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.
Files changed (116) hide show
  1. torchspin/__init__.py +363 -0
  2. torchspin/_cardamom_istos.py +327 -0
  3. torchspin/_cardamom_propagatedm.py +413 -0
  4. torchspin/_cardamom_utils.py +113 -0
  5. torchspin/_compile.py +100 -0
  6. torchspin/_linalg.py +140 -0
  7. torchspin/angmom.py +742 -0
  8. torchspin/autograd.py +557 -0
  9. torchspin/autoguess.py +342 -0
  10. torchspin/batch.py +225 -0
  11. torchspin/blochsteady.py +274 -0
  12. torchspin/cardamom.py +552 -0
  13. torchspin/chili.py +932 -0
  14. torchspin/chili_sle.py +1233 -0
  15. torchspin/constants.py +88 -0
  16. torchspin/convspec.py +148 -0
  17. torchspin/ctafft.py +91 -0
  18. torchspin/curry.py +241 -0
  19. torchspin/data/FourierSeriesCoefficients.txt +89 -0
  20. torchspin/data/GaussianCascadeCoefficients.txt +56 -0
  21. torchspin/data/isotopedata.txt +416 -0
  22. torchspin/data/spacegroups.txt +637 -0
  23. torchspin/dataproc.py +361 -0
  24. torchspin/dipbackground.py +69 -0
  25. torchspin/dipkernel.py +102 -0
  26. torchspin/diptensor.py +138 -0
  27. torchspin/endorfrq.py +198 -0
  28. torchspin/endorfrq_perturb.py +302 -0
  29. torchspin/eprload.py +1538 -0
  30. torchspin/eprsave.py +165 -0
  31. torchspin/esfit.py +2427 -0
  32. torchspin/evolve.py +274 -0
  33. torchspin/ewrls.py +124 -0
  34. torchspin/excitation.py +107 -0
  35. torchspin/exciteprofile.py +132 -0
  36. torchspin/experiment.py +323 -0
  37. torchspin/exponfit.py +177 -0
  38. torchspin/fastmotion.py +182 -0
  39. torchspin/fdaxis.py +57 -0
  40. torchspin/fitgui.py +350 -0
  41. torchspin/garlic.py +760 -0
  42. torchspin/ham.py +145 -0
  43. torchspin/ham_cf.py +115 -0
  44. torchspin/ham_ee.py +124 -0
  45. torchspin/ham_ez.py +119 -0
  46. torchspin/ham_ezho.py +403 -0
  47. torchspin/ham_hf.py +122 -0
  48. torchspin/ham_nn.py +108 -0
  49. torchspin/ham_nq.py +100 -0
  50. torchspin/ham_nz.py +136 -0
  51. torchspin/ham_oz.py +115 -0
  52. torchspin/ham_so.py +118 -0
  53. torchspin/ham_zf.py +254 -0
  54. torchspin/hamsymm.py +723 -0
  55. torchspin/initstate.py +158 -0
  56. torchspin/isotopologues.py +296 -0
  57. torchspin/levels.py +189 -0
  58. torchspin/levelsplot.py +363 -0
  59. torchspin/lineshape.py +715 -0
  60. torchspin/lpsvd.py +219 -0
  61. torchspin/makespec.py +87 -0
  62. torchspin/mdhmm.py +573 -0
  63. torchspin/mdload.py +788 -0
  64. torchspin/mdtraj2oripot.py +90 -0
  65. torchspin/ml.py +377 -0
  66. torchspin/mlpsvd.py +297 -0
  67. torchspin/nucdata.py +310 -0
  68. torchspin/nucfrq2d.py +228 -0
  69. torchspin/orca2torchspin.py +1198 -0
  70. torchspin/ordering.py +78 -0
  71. torchspin/oripotentialplot.py +246 -0
  72. torchspin/orisel.py +176 -0
  73. torchspin/pepper.py +2019 -0
  74. torchspin/pepper_autograd.py +613 -0
  75. torchspin/photoselect.py +174 -0
  76. torchspin/plegendre.py +131 -0
  77. torchspin/propint.py +127 -0
  78. torchspin/pulse.py +652 -0
  79. torchspin/py.typed +0 -0
  80. torchspin/rapidscan2spc.py +101 -0
  81. torchspin/resfields.py +337 -0
  82. torchspin/resfields_batch.py +446 -0
  83. torchspin/resfields_eig.py +250 -0
  84. torchspin/resfields_perturb.py +706 -0
  85. torchspin/resfreqs_matrix.py +344 -0
  86. torchspin/resfreqs_perturb.py +299 -0
  87. torchspin/resonator.py +310 -0
  88. torchspin/resonatorprofile.py +85 -0
  89. torchspin/rfmixer.py +147 -0
  90. torchspin/rotations.py +87 -0
  91. torchspin/rotutils.py +880 -0
  92. torchspin/saffron.py +1523 -0
  93. torchspin/saffron_pathways.py +92 -0
  94. torchspin/saffron_peaks.py +580 -0
  95. torchspin/saffron_thyme.py +472 -0
  96. torchspin/salt.py +680 -0
  97. torchspin/sigeq.py +71 -0
  98. torchspin/signalprocessing.py +226 -0
  99. torchspin/sitetransforms.py +254 -0
  100. torchspin/sphgrid.py +470 -0
  101. torchspin/spidyan.py +1508 -0
  102. torchspin/spinladder.py +161 -0
  103. torchspin/spinops.py +265 -0
  104. torchspin/spinsystem.py +1427 -0
  105. torchspin/stackplot.py +167 -0
  106. torchspin/stev.py +283 -0
  107. torchspin/stochtraj_diffusion.py +494 -0
  108. torchspin/stochtraj_jump.py +195 -0
  109. torchspin/strainwidth.py +727 -0
  110. torchspin/transmitter.py +89 -0
  111. torchspin/utils.py +605 -0
  112. torchspin-0.3.0.dist-info/METADATA +267 -0
  113. torchspin-0.3.0.dist-info/RECORD +116 -0
  114. torchspin-0.3.0.dist-info/WHEEL +5 -0
  115. torchspin-0.3.0.dist-info/licenses/LICENSE.md +22 -0
  116. 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