pytme 0.1.1__cp311-cp311-macosx_13_0_arm64.whl → 0.1.3__cp311-cp311-macosx_13_0_arm64.whl

tme/structure.py

@@ -189,11 +189,11 @@ class Structure:
 
         min_atom = np.min(self.atom_serial_number)
         max_atom = np.max(self.atom_serial_number)
-        n_atom = max_atom - min_atom + 1
+        n_atom = self.atom_serial_number.size
 
         min_residue = np.min(self.residue_sequence_number)
         max_residue = np.max(self.residue_sequence_number)
-        n_residue = max_residue - min_residue + 1
+        n_residue = self.residue_sequence_number.size
 
         repr_str = (
             f"Structure object at {id(self)}\n"
@@ -669,14 +669,16 @@ class Structure:
         Parameters
         ----------
         chain : str, optional
-            The chain identifier. If None, all chains are returned. Default is None.
+            The chain identifier. If multiple chains should be selected they need
+            to be a comma separated string, e.g. 'A,B,CE'. If chain None,
+            all chains are returned. Default is None.
 
         Returns
         -------
         Structure
             A subset of the original structure containing only the specified chain.
         """
-        chain = np.unique(self.chain_identifier) if chain is None else chain
+        chain = np.unique(self.chain_identifier) if chain is None else chain.split(",")
         keep = np.in1d(self.chain_identifier, chain)
         return self[keep]
 
@@ -698,7 +700,9 @@ class Structure:
             The ending residue sequence number.
 
         chain : str, optional
-            The chain identifier. If None, all chains are considered. Default is None.
+            The chain identifier. If multiple chains should be selected they need
+            to be a comma separated string, e.g. 'A,B,CE'. If chain None,
+            all chains are returned. Default is None.
 
         Returns
         -------
@@ -789,7 +793,6 @@ class Structure:
         shape: Tuple[int],
         sampling_rate: Tuple[float],
         origin: Tuple[float],
-        chain: str,
     ) -> (NDArray, Tuple[str], Tuple[int], float, Tuple[float]):
         """
         Converts coordinates to positions.
@@ -804,10 +807,6 @@ class Structure:
 
         origin : Tuple[float,]
             The origin of the coordinate system.
-
-        chain : str
-            The protein chain that should be considered.
-
         Returns
         -------
         Tuple[NDArray, List[str], Tuple[int, ], float, Tuple[float,]]
@@ -819,11 +818,6 @@ class Structure:
         # positions are in x, y, z map is z, y, x
         coordinates = coordinates[:, ::-1]
 
-        if chain is not None:
-            idx = np.in1d(self.chain_identifier, chain)
-            coordinates = coordinates[idx]
-            atom_types = atom_types[idx]
-
         sampling_rate = 1 if sampling_rate is None else sampling_rate
         adjust_origin = origin is not None and shape is None
         origin = coordinates.min(axis=0) if origin is None else origin
@@ -888,7 +882,6 @@ class Structure:
                 ).astype(int)
 
                 vdw_rad[atom_type] = atom_vdwr
-                print(atom_vdwr)
                 atom_slice = tuple(slice(-k, k + 1) for k in atom_vdwr)
                 distances = np.linalg.norm(
                     np.divide(
@@ -1039,7 +1032,9 @@ class Structure:
             Origin of the coordinate system. If origin is given its expected to be
             in z, y, x form.
         chain : str, optional
-            Protein chain that should be considered. Default None uses all.
+            The chain identifier. If multiple chains should be selected they need
+            to be a comma separated string, e.g. 'A,B,CE'. If chain None,
+            all chains are returned. Default is None.
         weight_type : str, optional
             Which weight should be given to individual atoms.
         scattering_args : dict, optional
@@ -1074,8 +1069,11 @@ class Structure:
         if "source" not in scattering_args:
             scattering_args["source"] = "peng1995"
 
+        temp = self.copy()
+        self = self.subset_by_chain(chain = chain)
+
         positions, atoms, shape, sampling_rate, origin = self._coordinate_to_position(
-            shape=shape, sampling_rate=sampling_rate, origin=origin, chain=chain
+            shape=shape, sampling_rate=sampling_rate, origin=origin
         )
         volume = np.zeros(shape, dtype=np.float32)
         if weight_type in ("atomic_weight", "atomic_number"):
@@ -1102,6 +1100,7 @@ class Structure:
                 **scattering_args,
             )
 
+        self = temp
         return volume, origin, sampling_rate
 
     @classmethod
@@ -1229,10 +1228,10 @@ class Structure:
         query = structure2.atom_coordinate.copy()
         if sampling_rate is not None:
             reference, atoms1, shape, _, _ = structure1._coordinate_to_position(
-                shape=None, sampling_rate=sampling_rate, origin=origin, chain=None
+                shape=None, sampling_rate=sampling_rate, origin=origin
             )
             query, atoms2, shape, _, _ = structure2._coordinate_to_position(
-                shape=None, sampling_rate=sampling_rate, origin=origin, chain=None
+                shape=None, sampling_rate=sampling_rate, origin=origin
             )
 
         reference_mean = reference.mean(axis=0)