From 2797b5c239340892ad13371b4c5718807f63acad Mon Sep 17 00:00:00 2001
From: Elias Carvalho <eliascarvdev@gmail.com>
Date: Wed, 23 Oct 2024 11:21:49 -0300
Subject: [PATCH] Apply suggestions

---
 src/fitting.jl | 35 +++++++++++++++++++++--------------
 1 file changed, 21 insertions(+), 14 deletions(-)

diff --git a/src/fitting.jl b/src/fitting.jl
index 81473fd..e0e7e6d 100644
--- a/src/fitting.jl
+++ b/src/fitting.jl
@@ -148,9 +148,18 @@ function fit_impl(
   # evaluate weights
   w = _weights(algo.weightfun, x, n)
 
+  # objective function
+  function J(θ)
+    γ = V(ball(θ[1]), sill=θ[2], nugget=θ[3])
+    sum(i -> w[i] * (γ(x′[i]) - y′[i])^2, eachindex(w, x′, y′))
+  end
+
   # linear constraint (sill ≥ nugget)
   L(θ) = θ[2] ≥ θ[3] ? 0.0 : θ[3] - θ[2]
 
+  # penalty for linear constraint (J + λL)
+  λ = sum(yᵢ -> yᵢ^2, y′)
+
   # maximum range, sill and nugget
   xmax = maximum(x′)
   ymax = maximum(y′)
@@ -173,7 +182,7 @@ function fit_impl(
   u = [rᵤ, sᵤ, nᵤ]
 
   # solve optimization problem
-  θ, ϵ = _optimize(θ -> V(ball(θ[1]), sill=θ[2], nugget=θ[3]), x′, y′, w, L, l, u, θₒ)
+  θ, ϵ = _optimize(J, L, λ, l, u, θₒ)
 
   # optimal variogram (with units)
   γ = V(ball(θ[1] * ux), sill=θ[2] * uy, nugget=θ[3] * uy)
@@ -218,11 +227,20 @@ function fit_impl(
   # evaluate weights
   w = _weights(algo.weightfun, x, n)
 
+  # objective function
+  function J(θ)
+    γ = V(scaling=θ[1], nugget=θ[2], exponent=θ[3])
+    sum(i -> w[i] * (γ(x′[i]) - y′[i])^2, eachindex(w, x′, y′))
+  end
+
   # linear constraints
   # 1. scaling ≥ 0
   # 2. 0 ≤ exponent ≤ 2
   L(θ) = θ[1] ≥ 0.0 ? 0.0 : -θ[1] + θ[3] ≥ 0.0 ? 0.0 : -θ[3] + 2.0 ≥ θ[3] ? 0.0 : θ[3] - 2.0
 
+  # penalty for linear constraint (J + λL)
+  λ = sum(yᵢ -> yᵢ^2, y′)
+
   # maximum scaling, nugget and exponent
   ymax = maximum(y′)
   smax = isnothing(maxscaling′) ? ymax : maxscaling′
@@ -244,7 +262,7 @@ function fit_impl(
   u = [sᵤ, nᵤ, eᵤ]
 
   # solve optimization problem
-  θ, ϵ = _optimize(θ -> V(scaling=θ[1], nugget=θ[2], exponent=θ[3]), x′, y′, w, L, l, u, θₒ)
+  θ, ϵ = _optimize(J, L, λ, l, u, θₒ)
 
   # optimal variogram (with units)
   γ = V(scaling=θ[1] * uy, nugget=θ[2] * uy, exponent=θ[3])
@@ -254,20 +272,9 @@ end
 
 _weights(f, x, n) = isnothing(f) ? n / sum(n) : map(xᵢ -> ustrip(f(xᵢ)), x)
 
-function _optimize(gamma, x, y, w, L, l, u, θₒ)
-  # objective function
-  function J(θ)
-    γ = gamma(θ)
-    sum(i -> w[i] * (γ(x[i]) - y[i])^2, eachindex(w, x, y))
-  end
-
-  # penalty for linear constraint (J + λL)
-  λ = sum(yᵢ -> yᵢ^2, y)
-
-  # solve optimization problem
+function _optimize(J, L, λ, l, u, θₒ)
   sol = Optim.optimize(θ -> J(θ) + λ * L(θ), l, u, θₒ)
   ϵ = Optim.minimum(sol)
   θ = Optim.minimizer(sol)
-
   θ, ϵ
 end