Anytime-Valid Neural Uncertainty Quantification for SPECT Imaging

Code for the thesis "Anytime-valid Neural Uncertainty Quantification for SPECT Imaging", implementing likelihood-based confidence sequences with classical estimators and neural predictors.

Installation

First, install micromamba, see https://mamba.readthedocs.io/en/latest/installation/micromamba-installation.html.

Then run

micromamba create -y -n anytime-nuq -f environment.cpu.yaml
# (optional) for cuda support, run the following command instead
# micromamba create -y -n anytime-nuq -f environment.cuda.yaml

and activate the environment via

micromamba activate anytime-nuq

In the following it is assumed that the PYTHONPATH environment variable includes a path to the directory in which this (README.md) file is. Assuming that your current working directory is in the directory of this file, executing

export PYTHONPATH=$PYTHONPATH:$(pwd)

sets PYTHONPATH appropriately.

Dataset

The dataset used in the thesis can not be made publicly accessible. However, it is possible to use different datasets without source code changes. They have to fulfill the following requirements:

• Location: data/dataset.h5
• Format: HDF5
• Top-level group/dataset names: must be one of tomogram_delta, images or data
• Dataset shape and data type: shape (1000, 64, 64) and data type float32
• Value range: [0, 1].

Note that anytime_nuq/data.py rotates images in the dataset by 30 degrees. If this is undesired behavior, comment out or remove

transforms.RandomAffine(
    (30, 30),
    (0, 0),
    (1.0, 1.0),
    interpolation=transforms.InterpolationMode.BILINEAR,
),

in anytime_nuq/data.py.

Training

To train the neural networks used in P-DMix, P-DMed, P-UMix, P-UMed, S-DMix, S-DMed, S-UMix and S-Med and some visualizations you may run

python3 training anytime_nuq/models/diffusion.py \
  --batch-size 32 --epochs 500 --model-size big --schedule linear \
  --learning-rate 0.0025 --weight-decay 0.0043 --dropout 0.37

to train the diffusion (DDPM) model and

for seed in {0..9}; do
  for aqct in 1 100 10000; do
    python3 anytime_nuq/models/unet.py \
      --model-size big --epochs 2 --batch-size 32 --weight-decay 0.042 \
      --dropout 0.29 --learning-rate 0.0003 --augment True \
      --acquisition-time $aqct --seed $seed
  done
done

to train the post-processing U-Nets. Since training the post-processing U-Nets serially in a loop like the above command might take a very long time to complete, running all training jobs in parallel on a cluster likely makes more sense. See anytime_nuq/cluster_scripts/run_unet_training.sh for a SLURM job script that does that.

Diffusion model checkpoints get saved in checkpoints/diffusion/. Post-processing U-Net checkpoints get saved in checkpoints/unet/acqt_$acqt/.

Confidence Sequences

To reproduce the experimental results with error level 0.05 and generate the confidence coefficients of all confidence sequences in the thesis, scripts

• scripts/run_priorlikemix_experiments.sh (local)
• scripts/run_seqlikemix_experiments.sh (local)

locally or

• cluster_scripts/run_priorlikemix_experiments.sh (SLURM cluster)
• cluster_scripts/run_seqlikemix_experiments.sh (SLURM cluster)

on a SLURM cluster.

Individual prior likelihood mixing confidence sequence runs can also be started via anytime_nuq/confidence_sequence/prior_likelihood_mixing.py. For example, to run a P-SN confidence sequence at level 5% with acqusition time 1 and starting time 1 for the first test set image, execute

python3 anytime_nuq/confidence_sequence/prior_likelihood_mixing.py \
  --idx 0 --acquisition-time 1.0 --prior standard-normal --t_0 1

The results of this run would get saved in confseqs/priorlikemix/acqt_1.0_t0_1_prior_standard-normal_idx_0/.

We have a similar CLI for sequential likelihood mixing: assuming that a diffusion model checkpoint is available, you may execute

python3 anytime_nuq/confidence_sequence/sequential_likelihood_mixing.py \
  --idx 1 --acquisition-time 100 --model diffusion-median  --t_0 60

to start a S-DMed sequence at level 5% with acquisition time 1 and starting time 60 for the second test set image. The results of this run would get saved in confseqs/seqlikemix/acqt_100.0_t0_60_model_diffusion-median_idx_1/.

Uncertainty Images

The Python script anytime_nuq/get_uncertainty.py can be used to generate uncertainty images. For example, to generate pixelwise uncertainty images for S-Med with MLEs with acquisition time 100 and starting time 60, execute

python3 anytime_nuq/get_uncertainty.py \
  --idx 0 --is-priorlikemix False --prior-model diffusion-median \
  --bound pixelwise --t -1 --t_0 1 --acquisition-time 100

Other uncertainty images can be generated by choosing a different parameter for bound, namely one of global, pixelwise, unet-ensemble, diffusion or l2-maximizer.

The --t flag controls for which acquisition steps to generate images. If --t -1 is supplied, it generates images for acquisition steps t_0, t_0 + 10, t_0 + 20, ....

A script that automates generating all pixelwise uncertainty images and average uncertainties can be bound at anytime_nuq/scripts/run_get_uncertainty.sh (cluster: anytime_nuq/cluster_scripts/run_get_uncertainty.sh).

The figures can be found in their corresponding confidence sequence directories confseqs/<priorlikemix or seqlikemix>/....

Visualizations

In the thesis are various figures. Some of them were generated using Python scripts. The scripts that generated them can be found in anytime_nuq/visualization/. Most work even if only a subset of all confidence sequence runs are available.

Figure 1.1

python3 anytime_nuq/visualization/plot_projections.py

Figure 4.3 and C.10 to C.13

python3 anytime_nuq/visualization/plot_unet_metrics_vs_step.py

Figure 4.6 and Figure C.1 to C.9

python3 anytime_nuq/visualization/plot_confseq_runs.py

Figure 4.7

python3 anytime_nuq/visualization/plot_violation_rates.py

Figure 4.8

python3 anytime_nuq/visualization/plot_uncertainties.py

Figure 4.9 to 4.15 and 4.17

These figures are generated with anytime_nuq/get_uncertainty.py and the resulting plots are saved in their corresponding confidence sequence directory confseqs/<priorlikemix or seqlikemix>/.....

Figure 4.16

python3 anytime_nuq/visualization/plot_mae_and_uncertainty_vs_time.py \
    confseqs/priorlikemix/acqt_1.0_t0_1_prior_diffusion-mixture_idx_0/uncertainty_visualization_metrics_unet-ensemble.json

Hyperparameter Optimization

We optimized hyperparameters via Bayesian optimization using Weights & Biases. The used sweep configs can be found in anytime_nuq/sweep_configs/. The SLURM script used to start one sweep agent can be found in cluster_scripts/run_unet_sweep.sh.