pyextal.jaxMultipole.intensity ============================== .. py:module:: pyextal.jaxMultipole.intensity .. autoapi-nested-parse:: 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_.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 ------- .. autoapisummary:: pyextal.jaxMultipole.intensity.IntensitySettings pyextal.jaxMultipole.intensity.MultipoleBridge pyextal.jaxMultipole.intensity.FullMatrix Functions --------- .. autoapisummary:: pyextal.jaxMultipole.intensity.warm_start pyextal.jaxMultipole.intensity.newton_x pyextal.jaxMultipole.intensity.curvature_esd pyextal.jaxMultipole.intensity.family_cost pyextal.jaxMultipole.intensity.setup_state pyextal.jaxMultipole.intensity.simulate_dataset pyextal.jaxMultipole.intensity.reflection_stats pyextal.jaxMultipole.intensity.save_three_way_png pyextal.jaxMultipole.intensity.draw_intensities pyextal.jaxMultipole.intensity.refine_sequence pyextal.jaxMultipole.intensity.refine_intensities Module Contents --------------- .. py:class:: IntensitySettings Options of `refine_intensities` beyond the jaxJoint ones. :ivar mode: ``'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``). :vartype mode: str :ivar newton_iters, newton_tol, newton_lambda: Newton iterations, the step size (in units of x) below which it stops, and the initial damping. :ivar hessian_every: Accepted Newton steps between exact Hessians (``len(x)`` Hessian-vector products, ~2-3 gradients each). :vartype hessian_every: int :ivar rest_steps: Adam steps on absorption + scalars after each Newton stage. :vartype rest_steps: int :ivar cycles: Newton / Adam alternations. :vartype cycles: int :ivar family_cost: Before fitting, report the loss cost of each model U_g in the joint solution (`family_cost`). :vartype family_cost: bool :ivar steps: 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. :vartype steps: int :ivar lr: Adam step (jaxGeo's ``lr``). :vartype lr: float :ivar lr_x: The U_g step of the multipole parameters, A^-2 (``lr`` when None). :vartype lr_x: float :ivar warmup: 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. :vartype warmup: int :ivar polish: After Adam, damped Newton on the multipole parameters at the final absorption and scalars. :vartype polish: bool :ivar lr: Adam step for absorption and per-dataset scalars (as jaxJoint). :vartype lr: float :ivar geometry_every: Geometry LM every this many steps (0: never). :vartype geometry_every: int :ivar esd: Evaluate the curvature esd at the end. :vartype esd: bool :ivar save_runs: Simulate and save a ``run_NNN`` in every dataset file. :vartype save_runs: bool :ivar out: Where the summary JSON goes. :vartype out: Path :ivar log_every: Print every this many steps. :vartype log_every: int .. py:attribute:: mode :type: str :value: 'newton' .. py:attribute:: newton_iters :type: int :value: 12 .. py:attribute:: newton_tol :type: float :value: 0.005 .. py:attribute:: newton_lambda :type: float :value: 0.001 .. py:attribute:: hessian_every :type: int :value: 3 .. py:attribute:: rest_steps :type: int :value: 0 .. py:attribute:: cycles :type: int :value: 1 .. py:attribute:: steps :type: int :value: 200 .. py:attribute:: lr :type: float :value: 0.0005 .. py:attribute:: lr_x :type: float :value: None .. py:attribute:: warmup :type: int :value: None .. py:attribute:: polish :type: bool :value: False .. py:attribute:: geometry_every :type: int :value: 0 .. py:attribute:: esd :type: bool :value: True .. py:attribute:: save_runs :type: bool :value: False .. py:attribute:: out :type: pathlib.Path :value: None .. py:attribute:: log_every :type: int :value: 5 .. py:attribute:: family_cost :type: bool :value: False .. py:class:: MultipoleBridge(model, pmap, keys, scale=None) Maps the multipole parameter tree onto jaxJoint's ``G``. :param model: Its crystal must be the one the datasets describe (same cell, setting and voltage). :type model: MultipoleModel :param pmap: Which multipole parameters are free. :type pmap: ParameterMap :param keys: The shared canonical keys, in ``G['shared']`` order. :type keys: list :param scale: Units of ``x`` (one per free parameter); 1 by default. :type scale: array, optional .. py:attribute:: keys .. py:attribute:: scale .. py:attribute:: geo .. py:attribute:: _p0 .. py:attribute:: _Ms .. py:method:: p(x) Physical parameters for scaled ``x`` (float32). .. py:method:: ug(x) Model U_g of the shared keys, complex64. .. py:method:: G(theta) jaxJoint's ``G`` for the tree ``theta``. .. py:method:: theta_from(G, x0=None) The tree for a joint ``G``: its absorption and scalars, ``x0`` (zeros). .. py:method:: 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. .. py:method:: physical(x) .. py:class:: 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). :param model: The multipole model (its crystal = the datasets'). :type model: MultipoleModel :param datasets: jaxJoint datasets (after `prepare`). :type datasets: list :param iam_model: 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. :type iam_model: MultipoleModel, optional .. py:attribute:: model .. py:attribute:: per .. py:attribute:: keys :value: [] .. py:attribute:: geo .. py:attribute:: iam_model .. py:attribute:: U_iam .. py:method:: dU(p) ``U_model - U_IAM`` of every family, complex (``p`` physical, float32). .. py:method:: delta(name, p=None, dU=None) The engine's ``ugh_delta`` of dataset ``name``: per layer ``(dr, di)``. .. py:method:: iam_matrix(name) The IAM elastic matrix this class assumes the static one holds, (nb, nb) complex. .. py:method:: 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. :rtype: dict .. py:function:: warm_start(datasets, G, joint_json, engine) Copies a finished joint run's geometry, scalars and absorption into ``G``. :param datasets: jaxJoint `Dataset` objects (after `prepare`). :type datasets: list :param G: The joint tree from `prepare`; modified in place. :type G: dict :param joint_json: ``joint_.json`` of the run to start from. :type joint_json: path :param engine: For the geometry rebuilds. :type engine: EngineSwitch :returns: What was copied, per dataset. :rtype: dict .. py:function:: 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)``. :rtype: tuple .. py:function:: 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``. :rtype: dict .. py:function:: 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. :rtype: list .. py:function:: 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). :rtype: dict .. py:function:: 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). :rtype: tuple .. py:function:: reflection_stats(sim, exp) Per reflection: R = sum|sim - exp| / sum|exp| and rms of (sim - exp) / max(exp), in %. .. py:function:: 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. .. py:function:: 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, ``__comparison.png`` (the jaxGeo layout: experiment, simulation, residual in counts and in %, line profile) and, with ``joint``, ``__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)``. :rtype: list .. py:function:: 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. :param stages: ``(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). :type stages: list :param start: Physical start values (names of ``model.layout``). :type start: dict :param final_refine: 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). :type final_refine: dict, optional :param scale_of: ``name -> unit of x`` (Newton is invariant to it except through the damping and the stopping step). :type scale_of: callable, optional :param out: Summary JSON (the `refine_intensities` keys for the end point, plus ``stages``); the tree goes next to it as ``.theta.pkl``. :type out: path :param full_matrix: Drive every coupling of the Bloch-wave matrix by the model (`FullMatrix`), not only the refined structure factors. :type full_matrix: bool :returns: ``stages`` (per stage: names, values, loss, per-dataset losses, time), ``values`` (final physical), ``esd`` (if ``final_refine``). :rtype: dict .. py:function:: 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. :param js: The datasets and common options, as for jaxJoint. :type js: JointSettings :param model_builder: ``() -> MultipoleModel``; called *before* the datasets are set up (building a `Crystal` re-parses the engine). :type model_builder: callable :param pmap_builder: ``model -> ParameterMap`` (free parameters and the physical start). :type pmap_builder: callable :param joint_json: Warm start from this finished joint run. :type joint_json: path, optional :param settings: Steps, learning rates, geometry, output. :type settings: IntensitySettings :param scale: Units of ``x``; e.g. the Option 1 esd. :type scale: array, optional :param theta_init: 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``). :type theta_init: dict, optional :returns: ``theta``, ``bridge``, ``datasets``, ``history``, ``per_dataset``, ``esd`` (if requested), ``summary`` (path). :rtype: dict