Implement l2 loss

Author: foldfelis

example/FlowOverCircle/src/FlowOverCircle.jl

@@ -31,7 +31,7 @@ function train()
         Dense(64, 1),
     ) |> device
 
-    loss(𝐱, 𝐲) = sum(abs2, 𝐲 .- m(𝐱)) / size(𝐱)[end]
+    loss(𝐱, 𝐲) = l₂loss(m(𝐱), 𝐲)
 
     opt = Flux.Optimiser(WeightDecay(1f-4), Flux.ADAM(1f-3))
 

src/NeuralOperators.jl

@@ -14,6 +14,7 @@ module NeuralOperators
 
     include("Transform/Transform.jl")
     include("operator_kernel.jl")
+    include("loss.jl")
     include("model.jl")
     include("DeepONet.jl")
     include("subnets.jl")

src/loss.jl

@@ -0,0 +1,8 @@
+export l₂loss
+
+function l₂loss(𝐲̂, 𝐲; agg=mean, is_grid_normalized=true)
+    feature_dims = 2:(ndims(𝐲)-1)
+    loss = agg(.√(sum(abs2, 𝐲̂-𝐲, dims=feature_dims)))
+
+    return is_grid_normalized ? loss/prod(feature_dims) : loss
+end