tomopt.optimisation.callbacks package

Submodules

tomopt.optimisation.callbacks.callback module

class tomopt.optimisation.callbacks.callback.Callback

Bases: object

Implements the base class from which all callback should inherit. Callbacks are used as part of the fitting, validation, and prediction methods of AbsVolumeWrapper. They can interject at various points, but by default do nothing. Please check in the AbsVolumeWrapper to see when exactly their interjections are called.

When writing new callbacks, the VolumeWrapper they are associated with will be their wrapper attribute. Their wrapper will have a fit_params attribute (FitParams) which is a data-class style object. It contains all the objects associated with the fit and predictions, including other callbacks. Callback interjections should read/write to wrapper.fit_params, rather than returning values.

Accounting for the interjection calls (on_*_begin & on_*_end), the full optimisation loop is:

1. Associate callbacks with wrapper (set_wrapper)

2. on_train_begin

3. for epoch in n_epochs:

   1. state = “train”

   2. on_epoch_begin

   3. for p, passive in enumerate(trn_passives):

      1. if p % passive_bs == 0:

         1. on_volume_batch_begin

         2. loss = 0

      2. load passive into passive volume

      3. on_volume_begin

      4. for muon_batch in range(n_mu_per_volume//mu_bs):

         1. on_mu_batch_begin

         2. Irradiate volume with mu_bs muons

         3. Infer scatter locations

         4. on_scatter_end

         5. Infer x0 and append to list of x0 predictions

         6. on_mu_batch_end

      5. on_x0_pred_begin

      6. Compute overall x0 prediction

      7. on_x0_pred_end

      8. Compute loss based on precision and cost, and add to loss

      9. if p`+1 % `passive_bs == 0:

         1. loss = loss/passive_bs

         2. on_volume_batch_end

         3. Zero parameter gradients

         4. on_backwards_begin

         5. Backpropagate loss and compute parameter gradients

         6. on_backwards_end

         7. Update detector parameters

         viii. Ensure detector parameters are within physical boundaries (AbsDetectorLayer.conform_detector) viv. loss = 0

      10. if len(trn_passives)-(p`+1) < `passive_bs:

          1. Break

   4. on_epoch_end

   5. state = “valid”

   6. on_epoch_begin

   7. for p, passive in enumerate(val_passives):

      1. if p % passive_bs == 0:

         1. on_volume_batch_begin

         2. loss = 0

      2. on_volume_begin

      3. for muon_batch in range(n_mu_per_volume//mu_bs):

         1. on_mu_batch_begin

         2. Irradiate volume with mu_bs muons

         3. Infer scatter locations

         4. on_scatter_end

         5. Infer x0 and append to list of x0 predictions

         6. on_mu_batch_end

      4. on_x0_pred_begin

      5. Compute overall x0 prediction

      6. on_x0_pred_end

      7. Compute loss based on precision and cost, and add to loss

      8. if p`+1 % `passive_bs == 0:

         1. loss = loss/passive_bs

         2. on_volume_batch_end

      9. if len(val_passives)-(p`+1) < `passive_bs:

         1. Break

   8. on_epoch_end

4. on_train_end

on_backwards_begin()

Runs when the loss for a batch of passive volumes has been computed, but not yet backpropagated.

Return type:

None

on_backwards_end()

Runs when the loss for a batch of passive volumes has been backpropagated, but parameters have not yet been updated.

Return type:

None

on_epoch_begin()

Runs when a new training or validations epoch begins.

Return type:

None

on_epoch_end()

Runs when a training or validations epoch ends.

Return type:

None

on_mu_batch_begin()

Runs when a new batch of muons begins.

Return type:

None

on_mu_batch_end()

Runs when a batch muons ends and scatters have been added to the volume inferrer.

Return type:

None

on_pred_begin()

Runs when the wrapper is about to begin in prediction mode.

Return type:

None

on_pred_end()

Runs when the wrapper has finished in prediction mode.

Return type:

None

on_scatter_end()

Runs when a scatters for the latest muon batch have been computed, but not yet added to the volume inferrer.

Return type:

None

on_step_end()

Runs when the parameters have been updated.

Return type:

None

on_train_begin()

Runs when detector fitting begins.

Return type:

None

on_train_end()

Runs when detector fitting ends.

Return type:

None

on_volume_batch_begin()

Runs when a new batch of passive volume layouts is begins.

Return type:

None

on_volume_batch_end()

Runs when a batch of passive volume layouts is ends.

Return type:

None

on_volume_begin()

Runs when a new passive volume layout is loaded.

Return type:

None

on_volume_end()

Runs when a passive volume layout has been predicted.

Return type:

None

on_x0_pred_begin()

Runs when the all the muons for a volume have propagated, and the volume inferrer is about to make its final prediction.

Return type:

None

on_x0_pred_end()

Runs after the volume inferrer has made its final prediction, but before the loss is computed.

Return type:

None

set_wrapper(wrapper)

Parameters:

wrapper (AbsVolumeWrapper) – Volume wrapper to associate with the callback

Return type:

None

wrapper: Optional[AbsVolumeWrapper] = None

tomopt.optimisation.callbacks.cyclic_callbacks module

class tomopt.optimisation.callbacks.cyclic_callbacks.CyclicCallback

Bases: Callback

tomopt.optimisation.callbacks.data_callbacks module

class tomopt.optimisation.callbacks.data_callbacks.MuonResampler

Bases: Callback

Resamples muons to only include those which will impact the passive volume at some point, even if they only hit the bottom layer.

static check_mu_batch(mu, volume)

Checks the provided muon batch to determine which muons will impact the passive volume at any point

Parameters:

• mu (MuonBatch) – incoming batch of muons

• volume (Volume) – Volume containing the passive volume to test against

Return type:

Tensor

Returns:

(muons) Boolean tensor where True indicates that the muon will hit the passive volume

on_mu_batch_begin()

Resamples muons prior to propagation through the volume such that all muons will hit the passive volume.

# TODO Add check for realistic validation

Return type:

None

static resample(mus, volume, gen)

Resamples muons until all muons will hit the passive volume.

Parameters:

• mus (Tensor) – xy_p_theta_phi tensor designed to initialise a MuonBatch

• volume (Volume) – Volume containing the passive volume to test against

• gen (AbsMuonGenerator) – Muon generator for sampling replacement muons

Return type:

Tensor

Returns:

xy_p_theta_phi tensor designed to initialise a MuonBatch

tomopt.optimisation.callbacks.detector_callbacks module

class tomopt.optimisation.callbacks.detector_callbacks.PanelCentring

Bases: Callback

Callback class for panel centring in the optimisation process.

This callback is used to centre the panels of PanelDetectorLayer objects by setting their xy coordinates to the mean xy value of all panels in the layer.

This update takes place after the panel positions have been updated in the optimisation process.

on_step_end()

Updates the xy coordinates of all panels in the PanelDetectorLayer objects after they have be updated, based on their current mean xy position.

Return type:

None

class tomopt.optimisation.callbacks.detector_callbacks.PanelUpdateLimiter(max_xy_step=None, max_z_step=None, max_xy_span_step=None)

Bases: Callback

Limits the maximum difference that optimisers can make to panel parameters, to prevent them from being affected by large updates from anomolous gradients. This is enacted by a hard-clamping based on the initial and final parameter values before/after each update step.

Parameters:

• max_xy_step (Optional[Tuple[float, float]]) – maximum update in xy position of panels

• max_z_step (Optional[float]) – maximum update in z position of panels

• max_xy_span_step (Optional[Tuple[float, float]]) – maximum update in xy_span position of panels

on_backwards_end()

Records the current paramaters of each panel before they are updated.

Return type:

None

on_step_end()

After the update step, goes through and hard-clamps parameter updates based on the difference between their current values and values before the update step.

Return type:

None

class tomopt.optimisation.callbacks.detector_callbacks.SigmoidPanelSmoothnessSchedule(smooth_range)

Bases: PostWarmupCallback

Creates an annealing schedule for the smooth attribute of SigmoidDetectorPanel. This can be used to move from smooth, unphysical panel with high sensitivity outside the physical panel boundaries, to one with sharper decrease in resolution | efficiency at the edge, and so more closely resembles a physical panel, whilst still being differentiable.

Parameters:

smooth_range (Tuple[float, float]) – tuple of initial and final values for the smooth attributes of all panels in the volume. A base-10 log schedule used over the number of epochs-total number of warmup epochs.

on_epoch_begin()

At the start of each training epoch, will anneal the SigmoidDetectorPanel s’ smooth attributes, if the callback is active.

Return type:

None

on_train_begin()

Sets all SigmoidDetectorPanel s to their initial smooth values.

Return type:

None

tomopt.optimisation.callbacks.diagnostic_callbacks module

class tomopt.optimisation.callbacks.diagnostic_callbacks.HitRecord

Bases: ScatterRecord

Records the hits of the muons. Once recorded, the hits can be retrieved via the get_record() method. plot_hit_density() may be used to plot the hit record.

Warning

Currently this callback makes no distinction between different volume layouts, and is designed to used over a single volume layout.

# TODO extend these to create one record per volume

on_scatter_end()

Saves the hits of the latest muon batch.

Return type:

None

class tomopt.optimisation.callbacks.diagnostic_callbacks.ScatterRecord

Bases: Callback

Records the PoCAs of the muons which are located inside the passive volume. Once recorded, the PoCAs can be retrieved via the get_record() method. plot_scatter_density() may be used to plot the scatter record.

Warning

Currently this callback makes no distinction between different volume layouts, and is designed to used over a single volume layout.

# TODO extend these to create one record per volume

get_record(as_df=False)

Access the recorded PoCAs.

Parameters:

as_df (bool) – if True, will return a Pandas DataFrame, otherwise will return a Tensor

Return type:

Union[Tensor, DataFrame]

Returns:

a Pandas DataFrame or a Tensor of recorded PoCAs

on_pred_begin()

Prepares to record scatters

Return type:

None

on_scatter_end()

Saves the PoCAs of the latest muon batch.

Return type:

None

on_train_begin()

Prepares to record scatters

Return type:

None

tomopt.optimisation.callbacks.eval_metric module

class tomopt.optimisation.callbacks.eval_metric.EvalMetric(lower_metric_better, name=None, main_metric=True)

Bases: Callback

Base class from which metric should inherit and implement the computation of their metric values. Inheriting classes will automatically be detected by MetricLogger and included in live feedback if it is the “main metric”

Parameters:

• lower_metric_better (bool) – if True, a lower value of the metric should be considered better than a higher value

• name (Optional[str]) – name to associate with the metric

• main_metric (bool) – whether this metric should be considered the “main metric”

get_metric()

This will be called by on_epoch_end()

Return type:

float

Returns:

metric value

on_train_begin()

Ensures that only one main metric is used

Return type:

None

tomopt.optimisation.callbacks.grad_callbacks module

class tomopt.optimisation.callbacks.grad_callbacks.NoMoreNaNs

Bases: Callback

Prior to parameter updates, this callback will check and set any NaN gradients to zero. Updates based on NaN gradients will set the parameter value to NaN.

Important

As new parameters are introduced, e.g. through new detector models, this callback will need to be updated.

on_backwards_end()

Prior to optimiser updates, parameter gradients are checked for NaNs.

Return type:

None

tomopt.optimisation.callbacks.heatmap_gif module

class tomopt.optimisation.callbacks.heatmap_gif.HeatMapGif(gif_filename='heatmap.gif')

Bases: Callback

Records a gif of the first heatmap in the first detector layer during training.

Parameters:

gif_filename (str) – savename for the gif (will be appended to the callback savepath)

on_epoch_begin()

When a new training epoch begins, saves an image of the current layout of the first heatmap in the first detector layer

Return type:

None

on_train_begin()

Prepares to record a new gif

Return type:

None

on_train_end()

When training, saves an image of the current layout of the first heatmap in the first detector layer and then combines all images into a gif

Return type:

None

tomopt.optimisation.callbacks.monitors module

class tomopt.optimisation.callbacks.monitors.MetricLogger(gif_filename='optimisation_history.gif', gif_length=10.0, show_plots=False)

Bases: Callback

Provides live feedback during training showing a variety of metrics to help highlight problems or test hyper-parameters without completing a full training. If show_plots is false, will instead print training and validation losses at the end of each epoch. The full history is available as a dictionary by calling get_loss_history(). Additionally, a gif of the optimisation can be saved.

Parameters:

• gif_filename (Optional[str]) – optional savename for recording a gif of the optimisation process (None -> no gif) The savename will be appended to the callback savepath

• gif_length (float) – If saving gifs, controls the total length in seconds

• show_plots (bool) – whether to provide live plots during optimisation in notebooks

cat_palette = 'tab10'

get_loss_history()

Get the current history of losses and metrics

Returns:

tuple of ordered dictionaries: first with losses, second with validation metrics

Return type:

history

get_results(loaded_best)

Return type:

Dict[str, float]

h_mid = 8

lbl_col = 'black'

lbl_sz = 24

leg_sz = 16

on_backwards_end()

Records the training loss for the latest volume batch

Return type:

None

on_epoch_begin()

Prepare to track new loss and snapshot the plots if training

Return type:

None

on_epoch_end()

If validation epoch finished, record validation losses, compute info and update plots

Return type:

None

on_train_begin()

Prepare for new training

Return type:

None

on_train_end()

Cleans up plots, and optionally creates a gif of the training history

Return type:

None

on_volume_batch_end()

Grabs the validation losses for the latest volume batch

Return type:

None

on_volume_end()

Grabs the validation sub-losses for the latest volume

Return type:

None

print_losses()

Print training and validation losses for the last epoch

Return type:

None

style = {'rc': {'patch.edgecolor': 'none'}, 'style': 'whitegrid'}

tk_col = 'black'

tk_sz = 16

update_plot()

Updates the plot(s).

Return type:

None

w_mid = 14.222222222222221

class tomopt.optimisation.callbacks.monitors.PanelMetricLogger(gif_filename='optimisation_history.gif', gif_length=10.0, show_plots=False)

Bases: MetricLogger

Logger for use with PanelDetectorLayer s

Parameters:

• gif_filename (Optional[str]) – optional savename for recording a gif of the optimisation process (None -> no gif) The savename will be appended to the callback savepath

• gif_length (float) – If saving gifs, controls the total length in seconds

• show_plots (bool) – whether to provide live plots during optimisation in notebooks

update_plot()

Updates the plot(s).

Return type:

None

tomopt.optimisation.callbacks.opt_callbacks module

class tomopt.optimisation.callbacks.opt_callbacks.EpochSave

Bases: Callback

Saves the state of the volume at the end of each training epoch to a unique file. This can be used to load a specifc state to either be used, or to resume training.

on_epoch_end()

Runs when a training or validations epoch ends.

Return type:

None

class tomopt.optimisation.callbacks.opt_callbacks.OneCycle(opt_name, warmup_length, init_lr=None, init_mom=None, mid_lr=None, mid_mom=None, final_lr=None, final_mom=None)

Bases: AbsOptSchedule

Callback implementing Smith 1-cycle evolution for lr and momentum (beta_1) https://arxiv.org/abs/1803.09820

In the warmup phase:

Learning rate is increased from init_lr to mid_lr, Momentum is decreased from init_mom to mid_mom, to stabilise the use of high LRs

In the convergence phase:

Learning rate is decreased from mid_lr to final_lr, Momentum is increased from mid_mom to final_mom

Setting the learning rate or momentum here will override the values specified when instantiating the VolumeWrapper. learning rate or momentum arguments can be None to avoid annealing or overriding their values.

Parameters:

• opt_name (str) – name of optimiser that should be affected by the scheduler

• warmup_length (int) – number of epochs to use for the warmup phase

• init_lr (Optional[float]) – initial learning rate (low)

• init_mom (Optional[float]) – initial momentum (high)

• mid_lr (Optional[float]) – nominal learning rate (high),

• mid_mom (Optional[float]) – nominal momentum (moderate),

• final_lr (Optional[float]) – final learning rate (low),

• final_mom (Optional[float]) – final momentum (high)

on_epoch_end()

Runs when a training or validations epoch ends.

Return type:

None

schedule()

Compute LR and momentum as a function of iter_cnt, according to defined ranges.

Return type:

Tuple[Optional[float], Optional[float]]

tomopt.optimisation.callbacks.pred_callbacks module

class tomopt.optimisation.callbacks.pred_callbacks.PredHandler

Bases: Callback

Default callback for predictions. Collects predictions and true voxelwise X0 pairs for a range of volumes and returns them as list of tuples of numpy arrays when get_preds() is called.

get_preds()

Return type:

List[Tuple[ndarray, ndarray]]

Returns:

List of predicted and target pairs

on_pred_begin()

Prepares to record predictions

Return type:

None

on_x0_pred_end()

Records predictions and true volume layout for the latest volume

Return type:

None

class tomopt.optimisation.callbacks.pred_callbacks.Save2HDF5PredHandler(path, use_volume_target, overwrite=False, x02id=None, compression='lzf')

Bases: VolumeTargetPredHandler

Saves predictions and targets to an HDF5 file, rather than caching and returning them. Samples are written incrementally. Can optionally save volume targets rather than voxel-wise X0 targets If an x02id lookup is provided, it transforms the target from an X0 value to a material class ID.

Parameters:

• path (Path) – savename of file to save predictions and targets

• use_volume_target (bool) – if True, saves the volume target value instead of the volume X0s

• overwrite (bool) – if True will overwrite existing files with the same path, otherwise will append to them

• x02id (Optional[Dict[float, int]]) – optional map from X0 values to class IDs

• compression (Optional[str]) – optional string representation of any compression to use when saving data

on_x0_pred_end()

Records predictions and true volume layout or target for the latest volume

Return type:

None

class tomopt.optimisation.callbacks.pred_callbacks.VolumeTargetPredHandler(x02id=None)

Bases: PredHandler

Returns the volume target as the target value, rather than the voxel-wise X0s. If an x02id lookup is provided, it transforms the target from an X0 value to a material class ID.

Parameters:

x02id (Optional[Dict[float, int]]) – optional map from X0 values to class IDs

on_x0_pred_end()

Records predictions and volume target for the latest volume

Return type:

None

tomopt.optimisation.callbacks.warmup_callbacks module

class tomopt.optimisation.callbacks.warmup_callbacks.CostCoefWarmup(n_warmup)

Bases: WarmupCallback

Sets a more stable cost coefficient in the AbsDetectorLoss by averaging the inference-error component for several epochs. During this warm-up monitoring phase, the detectors will be kept fixed.

Parameters:

n_warmup (int) – number of training epochs to wait before setting the cost coefficient

on_epoch_end()

If enough epochs have past, the overall median inference-error is computed and used to set the cost coefficient in the loss.

Return type:

None

on_volume_end()

If training, grabs the inference-error for the latest volume

Return type:

None

class tomopt.optimisation.callbacks.warmup_callbacks.OptConfig(n_warmup, rates)

Bases: WarmupCallback

Allows the user to specify the desired update steps for parameters in physical units. Over the course of several warm-up epochs the gradients on the parameters are monitored, after which suitable learning rates for the optimisers are set, such that the parameters will move by the desired amount every update. During the warm-up, the detectors will not be updated as optimiser learning rates will be set to zero.

The calculation here does not account for the effect of the optimiser’s momentum, nor scheduling and adaption of learning rates, and so the actual update rates may be different from the desired ones.

Parameters:

• n_warmup (int) – number of training epochs to wait before setting learning rates

• rates (Dict[str, float]) – dictionary of desired update rates for the parameters The keys are the names of the optimisers specified in the optimiser dictionary of the wrapper. The values are the desired update rates for the parameters in physical units. For example, if the optimiser is SGD, and the parameter is the xy position of a panel, then the update rate should be in metres. The parameters that are being optimisered are expected to be found in the zeroth parameter group of the optimiser, i.e. wrapper.opts[opt].param_groups[0][‘params’]. This implies that the optimiser is expected to have only one parameter group.

Example::

>>> OptConfig(n_warmup=2, rates={'xy_pos_opt':xy_pos_rate, 'z_pos_opt':z_pos_rate, 'xy_span_opt':xy_span_rate})

on_backwards_end()

Grabs training gradients from detector parameters

Return type:

None

on_epoch_end()

When enough training epochs have passed, sets suitable learning rates for the optimisers based on the median gradients and desired update rates

Return type:

None

class tomopt.optimisation.callbacks.warmup_callbacks.PostWarmupCallback

Bases: Callback

Callback class that waits for all WarmupCallback obejcts to finish their warmups before activating.

check_warmups()

When all WarmupCallbacks have finished, sets the callback to be active.

Return type:

None

on_epoch_begin()

Checks to see whether the callback should be active.

Return type:

None

on_train_begin()

Prepares for new training

Return type:

None

class tomopt.optimisation.callbacks.warmup_callbacks.WarmupCallback(n_warmup)

Bases: Callback

Warmup callbacks act at the start of training to track and set parameters based on the initial state of the detector. During warmup, optimisation of the detector is prevented, via a flag. If multiple warmup callbacks are present, they will wait to warmup according to the order they are provided in. Once the last warmup callback finished, the flag will be set to allow the detectors to be optimised. When a WarmupCallback is warming up, its warmup_active attribute will be True.

Important

When inheriting from WarmupCallback, the super methods of on_train_begin, on_epoch_begin, and on_epoch_end must be called.

Parameters:

n_warmup (int) – number of training epochs over-which to warmup

check_warmups()

If a WarmupCallback has finished, then its warmup_active is set to False, and the next WarmupCallback will have its warmup_active is set to True. If the finishing callback was the last WarmupCallback, then the “skip optimisation” flag is unset.

Return type:

None

on_epoch_begin()

Ensures that when one WarmupCallback has finished, either the next is called, or the detectors are set to be optimised.

Return type:

None

on_epoch_end()

After a training epoch is finished, increments the number of epochs that the callback has been warming up, provided it is active.

Return type:

None

on_train_begin()

Prepares to warmup

Return type:

None