Coverage for strongcoca/coupled_system.py: 100%
88 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-10-26 18:44 +0000
1from __future__ import annotations
2import logging
3from typing import Generator, List, Optional
4import numpy as np
6from ase import Atoms
7from ase.geometry import Cell
9from .polarizable_unit import PolarizableUnit
10from .types import Vector
11from .utilities import ClassFormatter
14logger = logging.getLogger(__name__)
17class CoupledSystem:
19 """Objects of this class represent collections of polarizable units.
21 Parameters
22 ----------
23 polarizable_units
24 Polarizable units comprising the coupled system.
25 cell
26 Cell metric (only used for periodic systems).
27 pbc
28 True if coupled system is periodic.
29 name
30 Name of coupled system.
32 Examples
33 --------
35 The following snippet illustrates how a coupled system can be constructed.
36 Here, an artificial response function is used. In practice
37 the response function would usually be imported from, e.g., TDDFT
38 calculations:
40 >>> from strongcoca import CoupledSystem, PolarizableUnit
41 >>> from strongcoca.response import build_random_casida
42 >>>
43 >>> # set up response function
44 >>> response = build_random_casida(n_states=3, name='Panda', random_seed=42)
45 >>>
46 >>> # set up polarizable units
47 >>> pu1 = PolarizableUnit(response, [0, 0, 0], name='Eats')
48 >>> pu2 = PolarizableUnit(response, [2, 0, 0], name='Shoots')
49 >>> pu3 = PolarizableUnit(response, [0, 2, 0], name='Leafs')
50 >>>
51 >>> # construct coupled system
52 >>> cs = CoupledSystem([pu1, pu2, pu3], name='Panda')
53 >>> print(cs)
54 ---------------------- CoupledSystem -----------------------
55 name : Panda
56 pbc : False
57 n_polarizable_units : 3
58 polarizable_unit 0 : Eats
59 polarizable_unit 1 : Shoots
60 polarizable_unit 2 : Leafs
62 CoupledSystem objects behave similar to lists with the polarizable units
63 as items as shown by the following examples:
65 >>> # loop over polarizable units
66 >>> for pu in cs:
67 ... print(pu)
68 --------------------- PolarizableUnit ----------------------
69 name : Eats
70 response : Panda
71 position (Å) : [0. 0. 0.]
72 ...
73 >>>
74 >>> # access a particular unit by index
75 >>> cs[0].name = 'Lime'
76 >>>
77 >>> # add another polarizable unit
78 >>> pu4 = PolarizableUnit(response, [0, 2, 2], name='Bamboo')
79 >>> cs.append(pu4)
80 >>>
81 >>> # add several polarizable units
82 >>> pu5 = PolarizableUnit(response, [-1, -2, 0], name='Oak')
83 >>> pu6 = PolarizableUnit(response, [3, -1, 0], name='Elm')
84 >>> cs.append([pu5, pu6])
85 """
87 def __init__(self,
88 polarizable_units: Optional[List[PolarizableUnit]] = None,
89 cell: Optional[Cell] = None,
90 pbc: bool = False,
91 name: str = 'CoupledSystem') -> None:
92 logger.debug('Setting up coupled system')
93 if polarizable_units is None:
94 self._polarizable_units = [] # type: List[PolarizableUnit]
95 else:
96 if not isinstance(polarizable_units, list) or \
97 not all([isinstance(pu, PolarizableUnit) for pu in polarizable_units]):
98 raise TypeError('polarizable_units must contain a list of PolarizableUnit objects')
99 self._polarizable_units = polarizable_units
100 if pbc is True:
101 logger.warning('Periodic boundary conditions are currently not supported.'
102 ' Proceed with caution.')
103 self._cell = cell
104 self._pbc = pbc
105 self._name = name
107 def __len__(self) -> int:
108 """Length method."""
109 return len(self._polarizable_units)
111 def __iadd__(self, other: CoupledSystem) -> CoupledSystem:
112 self.extend(other)
113 return self
115 def append(self, polarizable_unit: PolarizableUnit) -> None:
116 """Append polarizable unit to this coupled system.
118 Example
119 -------
121 >>> from strongcoca import CoupledSystem, PolarizableUnit
122 >>> from strongcoca.response import build_random_casida
123 >>>
124 >>> response = build_random_casida(n_states=3, name='Panda', random_seed=42)
125 >>>
126 >>> pu1 = PolarizableUnit(response, [0, 0, 0], name='Monkey')
127 >>> pu2 = PolarizableUnit(response, [2, 0, 0], name='Eats')
128 >>> pu3 = PolarizableUnit(response, [0, 2, 0], name='Banana')
129 >>>
130 >>> cs = CoupledSystem([pu1, pu2])
131 >>> print(len(cs))
132 2
133 >>> cs.append(pu3)
134 >>> print(len(cs))
135 3
136 """
137 self._polarizable_units.append(polarizable_unit)
139 def copy(self) -> CoupledSystem:
140 """Return a shallow copy of the coupled system."""
141 cs = self.__class__(polarizable_units=self._polarizable_units,
142 cell=self.cell, pbc=self.pbc, name=self.name)
143 return cs
145 def extend(self, other: CoupledSystem) -> None:
146 """Extend the current coupled system by adding the polarizable units
147 from another coupled system.
149 Example
150 -------
152 >>> from strongcoca import CoupledSystem, PolarizableUnit
153 >>> from strongcoca.response import build_random_casida
154 >>>
155 >>> response = build_random_casida(n_states=3, name='Panda', random_seed=42)
156 >>>
157 >>> pu1 = PolarizableUnit(response, [0, 0, 0], name='Donkey')
158 >>> pu2 = PolarizableUnit(response, [2, 0, 0], name='Camel')
159 >>> pu3 = PolarizableUnit(response, [0, 2, 0], name='Horse')
160 >>> pu4 = PolarizableUnit(response, [0, 0, 2], name='Zebra')
161 >>>
162 >>> cs1 = CoupledSystem([pu1, pu2])
163 >>> cs2 = CoupledSystem([pu3, pu4])
164 >>> print(len(cs1), len(cs2))
165 2 2
166 >>> cs1.extend(cs2)
167 >>> print(len(cs1), len(cs2))
168 4 2
169 """
170 for pu in other:
171 self.append(pu)
173 def __iter__(self) -> Generator[PolarizableUnit, None, None]:
174 for i in range(len(self)):
175 yield self[i]
177 def __getitem__(self, k) -> PolarizableUnit:
178 """Item getter method."""
179 pu = self._polarizable_units[k] # type: PolarizableUnit
180 return pu
182 def __add__(self, other: CoupledSystem) -> CoupledSystem:
183 coupled_system = self.copy()
184 coupled_system += other
185 return coupled_system
187 def __str__(self) -> str:
188 """String representation."""
189 fmt = ClassFormatter(self, pad=21)
190 fmt.append_class_name()
192 fmt.append_attr('name')
193 fmt.append_attr('pbc', self._pbc)
194 if self._pbc:
195 fmt.append_attr('cell', self._cell)
197 if len(self._polarizable_units) > 0:
198 fmt.append_attr('n_polarizable_units', len(self._polarizable_units))
199 for k, pu in enumerate(self._polarizable_units):
200 fmt.append_attr(f' polarizable_unit {k}', pu.name)
202 return fmt.to_string()
204 def to_atoms(self) -> Atoms:
205 """Return the combined atomic configurations of the underlying
206 polarizable units."""
207 atoms = Atoms(pbc=self._pbc)
208 if self._cell is not None:
209 atoms.set_cell(self._cell)
210 for k, pu in enumerate(self._polarizable_units):
211 atoms.extend(pu.atoms)
212 return atoms
214 @property
215 def positions(self) -> List[np.ndarray]:
216 """Positions of centers of mass of polarizable units."""
217 return [pu.position for pu in self._polarizable_units]
219 @positions.setter
220 def positions(self, new_positions: List[Vector]) -> None:
221 """Positions of centers of mass of polarizable units."""
222 if not isinstance(new_positions, list):
223 raise TypeError(f'Invalid positions; not a list: {new_positions}')
224 if len(new_positions) != len(self._polarizable_units):
225 raise ValueError('Invalid positions; does not match number of polarizable_units:'
226 f' {new_positions}')
227 for pu, pos in zip(self._polarizable_units, new_positions):
228 pu.position = pos # type: ignore
230 @property
231 def pbc(self) -> bool:
232 """True if coupled system is periodic."""
233 return self._pbc
235 @property
236 def cell(self) -> Optional[Cell]:
237 """Cell metric of coupled system."""
238 return self._cell
240 @property
241 def name(self) -> str:
242 """Name of coupled system."""
243 return self._name
245 @name.setter
246 def name(self, new_name: str) -> None:
247 """Name of coupled system."""
248 self._name = new_name