Keep track of the unit when fitting

MineralsCloud · Sep 30, 2019 · 9f61a51 · 9f61a51
1 parent 0290ea9
commit 9f61a51
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/src/NonlinearFitting.jl b/src/NonlinearFitting.jl
@@ -74,9 +74,15 @@ function lsqfit(
     kwargs...,
 )
     E = typeof(eos).name.wrapper
+    units = unit.(Collections.fieldvalues(eos))
     trial_params = map(ustrip, Collections.fieldvalues(upreferred(eos)))
     xdata, ydata = map(ustrip ∘ upreferred, xdata), map(ustrip ∘ upreferred, ydata)
-    return lsqfit(form, E(trial_params...), xdata, ydata, kwargs...)
+    result = lsqfit(form, E(trial_params...), xdata, ydata, kwargs...)
+    if result isa EquationOfState
+        E(
+            [uconvert(u, Collections.fieldvalues(result)[i] * upreferred(u)) for (i, u) in enumerate(units)]...
+        )
+    end
 end  # function lsqfit
 
 end