Coverage for cosmolayer/store/profiles.py: 99%
83 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-26 00:09 +0000
1"""Area-weighted sigma profiles, at atom or molecule level."""
3import math
4from dataclasses import dataclass, replace
5from typing import Literal
7import numpy as np
8from numpy.typing import NDArray
9from tqdm.auto import tqdm
11from .binning import (
12 AtomProfileAccumulator,
13 AtomTranslationBatch,
14 accumulate_atom_profiles,
15 accumulate_translated_profiles,
16 compute_per_molecule_properties,
17 row_indices_from_offsets,
18)
19from .grid import DEFAULT_SIGMA_GRID, SigmaGrid
20from .parallel import run_in_threads
23@dataclass(frozen=True)
24class SigmaProfileTable:
25 """Area-weighted sigma profiles, at atom or molecule level.
27 Build with ``from_segments`` (bin segment-level data) or ``aggregate``
28 (atom-level table to molecule-level). Direct construction wraps
29 already-computed arrays.
31 Parameters
32 ----------
33 areas : np.ndarray
34 Per-row area, shape ``(n,)``.
35 charges : np.ndarray
36 Per-row net (or smoothed equivalent) charge, shape ``(n,)``.
37 profiles : np.ndarray
38 Per-row area-fraction sigma profile, shape ``(n, len(grid))``.
39 grid : SigmaGrid
40 Grid these profiles are binned on.
41 centered : bool
42 Whether each row was centered on its mean charge density before
43 binning (zero first moment). A property of the data, not the grid.
44 atom_offsets : np.ndarray | None, optional
45 Global index of each molecule's first atom. ``None`` (default)
46 means molecule-level; when set, the table is atom-level and
47 ``aggregate`` can reassemble it.
48 """
50 areas: NDArray[np.float32]
51 charges: NDArray[np.float32]
52 profiles: NDArray[np.float32]
53 grid: SigmaGrid
54 centered: bool
55 atom_offsets: NDArray[np.int64] | None = None
57 @property
58 def level(self) -> Literal["atom", "molecule"]:
59 """Whether each row of this table describes one atom or one
60 molecule.
62 Returns
63 -------
64 Literal["atom", "molecule"]
65 ``"atom"`` if ``atom_offsets`` is set, else ``"molecule"``.
66 """
67 return "molecule" if self.atom_offsets is None else "atom"
69 @property
70 def sigma_values(self) -> NDArray[np.float64]:
71 """Sigma value at every profile column, in e/Ų.
73 Returns
74 -------
75 np.ndarray, shape (len(grid),)
76 ``grid.values``.
77 """
78 return self.grid.values
80 @classmethod
81 def from_segments( # noqa: PLR0913
82 cls,
83 sigmas: NDArray[np.float64],
84 areas: NDArray[np.float64],
85 segment_offsets: NDArray[np.int64],
86 *,
87 atom_indices: NDArray[np.int64] | None = None,
88 atom_offsets: NDArray[np.int64] | None = None,
89 num_rows: int | None = None,
90 grid: SigmaGrid = DEFAULT_SIGMA_GRID,
91 centered: bool = False,
92 num_threads: int | None = None,
93 progress: bool = False,
94 ) -> "SigmaProfileTable":
95 """Bin segment-level data into per-atom or per-molecule sigma
96 profiles.
98 ``sigmas`` is each segment's charge density: pass raw
99 ``charges / areas``, or an averaged density from
100 ``average_sigmas_by_molecule``.
102 Parameters
103 ----------
104 sigmas : np.ndarray
105 Segment charge density, in e/Ų (raw or averaged).
106 areas : np.ndarray
107 Segment areas.
108 segment_offsets : np.ndarray
109 Start index of each molecule's segments. Must describe
110 *exactly* the molecules present in the segment-level arrays;
111 slice offsets and arrays together.
112 atom_indices : np.ndarray | None, optional
113 Global atom index of each segment. ``None`` (default) builds
114 one profile per molecule. When given, builds one profile per
115 atom and ``atom_offsets`` is required.
116 atom_offsets : np.ndarray | None, optional
117 Start index of each molecule's atoms in the concatenated atom
118 table (including buried atoms that parent no segments).
119 Required when ``atom_indices`` is given; must not be inferred
120 from segment order or from ``min`` over parent atoms.
121 num_rows : int | None, optional
122 Number of profile rows. ``None`` (default) uses
123 ``len(segment_offsets)`` (molecule level) or
124 ``int(atom_indices.max()) + 1`` (atom level). Pass explicitly
125 when subsetting.
126 grid : SigmaGrid, optional
127 Grid to bin onto, by default ``DEFAULT_SIGMA_GRID``.
128 centered : bool, optional
129 If True, center each profile on its mean charge density
130 before binning.
131 num_threads : int | None, optional
132 Thread count. ``None`` (default) uses every CPU core.
133 progress : bool, optional
134 If True, show a tqdm bar over molecules. Default False.
136 Returns
137 -------
138 SigmaProfileTable
139 Atom-level if ``atom_indices`` is given, else molecule-level.
140 """
141 sigmas = np.asarray(sigmas)
142 areas = np.asarray(areas)
143 segment_offsets = np.asarray(segment_offsets)
145 if atom_indices is None:
146 if atom_offsets is not None:
147 raise ValueError("atom_offsets is only valid with atom_indices")
148 row_indices = row_indices_from_offsets(segment_offsets, len(sigmas))
149 num_rows = len(segment_offsets) if num_rows is None else num_rows
150 atom_offsets = None
151 else:
152 if atom_offsets is None:
153 raise ValueError(
154 "atom_offsets is required when atom_indices is given; "
155 "it must be the atom-table start of each molecule, not "
156 "inferred from segments (buried atoms parent none)."
157 )
158 row_indices = np.asarray(atom_indices)
159 num_rows = int(np.max(row_indices)) + 1 if num_rows is None else num_rows
160 atom_offsets = np.asarray(atom_offsets, dtype=np.int64)
162 areas_out = np.zeros(num_rows, dtype=np.float32)
163 charges_out = np.zeros(num_rows, dtype=np.float32)
164 profiles_out = np.zeros((num_rows, len(grid)), dtype=np.float32)
165 assert int(np.max(row_indices, initial=-1)) < num_rows, (
166 "atom_indices/segment_offsets reference a row index >= num_rows; "
167 "segment_offsets must describe exactly the molecules present in "
168 "the segment-level arrays"
169 )
171 num_segs = len(sigmas)
172 num_mols = len(segment_offsets)
173 accumulator = AtomProfileAccumulator(areas_out, charges_out, profiles_out)
175 with tqdm(
176 total=num_mols, desc="Binning sigma profiles", disable=not progress
177 ) as progress_bar:
179 def process_range(start_mol: int, stop_mol: int) -> None:
180 start_seg = segment_offsets[start_mol]
181 stop_seg = (
182 segment_offsets[stop_mol] if stop_mol < num_mols else num_segs
183 )
184 accumulate_atom_profiles(
185 accumulator,
186 sigmas[start_seg:stop_seg],
187 areas[start_seg:stop_seg],
188 row_indices[start_seg:stop_seg],
189 grid,
190 centered,
191 )
192 progress_bar.update(stop_mol - start_mol)
194 run_in_threads(process_range, num_mols, num_threads=num_threads)
196 return cls(
197 areas_out,
198 charges_out,
199 profiles_out,
200 grid,
201 centered,
202 atom_offsets=atom_offsets,
203 )
205 def aggregate(
206 self,
207 *,
208 grid: SigmaGrid | None = None,
209 normalize: bool = False,
210 num_threads: int | None = None,
211 progress: bool = False,
212 ) -> "SigmaProfileTable":
213 """Reassemble per-molecule profiles from these per-atom ones.
215 Area-weighted sum of atom profiles onto a shared molecule axis.
216 If this table is centered, each atom is translated back by its
217 mean charge density first. The result is always uncentered.
219 Parameters
220 ----------
221 grid : SigmaGrid | None, optional
222 Grid for the molecule profiles. ``None`` (default) uses
223 ``self.grid``. A different grid must share this table's
224 ``bin_width``.
225 normalize : bool, optional
226 If True, divide each molecule's profile by its total area so
227 it sums to 1.
228 num_threads : int | None, optional
229 Thread count. ``None`` (default) uses every CPU core.
230 progress : bool, optional
231 If True, show a tqdm bar over molecules. Default False.
233 Returns
234 -------
235 SigmaProfileTable
236 Molecule-level (``atom_offsets`` is None, ``centered`` False).
238 Raises
239 ------
240 ValueError
241 If this table is already molecule-level, or if ``grid`` does
242 not share this table's ``bin_width``.
243 """
244 if self.atom_offsets is None:
245 raise ValueError(
246 "This SigmaProfileTable has no atom_offsets -- it is "
247 "already at molecule level, so there is nothing to "
248 "aggregate."
249 )
250 output_grid = self.grid if grid is None else grid
251 if not math.isclose(output_grid.bin_width, self.grid.bin_width, rel_tol=1e-9):
252 raise ValueError(
253 f"Output grid bin width {output_grid.bin_width!r} does not "
254 f"match this table's {self.grid.bin_width!r}. Aggregation "
255 "aligns columns by point-count difference alone, which is "
256 "only valid between symmetric grids of equal bin width."
257 )
259 num_atoms = len(self.areas)
260 num_mols = len(self.atom_offsets)
261 molecule_num_points = len(output_grid)
263 molecule_indices = row_indices_from_offsets(self.atom_offsets, num_atoms)
265 translations = np.zeros(num_atoms, dtype=np.float64)
266 if self.centered:
267 has_area = self.areas > 0
268 translations[has_area] = self.charges[has_area] / self.areas[has_area]
270 molecule_profiles = np.zeros((num_mols, molecule_num_points), dtype=np.float64)
271 atom_offsets = self.atom_offsets
273 with tqdm(
274 total=num_mols, desc="Aggregating sigma profiles", disable=not progress
275 ) as progress_bar:
277 def process_range(start_mol: int, stop_mol: int) -> None:
278 start_atom = atom_offsets[start_mol]
279 stop_atom = atom_offsets[stop_mol] if stop_mol < num_mols else num_atoms
280 batch = AtomTranslationBatch(
281 self.areas[start_atom:stop_atom],
282 translations[start_atom:stop_atom],
283 self.profiles[start_atom:stop_atom],
284 self.grid,
285 )
286 accumulate_translated_profiles(
287 molecule_profiles,
288 molecule_indices[start_atom:stop_atom],
289 batch,
290 output_grid,
291 )
292 progress_bar.update(stop_mol - start_mol)
294 run_in_threads(process_range, num_mols, num_threads=num_threads)
296 if normalize:
297 molecule_profiles = molecule_profiles / molecule_profiles.sum(
298 axis=1, keepdims=True
299 )
301 molecule_areas = compute_per_molecule_properties(self.areas, atom_offsets)
302 molecule_charges = compute_per_molecule_properties(self.charges, atom_offsets)
304 return replace(
305 self,
306 areas=molecule_areas,
307 charges=molecule_charges,
308 profiles=molecule_profiles.astype(np.float32),
309 grid=output_grid,
310 centered=False,
311 atom_offsets=None,
312 )