epsi
/
msspec_python3
3
0
Fork 0
Code Issues 6 Pull Requests Releases Wiki Activity

Modified cut_cylinder and cut_plane functions.
epsi-builds/msspec_python3/pipeline/head There was a failure building this commit Details 

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.
This commit is contained in:
Sylvain Tricot 2022-02-08 15:20:32 +01:00
parent a657b1874e
commit ca1fd04163
1 changed files with 15 additions and 32 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

47
src/msspec/utils.py
View File

@ -311,34 +311,13 @@ def cut_cylinder(atoms, axis="z", radius=None):
:return: The modified atom cluster :return: The modified atom cluster
:rtype: ase.Atoms :rtype: ase.Atoms
""" """
if radius is None: if axis not in ('z',):
raise ValueError("radius not set") raise ValueError("axis value != 'z' is not supported yet.")
X, Y, Z = atoms.positions.T
new_atoms = atoms.copy() R = np.sqrt(X**2 + Y **2)
T = np.arctan2(Y, X)
dims = {"x": 0, "y": 1, "z": 2} i = np.where(R <= radius)[0]
if axis in dims: return atoms[i]
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
def cut_cone(atoms, radius, z=0): 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) dim_values = np.array(dim_values)
def constraint(coordinates): X, Y, Z = atoms.positions.T
return np.all(np.logical_and(coordinates >= dim_values[:, 0], i0 = np.where(X >= dim_values[0, 0])[0]
coordinates <= dim_values[:, 1])) 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] return atoms[indices]