Coverage for strongcoca / polarizable_unit.py: 100%
152 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-04-15 18:15 +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 def _repr_html_(self) -> str:
93 """HTML representation for Jupyter notebooks."""
94 pos = self.position
95 quat = self._orientation.as_quat()
96 rot_v = self._orientation.as_rotvec()
97 rows = [
98 f'<tr><td style="text-align: left;">name</td><td>{self._name}</td></tr>',
99 f'<tr><td style="text-align: left;">response</td><td>{self._response.name}</td></tr>',
100 f'<tr><td style="text-align: left;">position (Å)</td>'
101 f'<td>[{pos[0]:.4f}, {pos[1]:.4f}, {pos[2]:.4f}]</td></tr>',
102 f'<tr><td style="text-align: left;">quaternion</td>'
103 f'<td>[{quat[0]:.4f}, {quat[1]:.4f}, {quat[2]:.4f}, {quat[3]:.4f}]</td></tr>',
104 f'<tr><td style="text-align: left;">rotation angle</td>'
105 f'<td>{np.linalg.norm(rot_v):.4f}</td></tr>',
106 ]
107 return (
108 '<h4>PolarizableUnit</h4>'
109 '<table>'
110 '<thead><tr>'
111 '<th style="text-align: left;">field</th>'
112 '<th style="text-align: left;">value</th>'
113 '</tr></thead>'
114 '<tbody>' + ''.join(rows) + '</tbody>'
115 '</table>'
116 )
118 @property
119 def name(self) -> str:
120 """Name of polarizable unit."""
121 return self._name
123 @name.setter
124 def name(self, name: str) -> None:
125 """Set name of polarizable unit."""
126 self._name = name
128 @property
129 def position(self) -> np.ndarray:
130 """Position of center of mass of polarizable unit."""
131 return self._position * au_to_A # type: ignore
133 @position.setter
134 def position(self, new_position: Vector) -> None:
135 """Set position of center of mass of polarizable unit."""
136 _position = np.array(new_position) * A_to_au
137 if _position.shape != (3, ):
138 raise TypeError('Invalid position; not a three-dimensional vector of numbers:'
139 f' {new_position}, {type(new_position)}')
140 self._position = _position
142 @property
143 def orientation(self) -> Rotation:
144 """Orientation of the polarizable unit."""
145 return self._orientation
147 @orientation.setter
148 def orientation(self, orientation: Optional[Rotation]) -> None:
149 """Set the orientation of the polarizable unit."""
150 if orientation is None:
151 orientation = Rotation.identity()
152 if not isinstance(orientation, Rotation):
153 raise TypeError('Invalid rotation type')
154 if not orientation.single:
155 raise ValueError('Rotation object must contain exactly one rotation')
156 self._orientation = orientation
158 def set_orientation(self,
159 rotation: Optional[Rotation] = None,
160 quaternion: Optional[Vector] = None,
161 rotation_matrix: Optional[Matrix] = None,
162 rotation_vector: Optional[Vector] = None,
163 rotation_axis: Optional[Vector] = None,
164 rotation_angle: Optional[float] = None,
165 euler_seq: Optional[str] = None,
166 euler_angles: Optional[Vector] = None,
167 units: str = 'radians') -> None:
168 """Set the orientation of the polarizable unit.
170 The orientation can be set in multiple ways that are mutually exclusive.
172 Parameters
173 ----------
174 rotation:
175 Rotation object.
176 quaternion:
177 Quaternion as in :meth:`~scipy.spatial.transform.Rotation.from_quat`.
178 rotation_matrix:
179 Rotation matrix as in :meth:`~scipy.spatial.transform.Rotation.from_matrix`.
180 rotation_vector:
181 Rotation vector as in :meth:`~scipy.spatial.transform.Rotation.from_rotvec`;
182 its norm gives the angle of rotation in radians.
183 rotation_axis:
184 Rotation axis; the given vector is normalized;
185 use together with :attr:`rotation_angle`.
186 rotation_angle:
187 Rotation angle (see :attr:`units`);
188 use together with :attr:`rotation_axis`.
189 euler_seq:
190 Sequence of axes for Euler rotations
191 as in :meth:`~scipy.spatial.transform.Rotation.from_euler`;
192 use together with :attr:`euler_angles`.
193 euler_angles:
194 Euler angles (see :attr:`units`)
195 as in :meth:`~scipy.spatial.transform.Rotation.from_euler`;
196 use together with :attr:`euler_seq`.
197 units:
198 `radians` or `degrees`.
199 """
200 # Check units argument
201 if units not in ['radians', 'degrees']:
202 raise ValueError(f'unknown units {units}')
204 # Construct Rotation objects based on the arguments
205 count = 0
206 units_def = 'radians'
207 units_error_fmt = 'The units parameter is not supported with {}.'
208 if rotation is not None:
209 if units != units_def:
210 raise ValueError(units_error_fmt.format('rotation'))
211 orientation = rotation
212 count += 1
213 if quaternion is not None:
214 if units != units_def:
215 raise ValueError(units_error_fmt.format('quaternion'))
216 orientation = Rotation.from_quat(quaternion)
217 count += 1
218 if rotation_matrix is not None:
219 if units != units_def:
220 raise ValueError(units_error_fmt.format('rotation_matrix'))
221 orientation = Rotation.from_matrix(rotation_matrix)
222 count += 1
223 if rotation_vector is not None:
224 if units != units_def:
225 raise ValueError(units_error_fmt.format('rotation_vector'))
226 orientation = Rotation.from_rotvec(rotation_vector)
227 count += 1
228 if rotation_axis is not None or rotation_angle is not None:
229 if rotation_axis is None or rotation_angle is None:
230 raise ValueError('To use rotation axis and angle, specify both '
231 'rotation_axis and rotation_angle.')
232 rot_v = np.asarray(rotation_axis, dtype=float)
233 if units == 'degrees':
234 rotation_angle *= np.pi / 180.0
235 rot_v *= rotation_angle / np.linalg.norm(rot_v)
236 orientation = Rotation.from_rotvec(rot_v)
237 count += 1
238 if euler_seq is not None or euler_angles is not None:
239 if euler_seq is None or euler_angles is None:
240 raise ValueError('To use Euler angles, specify both '
241 'euler_seq and euler_angles.')
242 orientation = Rotation.from_euler(euler_seq, euler_angles, degrees=units == 'degrees')
243 count += 1
245 # Check if multiple orientations set
246 if count > 1:
247 raise ValueError('Multiple specifications for orientation.')
249 # Check if no orientations set
250 if count == 0:
251 orientation = Rotation.identity()
253 # Set orientation
254 self.orientation = orientation
256 @property
257 def response(self) -> BaseResponse:
258 """Dielectric response."""
259 return self._response
261 @property
262 def broadening(self) -> Broadening:
263 """Broadening of the response."""
264 return self._response.broadening
266 @property
267 def atoms(self) -> Atoms:
268 """Atomic structure associated with this polarizable unit (if available)."""
269 if self._response.atoms is None:
270 atoms = Atoms(pbc=False)
271 atoms.append(Atom('La', self.position))
272 else:
273 atoms = self._response.atoms.copy()
274 atoms.translate(-atoms.get_center_of_mass())
275 rotation_vector = self._orientation.as_rotvec()
276 rotation_angle = np.linalg.norm(rotation_vector)
277 if rotation_angle != 0.0:
278 atoms.rotate(a=180/np.pi*rotation_angle, v=rotation_vector)
279 atoms.translate(self.position)
280 return atoms
282 @property
283 def excitations(self) -> Excitations:
284 """Excitations as :class:`~strongcoca.response.excitations.Excitations`.
285 Available only if the underlying response can be represented in this form."""
286 if hasattr(self, '_excitations'):
287 return self._excitations # type: ignore
289 if isinstance(self._response, Excitations):
290 exc = self._response
291 elif isinstance(self._response, CasidaResponse):
292 exc = self._response.excitations
293 else:
294 raise RuntimeError(
295 f'{self.__class__.__name__} {self.name} has no response'
296 ' in the form of Excitations.')
297 rot = self._orientation
298 excitations = Excitations(exc._energies,
299 rot.apply(exc._transition_dipole_moments),
300 broadening=exc._broadening,
301 units='au')
302 self._excitations = excitations
303 return excitations
305 def _get_dynamic_polarizability(self, frequencies: Array) -> np.ndarray:
306 dm_wvv = self._response._get_dynamic_polarizability(frequencies)
307 # rotate
308 rot_vv = self._orientation.as_matrix()
309 dm_wvv = np.einsum('wab,xa,yb->wxy', dm_wvv, rot_vv, rot_vv, optimize=True)
310 return dm_wvv # type: ignore
312 def _get_dynamic_polarizability_imaginary_frequency(
313 self, frequencies: Array) -> np.ndarray:
314 dm_wvv = self._response._get_dynamic_polarizability_imaginary_frequency(
315 frequencies)
316 # rotate
317 rot_vv = self._orientation.as_matrix()
318 dm_wvv = np.einsum('wab,xa,yb->wxy', dm_wvv, rot_vv, rot_vv, optimize=True)
319 return dm_wvv # type: ignore