pyextal.jaxGeo.geometry ======================= .. py:module:: pyextal.jaxGeo.geometry .. autoapi-nested-parse:: 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 ---------- .. autoapisummary:: pyextal.jaxGeo.geometry.GEOMETRY pyextal.jaxGeo.geometry.NEUTRAL Classes ------- .. autoapisummary:: pyextal.jaxGeo.geometry.GeometryModel Functions --------- .. autoapisummary:: pyextal.jaxGeo.geometry.delta_to_x pyextal.jaxGeo.geometry.x_to_delta pyextal.jaxGeo.geometry.fmt_delta pyextal.jaxGeo.geometry.units_for pyextal.jaxGeo.geometry.rebuild_metadata Module Contents --------------- .. py:data:: GEOMETRY :value: ('gl', 'rotation', 'tiltx', 'tilty') .. py:data:: NEUTRAL .. py:function:: delta_to_x(delta) Geometry dict -> internal vector; ``gl`` goes in as its logarithm. :param delta: ``gl`` factor, ``rotation``/``tiltx``/``tilty`` in degrees. :type delta: dict :returns: (4,) ``[ln gl, rotation, tiltx, tilty]``. :rtype: np.ndarray .. py:function:: x_to_delta(x) Inverse of :func:`delta_to_x`. :param x: (4,) internal vector. :type x: np.ndarray :returns: Geometry dict. :rtype: dict .. py:function:: fmt_delta(delta) Formats a geometry dict for a log line. .. py:function:: units_for(step_gl, step_rotation, step_tilt) Size of one refinement unit of each entry of ``x``. :param step_gl: One unit of ``gl``, as a fraction (0.01 = 1 %). :type step_gl: float :param step_rotation: One unit of rotation, degrees. :type step_rotation: float :param step_tilt: One unit of either tilt, degrees. :type step_tilt: float :returns: (4,) in the units of ``x``. :rtype: np.ndarray .. py:class:: GeometryModel(centre, u, v, shape, gg, gh, gmxr, bigK) The affine pixel-to-tilt map of one scan, and its perturbations. :param centre: (3,) tilt at the ROI centre. :type centre: np.ndarray :param u: (3,) tilt step per column. :type u: np.ndarray :param v: (3,) tilt step per row. :type v: np.ndarray :param shape: ``(rows, cols)`` of the scan. :type shape: tuple :param gg: (3,) horizontal/vertical reciprocal axes, ``param.gg``/``gh``. :type gg: np.ndarray :param gh: (3,) horizontal/vertical reciprocal axes, ``param.gg``/``gh``. :type gh: np.ndarray :param gmxr: (3, 3) reciprocal metric. :type gmxr: np.ndarray :param bigK: ``|K|``; the wavelength is ``1/bigK``. :type bigK: float .. py:attribute:: centre .. py:attribute:: u .. py:attribute:: v .. py:attribute:: shape .. py:attribute:: gg .. py:attribute:: gh .. py:attribute:: gg_scale .. py:attribute:: wavel .. py:method:: from_simgrid(simgrid, param) :classmethod: 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. :param simgrid: (rows, cols, 3) tilt of each experimental pixel. :type simgrid: np.ndarray :param param: From ``bloch_parse``; supplies ``gg``, ``gh``, ``gmxr`` and ``bigK``. :type param: SimParams :returns: GeometryModel :raises ValueError: If the grid is not affine in the scan indices, which means it did not come from ``LARBEDROI.updateSimGrid``. .. py:method:: offsets() Row and column offsets of every pixel from the ROI centre. :returns: ``(dr, dc)``, each (rows, cols) float32. :rtype: tuple .. py:method:: grid(geometry=None) The tilt grid for a geometry. :param geometry: A geometry dict or the internal vector ``x``. Neutral when None. :type geometry: dict | np.ndarray, optional :returns: (rows, cols, 3) tilts. :rtype: 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. .. py:method:: 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. :param simgrid: (rows, cols, 3) the exported grid. :type simgrid: np.ndarray :returns: float .. py:method:: jax_fns(constants) JAX versions of the tilt grid and the tilt-dependent values. :param constants: Bloch constants (``hkl``, ``gmxr``, ``bigk``, ``zone``, ``snorm``, ``alpha``). :type constants: dict :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 :meth:`grid`. ``sg_bigkn(t)`` gives ``(sg (n, nbeams), bigkn (n,))`` -- the same formulas as `calculate_tilt_dependent_values`. :rtype: tuple .. py:method:: 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. :param fns: From :meth:`jax_fns`. :type fns: tuple :param metadata: Built for the neutral grid; supplies ``sg_batch`` and ``bigkn_batch``. :type metadata: dict :returns: ``(sg_rel, bigkn_rel)``, largest relative differences. :rtype: tuple .. py:function:: rebuild_metadata(tilts, constants, config, metadata, quiet=True) Recomputes the tilt-dependent metadata for a geometry (``beamcheck`` per tilt). :param tilts: (rows, cols, 3) or (n, 3) tilts. :type tilts: np.ndarray :param constants: Bloch constants. ``ndiag`` is updated in place by the engine, which is why the caller must keep it consistent. :type constants: dict :param config: The configuration. :type config: OptConfig :param metadata: The existing metadata, whose non-tilt entries are kept. :type metadata: dict :param quiet: Swallow the engine's progress output. :type quiet: bool :returns: Metadata with the tilt-dependent entries replaced. :rtype: dict