diff --git a/CodeEntropy/levels.py b/CodeEntropy/levels.py
index 64768a6..9d43670 100644
--- a/CodeEntropy/levels.py
+++ b/CodeEntropy/levels.py
@@ -80,7 +80,6 @@ def get_matrices(
         self,
         data_container,
         level,
-        number_frames,
         highest_level,
         force_matrix,
         torque_matrix,
@@ -92,7 +91,6 @@ def get_matrices(
         Parameters:
           data_container (MDAnalysis.Universe): Data for a molecule or residue.
           level (str): 'polymer', 'residue', or 'united_atom'.
-          number_frames (int): Number of frames being processed.
           highest_level (bool): Whether this is the top (largest bead size) level.
           force_matrix, torque_matrix (np.ndarray or None): Accumulated matrices to add
           to.
@@ -116,21 +114,23 @@ 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, rot_axes = self.get_axes(data_container, level, bead_index)
+            trans_axes = data_container.atoms.principal_axes()
+            rot_axes = np.real(bead.principal_axes())
 
             # Sort out coordinates, forces, and torques for each atom in the bead
             weighted_forces[bead_index] = self.get_weighted_forces(
                 data_container,
-                list_of_beads[bead_index],
+                bead,
                 trans_axes,
                 highest_level,
                 force_partitioning,
             )
             weighted_torques[bead_index] = self.get_weighted_torques(
-                data_container, list_of_beads[bead_index], rot_axes, force_partitioning
+                data_container, bead, rot_axes, force_partitioning
             )
 
         # Create covariance submatrices
@@ -233,255 +233,57 @@ def get_beads(self, data_container, level):
 
         return list_of_beads
 
-    def get_axes(self, data_container, level, index=0):
-        """
-        Function to set the translational and rotational axes.
-        The translational axes are based on the principal axes of the unit
-        one level larger than the level we are interested in (except for
-        the polymer level where there is no larger unit). The rotational
-        axes use the covalent links between residues or atoms where possible
-        to define the axes, or if the unit is not bonded to others of the
-        same level the prinicpal axes of the unit are used.
-
-        Args:
-          data_container (MDAnalysis.Universe): the molecule and trajectory data
-          level (str): the level (united atom, residue, or polymer) of interest
-          index (int): residue index
-
-        Returns:
-          trans_axes : translational axes
-          rot_axes : rotational axes
-        """
-        index = int(index)
-
-        if level == "polymer":
-            # for polymer use principle axis for both translation and rotation
-            trans_axes = data_container.atoms.principal_axes()
-            rot_axes = data_container.atoms.principal_axes()
-
-        elif level == "residue":
-            # Translation
-            # for residues use principal axes of whole molecule for translation
-            trans_axes = data_container.atoms.principal_axes()
-
-            # Rotation
-            # find bonds between atoms in residue of interest and other residues
-            # we are assuming bonds only exist between adjacent residues
-            # (linear chains of residues)
-            # 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}")
-
-            if len(atom_set) == 0:
-                # if no bonds to other residues use pricipal axes of residue
-                rot_axes = residue.atoms.principal_axes()
-
-            else:
-                # set center of rotation to center of mass of the residue
-                center = residue.atoms.center_of_mass()
-
-                # get vector for average position of bonded atoms
-                vector = self.get_avg_pos(atom_set, center)
-
-                # use spherical coordinates function to get rotational axes
-                rot_axes = self.get_sphCoord_axes(vector)
-
-        elif level == "united_atom":
-            # Translation
-            # for united atoms use principal axes of residue for translation
-            trans_axes = data_container.residues.principal_axes()
-
-            # Rotation
-            # for united atoms use heavy atoms bonded to the heavy atom
-            atom_set = data_container.select_atoms(
-                f"(prop mass > 1.1) and bonded index {index}"
-            )
-
-            if len(atom_set) == 0:
-                # if no bonds to other residues use pricipal axes of residue
-                rot_axes = data_container.residues.principal_axes()
-            else:
-                # center at position of heavy atom
-                atom_group = data_container.select_atoms(f"index {index}")
-                center = atom_group.positions[0]
-
-                # get vector for average position of bonded atoms
-                vector = self.get_avg_pos(atom_set, center)
-
-                # use spherical coordinates function to get rotational axes
-                rot_axes = self.get_sphCoord_axes(vector)
-
-        logger.debug(f"Translational Axes: {trans_axes}")
-        logger.debug(f"Rotational Axes: {rot_axes}")
-
-        return trans_axes, rot_axes
-
-    def get_avg_pos(self, atom_set, center):
-        """
-        Function to get the average position of a set of atoms.
-
-        Args:
-            atom_set : MDAnalysis atom group
-            center : position for center of rotation
-
-        Returns:
-            avg_position : three dimensional vector
-        """
-        # start with an empty vector
-        avg_position = np.zeros((3))
-
-        # get number of atoms
-        number_atoms = len(atom_set.names)
-
-        if number_atoms != 0:
-            # sum positions for all atoms in the given set
-            for atom_index in range(number_atoms):
-                atom_position = atom_set.atoms[atom_index].position
-
-                avg_position += atom_position
-
-            avg_position /= number_atoms  # divide by number of atoms to get average
-
-        else:
-            # if no atoms in set the unit has no bonds to restrict its rotational
-            # motion, so we can use a random vector to get spherical
-            # coordinate axes
-            avg_position = np.random.random(3)
-
-        # transform the average position to a coordinate system with the origin
-        # at center
-        avg_position = avg_position - center
-
-        logger.debug(f"Average Position: {avg_position}")
-
-        return avg_position
-
-    def get_sphCoord_axes(self, arg_r):
-        """
-        For a given vector in space, treat it is a radial vector rooted at
-        0,0,0 and derive a curvilinear coordinate system according to the
-        rules of polar spherical coordinates
-
-        Args:
-            arg_r: 3 dimensional vector
-
-        Returns:
-            spherical_basis: axes set (3 vectors)
-        """
-
-        x2y2 = arg_r[0] ** 2 + arg_r[1] ** 2
-        r2 = x2y2 + arg_r[2] ** 2
-
-        # Check for division by zero
-        if r2 == 0.0:
-            raise ValueError("r2 is zero, cannot compute spherical coordinates.")
-
-        if x2y2 == 0.0:
-            raise ValueError("x2y2 is zero, cannot compute sin_phi and cos_phi.")
-
-        # These conditions are mathematically unreachable for real-valued vectors.
-        # Marked as no cover to avoid false negatives in coverage reports.
-
-        # Check for non-negative values inside the square root
-        if x2y2 / r2 < 0:  # pragma: no cover
-            raise ValueError(
-                f"Negative value encountered for sin_theta calculation: {x2y2 / r2}. "
-                f"Cannot take square root."
-            )
-
-        if x2y2 < 0:  # pragma: no cover
-            raise ValueError(
-                f"Negative value encountered for sin_phi and cos_phi "
-                f"calculation: {x2y2}. "
-                f"Cannot take square root."
-            )
-
-        if x2y2 != 0.0:
-            sin_theta = np.sqrt(x2y2 / r2)
-            cos_theta = arg_r[2] / np.sqrt(r2)
-
-            sin_phi = arg_r[1] / np.sqrt(x2y2)
-            cos_phi = arg_r[0] / np.sqrt(x2y2)
-
-        else:  # pragma: no cover
-            sin_theta = 0.0
-            cos_theta = 1
-
-            sin_phi = 0.0
-            cos_phi = 1
-
-        # if abs(sin_theta) > 1 or abs(sin_phi) > 1:
-        #     print('Bad sine : T {} , P {}'.format(sin_theta, sin_phi))
-
-        # cos_theta = np.sqrt(1 - sin_theta*sin_theta)
-        # cos_phi = np.sqrt(1 - sin_phi*sin_phi)
-
-        # print('{} {} {}'.format(*arg_r))
-        # print('Sin T : {}, cos T : {}'.format(sin_theta, cos_theta))
-        # print('Sin P : {}, cos P : {}'.format(sin_phi, cos_phi))
-
-        spherical_basis = np.zeros((3, 3))
-
-        # r^
-        spherical_basis[0, :] = np.asarray(
-            [sin_theta * cos_phi, sin_theta * sin_phi, cos_theta]
-        )
-
-        # Theta^
-        spherical_basis[1, :] = np.asarray(
-            [cos_theta * cos_phi, cos_theta * sin_phi, -sin_theta]
-        )
-
-        # Phi^
-        spherical_basis[2, :] = np.asarray([-sin_phi, cos_phi, 0.0])
-
-        logger.debug(f"Spherical Basis: {spherical_basis}")
-
-        return spherical_basis
-
     def get_weighted_forces(
         self, data_container, bead, trans_axes, highest_level, force_partitioning
     ):
         """
-        Function to calculate the mass weighted forces for a given bead.
+        Compute mass-weighted translational forces for a bead.
 
-        Args:
-           data_container (MDAnalysis.Universe): Contains atomic positions and forces.
-           bead : The part of the molecule to be considered.
-           trans_axes (np.ndarray): The axes relative to which the forces are located.
-           highest_level (bool): Is this the largest level of the length scale hierarchy
-           force_partitioning (float): Factor to adjust force contributions to avoid
-           over counting correlated forces, default is 0.5.
+        The forces acting on all atoms belonging to the bead are first transformed
+        into the provided translational reference frame and summed. If this bead
+        corresponds to the highest level of a hierarchical coarse-graining scheme,
+        the total force is scaled by a force-partitioning factor to avoid double
+        counting forces from weakly correlated atoms.
 
-        Returns:
-            weighted_force (np.ndarray): The mass-weighted sum of the forces in the
-            bead.
-        """
+        The resulting force vector is then normalized by the square root of the
+        bead's total mass.
+
+        Parameters
+        ----------
+        data_container : MDAnalysis.Universe
+            Container holding atomic positions and forces.
+        bead : object
+            Molecular subunit whose atoms contribute to the force.
+        trans_axes : np.ndarray
+            Transformation matrix defining the translational reference frame.
+        highest_level : bool
+            Whether this bead is the highest level in the length-scale hierarchy.
+            If True, force partitioning is applied.
+        force_partitioning : float
+            Scaling factor applied to forces to avoid over-counting correlated
+            contributions (typically 0.5).
+
+        Returns
+        -------
+        weighted_force : np.ndarray
+            Mass-weighted translational force acting on the bead.
 
+        Raises
+        ------
+        ValueError
+            If the bead mass is zero or negative.
+        """
         forces_trans = np.zeros((3,))
 
-        # Sum forces from all atoms in the bead
         for atom in bead.atoms:
-            # update local forces in translational axes
             forces_local = np.matmul(trans_axes, data_container.atoms[atom.index].force)
             forces_trans += forces_local
 
         if highest_level:
-            # multiply by the force_partitioning parameter to avoid double counting
-            # of the forces on weakly correlated atoms
-            # the default value of force_partitioning is 0.5 (dividing by two)
             forces_trans = force_partitioning * forces_trans
 
-        # divide the sum of forces by the mass of the bead to get the weighted forces
         mass = bead.total_mass()
 
-        # Check that mass is positive to avoid division by 0 or negative values inside
-        # sqrt
         if mass <= 0:
             raise ValueError(
                 f"Invalid mass value: {mass}. Mass must be positive to compute the "
@@ -496,87 +298,83 @@ def get_weighted_forces(
 
     def get_weighted_torques(self, data_container, bead, rot_axes, force_partitioning):
         """
-        Function to calculate the moment of inertia weighted torques for a given bead.
+        Compute moment-of-inertia weighted torques for a bead.
+
+        Atomic coordinates and forces are transformed into the provided rotational
+        reference frame. Torques are computed as the cross product of position
+        vectors (relative to the bead center of mass) and forces, with a
+        force-partitioning factor applied to reduce over-counting of correlated
+        atomic contributions.
 
-        This function computes torques in a rotated frame and then weights them using
-        the moment of inertia tensor. To prevent numerical instability, it treats
-        extremely small diagonal elements of the moment of inertia tensor as zero
-        (since values below machine precision are effectively zero). This avoids
-        unnecessary use of extended precision (e.g., float128).
+        The total torque vector is then weighted by the square root of the bead's
+        principal moments of inertia. Weighting is performed component-wise using
+        the sorted eigenvalues of the moment of inertia tensor.
 
-        Additionally, if the computed torque is already zero, the function skips
-        the division step, reducing unnecessary computations and potential errors.
+        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.
 
         Parameters
         ----------
         data_container : object
-            Contains atomic positions and forces.
+            Container holding atomic positions and forces.
         bead : object
-            The part of the molecule to be considered.
+            Molecular subunit whose atoms contribute to the torque.
         rot_axes : np.ndarray
-            The axes relative to which the forces and coordinates are located.
-        force_partitioning : float, optional
-            Factor to adjust force contributions, default is 0.5.
+            Transformation matrix defining the rotational reference frame.
+        force_partitioning : float
+            Scaling factor applied to forces to avoid over-counting correlated
+            contributions (typically 0.5).
 
         Returns
         -------
         weighted_torque : np.ndarray
-            The mass-weighted sum of the torques in the bead.
-        """
+            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:
-
-            # update local coordinates in rotational axes
             coords_rot = (
                 data_container.atoms[atom.index].position - bead.center_of_mass()
             )
             coords_rot = np.matmul(rot_axes, coords_rot)
-            # update local forces in rotational frame
             forces_rot = np.matmul(rot_axes, data_container.atoms[atom.index].force)
 
-            # multiply by the force_partitioning parameter to avoid double counting
-            # of the forces on weakly correlated atoms
-            # the default value of force_partitioning is 0.5 (dividing by two)
             forces_rot = force_partitioning * forces_rot
 
-            # define torques (cross product of coordinates and forces) in rotational
-            # axes
             torques_local = np.cross(coords_rot, forces_rot)
             torques += torques_local
 
-        # divide by moment of inertia to get weighted torques
-        # moment of inertia is a 3x3 tensor
-        # the weighting is done in each dimension (x,y,z) using the diagonal
-        # elements of the moment of inertia tensor
-        moment_of_inertia = bead.moment_of_inertia()
+        eigenvalues, _ = np.linalg.eig(bead.moment_of_inertia())
+        moments_of_inertia = sorted(eigenvalues, reverse=True)
 
         for dimension in range(3):
-            # Skip calculation if torque is already zero
             if np.isclose(torques[dimension], 0):
                 weighted_torque[dimension] = 0
                 continue
 
-            # Check for zero moment of inertia
-            if np.isclose(moment_of_inertia[dimension, dimension], 0):
-                raise ZeroDivisionError(
-                    f"Attempted to divide by zero moment of inertia in dimension "
-                    f"{dimension}."
-                )
+            if np.isclose(moments_of_inertia[dimension], 0):
+                weighted_torque[dimension] = 0
+                logger.warning("Zero moment of inertia. Setting torque to 0")
+                continue
 
-            # Check for negative moment of inertia
-            if moment_of_inertia[dimension, dimension] < 0:
+            if moments_of_inertia[dimension] < 0:
                 raise ValueError(
                     f"Negative value encountered for moment of inertia: "
-                    f"{moment_of_inertia[dimension, dimension]} "
+                    f"{moment_of_inertia[dimension]} "
                     f"Cannot compute weighted torque."
                 )
 
-            # Compute weighted torque
             weighted_torque[dimension] = torques[dimension] / np.sqrt(
-                moment_of_inertia[dimension, dimension]
+                moments_of_inertia[dimension]
             )
 
         logger.debug(f"Weighted Torque: {weighted_torque}")
@@ -817,7 +615,6 @@ def update_force_torque_matrices(
                 f_mat, t_mat = self.get_matrices(
                     res,
                     level,
-                    num_frames,
                     highest,
                     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],
@@ -847,7 +644,6 @@ def update_force_torque_matrices(
             f_mat, t_mat = self.get_matrices(
                 mol,
                 level,
-                num_frames,
                 highest,
                 None if force_avg[key][group_id] is None else force_avg[key][group_id],
                 (
diff --git a/tests/test_CodeEntropy/test_levels.py b/tests/test_CodeEntropy/test_levels.py
index dd751de..98a2684 100644
--- a/tests/test_CodeEntropy/test_levels.py
+++ b/tests/test_CodeEntropy/test_levels.py
@@ -1,4 +1,4 @@
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock
 
 import numpy as np
 
@@ -60,18 +60,19 @@ def test_select_levels(self):
 
     def test_get_matrices(self):
         """
-        Test `get_matrices` with mocked internal methods and a simple trajectory.
-        Ensures that the method returns correctly shaped matrices after filtering.
+        Test `get_matrices` with mocked internal methods and a simple setup.
+        Ensures that the method returns correctly shaped force and torque matrices.
         """
-
-        # Create a mock UniverseOperations and LevelManager
         universe_operations = UniverseOperations()
-
         level_manager = LevelManager(universe_operations)
 
-        # Mock internal methods
-        level_manager.get_beads = MagicMock(return_value=["bead1", "bead2"])
-        level_manager.get_axes = MagicMock(return_value=("trans_axes", "rot_axes"))
+        bead1 = MagicMock()
+        bead1.principal_axes.return_value = np.ones(3)
+
+        bead2 = MagicMock()
+        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])
         )
@@ -79,39 +80,32 @@ def test_get_matrices(self):
             return_value=np.array([0.5, 1.5, 2.5])
         )
         level_manager.create_submatrix = MagicMock(return_value=np.identity(3))
-        level_manager.filter_zero_rows_columns = MagicMock(side_effect=lambda x: x)
 
-        # Mock data_container and trajectory
         data_container = MagicMock()
-        timestep1 = MagicMock()
-        timestep1.frame = 0
-        timestep2 = MagicMock()
-        timestep2.frame = 1
-        data_container.trajectory.__getitem__.return_value = [timestep1, timestep2]
+        data_container.atoms = MagicMock()
+        data_container.atoms.principal_axes.return_value = np.ones(3)
 
-        # Call the method
         force_matrix, torque_matrix = level_manager.get_matrices(
             data_container=data_container,
             level="residue",
-            number_frames=2,
             highest_level=True,
             force_matrix=None,
             torque_matrix=None,
             force_partitioning=0.5,
         )
 
-        # Assertions
-        self.assertTrue(isinstance(force_matrix, np.ndarray))
-        self.assertTrue(isinstance(torque_matrix, np.ndarray))
-        self.assertEqual(force_matrix.shape, (6, 6))  # 2 beads × 3D
+        self.assertIsInstance(force_matrix, np.ndarray)
+        self.assertIsInstance(torque_matrix, np.ndarray)
+
+        self.assertEqual(force_matrix.shape, (6, 6))
         self.assertEqual(torque_matrix.shape, (6, 6))
 
-        # Check that internal methods were called
-        self.assertEqual(level_manager.get_beads.call_count, 1)
-        self.assertEqual(level_manager.get_axes.call_count, 2)  # 2 beads
-        self.assertEqual(
-            level_manager.create_submatrix.call_count, 6
-        )  # 3 force + 3 torque
+        level_manager.get_beads.assert_called_once_with(data_container, "residue")
+
+        self.assertEqual(level_manager.get_weighted_forces.call_count, 2)
+        self.assertEqual(level_manager.get_weighted_torques.call_count, 2)
+
+        self.assertEqual(level_manager.create_submatrix.call_count, 6)
 
     def test_get_matrices_force_shape_mismatch(self):
         """
@@ -119,12 +113,16 @@ def test_get_matrices_force_shape_mismatch(self):
         has a shape mismatch with the computed force block matrix.
         """
         universe_operations = UniverseOperations()
-
         level_manager = LevelManager(universe_operations)
 
-        # Mock internal methods
-        level_manager.get_beads = MagicMock(return_value=["bead1", "bead2"])
-        level_manager.get_axes = MagicMock(return_value=("trans_axes", "rot_axes"))
+        bead1 = MagicMock()
+        bead1.principal_axes.return_value = np.ones(3)
+
+        bead2 = MagicMock()
+        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])
         )
@@ -134,8 +132,9 @@ def test_get_matrices_force_shape_mismatch(self):
         level_manager.create_submatrix = MagicMock(return_value=np.identity(3))
 
         data_container = MagicMock()
+        data_container.atoms = MagicMock()
+        data_container.atoms.principal_axes.return_value = np.ones(3)
 
-        # Incorrect shape for force matrix (should be 6x6 for 2 beads)
         bad_force_matrix = np.zeros((3, 3))
         correct_torque_matrix = np.zeros((6, 6))
 
@@ -143,7 +142,6 @@ def test_get_matrices_force_shape_mismatch(self):
             level_manager.get_matrices(
                 data_container=data_container,
                 level="residue",
-                number_frames=2,
                 highest_level=True,
                 force_matrix=bad_force_matrix,
                 torque_matrix=correct_torque_matrix,
@@ -158,12 +156,16 @@ def test_get_matrices_torque_shape_mismatch(self):
         has a shape mismatch with the computed torque block matrix.
         """
         universe_operations = UniverseOperations()
-
         level_manager = LevelManager(universe_operations)
 
-        # Mock internal methods
-        level_manager.get_beads = MagicMock(return_value=["bead1", "bead2"])
-        level_manager.get_axes = MagicMock(return_value=("trans_axes", "rot_axes"))
+        bead1 = MagicMock()
+        bead1.principal_axes.return_value = np.ones(3)
+
+        bead2 = MagicMock()
+        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])
         )
@@ -173,6 +175,8 @@ def test_get_matrices_torque_shape_mismatch(self):
         level_manager.create_submatrix = MagicMock(return_value=np.identity(3))
 
         data_container = MagicMock()
+        data_container.atoms = MagicMock()
+        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
@@ -181,7 +185,6 @@ def test_get_matrices_torque_shape_mismatch(self):
             level_manager.get_matrices(
                 data_container=data_container,
                 level="residue",
-                number_frames=2,
                 highest_level=True,
                 force_matrix=correct_force_matrix,
                 torque_matrix=bad_torque_matrix,
@@ -192,15 +195,20 @@ def test_get_matrices_torque_shape_mismatch(self):
 
     def test_get_matrices_torque_consistency(self):
         """
-        Test that get_matrices returns consistent torque and force matrices
+        Test that get_matrices returns consistent force and torque matrices
         when called multiple times with the same inputs.
         """
         universe_operations = UniverseOperations()
-
         level_manager = LevelManager(universe_operations)
 
-        level_manager.get_beads = MagicMock(return_value=["bead1", "bead2"])
-        level_manager.get_axes = MagicMock(return_value=("trans_axes", "rot_axes"))
+        bead1 = MagicMock()
+        bead1.principal_axes.return_value = np.ones(3)
+
+        bead2 = MagicMock()
+        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])
         )
@@ -210,6 +218,8 @@ def test_get_matrices_torque_consistency(self):
         level_manager.create_submatrix = MagicMock(return_value=np.identity(3))
 
         data_container = MagicMock()
+        data_container.atoms = MagicMock()
+        data_container.atoms.principal_axes.return_value = np.ones(3)
 
         initial_force_matrix = np.zeros((6, 6))
         initial_torque_matrix = np.zeros((6, 6))
@@ -217,7 +227,6 @@ def test_get_matrices_torque_consistency(self):
         force_matrix_1, torque_matrix_1 = level_manager.get_matrices(
             data_container=data_container,
             level="residue",
-            number_frames=2,
             highest_level=True,
             force_matrix=initial_force_matrix.copy(),
             torque_matrix=initial_torque_matrix.copy(),
@@ -227,16 +236,14 @@ def test_get_matrices_torque_consistency(self):
         force_matrix_2, torque_matrix_2 = level_manager.get_matrices(
             data_container=data_container,
             level="residue",
-            number_frames=2,
             highest_level=True,
             force_matrix=initial_force_matrix.copy(),
             torque_matrix=initial_torque_matrix.copy(),
             force_partitioning=0.5,
         )
 
-        # Check that repeated calls produce the same output
-        self.assertTrue(np.allclose(torque_matrix_1, torque_matrix_2, atol=1e-8))
-        self.assertTrue(np.allclose(force_matrix_1, force_matrix_2, atol=1e-8))
+        np.testing.assert_array_equal(force_matrix_1, force_matrix_2)
+        np.testing.assert_array_equal(torque_matrix_1, torque_matrix_2)
 
     def test_get_beads_polymer_level(self):
         """
@@ -328,260 +335,6 @@ def test_get_beads_hydrogen_molecule(self):
             data_container.select_atoms.call_count, 2
         )  # 1 for heavy_atoms + 1 beads
 
-    def test_get_axes_united_atom_no_bonds(self):
-        """
-        Test `get_axes` for 'united_atom' level when no bonded atoms are found.
-        Ensures that rotational axes fall back to residues' principal axes.
-        """
-        universe_operations = UniverseOperations()
-
-        level_manager = LevelManager(universe_operations)
-
-        data_container = MagicMock()
-
-        # Mock principal axes for translation and rotation
-        mock_rot_axes = MagicMock(name="rot_axes")
-
-        data_container.residues.principal_axes.return_value = mock_rot_axes
-        data_container.residues.principal_axes.return_value = mock_rot_axes
-        data_container.residues.principal_axes.return_value = mock_rot_axes  # fallback
-
-        # First select_atoms returns empty bonded atom set
-        atom_set = MagicMock()
-        atom_set.__len__.return_value = 0  # triggers fallback
-
-        data_container.select_atoms.side_effect = [atom_set]
-
-        trans_axes, rot_axes = level_manager.get_axes(
-            data_container=data_container, level="united_atom", index=5
-        )
-
-        # Assertions
-        self.assertEqual(trans_axes, mock_rot_axes)
-        self.assertEqual(rot_axes, mock_rot_axes)
-        data_container.residues.principal_axes.assert_called()
-        self.assertEqual(data_container.select_atoms.call_count, 1)
-
-    def test_get_axes_polymer_level(self):
-        """
-        Test `get_axes` for 'polymer' level.
-        Should return principal axes of the full system for both
-        translation and rotation.
-        """
-        universe_operations = UniverseOperations()
-
-        level_manager = LevelManager(universe_operations)
-
-        data_container = MagicMock()
-        principal_axes = np.identity(3)
-        data_container.atoms.principal_axes.return_value = principal_axes
-
-        trans_axes, rot_axes = level_manager.get_axes(data_container, level="polymer")
-
-        self.assertTrue((trans_axes == principal_axes).all())
-        self.assertTrue((rot_axes == principal_axes).all())
-
-    def test_get_axes_residue_level_with_bonds(self):
-        """
-        Test `get_axes` for 'residue' level with bonded neighbors.
-        Should use spherical coordinate axes for rotation.
-        """
-        universe_operations = UniverseOperations()
-
-        level_manager = LevelManager(universe_operations)
-
-        data_container = MagicMock()
-        data_container.atoms.principal_axes.return_value = "trans_axes"
-
-        atom_set = MagicMock()
-        atom_set.__len__.return_value = 1
-
-        residue = MagicMock()
-        residue.atoms.center_of_mass.return_value = "center"
-        residue.atoms.principal_axes.return_value = "fallback_rot_axes"
-
-        data_container.select_atoms.side_effect = [atom_set, residue]
-
-        level_manager.get_avg_pos = MagicMock(return_value="vector")
-        level_manager.get_sphCoord_axes = MagicMock(return_value="rot_axes")
-
-        trans_axes, rot_axes = level_manager.get_axes(
-            data_container, level="residue", index=2
-        )
-
-        self.assertEqual(trans_axes, "trans_axes")
-        self.assertEqual(rot_axes, "rot_axes")
-
-    def test_get_axes_residue_level_without_bonds(self):
-        """
-        Test `get_axes` for 'residue' level with no bonded neighbors.
-        Should use principal axes of the residue for rotation.
-        """
-        universe_operations = UniverseOperations()
-
-        level_manager = LevelManager(universe_operations)
-
-        data_container = MagicMock()
-        data_container.atoms.principal_axes.return_value = "trans_axes"
-
-        empty_atom_set = []
-        residue = MagicMock()
-        residue.atoms.principal_axes.return_value = "rot_axes"
-
-        data_container.select_atoms.side_effect = [empty_atom_set, residue]
-
-        trans_axes, rot_axes = level_manager.get_axes(
-            data_container, level="residue", index=2
-        )
-
-        self.assertEqual(trans_axes, "trans_axes")
-        self.assertEqual(rot_axes, "rot_axes")
-
-    def test_get_axes_united_atom_level(self):
-        """
-        Test `get_axes` for 'united_atom' level.
-        Should use residue principal axes for translation and spherical
-        axes for rotation.
-        """
-        universe_operations = UniverseOperations()
-
-        level_manager = LevelManager(universe_operations)
-
-        data_container = MagicMock()
-        data_container.residues.principal_axes.return_value = "trans_axes"
-
-        atom_set = MagicMock()
-        atom_set.__len__.return_value = 1
-
-        atom_group = MagicMock()
-        atom_group.positions = [[1.0, 2.0, 3.0]]
-
-        data_container.select_atoms.side_effect = [atom_set, atom_group]
-
-        level_manager.get_avg_pos = MagicMock(return_value="vector")
-        level_manager.get_sphCoord_axes = MagicMock(return_value="rot_axes")
-
-        trans_axes, rot_axes = level_manager.get_axes(
-            data_container, level="united_atom", index=5
-        )
-
-        self.assertEqual(trans_axes, "trans_axes")
-        self.assertEqual(rot_axes, "rot_axes")
-
-    def test_get_avg_pos_with_atoms(self):
-        """
-        Test `get_avg_pos` with a non-empty atom set.
-        Should return the average of atom positions minus the center.
-        """
-        universe_operations = UniverseOperations()
-
-        level_manager = LevelManager(universe_operations)
-
-        atom1 = MagicMock()
-        atom1.position = np.array([1.0, 2.0, 3.0])
-        atom2 = MagicMock()
-        atom2.position = np.array([4.0, 5.0, 6.0])
-
-        atom_set = MagicMock()
-        atom_set.names = ["A", "B"]
-        atom_set.atoms = [atom1, atom2]
-
-        center = np.array([1.0, 1.0, 1.0])
-        expected_avg = ((atom1.position + atom2.position) / 2) - center
-
-        result = level_manager.get_avg_pos(atom_set, center)
-        np.testing.assert_array_almost_equal(result, expected_avg)
-
-    @patch("numpy.random.random")
-    def test_get_avg_pos_empty(self, mock_random):
-        """
-        Test `get_avg_pos` with an empty atom set.
-        Should return a random vector minus the center.
-        """
-        universe_operations = UniverseOperations()
-
-        level_manager = LevelManager(universe_operations)
-
-        atom_set = MagicMock()
-        atom_set.names = []
-        atom_set.atoms = []
-
-        center = np.array([1.0, 1.0, 1.0])
-        mock_random.return_value = np.array([0.5, 0.5, 0.5])
-
-        result = level_manager.get_avg_pos(atom_set, center)
-        expected = np.array([0.5, 0.5, 0.5]) - center
-
-        np.testing.assert_array_almost_equal(result, expected)
-
-    def test_get_sphCoord_axes_valid_vector(self):
-        """
-        Test with a valid non-zero vector.
-        Should return a 3x3 orthonormal basis matrix.
-        """
-        universe_operations = UniverseOperations()
-
-        level_manager = LevelManager(universe_operations)
-
-        vector = np.array([1.0, 1.0, 1.0])
-        result = level_manager.get_sphCoord_axes(vector)
-
-        self.assertEqual(result.shape, (3, 3))
-        self.assertTrue(np.all(np.isfinite(result)))
-
-    def test_get_sphCoord_axes_vector_on_z_axis_raises(self):
-        """
-        Test with a vector along the z-axis (x2y2 == 0).
-        Should raise ValueError due to undefined phi.
-        """
-        universe_operations = UniverseOperations()
-
-        level_manager = LevelManager(universe_operations)
-
-        vector = np.array([0.0, 0.0, 1.0])
-        with self.assertRaises(ValueError):
-            level_manager.get_sphCoord_axes(vector)
-
-    def test_get_sphCoord_axes_negative_x2y2_div_r2(self):
-        """
-        Test with a vector that would cause x2y2 / r2 < 0.
-        """
-        universe_operations = UniverseOperations()
-
-        level_manager = LevelManager(universe_operations)
-
-        vector = np.array([1e-10, 1e-10, 1e10])  # x2y2 is tiny, r2 is huge
-        result = level_manager.get_sphCoord_axes(vector)
-        self.assertEqual(result.shape, (3, 3))
-
-    def test_get_sphCoord_axes_zero_vector_raises(self):
-        """
-        Test with a zero vector.
-        Should raise ValueError due to r2 == 0.
-        """
-        universe_operations = UniverseOperations()
-
-        level_manager = LevelManager(universe_operations)
-
-        vector = np.array([0.0, 0.0, 0.0])
-        with self.assertRaises(ValueError) as context:
-            level_manager.get_sphCoord_axes(vector)
-        self.assertIn("r2 is zero", str(context.exception))
-
-    def test_get_sphCoord_axes_x2y2_zero_raises(self):
-        """
-        Test with a vector along the z-axis (x2y2 == 0, r2 != 0).
-        Should raise ValueError due to undefined phi.
-        """
-        universe_operations = UniverseOperations()
-
-        level_manager = LevelManager(universe_operations)
-
-        vector = np.array([0.0, 0.0, 1.0])  # r2 = 1.0, x2y2 = 0.0
-        with self.assertRaises(ValueError) as context:
-            level_manager.get_sphCoord_axes(vector)
-        self.assertIn("x2y2 is zero", str(context.exception))
-
     def test_get_weighted_force_with_partitioning(self):
         """
         Test correct weighted force calculation with partitioning enabled.
@@ -765,9 +518,9 @@ def test_get_weighted_torques_zero_torque_skips_division(self):
         )
         np.testing.assert_array_almost_equal(result, np.zeros(3))
 
-    def test_get_weighted_torques_zero_moi_raises(self):
+    def test_get_weighted_torques_zero_moi(self):
         """
-        Should raise ZeroDivisionError when moment of inertia is zero in a dimension
+        Should set torque to 0 when moment of inertia is zero in a dimension
         and torque in that dimension is non-zero.
         """
         universe_operations = UniverseOperations()
@@ -782,7 +535,7 @@ def test_get_weighted_torques_zero_moi_raises(self):
         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.identity(3)
+        moi = np.zeros((3, 3))
         moi[2, 2] = 0.0
         bead.moment_of_inertia.return_value = moi
 
@@ -797,10 +550,11 @@ def test_get_weighted_torques_zero_moi_raises(self):
 
         force_partitioning = 0.5
 
-        with self.assertRaises(ZeroDivisionError):
-            level_manager.get_weighted_torques(
-                data_container, bead, rot_axes, force_partitioning
-            )
+        torque = level_manager.get_weighted_torques(
+            data_container, bead, rot_axes, force_partitioning
+        )
+
+        self.assertEqual(torque[2], 0)
 
     def test_get_weighted_torques_negative_moi_raises(self):
         """
@@ -820,7 +574,7 @@ def test_get_weighted_torques_negative_moi_raises(self):
 
         # Set moment of inertia with negative value in dimension 2
         moi = np.identity(3)
-        moi[2, 2] = -1.0
+        moi *= -1.0
         bead.moment_of_inertia.return_value = moi
 
         data_container = MagicMock()