pyextal.jaxGeo

Geometry-refining LARBED refinement on top of jaxEngine.

Adam refines the structure factors, thickness, scale, background and probe; between Adam segments the geometry (gl, rotation, tiltx, tilty) is refined by Levenberg-Marquardt on exact JAX derivatives, with everything else held fixed. On top of that:

  • complex U_g (Friedel-paired, imaginary parts free) – imag;

  • a specimen in layers stacked along the beam, each with its own thickness and structure factors, some of which may be held – layers;

  • an over-subtraction model for reflections whose integration background was overestimated and clipped at zero – oversub.

Nothing here is specific to one crystal: the beam list, the reflections, the starting geometry and the .dat all come from the dataset file written by the acom notebook (/raw + /pre_refinement).

Usage from Python:

from pyextal.jaxGeo import RefineSettings, refine
result = refine(RefineSettings(h5="data.h5", reflections="exported",
                               tile=16, steps=800, imag="both", layers=2))

or from a script, with dataset defaults and the full command line:

from pyextal.jaxGeo import cli
raise SystemExit(cli.main(h5=H5, tile=16, geo_max_block=12.0))

See jaxGeo.refine for the method.

Importing this package sets XLA_PYTHON_CLIENT_ALLOCATOR=cuda_async and XLA_PYTHON_CLIENT_PREALLOCATE=false unless already set, which only takes effect if JAX has not been imported yet – so import jaxGeo first. Under WSL, JAX’s default (BFC) and ‘platform’ allocators cannot get a single block over 4 GiB from CUDA however much of the card is free; the stream-ordered allocator pools memory and has no such limit, which lets the geometry derivatives run on the training tiling.

Submodules

Attributes

SF_MASKS

Classes

RefineSettings

Options for jaxGeo.refine.refine.

Functions

refine([settings])

Runs a refinement, or validates its setup with dry_run.

Package Contents

class pyextal.jaxGeo.RefineSettings

Options for jaxGeo.refine.refine.

Field names match the command-line options (geo_max_block is --geo-max-block); the command line’s --help describes each.

h5: pathlib.Path | None = None
dat: pathlib.Path | None = None
subsample: int = 2
reflections: str | int | tuple = 'review'
review_dir: pathlib.Path | None = None
symmetry: str | None = None
imag: str = 'none'
layers: int = 1
layer_fractions: Tuple[float, ...] | None = None
refine_layers: Tuple[int, ...] | None = None
diagonalise_measured: bool = True
steps: int = 800
warmup: int = 5
lr: float = 0.0001
tile: int = 32
thickness_window: float = 0.2
scale_window: float = 0.05
background_window: float = 0.05
kernel_radius: int = 4
voigt: bool = True
probe_sigma_start: float = 0.5
probe_gamma_start: float = 0.5
tilt_batch: int | None = None
sim_batch: int | None = None
sim_memory: float = 1.5
memory_warn: float = 8.0
figure: pathlib.Path | str | None = 'auto'
history: bool = True
logy: bool = False
notes: str = ''
dry_run: bool = False
cpu: bool = False
precision: str = 'highest'
geometry: bool = True
geo_start: int = 100
geo_every: int = 100
geo_params: Tuple[str, ...] = ('gl', 'rotation', 'tiltx', 'tilty')
geo_iters: int = 8
geo_hessian: str = 'auto'
geo_max_block: float | None = None
geo_tile: int | None = None
geo_step_gl: float = 0.01
geo_step_rotation: float = 0.2
geo_step_tilt: float = 0.02
geo_trust: float = 1.0
geo_refresh: float = 2.0
geo_lambda: float = 0.001
geo_tol: float = 1e-05
geo_check: bool = False
oversub: Tuple[Tuple[int, int, int], ...] = ()
oversub_start: int | None = None
oversub_leak: bool = True
resolved()

Checks the options and returns a canonical copy.

Returns:

reflections as ‘exported’/’review’/int/tuple of triples, symmetry decided, layer_fractions normalised to sum to one, refine_layers and geo_params as tuples, oversub as triples and oversub_start filled in.

Return type:

RefineSettings

Raises:

SetupError – On any inconsistent or unknown option.

property frozen_layers

0-based layers whose structure factors are held (OptConfig.frozen_layers).

classmethod field_names()
pyextal.jaxGeo.SF_MASKS
pyextal.jaxGeo.refine(settings=None, **overrides)

Runs a refinement, or validates its setup with dry_run.

Parameters:

  • settings (RefineSettings, optional) – The options.

  • **overrides – Fields to change on settings (or on the defaults).

Returns:

config, params, metadata, constants, geometry (the refined geometry dict) and, unless dry_run, run (the run key written to the file, or None if the simulation failed), loss_history and geometry_history.

Return type:

dict

Raises:

SetupError – If the dataset, the .dat or the options are unusable.