Coverage for strongcoca / coupled_system.py: 100%
92 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-04-15 18:15 +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=list(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 _repr_html_(self) -> str:
205 """HTML representation for Jupyter notebooks."""
206 summary_rows = [
207 f'<tr><td style="text-align: left;">name</td><td>{self._name}</td></tr>',
208 f'<tr><td style="text-align: left;">pbc</td><td>{self._pbc}</td></tr>',
209 f'<tr><td style="text-align: left;">n_polarizable_units</td><td>{len(self)}</td></tr>',
210 ]
211 html = (
212 '<h4>CoupledSystem</h4>'
213 '<table>'
214 '<thead><tr>'
215 '<th style="text-align: left;">field</th>'
216 '<th style="text-align: left;">value</th>'
217 '</tr></thead>'
218 '<tbody>' + ''.join(summary_rows) + '</tbody>'
219 '</table>'
220 )
221 return html
223 def to_atoms(self) -> Atoms:
224 """Return the combined atomic configurations of the underlying
225 polarizable units."""
226 atoms = Atoms(pbc=self._pbc)
227 if self._cell is not None:
228 atoms.set_cell(self._cell)
229 for k, pu in enumerate(self._polarizable_units):
230 atoms.extend(pu.atoms)
231 return atoms
233 @property
234 def positions(self) -> List[np.ndarray]:
235 """Positions of centers of mass of polarizable units."""
236 return [pu.position for pu in self._polarizable_units]
238 @positions.setter
239 def positions(self, new_positions: List[Vector]) -> None:
240 """Positions of centers of mass of polarizable units."""
241 if not isinstance(new_positions, list):
242 raise TypeError(f'Invalid positions; not a list: {new_positions}')
243 if len(new_positions) != len(self._polarizable_units):
244 raise ValueError('Invalid positions; does not match number of polarizable_units:'
245 f' {new_positions}')
246 for pu, pos in zip(self._polarizable_units, new_positions):
247 pu.position = pos # type: ignore
249 @property
250 def pbc(self) -> bool:
251 """True if coupled system is periodic."""
252 return self._pbc
254 @property
255 def cell(self) -> Optional[Cell]:
256 """Cell metric of coupled system."""
257 return self._cell
259 @property
260 def name(self) -> str:
261 """Name of coupled system."""
262 return self._name
264 @name.setter
265 def name(self, new_name: str) -> None:
266 """Name of coupled system."""
267 self._name = new_name