torchspin 0.3.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- torchspin/__init__.py +363 -0
- torchspin/_cardamom_istos.py +327 -0
- torchspin/_cardamom_propagatedm.py +413 -0
- torchspin/_cardamom_utils.py +113 -0
- torchspin/_compile.py +100 -0
- torchspin/_linalg.py +140 -0
- torchspin/angmom.py +742 -0
- torchspin/autograd.py +557 -0
- torchspin/autoguess.py +342 -0
- torchspin/batch.py +225 -0
- torchspin/blochsteady.py +274 -0
- torchspin/cardamom.py +552 -0
- torchspin/chili.py +932 -0
- torchspin/chili_sle.py +1233 -0
- torchspin/constants.py +88 -0
- torchspin/convspec.py +148 -0
- torchspin/ctafft.py +91 -0
- torchspin/curry.py +241 -0
- torchspin/data/FourierSeriesCoefficients.txt +89 -0
- torchspin/data/GaussianCascadeCoefficients.txt +56 -0
- torchspin/data/isotopedata.txt +416 -0
- torchspin/data/spacegroups.txt +637 -0
- torchspin/dataproc.py +361 -0
- torchspin/dipbackground.py +69 -0
- torchspin/dipkernel.py +102 -0
- torchspin/diptensor.py +138 -0
- torchspin/endorfrq.py +198 -0
- torchspin/endorfrq_perturb.py +302 -0
- torchspin/eprload.py +1538 -0
- torchspin/eprsave.py +165 -0
- torchspin/esfit.py +2427 -0
- torchspin/evolve.py +274 -0
- torchspin/ewrls.py +124 -0
- torchspin/excitation.py +107 -0
- torchspin/exciteprofile.py +132 -0
- torchspin/experiment.py +323 -0
- torchspin/exponfit.py +177 -0
- torchspin/fastmotion.py +182 -0
- torchspin/fdaxis.py +57 -0
- torchspin/fitgui.py +350 -0
- torchspin/garlic.py +760 -0
- torchspin/ham.py +145 -0
- torchspin/ham_cf.py +115 -0
- torchspin/ham_ee.py +124 -0
- torchspin/ham_ez.py +119 -0
- torchspin/ham_ezho.py +403 -0
- torchspin/ham_hf.py +122 -0
- torchspin/ham_nn.py +108 -0
- torchspin/ham_nq.py +100 -0
- torchspin/ham_nz.py +136 -0
- torchspin/ham_oz.py +115 -0
- torchspin/ham_so.py +118 -0
- torchspin/ham_zf.py +254 -0
- torchspin/hamsymm.py +723 -0
- torchspin/initstate.py +158 -0
- torchspin/isotopologues.py +296 -0
- torchspin/levels.py +189 -0
- torchspin/levelsplot.py +363 -0
- torchspin/lineshape.py +715 -0
- torchspin/lpsvd.py +219 -0
- torchspin/makespec.py +87 -0
- torchspin/mdhmm.py +573 -0
- torchspin/mdload.py +788 -0
- torchspin/mdtraj2oripot.py +90 -0
- torchspin/ml.py +377 -0
- torchspin/mlpsvd.py +297 -0
- torchspin/nucdata.py +310 -0
- torchspin/nucfrq2d.py +228 -0
- torchspin/orca2torchspin.py +1198 -0
- torchspin/ordering.py +78 -0
- torchspin/oripotentialplot.py +246 -0
- torchspin/orisel.py +176 -0
- torchspin/pepper.py +2019 -0
- torchspin/pepper_autograd.py +613 -0
- torchspin/photoselect.py +174 -0
- torchspin/plegendre.py +131 -0
- torchspin/propint.py +127 -0
- torchspin/pulse.py +652 -0
- torchspin/py.typed +0 -0
- torchspin/rapidscan2spc.py +101 -0
- torchspin/resfields.py +337 -0
- torchspin/resfields_batch.py +446 -0
- torchspin/resfields_eig.py +250 -0
- torchspin/resfields_perturb.py +706 -0
- torchspin/resfreqs_matrix.py +344 -0
- torchspin/resfreqs_perturb.py +299 -0
- torchspin/resonator.py +310 -0
- torchspin/resonatorprofile.py +85 -0
- torchspin/rfmixer.py +147 -0
- torchspin/rotations.py +87 -0
- torchspin/rotutils.py +880 -0
- torchspin/saffron.py +1523 -0
- torchspin/saffron_pathways.py +92 -0
- torchspin/saffron_peaks.py +580 -0
- torchspin/saffron_thyme.py +472 -0
- torchspin/salt.py +680 -0
- torchspin/sigeq.py +71 -0
- torchspin/signalprocessing.py +226 -0
- torchspin/sitetransforms.py +254 -0
- torchspin/sphgrid.py +470 -0
- torchspin/spidyan.py +1508 -0
- torchspin/spinladder.py +161 -0
- torchspin/spinops.py +265 -0
- torchspin/spinsystem.py +1427 -0
- torchspin/stackplot.py +167 -0
- torchspin/stev.py +283 -0
- torchspin/stochtraj_diffusion.py +494 -0
- torchspin/stochtraj_jump.py +195 -0
- torchspin/strainwidth.py +727 -0
- torchspin/transmitter.py +89 -0
- torchspin/utils.py +605 -0
- torchspin-0.3.0.dist-info/METADATA +267 -0
- torchspin-0.3.0.dist-info/RECORD +116 -0
- torchspin-0.3.0.dist-info/WHEEL +5 -0
- torchspin-0.3.0.dist-info/licenses/LICENSE.md +22 -0
- torchspin-0.3.0.dist-info/top_level.txt +1 -0
torchspin/pepper.py
ADDED
|
@@ -0,0 +1,2019 @@
|
|
|
1
|
+
"""Top-level CW EPR powder spectrum simulator for torchspin.
|
|
2
|
+
|
|
3
|
+
Port of EasySpin's ``pepper`` (matrix method, field-swept powder average).
|
|
4
|
+
|
|
5
|
+
Pipeline::
|
|
6
|
+
|
|
7
|
+
ham(sys) → H0, mux, muy, muz
|
|
8
|
+
sphgrid(...) → phi, theta, weights
|
|
9
|
+
for each orient:
|
|
10
|
+
resfields(...) → B_res, intensities, widths
|
|
11
|
+
makespec(...) → stick spectrum
|
|
12
|
+
convspec(...) → lineshape broadening (per-line if strain)
|
|
13
|
+
spec += stick * weight
|
|
14
|
+
convspec(...) → lineshape broadening (global if no strain)
|
|
15
|
+
return x, spec
|
|
16
|
+
"""
|
|
17
|
+
from __future__ import annotations
|
|
18
|
+
|
|
19
|
+
import functools
|
|
20
|
+
import math
|
|
21
|
+
from typing import Optional
|
|
22
|
+
|
|
23
|
+
import torch
|
|
24
|
+
|
|
25
|
+
import numpy as np
|
|
26
|
+
|
|
27
|
+
from torchspin.convspec import convspec
|
|
28
|
+
from torchspin.isotopologues import expand_components
|
|
29
|
+
from torchspin.experiment import auto_batch_size, Experiment, Options
|
|
30
|
+
from torchspin.ham import ham
|
|
31
|
+
from torchspin.hamsymm import hamsymm
|
|
32
|
+
from torchspin.rotations import erot
|
|
33
|
+
from torchspin.makespec import makespec
|
|
34
|
+
from torchspin.resfields import resfields
|
|
35
|
+
from torchspin.resfields_batch import resfields_batch
|
|
36
|
+
from torchspin.resfields_perturb import resfields_perturb_batch
|
|
37
|
+
from torchspin.resfreqs_matrix import resfreqs_matrix, resfreqs_batch
|
|
38
|
+
from torchspin.sphgrid import sphgrid, _d2h_triangulation, _triangle_areas, grid_triangulation, gridparam
|
|
39
|
+
from torchspin.spinsystem import SpinSystem
|
|
40
|
+
|
|
41
|
+
__all__ = ['pepper']
|
|
42
|
+
|
|
43
|
+
# Grid symmetries that use a closedPhi, nOct=1 structure (triangulatable by _d2h_triangulation)
|
|
44
|
+
_CLOSEDPHI_1OCT = frozenset({'D2h', 'D4h', 'D6h', 'D3d', 'Th', 'Oh'})
|
|
45
|
+
|
|
46
|
+
|
|
47
|
+
_INTERP_MATRIX_MAX = 8_000_000 # n_coarse × n_fine above which the scipy splines are used
|
|
48
|
+
|
|
49
|
+
|
|
50
|
+
@functools.lru_cache(maxsize=16)
|
|
51
|
+
def _d2h_interp_matrix(N_c: int, N_f: int) -> np.ndarray:
|
|
52
|
+
"""Dense ``(n_fine, n_coarse)`` matrix of EasySpin's G3 interpolation on the
|
|
53
|
+
triangular D2h SOPHE grid (row-wise not-a-knot cubic resampling along φ to
|
|
54
|
+
an N_c × N_c rectangle, then the tensor-product not-a-knot bicubic spline —
|
|
55
|
+
identical to ``RectBivariateSpline(kx=ky=3, s=0)`` — evaluated at the fine
|
|
56
|
+
knots). Both steps are linear in the values, so one matrix product per
|
|
57
|
+
transition slot replaces ~N_c spline constructions.
|
|
58
|
+
"""
|
|
59
|
+
from scipy.interpolate import CubicSpline
|
|
60
|
+
n_c = N_c * (N_c + 1) // 2
|
|
61
|
+
phi_rect = np.linspace(0.0, np.pi / 2, N_c)
|
|
62
|
+
theta_1d = np.linspace(0.0, np.pi / 2, N_c)
|
|
63
|
+
R = np.zeros((N_c * N_c, n_c))
|
|
64
|
+
idx = 0
|
|
65
|
+
for r in range(1, N_c + 1):
|
|
66
|
+
rows = slice((r - 1) * N_c, r * N_c)
|
|
67
|
+
if r == 1:
|
|
68
|
+
R[rows, idx] = 1.0
|
|
69
|
+
elif r == N_c:
|
|
70
|
+
R[rows, idx:idx + r] = np.eye(r)
|
|
71
|
+
else:
|
|
72
|
+
R[rows, idx:idx + r] = CubicSpline(np.linspace(0.0, np.pi / 2, r), np.eye(r), axis=0)(phi_rect)
|
|
73
|
+
idx += r
|
|
74
|
+
rows_f = np.arange(1, N_f + 1)
|
|
75
|
+
theta_all = np.repeat((rows_f - 1) / (N_f - 1) * np.pi / 2 if N_f > 1 else np.zeros(1), rows_f)
|
|
76
|
+
phi_all = np.concatenate([np.array([0.0])] + [np.linspace(0.0, np.pi / 2, r) for r in range(2, N_f + 1)])
|
|
77
|
+
S_th = CubicSpline(theta_1d, np.eye(N_c), axis=0)(theta_all) # (n_f, N_c)
|
|
78
|
+
S_ph = CubicSpline(phi_rect, np.eye(N_c), axis=0)(phi_all) # (n_f, N_c)
|
|
79
|
+
# value_k = Σ_ij S_th[k,i] Z[i,j] S_ph[k,j] with Z = (R v).reshape(N_c, N_c)
|
|
80
|
+
W = (S_th[:, :, None] * S_ph[:, None, :]).reshape(S_th.shape[0], N_c * N_c)
|
|
81
|
+
return W @ R
|
|
82
|
+
|
|
83
|
+
|
|
84
|
+
def _interp_sph(
|
|
85
|
+
coarse_vecs: torch.Tensor,
|
|
86
|
+
fine_vecs: torch.Tensor,
|
|
87
|
+
values: np.ndarray,
|
|
88
|
+
) -> np.ndarray:
|
|
89
|
+
"""Interpolate per-orientation scalar data from a coarse to a fine SOPHE grid.
|
|
90
|
+
|
|
91
|
+
For D2h-type triangular SOPHE grids, uses MATLAB's G3 structured interpolation:
|
|
92
|
+
first a row-by-row cubic spline along phi (converting triangular→rectangular),
|
|
93
|
+
then a 2-D bicubic spline in (theta, phi) space. This matches EasySpin's
|
|
94
|
+
``gridinterp(..., 'G3')`` and eliminates the quantization artefacts that arise
|
|
95
|
+
from unstructured 2-D scattered interpolation.
|
|
96
|
+
|
|
97
|
+
For Dinfh grids (1-D meridional), uses a cubic spline in cos(theta).
|
|
98
|
+
|
|
99
|
+
For full-sphere (C1) grids, falls back to CloughTocher 2-D scattered interpolation.
|
|
100
|
+
|
|
101
|
+
Parameters
|
|
102
|
+
----------
|
|
103
|
+
coarse_vecs : torch.Tensor, shape (3, N_coarse)
|
|
104
|
+
Unit vectors for the coarse grid (rows: x, y, z).
|
|
105
|
+
fine_vecs : torch.Tensor, shape (3, N_fine)
|
|
106
|
+
Unit vectors for the fine grid.
|
|
107
|
+
values : np.ndarray, shape (N_coarse,)
|
|
108
|
+
Scalar quantity to interpolate (resonance field, intensity, or width).
|
|
109
|
+
|
|
110
|
+
Returns
|
|
111
|
+
-------
|
|
112
|
+
np.ndarray, shape (N_fine,)
|
|
113
|
+
"""
|
|
114
|
+
x_c = coarse_vecs[0].numpy()
|
|
115
|
+
y_c = coarse_vecs[1].numpy()
|
|
116
|
+
z_c = coarse_vecs[2].numpy()
|
|
117
|
+
z_f = fine_vecs[2].numpy()
|
|
118
|
+
|
|
119
|
+
# ---- 1-D fallback for meridional grids (Dinfh: all phi = 0, y ≈ 0) ----
|
|
120
|
+
if np.std(y_c) < 1e-8:
|
|
121
|
+
from scipy.interpolate import CubicSpline
|
|
122
|
+
sort_idx = np.argsort(z_c)
|
|
123
|
+
zs = z_c[sort_idx]
|
|
124
|
+
vs = values[sort_idx]
|
|
125
|
+
if len(zs) >= 4:
|
|
126
|
+
cs = CubicSpline(zs, vs, extrapolate=False)
|
|
127
|
+
res = cs(z_f)
|
|
128
|
+
nan_mask = np.isnan(res)
|
|
129
|
+
if nan_mask.any():
|
|
130
|
+
res[nan_mask] = np.interp(z_f[nan_mask], zs, vs,
|
|
131
|
+
left=vs[0], right=vs[-1])
|
|
132
|
+
return res
|
|
133
|
+
return np.interp(z_f, zs, vs, left=vs[0], right=vs[-1])
|
|
134
|
+
|
|
135
|
+
n_c = len(values)
|
|
136
|
+
|
|
137
|
+
# ---- Detect D2h triangular SOPHE grid structure ----
|
|
138
|
+
# A triangular SOPHE grid with N knots along the meridian has N*(N+1)//2 points.
|
|
139
|
+
disc_c = 1 + 8 * n_c
|
|
140
|
+
sqrt_c = math.isqrt(disc_c)
|
|
141
|
+
is_d2h_structured = (sqrt_c * sqrt_c == disc_c) and ((sqrt_c - 1) % 2 == 0)
|
|
142
|
+
|
|
143
|
+
n_f = fine_vecs.shape[1]
|
|
144
|
+
disc_f = 1 + 8 * n_f
|
|
145
|
+
sqrt_f = math.isqrt(disc_f)
|
|
146
|
+
is_d2h_fine = (sqrt_f * sqrt_f == disc_f) and ((sqrt_f - 1) % 2 == 0)
|
|
147
|
+
|
|
148
|
+
if is_d2h_structured and is_d2h_fine:
|
|
149
|
+
from scipy.interpolate import CubicSpline, RectBivariateSpline
|
|
150
|
+
N_c = (sqrt_c - 1) // 2
|
|
151
|
+
N_f = (sqrt_f - 1) // 2
|
|
152
|
+
if n_c * n_f <= _INTERP_MATRIX_MAX:
|
|
153
|
+
# the G3 interpolation is linear in the values: cached matrix
|
|
154
|
+
return _d2h_interp_matrix(N_c, N_f) @ np.asarray(values, dtype=float)
|
|
155
|
+
|
|
156
|
+
# Step 1: Convert triangular coarse grid → rectangular array z_rect[N_c × N_c]
|
|
157
|
+
# by interpolating each theta row along phi via cubic spline.
|
|
158
|
+
theta_1d = np.linspace(0.0, np.pi / 2, N_c)
|
|
159
|
+
phi_rect = np.linspace(0.0, np.pi / 2, N_c) # target phi grid (N_c columns)
|
|
160
|
+
|
|
161
|
+
z_rect = np.zeros((N_c, N_c))
|
|
162
|
+
idx = 0
|
|
163
|
+
for r in range(1, N_c + 1):
|
|
164
|
+
row_vals = values[idx: idx + r]
|
|
165
|
+
idx += r
|
|
166
|
+
if r == 1:
|
|
167
|
+
z_rect[0, :] = row_vals[0] # north pole: constant in phi
|
|
168
|
+
elif r == N_c:
|
|
169
|
+
z_rect[N_c - 1, :] = row_vals # equator: already has N_c points
|
|
170
|
+
else:
|
|
171
|
+
phi_row = np.linspace(0.0, np.pi / 2, r)
|
|
172
|
+
cs = CubicSpline(phi_row, row_vals)
|
|
173
|
+
z_rect[r - 1, :] = cs(phi_rect)
|
|
174
|
+
|
|
175
|
+
# Step 2: 2-D bicubic spline on the rectangular coarse grid.
|
|
176
|
+
spline2d = RectBivariateSpline(theta_1d, phi_rect, z_rect, kx=3, ky=3, s=0)
|
|
177
|
+
|
|
178
|
+
# Step 3: Evaluate at the fine triangular grid positions (one call).
|
|
179
|
+
rows = np.arange(1, N_f + 1)
|
|
180
|
+
theta_all = np.repeat((rows - 1) / (N_f - 1) * np.pi / 2 if N_f > 1 else np.zeros(1), rows)
|
|
181
|
+
phi_all = np.concatenate([np.array([0.0])] + [np.linspace(0.0, np.pi / 2, r) for r in range(2, N_f + 1)])
|
|
182
|
+
return spline2d.ev(theta_all, phi_all)
|
|
183
|
+
|
|
184
|
+
# ---- Fallback: 2-D scattered interpolation (Ci / C1 / other) ----
|
|
185
|
+
from scipy.interpolate import CloughTocher2DInterpolator, NearestNDInterpolator
|
|
186
|
+
|
|
187
|
+
x_f = fine_vecs[0].numpy()
|
|
188
|
+
y_f = fine_vecs[1].numpy()
|
|
189
|
+
|
|
190
|
+
def _cubic2d(xc, yc, vals, xf, yf):
|
|
191
|
+
pts_c = np.column_stack([xc, yc])
|
|
192
|
+
pts_f = np.column_stack([xf, yf])
|
|
193
|
+
res = CloughTocher2DInterpolator(pts_c, vals)(pts_f)
|
|
194
|
+
nan_mask = np.isnan(res)
|
|
195
|
+
if nan_mask.any():
|
|
196
|
+
res[nan_mask] = NearestNDInterpolator(pts_c, vals)(pts_f[nan_mask])
|
|
197
|
+
return res
|
|
198
|
+
|
|
199
|
+
is_full_sphere = z_c.min() < -0.01
|
|
200
|
+
if not is_full_sphere:
|
|
201
|
+
return _cubic2d(x_c, y_c, values, x_f, y_f)
|
|
202
|
+
|
|
203
|
+
# Full sphere (C1): interpolate upper and lower hemispheres separately
|
|
204
|
+
result = np.empty(n_f)
|
|
205
|
+
y_c = coarse_vecs[1].numpy()
|
|
206
|
+
for c_mask, f_mask in [(z_c >= 0, z_f >= 0), (z_c <= 0, z_f < 0)]:
|
|
207
|
+
if c_mask.sum() < 3 or f_mask.sum() == 0:
|
|
208
|
+
continue
|
|
209
|
+
result[f_mask] = _cubic2d(
|
|
210
|
+
x_c[c_mask], y_c[c_mask], values[c_mask],
|
|
211
|
+
x_f[f_mask], y_f[f_mask],
|
|
212
|
+
)
|
|
213
|
+
return result
|
|
214
|
+
|
|
215
|
+
|
|
216
|
+
def _transition_slots(all_B, all_I, all_W, all_pairs, n_orient):
|
|
217
|
+
"""Group per-orientation resonances into transition slots across the grid.
|
|
218
|
+
|
|
219
|
+
Mirrors EasySpin's ``Pdat``/``Idat``/``Wdat`` (nTransitions × nOrientations)
|
|
220
|
+
bookkeeping: a slot is a level pair ``(u, v)`` plus its occurrence index
|
|
221
|
+
(looping transitions can resonate more than once), and every slot holds one
|
|
222
|
+
value per orientation with NaN where that transition has no resonance.
|
|
223
|
+
Interpolation and projection then act on physically the same transition at
|
|
224
|
+
every grid point. Without level pairs (perturbation path, whose slot
|
|
225
|
+
structure is fixed by construction) the arrays are matched by index.
|
|
226
|
+
|
|
227
|
+
Returns (B_arrs, I_arrs, W_arrs): lists of (n_orient,) float arrays; W_arrs
|
|
228
|
+
is None when no widths were supplied.
|
|
229
|
+
"""
|
|
230
|
+
have_widths = any(w is not None and w.numel() > 0 for w in all_W) if all_W is not None else False
|
|
231
|
+
if all_pairs is None or all(p is None for p in all_pairs):
|
|
232
|
+
n_slots = max((b.numel() if b is not None else 0) for b in all_B) if all_B else 0
|
|
233
|
+
B_arrs = [np.full(n_orient, np.nan) for _ in range(n_slots)]
|
|
234
|
+
I_arrs = [np.full(n_orient, np.nan) for _ in range(n_slots)]
|
|
235
|
+
W_arrs = [np.full(n_orient, np.nan) for _ in range(n_slots)] if have_widths else None
|
|
236
|
+
for k in range(n_orient):
|
|
237
|
+
B = all_B[k]
|
|
238
|
+
if B is None:
|
|
239
|
+
continue
|
|
240
|
+
Bn = B.detach().cpu().numpy(); In = all_I[k].detach().cpu().numpy()
|
|
241
|
+
Wn = all_W[k].detach().cpu().numpy() if (have_widths and all_W[k] is not None) else None
|
|
242
|
+
for i in range(Bn.shape[0]):
|
|
243
|
+
B_arrs[i][k] = Bn[i]; I_arrs[i][k] = In[i]
|
|
244
|
+
if Wn is not None and i < Wn.shape[0]:
|
|
245
|
+
W_arrs[i][k] = Wn[i]
|
|
246
|
+
_transition_slots.last_keys = None # index-matched slots (perturbation path)
|
|
247
|
+
return B_arrs, I_arrs, W_arrs
|
|
248
|
+
slots: dict = {}
|
|
249
|
+
B_arrs: list = []; I_arrs: list = []; W_arrs: list = []
|
|
250
|
+
for k in range(n_orient):
|
|
251
|
+
B = all_B[k]; pairs = all_pairs[k]
|
|
252
|
+
if B is None or pairs is None or B.numel() == 0:
|
|
253
|
+
continue
|
|
254
|
+
Bn = B.detach().cpu().numpy(); In = all_I[k].detach().cpu().numpy()
|
|
255
|
+
Wn = all_W[k].detach().cpu().numpy() if (have_widths and all_W[k] is not None) else None
|
|
256
|
+
pn = pairs.detach().cpu().numpy()
|
|
257
|
+
order = sorted(range(Bn.shape[0]), key=lambda i: (int(pn[i, 0]), int(pn[i, 1]), float(Bn[i])))
|
|
258
|
+
occ: dict = {}
|
|
259
|
+
for i in order:
|
|
260
|
+
uv = (int(pn[i, 0]), int(pn[i, 1]))
|
|
261
|
+
c = occ.get(uv, 0); occ[uv] = c + 1
|
|
262
|
+
key = (uv[0], uv[1], c)
|
|
263
|
+
if key not in slots:
|
|
264
|
+
slots[key] = len(B_arrs)
|
|
265
|
+
B_arrs.append(np.full(n_orient, np.nan))
|
|
266
|
+
I_arrs.append(np.full(n_orient, np.nan))
|
|
267
|
+
W_arrs.append(np.full(n_orient, np.nan))
|
|
268
|
+
j = slots[key]
|
|
269
|
+
B_arrs[j][k] = Bn[i]; I_arrs[j][k] = In[i]
|
|
270
|
+
if Wn is not None and i < Wn.shape[0]:
|
|
271
|
+
W_arrs[j][k] = Wn[i]
|
|
272
|
+
_transition_slots.last_keys = sorted(slots.keys(), key=lambda kk: slots[kk]) # (u, v, occurrence) per slot
|
|
273
|
+
return B_arrs, I_arrs, (W_arrs if have_widths else None)
|
|
274
|
+
|
|
275
|
+
|
|
276
|
+
def _filter_slots(B_arrs, I_arrs, W_arrs, threshold):
|
|
277
|
+
"""EasySpin transition pre-selection: keep a transition (slot) when its
|
|
278
|
+
maximum |intensity| over the grid reaches ``threshold`` × the global maximum.
|
|
279
|
+
Resonances are never dropped per orientation, so slots stay contiguous for
|
|
280
|
+
interpolation."""
|
|
281
|
+
if threshold <= 0 or not I_arrs:
|
|
282
|
+
return B_arrs, I_arrs, W_arrs
|
|
283
|
+
slot_max = np.array([np.nanmax(np.abs(I)) if np.isfinite(I).any() else 0.0 for I in I_arrs])
|
|
284
|
+
gmax = slot_max.max() if slot_max.size else 0.0
|
|
285
|
+
keep = [k for k in range(len(I_arrs)) if slot_max[k] >= threshold * gmax]
|
|
286
|
+
keys = getattr(_transition_slots, 'last_keys', None)
|
|
287
|
+
if keys is not None:
|
|
288
|
+
_transition_slots.last_keys = [keys[k] for k in keep]
|
|
289
|
+
return ([B_arrs[k] for k in keep], [I_arrs[k] for k in keep],
|
|
290
|
+
None if W_arrs is None else [W_arrs[k] for k in keep])
|
|
291
|
+
|
|
292
|
+
|
|
293
|
+
def _gaussian_bins(x: torch.Tensor, pos: np.ndarray, fwhm: np.ndarray, amp: np.ndarray,
|
|
294
|
+
min_fwhm: float, chunk: int = 256) -> torch.Tensor:
|
|
295
|
+
"""Sum of bin-integrated Gaussians (EasySpin lisum1i with a Gaussian template).
|
|
296
|
+
|
|
297
|
+
Each line contributes ``amp`` × [Φ(upper bin edge) − Φ(lower bin edge)], so
|
|
298
|
+
the sum over points equals ``amp`` (a spectral density after the final
|
|
299
|
+
1/ΔB). Widths below ``min_fwhm`` (ΔB/100, lisum1i's clip) are clipped.
|
|
300
|
+
|
|
301
|
+
Every line is evaluated only inside a window of ±7σ√2 around its centre
|
|
302
|
+
(erfc(7) ≈ 4e-23, i.e. the omitted bins are zero to double precision);
|
|
303
|
+
lines are sorted by width and processed in chunks of 256 (cache-sized:
|
|
304
|
+
larger chunks are 3-4× slower per line) with a tight common window; on the
|
|
305
|
+
CPU the chunks run concurrently in a thread pool
|
|
306
|
+
(``_linalg.pool_map``), on CUDA sequentially on the device.
|
|
307
|
+
"""
|
|
308
|
+
n = x.shape[0]
|
|
309
|
+
dev = x.device
|
|
310
|
+
if pos.size == 0:
|
|
311
|
+
return torch.zeros(n, dtype=torch.float64, device=dev)
|
|
312
|
+
dx = float(x[1] - x[0])
|
|
313
|
+
x0 = float(x[0])
|
|
314
|
+
sig = np.maximum(np.asarray(fwhm, dtype=float), min_fwhm) / math.sqrt(8.0 * math.log(2.0))
|
|
315
|
+
order = np.argsort(sig, kind='stable')
|
|
316
|
+
pos_t = torch.from_numpy(np.asarray(pos, dtype=float)[order]).to(dev)
|
|
317
|
+
sg_all = torch.from_numpy(sig[order] * math.sqrt(2.0)).to(dev)
|
|
318
|
+
amp_t = torch.from_numpy(np.asarray(amp, dtype=float)[order]).to(dev)
|
|
319
|
+
half_all = 7.0 * sg_all # window half-width (mT)
|
|
320
|
+
nw_np = ((2.0 * half_all / dx).ceil().to(torch.long) + 3).cpu().numpy() # bins per window
|
|
321
|
+
ar = torch.arange(n, dtype=torch.long, device=dev)
|
|
322
|
+
xr = x.to(dev)
|
|
323
|
+
|
|
324
|
+
def _chunk(k0: int) -> torch.Tensor:
|
|
325
|
+
sl = slice(k0, k0 + chunk)
|
|
326
|
+
p = pos_t[sl, None]
|
|
327
|
+
sg = sg_all[sl, None]
|
|
328
|
+
a = amp_t[sl, None]
|
|
329
|
+
nw = int(nw_np[sl].max())
|
|
330
|
+
if nw >= n:
|
|
331
|
+
hi = torch.erf((xr[None, :] + 0.5 * dx - p) / sg)
|
|
332
|
+
lo = torch.erf((xr[None, :] - 0.5 * dx - p) / sg)
|
|
333
|
+
return (0.5 * a * (hi - lo)).sum(dim=0)
|
|
334
|
+
start = ((p - half_all[sl, None] - x0) / dx).floor().to(torch.long) - 1
|
|
335
|
+
idx = start + ar[None, :nw]
|
|
336
|
+
idxc = idx.clamp(0, n - 1)
|
|
337
|
+
z = (xr[idxc] - p) / sg # bin centres
|
|
338
|
+
d = (0.5 * dx) / sg # half bin in σ√2 units
|
|
339
|
+
vals = (0.5 * a) * (torch.erf(z + d) - torch.erf(z - d))
|
|
340
|
+
vals = torch.where((idx >= 0) & (idx < n), vals, torch.zeros((), dtype=torch.float64, device=dev))
|
|
341
|
+
return torch.bincount(idxc.reshape(-1), weights=vals.reshape(-1), minlength=n)
|
|
342
|
+
|
|
343
|
+
starts = list(range(0, pos_t.shape[0], chunk))
|
|
344
|
+
if dev.type == 'cpu':
|
|
345
|
+
from torchspin._linalg import pool_map
|
|
346
|
+
parts = pool_map(_chunk, starts)
|
|
347
|
+
else:
|
|
348
|
+
parts = [_chunk(k0) for k0 in starts]
|
|
349
|
+
return torch.stack(parts).sum(dim=0) if len(parts) > 1 else parts[0]
|
|
350
|
+
|
|
351
|
+
|
|
352
|
+
def _interp_grid(values: np.ndarray, symmetry: str, N_c: int, N_f: int, linear: bool = False) -> np.ndarray:
|
|
353
|
+
"""EasySpin ``gridinterp`` for open-φ grids (Ci, C2h, C1, C4h, C6h, S6).
|
|
354
|
+
|
|
355
|
+
The coarse per-knot values are rectified into a (θ-row × φ-column) array
|
|
356
|
+
(each θ slice resampled along φ with a cubic spline, periodic in φ where
|
|
357
|
+
the grid wraps), then interpolated with a global bicubic spline at the
|
|
358
|
+
fine grid's (θ, φ) — EasySpin mode ``G3``. When values are missing (NaN,
|
|
359
|
+
e.g. a transition without a resonance at some knots) EasySpin switches to
|
|
360
|
+
linear interpolation (``L1``) so the NaNs propagate to the affected fine
|
|
361
|
+
knots only; the same is done here.
|
|
362
|
+
"""
|
|
363
|
+
from scipy.interpolate import CubicSpline, RectBivariateSpline
|
|
364
|
+
maxPhi, closedPhi, nOct = gridparam(symmetry)
|
|
365
|
+
periodic = not closedPhi
|
|
366
|
+
full = nOct == 8
|
|
367
|
+
y = np.asarray(values, dtype=float)
|
|
368
|
+
cubic = (not linear) and not np.isnan(y).any()
|
|
369
|
+
dlen = 4 if full else nOct
|
|
370
|
+
nr = 2 * N_c - 1 if full else N_c
|
|
371
|
+
nc = dlen * (N_c - 1) + 1
|
|
372
|
+
z = np.full((nr, nc), np.nan)
|
|
373
|
+
|
|
374
|
+
def _row(vals, is_periodic):
|
|
375
|
+
if is_periodic:
|
|
376
|
+
vals = np.concatenate([vals, vals[:1]])
|
|
377
|
+
n = vals.size
|
|
378
|
+
xx = np.linspace(1.0, float(n), nc)
|
|
379
|
+
if cubic:
|
|
380
|
+
if is_periodic:
|
|
381
|
+
return CubicSpline(np.arange(1, n + 1), vals, bc_type='periodic')(xx)
|
|
382
|
+
return CubicSpline(np.arange(1, n + 1), vals)(xx)
|
|
383
|
+
k = np.minimum(np.maximum(np.floor(xx - 1).astype(int), 0), n - 2)
|
|
384
|
+
return vals[k] + (xx - (k + 1)) * (vals[k + 1] - vals[k])
|
|
385
|
+
|
|
386
|
+
z[0, :] = y[0]
|
|
387
|
+
idx = 1
|
|
388
|
+
ln = dlen + 1 - int(periodic)
|
|
389
|
+
for ir in range(1, N_c - 1):
|
|
390
|
+
z[ir, :] = _row(y[idx:idx + ln], periodic)
|
|
391
|
+
idx += ln
|
|
392
|
+
ln += dlen
|
|
393
|
+
z[N_c - 1, :] = np.concatenate([y[idx:idx + ln], y[idx:idx + 1]]) if periodic else y[idx:idx + ln]
|
|
394
|
+
if full:
|
|
395
|
+
idx += ln
|
|
396
|
+
ln -= dlen
|
|
397
|
+
for ir in range(N_c, 2 * N_c - 2):
|
|
398
|
+
z[ir, :] = _row(y[idx:idx + ln], periodic)
|
|
399
|
+
idx += ln
|
|
400
|
+
ln += -dlen
|
|
401
|
+
z[2 * N_c - 2, :] = y[idx]
|
|
402
|
+
|
|
403
|
+
phi_f, theta_f, _, _ = sphgrid(symmetry, N_f)
|
|
404
|
+
phi_f = phi_f.numpy(); theta_f = theta_f.numpy()
|
|
405
|
+
if full:
|
|
406
|
+
iphi = 4 * (N_c - 1) * (phi_f / maxPhi)
|
|
407
|
+
ithe = 2 * (N_c - 1) * (theta_f / math.pi)
|
|
408
|
+
else:
|
|
409
|
+
iphi = nOct * (N_c - 1) * (phi_f / maxPhi)
|
|
410
|
+
ithe = (N_c - 1) * (theta_f / (math.pi / 2))
|
|
411
|
+
iphi = np.clip(iphi, 0, nc - 1); ithe = np.clip(ithe, 0, nr - 1)
|
|
412
|
+
if cubic:
|
|
413
|
+
spl = RectBivariateSpline(np.arange(nr), np.arange(nc), z, kx=3, ky=3, s=0)
|
|
414
|
+
return spl.ev(ithe, iphi)
|
|
415
|
+
# bilinear with NaN propagation (EasySpin interp2 'linear')
|
|
416
|
+
r0 = np.minimum(np.floor(ithe).astype(int), nr - 2); c0 = np.minimum(np.floor(iphi).astype(int), nc - 2)
|
|
417
|
+
fr = ithe - r0; fc = iphi - c0
|
|
418
|
+
return ((1 - fr) * (1 - fc) * z[r0, c0] + (1 - fr) * fc * z[r0, c0 + 1]
|
|
419
|
+
+ fr * (1 - fc) * z[r0 + 1, c0] + fr * fc * z[r0 + 1, c0 + 1])
|
|
420
|
+
|
|
421
|
+
|
|
422
|
+
def _interp_l3(y: np.ndarray, N_f: int) -> np.ndarray:
|
|
423
|
+
"""EasySpin gridinterp 'L3': local cubic Hermite interpolation in index space
|
|
424
|
+
with Fritsch-Carlson monotone tangents (1-D, axial grids). NaNs propagate."""
|
|
425
|
+
y = np.asarray(y, dtype=float)
|
|
426
|
+
n = y.size
|
|
427
|
+
factor = int(round((N_f - 1) / (n - 1)))
|
|
428
|
+
x = np.arange(factor) / factor
|
|
429
|
+
X = np.stack([x ** 3, x ** 2, x, np.ones_like(x)], axis=1) # (factor, 4)
|
|
430
|
+
H = np.array([[2, -2, 1, 1], [-3, 3, -2, -1], [0, 0, 1, 0], [1, 0, 0, 0]], dtype=float)
|
|
431
|
+
d = np.diff(y)
|
|
432
|
+
T = np.zeros(n)
|
|
433
|
+
with np.errstate(invalid='ignore'):
|
|
434
|
+
k = np.where(np.sign(d[:n - 2]) * np.sign(d[1:n - 1]) > 0)[0]
|
|
435
|
+
dmax = np.maximum(np.abs(d[k]), np.abs(d[k + 1]))
|
|
436
|
+
dmin = np.minimum(np.abs(d[k]), np.abs(d[k + 1]))
|
|
437
|
+
T[k + 1] = 2 * dmin * dmax / (d[k] + d[k + 1])
|
|
438
|
+
C = np.stack([y[:-1], y[1:], T[:-1], T[1:]], axis=0) # (4, n-1)
|
|
439
|
+
yii = X @ H @ C # (factor, n-1)
|
|
440
|
+
return np.concatenate([yii.reshape(-1, order='F'), y[-1:]])
|
|
441
|
+
|
|
442
|
+
|
|
443
|
+
def _interp_slot(vals: np.ndarray, symmetry: str, N_c: int, N_f: int,
|
|
444
|
+
vecs_c: torch.Tensor, vecs_f: torch.Tensor, mode: str = 'pos', any_nan: bool = False):
|
|
445
|
+
"""Interpolate one transition slot from the coarse to the fine grid.
|
|
446
|
+
|
|
447
|
+
Closed-φ one-octant grids and Dinfh use :func:`_interp_sph` (structured
|
|
448
|
+
G3 / meridional spline); open-φ grids use :func:`_interp_grid`. Returns
|
|
449
|
+
``None`` when fewer than three knots carry a value.
|
|
450
|
+
|
|
451
|
+
EasySpin pepper.m interpolation modes: positions ``G3`` (global cubic;
|
|
452
|
+
``L3``/``L1`` when knots are missing), intensities and widths ``L3``
|
|
453
|
+
(local cubic, axial grids) or ``L1`` (linear, all other grids). Pass
|
|
454
|
+
``mode='val'`` for intensities/widths.
|
|
455
|
+
"""
|
|
456
|
+
vals = np.asarray(vals, dtype=float)
|
|
457
|
+
valid = ~np.isnan(vals)
|
|
458
|
+
if valid.sum() < 3:
|
|
459
|
+
return None
|
|
460
|
+
if mode == 'val':
|
|
461
|
+
if symmetry == 'Dinfh':
|
|
462
|
+
return _interp_l3(vals, N_f)
|
|
463
|
+
return _interp_grid(vals, symmetry, N_c, N_f, linear=True)
|
|
464
|
+
# EasySpin: NaN_in_Pdat is a global flag — if any transition lacks a
|
|
465
|
+
# resonance somewhere, all positions are interpolated with L3 (axial) / L1
|
|
466
|
+
if not valid.all() or any_nan:
|
|
467
|
+
# EasySpin: with NaNs in Pdat the interpolation switches to linear
|
|
468
|
+
# ('L1'/'L3' → linear here) so that the NaNs propagate and the facets
|
|
469
|
+
# touching knots without a resonance are dropped, instead of being
|
|
470
|
+
# filled from neighbouring knots.
|
|
471
|
+
if symmetry == 'Dinfh':
|
|
472
|
+
return _interp_l3(vals, N_f) # EasySpin: 'L3' for axial grids with NaNs
|
|
473
|
+
return _interp_grid(vals, symmetry, N_c, N_f)
|
|
474
|
+
if symmetry in _CLOSEDPHI_1OCT or symmetry == 'Dinfh':
|
|
475
|
+
return _interp_sph(vecs_c, vecs_f, vals)
|
|
476
|
+
return _interp_grid(vals, symmetry, N_c, N_f)
|
|
477
|
+
|
|
478
|
+
|
|
479
|
+
def _facets(symmetry: str, N: int, theta_sym: torch.Tensor, ordering=None):
|
|
480
|
+
"""(tri_idx or None, facet weights summing to 4π) for projection/summation.
|
|
481
|
+
|
|
482
|
+
``ordering=(f, R_L2S)`` weights every facet by the orientational
|
|
483
|
+
distribution evaluated at the facet centre (EasySpin pepper.m
|
|
484
|
+
``orderingWeights`` for partially ordered samples).
|
|
485
|
+
"""
|
|
486
|
+
if symmetry == 'Dinfh':
|
|
487
|
+
if ordering is not None:
|
|
488
|
+
raise ValueError('Cannot use axial grid for partially ordered samples.')
|
|
489
|
+
th = theta_sym.cpu().numpy()
|
|
490
|
+
return None, (np.cos(th[:-1]) - np.cos(th[1:])) * 4.0 * math.pi
|
|
491
|
+
tri, areas = grid_triangulation(symmetry, N)
|
|
492
|
+
if ordering is not None:
|
|
493
|
+
from torchspin.ordering import orifun_M2L
|
|
494
|
+
fun, R_L2S = ordering
|
|
495
|
+
phi_g, theta_g, _, _ = sphgrid(symmetry, N)
|
|
496
|
+
c_phi = phi_g.detach().cpu().numpy()[tri].mean(axis=1)
|
|
497
|
+
c_theta = theta_g.detach().cpu().numpy()[tri].mean(axis=1)
|
|
498
|
+
ow = orifun_M2L(fun, R_L2S, c_phi, c_theta)
|
|
499
|
+
if np.any(ow < 0):
|
|
500
|
+
raise ValueError('User-supplied orientation distribution gives negative values.')
|
|
501
|
+
if np.all(ow == 0):
|
|
502
|
+
raise ValueError('User-supplied orientation distribution is all-zero.')
|
|
503
|
+
areas = areas * ow / ow.sum()
|
|
504
|
+
return tri, areas * (4.0 * math.pi / areas.sum())
|
|
505
|
+
|
|
506
|
+
|
|
507
|
+
def _central_difference(spec: torch.Tensor, dx: float) -> torch.Tensor:
|
|
508
|
+
"""EasySpin DerivHarmonic: mean of forward and backward differences."""
|
|
509
|
+
spec_der = torch.zeros_like(spec)
|
|
510
|
+
spec_der[1:-1] = (spec[2:] - spec[:-2]) / (2.0 * dx)
|
|
511
|
+
spec_der[0] = (spec[1] - spec[0]) / dx
|
|
512
|
+
spec_der[-1] = (spec[-1] - spec[-2]) / dx
|
|
513
|
+
return spec_der
|
|
514
|
+
|
|
515
|
+
|
|
516
|
+
def _lorentzian_bins(x: torch.Tensor, pos: np.ndarray, fwhm: np.ndarray, amp: np.ndarray,
|
|
517
|
+
min_fwhm: float, chunk: int = 256, phase: float = 0.0) -> torch.Tensor:
|
|
518
|
+
"""Sum of bin-integrated Lorentzians (EasySpin lisum1i with a Lorentzian template).
|
|
519
|
+
|
|
520
|
+
``phase`` mixes in the dispersion line shape (EasySpin ``Exp.mwPhase``; the
|
|
521
|
+
isotropic template ``lorentzian(xT,x0T,wT,-1,Exp.mwPhase)``)."""
|
|
522
|
+
n = x.shape[0]
|
|
523
|
+
out = torch.zeros(n, dtype=torch.float64)
|
|
524
|
+
if pos.size == 0:
|
|
525
|
+
return out
|
|
526
|
+
dx = float(x[1] - x[0])
|
|
527
|
+
hw = np.maximum(np.asarray(fwhm, dtype=float), min_fwhm) / 2.0
|
|
528
|
+
if phase:
|
|
529
|
+
from torchspin.lineshape import lorentzian as _lor
|
|
530
|
+
xn = x.detach().cpu().numpy()
|
|
531
|
+
for k in range(pos.size):
|
|
532
|
+
F_hi, _ = _lor(xn + 0.5 * dx, float(pos[k]), 2.0 * float(hw[k]), -1, float(phase))
|
|
533
|
+
F_lo, _ = _lor(xn - 0.5 * dx, float(pos[k]), 2.0 * float(hw[k]), -1, float(phase))
|
|
534
|
+
out += float(amp[k]) * torch.from_numpy(np.asarray(F_hi - F_lo, dtype=float))
|
|
535
|
+
return out
|
|
536
|
+
pos_t = torch.from_numpy(np.asarray(pos, dtype=float)); hw_t = torch.from_numpy(hw)
|
|
537
|
+
amp_t = torch.from_numpy(np.asarray(amp, dtype=float))
|
|
538
|
+
for k0 in range(0, pos_t.shape[0], chunk):
|
|
539
|
+
p = pos_t[k0:k0 + chunk, None]; g = hw_t[k0:k0 + chunk, None]; a = amp_t[k0:k0 + chunk, None]
|
|
540
|
+
hi = torch.atan((x[None, :] + 0.5 * dx - p) / g)
|
|
541
|
+
lo = torch.atan((x[None, :] - 0.5 * dx - p) / g)
|
|
542
|
+
out += (a * (hi - lo) / math.pi).sum(dim=0)
|
|
543
|
+
return out
|
|
544
|
+
|
|
545
|
+
|
|
546
|
+
def _finish_field_sweep(spec: torch.Tensor, x: torch.Tensor, dx: float, exp,
|
|
547
|
+
fwhm_g: float, fwhm_l: float, fd_deriv: bool,
|
|
548
|
+
phase: float = 0.0) -> torch.Tensor:
|
|
549
|
+
"""EasySpin's final broadening/harmonic bookkeeping for field sweeps.
|
|
550
|
+
|
|
551
|
+
With field modulation (``Exp.ModAmp``) the absorption spectrum is convolved
|
|
552
|
+
(ConvHarmonic 0) and the harmonic comes from the pseudo-modulation
|
|
553
|
+
(``fieldmod``); otherwise the harmonic is produced by the convolution
|
|
554
|
+
(ConvHarmonic) or, without residual line width, by differentiation
|
|
555
|
+
(DerivHarmonic: finite differences for template/strain spectra, spectral
|
|
556
|
+
derivative for stick spectra). ``phase`` (``Exp.mwPhase``) enters the
|
|
557
|
+
convolution kernel. Multi-row spectra are processed row-wise.
|
|
558
|
+
"""
|
|
559
|
+
mod_amp = float(getattr(exp, 'ModAmp', 0.0) or 0.0)
|
|
560
|
+
if mod_amp > 0 and exp.Harmonic < 1:
|
|
561
|
+
raise ValueError("With field modulation (Exp.ModAmp), Exp.Harmonic=0 does not work.")
|
|
562
|
+
conv_harm = 0 if mod_amp > 0 else int(exp.Harmonic)
|
|
563
|
+
rows = [spec] if spec.ndim == 1 else list(spec)
|
|
564
|
+
out = []
|
|
565
|
+
for row in rows:
|
|
566
|
+
if fwhm_g > 0 or fwhm_l > 0:
|
|
567
|
+
row = convspec(row, dx, fwhm_g=fwhm_g, fwhm_l=fwhm_l, deriv=conv_harm, phase=phase)
|
|
568
|
+
elif conv_harm > 0:
|
|
569
|
+
if fd_deriv:
|
|
570
|
+
for _h in range(conv_harm):
|
|
571
|
+
row = _central_difference(row, dx)
|
|
572
|
+
else:
|
|
573
|
+
row = convspec(row, dx, fwhm_g=0.0, fwhm_l=0.0, deriv=conv_harm)
|
|
574
|
+
if mod_amp > 0:
|
|
575
|
+
from torchspin.dataproc import fieldmod as _fieldmod
|
|
576
|
+
row = torch.tensor(_fieldmod(x.detach().cpu().numpy(), row.detach().cpu().numpy(),
|
|
577
|
+
mod_amp, int(exp.Harmonic)), dtype=torch.float64)
|
|
578
|
+
out.append(row)
|
|
579
|
+
return out[0] if spec.ndim == 1 else torch.stack(out)
|
|
580
|
+
|
|
581
|
+
|
|
582
|
+
def _preselect_pairs(H0, mux, muy, muz, sys, exp, opt, B_center, freq_sweep=False):
|
|
583
|
+
"""EasySpin resfields transition pre-selection: transition rates of all level
|
|
584
|
+
pairs at the centre field on a small D2h grid (TPSGridSize 4); keep the pairs
|
|
585
|
+
whose maximum rate reaches Opt.Threshold × the largest; pure nuclear
|
|
586
|
+
transitions are dropped when the hyperfine interaction is weak
|
|
587
|
+
(HFIStrength = max|A|·(I+1/2)/ν < 0.5). Returns a list of (u, v) pairs."""
|
|
588
|
+
from torchspin.excitation import excitation_geometry, exp_mw_mode
|
|
589
|
+
thr = float(opt.Threshold)
|
|
590
|
+
n = H0.shape[0]
|
|
591
|
+
if thr <= 0:
|
|
592
|
+
return None
|
|
593
|
+
phi_t, theta_t, _, _ = sphgrid('D2h', 4)
|
|
594
|
+
geom = excitation_geometry(exp_mw_mode(exp))
|
|
595
|
+
Ex, Ey, Ez = -mux, -muy, -muz
|
|
596
|
+
max_rate = torch.zeros(n, n, dtype=torch.float64)
|
|
597
|
+
for ph, th in zip(phi_t.tolist(), theta_t.tolist()):
|
|
598
|
+
st, ct, sp, cp = math.sin(th), math.cos(th), math.sin(ph), math.cos(ph)
|
|
599
|
+
muzL = st * (cp * mux + sp * muy) + ct * muz
|
|
600
|
+
_, V = torch.linalg.eigh(H0 - float(B_center) * muzL)
|
|
601
|
+
Exy = cp * Ex + sp * Ey
|
|
602
|
+
if geom.parallel:
|
|
603
|
+
EzL = st * Exy + ct * Ez
|
|
604
|
+
rate = (V.conj().T @ EzL @ V).abs() ** 2
|
|
605
|
+
else:
|
|
606
|
+
EyL = -sp * Ex + cp * Ey
|
|
607
|
+
ExL = ct * Exy - st * Ez
|
|
608
|
+
rate = ((V.conj().T @ ExL @ V).abs() ** 2 + (V.conj().T @ EyL @ V).abs() ** 2) / 2
|
|
609
|
+
max_rate = torch.maximum(max_rate, rate.real.to(torch.float64))
|
|
610
|
+
keep = torch.triu(torch.ones(n, n, dtype=torch.bool), diagonal=1)
|
|
611
|
+
if sys.nNuclei > 0:
|
|
612
|
+
nu = float(exp.mwFreq) * 1e3 if not freq_sweep else 1e9 # MHz (frequency sweeps: no exclusion)
|
|
613
|
+
A = sys.A.detach().cpu().numpy()
|
|
614
|
+
Amax = np.abs(A).reshape(sys.nNuclei, -1).max(axis=1)
|
|
615
|
+
hfi = Amax * (0.5 + np.asarray(sys.I, dtype=float)) / nu
|
|
616
|
+
if hfi.max() < 0.5:
|
|
617
|
+
n_el = int(np.prod([2 * float(S_) + 1 for S_ in sys.S]))
|
|
618
|
+
n_nuc = n // n_el
|
|
619
|
+
el_idx = torch.arange(n) // n_nuc
|
|
620
|
+
keep &= el_idx.unsqueeze(1) != el_idx.unsqueeze(0) # drop pure nuclear transitions
|
|
621
|
+
rates = max_rate[keep]
|
|
622
|
+
cutoff = thr * float(rates.max()) if rates.numel() else 0.0
|
|
623
|
+
sel = keep & (max_rate > cutoff)
|
|
624
|
+
ii, jj = torch.nonzero(sel, as_tuple=True)
|
|
625
|
+
return [(int(a), int(b)) for a, b in zip(ii.tolist(), jj.tolist())]
|
|
626
|
+
|
|
627
|
+
|
|
628
|
+
def _projectzones_loop(
|
|
629
|
+
pos: np.ndarray,
|
|
630
|
+
amp: np.ndarray,
|
|
631
|
+
seg_weights: np.ndarray,
|
|
632
|
+
x: np.ndarray,
|
|
633
|
+
) -> np.ndarray:
|
|
634
|
+
"""Port of EasySpin's ``projectzones.c`` — axial-grid projection.
|
|
635
|
+
|
|
636
|
+
Each segment between consecutive orientations (pos[i], pos[i+1]) contributes
|
|
637
|
+
a rectangular (uniform) density to the spectrum, scaled by seg_weights[i].
|
|
638
|
+
|
|
639
|
+
Parameters
|
|
640
|
+
----------
|
|
641
|
+
pos : (N,) resonance field positions in mT
|
|
642
|
+
amp : (N,) or (1,) intensities; per-segment mean of adjacent values is used
|
|
643
|
+
seg_weights : (N-1,) solid-angle weights (should sum to 4π)
|
|
644
|
+
x : (nPoints,) field axis (mT, uniform spacing)
|
|
645
|
+
|
|
646
|
+
Returns
|
|
647
|
+
-------
|
|
648
|
+
spec : (nPoints,) spectral density (same units as projecttriangles output)
|
|
649
|
+
"""
|
|
650
|
+
nPoints = int(len(x))
|
|
651
|
+
delta = float(x[1] - x[0])
|
|
652
|
+
x0 = float(x[0])
|
|
653
|
+
spec = np.zeros(nPoints)
|
|
654
|
+
|
|
655
|
+
anisoAmp = (len(amp) > 1)
|
|
656
|
+
if not anisoAmp:
|
|
657
|
+
fixed_meanAmp = float(amp[0]) / delta
|
|
658
|
+
|
|
659
|
+
for iSeg in range(len(seg_weights)):
|
|
660
|
+
left = float(pos[iSeg])
|
|
661
|
+
right = float(pos[iSeg + 1])
|
|
662
|
+
if np.isnan(left) or np.isnan(right):
|
|
663
|
+
continue
|
|
664
|
+
if left > right:
|
|
665
|
+
left, right = right, left
|
|
666
|
+
|
|
667
|
+
if anisoAmp:
|
|
668
|
+
meanAmp = (float(amp[iSeg]) + float(amp[iSeg + 1])) / 2.0 / delta
|
|
669
|
+
else:
|
|
670
|
+
meanAmp = fixed_meanAmp
|
|
671
|
+
|
|
672
|
+
sw = float(seg_weights[iSeg])
|
|
673
|
+
left_b = (left - x0) / delta
|
|
674
|
+
right_b = (right - x0) / delta
|
|
675
|
+
first = int(left_b)
|
|
676
|
+
last = int(right_b)
|
|
677
|
+
|
|
678
|
+
if first >= nPoints or last < 0:
|
|
679
|
+
continue
|
|
680
|
+
|
|
681
|
+
if first == last:
|
|
682
|
+
if 0 <= first < nPoints:
|
|
683
|
+
spec[first] += meanAmp * sw
|
|
684
|
+
else:
|
|
685
|
+
Height = meanAmp * sw / (right_b - left_b)
|
|
686
|
+
if first >= 0:
|
|
687
|
+
spec[first] += Height * (first + 1 - left_b)
|
|
688
|
+
else:
|
|
689
|
+
first = -1
|
|
690
|
+
if last < nPoints:
|
|
691
|
+
spec[last] += Height * (right_b - last)
|
|
692
|
+
else:
|
|
693
|
+
last = nPoints
|
|
694
|
+
# Interior bins — replace Python loop with NumPy slice
|
|
695
|
+
lo, hi = first + 1, last # bounds already clamped above
|
|
696
|
+
if lo < hi:
|
|
697
|
+
spec[lo:hi] += Height
|
|
698
|
+
|
|
699
|
+
return spec
|
|
700
|
+
|
|
701
|
+
|
|
702
|
+
def _projecttriangles_loop(
|
|
703
|
+
tri_idx: np.ndarray,
|
|
704
|
+
areas: np.ndarray,
|
|
705
|
+
fun: np.ndarray,
|
|
706
|
+
amp: np.ndarray,
|
|
707
|
+
x: np.ndarray,
|
|
708
|
+
) -> np.ndarray:
|
|
709
|
+
"""Port of EasySpin's ``projecttriangles.c`` — 2-D triangulation projection.
|
|
710
|
+
|
|
711
|
+
For each triangle the three resonance-field positions define a tent (hat)
|
|
712
|
+
function: linearly rising from the leftmost to the middle vertex, then
|
|
713
|
+
linearly falling to the rightmost vertex. This eliminates the quantization
|
|
714
|
+
noise of stick-spectrum binning and converges at much coarser grids.
|
|
715
|
+
|
|
716
|
+
Parameters
|
|
717
|
+
----------
|
|
718
|
+
tri_idx : (nTri, 3) 0-based vertex indices into *fun* / *amp*
|
|
719
|
+
areas : (nTri,) normalized solid-angle areas (must sum to 4π)
|
|
720
|
+
fun : (nPts,) resonance field positions (mT)
|
|
721
|
+
amp : (nPts,) transition intensities
|
|
722
|
+
x : (nPoints,) field axis (mT, uniform spacing)
|
|
723
|
+
|
|
724
|
+
Returns
|
|
725
|
+
-------
|
|
726
|
+
spec : (nPoints,) spectral density
|
|
727
|
+
"""
|
|
728
|
+
nPoints = int(len(x))
|
|
729
|
+
delta = float(x[1] - x[0])
|
|
730
|
+
x0 = float(x[0])
|
|
731
|
+
spec = np.zeros(nPoints)
|
|
732
|
+
# Pre-allocated float64 index buffer avoids per-triangle np.arange allocation
|
|
733
|
+
_idx_buf = np.arange(nPoints, dtype=np.float64)
|
|
734
|
+
|
|
735
|
+
for iTri in range(len(tri_idx)):
|
|
736
|
+
i1, i2, i3 = int(tri_idx[iTri, 0]), int(tri_idx[iTri, 1]), int(tri_idx[iTri, 2])
|
|
737
|
+
p1, p2, p3 = float(fun[i1]), float(fun[i2]), float(fun[i3])
|
|
738
|
+
if np.isnan(p1) or np.isnan(p2) or np.isnan(p3):
|
|
739
|
+
continue
|
|
740
|
+
|
|
741
|
+
# Sort vertices: left ≤ middle ≤ right (mT)
|
|
742
|
+
left, middle, right = sorted([p1, p2, p3])
|
|
743
|
+
|
|
744
|
+
area = float(areas[iTri])
|
|
745
|
+
amplitude = (float(amp[i1]) + float(amp[i2]) + float(amp[i3])) / 3.0
|
|
746
|
+
|
|
747
|
+
# Convert to fractional bin coordinates
|
|
748
|
+
left_b = (left - x0) / delta
|
|
749
|
+
middle_b = (middle - x0) / delta
|
|
750
|
+
right_b = (right - x0) / delta
|
|
751
|
+
|
|
752
|
+
Width = right_b - left_b
|
|
753
|
+
Width1 = middle_b - left_b # left sub-triangle
|
|
754
|
+
Width2 = right_b - middle_b # right sub-triangle
|
|
755
|
+
|
|
756
|
+
# Zero-width: all vertices at same field → delta function
|
|
757
|
+
if Width == 0.0:
|
|
758
|
+
fi = int(left_b)
|
|
759
|
+
if 0 <= fi < nPoints:
|
|
760
|
+
spec[fi] += amplitude * area / delta
|
|
761
|
+
continue
|
|
762
|
+
|
|
763
|
+
# ── LEFT sub-triangle: ramp up from left_b to middle_b ────────────
|
|
764
|
+
if Width1 > 0.0:
|
|
765
|
+
f0 = (2.0 * amplitude * area / Width) / Width1 / delta
|
|
766
|
+
first1 = int(left_b)
|
|
767
|
+
last1 = int(middle_b)
|
|
768
|
+
if first1 < nPoints and last1 >= 0:
|
|
769
|
+
if first1 == last1:
|
|
770
|
+
spec[first1] += f0 * Width1 * Width1 / 2.0
|
|
771
|
+
else:
|
|
772
|
+
if first1 >= 0:
|
|
773
|
+
spec[first1] += f0 * (first1 + 1 - left_b) ** 2 / 2.0
|
|
774
|
+
else:
|
|
775
|
+
first1 = -1
|
|
776
|
+
if last1 < nPoints:
|
|
777
|
+
spec[last1] += f0 * ((last1 + middle_b) / 2.0 - left_b) * (middle_b - last1)
|
|
778
|
+
else:
|
|
779
|
+
last1 = nPoints
|
|
780
|
+
lsh = left_b - 0.5
|
|
781
|
+
lo1, hi1 = first1 + 1, last1 # bounds already clamped
|
|
782
|
+
if lo1 < hi1:
|
|
783
|
+
spec[lo1:hi1] += f0 * (_idx_buf[lo1:hi1] - lsh)
|
|
784
|
+
|
|
785
|
+
# ── RIGHT sub-triangle: ramp down from middle_b to right_b ────────
|
|
786
|
+
if Width2 > 0.0:
|
|
787
|
+
f0 = (2.0 * amplitude * area / Width) / Width2 / delta
|
|
788
|
+
first2 = int(middle_b)
|
|
789
|
+
last2 = int(right_b)
|
|
790
|
+
if first2 < nPoints and last2 >= 0:
|
|
791
|
+
if first2 == last2:
|
|
792
|
+
spec[first2] += f0 * Width2 * Width2 / 2.0
|
|
793
|
+
else:
|
|
794
|
+
if first2 >= 0:
|
|
795
|
+
spec[first2] += f0 * (first2 + 1 - middle_b) * (right_b - (middle_b + first2 + 1) / 2.0)
|
|
796
|
+
else:
|
|
797
|
+
first2 = -1
|
|
798
|
+
if last2 < nPoints:
|
|
799
|
+
spec[last2] += f0 * (right_b - last2) ** 2 / 2.0
|
|
800
|
+
else:
|
|
801
|
+
last2 = nPoints
|
|
802
|
+
rsh = right_b - 0.5
|
|
803
|
+
lo2, hi2 = first2 + 1, last2 # bounds already clamped
|
|
804
|
+
if lo2 < hi2:
|
|
805
|
+
spec[lo2:hi2] += f0 * (rsh - _idx_buf[lo2:hi2])
|
|
806
|
+
|
|
807
|
+
return spec
|
|
808
|
+
|
|
809
|
+
|
|
810
|
+
def _ramp_bins(spec: np.ndarray, f0: np.ndarray, a: np.ndarray, b: np.ndarray, rising: bool) -> None:
|
|
811
|
+
"""Accumulate linear ramps ``f0·(u−a)`` (rising) or ``f0·(b−u)`` (falling) on
|
|
812
|
+
``[a, b]`` (fractional bin coordinates, a ≤ b) into ``spec`` — the vectorised
|
|
813
|
+
form of the sub-triangle branches of ``projecttriangles.c``.
|
|
814
|
+
|
|
815
|
+
Ramps that touch at most three bins are scattered bin by bin; wider ramps
|
|
816
|
+
(whose ``f0`` is bounded, so no cancellation) contribute their edge bins
|
|
817
|
+
directly and their interior bins through two cumulative sums (the interior
|
|
818
|
+
value is linear in the bin index).
|
|
819
|
+
"""
|
|
820
|
+
nPoints = spec.shape[0]
|
|
821
|
+
width = b - a
|
|
822
|
+
first = np.trunc(a).astype(np.int64) # C-style truncation (as in the .c code)
|
|
823
|
+
last = np.trunc(b).astype(np.int64)
|
|
824
|
+
ok = (width > 0.0) & (first < nPoints) & (last >= 0)
|
|
825
|
+
if not ok.any():
|
|
826
|
+
return
|
|
827
|
+
f0, a, b, width, first, last = f0[ok], a[ok], b[ok], width[ok], first[ok], last[ok]
|
|
828
|
+
one = first == last
|
|
829
|
+
if one.any():
|
|
830
|
+
idx = first[one]
|
|
831
|
+
inr = (idx >= 0) & (idx < nPoints)
|
|
832
|
+
np.add.at(spec, idx[inr], (f0[one] * width[one] * width[one] / 2.0)[inr])
|
|
833
|
+
multi = ~one
|
|
834
|
+
f0, a, b, first, last = f0[multi], a[multi], b[multi], first[multi], last[multi]
|
|
835
|
+
if f0.size == 0:
|
|
836
|
+
return
|
|
837
|
+
# edge bins
|
|
838
|
+
if rising:
|
|
839
|
+
v_first = f0 * (first + 1 - a) ** 2 / 2.0
|
|
840
|
+
v_last = f0 * ((last + b) / 2.0 - a) * (b - last)
|
|
841
|
+
else:
|
|
842
|
+
v_first = f0 * (first + 1 - a) * (b - (a + first + 1) / 2.0)
|
|
843
|
+
v_last = f0 * (b - last) ** 2 / 2.0
|
|
844
|
+
m = first >= 0
|
|
845
|
+
np.add.at(spec, first[m], v_first[m])
|
|
846
|
+
m = last < nPoints
|
|
847
|
+
np.add.at(spec, last[m], v_last[m])
|
|
848
|
+
lo = np.maximum(first + 1, 0) # interior bins lo .. hi-1
|
|
849
|
+
hi = np.minimum(last, nPoints)
|
|
850
|
+
has_int = lo < hi
|
|
851
|
+
if not has_int.any():
|
|
852
|
+
return
|
|
853
|
+
f0, a, b, lo, hi = f0[has_int], a[has_int], b[has_int], lo[has_int], hi[has_int]
|
|
854
|
+
# interior value at bin k: rising f0·(k − (a − ½)), falling f0·((b − ½) − k)
|
|
855
|
+
narrow = (hi - lo) <= 2
|
|
856
|
+
if narrow.any():
|
|
857
|
+
for off in (0, 1):
|
|
858
|
+
k = lo[narrow] + off
|
|
859
|
+
sel = k < hi[narrow]
|
|
860
|
+
kk = k[sel]
|
|
861
|
+
if rising:
|
|
862
|
+
val = f0[narrow][sel] * (kk - (a[narrow][sel] - 0.5))
|
|
863
|
+
else:
|
|
864
|
+
val = f0[narrow][sel] * ((b[narrow][sel] - 0.5) - kk)
|
|
865
|
+
np.add.at(spec, kk, val)
|
|
866
|
+
wide = ~narrow
|
|
867
|
+
if wide.any():
|
|
868
|
+
f0w, lo_w, hi_w = f0[wide], lo[wide], hi[wide]
|
|
869
|
+
if rising:
|
|
870
|
+
s1, s0 = f0w, -f0w * (a[wide] - 0.5)
|
|
871
|
+
else:
|
|
872
|
+
s1, s0 = -f0w, f0w * (b[wide] - 0.5)
|
|
873
|
+
S1 = np.zeros(nPoints + 1); S0 = np.zeros(nPoints + 1)
|
|
874
|
+
np.add.at(S1, lo_w, s1); np.add.at(S1, hi_w, -s1)
|
|
875
|
+
np.add.at(S0, lo_w, s0); np.add.at(S0, hi_w, -s0)
|
|
876
|
+
c1 = np.cumsum(S1)[:nPoints]; c0 = np.cumsum(S0)[:nPoints]
|
|
877
|
+
spec += c1 * np.arange(nPoints, dtype=np.float64) + c0
|
|
878
|
+
|
|
879
|
+
|
|
880
|
+
def _projectzones(
|
|
881
|
+
pos: np.ndarray,
|
|
882
|
+
amp: np.ndarray,
|
|
883
|
+
seg_weights: np.ndarray,
|
|
884
|
+
x: np.ndarray,
|
|
885
|
+
) -> np.ndarray:
|
|
886
|
+
"""Vectorised :func:`_projectzones_loop` (EasySpin ``projectzones.c``)."""
|
|
887
|
+
nPoints = int(len(x))
|
|
888
|
+
delta = float(x[1] - x[0])
|
|
889
|
+
x0 = float(x[0])
|
|
890
|
+
spec = np.zeros(nPoints)
|
|
891
|
+
pos = np.asarray(pos, dtype=float); amp = np.asarray(amp, dtype=float)
|
|
892
|
+
sw = np.asarray(seg_weights, dtype=float)
|
|
893
|
+
nSeg = sw.shape[0]
|
|
894
|
+
left = np.minimum(pos[:nSeg], pos[1:nSeg + 1]); right = np.maximum(pos[:nSeg], pos[1:nSeg + 1])
|
|
895
|
+
if amp.size > 1:
|
|
896
|
+
meanAmp = (amp[:nSeg] + amp[1:nSeg + 1]) / 2.0 / delta
|
|
897
|
+
else:
|
|
898
|
+
meanAmp = np.full(nSeg, float(amp[0]) / delta)
|
|
899
|
+
ok = ~(np.isnan(left) | np.isnan(right))
|
|
900
|
+
left, right, meanAmp, sw = left[ok], right[ok], meanAmp[ok], sw[ok]
|
|
901
|
+
left_b = (left - x0) / delta; right_b = (right - x0) / delta
|
|
902
|
+
first = np.trunc(left_b).astype(np.int64); last = np.trunc(right_b).astype(np.int64)
|
|
903
|
+
ok = (first < nPoints) & (last >= 0)
|
|
904
|
+
left_b, right_b, meanAmp, sw, first, last = left_b[ok], right_b[ok], meanAmp[ok], sw[ok], first[ok], last[ok]
|
|
905
|
+
one = first == last
|
|
906
|
+
if one.any():
|
|
907
|
+
idx = first[one]; inr = (idx >= 0) & (idx < nPoints)
|
|
908
|
+
np.add.at(spec, idx[inr], (meanAmp[one] * sw[one])[inr])
|
|
909
|
+
multi = ~one
|
|
910
|
+
if not multi.any():
|
|
911
|
+
return spec
|
|
912
|
+
left_b, right_b, meanAmp, sw, first, last = left_b[multi], right_b[multi], meanAmp[multi], sw[multi], first[multi], last[multi]
|
|
913
|
+
Height = meanAmp * sw / (right_b - left_b)
|
|
914
|
+
m = first >= 0
|
|
915
|
+
np.add.at(spec, first[m], (Height * (first + 1 - left_b))[m])
|
|
916
|
+
m = last < nPoints
|
|
917
|
+
np.add.at(spec, last[m], (Height * (right_b - last))[m])
|
|
918
|
+
lo = np.maximum(first + 1, 0); hi = np.minimum(last, nPoints)
|
|
919
|
+
has_int = lo < hi
|
|
920
|
+
if has_int.any():
|
|
921
|
+
S0 = np.zeros(nPoints + 1)
|
|
922
|
+
np.add.at(S0, lo[has_int], Height[has_int]); np.add.at(S0, hi[has_int], -Height[has_int])
|
|
923
|
+
spec += np.cumsum(S0)[:nPoints]
|
|
924
|
+
return spec
|
|
925
|
+
|
|
926
|
+
|
|
927
|
+
def _projecttriangles(
|
|
928
|
+
tri_idx: np.ndarray,
|
|
929
|
+
areas: np.ndarray,
|
|
930
|
+
fun: np.ndarray,
|
|
931
|
+
amp: np.ndarray,
|
|
932
|
+
x: np.ndarray,
|
|
933
|
+
) -> np.ndarray:
|
|
934
|
+
"""Vectorised :func:`_projecttriangles_loop` (EasySpin ``projecttriangles.c``):
|
|
935
|
+
every triangle's tent function is split into its rising and falling ramps,
|
|
936
|
+
which are accumulated by :func:`_ramp_bins`."""
|
|
937
|
+
nPoints = int(len(x))
|
|
938
|
+
delta = float(x[1] - x[0])
|
|
939
|
+
x0 = float(x[0])
|
|
940
|
+
spec = np.zeros(nPoints)
|
|
941
|
+
tri_idx = np.asarray(tri_idx, dtype=np.int64)
|
|
942
|
+
fun = np.asarray(fun, dtype=float); amp = np.asarray(amp, dtype=float)
|
|
943
|
+
if tri_idx.size == 0:
|
|
944
|
+
return spec
|
|
945
|
+
P = fun[tri_idx] # (nTri, 3)
|
|
946
|
+
ok = ~np.isnan(P).any(axis=1)
|
|
947
|
+
if not ok.any():
|
|
948
|
+
return spec
|
|
949
|
+
P = np.sort(P[ok], axis=1)
|
|
950
|
+
area = np.asarray(areas, dtype=float)[ok]
|
|
951
|
+
A = amp[tri_idx[ok]]
|
|
952
|
+
amplitude = (A[:, 0] + A[:, 1] + A[:, 2]) / 3.0
|
|
953
|
+
left_b = (P[:, 0] - x0) / delta; middle_b = (P[:, 1] - x0) / delta; right_b = (P[:, 2] - x0) / delta
|
|
954
|
+
Width = right_b - left_b
|
|
955
|
+
Width1 = middle_b - left_b
|
|
956
|
+
Width2 = right_b - middle_b
|
|
957
|
+
zero = Width == 0.0
|
|
958
|
+
if zero.any():
|
|
959
|
+
fi = np.trunc(left_b[zero]).astype(np.int64)
|
|
960
|
+
inr = (fi >= 0) & (fi < nPoints)
|
|
961
|
+
np.add.at(spec, fi[inr], (amplitude[zero] * area[zero] / delta)[inr])
|
|
962
|
+
nz = ~zero
|
|
963
|
+
if not nz.any():
|
|
964
|
+
return spec
|
|
965
|
+
amplitude, area, left_b, middle_b, right_b, Width, Width1, Width2 = (
|
|
966
|
+
amplitude[nz], area[nz], left_b[nz], middle_b[nz], right_b[nz], Width[nz], Width1[nz], Width2[nz])
|
|
967
|
+
base = 2.0 * amplitude * area / Width
|
|
968
|
+
with np.errstate(divide='ignore', invalid='ignore'):
|
|
969
|
+
f0L = base / Width1 / delta
|
|
970
|
+
f0R = base / Width2 / delta
|
|
971
|
+
sL = Width1 > 0.0
|
|
972
|
+
if sL.any():
|
|
973
|
+
_ramp_bins(spec, f0L[sL], left_b[sL], middle_b[sL], rising=True)
|
|
974
|
+
sR = Width2 > 0.0
|
|
975
|
+
if sR.any():
|
|
976
|
+
_ramp_bins(spec, f0R[sR], middle_b[sR], right_b[sR], rising=False)
|
|
977
|
+
return spec
|
|
978
|
+
|
|
979
|
+
|
|
980
|
+
def pepper(sys_or_list, exp: Experiment, opt: Optional[Options] = None):
|
|
981
|
+
"""Field-swept CW EPR powder spectrum.
|
|
982
|
+
|
|
983
|
+
Accepts either a single :class:`SpinSystem` or a list/tuple of systems
|
|
984
|
+
(multi-component fit). When a list is given each component is simulated
|
|
985
|
+
independently and the results summed, weighted by ``sys.weight``
|
|
986
|
+
(default 1.0).
|
|
987
|
+
|
|
988
|
+
Parameters
|
|
989
|
+
----------
|
|
990
|
+
sys_or_list : SpinSystem or list of SpinSystem
|
|
991
|
+
Spin system(s). ``sys.lw = [fwhm_g, fwhm_l]`` (mT) sets the global
|
|
992
|
+
linewidth. Strain parameters (``HStrain``, ``gStrain``, ``AStrain``,
|
|
993
|
+
``DStrain``) enable orientation-dependent broadening for each
|
|
994
|
+
transition. For multi-component, each ``sys.weight`` scales that
|
|
995
|
+
component's contribution.
|
|
996
|
+
exp : Experiment
|
|
997
|
+
Experimental parameters. Field-swept mode requires ``mwFreq`` (GHz)
|
|
998
|
+
and ``Range`` (mT). Frequency-swept mode (``Field`` and ``mwRange``
|
|
999
|
+
both set) dispatches automatically. ``Harmonic`` sets derivative
|
|
1000
|
+
order (0=absorption, 1=first derivative, 2=second derivative).
|
|
1001
|
+
``Temperature`` (K) enables Boltzmann populations.
|
|
1002
|
+
opt : Options, optional
|
|
1003
|
+
Computational options. Defaults to ``Options()`` if ``None``.
|
|
1004
|
+
Key fields: ``GridSize`` (int or [N_coarse, N_interp]),
|
|
1005
|
+
``GridSymmetry`` (``'auto'`` by default), ``device``, ``Verbosity``.
|
|
1006
|
+
|
|
1007
|
+
Returns
|
|
1008
|
+
-------
|
|
1009
|
+
B : torch.Tensor, shape (nPoints,)
|
|
1010
|
+
Field axis in mT.
|
|
1011
|
+
spec : torch.Tensor, shape (nPoints,)
|
|
1012
|
+
Spectrum in arbitrary units (consistent with EasySpin pepper).
|
|
1013
|
+
|
|
1014
|
+
Examples
|
|
1015
|
+
--------
|
|
1016
|
+
Nitroxide at X-band::
|
|
1017
|
+
|
|
1018
|
+
>>> from torchspin import SpinSystem, Experiment, Options, pepper
|
|
1019
|
+
>>> sys = SpinSystem(S=[0.5], g=[[2.009, 2.006, 2.002]],
|
|
1020
|
+
... Nucs=['14N'], A=[[10, 10, 95]], lw=[1.0, 0.0])
|
|
1021
|
+
>>> exp = Experiment(mwFreq=9.5, Range=[330, 350], nPoints=1024, Harmonic=1)
|
|
1022
|
+
>>> B, spc = pepper(sys, exp, Options(GridSize=50))
|
|
1023
|
+
|
|
1024
|
+
Two-component mixture weighted 70:30::
|
|
1025
|
+
|
|
1026
|
+
>>> sys_a = SpinSystem(S=[0.5], g=[[2.0, 2.1, 2.2]], lw=[1.0, 0.0], weight=0.7)
|
|
1027
|
+
>>> sys_b = SpinSystem(S=[0.5], g=2.0037, lw=[0.8, 0.0], weight=0.3)
|
|
1028
|
+
>>> B, spc = pepper([sys_a, sys_b], exp)
|
|
1029
|
+
|
|
1030
|
+
Notes
|
|
1031
|
+
-----
|
|
1032
|
+
* Only the matrix-diagonalization method is implemented (set via
|
|
1033
|
+
``Options.Method='matrix'``).
|
|
1034
|
+
* Powder averaging uses the SOPHE grid (Wang & Hanson, 1995).
|
|
1035
|
+
* Lineshape convolution is FFT-based.
|
|
1036
|
+
* Strain broadening adds per-transition widths in quadrature with
|
|
1037
|
+
``sys.lw``.
|
|
1038
|
+
* MATLAB-validated to cosine similarity > 0.999 on standard test cases.
|
|
1039
|
+
"""
|
|
1040
|
+
opt_ = opt if opt is not None else Options()
|
|
1041
|
+
# Components × isotopologues (EasySpin compisoloop): every component is
|
|
1042
|
+
# expanded into its isotopologues (natural abundance or Sys.Abund), each
|
|
1043
|
+
# simulated separately and summed with weight Sys.weight·abundance. With
|
|
1044
|
+
# Opt.separate='components' the isotopologue spectra are returned as rows.
|
|
1045
|
+
components = expand_components(sys_or_list, getattr(opt_, 'IsoCutoff', 1e-4))
|
|
1046
|
+
for _s in components:
|
|
1047
|
+
if getattr(_s, 'n', None) is not None and any(v > 1 for v in _s.n):
|
|
1048
|
+
raise ValueError(
|
|
1049
|
+
"pepper does not support sets of equivalent nuclei (SpinSystem.n > 1). "
|
|
1050
|
+
"List each nucleus separately in Nucs, or use garlic for isotropic spectra."
|
|
1051
|
+
)
|
|
1052
|
+
rows, B_out = [], None
|
|
1053
|
+
for sys in components:
|
|
1054
|
+
w = float(getattr(sys, 'weight', 1.0))
|
|
1055
|
+
B_out, spec = _pepper_single(sys, exp, opt_)
|
|
1056
|
+
rows.append(w * spec)
|
|
1057
|
+
if str(getattr(opt_, 'separate', '')) in ('components', 'transitions', 'orientations', 'sites'):
|
|
1058
|
+
# EasySpin compisoloop: any separate output concatenates the rows of all
|
|
1059
|
+
# components / isotopologues along the first dimension
|
|
1060
|
+
return B_out, torch.cat([r if r.ndim == 2 else r.unsqueeze(0) for r in rows], dim=0)
|
|
1061
|
+
return B_out, sum(rows)
|
|
1062
|
+
|
|
1063
|
+
|
|
1064
|
+
|
|
1065
|
+
def _pepper_single(
|
|
1066
|
+
sys: SpinSystem,
|
|
1067
|
+
exp: Experiment,
|
|
1068
|
+
opt: Optional[Options] = None,
|
|
1069
|
+
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
1070
|
+
"""Field-swept CW EPR powder spectrum.
|
|
1071
|
+
|
|
1072
|
+
Parameters
|
|
1073
|
+
----------
|
|
1074
|
+
sys:
|
|
1075
|
+
Spin system. ``sys.lw = [fwhm_g, fwhm_l]`` (mT) sets the global linewidth.
|
|
1076
|
+
Strain parameters (``HStrain``, ``gStrain``, ``AStrain``, ``DStrain``) enable
|
|
1077
|
+
orientation-dependent broadening for each transition.
|
|
1078
|
+
exp:
|
|
1079
|
+
Experimental parameters (mwFreq, Range, nPoints, Harmonic).
|
|
1080
|
+
opt:
|
|
1081
|
+
Computational options. Defaults to ``Options()`` if ``None``.
|
|
1082
|
+
|
|
1083
|
+
Returns
|
|
1084
|
+
-------
|
|
1085
|
+
x:
|
|
1086
|
+
Field axis in mT, shape ``(nPoints,)``.
|
|
1087
|
+
spec:
|
|
1088
|
+
Spectrum in a.u., shape ``(nPoints,)``.
|
|
1089
|
+
|
|
1090
|
+
Notes
|
|
1091
|
+
-----
|
|
1092
|
+
* Only the matrix diagonalization method is implemented.
|
|
1093
|
+
* Powder averaging uses the SOPHE grid (Wang & Hanson 1995).
|
|
1094
|
+
* Lineshape convolution is FFT-based.
|
|
1095
|
+
* Temperature-dependent Boltzmann populations: set ``exp.Temperature`` (K).
|
|
1096
|
+
If ``None`` or not provided, all levels have equal population (infinite T).
|
|
1097
|
+
* Strain broadening: If any strain parameters are present, per-transition
|
|
1098
|
+
broadening is applied. This is slower than global broadening but essential
|
|
1099
|
+
for realistic spectra. sys.lw is added in quadrature with strain widths.
|
|
1100
|
+
"""
|
|
1101
|
+
if opt is None:
|
|
1102
|
+
opt = Options()
|
|
1103
|
+
|
|
1104
|
+
# Dispatch to frequency-swept path when Exp.Field and Exp.mwRange are set
|
|
1105
|
+
# Frequency sweeps (Exp.Field + Exp.mwRange) run through the same grid /
|
|
1106
|
+
# interpolation / projection / summation machinery as field sweeps
|
|
1107
|
+
# (EasySpin pepper.m): the x axis is the frequency range in GHz and the
|
|
1108
|
+
# resonances come from resfreqs_batch instead of resfields_batch.
|
|
1109
|
+
freq_sweep = bool(exp.is_freq_swept) or (exp.Field is not None and exp.mwFreq is None)
|
|
1110
|
+
from dataclasses import replace as _dc_replace_f
|
|
1111
|
+
if not freq_sweep and exp.Range is None:
|
|
1112
|
+
# EasySpin pepper.m: automatic field range for S=1/2 from the g and A extremes
|
|
1113
|
+
if sys.nElectrons != 1 or float(sys.S[0]) != 0.5:
|
|
1114
|
+
raise ValueError('Cannot automatically determine field range. Please provide Exp.CenterSweep or Exp.Range.')
|
|
1115
|
+
from torchspin.constants import PLANCK, BMAGN
|
|
1116
|
+
if sys.nNuclei > 0:
|
|
1117
|
+
A = sys.A.detach().cpu().numpy()
|
|
1118
|
+
if sys.fullA:
|
|
1119
|
+
Amax = np.abs(A).reshape(sys.nNuclei, 3, -1).max(axis=(1, 2))
|
|
1120
|
+
else:
|
|
1121
|
+
Amax = np.abs(A).max(axis=1)
|
|
1122
|
+
hf = float(np.sum(np.asarray(sys.I, dtype=float) * Amax)) * 1e6 # Hz
|
|
1123
|
+
else:
|
|
1124
|
+
hf = 0.0
|
|
1125
|
+
g = sys.g.detach().cpu().numpy()
|
|
1126
|
+
gvals = np.linalg.eigvalsh(g) if sys.fullg else g.reshape(-1)
|
|
1127
|
+
gmax, gmin = float(np.max(gvals)), float(np.min(gvals))
|
|
1128
|
+
minB = PLANCK * (float(exp.mwFreq) * 1e9 - hf) / BMAGN / gmax / 1e-3
|
|
1129
|
+
maxB = PLANCK * (float(exp.mwFreq) * 1e9 + hf) / BMAGN / gmin / 1e-3
|
|
1130
|
+
lw_max = float(max(sys.get_lw()))
|
|
1131
|
+
center = (maxB + minB) / 2
|
|
1132
|
+
sweep = maxB - minB + 3 * lw_max
|
|
1133
|
+
if sweep == 0:
|
|
1134
|
+
sweep = 5 * lw_max
|
|
1135
|
+
if sweep == 0:
|
|
1136
|
+
sweep = 10.0
|
|
1137
|
+
sweep *= 1.25
|
|
1138
|
+
exp = _dc_replace_f(exp, Range=[center - sweep / 2, center + sweep / 2], CenterSweep=None)
|
|
1139
|
+
freq_auto_range = freq_sweep and exp.mwRange is None
|
|
1140
|
+
if freq_sweep:
|
|
1141
|
+
if float(getattr(exp, 'ModAmp', 0.0) or 0.0) > 0:
|
|
1142
|
+
raise ValueError('Exp.ModAmp cannot be used with frequency sweeps.')
|
|
1143
|
+
rng0 = [float(exp.mwRange[0]), float(exp.mwRange[1])] if not freq_auto_range else [0.0, 1.0]
|
|
1144
|
+
exp = _dc_replace_f(exp, Range=rng0, mwRange=rng0, CenterSweep=None, mwCenterSweep=None)
|
|
1145
|
+
|
|
1146
|
+
# -----------------------------------------------------------------------
|
|
1147
|
+
# Step 0: Determine grid symmetry (automatic or user-specified)
|
|
1148
|
+
# -----------------------------------------------------------------------
|
|
1149
|
+
grid_sym = opt.GridSymmetry
|
|
1150
|
+
R_sym2mol = torch.eye(3, dtype=torch.float64) # symmetry → molecular frame rotation
|
|
1151
|
+
if grid_sym == 'auto' or grid_sym == '':
|
|
1152
|
+
# Automatic symmetry detection
|
|
1153
|
+
grid_sym, R_sym2mol = hamsymm(sys)
|
|
1154
|
+
if opt.Verbosity >= 1:
|
|
1155
|
+
print(f"pepper: automatic symmetry detection → '{grid_sym}'")
|
|
1156
|
+
else:
|
|
1157
|
+
if opt.Verbosity >= 1:
|
|
1158
|
+
print(f"pepper: user-specified symmetry '{grid_sym}'")
|
|
1159
|
+
|
|
1160
|
+
# EasySpin p_sampletype/p_gridsetup: partially ordered samples (Exp.Ordering)
|
|
1161
|
+
# use a Ci grid in the molecular frame; a non-equilibrium state reduces an
|
|
1162
|
+
# axial (Dinfh) grid to D2h.
|
|
1163
|
+
from torchspin.ordering import ordering_function
|
|
1164
|
+
ordering_fun = ordering_function(getattr(exp, 'Ordering', None))
|
|
1165
|
+
if ordering_fun is not None:
|
|
1166
|
+
if getattr(exp, 'MolFrame', None) is not None or getattr(exp, 'CrystalSymmetry', None) is not None:
|
|
1167
|
+
raise ValueError('Exp.MolFrame/Exp.CrystalSymmetry cannot be used for partially ordered samples (Exp.Ordering given).')
|
|
1168
|
+
if opt.GridSymmetry in ('auto', ''):
|
|
1169
|
+
grid_sym, R_sym2mol = 'Ci', torch.eye(3, dtype=torch.float64)
|
|
1170
|
+
elif grid_sym != 'Ci':
|
|
1171
|
+
raise ValueError('For partially ordered samples, Ci grid symmetry is required.')
|
|
1172
|
+
if getattr(sys, 'initState', None) is not None and grid_sym == 'Dinfh' and opt.GridSymmetry in ('auto', ''):
|
|
1173
|
+
grid_sym = 'D2h'
|
|
1174
|
+
|
|
1175
|
+
# -----------------------------------------------------------------------
|
|
1176
|
+
# Step 1: Build field-independent Hamiltonian and moment operators
|
|
1177
|
+
# -----------------------------------------------------------------------
|
|
1178
|
+
H0, mux, muy, muz = ham(sys, B0=None)
|
|
1179
|
+
|
|
1180
|
+
# Non-equilibrium populations (Sys.initState) → density matrix in the
|
|
1181
|
+
# Hamiltonian basis; not available with perturbation theory (EasySpin).
|
|
1182
|
+
from torchspin.initstate import init_state_density
|
|
1183
|
+
init_state = init_state_density(sys, H0) if getattr(sys, 'initState', None) is not None else None
|
|
1184
|
+
if init_state is not None and str(opt.Method).startswith('perturb'):
|
|
1185
|
+
raise ValueError('Perturbation theory not available for systems with non-equilibrium populations.')
|
|
1186
|
+
# Photoselection (Exp.lightBeam, Sys.tdm): per-orientation weights
|
|
1187
|
+
light_beam = getattr(exp, 'lightBeam', None)
|
|
1188
|
+
light_scatter = float(getattr(exp, 'lightScatter', 0.0) or 0.0)
|
|
1189
|
+
use_photo = light_beam not in (None, '') and light_scatter < 1
|
|
1190
|
+
if use_photo:
|
|
1191
|
+
if getattr(sys, 'tdm', None) is None:
|
|
1192
|
+
raise ValueError('To include photoselection weights, Sys.tdm must be given.')
|
|
1193
|
+
if isinstance(light_beam, str):
|
|
1194
|
+
photo_k = [0.0, 1.0, 0.0] # beam along yL
|
|
1195
|
+
photo_alpha = {'perpendicular': -math.pi / 2, 'parallel': math.pi,
|
|
1196
|
+
'unpolarized': float('nan')}.get(light_beam)
|
|
1197
|
+
if photo_alpha is None:
|
|
1198
|
+
raise ValueError("Unknown string in Exp.lightBeam. Use '', 'perpendicular', 'parallel' or 'unpolarized'.")
|
|
1199
|
+
else:
|
|
1200
|
+
photo_k, photo_alpha = light_beam[0], float(light_beam[1])
|
|
1201
|
+
|
|
1202
|
+
def _photo_w(ori):
|
|
1203
|
+
from torchspin.photoselect import photoselect
|
|
1204
|
+
w = photoselect(sys.tdm, np.asarray(ori, dtype=float), photo_k, photo_alpha)
|
|
1205
|
+
return torch.tensor((1.0 - light_scatter) * w + light_scatter, dtype=torch.float64)
|
|
1206
|
+
# perturbation shortcuts are only taken for equilibrium, non-photoselected spectra
|
|
1207
|
+
_pt_ok = init_state is None and not freq_sweep
|
|
1208
|
+
from dataclasses import replace as _dc_replace_o
|
|
1209
|
+
_opt_rf = _dc_replace_o(opt, Threshold=0.0) # slot-based paths pre-select transitions globally (EasySpin)
|
|
1210
|
+
_pairs_sel = _preselect_pairs(H0, mux, muy, muz, sys, exp, opt,
|
|
1211
|
+
float(exp.Field) if freq_sweep else 0.5 * (exp.Range[0] + exp.Range[1]),
|
|
1212
|
+
freq_sweep=freq_sweep) if not str(opt.Method).startswith('perturb') else None
|
|
1213
|
+
|
|
1214
|
+
# Opt.separate='transitions': one spectrum row per transition slot (EasySpin
|
|
1215
|
+
# nSpectra = nTransitions). Rows are collected per slot by _acc and, for the
|
|
1216
|
+
# matrix path, ordered by descending maximum intensity like EasySpin's
|
|
1217
|
+
# transition pre-selection.
|
|
1218
|
+
sep_trans = str(getattr(opt, 'separate', '')) == 'transitions'
|
|
1219
|
+
_rows: dict = {}
|
|
1220
|
+
_row_max: dict = {}
|
|
1221
|
+
|
|
1222
|
+
def _rid(i_t):
|
|
1223
|
+
keys = getattr(_transition_slots, 'last_keys', None)
|
|
1224
|
+
return (keys[i_t][0], keys[i_t][1]) if keys is not None and i_t < len(keys) else i_t
|
|
1225
|
+
|
|
1226
|
+
def _row(i_t):
|
|
1227
|
+
rid = _rid(i_t)
|
|
1228
|
+
if rid not in _rows:
|
|
1229
|
+
_rows[rid] = torch.zeros(exp.nPoints, dtype=torch.float64)
|
|
1230
|
+
_row_max[rid] = 0.0
|
|
1231
|
+
return _rows[rid]
|
|
1232
|
+
|
|
1233
|
+
def _acc(i_t, contrib, imax=None):
|
|
1234
|
+
nonlocal spec
|
|
1235
|
+
if sep_trans:
|
|
1236
|
+
rid = _rid(i_t)
|
|
1237
|
+
_rows[rid] = _row(i_t) + contrib
|
|
1238
|
+
if imax is not None:
|
|
1239
|
+
_row_max[rid] = max(_row_max[rid], float(imax))
|
|
1240
|
+
else:
|
|
1241
|
+
spec = spec + contrib
|
|
1242
|
+
|
|
1243
|
+
def _resonances(H0_, mux_, muy_, muz_, phi_b, theta_b, exp_, opt_, sys_, **kw):
|
|
1244
|
+
opt_ = kw.pop('opt_override', opt_)
|
|
1245
|
+
if freq_sweep:
|
|
1246
|
+
return resfreqs_batch(H0_, mux_, muy_, muz_, phi_b, theta_b, exp_, opt_, sys_, **kw)
|
|
1247
|
+
return resfields_batch(H0_, mux_, muy_, muz_, phi_b, theta_b, exp_, opt_, sys_, **kw)
|
|
1248
|
+
|
|
1249
|
+
# Move operators to target device (GPU if requested)
|
|
1250
|
+
device = torch.device(opt.device)
|
|
1251
|
+
H0 = H0.to(device)
|
|
1252
|
+
mux = mux.to(device)
|
|
1253
|
+
muy = muy.to(device)
|
|
1254
|
+
muz = muz.to(device)
|
|
1255
|
+
|
|
1256
|
+
# -----------------------------------------------------------------------
|
|
1257
|
+
# Step 2: Spherical powder grid
|
|
1258
|
+
# -----------------------------------------------------------------------
|
|
1259
|
+
# Parse GridSize: int N → compute + accumulate at N knots (no interpolation)
|
|
1260
|
+
# [N1, Ni] → compute at coarse N1-knot grid, interpolate to
|
|
1261
|
+
# fine nfKnots = (N1-1)*Ni + 1 knot grid
|
|
1262
|
+
N_coarse = opt.grid_size_coarse
|
|
1263
|
+
N_interp = opt.grid_size_interp
|
|
1264
|
+
nfKnots = (N_coarse - 1) * N_interp + 1 if N_interp > 1 else N_coarse
|
|
1265
|
+
|
|
1266
|
+
phi_arr, theta_arr, weights_arr, vecs_orig = sphgrid(grid_sym, N_coarse)
|
|
1267
|
+
n_orient = phi_arr.shape[0]
|
|
1268
|
+
# Symmetry-frame polar angles: grid interpolation and zone/triangle weights
|
|
1269
|
+
# live in the symmetry frame; only the Hamiltonian is evaluated along the
|
|
1270
|
+
# rotated (molecular-frame) field directions below.
|
|
1271
|
+
theta_sym = theta_arr.clone()
|
|
1272
|
+
vecs_sym_c = vecs_orig if vecs_orig.shape[0] == 3 else vecs_orig.T # (3, n_orient)
|
|
1273
|
+
|
|
1274
|
+
# Rotate grid orientations from the symmetry frame to the molecular frame.
|
|
1275
|
+
# erot(gFrame) uses the passive convention: v_g = erot(gFrame) @ v_mol (mol→g).
|
|
1276
|
+
# hamsymm's _tensor_symmetry stores erot(euler_angles) directly as R, so
|
|
1277
|
+
# R maps mol→symmetry. To convert sphgrid unit vectors (in symmetry frame)
|
|
1278
|
+
# to the molecular frame: v_mol = R.T @ v_sym.
|
|
1279
|
+
if not torch.allclose(R_sym2mol, torch.eye(3, dtype=torch.float64), atol=1e-10):
|
|
1280
|
+
sin_t = torch.sin(theta_arr)
|
|
1281
|
+
v_sym = torch.stack(
|
|
1282
|
+
[sin_t * torch.cos(phi_arr), sin_t * torch.sin(phi_arr), torch.cos(theta_arr)],
|
|
1283
|
+
dim=1,
|
|
1284
|
+
) # (N, 3)
|
|
1285
|
+
v_mol = (R_sym2mol.T @ v_sym.T).T # (N, 3)
|
|
1286
|
+
theta_arr = torch.acos(v_mol[:, 2].clamp(-1.0, 1.0))
|
|
1287
|
+
phi_arr = torch.atan2(v_mol[:, 1], v_mol[:, 0])
|
|
1288
|
+
|
|
1289
|
+
# Move grid arrays to target device
|
|
1290
|
+
phi_arr = phi_arr.to(device)
|
|
1291
|
+
theta_arr = theta_arr.to(device)
|
|
1292
|
+
weights_arr = weights_arr.to(device)
|
|
1293
|
+
|
|
1294
|
+
if opt.Verbosity >= 1:
|
|
1295
|
+
print(f"pepper: {n_orient} orientations, grid '{grid_sym}', N={opt.GridSize}")
|
|
1296
|
+
|
|
1297
|
+
# -----------------------------------------------------------------------
|
|
1298
|
+
# Step 3: Accumulate spectrum over all orientations (with optional batching)
|
|
1299
|
+
# -----------------------------------------------------------------------
|
|
1300
|
+
# x and spec stay on CPU; only Hamiltonians/grid move to device for fast eigh
|
|
1301
|
+
if freq_auto_range:
|
|
1302
|
+
# EasySpin pepper.m: automatic frequency range from the resonance
|
|
1303
|
+
# frequencies of the coarse grid (± spread/5, ≥ 5×strain width, ≥ 5×Σlw)
|
|
1304
|
+
pw0 = _photo_w(torch.stack([phi_arr, theta_arr], dim=1).detach().cpu().numpy()) if use_photo else None
|
|
1305
|
+
P_l, _, W_l = _resonances(H0, mux, muy, muz, phi_arr, theta_arr, exp, opt, sys,
|
|
1306
|
+
init_state=init_state, photo_weights=pw0)
|
|
1307
|
+
pos = torch.cat([p_ for p_ in P_l if p_ is not None and p_.numel() > 0])
|
|
1308
|
+
f_min, f_max = float(pos.min()), float(pos.max())
|
|
1309
|
+
padding = (f_max - f_min) / 5.0
|
|
1310
|
+
if padding == 0:
|
|
1311
|
+
padding = 0.1
|
|
1312
|
+
w_all = [w_ for w_ in W_l if w_ is not None and w_.numel() > 0]
|
|
1313
|
+
if w_all:
|
|
1314
|
+
padding = max(padding, 5.0 * float(torch.cat(w_all).max()))
|
|
1315
|
+
padding = max(padding, 5.0 * float(sum(sys.get_lw())) / 1e3) # Sys.lw in MHz
|
|
1316
|
+
rng_auto = [max(0.0, f_min - padding), f_max + padding]
|
|
1317
|
+
exp = _dc_replace_f(exp, Range=rng_auto, mwRange=rng_auto)
|
|
1318
|
+
x = torch.linspace(exp.Range[0], exp.Range[1], exp.nPoints, dtype=torch.float64)
|
|
1319
|
+
spec = torch.zeros(exp.nPoints, dtype=torch.float64)
|
|
1320
|
+
dx = (exp.Range[1] - exp.Range[0]) / (exp.nPoints - 1) # mT per point
|
|
1321
|
+
# EasySpin pepper.m: the perturbation solver (resfields_perturb) searches
|
|
1322
|
+
# Exp.SearchRange = Range ± 20 % of the sweep width (clamped at 0) so that
|
|
1323
|
+
# lines just outside the window contribute their in-range part; the
|
|
1324
|
+
# matrix-diagonalisation solver (resfields) searches the sweep range only.
|
|
1325
|
+
from dataclasses import replace as _dc_replace
|
|
1326
|
+
_w = exp.Range[1] - exp.Range[0]
|
|
1327
|
+
exp_search = _dc_replace(exp, Range=[max(0.0, exp.Range[0] - 0.2 * _w), exp.Range[1] + 0.2 * _w],
|
|
1328
|
+
CenterSweep=None)
|
|
1329
|
+
|
|
1330
|
+
total_weight = weights_arr.sum().item() # should be ~4π
|
|
1331
|
+
|
|
1332
|
+
# Get effective linewidth (handles both lw and lwpp)
|
|
1333
|
+
lw_effective = sys.get_lw()
|
|
1334
|
+
mw_phase = float(getattr(exp, 'mwPhase', 0.0) or 0.0) # EasySpin Exp.mwPhase (dispersion)
|
|
1335
|
+
if freq_sweep:
|
|
1336
|
+
lw_effective = [float(v) / 1e3 for v in lw_effective] # Sys.lw in MHz → GHz axis
|
|
1337
|
+
mw_phase = -mw_phase # EasySpin: mwPhase negated for frequency sweeps
|
|
1338
|
+
# EasySpin auto-harmonic: with no broadening at all (no lw, no strain) the
|
|
1339
|
+
# spectrum is a stick/projection absorption spectrum; a derivative of it
|
|
1340
|
+
# is not meaningful, so Harmonic falls back to 0 (EasySpin errors when the
|
|
1341
|
+
# harmonic was requested explicitly; torchspin cannot tell that apart).
|
|
1342
|
+
|
|
1343
|
+
# Check if we need per-line broadening (strain present)
|
|
1344
|
+
has_strain = (
|
|
1345
|
+
(sys.HStrain is not None and torch.any(sys.HStrain > 0))
|
|
1346
|
+
or (sys.gStrain is not None and torch.any(sys.gStrain > 0))
|
|
1347
|
+
or (sys.AStrain is not None and torch.any(sys.AStrain > 0))
|
|
1348
|
+
or (sys.DStrain is not None and torch.any(sys.DStrain > 0))
|
|
1349
|
+
)
|
|
1350
|
+
if not has_strain and lw_effective[0] == 0 and lw_effective[1] == 0 and exp.Harmonic > 0:
|
|
1351
|
+
if opt.Verbosity >= 1:
|
|
1352
|
+
print(f"pepper: no broadening given; using Harmonic=0 instead of {exp.Harmonic}")
|
|
1353
|
+
from dataclasses import replace as _dc_replace0
|
|
1354
|
+
exp = _dc_replace0(exp, Harmonic=0, CenterSweep=None)
|
|
1355
|
+
|
|
1356
|
+
# photoselection weights for the (χ-averaged) powder orientations, molecular frame
|
|
1357
|
+
photo_w_all = _photo_w(torch.stack([phi_arr, theta_arr], dim=1).detach().cpu().numpy()) if use_photo else None
|
|
1358
|
+
_pw = (lambda a, b: None) if photo_w_all is None else (lambda a, b: photo_w_all[a:b])
|
|
1359
|
+
# ordering: sample orientation R_L2S from Exp.SampleFrame / Exp.SampleRotation (single)
|
|
1360
|
+
_ordering = None
|
|
1361
|
+
if ordering_fun is not None:
|
|
1362
|
+
from torchspin.rotutils import rotaxi2mat as _rotaxi2mat
|
|
1363
|
+
sf_ = getattr(exp, 'SampleFrame', None)
|
|
1364
|
+
sf_ = np.zeros((1, 3)) if sf_ is None else np.atleast_2d(np.asarray(sf_, dtype=float))
|
|
1365
|
+
if sf_.shape[0] != 1:
|
|
1366
|
+
raise ValueError('For partially ordered samples, only a single sample orientation (Exp.SampleFrame) can be used.')
|
|
1367
|
+
R_L2S = erot(sf_[0].tolist()).numpy()
|
|
1368
|
+
rot_ = getattr(exp, 'SampleRotation', None)
|
|
1369
|
+
if rot_ is not None:
|
|
1370
|
+
axis_, rho_ = rot_
|
|
1371
|
+
if isinstance(axis_, str):
|
|
1372
|
+
axis_ = {'x': [1, 0, 0], 'y': [0, 1, 0], 'z': [0, 0, 1]}[axis_.lower()]
|
|
1373
|
+
rho_ = np.atleast_1d(np.asarray(rho_, dtype=float))
|
|
1374
|
+
if rho_.size != 1:
|
|
1375
|
+
raise ValueError('For partially ordered samples, only a single sample rotation can be used.')
|
|
1376
|
+
R_L2S = R_L2S @ _rotaxi2mat(np.asarray(axis_, dtype=float), float(rho_[0]))
|
|
1377
|
+
_ordering = (ordering_fun, R_L2S)
|
|
1378
|
+
|
|
1379
|
+
# EasySpin p_sampletype: a crystal needs Exp.MolFrame or Exp.CrystalSymmetry;
|
|
1380
|
+
# Exp.SampleFrame/SampleRotation alone describe a (rotated) disordered sample,
|
|
1381
|
+
# which is orientation-independent and therefore simulated as a plain powder.
|
|
1382
|
+
crystal_sample = (getattr(exp, 'CrystalSymmetry', None) is not None
|
|
1383
|
+
or getattr(exp, 'MolFrame', None) is not None)
|
|
1384
|
+
if crystal_sample:
|
|
1385
|
+
if sep_trans:
|
|
1386
|
+
raise ValueError("Cannot return separate transitions for crystal spectra (Opt.separate='transitions').")
|
|
1387
|
+
# ---------------------------------------------------------------
|
|
1388
|
+
# Single crystal (EasySpin pepper.m crystalSample branch): explicit
|
|
1389
|
+
# molecule→lab orientations for every sample orientation × site,
|
|
1390
|
+
# intensities with B1 along lab x (no χ average), each line
|
|
1391
|
+
# accumulated with an exact template, weight 2π·4π/(nSites·nOri).
|
|
1392
|
+
# ---------------------------------------------------------------
|
|
1393
|
+
from torchspin.sitetransforms import crystal_orientations, sitetransforms
|
|
1394
|
+
from torchspin.rotutils import eulang, rotaxi2mat
|
|
1395
|
+
sf = getattr(exp, 'SampleFrame', None)
|
|
1396
|
+
sample_frames = np.zeros((1, 3)) if sf is None else np.atleast_2d(np.asarray(sf, dtype=float))
|
|
1397
|
+
rot = getattr(exp, 'SampleRotation', None)
|
|
1398
|
+
if rot is not None:
|
|
1399
|
+
axis, rho = rot
|
|
1400
|
+
if isinstance(axis, str):
|
|
1401
|
+
axis = {'x': [1, 0, 0], 'y': [0, 1, 0], 'z': [0, 0, 1]}[axis.lower()]
|
|
1402
|
+
rhos = np.atleast_1d(np.asarray(rho, dtype=float))
|
|
1403
|
+
frames = []
|
|
1404
|
+
for sfrow in sample_frames:
|
|
1405
|
+
R_L2S0 = erot(sfrow.tolist()).numpy()
|
|
1406
|
+
for r_ in rhos:
|
|
1407
|
+
frames.append(eulang(R_L2S0 @ rotaxi2mat(np.asarray(axis, dtype=float), float(r_))))
|
|
1408
|
+
sample_frames = np.array(frames).reshape(-1, 3)
|
|
1409
|
+
csym = getattr(exp, 'CrystalSymmetry', None)
|
|
1410
|
+
n_sites = len(sitetransforms(csym if csym not in (None, '') else 'P1'))
|
|
1411
|
+
angles = crystal_orientations(sample_frames, csym, getattr(exp, 'MolFrame', None)) # (nSamples*nSites, 3)
|
|
1412
|
+
n_samples = sample_frames.shape[0]
|
|
1413
|
+
site_of = np.tile(np.arange(n_sites), n_samples)
|
|
1414
|
+
sample_of = np.repeat(np.arange(n_samples), n_sites)
|
|
1415
|
+
if opt.Sites is not None:
|
|
1416
|
+
keep_s = np.isin(site_of + 1, np.atleast_1d(opt.Sites))
|
|
1417
|
+
angles, site_of, sample_of = angles[keep_s], site_of[keep_s], sample_of[keep_s]
|
|
1418
|
+
n_sites_eff = int(np.unique(site_of).size)
|
|
1419
|
+
else:
|
|
1420
|
+
n_sites_eff = n_sites
|
|
1421
|
+
photo_w_c = _photo_w(angles) if use_photo else None # crystal: no χ average
|
|
1422
|
+
Rs = np.stack([erot(a.tolist()).numpy() for a in angles]) # (N, 3, 3); rows = xLab_M, yLab_M, zLab_M
|
|
1423
|
+
zlab = Rs[:, 2, :]; xlab = Rs[:, 0, :]
|
|
1424
|
+
theta_c = torch.tensor(np.arccos(np.clip(zlab[:, 2], -1, 1)), dtype=torch.float64)
|
|
1425
|
+
phi_c = torch.tensor(np.arctan2(zlab[:, 1], zlab[:, 0]), dtype=torch.float64)
|
|
1426
|
+
B_l, I_l, W_l = _resonances(H0, mux, muy, muz, phi_c.to(device), theta_c.to(device), exp, opt, sys,
|
|
1427
|
+
R_batch=torch.tensor(Rs, dtype=torch.float64), init_state=init_state, photo_weights=photo_w_c)
|
|
1428
|
+
_min_fwhm0 = dx / 100.0
|
|
1429
|
+
fwhm_g0, fwhm_l0 = lw_effective[0], lw_effective[1]
|
|
1430
|
+
n_ori_tot = angles.shape[0]
|
|
1431
|
+
sep = str(opt.separate or '')
|
|
1432
|
+
n_spec = n_samples if sep == 'orientations' else (n_sites_eff if sep == 'sites' else 1)
|
|
1433
|
+
specs = torch.zeros(n_spec, exp.nPoints, dtype=torch.float64)
|
|
1434
|
+
for k in range(n_ori_tot):
|
|
1435
|
+
B0 = B_l[k].detach().cpu().numpy(); I0 = I_l[k].detach().cpu().numpy()
|
|
1436
|
+
if B0.size == 0:
|
|
1437
|
+
continue
|
|
1438
|
+
W0 = W_l[k].detach().cpu().numpy() if (W_l[k] is not None and W_l[k].numel() == B0.size) else np.zeros_like(B0)
|
|
1439
|
+
amp0 = I0 * (4.0 * math.pi) / n_sites_eff / n_samples
|
|
1440
|
+
if has_strain:
|
|
1441
|
+
line = _gaussian_bins(x, B0, W0, amp0, _min_fwhm0)
|
|
1442
|
+
elif fwhm_g0 > 0:
|
|
1443
|
+
line = _gaussian_bins(x, B0, np.full(B0.size, fwhm_g0), amp0, _min_fwhm0)
|
|
1444
|
+
elif fwhm_l0 > 0:
|
|
1445
|
+
line = _lorentzian_bins(x, B0, np.full(B0.size, fwhm_l0), amp0, _min_fwhm0, phase=mw_phase)
|
|
1446
|
+
else:
|
|
1447
|
+
line = _gaussian_bins(x, B0, np.zeros(B0.size), amp0, _min_fwhm0)
|
|
1448
|
+
idx_s = sample_of[k] if sep == 'orientations' else (int(np.searchsorted(np.unique(site_of), site_of[k])) if sep == 'sites' else 0)
|
|
1449
|
+
specs[idx_s] += line
|
|
1450
|
+
specs = specs / dx * (2.0 * math.pi)
|
|
1451
|
+
if has_strain:
|
|
1452
|
+
g_rem, l_rem, deriv_fd = fwhm_g0, fwhm_l0, (fwhm_g0 == 0 and fwhm_l0 == 0)
|
|
1453
|
+
elif fwhm_g0 > 0:
|
|
1454
|
+
g_rem, l_rem, deriv_fd = 0.0, fwhm_l0, (fwhm_l0 == 0)
|
|
1455
|
+
else:
|
|
1456
|
+
g_rem, l_rem, deriv_fd = 0.0, 0.0, True
|
|
1457
|
+
spec = torch.stack(list(specs)) if n_spec > 1 else specs[0]
|
|
1458
|
+
# phase enters the Lorentzian template (above) or the Lorentzian convolution
|
|
1459
|
+
spec = _finish_field_sweep(spec, x, dx, exp, g_rem, l_rem, deriv_fd,
|
|
1460
|
+
phase=mw_phase if l_rem > 0 else 0.0)
|
|
1461
|
+
return x, spec
|
|
1462
|
+
|
|
1463
|
+
if grid_sym == 'O3':
|
|
1464
|
+
# Isotropic system, one orientation (EasySpin ~anisotropicSpectrum):
|
|
1465
|
+
# every line is accumulated with an exact template at its exact
|
|
1466
|
+
# position — Gaussian (strain width, or Sys.lw(1)) or Lorentzian
|
|
1467
|
+
# (Sys.lw(2) when there is no Gaussian part) — instead of stick binning.
|
|
1468
|
+
_min_fwhm0 = dx / 100.0
|
|
1469
|
+
use_pt = (sys.nElectrons == 1 and not has_strain
|
|
1470
|
+
and not (sys.S[0] > 0.5 and sys.D is not None)
|
|
1471
|
+
and _pt_ok and str(opt.Method).startswith('perturb'))
|
|
1472
|
+
if use_pt:
|
|
1473
|
+
Bl, Il = resfields_perturb_batch(sys, phi_arr, theta_arr, exp_search, opt, return_full=True, photo_weights=_pw(0, n_orient))
|
|
1474
|
+
B0 = Bl[0].detach().cpu().numpy(); I0 = Il[0].detach().cpu().numpy()
|
|
1475
|
+
keep = np.isfinite(B0) & (I0 != 0)
|
|
1476
|
+
B0, I0 = B0[keep], I0[keep]; W0 = np.zeros_like(B0)
|
|
1477
|
+
else:
|
|
1478
|
+
Bl, Il, Wl = _resonances(H0, mux, muy, muz, phi_arr, theta_arr, exp, opt, sys, init_state=init_state, photo_weights=_pw(0, n_orient))
|
|
1479
|
+
B0 = Bl[0].detach().cpu().numpy(); I0 = Il[0].detach().cpu().numpy()
|
|
1480
|
+
W0 = Wl[0].detach().cpu().numpy() if (Wl[0] is not None and Wl[0].numel() == B0.size) else np.zeros_like(B0)
|
|
1481
|
+
amp0 = I0 * float(weights_arr[0]) # 4π
|
|
1482
|
+
fwhm_g0, fwhm_l0 = lw_effective[0], lw_effective[1]
|
|
1483
|
+
def _o3_template(sel):
|
|
1484
|
+
if has_strain:
|
|
1485
|
+
return _gaussian_bins(x, B0[sel], W0[sel], amp0[sel], _min_fwhm0)
|
|
1486
|
+
if fwhm_g0 > 0:
|
|
1487
|
+
return _gaussian_bins(x, B0[sel], np.full(B0[sel].size, fwhm_g0), amp0[sel], _min_fwhm0)
|
|
1488
|
+
if fwhm_l0 > 0:
|
|
1489
|
+
# EasySpin: Lorentzian template built with Exp.mwPhase
|
|
1490
|
+
return _lorentzian_bins(x, B0[sel], np.full(B0[sel].size, fwhm_l0), amp0[sel], _min_fwhm0, phase=mw_phase)
|
|
1491
|
+
return _gaussian_bins(x, B0[sel], np.zeros(B0[sel].size), amp0[sel], _min_fwhm0) # delta-like sticks
|
|
1492
|
+
if has_strain:
|
|
1493
|
+
g_rem, l_rem, deriv_fd = fwhm_g0, fwhm_l0, (fwhm_g0 == 0 and fwhm_l0 == 0)
|
|
1494
|
+
elif fwhm_g0 > 0:
|
|
1495
|
+
g_rem, l_rem, deriv_fd = 0.0, fwhm_l0, (fwhm_l0 == 0)
|
|
1496
|
+
else:
|
|
1497
|
+
g_rem, l_rem, deriv_fd = 0.0, 0.0, True
|
|
1498
|
+
if sep_trans:
|
|
1499
|
+
# one row per line: matrix path by descending intensity (EasySpin
|
|
1500
|
+
# transition pre-selection), perturbation path in line order
|
|
1501
|
+
order = np.arange(B0.size) if use_pt else np.argsort(-np.abs(I0), kind='stable')
|
|
1502
|
+
spec = torch.stack([_o3_template(slice(int(k_), int(k_) + 1)) for k_ in order]) if B0.size else spec.unsqueeze(0)
|
|
1503
|
+
else:
|
|
1504
|
+
spec = _o3_template(slice(None))
|
|
1505
|
+
spec = spec / dx * (2.0 * math.pi)
|
|
1506
|
+
# EasySpin disregards mwPhase for a Gaussian-only template spectrum
|
|
1507
|
+
spec = _finish_field_sweep(spec, x, dx, exp, g_rem, l_rem, deriv_fd,
|
|
1508
|
+
phase=mw_phase if l_rem > 0 else 0.0)
|
|
1509
|
+
return x, spec
|
|
1510
|
+
|
|
1511
|
+
# Strain on a projectable grid (D2h-type or Dinfh): EasySpin pepper.m's
|
|
1512
|
+
# "doSummation" — per transition slot the positions, intensities and strain
|
|
1513
|
+
# widths are (optionally) interpolated to the fine grid, then every facet
|
|
1514
|
+
# (triangle, or meridian segment for Dinfh) contributes one Gaussian with
|
|
1515
|
+
# the facet's mean width, inflated by the Lambda smoothing when the line is
|
|
1516
|
+
# narrower than the facet's spread on the field axis. The global linewidth
|
|
1517
|
+
# (Sys.lw) is applied afterwards by convolution, as in EasySpin.
|
|
1518
|
+
# O3 (isotropic, one orientation) has no facets: it uses the stick / per-line path.
|
|
1519
|
+
use_summation = has_strain and grid_sym != 'O3'
|
|
1520
|
+
|
|
1521
|
+
_inv_sqrt_2pi = 1.0 / math.sqrt(2.0 * math.pi)
|
|
1522
|
+
_sqrt_8log2 = math.sqrt(8.0 * math.log(2.0))
|
|
1523
|
+
_min_fwhm = dx / 100.0 # matches MATLAB lisum1i.c lower-bound clipping
|
|
1524
|
+
|
|
1525
|
+
# True when the SOPHE projection produced a spectral density already
|
|
1526
|
+
# (no 1/ΔB conversion needed at the end).
|
|
1527
|
+
used_projection = False
|
|
1528
|
+
|
|
1529
|
+
if use_summation:
|
|
1530
|
+
all_B_res = [None] * n_orient
|
|
1531
|
+
all_intens = [None] * n_orient
|
|
1532
|
+
all_widths = [None] * n_orient
|
|
1533
|
+
all_pairs = [None] * n_orient
|
|
1534
|
+
batch_size = auto_batch_size(opt.BatchSize, H0.shape[0], n_orient)
|
|
1535
|
+
for batch_idx in range((n_orient + batch_size - 1) // batch_size):
|
|
1536
|
+
s = batch_idx * batch_size
|
|
1537
|
+
e = min(s + batch_size, n_orient)
|
|
1538
|
+
B_list, I_list, W_list, P_list = _resonances(
|
|
1539
|
+
H0, mux, muy, muz, phi_arr[s:e], theta_arr[s:e], exp, opt, sys,
|
|
1540
|
+
return_pairs=True, init_state=init_state, photo_weights=_pw(s, e), opt_override=_opt_rf, pairs=_pairs_sel
|
|
1541
|
+
)
|
|
1542
|
+
for j, (B_res, intens, widths, prs) in enumerate(zip(B_list, I_list, W_list, P_list)):
|
|
1543
|
+
all_B_res[s + j] = B_res
|
|
1544
|
+
all_intens[s + j] = intens
|
|
1545
|
+
all_widths[s + j] = widths
|
|
1546
|
+
all_pairs[s + j] = prs
|
|
1547
|
+
slot_B, slot_I, slot_W = _transition_slots(all_B_res, all_intens, all_widths, all_pairs, n_orient)
|
|
1548
|
+
slot_B, slot_I, slot_W = _filter_slots(slot_B, slot_I, slot_W, float(opt.Threshold))
|
|
1549
|
+
_any_nan = any(np.isnan(b_).any() for b_ in slot_B)
|
|
1550
|
+
if slot_W is None:
|
|
1551
|
+
slot_W = [np.zeros(n_orient) for _ in slot_B]
|
|
1552
|
+
|
|
1553
|
+
do_interp_s = (N_interp > 1) and (nfKnots != N_coarse) and (n_orient > 1)
|
|
1554
|
+
if do_interp_s:
|
|
1555
|
+
_, theta_fine_s, _, vecs_fine_s = sphgrid(grid_sym, nfKnots)
|
|
1556
|
+
vecs_sym_f_s = vecs_fine_s if vecs_fine_s.shape[0] == 3 else vecs_fine_s.T
|
|
1557
|
+
theta_final = theta_fine_s
|
|
1558
|
+
vecs_final = vecs_fine_s
|
|
1559
|
+
N_final = nfKnots
|
|
1560
|
+
else:
|
|
1561
|
+
theta_final = theta_sym
|
|
1562
|
+
vecs_final = vecs_orig
|
|
1563
|
+
N_final = N_coarse
|
|
1564
|
+
tri_idx, facet_w = _facets(grid_sym, N_final, theta_final, ordering=_ordering)
|
|
1565
|
+
c1, c2 = (1.57246, 18.6348) if tri_idx is None else (2.8269, 42.6843)
|
|
1566
|
+
smoothing = 2.0 # EasySpin Opt.Smoothing default
|
|
1567
|
+
|
|
1568
|
+
for i_t in range(len(slot_B)):
|
|
1569
|
+
B_c, I_c, W_c = slot_B[i_t], slot_I[i_t], slot_W[i_t]
|
|
1570
|
+
valid = ~np.isnan(B_c)
|
|
1571
|
+
if do_interp_s:
|
|
1572
|
+
B_f = _interp_slot(B_c, grid_sym, N_coarse, nfKnots, vecs_sym_c, vecs_sym_f_s, any_nan=_any_nan)
|
|
1573
|
+
if B_f is None:
|
|
1574
|
+
continue
|
|
1575
|
+
I_f = _interp_slot(I_c, grid_sym, N_coarse, nfKnots, vecs_sym_c, vecs_sym_f_s, mode='val')
|
|
1576
|
+
W_f = _interp_slot(W_c, grid_sym, N_coarse, nfKnots, vecs_sym_c, vecs_sym_f_s, mode='val')
|
|
1577
|
+
else:
|
|
1578
|
+
B_f, I_f, W_f = B_c, I_c, W_c
|
|
1579
|
+
if tri_idx is None:
|
|
1580
|
+
fPosC = 0.5 * (B_f[:-1] + B_f[1:])
|
|
1581
|
+
fSpread = np.abs(B_f[1:] - B_f[:-1])
|
|
1582
|
+
fIntC = facet_w * 0.5 * (I_f[:-1] + I_f[1:])
|
|
1583
|
+
fWidM = 0.5 * (W_f[:-1] + W_f[1:])
|
|
1584
|
+
else:
|
|
1585
|
+
Bv = B_f[tri_idx] # (nTri, 3)
|
|
1586
|
+
fPosC = Bv.mean(axis=1)
|
|
1587
|
+
fSpread = Bv.max(axis=1) - Bv.min(axis=1)
|
|
1588
|
+
fIntC = facet_w * I_f[tri_idx].mean(axis=1)
|
|
1589
|
+
fWidM = W_f[tri_idx].mean(axis=1)
|
|
1590
|
+
ok = np.isfinite(fPosC) & np.isfinite(fIntC) & np.isfinite(fWidM) & (fIntC > 0)
|
|
1591
|
+
fPosC, fSpread, fIntC, fWidM = fPosC[ok], fSpread[ok], fIntC[ok], fWidM[ok]
|
|
1592
|
+
with np.errstate(divide='ignore', invalid='ignore'):
|
|
1593
|
+
Lambda = np.where(fSpread > 0, fWidM / np.where(fSpread > 0, fSpread, 1.0), np.inf)
|
|
1594
|
+
gam = 1.0 / np.sqrt(c1 * Lambda ** 2 + c2 * Lambda ** 4)
|
|
1595
|
+
gam[~np.isfinite(gam)] = 0.0
|
|
1596
|
+
fWidC = fWidM * (1.0 + smoothing * gam)
|
|
1597
|
+
_acc(i_t, _gaussian_bins(x, fPosC, fWidC, fIntC, _min_fwhm), np.nanmax(np.abs(I_c)) if np.isfinite(I_c).any() else 0.0)
|
|
1598
|
+
if opt.Verbosity >= 1:
|
|
1599
|
+
print(f"pepper: strain summation over {len(facet_w)} facets, {len(slot_B)} transition slots"
|
|
1600
|
+
f"{' (interpolated)' if do_interp_s else ''}")
|
|
1601
|
+
|
|
1602
|
+
else:
|
|
1603
|
+
# ---------------------------------------------------------------
|
|
1604
|
+
# Standard path: per-orientation accumulation, with optional
|
|
1605
|
+
# spherical interpolation from coarse to fine grid (N_interp > 1).
|
|
1606
|
+
# ---------------------------------------------------------------
|
|
1607
|
+
batch_size = auto_batch_size(opt.BatchSize, H0.shape[0], n_orient)
|
|
1608
|
+
n_batches = (n_orient + batch_size - 1) // batch_size
|
|
1609
|
+
|
|
1610
|
+
# Isotropic systems (O3, a single orientation) have nothing to interpolate
|
|
1611
|
+
do_interp = (N_interp > 1) and (nfKnots != N_coarse) and (n_orient > 1)
|
|
1612
|
+
|
|
1613
|
+
if do_interp:
|
|
1614
|
+
# -----------------------------------------------------------
|
|
1615
|
+
# Interpolation path: collect coarse data, interpolate to
|
|
1616
|
+
# fine grid (matching MATLAB's GridSize = [N_coarse, N_interp]).
|
|
1617
|
+
# -----------------------------------------------------------
|
|
1618
|
+
all_B_res = [None] * n_orient
|
|
1619
|
+
all_intens = [None] * n_orient
|
|
1620
|
+
all_widths = [None] * n_orient
|
|
1621
|
+
all_pairs = [None] * n_orient
|
|
1622
|
+
|
|
1623
|
+
# Use perturbation theory for single-electron systems without strain.
|
|
1624
|
+
# Falls back to matrix diagonalization for high-spin + ZFS:
|
|
1625
|
+
# perturbation theory cannot reproduce ZFS-split energy levels regardless
|
|
1626
|
+
# of temperature (kT comparison is not the only issue — the ZFS mixes
|
|
1627
|
+
# states and shifts resonance fields in ways perturbation theory misses).
|
|
1628
|
+
# S=1/2 + temperature is fine: Zeeman-only Boltzmann is exact for S=1/2.
|
|
1629
|
+
_needs_exact_boltz = (sys.S[0] > 0.5 and sys.D is not None)
|
|
1630
|
+
_use_perturb = (sys.nElectrons == 1
|
|
1631
|
+
and not has_strain
|
|
1632
|
+
and not _needs_exact_boltz
|
|
1633
|
+
and _pt_ok and str(opt.Method).startswith('perturb'))
|
|
1634
|
+
if _use_perturb:
|
|
1635
|
+
B_list, I_list = resfields_perturb_batch(
|
|
1636
|
+
sys, phi_arr, theta_arr, exp_search, opt, return_full=True, photo_weights=_pw(0, n_orient)
|
|
1637
|
+
)
|
|
1638
|
+
all_B_res = B_list
|
|
1639
|
+
all_intens = I_list
|
|
1640
|
+
all_widths = [None] * n_orient
|
|
1641
|
+
if opt.Verbosity >= 1:
|
|
1642
|
+
print(f"pepper: coarse grid {n_orient} orientations (perturb)")
|
|
1643
|
+
else:
|
|
1644
|
+
for batch_idx in range(n_batches):
|
|
1645
|
+
s = batch_idx * batch_size
|
|
1646
|
+
e = min(s + batch_size, n_orient)
|
|
1647
|
+
B_list, I_list, W_list, P_list = _resonances(
|
|
1648
|
+
H0, mux, muy, muz, phi_arr[s:e], theta_arr[s:e], exp, opt, sys,
|
|
1649
|
+
return_pairs=True, init_state=init_state, photo_weights=_pw(s, e), opt_override=_opt_rf, pairs=_pairs_sel
|
|
1650
|
+
)
|
|
1651
|
+
for j, (B_res, intens, widths, prs) in enumerate(zip(B_list, I_list, W_list, P_list)):
|
|
1652
|
+
all_B_res[s + j] = B_res
|
|
1653
|
+
all_intens[s + j] = intens
|
|
1654
|
+
all_widths[s + j] = widths
|
|
1655
|
+
all_pairs[s + j] = prs
|
|
1656
|
+
if opt.Verbosity >= 1:
|
|
1657
|
+
print(f"pepper: coarse grid {e}/{n_orient} orientations computed")
|
|
1658
|
+
|
|
1659
|
+
# Build the fine grid
|
|
1660
|
+
phi_fine, theta_fine, weights_fine, vecs_fine_orig = sphgrid(grid_sym, nfKnots)
|
|
1661
|
+
n_orient_fine = phi_fine.shape[0]
|
|
1662
|
+
total_weight = weights_fine.sum().item()
|
|
1663
|
+
theta_fine_sym = theta_fine.clone()
|
|
1664
|
+
vecs_sym_f = vecs_fine_orig if vecs_fine_orig.shape[0] == 3 else vecs_fine_orig.T
|
|
1665
|
+
|
|
1666
|
+
# Apply frame rotation to fine grid if needed
|
|
1667
|
+
if not torch.allclose(R_sym2mol, torch.eye(3, dtype=torch.float64), atol=1e-10):
|
|
1668
|
+
sin_t_f = torch.sin(theta_fine)
|
|
1669
|
+
v_sym_f = torch.stack([
|
|
1670
|
+
sin_t_f * torch.cos(phi_fine),
|
|
1671
|
+
sin_t_f * torch.sin(phi_fine),
|
|
1672
|
+
torch.cos(theta_fine),
|
|
1673
|
+
], dim=1)
|
|
1674
|
+
v_mol_f = (R_sym2mol.T @ v_sym_f.T).T
|
|
1675
|
+
theta_fine = torch.acos(v_mol_f[:, 2].clamp(-1.0, 1.0))
|
|
1676
|
+
phi_fine = torch.atan2(v_mol_f[:, 1], v_mol_f[:, 0])
|
|
1677
|
+
|
|
1678
|
+
if opt.Verbosity >= 1:
|
|
1679
|
+
print(f"pepper: interpolating to fine grid ({n_orient_fine} orientations)")
|
|
1680
|
+
|
|
1681
|
+
# Coarse unit vectors (already rotated to mol frame from phi_arr/theta_arr)
|
|
1682
|
+
sin_tc = torch.sin(theta_arr)
|
|
1683
|
+
vecs_c = torch.stack([
|
|
1684
|
+
sin_tc * torch.cos(phi_arr),
|
|
1685
|
+
sin_tc * torch.sin(phi_arr),
|
|
1686
|
+
torch.cos(theta_arr),
|
|
1687
|
+
], dim=0) # (3, n_orient)
|
|
1688
|
+
|
|
1689
|
+
# Fine unit vectors
|
|
1690
|
+
sin_tf = torch.sin(theta_fine)
|
|
1691
|
+
vecs_f = torch.stack([
|
|
1692
|
+
sin_tf * torch.cos(phi_fine),
|
|
1693
|
+
sin_tf * torch.sin(phi_fine),
|
|
1694
|
+
torch.cos(theta_fine),
|
|
1695
|
+
], dim=0) # (3, n_orient_fine)
|
|
1696
|
+
|
|
1697
|
+
# Determine consistent transition count
|
|
1698
|
+
n_trans_list = [b.numel() if b is not None else 0 for b in all_B_res]
|
|
1699
|
+
n_trans = max(n_trans_list) if n_trans_list else 0
|
|
1700
|
+
|
|
1701
|
+
if n_trans > 0:
|
|
1702
|
+
# SOPHE projection: for D2h/Dinfh grids without strain, use
|
|
1703
|
+
# projecttriangles / projectzones instead of per-orientation
|
|
1704
|
+
# stick binning. This eliminates quantization artifacts at
|
|
1705
|
+
# small linewidths (the root cause of the oscillation problem).
|
|
1706
|
+
x_np = x.numpy()
|
|
1707
|
+
# SOPHE projection for every grid symmetry (EasySpin: zones for
|
|
1708
|
+
# Dinfh, triangles otherwise — including the Delaunay-triangulated
|
|
1709
|
+
# open-φ grids Ci/C2h/C1).
|
|
1710
|
+
use_proj = (not has_strain) and grid_sym != 'O3'
|
|
1711
|
+
if use_proj:
|
|
1712
|
+
_tri_f, _areas_f = _facets(grid_sym, nfKnots, theta_fine_sym, ordering=_ordering)
|
|
1713
|
+
_seg_wts = _areas_f
|
|
1714
|
+
|
|
1715
|
+
slot_B, slot_I, slot_W = _transition_slots(
|
|
1716
|
+
all_B_res, all_intens, all_widths, all_pairs, n_orient
|
|
1717
|
+
)
|
|
1718
|
+
slot_B, slot_I, slot_W = _filter_slots(slot_B, slot_I, slot_W, float(opt.Threshold))
|
|
1719
|
+
_any_nan = any(np.isnan(b_).any() for b_ in slot_B)
|
|
1720
|
+
for i_t in range(len(slot_B)):
|
|
1721
|
+
B_c = slot_B[i_t]
|
|
1722
|
+
I_c = slot_I[i_t]
|
|
1723
|
+
valid = ~np.isnan(B_c)
|
|
1724
|
+
B_f = _interp_slot(B_c, grid_sym, N_coarse, nfKnots, vecs_sym_c, vecs_sym_f, any_nan=_any_nan)
|
|
1725
|
+
if B_f is None:
|
|
1726
|
+
continue
|
|
1727
|
+
I_f = _interp_slot(I_c, grid_sym, N_coarse, nfKnots, vecs_sym_c, vecs_sym_f, mode='val')
|
|
1728
|
+
|
|
1729
|
+
if use_proj:
|
|
1730
|
+
# SOPHE triangle/zone projection: each facet contributes a
|
|
1731
|
+
# tent/rect density → smooth spectrum without quantization noise.
|
|
1732
|
+
if grid_sym == 'Dinfh':
|
|
1733
|
+
contrib = _projectzones(B_f, I_f, _seg_wts, x_np)
|
|
1734
|
+
else:
|
|
1735
|
+
contrib = _projecttriangles(_tri_f, _areas_f, B_f, I_f, x_np)
|
|
1736
|
+
_acc(i_t, torch.from_numpy(contrib * (2.0 * math.pi)), np.nanmax(np.abs(I_c)) if np.isfinite(I_c).any() else 0.0)
|
|
1737
|
+
used_projection = True
|
|
1738
|
+
else:
|
|
1739
|
+
# Fallback: per-orientation accumulation.
|
|
1740
|
+
# Used for has_strain (spatial Gaussians) or grids without a
|
|
1741
|
+
# standard triangulation (e.g. Ci symmetry → stick binning).
|
|
1742
|
+
if has_strain and slot_W is not None:
|
|
1743
|
+
W_c = slot_W[i_t]
|
|
1744
|
+
W_f = _interp_slot(W_c, grid_sym, N_coarse, nfKnots, vecs_sym_c, vecs_sym_f, mode='val')
|
|
1745
|
+
else:
|
|
1746
|
+
W_f = np.zeros(n_orient_fine)
|
|
1747
|
+
|
|
1748
|
+
if not has_strain:
|
|
1749
|
+
# Vectorized stick binning for non-projectable no-strain grids
|
|
1750
|
+
# (e.g. Ci symmetry). Replaces O(n_fine) Python loop.
|
|
1751
|
+
_vld = (~np.isnan(B_f)) & (~np.isnan(I_f)) & (I_f > 0.0)
|
|
1752
|
+
if _vld.any():
|
|
1753
|
+
_Bv = torch.from_numpy(B_f[_vld])
|
|
1754
|
+
_Iv = torch.from_numpy(I_f[_vld]).to(spec.dtype)
|
|
1755
|
+
_wv = weights_fine[_vld].to(spec.dtype)
|
|
1756
|
+
_frac = (_Bv - exp.Range[0]) / dx
|
|
1757
|
+
_ilo = _frac.floor().long()
|
|
1758
|
+
_ihi = _ilo + 1
|
|
1759
|
+
_a = (_frac - _ilo.to(_frac.dtype)).to(spec.dtype)
|
|
1760
|
+
_amp = _wv * _Iv
|
|
1761
|
+
_mlo = (_ilo >= 0) & (_ilo < exp.nPoints)
|
|
1762
|
+
_mhi = (_ihi >= 0) & (_ihi < exp.nPoints)
|
|
1763
|
+
if _mlo.any():
|
|
1764
|
+
(_row(i_t) if sep_trans else spec).scatter_add_(0, _ilo[_mlo],
|
|
1765
|
+
_amp[_mlo] * (1.0 - _a[_mlo]))
|
|
1766
|
+
if _mhi.any():
|
|
1767
|
+
(_row(i_t) if sep_trans else spec).scatter_add_(0, _ihi[_mhi],
|
|
1768
|
+
_amp[_mhi] * _a[_mhi])
|
|
1769
|
+
else:
|
|
1770
|
+
for j in range(n_orient_fine):
|
|
1771
|
+
Bj = B_f[j]
|
|
1772
|
+
Ij = I_f[j]
|
|
1773
|
+
Wj = W_f[j]
|
|
1774
|
+
wj = weights_fine[j].item()
|
|
1775
|
+
|
|
1776
|
+
if np.isnan(Bj) or np.isnan(Ij) or Ij <= 0.0:
|
|
1777
|
+
continue
|
|
1778
|
+
|
|
1779
|
+
fwhm_strain = max(float(Wj), _min_fwhm)
|
|
1780
|
+
fwhm_g_total = fwhm_strain # global lw is convolved afterwards (EasySpin)
|
|
1781
|
+
fwhm_l = 0.0
|
|
1782
|
+
if fwhm_l == 0.0:
|
|
1783
|
+
if fwhm_g_total > 0.0:
|
|
1784
|
+
sigma = fwhm_g_total / _sqrt_8log2
|
|
1785
|
+
peak_amp = float(Ij) * dx * _inv_sqrt_2pi / sigma
|
|
1786
|
+
spec = spec + wj * peak_amp * torch.exp(
|
|
1787
|
+
-0.5 * ((x - Bj) / sigma) ** 2
|
|
1788
|
+
)
|
|
1789
|
+
else:
|
|
1790
|
+
frac = (Bj - exp.Range[0]) / dx
|
|
1791
|
+
ilo = int(frac)
|
|
1792
|
+
ihi = ilo + 1
|
|
1793
|
+
a = frac - ilo
|
|
1794
|
+
if 0 <= ilo < exp.nPoints:
|
|
1795
|
+
spec[ilo] = spec[ilo] + wj * float(Ij) * (1.0 - a)
|
|
1796
|
+
if 0 <= ihi < exp.nPoints:
|
|
1797
|
+
spec[ihi] = spec[ihi] + wj * float(Ij) * a
|
|
1798
|
+
else:
|
|
1799
|
+
Bt = torch.tensor([Bj], dtype=torch.float64)
|
|
1800
|
+
It = torch.tensor([float(Ij)], dtype=torch.float64)
|
|
1801
|
+
_, line_i = makespec(exp.Range, exp.nPoints, Bt, It)
|
|
1802
|
+
line_i = convspec(line_i, dx,
|
|
1803
|
+
fwhm_g=fwhm_g_total, fwhm_l=fwhm_l, deriv=0)
|
|
1804
|
+
_acc(i_t, wj * line_i)
|
|
1805
|
+
|
|
1806
|
+
else:
|
|
1807
|
+
# -----------------------------------------------------------
|
|
1808
|
+
# Direct path: no interpolation — accumulate at coarse grid.
|
|
1809
|
+
# For D2h/Dinfh without strain: two-phase collect → project.
|
|
1810
|
+
# For other grids or has_strain: existing per-orientation path.
|
|
1811
|
+
# -----------------------------------------------------------
|
|
1812
|
+
use_proj_d = (not has_strain) and grid_sym != 'O3'
|
|
1813
|
+
|
|
1814
|
+
if use_proj_d:
|
|
1815
|
+
# Phase 1: collect B_res and intensities for every orientation
|
|
1816
|
+
all_B_d = [None] * n_orient
|
|
1817
|
+
all_I_d = [None] * n_orient
|
|
1818
|
+
all_P_d = [None] * n_orient
|
|
1819
|
+
_needs_exact_boltz_d = (sys.S[0] > 0.5 and sys.D is not None)
|
|
1820
|
+
if (sys.nElectrons == 1
|
|
1821
|
+
and not has_strain
|
|
1822
|
+
and not _needs_exact_boltz_d
|
|
1823
|
+
and _pt_ok and _pt_ok and str(opt.Method).startswith('perturb')):
|
|
1824
|
+
all_B_d, all_I_d = resfields_perturb_batch(
|
|
1825
|
+
sys, phi_arr, theta_arr, exp_search, opt, return_full=True, photo_weights=_pw(0, n_orient)
|
|
1826
|
+
)
|
|
1827
|
+
else:
|
|
1828
|
+
for batch_idx in range(n_batches):
|
|
1829
|
+
s = batch_idx * batch_size
|
|
1830
|
+
e = min(s + batch_size, n_orient)
|
|
1831
|
+
B_list, I_list, W_list, P_list = _resonances(
|
|
1832
|
+
H0, mux, muy, muz, phi_arr[s:e], theta_arr[s:e], exp, opt, sys,
|
|
1833
|
+
return_pairs=True, init_state=init_state, photo_weights=_pw(s, e), opt_override=_opt_rf, pairs=_pairs_sel
|
|
1834
|
+
)
|
|
1835
|
+
for j, (B_res, intens, _, prs) in enumerate(zip(B_list, I_list, W_list, P_list)):
|
|
1836
|
+
all_B_d[s + j] = B_res
|
|
1837
|
+
all_I_d[s + j] = intens
|
|
1838
|
+
all_P_d[s + j] = prs
|
|
1839
|
+
|
|
1840
|
+
# Phase 2: project per-transition onto the spectrum
|
|
1841
|
+
x_np_d = x.numpy()
|
|
1842
|
+
n_td_list = [b.numel() if b is not None else 0 for b in all_B_d]
|
|
1843
|
+
n_td = max(n_td_list) if n_td_list else 0
|
|
1844
|
+
|
|
1845
|
+
_tri_c, _areas_c = _facets(grid_sym, N_coarse, theta_sym, ordering=_ordering)
|
|
1846
|
+
_seg_wts_d = _areas_c
|
|
1847
|
+
|
|
1848
|
+
slot_B_d, slot_I_d, _ = _transition_slots(all_B_d, all_I_d, None, all_P_d, n_orient)
|
|
1849
|
+
slot_B_d, slot_I_d, _ = _filter_slots(slot_B_d, slot_I_d, None, float(opt.Threshold))
|
|
1850
|
+
for i_t in range(len(slot_B_d)):
|
|
1851
|
+
B_t = slot_B_d[i_t]
|
|
1852
|
+
I_t = slot_I_d[i_t]
|
|
1853
|
+
if grid_sym == 'Dinfh':
|
|
1854
|
+
contrib = _projectzones(B_t, I_t, _seg_wts_d, x_np_d)
|
|
1855
|
+
else:
|
|
1856
|
+
contrib = _projecttriangles(_tri_c, _areas_c, B_t, I_t, x_np_d)
|
|
1857
|
+
_acc(i_t, torch.from_numpy(contrib * (2.0 * math.pi)), np.nanmax(np.abs(I_t)) if np.isfinite(I_t).any() else 0.0)
|
|
1858
|
+
|
|
1859
|
+
used_projection = True
|
|
1860
|
+
|
|
1861
|
+
if opt.Verbosity >= 1:
|
|
1862
|
+
print(f"pepper: direct projection over {n_orient} orientations, {n_td} transitions")
|
|
1863
|
+
|
|
1864
|
+
else:
|
|
1865
|
+
# Per-orientation accumulation (Ci/C2h/C1 grids, or has_strain).
|
|
1866
|
+
# Pre-collect (B_res, intens, widths, w_k) for all orientations.
|
|
1867
|
+
# For single-electron no-strain systems, use perturbation theory
|
|
1868
|
+
# (one vectorized call) instead of batched matrix eigh.
|
|
1869
|
+
_needs_exact_boltz_p = (sys.S[0] > 0.5 and sys.D is not None)
|
|
1870
|
+
_use_perturb_plain = (sys.nElectrons == 1
|
|
1871
|
+
and not has_strain
|
|
1872
|
+
and not _needs_exact_boltz_p
|
|
1873
|
+
and _pt_ok and str(opt.Method).startswith('perturb'))
|
|
1874
|
+
if _use_perturb_plain:
|
|
1875
|
+
# Vectorized accumulation: all orientations × transitions in one pass.
|
|
1876
|
+
# resfields_perturb_batch returns (M, n_slots) tensors (return_full=True)
|
|
1877
|
+
# with NaN for out-of-range slots. We scatter all valid sticks at once.
|
|
1878
|
+
_B_full, _I_full = resfields_perturb_batch(
|
|
1879
|
+
sys, phi_arr, theta_arr, exp_search, opt, return_full=True, photo_weights=_pw(0, n_orient)
|
|
1880
|
+
) # each is list of (n_slots,) tensors, length M
|
|
1881
|
+
_B_mat = torch.stack(_B_full, dim=0) # (M, n_slots)
|
|
1882
|
+
_I_mat = torch.stack(_I_full, dim=0) # (M, n_slots)
|
|
1883
|
+
_W_col = weights_arr.unsqueeze(1) # (M, 1)
|
|
1884
|
+
_WI = (_W_col * _I_mat).reshape(-1) # (M * n_slots,)
|
|
1885
|
+
if sep_trans:
|
|
1886
|
+
for _k in range(_I_mat.shape[1]):
|
|
1887
|
+
_Bk = _B_mat[:, _k]; _WIk = (_W_col[:, 0] * _I_mat[:, _k])
|
|
1888
|
+
_vk = ~torch.isnan(_Bk) & (_WIk != 0)
|
|
1889
|
+
if _vk.any():
|
|
1890
|
+
_bk = ((_Bk[_vk] - exp.Range[0]) / dx).floor().long()
|
|
1891
|
+
_ink = (_bk >= 0) & (_bk < exp.nPoints)
|
|
1892
|
+
_row(_k).scatter_add_(0, _bk[_ink].cpu(), _WIk[_vk][_ink].cpu().to(torch.float64))
|
|
1893
|
+
else:
|
|
1894
|
+
_B_flat = _B_mat.reshape(-1) # (M * n_slots,)
|
|
1895
|
+
# Bin valid (non-NaN, in-range) transitions
|
|
1896
|
+
_valid = ~torch.isnan(_B_flat) & (_WI != 0)
|
|
1897
|
+
if _valid.any():
|
|
1898
|
+
_B_v = _B_flat[_valid]
|
|
1899
|
+
_W_v = _WI[_valid]
|
|
1900
|
+
_bin = ((_B_v - exp.Range[0]) / dx).floor().long()
|
|
1901
|
+
_in = (_bin >= 0) & (_bin < exp.nPoints) # drop out-of-window lines
|
|
1902
|
+
_bin = _bin[_in]
|
|
1903
|
+
# scatter on whichever device the data lives, then move back to CPU
|
|
1904
|
+
spec_dev = spec.to(_bin.device)
|
|
1905
|
+
spec = spec_dev.scatter_add(0, _bin, _W_v[_in].to(spec_dev.dtype)).cpu()
|
|
1906
|
+
used_projection = False # will still apply normalization below
|
|
1907
|
+
else:
|
|
1908
|
+
_all_orient = []
|
|
1909
|
+
for batch_idx in range(n_batches):
|
|
1910
|
+
start_idx = batch_idx * batch_size
|
|
1911
|
+
end_idx = min(start_idx + batch_size, n_orient)
|
|
1912
|
+
phi_b = phi_arr[start_idx:end_idx]
|
|
1913
|
+
theta_b = theta_arr[start_idx:end_idx]
|
|
1914
|
+
weights_b = weights_arr[start_idx:end_idx]
|
|
1915
|
+
B_res_list, intens_list, widths_list = _resonances(
|
|
1916
|
+
H0, mux, muy, muz, phi_b, theta_b, exp, opt, sys, init_state=init_state, photo_weights=_pw(start_idx, end_idx)
|
|
1917
|
+
)
|
|
1918
|
+
for k_in_batch, (B_res, intens, widths) in enumerate(
|
|
1919
|
+
zip(B_res_list, intens_list, widths_list)
|
|
1920
|
+
):
|
|
1921
|
+
_all_orient.append(
|
|
1922
|
+
(B_res, intens, widths, weights_b[k_in_batch].item())
|
|
1923
|
+
)
|
|
1924
|
+
|
|
1925
|
+
_orient_seq = _all_orient if not _use_perturb_plain else []
|
|
1926
|
+
if sep_trans and _orient_seq:
|
|
1927
|
+
raise NotImplementedError("Opt.separate='transitions' is not available on the per-orientation accumulation path.")
|
|
1928
|
+
|
|
1929
|
+
if not has_strain and _orient_seq:
|
|
1930
|
+
# Vectorized scatter_add: collect all (B_res, w*intens) pairs across
|
|
1931
|
+
# all orientations in one pass, then bin in a single scatter operation.
|
|
1932
|
+
# Replaces O(n_orient) per-orientation makespec calls.
|
|
1933
|
+
_B_parts: list = []
|
|
1934
|
+
_WI_parts: list = []
|
|
1935
|
+
for B_res, intens, widths, w_k in _orient_seq:
|
|
1936
|
+
if B_res is not None and B_res.numel() > 0:
|
|
1937
|
+
_B_parts.append(B_res)
|
|
1938
|
+
_WI_parts.append((w_k * intens).to(spec.dtype))
|
|
1939
|
+
if _B_parts:
|
|
1940
|
+
_B_cat = torch.cat(_B_parts)
|
|
1941
|
+
_WI_cat = torch.cat(_WI_parts)
|
|
1942
|
+
_valid = ~torch.isnan(_B_cat)
|
|
1943
|
+
if _valid.any():
|
|
1944
|
+
_bin = ((_B_cat[_valid] - exp.Range[0]) / dx).floor().long()
|
|
1945
|
+
_in = (_bin >= 0) & (_bin < exp.nPoints) # drop out-of-window lines
|
|
1946
|
+
spec_dev = spec.to(_bin.device)
|
|
1947
|
+
spec = spec_dev.scatter_add_(0, _bin[_in], _WI_cat[_valid][_in]).cpu()
|
|
1948
|
+
else:
|
|
1949
|
+
for k_orient, (B_res, intens, widths, w_k) in enumerate(_orient_seq):
|
|
1950
|
+
if B_res.numel() == 0:
|
|
1951
|
+
continue
|
|
1952
|
+
|
|
1953
|
+
# has_strain: per-line spatial-domain Gaussian broadening
|
|
1954
|
+
for i in range(B_res.numel()):
|
|
1955
|
+
Bi = B_res[i].item()
|
|
1956
|
+
Ai = intens[i].item()
|
|
1957
|
+
fwhm_strain = max(widths[i].item(), _min_fwhm)
|
|
1958
|
+
fwhm_g_total = fwhm_strain # global lw is convolved afterwards (EasySpin)
|
|
1959
|
+
fwhm_l = 0.0
|
|
1960
|
+
|
|
1961
|
+
if fwhm_l == 0.0:
|
|
1962
|
+
if fwhm_g_total > 0.0:
|
|
1963
|
+
sigma = fwhm_g_total / _sqrt_8log2
|
|
1964
|
+
peak_amp = Ai * dx * _inv_sqrt_2pi / sigma
|
|
1965
|
+
spec = spec + w_k * peak_amp * torch.exp(
|
|
1966
|
+
-0.5 * ((x - Bi) / sigma) ** 2
|
|
1967
|
+
)
|
|
1968
|
+
else:
|
|
1969
|
+
frac = (Bi - exp.Range[0]) / dx
|
|
1970
|
+
ilo = int(frac)
|
|
1971
|
+
ihi = ilo + 1
|
|
1972
|
+
a = frac - ilo
|
|
1973
|
+
if 0 <= ilo < exp.nPoints:
|
|
1974
|
+
spec[ilo] = spec[ilo] + w_k * Ai * (1.0 - a)
|
|
1975
|
+
if 0 <= ihi < exp.nPoints:
|
|
1976
|
+
spec[ihi] = spec[ihi] + w_k * Ai * a
|
|
1977
|
+
else:
|
|
1978
|
+
_, line_i = makespec(
|
|
1979
|
+
exp.Range, exp.nPoints, B_res[i:i+1], intens[i:i+1],
|
|
1980
|
+
)
|
|
1981
|
+
line_i = convspec(
|
|
1982
|
+
line_i, dx, fwhm_g=fwhm_g_total, fwhm_l=fwhm_l, deriv=0
|
|
1983
|
+
)
|
|
1984
|
+
spec = spec + w_k * line_i
|
|
1985
|
+
|
|
1986
|
+
if opt.Verbosity >= 2:
|
|
1987
|
+
print(f" orient {k_orient+1}/{n_orient}: {B_res.numel()} transitions")
|
|
1988
|
+
|
|
1989
|
+
if sep_trans:
|
|
1990
|
+
if not _rows:
|
|
1991
|
+
spec = spec.unsqueeze(0)
|
|
1992
|
+
else:
|
|
1993
|
+
keys = sorted(_rows.keys(), key=lambda kk: (kk if isinstance(kk, tuple) else (-1, kk)))
|
|
1994
|
+
if any(v > 0 for v in _row_max.values()):
|
|
1995
|
+
keys = sorted(keys, key=lambda k_: -_row_max[k_]) # EasySpin: transitions by descending max rate
|
|
1996
|
+
spec = torch.stack([_rows[k_] for k_ in keys])
|
|
1997
|
+
if not used_projection:
|
|
1998
|
+
# Non-projection paths accumulate ∑ w·I per bin (stick), or per-line
|
|
1999
|
+
# shapes whose sum over points equals w·I. EasySpin pepper.m converts
|
|
2000
|
+
# these to a spectral density (spec/deltaX) and multiplies by 2π for
|
|
2001
|
+
# the χ integral; the orientation weights already sum to 4π, so no
|
|
2002
|
+
# further normalisation is applied. Together with the dBdE factor in
|
|
2003
|
+
# resfields this reproduces EasySpin's absolute intensity
|
|
2004
|
+
# (pepper_intensity_isopowder: ∫spec dB = 8π²·TransitionRate·dBdE).
|
|
2005
|
+
spec = spec / dx
|
|
2006
|
+
spec = spec * (2.0 * math.pi)
|
|
2007
|
+
|
|
2008
|
+
# -----------------------------------------------------------------------
|
|
2009
|
+
# Step 5: Global broadening / derivative
|
|
2010
|
+
# -----------------------------------------------------------------------
|
|
2011
|
+
fwhm_g = lw_effective[0] # Gaussian FWHM (mT) from sys.lw
|
|
2012
|
+
fwhm_l = lw_effective[1] # Lorentzian FWHM (mT) from sys.lw
|
|
2013
|
+
# EasySpin harmonic bookkeeping: with field modulation (Exp.ModAmp) the
|
|
2014
|
+
# absorption spectrum is convolved and the harmonic comes from the pseudo-
|
|
2015
|
+
# modulation; otherwise the harmonic is produced by the convolution
|
|
2016
|
+
# (ConvHarmonic) or, without lw, by finite differences (DerivHarmonic).
|
|
2017
|
+
# Exp.mwPhase (dispersion admixture) enters through the line-shape kernel.
|
|
2018
|
+
spec = _finish_field_sweep(spec, x, dx, exp, fwhm_g, fwhm_l, has_strain, phase=mw_phase)
|
|
2019
|
+
return x, spec
|