From 9920253652d1f21de0199af2dfd6ead5da7c2d55 Mon Sep 17 00:00:00 2001 From: ioana Date: Thu, 12 Mar 2026 13:41:12 +0000 Subject: [PATCH 01/35] 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 2d05fcb..5a5d773 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 dd56d18..48f62da 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/35] 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 1425078..7db30c7 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/35] 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 7db30c7..59bfdf5 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/35] 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 59bfdf5..e171be4 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/35] 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 e171be4..ccd4894 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/35] 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 ccd4894..e50aa0e 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/35] 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 e50aa0e..b52c014 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/35] 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 b52c014..b0d0aba 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/35] 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 b0d0aba..bb422c9 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/35] 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 bb422c9..e6bc77e 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/35] 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 e6bc77e..8c48b40 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/35] 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 8c48b40..4846c39 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/35] 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 4846c39..5383e3c 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/35] 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 fbe9a82..16e16ff 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/35] 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 5383e3c..8e4956a 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/35] 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 d6d0093..0eba6e0 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/35] 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 8e4956a..a597a4d 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 16e16ff..cbf2c3b 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/35] 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 a597a4d..2c99815 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/35] 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 0eba6e0..4fc0c56 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/35] 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 a597a4d..930da81 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/35] 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 f4fbf25..5077714 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 1b51006..1fdafd9 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/35] 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 1fdafd9..00233f1 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/35] 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 5077714..20ab150 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/35] 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 4fc0c56..c0c5423 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/35] 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 20ab150..6a5e66f 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/35] 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 c0c5423..2f369e9 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/35] 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 6a5e66f..7eff00f 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 2f369e9..40252f4 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/35] 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 7eff00f..a06c534 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 40252f4..4c4a894 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/35] 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 e867cee..5671e03 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 2cbef74..760c664 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 c35d221..a4dc5e6 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 2cbef74..760c664 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/35] 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 2cbef74..760c664 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/35] 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 a06c534..892ec4c 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/35] 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 892ec4c..3d82f8f 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/35] 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 3d82f8f..2138922 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/35] 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 2138922..24ad97a 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/35] 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 24ad97a..c3d1b4b 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 | |