jax_dna.energy.base =================== .. py:module:: jax_dna.energy.base .. autoapi-nested-parse:: Base classes for energy functions. Attributes ---------- .. autoapisummary:: jax_dna.energy.base.ERR_PARAM_NOT_FOUND jax_dna.energy.base.ERR_CALL_NOT_IMPLEMENTED jax_dna.energy.base.ERR_COMPOSED_ENERGY_FN_LEN_MISMATCH jax_dna.energy.base.ERR_COMPOSED_ENERGY_FN_TYPE_ENERGY_FNS Classes ------- .. autoapisummary:: jax_dna.energy.base.BaseEnergyFunction jax_dna.energy.base.ComposedEnergyFunction jax_dna.energy.base.BaseNucleotide Module Contents --------------- .. py:data:: ERR_PARAM_NOT_FOUND :value: "Parameter '{key}' not found in {class_name}" .. py:data:: ERR_CALL_NOT_IMPLEMENTED :value: 'Subclasses must implement this method' .. py:data:: ERR_COMPOSED_ENERGY_FN_LEN_MISMATCH :value: 'Weights must have the same length as energy functions' .. py:data:: ERR_COMPOSED_ENERGY_FN_TYPE_ENERGY_FNS :value: 'energy_fns must be a list of energy functions' .. py:class:: BaseEnergyFunction Base class for energy functions. This class should not be used directly. Subclasses should implement the __call__ method. :param displacement_fn: an instance of a displacement function from jax_md.space :type displacement_fn: Callable .. py:attribute:: displacement_fn :type: collections.abc.Callable .. py:property:: displacement_mapped :type: collections.abc.Callable Returns the displacement function mapped to the space. .. py:method:: __add__(other: BaseEnergyFunction) -> ComposedEnergyFunction Add two energy functions together to create a ComposedEnergyFunction. .. py:method:: __mul__(other: float) -> ComposedEnergyFunction Multiply an energy function by a scalar to create a ComposedEnergyFunction. .. py:method:: __call__(body: jax_md.rigid_body.RigidBody, seq: jax_dna.utils.types.Sequence, bonded_neighbors: jax_dna.utils.types.Arr_Bonded_Neighbors_2, unbonded_neighbors: jax_dna.utils.types.Arr_Unbonded_Neighbors_2) -> float :abstractmethod: Calculate the energy of the system. .. py:class:: ComposedEnergyFunction Represents a linear combination of energy functions. :param energy_fns: a list of energy functions :type energy_fns: list[BaseEnergyFunction] :param weights: optional, the weights of the energy functions :type weights: jnp.ndarray :param rigid_body_transform_fn: a function to transform the rigid body :type rigid_body_transform_fn: Callable :param to into something that can be used by the energy functions like a DNA1 nucleotide: .. py:attribute:: energy_fns :type: tuple[BaseEnergyFunction] .. py:attribute:: weights :type: jax.numpy.ndarray | None :value: None .. py:attribute:: rigid_body_transform_fn :type: collections.abc.Callable[[jax_md.rigid_body.RigidBody], Any] | None :value: None .. py:method:: __post_init__() -> None Check that the input is valid. .. py:method:: compute_terms(body: jax_md.rigid_body.RigidBody, seq: jax_dna.utils.types.Sequence, bonded_neighbors: jax_dna.utils.types.Arr_Bonded_Neighbors_2, unbonded_neighbors: jax_dna.utils.types.Arr_Unbonded_Neighbors_2) -> jax.numpy.ndarray Compute each of the energy terms in the energy function. .. py:method:: __call__(body: jax_md.rigid_body.RigidBody, seq: jax_dna.utils.types.Sequence, bonded_neighbors: jax_dna.utils.types.Arr_Bonded_Neighbors_2, unbonded_neighbors: jax_dna.utils.types.Arr_Unbonded_Neighbors_2) -> float Calculates the energy of the system using the all of the function in `energy_fns`. :param body: The rigid body(ies) of the system :type body: jax_md.rigid_body.RigidBody :param seq: the sequence of the system :type seq: typ.Sequence :param bonded_neighbors: the bonded neighbors :type bonded_neighbors: typ.Arr_Bonded_Neighbors_2 :param unbonded_neighbors: the unbonded neighbors :type unbonded_neighbors: typ.Arr_Unbonded_Neighbors_2 :returns: the energy of the system :rtype: float .. py:method:: add_energy_fn(energy_fn: BaseEnergyFunction, weight: float = 1.0) -> ComposedEnergyFunction Add an energy function to the list of energy functions. :param energy_fn: the energy function to add :type energy_fn: BaseEnergyFunction :param weight: the weight of the energy function :type weight: float :returns: a new ComposedEnergyFunction with the added energy function :rtype: ComposedEnergyFunction .. py:method:: add_composable_energy_fn(energy_fn: ComposedEnergyFunction) -> ComposedEnergyFunction Add a ComposedEnergyFunction to the list of energy functions. :param energy_fn: the ComposedEnergyFunction to add :type energy_fn: ComposedEnergyFunction :returns: a new ComposedEnergyFunction with the added energy function :rtype: ComposedEnergyFunction .. py:method:: __add__(other: Union[BaseEnergyFunction, ComposedEnergyFunction]) -> ComposedEnergyFunction Create a new ComposedEnergyFunction by adding another energy function. This is a convenience method for the add_energy_fn and add_composable_energy_fn methods. .. py:method:: __radd__(other: Union[BaseEnergyFunction, ComposedEnergyFunction]) -> ComposedEnergyFunction Create a new ComposedEnergyFunction by adding another energy function. This is a convenience method for the add_energy_fn and add_composable_energy_fn methods. .. py:class:: BaseNucleotide Bases: :py:obj:`jax_md.rigid_body.RigidBody`, :py:obj:`abc.ABC` Base nucleotide class. .. py:attribute:: center :type: jax_dna.utils.types.Arr_Nucleotide_3 .. py:attribute:: orientation :type: jax_dna.utils.types.Arr_Nucleotide_3 | jax_md.rigid_body.Quaternion .. py:attribute:: stack_sites :type: jax_dna.utils.types.Arr_Nucleotide_3 .. py:attribute:: back_sites :type: jax_dna.utils.types.Arr_Nucleotide_3 .. py:attribute:: base_sites :type: jax_dna.utils.types.Arr_Nucleotide_3 .. py:attribute:: back_base_vectors :type: jax_dna.utils.types.Arr_Nucleotide_3 .. py:attribute:: base_normals :type: jax_dna.utils.types.Arr_Nucleotide_3 .. py:attribute:: cross_prods :type: jax_dna.utils.types.Arr_Nucleotide_3 .. py:method:: from_rigid_body(rigid_body: jax_md.rigid_body.RigidBody, **kwargs) -> BaseNucleotide :staticmethod: :abstractmethod: Create an instance of the subclass from a RigidBody..