from __future__ import annotations
import logging
from typing import Generator, List, Optional
import numpy as np
from ase import Atoms
from ase.geometry import Cell
from .polarizable_unit import PolarizableUnit
from .types import Vector
from .utilities import ClassFormatter
logger = logging.getLogger(__name__)
[docs]
class CoupledSystem:
"""Objects of this class represent collections of polarizable units.
Parameters
----------
polarizable_units
Polarizable units comprising the coupled system.
cell
Cell metric (only used for periodic systems).
pbc
True if coupled system is periodic.
name
Name of coupled system.
Examples
--------
The following snippet illustrates how a coupled system can be constructed.
Here, an artificial response function is used. In practice
the response function would usually be imported from, e.g., TDDFT
calculations:
>>> from strongcoca import CoupledSystem, PolarizableUnit
>>> from strongcoca.response import build_random_casida
>>>
>>> # set up response function
>>> response = build_random_casida(n_states=3, name='Panda', random_seed=42)
>>>
>>> # set up polarizable units
>>> pu1 = PolarizableUnit(response, [0, 0, 0], name='Eats')
>>> pu2 = PolarizableUnit(response, [2, 0, 0], name='Shoots')
>>> pu3 = PolarizableUnit(response, [0, 2, 0], name='Leafs')
>>>
>>> # construct coupled system
>>> cs = CoupledSystem([pu1, pu2, pu3], name='Panda')
>>> print(cs)
---------------------- CoupledSystem -----------------------
name : Panda
pbc : False
n_polarizable_units : 3
polarizable_unit 0 : Eats
polarizable_unit 1 : Shoots
polarizable_unit 2 : Leafs
CoupledSystem objects behave similar to lists with the polarizable units
as items as shown by the following examples:
>>> # loop over polarizable units
>>> for pu in cs:
... print(pu)
--------------------- PolarizableUnit ----------------------
name : Eats
response : Panda
position (Å) : [0. 0. 0.]
...
>>>
>>> # access a particular unit by index
>>> cs[0].name = 'Lime'
>>>
>>> # add another polarizable unit
>>> pu4 = PolarizableUnit(response, [0, 2, 2], name='Bamboo')
>>> cs.append(pu4)
>>>
>>> # add several polarizable units
>>> pu5 = PolarizableUnit(response, [-1, -2, 0], name='Oak')
>>> pu6 = PolarizableUnit(response, [3, -1, 0], name='Elm')
>>> cs.append([pu5, pu6])
"""
def __init__(self,
polarizable_units: Optional[List[PolarizableUnit]] = None,
cell: Optional[Cell] = None,
pbc: bool = False,
name: str = 'CoupledSystem') -> None:
logger.debug('Setting up coupled system')
if polarizable_units is None:
self._polarizable_units = [] # type: List[PolarizableUnit]
else:
if not isinstance(polarizable_units, list) or \
not all([isinstance(pu, PolarizableUnit) for pu in polarizable_units]):
raise TypeError('polarizable_units must contain a list of PolarizableUnit objects')
self._polarizable_units = polarizable_units
if pbc is True:
logger.warning('Periodic boundary conditions are currently not supported.'
' Proceed with caution.')
self._cell = cell
self._pbc = pbc
self._name = name
def __len__(self) -> int:
"""Length method."""
return len(self._polarizable_units)
def __iadd__(self, other: CoupledSystem) -> CoupledSystem:
self.extend(other)
return self
[docs]
def append(self, polarizable_unit: PolarizableUnit) -> None:
"""Append polarizable unit to this coupled system.
Example
-------
>>> from strongcoca import CoupledSystem, PolarizableUnit
>>> from strongcoca.response import build_random_casida
>>>
>>> response = build_random_casida(n_states=3, name='Panda', random_seed=42)
>>>
>>> pu1 = PolarizableUnit(response, [0, 0, 0], name='Monkey')
>>> pu2 = PolarizableUnit(response, [2, 0, 0], name='Eats')
>>> pu3 = PolarizableUnit(response, [0, 2, 0], name='Banana')
>>>
>>> cs = CoupledSystem([pu1, pu2])
>>> print(len(cs))
2
>>> cs.append(pu3)
>>> print(len(cs))
3
"""
self._polarizable_units.append(polarizable_unit)
[docs]
def copy(self) -> CoupledSystem:
"""Return a shallow copy of the coupled system."""
cs = self.__class__(polarizable_units=self._polarizable_units,
cell=self.cell, pbc=self.pbc, name=self.name)
return cs
[docs]
def extend(self, other: CoupledSystem) -> None:
"""Extend the current coupled system by adding the polarizable units
from another coupled system.
Example
-------
>>> from strongcoca import CoupledSystem, PolarizableUnit
>>> from strongcoca.response import build_random_casida
>>>
>>> response = build_random_casida(n_states=3, name='Panda', random_seed=42)
>>>
>>> pu1 = PolarizableUnit(response, [0, 0, 0], name='Donkey')
>>> pu2 = PolarizableUnit(response, [2, 0, 0], name='Camel')
>>> pu3 = PolarizableUnit(response, [0, 2, 0], name='Horse')
>>> pu4 = PolarizableUnit(response, [0, 0, 2], name='Zebra')
>>>
>>> cs1 = CoupledSystem([pu1, pu2])
>>> cs2 = CoupledSystem([pu3, pu4])
>>> print(len(cs1), len(cs2))
2 2
>>> cs1.extend(cs2)
>>> print(len(cs1), len(cs2))
4 2
"""
for pu in other:
self.append(pu)
def __iter__(self) -> Generator[PolarizableUnit, None, None]:
for i in range(len(self)):
yield self[i]
def __getitem__(self, k) -> PolarizableUnit:
"""Item getter method."""
pu = self._polarizable_units[k] # type: PolarizableUnit
return pu
def __add__(self, other: CoupledSystem) -> CoupledSystem:
coupled_system = self.copy()
coupled_system += other
return coupled_system
def __str__(self) -> str:
"""String representation."""
fmt = ClassFormatter(self, pad=21)
fmt.append_class_name()
fmt.append_attr('name')
fmt.append_attr('pbc', self._pbc)
if self._pbc:
fmt.append_attr('cell', self._cell)
if len(self._polarizable_units) > 0:
fmt.append_attr('n_polarizable_units', len(self._polarizable_units))
for k, pu in enumerate(self._polarizable_units):
fmt.append_attr(f' polarizable_unit {k}', pu.name)
return fmt.to_string()
[docs]
def to_atoms(self) -> Atoms:
"""Return the combined atomic configurations of the underlying
polarizable units."""
atoms = Atoms(pbc=self._pbc)
if self._cell is not None:
atoms.set_cell(self._cell)
for k, pu in enumerate(self._polarizable_units):
atoms.extend(pu.atoms)
return atoms
@property
def positions(self) -> List[np.ndarray]:
"""Positions of centers of mass of polarizable units."""
return [pu.position for pu in self._polarizable_units]
@positions.setter
def positions(self, new_positions: List[Vector]) -> None:
"""Positions of centers of mass of polarizable units."""
if not isinstance(new_positions, list):
raise TypeError(f'Invalid positions; not a list: {new_positions}')
if len(new_positions) != len(self._polarizable_units):
raise ValueError('Invalid positions; does not match number of polarizable_units:'
f' {new_positions}')
for pu, pos in zip(self._polarizable_units, new_positions):
pu.position = pos # type: ignore
@property
def pbc(self) -> bool:
"""True if coupled system is periodic."""
return self._pbc
@property
def cell(self) -> Optional[Cell]:
"""Cell metric of coupled system."""
return self._cell
@property
def name(self) -> str:
"""Name of coupled system."""
return self._name
@name.setter
def name(self, new_name: str) -> None:
"""Name of coupled system."""
self._name = new_name
