from typing import Optional
import numpy as np
from scipy.spatial.transform import Rotation
from ase import Atom, Atoms
from .response.base import BaseResponse
from .response.casida import CasidaResponse
from .response.excitations import Excitations
from .response.utilities import Broadening
from .types import Array, Matrix, Vector
from .units import au_to_A, A_to_au
from .utilities import ClassFormatter
[docs]
class PolarizableUnit(BaseResponse):
"""Class for managing polarizable units in CoupledSystem objects along
with their position and orientation.
Parameters
----------
response
Object that handles response function representations.
position
Position of polarizable unit in Å.
orientation
Orientation of polarizable unit.
name
Name of polarizable unit.
**orientation_kwargs
Alternative keyword arguments for setting orientation;
these arguments are passed to :meth:`set_orientation`.
Examples
--------
The following snippet illustrates how polarizable unit objects can be
initialized. Here, an artificial response function is used. In
practice the response function would usually be imported from, e.g., TDDFT
calculations:
>>> from math import pi
>>> from strongcoca import PolarizableUnit
>>> from strongcoca.response import build_random_casida
>>>
>>> response = build_random_casida(n_states=3, name='Merkat', random_seed=42)
>>> pu = PolarizableUnit(response, [0, 2, -1], name='Apple',
... rotation_vector=[0, 0, pi])
>>>
>>> print(pu)
--------------------- PolarizableUnit ----------------------
name : Apple
response : Merkat
position (Å) : [ 0. 2. -1.]
quaternion : [0. 0. 1. 0.]
rotation_vector : [0. ... 0. ... 3.14...]
rotation_angle : 3.14...
"""
def __init__(self,
response: BaseResponse,
position: Vector,
orientation: Optional[Rotation] = None,
name: str = 'PolarizableUnit',
**orientation_kwargs):
super().__init__(broadening=Broadening(), pbc=False, name=name)
self._response = response
self.position = np.array(position)
if len(orientation_kwargs) > 0:
if orientation is not None:
raise ValueError('Both orientation and orientation keywords specified.')
self.set_orientation(**orientation_kwargs)
else:
self.orientation = orientation
def __str__(self) -> str:
"""String representation."""
fmt = ClassFormatter(self)
fmt.append_class_name()
fmt.append_attr('name')
fmt.append_attr('response', self._response.name)
fmt.append_attr('position', unit='Å')
fmt.append_attr('quaternion', np.around(self._orientation.as_quat(), 2))
rot_v = self._orientation.as_rotvec()
fmt.append_attr('rotation_vector', rot_v)
fmt.append_attr('rotation_angle', np.linalg.norm(rot_v))
return fmt.to_string()
@property
def name(self) -> str:
"""Name of polarizable unit."""
return self._name
@name.setter
def name(self, name: str) -> None:
"""Set name of polarizable unit."""
self._name = name
@property
def position(self) -> np.ndarray:
"""Position of center of mass of polarizable unit."""
return self._position * au_to_A # type: ignore
@position.setter
def position(self, new_position: Vector) -> None:
"""Set position of center of mass of polarizable unit."""
_position = np.array(new_position) * A_to_au
if _position.shape != (3, ):
raise TypeError('Invalid position; not a three-dimensional vector of numbers:'
f' {new_position}, {type(new_position)}')
self._position = _position
@property
def orientation(self) -> Rotation:
"""Orientation of the polarizable unit."""
return self._orientation
@orientation.setter
def orientation(self, orientation: Optional[Rotation]) -> None:
"""Set the orientation of the polarizable unit."""
if orientation is None:
orientation = Rotation.identity()
if not isinstance(orientation, Rotation):
raise TypeError('Invalid rotation type')
try:
# SciPy 1.6.0
is_single_rotation = orientation.single
except AttributeError:
# SciPy 1.5.4
is_single_rotation = len(orientation) == 1
if not is_single_rotation:
raise ValueError('Rotation object must contain exactly one rotation')
self._orientation = orientation
[docs]
def set_orientation(self,
rotation: Optional[Rotation] = None,
quaternion: Optional[Vector] = None,
rotation_matrix: Optional[Matrix] = None,
rotation_vector: Optional[Vector] = None,
rotation_axis: Optional[Vector] = None,
rotation_angle: Optional[float] = None,
euler_seq: Optional[str] = None,
euler_angles: Optional[Vector] = None,
units: str = 'radians') -> None:
"""Set the orientation of the polarizable unit.
The orientation can be set in multiple ways that are mutually exclusive.
Parameters
----------
rotation:
Rotation object.
quaternion:
Quaternion as in :meth:`~scipy.spatial.transform.Rotation.from_quat`.
rotation_matrix:
Rotation matrix as in :meth:`~scipy.spatial.transform.Rotation.from_matrix`.
rotation_vector:
Rotation vector as in :meth:`~scipy.spatial.transform.Rotation.from_rotvec`;
its norm gives the angle of rotation in radians.
rotation_axis:
Rotation axis; the given vector is normalized;
use together with :attr:`rotation_angle`.
rotation_angle:
Rotation angle (see :attr:`units`);
use together with :attr:`rotation_axis`.
euler_seq:
Sequence of axes for Euler rotations
as in :meth:`~scipy.spatial.transform.Rotation.from_euler`;
use together with :attr:`euler_angles`.
euler_angles:
Euler angles (see :attr:`units`)
as in :meth:`~scipy.spatial.transform.Rotation.from_euler`;
use together with :attr:`euler_seq`.
units:
`radians` or `degrees`.
"""
# Check units argument
if units not in ['radians', 'degrees']:
raise ValueError(f'unknown units {units}')
# Construct Rotation objects based on the arguments
count = 0
units_def = 'radians'
units_error_fmt = 'The units parameter is not supported with {}.'
if rotation is not None:
if units != units_def:
raise ValueError(units_error_fmt.format('rotation'))
orientation = rotation
count += 1
if quaternion is not None:
if units != units_def:
raise ValueError(units_error_fmt.format('quaternion'))
orientation = Rotation.from_quat(quaternion)
count += 1
if rotation_matrix is not None:
if units != units_def:
raise ValueError(units_error_fmt.format('rotation_matrix'))
orientation = Rotation.from_matrix(rotation_matrix)
count += 1
if rotation_vector is not None:
if units != units_def:
raise ValueError(units_error_fmt.format('rotation_vector'))
orientation = Rotation.from_rotvec(rotation_vector)
count += 1
if rotation_axis is not None or rotation_angle is not None:
if rotation_axis is None or rotation_angle is None:
raise ValueError('To use rotation axis and angle, specify both '
'rotation_axis and rotation_angle.')
rot_v = np.asarray(rotation_axis, dtype=float)
if units == 'degrees':
rotation_angle *= np.pi / 180.0
rot_v *= rotation_angle / np.linalg.norm(rot_v)
orientation = Rotation.from_rotvec(rot_v)
count += 1
if euler_seq is not None or euler_angles is not None:
if euler_seq is None or euler_angles is None:
raise ValueError('To use Euler angles, specify both '
'euler_seq and euler_angles.')
orientation = Rotation.from_euler(euler_seq, euler_angles, degrees=units == 'degrees')
count += 1
# Check if multiple orientations set
if count > 1:
raise ValueError('Multiple specifications for orientation.')
# Check if no orientations set
if count == 0:
orientation = Rotation.identity()
# Set orientation
self.orientation = orientation
@property
def response(self) -> BaseResponse:
"""Dielectric response."""
return self._response
@property
def broadening(self) -> Broadening:
"""Broadening of the response."""
return self._response.broadening
@property
def atoms(self) -> Atoms:
"""Atomic structure associated with this polarizable unit (if available)."""
if self._response.atoms is None:
atoms = Atoms(pbc=False)
atoms.append(Atom('La', self.position))
else:
atoms = self._response.atoms.copy()
atoms.translate(-atoms.get_center_of_mass())
rotation_vector = self._orientation.as_rotvec()
rotation_angle = np.linalg.norm(rotation_vector)
if rotation_angle != 0.0:
atoms.rotate(a=180/np.pi*rotation_angle, v=rotation_vector)
atoms.translate(self.position)
return atoms
@property
def excitations(self) -> Excitations:
"""Excitations as :class:`~strongcoca.response.excitations.Excitations`.
Available only if the underlying response can be represented in this form."""
if hasattr(self, '_excitations'):
return self._excitations # type: ignore
if isinstance(self._response, Excitations):
exc = self._response
elif isinstance(self._response, CasidaResponse):
exc = self._response.excitations
else:
raise RuntimeError(
f'{self.__class__.__name__} {self.name} has no response'
' in the form of Excitations.')
rot = self._orientation
excitations = Excitations(exc._energies,
rot.apply(exc._transition_dipole_moments),
broadening=exc._broadening,
units='au')
self._excitations = excitations
return excitations
def _get_dynamic_polarizability(self, frequencies: Array) -> np.ndarray:
dm_wvv = self._response._get_dynamic_polarizability(frequencies)
# rotate
rot_vv = self._orientation.as_matrix()
dm_wvv = np.einsum('wab,xa,yb->wxy', dm_wvv, rot_vv, rot_vv, optimize=True)
return dm_wvv # type: ignore
def _get_dynamic_polarizability_imaginary_frequency(
self, frequencies: Array) -> np.ndarray:
dm_wvv = self._response._get_dynamic_polarizability_imaginary_frequency(
frequencies)
# rotate
rot_vv = self._orientation.as_matrix()
dm_wvv = np.einsum('wab,xa,yb->wxy', dm_wvv, rot_vv, rot_vv, optimize=True)
return dm_wvv # type: ignore
