pyextal.jaxGeo ============== .. py:module:: pyextal.jaxGeo .. autoapi-nested-parse:: 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 ---------- .. toctree:: :maxdepth: 1 /autoapi/pyextal/jaxGeo/cli/index /autoapi/pyextal/jaxGeo/config/index /autoapi/pyextal/jaxGeo/figures/index /autoapi/pyextal/jaxGeo/geometry/index /autoapi/pyextal/jaxGeo/io/index /autoapi/pyextal/jaxGeo/lm/index /autoapi/pyextal/jaxGeo/objective/index /autoapi/pyextal/jaxGeo/oversub/index /autoapi/pyextal/jaxGeo/refine/index /autoapi/pyextal/jaxGeo/reflections/index /autoapi/pyextal/jaxGeo/settings/index /autoapi/pyextal/jaxGeo/simulate/index Attributes ---------- .. autoapisummary:: pyextal.jaxGeo.SF_MASKS Classes ------- .. autoapisummary:: pyextal.jaxGeo.RefineSettings Functions --------- .. autoapisummary:: pyextal.jaxGeo.refine Package Contents ---------------- .. py:class:: 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. .. py:attribute:: h5 :type: Optional[pathlib.Path] :value: None .. py:attribute:: dat :type: Optional[pathlib.Path] :value: None .. py:attribute:: subsample :type: int :value: 2 .. py:attribute:: reflections :type: Union[str, int, tuple] :value: 'review' .. py:attribute:: review_dir :type: Optional[pathlib.Path] :value: None .. py:attribute:: symmetry :type: Optional[str] :value: None .. py:attribute:: imag :type: str :value: 'none' .. py:attribute:: layers :type: int :value: 1 .. py:attribute:: layer_fractions :type: Optional[Tuple[float, ...]] :value: None .. py:attribute:: refine_layers :type: Optional[Tuple[int, ...]] :value: None .. py:attribute:: diagonalise_measured :type: bool :value: True .. py:attribute:: steps :type: int :value: 800 .. py:attribute:: warmup :type: int :value: 5 .. py:attribute:: lr :type: float :value: 0.0001 .. py:attribute:: tile :type: int :value: 32 .. py:attribute:: thickness_window :type: float :value: 0.2 .. py:attribute:: scale_window :type: float :value: 0.05 .. py:attribute:: background_window :type: float :value: 0.05 .. py:attribute:: kernel_radius :type: int :value: 4 .. py:attribute:: voigt :type: bool :value: True .. py:attribute:: probe_sigma_start :type: float :value: 0.5 .. py:attribute:: probe_gamma_start :type: float :value: 0.5 .. py:attribute:: tilt_batch :type: Optional[int] :value: None .. py:attribute:: sim_batch :type: Optional[int] :value: None .. py:attribute:: sim_memory :type: float :value: 1.5 .. py:attribute:: memory_warn :type: float :value: 8.0 .. py:attribute:: figure :type: Union[pathlib.Path, str, None] :value: 'auto' .. py:attribute:: history :type: bool :value: True .. py:attribute:: logy :type: bool :value: False .. py:attribute:: notes :type: str :value: '' .. py:attribute:: dry_run :type: bool :value: False .. py:attribute:: cpu :type: bool :value: False .. py:attribute:: precision :type: str :value: 'highest' .. py:attribute:: geometry :type: bool :value: True .. py:attribute:: geo_start :type: int :value: 100 .. py:attribute:: geo_every :type: int :value: 100 .. py:attribute:: geo_params :type: Tuple[str, ...] :value: ('gl', 'rotation', 'tiltx', 'tilty') .. py:attribute:: geo_iters :type: int :value: 8 .. py:attribute:: geo_hessian :type: str :value: 'auto' .. py:attribute:: geo_max_block :type: Optional[float] :value: None .. py:attribute:: geo_tile :type: Optional[int] :value: None .. py:attribute:: geo_step_gl :type: float :value: 0.01 .. py:attribute:: geo_step_rotation :type: float :value: 0.2 .. py:attribute:: geo_step_tilt :type: float :value: 0.02 .. py:attribute:: geo_trust :type: float :value: 1.0 .. py:attribute:: geo_refresh :type: float :value: 2.0 .. py:attribute:: geo_lambda :type: float :value: 0.001 .. py:attribute:: geo_tol :type: float :value: 1e-05 .. py:attribute:: geo_check :type: bool :value: False .. py:attribute:: oversub :type: Tuple[Tuple[int, int, int], ...] :value: () .. py:attribute:: oversub_start :type: Optional[int] :value: None .. py:attribute:: oversub_leak :type: bool :value: True .. py:method:: 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. :rtype: RefineSettings :raises SetupError: On any inconsistent or unknown option. .. py:property:: frozen_layers 0-based layers whose structure factors are held (``OptConfig.frozen_layers``). .. py:method:: field_names() :classmethod: .. py:data:: SF_MASKS .. py:function:: refine(settings=None, **overrides) Runs a refinement, or validates its setup with ``dry_run``. :param settings: The options. :type settings: RefineSettings, optional :param \*\*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``. :rtype: dict :raises SetupError: If the dataset, the ``.dat`` or the options are unusable.