KullbackLeibler gradient numba

[epapoutsellis]

The following code

https://github.com/TomographicImaging/CIL/blob/c133dc465ade1bcea5a8c0f597f426afe0b1b64a/Wrappers/Python/cil/optimisation/functions/KullbackLeibler.py#L456-L458

should be

if self.mask is not None: 
     kl_gradient_mask(x.as_array(), self.b.as_array(), out_np, self.eta.as_array(), self.mask)          
else:
     kl_gradient(x.as_array(), self.b.as_array(), out_np, self.eta.as_array()) 

Also, mask is passed and is used in the numba case only to determine if mask.flat[i]>0 and then compute the gradient, proximal, proximal conjugate of KL. I think is better if we extract indices of positive mask, e.g., self.ind = np.flatnonzero(self.mask)

and then

    @njit(parallel=True) 
    def kl_gradient_mask(x, b, out, eta, ind):
        for i in prange(len(ind)):
                out.flat[ind[i]] = 1 - b.flat[ind[i]]/(x.flat[ind[i]] + eta.flat[ind[i]])

We can also decorate the function above with @njit(void(float32[:], float32[:], float32[:], float32[:], int64[:]), parallel=True) provided that we have flatten our input arrays and we reshape at the final step out.fill(out_np.reshape(out_np.shape))

Also, parallel=True seems not to work with STIR but maybe is a problem with my SIRF installation @KrisThielemans

[KrisThielemans]

Also, parallel=True seems not to work with STIR but maybe is a problem with my SIRF installation @KrisThielemans

As far as I can see, this is all with numpy arrays, so has nothing to do with SIRF.

By the way, I'm amazed that this numba stuff would be better than

out.flat[ind] = 1 - b.flat[ind]/(x.flat[ind] + eta.flat[ind])

but I know nothing about numba really.

[epapoutsellis]

@paskino the error that I get is AsyncIOLoopKernelRestarter: restarting kernel (1/5), keep random ports

I use the notebook from SIRF-Exercises

First, you need to make some changes: 1) SIRF Prior to be compatible with CIL Function, in here 2) KullbackLeibler gradient method using numba, in here

Then, you continue with the notebook (download data etc) and you stopped after you estimate the random counts. No need to run all the nb.

Then, you run the following code

acq_model.set_up(acq_data, initial_image*0.)
mask = acq_model.direct(initial_image.get_uniform_copy(1.0))
tmp_f = KullbackLeibler(b = acq_data, eta = randoms, mask=mask, backend="numba") 

f1 = OperatorCompositionFunction(tmp_f, acq_model)

f2 = pet.RelativeDifferencePrior()
f2.set_epsilon(1e-5)
f2.set_penalisation_factor(10.)
f2.set_up(image)

# from cil.optimisation.utilities import ArmijoStepSize
# ss = ArmijoStepSize(initial=1., rho=0.5, c=0.5, iterations=100)
gd = ISTA(initial = initial_image*0., f = f1 + f2, g=IndicatorBox(lower=0.), step_size = 1e-5, 
             update_objective_interval=1, max_iteration=32)
gd.run(2, verbose=1)