mythos.energy.base

Base classes for energy functions.

Attributes

`ERR_PARAM_NOT_FOUND`
`ERR_CALL_NOT_IMPLEMENTED`
`ERR_COMPOSED_ENERGY_FN_LEN_MISMATCH`
`ERR_COMPOSED_ENERGY_FN_TYPE_ENERGY_FNS`

Classes

`EnergyFunction`	Abstract base class for energy functions.
`BaseNucleotide`	Base nucleotide class.
`BaseEnergyFunction`	Base class for energy functions.
`ComposedEnergyFunction`	Represents a linear combination of energy functions.
`QualifiedComposedEnergyFunction`	A ComposedEnergyFunction that qualifies parameters by their function.

Module Contents

mythos.energy.base.ERR_PARAM_NOT_FOUND = "Parameter '{key}' not found in {class_name}"

mythos.energy.base.ERR_CALL_NOT_IMPLEMENTED = 'Subclasses must implement this method'

mythos.energy.base.ERR_COMPOSED_ENERGY_FN_LEN_MISMATCH = 'Weights must have the same length as energy functions'

mythos.energy.base.ERR_COMPOSED_ENERGY_FN_TYPE_ENERGY_FNS = 'energy_fns must be a list of energy functions'

class mythos.energy.base.EnergyFunction[source]

Bases: 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.

abstractmethod __call__(body: jax_md.rigid_body.RigidBody) → float[source]: Calculate the energy of the system.

abstractmethod with_params(*repl_dicts: dict, **repl_kwargs: Any) → EnergyFunction[source]

Return a new energy function with updated parameters.

Parameters:

*repl_dicts (dict) – dictionaries of parameters to update. These must come first in the argument list and will be applied in order.
**repl_kwargs – keyword arguments of parameters to update. These are applied after any parameter dictionaries supplied as positional arguments.

abstractmethod with_props(**kwargs) → EnergyFunction[source]

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.

abstractmethod with_noopt(*params: str) → EnergyFunction[source]: Create a new energy function from this with specified parameters non-optimizable.

abstractmethod params_dict(*, include_dependent: bool = True, exclude_non_optimizable: bool = False) → dict[source]

Get the parameters as a dictionary.

Parameters:

include_dependent (bool) – whether to include dependent parameters
exclude_non_optimizable (bool) – whether to exclude non-optimizable parameters

abstractmethod opt_params() → dict[str, mythos.utils.types.Scalar][source]: Get the configured optimizable parameters.

map(body_sequence: jax.numpy.ndarray) → jax.numpy.ndarray[source]: Map the energy function over a sequence of rigid bodies.

class mythos.energy.base.BaseNucleotide[source]

Bases: jax_md.rigid_body.RigidBody, abc.ABC

Base nucleotide class.

center: mythos.utils.types.Arr_Nucleotide_3

orientation: mythos.utils.types.Arr_Nucleotide_3 | jax_md.rigid_body.Quaternion

stack_sites: mythos.utils.types.Arr_Nucleotide_3

back_sites: mythos.utils.types.Arr_Nucleotide_3

base_sites: mythos.utils.types.Arr_Nucleotide_3

back_base_vectors: mythos.utils.types.Arr_Nucleotide_3

base_normals: mythos.utils.types.Arr_Nucleotide_3

cross_prods: mythos.utils.types.Arr_Nucleotide_3

static from_rigid_body(rigid_body: jax_md.rigid_body.RigidBody, **kwargs) → BaseNucleotide[source]

Abstractmethod:

Create an instance of the subclass from a RigidBody..

class mythos.energy.base.BaseEnergyFunction[source]

Bases: EnergyFunction

Base class for energy functions.

This class should not be used directly. Subclasses should implement the __call__ method.

Parameters:: displacement_fn (Callable) – an instance of a displacement function from jax_md.space

params: mythos.energy.configuration.BaseConfiguration

displacement_fn: collections.abc.Callable

seq: mythos.utils.types.Sequence | None = None

bonded_neighbors: mythos.utils.types.Arr_Bonded_Neighbors_2 | None = None

unbonded_neighbors: mythos.utils.types.Arr_Unbonded_Neighbors_2 | None = None

topology: dataclasses.InitVar[mythos.input.topology.Topology | None] = None

transform_fn: collections.abc.Callable | None = None

__post_init__(topology: mythos.input.topology.Topology | None) → None[source]

classmethod create_from(other: EnergyFunction, **kwargs) → EnergyFunction[source]

Create a new energy function from another with updated properties.

Parameters:

other – the energy function to copy properties from
**kwargs – properties to update, overriding those from other

property displacement_mapped: collections.abc.Callable: Returns the displacement function mapped to the space.

__add__(other: BaseEnergyFunction) → ComposedEnergyFunction[source]: Add two energy functions together to create a ComposedEnergyFunction.

__mul__(other: float) → ComposedEnergyFunction[source]: Multiply an energy function by a scalar to create a ComposedEnergyFunction.

with_props(**kwargs: Any) → EnergyFunction[source]

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.

with_noopt(*params: str) → EnergyFunction[source]: Create a new energy function from this with specified parameters non-optimizable.

opt_params() → dict[str, mythos.utils.types.Scalar][source]: Get the configured optimizable parameters.

with_params(*repl_dicts: dict, **repl_kwargs: Any) → EnergyFunction[source]

Return a new energy function with updated parameters.

Parameters:

*repl_dicts (dict) – dictionaries of parameters to update. These must come first in the argument list and will be applied in order.
**repl_kwargs – keyword arguments of parameters to update. These are applied after any parameter dictionaries supplied as positional arguments.

params_dict(include_dependent: bool = True, exclude_non_optimizable: bool = False) → dict[source]

Get the parameters as a dictionary.

Parameters:

include_dependent (bool) – whether to include dependent parameters
exclude_non_optimizable (bool) – whether to exclude non-optimizable parameters

__call__(body: jax_md.rigid_body.RigidBody) → float[source]: Calculate the energy of the system.

abstractmethod compute_energy(nucleotide: BaseNucleotide) → float[source]: Compute the energy of the system given the nucleotide.

class mythos.energy.base.ComposedEnergyFunction[source]

Bases: 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.

Parameters:

energy_fns (list[BaseEnergyFunction]) – a list of energy functions
weights (jnp.ndarray) – optional, the weights of the energy functions

energy_fns: list[BaseEnergyFunction]

weights: jax.numpy.ndarray | None = None

__post_init__() → None[source]: Check that the input is valid.

with_props(**kwargs: Any) → ComposedEnergyFunction[source]

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.

_param_in_fn(param: str, fn: BaseEnergyFunction) → bool[source]: Helper for with_params to check if a param is in a given energy function.

_rename_param_for_fn(param: str, _fn: BaseEnergyFunction) → str[source]: Helper to rename a param for input to a given energy function.

_rename_param_from_fn(param: str, _fn: BaseEnergyFunction) → str[source]: Helper to rename a param for output from a given energy function.

with_noopt(*params: str) → ComposedEnergyFunction[source]: Create a new energy function from this with specified parameters non-optimizable.

opt_params(from_fns: list[type] | None = None) → dict[str, mythos.utils.types.Scalar][source]: Get the configured optimizable parameters.

with_params(*repl_dicts: dict, **repl_kwargs: Any) → ComposedEnergyFunction[source]

Return a new energy function with updated parameters.

Parameters:

*repl_dicts (dict) – dictionaries of parameters to update. These must come first in the argument list and will be applied in order.
**repl_kwargs – keyword arguments of parameters to update. These are applied after any parameter dictionaries supplied as positional arguments.

params_dict(*, include_dependent: bool = True, exclude_non_optimizable: bool = False) → dict[source]

Get the parameters as a dictionary.

Parameters:

include_dependent (bool) – whether to include dependent parameters
exclude_non_optimizable (bool) – whether to exclude non-optimizable parameters

compute_terms(body: jax_md.rigid_body.RigidBody) → jax.numpy.ndarray[source]: Compute each of the energy terms in the energy function.

__call__(body: jax_md.rigid_body.RigidBody) → float[source]: Calculate the energy of the system.

without_terms(*terms: list[str | type]) → ComposedEnergyFunction[source]

Create a new ComposedEnergyFunction without the specified terms.

Parameters:

*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

Return type:

ComposedEnergyFunction

add_energy_fn(energy_fn: BaseEnergyFunction, weight: float = 1.0) → ComposedEnergyFunction[source]

Add an energy function to the list of energy functions.

Parameters:

energy_fn (BaseEnergyFunction) – the energy function to add
weight (float) – the weight of the energy function

Returns:

a new ComposedEnergyFunction with the added energy function

Return type:

ComposedEnergyFunction

add_composable_energy_fn(energy_fn: ComposedEnergyFunction) → ComposedEnergyFunction[source]

Add a ComposedEnergyFunction to the list of energy functions.

Parameters:: energy_fn (ComposedEnergyFunction) – the ComposedEnergyFunction to add
Returns:: a new ComposedEnergyFunction with the added energy function
Return type:: ComposedEnergyFunction

__add__(other: BaseEnergyFunction | ComposedEnergyFunction) → ComposedEnergyFunction[source]

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.

__radd__(other: BaseEnergyFunction | ComposedEnergyFunction) → ComposedEnergyFunction[source]

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.

classmethod from_lists(energy_fns: list[BaseEnergyFunction], energy_configs: list[mythos.energy.configuration.BaseConfiguration], weights: list[float] | None = None, **kwargs) → ComposedEnergyFunction[source]

Create a ComposedEnergyFunction from lists of energy functions and weights.

Parameters:

energy_fns (list[BaseEnergyFunction]) – a list of energy functions
energy_configs (list[BaseConfiguration]) – a list of energy configurations
weights (list[float] | None) – optional, a list of weights for the energy functions
**kwargs – keyword arguments to pass to each energy function

Returns:

a new ComposedEnergyFunction

Return type:

ComposedEnergyFunction

class mythos.energy.base.QualifiedComposedEnergyFunction[source]

Bases: 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.

_param_in_fn(param: str, fn: BaseEnergyFunction) → bool[source]: Helper for with_params to check if a param is in a given energy function.

_rename_param_for_fn(param: str, fn: BaseEnergyFunction) → str[source]: Helper to rename a param for input to a given energy function.

_rename_param_from_fn(param: str, fn: BaseEnergyFunction) → str[source]: Helper to rename a param for output from a given energy function.