From 9920253652d1f21de0199af2dfd6ead5da7c2d55 Mon Sep 17 00:00:00 2001 From: ioana Date: Thu, 12 Mar 2026 13:41:12 +0000 Subject: [PATCH 01/47] data_container at UA level is now formed of a residue group and customised axes for rotation at residue level are introduced --- CodeEntropy/levels/axes.py | 247 ++++++++++++++++++++++--- CodeEntropy/levels/nodes/covariance.py | 42 ++++- 2 files changed, 253 insertions(+), 36 deletions(-) diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py index 2d05fcb3..5a5d7739 100644 --- a/CodeEntropy/levels/axes.py +++ b/CodeEntropy/levels/axes.py @@ -99,43 +99,109 @@ def get_residue_axes(self, data_container, index: int, residue=None): If the residue selection is empty. """ # TODO refine selection so that it will work for branched polymers + # match indexing to MDAnalysis indexing index_prev = index - 1 index_next = index + 1 - if residue is None: residue = data_container.select_atoms(f"resindex {index}") + # residue of interest if len(residue) == 0: raise ValueError(f"Empty residue selection for resindex={index}") - - center = residue.atoms.center_of_mass(unwrap=True) - atom_set = data_container.select_atoms( - f"(resindex {index_prev} or resindex {index_next}) and bonded resid {index}" + atom_set = data_container.atoms.select_atoms( + f"(resindex {index_prev} or " + f"resindex {index_next}) and " + f"bonded resindex {index}" ) + uas = residue.select_atoms("mass 2 to 999") + ua_masses = self.get_UA_masses(residue) 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 H. - uas = residue.select_atoms("mass 2 to 999") - ua_masses = self.get_UA_masses(residue) moi_tensor = self.get_moment_of_inertia_tensor( - center_of_mass=center, + center_of_mass=np.array(residue.center_of_mass()), positions=uas.positions, masses=ua_masses, dimensions=data_container.dimensions[:3], ) rot_axes, moment_of_inertia = self.get_custom_principal_axes(moi_tensor) trans_axes = rot_axes # per original convention + center = np.array(residue.center_of_mass()) else: - # If bonded to other residues, use default axes and MOI. + # If bonded to other residues, use local axes. make_whole(data_container.atoms) trans_axes = data_container.atoms.principal_axes() - rot_axes, moment_of_inertia = self.get_vanilla_axes(residue) - center = residue.center_of_mass(unwrap=True) - + backbone = residue.select_atoms("backbone") + if len(backbone) == 0: + # not a protein + center = np.array(residue.center_of_mass()) + else: + # protein backbone identified + center = np.array(backbone.center_of_mass()) + + if len(atom_set) == 1: + # only one neighbour + if index == 0: + # first residue + next_resid = data_container.select_atoms(f"resindex {index_next}") + next_backbone = next_resid.select_atoms("backbone") + if len(next_backbone) == 0: + anchor = np.array(next_resid.center_of_mass()) + else: + anchor = np.array(next_backbone.center_of_mass()) + # anchor = atom_set[0].position + else: + # last residue + prev_resid = data_container.select_atoms(f"resindex {index_prev}") + prev_backbone = prev_resid.select_atoms("backbone") + if len(prev_backbone) == 0: + anchor = np.array(prev_resid.center_of_mass()) + else: + anchor = np.array(prev_backbone.center_of_mass()) + # anchor = atom_set[0].position + rot_axes = self.get_custom_axes( + a=center, + b_list=anchor, + c=np.zeros(3), + dimensions=data_container.dimensions[:3], + ) + else: + # two neighbours + prev_resid = data_container.select_atoms(f"resindex {index_prev}") + next_resid = data_container.select_atoms(f"resindex {index_next}") + prev_backbone = prev_resid.select_atoms("backbone") + next_backbone = next_resid.select_atoms("backbone") + anchors = [] + # check separately in case we have a protein with a PTM + # or similar case + if len(prev_backbone) == 0: + anchors.append(np.array(prev_resid.center_of_mass())) + else: + anchors.append(np.array(prev_backbone.center_of_mass())) + if len(next_backbone) == 0: + anchors.append(np.array(next_resid.center_of_mass())) + else: + anchors.append(np.array(next_backbone.center_of_mass())) + # anchors = atom_set.positions + rot_axes = self.get_custom_axes( + a=center, + b_list=anchors, + c=anchors[1], + dimensions=data_container.dimensions[:3], + ) + # analogous to the UA case where a heavy atom is bound to >=2 heavy atoms + + moment_of_inertia = self.get_custom_residue_moment_of_inertia( + center_of_mass=center, + positions=uas.positions, + masses=ua_masses, + custom_rot_axes=rot_axes, + dimensions=data_container.dimensions[:3], + ) return trans_axes, rot_axes, center, moment_of_inertia - def get_UA_axes(self, data_container, index: int): + def get_UA_axes(self, data_container, index: int, res_position): """Compute united-atom-level translational and rotational axes. The translational and rotational axes at the united-atom level. @@ -176,20 +242,110 @@ def get_UA_axes(self, data_container, index: int): index = int(index) # bead index # use the same customPI trans axes as the residue level - heavy_atoms = data_container.select_atoms("prop mass > 1.1") - if len(heavy_atoms) > 1: - UA_masses = self.get_UA_masses(data_container.atoms) + if len(data_container.residues) == 1: + # only the one residue => use principal axes + residue = data_container center = data_container.atoms.center_of_mass(unwrap=True) - moment_of_inertia_tensor = self.get_moment_of_inertia_tensor( - center, heavy_atoms.positions, UA_masses, data_container.dimensions[:3] - ) - trans_axes, _moment_of_inertia = self.get_custom_principal_axes( - moment_of_inertia_tensor - ) + trans_axes = data_container.atoms.principal_axes else: - # use standard PA for UA not bonded to anything else - make_whole(data_container.atoms) - trans_axes = data_container.atoms.principal_axes() + # residue of interest has at least one neighbour + if res_position == -1: + residue = data_container.residues[0] + index_next = residue.resid + 1 + # atom_set = data_container.atoms.select_atoms( + # f"resindex {index_next-1} and " + # f"bonded resindex {residue.resid-1}" + # ) + # the .resid attribute gives 1-indexing + # substract 1 to match indexing later + next_resid = data_container.select_atoms(f"resindex {index_next - 1}") + next_backbone = next_resid.atoms.select_atoms("backbone") + if len(next_backbone) == 0: + anchor = np.array(next_resid.center_of_mass()) + else: + anchor = np.array(next_backbone.center_of_mass()) + # anchor = atom_set[0].position + backbone = residue.atoms.select_atoms("backbone") + if len(backbone) == 0: + # not a protein + center = np.array(residue.atoms.center_of_mass()) + else: + # protein backbone identified + center = np.array(backbone.center_of_mass()) + trans_axes = self.get_custom_axes( + a=center, + b_list=anchor, + c=np.zeros(3), + dimensions=data_container.dimensions[:3], + ) + + elif res_position == 0: + # between 2 residues + residue = data_container.residues[1] + index_prev = residue.resid - 1 + index_next = residue.resid + 1 + prev_resid = data_container.select_atoms(f"resindex {index_prev - 1}") + next_resid = data_container.select_atoms(f"resindex {index_next - 1}") + prev_backbone = prev_resid.atoms.select_atoms("backbone") + next_backbone = next_resid.atoms.select_atoms("backbone") + # atom_set = data_container.atoms.select_atoms( + # f"(resindex {index_prev-1} or " + # f"resindex {index_next-1}) and " + # f"bonded resindex {residue.resid-1}" + # ) + anchors = [] + if len(prev_backbone) == 0: + anchors.append(np.array(prev_resid.center_of_mass())) + else: + anchors.append(np.array(prev_backbone.center_of_mass())) + if len(next_backbone) == 0: + anchors.append(np.array(next_resid.center_of_mass())) + else: + anchors.append(np.array(next_backbone.center_of_mass())) + # anchors = atom_set.positions + backbone = residue.atoms.select_atoms("backbone") + if len(backbone) == 0: + # not a protein + center = np.array(residue.atoms.center_of_mass()) + else: + # protein backbone identified + center = np.array(backbone.center_of_mass()) + trans_axes = self.get_custom_axes( + a=center, + b_list=anchors, + c=anchors[1], + dimensions=data_container.dimensions[:3], + ) + + else: + # last resid + residue = data_container.residues[1] + index_prev = residue.resid - 1 + prev_resid = data_container.select_atoms(f"resindex {index_prev - 1}") + prev_backbone = prev_resid.atoms.select_atoms("backbone") + if len(prev_backbone) == 0: + anchor = np.array(prev_resid.center_of_mass()) + else: + anchor = np.array(prev_backbone.center_of_mass()) + # atom_set = data_container.atoms.select_atoms( + # f"resindex {index_prev-1} and " + # f"bonded resindex {residue.resid-1}" + # ) + # anchor = atom_set[0].position + backbone = residue.atoms.select_atoms("backbone") + if len(backbone) == 0: + # not a protein + center = np.array(residue.atoms.center_of_mass()) + else: + center = np.array(backbone.center_of_mass()) + trans_axes = self.get_custom_axes( + a=center, + b_list=anchor, + c=np.zeros(3), + dimensions=data_container.dimensions[:3], + ) + + heavy_atoms = residue.atoms.select_atoms("mass 2 to 999") # look for heavy atoms in residue of interest heavy_atom_indices = [] @@ -198,9 +354,9 @@ def get_UA_axes(self, data_container, index: int): # 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}") + heavy_atom = residue.atoms.select_atoms(f"index {heavy_atom_index}") - center = heavy_atom.positions[0] + rot_center = heavy_atom.positions[0] rot_axes, moment_of_inertia = self.get_bonded_axes( system=data_container, atom=heavy_atom[0], @@ -214,7 +370,7 @@ def get_UA_axes(self, data_container, index: int): logger.debug("Center: %s", center) logger.debug("Moment of Inertia: %s", moment_of_inertia) - return trans_axes, rot_axes, center, moment_of_inertia + return trans_axes, rot_axes, rot_center, moment_of_inertia def get_bonded_axes(self, system, atom, dimensions: np.ndarray): r"""Compute UA rotational axes from bonded topology around a heavy atom. @@ -446,6 +602,41 @@ def get_custom_axes( scaled_custom_axes = unscaled_custom_axes / mod[:, np.newaxis] return scaled_custom_axes + def get_custom_residue_moment_of_inertia( + self, + center_of_mass: np.ndarray, + positions: np.ndarray, + masses: np.ndarray, + custom_rot_axes: np.ndarray, + dimensions: np.ndarray, + ): + """ + Compute moment of inertia around custom axes for a bead + formed of multiple UAs. + + Args: + center_of_mass: (3, ) COM for bead + positions: (N,3) positions of the UAs in the bead + masses: (N,) masses of the UAs in the bead + custom_rot_axes: (3,3) array of residue rotation axes + dimensions: (3,) simulation_box_dimensions + + Returns: + np.ndarray: (3,) moment of inertia array. + + """ + + translated_coords = self.get_vector(center_of_mass, positions, dimensions) + custom_moment_of_inertia = np.zeros(3, dtype=float) + + for coord, mass in zip(translated_coords, masses, strict=True): + axis_component = np.sum( + np.cross(custom_rot_axes, coord) ** 2 * mass, axis=1 + ) + custom_moment_of_inertia += axis_component + + return custom_moment_of_inertia + def get_custom_moment_of_inertia( self, UA, diff --git a/CodeEntropy/levels/nodes/covariance.py b/CodeEntropy/levels/nodes/covariance.py index dd56d180..48f62dad 100644 --- a/CodeEntropy/levels/nodes/covariance.py +++ b/CodeEntropy/levels/nodes/covariance.py @@ -200,7 +200,25 @@ def _process_united_atom( Returns: None. Mutates out_force/out_torque and molcount in-place. """ + for local_res_i, res in enumerate(mol.residues): + # build residue group here + if local_res_i == 0: + # first residue + res_position = -1 + res_next = mol.residues[1] + residue_group = res + res_next + elif local_res_i == len(mol.residues) - 1: + # last residue + res_position = 1 + res_prev = mol.residues[-2] + residue_group = res + res_prev + else: + res_position = 0 + res_prev = mol.residues[local_res_i - 1] + res_next = mol.residues[local_res_i + 1] + residue_group = res_prev + res + res_next + bead_key = (mol_id, "united_atom", local_res_i) bead_idx_list = beads.get(bead_key, []) if not bead_idx_list: @@ -211,13 +229,14 @@ def _process_united_atom( continue force_vecs, torque_vecs = self._build_ua_vectors( - residue_atoms=res.atoms, + residue_group=residue_group.atoms, bead_groups=bead_groups, axes_manager=axes_manager, box=box, force_partitioning=force_partitioning, customised_axes=customised_axes, is_highest=is_highest, + res_position=res_position, ) F, T = self._ft.compute_frame_covariance(force_vecs, torque_vecs) @@ -413,23 +432,26 @@ def _build_ua_vectors( self, *, bead_groups: List[Any], - residue_atoms: Any, + residue_group: Any, axes_manager: Any, box: Optional[np.ndarray], force_partitioning: float, customised_axes: bool, is_highest: bool, + res_position: int, ) -> Tuple[List[np.ndarray], List[np.ndarray]]: """Build force/torque vectors for UA-level beads of one residue. Args: bead_groups: List of UA bead AtomGroups for the residue. - residue_atoms: AtomGroup for the residue atoms (used for axes when vanilla). + residue_group: AtomGroup for the residue group atoms. axes_manager: Axes manager used to determine axes/centers/MOI. box: Optional box vector used for PBC-aware displacements. force_partitioning: Force scaling factor applied at highest level. customised_axes: Whether to use customised axes methods when available. is_highest: Whether UA level is the highest level for the molecule. + res_position: Where the residue is in the residue group + Returns: A tuple (force_vecs, torque_vecs), each a list of (3,) vectors ordered @@ -437,17 +459,21 @@ def _build_ua_vectors( """ force_vecs: List[np.ndarray] = [] torque_vecs: List[np.ndarray] = [] - for ua_i, bead in enumerate(bead_groups): if customised_axes: trans_axes, rot_axes, center, moi = axes_manager.get_UA_axes( - residue_atoms, ua_i + residue_group, ua_i, res_position ) else: - make_whole(residue_atoms) + make_whole(residue_group) make_whole(bead) - - trans_axes = residue_atoms.principal_axes() + if res_position == -1: + # first residue in group + residue = residue_group.residues[0] + else: + # middle or last residue => second in group + residue = residue_group.residues[1] + trans_axes = residue.atoms.principal_axes() rot_axes, moi = axes_manager.get_vanilla_axes(bead) center = bead.center_of_mass(unwrap=True) From f1de766171a14a4b96cae65ec2909dcaec4cd8a9 Mon Sep 17 00:00:00 2001 From: ioana Date: Wed, 25 Mar 2026 10:33:19 +0000 Subject: [PATCH 02/47] added new functions to find custom backbone --- CodeEntropy/levels/axes.py | 309 ++++++++++++++++++++++--------------- 1 file changed, 185 insertions(+), 124 deletions(-) diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py index 14250785..7db30c79 100644 --- a/CodeEntropy/levels/axes.py +++ b/CodeEntropy/levels/axes.py @@ -67,7 +67,7 @@ def get_residue_axes(self, data_container, index: int, residue=None): The translational and rotational axes at the residue level. - Identify the residue (either provided or selected by `resindex index`). - - Determine whether the residue is bonded to neighboring residues + - Determine whether the residue is bonded to neighbouring residues (previous/next in sequence) using MDAnalysis bonded selections. - If there are *no* bonds to other residues: * Use a custom principal axes, from a moment-of-inertia (MOI) tensor @@ -107,15 +107,16 @@ def get_residue_axes(self, data_container, index: int, residue=None): # residue of interest if len(residue) == 0: raise ValueError(f"Empty residue selection for resindex={index}") - atom_set = data_container.atoms.select_atoms( + anchors = data_container.select_atoms( f"(resindex {index_prev} or " f"resindex {index_next}) and " f"bonded resindex {index}" ) + uas = residue.select_atoms("mass 2 to 999") ua_masses = self.get_UA_masses(residue) - if len(atom_set) == 0: + if len(anchors) == 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 H. @@ -129,65 +130,33 @@ def get_residue_axes(self, data_container, index: int, residue=None): trans_axes = rot_axes # per original convention center = np.array(residue.center_of_mass()) else: + print("-----RESIDUE LEVEL-----") # If bonded to other residues, use local axes. make_whole(data_container.atoms) trans_axes = data_container.atoms.principal_axes() - backbone = residue.select_atoms("backbone") - if len(backbone) == 0: - # not a protein - center = np.array(residue.center_of_mass()) - else: - # protein backbone identified - center = np.array(backbone.center_of_mass()) - - if len(atom_set) == 1: + backbone = self.get_backbone(data_container, index) + # get edge atoms of the residue + # for terminal residues, this will include the C/N terminus + # bond-derived rotation axes + print(f"The backbone: {backbone}") + center = np.zeros(3) + for atom in backbone: + center += atom.position + center /= len(backbone) + if len(anchors) == 1: # only one neighbour - if index == 0: - # first residue - next_resid = data_container.select_atoms(f"resindex {index_next}") - next_backbone = next_resid.select_atoms("backbone") - if len(next_backbone) == 0: - anchor = np.array(next_resid.center_of_mass()) - else: - anchor = np.array(next_backbone.center_of_mass()) - # anchor = atom_set[0].position - else: - # last residue - prev_resid = data_container.select_atoms(f"resindex {index_prev}") - prev_backbone = prev_resid.select_atoms("backbone") - if len(prev_backbone) == 0: - anchor = np.array(prev_resid.center_of_mass()) - else: - anchor = np.array(prev_backbone.center_of_mass()) - # anchor = atom_set[0].position rot_axes = self.get_custom_axes( a=center, - b_list=anchor, + b_list=anchors[0].position, c=np.zeros(3), dimensions=data_container.dimensions[:3], ) else: # two neighbours - prev_resid = data_container.select_atoms(f"resindex {index_prev}") - next_resid = data_container.select_atoms(f"resindex {index_next}") - prev_backbone = prev_resid.select_atoms("backbone") - next_backbone = next_resid.select_atoms("backbone") - anchors = [] - # check separately in case we have a protein with a PTM - # or similar case - if len(prev_backbone) == 0: - anchors.append(np.array(prev_resid.center_of_mass())) - else: - anchors.append(np.array(prev_backbone.center_of_mass())) - if len(next_backbone) == 0: - anchors.append(np.array(next_resid.center_of_mass())) - else: - anchors.append(np.array(next_backbone.center_of_mass())) - # anchors = atom_set.positions rot_axes = self.get_custom_axes( a=center, - b_list=anchors, - c=anchors[1], + b_list=anchors.positions, + c=anchors[1].position, dimensions=data_container.dimensions[:3], ) # analogous to the UA case where a heavy atom is bound to >=2 heavy atoms @@ -224,7 +193,8 @@ def get_UA_axes(self, data_container, index: int, res_position): Molecule and trajectory data. index (int): Bead index (ordinal among heavy atoms). - + res_position: where the residue of interest is + in data_container Returns: Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: - trans_axes: Translational axes (3, 3). @@ -245,36 +215,28 @@ def get_UA_axes(self, data_container, index: int, res_position): if len(data_container.residues) == 1: # only the one residue => use principal axes residue = data_container - center = data_container.atoms.center_of_mass(unwrap=True) + trans_center = data_container.atoms.center_of_mass(unwrap=True) trans_axes = data_container.atoms.principal_axes else: + print("-----UA LEVEL-----") # residue of interest has at least one neighbour if res_position == -1: residue = data_container.residues[0] index_next = residue.resid + 1 - # atom_set = data_container.atoms.select_atoms( - # f"resindex {index_next-1} and " - # f"bonded resindex {residue.resid-1}" - # ) + anchor = data_container.atoms.select_atoms( + f"resindex {index_next - 1} and bonded resindex {residue.resid - 1}" + ) # the .resid attribute gives 1-indexing # substract 1 to match indexing later - next_resid = data_container.select_atoms(f"resindex {index_next - 1}") - next_backbone = next_resid.atoms.select_atoms("backbone") - if len(next_backbone) == 0: - anchor = np.array(next_resid.center_of_mass()) - else: - anchor = np.array(next_backbone.center_of_mass()) - # anchor = atom_set[0].position - backbone = residue.atoms.select_atoms("backbone") - if len(backbone) == 0: - # not a protein - center = np.array(residue.atoms.center_of_mass()) - else: - # protein backbone identified - center = np.array(backbone.center_of_mass()) + backbone = self.get_backbone(data_container, 0) + print(f"The backbone: {backbone}") + trans_center = np.zeros(3) + for atom in backbone: + trans_center += atom.position + trans_center /= len(backbone) trans_axes = self.get_custom_axes( - a=center, - b_list=anchor, + a=trans_center, + b_list=anchor[0].position, c=np.zeros(3), dimensions=data_container.dimensions[:3], ) @@ -284,63 +246,39 @@ def get_UA_axes(self, data_container, index: int, res_position): residue = data_container.residues[1] index_prev = residue.resid - 1 index_next = residue.resid + 1 - prev_resid = data_container.select_atoms(f"resindex {index_prev - 1}") - next_resid = data_container.select_atoms(f"resindex {index_next - 1}") - prev_backbone = prev_resid.atoms.select_atoms("backbone") - next_backbone = next_resid.atoms.select_atoms("backbone") - # atom_set = data_container.atoms.select_atoms( - # f"(resindex {index_prev-1} or " - # f"resindex {index_next-1}) and " - # f"bonded resindex {residue.resid-1}" - # ) - anchors = [] - if len(prev_backbone) == 0: - anchors.append(np.array(prev_resid.center_of_mass())) - else: - anchors.append(np.array(prev_backbone.center_of_mass())) - if len(next_backbone) == 0: - anchors.append(np.array(next_resid.center_of_mass())) - else: - anchors.append(np.array(next_backbone.center_of_mass())) - # anchors = atom_set.positions - backbone = residue.atoms.select_atoms("backbone") - if len(backbone) == 0: - # not a protein - center = np.array(residue.atoms.center_of_mass()) - else: - # protein backbone identified - center = np.array(backbone.center_of_mass()) + anchors = data_container.atoms.select_atoms( + f"(resindex {index_prev - 1} or " + f"resindex {index_next - 1}) and " + f"bonded resindex {residue.resid - 1}" + ) + backbone = self.get_backbone(data_container, 1) + print(f"The backbone: {backbone}") + trans_center = np.zeros(3) + for atom in backbone: + trans_center += atom.position + trans_center /= len(backbone) trans_axes = self.get_custom_axes( - a=center, - b_list=anchors, - c=anchors[1], + a=trans_center, + b_list=anchors.positions, + c=anchors[1].position, dimensions=data_container.dimensions[:3], ) - else: # last resid residue = data_container.residues[1] index_prev = residue.resid - 1 - prev_resid = data_container.select_atoms(f"resindex {index_prev - 1}") - prev_backbone = prev_resid.atoms.select_atoms("backbone") - if len(prev_backbone) == 0: - anchor = np.array(prev_resid.center_of_mass()) - else: - anchor = np.array(prev_backbone.center_of_mass()) - # atom_set = data_container.atoms.select_atoms( - # f"resindex {index_prev-1} and " - # f"bonded resindex {residue.resid-1}" - # ) - # anchor = atom_set[0].position - backbone = residue.atoms.select_atoms("backbone") - if len(backbone) == 0: - # not a protein - center = np.array(residue.atoms.center_of_mass()) - else: - center = np.array(backbone.center_of_mass()) + anchor = data_container.atoms.select_atoms( + f"resindex {index_prev - 1} and bonded resindex {residue.resid - 1}" + ) + backbone = self.get_backbone(data_container, 1) + print(f"The backbone: {backbone}") + trans_center = np.zeros(3) + for atom in backbone: + trans_center += atom.position + trans_center /= len(backbone) trans_axes = self.get_custom_axes( - a=center, - b_list=anchor, + a=trans_center, + b_list=anchor[0].position, c=np.zeros(3), dimensions=data_container.dimensions[:3], ) @@ -365,10 +303,11 @@ def get_UA_axes(self, data_container, index: int, res_position): if rot_axes is None or moment_of_inertia is None: raise ValueError("Unable to compute bonded axes for UA bead.") - 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}") + logger.debug("Translational Axes: %s", trans_axes) + logger.debug("Rotational Axes: %s", rot_axes) + logger.debug("Translational center: %s", trans_center) + logger.debug("Rotational center: %s", rot_center) + logger.debug("Moment of Inertia: %s", moment_of_inertia) return trans_axes, rot_axes, rot_center, moment_of_inertia @@ -827,3 +766,125 @@ def get_UA_masses(self, molecule) -> list[float]: ua_mass += float(h.mass) ua_masses.append(ua_mass) return ua_masses + + def get_backbone(self, data_container, index): + """ + For a given residue, return AtomGroup corresponding to + its backbone. + This looks for heavy atoms between edge atoms of a residue. + Note: For a protein, this gives only the NCC atoms. + Meanwhile, the MDAnalysis "backbone" keyword includes + the O. + Args: + data_container (MDAnalysis.Universe or AtomGroup): + Molecule and trajectory data. + index: index of the residue of interest in + the data_container + + Returns: + backbone: Array containing + the backbone atoms. + + """ + # identify atoms with bind to neighbour residues + residue = data_container.residues[index] + print(f"Our residue of interest: {residue}") + print(f"Index of the residue: {residue.resid}") + edge_atom_set = data_container.atoms.select_atoms( + f" resindex {residue.resid - 1} and " + f"(bonded resindex {residue.resid - 2} or " + f"resindex {residue.resid})" + ) + print(f"The edge atoms are: {edge_atom_set}") + heavy_atoms = residue.atoms.select_atoms("mass 2 to 999") + if len(edge_atom_set) == 1: + # terminal residue + if index == 0: + # first residue + # assume first backbone atom will be first + backbone = self.get_chain(residue, residue.atoms[0], edge_atom_set[0]) + # add terminal atom to edge atom set + else: + # last residue + index = len(heavy_atoms) - 1 + last = None + while index > 0 and last is None: + # find a terminal atom + # look for last atom with only one bond to another heavy atom + heavy_atom = heavy_atoms[index] + bonded_atoms = residue.atoms.select_atoms( + f"(mass 2 to 999) and bonded index {heavy_atom.index}" + ) + if len(bonded_atoms) == 1: + last = heavy_atom + else: + index -= 1 + backbone = self.get_chain(residue, edge_atom_set[0], last) + else: + # will identify 2 edge atoms from linear neighbours + # disulfide bonds will be accounted for in the future + # not terminal residue + backbone = self.get_chain(residue, edge_atom_set[0], edge_atom_set[1]) + + return backbone + + def get_chain(self, residue, first, last): + """ + For a given MDAnalysis AtomGroup and two given heavy atoms + within that AtomGroup, return the + shortest path between the two atoms. + Args: + residue: MDAnalysis AtomGroup representing + the residue/monomer of interest. + first: First heavy atom in the chain + last: Last heavy atom in the chain + + Returns: + chain: array containing + the chain heavy atoms. + """ + chain = [] + # at the beggining we've only visited the first atom + visited_dict = {first: True} + # keep the previous atom to trace back the path + prev = {} + # queue of next heavy atoms to visit + next_to_visit = [first] + # all others heavy atoms in the residue, we have not yet visited + remaining_heavy_atoms = residue.atoms.select_atoms( + f"(mass 2 to 999) and not index {first.index}" + ) + for atom in remaining_heavy_atoms: + visited_dict[atom] = False + current = first + while not visited_dict[last]: + # we haven't found a path to the last residue + next_to_visit.pop(0) + # we're visiting the current atom => we remove it from the queue + bonded_atoms = residue.atoms.select_atoms( + f"(mass 2 to 999) and bonded index {current.index}" + ) + if last in bonded_atoms: + # we found a path to the last atom + visited_dict[last] = True + chain.append(last) + prev[last] = current + else: + for bonded_atom in bonded_atoms: + # look for unvisited bonded atoms to the current atom we're visiting + if not visited_dict[bonded_atom]: + # we're going to want to visit the atoms + next_to_visit.append(bonded_atom) + prev[bonded_atom] = current + # we visit the next atom in the queue + current = next_to_visit[0] + visited_dict[current] = True + + # we track the previous atom back to the first atom now + current = last + while chain[-1] != first: + # we haven't yet returned to the first atom + current = prev[current] + chain.append(current) + chain = np.flip(chain) + return chain From aa9a523b751d959ecbac2a4d46a192d1872aab24 Mon Sep 17 00:00:00 2001 From: ioanapapa Date: Tue, 31 Mar 2026 13:15:41 +0100 Subject: [PATCH 03/47] match backbone of last residue to that of previous residue --- CodeEntropy/levels/axes.py | 36 ++++++++++++++++++++++-------------- 1 file changed, 22 insertions(+), 14 deletions(-) diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py index 7db30c79..59bfdf58 100644 --- a/CodeEntropy/levels/axes.py +++ b/CodeEntropy/levels/axes.py @@ -788,14 +788,11 @@ def get_backbone(self, data_container, index): """ # identify atoms with bind to neighbour residues residue = data_container.residues[index] - print(f"Our residue of interest: {residue}") - print(f"Index of the residue: {residue.resid}") edge_atom_set = data_container.atoms.select_atoms( f" resindex {residue.resid - 1} and " f"(bonded resindex {residue.resid - 2} or " f"resindex {residue.resid})" ) - print(f"The edge atoms are: {edge_atom_set}") heavy_atoms = residue.atoms.select_atoms("mass 2 to 999") if len(edge_atom_set) == 1: # terminal residue @@ -808,17 +805,27 @@ def get_backbone(self, data_container, index): # last residue index = len(heavy_atoms) - 1 last = None - while index > 0 and last is None: - # find a terminal atom - # look for last atom with only one bond to another heavy atom - heavy_atom = heavy_atoms[index] - bonded_atoms = residue.atoms.select_atoms( - f"(mass 2 to 999) and bonded index {heavy_atom.index}" - ) - if len(bonded_atoms) == 1: - last = heavy_atom - else: - index -= 1 + # look for last heavy atom + # same type as terminal atom + # from previous residue + prev_terminal_atom = data_container.atoms.select_atoms( + f" resindex {residue.resid - 2} and " + f"bonded resindex {residue.resid - 1}" + ) + last_name = prev_terminal_atom.name + print("Name of last residue in chain") + last = residue.atoms.select_atoms(f"name {last_name}") + # while index > 0 and last is None: + # find the last atom of + # same + # heavy_atom = heavy_atoms[index] + # bonded_atoms = residue.atoms.select_atoms( + # f"(mass 2 to 999) and bonded index {heavy_atom.index}" + # ) + # if len(bonded_atoms) == 1: + # last = heavy_atom + # else: + # index -= 1 backbone = self.get_chain(residue, edge_atom_set[0], last) else: # will identify 2 edge atoms from linear neighbours @@ -887,4 +894,5 @@ def get_chain(self, residue, first, last): current = prev[current] chain.append(current) chain = np.flip(chain) + print(f"The chain is: {chain}") return chain From 138e770a6d4b980179de2751c22150171928fcaa Mon Sep 17 00:00:00 2001 From: ioanapapa Date: Tue, 31 Mar 2026 14:30:53 +0100 Subject: [PATCH 04/47] fix getting name of second to last terminal atom --- CodeEntropy/levels/axes.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py index 59bfdf58..e171be48 100644 --- a/CodeEntropy/levels/axes.py +++ b/CodeEntropy/levels/axes.py @@ -812,8 +812,9 @@ def get_backbone(self, data_container, index): f" resindex {residue.resid - 2} and " f"bonded resindex {residue.resid - 1}" ) - last_name = prev_terminal_atom.name - print("Name of last residue in chain") + print(f"Terminal atom of last resid: {prev_terminal_atom}") + last_name = prev_terminal_atom.names[0] + print(f"Name of last residue in chain: {last_name}") last = residue.atoms.select_atoms(f"name {last_name}") # while index > 0 and last is None: # find the last atom of From 9e1a3b0bff5c6c0802264e62866550a588557a8b Mon Sep 17 00:00:00 2001 From: ioanapapa Date: Tue, 31 Mar 2026 15:00:38 +0100 Subject: [PATCH 05/47] fix getting name of second to last terminal atom --- CodeEntropy/levels/axes.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py index e171be48..ccd48946 100644 --- a/CodeEntropy/levels/axes.py +++ b/CodeEntropy/levels/axes.py @@ -812,10 +812,10 @@ def get_backbone(self, data_container, index): f" resindex {residue.resid - 2} and " f"bonded resindex {residue.resid - 1}" ) - print(f"Terminal atom of last resid: {prev_terminal_atom}") last_name = prev_terminal_atom.names[0] print(f"Name of last residue in chain: {last_name}") last = residue.atoms.select_atoms(f"name {last_name}") + print(f"The last atom of the residue is: {last}") # while index > 0 and last is None: # find the last atom of # same From 3fba88e71769da2a97a4ac4a6ebe2246cd80196f Mon Sep 17 00:00:00 2001 From: ioanapapa Date: Tue, 31 Mar 2026 15:52:48 +0100 Subject: [PATCH 06/47] fix getting backbone for last residue --- CodeEntropy/levels/axes.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py index ccd48946..e50aa0e5 100644 --- a/CodeEntropy/levels/axes.py +++ b/CodeEntropy/levels/axes.py @@ -827,7 +827,7 @@ def get_backbone(self, data_container, index): # last = heavy_atom # else: # index -= 1 - backbone = self.get_chain(residue, edge_atom_set[0], last) + backbone = self.get_chain(residue, edge_atom_set[0], last[0]) else: # will identify 2 edge atoms from linear neighbours # disulfide bonds will be accounted for in the future From e3f4eafc72affaa23c14989142ae4f8554bef402 Mon Sep 17 00:00:00 2001 From: Ioana Papa Date: Wed, 1 Apr 2026 17:25:58 +0100 Subject: [PATCH 07/47] backbone is now returned as MDAnalysis atom group; backbone COM is used as center rather than average position of atoms --- CodeEntropy/levels/axes.py | 102 ++++++++++++++++++++----------------- 1 file changed, 55 insertions(+), 47 deletions(-) diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py index e50aa0e5..b52c0143 100644 --- a/CodeEntropy/levels/axes.py +++ b/CodeEntropy/levels/axes.py @@ -130,7 +130,6 @@ def get_residue_axes(self, data_container, index: int, residue=None): trans_axes = rot_axes # per original convention center = np.array(residue.center_of_mass()) else: - print("-----RESIDUE LEVEL-----") # If bonded to other residues, use local axes. make_whole(data_container.atoms) trans_axes = data_container.atoms.principal_axes() @@ -138,11 +137,11 @@ def get_residue_axes(self, data_container, index: int, residue=None): # get edge atoms of the residue # for terminal residues, this will include the C/N terminus # bond-derived rotation axes - print(f"The backbone: {backbone}") - center = np.zeros(3) - for atom in backbone: - center += atom.position - center /= len(backbone) + # center = np.zeros(3) + # for atom in backbone: + # center += atom.position + # center /= len(backbone) + center = np.array(backbone.center_of_mass()) if len(anchors) == 1: # only one neighbour rot_axes = self.get_custom_axes( @@ -218,7 +217,6 @@ def get_UA_axes(self, data_container, index: int, res_position): trans_center = data_container.atoms.center_of_mass(unwrap=True) trans_axes = data_container.atoms.principal_axes else: - print("-----UA LEVEL-----") # residue of interest has at least one neighbour if res_position == -1: residue = data_container.residues[0] @@ -229,11 +227,11 @@ def get_UA_axes(self, data_container, index: int, res_position): # the .resid attribute gives 1-indexing # substract 1 to match indexing later backbone = self.get_backbone(data_container, 0) - print(f"The backbone: {backbone}") - trans_center = np.zeros(3) - for atom in backbone: - trans_center += atom.position - trans_center /= len(backbone) + # trans_center = np.zeros(3) + # for atom in backbone: + # trans_center += atom.position + # trans_center /= len(backbone) + trans_center = np.array(backbone.center_of_mass()) trans_axes = self.get_custom_axes( a=trans_center, b_list=anchor[0].position, @@ -252,11 +250,11 @@ def get_UA_axes(self, data_container, index: int, res_position): f"bonded resindex {residue.resid - 1}" ) backbone = self.get_backbone(data_container, 1) - print(f"The backbone: {backbone}") - trans_center = np.zeros(3) - for atom in backbone: - trans_center += atom.position - trans_center /= len(backbone) + # trans_center = np.zeros(3) + # for atom in backbone: + # trans_center += atom.position + # trans_center /= len(backbone) + trans_center = np.array(backbone.center_of_mass()) trans_axes = self.get_custom_axes( a=trans_center, b_list=anchors.positions, @@ -271,11 +269,12 @@ def get_UA_axes(self, data_container, index: int, res_position): f"resindex {index_prev - 1} and bonded resindex {residue.resid - 1}" ) backbone = self.get_backbone(data_container, 1) - print(f"The backbone: {backbone}") - trans_center = np.zeros(3) - for atom in backbone: - trans_center += atom.position - trans_center /= len(backbone) + # trans_center = np.zeros(3) + # for atom in backbone: + # trans_center += atom.position + + # trans_center /= len(backbone) + trans_center = np.array(backbone.center_of_mass()) trans_axes = self.get_custom_axes( a=trans_center, b_list=anchor[0].position, @@ -782,7 +781,7 @@ def get_backbone(self, data_container, index): the data_container Returns: - backbone: Array containing + backbone: MDAnalysis AtomGroup containing the backbone atoms. """ @@ -806,28 +805,25 @@ def get_backbone(self, data_container, index): index = len(heavy_atoms) - 1 last = None # look for last heavy atom - # same type as terminal atom - # from previous residue - prev_terminal_atom = data_container.atoms.select_atoms( - f" resindex {residue.resid - 2} and " - f"bonded resindex {residue.resid - 1}" - ) - last_name = prev_terminal_atom.names[0] - print(f"Name of last residue in chain: {last_name}") - last = residue.atoms.select_atoms(f"name {last_name}") - print(f"The last atom of the residue is: {last}") - # while index > 0 and last is None: - # find the last atom of - # same - # heavy_atom = heavy_atoms[index] - # bonded_atoms = residue.atoms.select_atoms( - # f"(mass 2 to 999) and bonded index {heavy_atom.index}" - # ) - # if len(bonded_atoms) == 1: - # last = heavy_atom - # else: - # index -= 1 - backbone = self.get_chain(residue, edge_atom_set[0], last[0]) + # with only one bound to another + # prev_terminal_atom = data_container.atoms.select_atoms( + # f" resindex {residue.resid - 2} and " + # f"bonded resindex {residue.resid - 1}" + # ) + # last_name = prev_terminal_atom.names[0] + # print(f"Name of last residue in chain: {last_name}") + # last = residue.atoms.select_atoms(f"name {last_name}") + # print(f"The last atom of the residue is: {last}") + while index > 0 and last is None: + heavy_atom = heavy_atoms[index] + bonded_atoms = residue.atoms.select_atoms( + f"(mass 2 to 999) and bonded index {heavy_atom.index}" + ) + if len(bonded_atoms) == 1: + last = heavy_atom + else: + index -= 1 + backbone = self.get_chain(residue, edge_atom_set[0], last) else: # will identify 2 edge atoms from linear neighbours # disulfide bonds will be accounted for in the future @@ -848,10 +844,11 @@ def get_chain(self, residue, first, last): last: Last heavy atom in the chain Returns: - chain: array containing + chain: MDAnalysis AtomGroup containing the chain heavy atoms. """ chain = [] + chain_indices = [] # at the beggining we've only visited the first atom visited_dict = {first: True} # keep the previous atom to trace back the path @@ -890,10 +887,21 @@ def get_chain(self, residue, first, last): # we track the previous atom back to the first atom now current = last + chain = [last] + # subtract index of first atom in resid + # most likely will coincide with first + # but this will work even if it doesn't + # accout for in-residue index + chain_indices = [last.index - residue.atoms.indices[0]] + # start from last atom in chain while chain[-1] != first: # we haven't yet returned to the first atom current = prev[current] chain.append(current) - chain = np.flip(chain) + chain_indices.append(current.index - residue.atoms.indices[0]) + chain_indices = np.flip(chain_indices) + # accout for in-residue index + chain_AtomGroup = residue.atoms[chain_indices] + chain = chain_AtomGroup.atoms.select_atoms("all") print(f"The chain is: {chain}") return chain From a366d9a5a16823ead84f9fa3b9ea0912ec69609a Mon Sep 17 00:00:00 2001 From: Ioana Papa Date: Fri, 17 Apr 2026 12:23:18 +0100 Subject: [PATCH 08/47] draft new function to compute residue axes independent of residue neighbours --- CodeEntropy/levels/axes.py | 308 +++++++++++++++++-------------------- 1 file changed, 139 insertions(+), 169 deletions(-) diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py index b52c0143..b0d0aba9 100644 --- a/CodeEntropy/levels/axes.py +++ b/CodeEntropy/levels/axes.py @@ -107,17 +107,22 @@ def get_residue_axes(self, data_container, index: int, residue=None): # residue of interest if len(residue) == 0: raise ValueError(f"Empty residue selection for resindex={index}") - anchors = data_container.select_atoms( - f"(resindex {index_prev} or " - f"resindex {index_next}) and " - f"bonded resindex {index}" + # anchors = data_container.select_atoms( + # f"(resindex {index_prev} or " + # f"resindex {index_next}) and " + # f"bonded resindex {index}" + # ) + edge_atom_set = data_container.atoms.select_atoms( + f" resindex {index} and " + f"(bonded resindex {index_prev} or " + f"resindex {index_next})" ) uas = residue.select_atoms("mass 2 to 999") ua_masses = self.get_UA_masses(residue) - if len(anchors) == 0: - # No bonds to other residues. + if len(edge_atom_set) == 0: + # No UAS are bonded 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 H. moi_tensor = self.get_moment_of_inertia_tensor( @@ -133,32 +138,41 @@ def get_residue_axes(self, data_container, index: int, residue=None): # If bonded to other residues, use local axes. make_whole(data_container.atoms) trans_axes = data_container.atoms.principal_axes() - backbone = self.get_backbone(data_container, index) + residue = data_container.residues[index] + if len(edge_atom_set) == 1: + if index == 0: + # first residue + # use first heavy atom + edges = [residue.atoms[0], edge_atom_set[0]] + backbone = self.get_chain( + residue, residue.atoms[0], edge_atom_set[0] + ) + else: + # last residue + last_index = len(uas) - 1 + last = None + # look for last heavy atom + # with only one bond to another + while last_index > 0 and last is None: + heavy_atom = uas[last_index] + bonded_atoms = residue.atoms.select_atoms( + f"(mass 2 to 999) and bonded index {heavy_atom.index}" + ) + if len(bonded_atoms) == 1: + last = heavy_atom + else: + last_index -= 1 + edges = [edge_atom_set[0], last] + backbone = self.get_chain(residue, edge_atom_set[0], last) + else: + # residue has two bonds to other residues + edges = [edge_atom_set[0], edge_atom_set[1]] + backbone = self.get_chain(residue, edge_atom_set[0], edge_atom_set[1]) # get edge atoms of the residue # for terminal residues, this will include the C/N terminus - # bond-derived rotation axes - # center = np.zeros(3) - # for atom in backbone: - # center += atom.position - # center /= len(backbone) center = np.array(backbone.center_of_mass()) - if len(anchors) == 1: - # only one neighbour - rot_axes = self.get_custom_axes( - a=center, - b_list=anchors[0].position, - c=np.zeros(3), - dimensions=data_container.dimensions[:3], - ) - else: - # two neighbours - rot_axes = self.get_custom_axes( - a=center, - b_list=anchors.positions, - c=anchors[1].position, - dimensions=data_container.dimensions[:3], - ) - # analogous to the UA case where a heavy atom is bound to >=2 heavy atoms + print(f"Edges of the residue: {edges}") + rot_axes = self.get_residue_custom_axes(edges, center) moment_of_inertia = self.get_custom_residue_moment_of_inertia( center_of_mass=center, @@ -209,84 +223,73 @@ def get_UA_axes(self, data_container, index: int, res_position): """ index = int(index) # bead index + heavy_atoms = data_container.atoms.select_atoms("mass 2 to 999") # use the same customPI trans axes as the residue level - if len(data_container.residues) == 1: - # only the one residue => use principal axes - residue = data_container - trans_center = data_container.atoms.center_of_mass(unwrap=True) - trans_axes = data_container.atoms.principal_axes - else: - # residue of interest has at least one neighbour - if res_position == -1: - residue = data_container.residues[0] - index_next = residue.resid + 1 - anchor = data_container.atoms.select_atoms( - f"resindex {index_next - 1} and bonded resindex {residue.resid - 1}" - ) - # the .resid attribute gives 1-indexing - # substract 1 to match indexing later - backbone = self.get_backbone(data_container, 0) - # trans_center = np.zeros(3) - # for atom in backbone: - # trans_center += atom.position - # trans_center /= len(backbone) - trans_center = np.array(backbone.center_of_mass()) - trans_axes = self.get_custom_axes( - a=trans_center, - b_list=anchor[0].position, - c=np.zeros(3), - dimensions=data_container.dimensions[:3], - ) - - elif res_position == 0: - # between 2 residues - residue = data_container.residues[1] - index_prev = residue.resid - 1 - index_next = residue.resid + 1 - anchors = data_container.atoms.select_atoms( - f"(resindex {index_prev - 1} or " - f"resindex {index_next - 1}) and " - f"bonded resindex {residue.resid - 1}" - ) - backbone = self.get_backbone(data_container, 1) - # trans_center = np.zeros(3) - # for atom in backbone: - # trans_center += atom.position - # trans_center /= len(backbone) - trans_center = np.array(backbone.center_of_mass()) - trans_axes = self.get_custom_axes( - a=trans_center, - b_list=anchors.positions, - c=anchors[1].position, - dimensions=data_container.dimensions[:3], - ) + if len(heavy_atoms) > 1: + if len(data_container.residues) == 1: + # only the one residue => use principal axes + residue = data_container + trans_center = data_container.atoms.center_of_mass(unwrap=True) + trans_axes = data_container.atoms.principal_axes else: - # last resid - residue = data_container.residues[1] - index_prev = residue.resid - 1 - anchor = data_container.atoms.select_atoms( - f"resindex {index_prev - 1} and bonded resindex {residue.resid - 1}" - ) - backbone = self.get_backbone(data_container, 1) - # trans_center = np.zeros(3) - # for atom in backbone: - # trans_center += atom.position - - # trans_center /= len(backbone) + # residue of interest has at least one neighbour + if res_position == -1: + residue = data_container.residues[0] + index_next = residue.resid + 1 + # the .resid attribute gives 1-indexing + # substract 1 to match indexing later + second_edge = data_container.atoms.select_atoms( + f"resindex {residue.resid - 1} and " + f"bonded resindex {index_next - 1}" + ) + edges = [residue.atoms[0], second_edge] + backbone = self.get_chain(residue, residue.atoms[0], second_edge) + elif res_position == 0: + # between 2 residues + residue = data_container.residues[1] + index_prev = residue.resid - 1 + index_next = residue.resid + 1 + edge_set = data_container.atoms.select_atoms( + f"(resindex {residue.resid - 1} and " + f"(bonded resindex {index_next - 1} or " + f"bonded resindex {residue.resid - 1})" + ) + edges = [edge_set[0], edge_set[1]] + backbone = self.get_chain(residue, edge_set[0], edge_set[1]) + else: + # last resid + residue = data_container.residues[1] + index_prev = residue.resid - 1 + first_edge = data_container.atoms.select_atoms( + f"resindex {residue.resid - 1} and " + f"bonded resindex {index_prev - 1}" + ) + last_index = len(heavy_atoms) - 1 + last = None + # look for last heavy atom + # with only one bond to another + while last_index > 0 and last is None: + heavy_atom = heavy_atoms[last_index] + bonded_atoms = residue.atoms.select_atoms( + f"(mass 2 to 999) and bonded index {heavy_atom.index}" + ) + if len(bonded_atoms) == 1: + last = heavy_atom + else: + last_index -= 1 + edges = [first_edge, last] + backbone = self.get_chain(residue, first_edge, last) + print(f"The edges of the residue: {edges}") trans_center = np.array(backbone.center_of_mass()) - trans_axes = self.get_custom_axes( - a=trans_center, - b_list=anchor[0].position, - c=np.zeros(3), - dimensions=data_container.dimensions[:3], - ) - - heavy_atoms = residue.atoms.select_atoms("mass 2 to 999") - + trans_axes = self.get_residue_custom_axes(edges, trans_center) + else: + make_whole(data_container.atoms) + trans_axes = data_container.atoms.principal_axes() + residue_heavy_atoms = residue.atoms.select_atoms("mass 2 to 999") # look for heavy atoms in residue of interest heavy_atom_indices = [] - for atom in heavy_atoms: + for atom in residue_heavy_atoms: heavy_atom_indices.append(atom.index) # we find the nth heavy atom # where n is the bead index @@ -310,8 +313,42 @@ def get_UA_axes(self, data_container, index: int, res_position): return trans_axes, rot_axes, rot_center, moment_of_inertia + def get_residue_custom_axes(self, edges, center): + """ + Compute rotation axes at the residue level, given + two edge atoms of the residue (heavy atoms bonded + to neighbouring residues), and the rotation centre. + + E1----O + \ + \ + E2 + Args: + edges: (2,3) positions of two edge atoms + center: (3,) coordinates of the rotation centre + Returns: + rot_axes: (3,3) rotation axes of residue + """ + # x axis is O-E1 + x_axis = edges[0].position - center + # y axis is perpendicular to x + # in the same plane as E2 + # look for projection of E2-E1 on O-E1 + E1E2_vector = edges[1].position - edges[1].position + y_axis = np.dot(x_axis, E1E2_vector) / (np.linalg.norm(x_axis) ** 2) + y_axis = y_axis * x_axis + y_axis = edges[1].position - y_axis + print(f"We found the perpendicular: {np.dot(x_axis, y_axis)}") + z_axis = np.cross(x_axis, y_axis) + x_axis /= np.linalg.norm(x_axis) + y_axis /= np.linalg.norm(y_axis) + z_axis /= np.linalg.norm(z_axis) + rot_axes = np.array([x_axis, y_axis, z_axis]) + + return rot_axes + def get_bonded_axes(self, system, atom, dimensions: np.ndarray): - r"""Compute UA rotational axes from bonded topology around a heavy atom. + """Compute UA rotational axes from bonded topology around a heavy atom. For a given heavy atom, use its bonded atoms to get the axes for rotating forces around. Few cases for choosing united atom axes, which are dependent @@ -766,72 +803,6 @@ def get_UA_masses(self, molecule) -> list[float]: ua_masses.append(ua_mass) return ua_masses - def get_backbone(self, data_container, index): - """ - For a given residue, return AtomGroup corresponding to - its backbone. - This looks for heavy atoms between edge atoms of a residue. - Note: For a protein, this gives only the NCC atoms. - Meanwhile, the MDAnalysis "backbone" keyword includes - the O. - Args: - data_container (MDAnalysis.Universe or AtomGroup): - Molecule and trajectory data. - index: index of the residue of interest in - the data_container - - Returns: - backbone: MDAnalysis AtomGroup containing - the backbone atoms. - - """ - # identify atoms with bind to neighbour residues - residue = data_container.residues[index] - edge_atom_set = data_container.atoms.select_atoms( - f" resindex {residue.resid - 1} and " - f"(bonded resindex {residue.resid - 2} or " - f"resindex {residue.resid})" - ) - heavy_atoms = residue.atoms.select_atoms("mass 2 to 999") - if len(edge_atom_set) == 1: - # terminal residue - if index == 0: - # first residue - # assume first backbone atom will be first - backbone = self.get_chain(residue, residue.atoms[0], edge_atom_set[0]) - # add terminal atom to edge atom set - else: - # last residue - index = len(heavy_atoms) - 1 - last = None - # look for last heavy atom - # with only one bound to another - # prev_terminal_atom = data_container.atoms.select_atoms( - # f" resindex {residue.resid - 2} and " - # f"bonded resindex {residue.resid - 1}" - # ) - # last_name = prev_terminal_atom.names[0] - # print(f"Name of last residue in chain: {last_name}") - # last = residue.atoms.select_atoms(f"name {last_name}") - # print(f"The last atom of the residue is: {last}") - while index > 0 and last is None: - heavy_atom = heavy_atoms[index] - bonded_atoms = residue.atoms.select_atoms( - f"(mass 2 to 999) and bonded index {heavy_atom.index}" - ) - if len(bonded_atoms) == 1: - last = heavy_atom - else: - index -= 1 - backbone = self.get_chain(residue, edge_atom_set[0], last) - else: - # will identify 2 edge atoms from linear neighbours - # disulfide bonds will be accounted for in the future - # not terminal residue - backbone = self.get_chain(residue, edge_atom_set[0], edge_atom_set[1]) - - return backbone - def get_chain(self, residue, first, last): """ For a given MDAnalysis AtomGroup and two given heavy atoms @@ -903,5 +874,4 @@ def get_chain(self, residue, first, last): # accout for in-residue index chain_AtomGroup = residue.atoms[chain_indices] chain = chain_AtomGroup.atoms.select_atoms("all") - print(f"The chain is: {chain}") return chain From bd9d6f65115c6b1886698b134c4c5b6a03d6a959 Mon Sep 17 00:00:00 2001 From: Ioana Papa Date: Fri, 17 Apr 2026 17:08:31 +0100 Subject: [PATCH 09/47] residue rotation axes are determined from center of rotation + two edge atoms --- CodeEntropy/levels/axes.py | 45 +++++++++++++++++++++++--------------- 1 file changed, 27 insertions(+), 18 deletions(-) diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py index b0d0aba9..bb422c9b 100644 --- a/CodeEntropy/levels/axes.py +++ b/CodeEntropy/levels/axes.py @@ -243,17 +243,19 @@ def get_UA_axes(self, data_container, index: int, res_position): f"resindex {residue.resid - 1} and " f"bonded resindex {index_next - 1}" ) - edges = [residue.atoms[0], second_edge] - backbone = self.get_chain(residue, residue.atoms[0], second_edge) + edges = [residue.atoms[0], second_edge.atoms[0]] + backbone = self.get_chain( + residue, residue.atoms[0], second_edge.atoms[0] + ) elif res_position == 0: # between 2 residues residue = data_container.residues[1] index_prev = residue.resid - 1 index_next = residue.resid + 1 edge_set = data_container.atoms.select_atoms( - f"(resindex {residue.resid - 1} and " + f"resindex {residue.resid - 1} and " f"(bonded resindex {index_next - 1} or " - f"bonded resindex {residue.resid - 1})" + f"resindex {index_prev - 1})" ) edges = [edge_set[0], edge_set[1]] backbone = self.get_chain(residue, edge_set[0], edge_set[1]) @@ -278,9 +280,9 @@ def get_UA_axes(self, data_container, index: int, res_position): last = heavy_atom else: last_index -= 1 - edges = [first_edge, last] - backbone = self.get_chain(residue, first_edge, last) - print(f"The edges of the residue: {edges}") + edges = [first_edge.atoms[0], last] + backbone = self.get_chain(residue, first_edge.atoms[0], last) + trans_center = np.array(backbone.center_of_mass()) trans_axes = self.get_residue_custom_axes(edges, trans_center) else: @@ -319,10 +321,10 @@ def get_residue_custom_axes(self, edges, center): two edge atoms of the residue (heavy atoms bonded to neighbouring residues), and the rotation centre. - E1----O - \ - \ - E2 + E1---O + \ | + \ | + E2 Args: edges: (2,3) positions of two edge atoms center: (3,) coordinates of the rotation centre @@ -330,15 +332,21 @@ def get_residue_custom_axes(self, edges, center): rot_axes: (3,3) rotation axes of residue """ # x axis is O-E1 - x_axis = edges[0].position - center + E1O_vector = center - edges[0].position + x_axis = -E1O_vector # y axis is perpendicular to x # in the same plane as E2 - # look for projection of E2-E1 on O-E1 - E1E2_vector = edges[1].position - edges[1].position - y_axis = np.dot(x_axis, E1E2_vector) / (np.linalg.norm(x_axis) ** 2) - y_axis = y_axis * x_axis - y_axis = edges[1].position - y_axis - print(f"We found the perpendicular: {np.dot(x_axis, y_axis)}") + # look for projection of E1-E2 on E1-O + E1E2_vector = edges[1].position - edges[0].position + projection = ( + np.dot(E1O_vector, E1E2_vector) / (np.linalg.norm(E1O_vector) ** 2) + ) * E1O_vector + # get the perpendicular onto E1-O + perpendicular = E1E2_vector - projection + # get the perpendicular through O + diagonal = -(projection - E1O_vector) - perpendicular + # get the projection from this diagonal + y_axis = (projection - E1O_vector) + diagonal z_axis = np.cross(x_axis, y_axis) x_axis /= np.linalg.norm(x_axis) y_axis /= np.linalg.norm(y_axis) @@ -818,6 +826,7 @@ def get_chain(self, residue, first, last): chain: MDAnalysis AtomGroup containing the chain heavy atoms. """ + print(f"Last: {last} ") chain = [] chain_indices = [] # at the beggining we've only visited the first atom From 11b726c336c2ced8497f039e6149f2c17792bcd4 Mon Sep 17 00:00:00 2001 From: Ioana Papa Date: Wed, 22 Apr 2026 10:36:22 +0100 Subject: [PATCH 10/47] tidied up --- CodeEntropy/levels/axes.py | 39 +++++++++++++++++++++++++++++++++----- 1 file changed, 34 insertions(+), 5 deletions(-) diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py index bb422c9b..e6bc77e1 100644 --- a/CodeEntropy/levels/axes.py +++ b/CodeEntropy/levels/axes.py @@ -171,7 +171,6 @@ def get_residue_axes(self, data_container, index: int, residue=None): # get edge atoms of the residue # for terminal residues, this will include the C/N terminus center = np.array(backbone.center_of_mass()) - print(f"Edges of the residue: {edges}") rot_axes = self.get_residue_custom_axes(edges, center) moment_of_inertia = self.get_custom_residue_moment_of_inertia( @@ -247,6 +246,7 @@ def get_UA_axes(self, data_container, index: int, res_position): backbone = self.get_chain( residue, residue.atoms[0], second_edge.atoms[0] ) + elif res_position == 0: # between 2 residues residue = data_container.residues[1] @@ -259,6 +259,7 @@ def get_UA_axes(self, data_container, index: int, res_position): ) edges = [edge_set[0], edge_set[1]] backbone = self.get_chain(residue, edge_set[0], edge_set[1]) + else: # last resid residue = data_container.residues[1] @@ -285,7 +286,9 @@ def get_UA_axes(self, data_container, index: int, res_position): trans_center = np.array(backbone.center_of_mass()) trans_axes = self.get_residue_custom_axes(edges, trans_center) + else: + # only one heavy atom or hydrogen molecule make_whole(data_container.atoms) trans_axes = data_container.atoms.principal_axes() residue_heavy_atoms = residue.atoms.select_atoms("mass 2 to 999") @@ -343,9 +346,12 @@ def get_residue_custom_axes(self, edges, center): ) * E1O_vector # get the perpendicular onto E1-O perpendicular = E1E2_vector - projection - # get the perpendicular through O - diagonal = -(projection - E1O_vector) - perpendicular - # get the projection from this diagonal + # get the perpendicular through O (Q-O) + # first get P-Q diagonal through paralellogram rule + # P- Q = P-E2 + P-O + diagonal = -(projection - E1O_vector) + perpendicular + # get the parallel of P-E2 through O + # OQ = OP + PQ y_axis = (projection - E1O_vector) + diagonal z_axis = np.cross(x_axis, y_axis) x_axis /= np.linalg.norm(x_axis) @@ -826,7 +832,6 @@ def get_chain(self, residue, first, last): chain: MDAnalysis AtomGroup containing the chain heavy atoms. """ - print(f"Last: {last} ") chain = [] chain_indices = [] # at the beggining we've only visited the first atom @@ -884,3 +889,27 @@ def get_chain(self, residue, first, last): chain_AtomGroup = residue.atoms[chain_indices] chain = chain_AtomGroup.atoms.select_atoms("all") return chain + + def get_NCC_axes(self, residue): + """ + Return axes based on the NCC atoms of a + residue in a protein. This is based on Argo's implementation + and it's here to compare protein results with previous published results. + Will most likely be deleted in the future when merging into main. + """ + N_coords = residue.atoms.select_atoms("name N").positions[0] + C_alpha_coords = residue.atoms.select_atoms("name CA").positions[0] + C_coords = residue.atoms.select_atoms("name C").positions[0] + # get projection of CC_alpha onto CN + CCa_vector = C_alpha_coords - C_coords + CN_vector = N_coords - C_coords + center = np.dot(CCa_vector, CN_vector) / (np.linalg.norm(CN_vector) ** 2) + center = center * CN_vector + C_coords + x_axis = N_coords - center + x_axis /= np.linalg.norm(x_axis) + y_axis = C_alpha_coords - center + y_axis /= np.linalg.norm(y_axis) + z_axis = np.cross(x_axis, y_axis) + z_axis /= np.linalg.norm(z_axis) + rot_axes = np.array([x_axis, y_axis, z_axis]) + return center, rot_axes From d98963654a39680111220cea6fc0960a0edc407e Mon Sep 17 00:00:00 2001 From: Ioana Papa Date: Wed, 22 Apr 2026 10:40:57 +0100 Subject: [PATCH 11/47] tidied up --- CodeEntropy/levels/axes.py | 24 ------------------------ 1 file changed, 24 deletions(-) diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py index e6bc77e1..8c48b406 100644 --- a/CodeEntropy/levels/axes.py +++ b/CodeEntropy/levels/axes.py @@ -889,27 +889,3 @@ def get_chain(self, residue, first, last): chain_AtomGroup = residue.atoms[chain_indices] chain = chain_AtomGroup.atoms.select_atoms("all") return chain - - def get_NCC_axes(self, residue): - """ - Return axes based on the NCC atoms of a - residue in a protein. This is based on Argo's implementation - and it's here to compare protein results with previous published results. - Will most likely be deleted in the future when merging into main. - """ - N_coords = residue.atoms.select_atoms("name N").positions[0] - C_alpha_coords = residue.atoms.select_atoms("name CA").positions[0] - C_coords = residue.atoms.select_atoms("name C").positions[0] - # get projection of CC_alpha onto CN - CCa_vector = C_alpha_coords - C_coords - CN_vector = N_coords - C_coords - center = np.dot(CCa_vector, CN_vector) / (np.linalg.norm(CN_vector) ** 2) - center = center * CN_vector + C_coords - x_axis = N_coords - center - x_axis /= np.linalg.norm(x_axis) - y_axis = C_alpha_coords - center - y_axis /= np.linalg.norm(y_axis) - z_axis = np.cross(x_axis, y_axis) - z_axis /= np.linalg.norm(z_axis) - rot_axes = np.array([x_axis, y_axis, z_axis]) - return center, rot_axes From 9d7948b395aa1cd37940b7504e33d3dcc9e8712c Mon Sep 17 00:00:00 2001 From: Ioana Papa Date: Wed, 22 Apr 2026 12:13:39 +0100 Subject: [PATCH 12/47] corrections to get_residue_custom_axes --- CodeEntropy/levels/axes.py | 24 ++++++++++++++++++------ 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py index 8c48b406..4846c39b 100644 --- a/CodeEntropy/levels/axes.py +++ b/CodeEntropy/levels/axes.py @@ -324,10 +324,10 @@ def get_residue_custom_axes(self, edges, center): two edge atoms of the residue (heavy atoms bonded to neighbouring residues), and the rotation centre. - E1---O - \ | - \ | - E2 + Q --- E2 + | | + | | + E1 ---- O --- P Args: edges: (2,3) positions of two edge atoms center: (3,) coordinates of the rotation centre @@ -339,13 +339,18 @@ def get_residue_custom_axes(self, edges, center): x_axis = -E1O_vector # y axis is perpendicular to x # in the same plane as E2 - # look for projection of E1-E2 on E1-O + # look for projection of E1-E2 on E1-O (E1-P) E1E2_vector = edges[1].position - edges[0].position projection = ( np.dot(E1O_vector, E1E2_vector) / (np.linalg.norm(E1O_vector) ** 2) ) * E1O_vector - # get the perpendicular onto E1-O + # get the perpendicular onto E1-O (P-E2) + # P-E2 = P-E1 + E1-E2 perpendicular = E1E2_vector - projection + print( + f"The perpendicular is perpendicular to the projection: " + f"{np.dot(projection, perpendicular)}" + ) # get the perpendicular through O (Q-O) # first get P-Q diagonal through paralellogram rule # P- Q = P-E2 + P-O @@ -353,6 +358,11 @@ def get_residue_custom_axes(self, edges, center): # get the parallel of P-E2 through O # OQ = OP + PQ y_axis = (projection - E1O_vector) + diagonal + print( + f"The y-axis and perpendicular are " + f"parallel: {np.cross(perpendicular, y_axis)}" + ) + print(f"The x and y axis are parallel: {np.cross(x_axis, y_axis)}") z_axis = np.cross(x_axis, y_axis) x_axis /= np.linalg.norm(x_axis) y_axis /= np.linalg.norm(y_axis) @@ -361,6 +371,8 @@ def get_residue_custom_axes(self, edges, center): return rot_axes + return rot_axes + def get_bonded_axes(self, system, atom, dimensions: np.ndarray): """Compute UA rotational axes from bonded topology around a heavy atom. From 64a5187bc4fc46b44653ace121ef0272a83e8224 Mon Sep 17 00:00:00 2001 From: Ioana Papa Date: Wed, 22 Apr 2026 15:19:02 +0100 Subject: [PATCH 13/47] removed debugging print statements --- CodeEntropy/levels/axes.py | 9 --------- 1 file changed, 9 deletions(-) diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py index 4846c39b..5383e3c3 100644 --- a/CodeEntropy/levels/axes.py +++ b/CodeEntropy/levels/axes.py @@ -347,10 +347,6 @@ def get_residue_custom_axes(self, edges, center): # get the perpendicular onto E1-O (P-E2) # P-E2 = P-E1 + E1-E2 perpendicular = E1E2_vector - projection - print( - f"The perpendicular is perpendicular to the projection: " - f"{np.dot(projection, perpendicular)}" - ) # get the perpendicular through O (Q-O) # first get P-Q diagonal through paralellogram rule # P- Q = P-E2 + P-O @@ -358,11 +354,6 @@ def get_residue_custom_axes(self, edges, center): # get the parallel of P-E2 through O # OQ = OP + PQ y_axis = (projection - E1O_vector) + diagonal - print( - f"The y-axis and perpendicular are " - f"parallel: {np.cross(perpendicular, y_axis)}" - ) - print(f"The x and y axis are parallel: {np.cross(x_axis, y_axis)}") z_axis = np.cross(x_axis, y_axis) x_axis /= np.linalg.norm(x_axis) y_axis /= np.linalg.norm(y_axis) From 95f937bc5a39c59d317e7a16b22d228433fbf4ac Mon Sep 17 00:00:00 2001 From: ioana Date: Fri, 24 Apr 2026 13:57:31 +0100 Subject: [PATCH 14/47] fix for molecules with only one residue --- CodeEntropy/levels/nodes/covariance.py | 36 ++++++++++++++------------ 1 file changed, 20 insertions(+), 16 deletions(-) diff --git a/CodeEntropy/levels/nodes/covariance.py b/CodeEntropy/levels/nodes/covariance.py index fbe9a821..16e16ffd 100644 --- a/CodeEntropy/levels/nodes/covariance.py +++ b/CodeEntropy/levels/nodes/covariance.py @@ -202,23 +202,27 @@ def _process_united_atom( """ for local_res_i, res in enumerate(mol.residues): - # build residue group here - if local_res_i == 0: - # first residue - res_position = -1 - res_next = mol.residues[1] - residue_group = res + res_next - elif local_res_i == len(mol.residues) - 1: - # last residue - res_position = 1 - res_prev = mol.residues[-2] - residue_group = res + res_prev + if len(mol.residues) > 1: + # there are multiple residues in the molecule + # build residue group here + if local_res_i == 0: + # first residue + res_position = -1 + res_next = mol.residues[1] + residue_group = res + res_next + elif local_res_i == len(mol.residues) - 1: + # last residue + res_position = 1 + res_prev = mol.residues[-2] + residue_group = res + res_prev + else: + res_position = 0 + res_prev = mol.residues[local_res_i - 1] + res_next = mol.residues[local_res_i + 1] + residue_group = res_prev + res + res_next else: - res_position = 0 - res_prev = mol.residues[local_res_i - 1] - res_next = mol.residues[local_res_i + 1] - residue_group = res_prev + res + res_next - + # only one residue + res_position = None bead_key = (mol_id, "united_atom", local_res_i) bead_idx_list = beads.get(bead_key, []) if not bead_idx_list: From 98a3987be6b15f0505a7ed0e93248e32de492996 Mon Sep 17 00:00:00 2001 From: ioana Date: Fri, 24 Apr 2026 14:00:49 +0100 Subject: [PATCH 15/47] deleted commented out old code --- CodeEntropy/levels/axes.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py index 5383e3c3..8e4956ac 100644 --- a/CodeEntropy/levels/axes.py +++ b/CodeEntropy/levels/axes.py @@ -107,11 +107,6 @@ def get_residue_axes(self, data_container, index: int, residue=None): # residue of interest if len(residue) == 0: raise ValueError(f"Empty residue selection for resindex={index}") - # anchors = data_container.select_atoms( - # f"(resindex {index_prev} or " - # f"resindex {index_next}) and " - # f"bonded resindex {index}" - # ) edge_atom_set = data_container.atoms.select_atoms( f" resindex {index} and " f"(bonded resindex {index_prev} or " From 2b2b3a9ae4dfa8d31705d9799e5ac1211ec61b81 Mon Sep 17 00:00:00 2001 From: ioana Date: Fri, 24 Apr 2026 14:16:40 +0100 Subject: [PATCH 16/47] changed UA axes tests to include new parameter relevant for new UA axes --- tests/unit/CodeEntropy/levels/test_axes.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/unit/CodeEntropy/levels/test_axes.py b/tests/unit/CodeEntropy/levels/test_axes.py index d6d0093f..0eba6e05 100644 --- a/tests/unit/CodeEntropy/levels/test_axes.py +++ b/tests/unit/CodeEntropy/levels/test_axes.py @@ -109,7 +109,7 @@ def _select_atoms(q): lambda moi: (np.eye(3), np.array([3.0, 2.0, 1.0])), ) - trans, rot, center, moi = ax.get_residue_axes(u, index=7) + trans, rot, center, moi = ax.get_residue_axes(u, index=7, residue=None) assert np.allclose(trans, np.eye(3)) assert np.allclose(rot, np.eye(3)) @@ -144,7 +144,7 @@ def _select_atoms(q): ax, "get_vanilla_axes", lambda mol: (np.eye(3) * 2, np.array([9.0, 8.0, 7.0])) ) - trans, rot, center, moi = ax.get_residue_axes(u, index=10) + trans, rot, center, moi = ax.get_residue_axes(u, index=10, residue=None) assert np.allclose(trans, np.eye(3)) assert np.allclose(rot, np.eye(3) * 2) @@ -184,7 +184,7 @@ def _sel(q): lambda system, atom, dimensions: (np.eye(3), np.array([1.0, 2.0, 3.0])), ) - trans, rot, center, moi = ax.get_UA_axes(u, index=0) + trans, rot, center, moi = ax.get_UA_axes(u, index=0, res_position=None) assert np.allclose(trans, np.eye(3)) assert np.allclose(rot, np.eye(3)) From aebbd48473b33024c8cbe5c6fc9189b11b18ad8e Mon Sep 17 00:00:00 2001 From: ioana Date: Fri, 24 Apr 2026 15:46:38 +0100 Subject: [PATCH 17/47] fix for building ua bead when customised_axes False and UA trans axes when customised_axes True --- CodeEntropy/levels/axes.py | 2 +- CodeEntropy/levels/nodes/covariance.py | 6 +++++- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py index 8e4956ac..a597a4d2 100644 --- a/CodeEntropy/levels/axes.py +++ b/CodeEntropy/levels/axes.py @@ -225,7 +225,7 @@ def get_UA_axes(self, data_container, index: int, res_position): # only the one residue => use principal axes residue = data_container trans_center = data_container.atoms.center_of_mass(unwrap=True) - trans_axes = data_container.atoms.principal_axes + trans_axes = data_container.atoms.principal_axes() else: # residue of interest has at least one neighbour if res_position == -1: diff --git a/CodeEntropy/levels/nodes/covariance.py b/CodeEntropy/levels/nodes/covariance.py index 16e16ffd..cbf2c3b1 100644 --- a/CodeEntropy/levels/nodes/covariance.py +++ b/CodeEntropy/levels/nodes/covariance.py @@ -223,6 +223,7 @@ def _process_united_atom( else: # only one residue res_position = None + residue_group = res bead_key = (mol_id, "united_atom", local_res_i) bead_idx_list = beads.get(bead_key, []) if not bead_idx_list: @@ -475,9 +476,12 @@ def _build_ua_vectors( if res_position == -1: # first residue in group residue = residue_group.residues[0] - else: + elif res_position == 0 or res_position == 1: # middle or last residue => second in group residue = residue_group.residues[1] + else: + # res_position is None bc there is only one residue + residue = residue_group trans_axes = residue.atoms.principal_axes() rot_axes, moi = axes_manager.get_vanilla_axes(bead) center = bead.center_of_mass(unwrap=True) From 3052c6699c52ecb6800f24b5720086407c5e7423 Mon Sep 17 00:00:00 2001 From: ioanapapa Date: Mon, 27 Apr 2026 13:36:12 +0100 Subject: [PATCH 18/47] removed duplicate return statement --- CodeEntropy/levels/axes.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py index a597a4d2..2c998156 100644 --- a/CodeEntropy/levels/axes.py +++ b/CodeEntropy/levels/axes.py @@ -357,8 +357,6 @@ def get_residue_custom_axes(self, edges, center): return rot_axes - return rot_axes - def get_bonded_axes(self, system, atom, dimensions: np.ndarray): """Compute UA rotational axes from bonded topology around a heavy atom. From 4fe8f26cc34aed859147aee40f9638cca2676c58 Mon Sep 17 00:00:00 2001 From: ioana Date: Tue, 28 Apr 2026 11:05:27 +0100 Subject: [PATCH 19/47] changed tests to include res_position --- tests/unit/CodeEntropy/levels/test_axes.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/tests/unit/CodeEntropy/levels/test_axes.py b/tests/unit/CodeEntropy/levels/test_axes.py index 0eba6e05..4fc0c563 100644 --- a/tests/unit/CodeEntropy/levels/test_axes.py +++ b/tests/unit/CodeEntropy/levels/test_axes.py @@ -109,7 +109,7 @@ def _select_atoms(q): lambda moi: (np.eye(3), np.array([3.0, 2.0, 1.0])), ) - trans, rot, center, moi = ax.get_residue_axes(u, index=7, residue=None) + trans, rot, center, moi = ax.get_residue_axes(u, index=7) assert np.allclose(trans, np.eye(3)) assert np.allclose(rot, np.eye(3)) @@ -144,7 +144,7 @@ def _select_atoms(q): ax, "get_vanilla_axes", lambda mol: (np.eye(3) * 2, np.array([9.0, 8.0, 7.0])) ) - trans, rot, center, moi = ax.get_residue_axes(u, index=10, residue=None) + trans, rot, center, moi = ax.get_residue_axes(u, index=10) assert np.allclose(trans, np.eye(3)) assert np.allclose(rot, np.eye(3) * 2) @@ -217,7 +217,7 @@ def _sel(q): monkeypatch.setattr(ax, "get_bonded_axes", lambda **kwargs: (None, None)) with pytest.raises(ValueError): - ax.get_UA_axes(u, index=0) + ax.get_UA_axes(u, index=0, res_position=None) def test_get_custom_axes_degenerate_axis1_raises(): @@ -531,7 +531,7 @@ def _select_atoms(q): ax, "get_vanilla_axes", lambda mol: (np.eye(3) * 2, np.array([9.0, 8.0, 7.0])) ) - trans, rot, center, moi = ax.get_residue_axes(u, index=10) + trans, rot, center, moi = ax.get_residue_axes(u, index=10, residue=residue) assert np.allclose(trans, np.eye(3) * 2) assert np.allclose(rot, np.eye(3) * 2) @@ -664,7 +664,9 @@ def _select_atoms(q): got_custom_axes = MagicMock(return_value=(np.eye(3), np.array([3.0, 2.0, 1.0]))) monkeypatch.setattr(ax, "get_custom_principal_axes", got_custom_axes) - trans_axes, rot_axes, center, moi = ax.get_UA_axes(data_container, index=0) + trans_axes, rot_axes, center, moi = ax.get_UA_axes( + data_container, index=0, res_position=None + ) assert trans_axes.shape == (3, 3) assert rot_axes.shape == (3, 3) From a939741bb3711be02e7ea56a595fbfa5ab169567 Mon Sep 17 00:00:00 2001 From: ioana Date: Tue, 28 Apr 2026 11:52:22 +0100 Subject: [PATCH 20/47] added comments to described new axis method --- CodeEntropy/levels/axes.py | 26 +++++++++++++++++++++----- 1 file changed, 21 insertions(+), 5 deletions(-) diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py index a597a4d2..930da819 100644 --- a/CodeEntropy/levels/axes.py +++ b/CodeEntropy/levels/axes.py @@ -76,7 +76,10 @@ def get_residue_axes(self, data_container, index: int, residue=None): * Set translational axes equal to rotational axes (as per the original code convention). - If bonded to other residues: - * Use default axes and MOI (MDAnalysis principal axes / inertia). + Find edge heavy atoms (i.e. heavy atoms bonded to neighbour residues) + and find the shortest chain between them: the backbone. Edge + atoms + backbone COM are used to determine UA translational axes + (see get_residue_custom_axes) Args: data_container (MDAnalysis.Universe or AtomGroup): @@ -185,7 +188,12 @@ def get_UA_axes(self, data_container, index: int, res_position): This preserves the original behaviour and its rationale: - Translational axes: - Use the same custom principal-axes approach as residue level: + Use the same approach as residue level rotational. + Identify residue of interest and neighbours, then select + edge heavy atoms (i.e. heavy atoms bonded to neighbour residues) + and find the shortest chain between them: the backbone. Edge + atoms + backbone COM are used to determine UA translational axes + (see get_residue_custom_axes) compute a custom MOI tensor using heavy-atom coordinates but UA masses (heavy + bonded H masses), then compute the principal axes from it. @@ -213,7 +221,7 @@ def get_UA_axes(self, data_container, index: int, res_position): IndexError: If `index` does not correspond to an existing heavy atom. ValueError: - If bonded-axis construction fails. + If axis construction fails. """ index = int(index) # bead index @@ -296,12 +304,16 @@ def get_UA_axes(self, data_container, index: int, res_position): heavy_atom_index = heavy_atom_indices[index] heavy_atom = residue.atoms.select_atoms(f"index {heavy_atom_index}") + if trans_axes is None: + raise ValueError("Unable to compute translation axes for UA bead.") + rot_center = heavy_atom.positions[0] rot_axes, moment_of_inertia = self.get_bonded_axes( system=data_container, atom=heavy_atom[0], dimensions=data_container.dimensions[:3], ) + if rot_axes is None or moment_of_inertia is None: raise ValueError("Unable to compute bonded axes for UA bead.") @@ -316,8 +328,12 @@ def get_UA_axes(self, data_container, index: int, res_position): def get_residue_custom_axes(self, edges, center): """ Compute rotation axes at the residue level, given - two edge atoms of the residue (heavy atoms bonded - to neighbouring residues), and the rotation centre. + two edge atoms of the residue (E1+E2), + and the rotation centre (O). + - x axis is O-E1 + - y axis is O-Q (perpendicular to O-E1 in the + same plane as E2) + - z axis is perpendicular to the two other axes Q --- E2 | | From 9f15b7e71118079fd4cd2d5b9d812bafba59bf64 Mon Sep 17 00:00:00 2001 From: ioana Date: Thu, 30 Apr 2026 14:21:36 +0100 Subject: [PATCH 21/47] fix for 1 heavy atom case --- CodeEntropy/levels/axes.py | 2 ++ .../CodeEntropy/levels/nodes/test_frame_covariance_node.py | 6 +++--- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py index f4fbf258..50777141 100644 --- a/CodeEntropy/levels/axes.py +++ b/CodeEntropy/levels/axes.py @@ -293,7 +293,9 @@ def get_UA_axes(self, data_container, index: int, res_position): else: # only one heavy atom or hydrogen molecule make_whole(data_container.atoms) + residue = data_container trans_axes = data_container.atoms.principal_axes() + residue_heavy_atoms = residue.atoms.select_atoms("mass 2 to 999") # look for heavy atoms in residue of interest heavy_atom_indices = [] diff --git a/tests/unit/CodeEntropy/levels/nodes/test_frame_covariance_node.py b/tests/unit/CodeEntropy/levels/nodes/test_frame_covariance_node.py index 1b51006b..1fdafd98 100644 --- a/tests/unit/CodeEntropy/levels/nodes/test_frame_covariance_node.py +++ b/tests/unit/CodeEntropy/levels/nodes/test_frame_covariance_node.py @@ -401,7 +401,7 @@ def test_build_ua_vectors_customised_axes_true_calls_get_UA_axes(): node = FrameCovarianceNode() bead = _BeadGroup(1) - residue_atoms = MagicMock() + residue_group = MagicMock() axes_manager = MagicMock() axes_manager.get_UA_axes.return_value = ( @@ -416,7 +416,7 @@ def test_build_ua_vectors_customised_axes_true_calls_get_UA_axes(): force_vecs, torque_vecs = node._build_ua_vectors( bead_groups=[bead], - residue_atoms=residue_atoms, + residue_group=residue_group, axes_manager=axes_manager, box=np.array([10.0, 10.0, 10.0]), force_partitioning=1.0, @@ -451,7 +451,7 @@ def test_build_ua_vectors_vanilla_path_uses_principal_axes_and_vanilla_axes( force_vecs, torque_vecs = node._build_ua_vectors( bead_groups=[bead], - residue_atoms=residue_atoms, + residue_group=residue_atoms, axes_manager=axes_manager, box=np.array([10.0, 10.0, 10.0]), force_partitioning=1.0, From 5da1d6a32720031477731eebad3bfb5d526e74ef Mon Sep 17 00:00:00 2001 From: ioana Date: Thu, 30 Apr 2026 15:13:19 +0100 Subject: [PATCH 22/47] changed tests to match code changes --- .../unit/CodeEntropy/levels/nodes/test_frame_covariance_node.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/unit/CodeEntropy/levels/nodes/test_frame_covariance_node.py b/tests/unit/CodeEntropy/levels/nodes/test_frame_covariance_node.py index 1fdafd98..00233f1a 100644 --- a/tests/unit/CodeEntropy/levels/nodes/test_frame_covariance_node.py +++ b/tests/unit/CodeEntropy/levels/nodes/test_frame_covariance_node.py @@ -422,6 +422,7 @@ def test_build_ua_vectors_customised_axes_true_calls_get_UA_axes(): force_partitioning=1.0, customised_axes=True, is_highest=True, + res_position=None, ) axes_manager.get_UA_axes.assert_called_once() @@ -457,6 +458,7 @@ def test_build_ua_vectors_vanilla_path_uses_principal_axes_and_vanilla_axes( force_partitioning=1.0, customised_axes=False, is_highest=True, + res_position=None, ) axes_manager.get_vanilla_axes.assert_called_once() From af80571050c71ef19e025e1bf05538ece9d06854 Mon Sep 17 00:00:00 2001 From: Ioana Papa Date: Fri, 1 May 2026 15:22:52 +0100 Subject: [PATCH 23/47] fixed translation centre for single heavy atom case --- CodeEntropy/levels/axes.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py index 50777141..20ab150e 100644 --- a/CodeEntropy/levels/axes.py +++ b/CodeEntropy/levels/axes.py @@ -294,6 +294,8 @@ def get_UA_axes(self, data_container, index: int, res_position): # only one heavy atom or hydrogen molecule make_whole(data_container.atoms) residue = data_container + # trans_center is center of mass + trans_center = np.array(data_container.center_of_mass()) trans_axes = data_container.atoms.principal_axes() residue_heavy_atoms = residue.atoms.select_atoms("mass 2 to 999") From 9d1e3b78e5712e70fc247a93bf0c869cb9f791ff Mon Sep 17 00:00:00 2001 From: ioanapapa Date: Mon, 11 May 2026 15:08:49 +0100 Subject: [PATCH 24/47] changed test_get_UA_axes_raises_when_bonded_axes_fail --- tests/unit/CodeEntropy/levels/test_axes.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/CodeEntropy/levels/test_axes.py b/tests/unit/CodeEntropy/levels/test_axes.py index 4fc0c563..c0c54235 100644 --- a/tests/unit/CodeEntropy/levels/test_axes.py +++ b/tests/unit/CodeEntropy/levels/test_axes.py @@ -217,7 +217,7 @@ def _sel(q): monkeypatch.setattr(ax, "get_bonded_axes", lambda **kwargs: (None, None)) with pytest.raises(ValueError): - ax.get_UA_axes(u, index=0, res_position=None) + ax.get_UA_axes(u, index=5, res_position=None) def test_get_custom_axes_degenerate_axis1_raises(): From 9694ad1031d021d7b0040873f7e71a4e41d00541 Mon Sep 17 00:00:00 2001 From: Ioana Papa Date: Mon, 8 Jun 2026 12:38:43 +0100 Subject: [PATCH 25/47] correction for single heavy atom case --- CodeEntropy/levels/axes.py | 24 +++++++++++++----------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py index 20ab150e..6a5e66fb 100644 --- a/CodeEntropy/levels/axes.py +++ b/CodeEntropy/levels/axes.py @@ -223,7 +223,6 @@ def get_UA_axes(self, data_container, index: int, res_position): ValueError: If axis construction fails. """ - index = int(index) # bead index heavy_atoms = data_container.atoms.select_atoms("mass 2 to 999") @@ -298,16 +297,19 @@ def get_UA_axes(self, data_container, index: int, res_position): trans_center = np.array(data_container.center_of_mass()) trans_axes = data_container.atoms.principal_axes() - residue_heavy_atoms = residue.atoms.select_atoms("mass 2 to 999") - # look for heavy atoms in residue of interest - heavy_atom_indices = [] - for atom in residue_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 = residue.atoms.select_atoms(f"index {heavy_atom_index}") - + if len(heavy_atoms) > 1: + residue_heavy_atoms = residue.atoms.select_atoms("mass 2 to 999") + # look for heavy atoms in residue of interest + heavy_atom_indices = [] + for atom in residue_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 = residue.atoms.select_atoms(f"index {heavy_atom_index}") + else: + # only the one heavy atom + heavy_atom = heavy_atoms[0] if trans_axes is None: raise ValueError("Unable to compute translation axes for UA bead.") From f5d54317d7683fbdc653a940d0a9058da38082a1 Mon Sep 17 00:00:00 2001 From: ioana Date: Tue, 9 Jun 2026 13:30:49 +0100 Subject: [PATCH 26/47] changed residue level test for vanilla axes --- tests/unit/CodeEntropy/levels/test_axes.py | 32 ++++++++++++++++------ 1 file changed, 24 insertions(+), 8 deletions(-) diff --git a/tests/unit/CodeEntropy/levels/test_axes.py b/tests/unit/CodeEntropy/levels/test_axes.py index c0c54235..2f369e9b 100644 --- a/tests/unit/CodeEntropy/levels/test_axes.py +++ b/tests/unit/CodeEntropy/levels/test_axes.py @@ -117,13 +117,17 @@ def _select_atoms(q): assert np.allclose(moi, np.array([3.0, 2.0, 1.0])) -def test_get_residue_axes_with_bonds_uses_vanilla_axes(monkeypatch): +def test_get_residue_axes_uses_vanilla_axes(monkeypatch): ax = AxesCalculator() residue = MagicMock() residue.__len__.return_value = 1 residue.atoms.center_of_mass.return_value = np.array([1.0, 2.0, 3.0]) residue.center_of_mass.return_value = np.array([1.0, 2.0, 3.0]) + residue.select_atoms.return_value = MagicMock( + positions=[[1.0, 2.0, 3.0], [3.0, 2.0, 1.0]] + ) + uas = MagicMock(positions=np.zeros((2, 3))) u = MagicMock() u.dimensions = np.array([10.0, 10.0, 10.0, 90, 90, 90]) @@ -135,20 +139,25 @@ def _select_atoms(q): return [1] # non-empty if q.startswith("resindex "): return residue + if q == "mass 2 to 999": + return uas return [] + monkeypatch.setattr(ax, "get_UA_masses", lambda mol: [10.0, 12.0]) u.select_atoms.side_effect = _select_atoms + residue = u.select_atoms("resindex 10") monkeypatch.setattr("CodeEntropy.levels.axes.make_whole", lambda _ag: None) - monkeypatch.setattr( - ax, "get_vanilla_axes", lambda mol: (np.eye(3) * 2, np.array([9.0, 8.0, 7.0])) - ) + eigenvalues, eigenvectors = np.linalg.eig([[48, 0, 48], [0, 96, 0], [48, 0, 48]]) + transposed = np.transpose(eigenvectors) + axes = transposed[[2, 0, 1]] + axes[2] = -axes[2] - trans, rot, center, moi = ax.get_residue_axes(u, index=10) + trans, rot, center, moi = ax.get_residue_axes(u, index=10, residue=residue) - assert np.allclose(trans, np.eye(3)) - assert np.allclose(rot, np.eye(3) * 2) - assert np.allclose(moi, np.array([9.0, 8.0, 7.0])) + assert np.allclose(trans, axes) + assert np.allclose(rot, axes) + assert np.allclose(moi, np.array([96.0, 96.0, 0.0])) def test_get_UA_axes_uses_principal_axes_when_single_heavy(monkeypatch): @@ -638,6 +647,13 @@ def center_of_mass(self, *args, **kwargs): def __getitem__(self, idx): return system_atom + def _select_atoms(q): + if q == "prop mass > 1.1": + return heavy_atoms + if q.startswith("index "): + return heavy_atom_selection + return _FakeAtomGroup([]) + data_container = MagicMock() data_container.atoms = _Atoms() data_container.dimensions = np.array([10.0, 10.0, 10.0, 90, 90, 90], dtype=float) From 9a8aae0278006ca5b5308cb3b9898563539101df Mon Sep 17 00:00:00 2001 From: ioana Date: Wed, 10 Jun 2026 10:22:30 +0100 Subject: [PATCH 27/47] updated tests and removed redundant 1 heavy atom code --- CodeEntropy/levels/axes.py | 40 ++++++++++++---------- tests/unit/CodeEntropy/levels/test_axes.py | 27 +++++++-------- 2 files changed, 33 insertions(+), 34 deletions(-) diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py index 6a5e66fb..7eff00f6 100644 --- a/CodeEntropy/levels/axes.py +++ b/CodeEntropy/levels/axes.py @@ -123,6 +123,7 @@ def get_residue_axes(self, data_container, index: int, residue=None): # No UAS are bonded 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 H. + moi_tensor = self.get_moment_of_inertia_tensor( center_of_mass=np.array(residue.center_of_mass()), positions=uas.positions, @@ -225,7 +226,6 @@ def get_UA_axes(self, data_container, index: int, res_position): """ index = int(index) # bead index heavy_atoms = data_container.atoms.select_atoms("mass 2 to 999") - # use the same customPI trans axes as the residue level if len(heavy_atoms) > 1: if len(data_container.residues) == 1: @@ -289,15 +289,6 @@ def get_UA_axes(self, data_container, index: int, res_position): trans_center = np.array(backbone.center_of_mass()) trans_axes = self.get_residue_custom_axes(edges, trans_center) - else: - # only one heavy atom or hydrogen molecule - make_whole(data_container.atoms) - residue = data_container - # trans_center is center of mass - trans_center = np.array(data_container.center_of_mass()) - trans_axes = data_container.atoms.principal_axes() - - if len(heavy_atoms) > 1: residue_heavy_atoms = residue.atoms.select_atoms("mass 2 to 999") # look for heavy atoms in residue of interest heavy_atom_indices = [] @@ -307,19 +298,30 @@ def get_UA_axes(self, data_container, index: int, res_position): # where n is the bead index heavy_atom_index = heavy_atom_indices[index] heavy_atom = residue.atoms.select_atoms(f"index {heavy_atom_index}") + rot_center = heavy_atom.positions[0] + rot_axes, moment_of_inertia = self.get_bonded_axes( + system=data_container, + atom=heavy_atom[0], + dimensions=data_container.dimensions[:3], + ) + else: - # only the one heavy atom + # 1 heavy atom in the data_container heavy_atom = heavy_atoms[0] + # trans and rot centres are centre of mass + rot_center = data_container.center_of_mass() + rot_axes, moment_of_inertia = self.get_bonded_axes( + system=data_container, + atom=heavy_atom[0], + dimensions=data_container.dimensions[:3], + ) + trans_center = rot_center + # principal axes + trans_axes = rot_axes + if trans_axes is None: raise ValueError("Unable to compute translation axes for UA bead.") - rot_center = heavy_atom.positions[0] - rot_axes, moment_of_inertia = self.get_bonded_axes( - system=data_container, - atom=heavy_atom[0], - dimensions=data_container.dimensions[:3], - ) - if rot_axes is None or moment_of_inertia is None: raise ValueError("Unable to compute bonded axes for UA bead.") @@ -765,7 +767,6 @@ def get_moment_of_inertia_tensor( """ r = self.get_vector(center_of_mass, positions, dimensions) r2 = np.sum(r**2, axis=1) - masses_arr = np.asarray(list(masses), dtype=float) moment_of_inertia_tensor = np.eye(3) * np.sum(masses_arr * r2) moment_of_inertia_tensor -= np.einsum("i,ij,ik->jk", masses_arr, r, r) @@ -797,6 +798,7 @@ def get_custom_principal_axes( - principal_axes: (3, 3) principal axes (rows). - moment_of_inertia: (3,) principal moments. """ + eigenvalues, eigenvectors = np.linalg.eig(moment_of_inertia_tensor) order = np.abs(eigenvalues).argsort()[::-1] # descending order transposed = np.transpose(eigenvectors) # columns -> rows diff --git a/tests/unit/CodeEntropy/levels/test_axes.py b/tests/unit/CodeEntropy/levels/test_axes.py index 2f369e9b..40252f43 100644 --- a/tests/unit/CodeEntropy/levels/test_axes.py +++ b/tests/unit/CodeEntropy/levels/test_axes.py @@ -166,13 +166,15 @@ def test_get_UA_axes_uses_principal_axes_when_single_heavy(monkeypatch): u = MagicMock() u.dimensions = np.array([10.0, 10.0, 10.0, 90, 90, 90]) u.atoms.principal_axes.return_value = np.eye(3) + u.center_of_mass.return_value = np.array([[4.0, 0.0, 0.0]]) # heavy_atoms length <= 1 => principal_axes path - heavy_atom = MagicMock(index=5) + heavy_atom = MagicMock(index=5, position=np.array([4.0, 0.0, 0.0])) heavy_atoms = [heavy_atom] def _sel(q): - if q == "prop mass > 1.1": + if q == "mass 2 to 999": + # return heavy atoms group return heavy_atoms if q.startswith("index "): # return atom group with positions @@ -181,6 +183,7 @@ def _sel(q): ag.__getitem__.return_value = MagicMock( mass=12.0, position=np.array([4.0, 0.0, 0.0]), index=5 ) + return ag return [] @@ -513,18 +516,16 @@ def _select_atoms(q): assert np.allclose(moi, np.array([3.0, 2.0, 1.0])) -def test_get_residue_axes_with_bonds_vanilla_path(monkeypatch): +def test_get_residue_axes_vanilla_path(monkeypatch): ax = AxesCalculator() - residue = MagicMock() residue.__len__.return_value = 1 - residue.atoms.principal_axes.return_value = np.eye(3) * 2 residue.atoms.center_of_mass.return_value = np.array([1.0, 2.0, 3.0]) residue.center_of_mass.return_value = np.array([1.0, 2.0, 3.0]) + residue.select_atoms.return_value = MagicMock(positions=np.zeros((1, 3))) u = MagicMock() u.dimensions = np.array([10.0, 10.0, 10.0, 90, 90, 90]) - u.atoms.principal_axes.return_value = np.eye(3) * 2 def _select_atoms(q): if q.startswith("(resindex"): @@ -537,7 +538,9 @@ def _select_atoms(q): monkeypatch.setattr("CodeEntropy.levels.axes.make_whole", lambda _ag: None) monkeypatch.setattr( - ax, "get_vanilla_axes", lambda mol: (np.eye(3) * 2, np.array([9.0, 8.0, 7.0])) + ax, + "get_custom_principal_axes", + lambda mol: (np.eye(3) * 2, np.array([9.0, 8.0, 7.0])), ) trans, rot, center, moi = ax.get_residue_axes(u, index=10, residue=residue) @@ -647,19 +650,13 @@ def center_of_mass(self, *args, **kwargs): def __getitem__(self, idx): return system_atom - def _select_atoms(q): - if q == "prop mass > 1.1": - return heavy_atoms - if q.startswith("index "): - return heavy_atom_selection - return _FakeAtomGroup([]) - data_container = MagicMock() data_container.atoms = _Atoms() data_container.dimensions = np.array([10.0, 10.0, 10.0, 90, 90, 90], dtype=float) + _FakeAtomGroup.atoms = heavy_atom_selection def _select_atoms(q): - if q == "prop mass > 1.1": + if q == "mass 2 to 999": return heavy_atoms if q.startswith("index "): return heavy_atom_selection From 6b4a2d8c53b3be65a7a8b20bfc2a2207f8618319 Mon Sep 17 00:00:00 2001 From: ioana Date: Wed, 10 Jun 2026 11:23:17 +0100 Subject: [PATCH 28/47] updated unit tests --- CodeEntropy/levels/axes.py | 9 +++++---- tests/unit/CodeEntropy/levels/test_axes.py | 20 +++++++++++--------- 2 files changed, 16 insertions(+), 13 deletions(-) diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py index 7eff00f6..a06c5346 100644 --- a/CodeEntropy/levels/axes.py +++ b/CodeEntropy/levels/axes.py @@ -225,7 +225,7 @@ def get_UA_axes(self, data_container, index: int, res_position): If axis construction fails. """ index = int(index) # bead index - heavy_atoms = data_container.atoms.select_atoms("mass 2 to 999") + heavy_atoms = data_container.select_atoms("mass 2 to 999") # use the same customPI trans axes as the residue level if len(heavy_atoms) > 1: if len(data_container.residues) == 1: @@ -240,7 +240,7 @@ def get_UA_axes(self, data_container, index: int, res_position): index_next = residue.resid + 1 # the .resid attribute gives 1-indexing # substract 1 to match indexing later - second_edge = data_container.atoms.select_atoms( + second_edge = data_container.select_atoms( f"resindex {residue.resid - 1} and " f"bonded resindex {index_next - 1}" ) @@ -254,7 +254,7 @@ def get_UA_axes(self, data_container, index: int, res_position): residue = data_container.residues[1] index_prev = residue.resid - 1 index_next = residue.resid + 1 - edge_set = data_container.atoms.select_atoms( + edge_set = data_container.select_atoms( f"resindex {residue.resid - 1} and " f"(bonded resindex {index_next - 1} or " f"resindex {index_prev - 1})" @@ -266,7 +266,7 @@ def get_UA_axes(self, data_container, index: int, res_position): # last resid residue = data_container.residues[1] index_prev = residue.resid - 1 - first_edge = data_container.atoms.select_atoms( + first_edge = data_container.select_atoms( f"resindex {residue.resid - 1} and " f"bonded resindex {index_prev - 1}" ) @@ -283,6 +283,7 @@ def get_UA_axes(self, data_container, index: int, res_position): last = heavy_atom else: last_index -= 1 + edges = [first_edge.atoms[0], last] backbone = self.get_chain(residue, first_edge.atoms[0], last) diff --git a/tests/unit/CodeEntropy/levels/test_axes.py b/tests/unit/CodeEntropy/levels/test_axes.py index 40252f43..4c4a894a 100644 --- a/tests/unit/CodeEntropy/levels/test_axes.py +++ b/tests/unit/CodeEntropy/levels/test_axes.py @@ -213,7 +213,7 @@ def test_get_UA_axes_raises_when_bonded_axes_fail(monkeypatch): heavy_atoms = [heavy_atom] def _sel(q): - if q == "prop mass > 1.1": + if q == "mass 2 to 999": return heavy_atoms if q.startswith("index "): ag = MagicMock() @@ -650,8 +650,18 @@ def center_of_mass(self, *args, **kwargs): def __getitem__(self, idx): return system_atom + def principal_axes(self, *args, **kwargs): + return np.eye(3) + + def select_atoms(self, q): + if q == "mass 2 to 999": + return heavy_atoms + if q.startswith("index "): + return heavy_atom_selection + data_container = MagicMock() data_container.atoms = _Atoms() + data_container.residues.__len__.return_value = 1 data_container.dimensions = np.array([10.0, 10.0, 10.0, 90, 90, 90], dtype=float) _FakeAtomGroup.atoms = heavy_atom_selection @@ -671,12 +681,6 @@ def _select_atoms(q): ) monkeypatch.setattr(ax, "get_UA_masses", lambda _ag: [12.0, 12.0]) - got_tensor = MagicMock(return_value=np.eye(3)) - monkeypatch.setattr(ax, "get_moment_of_inertia_tensor", got_tensor) - - got_custom_axes = MagicMock(return_value=(np.eye(3), np.array([3.0, 2.0, 1.0]))) - monkeypatch.setattr(ax, "get_custom_principal_axes", got_custom_axes) - trans_axes, rot_axes, center, moi = ax.get_UA_axes( data_container, index=0, res_position=None ) @@ -685,8 +689,6 @@ def _select_atoms(q): assert rot_axes.shape == (3, 3) assert np.allclose(center, np.array([0.0, 0.0, 0.0])) assert moi.shape == (3,) - got_tensor.assert_called_once() - got_custom_axes.assert_called_once() def test_get_bonded_axes_returns_none_none_if_custom_axes_none(monkeypatch): From 620c22d0ad1fe8d9330436591da6a7499afaae1f Mon Sep 17 00:00:00 2001 From: ioana Date: Wed, 10 Jun 2026 13:54:07 +0100 Subject: [PATCH 29/47] modified dna regression tests to match results with new custom axes --- .../regression/baselines/dna/combined_forcetorque_off.json | 4 ++-- tests/regression/baselines/dna/default.json | 4 ++-- tests/regression/baselines/dna/frame_window.json | 6 +++--- tests/regression/baselines/dna/grouping_each.json | 4 ++-- 4 files changed, 9 insertions(+), 9 deletions(-) diff --git a/tests/regression/baselines/dna/combined_forcetorque_off.json b/tests/regression/baselines/dna/combined_forcetorque_off.json index e867cee1..5671e036 100644 --- a/tests/regression/baselines/dna/combined_forcetorque_off.json +++ b/tests/regression/baselines/dna/combined_forcetorque_off.json @@ -5,14 +5,14 @@ "united_atom:Transvibrational": 0.0, "united_atom:Rovibrational": 0.002160679012128457, "residue:Transvibrational": 0.0, - "residue:Rovibrational": 3.376800684085249, + "residue:Rovibrational": 6.832779629985204, "polymer:Transvibrational": 21.18266215491188, "polymer:Rovibrational": 12.837576042626923, "united_atom:Conformational": 0.0, "residue:Conformational": 0.0, "polymer:Orientational": 4.758905336627712 }, - "total": 42.1581048972639 + "total": 45.61408384316385 }, "1": { "components": { diff --git a/tests/regression/baselines/dna/default.json b/tests/regression/baselines/dna/default.json index 2cbef740..760c664d 100644 --- a/tests/regression/baselines/dna/default.json +++ b/tests/regression/baselines/dna/default.json @@ -5,14 +5,14 @@ "united_atom:Transvibrational": 0.0, "united_atom:Rovibrational": 0.002160679012128457, "residue:Transvibrational": 0.0, - "residue:Rovibrational": 3.376800684085249, + "residue:Rovibrational": 6.832779629985204, "polymer:FTmat-Transvibrational": 12.341104347192612, "polymer:FTmat-Rovibrational": 0.0, "united_atom:Conformational": 0.0, "residue:Conformational": 0.0, "polymer:Orientational": 4.758905336627712 }, - "total": 20.478971046917703 + "total": 23.934949992817657 }, "1": { "components": { diff --git a/tests/regression/baselines/dna/frame_window.json b/tests/regression/baselines/dna/frame_window.json index c35d2211..a4dc5e60 100644 --- a/tests/regression/baselines/dna/frame_window.json +++ b/tests/regression/baselines/dna/frame_window.json @@ -2,17 +2,17 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 0.0, + "united_atom:Transvibrational": 2.3245522868433963e-06, "united_atom:Rovibrational": 1.5821720528374943, "residue:Transvibrational": 0.0, - "residue:Rovibrational": 27.397449238560412, + "residue:Rovibrational": 36.126942876174176, "polymer:FTmat-Transvibrational": 48.62026970762269, "polymer:FTmat-Rovibrational": 0.0, "united_atom:Conformational": 10.584542990557836, "residue:Conformational": 0.0, "polymer:Orientational": 4.758905336627712 }, - "total": 92.94333932620614 + "total": 101.6728352883722 }, "1": { "components": { diff --git a/tests/regression/baselines/dna/grouping_each.json b/tests/regression/baselines/dna/grouping_each.json index 2cbef740..760c664d 100644 --- a/tests/regression/baselines/dna/grouping_each.json +++ b/tests/regression/baselines/dna/grouping_each.json @@ -5,14 +5,14 @@ "united_atom:Transvibrational": 0.0, "united_atom:Rovibrational": 0.002160679012128457, "residue:Transvibrational": 0.0, - "residue:Rovibrational": 3.376800684085249, + "residue:Rovibrational": 6.832779629985204, "polymer:FTmat-Transvibrational": 12.341104347192612, "polymer:FTmat-Rovibrational": 0.0, "united_atom:Conformational": 0.0, "residue:Conformational": 0.0, "polymer:Orientational": 4.758905336627712 }, - "total": 20.478971046917703 + "total": 23.934949992817657 }, "1": { "components": { From 246c72bcc8155e8a98eff8d36fb226d37ce88e66 Mon Sep 17 00:00:00 2001 From: ioana Date: Wed, 10 Jun 2026 13:58:41 +0100 Subject: [PATCH 30/47] changed dna selection subset regression test --- tests/regression/baselines/dna/selection_subset.json | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/regression/baselines/dna/selection_subset.json b/tests/regression/baselines/dna/selection_subset.json index 2cbef740..760c664d 100644 --- a/tests/regression/baselines/dna/selection_subset.json +++ b/tests/regression/baselines/dna/selection_subset.json @@ -5,14 +5,14 @@ "united_atom:Transvibrational": 0.0, "united_atom:Rovibrational": 0.002160679012128457, "residue:Transvibrational": 0.0, - "residue:Rovibrational": 3.376800684085249, + "residue:Rovibrational": 6.832779629985204, "polymer:FTmat-Transvibrational": 12.341104347192612, "polymer:FTmat-Rovibrational": 0.0, "united_atom:Conformational": 0.0, "residue:Conformational": 0.0, "polymer:Orientational": 4.758905336627712 }, - "total": 20.478971046917703 + "total": 23.934949992817657 }, "1": { "components": { From 6be94e63a1ab87015ef96898b6ec468d1be28048 Mon Sep 17 00:00:00 2001 From: Ioana Papa Date: Thu, 11 Jun 2026 16:25:15 +0100 Subject: [PATCH 31/47] changed residue axes to match NCC for proteins --- CodeEntropy/levels/axes.py | 74 ++++++++++++++++++++------------------ 1 file changed, 40 insertions(+), 34 deletions(-) diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py index a06c5346..892ec4ca 100644 --- a/CodeEntropy/levels/axes.py +++ b/CodeEntropy/levels/axes.py @@ -118,7 +118,7 @@ def get_residue_axes(self, data_container, index: int, residue=None): uas = residue.select_atoms("mass 2 to 999") ua_masses = self.get_UA_masses(residue) - + print(f"The edge atoms: {edge_atom_set}") if len(edge_atom_set) == 0: # No UAS are bonded to other residues # Use a custom principal axes, from a MOI tensor that uses positions of @@ -132,7 +132,7 @@ def get_residue_axes(self, data_container, index: int, residue=None): ) rot_axes, moment_of_inertia = self.get_custom_principal_axes(moi_tensor) trans_axes = rot_axes # per original convention - center = np.array(residue.center_of_mass()) + rot_center = np.array(residue.center_of_mass()) else: # If bonded to other residues, use local axes. make_whole(data_container.atoms) @@ -169,17 +169,20 @@ def get_residue_axes(self, data_container, index: int, residue=None): backbone = self.get_chain(residue, edge_atom_set[0], edge_atom_set[1]) # get edge atoms of the residue # for terminal residues, this will include the C/N terminus - center = np.array(backbone.center_of_mass()) - rot_axes = self.get_residue_custom_axes(edges, center) + backbone_center = np.zeros(3) + for heavy_atom in backbone: + backbone_center += heavy_atom.position + backbone_center = backbone_center / len(backbone) + rot_center, rot_axes = self.get_residue_custom_axes(edges, backbone_center) moment_of_inertia = self.get_custom_residue_moment_of_inertia( - center_of_mass=center, + center_of_mass=rot_center, positions=uas.positions, masses=ua_masses, custom_rot_axes=rot_axes, dimensions=data_container.dimensions[:3], ) - return trans_axes, rot_axes, center, moment_of_inertia + return trans_axes, rot_axes, rot_center, moment_of_inertia def get_UA_axes(self, data_container, index: int, res_position): """Compute united-atom-level translational and rotational axes. @@ -287,8 +290,13 @@ def get_UA_axes(self, data_container, index: int, res_position): edges = [first_edge.atoms[0], last] backbone = self.get_chain(residue, first_edge.atoms[0], last) - trans_center = np.array(backbone.center_of_mass()) - trans_axes = self.get_residue_custom_axes(edges, trans_center) + backbone_center = np.zeros(3) + for heavy_atom in backbone: + backbone_center += heavy_atom.position + backbone_center = backbone_center / len(backbone) + trans_center, trans_axes = self.get_residue_custom_axes( + edges, backbone_center + ) residue_heavy_atoms = residue.atoms.select_atoms("mass 2 to 999") # look for heavy atoms in residue of interest @@ -338,49 +346,45 @@ def get_residue_custom_axes(self, edges, center): """ Compute rotation axes at the residue level, given two edge atoms of the residue (E1+E2), - and the rotation centre (O). + and the centre of geometry of backbone atoms + that are not edges (C). - x axis is O-E1 - - y axis is O-Q (perpendicular to O-E1 in the + - y axis is O-C (perpendicular to O-E1 in the same plane as E2) - z axis is perpendicular to the two other axes - Q --- E2 - | | - | | - E1 ---- O --- P + C + | + | + E1 ---- O --- E2 Args: edges: (2,3) positions of two edge atoms - center: (3,) coordinates of the rotation centre + center: (3,) coordinates of the inner backbone + centre of geometry Returns: + rot_center: (3,) rotation centre -- + it lies on the E1-E2 vector rot_axes: (3,3) rotation axes of residue """ # x axis is O-E1 - E1O_vector = center - edges[0].position - x_axis = -E1O_vector - # y axis is perpendicular to x - # in the same plane as E2 - # look for projection of E1-E2 on E1-O (E1-P) + E1C_vector = center - edges[0].position + # look for projection of E1-O onto E1-E2 (E1-C) E1E2_vector = edges[1].position - edges[0].position - projection = ( - np.dot(E1O_vector, E1E2_vector) / (np.linalg.norm(E1O_vector) ** 2) - ) * E1O_vector - # get the perpendicular onto E1-O (P-E2) - # P-E2 = P-E1 + E1-E2 - perpendicular = E1E2_vector - projection - # get the perpendicular through O (Q-O) - # first get P-Q diagonal through paralellogram rule - # P- Q = P-E2 + P-O - diagonal = -(projection - E1O_vector) + perpendicular - # get the parallel of P-E2 through O - # OQ = OP + PQ - y_axis = (projection - E1O_vector) + diagonal + E1O_vector = ( + np.dot(E1E2_vector, E1C_vector) / (np.linalg.norm(E1E2_vector) ** 2) + ) * E1E2_vector + x_axis = -E1O_vector + # O-C = O-E1 + E1-C + OC_vector = -E1O_vector + E1C_vector + y_axis = OC_vector z_axis = np.cross(x_axis, y_axis) x_axis /= np.linalg.norm(x_axis) y_axis /= np.linalg.norm(y_axis) z_axis /= np.linalg.norm(z_axis) rot_axes = np.array([x_axis, y_axis, z_axis]) + rot_center = E1O_vector - edges[0].position - return rot_axes + return rot_center, rot_axes def get_bonded_axes(self, system, atom, dimensions: np.ndarray): """Compute UA rotational axes from bonded topology around a heavy atom. @@ -906,6 +910,8 @@ def get_chain(self, residue, first, last): chain.append(current) chain_indices.append(current.index - residue.atoms.indices[0]) chain_indices = np.flip(chain_indices) + # only get in between residues + chain_indices = chain_indices[1:-1] # accout for in-residue index chain_AtomGroup = residue.atoms[chain_indices] chain = chain_AtomGroup.atoms.select_atoms("all") From 1f47e6b51c648d29022acde2d5302ab45fc66ca6 Mon Sep 17 00:00:00 2001 From: ioana Date: Fri, 12 Jun 2026 10:41:58 +0100 Subject: [PATCH 32/47] axes change to make more similar to NCC --- CodeEntropy/levels/axes.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py index 892ec4ca..3d82f8f9 100644 --- a/CodeEntropy/levels/axes.py +++ b/CodeEntropy/levels/axes.py @@ -118,7 +118,7 @@ def get_residue_axes(self, data_container, index: int, residue=None): uas = residue.select_atoms("mass 2 to 999") ua_masses = self.get_UA_masses(residue) - print(f"The edge atoms: {edge_atom_set}") + if len(edge_atom_set) == 0: # No UAS are bonded to other residues # Use a custom principal axes, from a MOI tensor that uses positions of @@ -294,6 +294,7 @@ def get_UA_axes(self, data_container, index: int, res_position): for heavy_atom in backbone: backbone_center += heavy_atom.position backbone_center = backbone_center / len(backbone) + trans_center, trans_axes = self.get_residue_custom_axes( edges, backbone_center ) From 25d01505e0645862f06fd9cc0a15618b395252dc Mon Sep 17 00:00:00 2001 From: ioana Date: Fri, 12 Jun 2026 13:06:33 +0100 Subject: [PATCH 33/47] corrected center of mass --- CodeEntropy/levels/axes.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py index 3d82f8f9..2138922c 100644 --- a/CodeEntropy/levels/axes.py +++ b/CodeEntropy/levels/axes.py @@ -383,8 +383,7 @@ def get_residue_custom_axes(self, edges, center): y_axis /= np.linalg.norm(y_axis) z_axis /= np.linalg.norm(z_axis) rot_axes = np.array([x_axis, y_axis, z_axis]) - rot_center = E1O_vector - edges[0].position - + rot_center = E1O_vector + edges[0].position return rot_center, rot_axes def get_bonded_axes(self, system, atom, dimensions: np.ndarray): From e9d3cfc1826ec4bd3ebbdbdb0f2eb1a406f07adf Mon Sep 17 00:00:00 2001 From: ioana Date: Fri, 12 Jun 2026 15:29:43 +0100 Subject: [PATCH 34/47] centre of rotation correction --- CodeEntropy/levels/axes.py | 32 +++++++++++++++++--------------- 1 file changed, 17 insertions(+), 15 deletions(-) diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py index 2138922c..24ad97a5 100644 --- a/CodeEntropy/levels/axes.py +++ b/CodeEntropy/levels/axes.py @@ -349,23 +349,25 @@ def get_residue_custom_axes(self, edges, center): two edge atoms of the residue (E1+E2), and the centre of geometry of backbone atoms that are not edges (C). - - x axis is O-E1 - - y axis is O-C (perpendicular to O-E1 in the + x axis is O-E1 + y axis is O-C (perpendicular to O-E1 in the same plane as E2) - - z axis is perpendicular to the two other axes + z axis is perpendicular to the two other axes C | | E1 ---- O --- E2 - Args: - edges: (2,3) positions of two edge atoms - center: (3,) coordinates of the inner backbone - centre of geometry - Returns: - rot_center: (3,) rotation centre -- - it lies on the E1-E2 vector - rot_axes: (3,3) rotation axes of residue + + Args: + edges: (2,3) positions of two edge atoms + center: (3,) coordinates of the inner backbone + centre of geometry + + Returns: + rot_center: (3,) rotation centre, + lies on the E1-E2 vector + rot_axes: (3,3) rotation axes of residue """ # x axis is O-E1 E1C_vector = center - edges[0].position @@ -847,15 +849,15 @@ def get_chain(self, residue, first, last): For a given MDAnalysis AtomGroup and two given heavy atoms within that AtomGroup, return the shortest path between the two atoms. + Args: residue: MDAnalysis AtomGroup representing - the residue/monomer of interest. + the residue/monomer of interest. first: First heavy atom in the chain last: Last heavy atom in the chain - Returns: - chain: MDAnalysis AtomGroup containing - the chain heavy atoms. + Returns: + chain: MDAnalysis AtomGroup containing chain atoms. """ chain = [] chain_indices = [] From 828bde01bc769f95253f5f446f30d112ce1841e6 Mon Sep 17 00:00:00 2001 From: ioana Date: Fri, 12 Jun 2026 15:31:00 +0100 Subject: [PATCH 35/47] centre of rotation correction --- CodeEntropy/levels/axes.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py index 24ad97a5..c3d1b4bf 100644 --- a/CodeEntropy/levels/axes.py +++ b/CodeEntropy/levels/axes.py @@ -354,6 +354,8 @@ def get_residue_custom_axes(self, edges, center): same plane as E2) z axis is perpendicular to the two other axes + :: + C | | From d8fe0421bf01046dd0b681d8e1e02de5edde7cad Mon Sep 17 00:00:00 2001 From: ioana Date: Mon, 22 Jun 2026 14:16:38 +0100 Subject: [PATCH 36/47] fix atom is not subscriptable error --- CodeEntropy/levels/axes.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py index c3d1b4bf..f1c8bb75 100644 --- a/CodeEntropy/levels/axes.py +++ b/CodeEntropy/levels/axes.py @@ -322,7 +322,7 @@ def get_UA_axes(self, data_container, index: int, res_position): rot_center = data_container.center_of_mass() rot_axes, moment_of_inertia = self.get_bonded_axes( system=data_container, - atom=heavy_atom[0], + atom=heavy_atom, dimensions=data_container.dimensions[:3], ) trans_center = rot_center From ec0c03073a1c3954a214d5459f3360e8a9c60b84 Mon Sep 17 00:00:00 2001 From: ioana Date: Mon, 22 Jun 2026 14:22:15 +0100 Subject: [PATCH 37/47] updates baselines --- .../combined_forcetorque_false.json | 4 +- .../baselines/benzaldehyde/frame_window.json | 4 +- .../benzaldehyde/selection_subset.json | 4 +- .../benzene/combined_forcetorque_off.json | 4 +- .../baselines/benzene/frame_window.json | 4 +- .../baselines/benzene/selection_subset.json | 4 +- .../cyclohexane/combined_forcetorque_off.json | 4 +- .../baselines/cyclohexane/frame_window.json | 4 +- .../cyclohexane/selection_subset.json | 4 +- .../dna/combined_forcetorque_off.json | 6 +-- .../baselines/dna/frame_window.json | 10 ++-- .../baselines/dna/grouping_each.json | 6 +-- .../baselines/dna/selection_subset.json | 6 +-- .../combined_forcetorque_off.json | 4 +- .../baselines/ethyl-acetate/frame_window.json | 4 +- .../ethyl-acetate/selection_subset.json | 4 +- .../methane/combined_forcetorque_off.json | 6 +-- .../baselines/methane/frame_window.json | 4 +- .../baselines/methane/grouping_each.json | 44 +++++++++--------- .../baselines/methane/selection_subset.json | 6 +-- .../octonol/combined_forcetorque_off.json | 4 +- .../baselines/octonol/frame_window.json | 4 +- .../baselines/octonol/selection_subset.json | 4 +- .../water/combined_forcetorque_off.json | 6 +-- .../baselines/water/frame_window.json | 4 +- .../baselines/water/grouping_each.json | 46 +++++++++---------- .../baselines/water/selection_subset.json | 6 +-- .../regression/baselines/water/water_off.json | 6 +-- 28 files changed, 108 insertions(+), 108 deletions(-) diff --git a/tests/regression/baselines/benzaldehyde/combined_forcetorque_false.json b/tests/regression/baselines/benzaldehyde/combined_forcetorque_false.json index a5d1af0e..27025e62 100644 --- a/tests/regression/baselines/benzaldehyde/combined_forcetorque_false.json +++ b/tests/regression/baselines/benzaldehyde/combined_forcetorque_false.json @@ -2,7 +2,7 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 0.07119323721997475, + "united_atom:Transvibrational": 0.08982962903796131, "united_atom:Rovibrational": 49.68669738152346, "residue:Transvibrational": 69.48692941204929, "residue:Rovibrational": 68.46147102540942, @@ -10,7 +10,7 @@ "residue:Conformational": 0.0, "residue:Orientational": 20.481571492615355 }, - "total": 208.1878625488175 + "total": 208.2064989406355 } } } diff --git a/tests/regression/baselines/benzaldehyde/frame_window.json b/tests/regression/baselines/benzaldehyde/frame_window.json index 70bb2766..27225723 100644 --- a/tests/regression/baselines/benzaldehyde/frame_window.json +++ b/tests/regression/baselines/benzaldehyde/frame_window.json @@ -2,7 +2,7 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 40.30267601961045, + "united_atom:Transvibrational": 41.75648599130188, "united_atom:Rovibrational": 38.21906858443615, "residue:FTmat-Transvibrational": 73.41098578352612, "residue:FTmat-Rovibrational": 57.881504393660364, @@ -10,7 +10,7 @@ "residue:Conformational": 0.0, "residue:Orientational": 20.481571492615355 }, - "total": 230.29580627384846 + "total": 231.74961624553984 } } } diff --git a/tests/regression/baselines/benzaldehyde/selection_subset.json b/tests/regression/baselines/benzaldehyde/selection_subset.json index 6d50b1a7..c7750ebb 100644 --- a/tests/regression/baselines/benzaldehyde/selection_subset.json +++ b/tests/regression/baselines/benzaldehyde/selection_subset.json @@ -2,7 +2,7 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 0.07119323721997475, + "united_atom:Transvibrational": 0.08982962903796131, "united_atom:Rovibrational": 49.68669738152346, "residue:FTmat-Transvibrational": 87.43527331108173, "residue:FTmat-Rovibrational": 61.67126452972779, @@ -10,7 +10,7 @@ "residue:Conformational": 0.0, "residue:Orientational": 20.481571492615355 }, - "total": 219.34599995216834 + "total": 219.3646363439863 } } } diff --git a/tests/regression/baselines/benzene/combined_forcetorque_off.json b/tests/regression/baselines/benzene/combined_forcetorque_off.json index 16b5e375..ee2c41fc 100644 --- a/tests/regression/baselines/benzene/combined_forcetorque_off.json +++ b/tests/regression/baselines/benzene/combined_forcetorque_off.json @@ -2,7 +2,7 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 0.16311383522946216, + "united_atom:Transvibrational": 0.17824017761231503, "united_atom:Rovibrational": 40.695812407899396, "residue:Transvibrational": 57.59675197041957, "residue:Rovibrational": 52.85496348419672, @@ -10,7 +10,7 @@ "residue:Conformational": 0.0, "residue:Orientational": 0.0 }, - "total": 151.31064169774515 + "total": 151.325768040128 } } } diff --git a/tests/regression/baselines/benzene/frame_window.json b/tests/regression/baselines/benzene/frame_window.json index c1ffd000..0f692cbb 100644 --- a/tests/regression/baselines/benzene/frame_window.json +++ b/tests/regression/baselines/benzene/frame_window.json @@ -2,7 +2,7 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 11.46552036084548, + "united_atom:Transvibrational": 11.552833417413368, "united_atom:Rovibrational": 34.98492056748493, "residue:FTmat-Transvibrational": 71.83491878606151, "residue:FTmat-Rovibrational": 55.6098400281744, @@ -10,7 +10,7 @@ "residue:Conformational": 0.0, "residue:Orientational": 0.0 }, - "total": 173.89519974256632 + "total": 173.9825127991342 } } } diff --git a/tests/regression/baselines/benzene/selection_subset.json b/tests/regression/baselines/benzene/selection_subset.json index ff2d09b5..807d8c05 100644 --- a/tests/regression/baselines/benzene/selection_subset.json +++ b/tests/regression/baselines/benzene/selection_subset.json @@ -2,7 +2,7 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 0.16311383522946216, + "united_atom:Transvibrational": 0.17824017761231503, "united_atom:Rovibrational": 40.695812407899396, "residue:FTmat-Transvibrational": 71.15771063929333, "residue:FTmat-Rovibrational": 47.26253953880574, @@ -10,7 +10,7 @@ "residue:Conformational": 0.0, "residue:Orientational": 0.0 }, - "total": 159.27917642122793 + "total": 159.2943027636108 } } } diff --git a/tests/regression/baselines/cyclohexane/combined_forcetorque_off.json b/tests/regression/baselines/cyclohexane/combined_forcetorque_off.json index a548349d..56de2915 100644 --- a/tests/regression/baselines/cyclohexane/combined_forcetorque_off.json +++ b/tests/regression/baselines/cyclohexane/combined_forcetorque_off.json @@ -2,7 +2,7 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 0.5949308536734588, + "united_atom:Transvibrational": 0.6825851535582973, "united_atom:Rovibrational": 24.234154676578637, "residue:Transvibrational": 69.9872567772153, "residue:Rovibrational": 68.68521019685565, @@ -10,7 +10,7 @@ "residue:Conformational": 0.0, "residue:Orientational": 0.0 }, - "total": 163.50155250432303 + "total": 163.5892068042079 } } } diff --git a/tests/regression/baselines/cyclohexane/frame_window.json b/tests/regression/baselines/cyclohexane/frame_window.json index d21d437d..dd423267 100644 --- a/tests/regression/baselines/cyclohexane/frame_window.json +++ b/tests/regression/baselines/cyclohexane/frame_window.json @@ -2,7 +2,7 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 17.398382397359153, + "united_atom:Transvibrational": 16.605273867798914, "united_atom:Rovibrational": 73.96792995794405, "residue:FTmat-Transvibrational": 76.4267996157354, "residue:FTmat-Rovibrational": 63.30469126284744, @@ -10,7 +10,7 @@ "residue:Conformational": 0.0, "residue:Orientational": 0.0 }, - "total": 231.09780323388605 + "total": 230.3046947043258 } } } diff --git a/tests/regression/baselines/cyclohexane/selection_subset.json b/tests/regression/baselines/cyclohexane/selection_subset.json index b1a14443..a1a37145 100644 --- a/tests/regression/baselines/cyclohexane/selection_subset.json +++ b/tests/regression/baselines/cyclohexane/selection_subset.json @@ -2,7 +2,7 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 0.5949308536734588, + "united_atom:Transvibrational": 0.6825851535582973, "united_atom:Rovibrational": 24.234154676578637, "residue:FTmat-Transvibrational": 84.37184730911717, "residue:FTmat-Rovibrational": 59.52377096811085, @@ -10,7 +10,7 @@ "residue:Conformational": 0.0, "residue:Orientational": 0.0 }, - "total": 168.72470380748013 + "total": 168.81235810736496 } } } diff --git a/tests/regression/baselines/dna/combined_forcetorque_off.json b/tests/regression/baselines/dna/combined_forcetorque_off.json index 5671e036..5e724622 100644 --- a/tests/regression/baselines/dna/combined_forcetorque_off.json +++ b/tests/regression/baselines/dna/combined_forcetorque_off.json @@ -2,17 +2,17 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 0.0, + "united_atom:Transvibrational": 4.762281610623415e-20, "united_atom:Rovibrational": 0.002160679012128457, "residue:Transvibrational": 0.0, - "residue:Rovibrational": 6.832779629985204, + "residue:Rovibrational": 6.633599673254765, "polymer:Transvibrational": 21.18266215491188, "polymer:Rovibrational": 12.837576042626923, "united_atom:Conformational": 0.0, "residue:Conformational": 0.0, "polymer:Orientational": 4.758905336627712 }, - "total": 45.61408384316385 + "total": 45.41490388643341 }, "1": { "components": { diff --git a/tests/regression/baselines/dna/frame_window.json b/tests/regression/baselines/dna/frame_window.json index a4dc5e60..fd24e0e2 100644 --- a/tests/regression/baselines/dna/frame_window.json +++ b/tests/regression/baselines/dna/frame_window.json @@ -2,21 +2,21 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 2.3245522868433963e-06, + "united_atom:Transvibrational": 0.0, "united_atom:Rovibrational": 1.5821720528374943, "residue:Transvibrational": 0.0, - "residue:Rovibrational": 36.126942876174176, + "residue:Rovibrational": 35.89784662172834, "polymer:FTmat-Transvibrational": 48.62026970762269, "polymer:FTmat-Rovibrational": 0.0, "united_atom:Conformational": 10.584542990557836, "residue:Conformational": 0.0, "polymer:Orientational": 4.758905336627712 }, - "total": 101.6728352883722 + "total": 101.44373670937406 }, "1": { "components": { - "united_atom:Transvibrational": 0.0, + "united_atom:Transvibrational": 5.359713875687138e-10, "united_atom:Rovibrational": 2.5277936366208014, "residue:Transvibrational": 0.0, "residue:Rovibrational": 24.80670067454149, @@ -26,7 +26,7 @@ "residue:Conformational": 0.0, "polymer:Orientational": 4.758905336627712 }, - "total": 97.85965049646045 + "total": 97.85965049699642 } } } diff --git a/tests/regression/baselines/dna/grouping_each.json b/tests/regression/baselines/dna/grouping_each.json index 760c664d..7a2b6c3a 100644 --- a/tests/regression/baselines/dna/grouping_each.json +++ b/tests/regression/baselines/dna/grouping_each.json @@ -2,17 +2,17 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 0.0, + "united_atom:Transvibrational": 4.762281610623415e-20, "united_atom:Rovibrational": 0.002160679012128457, "residue:Transvibrational": 0.0, - "residue:Rovibrational": 6.832779629985204, + "residue:Rovibrational": 6.633599673254765, "polymer:FTmat-Transvibrational": 12.341104347192612, "polymer:FTmat-Rovibrational": 0.0, "united_atom:Conformational": 0.0, "residue:Conformational": 0.0, "polymer:Orientational": 4.758905336627712 }, - "total": 23.934949992817657 + "total": 23.735770036087217 }, "1": { "components": { diff --git a/tests/regression/baselines/dna/selection_subset.json b/tests/regression/baselines/dna/selection_subset.json index 760c664d..7a2b6c3a 100644 --- a/tests/regression/baselines/dna/selection_subset.json +++ b/tests/regression/baselines/dna/selection_subset.json @@ -2,17 +2,17 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 0.0, + "united_atom:Transvibrational": 4.762281610623415e-20, "united_atom:Rovibrational": 0.002160679012128457, "residue:Transvibrational": 0.0, - "residue:Rovibrational": 6.832779629985204, + "residue:Rovibrational": 6.633599673254765, "polymer:FTmat-Transvibrational": 12.341104347192612, "polymer:FTmat-Rovibrational": 0.0, "united_atom:Conformational": 0.0, "residue:Conformational": 0.0, "polymer:Orientational": 4.758905336627712 }, - "total": 23.934949992817657 + "total": 23.735770036087217 }, "1": { "components": { diff --git a/tests/regression/baselines/ethyl-acetate/combined_forcetorque_off.json b/tests/regression/baselines/ethyl-acetate/combined_forcetorque_off.json index 8c68cce7..499fca4c 100644 --- a/tests/regression/baselines/ethyl-acetate/combined_forcetorque_off.json +++ b/tests/regression/baselines/ethyl-acetate/combined_forcetorque_off.json @@ -2,7 +2,7 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 1.196367074472673, + "united_atom:Transvibrational": 0.8081103094129065, "united_atom:Rovibrational": 82.36455154278131, "residue:Transvibrational": 64.68750154134017, "residue:Rovibrational": 58.903938937880085, @@ -10,7 +10,7 @@ "residue:Conformational": 0.0, "residue:Orientational": 0.0 }, - "total": 215.05334534012698 + "total": 214.6650885750672 } } } diff --git a/tests/regression/baselines/ethyl-acetate/frame_window.json b/tests/regression/baselines/ethyl-acetate/frame_window.json index be566d07..911a2667 100644 --- a/tests/regression/baselines/ethyl-acetate/frame_window.json +++ b/tests/regression/baselines/ethyl-acetate/frame_window.json @@ -2,7 +2,7 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 28.312711247966337, + "united_atom:Transvibrational": 29.222699118677973, "united_atom:Rovibrational": 51.99036142818903, "residue:FTmat-Transvibrational": 67.80560626717748, "residue:FTmat-Rovibrational": 55.1631201009248, @@ -10,7 +10,7 @@ "residue:Conformational": 0.0, "residue:Orientational": 3.8984476469128193 }, - "total": 215.80571814986257 + "total": 216.7157060205742 } } } diff --git a/tests/regression/baselines/ethyl-acetate/selection_subset.json b/tests/regression/baselines/ethyl-acetate/selection_subset.json index 270dbbd5..dd0f2695 100644 --- a/tests/regression/baselines/ethyl-acetate/selection_subset.json +++ b/tests/regression/baselines/ethyl-acetate/selection_subset.json @@ -2,7 +2,7 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 1.196367074472673, + "united_atom:Transvibrational": 0.8081103094129065, "united_atom:Rovibrational": 82.36455154278131, "residue:FTmat-Transvibrational": 74.5939232239247, "residue:FTmat-Rovibrational": 55.42964192531005, @@ -10,7 +10,7 @@ "residue:Conformational": 0.0, "residue:Orientational": 0.0 }, - "total": 221.48547001014148 + "total": 221.09721324508172 } } } diff --git a/tests/regression/baselines/methane/combined_forcetorque_off.json b/tests/regression/baselines/methane/combined_forcetorque_off.json index 318c30c6..70be3cef 100644 --- a/tests/regression/baselines/methane/combined_forcetorque_off.json +++ b/tests/regression/baselines/methane/combined_forcetorque_off.json @@ -2,12 +2,12 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 42.51619525142699, - "united_atom:Rovibrational": 34.651932322640015, + "united_atom:Transvibrational": 42.516195251426986, + "united_atom:Rovibrational": 34.69255616696802, "united_atom:Conformational": 0.0, "united_atom:Orientational": 0.0 }, - "total": 77.168127574067 + "total": 77.208751418395 } } } diff --git a/tests/regression/baselines/methane/frame_window.json b/tests/regression/baselines/methane/frame_window.json index 9a45dc98..0f70346d 100644 --- a/tests/regression/baselines/methane/frame_window.json +++ b/tests/regression/baselines/methane/frame_window.json @@ -3,11 +3,11 @@ "0": { "components": { "united_atom:Transvibrational": 40.70376001258969, - "united_atom:Rovibrational": 33.778792396823484, + "united_atom:Rovibrational": 33.83818524319629, "united_atom:Conformational": 0.0, "united_atom:Orientational": 0.0 }, - "total": 74.48255240941317 + "total": 74.54194525578598 } } } diff --git a/tests/regression/baselines/methane/grouping_each.json b/tests/regression/baselines/methane/grouping_each.json index 3b31963b..e5f74df8 100644 --- a/tests/regression/baselines/methane/grouping_each.json +++ b/tests/regression/baselines/methane/grouping_each.json @@ -3,92 +3,92 @@ "0": { "components": { "united_atom:Transvibrational": 24.315699188266926, - "united_atom:Rovibrational": 21.0483901242056, + "united_atom:Rovibrational": 21.06129596642336, "united_atom:Conformational": 0.0, "united_atom:Orientational": 0.0 }, - "total": 45.36408931247253 + "total": 45.376995154690285 }, "1": { "components": { "united_atom:Transvibrational": 23.76608267046309, - "united_atom:Rovibrational": 16.769636839671296, + "united_atom:Rovibrational": 16.800440306419098, "united_atom:Conformational": 0.0, "united_atom:Orientational": 0.0 }, - "total": 40.535719510134385 + "total": 40.56652297688218 }, "2": { "components": { "united_atom:Transvibrational": 24.833896219788347, - "united_atom:Rovibrational": 16.996392060962357, + "united_atom:Rovibrational": 17.0032604785275, "united_atom:Conformational": 0.0, "united_atom:Orientational": 0.0 }, - "total": 41.83028828075071 + "total": 41.83715669831585 }, "3": { "components": { "united_atom:Transvibrational": 6.6723857079283215, - "united_atom:Rovibrational": 4.986376282798274, + "united_atom:Rovibrational": 4.958640584432368, "united_atom:Conformational": 0.0, "united_atom:Orientational": 0.0 }, - "total": 11.658761990726596 + "total": 11.63102629236069 }, "4": { "components": { - "united_atom:Transvibrational": 6.880461178899526, - "united_atom:Rovibrational": 4.441585994100298, + "united_atom:Transvibrational": 6.880461178899524, + "united_atom:Rovibrational": 4.47005523879151, "united_atom:Conformational": 0.0, "united_atom:Orientational": 0.0 }, - "total": 11.322047172999824 + "total": 11.350516417691033 }, "5": { "components": { "united_atom:Transvibrational": 15.473092273577388, - "united_atom:Rovibrational": 7.777083478334702, + "united_atom:Rovibrational": 7.7776511992146515, "united_atom:Conformational": 0.0, "united_atom:Orientational": 0.0 }, - "total": 23.25017575191209 + "total": 23.25074347279204 }, "6": { "components": { "united_atom:Transvibrational": 25.694168010993888, - "united_atom:Rovibrational": 6.731392483940261, + "united_atom:Rovibrational": 6.727950082765358, "united_atom:Conformational": 0.0, "united_atom:Orientational": 0.0 }, - "total": 32.425560494934146 + "total": 32.422118093759245 }, "7": { "components": { - "united_atom:Transvibrational": 10.533257828925997, - "united_atom:Rovibrational": 8.97563913301586, + "united_atom:Transvibrational": 10.533257828925999, + "united_atom:Rovibrational": 9.059251947808503, "united_atom:Conformational": 0.0, "united_atom:Orientational": 0.0 }, - "total": 19.508896961941858 + "total": 19.592509776734502 }, "8": { "components": { "united_atom:Transvibrational": 7.332978272264879, - "united_atom:Rovibrational": 3.0418502362906086, + "united_atom:Rovibrational": 3.084945359076361, "united_atom:Conformational": 0.0, "united_atom:Orientational": 0.0 }, - "total": 10.374828508555488 + "total": 10.417923631341239 }, "9": { "components": { "united_atom:Transvibrational": 4.023932761002867, - "united_atom:Rovibrational": 8.289324383636089, + "united_atom:Rovibrational": 8.271395156034108, "united_atom:Conformational": 0.0, "united_atom:Orientational": 0.0 }, - "total": 12.313257144638957 + "total": 12.295327917036975 } } } diff --git a/tests/regression/baselines/methane/selection_subset.json b/tests/regression/baselines/methane/selection_subset.json index 318c30c6..70be3cef 100644 --- a/tests/regression/baselines/methane/selection_subset.json +++ b/tests/regression/baselines/methane/selection_subset.json @@ -2,12 +2,12 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 42.51619525142699, - "united_atom:Rovibrational": 34.651932322640015, + "united_atom:Transvibrational": 42.516195251426986, + "united_atom:Rovibrational": 34.69255616696802, "united_atom:Conformational": 0.0, "united_atom:Orientational": 0.0 }, - "total": 77.168127574067 + "total": 77.208751418395 } } } diff --git a/tests/regression/baselines/octonol/combined_forcetorque_off.json b/tests/regression/baselines/octonol/combined_forcetorque_off.json index 920589a7..362cd1c2 100644 --- a/tests/regression/baselines/octonol/combined_forcetorque_off.json +++ b/tests/regression/baselines/octonol/combined_forcetorque_off.json @@ -2,7 +2,7 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 0.34382967782781193, + "united_atom:Transvibrational": 0.2907129992470817, "united_atom:Rovibrational": 19.154228264604278, "residue:Transvibrational": 74.72816183790984, "residue:Rovibrational": 60.84390144550801, @@ -10,7 +10,7 @@ "residue:Conformational": 0.0, "residue:Orientational": 25.79114991860739 }, - "total": 197.70076468713353 + "total": 197.64764800855278 } } } diff --git a/tests/regression/baselines/octonol/frame_window.json b/tests/regression/baselines/octonol/frame_window.json index 5c07dcc2..3fcdaef3 100644 --- a/tests/regression/baselines/octonol/frame_window.json +++ b/tests/regression/baselines/octonol/frame_window.json @@ -2,7 +2,7 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 65.86600728016141, + "united_atom:Transvibrational": 68.20033834493445, "united_atom:Rovibrational": 162.26463085986236, "residue:FTmat-Transvibrational": 74.84326467422125, "residue:FTmat-Rovibrational": 60.95710412746361, @@ -10,7 +10,7 @@ "residue:Conformational": 0.0, "residue:Orientational": 25.79114991860739 }, - "total": 406.37330124036134 + "total": 408.7076323051344 } } } diff --git a/tests/regression/baselines/octonol/selection_subset.json b/tests/regression/baselines/octonol/selection_subset.json index 5f9d18f1..e15b9555 100644 --- a/tests/regression/baselines/octonol/selection_subset.json +++ b/tests/regression/baselines/octonol/selection_subset.json @@ -2,7 +2,7 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 0.34382967782781193, + "united_atom:Transvibrational": 0.2907129992470817, "united_atom:Rovibrational": 19.154228264604278, "residue:FTmat-Transvibrational": 83.17889385114952, "residue:FTmat-Rovibrational": 57.480841206687685, @@ -10,7 +10,7 @@ "residue:Conformational": 0.0, "residue:Orientational": 25.79114991860739 }, - "total": 202.78843646155286 + "total": 202.73531978297217 } } } diff --git a/tests/regression/baselines/water/combined_forcetorque_off.json b/tests/regression/baselines/water/combined_forcetorque_off.json index 2e632568..b92ef4e3 100644 --- a/tests/regression/baselines/water/combined_forcetorque_off.json +++ b/tests/regression/baselines/water/combined_forcetorque_off.json @@ -2,12 +2,12 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 51.22145722965398, - "united_atom:Rovibrational": 17.232634718413053, + "united_atom:Transvibrational": 51.22145722965397, + "united_atom:Rovibrational": 17.307840626636313, "united_atom:Conformational": 0.0, "united_atom:Orientational": 0.0 }, - "total": 68.45409194806703 + "total": 68.52929785629028 } } } diff --git a/tests/regression/baselines/water/frame_window.json b/tests/regression/baselines/water/frame_window.json index c6ae550d..8b368026 100644 --- a/tests/regression/baselines/water/frame_window.json +++ b/tests/regression/baselines/water/frame_window.json @@ -3,11 +3,11 @@ "0": { "components": { "united_atom:Transvibrational": 46.684099714230925, - "united_atom:Rovibrational": 17.694014405444744, + "united_atom:Rovibrational": 17.66359738549709, "united_atom:Conformational": 0.0, "united_atom:Orientational": 0.0 }, - "total": 64.37811411967567 + "total": 64.34769709972801 } } } diff --git a/tests/regression/baselines/water/grouping_each.json b/tests/regression/baselines/water/grouping_each.json index acccd86b..676010dd 100644 --- a/tests/regression/baselines/water/grouping_each.json +++ b/tests/regression/baselines/water/grouping_each.json @@ -3,92 +3,92 @@ "0": { "components": { "united_atom:Transvibrational": 14.945907953805552, - "united_atom:Rovibrational": 1.8947243953648254, + "united_atom:Rovibrational": 1.94076792014105, "united_atom:Conformational": 0.0, "united_atom:Orientational": 0.0 }, - "total": 16.840632349170377 + "total": 16.8866758739466 }, "1": { "components": { - "united_atom:Transvibrational": 8.111522208287779, - "united_atom:Rovibrational": 0.9607747968334215, + "united_atom:Transvibrational": 8.11152220828778, + "united_atom:Rovibrational": 0.9798912337680555, "united_atom:Conformational": 0.0, "united_atom:Orientational": 0.0 }, - "total": 9.0722970051212 + "total": 9.091413442055837 }, "2": { "components": { - "united_atom:Transvibrational": 15.79497347968759, - "united_atom:Rovibrational": 4.606072530590255, + "united_atom:Transvibrational": 15.794973479687599, + "united_atom:Rovibrational": 4.4692860930919585, "united_atom:Conformational": 0.0, "united_atom:Orientational": 0.0 }, - "total": 20.401046010277845 + "total": 20.264259572779558 }, "3": { "components": { "united_atom:Transvibrational": 9.462487265119448, - "united_atom:Rovibrational": 1.6509513986984767, + "united_atom:Rovibrational": 1.6054615846390614, "united_atom:Conformational": 0.0, "united_atom:Orientational": 0.0 }, - "total": 11.113438663817924 + "total": 11.067948849758508 }, "4": { "components": { "united_atom:Transvibrational": 16.254161245740693, - "united_atom:Rovibrational": 1.367597374110767, + "united_atom:Rovibrational": 1.448700454979353, "united_atom:Conformational": 0.0, "united_atom:Orientational": 0.0 }, - "total": 17.62175861985146 + "total": 17.702861700720046 }, "5": { "components": { "united_atom:Transvibrational": 11.57779876203329, - "united_atom:Rovibrational": 7.351163568004953, + "united_atom:Rovibrational": 7.23849459157533, "united_atom:Conformational": 0.0, "united_atom:Orientational": 0.0 }, - "total": 18.928962330038242 + "total": 18.81629335360862 }, "6": { "components": { "united_atom:Transvibrational": 10.172061262201801, - "united_atom:Rovibrational": 3.671537296888502, + "united_atom:Rovibrational": 4.086982297438225, "united_atom:Conformational": 0.0, "united_atom:Orientational": 0.0 }, - "total": 13.843598559090303 + "total": 14.259043559640027 }, "7": { "components": { - "united_atom:Transvibrational": 13.49951045186194, - "united_atom:Rovibrational": 1.7029491215274863, + "united_atom:Transvibrational": 13.499510451861939, + "united_atom:Rovibrational": 1.6178142656646386, "united_atom:Conformational": 0.0, "united_atom:Orientational": 0.0 }, - "total": 15.202459573389428 + "total": 15.117324717526577 }, "8": { "components": { "united_atom:Transvibrational": 16.597759171924583, - "united_atom:Rovibrational": 1.6190289365617598, + "united_atom:Rovibrational": 1.6676178150275727, "united_atom:Conformational": 0.0, "united_atom:Orientational": 0.0 }, - "total": 18.21678810848634 + "total": 18.265376986952155 }, "9": { "components": { "united_atom:Transvibrational": 13.731375187064538, - "united_atom:Rovibrational": 5.650620990305761, + "united_atom:Rovibrational": 5.161816204002109, "united_atom:Conformational": 0.0, "united_atom:Orientational": 0.0 }, - "total": 19.381996177370297 + "total": 18.893191391066647 } } } diff --git a/tests/regression/baselines/water/selection_subset.json b/tests/regression/baselines/water/selection_subset.json index 2e632568..b92ef4e3 100644 --- a/tests/regression/baselines/water/selection_subset.json +++ b/tests/regression/baselines/water/selection_subset.json @@ -2,12 +2,12 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 51.22145722965398, - "united_atom:Rovibrational": 17.232634718413053, + "united_atom:Transvibrational": 51.22145722965397, + "united_atom:Rovibrational": 17.307840626636313, "united_atom:Conformational": 0.0, "united_atom:Orientational": 0.0 }, - "total": 68.45409194806703 + "total": 68.52929785629028 } } } diff --git a/tests/regression/baselines/water/water_off.json b/tests/regression/baselines/water/water_off.json index 2e632568..b92ef4e3 100644 --- a/tests/regression/baselines/water/water_off.json +++ b/tests/regression/baselines/water/water_off.json @@ -2,12 +2,12 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 51.22145722965398, - "united_atom:Rovibrational": 17.232634718413053, + "united_atom:Transvibrational": 51.22145722965397, + "united_atom:Rovibrational": 17.307840626636313, "united_atom:Conformational": 0.0, "united_atom:Orientational": 0.0 }, - "total": 68.45409194806703 + "total": 68.52929785629028 } } } From 950e411a6f2599d8cd9076a63b2aa41e0afcff1b Mon Sep 17 00:00:00 2001 From: ioana Date: Mon, 22 Jun 2026 14:39:06 +0100 Subject: [PATCH 38/47] updated baselines --- tests/regression/baselines/benzaldehyde/default.json | 4 ++-- tests/regression/baselines/benzaldehyde/rad.json | 4 ++-- tests/regression/baselines/benzene/default.json | 4 ++-- tests/regression/baselines/benzene/rad.json | 4 ++-- tests/regression/baselines/cyclohexane/default.json | 4 ++-- tests/regression/baselines/cyclohexane/rad.json | 4 ++-- tests/regression/baselines/dna/default.json | 6 +++--- tests/regression/baselines/ethyl-acetate/default.json | 4 ++-- tests/regression/baselines/ethyl-acetate/rad.json | 4 ++-- tests/regression/baselines/methane/default.json | 6 +++--- tests/regression/baselines/methane/rad.json | 6 +++--- tests/regression/baselines/octonol/default.json | 4 ++-- tests/regression/baselines/octonol/rad.json | 4 ++-- tests/regression/baselines/water/default.json | 4 ++-- tests/regression/baselines/water/rad.json | 6 +++--- 15 files changed, 34 insertions(+), 34 deletions(-) diff --git a/tests/regression/baselines/benzaldehyde/default.json b/tests/regression/baselines/benzaldehyde/default.json index c99a2d15..1ed1a4e7 100644 --- a/tests/regression/baselines/benzaldehyde/default.json +++ b/tests/regression/baselines/benzaldehyde/default.json @@ -2,7 +2,7 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 24.88518233240474, + "united_atom:Transvibrational": 24.869562693152986, "united_atom:Rovibrational": 27.950376507672583, "residue:FTmat-Transvibrational": 71.03412922724692, "residue:FTmat-Rovibrational": 59.44169664956799, @@ -10,7 +10,7 @@ "residue:Conformational": 0.0, "residue:Orientational": 59.43920971558428 }, - "total": 250.27323833217878 + "total": 250.25761869292705 } } } diff --git a/tests/regression/baselines/benzaldehyde/rad.json b/tests/regression/baselines/benzaldehyde/rad.json index 8e39e21b..02d56d59 100644 --- a/tests/regression/baselines/benzaldehyde/rad.json +++ b/tests/regression/baselines/benzaldehyde/rad.json @@ -2,7 +2,7 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 24.88518233240474, + "united_atom:Transvibrational": 24.869562693152986, "united_atom:Rovibrational": 27.950376507672583, "residue:FTmat-Transvibrational": 71.03412922724692, "residue:FTmat-Rovibrational": 59.44169664956799, @@ -10,7 +10,7 @@ "residue:Conformational": 0.0, "residue:Orientational": 25.502722133228936 }, - "total": 216.33675074982344 + "total": 216.32113111057168 } } } diff --git a/tests/regression/baselines/benzene/default.json b/tests/regression/baselines/benzene/default.json index 86ae25d0..de6eebc0 100644 --- a/tests/regression/baselines/benzene/default.json +++ b/tests/regression/baselines/benzene/default.json @@ -2,7 +2,7 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 9.081853431559052, + "united_atom:Transvibrational": 9.050489789631982, "united_atom:Rovibrational": 27.534380126250227, "residue:FTmat-Transvibrational": 72.66211800013413, "residue:FTmat-Rovibrational": 59.93761874375924, @@ -10,7 +10,7 @@ "residue:Conformational": 0.0, "residue:Orientational": 41.58427729275938 }, - "total": 210.800247594462 + "total": 210.76888395253496 } } } diff --git a/tests/regression/baselines/benzene/rad.json b/tests/regression/baselines/benzene/rad.json index c4aca55f..2ccf3306 100644 --- a/tests/regression/baselines/benzene/rad.json +++ b/tests/regression/baselines/benzene/rad.json @@ -2,7 +2,7 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 9.081853431559052, + "united_atom:Transvibrational": 9.050489789631982, "united_atom:Rovibrational": 27.534380126250227, "residue:FTmat-Transvibrational": 72.66211800013413, "residue:FTmat-Rovibrational": 59.93761874375924, @@ -10,7 +10,7 @@ "residue:Conformational": 0.0, "residue:Orientational": 12.386543820026672 }, - "total": 181.6025141217293 + "total": 181.57115047980224 } } } diff --git a/tests/regression/baselines/cyclohexane/default.json b/tests/regression/baselines/cyclohexane/default.json index c0233178..e7867e50 100644 --- a/tests/regression/baselines/cyclohexane/default.json +++ b/tests/regression/baselines/cyclohexane/default.json @@ -2,7 +2,7 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 11.726471858700231, + "united_atom:Transvibrational": 11.65610486557721, "united_atom:Rovibrational": 47.71088602161552, "residue:FTmat-Transvibrational": 70.3922327806972, "residue:FTmat-Rovibrational": 63.44920824648768, @@ -10,7 +10,7 @@ "residue:Conformational": 0.0, "residue:Orientational": 47.54362774624537 }, - "total": 242.11130586044152 + "total": 242.0409388673185 } } } diff --git a/tests/regression/baselines/cyclohexane/rad.json b/tests/regression/baselines/cyclohexane/rad.json index f65d9a4c..8de06e5b 100644 --- a/tests/regression/baselines/cyclohexane/rad.json +++ b/tests/regression/baselines/cyclohexane/rad.json @@ -2,7 +2,7 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 11.726471858700231, + "united_atom:Transvibrational": 11.65610486557721, "united_atom:Rovibrational": 47.71088602161552, "residue:FTmat-Transvibrational": 70.3922327806972, "residue:FTmat-Rovibrational": 63.44920824648768, @@ -10,7 +10,7 @@ "residue:Conformational": 0.0, "residue:Orientational": 14.049324060636614 }, - "total": 208.61700217483278 + "total": 208.54663518170975 } } } diff --git a/tests/regression/baselines/dna/default.json b/tests/regression/baselines/dna/default.json index 760c664d..7a2b6c3a 100644 --- a/tests/regression/baselines/dna/default.json +++ b/tests/regression/baselines/dna/default.json @@ -2,17 +2,17 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 0.0, + "united_atom:Transvibrational": 4.762281610623415e-20, "united_atom:Rovibrational": 0.002160679012128457, "residue:Transvibrational": 0.0, - "residue:Rovibrational": 6.832779629985204, + "residue:Rovibrational": 6.633599673254765, "polymer:FTmat-Transvibrational": 12.341104347192612, "polymer:FTmat-Rovibrational": 0.0, "united_atom:Conformational": 0.0, "residue:Conformational": 0.0, "polymer:Orientational": 4.758905336627712 }, - "total": 23.934949992817657 + "total": 23.735770036087217 }, "1": { "components": { diff --git a/tests/regression/baselines/ethyl-acetate/default.json b/tests/regression/baselines/ethyl-acetate/default.json index 077cdbe3..a98d4af8 100644 --- a/tests/regression/baselines/ethyl-acetate/default.json +++ b/tests/regression/baselines/ethyl-acetate/default.json @@ -2,7 +2,7 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 19.703248485425714, + "united_atom:Transvibrational": 19.67356252280246, "united_atom:Rovibrational": 49.35007067210558, "residue:FTmat-Transvibrational": 67.91350627765567, "residue:FTmat-Rovibrational": 60.042233035077246, @@ -10,7 +10,7 @@ "residue:Conformational": 0.0, "residue:Orientational": 65.33877737416202 }, - "total": 270.52337372549465 + "total": 270.4936877628714 } } } diff --git a/tests/regression/baselines/ethyl-acetate/rad.json b/tests/regression/baselines/ethyl-acetate/rad.json index ee7df624..51bed453 100644 --- a/tests/regression/baselines/ethyl-acetate/rad.json +++ b/tests/regression/baselines/ethyl-acetate/rad.json @@ -2,7 +2,7 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 19.703248485425714, + "united_atom:Transvibrational": 19.67356252280246, "united_atom:Rovibrational": 49.35007067210558, "residue:FTmat-Transvibrational": 67.91350627765567, "residue:FTmat-Rovibrational": 60.042233035077246, @@ -10,7 +10,7 @@ "residue:Conformational": 0.0, "residue:Orientational": 30.739736508673403 }, - "total": 235.92433286000602 + "total": 235.89464689738276 } } } diff --git a/tests/regression/baselines/methane/default.json b/tests/regression/baselines/methane/default.json index d08ac60f..c6716b00 100644 --- a/tests/regression/baselines/methane/default.json +++ b/tests/regression/baselines/methane/default.json @@ -2,12 +2,12 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 40.3239711717637, - "united_atom:Rovibrational": 33.60582165153992, + "united_atom:Transvibrational": 40.32397117176369, + "united_atom:Rovibrational": 33.617646190915714, "united_atom:Conformational": 0.0, "united_atom:Orientational": 3.596676624528629 }, - "total": 77.52646944783224 + "total": 77.53829398720804 } } } diff --git a/tests/regression/baselines/methane/rad.json b/tests/regression/baselines/methane/rad.json index df8fbbd7..2b506858 100644 --- a/tests/regression/baselines/methane/rad.json +++ b/tests/regression/baselines/methane/rad.json @@ -2,12 +2,12 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 40.3239711717637, - "united_atom:Rovibrational": 33.60582165153992, + "united_atom:Transvibrational": 40.32397117176369, + "united_atom:Rovibrational": 33.617646190915714, "united_atom:Conformational": 0.0, "united_atom:Orientational": 6.667209510980459 }, - "total": 80.59700233428407 + "total": 80.60882687365986 } } } diff --git a/tests/regression/baselines/octonol/default.json b/tests/regression/baselines/octonol/default.json index fdd22382..76b4b259 100644 --- a/tests/regression/baselines/octonol/default.json +++ b/tests/regression/baselines/octonol/default.json @@ -2,7 +2,7 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 37.860034170693694, + "united_atom:Transvibrational": 37.90014130683313, "united_atom:Rovibrational": 84.2608785308744, "residue:FTmat-Transvibrational": 66.13571365167344, "residue:FTmat-Rovibrational": 57.090651827515686, @@ -10,7 +10,7 @@ "residue:Conformational": 0.0, "residue:Orientational": 74.78356386019003 }, - "total": 354.36024401601327 + "total": 354.4003511521527 } } } diff --git a/tests/regression/baselines/octonol/rad.json b/tests/regression/baselines/octonol/rad.json index c2c20839..1d294585 100644 --- a/tests/regression/baselines/octonol/rad.json +++ b/tests/regression/baselines/octonol/rad.json @@ -2,7 +2,7 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 37.860034170693694, + "united_atom:Transvibrational": 37.90014130683313, "united_atom:Rovibrational": 84.2608785308744, "residue:FTmat-Transvibrational": 66.13571365167344, "residue:FTmat-Rovibrational": 57.090651827515686, @@ -10,7 +10,7 @@ "residue:Conformational": 0.0, "residue:Orientational": 28.084239413273323 }, - "total": 307.66091956909656 + "total": 307.701026705236 } } } diff --git a/tests/regression/baselines/water/default.json b/tests/regression/baselines/water/default.json index c98bc3eb..7f5baa09 100644 --- a/tests/regression/baselines/water/default.json +++ b/tests/regression/baselines/water/default.json @@ -3,11 +3,11 @@ "0": { "components": { "united_atom:Transvibrational": 43.725393337304425, - "united_atom:Rovibrational": 17.28911070147887, + "united_atom:Rovibrational": 17.42448032148262, "united_atom:Conformational": 0.0, "united_atom:Orientational": 32.193445593362114 }, - "total": 93.20794963214541 + "total": 93.34331925214916 } } } diff --git a/tests/regression/baselines/water/rad.json b/tests/regression/baselines/water/rad.json index cbda37a9..cf8510e3 100644 --- a/tests/regression/baselines/water/rad.json +++ b/tests/regression/baselines/water/rad.json @@ -2,12 +2,12 @@ "groups": { "0": { "components": { - "united_atom:Transvibrational": 79.20298312418278, - "united_atom:Rovibrational": 50.90260688502127, + "united_atom:Transvibrational": 79.2029831241828, + "united_atom:Rovibrational": 51.07316520629525, "united_atom:Conformational": 0.0, "united_atom:Orientational": 22.760377489372107 }, - "total": 152.86596749857614 + "total": 153.03652581985014 } } } From d8be74d7c29c86f9da4c0636e0759a16de5b214d Mon Sep 17 00:00:00 2001 From: ioana Date: Mon, 22 Jun 2026 14:41:19 +0100 Subject: [PATCH 39/47] will add unit tests for new lines of code --- tests/unit/CodeEntropy/levels/test_axes.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/unit/CodeEntropy/levels/test_axes.py b/tests/unit/CodeEntropy/levels/test_axes.py index 4c4a894a..aec1ff99 100644 --- a/tests/unit/CodeEntropy/levels/test_axes.py +++ b/tests/unit/CodeEntropy/levels/test_axes.py @@ -716,3 +716,22 @@ def test_get_bonded_axes_returns_none_none_if_custom_axes_none(monkeypatch): assert custom_axes is None assert moi is None + + +""" +def test_get_chain(monkeypatch): + ax = AxesCalculator() + + atom = _FakeAtom(index=7, mass=12.0, position=[0, 0, 0]) + system = MagicMock() + dimensions = np.array([10.0, 10.0, 10.0], dtype=float) + + + monkeypatch.setattr( + ax, "get_chain", lambda _idx, _sys: (first,middle,last) + ) + + + assert + assert +""" From 166937f4e2e431f7b284386922b831bbd49e8f45 Mon Sep 17 00:00:00 2001 From: ioanaapapa Date: Mon, 22 Jun 2026 15:06:53 +0100 Subject: [PATCH 40/47] updated tests --- tests/unit/CodeEntropy/levels/nodes/test_covariance_node.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unit/CodeEntropy/levels/nodes/test_covariance_node.py b/tests/unit/CodeEntropy/levels/nodes/test_covariance_node.py index 8398113c..132bf8e7 100644 --- a/tests/unit/CodeEntropy/levels/nodes/test_covariance_node.py +++ b/tests/unit/CodeEntropy/levels/nodes/test_covariance_node.py @@ -415,7 +415,7 @@ def test_build_ua_vectors_uses_customised_axes(): force_vecs, torque_vecs = node._build_ua_vectors( bead_groups=[FakeAtomGroup("ua")], - residue_atoms=FakeAtomGroup("res"), + residue_group=FakeAtomGroup("res"), axes_manager=axes_manager, box=None, force_partitioning=0.5, @@ -441,7 +441,7 @@ def test_build_ua_vectors_uses_vanilla_axes_when_not_customised(): with patch("CodeEntropy.levels.nodes.covariance.make_whole") as make_whole: node._build_ua_vectors( bead_groups=[FakeAtomGroup("ua")], - residue_atoms=FakeAtomGroup("res"), + residue_group=FakeAtomGroup("res"), axes_manager=axes_manager, box=None, force_partitioning=0.5, From c7eac7dc936e4471ef5d58b13453b0b2ba4279ad Mon Sep 17 00:00:00 2001 From: ioanaapapa Date: Mon, 22 Jun 2026 15:30:57 +0100 Subject: [PATCH 41/47] updated unit tests --- .../levels/nodes/test_covariance_node.py | 3 +++ tests/unit/CodeEntropy/levels/test_axes.py | 19 ------------------- 2 files changed, 3 insertions(+), 19 deletions(-) diff --git a/tests/unit/CodeEntropy/levels/nodes/test_covariance_node.py b/tests/unit/CodeEntropy/levels/nodes/test_covariance_node.py index 132bf8e7..94360764 100644 --- a/tests/unit/CodeEntropy/levels/nodes/test_covariance_node.py +++ b/tests/unit/CodeEntropy/levels/nodes/test_covariance_node.py @@ -421,6 +421,7 @@ def test_build_ua_vectors_uses_customised_axes(): force_partitioning=0.5, customised_axes=True, is_highest=True, + res_position=None, ) assert len(force_vecs) == 1 @@ -437,6 +438,7 @@ def test_build_ua_vectors_uses_vanilla_axes_when_not_customised(): ) node._ft.get_weighted_forces = MagicMock(return_value=np.array([1.0, 0.0, 0.0])) node._ft.get_weighted_torques = MagicMock(return_value=np.array([0.0, 1.0, 0.0])) + FakeAtomGroup.atoms = FakeAtomGroup with patch("CodeEntropy.levels.nodes.covariance.make_whole") as make_whole: node._build_ua_vectors( @@ -447,6 +449,7 @@ def test_build_ua_vectors_uses_vanilla_axes_when_not_customised(): force_partitioning=0.5, customised_axes=False, is_highest=False, + res_position=None, ) assert make_whole.call_count == 2 diff --git a/tests/unit/CodeEntropy/levels/test_axes.py b/tests/unit/CodeEntropy/levels/test_axes.py index aec1ff99..4c4a894a 100644 --- a/tests/unit/CodeEntropy/levels/test_axes.py +++ b/tests/unit/CodeEntropy/levels/test_axes.py @@ -716,22 +716,3 @@ def test_get_bonded_axes_returns_none_none_if_custom_axes_none(monkeypatch): assert custom_axes is None assert moi is None - - -""" -def test_get_chain(monkeypatch): - ax = AxesCalculator() - - atom = _FakeAtom(index=7, mass=12.0, position=[0, 0, 0]) - system = MagicMock() - dimensions = np.array([10.0, 10.0, 10.0], dtype=float) - - - monkeypatch.setattr( - ax, "get_chain", lambda _idx, _sys: (first,middle,last) - ) - - - assert - assert -""" From 9c1ec28fc47d3197c730260e8635adef2870c8c5 Mon Sep 17 00:00:00 2001 From: ioanaapapa Date: Mon, 22 Jun 2026 15:34:30 +0100 Subject: [PATCH 42/47] update covariance unit tests --- tests/unit/CodeEntropy/levels/nodes/test_covariance_node.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/unit/CodeEntropy/levels/nodes/test_covariance_node.py b/tests/unit/CodeEntropy/levels/nodes/test_covariance_node.py index 94360764..d27892a3 100644 --- a/tests/unit/CodeEntropy/levels/nodes/test_covariance_node.py +++ b/tests/unit/CodeEntropy/levels/nodes/test_covariance_node.py @@ -439,6 +439,7 @@ def test_build_ua_vectors_uses_vanilla_axes_when_not_customised(): node._ft.get_weighted_forces = MagicMock(return_value=np.array([1.0, 0.0, 0.0])) node._ft.get_weighted_torques = MagicMock(return_value=np.array([0.0, 1.0, 0.0])) FakeAtomGroup.atoms = FakeAtomGroup + FakeAtomGroup.atoms.principal_axes.return_value = np.eye(3) with patch("CodeEntropy.levels.nodes.covariance.make_whole") as make_whole: node._build_ua_vectors( From 5987746658908705c4c2c6464b65c2f4551dcbee Mon Sep 17 00:00:00 2001 From: ioanaapapa Date: Mon, 22 Jun 2026 16:51:03 +0100 Subject: [PATCH 43/47] fix test for building ua vectors with vanilla axes --- tests/unit/CodeEntropy/levels/nodes/test_covariance_node.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/unit/CodeEntropy/levels/nodes/test_covariance_node.py b/tests/unit/CodeEntropy/levels/nodes/test_covariance_node.py index d27892a3..d989069f 100644 --- a/tests/unit/CodeEntropy/levels/nodes/test_covariance_node.py +++ b/tests/unit/CodeEntropy/levels/nodes/test_covariance_node.py @@ -436,10 +436,11 @@ def test_build_ua_vectors_uses_vanilla_axes_when_not_customised(): np.eye(3), np.array([1.0, 2.0, 3.0]), ) + FakeAtomGroup_Atoms = MagicMock(positions=np.zeros((2, 3))) + FakeAtomGroup.atoms = FakeAtomGroup_Atoms + axes_manager.principal_axes.return_value = np.eye(3) node._ft.get_weighted_forces = MagicMock(return_value=np.array([1.0, 0.0, 0.0])) node._ft.get_weighted_torques = MagicMock(return_value=np.array([0.0, 1.0, 0.0])) - FakeAtomGroup.atoms = FakeAtomGroup - FakeAtomGroup.atoms.principal_axes.return_value = np.eye(3) with patch("CodeEntropy.levels.nodes.covariance.make_whole") as make_whole: node._build_ua_vectors( From 1e6cd563c1a42d1a81a5617311cf35665c8f362d Mon Sep 17 00:00:00 2001 From: ioanaapapa Date: Mon, 29 Jun 2026 16:04:35 +0100 Subject: [PATCH 44/47] updated documentation --- docs/science.rst | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/science.rst b/docs/science.rst index 1a98f977..16095e3e 100644 --- a/docs/science.rst +++ b/docs/science.rst @@ -72,9 +72,9 @@ For the polymer level, the translational and rotational axes are defined as the For the residue level, there are two situations. When the residue is not bonded to any other residues, the translational and rotational axes are the principal axes of the molecule. -When the residue is part of a larger polymer, the translational axes are the principal axes of the polymer, and the rotational axes are defined from the average position of the bonds to neighbouring residues. +When the residue is part of a larger polymer, the translational axes are the principal axes of the polymer, and the rotational axes are defined from the positions of edge atoms in a residue (i.e. heavy atoms bonded to atoms in a neighbouring residue) and average position of all inner backbone (defined as the shortest path between the edge atoms, not including them) heavy atoms. -For the united atom level, the translational axes are defined as the principal axes of the residue and the rotational axes are defined from the average position of the bonds to neighbouring heavy atoms. +For the united atom level, the translational axes are the same as the rotational axes at the residue level and the rotational axes are defined from the average position of the bonds to neighbouring heavy atoms. If there are no bonds to other heavy atoms, the principal axes of the molecule are used. Conformational Entropy From 3c4896ebe9f6dfd89f19a1526b75d3a06eda01fe Mon Sep 17 00:00:00 2001 From: ioanaapapa Date: Thu, 2 Jul 2026 10:36:15 +0100 Subject: [PATCH 45/47] Revert "updated documentation" This reverts commit 1e6cd563c1a42d1a81a5617311cf35665c8f362d. --- docs/science.rst | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/science.rst b/docs/science.rst index 16095e3e..1a98f977 100644 --- a/docs/science.rst +++ b/docs/science.rst @@ -72,9 +72,9 @@ For the polymer level, the translational and rotational axes are defined as the For the residue level, there are two situations. When the residue is not bonded to any other residues, the translational and rotational axes are the principal axes of the molecule. -When the residue is part of a larger polymer, the translational axes are the principal axes of the polymer, and the rotational axes are defined from the positions of edge atoms in a residue (i.e. heavy atoms bonded to atoms in a neighbouring residue) and average position of all inner backbone (defined as the shortest path between the edge atoms, not including them) heavy atoms. +When the residue is part of a larger polymer, the translational axes are the principal axes of the polymer, and the rotational axes are defined from the average position of the bonds to neighbouring residues. -For the united atom level, the translational axes are the same as the rotational axes at the residue level and the rotational axes are defined from the average position of the bonds to neighbouring heavy atoms. +For the united atom level, the translational axes are defined as the principal axes of the residue and the rotational axes are defined from the average position of the bonds to neighbouring heavy atoms. If there are no bonds to other heavy atoms, the principal axes of the molecule are used. Conformational Entropy From d2aaf9bd15dddd22244f0767b6ad05ae8321eacb Mon Sep 17 00:00:00 2001 From: ioanaapapa Date: Thu, 2 Jul 2026 10:52:59 +0100 Subject: [PATCH 46/47] Revert "Merge branch '26-residue-bonded-axes' of https://github.com/CCPBioSim/CodeEntropy into 26-residue-bonded-axes" This reverts commit 7b91b08609fecbacf3508fe16ebdcea3b7490c1a, reversing changes made to 3c4896ebe9f6dfd89f19a1526b75d3a06eda01fe. --- .github/workflows/coveralls-main.yaml | 54 -- .github/workflows/pr.yaml | 23 +- CITATION.cff | 4 +- CodeEntropy/__init__.py | 2 +- CodeEntropy/entropy/workflow.py | 2 +- CodeEntropy/levels/axes.py | 155 ------ .../levels/{graph => }/conformation_dag.py | 0 CodeEntropy/levels/dihedrals/topology.py | 123 +--- CodeEntropy/levels/execution/__init__.py | 0 CodeEntropy/levels/execution/scheduler.py | 2 +- CodeEntropy/levels/execution/tasks.py | 2 +- CodeEntropy/levels/{graph => }/frame_dag.py | 0 CodeEntropy/levels/graph/__init__.py | 0 CodeEntropy/levels/{graph => }/level_dag.py | 10 +- CodeEntropy/levels/nodes/axes_topology.py | 316 ----------- CodeEntropy/levels/nodes/covariance.py | 86 +-- CodeEntropy/trajectory/mda.py | 22 +- docs/api/CodeEntropy.config.rst | 13 +- docs/api/CodeEntropy.core.rst | 13 +- docs/api/CodeEntropy.entropy.nodes.rst | 13 +- docs/api/CodeEntropy.entropy.rst | 13 +- ...> CodeEntropy.levels.conformation_dag.rst} | 4 +- ...py.levels.dihedrals.angle_observations.rst | 7 - ...dihedrals.conformational_state_builder.rst | 7 - ...ntropy.levels.dihedrals.peak_detection.rst | 7 - docs/api/CodeEntropy.levels.dihedrals.rst | 20 +- ...ropy.levels.dihedrals.state_assignment.rst | 7 - .../CodeEntropy.levels.dihedrals.topology.rst | 7 - .../CodeEntropy.levels.execution.chunks.rst | 7 - .../CodeEntropy.levels.execution.policy.rst | 7 - .../CodeEntropy.levels.execution.reducers.rst | 7 - docs/api/CodeEntropy.levels.execution.rst | 22 - ...CodeEntropy.levels.execution.scheduler.rst | 7 - .../CodeEntropy.levels.execution.tasks.rst | 7 - docs/api/CodeEntropy.levels.frame_dag.rst | 7 + ...eEntropy.levels.graph.conformation_dag.rst | 7 - .../CodeEntropy.levels.graph.frame_dag.rst | 7 - .../CodeEntropy.levels.graph.level_dag.rst | 7 - docs/api/CodeEntropy.levels.graph.rst | 20 - docs/api/CodeEntropy.levels.level_dag.rst | 7 + ...CodeEntropy.levels.nodes.axes_topology.rst | 7 - docs/api/CodeEntropy.levels.nodes.rst | 14 +- docs/api/CodeEntropy.levels.rst | 20 +- docs/api/CodeEntropy.molecules.rst | 13 +- docs/api/CodeEntropy.results.rst | 13 +- docs/api/CodeEntropy.rst | 13 +- docs/api/CodeEntropy.trajectory.rst | 13 +- docs/getting_started.rst | 4 +- tests/regression/helpers.py | 2 +- .../levels/dihedrals/test_topology.py | 526 ++++-------------- .../levels/nodes/test_axes_topology.py | 415 -------------- .../levels/nodes/test_covariance_node.py | 104 +--- tests/unit/CodeEntropy/levels/test_axes.py | 497 ----------------- .../levels/test_conformation_dag.py | 2 +- .../CodeEntropy/levels/test_frame_graph.py | 2 +- .../unit/CodeEntropy/levels/test_level_dag.py | 33 +- .../levels/test_mda_universe_operations.py | 19 +- 57 files changed, 243 insertions(+), 2478 deletions(-) delete mode 100644 .github/workflows/coveralls-main.yaml rename CodeEntropy/levels/{graph => }/conformation_dag.py (100%) delete mode 100644 CodeEntropy/levels/execution/__init__.py rename CodeEntropy/levels/{graph => }/frame_dag.py (100%) delete mode 100644 CodeEntropy/levels/graph/__init__.py rename CodeEntropy/levels/{graph => }/level_dag.py (95%) delete mode 100644 CodeEntropy/levels/nodes/axes_topology.py rename docs/api/{CodeEntropy.levels.dihedrals.kernels.rst => CodeEntropy.levels.conformation_dag.rst} (50%) delete mode 100644 docs/api/CodeEntropy.levels.dihedrals.angle_observations.rst delete mode 100644 docs/api/CodeEntropy.levels.dihedrals.conformational_state_builder.rst delete mode 100644 docs/api/CodeEntropy.levels.dihedrals.peak_detection.rst delete mode 100644 docs/api/CodeEntropy.levels.dihedrals.state_assignment.rst delete mode 100644 docs/api/CodeEntropy.levels.dihedrals.topology.rst delete mode 100644 docs/api/CodeEntropy.levels.execution.chunks.rst delete mode 100644 docs/api/CodeEntropy.levels.execution.policy.rst delete mode 100644 docs/api/CodeEntropy.levels.execution.reducers.rst delete mode 100644 docs/api/CodeEntropy.levels.execution.rst delete mode 100644 docs/api/CodeEntropy.levels.execution.scheduler.rst delete mode 100644 docs/api/CodeEntropy.levels.execution.tasks.rst create mode 100644 docs/api/CodeEntropy.levels.frame_dag.rst delete mode 100644 docs/api/CodeEntropy.levels.graph.conformation_dag.rst delete mode 100644 docs/api/CodeEntropy.levels.graph.frame_dag.rst delete mode 100644 docs/api/CodeEntropy.levels.graph.level_dag.rst delete mode 100644 docs/api/CodeEntropy.levels.graph.rst create mode 100644 docs/api/CodeEntropy.levels.level_dag.rst delete mode 100644 docs/api/CodeEntropy.levels.nodes.axes_topology.rst delete mode 100644 tests/unit/CodeEntropy/levels/nodes/test_axes_topology.py diff --git a/.github/workflows/coveralls-main.yaml b/.github/workflows/coveralls-main.yaml deleted file mode 100644 index 6537127d..00000000 --- a/.github/workflows/coveralls-main.yaml +++ /dev/null @@ -1,54 +0,0 @@ -name: Coveralls Main - -on: - push: - branches: [main] - - workflow_dispatch: - -permissions: - contents: read - -concurrency: - group: coveralls-main-${{ github.ref }} - cancel-in-progress: true - -jobs: - coveralls: - name: Build main coverage baseline - runs-on: ubuntu-24.04 - timeout-minutes: 30 - - permissions: - contents: read - checks: write - statuses: write - - steps: - - name: Checkout - uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 - - - name: Set up Python 3.14 - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 - with: - python-version: "3.14" - - - name: Install testing dependencies - run: | - python -m pip install --upgrade pip - python -m pip install -e .[testing] - - - name: Run coverage - run: | - python -m pytest tests/unit \ - --cov CodeEntropy \ - --cov-report term-missing \ - --cov-report xml \ - -q - - - name: Upload coverage to Coveralls - uses: coverallsapp/github-action@5cbfd81b66ca5d10c19b062c04de0199c215fb6e - with: - github-token: ${{ secrets.GITHUB_TOKEN }} - file: coverage.xml - fail-on-error: false diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 795348dc..e2269cbe 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -2,17 +2,22 @@ name: Pull Request CI on: pull_request: + push: + branches: [main] permissions: contents: read + checks: write + pull-requests: write concurrency: - group: pr-ci-${{ github.ref }} + group: pr-${{ github.ref }} cancel-in-progress: true jobs: unit: name: Unit + if: github.event_name == 'pull_request' runs-on: ${{ matrix.os }} timeout-minutes: 25 strategy: @@ -35,11 +40,12 @@ jobs: python -m pip install --upgrade pip python -m pip install -e .[testing] - - name: Pytest unit • ${{ matrix.os }}, ${{ matrix.python-version }} + - name: Pytest (unit) • ${{ matrix.os }}, ${{ matrix.python-version }} run: python -m pytest tests/unit discover-systems: name: Discover regression systems + if: github.event_name == 'pull_request' runs-on: ubuntu-24.04 outputs: systems: ${{ steps.set-systems.outputs.systems }} @@ -62,10 +68,11 @@ jobs: id: set-systems run: | SYSTEMS=$(python -m tests.regression.list_systems) - echo "systems=$SYSTEMS" >> "$GITHUB_OUTPUT" + echo "systems=$SYSTEMS" >> $GITHUB_OUTPUT regression-quick: - name: Regression fast • ${{ matrix.system }} + name: Regression (fast) • ${{ matrix.system }} + if: github.event_name == 'pull_request' needs: discover-systems runs-on: ubuntu-24.04 timeout-minutes: 35 @@ -110,6 +117,7 @@ jobs: docs: name: Docs + if: github.event_name == 'pull_request' needs: unit runs-on: ubuntu-24.04 timeout-minutes: 25 @@ -142,6 +150,7 @@ jobs: pre-commit: name: Pre-commit + if: github.event_name == 'pull_request' needs: unit runs-on: ubuntu-24.04 timeout-minutes: 15 @@ -175,12 +184,6 @@ jobs: runs-on: ubuntu-24.04 timeout-minutes: 30 - permissions: - contents: read - checks: write - statuses: write - pull-requests: write - steps: - name: Checkout uses: actions/checkout@8e8c483db84b4bee98b60c0593521ed34d9990e8 diff --git a/CITATION.cff b/CITATION.cff index b98032b2..e4f657e8 100644 --- a/CITATION.cff +++ b/CITATION.cff @@ -76,5 +76,5 @@ keywords: - biomolecular simulations - protein flexibility license: MIT -version: 2.3.0 -date-released: '2026-06-26' +version: 2.2.4 +date-released: '2026-05-15' diff --git a/CodeEntropy/__init__.py b/CodeEntropy/__init__.py index bef57d77..b2301f7f 100644 --- a/CodeEntropy/__init__.py +++ b/CodeEntropy/__init__.py @@ -8,4 +8,4 @@ and statistical mechanics. """ -__version__ = "2.3.0" +__version__ = "2.2.4" diff --git a/CodeEntropy/entropy/workflow.py b/CodeEntropy/entropy/workflow.py index 9d5298e2..f8ddd0a8 100644 --- a/CodeEntropy/entropy/workflow.py +++ b/CodeEntropy/entropy/workflow.py @@ -26,8 +26,8 @@ from CodeEntropy.core.logging import LoggingConfig from CodeEntropy.entropy.graph import EntropyGraph from CodeEntropy.entropy.water import WaterEntropy -from CodeEntropy.levels.graph.level_dag import LevelDAG from CodeEntropy.levels.hierarchy import HierarchyBuilder +from CodeEntropy.levels.level_dag import LevelDAG from CodeEntropy.trajectory.frames import FrameSelection from CodeEntropy.trajectory.source import FrameSource diff --git a/CodeEntropy/levels/axes.py b/CodeEntropy/levels/axes.py index 87f27317..f1c8bb75 100644 --- a/CodeEntropy/levels/axes.py +++ b/CodeEntropy/levels/axes.py @@ -390,161 +390,6 @@ def get_residue_custom_axes(self, edges, center): rot_center = E1O_vector + edges[0].position return rot_center, rot_axes - def get_UA_axes_from_topology( - self, - *, - u, - residue_atoms, - topology, - box: np.ndarray | None, - ): - """Compute UA axes using cached static topology. - - This is the cached-index equivalent of ``get_UA_axes``. It preserves the - frame-dependent numerical calculations, but avoids repeated MDAnalysis - selection strings for heavy atoms, bonded atoms, and UA masses. - - Args: - u: Current-frame universe. - residue_atoms: AtomGroup for the parent residue in the current frame. - topology: Cached ``UAAxesTopology`` for this UA bead. - box: Current periodic box lengths. If omitted, ``u.dimensions`` is used. - - Returns: - Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: - - trans_axes: Translational axes, shape ``(3, 3)``. - - rot_axes: Rotational axes, shape ``(3, 3)``. - - center: Rotation centre, shape ``(3,)``. - - moment_of_inertia: Principal moments, shape ``(3,)``. - - Raises: - ValueError: If cached bonded-axis construction fails. - """ - dimensions = ( - np.asarray(box, dtype=float) - if box is not None - else np.asarray(u.dimensions[:3], dtype=float) - ) - - heavy_atoms = u.atoms[topology.residue_heavy_indices] - heavy_atom = u.atoms[int(topology.heavy_atom_index)] - - if len(heavy_atoms) > 1: - center = residue_atoms.center_of_mass(unwrap=True) - moment_of_inertia_tensor = self.get_moment_of_inertia_tensor( - center_of_mass=center, - positions=heavy_atoms.positions, - masses=topology.residue_ua_masses, - dimensions=dimensions, - ) - trans_axes, _moment_of_inertia = self.get_custom_principal_axes( - moment_of_inertia_tensor - ) - else: - make_whole(residue_atoms) - trans_axes = residue_atoms.principal_axes() - - center = heavy_atom.position - rot_axes, moment_of_inertia = self.get_bonded_axes_from_topology( - u=u, - heavy_atom=heavy_atom, - topology=topology, - dimensions=dimensions, - ) - if rot_axes is None or moment_of_inertia is None: - raise ValueError("Unable to compute bonded axes for cached UA bead.") - - logger.debug("Translational Axes: %s", trans_axes) - logger.debug("Rotational Axes: %s", rot_axes) - logger.debug("Center: %s", center) - logger.debug("Moment of Inertia: %s", moment_of_inertia) - - return trans_axes, rot_axes, center, moment_of_inertia - - def get_bonded_axes_from_topology( - self, - *, - u, - heavy_atom, - topology, - dimensions: np.ndarray, - ): - """Compute UA bonded axes using cached bonded atom indices. - - This mirrors ``get_bonded_axes`` but receives precomputed bonded atom - memberships from ``UAAxesTopology`` instead of rediscovering them with - MDAnalysis selection strings inside the frame loop. - - Args: - u: Current-frame universe. - heavy_atom: Current-frame heavy atom for the UA bead. - topology: Cached ``UAAxesTopology`` for the UA bead. - dimensions: Simulation box lengths, shape ``(3,)``. - - Returns: - Tuple[np.ndarray | None, np.ndarray | None]: - - custom_axes: Custom rotation axes, shape ``(3, 3)``, or ``None``. - - custom_moment_of_inertia: Principal moments, shape ``(3,)``, or - ``None``. - """ - if not heavy_atom.mass > 1.1: - return None, None - - custom_moment_of_inertia = None - custom_axes = None - - heavy_bonded = u.atoms[topology.bonded_heavy_indices] - light_bonded = u.atoms[topology.bonded_light_indices] - ua = u.atoms[topology.ua_atom_indices] - ua_all = u.atoms[topology.ua_all_atom_indices] - - if len(heavy_bonded) == 0: - custom_axes, custom_moment_of_inertia = self.get_vanilla_axes(ua_all) - - if len(heavy_bonded) == 1 and len(light_bonded) == 0: - custom_axes = self.get_custom_axes( - a=heavy_atom.position, - b_list=[heavy_bonded[0].position], - c=np.zeros(3), - dimensions=dimensions, - ) - - if len(heavy_bonded) == 1 and len(light_bonded) >= 1: - custom_axes = self.get_custom_axes( - a=heavy_atom.position, - b_list=[heavy_bonded[0].position], - c=light_bonded[0].position, - dimensions=dimensions, - ) - - if len(heavy_bonded) >= 2: - custom_axes = self.get_custom_axes( - a=heavy_atom.position, - b_list=heavy_bonded.positions, - c=heavy_bonded[1].position, - dimensions=dimensions, - ) - - if custom_axes is None: - return None, None - - if custom_moment_of_inertia is None: - custom_moment_of_inertia = self.get_custom_moment_of_inertia( - UA=ua, - custom_rotation_axes=custom_axes, - center_of_mass=heavy_atom.position, - dimensions=dimensions, - ) - - custom_axes = self.get_flipped_axes( - ua, - custom_axes, - heavy_atom.position, - dimensions, - ) - - return custom_axes, custom_moment_of_inertia - def get_bonded_axes(self, system, atom, dimensions: np.ndarray): """Compute UA rotational axes from bonded topology around a heavy atom. diff --git a/CodeEntropy/levels/graph/conformation_dag.py b/CodeEntropy/levels/conformation_dag.py similarity index 100% rename from CodeEntropy/levels/graph/conformation_dag.py rename to CodeEntropy/levels/conformation_dag.py diff --git a/CodeEntropy/levels/dihedrals/topology.py b/CodeEntropy/levels/dihedrals/topology.py index 31e35ffd..680639a1 100644 --- a/CodeEntropy/levels/dihedrals/topology.py +++ b/CodeEntropy/levels/dihedrals/topology.py @@ -59,7 +59,9 @@ def _discover_group_dihedral_topology( topologies: list[MoleculeDihedralTopology] = [] for molecule_order, molecule_id in enumerate(molecules): - mol = self._extract_topology_fragment(data_container, molecule_id) + mol = self._universe_operations.extract_fragment( + data_container, molecule_id + ) num_residues = len(mol.residues) ua_dihedrals_by_residue: dict[int, list[Any]] = {} residue_dihedrals: list[Any] = [] @@ -88,26 +90,6 @@ def _discover_group_dihedral_topology( return topologies - def _extract_topology_fragment(self, data_container: Any, molecule_id: Any) -> Any: - """Return a molecule fragment for topology discovery. - - This uses the lightweight AtomGroup extraction helper when available so - static conformational topology discovery does not create a standalone - in-memory universe or copy trajectory frames. The fallback preserves - compatibility with older ``UniverseOperations`` implementations. - - Args: - data_container: Source MDAnalysis universe or universe-like container. - molecule_id: Fragment index identifying the molecule to extract. - - Returns: - MDAnalysis AtomGroup for the selected molecule - """ - return self._universe_operations.extract_fragment_atomgroup( - data_container, - int(molecule_id), - ) - def _select_heavy_residue(self, mol: Any, res_id: int) -> Any: """Select heavy atoms in a residue by residue index. @@ -118,15 +100,13 @@ def _select_heavy_residue(self, mol: Any, res_id: int) -> Any: Returns: AtomGroup containing heavy atoms in the residue selection. """ - residue_atoms = mol.residues[int(res_id)].atoms - selection1 = residue_atoms.indices[0] - selection2 = residue_atoms.indices[-1] + selection1 = mol.residues[res_id].atoms.indices[0] + selection2 = mol.residues[res_id].atoms.indices[-1] - res_container = mol.select_atoms( - f"index {selection1}:{selection2}", - updating=False, + res_container = self._universe_operations.select_atoms( + mol, f"index {selection1}:{selection2}" ) - return res_container.select_atoms("prop mass > 1.1", updating=False) + return self._universe_operations.select_atoms(res_container, "prop mass > 1.1") def _get_dihedrals(self, data_container: Any, level: str) -> list[Any]: """Return dihedral AtomGroups for a container at a given level. @@ -141,93 +121,26 @@ def _get_dihedrals(self, data_container: Any, level: str) -> list[Any]: atom_groups: list[Any] = [] if level == "united_atom": - selected_indices = {int(index) for index in data_container.indices} - for dihedral in data_container.dihedrals: - dihedral_atoms = dihedral.atoms - dihedral_indices = {int(index) for index in dihedral_atoms.indices} - - if len(dihedral_atoms) == 4 and dihedral_indices.issubset( - selected_indices - ): - atom_groups.append(dihedral_atoms) + atom_groups.append(dihedral.atoms) if level == "residue": num_residues = len(data_container.residues) if num_residues >= 4: for residue in range(4, num_residues + 1): - residue1 = data_container.residues[residue - 4] - residue2 = data_container.residues[residue - 3] - residue3 = data_container.residues[residue - 2] - residue4 = data_container.residues[residue - 1] - - atom1 = self._atoms_in_source_bonded_to_target( - residue1, - residue2, + atom1 = data_container.select_atoms( + f"resindex {residue - 4} and bonded resindex {residue - 3}" ) - atom2 = self._atoms_in_source_bonded_to_target( - residue2, - residue1, + atom2 = data_container.select_atoms( + f"resindex {residue - 3} and bonded resindex {residue - 4}" ) - atom3 = self._atoms_in_source_bonded_to_target( - residue3, - residue4, + atom3 = data_container.select_atoms( + f"resindex {residue - 2} and bonded resindex {residue - 1}" ) - atom4 = self._atoms_in_source_bonded_to_target( - residue4, - residue3, + atom4 = data_container.select_atoms( + f"resindex {residue - 1} and bonded resindex {residue - 2}" ) - - dihedral_atoms = atom1 + atom2 + atom3 + atom4 - - if len(dihedral_atoms) == 4: - atom_groups.append(dihedral_atoms) - else: - logger.debug( - "Skipping residue-level dihedral for local residues " - "%s-%s-%s-%s because it produced %d atoms.", - residue - 4, - residue - 3, - residue - 2, - residue - 1, - len(dihedral_atoms), - ) + atom_groups.append(atom1 + atom2 + atom3 + atom4) logger.debug("Level: %s, Dihedrals: %s", level, atom_groups) return atom_groups - - @staticmethod - def _atoms_in_source_bonded_to_target( - source_residue: Any, - target_residue: Any, - ) -> Any: - """Return source-residue atoms bonded to atoms in a target residue. - - This helper is used when constructing residue-level dihedral definitions - from lightweight molecule AtomGroups. It selects atoms from the source - residue that are bonded to any atom in the target residue without using - global ``resindex`` selection strings. - - Args: - source_residue: Residue whose atoms should be tested for bonds. - target_residue: Adjacent residue providing the target bonded atoms. - - Returns: - MDAnalysis AtomGroup containing atoms from ``source_residue`` that are - bonded to at least one atom in ``target_residue``. If no matching - atoms are found, an empty AtomGroup is returned. - """ - source_atoms = source_residue.atoms - target_indices = {int(index) for index in target_residue.atoms.indices} - selected_indices: list[int] = [] - - for atom in source_atoms: - bonded_atoms = getattr(atom, "bonded_atoms", None) - if bonded_atoms is None: - continue - - bonded_indices = {int(index) for index in bonded_atoms.indices} - if bonded_indices.intersection(target_indices): - selected_indices.append(int(atom.index)) - - return source_atoms.universe.atoms[selected_indices] diff --git a/CodeEntropy/levels/execution/__init__.py b/CodeEntropy/levels/execution/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/CodeEntropy/levels/execution/scheduler.py b/CodeEntropy/levels/execution/scheduler.py index 0a5b5063..5f3d8455 100644 --- a/CodeEntropy/levels/execution/scheduler.py +++ b/CodeEntropy/levels/execution/scheduler.py @@ -16,7 +16,7 @@ execute_frame_map_output, make_frame_worker_shared_data, ) -from CodeEntropy.levels.graph.frame_dag import FrameGraph +from CodeEntropy.levels.frame_dag import FrameGraph from CodeEntropy.levels.neighbors import Neighbors from CodeEntropy.results.reporter import _RichProgressSink diff --git a/CodeEntropy/levels/execution/tasks.py b/CodeEntropy/levels/execution/tasks.py index 872d1ad1..a0997ccf 100644 --- a/CodeEntropy/levels/execution/tasks.py +++ b/CodeEntropy/levels/execution/tasks.py @@ -5,7 +5,7 @@ from dataclasses import dataclass, field from typing import Any -from CodeEntropy.levels.graph.frame_dag import FrameGraph +from CodeEntropy.levels.frame_dag import FrameGraph from CodeEntropy.levels.neighbors import Neighbors FRAME_WORKER_EXCLUDED_SHARED_KEYS = { diff --git a/CodeEntropy/levels/graph/frame_dag.py b/CodeEntropy/levels/frame_dag.py similarity index 100% rename from CodeEntropy/levels/graph/frame_dag.py rename to CodeEntropy/levels/frame_dag.py diff --git a/CodeEntropy/levels/graph/__init__.py b/CodeEntropy/levels/graph/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/CodeEntropy/levels/graph/level_dag.py b/CodeEntropy/levels/level_dag.py similarity index 95% rename from CodeEntropy/levels/graph/level_dag.py rename to CodeEntropy/levels/level_dag.py index e333c772..cc8129a1 100644 --- a/CodeEntropy/levels/graph/level_dag.py +++ b/CodeEntropy/levels/level_dag.py @@ -12,14 +12,13 @@ import networkx as nx from CodeEntropy.levels.axes import AxesCalculator +from CodeEntropy.levels.conformation_dag import ConformationDAG from CodeEntropy.levels.execution.policy import ExecutionPolicy from CodeEntropy.levels.execution.reducers import NeighborReducer from CodeEntropy.levels.execution.scheduler import FrameScheduler -from CodeEntropy.levels.graph.conformation_dag import ConformationDAG -from CodeEntropy.levels.graph.frame_dag import FrameGraph +from CodeEntropy.levels.frame_dag import FrameGraph from CodeEntropy.levels.neighbors import Neighbors from CodeEntropy.levels.nodes.accumulators import InitCovarianceAccumulatorsNode -from CodeEntropy.levels.nodes.axes_topology import BuildAxesTopologyNode from CodeEntropy.levels.nodes.beads import BuildBeadsNode from CodeEntropy.levels.nodes.detect_levels import DetectLevelsNode from CodeEntropy.levels.nodes.detect_molecules import DetectMoleculesNode @@ -50,11 +49,6 @@ def build(self) -> LevelDAG: self._add_static("detect_molecules", DetectMoleculesNode()) self._add_static("detect_levels", DetectLevelsNode(), deps=["detect_molecules"]) self._add_static("build_beads", BuildBeadsNode(), deps=["detect_levels"]) - self._add_static( - "build_axes_topology", - BuildAxesTopologyNode(), - deps=["build_beads"], - ) self._add_static( "init_covariance_accumulators", InitCovarianceAccumulatorsNode(), diff --git a/CodeEntropy/levels/nodes/axes_topology.py b/CodeEntropy/levels/nodes/axes_topology.py deleted file mode 100644 index 5197d752..00000000 --- a/CodeEntropy/levels/nodes/axes_topology.py +++ /dev/null @@ -1,316 +0,0 @@ -"""Build static axes-topology metadata for frame covariance calculations. - -This module caches topology-only atom-index relationships needed by customised -axes calculations. The cache avoids repeated MDAnalysis selection parsing inside -the frame-local covariance loop while preserving frame-dependent positions, -forces, centres, axes, torques, and moments of inertia. -""" - -from __future__ import annotations - -import logging -from dataclasses import dataclass, field -from typing import Any - -import numpy as np - -logger = logging.getLogger(__name__) - -UAKey = tuple[int, int, int] -ResidueKey = tuple[int, int] - - -@dataclass(frozen=True) -class UAAxesTopology: - """Static topology required to compute customised united-atom axes. - - Attributes: - heavy_atom_index: Reduced-universe atom index for the UA heavy atom. - ua_atom_indices: Atom indices for the UA heavy atom and its bonded - hydrogens/light atoms. - ua_all_atom_indices: Atom indices for the UA heavy atom, bonded heavy - atoms, and bonded hydrogens/light atoms. - bonded_heavy_indices: Heavy atoms bonded to the UA heavy atom. - bonded_light_indices: Hydrogens/light atoms bonded to the UA heavy atom. - residue_heavy_indices: Heavy atoms in the parent residue. - residue_ua_masses: UA masses for heavy atoms in the parent residue. - """ - - heavy_atom_index: int - ua_atom_indices: np.ndarray - ua_all_atom_indices: np.ndarray - bonded_heavy_indices: np.ndarray - bonded_light_indices: np.ndarray - residue_heavy_indices: np.ndarray - residue_ua_masses: np.ndarray - - -@dataclass(frozen=True) -class ResidueAxesTopology: - """Static topology required to compute customised residue axes. - - Attributes: - residue_heavy_indices: Heavy atom indices in the residue. - residue_ua_masses: UA masses for heavy atoms in the residue. - has_neighbor_bonds: Whether the residue is bonded to a neighbouring - residue according to the original customised residue-axis selection. - """ - - residue_heavy_indices: np.ndarray - residue_ua_masses: np.ndarray - has_neighbor_bonds: bool - - -@dataclass(frozen=True) -class AxesTopology: - """Cached axes topology for frame covariance calculations. - - Attributes: - ua: Mapping from ``(mol_id, local_residue_id, ua_id)`` to cached - united-atom axes topology. - residue: Mapping from ``(mol_id, local_residue_id)`` to cached - residue axes topology. - """ - - ua: dict[UAKey, UAAxesTopology] = field(default_factory=dict) - residue: dict[ResidueKey, ResidueAxesTopology] = field(default_factory=dict) - - -class BuildAxesTopologyNode: - """Build static customised-axes topology before frame covariance execution.""" - - def run(self, shared_data: dict[str, Any]) -> dict[str, Any]: - """Build cached axes topology and write it into shared data. - - The cache is only populated when ``args.customised_axes`` is true. When - customised axes are disabled, an empty cache is still written so later - stages can read ``shared_data["axes_topology"]`` safely. - - Args: - shared_data: Shared workflow data containing ``args`` and, when - customised axes are enabled, ``reduced_universe``, ``levels``, - and ``beads``. - - Returns: - Dict containing the cached ``axes_topology`` object. - """ - args = shared_data["args"] - topology = AxesTopology() - - if not bool(getattr(args, "customised_axes", False)): - shared_data["axes_topology"] = topology - return {"axes_topology": topology} - - u = shared_data["reduced_universe"] - levels = shared_data["levels"] - beads = shared_data["beads"] - - ua_topology: dict[UAKey, UAAxesTopology] = {} - residue_topology: dict[ResidueKey, ResidueAxesTopology] = {} - fragments = u.atoms.fragments - - for mol_id, level_list in enumerate(levels): - mol = fragments[mol_id] - - if "residue" in level_list: - self._add_residue_topology( - mol=mol, - mol_id=mol_id, - beads=beads, - out=residue_topology, - ) - - if "united_atom" in level_list: - self._add_ua_topology( - u=u, - mol=mol, - mol_id=mol_id, - beads=beads, - out=ua_topology, - ) - - topology = AxesTopology(ua=ua_topology, residue=residue_topology) - shared_data["axes_topology"] = topology - return {"axes_topology": topology} - - def _add_residue_topology( - self, - *, - mol: Any, - mol_id: int, - beads: dict[Any, list[np.ndarray]], - out: dict[ResidueKey, ResidueAxesTopology], - ) -> None: - """Cache static residue axes topology for one molecule. - - Args: - mol: Molecule AtomGroup. - mol_id: Molecule index. - beads: Bead-index mapping produced by ``BuildBeadsNode``. - out: Output residue topology mapping mutated in place. - """ - bead_key = (mol_id, "residue") - bead_idx_list = beads.get(bead_key, []) - if not bead_idx_list: - return - - for local_res_i, residue in enumerate(mol.residues): - if local_res_i >= len(bead_idx_list): - continue - - residue_atoms = residue.atoms - residue_heavy = residue_atoms.select_atoms("mass 2 to 999") - residue_heavy_indices = residue_heavy.indices.astype(int, copy=True) - residue_ua_masses = np.asarray( - self._get_ua_masses_from_topology(residue_atoms), - dtype=float, - ) - has_neighbor_bonds = self._has_neighbor_bonds( - mol=mol, - local_res_i=local_res_i, - ) - - out[(mol_id, local_res_i)] = ResidueAxesTopology( - residue_heavy_indices=residue_heavy_indices, - residue_ua_masses=residue_ua_masses, - has_neighbor_bonds=has_neighbor_bonds, - ) - - def _add_ua_topology( - self, - *, - u: Any, - mol: Any, - mol_id: int, - beads: dict[Any, list[np.ndarray]], - out: dict[UAKey, UAAxesTopology], - ) -> None: - """Cache static UA axes topology for one molecule. - - Args: - u: Reduced universe used to resolve bead atom-index arrays. - mol: Molecule AtomGroup. - mol_id: Molecule index. - beads: Bead-index mapping produced by ``BuildBeadsNode``. - out: Output UA topology mapping mutated in place. - """ - for local_res_i, residue in enumerate(mol.residues): - bead_key = (mol_id, "united_atom", local_res_i) - bead_idx_list = beads.get(bead_key, []) - - if not bead_idx_list: - continue - - residue_atoms = residue.atoms - residue_heavy = residue_atoms.select_atoms("prop mass > 1.1") - residue_heavy_indices = residue_heavy.indices.astype(int, copy=True) - residue_ua_masses = np.asarray( - self._get_ua_masses_from_topology(residue_atoms), - dtype=float, - ) - - for ua_i, bead_indices in enumerate(bead_idx_list): - bead = u.atoms[bead_indices] - heavy = bead.select_atoms("prop mass > 1.1") - - if len(heavy) == 0: - logger.warning( - "Skipping UA axes topology with no heavy atom: " - "mol=%s residue=%s ua=%s", - mol_id, - local_res_i, - ua_i, - ) - continue - - heavy_atom = heavy[0] - bonded_heavy, bonded_light = self._split_bonded_atoms(heavy_atom) - - heavy_index = np.asarray([int(heavy_atom.index)], dtype=int) - bonded_heavy_indices = bonded_heavy.indices.astype(int, copy=True) - bonded_light_indices = bonded_light.indices.astype(int, copy=True) - - ua_atom_indices = np.concatenate( - [heavy_index, bonded_light_indices], - axis=0, - ) - ua_all_atom_indices = np.concatenate( - [heavy_index, bonded_heavy_indices, bonded_light_indices], - axis=0, - ) - - out[(mol_id, local_res_i, ua_i)] = UAAxesTopology( - heavy_atom_index=int(heavy_atom.index), - ua_atom_indices=ua_atom_indices, - ua_all_atom_indices=ua_all_atom_indices, - bonded_heavy_indices=bonded_heavy_indices, - bonded_light_indices=bonded_light_indices, - residue_heavy_indices=residue_heavy_indices, - residue_ua_masses=residue_ua_masses, - ) - - @staticmethod - def _has_neighbor_bonds(*, mol: Any, local_res_i: int) -> bool: - """Return whether a residue is bonded to neighbouring residues. - - Args: - mol: Molecule AtomGroup used for the original bonded-neighbour - selection. - local_res_i: Residue index local to ``mol``. - - Returns: - True when the residue has bonded atoms in the previous or next - residue according to the original customised residue-axis query. - """ - index_prev = local_res_i - 1 - index_next = local_res_i + 1 - atom_set = mol.select_atoms( - f"(resindex {index_prev} or resindex {index_next}) " - f"and bonded resid {local_res_i}" - ) - return len(atom_set) > 0 - - @staticmethod - def _split_bonded_atoms(atom: Any) -> tuple[Any, Any]: - """Return bonded heavy and light atoms for one atom. - - Args: - atom: MDAnalysis Atom. - - Returns: - Tuple containing bonded heavy atoms and bonded hydrogens/light atoms. - """ - bonded_atoms = atom.bonded_atoms - bonded_heavy = bonded_atoms.select_atoms("mass 2 to 999") - bonded_light = bonded_atoms.select_atoms("mass 1 to 1.1") - return bonded_heavy, bonded_light - - @staticmethod - def _get_ua_masses_from_topology(atom_group: Any) -> list[float]: - """Return UA masses using static bonded atom relationships. - - Args: - atom_group: AtomGroup containing atoms from one residue. - - Returns: - List of UA masses, one for each heavy atom in ``atom_group``. - """ - ua_masses: list[float] = [] - - for atom in atom_group: - if atom.mass <= 1.1: - continue - - ua_mass = float(atom.mass) - bonded_atoms = getattr(atom, "bonded_atoms", None) - if bonded_atoms is None: - ua_masses.append(ua_mass) - continue - - bonded_h_atoms = bonded_atoms.select_atoms("mass 1 to 1.1") - for hydrogen in bonded_h_atoms: - ua_mass += float(hydrogen.mass) - - ua_masses.append(ua_mass) - - return ua_masses diff --git a/CodeEntropy/levels/nodes/covariance.py b/CodeEntropy/levels/nodes/covariance.py index 9588c050..38ec3d1c 100644 --- a/CodeEntropy/levels/nodes/covariance.py +++ b/CodeEntropy/levels/nodes/covariance.py @@ -79,7 +79,6 @@ def run(self, ctx: FrameCtx) -> dict[str, Any]: beads = shared["beads"] args = shared["args"] axes_manager = shared.get("axes_manager") - axes_topology = shared.get("axes_topology") fp = float(args.force_partitioning) combined = bool(getattr(args, "combined_forcetorque", False)) @@ -111,7 +110,6 @@ def run(self, ctx: FrameCtx) -> dict[str, Any]: group_id=group_id, beads=beads, axes_manager=axes_manager, - axes_topology=axes_topology, box=box, force_partitioning=fp, customised_axes=customised_axes, @@ -129,7 +127,6 @@ def run(self, ctx: FrameCtx) -> dict[str, Any]: group_id=group_id, beads=beads, axes_manager=axes_manager, - axes_topology=axes_topology, box=box, customised_axes=customised_axes, force_partitioning=fp, @@ -175,7 +172,6 @@ def _process_united_atom( group_id: int, beads: dict[Any, list[Any]], axes_manager: Any, - axes_topology: Any | None, box: np.ndarray | None, force_partitioning: float, customised_axes: bool, @@ -193,7 +189,6 @@ def _process_united_atom( group_id: Molecule-group identifier used for within-frame averaging. beads: Mapping of bead keys to reduced-universe atom-index arrays. axes_manager: Axes helper used to build translation and rotation axes. - axes_topology: Optional cached axes topology generated during static setup. box: Optional periodic box vector. force_partitioning: Force partitioning factor for highest-level vectors. customised_axes: Whether customised UA axes should be used. @@ -239,7 +234,6 @@ def _process_united_atom( residue_group=residue_group.atoms, bead_groups=bead_groups, axes_manager=axes_manager, - axes_topology=axes_topology, box=box, force_partitioning=force_partitioning, customised_axes=customised_axes, @@ -264,7 +258,6 @@ def _process_residue( group_id: int, beads: dict[Any, list[Any]], axes_manager: Any, - axes_topology: Any | None, box: np.ndarray | None, customised_axes: bool, force_partitioning: float, @@ -284,7 +277,6 @@ def _process_residue( group_id: Molecule-group identifier used for within-frame averaging. beads: Mapping of bead keys to reduced-universe atom-index arrays. axes_manager: Axes helper used to build translation and rotation axes. - axes_topology: Optional cached axes topology generated during static setup. box: Optional periodic box vector. customised_axes: Whether customised residue axes should be used. force_partitioning: Force partitioning factor for highest-level vectors. @@ -305,12 +297,9 @@ def _process_residue( return force_vecs, torque_vecs = self._build_residue_vectors( - u=u, mol=mol, - mol_id=mol_id, bead_groups=bead_groups, axes_manager=axes_manager, - axes_topology=axes_topology, box=box, customised_axes=customised_axes, force_partitioning=force_partitioning, @@ -423,13 +412,9 @@ def _process_polymer( def _build_ua_vectors( self, *, - u: Any, - mol_id: int, - local_res_i: int, bead_groups: list[Any], residue_group: Any, axes_manager: Any, - axes_topology: Any | None, box: np.ndarray | None, force_partitioning: float, customised_axes: bool, @@ -439,19 +424,16 @@ def _build_ua_vectors( """Build force and torque vectors for united-atom beads. Args: - u: Universe-like object used to resolve cached atom indices. - mol_id: Molecule index used in axes-topology lookup keys. - local_res_i: Local residue index used in axes-topology lookup keys. - bead_groups: Atom groups representing UA beads in a residue. + bead_groups: List of UA bead AtomGroups for the residue. residue_group: AtomGroup for the residue group atoms. - axes_manager: Axes helper used to select axes, centres, and moments. - axes_topology: Optional cached axes topology generated during static setup. - box: Optional periodic box vector. - force_partitioning: Force partitioning factor for highest-level vectors. - customised_axes: Whether customised UA axes should be used. - is_highest: Whether UA is the highest active level. + axes_manager: Axes manager used to determine axes/centers/MOI. + box: Optional box vector used for PBC-aware displacements. + force_partitioning: Force scaling factor applied at highest level. + customised_axes: Whether to use customised axes methods when available. + is_highest: Whether UA level is the highest level for the molecule. res_position: Where the residue is in the residue group + Returns: A tuple containing lists of force vectors and torque vectors. """ @@ -460,26 +442,9 @@ def _build_ua_vectors( for ua_i, bead in enumerate(bead_groups): if customised_axes: - ua_topology = None - if axes_topology is not None: - ua_topology = axes_topology.ua.get((mol_id, local_res_i, - ua_i, - res_position)) - - if ua_topology is not None: - trans_axes, rot_axes, center, moi = ( - axes_manager.get_UA_axes_from_topology( - u=u, - residue_group=residue_group, - topology=ua_topology, - box=box, - res_position=res_position - ) - ) - else: - trans_axes, rot_axes, center, moi = axes_manager.get_UA_axes( - residue_atoms, ua_i - ) + trans_axes, rot_axes, center, moi = axes_manager.get_UA_axes( + residue_group, ua_i, res_position + ) else: make_whole(residue_group) make_whole(bead) @@ -521,12 +486,9 @@ def _build_ua_vectors( def _build_residue_vectors( self, *, - u: Any, mol: Any, - mol_id: int, bead_groups: list[Any], axes_manager: Any, - axes_topology: Any | None, box: np.ndarray | None, customised_axes: bool, force_partitioning: float, @@ -535,12 +497,9 @@ def _build_residue_vectors( """Build force and torque vectors for residue beads. Args: - u: Universe-like object used to resolve cached atom indices. mol: Molecule fragment containing residues and atoms. - mol_id: Molecule index used in axes-topology lookup keys. bead_groups: Atom groups representing residue beads. axes_manager: Axes helper used to select axes, centres, and moments. - axes_topology: Optional cached axes topology generated during static setup. box: Optional periodic box vector. customised_axes: Whether customised residue axes should be used. force_partitioning: Force partitioning factor for highest-level vectors. @@ -554,14 +513,10 @@ def _build_residue_vectors( for local_res_i, bead in enumerate(bead_groups): trans_axes, rot_axes, center, moi = self._get_residue_axes( - u=u, mol=mol, - mol_id=mol_id, bead=bead, local_res_i=local_res_i, axes_manager=axes_manager, - axes_topology=axes_topology, - box=box, customised_axes=customised_axes, ) @@ -590,27 +545,19 @@ def _build_residue_vectors( def _get_residue_axes( self, *, - u: Any, mol: Any, - mol_id: int, bead: Any, local_res_i: int, axes_manager: Any, - axes_topology: Any | None, - box: np.ndarray | None, customised_axes: bool, ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: """Return axes, centre, and inertia data for a residue bead. Args: - u: Universe-like object used to resolve cached atom indices. mol: Molecule fragment containing residues and atoms. - mol_id: Molecule index used in axes-topology lookup keys. bead: Atom group representing the residue bead. local_res_i: Residue index local to ``mol``. axes_manager: Axes helper used to select axes, centres, and moments. - axes_topology: Optional cached axes topology generated during static setup. - box: Optional periodic box vector. customised_axes: Whether customised residue axes should be used. Returns: @@ -618,19 +565,6 @@ def _get_residue_axes( """ if customised_axes: res = mol.residues[local_res_i] - residue_topology = None - if axes_topology is not None: - residue_topology = axes_topology.residue.get((mol_id, local_res_i)) - - if residue_topology is not None: - return axes_manager.get_residue_axes_from_topology( - u=u, - mol=mol, - residue_atoms=res.atoms, - topology=residue_topology, - box=box, - ) - return axes_manager.get_residue_axes(mol, local_res_i, residue=res.atoms) make_whole(mol.atoms) diff --git a/CodeEntropy/trajectory/mda.py b/CodeEntropy/trajectory/mda.py index 1cdb6e75..04943a4c 100644 --- a/CodeEntropy/trajectory/mda.py +++ b/CodeEntropy/trajectory/mda.py @@ -25,8 +25,7 @@ class UniverseOperations: This helper provides methods to: - Build reduced universes by selecting subsets of frames or atoms. - Extract a single fragment (molecule) into a standalone universe. - - Merge coordinates from one trajectory with forces sourced from another - trajectory. + - Merge coordinates from one trajectory with forces from another trajectory. """ def __init__(self) -> None: @@ -187,25 +186,6 @@ def extract_fragment( selection_string = f"index {frag.indices[0]}:{frag.indices[-1]}" return self.select_atoms(universe, selection_string) - def extract_fragment_atomgroup(self, universe: mda.Universe, molecule_id: int): - """Return a molecule fragment as an AtomGroup. - - This helper mirrors the atom-index range used by ``extract_fragment`` but - avoids building a standalone in-memory universe. It is intended for - topology discovery paths where only the static atom selection is needed - and trajectory coordinates do not need to be copied. - - Args: - universe: Source MDAnalysis universe. - molecule_id: Fragment index in ``universe.atoms.fragments``. - - Returns: - MDAnalysis AtomGroup containing the atoms in the selected fragment. - """ - frag = universe.atoms.fragments[int(molecule_id)] - selection_string = f"index {frag.indices[0]}:{frag.indices[-1]}" - return universe.select_atoms(selection_string, updating=False) - def convert_lammps( self, tprfile: str, diff --git a/docs/api/CodeEntropy.config.rst b/docs/api/CodeEntropy.config.rst index d7945e83..c85bcb89 100644 --- a/docs/api/CodeEntropy.config.rst +++ b/docs/api/CodeEntropy.config.rst @@ -1,6 +1,11 @@ CodeEntropy.config package ========================== +.. automodule:: CodeEntropy.config + :members: + :show-inheritance: + :undoc-members: + Submodules ---------- @@ -9,11 +14,3 @@ Submodules CodeEntropy.config.argparse CodeEntropy.config.runtime - -Module contents ---------------- - -.. automodule:: CodeEntropy.config - :members: - :show-inheritance: - :undoc-members: diff --git a/docs/api/CodeEntropy.core.rst b/docs/api/CodeEntropy.core.rst index 86ebc99c..695ee31d 100644 --- a/docs/api/CodeEntropy.core.rst +++ b/docs/api/CodeEntropy.core.rst @@ -1,6 +1,11 @@ CodeEntropy.core package ======================== +.. automodule:: CodeEntropy.core + :members: + :show-inheritance: + :undoc-members: + Submodules ---------- @@ -9,11 +14,3 @@ Submodules CodeEntropy.core.dask_clusters CodeEntropy.core.logging - -Module contents ---------------- - -.. automodule:: CodeEntropy.core - :members: - :show-inheritance: - :undoc-members: diff --git a/docs/api/CodeEntropy.entropy.nodes.rst b/docs/api/CodeEntropy.entropy.nodes.rst index 4ff7516c..2fa9338e 100644 --- a/docs/api/CodeEntropy.entropy.nodes.rst +++ b/docs/api/CodeEntropy.entropy.nodes.rst @@ -1,6 +1,11 @@ CodeEntropy.entropy.nodes package ================================= +.. automodule:: CodeEntropy.entropy.nodes + :members: + :show-inheritance: + :undoc-members: + Submodules ---------- @@ -11,11 +16,3 @@ Submodules CodeEntropy.entropy.nodes.configurational CodeEntropy.entropy.nodes.orientational CodeEntropy.entropy.nodes.vibrational - -Module contents ---------------- - -.. automodule:: CodeEntropy.entropy.nodes - :members: - :show-inheritance: - :undoc-members: diff --git a/docs/api/CodeEntropy.entropy.rst b/docs/api/CodeEntropy.entropy.rst index f4322465..99ed6b3f 100644 --- a/docs/api/CodeEntropy.entropy.rst +++ b/docs/api/CodeEntropy.entropy.rst @@ -1,6 +1,11 @@ CodeEntropy.entropy package =========================== +.. automodule:: CodeEntropy.entropy + :members: + :show-inheritance: + :undoc-members: + Subpackages ----------- @@ -21,11 +26,3 @@ Submodules CodeEntropy.entropy.vibrational CodeEntropy.entropy.water CodeEntropy.entropy.workflow - -Module contents ---------------- - -.. automodule:: CodeEntropy.entropy - :members: - :show-inheritance: - :undoc-members: diff --git a/docs/api/CodeEntropy.levels.dihedrals.kernels.rst b/docs/api/CodeEntropy.levels.conformation_dag.rst similarity index 50% rename from docs/api/CodeEntropy.levels.dihedrals.kernels.rst rename to docs/api/CodeEntropy.levels.conformation_dag.rst index 3e51f72a..ffea54eb 100644 --- a/docs/api/CodeEntropy.levels.dihedrals.kernels.rst +++ b/docs/api/CodeEntropy.levels.conformation_dag.rst @@ -1,7 +1,7 @@ -CodeEntropy.levels.dihedrals.kernels module +CodeEntropy.levels.conformation\_dag module =========================================== -.. automodule:: CodeEntropy.levels.dihedrals.kernels +.. automodule:: CodeEntropy.levels.conformation_dag :members: :show-inheritance: :undoc-members: diff --git a/docs/api/CodeEntropy.levels.dihedrals.angle_observations.rst b/docs/api/CodeEntropy.levels.dihedrals.angle_observations.rst deleted file mode 100644 index 16bf9470..00000000 --- a/docs/api/CodeEntropy.levels.dihedrals.angle_observations.rst +++ /dev/null @@ -1,7 +0,0 @@ -CodeEntropy.levels.dihedrals.angle\_observations module -======================================================= - -.. automodule:: CodeEntropy.levels.dihedrals.angle_observations - :members: - :show-inheritance: - :undoc-members: diff --git a/docs/api/CodeEntropy.levels.dihedrals.conformational_state_builder.rst b/docs/api/CodeEntropy.levels.dihedrals.conformational_state_builder.rst deleted file mode 100644 index 449f21c8..00000000 --- a/docs/api/CodeEntropy.levels.dihedrals.conformational_state_builder.rst +++ /dev/null @@ -1,7 +0,0 @@ -CodeEntropy.levels.dihedrals.conformational\_state\_builder module -================================================================== - -.. automodule:: CodeEntropy.levels.dihedrals.conformational_state_builder - :members: - :show-inheritance: - :undoc-members: diff --git a/docs/api/CodeEntropy.levels.dihedrals.peak_detection.rst b/docs/api/CodeEntropy.levels.dihedrals.peak_detection.rst deleted file mode 100644 index 11c89969..00000000 --- a/docs/api/CodeEntropy.levels.dihedrals.peak_detection.rst +++ /dev/null @@ -1,7 +0,0 @@ -CodeEntropy.levels.dihedrals.peak\_detection module -=================================================== - -.. automodule:: CodeEntropy.levels.dihedrals.peak_detection - :members: - :show-inheritance: - :undoc-members: diff --git a/docs/api/CodeEntropy.levels.dihedrals.rst b/docs/api/CodeEntropy.levels.dihedrals.rst index 30aa5a67..a69952a9 100644 --- a/docs/api/CodeEntropy.levels.dihedrals.rst +++ b/docs/api/CodeEntropy.levels.dihedrals.rst @@ -1,21 +1,5 @@ -CodeEntropy.levels.dihedrals package -==================================== - -Submodules ----------- - -.. toctree:: - :maxdepth: 4 - - CodeEntropy.levels.dihedrals.angle_observations - CodeEntropy.levels.dihedrals.conformational_state_builder - CodeEntropy.levels.dihedrals.kernels - CodeEntropy.levels.dihedrals.peak_detection - CodeEntropy.levels.dihedrals.state_assignment - CodeEntropy.levels.dihedrals.topology - -Module contents ---------------- +CodeEntropy.levels.dihedrals module +=================================== .. automodule:: CodeEntropy.levels.dihedrals :members: diff --git a/docs/api/CodeEntropy.levels.dihedrals.state_assignment.rst b/docs/api/CodeEntropy.levels.dihedrals.state_assignment.rst deleted file mode 100644 index aeed41ce..00000000 --- a/docs/api/CodeEntropy.levels.dihedrals.state_assignment.rst +++ /dev/null @@ -1,7 +0,0 @@ -CodeEntropy.levels.dihedrals.state\_assignment module -===================================================== - -.. automodule:: CodeEntropy.levels.dihedrals.state_assignment - :members: - :show-inheritance: - :undoc-members: diff --git a/docs/api/CodeEntropy.levels.dihedrals.topology.rst b/docs/api/CodeEntropy.levels.dihedrals.topology.rst deleted file mode 100644 index 30074514..00000000 --- a/docs/api/CodeEntropy.levels.dihedrals.topology.rst +++ /dev/null @@ -1,7 +0,0 @@ -CodeEntropy.levels.dihedrals.topology module -============================================ - -.. automodule:: CodeEntropy.levels.dihedrals.topology - :members: - :show-inheritance: - :undoc-members: diff --git a/docs/api/CodeEntropy.levels.execution.chunks.rst b/docs/api/CodeEntropy.levels.execution.chunks.rst deleted file mode 100644 index 6f7924f8..00000000 --- a/docs/api/CodeEntropy.levels.execution.chunks.rst +++ /dev/null @@ -1,7 +0,0 @@ -CodeEntropy.levels.execution.chunks module -========================================== - -.. automodule:: CodeEntropy.levels.execution.chunks - :members: - :show-inheritance: - :undoc-members: diff --git a/docs/api/CodeEntropy.levels.execution.policy.rst b/docs/api/CodeEntropy.levels.execution.policy.rst deleted file mode 100644 index 3dd1f5aa..00000000 --- a/docs/api/CodeEntropy.levels.execution.policy.rst +++ /dev/null @@ -1,7 +0,0 @@ -CodeEntropy.levels.execution.policy module -========================================== - -.. automodule:: CodeEntropy.levels.execution.policy - :members: - :show-inheritance: - :undoc-members: diff --git a/docs/api/CodeEntropy.levels.execution.reducers.rst b/docs/api/CodeEntropy.levels.execution.reducers.rst deleted file mode 100644 index 0beef239..00000000 --- a/docs/api/CodeEntropy.levels.execution.reducers.rst +++ /dev/null @@ -1,7 +0,0 @@ -CodeEntropy.levels.execution.reducers module -============================================ - -.. automodule:: CodeEntropy.levels.execution.reducers - :members: - :show-inheritance: - :undoc-members: diff --git a/docs/api/CodeEntropy.levels.execution.rst b/docs/api/CodeEntropy.levels.execution.rst deleted file mode 100644 index 2932017a..00000000 --- a/docs/api/CodeEntropy.levels.execution.rst +++ /dev/null @@ -1,22 +0,0 @@ -CodeEntropy.levels.execution package -==================================== - -Submodules ----------- - -.. toctree:: - :maxdepth: 4 - - CodeEntropy.levels.execution.chunks - CodeEntropy.levels.execution.policy - CodeEntropy.levels.execution.reducers - CodeEntropy.levels.execution.scheduler - CodeEntropy.levels.execution.tasks - -Module contents ---------------- - -.. automodule:: CodeEntropy.levels.execution - :members: - :show-inheritance: - :undoc-members: diff --git a/docs/api/CodeEntropy.levels.execution.scheduler.rst b/docs/api/CodeEntropy.levels.execution.scheduler.rst deleted file mode 100644 index b8a95196..00000000 --- a/docs/api/CodeEntropy.levels.execution.scheduler.rst +++ /dev/null @@ -1,7 +0,0 @@ -CodeEntropy.levels.execution.scheduler module -============================================= - -.. automodule:: CodeEntropy.levels.execution.scheduler - :members: - :show-inheritance: - :undoc-members: diff --git a/docs/api/CodeEntropy.levels.execution.tasks.rst b/docs/api/CodeEntropy.levels.execution.tasks.rst deleted file mode 100644 index 9f2bda92..00000000 --- a/docs/api/CodeEntropy.levels.execution.tasks.rst +++ /dev/null @@ -1,7 +0,0 @@ -CodeEntropy.levels.execution.tasks module -========================================= - -.. automodule:: CodeEntropy.levels.execution.tasks - :members: - :show-inheritance: - :undoc-members: diff --git a/docs/api/CodeEntropy.levels.frame_dag.rst b/docs/api/CodeEntropy.levels.frame_dag.rst new file mode 100644 index 00000000..9a747a87 --- /dev/null +++ b/docs/api/CodeEntropy.levels.frame_dag.rst @@ -0,0 +1,7 @@ +CodeEntropy.levels.frame\_dag module +==================================== + +.. automodule:: CodeEntropy.levels.frame_dag + :members: + :show-inheritance: + :undoc-members: diff --git a/docs/api/CodeEntropy.levels.graph.conformation_dag.rst b/docs/api/CodeEntropy.levels.graph.conformation_dag.rst deleted file mode 100644 index ee59e581..00000000 --- a/docs/api/CodeEntropy.levels.graph.conformation_dag.rst +++ /dev/null @@ -1,7 +0,0 @@ -CodeEntropy.levels.graph.conformation\_dag module -================================================= - -.. automodule:: CodeEntropy.levels.graph.conformation_dag - :members: - :show-inheritance: - :undoc-members: diff --git a/docs/api/CodeEntropy.levels.graph.frame_dag.rst b/docs/api/CodeEntropy.levels.graph.frame_dag.rst deleted file mode 100644 index 361e5502..00000000 --- a/docs/api/CodeEntropy.levels.graph.frame_dag.rst +++ /dev/null @@ -1,7 +0,0 @@ -CodeEntropy.levels.graph.frame\_dag module -========================================== - -.. automodule:: CodeEntropy.levels.graph.frame_dag - :members: - :show-inheritance: - :undoc-members: diff --git a/docs/api/CodeEntropy.levels.graph.level_dag.rst b/docs/api/CodeEntropy.levels.graph.level_dag.rst deleted file mode 100644 index a00e69b5..00000000 --- a/docs/api/CodeEntropy.levels.graph.level_dag.rst +++ /dev/null @@ -1,7 +0,0 @@ -CodeEntropy.levels.graph.level\_dag module -========================================== - -.. automodule:: CodeEntropy.levels.graph.level_dag - :members: - :show-inheritance: - :undoc-members: diff --git a/docs/api/CodeEntropy.levels.graph.rst b/docs/api/CodeEntropy.levels.graph.rst deleted file mode 100644 index d4f75dbb..00000000 --- a/docs/api/CodeEntropy.levels.graph.rst +++ /dev/null @@ -1,20 +0,0 @@ -CodeEntropy.levels.graph package -================================ - -Submodules ----------- - -.. toctree:: - :maxdepth: 4 - - CodeEntropy.levels.graph.conformation_dag - CodeEntropy.levels.graph.frame_dag - CodeEntropy.levels.graph.level_dag - -Module contents ---------------- - -.. automodule:: CodeEntropy.levels.graph - :members: - :show-inheritance: - :undoc-members: diff --git a/docs/api/CodeEntropy.levels.level_dag.rst b/docs/api/CodeEntropy.levels.level_dag.rst new file mode 100644 index 00000000..78c08786 --- /dev/null +++ b/docs/api/CodeEntropy.levels.level_dag.rst @@ -0,0 +1,7 @@ +CodeEntropy.levels.level\_dag module +==================================== + +.. automodule:: CodeEntropy.levels.level_dag + :members: + :show-inheritance: + :undoc-members: diff --git a/docs/api/CodeEntropy.levels.nodes.axes_topology.rst b/docs/api/CodeEntropy.levels.nodes.axes_topology.rst deleted file mode 100644 index af80bc9e..00000000 --- a/docs/api/CodeEntropy.levels.nodes.axes_topology.rst +++ /dev/null @@ -1,7 +0,0 @@ -CodeEntropy.levels.nodes.axes\_topology module -============================================== - -.. automodule:: CodeEntropy.levels.nodes.axes_topology - :members: - :show-inheritance: - :undoc-members: diff --git a/docs/api/CodeEntropy.levels.nodes.rst b/docs/api/CodeEntropy.levels.nodes.rst index 048f7fcf..4499f939 100644 --- a/docs/api/CodeEntropy.levels.nodes.rst +++ b/docs/api/CodeEntropy.levels.nodes.rst @@ -1,6 +1,11 @@ CodeEntropy.levels.nodes package ================================ +.. automodule:: CodeEntropy.levels.nodes + :members: + :show-inheritance: + :undoc-members: + Submodules ---------- @@ -8,16 +13,7 @@ Submodules :maxdepth: 4 CodeEntropy.levels.nodes.accumulators - CodeEntropy.levels.nodes.axes_topology CodeEntropy.levels.nodes.beads CodeEntropy.levels.nodes.covariance CodeEntropy.levels.nodes.detect_levels CodeEntropy.levels.nodes.detect_molecules - -Module contents ---------------- - -.. automodule:: CodeEntropy.levels.nodes - :members: - :show-inheritance: - :undoc-members: diff --git a/docs/api/CodeEntropy.levels.rst b/docs/api/CodeEntropy.levels.rst index 9438f461..cca6398d 100644 --- a/docs/api/CodeEntropy.levels.rst +++ b/docs/api/CodeEntropy.levels.rst @@ -1,15 +1,17 @@ CodeEntropy.levels package ========================== +.. automodule:: CodeEntropy.levels + :members: + :show-inheritance: + :undoc-members: + Subpackages ----------- .. toctree:: :maxdepth: 4 - CodeEntropy.levels.dihedrals - CodeEntropy.levels.execution - CodeEntropy.levels.graph CodeEntropy.levels.nodes Submodules @@ -19,16 +21,12 @@ Submodules :maxdepth: 4 CodeEntropy.levels.axes + CodeEntropy.levels.conformation_dag + CodeEntropy.levels.dihedrals CodeEntropy.levels.forces + CodeEntropy.levels.frame_dag CodeEntropy.levels.hierarchy + CodeEntropy.levels.level_dag CodeEntropy.levels.linalg CodeEntropy.levels.neighbors CodeEntropy.levels.search - -Module contents ---------------- - -.. automodule:: CodeEntropy.levels - :members: - :show-inheritance: - :undoc-members: diff --git a/docs/api/CodeEntropy.molecules.rst b/docs/api/CodeEntropy.molecules.rst index 3dd473be..af25747e 100644 --- a/docs/api/CodeEntropy.molecules.rst +++ b/docs/api/CodeEntropy.molecules.rst @@ -1,6 +1,11 @@ CodeEntropy.molecules package ============================= +.. automodule:: CodeEntropy.molecules + :members: + :show-inheritance: + :undoc-members: + Submodules ---------- @@ -8,11 +13,3 @@ Submodules :maxdepth: 4 CodeEntropy.molecules.grouping - -Module contents ---------------- - -.. automodule:: CodeEntropy.molecules - :members: - :show-inheritance: - :undoc-members: diff --git a/docs/api/CodeEntropy.results.rst b/docs/api/CodeEntropy.results.rst index 2606aaf4..40d4eee2 100644 --- a/docs/api/CodeEntropy.results.rst +++ b/docs/api/CodeEntropy.results.rst @@ -1,6 +1,11 @@ CodeEntropy.results package =========================== +.. automodule:: CodeEntropy.results + :members: + :show-inheritance: + :undoc-members: + Submodules ---------- @@ -8,11 +13,3 @@ Submodules :maxdepth: 4 CodeEntropy.results.reporter - -Module contents ---------------- - -.. automodule:: CodeEntropy.results - :members: - :show-inheritance: - :undoc-members: diff --git a/docs/api/CodeEntropy.rst b/docs/api/CodeEntropy.rst index 2fb3c5bb..ae75c621 100644 --- a/docs/api/CodeEntropy.rst +++ b/docs/api/CodeEntropy.rst @@ -1,6 +1,11 @@ CodeEntropy package =================== +.. automodule:: CodeEntropy + :members: + :show-inheritance: + :undoc-members: + Subpackages ----------- @@ -22,11 +27,3 @@ Submodules :maxdepth: 4 CodeEntropy.cli - -Module contents ---------------- - -.. automodule:: CodeEntropy - :members: - :show-inheritance: - :undoc-members: diff --git a/docs/api/CodeEntropy.trajectory.rst b/docs/api/CodeEntropy.trajectory.rst index 2ac5a826..14005c6c 100644 --- a/docs/api/CodeEntropy.trajectory.rst +++ b/docs/api/CodeEntropy.trajectory.rst @@ -1,6 +1,11 @@ CodeEntropy.trajectory package ============================== +.. automodule:: CodeEntropy.trajectory + :members: + :show-inheritance: + :undoc-members: + Submodules ---------- @@ -10,11 +15,3 @@ Submodules CodeEntropy.trajectory.frames CodeEntropy.trajectory.mda CodeEntropy.trajectory.source - -Module contents ---------------- - -.. automodule:: CodeEntropy.trajectory - :members: - :show-inheritance: - :undoc-members: diff --git a/docs/getting_started.rst b/docs/getting_started.rst index 192323cc..6ddc6a21 100644 --- a/docs/getting_started.rst +++ b/docs/getting_started.rst @@ -515,7 +515,7 @@ Example 1: DNA Fragment (Smallest / Fastest) Data files: -`DNA fragment example (~1MB) `_ +`DNA fragment example (~1MB) `_ Create or edit ``config.yaml`` in your working directory: @@ -551,7 +551,7 @@ Example 2: Lysozyme (Larger / Slower) Data files: -`Lysozyme example (~1.2GB) `_ +`Lysozyme example (~1.2GB) `_ Create or edit ``config.yaml`` in your working directory: diff --git a/tests/regression/helpers.py b/tests/regression/helpers.py index d5c09a12..9f27b41b 100644 --- a/tests/regression/helpers.py +++ b/tests/regression/helpers.py @@ -11,7 +11,7 @@ import yaml -DEFAULT_TESTDATA_BASE_URL = "https://file-store.ccpbiosim.org/codeentropy-testing" +DEFAULT_TESTDATA_BASE_URL = "https://www.ccpbiosim.ac.uk/file-store/codeentropy-testing" @dataclass(frozen=True) diff --git a/tests/unit/CodeEntropy/levels/dihedrals/test_topology.py b/tests/unit/CodeEntropy/levels/dihedrals/test_topology.py index cbc38b19..60affa71 100644 --- a/tests/unit/CodeEntropy/levels/dihedrals/test_topology.py +++ b/tests/unit/CodeEntropy/levels/dihedrals/test_topology.py @@ -1,472 +1,146 @@ from __future__ import annotations -import logging -from dataclasses import dataclass -from types import SimpleNamespace -from typing import Any -from unittest.mock import Mock, call +from unittest.mock import MagicMock -import pytest +import numpy as np -from CodeEntropy.levels.dihedrals.topology import ( - DihedralTopologyDiscovery, - MoleculeDihedralTopology, -) +from CodeEntropy.levels.dihedrals.topology import DihedralTopologyDiscovery -@dataclass -class FakeAtom: - """Small test double for an MDAnalysis Atom.""" +class _AddableAG: + """Minimal addable AtomGroup test double.""" - index: int - bonded_atoms: Any = None + def __init__(self, name: str) -> None: + """Initialize the fake AtomGroup. + Args: + name: Human-readable identifier used in composed names. + """ + self.name = name -class FakeUniverseAtoms: - """Indexable fake universe atom container.""" + def __add__(self, other: _AddableAG) -> _AddableAG: + """Return a composed fake AtomGroup. - def __init__(self, universe: FakeUniverse) -> None: - self._universe = universe + Args: + other: Fake AtomGroup to combine with this object. - def __getitem__(self, indices: int | list[int] | tuple[int, ...]) -> Any: - if isinstance(indices, int): - return self._universe.atom_by_index[int(indices)] + Returns: + New fake AtomGroup containing a composed name. + """ + return _AddableAG(f"({self.name}+{other.name})") - return FakeAtomGroup( - [self._universe.atom_by_index[int(index)] for index in indices], - universe=self._universe, - ) +class _TopologyDiscovery(DihedralTopologyDiscovery): + """Concrete topology-discovery helper for unit tests.""" -class FakeUniverse: - """Small fake universe supporting ``universe.atoms[indices]``.""" + def __init__(self, universe_operations: MagicMock) -> None: + """Initialize the test helper. - def __init__(self, atoms: list[FakeAtom]) -> None: - self.atom_by_index = {int(atom.index): atom for atom in atoms} - self.atoms = FakeUniverseAtoms(self) + Args: + universe_operations: Mock universe-operation adapter. + """ + self._universe_operations = universe_operations -class FakeAtomGroup: - """Small AtomGroup-like test double for topology discovery tests.""" +def test_select_heavy_residue_builds_expected_selections(): + uops = MagicMock() + helper = _TopologyDiscovery(universe_operations=uops) - def __init__( - self, - atoms: list[FakeAtom], - *, - residues: list[Any] | None = None, - dihedrals: list[Any] | None = None, - select_map: dict[str, Any] | None = None, - universe: FakeUniverse | None = None, - ) -> None: - self._atoms = list(atoms) - self.residues = list(residues or []) - self.dihedrals = list(dihedrals or []) - self._select_map = dict(select_map or {}) - self.universe = universe + mol = MagicMock() + mol.residues = [MagicMock()] + mol.residues[0].atoms.indices = np.array([10, 11, 12], dtype=int) + uops.select_atoms.side_effect = ["residue_atoms", "heavy_atoms"] - @property - def atoms(self) -> FakeAtomGroup: - return self - - @property - def indices(self) -> list[int]: - return [int(atom.index) for atom in self._atoms] - - def __iter__(self): - return iter(self._atoms) - - def __len__(self) -> int: - return len(self._atoms) - - def __add__(self, other: FakeAtomGroup) -> FakeAtomGroup: - return FakeAtomGroup( - self._atoms + other._atoms, - universe=self.universe or other.universe, - ) - - def select_atoms(self, select_string: str, updating: bool = False) -> Any: - if select_string not in self._select_map: - raise AssertionError(f"Unexpected selection: {select_string!r}") - return self._select_map[select_string] - - -@dataclass -class FakeResidue: - """Small residue test double.""" - - atoms: FakeAtomGroup - - -@dataclass -class FakeDihedral: - """Small dihedral topology object.""" - - atoms: FakeAtomGroup - - -def _make_discovery( - universe_operations: Any | None = None, -) -> DihedralTopologyDiscovery: - discovery = DihedralTopologyDiscovery() - discovery._universe_operations = universe_operations or Mock() - return discovery - - -def test_molecule_dihedral_topology_stores_expected_fields() -> None: - topology = MoleculeDihedralTopology( - group_id=1, - molecule_id=2, - molecule_order=3, - num_residues=4, - ua_dihedrals_by_residue={0: ["ua"]}, - residue_dihedrals=["res"], - ) - - assert topology.group_id == 1 - assert topology.molecule_id == 2 - assert topology.molecule_order == 3 - assert topology.num_residues == 4 - assert topology.ua_dihedrals_by_residue == {0: ["ua"]} - assert topology.residue_dihedrals == ["res"] - - -def test_extract_topology_fragment_uses_lightweight_atomgroup_helper() -> None: - universe_operations = Mock() - universe_operations.extract_fragment_atomgroup.return_value = "fragment_atomgroup" - universe_operations.extract_fragment.return_value = "heavy_fragment" - - discovery = _make_discovery(universe_operations) - - result = discovery._extract_topology_fragment("universe", 5) - - assert result == "fragment_atomgroup" - universe_operations.extract_fragment_atomgroup.assert_called_once_with( - "universe", - 5, - ) - universe_operations.extract_fragment.assert_not_called() - - -def test_select_heavy_residue_builds_expected_selections() -> None: - discovery = _make_discovery() - - residue_atoms = Mock() - residue_atoms.indices = [10, 11, 12, 13] - residue = SimpleNamespace(atoms=residue_atoms) - - heavy_atoms = object() - residue_container = Mock() - residue_container.select_atoms.return_value = heavy_atoms - - mol = Mock() - mol.residues = [residue] - mol.select_atoms.return_value = residue_container - - result = discovery._select_heavy_residue(mol, 0) - - assert result is heavy_atoms - mol.select_atoms.assert_called_once_with("index 10:13", updating=False) - residue_container.select_atoms.assert_called_once_with( - "prop mass > 1.1", - updating=False, - ) - discovery._universe_operations.select_atoms.assert_not_called() - - -def test_get_dihedrals_united_atom_collects_atoms_from_dihedral_objects() -> None: - discovery = _make_discovery() - - atoms = [FakeAtom(index) for index in range(1, 8)] - universe = FakeUniverse(atoms) - - valid_dihedral_atoms = FakeAtomGroup( - [atoms[0], atoms[1], atoms[2], atoms[3]], - universe=universe, - ) - outside_selection_atoms = FakeAtomGroup( - [atoms[0], atoms[1], atoms[2], atoms[6]], - universe=universe, - ) - wrong_size_atoms = FakeAtomGroup( - [atoms[0], atoms[1], atoms[2]], - universe=universe, - ) - - selected_heavy_atoms = FakeAtomGroup( - [atoms[0], atoms[1], atoms[2], atoms[3], atoms[4]], - dihedrals=[ - FakeDihedral(valid_dihedral_atoms), - FakeDihedral(outside_selection_atoms), - FakeDihedral(wrong_size_atoms), - ], - universe=universe, - ) - - result = discovery._get_dihedrals(selected_heavy_atoms, "united_atom") - - assert result == [valid_dihedral_atoms] - - -def test_get_dihedrals_united_atom_returns_empty_when_no_valid_dihedrals() -> None: - discovery = _make_discovery() - - atoms = [FakeAtom(index) for index in range(1, 5)] - universe = FakeUniverse(atoms) - - invalid_dihedral_atoms = FakeAtomGroup( - [atoms[0], atoms[1], atoms[2]], - universe=universe, - ) - selected_heavy_atoms = FakeAtomGroup( - atoms, - dihedrals=[FakeDihedral(invalid_dihedral_atoms)], - universe=universe, - ) - - assert discovery._get_dihedrals(selected_heavy_atoms, "united_atom") == [] - - -def test_get_dihedrals_residue_builds_one_dihedral_when_four_residues() -> None: - discovery = _make_discovery() - - atom1 = FakeAtom(10) - atom2 = FakeAtom(20) - atom3 = FakeAtom(30) - atom4 = FakeAtom(40) - - universe = FakeUniverse([atom1, atom2, atom3, atom4]) - - group1 = FakeAtomGroup([atom1], universe=universe) - group2 = FakeAtomGroup([atom2], universe=universe) - group3 = FakeAtomGroup([atom3], universe=universe) - group4 = FakeAtomGroup([atom4], universe=universe) - - atom1.bonded_atoms = group2 - atom2.bonded_atoms = group1 - atom3.bonded_atoms = group4 - atom4.bonded_atoms = group3 - - mol = FakeAtomGroup( - [atom1, atom2, atom3, atom4], - residues=[ - FakeResidue(group1), - FakeResidue(group2), - FakeResidue(group3), - FakeResidue(group4), - ], - universe=universe, - ) - - result = discovery._get_dihedrals(mol, "residue") - - assert len(result) == 1 - assert result[0].indices == [10, 20, 30, 40] - - -def test_get_dihedrals_residue_skips_invalid_four_residue_window( - caplog: pytest.LogCaptureFixture, -) -> None: - caplog.set_level(logging.DEBUG) - - discovery = _make_discovery() - - atom1 = FakeAtom(10) - atom2 = FakeAtom(20) - atom3 = FakeAtom(30) - atom4 = FakeAtom(40) - - universe = FakeUniverse([atom1, atom2, atom3, atom4]) - - group1 = FakeAtomGroup([atom1], universe=universe) - group2 = FakeAtomGroup([atom2], universe=universe) - group3 = FakeAtomGroup([atom3], universe=universe) - group4 = FakeAtomGroup([atom4], universe=universe) - - atom1.bonded_atoms = group2 - atom2.bonded_atoms = group1 - atom3.bonded_atoms = FakeAtomGroup([], universe=universe) - atom4.bonded_atoms = FakeAtomGroup([], universe=universe) - - mol = FakeAtomGroup( - [atom1, atom2, atom3, atom4], - residues=[ - FakeResidue(group1), - FakeResidue(group2), - FakeResidue(group3), - FakeResidue(group4), - ], - universe=universe, - ) - - result = discovery._get_dihedrals(mol, "residue") - - assert result == [] - assert "Skipping residue-level dihedral" in caplog.text - - -def test_get_dihedrals_residue_returns_empty_when_fewer_than_four_residues() -> None: - discovery = _make_discovery() - - atoms = [FakeAtom(1), FakeAtom(2), FakeAtom(3)] - universe = FakeUniverse(atoms) - groups = [FakeAtomGroup([atom], universe=universe) for atom in atoms] - - mol = FakeAtomGroup( - atoms, - residues=[FakeResidue(group) for group in groups], - universe=universe, - ) - - assert discovery._get_dihedrals(mol, "residue") == [] + out = helper._select_heavy_residue(mol, res_id=0) + assert out == "heavy_atoms" + assert uops.select_atoms.call_args_list == [ + ((mol, "index 10:12"),), + (("residue_atoms", "prop mass > 1.1"),), + ] -def test_atoms_in_source_bonded_to_target_returns_matching_source_atoms() -> None: - atom1 = FakeAtom(1) - atom2 = FakeAtom(2) - atom3 = FakeAtom(3) - universe = FakeUniverse([atom1, atom2, atom3]) +def test_get_dihedrals_united_atom_collects_atoms_from_dihedral_objects(): + helper = _TopologyDiscovery(universe_operations=MagicMock()) - source_group = FakeAtomGroup([atom1, atom3], universe=universe) - target_group = FakeAtomGroup([atom2], universe=universe) + d0 = MagicMock() + d0.atoms = "A0" + d1 = MagicMock() + d1.atoms = "A1" - atom1.bonded_atoms = target_group - atom3.bonded_atoms = FakeAtomGroup([], universe=universe) + container = MagicMock() + container.dihedrals = [d0, d1] - source_residue = FakeResidue(source_group) - target_residue = FakeResidue(target_group) + assert helper._get_dihedrals(container, level="united_atom") == ["A0", "A1"] - result = DihedralTopologyDiscovery._atoms_in_source_bonded_to_target( - source_residue, - target_residue, - ) - assert result.indices == [1] +def test_get_dihedrals_residue_returns_empty_when_less_than_four_residues(): + helper = _TopologyDiscovery(universe_operations=MagicMock()) + mol = MagicMock() + mol.residues = [MagicMock(), MagicMock(), MagicMock()] + mol.select_atoms = MagicMock() -def test_atoms_in_source_bonded_to_target_returns_empty_when_atom_has_no_bonds() -> ( - None -): - atom1 = FakeAtom(1) - atom2 = FakeAtom(2) + assert helper._get_dihedrals(mol, level="residue") == [] + mol.select_atoms.assert_not_called() - universe = FakeUniverse([atom1, atom2]) - source_residue = FakeResidue(FakeAtomGroup([atom1], universe=universe)) - target_residue = FakeResidue(FakeAtomGroup([atom2], universe=universe)) +def test_get_dihedrals_residue_builds_one_dihedral_when_four_residues(): + helper = _TopologyDiscovery(universe_operations=MagicMock()) - result = DihedralTopologyDiscovery._atoms_in_source_bonded_to_target( - source_residue, - target_residue, + mol = MagicMock() + mol.residues = [MagicMock(), MagicMock(), MagicMock(), MagicMock()] + mol.select_atoms = MagicMock( + side_effect=[ + _AddableAG("a1"), + _AddableAG("a2"), + _AddableAG("a3"), + _AddableAG("a4"), + ] ) - assert result.indices == [] - + out = helper._get_dihedrals(mol, level="residue") -def test_discover_group_dihedral_topology_builds_one_entry_per_molecule( - monkeypatch: pytest.MonkeyPatch, -) -> None: - universe_operations = Mock() + assert len(out) == 1 + assert isinstance(out[0], _AddableAG) + assert mol.select_atoms.call_count == 4 - molecule0 = SimpleNamespace(label="molecule0", residues=[object(), object()]) - molecule1 = SimpleNamespace(label="molecule1", residues=[object()]) - universe_operations.extract_fragment_atomgroup.side_effect = [ - molecule0, - molecule1, - ] - - discovery = _make_discovery(universe_operations) - - def fake_select_heavy_residue(mol: Any, res_id: int) -> str: - return f"{mol.label}-heavy-{res_id}" +def test_discover_group_dihedral_topology_builds_one_entry_per_molecule(): + uops = MagicMock() + helper = _TopologyDiscovery(universe_operations=uops) - def fake_get_dihedrals(data_container: Any, level: str) -> list[str]: - if level == "united_atom": - return [f"ua-{data_container}"] - return [f"res-{data_container.label}"] + mol0 = MagicMock() + mol0.residues = [MagicMock(), MagicMock()] + mol1 = MagicMock() + mol1.residues = [MagicMock(), MagicMock()] + uops.extract_fragment.side_effect = [mol0, mol1] - monkeypatch.setattr( - discovery, - "_select_heavy_residue", - fake_select_heavy_residue, + helper._select_heavy_residue = MagicMock( + side_effect=["heavy0", "heavy1", "heavy2", "heavy3"] ) - monkeypatch.setattr( - discovery, - "_get_dihedrals", - fake_get_dihedrals, + helper._get_dihedrals = MagicMock( + side_effect=[ + ["ua0r0"], + ["ua0r1"], + ["res0"], + ["ua1r0"], + ["ua1r1"], + ["res1"], + ] ) - topologies = discovery._discover_group_dihedral_topology( + topologies = helper._discover_group_dihedral_topology( data_container="universe", - group_id=9, - molecules=[100, 200], + group_id=3, + molecules=[7, 8], level_list=["united_atom", "residue"], ) - assert topologies == [ - MoleculeDihedralTopology( - group_id=9, - molecule_id=100, - molecule_order=0, - num_residues=2, - ua_dihedrals_by_residue={ - 0: ["ua-molecule0-heavy-0"], - 1: ["ua-molecule0-heavy-1"], - }, - residue_dihedrals=["res-molecule0"], - ), - MoleculeDihedralTopology( - group_id=9, - molecule_id=200, - molecule_order=1, - num_residues=1, - ua_dihedrals_by_residue={ - 0: ["ua-molecule1-heavy-0"], - }, - residue_dihedrals=["res-molecule1"], - ), - ] - - assert universe_operations.extract_fragment_atomgroup.call_args_list == [ - call("universe", 100), - call("universe", 200), - ] - - -def test_discover_group_dihedral_topology_respects_enabled_levels( - monkeypatch: pytest.MonkeyPatch, -) -> None: - universe_operations = Mock() - - molecule = SimpleNamespace(label="molecule", residues=[object(), object()]) - universe_operations.extract_fragment_atomgroup.return_value = molecule - - discovery = _make_discovery(universe_operations) - - select_heavy_residue = Mock(return_value="heavy-residue") - get_dihedrals = Mock(return_value=["residue-dihedral"]) - - monkeypatch.setattr(discovery, "_select_heavy_residue", select_heavy_residue) - monkeypatch.setattr(discovery, "_get_dihedrals", get_dihedrals) - - topologies = discovery._discover_group_dihedral_topology( - data_container="universe", - group_id=1, - molecules=[0], - level_list=["residue"], - ) - - assert topologies == [ - MoleculeDihedralTopology( - group_id=1, - molecule_id=0, - molecule_order=0, - num_residues=2, - ua_dihedrals_by_residue={}, - residue_dihedrals=["residue-dihedral"], - ) - ] - - select_heavy_residue.assert_not_called() - get_dihedrals.assert_called_once_with(molecule, "residue") + assert [topology.molecule_id for topology in topologies] == [7, 8] + assert [topology.molecule_order for topology in topologies] == [0, 1] + assert topologies[0].group_id == 3 + assert topologies[0].ua_dihedrals_by_residue == {0: ["ua0r0"], 1: ["ua0r1"]} + assert topologies[0].residue_dihedrals == ["res0"] + assert topologies[1].ua_dihedrals_by_residue == {0: ["ua1r0"], 1: ["ua1r1"]} + assert topologies[1].residue_dihedrals == ["res1"] diff --git a/tests/unit/CodeEntropy/levels/nodes/test_axes_topology.py b/tests/unit/CodeEntropy/levels/nodes/test_axes_topology.py deleted file mode 100644 index 2c319985..00000000 --- a/tests/unit/CodeEntropy/levels/nodes/test_axes_topology.py +++ /dev/null @@ -1,415 +0,0 @@ -from __future__ import annotations - -from types import SimpleNamespace - -import numpy as np - -from CodeEntropy.levels.nodes.axes_topology import ( - AxesTopology, - BuildAxesTopologyNode, - ResidueAxesTopology, - UAAxesTopology, -) - - -class FakeAtomGroup: - """Small AtomGroup-like object for axes-topology tests.""" - - def __init__(self, atoms=None, *, name="ag"): - self._atoms = list(atoms or []) - self.name = name - self.indices = np.asarray([atom.index for atom in self._atoms], dtype=int) - - def __iter__(self): - return iter(self._atoms) - - def __len__(self): - return len(self._atoms) - - def __getitem__(self, index): - if isinstance(index, (list, tuple, np.ndarray)): - return FakeAtomGroup([self._atoms[int(i)] for i in index]) - return self._atoms[int(index)] - - def select_atoms(self, selection): - """Return atoms matching the small mass selections used by the node.""" - if selection == "prop mass > 1.1": - return FakeAtomGroup([atom for atom in self._atoms if atom.mass > 1.1]) - if selection == "mass 2 to 999": - return FakeAtomGroup( - [atom for atom in self._atoms if 2.0 <= atom.mass <= 999.0] - ) - if selection == "mass 1 to 1.1": - return FakeAtomGroup( - [atom for atom in self._atoms if 1.0 <= atom.mass <= 1.1] - ) - raise AssertionError(f"Unexpected selection: {selection}") - - -class FakeAtom: - """Small Atom-like object with mass, index, and bonded atoms.""" - - def __init__(self, index, mass, bonded_atoms=None): - self.index = index - self.mass = mass - self.bonded_atoms = bonded_atoms or FakeAtomGroup([]) - - -class FakeResidue: - """Small residue-like object.""" - - def __init__(self, atoms): - self.atoms = atoms - - -class FakeMolecule: - """Small molecule-like object with residues and selections.""" - - def __init__(self, residues, select_map=None): - self.residues = residues - self._select_map = dict(select_map or {}) - - def select_atoms(self, selection): - return self._select_map.get(selection, FakeAtomGroup([])) - - -class FakeAtoms: - """Container supporting fragments and u.atoms[index_array].""" - - def __init__(self, fragments, atom_map): - self.fragments = fragments - self._atom_map = dict(atom_map) - - def __getitem__(self, indices): - if isinstance(indices, np.ndarray): - return FakeAtomGroup([self._atom_map[int(index)] for index in indices]) - if isinstance(indices, (list, tuple)): - return FakeAtomGroup([self._atom_map[int(index)] for index in indices]) - return self._atom_map[int(indices)] - - -class FakeUniverse: - """Small universe-like object.""" - - def __init__(self, fragments, atom_map): - self.atoms = FakeAtoms(fragments=fragments, atom_map=atom_map) - - -def _args(*, customised_axes): - return SimpleNamespace(customised_axes=customised_axes) - - -def _single_molecule_universe(): - """Build a small molecule containing one residue and one UA bead.""" - hydrogen = FakeAtom(index=2, mass=1.0) - bonded_heavy = FakeAtom(index=3, mass=12.0) - heavy = FakeAtom( - index=1, - mass=12.0, - bonded_atoms=FakeAtomGroup([hydrogen, bonded_heavy]), - ) - other_residue_heavy = FakeAtom(index=4, mass=14.0) - - residue_atoms = FakeAtomGroup([heavy, hydrogen, bonded_heavy, other_residue_heavy]) - residue = FakeResidue(residue_atoms) - molecule = FakeMolecule([residue]) - - atom_map = { - heavy.index: heavy, - hydrogen.index: hydrogen, - bonded_heavy.index: bonded_heavy, - other_residue_heavy.index: other_residue_heavy, - } - universe = FakeUniverse([molecule], atom_map) - - return universe, molecule, heavy, hydrogen, bonded_heavy, other_residue_heavy - - -def test_ua_axes_topology_dataclass_preserves_arrays(): - topology = UAAxesTopology( - heavy_atom_index=1, - ua_atom_indices=np.array([1, 2]), - ua_all_atom_indices=np.array([1, 3, 2]), - bonded_heavy_indices=np.array([3]), - bonded_light_indices=np.array([2]), - residue_heavy_indices=np.array([1, 3, 4]), - residue_ua_masses=np.array([13.0, 12.0, 14.0]), - ) - - assert topology.heavy_atom_index == 1 - np.testing.assert_array_equal(topology.ua_atom_indices, np.array([1, 2])) - np.testing.assert_array_equal(topology.ua_all_atom_indices, np.array([1, 3, 2])) - np.testing.assert_array_equal(topology.bonded_heavy_indices, np.array([3])) - np.testing.assert_array_equal(topology.bonded_light_indices, np.array([2])) - np.testing.assert_array_equal(topology.residue_heavy_indices, np.array([1, 3, 4])) - np.testing.assert_allclose(topology.residue_ua_masses, np.array([13.0, 12.0, 14.0])) - - -def test_residue_axes_topology_dataclass_preserves_arrays(): - topology = ResidueAxesTopology( - residue_heavy_indices=np.array([1, 3, 4]), - residue_ua_masses=np.array([13.0, 12.0, 14.0]), - has_neighbor_bonds=True, - ) - - np.testing.assert_array_equal(topology.residue_heavy_indices, np.array([1, 3, 4])) - np.testing.assert_allclose(topology.residue_ua_masses, np.array([13.0, 12.0, 14.0])) - assert topology.has_neighbor_bonds is True - - -def test_axes_topology_defaults_to_empty_mappings(): - topology = AxesTopology() - - assert topology.ua == {} - assert topology.residue == {} - - -def test_run_writes_empty_topology_when_customised_axes_disabled(): - node = BuildAxesTopologyNode() - shared_data = {"args": _args(customised_axes=False)} - - result = node.run(shared_data) - - assert isinstance(result["axes_topology"], AxesTopology) - assert result["axes_topology"].ua == {} - assert result["axes_topology"].residue == {} - assert shared_data["axes_topology"] is result["axes_topology"] - - -def test_run_builds_residue_topology_for_residue_levels(): - node = BuildAxesTopologyNode() - universe, molecule, heavy, hydrogen, bonded_heavy, other_residue_heavy = ( - _single_molecule_universe() - ) - neighbor_query = "(resindex -1 or resindex 1) and bonded resid 0" - molecule._select_map[neighbor_query] = FakeAtomGroup([bonded_heavy]) - shared_data = { - "args": _args(customised_axes=True), - "reduced_universe": universe, - "levels": [["residue"]], - "beads": {(0, "residue"): [np.array([1, 2, 3, 4])]}, - } - - result = node.run(shared_data) - - axes_topology = result["axes_topology"] - assert axes_topology.ua == {} - assert set(axes_topology.residue) == {(0, 0)} - - residue_topology = axes_topology.residue[(0, 0)] - np.testing.assert_array_equal( - residue_topology.residue_heavy_indices, - np.array([heavy.index, bonded_heavy.index, other_residue_heavy.index]), - ) - np.testing.assert_allclose( - residue_topology.residue_ua_masses, - np.array([13.0, 12.0, 14.0]), - ) - assert residue_topology.has_neighbor_bonds is True - - -def test_add_residue_topology_skips_when_no_residue_beads(): - node = BuildAxesTopologyNode() - _, molecule, _, _, _, _ = _single_molecule_universe() - out = {} - - node._add_residue_topology( - mol=molecule, - mol_id=0, - beads={}, - out=out, - ) - - assert out == {} - - -def test_has_neighbor_bonds_uses_original_residue_selection(): - node = BuildAxesTopologyNode() - query = "(resindex 0 or resindex 2) and bonded resid 1" - molecule = FakeMolecule([], select_map={query: FakeAtomGroup([FakeAtom(1, 12.0)])}) - - assert node._has_neighbor_bonds(mol=molecule, local_res_i=1) is True - - -def test_has_neighbor_bonds_returns_false_for_empty_selection(): - node = BuildAxesTopologyNode() - molecule = FakeMolecule([]) - - assert node._has_neighbor_bonds(mol=molecule, local_res_i=1) is False - - -def test_run_builds_ua_topology_for_united_atom_levels(): - node = BuildAxesTopologyNode() - universe, _, heavy, hydrogen, bonded_heavy, other_residue_heavy = ( - _single_molecule_universe() - ) - shared_data = { - "args": _args(customised_axes=True), - "reduced_universe": universe, - "levels": [["united_atom"]], - "beads": {(0, "united_atom", 0): [np.array([1, 2])]}, - } - - result = node.run(shared_data) - - axes_topology = result["axes_topology"] - assert shared_data["axes_topology"] is axes_topology - assert set(axes_topology.ua) == {(0, 0, 0)} - - ua_topology = axes_topology.ua[(0, 0, 0)] - assert ua_topology.heavy_atom_index == heavy.index - np.testing.assert_array_equal(ua_topology.ua_atom_indices, np.array([1, 2])) - np.testing.assert_array_equal(ua_topology.ua_all_atom_indices, np.array([1, 3, 2])) - np.testing.assert_array_equal(ua_topology.bonded_heavy_indices, np.array([3])) - np.testing.assert_array_equal(ua_topology.bonded_light_indices, np.array([2])) - np.testing.assert_array_equal( - ua_topology.residue_heavy_indices, - np.array([heavy.index, bonded_heavy.index, other_residue_heavy.index]), - ) - np.testing.assert_allclose( - ua_topology.residue_ua_masses, - np.array([13.0, 12.0, 14.0]), - ) - - -def test_run_ignores_molecules_without_united_atom_level(): - node = BuildAxesTopologyNode() - universe, _, _, _, _, _ = _single_molecule_universe() - shared_data = { - "args": _args(customised_axes=True), - "reduced_universe": universe, - "levels": [["residue"]], - "beads": {(0, "united_atom", 0): [np.array([1, 2])]}, - } - - result = node.run(shared_data) - - assert result["axes_topology"].ua == {} - assert shared_data["axes_topology"].ua == {} - - -def test_add_ua_topology_skips_residues_without_beads(): - node = BuildAxesTopologyNode() - universe, molecule, _, _, _, _ = _single_molecule_universe() - out = {} - - node._add_ua_topology( - u=universe, - mol=molecule, - mol_id=0, - beads={}, - out=out, - ) - - assert out == {} - - -def test_add_ua_topology_skips_ua_beads_without_heavy_atoms(caplog): - node = BuildAxesTopologyNode() - hydrogen = FakeAtom(index=2, mass=1.0) - residue = FakeResidue(FakeAtomGroup([hydrogen])) - molecule = FakeMolecule([residue]) - universe = FakeUniverse([molecule], {2: hydrogen}) - out = {} - - node._add_ua_topology( - u=universe, - mol=molecule, - mol_id=5, - beads={(5, "united_atom", 0): [np.array([2])]}, - out=out, - ) - - assert out == {} - assert "Skipping UA axes topology with no heavy atom" in caplog.text - - -def test_add_ua_topology_handles_multiple_residues_and_ua_beads(): - node = BuildAxesTopologyNode() - - h0 = FakeAtom(index=10, mass=1.0) - c0 = FakeAtom(index=11, mass=12.0, bonded_atoms=FakeAtomGroup([h0])) - residue0 = FakeResidue(FakeAtomGroup([c0, h0])) - - h1 = FakeAtom(index=20, mass=1.0) - c1 = FakeAtom(index=21, mass=12.0, bonded_atoms=FakeAtomGroup([h1])) - residue1 = FakeResidue(FakeAtomGroup([c1, h1])) - - molecule = FakeMolecule([residue0, residue1]) - universe = FakeUniverse( - [molecule], - { - h0.index: h0, - c0.index: c0, - h1.index: h1, - c1.index: c1, - }, - ) - out = {} - - node._add_ua_topology( - u=universe, - mol=molecule, - mol_id=3, - beads={ - (3, "united_atom", 0): [np.array([11, 10])], - (3, "united_atom", 1): [np.array([21, 20])], - }, - out=out, - ) - - assert set(out) == {(3, 0, 0), (3, 1, 0)} - assert out[(3, 0, 0)].heavy_atom_index == 11 - assert out[(3, 1, 0)].heavy_atom_index == 21 - - -def test_split_bonded_atoms_returns_heavy_and_light_atom_groups(): - hydrogen = FakeAtom(index=2, mass=1.0) - heavy_bonded = FakeAtom(index=3, mass=12.0) - atom = FakeAtom( - index=1, - mass=12.0, - bonded_atoms=FakeAtomGroup([hydrogen, heavy_bonded]), - ) - - bonded_heavy, bonded_light = BuildAxesTopologyNode._split_bonded_atoms(atom) - - np.testing.assert_array_equal(bonded_heavy.indices, np.array([3])) - np.testing.assert_array_equal(bonded_light.indices, np.array([2])) - - -def test_get_ua_masses_from_topology_adds_bonded_hydrogen_masses(): - hydrogen = FakeAtom(index=2, mass=1.0) - heavy = FakeAtom(index=1, mass=12.0, bonded_atoms=FakeAtomGroup([hydrogen])) - other_heavy = FakeAtom(index=3, mass=14.0) - atom_group = FakeAtomGroup([heavy, hydrogen, other_heavy]) - - masses = BuildAxesTopologyNode._get_ua_masses_from_topology(atom_group) - - assert masses == [13.0, 14.0] - - -def test_get_ua_masses_from_topology_handles_atoms_without_bonded_atoms_attribute(): - class AtomWithoutBonds: - """Small atom-like object without bonded_atoms.""" - - def __init__(self, index, mass): - self.index = index - self.mass = mass - - heavy = AtomWithoutBonds(index=1, mass=12.0) - hydrogen = AtomWithoutBonds(index=2, mass=1.0) - atom_group = FakeAtomGroup([heavy, hydrogen]) - - masses = BuildAxesTopologyNode._get_ua_masses_from_topology(atom_group) - - assert masses == [12.0] - - -def test_get_ua_masses_from_topology_returns_empty_list_for_light_atoms_only(): - hydrogen = FakeAtom(index=2, mass=1.0) - atom_group = FakeAtomGroup([hydrogen]) - - masses = BuildAxesTopologyNode._get_ua_masses_from_topology(atom_group) - - assert masses == [] diff --git a/tests/unit/CodeEntropy/levels/nodes/test_covariance_node.py b/tests/unit/CodeEntropy/levels/nodes/test_covariance_node.py index b0fb45f1..d989069f 100644 --- a/tests/unit/CodeEntropy/levels/nodes/test_covariance_node.py +++ b/tests/unit/CodeEntropy/levels/nodes/test_covariance_node.py @@ -1,3 +1,5 @@ +"""Atomic unit tests for frame-local covariance construction.""" + from __future__ import annotations from types import SimpleNamespace @@ -86,7 +88,6 @@ def test_run_processes_all_levels_and_writes_frame_covariance(): mol = FakeMolecule() universe = FakeUniverse([mol], dimensions=np.array([10.0, 20.0, 30.0, 90.0])) axes_manager = object() - axes_topology = object() ctx = { "shared": { @@ -96,7 +97,6 @@ def test_run_processes_all_levels_and_writes_frame_covariance(): "beads": {}, "args": _args(combined_forcetorque=True, customised_axes=True), "axes_manager": axes_manager, - "axes_topology": axes_topology, } } @@ -115,14 +115,10 @@ def test_run_processes_all_levels_and_writes_frame_covariance(): assert ua_kwargs["mol_id"] == 0 assert ua_kwargs["group_id"] == 7 assert ua_kwargs["axes_manager"] is axes_manager - assert ua_kwargs["axes_topology"] is axes_topology assert ua_kwargs["force_partitioning"] == 0.5 assert ua_kwargs["customised_axes"] is True assert ua_kwargs["is_highest"] is False - res_kwargs = node._process_residue.call_args.kwargs - assert res_kwargs["axes_topology"] is axes_topology - def test_run_omits_forcetorque_when_combined_is_false(): node = FrameCovarianceNode() @@ -169,7 +165,6 @@ def test_process_united_atom_updates_outputs_and_molcount(): group_id=7, beads={(0, "united_atom", 0): [np.array([0])]}, axes_manager="axes", - axes_topology=None, box=None, force_partitioning=0.5, customised_axes=False, @@ -197,7 +192,6 @@ def test_process_united_atom_returns_when_no_beads_or_empty_atom_groups(): group_id=7, beads={}, axes_manager=None, - axes_topology=None, box=None, force_partitioning=0.5, customised_axes=False, @@ -219,7 +213,6 @@ def test_process_united_atom_returns_when_no_beads_or_empty_atom_groups(): group_id=7, beads={(0, "united_atom", 0): [np.array([0])]}, axes_manager=None, - axes_topology=None, box=None, force_partitioning=0.5, customised_axes=False, @@ -256,7 +249,6 @@ def test_process_residue_updates_outputs_and_combined_ft(): group_id=7, beads={(0, "residue"): [np.array([0])]}, axes_manager="axes", - axes_topology=None, box=None, customised_axes=True, force_partitioning=0.5, @@ -286,7 +278,6 @@ def test_process_residue_returns_when_no_beads_or_empty_groups(): group_id=7, beads={}, axes_manager=None, - axes_topology=None, box=None, customised_axes=False, force_partitioning=0.5, @@ -308,7 +299,6 @@ def test_process_residue_returns_when_no_beads_or_empty_groups(): group_id=7, beads={(0, "residue"): [np.array([0])]}, axes_manager=None, - axes_topology=None, box=None, customised_axes=False, force_partitioning=0.5, @@ -424,13 +414,9 @@ def test_build_ua_vectors_uses_customised_axes(): node._ft.get_weighted_torques = MagicMock(return_value=np.array([0.0, 1.0, 0.0])) force_vecs, torque_vecs = node._build_ua_vectors( - u=FakeUniverse([]), - mol_id=0, - local_res_i=0, bead_groups=[FakeAtomGroup("ua")], residue_group=FakeAtomGroup("res"), axes_manager=axes_manager, - axes_topology=None, box=None, force_partitioning=0.5, customised_axes=True, @@ -443,49 +429,6 @@ def test_build_ua_vectors_uses_customised_axes(): axes_manager.get_UA_axes.assert_called_once() -def test_build_ua_vectors_uses_cached_axes_topology_when_available(): - node = FrameCovarianceNode() - axes_manager = MagicMock() - - u = FakeUniverse([]) - ua_topology = object() - axes_topology = SimpleNamespace(ua={(3, 4, 0): ua_topology}) - - axes_manager.get_UA_axes_from_topology.return_value = ( - np.eye(3), - 2.0 * np.eye(3), - np.ones(3), - np.array([1.0, 2.0, 3.0]), - ) - node._ft.get_weighted_forces = MagicMock(return_value=np.array([1.0, 0.0, 0.0])) - node._ft.get_weighted_torques = MagicMock(return_value=np.array([0.0, 1.0, 0.0])) - - force_vecs, torque_vecs = node._build_ua_vectors( - u=u, - mol_id=3, - local_res_i=4, - bead_groups=[FakeAtomGroup("ua")], - residue_atoms=FakeAtomGroup("res"), - axes_manager=axes_manager, - axes_topology=axes_topology, - box=None, - force_partitioning=0.5, - customised_axes=True, - is_highest=True, - ) - - assert len(force_vecs) == 1 - assert len(torque_vecs) == 1 - - called_kwargs = axes_manager.get_UA_axes_from_topology.call_args.kwargs - assert called_kwargs["u"] is u - assert called_kwargs["topology"] is ua_topology - assert called_kwargs["box"] is None - assert called_kwargs["residue_atoms"].name == "res" - - axes_manager.get_UA_axes.assert_not_called() - - def test_build_ua_vectors_uses_vanilla_axes_when_not_customised(): node = FrameCovarianceNode() axes_manager = MagicMock() @@ -501,13 +444,9 @@ def test_build_ua_vectors_uses_vanilla_axes_when_not_customised(): with patch("CodeEntropy.levels.nodes.covariance.make_whole") as make_whole: node._build_ua_vectors( - u=FakeUniverse([]), - mol_id=0, - local_res_i=0, bead_groups=[FakeAtomGroup("ua")], residue_group=FakeAtomGroup("res"), axes_manager=axes_manager, - axes_topology=None, box=None, force_partitioning=0.5, customised_axes=False, @@ -531,12 +470,9 @@ def test_build_residue_vectors_uses_residue_axes(): node._ft.get_weighted_torques = MagicMock(return_value=np.array([0.0, 1.0, 0.0])) force_vecs, torque_vecs = node._build_residue_vectors( - u=FakeUniverse([mol]), mol=mol, - mol_id=0, bead_groups=[FakeAtomGroup("res")], axes_manager=axes_manager, - axes_topology=None, box=None, customised_axes=True, force_partitioning=0.5, @@ -548,34 +484,6 @@ def test_build_residue_vectors_uses_residue_axes(): node._get_residue_axes.assert_called_once() -def test_get_residue_axes_customised_uses_cached_topology_when_available(): - node = FrameCovarianceNode() - mol = FakeMolecule(n_residues=1) - axes_manager = MagicMock() - expected = (np.eye(3), np.eye(3) * 2.0, np.zeros(3), np.ones(3)) - residue_topology = object() - axes_topology = SimpleNamespace(residue={(3, 0): residue_topology}) - axes_manager.get_residue_axes_from_topology.return_value = expected - - result = node._get_residue_axes( - u=FakeUniverse([mol]), - mol=mol, - mol_id=3, - bead=FakeAtomGroup("res"), - local_res_i=0, - axes_manager=axes_manager, - axes_topology=axes_topology, - box=None, - customised_axes=True, - ) - - assert result == expected - called_kwargs = axes_manager.get_residue_axes_from_topology.call_args.kwargs - assert called_kwargs["topology"] is residue_topology - assert called_kwargs["residue_atoms"] is mol.residues[0].atoms - axes_manager.get_residue_axes.assert_not_called() - - def test_get_residue_axes_customised_delegates_to_axes_manager(): node = FrameCovarianceNode() mol = FakeMolecule(n_residues=1) @@ -585,14 +493,10 @@ def test_get_residue_axes_customised_delegates_to_axes_manager(): assert ( node._get_residue_axes( - u=FakeUniverse([mol]), mol=mol, - mol_id=0, bead=FakeAtomGroup("res"), local_res_i=0, axes_manager=axes_manager, - axes_topology=None, - box=None, customised_axes=True, ) == expected @@ -617,14 +521,10 @@ def test_get_residue_axes_vanilla_uses_make_whole_and_vanilla_axes(): with patch("CodeEntropy.levels.nodes.covariance.make_whole") as make_whole: trans_axes, rot_axes, center, moi = node._get_residue_axes( - u=FakeUniverse([mol]), mol=mol, - mol_id=0, bead=bead, local_res_i=0, axes_manager=axes_manager, - axes_topology=None, - box=None, customised_axes=False, ) diff --git a/tests/unit/CodeEntropy/levels/test_axes.py b/tests/unit/CodeEntropy/levels/test_axes.py index 2e42663f..4c4a894a 100644 --- a/tests/unit/CodeEntropy/levels/test_axes.py +++ b/tests/unit/CodeEntropy/levels/test_axes.py @@ -4,7 +4,6 @@ import pytest from CodeEntropy.levels.axes import AxesCalculator -from CodeEntropy.levels.nodes.axes_topology import ResidueAxesTopology, UAAxesTopology class _FakeAtom: @@ -717,499 +716,3 @@ def test_get_bonded_axes_returns_none_none_if_custom_axes_none(monkeypatch): assert custom_axes is None assert moi is None - - -class _FakeIndexedAtoms: - """Container supporting ``u.atoms[index]`` and ``u.atoms[index_array]``.""" - - def __init__(self, atom_map): - self._atom_map = dict(atom_map) - - def __getitem__(self, index): - if isinstance(index, np.ndarray): - return _FakeAtomGroup([self._atom_map[int(i)] for i in index]) - if isinstance(index, (list, tuple)): - return _FakeAtomGroup([self._atom_map[int(i)] for i in index]) - return self._atom_map[int(index)] - - -class _FakeUniverse: - """Small universe-like object with indexed atoms and dimensions.""" - - def __init__(self, atom_map, dimensions=None): - self.atoms = _FakeIndexedAtoms(atom_map) - self.dimensions = np.asarray( - dimensions - if dimensions is not None - else [10.0, 10.0, 10.0, 90.0, 90.0, 90.0], - dtype=float, - ) - - -def _ua_topology( - *, - heavy_atom_index=1, - ua_atom_indices=(1,), - ua_all_atom_indices=(1,), - bonded_heavy_indices=(), - bonded_light_indices=(), - residue_heavy_indices=(1,), - residue_ua_masses=(12.0,), -): - """Build a small cached UA topology fixture.""" - return UAAxesTopology( - heavy_atom_index=int(heavy_atom_index), - ua_atom_indices=np.asarray(ua_atom_indices, dtype=int), - ua_all_atom_indices=np.asarray(ua_all_atom_indices, dtype=int), - bonded_heavy_indices=np.asarray(bonded_heavy_indices, dtype=int), - bonded_light_indices=np.asarray(bonded_light_indices, dtype=int), - residue_heavy_indices=np.asarray(residue_heavy_indices, dtype=int), - residue_ua_masses=np.asarray(residue_ua_masses, dtype=float), - ) - - -def _residue_topology( - *, - residue_heavy_indices=(1,), - residue_ua_masses=(12.0,), - has_neighbor_bonds=False, -): - """Build a small cached residue topology fixture.""" - return ResidueAxesTopology( - residue_heavy_indices=np.asarray(residue_heavy_indices, dtype=int), - residue_ua_masses=np.asarray(residue_ua_masses, dtype=float), - has_neighbor_bonds=bool(has_neighbor_bonds), - ) - - -def test_get_residue_axes_from_topology_no_neighbor_bonds_uses_cached_indices( - monkeypatch, -): - ax = AxesCalculator() - - heavy_atom = _FakeAtom(1, 12.0, [1.0, 2.0, 3.0]) - other_heavy = _FakeAtom(3, 14.0, [4.0, 5.0, 6.0]) - universe = _FakeUniverse({1: heavy_atom, 3: other_heavy}) - mol = MagicMock() - residue_atoms = MagicMock() - residue_atoms.center_of_mass.return_value = np.array([9.0, 8.0, 7.0]) - topology = _residue_topology( - residue_heavy_indices=(1, 3), - residue_ua_masses=(13.0, 14.0), - has_neighbor_bonds=False, - ) - - get_tensor = MagicMock(return_value=np.eye(3)) - get_principal = MagicMock(return_value=(np.eye(3) * 2.0, np.array([3.0, 2.0, 1.0]))) - - monkeypatch.setattr(ax, "get_moment_of_inertia_tensor", get_tensor) - monkeypatch.setattr(ax, "get_custom_principal_axes", get_principal) - - box = np.array([20.0, 30.0, 40.0]) - trans_axes, rot_axes, center, moi = ax.get_residue_axes_from_topology( - u=universe, - mol=mol, - residue_atoms=residue_atoms, - topology=topology, - box=box, - ) - - np.testing.assert_allclose(trans_axes, np.eye(3) * 2.0) - np.testing.assert_allclose(rot_axes, np.eye(3) * 2.0) - np.testing.assert_allclose(center, np.array([9.0, 8.0, 7.0])) - np.testing.assert_allclose(moi, np.array([3.0, 2.0, 1.0])) - - tensor_kwargs = get_tensor.call_args.kwargs - np.testing.assert_allclose( - tensor_kwargs["center_of_mass"], - np.array([9.0, 8.0, 7.0]), - ) - np.testing.assert_allclose( - tensor_kwargs["positions"], - np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]), - ) - np.testing.assert_allclose(tensor_kwargs["masses"], np.array([13.0, 14.0])) - np.testing.assert_allclose(tensor_kwargs["dimensions"], box) - get_principal.assert_called_once() - - -def test_get_residue_axes_from_topology_neighbor_bonds_uses_vanilla_axes( - monkeypatch, -): - ax = AxesCalculator() - - universe = _FakeUniverse( - {}, - dimensions=[11.0, 12.0, 13.0, 90.0, 90.0, 90.0], - ) - mol = MagicMock() - mol.atoms.principal_axes.return_value = np.eye(3) * 5.0 - residue_atoms = MagicMock() - residue_atoms.center_of_mass.return_value = np.array([1.0, 2.0, 3.0]) - topology = _residue_topology(has_neighbor_bonds=True) - - make_whole = MagicMock() - get_vanilla = MagicMock(return_value=(np.eye(3) * 6.0, np.array([6.0, 5.0, 4.0]))) - - monkeypatch.setattr("CodeEntropy.levels.axes.make_whole", make_whole) - monkeypatch.setattr(ax, "get_vanilla_axes", get_vanilla) - - trans_axes, rot_axes, center, moi = ax.get_residue_axes_from_topology( - u=universe, - mol=mol, - residue_atoms=residue_atoms, - topology=topology, - box=None, - ) - - make_whole.assert_called_once_with(mol.atoms) - mol.atoms.principal_axes.assert_called_once() - get_vanilla.assert_called_once_with(residue_atoms) - np.testing.assert_allclose(trans_axes, np.eye(3) * 5.0) - np.testing.assert_allclose(rot_axes, np.eye(3) * 6.0) - np.testing.assert_allclose(center, np.array([1.0, 2.0, 3.0])) - np.testing.assert_allclose(moi, np.array([6.0, 5.0, 4.0])) - - -def test_get_UA_axes_from_topology_multiple_heavy_uses_cached_indices_and_box( - monkeypatch, -): - ax = AxesCalculator() - - heavy_atom = _FakeAtom(1, 12.0, [1.0, 2.0, 3.0]) - other_heavy = _FakeAtom(3, 14.0, [4.0, 5.0, 6.0]) - universe = _FakeUniverse({1: heavy_atom, 3: other_heavy}) - residue_atoms = MagicMock() - residue_atoms.center_of_mass.return_value = np.array([9.0, 8.0, 7.0]) - - topology = _ua_topology( - heavy_atom_index=1, - residue_heavy_indices=(1, 3), - residue_ua_masses=(13.0, 14.0), - ) - - get_tensor = MagicMock(return_value=np.eye(3)) - get_principal = MagicMock(return_value=(np.eye(3) * 2.0, np.array([3.0, 2.0, 1.0]))) - get_bonded = MagicMock(return_value=(np.eye(3) * 4.0, np.array([1.0, 1.0, 1.0]))) - - monkeypatch.setattr(ax, "get_moment_of_inertia_tensor", get_tensor) - monkeypatch.setattr(ax, "get_custom_principal_axes", get_principal) - monkeypatch.setattr(ax, "get_bonded_axes_from_topology", get_bonded) - - box = np.array([20.0, 30.0, 40.0]) - trans_axes, rot_axes, center, moi = ax.get_UA_axes_from_topology( - u=universe, - residue_atoms=residue_atoms, - topology=topology, - box=box, - ) - - np.testing.assert_allclose(trans_axes, np.eye(3) * 2.0) - np.testing.assert_allclose(rot_axes, np.eye(3) * 4.0) - np.testing.assert_allclose(center, heavy_atom.position) - np.testing.assert_allclose(moi, np.array([1.0, 1.0, 1.0])) - - get_tensor.assert_called_once() - tensor_kwargs = get_tensor.call_args.kwargs - np.testing.assert_allclose( - tensor_kwargs["center_of_mass"], np.array([9.0, 8.0, 7.0]) - ) - np.testing.assert_allclose( - tensor_kwargs["positions"], np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) - ) - np.testing.assert_allclose(tensor_kwargs["masses"], np.array([13.0, 14.0])) - np.testing.assert_allclose(tensor_kwargs["dimensions"], box) - - get_principal.assert_called_once() - np.testing.assert_allclose(get_principal.call_args.args[0], np.eye(3)) - get_bonded.assert_called_once_with( - u=universe, - heavy_atom=heavy_atom, - topology=topology, - dimensions=box, - ) - - -def test_get_UA_axes_from_topology_single_heavy_uses_residue_principal_axes( - monkeypatch, -): - ax = AxesCalculator() - - heavy_atom = _FakeAtom(1, 12.0, [1.0, 0.0, 0.0]) - universe = _FakeUniverse( - {1: heavy_atom}, dimensions=[11.0, 12.0, 13.0, 90.0, 90.0, 90.0] - ) - residue_atoms = MagicMock() - residue_atoms.principal_axes.return_value = np.eye(3) * 5.0 - - topology = _ua_topology(heavy_atom_index=1, residue_heavy_indices=(1,)) - - make_whole = MagicMock() - get_bonded = MagicMock(return_value=(np.eye(3) * 6.0, np.array([6.0, 5.0, 4.0]))) - - monkeypatch.setattr("CodeEntropy.levels.axes.make_whole", make_whole) - monkeypatch.setattr(ax, "get_bonded_axes_from_topology", get_bonded) - - trans_axes, rot_axes, center, moi = ax.get_UA_axes_from_topology( - u=universe, - residue_atoms=residue_atoms, - topology=topology, - box=None, - ) - - make_whole.assert_called_once_with(residue_atoms) - residue_atoms.principal_axes.assert_called_once() - np.testing.assert_allclose(trans_axes, np.eye(3) * 5.0) - np.testing.assert_allclose(rot_axes, np.eye(3) * 6.0) - np.testing.assert_allclose(center, heavy_atom.position) - np.testing.assert_allclose(moi, np.array([6.0, 5.0, 4.0])) - - called_kwargs = get_bonded.call_args.kwargs - np.testing.assert_allclose( - called_kwargs["dimensions"], np.array([11.0, 12.0, 13.0]) - ) - - -def test_get_UA_axes_from_topology_raises_when_cached_bonded_axes_fail(monkeypatch): - ax = AxesCalculator() - - heavy_atom = _FakeAtom(1, 12.0, [1.0, 0.0, 0.0]) - universe = _FakeUniverse({1: heavy_atom}) - residue_atoms = MagicMock() - residue_atoms.principal_axes.return_value = np.eye(3) - topology = _ua_topology(heavy_atom_index=1, residue_heavy_indices=(1,)) - - monkeypatch.setattr("CodeEntropy.levels.axes.make_whole", lambda _ag: None) - monkeypatch.setattr( - ax, "get_bonded_axes_from_topology", lambda **kwargs: (None, None) - ) - - with pytest.raises(ValueError, match="cached UA bead"): - ax.get_UA_axes_from_topology( - u=universe, - residue_atoms=residue_atoms, - topology=topology, - box=None, - ) - - -def test_get_bonded_axes_from_topology_non_heavy_returns_none_none(): - ax = AxesCalculator() - light_atom = _FakeAtom(1, 1.0, [0.0, 0.0, 0.0]) - - custom_axes, moi = ax.get_bonded_axes_from_topology( - u=MagicMock(), - heavy_atom=light_atom, - topology=_ua_topology(heavy_atom_index=1), - dimensions=np.array([10.0, 10.0, 10.0]), - ) - - assert custom_axes is None - assert moi is None - - -def test_get_bonded_axes_from_topology_no_bonded_heavy_uses_vanilla_axes( - monkeypatch, -): - ax = AxesCalculator() - - heavy_atom = _FakeAtom(1, 12.0, [0.0, 0.0, 0.0]) - hydrogen = _FakeAtom(2, 1.0, [1.0, 0.0, 0.0]) - universe = _FakeUniverse({1: heavy_atom, 2: hydrogen}) - topology = _ua_topology( - heavy_atom_index=1, - ua_atom_indices=(1, 2), - ua_all_atom_indices=(1, 2), - bonded_heavy_indices=(), - bonded_light_indices=(2,), - ) - - get_vanilla = MagicMock(return_value=(np.eye(3) * 7.0, np.array([7.0, 8.0, 9.0]))) - get_custom_moi = MagicMock() - get_flipped = MagicMock(return_value=np.eye(3) * -7.0) - - monkeypatch.setattr(ax, "get_vanilla_axes", get_vanilla) - monkeypatch.setattr(ax, "get_custom_moment_of_inertia", get_custom_moi) - monkeypatch.setattr(ax, "get_flipped_axes", get_flipped) - - custom_axes, moi = ax.get_bonded_axes_from_topology( - u=universe, - heavy_atom=heavy_atom, - topology=topology, - dimensions=np.array([10.0, 10.0, 10.0]), - ) - - np.testing.assert_allclose(custom_axes, np.eye(3) * -7.0) - np.testing.assert_allclose(moi, np.array([7.0, 8.0, 9.0])) - get_vanilla.assert_called_once() - get_custom_moi.assert_not_called() - get_flipped.assert_called_once() - - -def test_get_bonded_axes_from_topology_one_heavy_no_light_uses_custom_axes( - monkeypatch, -): - ax = AxesCalculator() - - heavy_atom = _FakeAtom(1, 12.0, [0.0, 0.0, 0.0]) - bonded_heavy = _FakeAtom(3, 12.0, [1.0, 0.0, 0.0]) - universe = _FakeUniverse({1: heavy_atom, 3: bonded_heavy}) - topology = _ua_topology( - heavy_atom_index=1, - ua_atom_indices=(1,), - ua_all_atom_indices=(1, 3), - bonded_heavy_indices=(3,), - bonded_light_indices=(), - ) - - get_custom_axes = MagicMock(return_value=np.eye(3) * 2.0) - get_custom_moi = MagicMock(return_value=np.array([2.0, 3.0, 4.0])) - get_flipped = MagicMock(return_value=np.eye(3) * 3.0) - - monkeypatch.setattr(ax, "get_custom_axes", get_custom_axes) - monkeypatch.setattr(ax, "get_custom_moment_of_inertia", get_custom_moi) - monkeypatch.setattr(ax, "get_flipped_axes", get_flipped) - - custom_axes, moi = ax.get_bonded_axes_from_topology( - u=universe, - heavy_atom=heavy_atom, - topology=topology, - dimensions=np.array([10.0, 10.0, 10.0]), - ) - - np.testing.assert_allclose(custom_axes, np.eye(3) * 3.0) - np.testing.assert_allclose(moi, np.array([2.0, 3.0, 4.0])) - - kwargs = get_custom_axes.call_args.kwargs - np.testing.assert_allclose(kwargs["a"], heavy_atom.position) - np.testing.assert_allclose(kwargs["b_list"][0], bonded_heavy.position) - np.testing.assert_allclose(kwargs["c"], np.zeros(3)) - get_custom_moi.assert_called_once() - get_flipped.assert_called_once() - - -def test_get_bonded_axes_from_topology_one_heavy_with_light_uses_light_as_c( - monkeypatch, -): - ax = AxesCalculator() - - heavy_atom = _FakeAtom(1, 12.0, [0.0, 0.0, 0.0]) - bonded_heavy = _FakeAtom(3, 12.0, [1.0, 0.0, 0.0]) - bonded_light = _FakeAtom(2, 1.0, [0.0, 1.0, 0.0]) - universe = _FakeUniverse({1: heavy_atom, 2: bonded_light, 3: bonded_heavy}) - topology = _ua_topology( - heavy_atom_index=1, - ua_atom_indices=(1, 2), - ua_all_atom_indices=(1, 3, 2), - bonded_heavy_indices=(3,), - bonded_light_indices=(2,), - ) - - get_custom_axes = MagicMock(return_value=np.eye(3)) - monkeypatch.setattr(ax, "get_custom_axes", get_custom_axes) - monkeypatch.setattr( - ax, - "get_custom_moment_of_inertia", - lambda **kwargs: np.array([1.0, 2.0, 3.0]), - ) - monkeypatch.setattr(ax, "get_flipped_axes", lambda ua, axes, com, dims: axes) - - custom_axes, moi = ax.get_bonded_axes_from_topology( - u=universe, - heavy_atom=heavy_atom, - topology=topology, - dimensions=np.array([10.0, 10.0, 10.0]), - ) - - np.testing.assert_allclose(custom_axes, np.eye(3)) - np.testing.assert_allclose(moi, np.array([1.0, 2.0, 3.0])) - np.testing.assert_allclose( - get_custom_axes.call_args.kwargs["c"], bonded_light.position - ) - - -def test_get_bonded_axes_from_topology_two_heavy_uses_heavy_positions_as_b_list( - monkeypatch, -): - ax = AxesCalculator() - - heavy_atom = _FakeAtom(1, 12.0, [0.0, 0.0, 0.0]) - bonded_heavy_0 = _FakeAtom(3, 12.0, [1.0, 0.0, 0.0]) - bonded_heavy_1 = _FakeAtom(4, 12.0, [0.0, 1.0, 0.0]) - universe = _FakeUniverse( - { - 1: heavy_atom, - 3: bonded_heavy_0, - 4: bonded_heavy_1, - } - ) - topology = _ua_topology( - heavy_atom_index=1, - ua_atom_indices=(1,), - ua_all_atom_indices=(1, 3, 4), - bonded_heavy_indices=(3, 4), - bonded_light_indices=(), - ) - - get_custom_axes = MagicMock(return_value=np.eye(3) * 4.0) - monkeypatch.setattr(ax, "get_custom_axes", get_custom_axes) - monkeypatch.setattr( - ax, - "get_custom_moment_of_inertia", - lambda **kwargs: np.array([4.0, 5.0, 6.0]), - ) - monkeypatch.setattr(ax, "get_flipped_axes", lambda ua, axes, com, dims: axes) - - custom_axes, moi = ax.get_bonded_axes_from_topology( - u=universe, - heavy_atom=heavy_atom, - topology=topology, - dimensions=np.array([10.0, 10.0, 10.0]), - ) - - np.testing.assert_allclose(custom_axes, np.eye(3) * 4.0) - np.testing.assert_allclose(moi, np.array([4.0, 5.0, 6.0])) - np.testing.assert_allclose( - get_custom_axes.call_args.kwargs["b_list"], - np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]), - ) - np.testing.assert_allclose( - get_custom_axes.call_args.kwargs["c"], - bonded_heavy_1.position, - ) - - -def test_get_bonded_axes_from_topology_returns_none_when_custom_axes_none( - monkeypatch, -): - ax = AxesCalculator() - - heavy_atom = _FakeAtom(1, 12.0, [0.0, 0.0, 0.0]) - bonded_heavy = _FakeAtom(3, 12.0, [1.0, 0.0, 0.0]) - universe = _FakeUniverse({1: heavy_atom, 3: bonded_heavy}) - topology = _ua_topology( - heavy_atom_index=1, - ua_atom_indices=(1,), - ua_all_atom_indices=(1, 3), - bonded_heavy_indices=(3,), - bonded_light_indices=(), - ) - - get_custom_moi = MagicMock() - get_flipped = MagicMock() - - monkeypatch.setattr(ax, "get_custom_axes", lambda **kwargs: None) - monkeypatch.setattr(ax, "get_custom_moment_of_inertia", get_custom_moi) - monkeypatch.setattr(ax, "get_flipped_axes", get_flipped) - - custom_axes, moi = ax.get_bonded_axes_from_topology( - u=universe, - heavy_atom=heavy_atom, - topology=topology, - dimensions=np.array([10.0, 10.0, 10.0]), - ) - - assert custom_axes is None - assert moi is None - get_custom_moi.assert_not_called() - get_flipped.assert_not_called() diff --git a/tests/unit/CodeEntropy/levels/test_conformation_dag.py b/tests/unit/CodeEntropy/levels/test_conformation_dag.py index cb634fe8..92568ca9 100644 --- a/tests/unit/CodeEntropy/levels/test_conformation_dag.py +++ b/tests/unit/CodeEntropy/levels/test_conformation_dag.py @@ -3,7 +3,7 @@ from types import SimpleNamespace from unittest.mock import MagicMock -from CodeEntropy.levels.graph.conformation_dag import ConformationDAG +from CodeEntropy.levels.conformation_dag import ConformationDAG from CodeEntropy.trajectory.frames import FrameSelection diff --git a/tests/unit/CodeEntropy/levels/test_frame_graph.py b/tests/unit/CodeEntropy/levels/test_frame_graph.py index 59a9e757..e1134772 100644 --- a/tests/unit/CodeEntropy/levels/test_frame_graph.py +++ b/tests/unit/CodeEntropy/levels/test_frame_graph.py @@ -2,7 +2,7 @@ import pytest -from CodeEntropy.levels.graph.frame_dag import FrameGraph +from CodeEntropy.levels.frame_dag import FrameGraph def test_make_frame_ctx_has_required_keys(): diff --git a/tests/unit/CodeEntropy/levels/test_level_dag.py b/tests/unit/CodeEntropy/levels/test_level_dag.py index 51674c73..1b9f4d4f 100644 --- a/tests/unit/CodeEntropy/levels/test_level_dag.py +++ b/tests/unit/CodeEntropy/levels/test_level_dag.py @@ -4,17 +4,16 @@ from unittest.mock import MagicMock, call, patch -from CodeEntropy.levels.graph.level_dag import LevelDAG +from CodeEntropy.levels.level_dag import LevelDAG def test_build_registers_static_nodes_and_builds_stage_dags(): with ( - patch("CodeEntropy.levels.graph.level_dag.DetectMoleculesNode"), - patch("CodeEntropy.levels.graph.level_dag.DetectLevelsNode"), - patch("CodeEntropy.levels.graph.level_dag.BuildBeadsNode"), - patch("CodeEntropy.levels.graph.level_dag.BuildAxesTopologyNode"), - patch("CodeEntropy.levels.graph.level_dag.InitCovarianceAccumulatorsNode"), - patch("CodeEntropy.levels.graph.level_dag.ConformationDAG"), + patch("CodeEntropy.levels.level_dag.DetectMoleculesNode"), + patch("CodeEntropy.levels.level_dag.DetectLevelsNode"), + patch("CodeEntropy.levels.level_dag.BuildBeadsNode"), + patch("CodeEntropy.levels.level_dag.InitCovarianceAccumulatorsNode"), + patch("CodeEntropy.levels.level_dag.ConformationDAG"), ): universe_operations = MagicMock() dag = LevelDAG(universe_operations=universe_operations) @@ -28,7 +27,6 @@ def test_build_registers_static_nodes_and_builds_stage_dags(): "detect_molecules", "detect_levels", "build_beads", - "build_axes_topology", "init_covariance_accumulators", } assert "find_neighbors" not in dag._static_nodes @@ -36,7 +34,6 @@ def test_build_registers_static_nodes_and_builds_stage_dags(): assert ("detect_molecules", "detect_levels") in dag._static_graph.edges assert ("detect_levels", "build_beads") in dag._static_graph.edges - assert ("build_beads", "build_axes_topology") in dag._static_graph.edges assert ("detect_levels", "init_covariance_accumulators") in dag._static_graph.edges assert ( "detect_levels", @@ -59,12 +56,8 @@ def test_execute_sets_default_axes_manager_and_runs_workflow_stages(): dag._run_frame_stage = MagicMock() with ( - patch( - "CodeEntropy.levels.graph.level_dag.NeighborReducer.initialise" - ) as initialise, - patch( - "CodeEntropy.levels.graph.level_dag.NeighborReducer.finalise" - ) as finalise, + patch("CodeEntropy.levels.level_dag.NeighborReducer.initialise") as initialise, + patch("CodeEntropy.levels.level_dag.NeighborReducer.finalise") as finalise, ): out = dag.execute(shared_data, progress=progress) @@ -168,7 +161,7 @@ def test_run_frame_stage_collects_frame_indices_and_delegates_to_scheduler(): shared_data = {"frame_source": frame_source} - with patch("CodeEntropy.levels.graph.level_dag.FrameScheduler") as Scheduler: + with patch("CodeEntropy.levels.level_dag.FrameScheduler") as Scheduler: scheduler = Scheduler.return_value dag._run_frame_stage(shared_data, progress=progress) @@ -193,7 +186,7 @@ def test_initialise_neighbor_metadata_writes_symmetry_and_linearity(): groups = {0: [0], 1: [1]} shared_data = {"reduced_universe": universe, "groups": groups} - with patch("CodeEntropy.levels.graph.level_dag.Neighbors") as Neighbors: + with patch("CodeEntropy.levels.level_dag.Neighbors") as Neighbors: helper = Neighbors.return_value helper.get_symmetry.return_value = ({0: 12, 1: 2}, {0: False, 1: True}) @@ -208,7 +201,7 @@ def test_initialise_neighbor_metadata_falls_back_to_universe_key(): universe = object() shared_data = {"universe": universe, "groups": {0: [0]}} - with patch("CodeEntropy.levels.graph.level_dag.Neighbors") as Neighbors: + with patch("CodeEntropy.levels.level_dag.Neighbors") as Neighbors: helper = Neighbors.return_value helper.get_symmetry.return_value = ({0: 1}, {0: False}) @@ -243,11 +236,11 @@ def test_level_dag_runs_static_conformation_then_frame(monkeypatch): ) monkeypatch.setattr( - "CodeEntropy.levels.graph.level_dag.NeighborReducer.initialise", + "CodeEntropy.levels.level_dag.NeighborReducer.initialise", lambda shared_data: calls.append("neighbor_initialise"), ) monkeypatch.setattr( - "CodeEntropy.levels.graph.level_dag.NeighborReducer.finalise", + "CodeEntropy.levels.level_dag.NeighborReducer.finalise", lambda shared_data: calls.append("neighbor_finalise"), ) diff --git a/tests/unit/CodeEntropy/levels/test_mda_universe_operations.py b/tests/unit/CodeEntropy/levels/test_mda_universe_operations.py index d298c486..e8d92329 100644 --- a/tests/unit/CodeEntropy/levels/test_mda_universe_operations.py +++ b/tests/unit/CodeEntropy/levels/test_mda_universe_operations.py @@ -1,5 +1,4 @@ -from types import SimpleNamespace -from unittest.mock import MagicMock, Mock +from unittest.mock import MagicMock import numpy as np import pytest @@ -408,19 +407,3 @@ def test_select_frame_indices_raises_when_frame_indices_empty(): ops.select_frame_indices(u, frame_indices=[]) u.select_atoms.assert_not_called() - - -def test_extract_fragment_atomgroup_returns_lightweight_range_selection() -> None: - universe = Mock() - universe.atoms.fragments = [ - SimpleNamespace(indices=[10, 11, 12, 13]), - ] - universe.select_atoms.return_value = "fragment_atomgroup" - - result = UniverseOperations().extract_fragment_atomgroup(universe, 0) - - assert result == "fragment_atomgroup" - universe.select_atoms.assert_called_once_with( - "index 10:13", - updating=False, - ) From 7f85a4f5dc5b33a792ffd3ee58a51dbda83aa45d Mon Sep 17 00:00:00 2001 From: ioanaapapa Date: Thu, 2 Jul 2026 11:07:43 +0100 Subject: [PATCH 47/47] added new unit tests --- tests/unit/CodeEntropy/levels/test_axes.py | 36 ++++++++++++++++++++++ 1 file changed, 36 insertions(+) diff --git a/tests/unit/CodeEntropy/levels/test_axes.py b/tests/unit/CodeEntropy/levels/test_axes.py index 4c4a894a..93740467 100644 --- a/tests/unit/CodeEntropy/levels/test_axes.py +++ b/tests/unit/CodeEntropy/levels/test_axes.py @@ -716,3 +716,39 @@ def test_get_bonded_axes_returns_none_none_if_custom_axes_none(monkeypatch): assert custom_axes is None assert moi is None + + +def test_get_residue_axes_custom_path(monkeypatch): + ax = AxesCalculator() + + edge_atoms = _FakeAtomGroup( + [_FakeAtom(8, 12.0, [1, 0, 0]), _FakeAtom(10, 14.0, [0, 0, 0])], + positions=np.array([[1.0, 0.0, 0.0], [0.0, 0.0, 0.0]], dtype=float), + ) + + backbone_center = np.array([0.0, 1.0, 0.0]) + rot_center, rot_axes = ax.get_residue_custom_axes( + [edge_atoms[0], edge_atoms[1]], backbone_center + ) + + assert rot_center.shape == (3,) + assert rot_axes.shape == (3, 3) + + +def test_get_custom_residue_moment_of_inertia(monkeypatch): + ax = AxesCalculator() + heavy_atoms = _FakeAtomGroup( + [_FakeAtom(0, 12.0, [1, 2, 1]), _FakeAtom(1, 12.0, [2, 1, 1])], + positions=np.array([[1, 2, 1], [2, 1, 1]], dtype=float), + ) + dimensions = np.array([10.0, 10.0, 10.0], dtype=float) + + moi = ax.get_custom_residue_moment_of_inertia( + center_of_mass=np.array([1, 1, 1]), + positions=heavy_atoms.positions, + masses=heavy_atoms.masses, + custom_rot_axes=np.eye(3), + dimensions=dimensions, + ) + + assert moi.shape == (3,)