pyextal.jaxMultipole.intensity¶

Option 2: multipole parameters refined directly against LARBED intensities.

The joint refinement (jaxJoint) refines one table of shared structure factors – G['shared'], one {xfre, xfim, afre, afim} per canonical family – that every dataset reads through its own sign/phase links. Here that table is no longer free: its elastic part is computed from the multipole model,

xfre + i xfim = U_g(p), p = p0 + M (s * x),

for every shared family, in float32, inside the differentiable forward model. The free parameters are

x: the multipole parameters chosen by a ParameterMap (populations, kappa, P_lm, … – whatever it frees), scaled by s (typically their Option 1 esd) so that one unit is a comparable step for Adam;
absorption: the absorptive U’_g of each shared family, still free (afre; afim as the joint run’s imag option says), as in the joint refinement;
datasets: each dataset’s thickness, scale, background and probe, unchanged.

Everything else is jaxJoint’s: the same per-dataset losses, links, chunking and pooled weighting, and between Adam segments the same Levenberg-Marquardt geometry refinement of each dataset.

Scope (step 2A)¶

Only the shared families come from the model. The other Bloch-matrix elements – the differences between beams that no dataset refines – keep the engine’s tabulated IAM (Cromer-Mann, the .dat Debye-Waller factors), as they do in the joint refinement; so do the beam selection (beamcheck) and U_0.

Warm start¶

warm_start copies a finished joint run into the tree: each dataset’s refined geometry (from joint_<stamp>.json), its thickness, scale, background and probe (from its run_NNN in the dataset file) and the shared absorption. The multipole start comes from the Option 1 fit. The first loss is then the joint run’s, except that the elastic U_g are now the model’s.

Uncertainties¶

curvature_esd evaluates the Hessian of the pooled loss with respect to x. The pooled loss is the mean chi^2 per data point, so chi^2_total = N L and cov(x) = 2 (N H)^-1; scaled by L (the reduced chi^2, well above 1 for these data) it is the conventional esd. Correlated pixels (the probe convolution, the geometry) make even that optimistic.

Classes¶

`IntensitySettings`	Options of refine_intensities beyond the jaxJoint ones.
`MultipoleBridge`	Maps the multipole parameter tree onto jaxJoint's `G`.
`FullMatrix`	The multipole model on every coupling of the Bloch-wave matrix (Option 2B).

Functions¶

`warm_start`(datasets, G, joint_json, engine)	Copies a finished joint run's geometry, scalars and absorption into `G`.
`newton_x`(theta, work[, iters, tol, lam, ...])	Damped Newton on `theta['x']` with everything else held.
`curvature_esd`(theta, work, n_points)	Curvature and dataset-jackknife esd of `x` at the optimum.
`family_cost`(datasets, G, bridge, x[, log])	Loss cost of each multipole U_g in the joint run's own solution.
`setup_state`(js, model_builder, pmap_builder[, ...])	Datasets, the warm-started joint tree and the multipole bridge.
`simulate_dataset`(ds, G[, ugh_delta])	One dataset's simulated and measured rocking curves for the tree `G`.
`reflection_stats`(sim, exp)	Per reflection: R = sum\|sim - exp\| / sum\|exp\| and rms of (sim - exp) / max(exp), in %.
`save_three_way_png`(exp, sim_ref, sim_new, labels, path)	Experiment \| reference simulation \| new simulation \| new - reference, per reflection.
`draw_intensities`(state, x, outdir, tag[, names, ...])	Simulated rocking curves of the multipole solution `x` for each dataset.
`refine_sequence`(js, model_builder, stages, start[, ...])	Option 2 in stages: each stage refines its own parameters, the rest frozen.
`refine_intensities`(js, model_builder, pmap_builder[, ...])	Option 2: the multipole model refined against the joint LARBED intensities.

Module Contents¶

class pyextal.jaxMultipole.intensity.IntensitySettings¶

Options of refine_intensities beyond the jaxJoint ones.

Variables:

mode (str) – 'newton' (default) – damped Newton on the multipole parameters with everything else held, using the exact Hessian of the pooled loss (a few parameters: a few gradient evaluations per iteration), optionally followed by rest_steps Adam steps on the absorption and per-dataset scalars, repeated cycles times; 'adam' – Adam on everything at once for steps steps, with jaxGeo’s scaling and warm-up (see steps).
newton_lambda (newton_iters, newton_tol,) – Newton iterations, the step size (in units of x) below which it stops, and the initial damping.
hessian_every (int) – Accepted Newton steps between exact Hessians (len(x) Hessian-vector products, ~2-3 gradients each).
rest_steps (int) – Adam steps on absorption + scalars after each Newton stage.
cycles (int) – Newton / Adam alternations.
family_cost (bool) – Before fitting, report the loss cost of each model U_g in the joint solution (family_cost).
steps (int) – Adam steps (mode='adam'): the multipole parameters, the absorption U’_g and the per-dataset scalars together, scaled as jaxGeo scales them – one lr; the scalars normalised to their windows, the absorption in A^-2, and each multipole parameter’s step scaled (MultipoleBridge.ug_units) so that one step moves the model U_g by at most lr_x A^-2, as one step moves a free U_g in jaxGeo.
lr (float) – Adam step (jaxGeo’s lr).
lr_x (float) – The U_g step of the multipole parameters, A^-2 (lr when None).
warmup (int) – Steps at the start with the multipole parameters and the absorption frozen, then a fresh Adam state – jaxGeo’s warmup; the joint run’s value when None.
polish (bool) – After Adam, damped Newton on the multipole parameters at the final absorption and scalars.
lr – Adam step for absorption and per-dataset scalars (as jaxJoint).
geometry_every (int) – Geometry LM every this many steps (0: never).
esd (bool) – Evaluate the curvature esd at the end.
save_runs (bool) – Simulate and save a run_NNN in every dataset file.
out (Path) – Where the summary JSON goes.
log_every (int) – Print every this many steps.

mode: str = 'newton'¶

newton_iters: int = 12¶

newton_tol: float = 0.005¶

newton_lambda: float = 0.001¶

hessian_every: int = 3¶

rest_steps: int = 0¶

cycles: int = 1¶

steps: int = 200¶

lr: float = 0.0005¶

lr_x: float = None¶

warmup: int = None¶

polish: bool = False¶

geometry_every: int = 0¶

esd: bool = True¶

save_runs: bool = False¶

out: pathlib.Path = None¶

log_every: int = 5¶

family_cost: bool = False¶

class pyextal.jaxMultipole.intensity.MultipoleBridge(model, pmap, keys, scale=None)¶

Maps the multipole parameter tree onto jaxJoint’s G.

Parameters:

model (MultipoleModel) – Its crystal must be the one the datasets describe (same cell, setting and voltage).
pmap (ParameterMap) – Which multipole parameters are free.
keys (list) – The shared canonical keys, in G['shared'] order.
scale (array, optional) – Units of x (one per free parameter); 1 by default.

keys¶

scale¶

geo¶

_p0¶

_Ms¶

p(x)¶: Physical parameters for scaled x (float32).

ug(x)¶: Model U_g of the shared keys, complex64.

G(theta)¶: jaxJoint’s G for the tree theta.

theta_from(G, x0=None)¶: The tree for a joint G: its absorption and scalars, x0 (zeros).

ug_units(x=None)¶

Per free parameter: the step in x that moves the shared U_g by at most 1 A^-2.

1 / max_g |dU_g / dx_i| at x (zeros by default). Multiplying an Adam update by lr times this makes lr mean what it means for a free U_g in jaxGeo / jaxJoint: the largest U_g change per step, in A^-2.

physical(x)¶

class pyextal.jaxMultipole.intensity.FullMatrix(model, datasets, iam_model=None, log=print)¶

The multipole model on every coupling of the Bloch-wave matrix (Option 2B).

MultipoleBridge drives only the shared, refined structure factors; every other U_(g-h) in a dataset’s matrix – the beams that are coupled but not refined, upper-layer (HOLZ) couplings such as (111), (211) included, and the weak beams of the Bethe sums – stays at the engine’s static IAM (Cromer-Mann with the .dat Debye-Waller factors). This class adds, at every one of those positions, the change U_model(g) - U_IAM(g) as the engine’s ugh_delta (jaxEngine.model.construct_ugh), so the whole matrix follows the multipole parameters and the refined ADPs.

Set-up, per dataset: the distinct difference vectors g_i - g_j of its beam list, each mapped to its canonical family with the symmetry link (U(g) = c U(key) or c U(key)*, symmetry.link_phase); the positions not covered by a refined reflection (metadata['mask']). The model is evaluated once per call for the union of families of all datasets.

The absorptive part of these couplings stays the engine’s Bird-King (.dat B); the refined reflections keep their own route (model elastic part, refined U’_g).

Parameters:

model (MultipoleModel) – The multipole model (its crystal = the datasets’).
datasets (list) – jaxJoint datasets (after prepare).
iam_model (MultipoleModel, optional) – The IAM the static matrices were built from; by default the model’s crystal in 'cm' mode, which reproduces the engine’s defaultsf. check compares it with the static matrices.

model¶

per¶

keys = []¶

geo¶

iam_model¶

U_iam¶

dU(p)¶: U_model - U_IAM of every family, complex (p physical, float32).

delta(name, p=None, dU=None)¶: The engine’s ugh_delta of dataset name: per layer (dr, di).

iam_matrix(name)¶: The IAM elastic matrix this class assumes the static one holds, (nb, nb) complex.

check(datasets, log=print)¶

Compares the assumed IAM with each dataset’s static matrix (real part, free positions).

Returns:: name -> (max |diff|, max |static|) in A^-2.
Return type:: dict

pyextal.jaxMultipole.intensity.warm_start(datasets, G, joint_json, engine)¶

Copies a finished joint run’s geometry, scalars and absorption into G.

Parameters:

datasets (list) – jaxJoint Dataset objects (after prepare).
G (dict) – The joint tree from prepare; modified in place.
joint_json (path) – joint_<stamp>.json of the run to start from.
engine (EngineSwitch) – For the geometry rebuilds.

Returns:

What was copied, per dataset.

Return type:

dict

pyextal.jaxMultipole.intensity.newton_x(theta, work, iters=12, tol=0.005, lam=0.001, hessian_every=3, log=print, label='', bounds=None)¶

Damped Newton on theta['x'] with everything else held.

Solves (H + lam diag|H|) dx = -g; a step is accepted when the pooled loss falls (then lam shrinks) and otherwise lam grows. The exact Hessian (len(x) Hessian-vector products per chunk) is recomputed every hessian_every accepted steps, or when a step with a stale one fails; between, only the gradient is. Stops when the accepted step is below tol (units of x) or after iters iterations.

bounds ((lower, upper) in units of x) clips every trial step, so a radial parameter cannot leave its physical range (e.g. kappa’ < 0.3).

Returns:: (theta, history[(loss, x, per-dataset)], H).
Return type:: tuple

pyextal.jaxMultipole.intensity.curvature_esd(theta, work, n_points)¶

Curvature and dataset-jackknife esd of x at the optimum.

Two estimates:

formal: cov = 2 (N H)^-1 from the Hessian H of the pooled loss (a mean over N pixels). It treats every pixel as an independent measurement; the probe convolution, the tiling and the model error make neighbouring pixels strongly correlated, so it is far too small;
jackknife over datasets: each dataset’s own Newton step from the optimum gives the delete-one estimates x_(-d) = x - (H - w_d H_d)^-1 (g - w_d g_d); their spread is the esd that the scan-to-scan reproducibility supports (the analogue of the per-scan scatter used for the Option 1 sigma). Needs >= 3 datasets.

Returns:: hessian, loss, cov, esd (formal), esd_scaled (formal times sqrt L), corr, jackknife (esd or None), per_dataset_x (each dataset’s own optimum, one Newton step), all in the units of x.
Return type:: dict

pyextal.jaxMultipole.intensity.family_cost(datasets, G, bridge, x, log=print)¶

Loss cost of each multipole U_g in the joint run’s own solution.

For every shared family, the joint run’s refined U_g is replaced by the model’s (all else held at the joint solution) and the pooled loss re-evaluated: the increase is what the intensities say against the model’s value of that family. Costs one loss evaluation per family and dataset.

Returns:: (key, U_joint, U_model, cost) sorted by cost, largest first.
Return type:: list

pyextal.jaxMultipole.intensity.setup_state(js, model_builder, pmap_builder, joint_json=None, scale=None, log=print)¶

Datasets, the warm-started joint tree and the multipole bridge.

Returns:: model, pmap, datasets, G (the joint run’s solution, its own refined U_g), table, engine, bridge, theta (x = 0: the pmap start), keys, js (resolved).
Return type:: dict

pyextal.jaxMultipole.intensity.simulate_dataset(ds, G, ugh_delta=None)¶

One dataset’s simulated and measured rocking curves for the tree G.

As jaxJoint.refine._finish draws them (over-subtraction applied, the experiment clipped to match), but nothing is written to the dataset’s HDF5.

Returns:: (sim, exp, probe) – arrays (rows, cols, n) and the Voigt probe dict for save_comparison_png (or None).
Return type:: tuple

pyextal.jaxMultipole.intensity.reflection_stats(sim, exp)¶: Per reflection: R = sum|sim - exp| / sum|exp| and rms of (sim - exp) / max(exp), in %.

pyextal.jaxMultipole.intensity.save_three_way_png(exp, sim_ref, sim_new, labels, path, names=('joint (free U_g)', 'multipole'), stats=None)¶

Experiment | reference simulation | new simulation | new - reference, per reflection.

The last column, in % of the experiment’s peak on a shared colour scale per reflection, shows where the multipole U_g change the intensities relative to the joint run’s own refined U_g, with every other parameter identical.

pyextal.jaxMultipole.intensity.draw_intensities(state, x, outdir, tag, names=None, joint=True, family_of=None, log=print)¶

Simulated rocking curves of the multipole solution x for each dataset.

Writes, per dataset, <name>_<tag>_comparison.png (the jaxGeo layout: experiment, simulation, residual in counts and in %, line profile) and, with joint, <name>_<tag>_vs_joint.png (experiment, the joint run’s own simulation, the multipole one, and their difference). Nothing is written to the datasets’ HDF5 files. With family_of (hkl -> family key) the figures show one reflection per family – its member with the most counts; the statistics always cover every reflection.

Returns:: per dataset dict(name, labels, R_joint, rms_joint, R, rms).
Return type:: list

pyextal.jaxMultipole.intensity.refine_sequence(js, model_builder, stages, start, joint_json=None, settings=None, theta_init=None, final_refine=None, scale_of=None, out=None, full_matrix=False, log=print)¶

Option 2 in stages: each stage refines its own parameters, the rest frozen.

The datasets are set up once; every stage builds a parameter map from the current physical values (so frozen parameters stay exactly where the previous stages left them) and runs damped Newton (newton_x) on the joint intensities, with absorption and per-dataset scalars held.

Parameters:

stages (list) – (name, refine) with refine a build_parameter_map refine dict; a stage whose map has no free parameter is logged and skipped (e.g. l = 1, 3 on centrosymmetric sites).
start (dict) – Physical start values (names of model.layout).
final_refine (dict, optional) – Everything free at once, evaluated at the end point without steps: the marginal jackknife esds (the per-stage ones are conditional on the frozen parameters, hence too small).
scale_of (callable, optional) – name -> unit of x (Newton is invariant to it except through the damping and the stopping step).
out (path) – Summary JSON (the refine_intensities keys for the end point, plus stages); the tree goes next to it as .theta.pkl.
full_matrix (bool) – Drive every coupling of the Bloch-wave matrix by the model (FullMatrix), not only the refined structure factors.

Returns:

stages (per stage: names, values, loss, per-dataset losses, time), values (final physical), esd (if final_refine).

Return type:

dict

pyextal.jaxMultipole.intensity.refine_intensities(js, model_builder, pmap_builder, joint_json=None, settings=None, scale=None, theta_init=None, full_matrix=False, log=print)¶

Option 2: the multipole model refined against the joint LARBED intensities.

Parameters:

js (JointSettings) – The datasets and common options, as for jaxJoint.
model_builder (callable) – () -> MultipoleModel; called before the datasets are set up (building a Crystal re-parses the engine).
pmap_builder (callable) – model -> ParameterMap (free parameters and the physical start).
joint_json (path, optional) – Warm start from this finished joint run.
settings (IntensitySettings) – Steps, learning rates, geometry, output.
scale (array, optional) – Units of x; e.g. the Option 1 esd.
theta_init (dict, optional) – A previous run’s final tree (its .theta.pkl): its absorption and per-dataset scalars replace the joint run’s; its x is ignored (the physical start comes from pmap_builder).

Returns:

theta, bridge, datasets, history, per_dataset, esd (if requested), summary (path).

Return type:

dict