Preconditioning Explicit PDHG

The following code initially implemented here fails due to possible numba implementation in MixedL21Norm.

# Explicit PDHG
f1 = 0.5 * L2NormSquared(b=data2D)
f2 = alpha * MixedL21Norm()
F = BlockFunction(f1, f2)
Grad = GradientOperator(ig2D)
K = BlockOperator(A, Grad)
G = IndicatorBox(lower=0)
normK = K.norm()

tmp_sigma = K.range_geometry().allocate()

tmp_sigma.get_item(0).fill(A.direct(A.domain_geometry().allocate(1)))
tmp_sigma.get_item(1).get_item(0).fill(Grad.domain_geometry().allocate(2))
tmp_sigma.get_item(1).get_item(1).fill(Grad.domain_geometry().allocate(2))
sigma = 1./tmp_sigma

tmp_tau = K.domain_geometry().allocate()
tmp_tau.fill(A.adjoint(A.range_geometry().allocate(1))
             + Grad.domain_geometry().allocate(2) + Grad.domain_geometry().allocate(2))
tau = 1./tmp_tau

pdhg_explicit_tv_precond = PDHG(f=F, g=G, operator=K, sigma=sigma, tau=tau,
            max_iteration = 10000,
            update_objective_interval=1)    
pdhg_explicit_tv_precond.run(verbose=1)

File ~/miniconda3/envs/cil_sirf/lib/python3.10/site-packages/cil/optimisation/functions/MixedL21Norm.py:73, in _proximal_step_numpy(tmp, tau)
     59 '''Numpy implementation of a step in the calculation of the proximal of MixedL21Norm
     60 
     61 Parameters:
   (...)
     69 A DataContainer where we have substituted nan with 0.
     70 '''
     71 # Note: we divide x by tau so the cases of tau both scalar and 
     72 # DataContainers run
---> 73 tmp /= np.abs(tau, dtype=np.float32)
     74 res = tmp - 1
     75 res.maximum(0.0, out=res)

UFuncTypeError: Cannot cast ufunc 'absolute' input from dtype('O') to dtype('float32') with casting rule 'same_kind'

Implicit PDHG with preconditioning works with no problems.

# implicit PDHG
F1 = 0.5 * L2NormSquared(b=data2D)    
K = A
sigma1 = 1./A.direct(A.domain_geometry().allocate(1.))
tau1 = 1./A.adjoint(A.range_geometry().allocate(1.))
G = alpha * TotalVariation()
pdhg_implicit_tv_precond = PDHG(f=F1, g=G, operator=K, sigma=sigma1, tau=tau1,
            max_iteration = 1000,
            update_objective_interval=1)    
pdhg_implicit_tv_precond.run(very_verbose=True)

Surely numba has nothing to do with it as it is the _proximal_numpy method that fails?

If I follow the code correctly, it fails at a call to proximal of MixedL21Norm. This in turn calls _proximal_numpy.

https://github.com/TomographicImaging/CIL/blob/2b3e2cd393c7ae15e27a7539c144c5197afbacef/Wrappers/Python/cil/optimisation/functions/MixedL21Norm.py#L163

where

https://github.com/TomographicImaging/CIL/blob/2b3e2cd393c7ae15e27a7539c144c5197afbacef/Wrappers/Python/cil/optimisation/functions/MixedL21Norm.py#L149

I suppose this fails during an update of PDHG, and I presume the error creeps up at this line

https://github.com/TomographicImaging/CIL/blob/2b3e2cd393c7ae15e27a7539c144c5197afbacef/Wrappers/Python/cil/optimisation/algorithms/PDHG.py#L382

So, if I understand correctly the thing that fails is

np.abs(pdhg_explicit_tv_precond.sigma, dtype=np.float32)

because sigma is a BlockDataContainer and np.abs doesn't know what to do with it, and our containers are not NEP-13 and 18 compliant, see https://github.com/TomographicImaging/CIL/issues/1156#issuecomment-1483037384.

Lastly, which version of CIL are you running?

Thanks @epapoutsellis . I think I now fixed my test: I was not passing to the method the right objects, also in view of where the code fails

https://github.com/TomographicImaging/CIL/blob/2b3e2cd393c7ae15e27a7539c144c5197afbacef/Wrappers/Python/cil/optimisation/algorithms/PDHG.py#L382

Now my test code is the following:

from cil.optimisation.functions.MixedL21Norm import _proximal_step_numpy
from cil.framework import ImageGeometry, AcquisitionGeometry
from cil.optimisation.operators import LinearOperator, GradientOperator

ag = AcquisitionGeometry.create_Parallel2D()\
    .set_panel(3)\
    .set_angles(np.linspace(0,360,4), angle_unit='degree')

ig = ag.get_ImageGeometry()

L = LinearOperator(domain_geometry=ig, range_geometry=ag)
L.set_norm(1)
Grad = GradientOperator(ig)
K = BlockOperator(L, Grad)

tmp_sigma = K.range_geometry().allocate(1)

sigma1 = 1./tmp_sigma

tmp = K.range_geometry().allocate(1.5)

a = _proximal_step_numpy(tmp, sigma1)

It does fail at calculating the pnorm(2) of the sigma1 BlockDataContainer, see this new unit test which fails with ValueError: Incompatible for operation add. BTW if the containers in the block have the same shape, this works.

`pnorm`: the new (old) bug

Looking at the definition of pnorm in BlockDataContainer it is not clear what it is supposed to return. The code is making a reduction for both pnorm(1) and pnorm(2).

https://github.com/TomographicImaging/CIL/blob/c133dc465ade1bcea5a8c0f597f426afe0b1b64a/Wrappers/Python/cil/framework/BlockDataContainer.py#L467-L475

pnorm(1) is accumulating in a float
pnorm(2) is accumulating in a container of the same type as the first element in the BlockDataContainer

In practice, in this case both methods will accumulate in a BlockDataContainer of the same type the range of the GradientOperator because the BlockDataContainer algebra takes precedence w.r.t. DataContainer algebra, thanks to the __container_priority__ class variable.

As we can use add with different types of DataContainers, the only important thing is that the shape of the containers are consistent for the pnorm to return something and not to crash.

There is no docstring attached to the pnorm so the question is: what is pnorm supposed to return?

This problem has already arisen https://github.com/TomographicImaging/CIL/issues/474

TomographicImaging / CIL