spine.model.full_chain

Module with the core full reconstruction chain.

Functions

process_chain_config(self[, dump_config])

Process the full chain configuration and dump it.

Classes

FullChain(*args, **kwargs)	Full reconstruction in all its glory.
FullChainLoss(*args, **kwargs)	Loss function for the full chain.

class spine.model.full_chain.FullChain(*args: Any, **kwargs: Any)

Full reconstruction in all its glory.

Modular, end-to-end particle imaging detector reconstruction chain: - Deghosting for 3D tomographic reconstruction artifiact removal - Voxel-wise semantic segmentation - Point proposal - Particle clustering - Shower primary identification - Interaction clustering - Particle type classification - Primary identification - Track orientation

Typical configuration can look like this:

model:
  name: grappa
  modules:
    chain:
       <dictionary of arguments to specify chain-wide configuration>
    uresnet_deghost:
       name: <name of the model used to deghost
       <dictionary of arguments to specify the deghosting module>
    uresnet_segmentation:
       name: <name of the model used to do segmentation>
       <dictionary of arguments to specify the segmentation module>
    dbscan:
       TODO
    graph_spice:
       name: <name of the model used to do segmentation>
       <dictionary of arguments to specify the segmentation module>
    grappa_shower:
       name: <name of the model used to do segmentation>
       <dictionary of arguments to specify the segmentation module>
    grappa_track:
       name: <name of the model used to do segmentation>
       <dictionary of arguments to specify the segmentation module>
    grappa_inter:
      TODO
    calibration:
      TODO

See configuration file(s) prefixed with full_chain_ under the config directory for detailed examples of working configurations.

The chain section enables or disables specific stages of the full chain. When a module is disabled through this section, it will not be constructed. The configuration blocks for each enabled module should also live under the modules section of the configuration.

Attributes:

modes

Dictionary of (stage, modes) pairs which determine which options are available to each of the the reconstruction stage.

Methods

__call__(*args, **kwargs)	Call self as a function.
build_groups(clusts, clust_shapes, group_pred)	Use groups predictions from GrapPA to build superstructures.
forward(data[, sources, seg_label, ...])	Run a batch of data through the full chain.
group_labels(data, clusts, clust_shapes[, ...])	Aggregate particles using labels.
prepare_grappa_input(model, data, clusts, ...)	Prepares the input to a GrapPA model.
restrict_clusts(clusts, clust_shapes, shapes)	Restricts a cluster list against a list of shapes.
run_calibration(data[, sources, ...])	Run the calibration algorithm.
run_deghosting(data[, sources, seg_label, ...])	Run the deghosting algorithm.
run_fragmentation(data[, clust_label])	Run the fragmentation algorithm, i.e. the dense clustering step.
run_grappa(prefix, model, data, clusts, ...)	Run the GNN-based particle aggregator.
run_inter_aggregation(data[, clust_label, ...])	Run the interaction aggreation step.
run_part_aggregation(data[, clust_label, ...])	Run the particle aggreation step.
run_segmentation_ppn(data[, seg_label, ...])	Run the semantic segmentation and the point proposal algorithms.

MODULES = ['grappa_shower', 'grappa_track', 'grappa_inter', 'grappa_shower_loss', 'grappa_track_loss', 'grappa_inter_loss', 'full_chain_loss', 'graph_spice', 'graph_spice_loss', 'fragment_clustering', 'chain', 'dbscan_frag', ('uresnet_ppn', ['uresnet', 'ppn'])]

property modes

Dictionary of (stage, modes) pairs which determine which options are available to each of the the reconstruction stage.

Returns:

Dictionary of (stage, modes) pairs to be used

Return type:

Dict[str, Tuple(str)]

forward(data, sources=None, seg_label=None, clust_label=None, orig_index=None, coord_label=None, energy_label=None, meta=None, run_info=None)

Run a batch of data through the full chain.

Parameters:

• data (TensorBatch) – (N, 1 + D + N_f) tensor of voxel/value pairs - N is the the total number of voxels in the image - 1 is the batch ID - D is the number of dimensions in the input image - N_f is the number of features per voxel

• sources (TensorBatch, optional) – (N, 2) tensor of module/tpc pair for each voxel

• seg_label (TensorBatch, optional) – (N, 1 + D + 1) Tensor of segmentation labels - 1 is the segmentation label

• clust_label (TensorBatch, optional) – (N, 1 + D + N_c) Tensor of cluster labels - N_c is is the number of cluster labels

• orig_index (IndexBatch, optional) – (N_deghost) Index of the deghosted voxels in the original input voxel ordering. This is used to adapt cluster labels on the fly when deghosting was cached offline.

• coord_label (TensorBatch, optional) – (N, 1 + D + N_p) Tensor of point of interest labels - N_p is the number point labels

• energy_label (TensorBatch, optional) – (N, 1 + D + 1) Tensor of true energy deposition values - 1 is the energy deposition value in each voxel

• meta (Meta, optional) – Image metadata information

• run_info (List[RunInfo], optional) – Object containing information about the run, subrun and event

Return type:

TODO

run_deghosting(data, sources=None, seg_label=None, clust_label=None)

Run the deghosting algorithm.

This removes points that are artifacts of the tomographic reconstruction. This is only relevant for detectors producing 2D projections of an event. If requested, charge is rescaled according to the deghosting mask.

Parameters:

• data (TensorBatch) – (N, 1 + D + N_f) tensor of voxel/value pairs

• sources (TensorBatch, optional) – (N, 2) tensor of module/tpc pair for each voxel

• seg_label (TensorBatch, optional) – (N, 1 + D + 1) Tensor of segmentation labels

• clust_label (TensorBatch, optional) – (N, 1 + D + N_c) Tensor of cluster labels

Returns:

(N, 1 + D + N_f) tensor of deghosted voxel/value pairs

Return type:

TensorBatch

run_segmentation_ppn(data, seg_label=None, clust_label=None, orig_index=None)

Run the semantic segmentation and the point proposal algorithms.

This classifies each individual voxel in the image into different particle topological categories and identifies poins of interest, namely track end points and shower fragment start points.

Parameters:

• data (TensorBatch) – (N, 1 + D + N_f) tensor of voxel/value pairs

• seg_label (TensorBatch, optional) – (N, 1 + D + 1) Tensor of segmentation labels

• clust_label (TensorBatch, optional) – (N, 1 + D + N_c) Tensor of cluster labels

• orig_index (IndexBatch, optional) – (N_deghost) Index of the adapted voxels in the original input voxel ordering. This is used to adapt cluster labels on the fly when deghosting was cached offline.

run_fragmentation(data, clust_label=None)

Run the fragmentation algorithm, i.e. the dense clustering step.

This breaks down each topological class individually into a set of fragments, each composed of contiguous voxels which belong to a single particle instance, but do not necessarily make up the whole instance.

Parameters:

• data (TensorBatch) – (N, 1 + D + N_f) tensor of voxel/value pairs

• clust_label (TensorBatch, optional) – (N, 1 + D + N_c) Tensor of cluster labels

run_part_aggregation(data, clust_label=None, coord_label=None)

Run the particle aggreation step.

This step gathers particle fragments into complete particle instances. It either aggregates shower and track fragments independently or jointly into a single step.

In the process of shower aggregation, shower primaries can be identified.

Parameters:

• data (TensorBatch) – (N, 1 + D + N_f) tensor of voxel/value pairs

• clust_label (TensorBatch, optional) – (N, 1 + D + N_c) Tensor of cluster labels

• coord_label (TensorBatch, optional) – (N, 1 + D + 6) Array of label particle end points

run_inter_aggregation(data, clust_label=None, coord_label=None)

Run the interaction aggreation step.

This step gathers particles into complete interaction instances.

In the process of interaction aggregation, other tasks may be performed: - Particle identification - Primary tagging (particle coming from the interaction vertex) - Orientation tagging (order start/end points of particles) - Vertex reconstruction

Parameters:

• data (TensorBatch) – (N, 1 + D + N_f) tensor of voxel/value pairs

• clust_label (TensorBatch, optional) – (N, 1 + D + N_c) Tensor of cluster labels

• coord_label (TensorBatch, optional) – (N, 1 + D + 6) Array of label particle end points

run_calibration(data, sources=None, energy_label=None, meta=None, run_info=None)

Run the calibration algorithm.

This converts the raw charge values in ADC to energy depositions expressed in MeV. It applies gain, recombination, transparency and electron lifetime corrections.

Parameters:

• data (TensorBatch) – (N, 1 + D + N_f) tensor of voxel/value pairs

• sources (TensorBatch, optional) – (N, 2) tensor of module/tpc pair for each voxel

• energy_label (TensorBatch, optional) – (N, 1 + D + 1) Tensor of true energy deposition values - 1 is the energy deposition value in each voxel

• meta (Meta, optional) – Image metadata information

• run_info (List[RunInfo], optional) – Object containing information about the run, subrun and event

Returns:

(N, 1 + D + N_f) tensor of calibrated voxel/value pairs

Return type:

TensorBatch

run_grappa(prefix, model, data, clusts, clust_shapes, clust_primaries=None, coord_label=None, aggregate_shapes=False, shape_use_primary=False, point_use_primary=False, retain_primaries=False)

Run the GNN-based particle aggregator.

Parameters:

• prefix (str) – Name of the aggregation step

• model (GraPA) – GraPA model to execute for this aggregation step

• data (TensorBatch) – (N, 1 + D + N_f) tensor of voxel/value pairs

• clusts (IndexBatch) – List of clusters to aggregate using GrapPA

• clust_shapes (TensorBatch) – Semantic type of each of the clusters

• clust_primaries (IndexBatch) – List of primary fragments associated with each input cluster

• coord_label (TensorBatch, optional) – (N, 1 + D + 6) Array of label particle end points

• aggregate_shapes (bool, default False) – Combine shapes to give a shape to the aggregated object

• shape_use_primary (bool, default False) – Use primary shape as the group shape

• point_use_primary (bool, default False) – Use the primary fragment to get the points

• retain_primaries (bool, default False) – Retain the primary cluster in the aggregated group

Returns:

• groups (IndexBatch) – List of cluster groups aggregated using GrapPA

• group_shapes (TensorBatch) – Semantic type of each of the cluster groups

• group_primaries (IndexBatch) – List of primary clusters for each group

• shape_index (np.ndarray) – List of indexes used to restrict the original cluster list

group_labels(data, clusts, clust_shapes, aggregate_shapes=False, shape_use_primary=False, retain_primaries=False)

Aggregate particles using labels.

Parameters:

• data (TensorBatch) – (N, 1 + D + N_f) tensor of voxel/value pairs

• clusts (IndexBatch) – List of clusters to aggregate using GrapPA

• clust_shapes (TensorBatch) – Semantic type of each of the clusters

• aggregate_shapes (bool, default False) – Combine shapes to give a shape to the aggregated object

• shape_use_primary (bool, default False) – Use primary shape as the group shape

• retain_primaries (bool, default False) – Retain the primary cluster

Returns:

• groups (IndexBatch) – List of cluster groups aggregated using labels

• group_shapes (TensorBatch) – Semantic type of each of the cluster groups

• group_primaries (IndexBatch) – List of primary clusters for each group

• shape_index (np.ndarray) – List of indexes used to restrict the original cluster list

restrict_clusts(clusts, clust_shapes, shapes)

Restricts a cluster list against a list of shapes.

Parameters:

• clusts (IndexBatch) – List of clusters to aggregate using GrapPA

• clust_shapes (TensorBatch) – Semantic type of each of the clusters

• shapes (List[int]) – List of semantic shapes to restrict to

Returns:

• clusts (IndexBatch) – Restricted list of clusters

• clust_shapes (TensorBatch) – Restricted list of semantic types

• shape_index (np.ndarray) – List of indexes used to restrict the original cluster list

prepare_grappa_input(model, data, clusts, clust_shapes, clust_primaries=None, coord_label=None, point_use_primaries=False)

Prepares the input to a GrapPA model.

It builds the following input to GrpPA: - points: end points of fragments/particles - value: mean/std of charge distribution in each cluster - shape: shape of each fragment/particle

Parameters:

• model (torch.nn.Module) – GrapPA model to feed information to

• data (TensorBatch) – (N, 1 + D + N_f) tensor of voxel/value pairs

• clusts (IndexBatch) – List of clusters to aggregate using GrapPA

• clust_shapes (TensorBatch) – Semantic type of each of the clusters

• clust_primaries (IndexBatch, optional) – List of primary fragment within each cluster to aggregate

• coord_label (TensorBatch, optional) – (N, 1 + D + 6) Array of label particle end points

• point_use_primaries – Use the primary fragment only to infer primaries

Returns:

• data (TensorBatch) – (N, 1 + D + N_f) tensor of voxel/value pairs

• clusts (IndexBatch) – Input clusters to the model

• shapes (TensorBatch) – List of semantic type of each clusters

• points (TensorBatch) – List of start/end points associated with each cluster

• extra (TensorBatch) – List of additional features to pass to the GrapPA model

build_groups(clusts, clust_shapes, group_pred, primary_mask=None, aggregate_shapes=False, shape_use_primary=False, retain_primaries=False)

Use groups predictions from GrapPA to build superstructures.

Parameters:

• clusts (IndexBatch) – List of clusters to aggregate using GrapPA

• clust_shapes (TensorBatch) – Semantic type of each of the clusters

• group_pred (TensorBatch) – Group ID of each node in the GraPA output

• primary_mask (TensorBatch) – Binary mask as to whether a node is a group primary or not

• aggregate_shapes (bool, default False) – Combine shapes to give a shape to the aggregated object

• shape_use_primary (bool, default False) – Use primary shape as the group shape

• retain_primaries (bool, default False) – Retain the primary cluster in the aggregated group

Returns:

• groups (IndexBatch) – List of cluster groups aggregated using GrapPA

• group_shapes (TensorBatch) – Semantic type of each of the cluster groups

• group_primaries (IndexBatch) – List of primary clusters for each group

class spine.model.full_chain.FullChainLoss(*args: Any, **kwargs: Any)

Loss function for the full chain.

See also

FullChain

Attributes:

modes

Dictionary of (stage, modes) pairs which determine which options are available to each of the the reconstruction stage.

Methods

__call__(*args, **kwargs)	Call self as a function.
forward([seg_label, ppn_label, clust_label, ...])	Run the full chain output through the full chain loss.
restrict_segmentation_loss_input(seg_label, ...)	Align segmentation loss inputs to the deghosted voxel subset.
update_result(result[, prefix])	Update loss and accuracy using the output of one of the module.

property modes

Dictionary of (stage, modes) pairs which determine which options are available to each of the the reconstruction stage.

Returns:

Dictionary of (stage, modes) pairs to be used

Return type:

Dict[str, Tuple(str)]

static restrict_segmentation_loss_input(seg_label, segmentation, orig_index=None)

Align segmentation loss inputs to the deghosted voxel subset.

Parameters:

• seg_label (TensorBatch) – Segmentation labels defined on the original voxel set.

• segmentation (TensorBatch or None) – Segmentation logits defined on the effective voxel set.

• orig_index (IndexBatch, optional) – Original-voxel indices corresponding to the effective voxel set.

Returns:

Segmentation labels and logits restricted to true non-ghosts.

Return type:

tuple[TensorBatch, TensorBatch | None]

forward(seg_label=None, ppn_label=None, clust_label=None, clust_label_adapt=None, coord_label=None, graph_label=None, meta=None, orig_index=None, ghost=None, segmentation=None, seg_pred=None, **output)

Run the full chain output through the full chain loss.

Parameters:

• seg_label (TensorBatch, optional) – (N, 1 + D + 1) Tensor of segmentation labels - 1 is the segmentation label

• ppn_label (TensorBatch, optional) – (N, 1 + D + N_l) Tensor of PPN labels

• clust_label (TensorBatch, optional) – (N, 1 + D + N_c) Tensor of cluster labels - N_c is is the number of cluster labels

• clust_label_adapt (TensorBatch, optional) – (N, 1 + D + N_c) Tensor of cluster labels adapted to seg predictions - N_c is is the number of cluster labels

• coord_label (TensorBatch, optional) – (P, 1 + D + 8) Tensor of start/end point labels for each true particle in the image

• graph_label (EdgeIndexTensor, optional) – (2, E) Tensor of edges that correspond to physical connections between true particle in the image

• meta (Meta, optional) – Image metadata information

• orig_index (IndexBatch, optional) – Original-voxel indices corresponding to the effective voxel set.

• ghost (TensorBatch, optional) – (N, 2) Tensor of logits from the deghosting model

• ghost_pred (TensorBatch, optional) – (N,) Tensor of ghost predictions

• segmentation (TensorBatch, optional) – (N, N_c) Tensor of logits from the segmentation model

• seg_pred (TensorBatch, optional) –

  14. Semantic prediction for each point

• **output (dict, optional) – Additional outputs of the reconstruction chain

update_result(result, prefix=None)

Update loss and accuracy using the output of one of the module.

Parameters:

• result (dict) – Dictionary output of the module

• prefix (str, optional) – Prefix to preface the loss output keys with

spine.model.full_chain.process_chain_config(self, dump_config=False, **parameters)

Process the full chain configuration and dump it.

Parameters:

• dump_config (bool, default False) – Whether to dump the chain configuration in the log file or not

• **parameters (dict) – Dictionary of chain configuration parameters