one example notebook about petric HOT 11 CLOSED

import os
from sirf import STIR as pet
from sirf.contrib.partitioner import partitioner
from cil.optimisation.functions import SGFunction
from cil.optimisation.algorithms import GD
from cil.utilities.display import show2D 
from cil.optimisation.utilities import Sampler
from cil.optimisation.utilities.callbacks import TextProgressCallback, ProgressCallback

# engine's messages go to files, except error messages, which go to stdout
_ = pet.MessageRedirector('info.txt', 'warn.txt')
# Needed for get_subsets()
pet.AcquisitionData.set_storage_scheme('memory')
# fewer message from STIR and SIRF
pet.set_verbosity(0)

os.chdir('/home/jovyan/work/Challenge24/data')

acquired_data = pet.AcquisitionData('prompts.hs')

additive_term = pet.AcquisitionData('additive.hs')

mult_factors = pet.AcquisitionData('multfactors.hs')

# somewhat crazy initialisation, currently hand-tuned scale
initial_image = pet.ImageData('20170809_NEMA_MUMAP_UCL.hv')+.5

initial_image=initial_image.zoom_image(zooms=(1,1,1), offsets_in_mm=(0,0,0), size=(-1,200,200))

num_subsets = 7
data, acq_models, obj_funs = partitioner.data_partition(acquired_data, additive_term, mult_factors, num_subsets, mode='staggered', initial_image=initial_image)

osem_sol = initial_OSEM(acquired_data, additive_term, mult_factors, initial_image)
# add RDP prior to the objective functions
step_size = 1e-7
add_regulariser = True
if add_regulariser:
    alpha = 5
    epsilon = initial_image.max()*1e-4
    prior = add_RDP(alpha, epsilon, obj_funs)
    step_size = 1e-10

#set up and run the stochastic gradient descent algorithm
sampler = Sampler.random_without_replacement(len(obj_funs))
F = - SGFunction(obj_funs, sampler=sampler)
alg = GD(initial=osem_sol, objective_function=F, step_size=step_size)

alg.update_objective_interval = num_subsets
alg.run(num_subsets * 1, callbacks=[TextProgressCallback()])

cmax = .15
im_slice = 70

show2D([osem_sol.as_array()[im_slice,:,:], 
        alg.solution.as_array()[im_slice,:,:]], 
       title=['OSEM',f"SGD epoch {alg.iteration/num_subsets}"], cmap="PuRd", fix_range=[(0, 0.2),(0,0.2)])

# somewhat crazy initialisation, currently hand-tuned scale
initial_image = pet.ImageData('20170809_NEMA_MUMAP_UCL.hv')+.5

initial_image=initial_image.zoom_image(zooms=(1,1,1), offsets_in_mm=(0,0,0), size=(-1,200,200))
osem_sol = initial_OSEM(acquired_data, additive_term, mult_factors, initial_image)

initial_image = pet.ImageData('OSEM_image.hv')

kappa_image = pet.ImageData('penalty_weights.hv')

def construct_RDP(penalty_strength, initial_image, kappa):
    '''
    construct RDP prior, set epsilon etc.
    '''
    prior = sirf.STIR.RelativeDifferencePrior()
    # make it differentiable
    epsilon = initial_image.max()*1e-4
    prior.set_epsilon(epsilon)
    prior.set_penalisation_factor(penalty_strength)
    prior.set_kappa(kappa)
    prior.set_up(initial_image)
    return prior

def add_prior(prior, obj_funs):
    '''
    add prior evenly to every objective function.
    
    WARNING: modifies the prior to divide the penalty_strength by the number of `obj_funs`
    '''
    prior.set_penalisation_factor(pror.get_penalisation_factor ()/ len(obj_funs));
    prior.set_up(initial_image)
    for f in obj_funs:
        f.set_prior(prior)

num_subsets = 7
data, acq_models, obj_funs = partitioner.data_partition(acquired_data, additive_term, mult_factors, num_subsets, mode='staggered', initial_image=initial_image)
add_RDP(prior, obj_funs)
step_size = 1e-10

#set up and run the stochastic gradient descent algorithm
sampler = Sampler.random_without_replacement(len(obj_funs))
F = - SGFunction(obj_funs, sampler=sampler)
alg = GD(initial=osem_sol, objective_function=F, step_size=step_size)
# optional, you might want to disable this
alg.update_objective_interval = num_subsets
# This would need changing for the actual timing
alg.run(num_subsets * 1, callbacks=[TextProgressCallback()])