jax_dna.simulators.jax_md.jaxmd
===============================

.. py:module:: jax_dna.simulators.jax_md.jaxmd

.. autoapi-nested-parse::

   A sampler based on running a jax_md simulation routine.



Attributes
----------

.. autoapisummary::

   jax_dna.simulators.jax_md.jaxmd.REQUIRED_KEYS
   jax_dna.simulators.jax_md.jaxmd.ERR_MISSING_REQUIRED_KEYS
   jax_dna.simulators.jax_md.jaxmd.SIM_STATE


Classes
-------

.. autoapisummary::

   jax_dna.simulators.jax_md.jaxmd.JaxMDSimulator


Functions
---------

.. autoapisummary::

   jax_dna.simulators.jax_md.jaxmd.build_run_fn


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

.. py:data:: REQUIRED_KEYS

.. py:data:: ERR_MISSING_REQUIRED_KEYS
   :value: 'Missing required keys: {}'


.. py:data:: SIM_STATE

.. py:class:: JaxMDSimulator

   Bases: :py:obj:`jax_dna.simulators.base.BaseSimulation`


   A sampler based on running a jax_md simulation routine.


   .. py:attribute:: energy_configs
      :type:  list[jax_dna.energy.configuration.BaseConfiguration]


   .. py:attribute:: energy_fns
      :type:  list[jax_dna.energy.base.BaseEnergyFunction]


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


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


   .. py:attribute:: transform_fn
      :type:  collections.abc.Callable


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


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


   .. py:attribute:: topology
      :type:  jax_dna.input.topology.Topology


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

      Builds the run function using the provided parameters.



.. py:function:: build_run_fn(energy_configs: list[jax_dna.energy.configuration.BaseConfiguration], energy_fns: list[jax_dna.energy.base.BaseEnergyFunction], simulator_params: jax_dna.simulators.jax_md.utils.StaticSimulatorParams, space: jax_md.space.Space, transform_fn: collections.abc.Callable, simulator_init: collections.abc.Callable[[collections.abc.Callable, collections.abc.Callable], jax_md.simulate.Simulator], neighbors: jax_dna.simulators.jax_md.utils.NeighborHelper) -> collections.abc.Callable[[dict[str, float], jax_md.rigid_body.RigidBody, int, jax.random.PRNGKey], jax_dna.input.trajectory.Trajectory]

   Builds the run function for the jax_md simulation.