add epsilon

MurrellGroup · Mar 2, 2024 · 7bb4af6 · 7bb4af6
1 parent bcf9088
commit 7bb4af6
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/src/MessagePassingIPA.jl b/src/MessagePassingIPA.jl
@@ -293,12 +293,14 @@ end
 
 # Normalization for vector features
 struct VectorNorm
-    #eps::Float32
+    ϵ::Float32
 end
 
+VectorNorm(; eps::Real = 1f-5) = VectorNorm(eps)
+
 function (norm::VectorNorm)(V::AbstractArray{T, 3}) where T
     @assert size(V, 1) == 3
-    V ./ sqrt.(mean(sum(abs2, V, dims = 1), dims = 2))
+    V ./ (sqrt.(mean(sum(abs2, V, dims = 1), dims = 2)) .+ norm.ϵ)
 end
 
 # L2 norm along the first dimension

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -89,10 +89,14 @@ using Test
 
     @testset "VectorNorm" begin
         norm = VectorNorm()
-        V = randn(Float32, 3, 8, 128)
+        V = randn(Float32, 3, 5, 10)
         @test norm(V) isa Array{Float32, 3}
         @test size(norm(V)) == size(V)
         @test norm(V) ≈ norm(100 * V)
         @test all(sqrt.(mean(sum(abs2, norm(V), dims = 1), dims = 2)) .≈ 1)
+
+        # zero values
+        V = zeros(Float32, 3, 5, 10)
+        @test all(!isnan, norm(V))
     end
 end