Coverage for cosmolayer/store/averaging.py: 100%
48 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-26 00:09 +0000
1"""COSMO-SAC-style distance-weighted averaging of segment charge density.
3In COSMO-RS and COSMO-SAC, "sigma" is this averaged charge density, not
4the raw per-segment ``charge / area``.
6``segment_offsets`` must describe *exactly* the molecules present in the
7accompanying segment-level arrays: the last molecule's segments run to
8the end of those arrays. Slice ``segment_offsets`` and the arrays
9together, or leftover segments are attributed to the wrong molecule.
10"""
12from collections.abc import Sequence
13from dataclasses import dataclass
15import numpy as np
16from numpy.typing import NDArray
17from tqdm.auto import tqdm
19from cosmolayer.cosmosac.constants import (
20 COSMO_SAC_2002_AVERAGING_RADIUS,
21 COSMO_SAC_2002_F_DECAY,
22 COSMO_SAC_2010_AVERAGING_RADIUS,
23 COSMO_SAC_2010_F_DECAY,
24)
26from .parallel import run_in_threads
29@dataclass(frozen=True)
30class AveragingScheme:
31 """Named averaging scheme (radius and decay) for segment charge
32 densities.
34 Parameters
35 ----------
36 name : str
37 Identifier used as the ``<name>.npy`` stem when a store saves
38 averaged sigmas. Must not collide with reserved store filenames
39 (``data``, ``atom_indices``).
40 averaging_radius : float
41 Effective averaging radius ``r_av``, in Å.
42 f_decay : float
43 Exponential decay factor.
44 """
46 name: str
47 averaging_radius: float
48 f_decay: float
51COSMO_RS = AveragingScheme("cosmo-rs", averaging_radius=0.5, f_decay=1.0)
52COSMO_SAC_2002 = AveragingScheme(
53 "cosmo-sac-2002",
54 averaging_radius=COSMO_SAC_2002_AVERAGING_RADIUS,
55 f_decay=COSMO_SAC_2002_F_DECAY,
56)
57COSMO_SAC_2010 = AveragingScheme(
58 "cosmo-sac-2010",
59 averaging_radius=COSMO_SAC_2010_AVERAGING_RADIUS,
60 f_decay=COSMO_SAC_2010_F_DECAY,
61)
63AVERAGING_SCHEMES: tuple[AveragingScheme, ...] = (
64 COSMO_RS,
65 COSMO_SAC_2002,
66 COSMO_SAC_2010,
67)
68"""Built-in schemes: COSMO-RS, COSMO-SAC 2002, and COSMO-SAC 2010."""
71def average_sigmas(
72 coords: NDArray[np.float64],
73 charges: NDArray[np.float64],
74 areas: NDArray[np.float64],
75 schemes: Sequence[AveragingScheme],
76) -> NDArray[np.float64]:
77 """Distance-weighted average of one molecule's segment charge
78 densities, under one or more averaging schemes.
80 For every segment ``m``, replaces its raw charge density
81 ``sigma_m = q_m / A_m`` with a weighted average over every segment
82 ``n`` in the same molecule, including itself::
84 sigma_avg[m] = sum_n(sigma[n] * w[m, n]) / sum_n(w[m, n])
85 w[m, n] = (r_n^2 * r_av^2 / (r_n^2 + r_av^2))
86 * exp(-f_decay * d_mn^2 / (r_n^2 + r_av^2))
88 where ``r_n = sqrt(A_n / pi)`` is the *neighbor* segment's effective
89 radius and ``d_mn`` is the distance between centroids, so ``w`` is
90 asymmetric even though ``d_mn`` is not. Pass segments of a single
91 molecule only; use ``average_sigmas_by_molecule`` for a dataset.
93 Parameters
94 ----------
95 coords : np.ndarray
96 Segment centroid coordinates for one molecule, shape
97 ``(n_segs, 3)``.
98 charges : np.ndarray
99 Segment charges for the same molecule, shape ``(n_segs,)``.
100 areas : np.ndarray
101 Segment areas for the same molecule, shape ``(n_segs,)``.
102 schemes : Sequence[AveragingScheme]
103 Schemes to apply, in the order of the result rows.
105 Returns
106 -------
107 np.ndarray, shape (len(schemes), n_segs)
108 Averaged charge density for each segment under each scheme, in
109 e/Ų, row ``i`` matching ``schemes[i]``.
110 """
111 # float64: the Gram-matrix distance expansion loses precision in float32.
112 coords = coords.astype(np.float64, copy=False)
113 charges = charges.astype(np.float64, copy=False)
114 areas = areas.astype(np.float64, copy=False)
116 sigmas = charges / areas
117 squared_norms = np.sum(np.square(coords), axis=1)
118 squared_distances = (
119 squared_norms[:, None] + squared_norms[None, :] - 2.0 * (coords @ coords.T)
120 )
121 np.clip(squared_distances, 0.0, None, out=squared_distances)
122 squared_radii = areas / np.pi
124 results = np.empty((len(schemes), len(charges)), dtype=np.float64)
125 for i, scheme in enumerate(schemes):
126 r_av_sq = scheme.averaging_radius**2
127 sums = squared_radii + r_av_sq
128 prods = squared_radii * r_av_sq
129 weights = np.exp(-scheme.f_decay * squared_distances / sums) * prods / sums
130 results[i] = np.sum(weights * sigmas, axis=1) / np.sum(weights, axis=1)
132 return results
135def average_sigmas_by_molecule( # noqa: PLR0913
136 coords: NDArray[np.float64],
137 charges: NDArray[np.float64],
138 areas: NDArray[np.float64],
139 segment_offsets: NDArray[np.int64],
140 schemes: Sequence[AveragingScheme],
141 num_threads: int | None = None,
142 *,
143 progress: bool = False,
144) -> NDArray[np.float64]:
145 """Apply one or more averaging schemes to every molecule in a dataset.
147 Each thread handles a disjoint range of whole molecules. Row ``i`` of
148 the result matches ``schemes[i]`` and can be passed as ``sigmas`` to
149 ``SigmaProfileTable.from_segments``.
151 Parameters
152 ----------
153 coords : np.ndarray
154 Segment centroid coordinates, shape ``(n_segs_total, 3)``.
155 charges : np.ndarray
156 Segment charges, shape ``(n_segs_total,)``.
157 areas : np.ndarray
158 Segment areas, shape ``(n_segs_total,)``.
159 segment_offsets : np.ndarray
160 Start index of each molecule's segments. Must describe exactly
161 the molecules present in the segment-level arrays.
162 schemes : Sequence[AveragingScheme]
163 Schemes to apply, in the order of the result rows.
164 num_threads : int | None, optional
165 Thread count. ``None`` (default) uses every CPU core.
166 progress : bool, optional
167 If True, show a tqdm bar over molecules. Default False.
169 Returns
170 -------
171 np.ndarray, shape (len(schemes), n_segs_total)
172 Averaged charge density for every segment under every scheme, in
173 e/Ų, row ``i`` matching ``schemes[i]``.
174 """
175 num_segs = len(charges)
176 num_mols = len(segment_offsets)
177 averaged_sigmas = np.empty((len(schemes), num_segs), dtype=np.float64)
179 with tqdm(
180 total=num_mols, desc="Averaging sigmas", disable=not progress
181 ) as progress_bar:
183 def process_range(start_mol: int, stop_mol: int) -> None:
184 for mol in range(start_mol, stop_mol):
185 start_seg = segment_offsets[mol]
186 stop_seg = segment_offsets[mol + 1] if mol + 1 < num_mols else num_segs
187 if stop_seg != start_seg:
188 averaged_sigmas[:, start_seg:stop_seg] = average_sigmas(
189 coords[start_seg:stop_seg],
190 charges[start_seg:stop_seg],
191 areas[start_seg:stop_seg],
192 schemes,
193 )
194 progress_bar.update(1)
196 run_in_threads(
197 process_range, num_mols, num_threads=num_threads, limit_blas=True
198 )
199 return averaged_sigmas