pytme 0.1.9__cp311-cp311-macosx_14_0_arm64.whl → 0.2.0__cp311-cp311-macosx_14_0_arm64.whl

tme/preprocessor.py

@@ -47,7 +47,7 @@ class Preprocessor:
 
     def apply_method(self, method: str, parameters: Dict):
         """
-        Apply a method on the atomic structure.
+        Invoke ``Preprocessor.method`` using ``parameters``.
 
         Parameters
         ----------
@@ -654,12 +654,9 @@ class Preprocessor:
         array = template.copy()
         interpolation_box = array.shape
 
-        print(array.shape)
-
         for k in range(template.ndim):
             array = decimate(array, q=level, axis=k)
 
-        print(array.shape)
         template = zoom(array, np.divide(template.shape, array.shape))
         template = self.interpolate_box(box=interpolation_box, arr=template)
 
@@ -768,21 +765,24 @@ class Preprocessor:
         sigma = sigma_factor * resolution
         sigma_grid = sigma / sampling_rate
         sigma_grid2 = sigma_grid * sigma_grid
-        for index, point in enumerate(np.rollaxis(positions, 0)):
-            starts = np.maximum(np.ceil(point - cutoff_value * sigma_grid), 0).astype(
-                int
-            )
-            stops = np.minimum(
-                np.floor(point + cutoff_value * sigma_grid), shape
-            ).astype(int)
 
-            grid_index = np.meshgrid(
-                *[range(start, stop) for start, stop in zip(starts, stops)]
-            )
-            distances = np.einsum(
-                "aijk->ijk",
-                np.array([(grid_index[i] - point[i]) ** 2 for i in range(len(point))]),
-                dtype=np.float64,
+        starts = np.maximum(np.ceil(positions - cutoff_value * sigma_grid), 0).astype(
+            int
+        )
+        stops = np.minimum(
+            np.floor(positions + cutoff_value * sigma_grid), shape
+        ).astype(int)
+        ranges = tuple(tuple(zip(start, stop)) for start, stop in zip(starts, stops))
+
+        positions = positions.reshape(
+            *positions.shape, *tuple(1 for _ in range(positions.shape[1]))
+        )
+        for index in range(positions.shape[0]):
+            grid_index = np.meshgrid(*[range(*coord) for coord in ranges[index]])
+            distances = np.sum(
+                np.square(np.subtract(grid_index, positions[index])),
+                dtype=np.float32,
+                axis=0,
             )
             np.add.at(
                 out,
@@ -1131,9 +1131,11 @@ class Preprocessor:
         stop_tilt: float,
         tilt_step: float,
         shape: Tuple[int],
+        tilt_angles: Tuple[float] = None,
         opening_axis: int = 0,
         tilt_axis: int = 2,
         sigma: float = 0,
+        weights: float = 1,
         omit_negative_frequencies: bool = True,
     ) -> NDArray:
         """
@@ -1164,6 +1166,8 @@ class Preprocessor:
             - 2 for X-axis
         sigma : float, optional
             Standard deviation for Gaussian kernel used for smoothing the wedge.
+        weights : float, tuple of float
+            Weight of each element in the wedge. Defaults to one.
         omit_negative_frequencies : bool, optional
             Whether the wedge mask should omit negative frequencies, i.e. be
             applicable to symmetric Fourier transforms (see :obj:`numpy.fft.fftn`)
@@ -1184,17 +1188,20 @@ class Preprocessor:
         :py:meth:`Preprocessor.wedge_mask`
         :py:meth:`Preprocessor.continuous_wedge_mask`
         """
-        tilt_angles = np.arange(-start_tilt, stop_tilt + tilt_step, tilt_step)
+        if tilt_angles is None:
+            tilt_angles = np.arange(-start_tilt, stop_tilt + tilt_step, tilt_step)
+
+        weights = np.asarray(weights)
+        weights = np.repeat(weights, tilt_angles.size // weights.size)
         plane = np.zeros((shape[opening_axis], shape[tilt_axis]), dtype=np.float32)
         subset = tuple(
             slice(None) if i != 0 else slice(x // 2, x // 2 + 1)
             for i, x in enumerate(plane.shape)
         )
-        plane[subset] = 1
         plane_rotated, wedge_volume = np.zeros_like(plane), np.zeros_like(plane)
         for index in range(tilt_angles.shape[0]):
             plane_rotated.fill(0)
-
+            plane[subset] = weights[index]
             rotation_matrix = euler_to_rotationmatrix((tilt_angles[index], 0))
             rotation_matrix = rotation_matrix[np.ix_((0, 1), (0, 1))]
 
@@ -1207,10 +1214,13 @@ class Preprocessor:
             )
             wedge_volume += plane_rotated
 
-        wedge_volume = self.gaussian_filter(
-            template=wedge_volume, sigma=sigma, fourier=False
-        )
-        wedge_volume = np.where(wedge_volume > np.exp(-2), 1, 0)
+        # Ramp filtering would be more accurate
+        np.fmin(wedge_volume, np.max(weights), wedge_volume)
+
+        if sigma > 0:
+            wedge_volume = self.gaussian_filter(
+                template=wedge_volume, sigma=sigma, fourier=False
+            )
 
         if opening_axis > tilt_axis:
             wedge_volume = np.moveaxis(wedge_volume, 1, 0)

tme/structure.py

@@ -26,45 +26,10 @@ from .types import NDArray
 
 @dataclass(repr=False)
 class Structure:
-    """Represents atomic structures in accordance with the Protein Data Bank (PDB)
+    """
+    Represents atomic structures in accordance with the Protein Data Bank (PDB)
     format specification.
 
-    Attributes
-    ----------
-    record_type : NDArray
-        Type of the record, e.g., ATOM, HETATM. Array shape = (n,)
-    atom_serial_number : NDArray
-        Serial number assigned to each atom.  Array shape = (n,)
-    atom_name : NDArray
-        Standardized names for each atom.  Array shape = (n,)
-    atom_coordinate : NDArray
-        The 3D Cartesian coordinates of each atom in x, y, z.  Array shape = (n,3 )
-    alternate_location_indicator : NDArray
-        Indicator for alternate locations of an atom if it exists in multiple places.
-        Array shape = (n,)
-    residue_name : NDArray
-        Standard residue names where each atom belongs. Array shape = (n,)
-    chain_identifier : NDArray
-        Identifier for the chain where each atom is located. Array shape = (n,)
-    residue_sequence_number : NDArray
-        Sequence number of the residue in the protein chain for each atom.
-        Array shape = (n,)
-    code_for_residue_insertion : NDArray
-        Code to denote any residue insertion. Array shape = (n,)
-    occupancy : NDArray
-        Occupancy factor of each atom, indicating the fraction of time the atom
-        is located at its position. Array shape = (n,)
-    temperature_factor : NDArray
-        Measure of the atomic displacement or B-factor for each atom. Array shape = (n,)
-    segment_identifier : NDArray
-        Identifier for the segment where each atom belongs. Array shape = (n,)
-    element_symbol : NDArray
-        Atomic element symbol for each atom. Array shape = (n,)
-    charge : NDArray
-        Charge on the atom. Array shape = (n,)
-    details : dict
-        Any additional or auxiliary details. Array shape = (n,)
-
     References
     ----------
     .. [1]  https://www.cgl.ucsf.edu/chimera/docs/UsersGuide/tutorials/pdbintro.html