AbsFunctions added (#1976)

* AbsFunctions added
* Documentation
* set default precision to double, as necessary for SAR which is currently the only use case for FunctionOfAbs
returned type hints
* added output type hints
* convex_conjugate returns 0. when mathematically undefined to ensure compatibility with optimisation algorithms.
precision default set to double to ensure sufficient accuracy with e.g. SAR imaging
docstring updated to reflect these changes
* made decorators separate from class
* Get precision from inputted x
* Update NOTICE.txt
* Changelog

---------

Signed-off-by: Margaret Duff <[email protected]>
Signed-off-by: Hannah Robarts <[email protected]>
Co-authored-by: Margaret Duff <[email protected]>
Co-authored-by: Laura Murgatroyd <[email protected]>
Co-authored-by: hrobarts <[email protected]>

Author: 4 people

CHANGELOG.md

@@ -1,4 +1,6 @@
 * XX.X
+  - New features:
+    - Added `FunctionOfAbs` class (#1976)
   - Bug fixes:
     - Fix deprecation warning for rtol and atol in GD (#2056)
     - Removed the deprecated usage of run method in test_SIRF.py (#2070)

NOTICE.txt

@@ -75,6 +75,7 @@ Emmanuel Ferdman (2025) - 14
 Mariam Demir (2025) - 1
 Satwik Pani (2025) - 15
 Hok Shing Wong (2025) - 7
+Francis M Watson (2025) - 3 
 
 CIL Advisory Board:
 Llion Evans - 9

Wrappers/Python/cil/optimisation/functions/AbsFunction.py

@@ -0,0 +1,222 @@
+#  Copyright 2024 United Kingdom Research and Innovation
+#  Copyright 2024 The University of Manchester
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+# Authors:
+# CIL Developers, listed at: https://github.com/TomographicImaging/CIL/blob/master/NOTICE.txt
+#
+# This work has been supported by the Royal Academy of Engineering and the
+# Office of the Chief Science Adviser for National Security under the UK
+# Intelligence Community Postdoctoral Research Fellowship programme.
+#
+# Francis M Watson, University of Manchester 2024
+#
+
+
+import numpy as np
+from cil.optimisation.functions import Function
+from cil.framework import DataContainer
+from typing import Optional
+import warnings
+import logging
+
+log = logging.getLogger(__name__)
+
+class FunctionOfAbs(Function):
+    r'''A function which acts on the absolute value of the complex-valued input, 
+
+    .. math::  G(z) = H(abs(z))
+
+    This function is initialised with another CIL function, :math:`H` in the above formula. When this function is called, first the absolute value of the input is taken, and then the input is passed to the provided function.
+    
+    Included in this class is the proximal map for FunctionOfAbs.  From this, the proximal conjugate is also available from the parent CIL Function class, which is valid for this function. 
+    In the case that :math:`H` is lower semi-continuous, convex, non-decreasing and finite at the origin, (and thus `assume_lower_semi` is set to `True` in the `init`) the convex conjugate is also defined.  
+    The gradient is not defined for this function.
+    
+    For reference: see https://doi.org/10.48550/arXiv.2410.22161
+
+    Parameters
+    ----------
+    function : Function
+        Function acting on a real input, :math:`H` in the above formula.
+    assume_lower_semi : bool, default False
+        If True, assume that the function is lower semi-continuous, convex, non-decreasing and finite at the origin.
+        This allows the convex conjugate to be calculated as the monotone conjugate, which is less than or equal to the convex conjugate.
+        If False, the convex conjugate returned as 0. This is to ensure compatibility with Algorithms such as PDHG.
+
+    '''
+
+    def __init__(self, function: Function, assume_lower_semi: bool=False):
+        self._function = function
+        self._lower_semi = assume_lower_semi
+                
+        super().__init__(L=function.L)
+
+    def __call__(self, x: DataContainer) -> float:
+        call_abs = _take_abs_input(self._function.__call__)
+        return call_abs(self._function, x)
+
+    def proximal(self, x: DataContainer, tau: float, out: Optional[DataContainer]=None) -> DataContainer:
+        r'''Returns the proximal map of function :math:`\tau G`  evaluated at x
+
+        .. math:: \text{prox}_{\tau G}(x) = \underset{z}{\text{argmin}} \frac{1}{2}\|z - x\|^{2} + \tau G(z)
+
+        This is accomplished by calculating a bounded proximal map and making a change of phase,
+        :math:`prox_G(z) = prox^+_H(r) \circ \Phi` where  :math:`z = r \circ \Phi`, :math:`r = abs(z)`, :math:`\Phi = \exp(i \angle(z))`,
+        and :math:`\circ` is element-wise product.  Also define :math:`prox^+` to be the proximal map of :math:`H`  in which the minimisation carried out over the positive orthant.
+
+
+        Parameters
+        ----------
+        x : DataContainer
+            The input to the function
+        tau: scalar
+            The scalar multiplying the function in the proximal map
+        out: return DataContainer, if None a new DataContainer is returned, default None.
+            DataContainer to store the result of the proximal map
+        
+
+        Returns
+        -------
+        DataContainer, the proximal map of the function at x with scalar :math:`\tau`.
+
+        '''
+        prox_abs = _abs_and_project(self._function.proximal)
+        return prox_abs(self._function, x, tau=tau, out=out)
+
+    def convex_conjugate(self, x: DataContainer) -> float:
+        r'''
+        Evaluation of the function G* at x, where G* is the convex conjugate of function G,
+
+        .. math:: G^{*}(x^{*}) = \underset{x}{\sup} \langle x^{*}, x \rangle - G(x)
+
+        If :math:`H` is lower semi-continuous, convex, non-decreasing 
+        finite at the origin, then :math:`G^*(z*) = H^+(|z*|)`, where the monotone conjugate :math:`g^+` is
+
+        .. math:: H^+(z^*) =sup {(z, z^*) - H(z) : z >= O}
+
+        The monotone conjugate will therefore be less than or equal to the convex conjugate, 
+        since it is taken over a smaller set.  It is not available directly, but may coincide with
+        the convex conjugate, which is therefore the best estimate we have.  This is only valid for
+        real x. In other cases, a general convex conjugate is not available or defined.      
+
+
+        For reference see:  Convex Analysis, R. Tyrrell Rocakfellar, pp110-111.
+        
+        
+        Parameters
+        ----------
+        x : DataContainer
+            The input to the function
+
+        Returns
+        -------
+        float:
+        The value of the convex conjugate of the function at x.
+        '''
+
+        if self._lower_semi:
+            conv_abs = _take_abs_input(self._function.convex_conjugate)
+            return conv_abs(self._function, x)
+        else:
+            warnings.warn('Convex conjugate is not defined for this function, returning 0 for compatibility with optimisation algorithms')
+            return 0.0
+
+    def gradient(self, x):
+        '''Gradient of the function at x is not defined for this function.
+        '''
+        raise NotImplementedError('Gradient not available for this function')
+
+
+
+
+def _take_abs_input(func):
+    '''Decorator for function to act on abs of input of a method'''
+
+    def _take_abs_decorator(self, x: DataContainer, *args, **kwargs):
+        real_dtype, _ = _get_real_complex_dtype(x)
+        
+        rgeo = x.geometry.copy()
+        rgeo.dtype = real_dtype
+        r = rgeo.allocate(0)
+        r.fill(np.abs(x.as_array()).astype(real_dtype))
+        fval = func(r, *args, **kwargs)
+        return fval
+    return _take_abs_decorator
+
+
+def _abs_and_project(func):
+    '''Decorator for function to act on abs of input, 
+    with return being projected to the angle of the input.
+    Requires function return to have the same shape as input,
+    such as prox.'''
+    
+        
+    def _abs_project_decorator(self, x: DataContainer, *args, **kwargs):
+        
+        real_dtype, complex_dtype = _get_real_complex_dtype(x)
+            
+                    
+        rgeo = x.geometry.copy()
+        rgeo.dtype = real_dtype
+        r = rgeo.allocate(None)
+        r.fill(np.abs(x.as_array()).astype(real_dtype))
+        Phi = np.exp((1j*np.angle(x.array)))
+        out = kwargs.pop('out', None)
+        
+        fvals = func(r, *args, **kwargs)
+        fvals_np = fvals.as_array()
+        
+        # Douglas-Rachford splitting to find solution in positive orthant
+        if np.any(fvals_np < 0):
+            log.info('AbsFunctions: projection to +ve orthant triggered')
+            cts = 0
+            y = r.copy()
+            fvals_np = fvals.as_array()
+            while np.any(fvals_np < 0):
+                tmp = fvals_np  - 0.5*y.as_array() + 0.5*r.as_array()
+                tmp[tmp < 0] = 0.
+                y += DataContainter(tmp, y.geometry) - fvals
+                fvals = func(y, *args, **kwargs)
+                cts += 1
+                if cts > 10:
+                    fvals_np = fvals.as_array()
+                    fvals_np[fvals_np < 0] = 0.
+                    break
+
+        if out is None: 
+            out_geom = x.geometry.copy()
+            out = out_geom.allocate(None)
+        if np.isreal(x.as_array()).all():
+            out.fill( np.real(fvals_np.astype(complex_dtype)*Phi))
+        else:
+            out.fill( fvals_np.astype(complex_dtype)*Phi)
+        return out
+    return _abs_project_decorator
+
+
+def _get_real_complex_dtype(x: DataContainer):
+    '''An internal function to find the type of x and set the corresponding real and complex data types '''
+    
+    x_dtype = x.as_array().dtype
+    if np.issubdtype(x_dtype, np.complexfloating):
+        complex_dtype = x_dtype
+        complex_example = np.array([1 + 1j], dtype=x_dtype)
+        real_dtype = np.real(complex_example).dtype
+    else: 
+        real_dtype = x_dtype
+        complex_example = 1j*np.array([1], dtype=x_dtype)
+        complex_dtype = complex_example.dtype
+    return real_dtype, complex_dtype            
+            

Wrappers/Python/cil/optimisation/functions/__init__.py

@@ -39,4 +39,5 @@
 from .SGFunction import SGFunction
 from .SVRGFunction import SVRGFunction, LSVRGFunction
 from .SAGFunction import SAGFunction, SAGAFunction
+from .AbsFunction import FunctionOfAbs
 

Wrappers/Python/test/test_functions.py

@@ -32,7 +32,7 @@
                                          L1Norm, MixedL21Norm, LeastSquares, \
                                          SmoothMixedL21Norm, OperatorCompositionFunction,\
                                          Rosenbrock, IndicatorBox, TotalVariation, ScaledFunction, SumFunction, SumScalarFunction, \
-                                         WeightedL2NormSquared, MixedL11Norm, ZeroFunction, L1Sparsity
+                                         WeightedL2NormSquared, MixedL11Norm, ZeroFunction, L1Sparsity, FunctionOfAbs
 
 from cil.optimisation.functions import BlockFunction
 
@@ -2124,3 +2124,105 @@ def test_set_num_threads(self):
             N = 10
             ib.set_num_threads(N)
             assert ib.num_threads == N
+
+
+class MockFunction(Function):
+    """Mock function to test FunctionOfAbs"""
+    def __init__(self, L=1.0):
+        super().__init__(L=L)
+    
+    def __call__(self, x):
+        return x.array**2  # Simple squared function
+    
+    def proximal(self, x, tau, out=None):
+        if out is None:
+            out = x.geometry.allocate(None)
+        out.array = x.array / (1 + tau)  # Simple proximal operator
+        return out
+    
+    def convex_conjugate(self, x):
+        return 0.5 * x.array**2  # Quadratic conjugate function
+
+
+class TestFunctionOfAbs(unittest.TestCase):
+    def setUp(self):
+        self.data_complex64 = VectorData(np.array([1+2j, -3+4j, -5-6j], dtype=np.complex64))
+        self.data_complex128 = VectorData(np.array([1+2j, -3+4j, -5-6j], dtype=np.complex128))
+        self.data_real32 = VectorData(np.array([1, 2, 3], dtype=np.float32))
+        self.mock_function = MockFunction()
+        self.abs_function = FunctionOfAbs(self.mock_function)
+        
+    
+    def test_initialization(self):
+        self.assertIsInstance(self.abs_function._function, MockFunction)
+        self.assertFalse(self.abs_function._lower_semi)
+
+
+
+    def test_call(self):
+        result = self.abs_function(self.data_complex64)
+        expected = np.abs(self.data_complex64.array)**2
+        np.testing.assert_array_almost_equal(result, expected, decimal=4)
+        
+        result = self.abs_function(self.data_complex128)
+        expected = np.abs(self.data_complex128.array)**2
+        np.testing.assert_array_almost_equal(result, expected, decimal=6)
+        
+        result = self.abs_function(self.data_real32)
+        expected = np.abs(self.data_real32.array)**2
+        np.testing.assert_array_almost_equal(result, expected, decimal=6)
+        
+    def test_get_real_complex_dtype(self):
+        from cil.optimisation.functions.AbsFunction import _get_real_complex_dtype
+        real, compl = _get_real_complex_dtype(self.data_complex128)
+        self.assertEqual( real, np.float64)
+        self.assertEqual(compl, np.complex128)
+        
+        real, compl = _get_real_complex_dtype(self.data_complex64)
+        self.assertEqual( real, np.float32)
+        self.assertEqual(compl, np.complex64)
+        
+        real, compl = _get_real_complex_dtype(self.data_real32)
+        self.assertEqual( real, np.float32)
+        self.assertEqual(compl, np.complex64)
+        
+        
+        
+    
+    def test_proximal(self):
+        tau = 0.5
+        result = self.abs_function.proximal(self.data_complex128, tau)
+        expected = (np.abs(self.data_complex128.array) / (1 + tau)) * np.exp(1j * np.angle(self.data_complex128.array))
+        np.testing.assert_array_almost_equal(result.array, expected, decimal=6)
+    
+        result = self.abs_function.proximal(self.data_complex64, tau)
+        expected = (np.abs(self.data_complex64.array) / (1 + tau)) * np.exp(1j * np.angle(self.data_complex64.array))
+        np.testing.assert_array_almost_equal(result.array, expected, decimal=4)
+        
+        result = self.abs_function.proximal(self.data_real32, tau)
+        expected = (np.abs(self.data_real32.array) / (1 + tau)) * np.exp(1j * np.angle(self.data_real32.array))
+        np.testing.assert_array_almost_equal(result.array, expected, decimal=4)
+        
+        
+        
+    def test_convex_conjugate_lower_semi(self):
+        self.abs_function._lower_semi = True
+        result = self.abs_function.convex_conjugate(self.data_complex128)
+        expected = 0.5 * np.abs(self.data_complex128.array) ** 2
+        np.testing.assert_array_almost_equal(result, expected, decimal=6)
+        
+        result = self.abs_function.convex_conjugate(self.data_complex64)
+        expected = 0.5 * np.abs(self.data_complex64.array) ** 2
+        np.testing.assert_array_almost_equal(result, expected, decimal=4)
+        
+        result = self.abs_function.convex_conjugate(self.data_real32)
+        expected = 0.5 * np.abs(self.data_real32.array) ** 2
+        np.testing.assert_array_almost_equal(result, expected, decimal=4)
+        
+    
+    def test_convex_conjugate_not_implemented(self):
+        self.abs_function._lower_semi = False
+        
+        self.assertEqual(self.abs_function.convex_conjugate(self.data_real32), 0.)
+
+

Wrappers/Python/test/test_out_in_place.py

@@ -38,7 +38,7 @@
     L1Norm, L2NormSquared, MixedL21Norm, LeastSquares, \
     SmoothMixedL21Norm, OperatorCompositionFunction, \
      IndicatorBox, TotalVariation,  SumFunction, SumScalarFunction, \
-    WeightedL2NormSquared, MixedL11Norm, ZeroFunction
+    WeightedL2NormSquared, MixedL11Norm, ZeroFunction, FunctionOfAbs
 
 from cil.processors import AbsorptionTransmissionConverter, Binner, CentreOfRotationCorrector, MaskGenerator, Masker, Normaliser, Padder, \
 RingRemover, Slicer, TransmissionAbsorptionConverter, PaganinProcessor, FluxNormaliser
@@ -87,6 +87,9 @@ def setUp(self):
         ag.set_panel(10)
 
         ig = ag.get_ImageGeometry()
+        
+        ig_complex = ig.copy()
+        ig_complex.dtype = np.complex128
 
         scalar = 4
 
@@ -127,13 +130,13 @@ def setUp(self):
             (MixedL11Norm(), bg, True, True, False),
             (BlockFunction(L1Norm(),L2NormSquared()), bg, True, True, False),
             (BlockFunction(L2NormSquared(),L2NormSquared()), bg, True, True, True),
-            (L1Sparsity(WaveletOperator(ig)), ig, True, True, False)
-
-
+            (L1Sparsity(WaveletOperator(ig)), ig, True, True, False),
+            (FunctionOfAbs(TotalVariation(backend='numpy', warm_start= False), assume_lower_semi=True), ig , True, True, False),
+            (FunctionOfAbs(TotalVariation(backend='numpy', warm_start= False), assume_lower_semi=True), ig_complex , True, True, False)             
         ]
 
         np.random.seed(5)
-        self.data_arrays=[np.random.normal(0,1, (10,10)).astype(np.float32),  np.array(range(0,65500, 655), dtype=np.uint16).reshape((10,10)), np.random.uniform(-0.1,1,(10,10)).astype(np.float32)]
+        self.data_arrays=[np.random.normal(0,1, (10,10)).astype(np.float32),  np.array(range(0,65500, 655), dtype=np.uint16).reshape((10,10)), np.random.uniform(-0.1,1,(10,10)).astype(np.float32) ]
 
     def get_result(self, function, method, x, *args):
         try: