diff --git a/src/msspec/utils.py b/src/msspec/utils.py index 1355e11..b9eedd3 100644 --- a/src/msspec/utils.py +++ b/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]