Mercurial > repos > AlgTools / changeset

--- a/src/Util.jl	Wed Dec 18 21:32:27 2019 +0200
+++ b/src/Util.jl	Fri Dec 27 22:04:49 2019 +0200
@@ -136,12 +136,26 @@
 
 @inline function proj_norm₂₁ball!(y, α)
     α²=α*α
-    y′=reshape(y, (size(y, 1), prod(size(y)[2:end])))
 
-    @inbounds @simd for i=1:size(y′, 2)# in CartesianIndices(size(y)[2:end])
-        n² = norm₂²(@view(y′[:, i]))
-        if n²>α²
-            y′[:, i] .*= (α/√n²)
+    if ndims(y)==3 && size(y, 1)==2
+        @inbounds for i=1:size(y, 2)
+            @simd for j=1:size(y, 3)
+                n² = y[1,i,j]*y[1,i,j]+y[2,i,j]*y[2,i,j]
+                if n²>α²
+                    v = α/√n²
+                    y[1, i, j] *= v
+                    y[2, i, j] *= v
+                end
+            end
+        end
+    else
+        y′=reshape(y, (size(y, 1), prod(size(y)[2:end])))
+
+        @inbounds @simd for i=1:size(y′, 2)# in CartesianIndices(size(y)[2:end])
+            n² = norm₂²(@view(y′[:, i]))
+            if n²>α²
+                y′[:, i] .*= (α/√n²)
+            end
         end
     end
 end