ensemble.sampling.TrajectoryState

ensemble.sampling.TrajectoryState#

class TrajectoryState(sim_state, trajectory, overflow=False, dynamic_kwargs=None, static_kwargs=None, aux=None, key=None, energy_params=None, entropy_diff=0.0, free_energy_diff=0.0)[source]#

A dataclass storing information of a generated trajectory.

Variables:

  • sim_state (chemtrain.ensemble.sampling.SimulatorState) – Last simulation state, a tuple of last state and nbrs

  • trajectory (chemtrain.ensemble.evaluation.State) – Generated trajectory

  • overflow (jax.Array) – True if neighbor list overflowed during trajectory generation

  • dynamic_kwargs (Dict[str, jax.Array]) – Additional information passed to the simulator and energy function, e.g., species, thermostat / barostat targets

  • static_kwargs (Dict[str, jax.Array]) – Same as dynamic_kwargs but constant for the trajectory.

  • aux (Dict[str, Any]) – Dict of auxilary per-snapshot quantities as defined by quantities in trajectory generator.

  • key (jax.Array) – PRNGKey of the trajectory state.

  • energy_params (Any) – Energy parameters used to generate the trajectory.

  • entropy_diff (jax.Array) – Entropy difference estimated for the trajectory, e.g., via DiffTRe optimization

  • free_energy_diff (jax.Array) – Free energy difference estimated for the trajectory, e.g., via DiffTRe optimization

Methods

__init__(sim_state, trajectory[, overflow, ...])

from_tuple()

get(k[,d])

items()

keys()

replace(**kwargs)

to_tuple()

values()

Attributes

aux: Dict[str, Any] = None#
barostat_press#

Target barostat pressure at time of respective snapshots.

dynamic_kwargs: Dict[str, Array] = None#
energy_params: Any = None#
entropy_diff: Array = 0.0#
free_energy_diff: Array = 0.0#
key: Array = None#
overflow: Array = False#
reference_nbrs#

Returns a single neighbor list.

static_kwargs: Dict[str, Array] = None#
thermostat_kbt#

Target thermostat kbT at time of respective snapshots.

sim_state: SimulatorState#
trajectory: State#
Contents