From 2de7bc3de29fdef5189a661274eb7deaa2b543dd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Luis=20Eduardo=20Ram=C3=ADrez=20Montoya?=
 <luiseduardoramirezmontoya@ciencias.unam.mx>
Date: Fri, 20 Dec 2024 10:24:48 -0600
Subject: [PATCH] Two stages in orbitdetermination(::ODProblem, ::NEOSolution,
 ::NEOParameters)

---
 src/orbitdetermination/neosolution.jl        |  3 ++
 src/orbitdetermination/orbitdetermination.jl | 30 ++++++++++++--------
 2 files changed, 21 insertions(+), 12 deletions(-)

diff --git a/src/orbitdetermination/neosolution.jl b/src/orbitdetermination/neosolution.jl
index 83844245..6cf4b2d6 100644
--- a/src/orbitdetermination/neosolution.jl
+++ b/src/orbitdetermination/neosolution.jl
@@ -152,6 +152,9 @@ end
 
 iszero(x::NEOSolution{T, U}) where {T <: Real, U <: Number} = x == zero(NEOSolution{T, U})
 
+# Number of observations
+nobs(x::NEOSolution) = length(x.res)
+
 # Override Base.min
 function min(x::NEOSolution{T, T}, y::NEOSolution{T, T}) where {T <: Real}
     nrms(x) <= nrms(y) && return x
diff --git a/src/orbitdetermination/orbitdetermination.jl b/src/orbitdetermination/orbitdetermination.jl
index 75cde7b5..03e759ea 100644
--- a/src/orbitdetermination/orbitdetermination.jl
+++ b/src/orbitdetermination/orbitdetermination.jl
@@ -295,14 +295,9 @@ function iodinitcond(A::AdmissibleRegion{T}) where {T <: Real}
 end
 
 # Update `sol` iff `_sol_` is complete and has a lower nrms
-function updatesol(sol::NEOSolution{T, T}, _sol_::NEOSolution{T, T},
-    radec::Vector{RadecMPC{T}}) where {T <: Real}
-    if length(_sol_.res) == length(radec)
-        return min(sol, _sol_)
-    else
-        return sol
-    end
-end
+updatesol(sol::NEOSolution{T, T}, _sol_::NEOSolution{T, T},
+    radec::Vector{RadecMPC{T}}) where {T <: Real} =
+    nobs(_sol_) == length(radec) ? min(sol, _sol_) : sol
 
 include("tooshortarc.jl")
 include("gaussinitcond.jl")
@@ -380,6 +375,8 @@ Fit a least squares orbit to `od` using `sol` as an initial condition.
 """
 function orbitdetermination(od::ODProblem{D, T}, sol::NEOSolution{T, T},
     params::NEOParameters{T}) where {D, T <: Real}
+    # Unfold parameters
+    varorder, significance, mode = params.jtlsorder, params.significance, params.adammode
     # Reference epoch [TDB]
     t = epoch(sol)
     jd0 = t + PE.J2000
@@ -388,12 +385,21 @@ function orbitdetermination(od::ODProblem{D, T}, sol::NEOSolution{T, T},
     # Scaling factors
     scalings = abs.(q0) ./ 10^6
     # Jet transport variables
-    dq = scaled_variables("dx", scalings; order = params.jtlsorder)
+    dq = scaled_variables("dx", scalings; order = varorder)
     # Jet Transport initial condition
     q = q0 + dq
+    # Jet Transport Least Squares
+    sol1 = jtls(od, jd0, q, sol.tracklets, params, true)
+    # Termination condition
+    (nobs(sol1) == nobs(od) && critical_value(sol1) < significance) && return sol1
     # ADAM refinement
-    _, j = findmin(@. abs(t - dtutc2days(date(od.tracklets))))
-    jd0 = _adam!(od, j, q, jd0, params)
+    _, i = findmin(@. abs(t - dtutc2days(date(od.tracklets))))
+    jd0 = _adam!(od, i, q, jd0, params)
     # Jet Transport Least Squares
-    return jtls(od, jd0, q, sol.tracklets, params, true)
+    trks = mode ? od.tracklets[:] : od.tracklets[i:i]
+    sol2 = jtls(od, jd0, q, trks, params, true)
+    # Update solution
+    sol1 = updatesol(sol1, sol2, od.radec)
+
+    return sol1
 end
\ No newline at end of file