adeed nonlinear-solve wrappers code split out of OptimizationSciPy into NonlinearSolve

Signed-off-by: AdityaPandeyCN <[email protected]>

Author: AdityaPandeyCN

Project.toml

@@ -44,6 +44,7 @@ PETSc = "ace2c81b-2b5f-4b1e-a30d-d662738edfe0"
 SIAMFANLEquations = "084e46ad-d928-497d-ad5e-07fa361a48c4"
 SpeedMapping = "f1835b91-879b-4a3f-a438-e4baacf14412"
 Sundials = "c3572dad-4567-51f8-b174-8c6c989267f4"
+PythonCall = "9d6bc502-a68f-5bec-b839-aff6a2a0a179"
 
 [sources.BracketingNonlinearSolve]
 path = "lib/BracketingNonlinearSolve"
@@ -74,6 +75,7 @@ NonlinearSolvePETScExt = ["PETSc", "MPI"]
 NonlinearSolveSIAMFANLEquationsExt = "SIAMFANLEquations"
 NonlinearSolveSpeedMappingExt = "SpeedMapping"
 NonlinearSolveSundialsExt = "Sundials"
+NonlinearSolveSciPyExt = "PythonCall"
 
 [compat]
 ADTypes = "1.9"
@@ -135,6 +137,7 @@ SymbolicIndexingInterface = "0.3.36"
 Test = "1.10"
 Zygote = "0.6.69, 0.7"
 julia = "1.10"
+PythonCall = "0.9"
 
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"

ext/NonlinearSolveSciPyExt.jl

@@ -0,0 +1,140 @@
+module NonlinearSolveSciPyExt
+
+# This file is loaded as an extension when PythonCall is available
+using PythonCall
+const scipy_optimize = try
+    pyimport("scipy.optimize")
+catch err
+    error("Python package `scipy` could not be imported. Install it in the Python environment used by PythonCall.")
+end
+
+using SciMLBase
+using NonlinearSolve
+
+# Re-export algorithm type so that `using NonlinearSolve` brings it in when the
+# extension is loaded.  (Matches convention in other extensions.)
+import ..NonlinearSolve: SciPyLeastSquares, SciPyRoot, SciPyRootScalar
+using NonlinearSolveBase: construct_extension_function_wrapper
+
+""" Internal: wrap a Julia residual function into a Python callable """
+function _make_py_residual(f, p)
+    return pyfunc(x_py -> begin
+        x = Vector{Float64}(x_py)     # convert NumPy array → Julia Vector
+        r  = f(x, p)
+        return r                      # auto-convert back to NumPy
+    end)
+end
+
+""" Internal: wrap a Julia scalar function into a Python callable """
+function _make_py_scalar(f, p)
+    return pyfunc(x_py -> begin
+        x = Float64(x_py)
+        return f(x, p)
+    end)
+end
+
+function SciMLBase.__solve(prob::SciMLBase.NonlinearLeastSquaresProblem, alg::SciPyLeastSquares;
+                           abstol = nothing, maxiters = 10_000, alias_u0::Bool = false,
+                           kwargs...)
+    # Construct Python residual
+    py_f = _make_py_residual(prob.f, prob.p)
+
+    # Bounds handling (lb/ub may be missing)
+    has_lb = hasproperty(prob, :lb)
+    has_ub = hasproperty(prob, :ub)
+    if has_lb || has_ub
+        lb = has_lb ? getproperty(prob, :lb) : fill(-Inf, length(prob.u0))
+        ub = has_ub ? getproperty(prob, :ub) : fill( Inf, length(prob.u0))
+        bounds = (lb, ub)
+    else
+        bounds = nothing
+    end
+
+    # Call SciPy
+    res = scipy_optimize.least_squares(py_f, collect(prob.u0);
+                                       method = alg.method,
+                                       loss   = alg.loss,
+                                       max_nfev = maxiters,
+                                       bounds = bounds === nothing ? py_none : bounds,
+                                       kwargs...)
+
+    # Extract solution
+    u_vec = Vector{Float64}(res.x)
+    resid = Vector{Float64}(res.fun)
+
+    u = prob.u0 isa Number ? u_vec[1] : reshape(u_vec, size(prob.u0))
+
+    ret = res.success ? SciMLBase.ReturnCode.Success : SciMLBase.ReturnCode.Failure
+    njev = try
+        Int(res.njev)
+    catch
+        0
+    end
+    stats = SciMLBase.NLStats(res.nfev, njev, 0, 0, res.nfev)
+
+    return SciMLBase.build_solution(prob, alg, u, resid; retcode = ret,
+                                    original = res, stats = stats)
+end
+
+function SciMLBase.__solve(prob::SciMLBase.NonlinearProblem, alg::SciPyRoot;
+                           abstol = nothing, maxiters = 10_000, alias_u0::Bool = false,
+                           kwargs...)
+    # Get in-place residual wrapper from NonlinearSolveBase.
+    f!, u0, resid = construct_extension_function_wrapper(prob; alias_u0)
+
+    py_f = pyfunc(x_py -> begin
+        x = Vector{Float64}(x_py)
+        f!(resid, x)
+        return resid
+    end)
+
+    tol = abstol === nothing ? nothing : abstol
+
+    res = scipy_optimize.root(py_f, collect(u0);
+                              method = alg.method,
+                              tol = tol,
+                              options = Dict("maxiter" => maxiters),
+                              kwargs...)
+
+    u_vec = Vector{Float64}(res.x)
+    f!(resid, u_vec)  # update residual
+
+    u_out = prob.u0 isa Number ? u_vec[1] : reshape(u_vec, size(prob.u0))
+
+    ret = res.success ? SciMLBase.ReturnCode.Success : SciMLBase.ReturnCode.Failure
+    nfev = try Int(res.nfev) catch; 0 end
+    niter = try Int(res.nit) catch; 0 end
+    stats = SciMLBase.NLStats(nfev, 0, 0, 0, niter)
+
+    return SciMLBase.build_solution(prob, alg, u_out, resid; retcode = ret,
+                                    original = res, stats = stats)
+end
+
+function SciMLBase.__solve(prob::SciMLBase.IntervalNonlinearProblem, alg::SciPyRootScalar;
+                           abstol = nothing, maxiters = 10_000, kwargs...)
+    f = prob.f
+    p = prob.p
+    py_f = _make_py_scalar(f, p)
+
+    a, b = prob.tspan
+
+    res = scipy_optimize.root_scalar(py_f;
+                                     method = alg.method,
+                                     bracket = (a, b),
+                                     maxiter = maxiters,
+                                     xtol = abstol,
+                                     kwargs...)
+
+    u_root = res.root
+    resid = f(u_root, p)
+
+    ret = res.converged ? SciMLBase.ReturnCode.Success : SciMLBase.ReturnCode.Failure
+    nfev = try Int(res.function_calls) catch; 0 end
+    niter = try Int(res.iterations) catch; 0 end
+    stats = SciMLBase.NLStats(nfev, 0, 0, 0, niter)
+
+    return SciMLBase.build_solution(prob, alg, u_root, resid; retcode = ret,
+                                    original = res, stats = stats)
+end
+
+end # module 

src/NonlinearSolve.jl

@@ -117,7 +117,9 @@ export NonlinearSolvePolyAlgorithm, FastShortcutNonlinearPolyalg, FastShortcutNL
 
 # Extension Algorithms
 export LeastSquaresOptimJL, FastLevenbergMarquardtJL, NLsolveJL, NLSolversJL,
-       FixedPointAccelerationJL, SpeedMappingJL, SIAMFANLEquationsJL
+       FixedPointAccelerationJL, SpeedMappingJL, SIAMFANLEquationsJL, SciPyLeastSquares,
+       SciPyLeastSquaresTRF, SciPyLeastSquaresDogbox, SciPyLeastSquaresLM,
+       SciPyRoot, SciPyRootScalar
 export PETScSNES, CMINPACK
 
 end

src/extension_algs.jl

@@ -488,3 +488,80 @@ function PETScSNES(; petsclib = missing, autodiff = nothing, mpi_comm = missing,
     end
     return PETScSNES(petsclib, mpi_comm, autodiff, kwargs)
 end
+
+"""
+    SciPyLeastSquares(; method = "trf", loss = "linear")
+
+Wrapper over `scipy.optimize.least_squares` (via PythonCall) for solving
+`NonlinearLeastSquaresProblem`s.  The keyword arguments correspond to the
+`method` ("trf", "dogbox", "lm") and the robust loss function ("linear",
+"soft_l1", "huber", "cauchy", "arctan").
+
+!!! note
+
+    This algorithm is only available if `PythonCall.jl` is installed and the
+    Python package `scipy` is import-able.  Attempting to construct or use the
+    algorithm otherwise will throw an informative error.
+"""
+@concrete struct SciPyLeastSquares <: AbstractNonlinearSolveAlgorithm
+    method::String
+    loss::String
+    name::Symbol
+end
+
+function SciPyLeastSquares(; method::String = "trf", loss::String = "linear")
+    if Base.get_extension(@__MODULE__, :NonlinearSolveSciPyExt) === nothing
+        error("`SciPyLeastSquares` requires `PythonCall.jl` (and SciPy) to be loaded")
+    end
+    valid_methods = ("trf", "dogbox", "lm")
+    valid_losses  = ("linear", "soft_l1", "huber", "cauchy", "arctan")
+    method in valid_methods ||
+        throw(ArgumentError("Invalid method: $method. Valid methods are: $(join(valid_methods, ", "))"))
+    loss in valid_losses ||
+        throw(ArgumentError("Invalid loss: $loss. Valid loss functions are: $(join(valid_losses, ", "))"))
+    return SciPyLeastSquares(method, loss, :SciPyLeastSquares)
+end
+
+SciPyLeastSquaresTRF()    = SciPyLeastSquares(method = "trf")
+SciPyLeastSquaresDogbox() = SciPyLeastSquares(method = "dogbox")
+SciPyLeastSquaresLM()     = SciPyLeastSquares(method = "lm")
+
+"""
+    SciPyRoot(; method = "hybr")
+
+Wrapper over `scipy.optimize.root` for solving `NonlinearProblem`s.  Available
+methods include "hybr" (default), "lm", "broyden1", "broyden2", "anderson",
+"diagbroyden", "linearmixing", "excitingmixing", "krylov", "df-sane" – any
+method accepted by SciPy.
+"""
+@concrete struct SciPyRoot <: AbstractNonlinearSolveAlgorithm
+    method::String
+    name::Symbol
+end
+
+function SciPyRoot(; method::String = "hybr")
+    if Base.get_extension(@__MODULE__, :NonlinearSolveSciPyExt) === nothing
+        error("`SciPyRoot` requires `PythonCall.jl` (and SciPy) to be loaded")
+    end
+    return SciPyRoot(method, :SciPyRoot)
+end
+
+"""
+    SciPyRootScalar(; method = "brentq")
+
+Wrapper over `scipy.optimize.root_scalar` for scalar `IntervalNonlinearProblem`s
+(bracketing problems).  The default method uses Brent's algorithm ( "brentq").
+Other valid options: "bisect", "brentq", "brenth", "ridder", "toms748",
+"secant", "newton" (secant/Newton do not require brackets).
+"""
+@concrete struct SciPyRootScalar <: AbstractNonlinearSolveAlgorithm
+    method::String
+    name::Symbol
+end
+
+function SciPyRootScalar(; method::String = "brentq")
+    if Base.get_extension(@__MODULE__, :NonlinearSolveSciPyExt) === nothing
+        error("`SciPyRootScalar` requires `PythonCall.jl` (and SciPy) to be loaded")
+    end
+    return SciPyRootScalar(method, :SciPyRootScalar)
+end

test/wrappers/least_squares_tests.jl

@@ -113,3 +113,30 @@ end
     sol = solve(prob_sa, FastLevenbergMarquardtJL())
     @test maximum(abs, sol.resid) < 1e-6
 end
+
+@testitem "SciPyLeastSquares" setup=[WrapperNLLSSetup] tags=[:wrappers] begin
+    success = false
+    try
+        import PythonCall
+        spopt = PythonCall.pyimport("scipy.optimize")
+        success = true
+    catch
+    end
+    if success
+        xdata = collect(0:0.1:1)
+        ydata = 2.0 .* xdata .+ 1.0 .+ 0.1 .* randn(length(xdata))
+        function residuals(params, p=nothing)
+            a, b = params
+            return ydata .- (a .* xdata .+ b)
+        end
+        x0_ls = [1.0, 0.0]
+        prob = NonlinearLeastSquaresProblem(residuals, x0_ls)
+        sol = solve(prob, SciPyLeastSquaresTRF())
+        @test SciMLBase.successful_retcode(sol)
+        prob_bounded = NonlinearLeastSquaresProblem(residuals, x0_ls; lb = [0.0,-2.0], ub = [5.0,3.0])
+        sol2 = solve(prob_bounded, SciPyLeastSquares(method="trf"))
+        @test SciMLBase.successful_retcode(sol2)
+    else
+        @test true # skip: SciPy not present
+    end
+end

test/wrappers/rootfind_tests.jl

@@ -185,3 +185,33 @@ end
 
     @test_throws AssertionError solve(probN, PETScSNES(); abstol = 1e-8)
 end
+
+@testitem "SciPyRoot + SciPyRootScalar" tags=[:wrappers] begin
+    success = false
+    try
+        import PythonCall
+        PythonCall.pyimport("scipy.optimize")
+        success = true
+    catch
+    end
+    if success
+        # Vector root example
+        function fvec(u, p)
+            return [2 - 2u[1]; u[1] - 4u[2]]
+        end
+        u0 = zeros(2)
+        prob_vec = NonlinearProblem(fvec, u0)
+        sol_vec = solve(prob_vec, SciPyRoot())
+        @test SciMLBase.successful_retcode(sol_vec)
+        @test maximum(abs, sol_vec.resid) < 1e-6
+
+        # Scalar bracketing root example
+        fscalar(x, p) = x^2 - 2
+        prob_interval = IntervalNonlinearProblem(fscalar, (1.0, 2.0))
+        sol_scalar = solve(prob_interval, SciPyRootScalar())
+        @test SciMLBase.successful_retcode(sol_scalar)
+        @test abs(sol_scalar.u - sqrt(2)) < 1e-6
+    else
+        @test true
+    end
+end