pyextal.jaxGeo.geometry¶

The geometry model: how gl, rotation and tilt map scan pixels to tilts.

Four parameters, all applied to the exported simgrid – the tilt attached to each experimental pixel:

gl scales the tilt step, so the tilt spacing between pixels goes as: 1/gl while the centre of the ROI stays put.

rotation rotates the scan axes against the simulation’s gg/gh. tiltx shifts the whole net along gg, in degrees. tilty shifts it along gh, in degrees, following the same convention

as callBloch.tilt.

The map from pixel to tilt is exactly affine, so it is recovered from the exported grid by least squares (GeometryModel.from_simgrid) and rebuilt for each geometry. With all four at their neutral values the rebuild reproduces the exported grid to rounding; GeometryModel.reproduction_error measures that, and nothing should be refined until it is small.

Internally the geometry is a vector x = [ln gl, rotation, tiltx, tilty] (delta_to_x), gl as a logarithm so a step in it is a relative change.

Why there is a gradient¶

The tilt reaches the loss through calculate_tilt_dependent_values, which calls the Fortran beamcheck per tilt to decide which beams are excited. That beam partition – ig and the output indices derived from it – is an integer set. Nothing else the tilt produces is: the excitation errors sg and the normal wavevector bigkn are closed-form smooth functions of the tilt (sg = 2 kz (h.zone) - (2t + h).G.h with kz = sqrt(K^2 - t.G.t)), and the forward model is already differentiable in both. GeometryModel.jax_fns gives those formulas in JAX, so with ig frozen the loss is a smooth function of the geometry (see jaxGeo.objective).

Attributes¶

`GEOMETRY`
`NEUTRAL`

Classes¶

GeometryModel

The affine pixel-to-tilt map of one scan, and its perturbations.

Functions¶

`delta_to_x`(delta)	Geometry dict -> internal vector; `gl` goes in as its logarithm.
`x_to_delta`(x)	Inverse of `delta_to_x()`.
`fmt_delta`(delta)	Formats a geometry dict for a log line.
`units_for`(step_gl, step_rotation, step_tilt)	Size of one refinement unit of each entry of `x`.
`rebuild_metadata`(tilts, constants, config, metadata[, ...])	Recomputes the tilt-dependent metadata for a geometry (`beamcheck` per tilt).

Module Contents¶

pyextal.jaxGeo.geometry.GEOMETRY = ('gl', 'rotation', 'tiltx', 'tilty')¶

pyextal.jaxGeo.geometry.NEUTRAL¶

pyextal.jaxGeo.geometry.delta_to_x(delta)¶

Geometry dict -> internal vector; gl goes in as its logarithm.

Parameters:: delta (dict) – gl factor, rotation/tiltx/tilty in degrees.
Returns:: (4,) [ln gl, rotation, tiltx, tilty].
Return type:: np.ndarray

pyextal.jaxGeo.geometry.x_to_delta(x)¶

Inverse of delta_to_x().

Parameters:: x (np.ndarray) – (4,) internal vector.
Returns:: Geometry dict.
Return type:: dict

pyextal.jaxGeo.geometry.fmt_delta(delta)¶: Formats a geometry dict for a log line.

pyextal.jaxGeo.geometry.units_for(step_gl, step_rotation, step_tilt)¶

Size of one refinement unit of each entry of x.

Parameters:

step_gl (float) – One unit of gl, as a fraction (0.01 = 1 %).
step_rotation (float) – One unit of rotation, degrees.
step_tilt (float) – One unit of either tilt, degrees.

Returns:

(4,) in the units of x.

Return type:

np.ndarray

class pyextal.jaxGeo.geometry.GeometryModel(centre, u, v, shape, gg, gh, gmxr, bigK)¶

The affine pixel-to-tilt map of one scan, and its perturbations.

Parameters:

centre (np.ndarray) – (3,) tilt at the ROI centre.
u (np.ndarray) – (3,) tilt step per column.
v (np.ndarray) – (3,) tilt step per row.
shape (tuple) – (rows, cols) of the scan.
gg (np.ndarray) – (3,) horizontal/vertical reciprocal axes, param.gg/gh.
gh (np.ndarray) – (3,) horizontal/vertical reciprocal axes, param.gg/gh.
gmxr (np.ndarray) – (3, 3) reciprocal metric.
bigK (float) – |K|; the wavelength is 1/bigK.

centre¶

u¶

v¶

shape¶

gg¶

gh¶

gg_scale¶

wavel¶

classmethod from_simgrid(simgrid, param)¶

Recovers the pixel-to-tilt map from an exported simulation grid.

The ROI builds simgrid as an affine function of the scan indices, so the fit is exact rather than approximate – which is what makes it safe to rebuild the grid for a different geometry.

Parameters:

simgrid (np.ndarray) – (rows, cols, 3) tilt of each experimental pixel.
param (SimParams) – From bloch_parse; supplies gg, gh, gmxr and bigK.

Returns:

GeometryModel

Raises:

ValueError – If the grid is not affine in the scan indices, which means it did not come from LARBEDROI.updateSimGrid.

offsets()¶

Row and column offsets of every pixel from the ROI centre.

Returns:: (dr, dc), each (rows, cols) float32.
Return type:: tuple

grid(geometry=None)¶

The tilt grid for a geometry.

Parameters:: geometry (dict | np.ndarray, optional) – A geometry dict or the internal vector x. Neutral when None.
Returns:: (rows, cols, 3) tilts.
Return type:: np.ndarray

Note

The tilt shift follows callBloch.tilt exactly – -deg2rad(t) * g / |gg| / wavel – so a tilt applied here means the same thing it does everywhere else in the package.

reproduction_error(simgrid)¶

Largest relative difference between the neutral grid and simgrid.

The file holds the grid in float64 and it is affine there to 1e-15, but load_for_refinement passes it through jnp.array and JAX defaults to float32 – so the achievable agreement is float32 epsilon, ~1.2e-7. Anything above ~1e-5 means the model does not describe the dataset.

Parameters:: simgrid (np.ndarray) – (rows, cols, 3) the exported grid.
Returns:: float

jax_fns(constants)¶

JAX versions of the tilt grid and the tilt-dependent values.

Parameters:: constants (dict) – Bloch constants (hkl, gmxr, bigk, zone, snorm, alpha).
Returns:: (tilts_at, sg_bigkn). tilts_at(x, dr, dc) gives the (n, 3) tilts of pixels at row/column offsets dr/dc from the ROI centre, for geometry x – the same map as grid(). sg_bigkn(t) gives (sg (n, nbeams), bigkn (n,)) – the same formulas as calculate_tilt_dependent_values.
Return type:: tuple

engine_agreement(fns, metadata)¶

How closely the JAX sg/bigkn reproduce the engine’s.

This is the sign and convention check: the derivatives mean nothing unless the JAX formulas give the engine’s values at the neutral geometry.

Parameters:

fns (tuple) – From jax_fns().
metadata (dict) – Built for the neutral grid; supplies sg_batch and bigkn_batch.

Returns:

(sg_rel, bigkn_rel), largest relative differences.

Return type:

tuple

pyextal.jaxGeo.geometry.rebuild_metadata(tilts, constants, config, metadata, quiet=True)¶

Recomputes the tilt-dependent metadata for a geometry (beamcheck per tilt).

Parameters:

tilts (np.ndarray) – (rows, cols, 3) or (n, 3) tilts.
constants (dict) – Bloch constants. ndiag is updated in place by the engine, which is why the caller must keep it consistent.
config (OptConfig) – The configuration.
metadata (dict) – The existing metadata, whose non-tilt entries are kept.
quiet (bool) – Swallow the engine’s progress output.

Returns:

Metadata with the tilt-dependent entries replaced.

Return type:

dict