work towards 3D OI

proxtoolbox/ProxOperators/P_M.py

@@ -14,7 +14,7 @@ class P_M(ProxOperator):
         Parameters
         ----------
         config : dict - Dictionary containing the problem configuration. It must contain the following mapping:
-        'M' : array_like - nonegative real FOURIER DOMAIN CONSTRAINT
+        'M' : array_like - non-negative real FOURIER DOMAIN CONSTRAINT
         """
         self.M = config['M']
 
@@ -31,6 +31,6 @@ class P_M(ProxOperator):
         array_like - [p_M,phat_M] ,where p_M = the projection IN THE PHYSICAL (time) DOMAIN and phat_M = projection IN THE FOURIER DOMAIN 
         """
         m = self.M
-        a = np.fft.fft2(u,axes=None)
+        a = np.fft.fft2(u, axes=None)
         b = magproj(a, m)
-        return np.fft.ifft2(b,axes=None)
+        return np.fft.ifft2(b, axes=None)

proxtoolbox/ProxOperators/P_Sparsity.py

@@ -9,6 +9,7 @@ class P_Sparsity(ProxOperator):
     USAGE: p_Sparsity = P_Sparsity(config)
            p_Sparsity.work(u)
     """
+
     def __init__(self, config):
         """
         Initialization
@@ -53,7 +54,7 @@ class P_Sparsity(ProxOperator):
         -------
         p_Sparsity : array_like, the projection
         """
-        u *= self.support # apply support
+        u *= self.support  # apply support
         p_Sparsity = 0 * u
         sorting = np.argsort(abs(u), axis=None)  # gives indices of sorted array in ascending order
         indices = np.asarray([divmod(sorting[i], self.ny) for i in range(-1 * self.sparsity_parameter, 0)])
@@ -68,6 +69,7 @@ class P_Sparsity_real(P_Sparsity):
 
     order: realness, sparsity
     """
+
     def __init__(self, config):
         super(P_Sparsity_real, self).__init__(config)
 
@@ -81,9 +83,10 @@ class P_Sparsity_Symmetric(P_Sparsity):
 
     order: symmetry, sparsity
     """
+
     def __init__(self, config):
         super(P_Sparsity_Symmetric, self).__init__(config=config)
-        self.symmetry = config['symmetry_type'] # antisymmetric = -1, symmetric = 1
+        self.symmetry = config['symmetry_type']  # antisymmetric = -1, symmetric = 1
         self.symmetry_axis = config['symmetry_axis']  # -1 for last, 0 for first, 'both', 'all' or None for full flip
 
     def work(self, u):
@@ -102,9 +105,71 @@ class P_Sparsity_Symmetric_real(P_Sparsity_Symmetric):
 
     order: realness, symmetry, sparsity
     """
+
     def __init__(self, config):
         super(P_Sparsity_Symmetric_real, self).__init__(config=config)
 
     def work(self, u):
         out = super(P_Sparsity_Symmetric_real, self).work(u.real)
         return out
+
+
+class P_Sparsity_Line(P_Sparsity):
+    def __init__(self, config):
+        """
+        Sparsity in n-1 dimensions, where n is the number of dimensions in the data u
+        example: data I(x,y,z) projected to Ip(x,y) -> I(x,y,z) = 0 if Ip(x,y,z) is lower than the threshold set by sparsity
+        default projection = linear sum
+
+        Args:
+            config: identical to P_Sparse, with extra keys:
+                'unconstrained dimension': in which dimension the data is not constrained, defaults to -1
+                'dimension reduction': method for dimension reducing projection e.g.:
+                    'linear' for sum(u) (default),
+                    'abs' for sum(abs(u)),
+                    'rms' for sqrt(sum(abs(u)**2))
+
+        """
+        super(P_Sparsity_Line, self).__init__(config=config)
+        if 'unconstrained dimension' not in config:
+            self.unconstrained_axis = -1
+        else:
+            self.unconstrained_axis = config['unconstrained dimension']
+
+        if 'dimension reduction' in config and config['3d to plane projection'] == 'abs':
+            def dim_proj(arr):
+                return np.sum(abs(arr), axis=config['unconstrained dimension'])
+
+        elif 'dimension reduction' in config and config['3d to plane projection'] == 'rms':
+            def dim_proj(arr):
+                return np.sqrt(np.sum(abs(arr) ** 2, axis=config['unconstrained dimension']))
+
+        else: # default to linear
+            def dim_proj(arr):
+                return np.sum(arr, axis=config['unconstrained dimension'])
+
+        self.dim_proj = dim_proj
+
+    def work(self, u):
+        # projection
+        sparsity_input = self.dim_proj(u)
+        # calculate sparse map
+        out = super(P_Sparsity_Line, self).work(sparsity_input)
+        # broadcast to original dimensions, keeping original values
+        rolled_u = np.moveaxis(u, self.unconstrained_axis, 0)
+        new = np.moveaxis(np.where(out != 0, rolled_u, np.zeros_like(rolled_u)),
+                          0, self.unconstrained_axis)
+        return new.astype(u.dtype)
+
+
+class P_Sparsity_Line_Real(P_Sparsity_Line):
+    """
+    Combining P_Sparsity_Line with a realness constraint
+    """
+
+    def __init__(self, config):
+        super(P_Sparsity_Line_Real, self).__init__(config=config)
+
+    def work(self, u):
+        out = super(P_Sparsity_Line_Real, self).work(u.real)
+        return out