Coverage for strongcoca/polarizable_unit.py: 99%
150 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-10-26 18:44 +0000
1from typing import Optional
3import numpy as np
4from scipy.spatial.transform import Rotation
5from ase import Atom, Atoms
7from .response.base import BaseResponse
8from .response.casida import CasidaResponse
9from .response.excitations import Excitations
10from .response.utilities import Broadening
11from .types import Array, Matrix, Vector
12from .units import au_to_A, A_to_au
13from .utilities import ClassFormatter
16class PolarizableUnit(BaseResponse):
17 """Class for managing polarizable units in CoupledSystem objects along
18 with their position and orientation.
20 Parameters
21 ----------
22 response
23 Object that handles response function representations.
24 position
25 Position of polarizable unit in Å.
26 orientation
27 Orientation of polarizable unit.
28 name
29 Name of polarizable unit.
30 **orientation_kwargs
31 Alternative keyword arguments for setting orientation;
32 these arguments are passed to :meth:`set_orientation`.
34 Examples
35 --------
37 The following snippet illustrates how polarizable unit objects can be
38 initialized. Here, an artificial response function is used. In
39 practice the response function would usually be imported from, e.g., TDDFT
40 calculations:
42 >>> from math import pi
43 >>> from strongcoca import PolarizableUnit
44 >>> from strongcoca.response import build_random_casida
45 >>>
46 >>> response = build_random_casida(n_states=3, name='Merkat', random_seed=42)
47 >>> pu = PolarizableUnit(response, [0, 2, -1], name='Apple',
48 ... rotation_vector=[0, 0, pi])
49 >>>
50 >>> print(pu)
51 --------------------- PolarizableUnit ----------------------
52 name : Apple
53 response : Merkat
54 position (Å) : [ 0. 2. -1.]
55 quaternion : [0. 0. 1. 0.]
56 rotation_vector : [0. ... 0. ... 3.14...]
57 rotation_angle : 3.14...
58 """
60 def __init__(self,
61 response: BaseResponse,
62 position: Vector,
63 orientation: Optional[Rotation] = None,
64 name: str = 'PolarizableUnit',
65 **orientation_kwargs):
66 super().__init__(broadening=Broadening(), pbc=False, name=name)
67 self._response = response
68 self.position = np.array(position)
69 if len(orientation_kwargs) > 0:
70 if orientation is not None:
71 raise ValueError('Both orientation and orientation keywords specified.')
72 self.set_orientation(**orientation_kwargs)
73 else:
74 self.orientation = orientation
76 def __str__(self) -> str:
77 """String representation."""
78 fmt = ClassFormatter(self)
79 fmt.append_class_name()
81 fmt.append_attr('name')
82 fmt.append_attr('response', self._response.name)
84 fmt.append_attr('position', unit='Å')
85 fmt.append_attr('quaternion', np.around(self._orientation.as_quat(), 2))
86 rot_v = self._orientation.as_rotvec()
87 fmt.append_attr('rotation_vector', rot_v)
88 fmt.append_attr('rotation_angle', np.linalg.norm(rot_v))
90 return fmt.to_string()
92 @property
93 def name(self) -> str:
94 """Name of polarizable unit."""
95 return self._name
97 @name.setter
98 def name(self, name: str) -> None:
99 """Set name of polarizable unit."""
100 self._name = name
102 @property
103 def position(self) -> np.ndarray:
104 """Position of center of mass of polarizable unit."""
105 return self._position * au_to_A # type: ignore
107 @position.setter
108 def position(self, new_position: Vector) -> None:
109 """Set position of center of mass of polarizable unit."""
110 _position = np.array(new_position) * A_to_au
111 if _position.shape != (3, ):
112 raise TypeError('Invalid position; not a three-dimensional vector of numbers:'
113 f' {new_position}, {type(new_position)}')
114 self._position = _position
116 @property
117 def orientation(self) -> Rotation:
118 """Orientation of the polarizable unit."""
119 return self._orientation
121 @orientation.setter
122 def orientation(self, orientation: Optional[Rotation]) -> None:
123 """Set the orientation of the polarizable unit."""
124 if orientation is None:
125 orientation = Rotation.identity()
126 if not isinstance(orientation, Rotation):
127 raise TypeError('Invalid rotation type')
128 try:
129 # SciPy 1.6.0
130 is_single_rotation = orientation.single
131 except AttributeError:
132 # SciPy 1.5.4
133 is_single_rotation = len(orientation) == 1
134 if not is_single_rotation:
135 raise ValueError('Rotation object must contain exactly one rotation')
136 self._orientation = orientation
138 def set_orientation(self,
139 rotation: Optional[Rotation] = None,
140 quaternion: Optional[Vector] = None,
141 rotation_matrix: Optional[Matrix] = None,
142 rotation_vector: Optional[Vector] = None,
143 rotation_axis: Optional[Vector] = None,
144 rotation_angle: Optional[float] = None,
145 euler_seq: Optional[str] = None,
146 euler_angles: Optional[Vector] = None,
147 units: str = 'radians') -> None:
148 """Set the orientation of the polarizable unit.
150 The orientation can be set in multiple ways that are mutually exclusive.
152 Parameters
153 ----------
154 rotation:
155 Rotation object.
156 quaternion:
157 Quaternion as in :meth:`~scipy.spatial.transform.Rotation.from_quat`.
158 rotation_matrix:
159 Rotation matrix as in :meth:`~scipy.spatial.transform.Rotation.from_matrix`.
160 rotation_vector:
161 Rotation vector as in :meth:`~scipy.spatial.transform.Rotation.from_rotvec`;
162 its norm gives the angle of rotation in radians.
163 rotation_axis:
164 Rotation axis; the given vector is normalized;
165 use together with :attr:`rotation_angle`.
166 rotation_angle:
167 Rotation angle (see :attr:`units`);
168 use together with :attr:`rotation_axis`.
169 euler_seq:
170 Sequence of axes for Euler rotations
171 as in :meth:`~scipy.spatial.transform.Rotation.from_euler`;
172 use together with :attr:`euler_angles`.
173 euler_angles:
174 Euler angles (see :attr:`units`)
175 as in :meth:`~scipy.spatial.transform.Rotation.from_euler`;
176 use together with :attr:`euler_seq`.
177 units:
178 `radians` or `degrees`.
179 """
180 # Check units argument
181 if units not in ['radians', 'degrees']:
182 raise ValueError(f'unknown units {units}')
184 # Construct Rotation objects based on the arguments
185 count = 0
186 units_def = 'radians'
187 units_error_fmt = 'The units parameter is not supported with {}.'
188 if rotation is not None:
189 if units != units_def:
190 raise ValueError(units_error_fmt.format('rotation'))
191 orientation = rotation
192 count += 1
193 if quaternion is not None:
194 if units != units_def:
195 raise ValueError(units_error_fmt.format('quaternion'))
196 orientation = Rotation.from_quat(quaternion)
197 count += 1
198 if rotation_matrix is not None:
199 if units != units_def:
200 raise ValueError(units_error_fmt.format('rotation_matrix'))
201 orientation = Rotation.from_matrix(rotation_matrix)
202 count += 1
203 if rotation_vector is not None:
204 if units != units_def:
205 raise ValueError(units_error_fmt.format('rotation_vector'))
206 orientation = Rotation.from_rotvec(rotation_vector)
207 count += 1
208 if rotation_axis is not None or rotation_angle is not None:
209 if rotation_axis is None or rotation_angle is None:
210 raise ValueError('To use rotation axis and angle, specify both '
211 'rotation_axis and rotation_angle.')
212 rot_v = np.asarray(rotation_axis, dtype=float)
213 if units == 'degrees':
214 rotation_angle *= np.pi / 180.0
215 rot_v *= rotation_angle / np.linalg.norm(rot_v)
216 orientation = Rotation.from_rotvec(rot_v)
217 count += 1
218 if euler_seq is not None or euler_angles is not None:
219 if euler_seq is None or euler_angles is None:
220 raise ValueError('To use Euler angles, specify both '
221 'euler_seq and euler_angles.')
222 orientation = Rotation.from_euler(euler_seq, euler_angles, degrees=units == 'degrees')
223 count += 1
225 # Check if multiple orientations set
226 if count > 1:
227 raise ValueError('Multiple specifications for orientation.')
229 # Check if no orientations set
230 if count == 0:
231 orientation = Rotation.identity()
233 # Set orientation
234 self.orientation = orientation
236 @property
237 def response(self) -> BaseResponse:
238 """Dielectric response."""
239 return self._response
241 @property
242 def broadening(self) -> Broadening:
243 """Broadening of the response."""
244 return self._response.broadening
246 @property
247 def atoms(self) -> Atoms:
248 """Atomic structure associated with this polarizable unit (if available)."""
249 if self._response.atoms is None:
250 atoms = Atoms(pbc=False)
251 atoms.append(Atom('La', self.position))
252 else:
253 atoms = self._response.atoms.copy()
254 atoms.translate(-atoms.get_center_of_mass())
255 rotation_vector = self._orientation.as_rotvec()
256 rotation_angle = np.linalg.norm(rotation_vector)
257 if rotation_angle != 0.0:
258 atoms.rotate(a=180/np.pi*rotation_angle, v=rotation_vector)
259 atoms.translate(self.position)
260 return atoms
262 @property
263 def excitations(self) -> Excitations:
264 """Excitations as :class:`~strongcoca.response.excitations.Excitations`.
265 Available only if the underlying response can be represented in this form."""
266 if hasattr(self, '_excitations'):
267 return self._excitations # type: ignore
269 if isinstance(self._response, Excitations):
270 exc = self._response
271 elif isinstance(self._response, CasidaResponse):
272 exc = self._response.excitations
273 else:
274 raise RuntimeError(
275 f'{self.__class__.__name__} {self.name} has no response'
276 ' in the form of Excitations.')
277 rot = self._orientation
278 excitations = Excitations(exc._energies,
279 rot.apply(exc._transition_dipole_moments),
280 broadening=exc._broadening,
281 units='au')
282 self._excitations = excitations
283 return excitations
285 def _get_dynamic_polarizability(self, frequencies: Array) -> np.ndarray:
286 dm_wvv = self._response._get_dynamic_polarizability(frequencies)
287 # rotate
288 rot_vv = self._orientation.as_matrix()
289 dm_wvv = np.einsum('wab,xa,yb->wxy', dm_wvv, rot_vv, rot_vv, optimize=True)
290 return dm_wvv # type: ignore
292 def _get_dynamic_polarizability_imaginary_frequency(
293 self, frequencies: Array) -> np.ndarray:
294 dm_wvv = self._response._get_dynamic_polarizability_imaginary_frequency(
295 frequencies)
296 # rotate
297 rot_vv = self._orientation.as_matrix()
298 dm_wvv = np.einsum('wab,xa,yb->wxy', dm_wvv, rot_vv, rot_vv, optimize=True)
299 return dm_wvv # type: ignore