wepy2: Refactor and test suite

[salotz]

This is a beta release of v2 of wepy which introduces breaking changes to APIs.

For a complete and realistic example for users writing scripts/applications see: tests/integration/test_openmm/test_realistic.py

The primary short term goals of these changes were:

1. Increase robustness of REVOResampler and work mappers for running OpenMM simulations
2. Refactor code to make code testable, and write tests for all core code in project to make the project maintainable over time.

Some changes in focus were also made. Primarily, that current maintainence effort in core wepy (and this repo) will be on MD simulations and OpenMM in specific.

To that end the generic work mappers have been removed and replaced by purpose built implementations that work very well for OpenMM specifically.

See below for a breakdown of the breaking changes as well as some non-breaking improvements.

A key non-interface improvement is the inclusion of a fairly comprehensive test suite of the core components. This will expand to the analysis modules in the future.

Breaking Changes

General Changes

• many classes are now defined using attrs
• removes the pytest_wepy module
• removed orchestration and util/kv.py modules removed
• addition and refactoring of test systems under wepy_tools. wepy_tools was never intended to follow the same API stability guarantees but is worth mentioning that there are breakages.

Simulation Manager

Previously the interfaces between the simulation manager and the components (resamplers, runners, and boundary conditions) were only ambiently defined and difficult to understand for new implementations of components. These interfaces have been completely reified and typed to make them explicit and clear. This includes the CycleReportDict which is handed off to reporters and the RunSegmentData which is produced by runners.

Because the simulation manager is effectively a fairly complicated state machine this was made explicit with the introduction of a ManagerStateMachine, a state transition table MANAGER_STATE_TRANSITION_TABLE, and associated ManagerStatus and ManagerEvent types. The state machine does not drive actual state transitions but does track them and prevent invalid transitions, which makes it much easier to validate and test.

The Manager class itself is now explicitly generic over the WalkerState being used in the simulation from a typing perspective.

Another pattern was also adopted to make polymorphic configuration/construction of components possible without littering the simulation manager with specific arguments for a component (e.g. resampler) that are only used in certain implementations. The simulation manager now accepts "factories" for runners, resamplers, and work mappers (and will soon accept for boundary conditions, monitors, and reporters). A factory is a function that generates a component conditioned only on a tighter interface inside the simulation manager. Additional arguments must be folded into the factory creation. Factories are simply functions and you are free to implement them however you please but most components provide factories for themselves. The factory objects themselves should be plain data and amenable to simple serialization. The "runtime" objects generated by them are then completely managed by the Manager object.

Concretely, previously you would construct a stateful Resampler and then pass it to the Manager:

resampler = Resampler(...)

manager = Manager(resampler=resampler)

The state of the resampler was then not owned by the manager. Now the manager is responsible for creating and calling hooks during the simulation lifecycle based only upon a static, stateless, factory input. E.g.:

manager = Manager(resampler_factory=ResamplerFactory(...))

Then manager.resampler is the stateful runtime object.

This architecture (and implementation) thus does not have problems with ownership requiring judicious use of deepcopy for Manager inputs, now Manager will always have complete ownership what it needs and will not effect the application calling environment.

Runners

The base Runner class is now a generic protocol over the WalkerState and the segment data.

All runners must implement this interface.

Runners are now split up into "runtime" and "static" classes in accordance with the object lifecycles in the simulation manager.

The "runtime" component manages the state of the runs in a simulation while the static version (often a "factory" class) is more decoupled and should contain data only that can be used in serialization, etc.

To make managing state for the runners a state machine was created to aid in implementation: RunnerStateMachine. This is not required for implementation and only provides checks on explicit opt-in. However the new Runner protocol does require a status property to return the correct state as defined in RunnerStatus.

The NoRunner shows a minimal implementation. When transitioning between states via self.state_machine.send(RunnerEvent.INIT) etc. the state machine will verify that the transition is valid. This also helps in validation that new implementations are correct.

The randomwalk.py module was removed.

A MockRunner and associated classes are defined in runners/mock.py. This is useful in testing and for a reference implementation.

OpenMM Runner

The OpenMM runner was refactored heavily. Firstly it is split up into multiple sub-modules.

OpenMM Simulation State

Previously the OpenMMState was used as a WalkerState in simulations. This class was a wrapper around the openmm.State object to avoid unnecessary copies. However, there was a number of issues surrounding the statefulness of openmm.State objects and their serializability that was causing subtle issues with threads and processes leading to instabilities in the simulations.

In this release the original OpenMMState has now been replaced by both an OpenMMStateWrapper which is used to wrap the openmm.State and facilitate Pythonic and simple access to the underlying data and conversion to simpler pure-data representations better suited for serialization etc. The new OpenMMState is such a pure-data implementation and does not wrap the openmm.State. There are a variety of methods available for interconverting between these two kinds of objects.

In general the workflow would be to:

1. Generate an OpenMMState from arbitrary input data (i.e. the from_dwim "do what I mean") (skip this if you already have a State XML file)
2. Convert to an OpenMMStateWrapper with OpenMMState.to_state_wrapper
3. Use the OpenMMStateWrapper.state for feeding openmm.State to an openmm.Context simulation
4. Wrap outputs of an openmm.Context in a new OpenMMStateWrapper()
5. OpenMMState.from_state_wrapper to convert to a plain-data representation.

The OpenMMRunner itself handles all of this and only accepts OpenMMState objects, but this pattern is useful for generic data conversion workflows.

As generation of openmm.State objects from scratch has always been a persistent problem, requiring an openmm.Context, we have provided a novel implementation that serializes the OpenMMState to the State XML format supported by OpenMM (state_to_xml). This avoids generating openmm.Contexts which can sometimes cause stateful corruption within a process which causes issues if the same process then later runs a real simulation. The OpenMMStateWrapper.from_xml and serialize_xml provide conveniences for interconversion.

Note, that all fields on these objects are now openmm.unit.Quantity objects rather than there being both quantity and plain numpy versions.

OpenMM Logging Reporters

Support for logging of MD simulations directly has been enabled.

Support for logging OpenMM reporters was added directly to the OpenMM simulations themselves (this is distinct) from the wepy reporters, via the openmm_reporter_factories argument to the OpenMMRunner constructor. These factories must generate concrete reporters that implement the interface defined for openmm.app.Simulation. The wepy/runners/openmm/reporter.py module provides some concrete definitions of those interfaces to aid in implementation.

By default no OpenMM reporter factories are configured by the OpenMMRunner class. However, for convenience the DEFAULT_OPENMM_REPORTER_FACTORIES provides some preconfigured logging reporters that can be used. Be aware that there really isn't a "good default" for a wide range of applications so you should consider configuring these yourself for your application.

To aid in this the wepy/runners/openmm/logger.py module provides both some base classes for aiding in implementing useful openmm.app.Simulation reporters. The LoggingReporter class provides a base class that accepts a logging.Logger object, a logging callback function, and the standard state_includes for getting state fields. There is also two more specialized generic LoggingReporter classes:

1. StepIntervalLoggingReporter which configures the interval of MD steps for which it should be called.
2. SamplingTimeIntervalLoggingReporter which configures a sampling time interval for which logging callback should be executed.

There are three concrete loggers that hardcode a useful logging callback of general interest for MD simulations available to use as well:

1. HeartBeatLoggingReporter running on a step-interval that reports: clock time, elapsed time, sim time (in picoseconds), and sim steps (count).
2. EnergyLoggingReporter running on a sampling time interval that reports sim steps and time as well as the kinetic, potential, and total energy of the system.
3. UnitCellLoggingReporter running on a sampling time interval that reports sim steps and time as well as the box volume and vectors.

Resampling

Decision

For applications there shouldn't be anything needed to update.

The base Decision class has been replaced with BaseDecisionABC which is a generic ABC class.

The "decision records" which were informally defined dictionaries have now been reified as a protocol DecisionRecord and a base record BaseDecisionRecord which has the minimal data needed. Each decision record supports conversion to a plain dictionary (to_dict) for compatibility with the reporting interfaces.

The Decision.parents method was removed. Implementations of Decisions must implement this now.

The NoDecision class was moved to the no_decision.py module.

The clone_merge.py module was refactored and additional runtime checks were added on construction.

Distances

The Distance class is now a generic protocol wrt the Image type and WalkerState type.

In practice now designers of distances should explicitly define their Image type specifically for their distance, rather than relying on dynamic typing. Currently, this is completely unconstrained, but will likely add some constraints to support data output.

There is now a DistanceABC that has the same default implementations for image and distance, only requiring implementation of image_distance, if you want the previous behavior of Distance.

The example and mock distance metrics were moved into their own submodules mock.py and simple.py.

The mock.py is written for the new wepy.runners.mock components.

The simple.py distance shows a more standalone example that does not integrate to any of the runners.

The randomwalk.py distances were removed.

Resamplers

There is now a core Resampler generic protocol which defines the interface. All class constants, methods, etc. are type annotated.

There is now a ResamplerABC for inheriting the "standard" settings and methods.

Adds an additional check that the max and min walkers do not conflict for ResamplerABC construction.

Removed debug_mode as a constructor argument.

The NoResampler was moved to the noresampler.py submodule.

REVO

The default optional constructor arguments were refined to not be optional and use explicit input types.
Removed acceptance of **kwargs.

The paralellization via multiprocessing in the _all_to_all_distance function was updated. No APIs were changed but you should be aware of the new behavior. The 'spawn' process start method is now required and hard coded (rather than 'fork' as the default).

This should be significantly more robust than in the past in which the Pool would deadlock. For each process only a finite number of calculations will be scheduled to avoid potential memory leaks. Additionally logging is configured in the computed image functions.

As before n_proc=1 will not use multiprocessing at all.

Adds the REVOResamplerFactory in accord with the changes to simulation component lifecycles are used in the SimManager (see above).

Walker

• The Walker class is explicitly a generic class type. This interface is used throughout the code to indicate generic walkers. It is a generic class so that applications can specify particular State types as being contained in a Walker.
• The WalkerState class is now a protocol to support independent implementations rather than confusing inheritance.
• To get a "black box" dynamically typed walker with roughly the same utility as the old Walker class use WalkerStateBox.
• The cloning and merging methods on the Walker class have been pulled out into standalone functions.

Work Mappers

To align with the goal of focused support on OpenMM MD the old work generic WorkerMapper and TaskMapper have been removed. These were designed to be generic to runners which made them onerous to configure properly and had numerous robustness problems that would cause simulation crashes. These have been replaced by OpenMM specific work mappers in the work_mapper.openmm submodule.

The generic SerialMapper remains as its implementation was fine.

In addition the work mappers in general were updated to fit the new manner of component lifecycle in the simulation manager with a focus on factories generated from static configurations to live runtime state managing objects.

A WorkMapper protocol interface was defined which specifies exactly what a work mapper must implement.

As for the openmm work mappers the primary goal was to provide an integrated approach to GPU device and other platform settings baked into their configuration. The OpenMMSerialWorkMapper provides a serial implementation and the OpenMMProcPoolWorkMapper provides a parallel implementation via multiprocessing.Pool.

The OpenMMProcPoolWorkMapper is the primary production work mapper that should be used in realistic simulations. Its behavior is similar to the previous WorkerMapper in that worker processes receive serialized inputs to run (rather than relying on shared process memory from forking.) This implementation uses the same util.multiprocessing infrastructure to coordinate handling of logs from runner processes which enables log integration of sub-processes for enhanced monitoring.

Currently, performance of the OpenMMProcPoolWorkMapper is worse than the previous TaskMapper but has significant improvements in robustness and multiple deadlocking scenarios were eliminated. Only the 'spawn' method is supported as forking resulted in bugs related to OpenMM Context memory problems. Additionally, in the process pool only one simulation segment is run per process to avoid similar memory corruption problems and leaks.

Overall the new implementation is much simpler and outsources most coordination to multiprocessing.Pool rather than using lower level primitives.

Over time performance improvements can be made alongside more extensive testing and evaluation of robustness.

HDF5 and Storage

The HDF5 on-disk format was not substantially changed to support backwards compatibility, and because it was unnecessary.

One change is the serialization format for dtypes, which was updated and made more robust according to more modern numpy conventions.

An additional wepy.storage submodule was added as a place to define common interfaces for data that can be use both by the generators of data (runners, resamplers, boundary conditions, etc.) and the consumers (reporters and especially HDF5 reporter).

Ultimately, the goal is to have a complete interface for the WepyHDF5 format, but agnostic to the underlying storage technology, in this case HDF5. This will allow for implementation of different storage backends such as xarray, zarr, or TileDB.

WepyHDF5Reporter

The interface for this has changed to incorporate an explicit __init__ constructor with only the needed information and NOT the components that are typically used where this information can be pulled from. Use the alternate constructor WepyHDF5Reporter.from_components to have a similar ergonomic experience as the old constructor.

Non-Breaking Changes

Misc

• Types have been added to many function signatures to aid in refactoring and increasing maintainability. However, the project is still not type checked and these should be regarded as comments and not checked types.

• Added much more extensive logging to the components.

• A simulation monitor interface was defined in wepy/monitor.py, which will be for observability metrics in e.g. Prometheus.

• Decision class hierarchy was refactored [...]

Reporters

• Created Reporter protocol type with explicit interfaces for construction arguments (SimComponentArgs), the CycleReportDict (the data that each reporter receives from the sim manager).
• Dashboard
  • Fixed messy jinja templating in some of the dashboard sections
  • Typed the various data structures used in dashboard sections
• Moved file based reporters to their own module reporter.file and created an explicit ABC base class for inheritance.

HDF5

• Fixed some incorrect doc string descriptions of arguments
• The WepyHDF5 constructor methods were refactored to be much cleaner and easier to test
• several bugs wrt new versions of h5py were identified and fixed.
• Switched to using non-deprecated special types for variable length and string types.

The test suite is currently only partial but covers the key creation and writing methods.

WepyHDF5Reporter

• Fixes behavior surrounding alt reps for topologies and those actually saved in alt_reps trajectory fields.
• If no explicit units are given they will be automatically determined from the initial walker fields.
• The internal code was refactored to be much cleaner and remove a lot of the unnecessary object state used only for initialization.
• Fixed bug with closing h5py.File handles.

Multiprocessing Paralellization

In multiple places the use of process parallel computation was needed. This caused a number of problems when used naively. In response a centralized version of a "process pool" has been provided in the wepy.util.multiprocessing module. This provides a context manager queue_listener_context that implements forwarding of logs from subprocesses in a safe manner that also integrates cleanly with the configured logging settings.

The proc_pool_worker_setup function is meant to be used in configuring a Pool so as to do initial logging configuration in workers.

For example:

            mp_ctx = mp.get_context(method="spawn")

            with (
                queue_listener_context(mp_ctx) as log_queue,
                mp_ctx.Pool(
                    self.num_proc,
                    initializer=proc_pool_worker_setup,
                    initargs=(log_queue,),
                    maxtasksperchild=4,
                ) as pool,
            ):

This infrastructure is meant only to work with spawn process start method.

For Maintainer's

• code base was formatted using black and unnecessary imports were cleaned out repo wide
• Switch to uv for project mgmt.
• Fix wepy.analysis.network_layouts.layout_graph.RGBA_to_hex function
• removal of the tasks and templates directories
• removal of junk integration tests
• Implementation of a large number of unit and integration tests
• removal of the tests/utils folder