pyextal.jaxJoint.optimize

Adam on the joint parameter tree.

jaxEngine.optimize.optimize for several datasets: each step evaluates every chunk of every dataset, averages each dataset’s chunk losses and gradients as the single-dataset per_epoch mode does, and combines the datasets with their weights into one loss and one gradient with respect to G. A dataset’s gradient reaches the shared structure factors through local_params (its links) and touches only its own scalars, so one Adam update moves the shared U_g by the pooled evidence and each dataset’s thickness, scale, background and probe by its own.

The schedule mirrors optimize exactly – the Adam state is reset at step warmup and snapshots are kept at the same steps – so a joint run of a single dataset is the jaxGeo run (jaxEngine.hdf5io.snapshot_steps applies).

Functions

_dataset_step_fn(ds)

Jitted value-and-gradient, with respect to G, of one chunk of ds.

_chunks(ds)

The dataset's training chunks, as optimize's ChunkingGenerator cuts them.

build_optimizer(G, datasets, lr, sf_mask, frozen_layers)

Adam with the joint freeze mask.

joint_optimize(G, datasets, *, steps, lr, sf_mask, ...)

Runs steps Adam steps on the pooled loss.

Module Contents

pyextal.jaxJoint.optimize._dataset_step_fn(ds)

Jitted value-and-gradient, with respect to G, of one chunk of ds.

pyextal.jaxJoint.optimize._chunks(ds)

The dataset’s training chunks, as optimize’s ChunkingGenerator cuts them.

pyextal.jaxJoint.optimize.build_optimizer(G, datasets, lr, sf_mask, frozen_layers)

Adam with the joint freeze mask.

pyextal.jaxJoint.optimize.joint_optimize(G, datasets, *, steps, lr, sf_mask, frozen_layers, warmup=0, offset=0, opt_state=None)

Runs steps Adam steps on the pooled loss.

Parameters:

  • G (dict) – The joint tree.

  • datasets (list) – Dataset objects with current config, metadata and tilts; each one’s weight scales its loss.

  • steps (int) – Steps to run.

  • lr (float) – Adam step size.

  • sf_mask (tuple) – Free components of the structure factors.

  • frozen_layers (tuple) – Layers whose structure factors are held.

  • warmup (int) – The Adam state is reset at this (absolute) step, as optimize does; 0 for never.

  • offset (int) – Steps already run, for the warmup and the log.

  • opt_state – Adam state to continue from; a fresh one when None.

Returns:

(G, loss_history, per_dataset_history, snapshots, opt_state). per_dataset_history holds each dataset’s own (unweighted) loss.

Return type:

tuple