pyextal.jaxGeo.objective ======================== .. py:module:: pyextal.jaxGeo.objective .. autoapi-nested-parse:: The epoch loss as a smooth function of the geometry, and how to tile it. With the beam partition (``ig``) frozen, the loss is differentiable end to end: geometry -> tilt grid -> ``sg``/``bigkn`` -> structure matrix -> ``expm`` -> Voigt -> chi^2 (see `jaxGeo.geometry`). `GeometryObjective` is that function, evaluated chunk by chunk with its value, gradient and Hessian jitted. Memory and tiling ----------------- Differentiating with respect to ``sg`` is heavier than with respect to the structure factors, because the gradient runs through the Bethe-perturbation sum over every weak beam. Measured on tilted SrTiO3 data (565 beams): 3.3 MB per tilt for a forward pass, 10.2 MB for the gradient, 17.7 MB for the exact Hessian. XLA allocates a program's scratch as a single block. Under WSL, JAX's default (BFC) and 'platform' allocators cannot get a block over 4 GiB from CUDA however much of the card is free, which is why importing `jaxGeo` selects the stream-ordered allocator (``cuda_async``) unless told otherwise. The geometry objective still gets its own tiling -- the largest, starting from the training tile, whose derivative program fits under a block limit, read from the compiled program (`fit_geometry_tile`) -- because a smaller card or another allocator may need it. The Voigt margin (``kernel_radius`` each side) is kept, so small tiles waste pixels: tile 4 is 12x12 chunks of which 16 pixels count, tile 16 is 24x24 with 256. Measured, one Hessian pass on the SrTiO3 data with the previous all-tilts-at-once forward model: tile 16, exact (8.5 GB scratch) 13 s <- chosen under cuda_async tile 16, fd (5.8 GB) 30 s tile 8, fd (2.6 GB) 13-50 s (allocator-dependent) tile 4, exact (2.6 GB) 34 s Since `jaxEngine.model.batched_forward` runs a chunk's tilts in rematerialised sub-batches, the scratch no longer grows with the tile: the exact Hessian of the whole 160x160 scan (507 beams, 99 diagonalised) as one chunk is 5.3 GB, so the search usually stops at the training tile. ``hessian='auto'`` takes the exact Hessian where it fits and differences of the gradient otherwise. At equal tiling the objective and the training loss are the same function. Classes ------- .. autoapisummary:: pyextal.jaxGeo.objective.GeometryObjective Functions --------- .. autoapisummary:: pyextal.jaxGeo.objective.geometry_config pyextal.jaxGeo.objective.fit_geometry_tile pyextal.jaxGeo.objective.default_max_block Module Contents --------------- .. py:class:: GeometryObjective(params, metadata, config, constants, data, fns, x_ref, active, units) The epoch loss as a smooth function of the geometry, partition frozen. The variable is ``z``: the offset from the reference geometry ``x_ref`` (the one the metadata was built for), in units of ``units``, over the active parameters only. ``sg`` and ``bigkn`` enter as the stored values plus the change the JAX formulas predict, so at ``z = 0`` the loss is exactly the stored one and float32 rounding in the formulas cannot offset it. :param params: Current normalised parameters, held fixed. :type params: dict :param metadata: Metadata built for ``x_ref``. :type metadata: dict :param config: The configuration (its tiling is the one used). :type config: OptConfig :param constants: Bloch constants, with ``ndiag`` matching ``metadata``. :type constants: dict :param data: Holds ``experimental_data`` and ``variance``. :type data: dict :param fns: From `GeometryModel.jax_fns`. :type fns: tuple :param x_ref: (4,) reference geometry. :type x_ref: np.ndarray :param active: Indices into ``x`` of the refined parameters. :type active: np.ndarray :param units: (4,) size of one unit of each parameter. :type units: np.ndarray .. py:attribute:: x_ref .. py:attribute:: active .. py:attribute:: units .. py:attribute:: evals .. py:attribute:: chunks :value: [] .. py:attribute:: _value .. py:attribute:: _vg .. py:attribute:: _vgh .. py:property:: k .. py:method:: scratch_bytes(mode) XLA's scratch allocation for the derivative program, largest chunk. XLA allocates a program's temporaries as one block, so this is the single allocation the device must grant. :param mode: ``'exact'`` measures the Hessian program, ``'fd'`` the gradient one. :type mode: str :returns: Bytes. :rtype: int .. py:method:: x_of(z) Full geometry vector for offset ``z``. .. py:method:: _chunk_mean(fn, z) Mean of ``fn(z, chunk)`` over the chunks, one host transfer. The per-chunk results are summed on the device, so the chunks are dispatched back to back instead of each waiting for the previous one's result to reach the host. :param fn: A jitted ``fn(z, arrays, shape)``; any pytree output. :type fn: callable :param z: (k,) offset. :type z: np.ndarray :returns: Same pytree as ``fn``, as numpy float64. .. py:method:: value(z) Mean loss over the chunks at ``z``. .. py:method:: value_grad(z) Mean loss and gradient at ``z``. .. py:method:: value_grad_hess(z, mode='exact', fd_step=0.02) Mean loss, gradient and Hessian at ``z``. :param z: (k,) offset. :type z: np.ndarray :param mode: ``'exact'`` for forward-over-reverse HVPs, ``'fd'`` for central differences of the gradient. :type mode: str :param fd_step: Difference step, in units, for ``'fd'``. :type fd_step: float :returns: ``(loss, grad (k,), hess (k, k))``, Hessian symmetrised. :rtype: tuple .. py:function:: geometry_config(config, tile) The configuration with the scan tiled into ``tile``-pixel chunks. The Voigt margin is kept. :param config: The training configuration. :type config: OptConfig :param tile: Useful pixels per chunk side. :type tile: int :returns: Same configuration, re-tiled. :rtype: OptConfig .. py:function:: fit_geometry_tile(make_obj, config, start, mode, max_block) Chooses the geometry tile and Hessian method that fit. The limit is on the single scratch block XLA asks for, read from the compiled program, so nothing is run to find out. With ``mode='auto'`` each tile tries the exact Hessian first and then differences of the gradient: at equal tile the exact one is faster (one gradient plus k HVPs against 2k + 1 gradients), and a larger tile beats either at a smaller one, because the Voigt margin wastes a smaller share of each chunk. :param make_obj: ``make_obj(geo_config) -> GeometryObjective``. :type make_obj: callable :param config: The training configuration. :type config: OptConfig :param start: First tile to try; halved until something fits. :type start: int :param mode: ``'auto'``, ``'exact'`` or ``'fd'``. :type mode: str :param max_block: Largest scratch block to allow, bytes. :type max_block: float :returns: ``(tile, mode)``, ``mode`` resolved to ``'exact'`` or ``'fd'``. :rtype: tuple :raises RuntimeError: If nothing fits even at tile 2. .. py:function:: default_max_block() Largest single scratch block to plan for, bytes, given the allocator. :returns: 3.5 GB under the BFC/platform allocators (the WSL single-block limit is 4 GiB); otherwise 80 % of the device memory. On a large matrix that can still be too much once the training state is counted -- under WSL an overflow spills into shared system memory rather than failing -- so pass ``geo_max_block`` for such datasets. :rtype: float