Coverage for cosmolayer/store/binning.py: 100%
55 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-26 00:09 +0000
1"""Linear binning of segment- and atom-level data onto a ``SigmaGrid``.
3Every accumulation function here uses the same two-tap linear
4interpolation: a value is split between the two nearest grid points in
5proportion to its distance from each, so mass is conserved exactly except
6for values outside the grid's range, which are folded into the nearest
7boundary point.
8"""
10from dataclasses import dataclass
11from typing import Any, TypeVar
13import numpy as np
14from numpy.typing import NDArray
16from .grid import SigmaGrid
18_Float = TypeVar("_Float", bound=np.floating[Any])
21def row_indices_from_offsets(
22 offsets: NDArray[np.int64], total_num_rows: int
23) -> NDArray[np.int64]:
24 """Expand a table of row-start offsets into a row index per element.
26 Parameters
27 ----------
28 offsets : np.ndarray
29 Start index of each row's elements within the element-level
30 arrays. The last row's elements are assumed to run to
31 ``total_num_rows``.
32 total_num_rows : int
33 Total number of elements described by ``offsets``.
35 Returns
36 -------
37 np.ndarray
38 Row index of every element, shape ``(total_num_rows,)``.
40 Examples
41 --------
42 >>> row_indices_from_offsets(np.array([0, 2, 3]), 5)
43 array([0, 0, 1, 2, 2])
44 """
45 return np.repeat(
46 np.arange(len(offsets), dtype=np.int64),
47 np.diff(np.append(offsets, total_num_rows)),
48 )
51def compute_per_atom_properties(
52 properties: NDArray[_Float],
53 atom_indices: NDArray[np.int64],
54 total_num_atoms: int,
55) -> NDArray[_Float]:
56 """Sum a segment-level property into per-atom totals.
58 Parameters
59 ----------
60 properties : np.ndarray
61 Segment-level values to sum, e.g. charges or areas.
62 atom_indices : np.ndarray
63 Atom indices for the segments.
64 total_num_atoms : int
65 Total number of atoms in the dataset.
67 Returns
68 -------
69 np.ndarray
70 Per-atom sum of ``properties``, of shape ``(total_num_atoms,)``.
71 """
72 atom_properties = np.zeros(total_num_atoms, dtype=properties.dtype)
73 np.add.at(atom_properties, atom_indices, properties)
74 return atom_properties
77def compute_per_molecule_properties(
78 properties: NDArray[_Float], atom_offsets: NDArray[np.int64]
79) -> NDArray[_Float]:
80 """Compute the per-molecule sum of an atom-level property.
82 Parameters
83 ----------
84 properties : np.ndarray
85 Atom-level property values, of shape ``(total_num_atoms,)``.
86 atom_offsets : np.ndarray
87 Global index of each molecule's first atom (cumulative sum of
88 ``num_atoms`` for preceding molecules) -- *not* the
89 ``segment_offsets`` column of the ``molecules`` table, which
90 indexes segment-level arrays instead.
92 Returns
93 -------
94 np.ndarray
95 Per-molecule sum of ``properties``, of shape ``(n_molecules,)``.
96 """
97 return np.add.reduceat(properties, atom_offsets)
100@dataclass
101class AtomProfileAccumulator:
102 """Per-atom areas, charges, and profiles updated in place by
103 ``accumulate_atom_profiles``.
105 Parameters
106 ----------
107 areas : np.ndarray
108 Per-atom areas, shape ``(num_atoms,)``.
109 charges : np.ndarray
110 Per-atom charges, shape ``(num_atoms,)``.
111 profiles : np.ndarray
112 Per-atom sigma profiles, shape ``(num_atoms, num_points)``.
113 """
115 areas: NDArray[np.float32]
116 charges: NDArray[np.float32]
117 profiles: NDArray[np.float32]
120def accumulate_atom_profiles(
121 accumulator: AtomProfileAccumulator,
122 sigmas: NDArray[np.float64],
123 areas: NDArray[np.float64],
124 atom_indices: NDArray[np.int64],
125 grid: SigmaGrid,
126 centered: bool,
127) -> None:
128 """Accumulate a batch of segments into per-atom sigma profiles.
130 Each segment's charge density is linearly interpolated between the
131 two nearest grid points, and its area split between those points.
132 Values outside ``grid`` fold into the nearest boundary. Each atom's
133 profile sums to 1 (all-zero if it has no surface segments).
135 ``sigmas`` is charge density (raw ``charges / areas``, or averaged).
136 ``accumulator.charges`` accumulates ``sigmas * areas``.
138 Safe concurrently on disjoint ranges split on *molecule* boundaries:
139 every segment of an atom must land in the same batch.
141 Parameters
142 ----------
143 accumulator : AtomProfileAccumulator
144 Per-atom arrays to update in place.
145 sigmas : np.ndarray
146 Segment charge density in this batch, in e/Ų.
147 areas : np.ndarray
148 Segment areas in this batch.
149 atom_indices : np.ndarray
150 Global atom index of each segment in this batch.
151 grid : SigmaGrid
152 Grid to bin onto.
153 centered : bool
154 If True, center each atom's profile on its mean charge density
155 before binning.
156 """
157 atom_areas, atom_charges, sigma_profiles = (
158 accumulator.areas,
159 accumulator.charges,
160 accumulator.profiles,
161 )
162 np.add.at(atom_areas, atom_indices, areas)
163 np.add.at(atom_charges, atom_indices, sigmas * areas)
165 num_points = len(grid)
166 summed_areas = atom_areas[atom_indices]
168 if centered:
169 sigmas = sigmas - atom_charges[atom_indices] / summed_areas
171 fractional_bins = (sigmas - (-grid.max_abs_sigma)) / grid.bin_width
173 points_at_left = np.floor(fractional_bins).astype(int)
174 points_at_right = points_at_left + 1
176 normalized_areas = areas / summed_areas
177 contributions_at_left = normalized_areas * (points_at_right - fractional_bins)
178 contributions_at_right = normalized_areas * (fractional_bins - points_at_left)
180 np.add.at(
181 sigma_profiles,
182 (atom_indices, points_at_left.clip(0, num_points - 1)),
183 contributions_at_left,
184 )
185 np.add.at(
186 sigma_profiles,
187 (atom_indices, points_at_right.clip(0, num_points - 1)),
188 contributions_at_right,
189 )
192@dataclass
193class AtomTranslationBatch:
194 """Atom profiles to translate and sum onto a molecule grid.
196 Parameters
197 ----------
198 areas : np.ndarray
199 Areas of the atoms in this batch, shape ``(n,)``.
200 translations : np.ndarray
201 Shift of each atom's profile, in sigma units (positive toward
202 the positive end of the grid). Pass zeros to leave untranslated.
203 profiles : np.ndarray
204 Per-atom sigma profiles, shape ``(n, len(grid))``.
205 grid : SigmaGrid
206 Grid ``profiles`` are binned on.
207 """
209 areas: NDArray[np.float32]
210 translations: NDArray[np.float64]
211 profiles: NDArray[np.float32]
212 grid: SigmaGrid
215def accumulate_translated_profiles(
216 molecule_profiles: NDArray[np.float64],
217 molecule_indices: NDArray[np.int64],
218 batch: AtomTranslationBatch,
219 molecule_grid: SigmaGrid,
220) -> None:
221 """Accumulate translated, area-weighted atom profiles into molecules.
223 Both grids must be symmetric about zero and share a ``bin_width``.
224 Zero translation places atom column ``k`` at molecule column
225 ``k + (len(molecule_grid) - len(batch.grid)) / 2``. Out-of-range
226 destinations fold into the nearest boundary column.
228 Safe concurrently on disjoint atom ranges that share no molecule.
230 Parameters
231 ----------
232 molecule_profiles : np.ndarray
233 Per-molecule profiles to update in place, shape
234 ``(num_molecules, len(molecule_grid))``.
235 molecule_indices : np.ndarray
236 Global molecule index of each atom in this batch.
237 batch : AtomTranslationBatch
238 Atom areas, translations, and profiles to accumulate.
239 molecule_grid : SigmaGrid
240 Grid for the output ``molecule_profiles``.
241 """
242 atom_areas, translations, atom_profiles, atom_grid = (
243 batch.areas,
244 batch.translations,
245 batch.profiles,
246 batch.grid,
247 )
248 bin_width = atom_grid.bin_width
249 num_points = len(atom_grid)
250 molecule_num_points = len(molecule_grid)
251 grid_offset = (molecule_num_points - num_points) / 2
253 fractional_translation = grid_offset + translations / bin_width
254 points_translation = np.floor(fractional_translation).astype(np.int64)
255 weight_right = (fractional_translation - points_translation)[:, None]
257 contributions = atom_areas[:, None].astype(np.float64) * atom_profiles.astype(
258 np.float64
259 )
260 points_at_left = np.arange(num_points)[None, :] + points_translation[:, None]
261 points_at_right = points_at_left + 1
263 np.add.at(
264 molecule_profiles,
265 (
266 molecule_indices[:, None],
267 points_at_left.clip(0, molecule_num_points - 1),
268 ),
269 contributions * (1.0 - weight_right),
270 )
271 np.add.at(
272 molecule_profiles,
273 (
274 molecule_indices[:, None],
275 points_at_right.clip(0, molecule_num_points - 1),
276 ),
277 contributions * weight_right,
278 )