pyextal.jaxJoint

Joint LARBED refinement of several datasets on top of jaxGeo.

Several datasets of one crystal – other scans, other camera lengths, other zone axes – are refined together. Their structure factors are one shared set of parameters, while each dataset keeps its own geometry (gl, rotation, tiltx, tilty), thickness, scale, background, probe and over-subtraction model. Adam minimises the pooled chi^2 of all datasets; between Adam segments each dataset’s geometry is refined by jaxGeo’s Levenberg-Marquardt, on its own loss, with everything else held fixed.

Usage from Python:

from pyextal.jaxJoint import JointSettings, refine
from pyextal.jaxGeo import RefineSettings
result = refine(JointSettings(
    common=RefineSettings(steps=800, tile=180, reflections='exported'),
    datasets=[dict(h5=R5), dict(h5=R4, subsample=2)]))

or from a script, with the full command line:

from pyextal.jaxJoint import cli
raise SystemExit(cli.main(doc=__doc__, datasets=[dict(h5=R5), dict(h5=R4)],
                          steps=800, tile=180))

See jaxJoint.refine for the method.

Importing this package sets the same JAX allocator defaults as jaxGeo (it imports jaxGeo first), so import jaxJoint before JAX.

Submodules

• pyextal.jaxJoint.cli
• pyextal.jaxJoint.datasets
• pyextal.jaxJoint.engine
• pyextal.jaxJoint.link
• pyextal.jaxJoint.optimize
• pyextal.jaxJoint.params
• pyextal.jaxJoint.refine
• pyextal.jaxJoint.report
• pyextal.jaxJoint.settings

Attributes

JOINT_FIELDS

Classes

JointSettings

Options for jaxJoint.refine.refine.

Functions

refine([settings])

Runs a joint refinement, or validates its setup with dry_run.

Package Contents

class pyextal.jaxJoint.JointSettings

Options for jaxJoint.refine.refine.

Parameters:

• common (RefineSettings) – Options every dataset starts from, including all of JOINT_FIELDS.

• datasets (sequence) – One entry per dataset: a dict of the fields that differ from common (at least h5), a full RefineSettings, or just the dataset path.

• names (sequence, optional) – Display names; default display_names.

• weighting – 'pixels' weights each dataset’s mean chi^2 by its number of data points, so the loss is the chi^2 of all of them pooled; 'equal' averages the datasets’ chi^2; a sequence gives explicit relative weights.

• out_dir (Path, optional) – Where the joint summary goes; default the first dataset’s folder.

common: pyextal.jaxGeo.settings.RefineSettings

datasets: Tuple = ()

names: Tuple[str, ...] | None = None

weighting: str | Tuple[float, ...] = 'pixels'

out_dir: pathlib.Path | None = None

resolved()

Checks the options and returns a canonical copy.

Returns:

common resolved, datasets as a tuple of resolved RefineSettings, names filled in and weighting as ‘pixels’, ‘equal’ or a tuple of floats.

Return type:

JointSettings

Raises:

SetupError – On an inconsistent option, a dataset setting a joint field, or the same file twice.

pyextal.jaxJoint.JOINT_FIELDS = ('symmetry', 'imag', 'layers', 'refine_layers', 'steps', 'warmup', 'lr', 'geometry',...

pyextal.jaxJoint.refine(settings=None, **overrides)

Runs a joint refinement, or validates its setup with dry_run.

Parameters:

• settings (JointSettings, optional) – The options.

• **overrides – Fields to change on settings.

Returns:

G (the joint tree), datasets (the Dataset objects), table (shared structure factors), engine and, unless dry_run, loss_history, runs (dataset name -> run key) and summary (the JSON path).

Return type:

dict

Raises:

SetupError – If the datasets cannot be refined together.