Avoid underflow and overflow in norm()

src/linalg.jl

@@ -217,17 +217,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)
@@ -236,7 +261,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
 
@@ -245,11 +273,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)
@@ -263,17 +311,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
 
@@ -466,4 +510,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

@@ -208,6 +208,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