mythos.energy.base ================== .. py:module:: mythos.energy.base .. autoapi-nested-parse:: Base classes for energy functions. Attributes ---------- .. autoapisummary:: mythos.energy.base.ERR_PARAM_NOT_FOUND mythos.energy.base.ERR_CALL_NOT_IMPLEMENTED mythos.energy.base.ERR_COMPOSED_ENERGY_FN_LEN_MISMATCH mythos.energy.base.ERR_COMPOSED_ENERGY_FN_TYPE_ENERGY_FNS Classes ------- .. autoapisummary:: mythos.energy.base.EnergyFunction mythos.energy.base.BaseNucleotide mythos.energy.base.BaseEnergyFunction mythos.energy.base.ComposedEnergyFunction mythos.energy.base.QualifiedComposedEnergyFunction 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:: EnergyFunction Bases: :py:obj:`abc.ABC` Abstract base class for energy functions. These are a class of callable-classes that take in a RigidBody and return the energy of the system as a scalar float. .. py:method:: __call__(body: jax_md.rigid_body.RigidBody) -> float :abstractmethod: Calculate the energy of the system. .. py:method:: with_params(*repl_dicts: dict, **repl_kwargs: Any) -> EnergyFunction :abstractmethod: Return a new energy function with updated parameters. :param \*repl_dicts: dictionaries of parameters to update. These must come first in the argument list and will be applied in order. :type \*repl_dicts: dict :param \*\*repl_kwargs: keyword arguments of parameters to update. These are applied after any parameter dictionaries supplied as positional arguments. .. py:method:: with_props(**kwargs) -> EnergyFunction :abstractmethod: Create a new energy function from this with updated properties. Properties are those that are defined at the energy function class level and not the parameters that are defined therein. For example, the `displacement_fn` can be modified using this method. .. py:method:: with_noopt(*params: str) -> EnergyFunction :abstractmethod: Create a new energy function from this with specified parameters non-optimizable. .. py:method:: params_dict(*, include_dependent: bool = True, exclude_non_optimizable: bool = False) -> dict :abstractmethod: Get the parameters as a dictionary. :param include_dependent: whether to include dependent parameters :type include_dependent: bool :param exclude_non_optimizable: whether to exclude non-optimizable parameters :type exclude_non_optimizable: bool .. py:method:: opt_params() -> dict[str, mythos.utils.types.Scalar] :abstractmethod: Get the configured optimizable parameters. .. py:method:: map(body_sequence: jax.numpy.ndarray) -> jax.numpy.ndarray Map the energy function over a sequence of rigid bodies. .. py:class:: BaseNucleotide Bases: :py:obj:`jax_md.rigid_body.RigidBody`, :py:obj:`abc.ABC` Base nucleotide class. .. py:attribute:: center :type: mythos.utils.types.Arr_Nucleotide_3 .. py:attribute:: orientation :type: mythos.utils.types.Arr_Nucleotide_3 | jax_md.rigid_body.Quaternion .. py:attribute:: stack_sites :type: mythos.utils.types.Arr_Nucleotide_3 .. py:attribute:: back_sites :type: mythos.utils.types.Arr_Nucleotide_3 .. py:attribute:: base_sites :type: mythos.utils.types.Arr_Nucleotide_3 .. py:attribute:: back_base_vectors :type: mythos.utils.types.Arr_Nucleotide_3 .. py:attribute:: base_normals :type: mythos.utils.types.Arr_Nucleotide_3 .. py:attribute:: cross_prods :type: mythos.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.. .. py:class:: BaseEnergyFunction Bases: :py:obj:`EnergyFunction` 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:: params :type: mythos.energy.configuration.BaseConfiguration .. py:attribute:: displacement_fn :type: collections.abc.Callable .. py:attribute:: seq :type: mythos.utils.types.Sequence | None :value: None .. py:attribute:: bonded_neighbors :type: mythos.utils.types.Arr_Bonded_Neighbors_2 | None :value: None .. py:attribute:: unbonded_neighbors :type: mythos.utils.types.Arr_Unbonded_Neighbors_2 | None :value: None .. py:attribute:: topology :type: dataclasses.InitVar[mythos.input.topology.Topology | None] :value: None .. py:attribute:: transform_fn :type: collections.abc.Callable | None :value: None .. py:method:: __post_init__(topology: mythos.input.topology.Topology | None) -> None .. py:method:: create_from(other: EnergyFunction, **kwargs) -> EnergyFunction :classmethod: Create a new energy function from another with updated properties. :param other: the energy function to copy properties from :param \*\*kwargs: properties to update, overriding those from other .. 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:: with_props(**kwargs: Any) -> EnergyFunction Create a new energy function from this with updated properties. Properties are those that are defined at the energy function class level and not the parameters that are defined therein. For example, the `displacement_fn` can be modified using this method. .. py:method:: with_noopt(*params: str) -> EnergyFunction Create a new energy function from this with specified parameters non-optimizable. .. py:method:: opt_params() -> dict[str, mythos.utils.types.Scalar] Get the configured optimizable parameters. .. py:method:: with_params(*repl_dicts: dict, **repl_kwargs: Any) -> EnergyFunction Return a new energy function with updated parameters. :param \*repl_dicts: dictionaries of parameters to update. These must come first in the argument list and will be applied in order. :type \*repl_dicts: dict :param \*\*repl_kwargs: keyword arguments of parameters to update. These are applied after any parameter dictionaries supplied as positional arguments. .. py:method:: params_dict(include_dependent: bool = True, exclude_non_optimizable: bool = False) -> dict Get the parameters as a dictionary. :param include_dependent: whether to include dependent parameters :type include_dependent: bool :param exclude_non_optimizable: whether to exclude non-optimizable parameters :type exclude_non_optimizable: bool .. py:method:: __call__(body: jax_md.rigid_body.RigidBody) -> float Calculate the energy of the system. .. py:method:: compute_energy(nucleotide: BaseNucleotide) -> float :abstractmethod: Compute the energy of the system given the nucleotide. .. py:class:: ComposedEnergyFunction Bases: :py:obj:`EnergyFunction` Represents a linear combination of energy functions. The parameters of all composite energy functions are treated as sharing a global namespace in all setting and retrieval methods. For example, calling `with_params(kt=0.1)` will set the parameter `kt` in all those energy functions that contain a parameter name `kt`. :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 .. py:attribute:: energy_fns :type: list[BaseEnergyFunction] .. py:attribute:: weights :type: jax.numpy.ndarray | None :value: None .. py:method:: __post_init__() -> None Check that the input is valid. .. py:method:: with_props(**kwargs: Any) -> ComposedEnergyFunction Create a new energy function from this with updated properties. Properties are those that are defined at the energy function class level and not the parameters that are defined therein. For example, the `displacement_fn` can be modified using this method. .. py:method:: _param_in_fn(param: str, fn: BaseEnergyFunction) -> bool Helper for with_params to check if a param is in a given energy function. .. py:method:: _rename_param_for_fn(param: str, _fn: BaseEnergyFunction) -> str Helper to rename a param for input to a given energy function. .. py:method:: _rename_param_from_fn(param: str, _fn: BaseEnergyFunction) -> str Helper to rename a param for output from a given energy function. .. py:method:: with_noopt(*params: str) -> ComposedEnergyFunction Create a new energy function from this with specified parameters non-optimizable. .. py:method:: opt_params(from_fns: list[type] | None = None) -> dict[str, mythos.utils.types.Scalar] Get the configured optimizable parameters. .. py:method:: with_params(*repl_dicts: dict, **repl_kwargs: Any) -> ComposedEnergyFunction Return a new energy function with updated parameters. :param \*repl_dicts: dictionaries of parameters to update. These must come first in the argument list and will be applied in order. :type \*repl_dicts: dict :param \*\*repl_kwargs: keyword arguments of parameters to update. These are applied after any parameter dictionaries supplied as positional arguments. .. py:method:: params_dict(*, include_dependent: bool = True, exclude_non_optimizable: bool = False) -> dict Get the parameters as a dictionary. :param include_dependent: whether to include dependent parameters :type include_dependent: bool :param exclude_non_optimizable: whether to exclude non-optimizable parameters :type exclude_non_optimizable: bool .. py:method:: compute_terms(body: jax_md.rigid_body.RigidBody) -> jax.numpy.ndarray Compute each of the energy terms in the energy function. .. py:method:: __call__(body: jax_md.rigid_body.RigidBody) -> float Calculate the energy of the system. .. py:method:: without_terms(*terms: list[str | type]) -> ComposedEnergyFunction Create a new ComposedEnergyFunction without the specified terms. :param \*terms: all positional arguments should be either a type or a string which is the name of the type to exclude. :returns: a new ComposedEnergyFunction without the specified terms :rtype: ComposedEnergyFunction .. 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:method:: from_lists(energy_fns: list[BaseEnergyFunction], energy_configs: list[mythos.energy.configuration.BaseConfiguration], weights: list[float] | None = None, **kwargs) -> ComposedEnergyFunction :classmethod: Create a ComposedEnergyFunction from lists of energy functions and weights. :param energy_fns: a list of energy functions :type energy_fns: list[BaseEnergyFunction] :param energy_configs: a list of energy configurations :type energy_configs: list[BaseConfiguration] :param weights: optional, a list of weights for the energy functions :type weights: list[float] | None :param \*\*kwargs: keyword arguments to pass to each energy function :returns: a new ComposedEnergyFunction :rtype: ComposedEnergyFunction .. py:class:: QualifiedComposedEnergyFunction Bases: :py:obj:`ComposedEnergyFunction` A ComposedEnergyFunction that qualifies parameters by their function. Parameters for composite functions do not share a global namespace, but instead are qualified by the function they belong to in all setting and retrieval methods. For example, parameter `eps_backbone` in Fene energy function would be referred to as `Fene.eps_backbone` in the this energy function. This is useful for isolating parameters from a specific energy function for optimization, however note that not all simulations will support this functionality - for example oxDNA simulations write only one value per parameter. .. py:method:: _param_in_fn(param: str, fn: BaseEnergyFunction) -> bool Helper for with_params to check if a param is in a given energy function. .. py:method:: _rename_param_for_fn(param: str, fn: BaseEnergyFunction) -> str Helper to rename a param for input to a given energy function. .. py:method:: _rename_param_from_fn(param: str, fn: BaseEnergyFunction) -> str Helper to rename a param for output from a given energy function.