jaxdem.utils.thermal#

Utility functions to compute thermodynamic quantities.

Functions

`compute_energy`(state, system)	Compute the total mechanical energy of the system.
`compute_potential_energy`(state, system)	Compute the total potential energy of the system.
`compute_potential_energy_per_particle`(state, ...)	Compute the potential energy per particle based on system interactions.
`compute_rotational_kinetic_energy`(state)	Compute the total rotational kinetic energy of the system.
`compute_rotational_kinetic_energy_per_particle`(state)	Compute the rotational kinetic energy per particle.
`compute_temperature`(state, can_rotate, ...)	Compute the temperature for a state.
`compute_translational_kinetic_energy`(state)	Compute the total translational kinetic energy of the system.
`compute_translational_kinetic_energy_per_particle`(state)	Compute the translational kinetic energy per particle.
`count_dynamic_dofs`(state, can_rotate, ...)	Count the number of degrees of freedom for the dynamics.
`scale_to_temperature`(state, ...[, k_B])	Scale the velocities of a state to a desired temperature state: State target_temperature: float - desired target temperature can_rotate: bool - whether to include the rigid body rotations subtract_drift: bool - whether to remove center of mass drift (usually only relevant for small systems) k_B: Optional[float] - boltzmanns constant, default is 1.0
`set_temperature`(state, target_temperature, ...)	Randomize the velocities of a state according to a desired temperature.

jaxdem.utils.thermal.compute_translational_kinetic_energy_per_particle(state: State) → jax.Array[source][source]#

Compute the translational kinetic energy per particle.

\[E_{trans} = \frac{1}{2} m |v|^2\]

Notes

The energy of clump members is divided by the number of spheres in the clump.

Parameters:: state (State) – The current state of the system containing particle masses and velocities.
Returns:: An array containing the translational kinetic energy for each particle.
Return type:: jax.Array

jaxdem.utils.thermal.compute_rotational_kinetic_energy_per_particle(state: State) → jax.Array[source][source]#

Compute the rotational kinetic energy per particle.

\[E_{rot} = \frac{1}{2} \vec{\omega}^T I \vec{\omega}\]

Notes

The energy of clump members is divided by the number of spheres in the clump.

Parameters:: state (State) – The current state of the system containing inertia, orientation, and angular velocity.
Returns:: An array containing the rotational kinetic energy for each particle.
Return type:: jax.Array

jaxdem.utils.thermal.compute_translational_kinetic_energy(state: State) → jax.Array[source][source]#

Compute the total translational kinetic energy of the system.

\[E_{trans, total} = \sum_{i} \frac{1}{2} m_i |v_i|^2\]

Parameters:: state (State) – The current state of the system.
Returns:: The scalar sum of translational kinetic energy across all particles.
Return type:: jax.Array

jaxdem.utils.thermal.compute_rotational_kinetic_energy(state: State) → jax.Array[source][source]#

Compute the total rotational kinetic energy of the system.

\[E_{rot, total} = \sum_{i} \frac{1}{2} \vec{\omega}_i^T I_i \vec{\omega}_i\]

Parameters:: state (State) – The current state of the system.
Returns:: The scalar sum of rotational kinetic energy across all particles.
Return type:: jax.Array

jaxdem.utils.thermal.compute_potential_energy_per_particle(state: State, system: System) → jax.Array[source][source]#

Compute the potential energy per particle based on system interactions. Energy is computed from the force models in the collider, and gravity and force functions that have potential energy associated with them in the force manager.

Parameters:

state (State) – The current state of the system.
system (System) – The system definition containing the collider and potential energy functions.

Returns:

An array containing the potential energy for each particle.

Return type:

jax.Array

jaxdem.utils.thermal.compute_potential_energy(state: State, system: System) → jax.Array[source][source]#

Compute the total potential energy of the system. Energy is computed from the force models in the collider, and gravity and force functions that have potential energy associated with them in the force manager.

\[E_{pot, total} = \sum_{i} U(r_i)\]

Parameters:

state (State) – The current state of the system.
system (System) – The system definition containing the collider.

Returns:

The scalar sum of potential energy across all particles.

Return type:

jax.Array

jaxdem.utils.thermal.compute_energy(state: State, system: System) → jax.Array[source][source]#

Compute the total mechanical energy of the system.

\[E_{total} = E_{pot, total} + E_{trans, total} + E_{rot, total}\]

Parameters:

state (State) – The current state of the system.
system (System) – The system definition containing physics parameters and colliders.

Returns:

The total energy (scalar) of the system.

Return type:

jax.Array

jaxdem.utils.thermal.count_dynamic_dofs(state: State, can_rotate: bool, subtract_drift: bool) → Tuple[jax.Array, jax.Array, jax.Array][source][source]#

Count the number of degrees of freedom for the dynamics.

Parameters:

state (State) – Current simulation state.
can_rotate (bool) – Whether to include rigid body rotations.
subtract_drift (bool) – Whether to include center-of-mass drift (usually only relevant for small systems).

jaxdem.utils.thermal.compute_temperature(state: State, can_rotate: bool, subtract_drift: bool, k_B: float = 1.0) → float[source][source]#

Compute the temperature for a state.

Parameters:

state (State) – Current simulation state.
can_rotate (bool) – Whether to include rigid body rotations.
subtract_drift (bool) – Whether to remove center-of-mass drift (usually only relevant for small systems).
k_B (float, optional) – Boltzmann constant (default is 1.0).

jaxdem.utils.thermal.set_temperature(state: State, target_temperature: float, can_rotate: bool, subtract_drift: bool, seed: int | None = 0, k_B: float = 1.0) → State[source][source]#

Randomize the velocities of a state according to a desired temperature.

Parameters:

state (State) – Current simulation state.
target_temperature (float) – Desired target temperature.
can_rotate (bool) – Whether to include rigid body rotations.
subtract_drift (bool) – Whether to remove center-of-mass drift (usually only relevant for small systems).
seed (int, optional) – RNG seed.
k_B (float, optional) – Boltzmann constant (default is 1.0).

jaxdem.utils.thermal.scale_to_temperature(state: State, target_temperature: float, can_rotate: bool, subtract_drift: bool, k_B: float = 1.0) → State[source][source]#: Scale the velocities of a state to a desired temperature state: State target_temperature: float - desired target temperature can_rotate: bool - whether to include the rigid body rotations subtract_drift: bool - whether to remove center of mass drift (usually only relevant for small systems) k_B: Optional[float] - boltzmanns constant, default is 1.0