mythos.simulators.jax_md.jaxmd
==============================

.. py:module:: mythos.simulators.jax_md.jaxmd

.. autoapi-nested-parse::

   A sampler based on running a jax_md simulation routine.



Attributes
----------

.. autoapisummary::

   mythos.simulators.jax_md.jaxmd.SIM_STATE


Classes
-------

.. autoapisummary::

   mythos.simulators.jax_md.jaxmd.JaxMDSimulator


Functions
---------

.. autoapisummary::

   mythos.simulators.jax_md.jaxmd.build_run_fn


Module Contents
---------------

.. py:data:: SIM_STATE

.. py:class:: JaxMDSimulator

   Bases: :py:obj:`mythos.simulators.base.Simulator`


   A sampler based on running a jax_md simulation routine.


   .. py:attribute:: energy_fn
      :type:  mythos.energy.base.EnergyFunction


   .. py:attribute:: simulator_params
      :type:  mythos.simulators.jax_md.utils.StaticSimulatorParams


   .. py:attribute:: space
      :type:  jax_md.space.Space


   .. py:attribute:: simulator_init
      :type:  collections.abc.Callable[[collections.abc.Callable, collections.abc.Callable], jax_md.simulate.Simulator]


   .. py:attribute:: neighbors
      :type:  mythos.simulators.jax_md.utils.NeighborHelper


   .. py:method:: __post_init__() -> None

      Builds the run function using the provided parameters.



.. py:function:: 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]

   Builds the run function for the jax_md simulation.