jax_dna.optimization.simulator

Simulation actors for use in an jax_dna.optimization.ray_optimization.Optimization loop.

Classes

`BaseSimulator`	A base class for a simulator actor.
`SimulatorActor`	A ray actor that runs a simulation and exposes observables.

Module Contents

class jax_dna.optimization.simulator.BaseSimulator(name: str, fn: Callable[[jax_dna.utils.types.Params, jax_dna.utils.types.MetaData], tuple[str, Ellipsis]], exposes: list[str], meta_data: jax_dna.utils.types.MetaData)[source]

A base class for a simulator actor.

The class is split this way to make testing easier.

_name

_fn

_exposes

_meta_data

name() → str[source]: Returns the name of the simulation.

exposes() → list[str][source]: Returns the list of observables exposed by the simulation.

meta_data() → jax_dna.utils.types.MetaData[source]: Returns the metadata used by the simulation.

run(params: jax_dna.utils.types.Params) → tuple[str, Ellipsis][source]: Runs the simulation using the given params and returns the observables and metadata.

class jax_dna.optimization.simulator.SimulatorActor(name: str, fn: Callable[[jax_dna.utils.types.Params, jax_dna.utils.types.MetaData], tuple[str, Ellipsis]], exposes: list[str], meta_data: jax_dna.utils.types.MetaData)[source]

Bases: BaseSimulator

A ray actor that runs a simulation and exposes observables.

The simulator actor is wrapper around a simulator function so to be used in a jax_dna.optimization.ray_optimization.Optimization. Because a simulation trajectory and derived observables can be large, the simulation function should write the trajectory to a file and return the path to the file.