Added check function for weighted data

JuliaNLSolvers · May 17, 2021 · 54c4636 · 54c4636
1 parent 2d3611f
commit 54c4636
Showing 1 changed file with 13 additions and 4 deletions.
diff --git a/src/curve_fit.jl b/src/curve_fit.jl
@@ -24,6 +24,15 @@ function check_data_health(xdata, ydata)
     end
 end
 
+function check_data_health(xdata, ydata, wt)
+    if any(ismissing, xdata) || any(ismissing, ydata) || any(ismissing, wt)
+        error("Data contains `missing` values and a fit cannot be performed")
+    end
+    if any(isinf, xdata) || any(isinf, ydata) || any(isinf, wt) || any(isnan, xdata) || any(isnan, ydata) || any(isnan, wt)
+        error("Data contains `Inf` or `NaN` values and a fit cannot be performed")
+    end
+end
+
 # provide a method for those who have their own Jacobian function
 function lmfit(f, g, p0::AbstractArray, wt::AbstractArray; kwargs...)
     r = f(p0)
@@ -134,7 +143,7 @@ function curve_fit(model, jacobian_model,
 end
 
 function curve_fit(model, xdata::AbstractArray, ydata::AbstractArray, wt::AbstractArray{T}, p0::AbstractArray; inplace = false, kwargs...) where T
-    check_data_health(xdata, ydata)
+    check_data_health(xdata, ydata, wt)
     # construct a weighted cost function, with a vector weight for each ydata
     # for example, this might be wt = 1/sigma where sigma is some error term
     u = sqrt.(wt) # to be consistant with the matrix form
@@ -150,7 +159,7 @@ end
 
 function curve_fit(model, jacobian_model,
             xdata::AbstractArray, ydata::AbstractArray, wt::AbstractArray{T}, p0; inplace = false, kwargs...) where T
-    check_data_health(xdata, ydata)
+    check_data_health(xdata, ydata, wt)
     u = sqrt.(wt) # to be consistant with the matrix form
 
     if inplace
@@ -165,7 +174,7 @@ function curve_fit(model, jacobian_model,
 end
 
 function curve_fit(model, xdata::AbstractArray, ydata::AbstractArray, wt::AbstractArray{T,2}, p0; kwargs...) where T
-    check_data_health(xdata, ydata)
+    check_data_health(xdata, ydata, wt)
 
     # as before, construct a weighted cost function with where this
     # method uses a matrix weight.
@@ -182,7 +191,7 @@ end
 
 function curve_fit(model, jacobian_model,
             xdata::AbstractArray, ydata::AbstractArray, wt::AbstractArray{T,2}, p0; kwargs...) where T
-    check_data_health(xdata, ydata)
+    check_data_health(xdata, ydata, wt)
 
     u = cholesky(wt).U