Modified cut_cylinder and cut_plane functions.

cut_cylinder and cut_plane functions were too slow due to
lists modifications in for loops. Implementation was modified
a bit while keeping the same API.
For now, cutting a cylinder for an axis 'x' or 'y' is not supported
anymore.

src/msspec/utils.py

@@ -311,34 +311,13 @@ def cut_cylinder(atoms, axis="z", radius=None):
     :return: The modified atom cluster
     :rtype: ase.Atoms
     """
-    if radius is None:
-        raise ValueError("radius not set")
-
-    new_atoms = atoms.copy()
-
-    dims = {"x": 0, "y": 1, "z": 2}
-    if axis in dims:
-        axis = dims[axis]
-    else:
-        raise ValueError("axis not valid, must be 'x','y', or 'z'")
-
-    del_list = []
-    for index, position in enumerate(new_atoms.positions):
-        # calculating the distance of the atom to the given axis
-        r = 0
-        for dim in range(3):
-            if dim != axis:
-                r = r + position[dim]**2
-        r = np.sqrt(r)
-
-        if r > radius:
-            del_list.append(index)
-
-    del_list.reverse()
-    for index in del_list:
-        del new_atoms[index]
-
-    return new_atoms
+    if axis not in ('z',):
+        raise ValueError("axis value != 'z' is not supported yet.")
+    X, Y, Z = atoms.positions.T
+    R = np.sqrt(X**2 + Y **2)
+    T = np.arctan2(Y, X)
+    i = np.where(R <= radius)[0]
+    return atoms[i]
 
 
 def cut_cone(atoms, radius, z=0):
@@ -426,11 +405,15 @@ def cut_plane(atoms, x=None, y=None, z=None):
 
     dim_values = np.array(dim_values)
 
-    def constraint(coordinates):
-        return np.all(np.logical_and(coordinates >= dim_values[:, 0],
-                                     coordinates <= dim_values[:, 1]))
+    X, Y, Z = atoms.positions.T
+    i0 = np.where(X                     >= dim_values[0, 0])[0]
+    i1 = np.where(X[i0]                 <= dim_values[0, 1])[0]
+    i2 = np.where(Y[i0][i1]             >= dim_values[1, 0])[0]
+    i3 = np.where(Y[i0][i1][i2]         <= dim_values[1, 1])[0]
+    i4 = np.where(Z[i0][i1][i2][i3]     >= dim_values[2, 0])[0]
+    i5 = np.where(Z[i0][i1][i2][i3][i4] <= dim_values[2, 1])[0]
+    indices = np.arange(len(atoms))[i0][i1][i2][i3][i4][i5]
 
-    indices = np.where(list(map(constraint, atoms.positions)))[0]
     return atoms[indices]