RFC: Add norm(A, p; dims)

stdlib/LinearAlgebra/src/generic.jl

@@ -657,6 +657,136 @@ julia> norm(-2, Inf)
 end
 norm(::Missing, p::Real=2) = missing
 
+# With dims keyword
+norm0_dims!(B, A) = count!(!iszero, B, A)
+norm1_dims!(B, A) = Base.mapreducedim!(norm, +, B, A)
+normInf_dims!(B, A) = Base.mapreducedim!(norm, max, B, A)
+normMinusInf_dims!(B, A) = Base.mapreducedim!(norm, min, B, A)
+
+function norm2_dims!(B::AbstractArray, A::AbstractArray, dims)
+    sum!(LinearAlgebra.norm_sqr, B, A)
+    map!(sqrt, B, B)
+    # Checking whether `A` is safe for the fast path is slower than taking it,
+    # so check and fix any zero/infinite answers afterwards:
+    _norm_dims_check!(B, A, dims, LinearAlgebra.norm2)
+    B
+end
+
+function normp_dims!(B::AbstractArray, A::AbstractArray, p::Real, dims)
+    if p == 0.5
+        sum!(sqrt ∘ norm, B, A)
+        map!(abs2, B, B)
+    elseif p == 3
+        sum!(x -> norm(x)^3, B, A)
+        map!(cbrt, B, B)
+    else
+        sum!(x -> norm(x)^p, B, A)
+        invp = inv(p)
+        map!(x -> x^invp, B, B)
+    end
+    _norm_dims_check!(B, A, dims, LinearAlgebra.normp, p)
+    B
+end
+
+function _norm_dims_check!(B, A, dims, norm, args...)
+    if A isa AbstractVecOrMat && dims == 1
+        for (i,x) in zip(eachindex(B), eachcol(A))
+            !iszero(B[i]) && isfinite(B[i]) && continue
+            B[i] = norm(x, args...)
+        end
+    elseif A isa AbstractVecOrMat && dims == 2
+        for (i,x) in zip(eachindex(B), eachrow(A))
+            !iszero(B[i]) && isfinite(B[i]) && continue
+            B[i] = norm(x, args...)
+        end
+    elseif all(y -> !iszero(y) && isfinite(y), B)
+        for I in CartesianIndices(B)
+            !iszero(B[I]) && isfinite(B[I]) && continue
+            # Unfortunately `eachslice(A; dims)` is not what we need here.
+            # This path is not type-stable, so quite slow, but hopefully rare.
+            J = ntuple(d -> d in dims ? Colon() : I[d], ndims(A))
+            B[I] = norm(view(A, J...), args...)
+        end
+    end
+    B
+end
+
+"""
+    norm(A::AbstractArray, [p]; dims)
+
+Find the vector `norm`s of slices of a given array.
+
+The result has the same size as `sum(A; dims)`, containing `norm(A[i,:,j,...], p)`
+for each possible `i,j,...`, with colons at dimensions `d ∈ dims`.
+
+The result is always a floating point array. To avoid this when `eltype(A) <: Integer`,
+0-norm is `count(!iszero, A; dims)`, 1-norm is `sum(abs, A; dims)`,
+and the `Inf`-norm is `maximum(abs, A; dims)`.
+
+!!! compat "Julia 1.10"
+    Methods taking keyword `dims` require at least Julia 1.10.
+
+# Examples
+```jldoctest
+julia> v = [3, -2, 6]; m = hcat(v, -v, [3,0,0], [4,4,4])
+3×4 Matrix{Int64}:
+  3  -3  3  4
+ -2   2  0  4
+  6  -6  0  4
+
+julia> norm(v)
+7.0
+
+julia> norm(m; dims=1)
+1×4 Matrix{Float64}:
+ 7.0  7.0  3.0  6.9282
+
+julia> map(norm, eachcol(m))  # same contents as dims=1
+4-element Vector{Float64}:
+ 7.0
+ 7.0
+ 3.0
+ 6.928203230275509
+
+julia> norm(v, 1), norm(m, 1; dims=1), sum(abs, m; dims=1)
+(11.0, [11.0 11.0 3.0 12.0], [11 11 3 12])
+
+julia> norm(m, 1; dims=2)
+3×1 Matrix{Float64}:
+ 13.0
+  8.0
+ 16.0
+
+julia> norm(v, Inf), norm(m, Inf; dims=1), maximum(abs, m; dims=1)
+(6.0, [6.0 6.0 3.0 4.0], [6 6 3 4])
+
+julia> norm(m, 0; dims=2), count(!iszero, m; dims=2)
+([4.0; 3.0; 3.0;;], [4; 3; 3;;])
+
+julia> norm([1e-200, 0, 1e-300]; dims=1)  # avoids underflow & overflow
+1-element Vector{Float64}:
+ 1.0e-200
+```
+"""
+function norm(A::AbstractArray, p::Real=2; dims=:)
+    dims isa Colon && return invoke(norm, Tuple{Any, Real}, A, p)
+    B = Base.reducedim_init(norm, +, A, dims)
+    if p == 2
+        norm2_dims!(B, A, dims)
+    elseif p == 1
+        norm1_dims!(B, A)
+    elseif p == Inf
+        normInf_dims!(B, A)
+    elseif p == 0
+        norm0_dims!(B, A)
+    elseif p == -Inf
+        normMinusInf_dims!(B, A)
+    else
+        normp_dims!(B, A, p, dims)
+    end
+    B
+end
+
 # special cases of opnorm
 function opnorm1(A::AbstractMatrix{T}) where T
     require_one_based_indexing(A)
@@ -1816,7 +1946,11 @@ end
 
 Normalize the array `a` in-place so that its `p`-norm equals unity,
 i.e. `norm(a, p) == 1`.
+
 See also [`normalize`](@ref) and [`norm`](@ref).
+
+Note that `normalize!(a)` does not take a `dims` keyword. You can write
+`a ./= norm(a, p; dims)`, although this omits some steps for avoiding overflow.
 """
 function normalize!(a::AbstractArray, p::Real=2)
     nrm = norm(a, p)
@@ -1837,16 +1971,40 @@ end
     return a
 end
 
+@inline function __normalize!(a::AbstractArray, nrm::AbstractArray)
+    δ = inv(prevfloat(typemax(eltype(nrm))))
+    if all(≥(δ), nrm)
+        nrm = inv.(nrm)  # know this is mutable as `norm2_dims!` etc wrote into it
+        a .= a .* nrm
+    else
+        εδ = eps(one(nrm))/δ
+        nrm .= inv.(nrm .* εδ)
+        a .= (a .* εδ) .* nrm
+    end
+    a
+end
+
+
 """
-    normalize(a, p::Real=2)
+    normalize(a, p::Real=2; [dims])
 
 Normalize `a` so that its `p`-norm equals unity,
-i.e. `norm(a, p) == 1`. For scalars, this is similar to sign(a),
-except normalize(0) = NaN.
+i.e. `norm(a, p) == 1`. For scalars, this is similar to `sign(a)`,
+except `normalize(0) == NaN`.
+
 See also [`normalize!`](@ref), [`norm`](@ref), and [`sign`](@ref).
 
 # Examples
 ```jldoctest
+julia> normalize(3, 1)
+1.0
+
+julia> normalize(-8, 1)
+-1.0
+
+julia> normalize(0, 1)
+NaN
+
 julia> a = [1,2,4];
 
 julia> b = normalize(a)
@@ -1866,43 +2024,62 @@ julia> c = normalize(a, 1)
 
 julia> norm(c, 1)
 1.0
+```
+"""
+normalize(x) = x / norm(x)
+normalize(x, p::Real) = x / norm(x, p)
 
-julia> a = [1 2 4 ; 1 2 4]
-2×3 Matrix{Int64}:
- 1  2  4
- 1  2  4
+"""
+    normalize(A::AbstractArray, p::Real=2; dims)
 
-julia> norm(a)
-6.48074069840786
+Return a similar array for which `norm(result, p; dims)` is everywhere 1.
+Equivalent to `mapslices(x -> normalize(x, p), mat; dims)`, but usually more efficient.
 
-julia> normalize(a)
+!!! compat "Julia 1.11"
+    The `dims` keyword requires at least Julia 1.11.
+
+# Examples
+```jldoctest
+julia> mat = [1 2 5 ; 1 2 5]
+2×3 Matrix{Int64}:
+ 1  2  5
+ 1  2  5
+
+julia> cols2 = normalize(mat; dims=1)
 2×3 Matrix{Float64}:
- 0.154303  0.308607  0.617213
- 0.154303  0.308607  0.617213
+ 0.707107  0.707107  0.707107
+ 0.707107  0.707107  0.707107
 
-julia> normalize(3, 1)
-1.0
+julia> norm(cols2; dims=1)
+1×3 Matrix{Float64}:
+ 1.0  1.0  1.0
 
-julia> normalize(-8, 1)
--1.0
+julia> rows1 = normalize(mat, 1; dims=2)
+2×3 Matrix{Float64}:
+ 0.125  0.25  0.625
+ 0.125  0.25  0.625
 
-julia> normalize(0, 1)
-NaN
+julia> sum(rows1, dims=2)
+2×1 Matrix{Float64}:
+ 1.0
+ 1.0
+
+julia> rows1 == mapslices(x -> normalize(x, 1), mat; dims=2)
+true
 ```
 """
-function normalize(a::AbstractArray, p::Real = 2)
-    nrm = norm(a, p)
+function normalize(a::AbstractArray, p::Real = 2; dims=:)
+    nrm = norm(a, p; dims)
     if !isempty(a)
-        aa = copymutable_oftype(a, typeof(first(a)/nrm))
+        aa = copymutable_oftype(a, typeof(first(a)/first(nrm)))
         return __normalize!(aa, nrm)
     else
+        @assert dims isa Colon  # else `norm(a; dims)` is already an error
         T = typeof(zero(eltype(a))/nrm)
         return T[]
     end
 end
 
-normalize(x) = x / norm(x)
-normalize(x, p::Real) = x / norm(x, p)
 
 """
     copytrito!(B, A, uplo) -> B

stdlib/LinearAlgebra/test/dense.jl

@@ -415,6 +415,45 @@ end
         @test norm(Float64[1e-300, 1], -3)*1e300 ≈ 1
         @test norm(Float64[1e300, 1], 3)*1e-300 ≈ 1
     end
+
+    @testset "norm(array; dims)" begin
+        _norm(A, p=2; dims) = mapslices(x -> norm(x, p), A; dims)
+
+        A = randn(5, 7)
+        @test norm(A; dims=2) ≈ _norm(A; dims=2)
+
+        B = zeros(3, 4, 2)
+        B[1] = pi
+        B[2] = -pi
+        B[2,2,2] = Inf
+        B[3,3,1] = NaN
+
+        for p in [-Inf,-2,-1.5,-1,-0.5,0.5,1,1.5,2,Inf]
+            @test norm(A, p; dims=1) ≈ _norm(A, p; dims=1) broken=(p==-Inf)  # TODO fix -Inf
+            @test norm(B, p; dims=3) ≈ _norm(B, p; dims=3) nans=true
+            @test norm(B, p; dims=(1,3)) ≈ _norm(B, p; dims=(1,3)) nans=true
+        end
+
+        for p in [-Inf,-1,1,2,3,Inf]  #10234 above
+            @test isnan(only(norm([0, NaN], p; dims=1)))
+            @test isnan(only(norm([NaN32, 0], p; dims=1)))
+            @test isnan(only(norm([im, NaN], p; dims=1)))
+        end
+
+        for p in [-2,-1,1,1.5,2,Inf] #10234 above
+            @test only(norm([Inf], p; dims=1)) == Inf
+        end
+        @test only(norm([Inf], 0; dims=1)) == norm([Inf], 0)
+        @test_broken only(norm([Inf], -Inf; dims=1)) == norm([Inf], -Inf)  # TODO fix this!
+
+        @test only(norm([0, 1e300]; dims=1)) ≈ norm([0, 1e300])
+        @test only(norm([0, 1e-300]; dims=1)) ≈ norm([0, 1e-300])
+        @test only(norm([0f0, 1f30]; dims=1)) ≈ norm([0f0, 1f30])
+        @test only(norm([0f0, 1f-30]; dims=1)) ≈ norm([0f0, 1f-30])
+
+        @test only(norm([1, 1e300], 3; dims=1))*1e-300 ≈ norm([1, 1e300], 3)*1e-300
+        @test only(norm([1, 1e-300], -3; dims=1))*1e300 ≈ norm([1, 1e-300], -3)*1e300
+    end
 end
 
 ## Issue related tests

stdlib/LinearAlgebra/test/generic.jl

@@ -425,6 +425,33 @@ end
                        [0.0 + 0.0im, 0.0 - 0.0im, NaN + NaN*im]))
 end
 
+@testset "normalize(array; dims)" begin
+    v = randn(10)
+    @test normalize(v) ≈ normalize(v; dims=1)
+    @test sign.(v) ≈ normalize(v; dims=2)
+
+    @testset for p in [-Inf,-1.5,-1,0,1,1.5,2,3,Inf]
+        @test normalize(v, p) ≈ normalize(v, p; dims=1)  broken=(p==-Inf)
+        @test_broken normalize(fill(-2pi), p) ≈ normalize(fill(-2pi), p; dims=1)  broken=(p==-Inf)
+        # TODO fix zero-arrays
+    end
+
+    x = randn(3, 5, 2)
+    @testset for p in [-1.5,-1,0,1,1.5,2,3,Inf]
+        for dims in (2, (1,3))
+            y = normalize(x, p; dims)  # error for p = -Inf
+            z = norm(y, p; dims)
+            if p==0
+                @test z ≈ zero.(z) .+ prod(size(x, d) for d in dims)
+            else
+                @test z ≈ one.(z)
+            end
+        end
+    end
+
+    # TODO test Inf, NaN, overflow cases
+end
+
 @testset "Issue 14657" begin
     @test det([true false; false true]) == det(Matrix(1I, 2, 2))
 end