CodeEntropy.levels.axes module

Axes utilities for entropy calculations.

This module contains the AxesCalculator, a geometry-focused helper used by the entropy pipeline to compute translational and rotational axes, centres, and moments of inertia at different hierarchy levels (residue / united-atom).

class CodeEntropy.levels.axes.AxesCalculator[source]

Bases: object

Compute translation/rotation axes and inertia utilities used by entropy.

Manages the structural and dynamic levels involved in entropy calculations. This includes selecting relevant levels, computing axes for translation and rotation, and handling bead-based representations of molecular systems.

Provides utility methods to:

  • Extract averaged positions.

  • Convert coordinates to spherical systems (future/legacy scope).

  • Compute axes used to rotate forces around.

  • Compute custom moments of inertia.

  • Manipulate vectors under periodic boundary conditions (PBC).

  • Construct custom moment-of-inertia tensors and principal axes.

Notes

This class deliberately does not:

  • Compute weighted forces/torques (that belongs in ForceTorqueCalculator).

  • Build covariances.

  • Compute entropies.

find_bonded_atoms(atom_idx: int, system)[source]

Find bonded heavy and hydrogen atoms for a given atom.

Parameters:

  • atom_idx – Atom index to find bonded atoms for.

  • system – MDAnalysis selection containing all atoms in current frame.

Returns:

  • bonded_heavy_atoms: bonded heavy atoms (mass 2 to 999)

  • bonded_H_atoms: bonded hydrogen atoms (mass 1 to 1.1)

Return type:

Tuple[AtomGroup, AtomGroup]

get_UA_axes(data_container, index: int)[source]

Compute united-atom-level translational and rotational axes.

The translational and rotational axes at the united-atom level.

This preserves the original behaviour and its rationale:

  • Translational axes:

    Use the same custom principal-axes approach as residue level: compute a custom MOI tensor using heavy-atom coordinates but UA masses (heavy + bonded H masses), then compute the principal axes from it.

  • Rotational axes:

    Identify heavy atoms in the residue/molecule of interest and choose the index-th heavy atom (where index corresponds to the bead index). Use bonded topology around that heavy atom to determine UA rotational axes (see get_bonded_axes()).

Parameters:

  • data_container (MDAnalysis.Universe or AtomGroup) – Molecule and trajectory data.

  • index (int) – Bead index (ordinal among heavy atoms).

Returns:

  • trans_axes: Translational axes (3, 3).

  • rot_axes: Rotational axes (3, 3).

  • center: Rotation centre (3,) (heavy atom position).

  • moment_of_inertia: (3,) moments for the UA around rot_axes.

Return type:

Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]

Raises:

  • IndexError – If index does not correspond to an existing heavy atom.

  • ValueError – If bonded-axis construction fails.

get_UA_masses(molecule) → list[float][source]

Return united-atom (UA) masses for a molecule.

For a given molecule, return a list of masses of UAs (combination of the heavy atoms + bonded hydrogen atoms). This list is used to get the moment of inertia tensor for molecules larger than one UA.

Parameters:

molecule – MDAnalysis AtomGroup representing the molecule.

Returns:

UA masses for each heavy atom.

Return type:

list[float]

get_bonded_axes(system, atom, dimensions: ndarray)[source]

Compute UA rotational axes from bonded topology around a heavy atom.

For a given heavy atom, use its bonded atoms to get the axes for rotating forces around. Few cases for choosing united atom axes, which are dependent on the bonds to the atom:

X -- H = bonded to zero or more light atom/s (case1)

X -- R = bonded to one heavy atom (case2)

R -- X -- H = bonded to one heavy and at least one light atom (case3)

R1 -- X -- R2 = bonded to two heavy atoms (case4)

R1 -- X -- R2 = bonded to more than two heavy atoms (case5)
      |
      R3

Note that axis2 is calculated by taking the cross product between axis1 and the vector chosen for each case, dependent on bonding:

  • case1: if all the bonded atoms are hydrogens, use the principal axes.

  • case2: use XR vector as axis1, arbitrary axis2.

  • case3: use XR vector as axis1, vector XH to calculate axis2

  • case4: use vector XR1 as axis1, and XR2 to calculate axis2

  • case5: get the sum of all XR normalised vectors as axis1, then use vector R1R2 to calculate axis2

axis3 is always the cross product of axis1 and axis2.

Parameters:

  • system – MDAnalysis selection containing all atoms in current frame.

  • atom – MDAnalysis Atom for the heavy atom.

  • dimensions – Simulation box dimensions (3,).

Returns:

  • custom_axes: Custom axes (3, 3), or None if atom is not heavy.

  • custom_moment_of_inertia: (3,) moment of inertia around axes.

Return type:

Tuple[np.ndarray | None, np.ndarray | None]

Notes

If custom_moment_of_inertia is not produced by the chosen method, it is computed using get_custom_moment_of_inertia() with the heavy atom as COM (matching original behaviour).

get_custom_axes(a: ndarray, b_list: Sequence[ndarray], c: ndarray, dimensions: ndarray) → ndarray[source]

Compute custom rotation axes from bonded vectors (PBC-aware).

For atoms a, b_list and c, calculate the axis to rotate forces around:

  • axis1: use the normalised vector ab as axis1. If there is more than one bonded heavy atom (HA), average over all the normalised vectors calculated from b_list and use this as axis1. b_list contains all the bonded heavy atom coordinates.

  • axis2: use the cross product of normalised vector ac and axis1 as axis2. If there are more than two bonded heavy atoms, then use normalised vector b[0]c to cross product with axis1. This gives the axis perpendicular to axis1.

  • axis3: the cross product of axis1 and axis2, which is perpendicular to axis1 and axis2.

Parameters:

  • a – Central united-atom coordinates (3,).

  • b_list – Positions of heavy bonded atoms.

  • c – Coordinates of a second heavy atom or a hydrogen atom.

  • dimensions – Simulation box dimensions (3,).

   a          1 = norm_ab
  / \         2 = perpendicular to norm_ab and norm_ac (or bc if >2 HAs)
 /   \        3 = perpendicular to 1 and 2
b     c
Returns:

(3, 3) array of the axes used to rotate forces.

Return type:

np.ndarray

Raises:

ValueError – If axes cannot be normalized due to degeneracy.

get_custom_moment_of_inertia(UA, custom_rotation_axes: ndarray, center_of_mass: ndarray, dimensions: ndarray) → ndarray[source]

Compute moment of inertia around custom axes for a UA.

Get the moment of inertia (specifically used for the united atom level) from a set of rotation axes and a given center of mass (COM is usually the heavy atom position in a UA).

Original behaviour preserved:

  • Uses PBC-aware translated coordinates.

  • Sums contributions from each atom using the squared norm of (axis × r) multiplied by mass.

  • Removes the lowest MOI degree of freedom if the UA only has a single bonded H (i.e. UA has 2 atoms total).

Parameters:

  • UA – MDAnalysis AtomGroup for the UA (heavy + bonded H atoms).

  • custom_rotation_axes – (3, 3) array of rotation axes.

  • center_of_mass – (3,) COM for the UA (typically HA position).

  • dimensions – (3,) simulation box dimensions.

Returns:

(3,) moment of inertia array.

Return type:

np.ndarray

get_custom_principal_axes(moment_of_inertia_tensor: ndarray) → tuple[ndarray, ndarray][source]

Compute principal axes and moments from a custom MOI tensor.

Principal axes and centre of axes from the ordered eigenvalues and eigenvectors of a moment of inertia tensor. This function allows for a custom moment of inertia tensor to be used, which isn’t possible with the built-in MDAnalysis principal_axes() function.

Original behaviour preserved:

  • Eigenvalues are sorted by descending absolute magnitude.

  • Eigenvectors are transposed so axes are returned as rows.

  • Z axis is flipped to enforce the same handedness convention as the original implementation.

Parameters:

moment_of_inertia_tensor – (3, 3) custom inertia tensor.

Returns:

  • principal_axes: (3, 3) principal axes (rows).

  • moment_of_inertia: (3,) principal moments.

Return type:

Tuple[np.ndarray, np.ndarray]

get_flipped_axes(UA, custom_axes: ndarray, center_of_mass: ndarray, dimensions: ndarray)[source]

Flip custom axes to a consistent direction with respect to the UA.

For a given set of custom axes, ensure the axes are pointing in the correct direction with respect to the heavy atom position and the chosen center of mass.

Parameters:

  • UA – MDAnalysis AtomGroup for the UA.

  • custom_axes – (3, 3) array of rotation axes.

  • center_of_mass – (3,) COM reference (usually HA position).

  • dimensions – (3,) simulation box dimensions.

Returns:

(3, 3) array of flipped/normalized axes.

Return type:

np.ndarray

get_moment_of_inertia_tensor(center_of_mass: ndarray, positions: ndarray, masses: Sequence[float], dimensions: ndarray) → ndarray[source]

Compute a custom moment of inertia tensor.

Calculate a custom moment of inertia tensor. E.g., for cases where the mass list will contain masses of UAs rather than individual atoms and the positions will be those for the UAs only (excluding the H atoms coordinates).

Parameters:

  • center_of_mass – (3,) chosen centre for the tensor.

  • positions – (N, 3) point positions.

  • masses – (N,) point masses corresponding to positions.

  • dimensions – (3,) simulation box dimensions.

Returns:

(3, 3) moment of inertia tensor.

Return type:

np.ndarray

get_residue_axes(data_container, index: int, residue=None)[source]

Compute residue-level translational and rotational axes.

The translational and rotational axes at the residue level.

  • Identify the residue (either provided or selected by resindex index).

  • Determine whether the residue is bonded to neighboring residues (previous/next in sequence) using MDAnalysis bonded selections.

  • If there are no bonds to other residues:

    • Use a custom principal axes, from a moment-of-inertia (MOI) tensor that uses positions of heavy atoms only, but including masses of heavy atom + bonded hydrogens.

    • Set translational axes equal to rotational axes (as per the original code convention).

  • If bonded to other residues:

    • Use default axes and MOI (MDAnalysis principal axes / inertia).

Parameters:

  • data_container (MDAnalysis.Universe or AtomGroup) – Molecule and trajectory data (the fragment/molecule container).

  • index (int) – Residue index (resindex) within data_container.

  • residue (MDAnalysis.AtomGroup, optional) – If provided, this residue selection will be used rather than selecting again.

Returns:

  • trans_axes: Translational axes array of shape (3, 3).

  • rot_axes: Rotational axes array of shape (3, 3).

  • center: Center of mass array of shape (3,).

  • moment_of_inertia: Principal moments array of shape (3,).

Return type:

Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]

Raises:

ValueError – If the residue selection is empty.

get_vanilla_axes(molecule)[source]

Get principal axes and sorted principal moments (vanilla method).

Compute the principal axes and moments of inertia for a molecule using MDAnalysis built-in functionality.

The original description is preserved: - The molecule is made whole to ensure correct handling of PBC. - The moments are obtained by diagonalising the moment of inertia tensor. - Eigenvalues are returned sorted from largest to smallest magnitude.

Parameters:

molecule (MDAnalysis.core.groups.AtomGroup) – AtomGroup representing the molecule/bead.

Returns:

  • principal_axes: (3, 3) axes.

  • moment_of_inertia: (3,) moments sorted descending by absolute value.

Return type:

Tuple[np.ndarray, np.ndarray]

get_vector(a: ndarray, b: ndarray, dimensions: ndarray)[source]

Compute PBC-wrapped displacement vector(s).

For vector of two coordinates over periodic boundary conditions (PBCs).

Parameters:

  • a – (3,) or (N, 3) array of coordinates.

  • b – (3,) or (N, 3) array of coordinates.

  • dimensions – (3,) simulation box dimensions.

Returns:

Wrapped displacement vector(s) with broadcasted shape.

Return type:

np.ndarray