_proximal_numpy accepts BDC

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

@@ -79,37 +79,44 @@ def _proximal_step_numpy(arr, tau):
         # DataContainer or BlockDataContainer
         tmp = tau.abs()
 
-    try:
-        # if tau is a DataContainer
-        arr /= tmp
-        res = arr - 1
-        res.maximum(0.0, out=res)
-        res /= arr
-
-        arr *= tmp
-
-        resarray = res.as_array()
-        resarray[np.isnan(resarray)] = 0
-        res.fill(resarray)
-        return res
-    except AttributeError:
+    if isinstance(tau, BlockDataContainer):
         # if tau is a BlockDataContainer
         # This is the old 21.3.1 CIL code for the Proximal of MixedL21Norm
         # https://github.com/TomographicImaging/CIL/blob/6edf48aee6e1dcf81a933d4d57dea661509fcba3/Wrappers/Python/cil/optimisation/functions/MixedL21Norm.py#L92-L104
 
         tmp = (arr/tau).pnorm(2)
         res = (tmp - 1)
         res.maximum(0.0, out=res)
-        res.multiply(arr, out=res)
+        res = arr * res
         res.divide(tmp, out=res)
 
-        for el in res.containers:
+        _remove_nans_bdc(res, 0)
 
-            elarray = el.as_array()
-            elarray[np.isnan(elarray)]=0
-            el.fill(elarray)
+        return res
+    else:
+        # if tau is a DataContainer
+        arr /= tmp
+        res = arr - 1
+        res.maximum(0.0, out=res)
+        res /= arr
+
+        arr *= tmp
 
+        _remove_nans(res, 0)
         return res
+
+def _remove_nans(res, value):
+    resarray = res.as_array()
+    resarray[np.isnan(resarray)] = value
+    res.fill(resarray)
+    # return res
+def _remove_nans_bdc(bdc, value):
+    for el in bdc.containers:
+        if isinstance(el, BlockDataContainer):
+            # recursive call
+            _remove_nans_bdc(el, value)
+        else:
+            _remove_nans(el, value)
 
 class MixedL21Norm(Function):