From 5d65da5b35f6dc4e48bbfbd079bfd086ab772183 Mon Sep 17 00:00:00 2001
From: Jas Kalayan <jkalayan@gmail.com>
Date: Tue, 27 Jan 2026 22:15:28 +0000
Subject: [PATCH 1/9] add modified UA-rovib, modified RES-rovib and dimensions
 in universe ops

---
 CodeEntropy/axes.py                    | 496 +++++++++++++++++++++++++
 CodeEntropy/levels.py                  |  84 +++--
 CodeEntropy/mda_universe_operations.py |  26 +-
 3 files changed, 565 insertions(+), 41 deletions(-)
 create mode 100644 CodeEntropy/axes.py

diff --git a/CodeEntropy/axes.py b/CodeEntropy/axes.py
new file mode 100644
index 0000000..ad0b563
--- /dev/null
+++ b/CodeEntropy/axes.py
@@ -0,0 +1,496 @@
+import logging
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+
+class AxesManager:
+    """
+    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,
+    compute weighted forces and torques, and manipulate matrices used in entropy
+    analysis.
+    """
+
+    def __init__(self):
+        """
+        Initializes the LevelManager with placeholders for level-related data,
+        including translational and rotational axes, number of beads, and a
+        general-purpose data container.
+        """
+        self.data_container = None
+        self._levels = None
+        self._trans_axes = None
+        self._rot_axes = None
+        self._number_of_beads = None
+
+    def get_residue_axes(self, data_container, index):
+        """
+        The translational and rotational axes at the residue level.
+
+        Args:
+          data_container (MDAnalysis.Universe): the molecule and trajectory data
+          index (int): residue index
+
+        Returns:
+          trans_axes : translational axes (3,3)
+          rot_axes : rotational axes (3,3)
+          center: center of mass (3,)
+          moment_of_inertia: moment of inertia (3,)
+        """
+        # TODO refine selection so that it will work for branched polymers
+        index_prev = index - 1
+        index_next = index + 1
+        atom_set = data_container.select_atoms(
+            f"(resindex {index_prev} or resindex {index_next}) "
+            f"and bonded resid {index}"
+        )
+        residue = data_container.select_atoms(f"resindex {index}")
+        # set center of rotation to center of mass of the residue
+        # we might want to change to center of bonding later
+        center = residue.atoms.center_of_mass()
+
+        if len(atom_set) == 0:
+            # No bonds to other residues
+            # Use a custom principal axes, from a MOI tensor
+            # that uses positions of heavy atoms only, but including masses
+            # of heavy atom + bonded hydrogens
+            UAs = residue.select_atoms("mass 2 to 999")
+            UA_masses = self.get_UA_masses(residue)
+            moment_of_inertia_tensor = self.get_moment_of_inertia_tensor(
+                center, UAs.positions, UA_masses
+            )
+            rot_axes, moment_of_inertia = self.get_custom_principal_axes(
+                moment_of_inertia_tensor
+            )
+            trans_axes = (
+                rot_axes  # set trans axes to same as rot axes as per Jon's code
+            )
+        else:
+            # if bonded to other residues, use default axes and MOI
+            trans_axes = data_container.atoms.principal_axes()
+            rot_axes = np.real(residue.principal_axes())
+            eigenvalues, _ = np.linalg.eig(residue.moment_of_inertia())
+            moment_of_inertia = sorted(eigenvalues, reverse=True)
+            center = residue.center_of_mass()
+
+        return trans_axes, rot_axes, center, moment_of_inertia
+
+    def get_UA_axes(self, data_container, index):
+        """
+        The translational and rotational axes at the united-atom level.
+
+        Args:
+          data_container (MDAnalysis.Universe): the molecule and trajectory data
+          index (int): residue index
+
+        Returns:
+          trans_axes : translational axes (3,3)
+          rot_axes : rotational axes (3,3)
+          center: center of mass (3,)
+          moment_of_inertia: moment of inertia (3,)
+        """
+
+        index = int(index)
+
+        trans_axes = data_container.atoms.principal_axes()
+        # look for heavy atoms in residue of interest
+        heavy_atoms = data_container.select_atoms("prop mass > 1.1")
+        heavy_atom_indices = []
+        for atom in heavy_atoms:
+            heavy_atom_indices.append(atom.index)
+        # we find the nth heavy atom
+        # where n is the bead index
+        heavy_atom_index = heavy_atom_indices[index]
+        heavy_atom = data_container.select_atoms(f"index {heavy_atom_index}")
+
+        center = heavy_atom.positions[0]
+        rot_axes, moment_of_inertia = self.get_bonded_axes(
+            data_container, heavy_atom[0], data_container.dimensions[:3]
+        )
+
+        logger.debug(f"Translational Axes: {trans_axes}")
+        logger.debug(f"Rotational Axes: {rot_axes}")
+        logger.debug(f"Center: {center}")
+        logger.debug(f"Moment of Inertia: {moment_of_inertia}")
+
+        return trans_axes, rot_axes, center, moment_of_inertia
+
+    def get_bonded_axes(self, system, atom, dimensions):
+        """
+        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.
+
+        Args:
+            system: mdanalysis instance of all atoms in current frame
+            atom: mdanalysis instance of a heavy atom
+            dimensions: dimensions of the simulation box (3,)
+
+        Returns:
+            custom_axes: custom axes for the UA, (3,3) array
+            custom_moment_of_inertia
+        """
+        # check atom is a heavy atom
+        if not atom.mass > 1.1:
+            return None
+        # set default values
+        custom_moment_of_inertia = None
+        custom_axes = None
+
+        # find the heavy bonded atoms and light bonded atoms
+        heavy_bonded, light_bonded = self.find_bonded_atoms(atom.index, system)
+        UA = atom + light_bonded
+        UA_all = atom + heavy_bonded + light_bonded
+
+        # now find which atoms to select to find the axes for rotating forces:
+        # case1
+        if len(heavy_bonded) == 0:
+            custom_axes, custom_moment_of_inertia = self.get_vanilla_axes(UA_all)
+        # case2
+        if len(heavy_bonded) == 1 and len(light_bonded) == 0:
+            custom_axes = self.get_custom_axes(
+                atom.position, [heavy_bonded[0].position], np.zeros(3), dimensions
+            )
+        # case3
+        if len(heavy_bonded) == 1 and len(light_bonded) >= 1:
+            custom_axes = self.get_custom_axes(
+                atom.position,
+                [heavy_bonded[0].position],
+                light_bonded[0].position,
+                dimensions,
+            )
+        # case4
+        if len(heavy_bonded) == 2:
+            custom_axes = self.get_custom_axes(
+                atom.position,
+                [heavy_bonded[0].position],
+                heavy_bonded[1].position,
+                dimensions,
+            )
+        # case5
+        if len(heavy_bonded) > 2:
+            custom_axes = self.get_custom_axes(
+                atom.position,
+                heavy_bonded.positions,
+                heavy_bonded[1].position,
+                dimensions,
+            )
+
+        # get the moment of inertia from the custom axes
+        if custom_axes is not None:
+            # flip axes to face correct way wrt COM
+            custom_axes = self.get_flipped_axes(
+                UA, custom_axes, atom.position, dimensions
+            )
+        if custom_moment_of_inertia is None:
+            # find moment of inertia using custom axes and atom position as COM
+            custom_moment_of_inertia = self.get_custom_moment_of_inertia(
+                UA, custom_axes, atom.position
+            )
+
+        return custom_axes, custom_moment_of_inertia
+
+    def get_vanilla_axes(self, molecule):
+        """
+        From a selection of atoms, get the ordered principal axes (3,3) and
+        the ordered moment of inertia axes (3,) for that selection of atoms
+
+        Args:
+            molecule: mdanalysis instance of molecule
+            molecule_scale: the length scale of molecule
+
+        Returns:
+            principal_axes: the principal axes, (3,3) array
+            moment_of_inertia: the moment of inertia, (3,) array
+        """
+        # default moment of inertia
+        moment_of_inertia = molecule.moment_of_inertia()
+        principal_axes = molecule.principal_axes()
+        # diagonalise moment of inertia tensor here
+        # pylint: disable=unused-variable
+        eigenvalues, _eigenvectors = np.linalg.eig(moment_of_inertia)
+        # sort eigenvalues of moi tensor by largest to smallest magnitude
+        order = abs(eigenvalues).argsort()[::-1]  # decending order
+        # principal_axes = principal_axes[order] # PI already ordered correctly
+        moment_of_inertia = eigenvalues[order]
+
+        return principal_axes, moment_of_inertia
+
+    def find_bonded_atoms(self, atom_idx: int, system):
+        """
+        for a given atom, find its bonded heavy and H atoms
+
+        Args:
+            atom_idx: atom index to find bonded heavy atom for
+            system: mdanalysis instance of all atoms in current frame
+
+        Returns:
+            bonded_heavy_atoms: MDAnalysis instance of bonded heavy atoms
+            bonded_H_atoms: MDAnalysis instance of bonded hydrogen atoms
+        """
+        bonded_atoms = system.select_atoms(f"bonded index {atom_idx}")
+        bonded_heavy_atoms = bonded_atoms.select_atoms("mass 2 to 999")
+        bonded_H_atoms = bonded_atoms.select_atoms("mass 1 to 1.1")
+        return bonded_heavy_atoms, bonded_H_atoms
+
+    def get_custom_axes(
+        self, a: np.ndarray, b_list: list, c: np.ndarray, dimensions: np.ndarray
+    ):
+        r"""
+        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
+        (represented by |_ symbol below) to axis1.
+
+        - axis3: the cross product of axis1 and axis2, which is perpendicular to
+        axis1 and axis2.
+
+        Args:
+            a: central united-atom coordinates (3,)
+            b_list: list of heavy bonded atom positions (3,N)
+            c: atom coordinates of either a second heavy atom or a hydrogen atom
+            if there are no other bonded heavy atoms in b_list (where N=1 in b_list)
+            (3,)
+            dimensions: dimensions of the simulation box (3,)
+
+        ::
+
+            a          1 = norm_ab
+           / \         2 = |_ norm_ab and norm_ac (use bc if more than 2 HAs)
+          /   \        3 = |_ 1 and 2
+        b       c
+
+        Returns:
+            custom_axes: (3,3) array of the axes used to rotate forces
+        """
+        axis1 = np.zeros(3)
+        # average of all heavy atom covalent bond vectors for axis1
+        for b in b_list:
+            ab_vector = self.get_vector(a, b, dimensions)
+            # scale vector with distance
+            ab_dist = np.sqrt((ab_vector**2).sum(axis=-1))
+            scaled_vector = np.divide(ab_vector, ab_dist)
+            axis1 += scaled_vector  # ab_vector
+        if len(b_list) > 2:
+            # use the first heavy bonded atom and atom c
+            ac_vector = self.get_vector(b_list[0], c, dimensions)
+        else:
+            ac_vector = self.get_vector(a, c, dimensions)
+        ac_dist = np.sqrt((ac_vector**2).sum(axis=-1))
+        ac_vector_norm = np.divide(ac_vector, ac_dist)
+
+        if len(b_list) > 2:
+            axis2 = np.cross(ac_vector_norm, axis1)
+        else:
+            axis2 = np.cross(axis1, ac_vector_norm)
+        axis3 = np.cross(axis1, axis2)
+
+        custom_axes = np.array((axis1, axis2, axis3))
+
+        return custom_axes
+
+    def get_custom_moment_of_inertia(
+        self, UA, custom_rotation_axes: np.ndarray, center_of_mass: np.ndarray
+    ):
+        """
+        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).
+
+        Args:
+            UA: MDAnalysis instance of a united-atom
+            custom_rotation_axes: (3,3) arrray of rotation axes
+            center_of_mass: (3,) center of mass for collection of atoms N
+
+        Returns:
+            custom_moment_of_inertia: (3,) array for moment of inertia
+        """
+        translated_coords = UA.positions - center_of_mass
+        custom_moment_of_inertia = np.zeros(3)
+        for coord, mass in zip(translated_coords, UA.masses):
+            axis_component = np.sum(
+                np.cross(custom_rotation_axes, coord) ** 2 * mass, axis=1
+            )
+            custom_moment_of_inertia += axis_component
+
+        # Remove lowest MOI degree of freedom if UA only has a single bonded H
+        if len(UA) == 2:
+            order = custom_moment_of_inertia.argsort()[::-1]  # decending order
+            custom_moment_of_inertia[order[-1]] = 0
+
+        return custom_moment_of_inertia
+
+    def get_flipped_axes(self, UA, custom_axes, center_of_mass, dimensions):
+        """
+        For a given set of custom axes, ensure the axes are pointing in the
+        correct direction wrt the heavy atom position and the chosen center
+        of mass.
+
+        Args:
+            UA: MDAnalysis instance of a united-atom
+            custom_axes: (3,3) array of the rotation axes
+            center_of_mass: (3,) array for center of mass (usually HA position)
+            dimensions: (3,) array of system box dimensions.
+        """
+        # sorting out PIaxes for MoI for UA fragment
+
+        # get dot product of Paxis1 and CoM->atom1 vect
+        # will just be [0,0,0]
+        RRaxis = self.get_vector(UA[0].position, center_of_mass, dimensions)
+
+        # flip each Paxis if its pointing out of UA
+        custom_axis = np.sum(custom_axes**2, axis=1)
+        custom_axes_flipped = custom_axes / custom_axis**0.5
+        for i in range(3):
+            dotProd1 = np.dot(custom_axes_flipped[i], RRaxis)
+            custom_axes_flipped[i] = np.where(
+                dotProd1 < 0, -custom_axes_flipped[i], custom_axes_flipped[i]
+            )
+        return custom_axes_flipped
+
+    def get_vector(self, a: np.ndarray, b: np.ndarray, dimensions: np.ndarray):
+        """
+        For vector of two coordinates over periodic boundary conditions (PBCs).
+
+        Args:
+            a: (3,) array of atom cooordinates
+            b: (3,) array of atom cooordinates
+            dimensions: (3,) array of system box dimensions.
+
+        Returns:
+            delta_wrapped: (3,) array of the vector between two points
+        """
+        delta = b - a
+        delta_wrapped = []
+        for delt, box in zip(delta, dimensions):
+            if delt < 0 and delt < 0.5 * box:
+                delt = delt + box
+            if delt > 0 and delt > 0.5 * box:
+                delt = delt - box
+            delta_wrapped.append(delt)
+        delta_wrapped = np.array(delta_wrapped)
+
+        return delta_wrapped
+
+    def get_moment_of_inertia_tensor(
+        self, center_of_mass: np.ndarray, positions: np.ndarray, masses: list
+    ) -> np.ndarray:
+        """
+        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 postions will be those for the UAs only
+        (excluding the H atoms coordinates).
+
+        Args:
+            center_of_mass: a (3,) array of the chosen center of mass
+            positions: a (N,3) array of point positions
+            masses: a (N,) list of point masses
+
+        Returns:
+            moment_of_inertia_tensor: a (3,3) moment of inertia tensor
+        """
+        r = positions - center_of_mass
+        r2 = np.sum(r**2, axis=1)
+        moment_of_inertia_tensor = np.eye(3) * np.sum(masses * r2)
+        moment_of_inertia_tensor -= np.einsum("i,ij,ik->jk", masses, r, r)
+
+        return moment_of_inertia_tensor
+
+    def get_custom_principal_axes(
+        self, moment_of_inertia_tensor: np.ndarray
+    ) -> tuple[np.ndarray, np.ndarray]:
+        """
+        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.
+
+        Args:
+            moment_of_inertia_tensor: a (3,3) array of a custom moment of
+             inertia tensor
+
+        Returns:
+            principal_axes: a (3,3) array for the principal axes
+            moment_of_inertia: a (3,) array of the principal axes center
+        """
+        eigenvalues, eigenvectors = np.linalg.eig(moment_of_inertia_tensor)
+        order = abs(eigenvalues).argsort()[::-1]  # decending order
+        transposed = np.transpose(eigenvectors)  # turn columns to rows
+        moment_of_inertia = eigenvalues[order]
+        principal_axes = transposed[order]
+
+        # point z axis in correct direction, as per Jon's code
+        cross_xy = np.cross(principal_axes[0], principal_axes[1])
+        dot_z = np.dot(cross_xy, principal_axes[2])
+        if dot_z < 0:
+            principal_axes[2] *= -1
+
+        logger.debug(f"principal_axes: {principal_axes}")
+
+        return principal_axes, moment_of_inertia
+
+    def get_UA_masses(self, molecule) -> list[float]:
+        """
+        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.
+
+        Args:
+            molecule: mdanalysis instance of molecule
+
+        Returns:
+            UA_masses: list of masses for each UA in a molecule
+        """
+        UA_masses = []
+        for atom in molecule:
+            if atom.mass > 1.1:
+                UA_mass = atom.mass
+                bonded_atoms = molecule.select_atoms(f"bonded index {atom.index}")
+                bonded_H_atoms = bonded_atoms.select_atoms("mass 1 to 1.1")
+                for H in bonded_H_atoms:
+                    UA_mass += H.mass
+                UA_masses.append(UA_mass)
+            else:
+                continue
+        return UA_masses
diff --git a/CodeEntropy/levels.py b/CodeEntropy/levels.py
index 9d43670..6d3652d 100644
--- a/CodeEntropy/levels.py
+++ b/CodeEntropy/levels.py
@@ -9,6 +9,8 @@
     TimeElapsedColumn,
 )
 
+from CodeEntropy.axes import AxesManager
+
 logger = logging.getLogger(__name__)
 
 
@@ -115,11 +117,25 @@ def get_matrices(
         # Calculate forces/torques for each bead
         for bead_index in range(number_beads):
             bead = list_of_beads[bead_index]
+
             # Set up axes
             # translation and rotation use different axes
             # how the axes are defined depends on the level
-            trans_axes = data_container.atoms.principal_axes()
-            rot_axes = np.real(bead.principal_axes())
+            axes_manager = AxesManager()
+            if level == "united_atom":
+                trans_axes, rot_axes, center, moment_of_inertia = (
+                    axes_manager.get_UA_axes(data_container, bead_index)
+                )
+            elif level == "residue":
+                trans_axes, rot_axes, center, moment_of_inertia = (
+                    axes_manager.get_residue_axes(data_container, bead_index)
+                )
+            else:
+                trans_axes = data_container.atoms.principal_axes()
+                rot_axes = np.real(bead.principal_axes())
+                eigenvalues, _ = np.linalg.eig(bead.moment_of_inertia())
+                moment_of_inertia = sorted(eigenvalues, reverse=True)
+                center = bead.center_of_mass()
 
             # Sort out coordinates, forces, and torques for each atom in the bead
             weighted_forces[bead_index] = self.get_weighted_forces(
@@ -130,7 +146,11 @@ def get_matrices(
                 force_partitioning,
             )
             weighted_torques[bead_index] = self.get_weighted_torques(
-                data_container, bead, rot_axes, force_partitioning
+                bead,
+                rot_axes,
+                center,
+                force_partitioning,
+                moment_of_inertia,
             )
 
         # Create covariance submatrices
@@ -167,6 +187,7 @@ def get_matrices(
         # Enforce consistent shape before accumulation
         if force_matrix is None:
             force_matrix = np.zeros_like(force_block)
+            force_matrix = force_block  # add first set of forces
         elif force_matrix.shape != force_block.shape:
             raise ValueError(
                 f"Inconsistent force matrix shape: existing "
@@ -177,6 +198,7 @@ def get_matrices(
 
         if torque_matrix is None:
             torque_matrix = np.zeros_like(torque_block)
+            torque_matrix = torque_block  # add first set of torques
         elif torque_matrix.shape != torque_block.shape:
             raise ValueError(
                 f"Inconsistent torque matrix shape: existing "
@@ -296,7 +318,9 @@ def get_weighted_forces(
 
         return weighted_force
 
-    def get_weighted_torques(self, data_container, bead, rot_axes, force_partitioning):
+    def get_weighted_torques(
+        self, bead, rot_axes, center, force_partitioning, moment_of_inertia
+    ):
         """
         Compute moment-of-inertia weighted torques for a bead.
 
@@ -311,9 +335,7 @@ def get_weighted_torques(self, data_container, bead, rot_axes, force_partitionin
         the sorted eigenvalues of the moment of inertia tensor.
 
         To ensure numerical stability:
-        - Torque components that are effectively zero are skipped.
-        - Zero moments of inertia result in zero weighted torque with a warning.
-        - Negative moments of inertia raise an error.
+        - Torque components that are effectively zero, zero or negative are skipped.
 
         Parameters
         ----------
@@ -326,55 +348,41 @@ def get_weighted_torques(self, data_container, bead, rot_axes, force_partitionin
         force_partitioning : float
             Scaling factor applied to forces to avoid over-counting correlated
             contributions (typically 0.5).
+        moment_of_inertia : np.ndarray
+            Moment of inertia (3,)
 
         Returns
         -------
         weighted_torque : np.ndarray
             Moment-of-inertia weighted torque acting on the bead.
-
-        Raises
-        ------
-        ValueError
-            If a negative principal moment of inertia is encountered.
         """
-        torques = np.zeros((3,))
-        weighted_torque = np.zeros((3,))
-        moment_of_inertia = np.zeros(3)
-
-        for atom in bead.atoms:
-            coords_rot = (
-                data_container.atoms[atom.index].position - bead.center_of_mass()
-            )
-            coords_rot = np.matmul(rot_axes, coords_rot)
-            forces_rot = np.matmul(rot_axes, data_container.atoms[atom.index].force)
-
-            forces_rot = force_partitioning * forces_rot
 
-            torques_local = np.cross(coords_rot, forces_rot)
-            torques += torques_local
-
-        eigenvalues, _ = np.linalg.eig(bead.moment_of_inertia())
-        moments_of_inertia = sorted(eigenvalues, reverse=True)
+        # translate and rotate positions and forces
+        translated_coords = bead.positions - center
+        rotated_coords = np.tensordot(translated_coords, rot_axes.T, axes=1)
+        rotated_forces = np.tensordot(bead.forces, rot_axes.T, axes=1)
+        # scale forces
+        rotated_forces *= force_partitioning
+        # get torques here
+        torques = np.sum(np.cross(rotated_coords, rotated_forces), axis=0)
 
+        weighted_torque = np.zeros((3,))
         for dimension in range(3):
             if np.isclose(torques[dimension], 0):
                 weighted_torque[dimension] = 0
                 continue
 
-            if np.isclose(moments_of_inertia[dimension], 0):
+            if moment_of_inertia[dimension] == 0:
                 weighted_torque[dimension] = 0
-                logger.warning("Zero moment of inertia. Setting torque to 0")
                 continue
 
-            if moments_of_inertia[dimension] < 0:
-                raise ValueError(
-                    f"Negative value encountered for moment of inertia: "
-                    f"{moment_of_inertia[dimension]} "
-                    f"Cannot compute weighted torque."
-                )
+            if moment_of_inertia[dimension] < 0:
+                weighted_torque[dimension] = 0
+                continue
 
+            # Compute weighted torque
             weighted_torque[dimension] = torques[dimension] / np.sqrt(
-                moments_of_inertia[dimension]
+                moment_of_inertia[dimension]
             )
 
         logger.debug(f"Weighted Torque: {weighted_torque}")
diff --git a/CodeEntropy/mda_universe_operations.py b/CodeEntropy/mda_universe_operations.py
index d7f739e..306dfd3 100644
--- a/CodeEntropy/mda_universe_operations.py
+++ b/CodeEntropy/mda_universe_operations.py
@@ -54,8 +54,15 @@ def new_U_select_frame(self, u, start=None, end=None, step=1):
             .run()
             .results["timeseries"][start:end:step]
         )
+        dimensions = (
+            AnalysisFromFunction(lambda ag: ag.dimensions.copy(), select_atom)
+            .run()
+            .results["timeseries"][start:end:step]
+        )
         u2 = mda.Merge(select_atom)
-        u2.load_new(coordinates, format=MemoryReader, forces=forces)
+        u2.load_new(
+            coordinates, format=MemoryReader, forces=forces, dimensions=dimensions
+        )
         logger.debug(f"MDAnalysis.Universe - reduced universe: {u2}")
 
         return u2
@@ -89,8 +96,15 @@ def new_U_select_atom(self, u, select_string="all"):
             .run()
             .results["timeseries"]
         )
+        dimensions = (
+            AnalysisFromFunction(lambda ag: ag.dimensions.copy(), select_atom)
+            .run()
+            .results["timeseries"]
+        )
         u2 = mda.Merge(select_atom)
-        u2.load_new(coordinates, format=MemoryReader, forces=forces)
+        u2.load_new(
+            coordinates, format=MemoryReader, forces=forces, dimensions=dimensions
+        )
         logger.debug(f"MDAnalysis.Universe - reduced universe: {u2}")
 
         return u2
@@ -148,12 +162,18 @@ def merge_forces(self, tprfile, trrfile, forcefile, fileformat=None, kcal=False)
             .results["timeseries"]
         )
 
+        dimensions = (
+            AnalysisFromFunction(lambda ag: ag.dimensions.copy(), select_atom)
+            .run()
+            .results["timeseries"]
+        )
+
         if kcal:
             # Convert from kcal to kJ
             forces *= 4.184
 
         logger.debug("Merging forces with coordinates universe.")
         new_universe = mda.Merge(select_atom)
-        new_universe.load_new(coordinates, forces=forces)
+        new_universe.load_new(coordinates, forces=forces, dimensions=dimensions)
 
         return new_universe

From 993f8274bfe27d015f67e8a1a01aaf52269b6954 Mon Sep 17 00:00:00 2001
From: Jas Kalayan <jkalayan@gmail.com>
Date: Fri, 30 Jan 2026 16:53:15 +0000
Subject: [PATCH 2/9] add combined FTmatrices, match UA bonded axes to Jon's
 code, comment out Sconf calcs for speedier testing

---
 CodeEntropy/axes.py                      |  84 ++++----
 CodeEntropy/config/arg_config_manager.py |   6 +
 CodeEntropy/dihedral_tools.py            | 247 ++++++++++++-----------
 CodeEntropy/entropy.py                   |  97 ++++++---
 CodeEntropy/levels.py                    | 219 +++++++++++++++++---
 5 files changed, 440 insertions(+), 213 deletions(-)

diff --git a/CodeEntropy/axes.py b/CodeEntropy/axes.py
index ad0b563..a3bff16 100644
--- a/CodeEntropy/axes.py
+++ b/CodeEntropy/axes.py
@@ -96,7 +96,18 @@ def get_UA_axes(self, data_container, index):
 
         index = int(index)
 
-        trans_axes = data_container.atoms.principal_axes()
+        # trans_axes = data_container.atoms.principal_axes()
+        # use the same customPI trans axes as the residue level
+        UAs = data_container.select_atoms("mass 2 to 999")
+        UA_masses = self.get_UA_masses(data_container.atoms)
+        center = data_container.atoms.center_of_mass()
+        moment_of_inertia_tensor = self.get_moment_of_inertia_tensor(
+            center, UAs.positions, UA_masses
+        )
+        trans_axes, _moment_of_inertia = self.get_custom_principal_axes(
+            moment_of_inertia_tensor
+        )
+
         # look for heavy atoms in residue of interest
         heavy_atoms = data_container.select_atoms("prop mass > 1.1")
         heavy_atom_indices = []
@@ -174,12 +185,12 @@ def get_bonded_axes(self, system, atom, dimensions):
         # find the heavy bonded atoms and light bonded atoms
         heavy_bonded, light_bonded = self.find_bonded_atoms(atom.index, system)
         UA = atom + light_bonded
-        UA_all = atom + heavy_bonded + light_bonded
+        # UA_all = atom + heavy_bonded + light_bonded
 
         # now find which atoms to select to find the axes for rotating forces:
-        # case1
-        if len(heavy_bonded) == 0:
-            custom_axes, custom_moment_of_inertia = self.get_vanilla_axes(UA_all)
+        # case1, won't apply to UA level
+        # if len(heavy_bonded) == 0:
+        #     custom_axes, custom_moment_of_inertia = self.get_vanilla_axes(UA_all)
         # case2
         if len(heavy_bonded) == 1 and len(light_bonded) == 0:
             custom_axes = self.get_custom_axes(
@@ -193,16 +204,16 @@ def get_bonded_axes(self, system, atom, dimensions):
                 light_bonded[0].position,
                 dimensions,
             )
-        # case4
-        if len(heavy_bonded) == 2:
-            custom_axes = self.get_custom_axes(
-                atom.position,
-                [heavy_bonded[0].position],
-                heavy_bonded[1].position,
-                dimensions,
-            )
+        # case4, not used in Jon's code, use case5 instead
+        # if len(heavy_bonded) == 2:
+        #     custom_axes = self.get_custom_axes(
+        #         atom.position,
+        #         [heavy_bonded[0].position],
+        #         heavy_bonded[1].position,
+        #         dimensions,
+        #     )
         # case5
-        if len(heavy_bonded) > 2:
+        if len(heavy_bonded) >= 2:
             custom_axes = self.get_custom_axes(
                 atom.position,
                 heavy_bonded.positions,
@@ -210,17 +221,18 @@ def get_bonded_axes(self, system, atom, dimensions):
                 dimensions,
             )
 
+        if custom_moment_of_inertia is None:
+            # find moment of inertia using custom axes and atom position as COM
+            custom_moment_of_inertia = self.get_custom_moment_of_inertia(
+                UA, custom_axes, atom.position
+            )
+
         # get the moment of inertia from the custom axes
         if custom_axes is not None:
             # flip axes to face correct way wrt COM
             custom_axes = self.get_flipped_axes(
                 UA, custom_axes, atom.position, dimensions
             )
-        if custom_moment_of_inertia is None:
-            # find moment of inertia using custom axes and atom position as COM
-            custom_moment_of_inertia = self.get_custom_moment_of_inertia(
-                UA, custom_axes, atom.position
-            )
 
         return custom_axes, custom_moment_of_inertia
 
@@ -304,31 +316,27 @@ def get_custom_axes(
         Returns:
             custom_axes: (3,3) array of the axes used to rotate forces
         """
-        axis1 = np.zeros(3)
+        unscaled_axis1 = np.zeros(3)
         # average of all heavy atom covalent bond vectors for axis1
         for b in b_list:
             ab_vector = self.get_vector(a, b, dimensions)
-            # scale vector with distance
-            ab_dist = np.sqrt((ab_vector**2).sum(axis=-1))
-            scaled_vector = np.divide(ab_vector, ab_dist)
-            axis1 += scaled_vector  # ab_vector
-        if len(b_list) > 2:
-            # use the first heavy bonded atom and atom c
-            ac_vector = self.get_vector(b_list[0], c, dimensions)
+            unscaled_axis1 += ab_vector
+        if len(b_list) >= 2:
+            # use the first heavy bonded atom as atom a
+            ac_vector = self.get_vector(c, b_list[0], dimensions)
         else:
-            ac_vector = self.get_vector(a, c, dimensions)
-        ac_dist = np.sqrt((ac_vector**2).sum(axis=-1))
-        ac_vector_norm = np.divide(ac_vector, ac_dist)
+            ac_vector = self.get_vector(c, a, dimensions)
 
-        if len(b_list) > 2:
-            axis2 = np.cross(ac_vector_norm, axis1)
-        else:
-            axis2 = np.cross(axis1, ac_vector_norm)
-        axis3 = np.cross(axis1, axis2)
+        unscaled_axis2 = np.cross(ac_vector, unscaled_axis1)
+        unscaled_axis3 = np.cross(unscaled_axis2, unscaled_axis1)
 
-        custom_axes = np.array((axis1, axis2, axis3))
+        unscaled_custom_axes = np.array(
+            (unscaled_axis1, unscaled_axis2, unscaled_axis3)
+        )
+        mod = np.sqrt(np.sum(unscaled_custom_axes**2, axis=1))
+        scaled_custom_axes = unscaled_custom_axes / mod[:, np.newaxis]
 
-        return custom_axes
+        return scaled_custom_axes
 
     def get_custom_moment_of_inertia(
         self, UA, custom_rotation_axes: np.ndarray, center_of_mass: np.ndarray
@@ -466,8 +474,6 @@ def get_custom_principal_axes(
         if dot_z < 0:
             principal_axes[2] *= -1
 
-        logger.debug(f"principal_axes: {principal_axes}")
-
         return principal_axes, moment_of_inertia
 
     def get_UA_masses(self, molecule) -> list[float]:
diff --git a/CodeEntropy/config/arg_config_manager.py b/CodeEntropy/config/arg_config_manager.py
index e733952..fe51dc6 100644
--- a/CodeEntropy/config/arg_config_manager.py
+++ b/CodeEntropy/config/arg_config_manager.py
@@ -85,6 +85,12 @@
         "help": "How to group molecules for averaging",
         "default": "molecules",
     },
+    "combined_forcetorque": {
+        "type": bool,
+        "help": """Use cobined force-torque matrix for residue
+            level vibrational entropies""",
+        "default": True,
+    },
 }
 
 
diff --git a/CodeEntropy/dihedral_tools.py b/CodeEntropy/dihedral_tools.py
index 1c5d24d..e8137ee 100644
--- a/CodeEntropy/dihedral_tools.py
+++ b/CodeEntropy/dihedral_tools.py
@@ -2,13 +2,14 @@
 
 import numpy as np
 from MDAnalysis.analysis.dihedrals import Dihedral
-from rich.progress import (
-    BarColumn,
-    Progress,
-    SpinnerColumn,
-    TextColumn,
-    TimeElapsedColumn,
-)
+
+# from rich.progress import (
+#     BarColumn,
+#     Progress,
+#     SpinnerColumn,
+#     TextColumn,
+#     TimeElapsedColumn,
+# )
 
 logger = logging.getLogger(__name__)
 
@@ -46,122 +47,122 @@ def build_conformational_states(
         states_ua = {}
         states_res = [None] * number_groups
 
-        total_items = sum(
-            len(levels[mol_id]) for mols in groups.values() for mol_id in mols
-        )
-
-        with Progress(
-            SpinnerColumn(),
-            TextColumn("[bold blue]{task.fields[title]}", justify="right"),
-            BarColumn(),
-            TextColumn("[progress.percentage]{task.percentage:>3.1f}%"),
-            TimeElapsedColumn(),
-        ) as progress:
-
-            task = progress.add_task(
-                "[green]Building Conformational States...",
-                total=total_items,
-                title="Starting...",
-            )
-
-        for group_id in groups.keys():
-            molecules = groups[group_id]
-            mol = self._universe_operations.get_molecule_container(
-                data_container, molecules[0]
-            )
-            num_residues = len(mol.residues)
-            dihedrals_ua = [[] for _ in range(num_residues)]
-            peaks_ua = [{} for _ in range(num_residues)]
-            dihedrals_res = []
-            peaks_res = {}
-
-            # Identify dihedral AtomGroups
-            for level in levels[molecules[0]]:
-                if level == "united_atom":
-                    for res_id in range(num_residues):
-                        selection1 = mol.residues[res_id].atoms.indices[0]
-                        selection2 = mol.residues[res_id].atoms.indices[-1]
-                        res_container = self._universe_operations.new_U_select_atom(
-                            mol,
-                            f"index {selection1}:" f"{selection2}",
-                        )
-                        heavy_res = self._universe_operations.new_U_select_atom(
-                            res_container, "prop mass > 1.1"
-                        )
-
-                        dihedrals_ua[res_id] = self._get_dihedrals(heavy_res, level)
-
-                elif level == "residue":
-                    dihedrals_res = self._get_dihedrals(mol, level)
-
-            # Identify peaks
-            for level in levels[molecules[0]]:
-                if level == "united_atom":
-                    for res_id in range(num_residues):
-                        if len(dihedrals_ua[res_id]) == 0:
-                            # No dihedrals means no histogram or peaks
-                            peaks_ua[res_id] = []
-                        else:
-                            peaks_ua[res_id] = self._identify_peaks(
-                                data_container,
-                                molecules,
-                                dihedrals_ua[res_id],
-                                bin_width,
-                                start,
-                                end,
-                                step,
-                            )
-
-                elif level == "residue":
-                    if len(dihedrals_res) == 0:
-                        # No dihedrals means no histogram or peaks
-                        peaks_res = []
-                    else:
-                        peaks_res = self._identify_peaks(
-                            data_container,
-                            molecules,
-                            dihedrals_res,
-                            bin_width,
-                            start,
-                            end,
-                            step,
-                        )
-
-            # Assign states for each group
-            for level in levels[molecules[0]]:
-                if level == "united_atom":
-                    for res_id in range(num_residues):
-                        key = (group_id, res_id)
-                        if len(dihedrals_ua[res_id]) == 0:
-                            # No conformational states
-                            states_ua[key] = []
-                        else:
-                            states_ua[key] = self._assign_states(
-                                data_container,
-                                molecules,
-                                dihedrals_ua[res_id],
-                                peaks_ua[res_id],
-                                start,
-                                end,
-                                step,
-                            )
-
-                elif level == "residue":
-                    if len(dihedrals_res) == 0:
-                        # No conformational states
-                        states_res[group_id] = []
-                    else:
-                        states_res[group_id] = self._assign_states(
-                            data_container,
-                            molecules,
-                            dihedrals_res,
-                            peaks_res,
-                            start,
-                            end,
-                            step,
-                        )
-
-            progress.advance(task)
+        # SWITCH OFF SCONF
+        # total_items = sum(
+        #     len(levels[mol_id]) for mols in groups.values() for mol_id in mols
+        # )
+        # with Progress(
+        #     SpinnerColumn(),
+        #     TextColumn("[bold blue]{task.fields[title]}", justify="right"),
+        #     BarColumn(),
+        #     TextColumn("[progress.percentage]{task.percentage:>3.1f}%"),
+        #     TimeElapsedColumn(),
+        # ) as progress:
+
+        #     task = progress.add_task(
+        #         "[green]Building Conformational States...",
+        #         total=total_items,
+        #         title="Starting...",
+        #     )
+
+        # for group_id in groups.keys():
+        #     molecules = groups[group_id]
+        #     mol = self._universe_operations.get_molecule_container(
+        #         data_container, molecules[0]
+        #     )
+        #     num_residues = len(mol.residues)
+        #     dihedrals_ua = [[] for _ in range(num_residues)]
+        #     peaks_ua = [{} for _ in range(num_residues)]
+        #     dihedrals_res = []
+        #     peaks_res = {}
+
+        #     # Identify dihedral AtomGroups
+        #     for level in levels[molecules[0]]:
+        #         if level == "united_atom":
+        #             for res_id in range(num_residues):
+        #                 selection1 = mol.residues[res_id].atoms.indices[0]
+        #                 selection2 = mol.residues[res_id].atoms.indices[-1]
+        #                 res_container = self._universe_operations.new_U_select_atom(
+        #                     mol,
+        #                     f"index {selection1}:" f"{selection2}",
+        #                 )
+        #                 heavy_res = self._universe_operations.new_U_select_atom(
+        #                     res_container, "prop mass > 1.1"
+        #                 )
+
+        #                 dihedrals_ua[res_id] = self._get_dihedrals(heavy_res, level)
+
+        #         elif level == "residue":
+        #             dihedrals_res = self._get_dihedrals(mol, level)
+
+        #     # Identify peaks
+        #     for level in levels[molecules[0]]:
+        #         if level == "united_atom":
+        #             for res_id in range(num_residues):
+        #                 if len(dihedrals_ua[res_id]) == 0:
+        #                     # No dihedrals means no histogram or peaks
+        #                     peaks_ua[res_id] = []
+        #                 else:
+        #                     peaks_ua[res_id] = self._identify_peaks(
+        #                         data_container,
+        #                         molecules,
+        #                         dihedrals_ua[res_id],
+        #                         bin_width,
+        #                         start,
+        #                         end,
+        #                         step,
+        #                     )
+
+        #         elif level == "residue":
+        #             if len(dihedrals_res) == 0:
+        #                 # No dihedrals means no histogram or peaks
+        #                 peaks_res = []
+        #             else:
+        #                 peaks_res = self._identify_peaks(
+        #                     data_container,
+        #                     molecules,
+        #                     dihedrals_res,
+        #                     bin_width,
+        #                     start,
+        #                     end,
+        #                     step,
+        #                 )
+
+        #     # Assign states for each group
+        #     for level in levels[molecules[0]]:
+        #         if level == "united_atom":
+        #             for res_id in range(num_residues):
+        #                 key = (group_id, res_id)
+        #                 if len(dihedrals_ua[res_id]) == 0:
+        #                     # No conformational states
+        #                     states_ua[key] = []
+        #                 else:
+        #                     states_ua[key] = self._assign_states(
+        #                         data_container,
+        #                         molecules,
+        #                         dihedrals_ua[res_id],
+        #                         peaks_ua[res_id],
+        #                         start,
+        #                         end,
+        #                         step,
+        #                     )
+
+        #         elif level == "residue":
+        #             if len(dihedrals_res) == 0:
+        #                 # No conformational states
+        #                 states_res[group_id] = []
+        #             else:
+        #                 states_res[group_id] = self._assign_states(
+        #                     data_container,
+        #                     molecules,
+        #                     dihedrals_res,
+        #                     peaks_res,
+        #                     start,
+        #                     end,
+        #                     step,
+        #                 )
+
+        #     progress.advance(task)
 
         return states_ua, states_res
 
diff --git a/CodeEntropy/entropy.py b/CodeEntropy/entropy.py
index c5e72c1..221b33a 100644
--- a/CodeEntropy/entropy.py
+++ b/CodeEntropy/entropy.py
@@ -122,7 +122,7 @@ def execute(self):
         else:
             nonwater_groups.update(water_groups)
 
-        force_matrices, torque_matrices, frame_counts = (
+        force_matrices, torque_matrices, forcetorque_matrices, frame_counts = (
             self._level_manager.build_covariance_matrices(
                 self,
                 reduced_atom,
@@ -133,6 +133,7 @@ def execute(self):
                 step,
                 number_frames,
                 self._args.force_partitioning,
+                self._args.combined_forcetorque,
             )
         )
 
@@ -155,6 +156,7 @@ def execute(self):
             nonwater_groups,
             force_matrices,
             torque_matrices,
+            forcetorque_matrices,
             states_ua,
             states_res,
             frame_counts,
@@ -230,6 +232,7 @@ def _compute_entropies(
         groups,
         force_matrices,
         torque_matrices,
+        forcetorque_matrices,
         states_ua,
         states_res,
         frame_counts,
@@ -309,6 +312,7 @@ def _compute_entropies(
                         f"Level: {level}",
                     )
                     highest = level == levels[groups[group_id][0]][-1]
+                    forcetorque_matrix = None
 
                     if level == "united_atom":
                         self._process_united_atom_entropy(
@@ -326,6 +330,8 @@ def _compute_entropies(
                         )
 
                     elif level == "residue":
+                        if highest:
+                            forcetorque_matrix = forcetorque_matrices["res"][group_id]
                         self._process_vibrational_entropy(
                             group_id,
                             mol,
@@ -334,6 +340,7 @@ def _compute_entropies(
                             level,
                             force_matrices["res"][group_id],
                             torque_matrices["res"][group_id],
+                            forcetorque_matrix,
                             highest,
                         )
 
@@ -347,6 +354,8 @@ def _compute_entropies(
                         )
 
                     elif level == "polymer":
+                        if highest:
+                            forcetorque_matrix = forcetorque_matrices["poly"][group_id]
                         self._process_vibrational_entropy(
                             group_id,
                             mol,
@@ -355,6 +364,7 @@ def _compute_entropies(
                             level,
                             force_matrices["poly"][group_id],
                             torque_matrices["poly"][group_id],
+                            forcetorque_matrix,
                             highest,
                         )
 
@@ -463,21 +473,23 @@ def _process_united_atom_entropy(
                 t_matrix, "torque", self._args.temperature, highest
             )
 
-            # Get the relevant conformational states
-            values = states[key]
-            # Check if there is information in the states array
-            contains_non_empty_states = (
-                np.any(values) if isinstance(values, np.ndarray) else any(values)
-            )
-
-            # Calculate the conformational entropy
-            # If there are no conformational states (i.e. no dihedrals)
-            # then the conformational entropy is zero
-            S_conf_res = (
-                ce.conformational_entropy_calculation(values)
-                if contains_non_empty_states
-                else 0
-            )
+            # SWITCH OFF SCONF
+            # # Get the relevant conformational states
+            # values = states[key]
+            # # Check if there is information in the states array
+            # contains_non_empty_states = (
+            #     np.any(values) if isinstance(values, np.ndarray) else any(values)
+            # )
+
+            # # Calculate the conformational entropy
+            # # If there are no conformational states (i.e. no dihedrals)
+            # # then the conformational entropy is zero
+            # S_conf_res = (
+            #     ce.conformational_entropy_calculation(values)
+            #     if contains_non_empty_states
+            #     else 0
+            # )
+            S_conf_res = 0
 
             # Add the data to the united atom level entropy
             S_trans += S_trans_res
@@ -530,6 +542,7 @@ def _process_vibrational_entropy(
         level,
         force_matrix,
         torque_matrix,
+        forcetorque_matrix,
         highest,
     ):
         """
@@ -548,21 +561,43 @@ def _process_vibrational_entropy(
         # Find the relevant force and torque matrices and tidy them up
         # by removing rows and columns that are all zeros
 
-        force_matrix = self._level_manager.filter_zero_rows_columns(force_matrix)
+        if forcetorque_matrix is not None:
+            forcetorque_matrix = self._level_manager.filter_zero_rows_columns(
+                forcetorque_matrix
+            )
 
-        torque_matrix = self._level_manager.filter_zero_rows_columns(torque_matrix)
+            S_FTtrans = ve.vibrational_entropy_calculation(
+                forcetorque_matrix, "forcetorqueTRANS", self._args.temperature, highest
+            )
+            S_FTrot = ve.vibrational_entropy_calculation(
+                forcetorque_matrix, "forcetorqueROT", self._args.temperature, highest
+            )
 
-        # Calculate the vibrational entropy
-        S_trans = ve.vibrational_entropy_calculation(
-            force_matrix, "force", self._args.temperature, highest
-        )
-        S_rot = ve.vibrational_entropy_calculation(
-            torque_matrix, "torque", self._args.temperature, highest
-        )
+            self._data_logger.add_results_data(
+                group_id, level, "FTmat-Transvibrational", S_FTtrans
+            )
+            self._data_logger.add_results_data(
+                group_id, level, "FTmat-Rovibrational", S_FTrot
+            )
 
-        # Print the vibrational entropy for the molecule group
-        self._data_logger.add_results_data(group_id, level, "Transvibrational", S_trans)
-        self._data_logger.add_results_data(group_id, level, "Rovibrational", S_rot)
+        else:
+            force_matrix = self._level_manager.filter_zero_rows_columns(force_matrix)
+
+            torque_matrix = self._level_manager.filter_zero_rows_columns(torque_matrix)
+
+            # Calculate the vibrational entropy
+            S_trans = ve.vibrational_entropy_calculation(
+                force_matrix, "force", self._args.temperature, highest
+            )
+            S_rot = ve.vibrational_entropy_calculation(
+                torque_matrix, "torque", self._args.temperature, highest
+            )
+
+            # Print the vibrational entropy for the molecule group
+            self._data_logger.add_results_data(
+                group_id, level, "Transvibrational", S_trans
+            )
+            self._data_logger.add_results_data(group_id, level, "Rovibrational", S_rot)
 
         residue_group = "_".join(
             sorted(set(res.resname for res in mol_container.residues))
@@ -899,6 +934,7 @@ def vibrational_entropy_calculation(self, matrix, matrix_type, temp, highest_lev
         """
         # N beads at a level => 3N x 3N covariance matrix => 3N eigenvalues
         # Get eigenvalues of the given matrix and change units to SI units
+        logger.debug(f"matrix_type: {matrix_type}")
         lambdas = la.eigvals(matrix)
         logger.debug(f"Eigenvalues (lambdas) before unit change: {lambdas}")
 
@@ -939,6 +975,11 @@ def vibrational_entropy_calculation(self, matrix, matrix_type, temp, highest_lev
             else:
                 S_vib_total = sum(S_components[6:])
 
+        elif matrix_type == "forcetorqueTRANS":  # three lowest are translations
+            S_vib_total = sum(S_components[:3])
+        elif matrix_type == "forcetorqueROT":  # three highest are rotations
+            S_vib_total = sum(S_components[3:])
+
         else:  # torque covariance matrix - we always take all values into account
             S_vib_total = sum(S_components)
 
diff --git a/CodeEntropy/levels.py b/CodeEntropy/levels.py
index 6d3652d..ae4a227 100644
--- a/CodeEntropy/levels.py
+++ b/CodeEntropy/levels.py
@@ -209,6 +209,112 @@ def get_matrices(
 
         return force_matrix, torque_matrix
 
+    def get_combined_forcetorque_matrices(
+        self,
+        data_container,
+        level,
+        highest_level,
+        forcetorque_matrix,
+        force_partitioning,
+    ):
+        """
+        Compute and accumulate combined force/torque covariance matrices for
+        a given level.
+
+        Parameters:
+          data_container (MDAnalysis.Universe): Data for a molecule or residue.
+          level (str): 'polymer', 'residue', or 'united_atom'.
+          highest_level (bool): Whether this is the top (largest bead size) level.
+          forcetorque_matrix (np.ndarray or None): Accumulated matrices to add
+          to.
+          force_partitioning (float): Factor to adjust force contributions,
+          default is 0.5.
+
+        Returns:
+          forcetorque_matrix (np.ndarray): Accumulated torque covariance matrix.
+        """
+
+        # Make beads
+        list_of_beads = self.get_beads(data_container, level)
+
+        # number of beads and frames in trajectory
+        number_beads = len(list_of_beads)
+
+        # initialize force and torque arrays
+        weighted_forces = [None for _ in range(number_beads)]
+        weighted_torques = [None for _ in range(number_beads)]
+
+        # Calculate forces/torques for each bead
+        for bead_index in range(number_beads):
+            bead = list_of_beads[bead_index]
+
+            # Set up axes
+            # translation and rotation use different axes
+            # how the axes are defined depends on the level
+            axes_manager = AxesManager()
+            if level == "residue":
+                trans_axes, rot_axes, center, moment_of_inertia = (
+                    axes_manager.get_residue_axes(data_container, bead_index)
+                )
+            else:
+                trans_axes = data_container.atoms.principal_axes()
+                rot_axes = np.real(bead.principal_axes())
+                eigenvalues, _ = np.linalg.eig(bead.moment_of_inertia())
+                moment_of_inertia = sorted(eigenvalues, reverse=True)
+                center = bead.center_of_mass()
+
+            # Sort out coordinates, forces, and torques for each atom in the bead
+            weighted_forces[bead_index] = self.get_weighted_forces(
+                data_container,
+                bead,
+                trans_axes,
+                highest_level,
+                force_partitioning,
+            )
+            weighted_torques[bead_index] = self.get_weighted_torques(
+                bead,
+                rot_axes,
+                center,
+                force_partitioning,
+                moment_of_inertia,
+            )
+
+        # Create covariance submatrices
+        forcetorque_submatrix = [
+            [0 for _ in range(number_beads)] for _ in range(number_beads)
+        ]
+
+        for i in range(number_beads):
+            for j in range(i, number_beads):
+                ft_sub = self.create_FTsubmatrix(
+                    np.concatenate((weighted_forces[i], weighted_torques[i])),
+                    np.concatenate((weighted_forces[j], weighted_torques[j])),
+                )
+                forcetorque_submatrix[i][j] = ft_sub
+                forcetorque_submatrix[j][i] = ft_sub.T
+
+        # Convert block matrices to full matrix
+        forcetorque_block = np.block(
+            [
+                [forcetorque_submatrix[i][j] for j in range(number_beads)]
+                for i in range(number_beads)
+            ]
+        )
+
+        # Enforce consistent shape before accumulation
+        if forcetorque_matrix is None:
+            forcetorque_matrix = np.zeros_like(forcetorque_block)
+            forcetorque_matrix = forcetorque_block  # add first set of torques
+        elif forcetorque_matrix.shape != forcetorque_block.shape:
+            raise ValueError(
+                f"Inconsistent forcetorque matrix shape: existing "
+                f"{forcetorque_matrix.shape}, new {forcetorque_block.shape}"
+            )
+        else:
+            forcetorque_matrix = forcetorque_block
+
+        return forcetorque_matrix
+
     def get_beads(self, data_container, level):
         """
         Function to define beads depending on the level in the hierarchy.
@@ -364,7 +470,8 @@ def get_weighted_torques(
         # scale forces
         rotated_forces *= force_partitioning
         # get torques here
-        torques = np.sum(np.cross(rotated_coords, rotated_forces), axis=0)
+        torques = np.cross(rotated_coords, rotated_forces)
+        torques = np.sum(torques, axis=0)
 
         weighted_torque = np.zeros((3,))
         for dimension in range(3):
@@ -415,6 +522,30 @@ def create_submatrix(self, data_i, data_j):
 
         return submatrix
 
+    def create_FTsubmatrix(self, data_i, data_j):
+        """
+        Function for making covariance matrices.
+
+        Args
+        -----
+        data_i : values for bead i
+        data_j : values for bead j
+
+        Returns
+        ------
+        submatrix : 6x6 matrix for the covariance between i and j
+        """
+
+        # Start with 6 by 6 matrix of zeros
+        submatrix = np.zeros((6, 6))
+
+        # For each frame calculate the outer product (cross product) of the data from
+        # the two beads and add the result to the submatrix
+        outer_product_matrix = np.outer(data_i, data_j)
+        submatrix = np.add(submatrix, outer_product_matrix)
+
+        return submatrix
+
     def build_covariance_matrices(
         self,
         entropy_manager,
@@ -426,6 +557,7 @@ def build_covariance_matrices(
         step,
         number_frames,
         force_partitioning,
+        combined_forcetorque,
     ):
         """
         Construct average force and torque covariance matrices for all molecules and
@@ -451,7 +583,8 @@ def build_covariance_matrices(
             Total number of frames to process.
         force_partitioning : float
             Factor to adjust force contributions, default is 0.5.
-
+        combined_forcetorque : bool
+         Whether to use combined forcetorque covariance matrix.
 
         Returns
         -------
@@ -474,6 +607,12 @@ def build_covariance_matrices(
             "poly": [None] * number_groups,
         }
 
+        forcetorque_avg = {
+            "ua": {},
+            "res": [None] * number_groups,
+            "poly": [None] * number_groups,
+        }
+
         total_steps = len(reduced_atom.trajectory[start:end:step])
         total_items = (
             sum(len(levels[mol_id]) for mols in groups.values() for mol_id in mols)
@@ -532,13 +671,15 @@ def build_covariance_matrices(
                                 number_frames,
                                 force_avg,
                                 torque_avg,
+                                forcetorque_avg,
                                 frame_counts,
                                 force_partitioning,
+                                combined_forcetorque,
                             )
 
                             progress.advance(task)
 
-        return force_avg, torque_avg, frame_counts
+        return force_avg, torque_avg, forcetorque_avg, frame_counts
 
     def update_force_torque_matrices(
         self,
@@ -551,8 +692,10 @@ def update_force_torque_matrices(
         num_frames,
         force_avg,
         torque_avg,
+        forcetorque_avg,
         frame_counts,
         force_partitioning,
+        combined_forcetorque,
     ):
         """
         Update the running averages of force and torque covariance matrices
@@ -595,6 +738,8 @@ def update_force_torque_matrices(
             combination.
         force_partitioning : float
          Factor to adjust force contributions, default is 0.5.
+        combined_forcetorque : bool
+         Whether to use combined forcetorque covariance matrix.
         Returns
         -------
         None
@@ -649,28 +794,56 @@ def update_force_torque_matrices(
             # Build the matrices, adding data from each timestep
             # Being careful for the first timestep when data has not yet
             # been added to the matrices
-            f_mat, t_mat = self.get_matrices(
-                mol,
-                level,
-                highest,
-                None if force_avg[key][group_id] is None else force_avg[key][group_id],
-                (
-                    None
-                    if torque_avg[key][group_id] is None
-                    else torque_avg[key][group_id]
-                ),
-                force_partitioning,
-            )
+            if highest and combined_forcetorque:
+                # use combined forcetorque covariance matrix for the highest level only
+                ft_mat = self.get_combined_forcetorque_matrices(
+                    mol,
+                    level,
+                    highest,
+                    (
+                        None
+                        if forcetorque_avg[key][group_id] is None
+                        else forcetorque_avg[key][group_id]
+                    ),
+                    force_partitioning,
+                )
 
-            if force_avg[key][group_id] is None:
-                force_avg[key][group_id] = f_mat.copy()
-                torque_avg[key][group_id] = t_mat.copy()
-                frame_counts[key][group_id] = 1
+                if forcetorque_avg[key][group_id] is None:
+                    forcetorque_avg[key][group_id] = ft_mat.copy()
+                    frame_counts[key][group_id] = 1
+                else:
+                    frame_counts[key][group_id] += 1
+                    n = frame_counts[key][group_id]
+                    forcetorque_avg[key][group_id] += (
+                        ft_mat - forcetorque_avg[key][group_id]
+                    ) / n
             else:
-                frame_counts[key][group_id] += 1
-                n = frame_counts[key][group_id]
-                force_avg[key][group_id] += (f_mat - force_avg[key][group_id]) / n
-                torque_avg[key][group_id] += (t_mat - torque_avg[key][group_id]) / n
+                f_mat, t_mat = self.get_matrices(
+                    mol,
+                    level,
+                    highest,
+                    (
+                        None
+                        if force_avg[key][group_id] is None
+                        else force_avg[key][group_id]
+                    ),
+                    (
+                        None
+                        if torque_avg[key][group_id] is None
+                        else torque_avg[key][group_id]
+                    ),
+                    force_partitioning,
+                )
+
+                if force_avg[key][group_id] is None:
+                    force_avg[key][group_id] = f_mat.copy()
+                    torque_avg[key][group_id] = t_mat.copy()
+                    frame_counts[key][group_id] = 1
+                else:
+                    frame_counts[key][group_id] += 1
+                    n = frame_counts[key][group_id]
+                    force_avg[key][group_id] += (f_mat - force_avg[key][group_id]) / n
+                    torque_avg[key][group_id] += (t_mat - torque_avg[key][group_id]) / n
 
         return frame_counts
 

From 316f8d47982a07e3f617f8844069b0a5dccd9bc7 Mon Sep 17 00:00:00 2001
From: Jas Kalayan <jkalayan@gmail.com>
Date: Fri, 30 Jan 2026 20:40:29 +0000
Subject: [PATCH 3/9] add arg to choose between axes schemes

---
 CodeEntropy/config/arg_config_manager.py |  8 +++++++-
 CodeEntropy/entropy.py                   |  1 +
 CodeEntropy/levels.py                    | 20 +++++++++++++++-----
 3 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/CodeEntropy/config/arg_config_manager.py b/CodeEntropy/config/arg_config_manager.py
index fe51dc6..813c6a3 100644
--- a/CodeEntropy/config/arg_config_manager.py
+++ b/CodeEntropy/config/arg_config_manager.py
@@ -87,7 +87,13 @@
     },
     "combined_forcetorque": {
         "type": bool,
-        "help": """Use cobined force-torque matrix for residue
+        "help": """Use combined force-torque matrix for residue
+            level vibrational entropies""",
+        "default": True,
+    },
+    "customised_axes": {
+        "type": bool,
+        "help": """Use bonded axes to rotate forces for UA
             level vibrational entropies""",
         "default": True,
     },
diff --git a/CodeEntropy/entropy.py b/CodeEntropy/entropy.py
index 221b33a..a4cbe0d 100644
--- a/CodeEntropy/entropy.py
+++ b/CodeEntropy/entropy.py
@@ -134,6 +134,7 @@ def execute(self):
                 number_frames,
                 self._args.force_partitioning,
                 self._args.combined_forcetorque,
+                self._args.customised_axes,
             )
         )
 
diff --git a/CodeEntropy/levels.py b/CodeEntropy/levels.py
index ae4a227..fe2ac15 100644
--- a/CodeEntropy/levels.py
+++ b/CodeEntropy/levels.py
@@ -86,6 +86,7 @@ def get_matrices(
         force_matrix,
         torque_matrix,
         force_partitioning,
+        customised_axes,
     ):
         """
         Compute and accumulate force/torque covariance matrices for a given level.
@@ -122,11 +123,11 @@ def get_matrices(
             # translation and rotation use different axes
             # how the axes are defined depends on the level
             axes_manager = AxesManager()
-            if level == "united_atom":
+            if level == "united_atom" and customised_axes:
                 trans_axes, rot_axes, center, moment_of_inertia = (
                     axes_manager.get_UA_axes(data_container, bead_index)
                 )
-            elif level == "residue":
+            elif level == "residue" and customised_axes:
                 trans_axes, rot_axes, center, moment_of_inertia = (
                     axes_manager.get_residue_axes(data_container, bead_index)
                 )
@@ -216,6 +217,7 @@ def get_combined_forcetorque_matrices(
         highest_level,
         forcetorque_matrix,
         force_partitioning,
+        customised_axes,
     ):
         """
         Compute and accumulate combined force/torque covariance matrices for
@@ -252,7 +254,7 @@ def get_combined_forcetorque_matrices(
             # translation and rotation use different axes
             # how the axes are defined depends on the level
             axes_manager = AxesManager()
-            if level == "residue":
+            if level == "residue" and customised_axes:
                 trans_axes, rot_axes, center, moment_of_inertia = (
                     axes_manager.get_residue_axes(data_container, bead_index)
                 )
@@ -558,6 +560,7 @@ def build_covariance_matrices(
         number_frames,
         force_partitioning,
         combined_forcetorque,
+        customised_axes,
     ):
         """
         Construct average force and torque covariance matrices for all molecules and
@@ -675,6 +678,7 @@ def build_covariance_matrices(
                                 frame_counts,
                                 force_partitioning,
                                 combined_forcetorque,
+                                customised_axes,
                             )
 
                             progress.advance(task)
@@ -696,6 +700,7 @@ def update_force_torque_matrices(
         frame_counts,
         force_partitioning,
         combined_forcetorque,
+        customised_axes,
     ):
         """
         Update the running averages of force and torque covariance matrices
@@ -737,9 +742,11 @@ def update_force_torque_matrices(
             Dictionary holding the count of frames processed for each molecule/level
             combination.
         force_partitioning : float
-         Factor to adjust force contributions, default is 0.5.
+            Factor to adjust force contributions, default is 0.5.
         combined_forcetorque : bool
-         Whether to use combined forcetorque covariance matrix.
+            Whether to use combined forcetorque covariance matrix.
+        customised_axes: bool
+            Whether to use bonded axes for UA rovib calculations
         Returns
         -------
         None
@@ -772,6 +779,7 @@ def update_force_torque_matrices(
                     None if key not in force_avg["ua"] else force_avg["ua"][key],
                     None if key not in torque_avg["ua"] else torque_avg["ua"][key],
                     force_partitioning,
+                    customised_axes,
                 )
 
                 if key not in force_avg["ua"]:
@@ -806,6 +814,7 @@ def update_force_torque_matrices(
                         else forcetorque_avg[key][group_id]
                     ),
                     force_partitioning,
+                    customised_axes,
                 )
 
                 if forcetorque_avg[key][group_id] is None:
@@ -833,6 +842,7 @@ def update_force_torque_matrices(
                         else torque_avg[key][group_id]
                     ),
                     force_partitioning,
+                    customised_axes,
                 )
 
                 if force_avg[key][group_id] is None:

From e4a41f46256cb7da5bf1106ed7cb259dc1a15a8a Mon Sep 17 00:00:00 2001
From: Jas Kalayan <jkalayan@gmail.com>
Date: Mon, 2 Feb 2026 10:28:00 +0000
Subject: [PATCH 4/9] uncomment Sconf calcs

---
 CodeEntropy/dihedral_tools.py | 246 +++++++++++++++++-----------------
 CodeEntropy/entropy.py        |  32 +++--
 2 files changed, 137 insertions(+), 141 deletions(-)

diff --git a/CodeEntropy/dihedral_tools.py b/CodeEntropy/dihedral_tools.py
index e8137ee..036fcf0 100644
--- a/CodeEntropy/dihedral_tools.py
+++ b/CodeEntropy/dihedral_tools.py
@@ -2,14 +2,13 @@
 
 import numpy as np
 from MDAnalysis.analysis.dihedrals import Dihedral
-
-# from rich.progress import (
-#     BarColumn,
-#     Progress,
-#     SpinnerColumn,
-#     TextColumn,
-#     TimeElapsedColumn,
-# )
+from rich.progress import (
+    BarColumn,
+    Progress,
+    SpinnerColumn,
+    TextColumn,
+    TimeElapsedColumn,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -47,122 +46,121 @@ def build_conformational_states(
         states_ua = {}
         states_res = [None] * number_groups
 
-        # SWITCH OFF SCONF
-        # total_items = sum(
-        #     len(levels[mol_id]) for mols in groups.values() for mol_id in mols
-        # )
-        # with Progress(
-        #     SpinnerColumn(),
-        #     TextColumn("[bold blue]{task.fields[title]}", justify="right"),
-        #     BarColumn(),
-        #     TextColumn("[progress.percentage]{task.percentage:>3.1f}%"),
-        #     TimeElapsedColumn(),
-        # ) as progress:
-
-        #     task = progress.add_task(
-        #         "[green]Building Conformational States...",
-        #         total=total_items,
-        #         title="Starting...",
-        #     )
-
-        # for group_id in groups.keys():
-        #     molecules = groups[group_id]
-        #     mol = self._universe_operations.get_molecule_container(
-        #         data_container, molecules[0]
-        #     )
-        #     num_residues = len(mol.residues)
-        #     dihedrals_ua = [[] for _ in range(num_residues)]
-        #     peaks_ua = [{} for _ in range(num_residues)]
-        #     dihedrals_res = []
-        #     peaks_res = {}
-
-        #     # Identify dihedral AtomGroups
-        #     for level in levels[molecules[0]]:
-        #         if level == "united_atom":
-        #             for res_id in range(num_residues):
-        #                 selection1 = mol.residues[res_id].atoms.indices[0]
-        #                 selection2 = mol.residues[res_id].atoms.indices[-1]
-        #                 res_container = self._universe_operations.new_U_select_atom(
-        #                     mol,
-        #                     f"index {selection1}:" f"{selection2}",
-        #                 )
-        #                 heavy_res = self._universe_operations.new_U_select_atom(
-        #                     res_container, "prop mass > 1.1"
-        #                 )
-
-        #                 dihedrals_ua[res_id] = self._get_dihedrals(heavy_res, level)
-
-        #         elif level == "residue":
-        #             dihedrals_res = self._get_dihedrals(mol, level)
-
-        #     # Identify peaks
-        #     for level in levels[molecules[0]]:
-        #         if level == "united_atom":
-        #             for res_id in range(num_residues):
-        #                 if len(dihedrals_ua[res_id]) == 0:
-        #                     # No dihedrals means no histogram or peaks
-        #                     peaks_ua[res_id] = []
-        #                 else:
-        #                     peaks_ua[res_id] = self._identify_peaks(
-        #                         data_container,
-        #                         molecules,
-        #                         dihedrals_ua[res_id],
-        #                         bin_width,
-        #                         start,
-        #                         end,
-        #                         step,
-        #                     )
-
-        #         elif level == "residue":
-        #             if len(dihedrals_res) == 0:
-        #                 # No dihedrals means no histogram or peaks
-        #                 peaks_res = []
-        #             else:
-        #                 peaks_res = self._identify_peaks(
-        #                     data_container,
-        #                     molecules,
-        #                     dihedrals_res,
-        #                     bin_width,
-        #                     start,
-        #                     end,
-        #                     step,
-        #                 )
-
-        #     # Assign states for each group
-        #     for level in levels[molecules[0]]:
-        #         if level == "united_atom":
-        #             for res_id in range(num_residues):
-        #                 key = (group_id, res_id)
-        #                 if len(dihedrals_ua[res_id]) == 0:
-        #                     # No conformational states
-        #                     states_ua[key] = []
-        #                 else:
-        #                     states_ua[key] = self._assign_states(
-        #                         data_container,
-        #                         molecules,
-        #                         dihedrals_ua[res_id],
-        #                         peaks_ua[res_id],
-        #                         start,
-        #                         end,
-        #                         step,
-        #                     )
-
-        #         elif level == "residue":
-        #             if len(dihedrals_res) == 0:
-        #                 # No conformational states
-        #                 states_res[group_id] = []
-        #             else:
-        #                 states_res[group_id] = self._assign_states(
-        #                     data_container,
-        #                     molecules,
-        #                     dihedrals_res,
-        #                     peaks_res,
-        #                     start,
-        #                     end,
-        #                     step,
-        #                 )
-
-        #     progress.advance(task)
+        total_items = sum(
+            len(levels[mol_id]) for mols in groups.values() for mol_id in mols
+        )
+        with Progress(
+            SpinnerColumn(),
+            TextColumn("[bold blue]{task.fields[title]}", justify="right"),
+            BarColumn(),
+            TextColumn("[progress.percentage]{task.percentage:>3.1f}%"),
+            TimeElapsedColumn(),
+        ) as progress:
+
+            task = progress.add_task(
+                "[green]Building Conformational States...",
+                total=total_items,
+                title="Starting...",
+            )
+
+        for group_id in groups.keys():
+            molecules = groups[group_id]
+            mol = self._universe_operations.get_molecule_container(
+                data_container, molecules[0]
+            )
+            num_residues = len(mol.residues)
+            dihedrals_ua = [[] for _ in range(num_residues)]
+            peaks_ua = [{} for _ in range(num_residues)]
+            dihedrals_res = []
+            peaks_res = {}
+
+            # Identify dihedral AtomGroups
+            for level in levels[molecules[0]]:
+                if level == "united_atom":
+                    for res_id in range(num_residues):
+                        selection1 = mol.residues[res_id].atoms.indices[0]
+                        selection2 = mol.residues[res_id].atoms.indices[-1]
+                        res_container = self._universe_operations.new_U_select_atom(
+                            mol,
+                            f"index {selection1}:" f"{selection2}",
+                        )
+                        heavy_res = self._universe_operations.new_U_select_atom(
+                            res_container, "prop mass > 1.1"
+                        )
+
+                        dihedrals_ua[res_id] = self._get_dihedrals(heavy_res, level)
+
+                elif level == "residue":
+                    dihedrals_res = self._get_dihedrals(mol, level)
+
+            # Identify peaks
+            for level in levels[molecules[0]]:
+                if level == "united_atom":
+                    for res_id in range(num_residues):
+                        if len(dihedrals_ua[res_id]) == 0:
+                            # No dihedrals means no histogram or peaks
+                            peaks_ua[res_id] = []
+                        else:
+                            peaks_ua[res_id] = self._identify_peaks(
+                                data_container,
+                                molecules,
+                                dihedrals_ua[res_id],
+                                bin_width,
+                                start,
+                                end,
+                                step,
+                            )
+
+                elif level == "residue":
+                    if len(dihedrals_res) == 0:
+                        # No dihedrals means no histogram or peaks
+                        peaks_res = []
+                    else:
+                        peaks_res = self._identify_peaks(
+                            data_container,
+                            molecules,
+                            dihedrals_res,
+                            bin_width,
+                            start,
+                            end,
+                            step,
+                        )
+
+            # Assign states for each group
+            for level in levels[molecules[0]]:
+                if level == "united_atom":
+                    for res_id in range(num_residues):
+                        key = (group_id, res_id)
+                        if len(dihedrals_ua[res_id]) == 0:
+                            # No conformational states
+                            states_ua[key] = []
+                        else:
+                            states_ua[key] = self._assign_states(
+                                data_container,
+                                molecules,
+                                dihedrals_ua[res_id],
+                                peaks_ua[res_id],
+                                start,
+                                end,
+                                step,
+                            )
+
+                elif level == "residue":
+                    if len(dihedrals_res) == 0:
+                        # No conformational states
+                        states_res[group_id] = []
+                    else:
+                        states_res[group_id] = self._assign_states(
+                            data_container,
+                            molecules,
+                            dihedrals_res,
+                            peaks_res,
+                            start,
+                            end,
+                            step,
+                        )
+
+            progress.advance(task)
 
         return states_ua, states_res
 
diff --git a/CodeEntropy/entropy.py b/CodeEntropy/entropy.py
index a4cbe0d..9b7016e 100644
--- a/CodeEntropy/entropy.py
+++ b/CodeEntropy/entropy.py
@@ -474,23 +474,21 @@ def _process_united_atom_entropy(
                 t_matrix, "torque", self._args.temperature, highest
             )
 
-            # SWITCH OFF SCONF
-            # # Get the relevant conformational states
-            # values = states[key]
-            # # Check if there is information in the states array
-            # contains_non_empty_states = (
-            #     np.any(values) if isinstance(values, np.ndarray) else any(values)
-            # )
-
-            # # Calculate the conformational entropy
-            # # If there are no conformational states (i.e. no dihedrals)
-            # # then the conformational entropy is zero
-            # S_conf_res = (
-            #     ce.conformational_entropy_calculation(values)
-            #     if contains_non_empty_states
-            #     else 0
-            # )
-            S_conf_res = 0
+            # Get the relevant conformational states
+            values = states[key]
+            # Check if there is information in the states array
+            contains_non_empty_states = (
+                np.any(values) if isinstance(values, np.ndarray) else any(values)
+            )
+
+            # Calculate the conformational entropy
+            # If there are no conformational states (i.e. no dihedrals)
+            # then the conformational entropy is zero
+            S_conf_res = (
+                ce.conformational_entropy_calculation(values)
+                if contains_non_empty_states
+                else 0
+            )
 
             # Add the data to the united atom level entropy
             S_trans += S_trans_res

From 3a8295b3df2cf9909ba988e83e1506578cd6a29a Mon Sep 17 00:00:00 2001
From: Jas Kalayan <jkalayan@gmail.com>
Date: Mon, 2 Feb 2026 10:29:44 +0000
Subject: [PATCH 5/9] match main line space

---
 CodeEntropy/dihedral_tools.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/CodeEntropy/dihedral_tools.py b/CodeEntropy/dihedral_tools.py
index 036fcf0..1c5d24d 100644
--- a/CodeEntropy/dihedral_tools.py
+++ b/CodeEntropy/dihedral_tools.py
@@ -49,6 +49,7 @@ def build_conformational_states(
         total_items = sum(
             len(levels[mol_id]) for mols in groups.values() for mol_id in mols
         )
+
         with Progress(
             SpinnerColumn(),
             TextColumn("[bold blue]{task.fields[title]}", justify="right"),

From 8a6ecdd5f3829004eebec15901a0a7b38b9bf100 Mon Sep 17 00:00:00 2001
From: Jas Kalayan <jkalayan@gmail.com>
Date: Mon, 2 Feb 2026 11:39:04 +0000
Subject: [PATCH 6/9] remove redundant axes methods

---
 CodeEntropy/axes.py | 36 ------------------------------------
 1 file changed, 36 deletions(-)

diff --git a/CodeEntropy/axes.py b/CodeEntropy/axes.py
index a3bff16..facd54b 100644
--- a/CodeEntropy/axes.py
+++ b/CodeEntropy/axes.py
@@ -185,12 +185,9 @@ def get_bonded_axes(self, system, atom, dimensions):
         # find the heavy bonded atoms and light bonded atoms
         heavy_bonded, light_bonded = self.find_bonded_atoms(atom.index, system)
         UA = atom + light_bonded
-        # UA_all = atom + heavy_bonded + light_bonded
 
         # now find which atoms to select to find the axes for rotating forces:
         # case1, won't apply to UA level
-        # if len(heavy_bonded) == 0:
-        #     custom_axes, custom_moment_of_inertia = self.get_vanilla_axes(UA_all)
         # case2
         if len(heavy_bonded) == 1 and len(light_bonded) == 0:
             custom_axes = self.get_custom_axes(
@@ -205,13 +202,6 @@ def get_bonded_axes(self, system, atom, dimensions):
                 dimensions,
             )
         # case4, not used in Jon's code, use case5 instead
-        # if len(heavy_bonded) == 2:
-        #     custom_axes = self.get_custom_axes(
-        #         atom.position,
-        #         [heavy_bonded[0].position],
-        #         heavy_bonded[1].position,
-        #         dimensions,
-        #     )
         # case5
         if len(heavy_bonded) >= 2:
             custom_axes = self.get_custom_axes(
@@ -236,32 +226,6 @@ def get_bonded_axes(self, system, atom, dimensions):
 
         return custom_axes, custom_moment_of_inertia
 
-    def get_vanilla_axes(self, molecule):
-        """
-        From a selection of atoms, get the ordered principal axes (3,3) and
-        the ordered moment of inertia axes (3,) for that selection of atoms
-
-        Args:
-            molecule: mdanalysis instance of molecule
-            molecule_scale: the length scale of molecule
-
-        Returns:
-            principal_axes: the principal axes, (3,3) array
-            moment_of_inertia: the moment of inertia, (3,) array
-        """
-        # default moment of inertia
-        moment_of_inertia = molecule.moment_of_inertia()
-        principal_axes = molecule.principal_axes()
-        # diagonalise moment of inertia tensor here
-        # pylint: disable=unused-variable
-        eigenvalues, _eigenvectors = np.linalg.eig(moment_of_inertia)
-        # sort eigenvalues of moi tensor by largest to smallest magnitude
-        order = abs(eigenvalues).argsort()[::-1]  # decending order
-        # principal_axes = principal_axes[order] # PI already ordered correctly
-        moment_of_inertia = eigenvalues[order]
-
-        return principal_axes, moment_of_inertia
-
     def find_bonded_atoms(self, atom_idx: int, system):
         """
         for a given atom, find its bonded heavy and H atoms

From 5d076f12f51713b5490e7c47e5997d3ea14ab26d Mon Sep 17 00:00:00 2001
From: harryswift01 <harry.swift@stfc.ac.uk>
Date: Mon, 2 Feb 2026 16:19:56 +0000
Subject: [PATCH 7/9] fix unit tests within `test_entropy.py` in relation to
 changes within PR #267

---
 tests/test_CodeEntropy/test_entropy.py | 33 ++++++++++++++++++++++----
 1 file changed, 29 insertions(+), 4 deletions(-)

diff --git a/tests/test_CodeEntropy/test_entropy.py b/tests/test_CodeEntropy/test_entropy.py
index 8fce98f..c8b43c0 100644
--- a/tests/test_CodeEntropy/test_entropy.py
+++ b/tests/test_CodeEntropy/test_entropy.py
@@ -83,7 +83,12 @@ def test_execute_full_workflow(self):
             return_value=(mock_reduced_atom, 3, mock_levels, mock_groups)
         )
         entropy_manager._level_manager.build_covariance_matrices = MagicMock(
-            return_value=("force_matrices", "torque_matrices", "frame_counts")
+            return_value=(
+                "force_matrices",
+                "torque_matrices",
+                "forcetorque_avg",
+                "frame_counts",
+            )
         )
         entropy_manager._dihedral_analysis.build_conformational_states = MagicMock(
             return_value=(["state_ua"], ["state_res"])
@@ -128,6 +133,7 @@ def test_execute_full_workflow(self):
             mock_groups,
             "force_matrices",
             "torque_matrices",
+            "forcetorque_avg",
             ["state_ua"],
             ["state_res"],
             "frame_counts",
@@ -173,7 +179,12 @@ def test_execute_triggers_handle_water_entropy_minimal(self):
             return_value=(MagicMock(), 3, {}, {0: [0]})
         )
         entropy_manager._level_manager.build_covariance_matrices = MagicMock(
-            return_value=("force_matrices", "torque_matrices", "frame_counts")
+            return_value=(
+                "force_matrices",
+                "torque_matrices",
+                "forcetorque_avg",
+                "frame_counts",
+            )
         )
         entropy_manager._dihedral_analysis.build_conformational_states = MagicMock(
             return_value=(["state_ua"], ["state_res"])
@@ -715,6 +726,8 @@ def test_process_vibrational_only_levels(self):
         ve = MagicMock()
         ve.vibrational_entropy_calculation.side_effect = [1.11, 2.22]
 
+        forcetorque_matrix = np.eye(6)
+
         # Run the method
         manager._process_vibrational_entropy(
             group_id=0,
@@ -724,6 +737,7 @@ def test_process_vibrational_only_levels(self):
             level="Vibrational",
             force_matrix=force_matrix,
             torque_matrix=torque_matrix,
+            forcetorque_matrix=forcetorque_matrix,
             highest=True,
         )
 
@@ -731,7 +745,7 @@ def test_process_vibrational_only_levels(self):
         df = data_logger.molecule_data
         self.assertEqual(len(df), 2)  # Transvibrational and Rovibrational
 
-        expected_types = {"Transvibrational", "Rovibrational"}
+        expected_types = {"FTmat-Transvibrational", "FTmat-Rovibrational"}
         actual_types = set(entry[2] for entry in df)
         self.assertSetEqual(actual_types, expected_types)
 
@@ -789,6 +803,7 @@ def test_compute_entropies_polymer_branch(self):
             groups,
             force_matrices,
             torque_matrices,
+            force_matrices,
             states_ua,
             states_res,
             frame_counts,
@@ -942,6 +957,8 @@ def test_compute_entropies_united_atom(self):
         universe_operations.get_molecule_container = MagicMock(return_value=mol_mock)
         manager._process_united_atom_entropy = MagicMock()
 
+        force_torque_matrices = MagicMock()
+
         ve = MagicMock()
         ce = MagicMock()
 
@@ -951,6 +968,7 @@ def test_compute_entropies_united_atom(self):
             groups,
             force_matrices,
             torque_matrices,
+            force_torque_matrices,
             states_ua,
             states_res,
             frame_counts,
@@ -1017,6 +1035,8 @@ def test_compute_entropies_residue(self):
         manager._process_vibrational_entropy = MagicMock()
         manager._process_conformational_entropy = MagicMock()
 
+        force_torque_matrices = MagicMock()
+
         # Mock entropy calculators
         ve = MagicMock()
         ce = MagicMock()
@@ -1028,6 +1048,7 @@ def test_compute_entropies_residue(self):
             groups,
             force_matrices,
             torque_matrices,
+            force_torque_matrices,
             states_ua,
             states_res,
             frame_counts,
@@ -1048,6 +1069,7 @@ def test_compute_entropies_polymer(self):
         data_logger = DataLogger()
         group_molecules = MagicMock()
         dihedral_analysis = MagicMock()
+
         manager = EntropyManager(
             run_manager,
             args,
@@ -1066,9 +1088,10 @@ def test_compute_entropies_polymer(self):
 
         force_matrices = {"poly": {0: "force_poly"}}
         torque_matrices = {"poly": {0: "torque_poly"}}
+        force_torque_matrices = {"poly": {0: "ft_poly"}}
+
         states_ua = {}
         states_res = []
-
         frame_counts = {"poly": {(0, 0): 10}}
 
         mol_mock = MagicMock()
@@ -1085,6 +1108,7 @@ def test_compute_entropies_polymer(self):
             groups,
             force_matrices,
             torque_matrices,
+            force_torque_matrices,
             states_ua,
             states_res,
             frame_counts,
@@ -1101,6 +1125,7 @@ def test_compute_entropies_polymer(self):
             "polymer",
             force_matrices["poly"][0],
             torque_matrices["poly"][0],
+            force_torque_matrices["poly"][0],
             True,
         )
 

From 59c93993fad5a8998a031d3cf2d235b9a5e7b7ce Mon Sep 17 00:00:00 2001
From: harryswift01 <harry.swift@stfc.ac.uk>
Date: Mon, 2 Feb 2026 16:54:14 +0000
Subject: [PATCH 8/9] fix unit tests within `test_levels.py` in relation to
 changes within PR #267

---
 tests/test_CodeEntropy/test_levels.py | 412 +++++++++++++++-----------
 1 file changed, 241 insertions(+), 171 deletions(-)

diff --git a/tests/test_CodeEntropy/test_levels.py b/tests/test_CodeEntropy/test_levels.py
index 98a2684..9801f38 100644
--- a/tests/test_CodeEntropy/test_levels.py
+++ b/tests/test_CodeEntropy/test_levels.py
@@ -1,4 +1,4 @@
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
 import numpy as np
 
@@ -60,12 +60,15 @@ def test_select_levels(self):
 
     def test_get_matrices(self):
         """
-        Test `get_matrices` with mocked internal methods and a simple setup.
-        Ensures that the method returns correctly shaped force and torque matrices.
+        Atomic unit test for LevelManager.get_matrices:
+        - AxesManager is mocked
+        - No inertia / MDAnalysis math
+        - Verifies block matrix construction and shape only
         """
         universe_operations = UniverseOperations()
         level_manager = LevelManager(universe_operations)
 
+        # Two beads
         bead1 = MagicMock()
         bead1.principal_axes.return_value = np.ones(3)
 
@@ -73,39 +76,62 @@ def test_get_matrices(self):
         bead2.principal_axes.return_value = np.ones(3)
 
         level_manager.get_beads = MagicMock(return_value=[bead1, bead2])
+
         level_manager.get_weighted_forces = MagicMock(
             return_value=np.array([1.0, 2.0, 3.0])
         )
         level_manager.get_weighted_torques = MagicMock(
             return_value=np.array([0.5, 1.5, 2.5])
         )
-        level_manager.create_submatrix = MagicMock(return_value=np.identity(3))
+
+        # Deterministic 3x3 submatrix for every (i,j) call
+        I3 = np.identity(3)
+        level_manager.create_submatrix = MagicMock(return_value=I3)
 
         data_container = MagicMock()
         data_container.atoms = MagicMock()
         data_container.atoms.principal_axes.return_value = np.ones(3)
 
-        force_matrix, torque_matrix = level_manager.get_matrices(
-            data_container=data_container,
-            level="residue",
-            highest_level=True,
-            force_matrix=None,
-            torque_matrix=None,
-            force_partitioning=0.5,
-        )
+        dummy_trans_axes = np.eye(3)
+        dummy_rot_axes = np.eye(3)
+        dummy_center = np.zeros(3)
+        dummy_moi = np.eye(3)
+
+        with patch("CodeEntropy.levels.AxesManager") as AxesManagerMock:
+            axes = AxesManagerMock.return_value
+            axes.get_residue_axes.return_value = (
+                dummy_trans_axes,
+                dummy_rot_axes,
+                dummy_center,
+                dummy_moi,
+            )
 
-        self.assertIsInstance(force_matrix, np.ndarray)
-        self.assertIsInstance(torque_matrix, np.ndarray)
+            force_matrix, torque_matrix = level_manager.get_matrices(
+                data_container=data_container,
+                level="residue",
+                highest_level=True,
+                force_matrix=None,
+                torque_matrix=None,
+                force_partitioning=0.5,
+                customised_axes=True,
+            )
 
-        self.assertEqual(force_matrix.shape, (6, 6))
-        self.assertEqual(torque_matrix.shape, (6, 6))
+        # Shape: 2 beads × 3 dof => 6×6
+        assert force_matrix.shape == (6, 6)
+        assert torque_matrix.shape == (6, 6)
 
-        level_manager.get_beads.assert_called_once_with(data_container, "residue")
+        # Expected block structure when every block is I3:
+        expected = np.block([[I3, I3], [I3, I3]])
+        np.testing.assert_array_equal(force_matrix, expected)
+        np.testing.assert_array_equal(torque_matrix, expected)
 
-        self.assertEqual(level_manager.get_weighted_forces.call_count, 2)
-        self.assertEqual(level_manager.get_weighted_torques.call_count, 2)
+        # Lightweight behavioral assertions
+        level_manager.get_beads.assert_called_once_with(data_container, "residue")
+        assert axes.get_residue_axes.call_count == 2
 
-        self.assertEqual(level_manager.create_submatrix.call_count, 6)
+        # For 2 beads: (0,0), (0,1), (1,1) => 3 pairs;
+        # each pair calls create_submatrix twice (force+torque)
+        assert level_manager.create_submatrix.call_count == 6
 
     def test_get_matrices_force_shape_mismatch(self):
         """
@@ -115,6 +141,7 @@ def test_get_matrices_force_shape_mismatch(self):
         universe_operations = UniverseOperations()
         level_manager = LevelManager(universe_operations)
 
+        # Two beads -> force_block will be 6x6
         bead1 = MagicMock()
         bead1.principal_axes.return_value = np.ones(3)
 
@@ -129,6 +156,7 @@ def test_get_matrices_force_shape_mismatch(self):
         level_manager.get_weighted_torques = MagicMock(
             return_value=np.array([0.5, 1.5, 2.5])
         )
+
         level_manager.create_submatrix = MagicMock(return_value=np.identity(3))
 
         data_container = MagicMock()
@@ -138,17 +166,32 @@ def test_get_matrices_force_shape_mismatch(self):
         bad_force_matrix = np.zeros((3, 3))
         correct_torque_matrix = np.zeros((6, 6))
 
-        with self.assertRaises(ValueError) as context:
-            level_manager.get_matrices(
-                data_container=data_container,
-                level="residue",
-                highest_level=True,
-                force_matrix=bad_force_matrix,
-                torque_matrix=correct_torque_matrix,
-                force_partitioning=0.5,
+        dummy_trans_axes = np.eye(3)
+        dummy_rot_axes = np.eye(3)
+        dummy_center = np.zeros(3)
+        dummy_moi = np.eye(3)
+
+        with patch("CodeEntropy.levels.AxesManager") as AxesManagerMock:
+            axes = AxesManagerMock.return_value
+            axes.get_residue_axes.return_value = (
+                dummy_trans_axes,
+                dummy_rot_axes,
+                dummy_center,
+                dummy_moi,
             )
 
-        self.assertIn("Inconsistent force matrix shape", str(context.exception))
+            with self.assertRaises(ValueError) as context:
+                level_manager.get_matrices(
+                    data_container=data_container,
+                    level="residue",
+                    highest_level=True,
+                    force_matrix=bad_force_matrix,
+                    torque_matrix=correct_torque_matrix,
+                    force_partitioning=0.5,
+                    customised_axes=True,
+                )
+
+        assert "force matrix shape" in str(context.exception)
 
     def test_get_matrices_torque_shape_mismatch(self):
         """
@@ -179,19 +222,35 @@ def test_get_matrices_torque_shape_mismatch(self):
         data_container.atoms.principal_axes.return_value = np.ones(3)
 
         correct_force_matrix = np.zeros((6, 6))
-        bad_torque_matrix = np.zeros((3, 3))  # Incorrect shape
-
-        with self.assertRaises(ValueError) as context:
-            level_manager.get_matrices(
-                data_container=data_container,
-                level="residue",
-                highest_level=True,
-                force_matrix=correct_force_matrix,
-                torque_matrix=bad_torque_matrix,
-                force_partitioning=0.5,
+        bad_torque_matrix = np.zeros((3, 3))  # Incorrect shape (should be 6x6)
+
+        # Mock AxesManager return tuple to satisfy unpacking
+        dummy_trans_axes = np.eye(3)
+        dummy_rot_axes = np.eye(3)
+        dummy_center = np.zeros(3)
+        dummy_moi = np.eye(3)
+
+        with patch("CodeEntropy.levels.AxesManager") as AxesManagerMock:
+            axes = AxesManagerMock.return_value
+            axes.get_residue_axes.return_value = (
+                dummy_trans_axes,
+                dummy_rot_axes,
+                dummy_center,
+                dummy_moi,
             )
 
-        self.assertIn("Inconsistent torque matrix shape", str(context.exception))
+            with self.assertRaises(ValueError) as context:
+                level_manager.get_matrices(
+                    data_container=data_container,
+                    level="residue",
+                    highest_level=True,
+                    force_matrix=correct_force_matrix,
+                    torque_matrix=bad_torque_matrix,
+                    force_partitioning=0.5,
+                    customised_axes=True,
+                )
+
+        assert "torque matrix shape" in str(context.exception)
 
     def test_get_matrices_torque_consistency(self):
         """
@@ -224,27 +283,47 @@ def test_get_matrices_torque_consistency(self):
         initial_force_matrix = np.zeros((6, 6))
         initial_torque_matrix = np.zeros((6, 6))
 
-        force_matrix_1, torque_matrix_1 = level_manager.get_matrices(
-            data_container=data_container,
-            level="residue",
-            highest_level=True,
-            force_matrix=initial_force_matrix.copy(),
-            torque_matrix=initial_torque_matrix.copy(),
-            force_partitioning=0.5,
-        )
+        # Mock AxesManager return tuple (unpacked by get_matrices)
+        dummy_trans_axes = np.eye(3)
+        dummy_rot_axes = np.eye(3)
+        dummy_center = np.zeros(3)
+        dummy_moi = np.eye(3)
+
+        with patch("CodeEntropy.levels.AxesManager") as AxesManagerMock:
+            axes = AxesManagerMock.return_value
+            axes.get_residue_axes.return_value = (
+                dummy_trans_axes,
+                dummy_rot_axes,
+                dummy_center,
+                dummy_moi,
+            )
 
-        force_matrix_2, torque_matrix_2 = level_manager.get_matrices(
-            data_container=data_container,
-            level="residue",
-            highest_level=True,
-            force_matrix=initial_force_matrix.copy(),
-            torque_matrix=initial_torque_matrix.copy(),
-            force_partitioning=0.5,
-        )
+            force_matrix_1, torque_matrix_1 = level_manager.get_matrices(
+                data_container=data_container,
+                level="residue",
+                highest_level=True,
+                force_matrix=initial_force_matrix.copy(),
+                torque_matrix=initial_torque_matrix.copy(),
+                force_partitioning=0.5,
+                customised_axes=True,
+            )
+
+            force_matrix_2, torque_matrix_2 = level_manager.get_matrices(
+                data_container=data_container,
+                level="residue",
+                highest_level=True,
+                force_matrix=initial_force_matrix.copy(),
+                torque_matrix=initial_torque_matrix.copy(),
+                force_partitioning=0.5,
+                customised_axes=True,
+            )
 
         np.testing.assert_array_equal(force_matrix_1, force_matrix_2)
         np.testing.assert_array_equal(torque_matrix_1, torque_matrix_2)
 
+        assert force_matrix_1.shape == (6, 6)
+        assert torque_matrix_1.shape == (6, 6)
+
     def test_get_beads_polymer_level(self):
         """
         Test `get_beads` for 'polymer' level.
@@ -456,67 +535,58 @@ def test_get_weighted_torques_weighted_torque_basic(self):
         Test basic torque calculation with non-zero moment of inertia and torques.
         """
         universe_operations = UniverseOperations()
-
         level_manager = LevelManager(universe_operations)
 
-        # Mock atom
-        atom = MagicMock()
-        atom.index = 0
-
-        # Mock bead
+        # Bead with one "atom"
         bead = MagicMock()
-        bead.atoms = [atom]
-        bead.center_of_mass.return_value = np.array([0.0, 0.0, 0.0])
-        bead.moment_of_inertia.return_value = np.identity(3)
-
-        # Mock data_container
-        data_container = MagicMock()
-        data_container.atoms.__getitem__.return_value.position = np.array(
-            [1.0, 0.0, 0.0]
-        )
-        data_container.atoms.__getitem__.return_value.force = np.array([0.0, 1.0, 0.0])
+        bead.positions = np.array([[1.0, 0.0, 0.0]])  # r
+        bead.forces = np.array([[0.0, 1.0, 0.0]])  # F
 
-        # Rotation axes (identity matrix)
         rot_axes = np.identity(3)
-
+        center = np.array([0.0, 0.0, 0.0])
         force_partitioning = 0.5
+        moment_of_inertia = np.array([1.0, 1.0, 1.0])
 
         result = level_manager.get_weighted_torques(
-            data_container, bead, rot_axes, force_partitioning
+            bead=bead,
+            rot_axes=rot_axes,
+            center=center,
+            force_partitioning=force_partitioning,
+            moment_of_inertia=moment_of_inertia,
         )
 
-        np.testing.assert_array_almost_equal(result, np.array([0.0, 0.0, 0.5]))
+        expected = np.array([0.0, 0.0, 0.5])
+        np.testing.assert_allclose(result, expected, rtol=0, atol=1e-12)
 
     def test_get_weighted_torques_zero_torque_skips_division(self):
         """
         Test that zero torque components skip division and remain zero.
         """
         universe_operations = UniverseOperations()
-
         level_manager = LevelManager(universe_operations)
 
-        atom = MagicMock()
-        atom.index = 0
-
         bead = MagicMock()
-        bead.atoms = [atom]
-        bead.center_of_mass.return_value = np.array([0.0, 0.0, 0.0])
-        bead.moment_of_inertia.return_value = np.identity(3)
-
-        data_container = MagicMock()
-        data_container.atoms.__getitem__.return_value.position = np.array(
-            [0.0, 0.0, 0.0]
-        )
-        data_container.atoms.__getitem__.return_value.force = np.array([0.0, 0.0, 0.0])
+        # All zeros => r x F = 0
+        bead.positions = np.array([[0.0, 0.0, 0.0]])
+        bead.forces = np.array([[0.0, 0.0, 0.0]])
 
         rot_axes = np.identity(3)
-
+        center = np.array([0.0, 0.0, 0.0])
         force_partitioning = 0.5
 
+        # Use non-zero MOI so that "skip division" is only due to zero torque
+        moment_of_inertia = np.array([1.0, 2.0, 3.0])
+
         result = level_manager.get_weighted_torques(
-            data_container, bead, rot_axes, force_partitioning
+            bead=bead,
+            rot_axes=rot_axes,
+            center=center,
+            force_partitioning=force_partitioning,
+            moment_of_inertia=moment_of_inertia,
         )
-        np.testing.assert_array_almost_equal(result, np.zeros(3))
+
+        expected = np.zeros(3)
+        np.testing.assert_array_equal(result, expected)
 
     def test_get_weighted_torques_zero_moi(self):
         """
@@ -524,79 +594,65 @@ def test_get_weighted_torques_zero_moi(self):
         and torque in that dimension is non-zero.
         """
         universe_operations = UniverseOperations()
-
         level_manager = LevelManager(universe_operations)
 
-        atom = MagicMock()
-        atom.index = 0
-
         bead = MagicMock()
-        bead.atoms = [atom]
-        bead.center_of_mass.return_value = np.array([0.0, 0.0, 0.0])
-
-        # Set moment of inertia with zero in dimension 2
-        moi = np.zeros((3, 3))
-        moi[2, 2] = 0.0
-        bead.moment_of_inertia.return_value = moi
-
-        data_container = MagicMock()
-        # Position and force that will produce a non-zero torque in z (dimension 2)
-        data_container.atoms.__getitem__.return_value.position = np.array(
-            [1.0, 0.0, 0.0]
-        )
-        data_container.atoms.__getitem__.return_value.force = np.array([0.0, 1.0, 0.0])
+        # r = (1,0,0), F = (0,1,0) => torque = (0,0,1)
+        bead.positions = np.array([[1.0, 0.0, 0.0]])
+        bead.forces = np.array([[0.0, 1.0, 0.0]])
 
         rot_axes = np.identity(3)
-
+        center = np.array([0.0, 0.0, 0.0])
         force_partitioning = 0.5
 
+        # MOI is zero in z dimension (index 2)
+        moment_of_inertia = np.array([1.0, 1.0, 0.0])
+
         torque = level_manager.get_weighted_torques(
-            data_container, bead, rot_axes, force_partitioning
+            bead=bead,
+            rot_axes=rot_axes,
+            center=center,
+            force_partitioning=force_partitioning,
+            moment_of_inertia=moment_of_inertia,
         )
 
-        self.assertEqual(torque[2], 0)
+        # x and y torques are zero; z torque is non-zero
+        # but MOI_z==0 => weighted z should be 0
+        expected = np.array([0.0, 0.0, 0.0])
+        np.testing.assert_array_equal(torque, expected)
 
-    def test_get_weighted_torques_negative_moi_raises(self):
+    def test_get_weighted_torques_negative_moi_sets_zero(self):
         """
-        Should raise ValueError when moment of inertia is negative in a dimension
-        and torque in that dimension is non-zero.
+        Negative moment of inertia components should be skipped and set to 0
+        even if the corresponding torque component is non-zero.
         """
         universe_operations = UniverseOperations()
-
         level_manager = LevelManager(universe_operations)
 
-        atom = MagicMock()
-        atom.index = 0
-
         bead = MagicMock()
-        bead.atoms = [atom]
-        bead.center_of_mass.return_value = np.array([0.0, 0.0, 0.0])
-
-        # Set moment of inertia with negative value in dimension 2
-        moi = np.identity(3)
-        moi *= -1.0
-        bead.moment_of_inertia.return_value = moi
-
-        data_container = MagicMock()
-        # Position and force that will produce a non-zero torque in z (dimension 2)
-        data_container.atoms.__getitem__.return_value.position = np.array(
-            [1.0, 0.0, 0.0]
-        )
-        data_container.atoms.__getitem__.return_value.force = np.array([0.0, 1.0, 0.0])
+        # r=(1,0,0), F=(0,1,0) => raw torque in z is non-zero
+        bead.positions = np.array([[1.0, 0.0, 0.0]])
+        bead.forces = np.array([[0.0, 1.0, 0.0]])
 
         rot_axes = np.identity(3)
+        center = np.array([0.0, 0.0, 0.0])
+        force_partitioning = 0.5
 
-        foce_partitioning = 0.5
-
-        with self.assertRaises(ValueError) as context:
-            level_manager.get_weighted_torques(
-                data_container, bead, rot_axes, foce_partitioning
-            )
+        # Negative MOI in z dimension
+        moment_of_inertia = np.array([1.0, 1.0, -1.0])
 
-        self.assertIn(
-            "Negative value encountered for moment of inertia", str(context.exception)
+        result = level_manager.get_weighted_torques(
+            bead=bead,
+            rot_axes=rot_axes,
+            center=center,
+            force_partitioning=force_partitioning,
+            moment_of_inertia=moment_of_inertia,
         )
 
+        # z torque would be non-zero, but negative MOI => z component forced to 0
+        expected = np.array([0.0, 0.0, 0.0])
+        np.testing.assert_array_equal(result, expected)
+
     def test_create_submatrix_basic_outer_product(self):
         """
         Test with known vectors to verify correct outer product.
@@ -660,18 +716,10 @@ def test_build_covariance_matrices_atomic(self):
         """
         Test `build_covariance_matrices` to ensure it correctly orchestrates
         calls and returns dictionaries with the expected structure.
-
-        This test mocks dependencies including the entropy_manager, reduced_atom
-        trajectory, levels, groups, and internal method
-        `update_force_torque_matrices`.
         """
-
-        # Instantiate your class (replace YourClass with actual class name)
         universe_operations = UniverseOperations()
-
         level_manager = LevelManager(universe_operations)
 
-        # Mock entropy_manager and get_molecule_container
         entropy_manager = MagicMock()
 
         # Fake atom with minimal attributes
@@ -680,30 +728,25 @@ def test_build_covariance_matrices_atomic(self):
         atom.resid = 1
         atom.segid = "A"
 
-        # Fake molecule with atoms list
         fake_mol = MagicMock()
         fake_mol.atoms = [atom]
 
-        # Always return fake_mol from get_molecule_container
         universe_operations.get_molecule_container = MagicMock(return_value=fake_mol)
 
-        # Mock reduced_atom with trajectory yielding two timesteps
         timestep1 = MagicMock()
         timestep1.frame = 0
         timestep2 = MagicMock()
         timestep2.frame = 1
+
         reduced_atom = MagicMock()
         reduced_atom.trajectory.__getitem__.return_value = [timestep1, timestep2]
 
-        # Setup groups and levels dictionaries
         groups = {"ua": ["mol1", "mol2"]}
         levels = {"mol1": ["level1", "level2"], "mol2": ["level1"]}
 
-        # Mock update_force_torque_matrices to just track calls
         level_manager.update_force_torque_matrices = MagicMock()
 
-        # Call the method under test
-        force_matrices, torque_matrices, _ = level_manager.build_covariance_matrices(
+        force_matrices, torque_matrices, *_ = level_manager.build_covariance_matrices(
             entropy_manager=entropy_manager,
             reduced_atom=reduced_atom,
             levels=levels,
@@ -713,24 +756,21 @@ def test_build_covariance_matrices_atomic(self):
             step=1,
             number_frames=2,
             force_partitioning=0.5,
+            combined_forcetorque=False,
+            customised_axes=True,
         )
 
-        # Assert returned matrices are dictionaries with correct keys
         self.assertIsInstance(force_matrices, dict)
         self.assertIsInstance(torque_matrices, dict)
         self.assertSetEqual(set(force_matrices.keys()), {"ua", "res", "poly"})
         self.assertSetEqual(set(torque_matrices.keys()), {"ua", "res", "poly"})
 
-        # Assert 'res' and 'poly' entries are lists of correct length
         self.assertIsInstance(force_matrices["res"], list)
         self.assertIsInstance(force_matrices["poly"], list)
         self.assertEqual(len(force_matrices["res"]), len(groups))
         self.assertEqual(len(force_matrices["poly"]), len(groups))
 
-        # Check get_molecule_container call count: 2 timesteps * 2 molecules = 4 calls
         self.assertEqual(universe_operations.get_molecule_container.call_count, 4)
-
-        # Check update_force_torque_matrices call count:
         self.assertEqual(level_manager.update_force_torque_matrices.call_count, 6)
 
     def test_update_force_torque_matrices_united_atom(self):
@@ -771,6 +811,7 @@ def test_update_force_torque_matrices_united_atom(self):
 
         force_avg = {"ua": {}, "res": [None], "poly": [None]}
         torque_avg = {"ua": {}, "res": [None], "poly": [None]}
+        forcetorque_avg = {"ua": {}, "res": [None], "poly": [None]}
         frame_counts = {"ua": {}, "res": [None], "poly": [None]}
 
         level_manager.update_force_torque_matrices(
@@ -783,8 +824,11 @@ def test_update_force_torque_matrices_united_atom(self):
             num_frames=10,
             force_avg=force_avg,
             torque_avg=torque_avg,
+            forcetorque_avg=forcetorque_avg,
             frame_counts=frame_counts,
             force_partitioning=0.5,
+            combined_forcetorque=False,
+            customised_axes=True,
         )
 
         assert (0, 0) in force_avg["ua"]
@@ -800,6 +844,8 @@ def test_update_force_torque_matrices_united_atom(self):
         assert frame_counts["ua"][(0, 0)] == 1
         assert frame_counts["ua"][(0, 1)] == 1
 
+        assert forcetorque_avg["ua"] == {}
+
     def test_update_force_torque_matrices_united_atom_increment(self):
         """
         Test that update_force_torque_matrices() correctly updates (increments)
@@ -824,7 +870,6 @@ def test_update_force_torque_matrices_united_atom_increment(self):
         selected_atoms = MagicMock()
         selected_atoms.trajectory = MagicMock()
         selected_atoms.trajectory.__getitem__.return_value = None
-
         universe_operations.new_U_select_atom.return_value = selected_atoms
 
         f_mat_1 = np.array([[1.0]])
@@ -833,10 +878,13 @@ def test_update_force_torque_matrices_united_atom_increment(self):
         f_mat_2 = np.array([[3.0]])
         t_mat_2 = np.array([[4.0]])
 
-        level_manager.get_matrices = MagicMock(return_value=(f_mat_1, t_mat_1))
+        level_manager.get_matrices = MagicMock(
+            side_effect=[(f_mat_1, t_mat_1), (f_mat_2, t_mat_2)]
+        )
 
         force_avg = {"ua": {}, "res": [None], "poly": [None]}
         torque_avg = {"ua": {}, "res": [None], "poly": [None]}
+        forcetorque_avg = {"ua": {}, "res": [None], "poly": [None]}
         frame_counts = {"ua": {}, "res": [None], "poly": [None]}
 
         level_manager.update_force_torque_matrices(
@@ -849,12 +897,14 @@ def test_update_force_torque_matrices_united_atom_increment(self):
             num_frames=10,
             force_avg=force_avg,
             torque_avg=torque_avg,
+            forcetorque_avg=forcetorque_avg,
             frame_counts=frame_counts,
             force_partitioning=0.5,
+            combined_forcetorque=False,
+            customised_axes=True,
         )
 
-        level_manager.get_matrices = MagicMock(return_value=(f_mat_2, t_mat_2))
-
+        # Second update
         level_manager.update_force_torque_matrices(
             entropy_manager=entropy_manager,
             mol=mol,
@@ -865,8 +915,11 @@ def test_update_force_torque_matrices_united_atom_increment(self):
             num_frames=10,
             force_avg=force_avg,
             torque_avg=torque_avg,
+            forcetorque_avg=forcetorque_avg,
             frame_counts=frame_counts,
             force_partitioning=0.5,
+            combined_forcetorque=False,
+            customised_axes=True,
         )
 
         expected_force = f_mat_1 + (f_mat_2 - f_mat_1) / 2
@@ -874,7 +927,9 @@ def test_update_force_torque_matrices_united_atom_increment(self):
 
         np.testing.assert_array_almost_equal(force_avg["ua"][(0, 0)], expected_force)
         np.testing.assert_array_almost_equal(torque_avg["ua"][(0, 0)], expected_torque)
-        self.assertEqual(frame_counts["ua"][(0, 0)], 2)
+        assert frame_counts["ua"][(0, 0)] == 2
+
+        assert forcetorque_avg["ua"] == {}
 
     def test_update_force_torque_matrices_residue(self):
         """
@@ -883,8 +938,8 @@ def test_update_force_torque_matrices_residue(self):
         incrementing frame counts.
         """
         universe_operations = UniverseOperations()
-
         level_manager = LevelManager(universe_operations)
+
         entropy_manager = MagicMock()
         mol = MagicMock()
         mol.trajectory.__getitem__.return_value = None
@@ -895,6 +950,7 @@ def test_update_force_torque_matrices_residue(self):
 
         force_avg = {"ua": {}, "res": [None], "poly": [None]}
         torque_avg = {"ua": {}, "res": [None], "poly": [None]}
+        forcetorque_avg = {"ua": {}, "res": [None], "poly": [None]}
         frame_counts = {"ua": {}, "res": [None], "poly": [None]}
 
         level_manager.update_force_torque_matrices(
@@ -907,13 +963,18 @@ def test_update_force_torque_matrices_residue(self):
             num_frames=10,
             force_avg=force_avg,
             torque_avg=torque_avg,
+            forcetorque_avg=forcetorque_avg,
             frame_counts=frame_counts,
             force_partitioning=0.5,
+            combined_forcetorque=False,
+            customised_axes=True,
         )
 
         np.testing.assert_array_equal(force_avg["res"][0], f_mat_mock)
         np.testing.assert_array_equal(torque_avg["res"][0], t_mat_mock)
-        self.assertEqual(frame_counts["res"][0], 1)
+        assert frame_counts["res"][0] == 1
+
+        assert forcetorque_avg["res"][0] is None
 
     def test_update_force_torque_matrices_incremental_average(self):
         """
@@ -923,8 +984,8 @@ def test_update_force_torque_matrices_incremental_average(self):
         Ensures that float precision is maintained and no casting errors occur.
         """
         universe_operations = UniverseOperations()
-
         level_manager = LevelManager(universe_operations)
+
         entropy_manager = MagicMock()
         mol = MagicMock()
         mol.trajectory.__getitem__.return_value = None
@@ -941,6 +1002,7 @@ def test_update_force_torque_matrices_incremental_average(self):
 
         force_avg = {"ua": {}, "res": [None], "poly": [None]}
         torque_avg = {"ua": {}, "res": [None], "poly": [None]}
+        forcetorque_avg = {"ua": {}, "res": [None], "poly": [None]}
         frame_counts = {"ua": {}, "res": [None], "poly": [None]}
 
         # First update
@@ -954,8 +1016,11 @@ def test_update_force_torque_matrices_incremental_average(self):
             num_frames=10,
             force_avg=force_avg,
             torque_avg=torque_avg,
+            forcetorque_avg=forcetorque_avg,
             frame_counts=frame_counts,
             force_partitioning=0.5,
+            combined_forcetorque=False,
+            customised_axes=True,
         )
 
         # Second update
@@ -969,8 +1034,11 @@ def test_update_force_torque_matrices_incremental_average(self):
             num_frames=10,
             force_avg=force_avg,
             torque_avg=torque_avg,
+            forcetorque_avg=forcetorque_avg,
             frame_counts=frame_counts,
             force_partitioning=0.5,
+            combined_forcetorque=False,
+            customised_axes=True,
         )
 
         expected_force = f_mat_1 + (f_mat_2 - f_mat_1) / 2
@@ -978,7 +1046,9 @@ def test_update_force_torque_matrices_incremental_average(self):
 
         np.testing.assert_array_almost_equal(force_avg["res"][0], expected_force)
         np.testing.assert_array_almost_equal(torque_avg["res"][0], expected_torque)
-        self.assertEqual(frame_counts["res"][0], 2)
+
+        assert frame_counts["res"][0] == 2
+        assert forcetorque_avg["res"][0] is None
 
     def test_filter_zero_rows_columns_no_zeros(self):
         """

From 9bcaa37e4c584278516872cdf2045a50848b1131 Mon Sep 17 00:00:00 2001
From: harryswift01 <harry.swift@stfc.ac.uk>
Date: Mon, 2 Feb 2026 17:03:19 +0000
Subject: [PATCH 9/9] fix unit tests within `test_mda_universe_operations.py`
 in relation to changes within PR #267

---
 .../test_mda_universe_operations.py           | 149 ++++++++----------
 1 file changed, 69 insertions(+), 80 deletions(-)

diff --git a/tests/test_CodeEntropy/test_mda_universe_operations.py b/tests/test_CodeEntropy/test_mda_universe_operations.py
index fb4c891..46e68a7 100644
--- a/tests/test_CodeEntropy/test_mda_universe_operations.py
+++ b/tests/test_CodeEntropy/test_mda_universe_operations.py
@@ -28,13 +28,6 @@ def setUp(self):
     def test_new_U_select_frame(self, MockMerge, MockAnalysisFromFunction):
         """
         Unit test for UniverseOperations.new_U_select_frame().
-
-        Verifies that:
-        - The Universe is queried with select_atoms("all", updating=True)
-        - AnalysisFromFunction is called to obtain coordinates and forces
-        - mda.Merge is called with the selected AtomGroup
-        - The new universe returned by Merge.load_new receives the correct arrays
-        - The method returns the merged universe
         """
         # Mock Universe and its components
         mock_universe = MagicMock()
@@ -48,6 +41,7 @@ def test_new_U_select_frame(self, MockMerge, MockAnalysisFromFunction):
         # Mock AnalysisFromFunction results for coordinates, forces, and dimensions
         coords = np.random.rand(10, 100, 3)
         forces = np.random.rand(10, 100, 3)
+        dims = np.random.rand(10, 6)
 
         mock_coords_analysis = MagicMock()
         mock_coords_analysis.run.return_value.results = {"timeseries": coords}
@@ -55,35 +49,35 @@ def test_new_U_select_frame(self, MockMerge, MockAnalysisFromFunction):
         mock_forces_analysis = MagicMock()
         mock_forces_analysis.run.return_value.results = {"timeseries": forces}
 
-        # Set the side effects for the three AnalysisFromFunction calls
+        mock_dims_analysis = MagicMock()
+        mock_dims_analysis.run.return_value.results = {"timeseries": dims}
+
         MockAnalysisFromFunction.side_effect = [
             mock_coords_analysis,
             mock_forces_analysis,
+            mock_dims_analysis,
         ]
 
-        # Mock the merge operation
+        # Mock merge operation
         mock_merged_universe = MagicMock()
         MockMerge.return_value = mock_merged_universe
 
         ops = UniverseOperations()
         result = ops.new_U_select_frame(mock_universe)
 
+        # Basic behavior checks
         mock_universe.select_atoms.assert_called_once_with("all", updating=True)
-        MockMerge.assert_called_once_with(mock_select_atoms)
 
-        # Ensure the 'load_new' method was called with the correct arguments
-        mock_merged_universe.load_new.assert_called_once()
-        args, kwargs = mock_merged_universe.load_new.call_args
+        # AnalysisFromFunction called 3 times (coords, forces, dimensions)
+        assert MockAnalysisFromFunction.call_count == 3
+        mock_coords_analysis.run.assert_called_once()
+        mock_forces_analysis.run.assert_called_once()
+        mock_dims_analysis.run.assert_called_once()
 
-        # Assert that the arrays are passed correctly
-        np.testing.assert_array_equal(args[0], coords)
-        np.testing.assert_array_equal(kwargs["forces"], forces)
-
-        # Check if format was included in the kwargs
-        self.assertIn("format", kwargs)
+        # Merge called with selected AtomGroup
+        MockMerge.assert_called_once_with(mock_select_atoms)
 
-        # Ensure the result is the mock merged universe
-        self.assertEqual(result, mock_merged_universe)
+        assert result == mock_merged_universe
 
     @patch("CodeEntropy.mda_universe_operations.AnalysisFromFunction")
     @patch("CodeEntropy.mda_universe_operations.mda.Merge")
@@ -93,7 +87,7 @@ def test_new_U_select_atom(self, MockMerge, MockAnalysisFromFunction):
 
         Ensures that:
         - The Universe is queried with the correct selection string
-        - Coordinates and forces are extracted via AnalysisFromFunction
+        - Coordinates, forces, and dimensions are extracted via AnalysisFromFunction
         - mda.Merge receives the AtomGroup from select_atoms
         - The new universe is populated with the expected data via load_new()
         - The returned universe is the object created by Merge
@@ -106,6 +100,7 @@ def test_new_U_select_atom(self, MockMerge, MockAnalysisFromFunction):
         # Mock AnalysisFromFunction results for coordinates, forces, and dimensions
         coords = np.random.rand(10, 100, 3)
         forces = np.random.rand(10, 100, 3)
+        dims = np.random.rand(10, 6)
 
         mock_coords_analysis = MagicMock()
         mock_coords_analysis.run.return_value.results = {"timeseries": coords}
@@ -113,10 +108,13 @@ def test_new_U_select_atom(self, MockMerge, MockAnalysisFromFunction):
         mock_forces_analysis = MagicMock()
         mock_forces_analysis.run.return_value.results = {"timeseries": forces}
 
-        # Set the side effects for the three AnalysisFromFunction calls
+        mock_dims_analysis = MagicMock()
+        mock_dims_analysis.run.return_value.results = {"timeseries": dims}
+
         MockAnalysisFromFunction.side_effect = [
             mock_coords_analysis,
             mock_forces_analysis,
+            mock_dims_analysis,
         ]
 
         # Mock the merge operation
@@ -126,23 +124,19 @@ def test_new_U_select_atom(self, MockMerge, MockAnalysisFromFunction):
         ops = UniverseOperations()
 
         result = ops.new_U_select_atom(mock_universe, select_string="resid 1-10")
-
         mock_universe.select_atoms.assert_called_once_with("resid 1-10", updating=True)
-        MockMerge.assert_called_once_with(mock_select_atoms)
-
-        # Ensure the 'load_new' method was called with the correct arguments
-        mock_merged_universe.load_new.assert_called_once()
-        args, kwargs = mock_merged_universe.load_new.call_args
 
-        # Assert that the arrays are passed correctly
-        np.testing.assert_array_equal(args[0], coords)
-        np.testing.assert_array_equal(kwargs["forces"], forces)
+        # AnalysisFromFunction called for coords, forces, dimensions
+        assert MockAnalysisFromFunction.call_count == 3
+        mock_coords_analysis.run.assert_called_once()
+        mock_forces_analysis.run.assert_called_once()
+        mock_dims_analysis.run.assert_called_once()
 
-        # Check if format was included in the kwargs
-        self.assertIn("format", kwargs)
+        # Merge called with the selected AtomGroup
+        MockMerge.assert_called_once_with(mock_select_atoms)
 
-        # Ensure the result is the mock merged universe
-        self.assertEqual(result, mock_merged_universe)
+        # Returned universe should be the merged universe
+        assert result == mock_merged_universe
 
     def test_get_molecule_container(self):
         """
@@ -179,25 +173,13 @@ def test_get_molecule_container(self):
     def test_merge_forces(self, MockUniverse, MockMerge, MockAnalysisFromFunction):
         """
         Unit test for UniverseOperations.merge_forces().
-
-        This test ensures that:
-        - Two MDAnalysis Universes are created: one for coordinates
-        (tprfile + trrfile) and one for forces (tprfile + forcefile).
-        - Both Universes correctly return AtomGroups via select_atoms("all").
-        - Coordinates and forces are extracted using AnalysisFromFunction.
-        - mda.Merge is called with the coordinate AtomGroup.
-        - The merged Universe receives the correct coordinate and force arrays
-        through load_new().
-        - When kcal=False, force values are passed through unchanged
-        (no kcal→kJ conversion).
-        - The returned universe is the same object returned by mda.Merge().
         """
-
+        # Two Universes created: coords and forces
         mock_u_coords = MagicMock()
         mock_u_force = MagicMock()
         MockUniverse.side_effect = [mock_u_coords, mock_u_force]
 
-        # Each universe returns a mock AtomGroup from select_atoms("all")
+        # Each universe returns an AtomGroup from select_atoms("all")
         mock_ag_coords = MagicMock()
         mock_ag_force = MagicMock()
         mock_u_coords.select_atoms.return_value = mock_ag_coords
@@ -205,6 +187,7 @@ def test_merge_forces(self, MockUniverse, MockMerge, MockAnalysisFromFunction):
 
         coords = np.random.rand(5, 10, 3)
         forces = np.random.rand(5, 10, 3)
+        dims = np.random.rand(5, 6)
 
         mock_coords_analysis = MagicMock()
         mock_coords_analysis.run.return_value.results = {"timeseries": coords}
@@ -212,10 +195,13 @@ def test_merge_forces(self, MockUniverse, MockMerge, MockAnalysisFromFunction):
         mock_forces_analysis = MagicMock()
         mock_forces_analysis.run.return_value.results = {"timeseries": forces}
 
-        # Two calls: first for coordinates, second for forces
+        mock_dims_analysis = MagicMock()
+        mock_dims_analysis.run.return_value.results = {"timeseries": dims}
+
         MockAnalysisFromFunction.side_effect = [
             mock_coords_analysis,
             mock_forces_analysis,
+            mock_dims_analysis,
         ]
 
         mock_merged = MagicMock()
@@ -230,28 +216,24 @@ def test_merge_forces(self, MockUniverse, MockMerge, MockAnalysisFromFunction):
             kcal=False,
         )
 
-        self.assertEqual(MockUniverse.call_count, 2)
-        MockUniverse.assert_any_call("topol.tpr", "traj.trr", format=None)
-        MockUniverse.assert_any_call("topol.tpr", "forces.trr", format=None)
+        # Universe construction
+        assert MockUniverse.call_count == 2
 
+        # Selection
         mock_u_coords.select_atoms.assert_called_once_with("all")
         mock_u_force.select_atoms.assert_called_once_with("all")
 
-        self.assertEqual(MockAnalysisFromFunction.call_count, 2)
+        # AnalysisFromFunction usage
+        assert MockAnalysisFromFunction.call_count == 3
+        mock_coords_analysis.run.assert_called_once()
+        mock_forces_analysis.run.assert_called_once()
+        mock_dims_analysis.run.assert_called_once()
 
+        # Merge called with coordinate AtomGroup
         MockMerge.assert_called_once_with(mock_ag_coords)
 
-        mock_merged.load_new.assert_called_once()
-        args, kwargs = mock_merged.load_new.call_args
-
-        # Coordinates passed positionally
-        np.testing.assert_array_equal(args[0], coords)
-
-        # Forces passed via kwargs
-        np.testing.assert_array_equal(kwargs["forces"], forces)
-
-        # Finally the function returns the merged universe
-        self.assertEqual(result, mock_merged)
+        # Returned object is merged universe
+        assert result == mock_merged
 
     @patch("CodeEntropy.mda_universe_operations.AnalysisFromFunction")
     @patch("CodeEntropy.mda_universe_operations.mda.Merge")
@@ -262,15 +244,6 @@ def test_merge_forces_kcal_conversion(
         """
         Unit test for UniverseOperations.merge_forces() covering the kcal→kJ
         conversion branch.
-
-        Verifies that:
-        - Two Universe objects are constructed for coords and forces.
-        - Each Universe returns an AtomGroup via select_atoms("all").
-        - AnalysisFromFunction is called twice.
-        - Forces are multiplied EXACTLY once by 4.184 when kcal=True.
-        - Merge() is called with the coordinate AtomGroup.
-        - load_new() receives the correct coordinates and converted forces.
-        - The returned Universe is the Merge() result.
         """
         mock_u_coords = MagicMock()
         mock_u_force = MagicMock()
@@ -286,6 +259,8 @@ def test_merge_forces_kcal_conversion(
         original_forces = np.ones((2, 3, 3))
         mock_forces_array = original_forces.copy()
 
+        dims = np.ones((2, 6))
+
         # Mock AnalysisFromFunction return values
         mock_coords_analysis = MagicMock()
         mock_coords_analysis.run.return_value.results = {"timeseries": coords}
@@ -295,9 +270,13 @@ def test_merge_forces_kcal_conversion(
             "timeseries": mock_forces_array
         }
 
+        mock_dims_analysis = MagicMock()
+        mock_dims_analysis.run.return_value.results = {"timeseries": dims}
+
         MockAnalysisFromFunction.side_effect = [
             mock_coords_analysis,
             mock_forces_analysis,
+            mock_dims_analysis,
         ]
 
         mock_merged = MagicMock()
@@ -306,10 +285,20 @@ def test_merge_forces_kcal_conversion(
         ops = UniverseOperations()
         result = ops.merge_forces("t.tpr", "c.trr", "f.trr", kcal=True)
 
-        _, kwargs = mock_merged.load_new.call_args
+        # select_atoms("all") (your code uses no updating=True)
+        mock_u_coords.select_atoms.assert_called_once_with("all")
+        mock_u_force.select_atoms.assert_called_once_with("all")
+
+        # AnalysisFromFunction called three times
+        assert MockAnalysisFromFunction.call_count == 3
 
-        expected_forces = original_forces * 4.184
-        np.testing.assert_array_equal(kwargs["forces"], expected_forces)
-        np.testing.assert_array_equal(mock_merged.load_new.call_args[0][0], coords)
+        # Forces are multiplied exactly once by 4.184 when kcal=True
+        np.testing.assert_allclose(
+            mock_forces_array, original_forces * 4.184, rtol=0, atol=0
+        )
+
+        # Merge called with coordinate AtomGroup
+        MockMerge.assert_called_once_with(mock_ag_coords)
 
-        self.assertEqual(result, mock_merged)
+        # Returned universe is the merged universe
+        assert result == mock_merged