mythos.simulators.jax_md.jaxmd

A sampler based on running a jax_md simulation routine.

Attributes

Classes

JaxMDSimulator

A sampler based on running a jax_md simulation routine.

Functions

build_run_fn(→ collections.abc.Callable[[dict[str, ...)

Builds the run function for the jax_md simulation.

Module Contents

mythos.simulators.jax_md.jaxmd.SIM_STATE

class mythos.simulators.jax_md.jaxmd.JaxMDSimulator[source]

Bases: mythos.simulators.base.Simulator

A sampler based on running a jax_md simulation routine.

energy_fn: mythos.energy.base.EnergyFunction

simulator_params: mythos.simulators.jax_md.utils.StaticSimulatorParams

space: jax_md.space.Space

simulator_init: collections.abc.Callable[[collections.abc.Callable, collections.abc.Callable], jax_md.simulate.Simulator]

neighbors: mythos.simulators.jax_md.utils.NeighborHelper

__post_init__() → None[source]: Builds the run function using the provided parameters.

mythos.simulators.jax_md.jaxmd.build_run_fn(energy_fn: mythos.energy.base.EnergyFunction, simulator_params: mythos.simulators.jax_md.utils.StaticSimulatorParams, space: jax_md.space.Space, simulator_init: collections.abc.Callable[[collections.abc.Callable, collections.abc.Callable], jax_md.simulate.Simulator], neighbors: mythos.simulators.jax_md.utils.NeighborHelper) → collections.abc.Callable[[dict[str, float], jax_md.rigid_body.RigidBody, int, jax.random.PRNGKey], mythos.input.trajectory.Trajectory][source]: Builds the run function for the jax_md simulation.