pyextal.jaxEngine.model¶

Attributes¶

`TILT_BATCH`
`_EXPM_THETA`
`_EXPM_MAX_SQUARINGS`
`_PADE7`
`vmapped_forward`
`vmapped_forward_reference`
`vmapped_grad_forward`

Functions¶

`construct_ugh`(→ Tuple[jax.numpy.ndarray, ...)	Update Ugh matrices with optimizable structure factors.
`select_outputs`(→ jax.numpy.ndarray)	Gathers the measured beams from the diagonalised ones; 0 where a beam was not computed.
`build_structure_matrix_single`(→ jax.numpy.ndarray)	Build structure matrix A for a single tilt.
`propagate`(→ jax.numpy.ndarray)	Propagate wave through scattering matrix
`forward_model_single`(→ jax.numpy.ndarray)	Forward model for a SINGLE tilt point.
`structure_matrix_slab`(→ jax.numpy.ndarray)	build_structure_matrix_single, gathering only the blocks it uses.
`_dot`(a, b)
`expm_batch`(→ jax.numpy.ndarray)	exp of each matrix in `X` (b, n, n), sharing one squaring count.
`_forward_sub_batch`(sg, ig, bigkn, ughs, thicknesses, ...)	\|phi\|^2 for one sub-batch of tilts, (b, ndiag).
`batched_forward`(sg_batch, ig_batch, bigkn_batch, ...)	forward_model_single over every tilt, in rematerialised sub-batches.
`forward_model`(→ jax.numpy.ndarray)	Vectorized forward model for multiple tilts.
`forward_model_crystalparams`(→ jax.numpy.ndarray)	Forward model for refining crystal parameters.

Module Contents¶

pyextal.jaxEngine.model.construct_ugh(params: Dict, metadata: Dict, layer_idx: int = 0, delta=None) → Tuple[jax.numpy.ndarray, jax.numpy.ndarray]¶

Update Ugh matrices with optimizable structure factors.

delta – optional (dr, di), each (nbeams, nbeams), added to the static (non-refined) part before the refined reflections: a model-driven change of every coupling the static matrix holds (e.g. a multipole model’s U_g against the IAM the static matrix was built from). None leaves the static matrix as it is.

pyextal.jaxEngine.model.select_outputs(phi: jax.numpy.ndarray, indices: jax.numpy.ndarray, axis: int = 1) → jax.numpy.ndarray¶

Gathers the measured beams from the diagonalised ones; 0 where a beam was not computed.

phi holds the ndiag beams diagonalised at each tilt. A measured beam that beamcheck ranks as weak at some tilt sits at a position >= ndiag there, so it has no amplitude, and an out-of-range gather would return NaN – which also poisons every gradient. The index is clipped instead and the value set to zero; see postProcess.excluded_outputs for keeping those entries out of the loss.

Parameters:

phi – (…, ndiag) intensities along axis.
indices – Positions of the measured beams in the diagonalised order.
axis – The beam axis.

Returns:

Intensities of the measured beams, shaped like indices.

pyextal.jaxEngine.model.build_structure_matrix_single(sg: jax.numpy.ndarray, ig: jax.numpy.ndarray, ughr: jax.numpy.ndarray, ughi: jax.numpy.ndarray, constants: Dict) → jax.numpy.ndarray¶

Build structure matrix A for a single tilt.

Parameters:

sg – Pre-calculated excitation errors for this tilt, shape (nbeams,)
ig – Pre-calculated ig indices for this tilt, shape (nbeams,)
ughr – Real part of Ugh matrix, shape (nbeams, nbeams)
ughi – Imaginary part of Ugh matrix, shape (nbeams, nbeams)
constants – Dictionary containing Bloch constants

Returns:

Structure matrix, shape (ndiag, ndiag)

Return type:

pyextal.jaxEngine.model.propagate(bigkn: float, A: jax.numpy.ndarray, phi_0: jax.numpy.ndarray, thickness: float) → jax.numpy.ndarray¶: Propagate wave through scattering matrix

pyextal.jaxEngine.model.forward_model_single(sg: jax.numpy.ndarray, ig: jax.numpy.ndarray, bigkn: float, sf_params: Dict, thickness_params: Dict, metadata: Dict, constants: Dict) → jax.numpy.ndarray¶

Forward model for a SINGLE tilt point.

Parameters:

sg – Excitation errors for this tilt, shape (nbeams,)
ig – ig indices for this tilt, shape (nbeams,)
bigkn – bigkn value for this tilt, scalar
params – Optimization parameters
metadata – Metadata containing pre-calculated values
constants – Bloch wave constants
phi_0 – Initial state, shape (ndiag,)

Returns:

Output wave function, shape (ndiag,)

Return type:

phi_t

pyextal.jaxEngine.model.TILT_BATCH = 1024¶

pyextal.jaxEngine.model._EXPM_THETA = 3.92572478313866¶

pyextal.jaxEngine.model._EXPM_MAX_SQUARINGS = 16¶

pyextal.jaxEngine.model._PADE7 = (17297280.0, 8648640.0, 1995840.0, 277200.0, 25200.0, 1512.0, 56.0, 1.0)¶

pyextal.jaxEngine.model.structure_matrix_slab(sg: jax.numpy.ndarray, ig: jax.numpy.ndarray, ughr: jax.numpy.ndarray, ughi: jax.numpy.ndarray, constants: Dict) → jax.numpy.ndarray¶

build_structure_matrix_single, gathering only the blocks it uses.

The Bethe sum reads the diagonalised rows against every beam and the weak rows against the diagonalised columns; the weak-weak block, most of the matrix, is never needed.

pyextal.jaxEngine.model._dot(a, b)¶

pyextal.jaxEngine.model.expm_batch(X: jax.numpy.ndarray) → jax.numpy.ndarray¶

exp of each matrix in X (b, n, n), sharing one squaring count.

The count is the largest any matrix in the batch needs; a matrix needing fewer is scaled further, which Pade-7 handles at least as accurately. The branch on it is real (it is one scalar), so unused squarings cost nothing. Returns NaN if a matrix needs more than 16 squarings, as expm does.

pyextal.jaxEngine.model._forward_sub_batch(sg, ig, bigkn, ughs, thicknesses, constants)¶: |phi|^2 for one sub-batch of tilts, (b, ndiag).

pyextal.jaxEngine.model.batched_forward(sg_batch, ig_batch, bigkn_batch, sf_params, thickness_params, metadata, constants, ugh_delta=None)¶

forward_model_single over every tilt, in rematerialised sub-batches.

Same signature and result as jax.vmap(forward_model_single, ...); see the notes above TILT_BATCH. ugh_delta – optional per-layer (dr, di) changes of the static Ugh matrix (construct_ugh).

pyextal.jaxEngine.model.vmapped_forward¶

pyextal.jaxEngine.model.vmapped_forward_reference¶

pyextal.jaxEngine.model.vmapped_grad_forward¶

pyextal.jaxEngine.model.forward_model(params: Dict, metadata: Dict, constants: Dict) → jax.numpy.ndarray¶

Vectorized forward model for multiple tilts.

Uses vmap over forward_model_single for efficient batch processing.

pyextal.jaxEngine.model.forward_model_crystalparams(params: Dict, metadata: Dict, constants: Dict, model_fn: callable) → jax.numpy.ndarray¶

Forward model for refining crystal parameters.

Parameters:

params – Parameters including ‘model_param’ and ‘beams’
metadata – Metadata
constants – Bloch wave constants
phi_0_batch – Initial states
model_fn – Model function that computes structure factors from crystal params

Returns:

Output wave functions, shape (n_tilts, ndiag)

Return type:

phi_batch