19
19
if TYPE_CHECKING :
20
20
from typing import Any
21
21
22
- from numpy .typing import ArrayLike
22
+ from numpy .typing import ArrayLike , NDArray
23
23
from typing_extensions import Self
24
24
25
25
__author__ = "Shyue Ping Ong, Shyam Dwaraknath, Matthew Horton"
@@ -31,7 +31,7 @@ class SymmOp(MSONable):
31
31
for efficiency. Read: https://wikipedia.org/wiki/Affine_transformation.
32
32
33
33
Attributes:
34
- affine_matrix (np.ndarray ): A 4x4 array representing the symmetry operation.
34
+ affine_matrix (NDArray ): A 4x4 array representing the symmetry operation.
35
35
"""
36
36
37
37
def __init__ (
@@ -116,7 +116,7 @@ def from_rotation_and_translation(
116
116
affine_matrix [:3 ][:, 3 ] = translation_vec
117
117
return cls (affine_matrix , tol )
118
118
119
- def operate (self , point : ArrayLike ) -> np . ndarray :
119
+ def operate (self , point : ArrayLike ) -> NDArray :
120
120
"""Apply the operation on a point.
121
121
122
122
Args:
@@ -128,7 +128,7 @@ def operate(self, point: ArrayLike) -> np.ndarray:
128
128
affine_point = np .asarray ([* point , 1 ])
129
129
return np .dot (self .affine_matrix , affine_point )[:3 ]
130
130
131
- def operate_multi (self , points : ArrayLike ) -> np . ndarray :
131
+ def operate_multi (self , points : ArrayLike ) -> NDArray :
132
132
"""Apply the operation on a list of points.
133
133
134
134
Args:
@@ -141,7 +141,7 @@ def operate_multi(self, points: ArrayLike) -> np.ndarray:
141
141
affine_points = np .concatenate ([points , np .ones (points .shape [:- 1 ] + (1 ,))], axis = - 1 )
142
142
return np .inner (affine_points , self .affine_matrix )[..., :- 1 ]
143
143
144
- def apply_rotation_only (self , vector : ArrayLike ) -> np . ndarray :
144
+ def apply_rotation_only (self , vector : ArrayLike ) -> NDArray :
145
145
"""Vectors should only be operated by the rotation matrix and not the
146
146
translation vector.
147
147
@@ -150,7 +150,7 @@ def apply_rotation_only(self, vector: ArrayLike) -> np.ndarray:
150
150
"""
151
151
return np .dot (self .rotation_matrix , vector )
152
152
153
- def transform_tensor (self , tensor : np . ndarray ) -> np . ndarray :
153
+ def transform_tensor (self , tensor : NDArray ) -> NDArray :
154
154
"""Apply rotation portion to a tensor. Note that tensor has to be in
155
155
full form, not the Voigt form.
156
156
@@ -239,12 +239,12 @@ def are_symmetrically_related_vectors(
239
239
return False , False
240
240
241
241
@property
242
- def rotation_matrix (self ) -> np . ndarray :
242
+ def rotation_matrix (self ) -> NDArray :
243
243
"""A 3x3 numpy.array representing the rotation matrix."""
244
244
return self .affine_matrix [:3 ][:, :3 ]
245
245
246
246
@property
247
- def translation_vector (self ) -> np . ndarray :
247
+ def translation_vector (self ) -> NDArray :
248
248
"""A rank 1 numpy.array of dim 3 representing the translation vector."""
249
249
return self .affine_matrix [:3 ][:, 3 ]
250
250
@@ -462,16 +462,17 @@ def as_xyz_str(self) -> str:
462
462
def from_xyz_str (cls , xyz_str : str ) -> Self :
463
463
"""
464
464
Args:
465
- xyz_str: string of the form ' x, y, z', ' -x, -y, z', ' -2y+1/2, 3x+1/2, z-y+1/2' , etc.
465
+ xyz_str (str): " x, y, z", " -x, -y, z", " -2y+1/2, 3x+1/2, z-y+1/2" , etc.
466
466
467
467
Returns:
468
468
SymmOp
469
469
"""
470
- rot_matrix = np .zeros ((3 , 3 ))
471
- trans = np .zeros (3 )
472
- tokens = xyz_str .strip ().replace (" " , "" ).lower ().split ("," )
470
+ rot_matrix : NDArray = np .zeros ((3 , 3 ))
471
+ trans : NDArray = np .zeros (3 )
472
+ tokens : list [ str ] = xyz_str .strip ().replace (" " , "" ).lower ().split ("," )
473
473
re_rot = re .compile (r"([+-]?)([\d\.]*)/?([\d\.]*)([x-z])" )
474
474
re_trans = re .compile (r"([+-]?)([\d\.]+)/?([\d\.]*)(?![x-z])" )
475
+
475
476
for idx , tok in enumerate (tokens ):
476
477
# Build the rotation matrix
477
478
for match in re_rot .finditer (tok ):
@@ -480,11 +481,13 @@ def from_xyz_str(cls, xyz_str: str) -> Self:
480
481
factor *= float (match [2 ]) / float (match [3 ]) if match [3 ] != "" else float (match [2 ])
481
482
j = ord (match [4 ]) - 120
482
483
rot_matrix [idx , j ] = factor
484
+
483
485
# Build the translation vector
484
486
for match in re_trans .finditer (tok ):
485
487
factor = - 1 if match [1 ] == "-" else 1
486
488
num = float (match [2 ]) / float (match [3 ]) if match [3 ] != "" else float (match [2 ])
487
489
trans [idx ] = num * factor
490
+
488
491
return cls .from_rotation_and_translation (rot_matrix , trans )
489
492
490
493
@classmethod
0 commit comments