Avoid underflow and overflow in norm() (#975)

* Avoid underflow in norm()

* Perform return-type assertion at last step

* Make _init_zero() take AbstractArray

* Define maxabs_nested() to find maximum of absolute values of entries in nested array

* Remove _norm() and define norm() directly for type stability

* Use maxabs_nested() to define norm() stably without underflow

* Call _init_zero() without qualification by StaticArrays

* Define norm_scaled() and call it inside norm()

* norm(a,p) does not call norm_scaled(a,p) if p == 0, 2, Inf

* Make code more succinct

* Use normInf() instead of maxabs_nested()

* Rename aₘ to scale

* Calculate scale with mapreduce()

* Re-introduce maxabs_nested to avoid type instability

* norm() calls _norm() and _norm_scaled()

* _norm_scaled(Size(a), a, p) does not need to deal with p==2

* Add test for overflow

* Add tests for underflow and overflow for p ≥ 2

* Avoid using iszero(l) and isinf(l) for speedup

* Avoid using iszero(scale) for speedup

Author: wsshin

src/linalg.jl

@@ -219,17 +219,42 @@ end
 # Norms
 _inner_eltype(v::AbstractArray) = isempty(v) ? eltype(v) : _inner_eltype(first(v))
 _inner_eltype(x::Number) = typeof(x)
-@inline _init_zero(v::StaticArray) = float(norm(zero(_inner_eltype(v))))
+@inline _init_zero(v::AbstractArray) = float(norm(zero(_inner_eltype(v))))
 
 @inline function LinearAlgebra.norm_sqr(v::StaticArray)
     return mapreduce(LinearAlgebra.norm_sqr, +, v; init=_init_zero(v))
 end
 
+@inline maxabs_nested(a::Number) = abs(a)
+function maxabs_nested(a::AbstractArray)
+    prod(size(a)) == 0 && (return _init_zero(a))
+
+    m = maxabs_nested(a[1])
+    for j = 2:prod(size(a))
+        m = @fastmath max(m, maxabs_nested(a[j]))
+    end
+
+    return m
+end
+
+@generated function _norm_scaled(::Size{S}, a::StaticArray) where {S}
+    expr = :(LinearAlgebra.norm_sqr(a[1]/scale))
+    for j = 2:prod(S)
+        expr = :($expr + LinearAlgebra.norm_sqr(a[$j]/scale))
+    end
+
+    return quote
+        $(Expr(:meta, :inline))
+        scale = maxabs_nested(a)
+
+        scale==0 && return _init_zero(a)
+        return @inbounds scale * sqrt($expr)
+    end
+end
+
 @inline norm(a::StaticArray) = _norm(Size(a), a)
 @generated function _norm(::Size{S}, a::StaticArray) where {S}
-    if prod(S) == 0
-        return :(_init_zero(a))
-    end
+    prod(S) == 0 && return :(_init_zero(a))
 
     expr = :(LinearAlgebra.norm_sqr(a[1]))
     for j = 2:prod(S)
@@ -238,7 +263,10 @@ end
 
     return quote
         $(Expr(:meta, :inline))
-        @inbounds return sqrt($expr)
+        l = @inbounds sqrt($expr)
+
+        0<l<Inf && return l
+        return _norm_scaled(Size(a), a)
     end
 end
 
@@ -247,11 +275,31 @@ function _norm_p0(x)
     return float(norm(iszero(x) ? zero(T) : one(T)))
 end
 
+# Do not need to deal with p == 0, 2, Inf; see norm(a, p).
+@generated function _norm_scaled(::Size{S}, a::StaticArray, p::Real) where {S}
+    expr = :(norm(a[1]/scale)^p)
+    for j = 2:prod(S)
+        expr = :($expr + norm(a[$j]/scale)^p)
+    end
+
+    expr_p1 = :(norm(a[1]/scale))
+    for j = 2:prod(S)
+        expr_p1 = :($expr_p1 + norm(a[$j]/scale))
+    end
+
+    return quote
+        $(Expr(:meta, :inline))
+        scale = maxabs_nested(a)
+
+        scale==0 && return _init_zero(a)
+        p == 1 && return @inbounds scale * $expr_p1
+        return @inbounds scale * ($expr)^(inv(p))
+    end
+end
+
 @inline norm(a::StaticArray, p::Real) = _norm(Size(a), a, p)
 @generated function _norm(::Size{S}, a::StaticArray, p::Real) where {S}
-    if prod(S) == 0
-        return :(_init_zero(a))
-    end
+    prod(S) == 0 && return :(_init_zero(a))
 
     expr = :(norm(a[1])^p)
     for j = 2:prod(S)
@@ -265,17 +313,13 @@ end
 
     return quote
         $(Expr(:meta, :inline))
-        if p == Inf
-            return mapreduce(norm, max, a)
-        elseif p == 1
-            @inbounds return $expr_p1
-        elseif p == 2
-            return norm(a)
-        elseif p == 0
-            return mapreduce(_norm_p0, +, a)
-        else
-            @inbounds return ($expr)^(inv(p))
-        end
+        p == 0 && return mapreduce(_norm_p0, +, a)  # no need for scaling
+        p == 2 && return norm(a)  # norm(a) takes care of scaling
+        p == Inf && return mapreduce(norm, max, a)  # no need for scaling
+
+        l = p==1 ? @inbounds($expr_p1) : @inbounds(($expr)^(inv(p)))
+        0<l<Inf && return l
+        return _norm_scaled(Size(a), a, p)  # p != 0, 2, Inf
     end
 end
 
@@ -468,4 +512,3 @@ end
 # Some shimming for special linear algebra matrix types
 @inline LinearAlgebra.Symmetric(A::StaticMatrix, uplo::Char='U') = (checksquare(A); Symmetric{eltype(A),typeof(A)}(A, uplo))
 @inline LinearAlgebra.Hermitian(A::StaticMatrix, uplo::Char='U') = (checksquare(A); Hermitian{eltype(A),typeof(A)}(A, uplo))
-

test/linalg.jl

@@ -236,6 +236,10 @@ StaticArrays.similar_type(::Union{RotMat2,Type{RotMat2}}) = SMatrix{2,2,Float64,
     end
 
     @testset "normalization" begin
+        @test norm(SVector(0.0,1e-180)) == 1e-180  # avoid underflow
+        @test norm(SVector(0.0,1e155)) == 1e155  # avoid overflow
+        @test all([norm(SVector(0.0,1e-180), p) == 1e-180 for p = [2,3,Inf]])  # avoid underflow
+        @test all([norm(SVector(0.0,1e155), p) == 1e155 for p = [2,3,Inf]])  # avoid overflow
         @test norm(SVector(1.0,2.0,2.0)) ≈ 3.0
         @test norm(SVector(1.0,2.0,2.0),2) ≈ 3.0
         @test norm(SVector(1.0,2.0,2.0),Inf) ≈ 2.0