More stats fixes

Project.toml

@@ -1,6 +1,6 @@
 name = "CurveFit"
 uuid = "5a033b19-8c74-5913-a970-47c3779ef25c"
-version = "1.9.4"
+version = "1.10.0"
 authors = ["Paulo José Saiz Jabardo <pjabardo@gmail.com>, Avik Pal <avikpal@mit.edu> and contributors"]
 
 [deps]

docs/src/_changelog.md

@@ -7,7 +7,26 @@ CurrentModule = CurveFit
 This documents notable changes in CurveFit.jl. The format is based on [Keep a
 Changelog](https://keepachangelog.com).
 
-## [v1.9.4] - 2026-06-27
+## [v1.10.0] - 2026-06-28
+
+### Added
+- Added support for passing arrays to solutions of
+  [`KingCurveFitAlgorithm`](@ref), [`ModifiedKingCurveFitAlgorithm`](@ref), and
+  [`RationalPolynomialFitAlgorithm`](@ref) ([#115]).
+
+### Fixed
+- Fixed the statistics functions for [`ExpSumFitAlgorithm`](@ref) to handle
+  `withconst=true` correctly ([#115]).
+- Fixed `sol(x::Number)` of [`ExpSumFitAlgorithm`](@ref) to return scalars for
+  consistency with the other solutions ([#115]).
+- Previously the original [`CurveFitProblem`](@ref) from a nonlinear fit was
+  always copied into the solution, even after calling `reinit!(cache,
+  ...)`. This meant that the statistics functions like [`margin_error()`](@ref)
+  etc would incorrectly return values for the original problem rather than the
+  one actually solved. Now `sol.prob` is reconstructed using the correct inputs
+  ([#115]).
+
+## [v1.9.4] - 2026-06-28
 
 ### Changed
 - Previously nonlinear fits would compute the residuals as `ŷ − y`, they are now

src/expsumfit.jl

@@ -222,8 +222,15 @@ function CommonSolve.solve!(cache::ExpSumFitCache)
     )
 end
 
-function (sol::CurveFitSolution{<:ExpSumFitAlgorithm})(x)
+function __eval_expsum(sol::CurveFitSolution{<:ExpSumFitAlgorithm}, x)
     (; k, p, λ) = sol.u
-    y = k .+ sum(exp.(x * λ') .* p'; dims = 2)
-    return real.(vec(y))
+    y = k[]
+    @inbounds for i in eachindex(p, λ)
+        y += p[i] * exp(λ[i] * x)
+    end
+    return real(y)
+end
+
+function (sol::CurveFitSolution{<:ExpSumFitAlgorithm})(x)
+    return __eval_expsum.(Ref(sol), x)
 end

src/king.jl

@@ -31,9 +31,9 @@ function CommonSolve.solve!(cache::KingFitCache)
     return CurveFitSolution(cache.alg, (A, B), resid, cache.prob, ReturnCode.Success)
 end
 
-function (sol::CurveFitSolution{<:KingCurveFitAlgorithm})(x::Number)
+function (sol::CurveFitSolution{<:KingCurveFitAlgorithm})(x)
     A, B = sol.u
-    return ((x^2 - A) / B)^2
+    return @. ((x^2 - A) / B)^2
 end
 
 # Common Solve Interface for ModifiedKingCurveFitAlgorithm
@@ -108,6 +108,6 @@ function CommonSolve.solve!(cache::ModifiedKingFitCache)
     return CurveFitSolution(cache.alg, sol.u, sol.resid, cache.prob, sol.retcode, sol.original)
 end
 
-function (sol::CurveFitSolution{<:ModifiedKingCurveFitAlgorithm})(x::Number)
-    return ((x .^ 2 .- sol.u[1]) ./ sol.u[2]) .^ (1 ./ sol.u[3])
+function (sol::CurveFitSolution{<:ModifiedKingCurveFitAlgorithm})(x)
+    return @. ((x^2 - sol.u[1]) / sol.u[2])^(1 / sol.u[3])
 end

src/nonlinfit.jl

@@ -8,6 +8,8 @@ SciMLBase.isinplace(::NonlinearFunctionWrapper{iip}) where {iip} = iip
 
 _unwrap_nonlinear_function(f::NonlinearFunctionWrapper) = f
 _unwrap_nonlinear_function(f::NonlinearSolveBase.AutoSpecializeCallable) = _unwrap_nonlinear_function(f.orig)
+_unwrap_nonlinear_function(f::NonlinearSolveBase.BoundedWrapper) = _unwrap_nonlinear_function(f.f.f)
+_unwrap_nonlinear_function(f) = f
 
 # If `target` is nothing then we can completely ignore sigma
 __wrap_nonlinear_function(f::NonlinearFunction, ::Nothing, _) = f
@@ -45,13 +47,15 @@ end
 @concrete struct GenericNonlinearCurveFitCache <: AbstractCurveFitCache
     prob <: CurveFitProblem
     cache
+    u0
     alg
     kwargs
 end
 
 function SciMLBase.reinit!(cache::GenericNonlinearCurveFitCache; u0 = nothing, x = nothing, y = nothing, sigma = nothing, kwargs...)
     if !isnothing(u0)
         kwargs = (; kwargs..., u0)
+        copyto!(cache.u0, u0)
     end
 
     # x becomes `p` (parameter) in the NonlinearLeastSquaresProblem
@@ -92,31 +96,59 @@ function CommonSolve.init(
             alg.alg;
             kwargs...
         ),
+        copy(prob.u0),
         alg,
         kwargs
     )
 end
 
 function CommonSolve.solve!(cache::GenericNonlinearCurveFitCache)
+    inner = _get_cache(cache)
+    x = inner.p
     sol = solve!(cache.cache)
-    return CurveFitSolution(cache.alg, sol.u, sol.resid, cache.prob, sol.retcode, sol)
+
+    y = cache.prob.y
+    sigma = cache.prob.sigma
+
+    wrapped_f = _unwrap_nonlinear_function(inner.prob.f.f)
+    if wrapped_f isa NonlinearFunctionWrapper
+        y = wrapped_f.target
+        sigma = wrapped_f.sigma
+    end
+
+    # Reconstruct the problem with the current settings. We can't copy
+    # cache.prob because the cache may have been reinit()'d in which case
+    # cache.prob will be out of date and will give wrong results for the stats
+    # functions that use it.
+    prob = CurveFitProblem(
+        x,
+        y,
+        sigma,
+        cache.prob.nlfunc,
+        cache.u0,
+        cache.prob.lb,
+        cache.prob.ub
+    )
+    return CurveFitSolution(cache.alg, sol.u, sol.resid, prob, sol.retcode, sol)
 end
 
 function (sol::CurveFitSolution{<:__FallbackNonlinearFitAlgorithm})(x)
     return sol.prob.nlfunc(sol.u, x)
 end
 
-function Base.show(io::IO, ::MIME"text/plain", cache::GenericNonlinearCurveFitCache)
+function _get_cache(cache::GenericNonlinearCurveFitCache)
     inner = cache.cache
-
-    # Get the actual working cache
-    is_polyalg = inner isa NonlinearSolveBase.NonlinearSolvePolyAlgorithmCache
-    current_cache = if is_polyalg
+    return if inner isa NonlinearSolveBase.NonlinearSolvePolyAlgorithmCache
         inner.caches[inner.current]
     else
         inner
     end
+end
 
+function Base.show(io::IO, ::MIME"text/plain", cache::GenericNonlinearCurveFitCache)
+    inner = cache.cache
+    is_polyalg = inner isa NonlinearSolveBase.NonlinearSolvePolyAlgorithmCache
+    current_cache = _get_cache(cache)
 
     context = (:compact => true, :limit => true)
 

src/rationalfit.jl

@@ -150,12 +150,13 @@ function CommonSolve.solve!(cache::NonlinearRationalFitCache)
     )
 end
 
-function (sol::CurveFitSolution{<:RationalPolynomialFitAlgorithm})(x::Number)
-    return RationalPolynomial(
+function (sol::CurveFitSolution{<:RationalPolynomialFitAlgorithm})(x)
+    rpoly = RationalPolynomial(
         view(sol.u, 1:(sol.alg.num_degree + 1)),
         vcat(
             one(eltype(sol.u)),
             view(sol.u, (sol.alg.num_degree + 2):(length(sol.u)))
         )
-    )(x)
+    )
+    return rpoly.(x)
 end

src/stats.jl

@@ -56,13 +56,27 @@ function StatsAPI.nobs(sol::CurveFitSolution)
     return length(sol.prob.y)
 end
 
+# Return the positions in `sol.u` that are held fixed. These coefficients
+# are still reported by `coef` for a consistent parameter layout, but they don't
+# have a degree of freedom and contribute a zero row/column to the covariance.
+fixed_param_indices(::CurveFitSolution) = ()
+function fixed_param_indices(sol::CurveFitSolution{<:ExpSumFitAlgorithm})
+    # `k` is always stored as the first coefficient, but it's only fitted when
+    # `withconst` is true; otherwise it's held at 0.
+    return sol.alg.withconst ? () : (1,)
+end
+
 """
     dof(sol::CurveFitSolution)
 
 Return the number of degrees of freedom of the model.
+
+Note that this counts the fitted coefficients. Coefficients that are held fixed
+(e.g. the constant of [`ExpSumFitAlgorithm`](@ref) with `withconst = false`) are
+excluded.
 """
 function StatsAPI.dof(sol::CurveFitSolution)
-    return length(sol.u)
+    return length(sol.u) - length(fixed_param_indices(sol))
 end
 
 """
@@ -193,29 +207,13 @@ function jacobian(sol::CurveFitSolution{<:ExpSumFitAlgorithm})
     # We know the sizes from sol.alg (n, m is irrelevant here).
 
     n = sol.alg.n
-    withconst = sol.alg.withconst
 
     function model_expsum(u_curr, x_val)
-        # Extract parameters from flat vector u_curr
-        # Layout: k (if withconst), p (n), λ (n)
-        # Check src/expsumfit.jl backing: (; k, p, λ)
-        # NamedArrayPartition stores them sequentially.
-
-        idx = 1
-        if withconst
-            k = u_curr[idx]
-            idx += 1
-        else
-            k = zero(eltype(u_curr))
-            # k doesn't advance idx
-        end
-
-        # p is next n
-        p = view(u_curr, idx:(idx + n - 1))
-        idx += n
-
-        # λ is next n
-        λ = view(u_curr, idx:(idx + n - 1))
+        # `sol.u` always stores the constant first, regardless of `withconst`
+        # (when `withconst` is false it's just held at 0). Layout: k, p (n), λ (n).
+        k = u_curr[1]
+        p = view(u_curr, 2:(n + 1))
+        λ = view(u_curr, (n + 2):(2n + 1))
 
         # Computation: k + sum(p .* exp.(λ .* x))
         # Use sum generator to avoid allocation
@@ -344,25 +342,32 @@ when `sigma` carries absolute physical uncertainties (analogous to scipy's
     original `y`-space as fit diagnostics, so for these fits `mse(sol)` is *not*
     the variance scaling behind `vcov`.
 """
+# (JᵀJ)⁻¹ via QR, which is more numerically stable than inv(J'J).
+function _vcov_from_jacobian(J)
+    R = LinearAlgebra.qr(J).R
+    Rinv = inv(R)
+    return Rinv * Rinv'
+end
+
 function StatsAPI.vcov(sol::CurveFitSolution; absolute_sigma::Bool = false)
     J = jacobian(sol)
 
     if !isnothing(sol.prob.sigma)
         J ./= sol.prob.sigma
     end
 
-    # Compute the covariance matrix from the QR decomposition
-    # This is numerically more stable than inv(J'J)
-    Q, R = LinearAlgebra.qr(J)
-
-    # Check for rank deficiency or other issues?
-    # LinearAlgebra.qr usually handles full rank.
-    # R is upper triangular. Rinv = inv(R)
-
-    # Ideally checking rank(R) would be good, but assuming J is full rank for now.
-
-    Rinv = inv(R)
-    covar = Rinv * Rinv'
+    fixed = fixed_param_indices(sol)
+    if isempty(fixed)
+        covar = _vcov_from_jacobian(J)
+    else
+        # Fixed coefficients are known exactly: drop them from the inversion
+        # (their Jacobian columns would make JᵀJ singular) and scatter the
+        # free-parameter covariance back, leaving zero rows/columns for them.
+        free = setdiff(axes(J, 2), fixed)
+        covar_free = _vcov_from_jacobian(J[:, free])
+        covar = zeros(eltype(covar_free), size(J, 2), size(J, 2))
+        covar[free, free] .= covar_free
+    end
 
     if !absolute_sigma
         covar .*= mse(sol)

test/expsumfit.jl

@@ -24,6 +24,8 @@ using Test
     @test sol.u.λ ≈ [-3, -2, 0.5] rtol = 5.0e-7
     @test sol.u.p ≈ [4, 2, 5] rtol = 9.0e-6
     @test sol.u.k[] ≈ -1 rtol = 2.0e-6
+    @test sol(x) ≈ y rtol = 2.0e-6
+    @test sol(x[1]) ≈ y[1] rtol = 2.0e-6
 
     # decay curve
     fs, ts, ω₀, τ = 20.0e3, 0.2, 6283.2, 0.0322

test/king_law.jl

@@ -20,6 +20,8 @@ using Test
         @test sol(val) ≈ fn(val)
     end
 
+    @test sol(E) ≈ U
+
     # Sigma not supported
     prob_sigma = CurveFitProblem(E, U; sigma = ones(length(U)))
     @test_throws AssertionError solve(prob_sigma, KingCurveFitAlgorithm())

test/king_modified.jl

@@ -28,6 +28,8 @@ using NonlinearSolveFirstOrder
         @testset for val in range(minimum(E), stop = maximum(E), length = 10)
             @test sol(val) ≈ fn(val) atol = 1.0e-8
         end
+
+        @test sol(E) ≈ U atol = 1.0e-8
     end
 
     # Test bounds: constrain n to [0.0, 0.3], which excludes the true value of 0.42

test/nonlinfit_reinit.jl

@@ -10,18 +10,37 @@ using NonlinearSolveBase: NonlinearSolveBase
 
     fn(a, x) = @. a[1] + a[2] * x^a[3]
     y = fn(a0, x)
+    sigma = ones(length(y))
+    lb = [0.0, 0.0, 0.0]
+    ub = [10.0, 10.0, 1.0]
 
-    prob = NonlinearCurveFitProblem(fn, [0.5, 0.5, 0.5], x, y)
+    prob = NonlinearCurveFitProblem(fn, [0.5, 0.5, 0.5], x, y, sigma; lb, ub)
     cache = CurveFit.init(prob)
     @test solve!(cache).u ≈ a0 atol = 1.0e-7
 
+    # Without bounds, reinit!() should preserve the user-provided u0 in the
+    # returned solution metadata.
+    cache_unbounded = CurveFit.init(NonlinearCurveFitProblem(fn, [0.5, 0.5, 0.5], x, y))
+    u0 = [1.0, 1.0, 1.0]
+    CurveFit.reinit!(cache_unbounded; u0, x, y)
+    @test solve!(cache_unbounded).prob.u0 == u0
+
     # reinit!() the cache with different parameters and recheck the solve
     a0 = [4.0, 5.0, 0.2]
     x = 11.0:20.0
     y = fn(a0, x)
+    sigma = collect(range(0.5, 1.5; length = length(y)))
 
-    CurveFit.reinit!(cache; u0 = [1.0, 1.0, 1.0], x, y)
-    @test solve!(cache).u ≈ a0 atol = 1.0e-7
+    CurveFit.reinit!(cache; u0, x, y, sigma)
+    sol = solve!(cache)
+    @test sol.u ≈ a0 atol = 1.0e-7
+    @test sol.prob.x == x
+    @test sol.prob.y == y
+    @test sol.prob.sigma == sigma
+    @test sol.prob.lb == lb
+    @test sol.prob.ub == ub
+    @test CurveFit.fitted(sol) ≈ y atol = 1.0e-7
+    @test CurveFit.residuals(sol; weighted = false) ≈ CurveFit.residuals(sol) .* sigma
 
     # Repeat with an in-place model: NonlinearSolve wraps in-place Float64
     # problems in `AutoSpecializeCallable`, which `reinit!` must unwrap.

test/qa/qa.jl

@@ -9,7 +9,8 @@ run_qa(
         ),
         all_qualified_accesses_are_public = (;
             ignore = (
-                :AutoSpecializeCallable, :NonlinearSolvePolyAlgorithmCache,  # NonlinearSolveBase
+                :AutoSpecializeCallable, :BoundedWrapper,                    # NonlinearSolveBase
+                :NonlinearSolvePolyAlgorithmCache,                           # NonlinearSolveBase
                 :Utils, :get_fu, :clean_sprint_struct,                       # NonlinearSolveBase(.Utils)
                 :rtoldefault,                                                # Base
             ),

test/rationalfit_linear.jl

@@ -15,4 +15,6 @@ using LinearSolve, LinearAlgebra
     @testset for val in (0.0, 1.5, 4.5, 10.0)
         @test sol(val) ≈ r(val) atol = 1.0e-8
     end
+
+    @test sol(x) ≈ y atol = 1.0e-8
 end

test/stats.jl

@@ -180,6 +180,19 @@ using Distributions: TDist, quantile
         @test size(vcov(sol)) == (3, 3) # k, p, lam
         @test all(stderror(sol) .> 0)
 
+        # Test with `withconst=false`
+        sol_nc = solve(prob, ExpSumFitAlgorithm(; n = 1, withconst = false))
+        @test length(coef(sol_nc)) == 3      # k, p, lam still reported
+        @test coef(sol_nc)[1] == 0           # k held at 0
+        @test dof(sol_nc) == 2               # only p, lam are fitted
+        @test dof_residual(sol_nc) == nobs(sol_nc) - 2
+        @test size(vcov(sol_nc)) == (3, 3)
+        @test vcov(sol_nc)[1, :] == zeros(3) # zero row/column for fixed k
+        @test vcov(sol_nc)[:, 1] == zeros(3)
+        @test stderror(sol_nc)[1] == 0
+        @test all(stderror(sol_nc)[2:end] .> 0)
+        @test confint(sol_nc)[1] == (0.0, 0.0)
+
         # Modified King Fit (E^2 = A + B * U^n)
         # x corresponds to E (Voltage) in Jacobian logic, but input data order is (U, E^2)?
         # User creates CurveFitProblem(x, y).