diff --git a/.JuliaFormatter.toml b/.JuliaFormatter.toml index 66c13bae3..1768a1a7f 100644 --- a/.JuliaFormatter.toml +++ b/.JuliaFormatter.toml @@ -2,4 +2,3 @@ style = "sciml" format_markdown = true annotate_untyped_fields_with_any = false format_docstrings = true -join_lines_based_on_source = false diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index b9a145400..8ddf4f4ff 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -1,15 +1,25 @@ steps: - - label: "Julia 1" + - label: "Julia 1 (NonlinearSolve)" plugins: - JuliaCI/julia#v1: version: "1" - - JuliaCI/julia-test#v1: - coverage: true - JuliaCI/julia-coverage#v1: codecov: true dirs: - src - ext + command: | + julia --color=yes --code-coverage=user --depwarn=yes --project=. -e ' + import Pkg; + Pkg.Registry.update(); + # Install packages present in subdirectories + dev_pks = Pkg.PackageSpec[]; + for path in ("lib/SciMLJacobianOperators",) + push!(dev_pks, Pkg.PackageSpec(; path)); + end + Pkg.develop(dev_pks); + Pkg.instantiate(); + Pkg.test(; coverage=true)' agents: queue: "juliagpu" cuda: "*" @@ -20,6 +30,4 @@ steps: env: GROUP: CUDA JULIA_PKG_SERVER: "" # it often struggles with our large artifacts - RETESTITEMS_NWORKERS: 4 - RETESTITEMS_NWORKER_THREADS: 2 SECRET_CODECOV_TOKEN: "HC7K/ymhi62KUQ5OLU4DOl+11gaQt4JhXX/2nfTGlTsBB8mEMxQ8R+sHIp/2HjEup5eSXAN2IWQDQ7RDBuQvVp0T1UVtr2e4YNZFztKnsJXrFO15hXxYShJodI//X/8DzhlQd/lyTDOAOJu3eznsc3sC2CUgJzXZxLUtQN9YaZ1i3a+NoN1mO5UpkkHVhXigwF5gjy+0tei8fCdcP+SIhG0EanS5yd9q/SurtCpMHsHyUG97+ZVPglSKgdaqr31+PdmiPJ+ynp4+Hnc/esosxUSHSIL+ryRTO+28RNwPTiNf99J51RJLQmz1knWTR1ky6tiYIZ5218O6wvNil0SqNw==;U2FsdGVkX18nBY3t4LZYlEIz3EVKjpqCd994JNeJGt006up+sAjXEssI0tgCVXnfXsenVsP3NCCEoOS1GXc44g==" diff --git a/.github/workflows/CI_NonlinearSolve.yml b/.github/workflows/CI_NonlinearSolve.yml new file mode 100644 index 000000000..a4c14820e --- /dev/null +++ b/.github/workflows/CI_NonlinearSolve.yml @@ -0,0 +1,80 @@ +name: CI (NonlinearSolve) + +on: + pull_request: + branches: + - master + paths: + - "src/**" + - "ext/**" + - "test/**" + - "Project.toml" + - ".github/workflows/CI_NonlinearSolve.yml" + - "lib/SciMLNonlinearOperators/**" + push: + branches: + - master + +concurrency: + # Skip intermediate builds: always. + # Cancel intermediate builds: only if it is a pull request build. + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }} + +jobs: + test: + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + group: + - Core + - Downstream + - Misc + - Wrappers + version: + - "min" + - "1" + os: + - ubuntu-latest + - macos-latest + - windows-latest + steps: + - uses: actions/checkout@v4 + - uses: julia-actions/setup-julia@v2 + with: + version: ${{ matrix.version }} + - uses: actions/cache@v4 + env: + cache-name: cache-artifacts + with: + path: ~/.julia/artifacts + key: ${{ runner.os }}-test-${{ env.cache-name }}-${{ hashFiles('**/Project.toml') }} + restore-keys: | + ${{ runner.os }}-test-${{ env.cache-name }}- + ${{ runner.os }}-test- + ${{ runner.os }}- + - name: "Install Dependencies and Run Tests" + run: | + import Pkg + Pkg.Registry.update() + # Install packages present in subdirectories + dev_pks = Pkg.PackageSpec[] + for path in ("lib/SciMLJacobianOperators",) + push!(dev_pks, Pkg.PackageSpec(; path)) + end + Pkg.develop(dev_pks) + Pkg.instantiate() + Pkg.test(; coverage=true) + shell: julia --color=yes --code-coverage=user --depwarn=yes --project=. {0} + env: + GROUP: ${{ matrix.group }} + - uses: julia-actions/julia-processcoverage@v1 + with: + directories: src,ext,lib/SciMLJacobianOperators/src + - uses: codecov/codecov-action@v4 + with: + file: lcov.info + token: ${{ secrets.CODECOV_TOKEN }} + verbose: true + fail_ci_if_error: true diff --git a/.github/workflows/CI.yml b/.github/workflows/CI_SciMLJacobianOperators.yml similarity index 71% rename from .github/workflows/CI.yml rename to .github/workflows/CI_SciMLJacobianOperators.yml index 21263f5f3..92daf23e9 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI_SciMLJacobianOperators.yml @@ -1,28 +1,30 @@ -name: CI +name: CI (SciMLJacobianOperators) + on: pull_request: branches: - master + paths: + - "lib/SciMLJacobianOperators/**" + - ".github/workflows/CI_SciMLJacobianOperators.yml" push: branches: - master + concurrency: # Skip intermediate builds: always. # Cancel intermediate builds: only if it is a pull request build. group: ${{ github.workflow }}-${{ github.ref }} cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }} + jobs: test: runs-on: ${{ matrix.os }} strategy: fail-fast: false matrix: - group: - - Core - - Downstream - - Misc - - Wrappers version: + - "min" - "1" os: - ubuntu-latest @@ -43,16 +45,16 @@ jobs: ${{ runner.os }}-test-${{ env.cache-name }}- ${{ runner.os }}-test- ${{ runner.os }}- - - uses: julia-actions/julia-buildpkg@v1 - - uses: julia-actions/julia-runtest@v1 - env: - GROUP: ${{ matrix.group }} - JULIA_NUM_THREADS: 11 - RETESTITEMS_NWORKERS: 4 - RETESTITEMS_NWORKER_THREADS: 2 + - name: "Install Dependencies and Run Tests" + run: | + import Pkg + Pkg.Registry.update() + Pkg.instantiate() + Pkg.test(; coverage=true) + shell: julia --color=yes --code-coverage=user --depwarn=yes --project=lib/SciMLJacobianOperators {0} - uses: julia-actions/julia-processcoverage@v1 with: - directories: src,ext + directories: lib/SciMLJacobianOperators/src - uses: codecov/codecov-action@v4 with: file: lcov.info diff --git a/.github/workflows/Documentation.yml b/.github/workflows/Documentation.yml index 0f767ec10..008cd511e 100644 --- a/.github/workflows/Documentation.yml +++ b/.github/workflows/Documentation.yml @@ -17,7 +17,17 @@ jobs: with: version: '1' - name: Install dependencies - run: julia --project=docs/ -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()' + run: | + import Pkg + Pkg.Registry.update() + # Install packages present in subdirectories + dev_pks = Pkg.PackageSpec[] + for path in ("lib/SciMLJacobianOperators", ".") + push!(dev_pks, Pkg.PackageSpec(; path)) + end + Pkg.develop(dev_pks) + Pkg.instantiate() + shell: julia --color=yes --project=docs/ {0} - name: Build and deploy env: JULIA_DEBUG: "Documenter" diff --git a/.github/workflows/Downgrade.yml b/.github/workflows/Downgrade.yml index 3d3c9a2cd..fa9be3c7a 100644 --- a/.github/workflows/Downgrade.yml +++ b/.github/workflows/Downgrade.yml @@ -26,7 +26,3 @@ jobs: skip: Pkg,TOML - uses: julia-actions/julia-buildpkg@v1 - uses: julia-actions/julia-runtest@v1 - env: - JULIA_NUM_THREADS: 11 - RETESTITEMS_NWORKERS: 4 - RETESTITEMS_NWORKER_THREADS: 2 diff --git a/.github/workflows/Downstream.yml b/.github/workflows/Downstream.yml index 3e5a3e4f5..9e54283aa 100644 --- a/.github/workflows/Downstream.yml +++ b/.github/workflows/Downstream.yml @@ -55,9 +55,6 @@ jobs: @info "Not compatible with this release. No problem." exception=err exit(0) # Exit immediately, as a success end - env: - RETESTITEMS_NWORKERS: 4 - RETESTITEMS_NWORKER_THREADS: 2 - uses: julia-actions/julia-processcoverage@v1 with: directories: src,ext diff --git a/Project.toml b/Project.toml index b706d37db..ab35df1e2 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "NonlinearSolve" uuid = "8913a72c-1f9b-4ce2-8d82-65094dcecaec" authors = ["SciML"] -version = "3.14.0" +version = "3.15.0-DEV" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" @@ -23,6 +23,7 @@ Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" Reexport = "189a3867-3050-52da-a836-e630ba90ab69" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" +SciMLJacobianOperators = "19f34311-ddf3-4b8b-af20-060888a46c0e" SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804" @@ -73,12 +74,14 @@ FastLevenbergMarquardt = "0.1" FiniteDiff = "2.22" FixedPointAcceleration = "0.3" ForwardDiff = "0.10.36" +Hwloc = "3" +InteractiveUtils = "<0.0.1, 1" LazyArrays = "1.8.2, 2" LeastSquaresOptim = "0.8.5" LineSearches = "7.2" LinearAlgebra = "1.10" LinearSolve = "2.30" -MINPACK = "=1.2" +MINPACK = "1.2" MaybeInplace = "0.1.3" ModelingToolkit = "9.15.0" NLSolvers = "0.5" @@ -96,6 +99,7 @@ RecursiveArrayTools = "3.8" Reexport = "1.2" SIAMFANLEquations = "1.0.1" SciMLBase = "2.34.0" +SciMLJacobianOperators = "0.1" SimpleNonlinearSolve = "1.8" SparseArrays = "1.10" SparseDiffTools = "2.19" @@ -120,6 +124,8 @@ Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7" FastLevenbergMarquardt = "7a0df574-e128-4d35-8cbd-3d84502bf7ce" FixedPointAcceleration = "817d07cb-a79a-5c30-9a31-890123675176" +Hwloc = "0e44f5e4-bd66-52a0-8798-143a42290a1d" +InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" LeastSquaresOptim = "0fc2ff8b-aaa3-5acd-a817-1944a5e08891" MINPACK = "4854310b-de5a-5eb6-a2a5-c1dee2bd17f9" ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78" @@ -141,4 +147,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [targets] -test = ["Aqua", "BandedMatrices", "BenchmarkTools", "CUDA", "Enzyme", "ExplicitImports", "FastLevenbergMarquardt", "FixedPointAcceleration", "LeastSquaresOptim", "MINPACK", "ModelingToolkit", "NLSolvers", "NLsolve", "NaNMath", "NonlinearProblemLibrary", "OrdinaryDiffEq", "Pkg", "Random", "ReTestItems", "SIAMFANLEquations", "SpeedMapping", "StableRNGs", "StaticArrays", "Sundials", "Symbolics", "Test", "Zygote"] +test = ["Aqua", "BandedMatrices", "BenchmarkTools", "CUDA", "Enzyme", "ExplicitImports", "FastLevenbergMarquardt", "FixedPointAcceleration", "Hwloc", "InteractiveUtils", "LeastSquaresOptim", "MINPACK", "ModelingToolkit", "NLSolvers", "NLsolve", "NaNMath", "NonlinearProblemLibrary", "OrdinaryDiffEq", "Pkg", "Random", "ReTestItems", "SIAMFANLEquations", "SpeedMapping", "StableRNGs", "StaticArrays", "Sundials", "Symbolics", "Test", "Zygote"] diff --git a/docs/Project.toml b/docs/Project.toml index 780cec74a..63d1803d4 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -14,6 +14,7 @@ OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" +SciMLJacobianOperators = "19f34311-ddf3-4b8b-af20-060888a46c0e" SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7" SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" @@ -36,6 +37,7 @@ OrdinaryDiffEq = "6" Plots = "1" Random = "<0.0.1, 1" SciMLBase = "2.4" +SciMLJacobianOperators = "0.1" SimpleNonlinearSolve = "1" SparseDiffTools = "2.14" StaticArrays = "1" diff --git a/docs/make.jl b/docs/make.jl index 622cebbfe..e11a04149 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -1,6 +1,7 @@ using Documenter, DocumenterCitations using NonlinearSolve, SimpleNonlinearSolve, Sundials, SteadyStateDiffEq, SciMLBase, DiffEqBase +using SciMLJacobianOperators cp(joinpath(@__DIR__, "Manifest.toml"), joinpath(@__DIR__, "src/assets/Manifest.toml"), force = true) @@ -13,8 +14,8 @@ bib = CitationBibliography(joinpath(@__DIR__, "src", "refs.bib")) makedocs(; sitename = "NonlinearSolve.jl", authors = "Chris Rackauckas", - modules = [NonlinearSolve, SimpleNonlinearSolve, - SteadyStateDiffEq, Sundials, DiffEqBase, SciMLBase], + modules = [NonlinearSolve, SimpleNonlinearSolve, SteadyStateDiffEq, + Sundials, DiffEqBase, SciMLBase, SciMLJacobianOperators], clean = true, doctest = false, linkcheck = true, diff --git a/docs/pages.jl b/docs/pages.jl index 0d548cd60..cfb2754c0 100644 --- a/docs/pages.jl +++ b/docs/pages.jl @@ -2,25 +2,61 @@ pages = ["index.md", "Getting Started with Nonlinear Rootfinding in Julia" => "tutorials/getting_started.md", - "Tutorials" => Any["tutorials/code_optimization.md", "tutorials/large_systems.md", - "tutorials/modelingtoolkit.md", "tutorials/small_compile.md", - "tutorials/iterator_interface.md", "tutorials/optimizing_parameterized_ode.md"], - "Basics" => Any["basics/nonlinear_problem.md", "basics/nonlinear_functions.md", - "basics/solve.md", "basics/nonlinear_solution.md", "basics/autodiff.md", - "basics/termination_condition.md", "basics/diagnostics_api.md", - "basics/sparsity_detection.md", "basics/faq.md"], - "Solver Summaries and Recommendations" => Any["solvers/nonlinear_system_solvers.md", - "solvers/bracketing_solvers.md", "solvers/steady_state_solvers.md", - "solvers/nonlinear_least_squares_solvers.md", "solvers/fixed_point_solvers.md"], - "Native Functionalities" => Any["native/solvers.md", "native/simplenonlinearsolve.md", - "native/steadystatediffeq.md", "native/descent.md", - "native/globalization.md", "native/diagnostics.md"], + "Tutorials" => Any[ + "tutorials/code_optimization.md", + "tutorials/large_systems.md", + "tutorials/modelingtoolkit.md", + "tutorials/small_compile.md", + "tutorials/iterator_interface.md", + "tutorials/optimizing_parameterized_ode.md" + ], + "Basics" => Any[ + "basics/nonlinear_problem.md", + "basics/nonlinear_functions.md", + "basics/solve.md", + "basics/nonlinear_solution.md", + "basics/autodiff.md", + "basics/termination_condition.md", + "basics/diagnostics_api.md", + "basics/sparsity_detection.md", + "basics/faq.md" + ], + "Solver Summaries and Recommendations" => Any[ + "solvers/nonlinear_system_solvers.md", + "solvers/bracketing_solvers.md", + "solvers/steady_state_solvers.md", + "solvers/nonlinear_least_squares_solvers.md", + "solvers/fixed_point_solvers.md" + ], + "Native Functionalities" => Any[ + "native/solvers.md", + "native/simplenonlinearsolve.md", + "native/steadystatediffeq.md", + "native/descent.md", + "native/globalization.md", + "native/diagnostics.md" + ], "Wrapped Solver APIs" => Any[ - "api/fastlevenbergmarquardt.md", "api/fixedpointacceleration.md", - "api/leastsquaresoptim.md", "api/minpack.md", "api/nlsolve.md", "api/nlsolvers.md", - "api/siamfanlequations.md", "api/speedmapping.md", "api/sundials.md"], + "api/fastlevenbergmarquardt.md", + "api/fixedpointacceleration.md", + "api/leastsquaresoptim.md", + "api/minpack.md", + "api/nlsolve.md", + "api/nlsolvers.md", + "api/siamfanlequations.md", + "api/speedmapping.md", + "api/sundials.md" + ], + "Sub-Packages" => Any[ + "api/SciMLJacobianOperators.md", + ], "Development Documentation" => [ - "devdocs/internal_interfaces.md", "devdocs/linear_solve.md", - "devdocs/jacobian.md", "devdocs/operators.md", "devdocs/algorithm_helpers.md"], + "devdocs/internal_interfaces.md", + "devdocs/linear_solve.md", + "devdocs/jacobian.md", + "devdocs/operators.md", + "devdocs/algorithm_helpers.md" + ], "Release Notes" => "release_notes.md", - "References" => "references.md"] + "References" => "references.md" +] diff --git a/docs/src/api/SciMLJacobianOperators.md b/docs/src/api/SciMLJacobianOperators.md new file mode 100644 index 000000000..91b7f61df --- /dev/null +++ b/docs/src/api/SciMLJacobianOperators.md @@ -0,0 +1,29 @@ +```@meta +CurrentModule = SciMLJacobianOperators +``` + +# SciMLJacobianOperators.jl + +This is a subpackage on NonlinearSolve providing a general purpose JacVec and VecJac +operator built on top on DifferentiationInterface.jl. + +```julia +import Pkg +Pkg.add("SciMLJacobianOperators") +using SciMLJacobianOperators +``` + +## Jacobian API + +```@docs +JacobianOperator +VecJacOperator +JacVecOperator +``` + +## Stateful Operators + +```@docs +StatefulJacobianOperator +StatefulJacobianNormalFormOperator +``` diff --git a/docs/src/devdocs/operators.md b/docs/src/devdocs/operators.md index b96a63f8c..15d00093a 100644 --- a/docs/src/devdocs/operators.md +++ b/docs/src/devdocs/operators.md @@ -6,21 +6,6 @@ NonlinearSolve.AbstractNonlinearSolveOperator ``` -## Jacobian Operators - -```@docs -NonlinearSolve.JacobianOperator -NonlinearSolve.VecJacOperator -NonlinearSolve.JacVecOperator -``` - -### Stateful Jacobian Operators - -```@docs -NonlinearSolve.StatefulJacobianOperator -NonlinearSolve.StatefulJacobianNormalFormOperator -``` - ## Low-Rank Jacobian Operators ```@docs diff --git a/docs/src/tutorials/large_systems.md b/docs/src/tutorials/large_systems.md index 4e732b9c0..368535287 100644 --- a/docs/src/tutorials/large_systems.md +++ b/docs/src/tutorials/large_systems.md @@ -142,7 +142,7 @@ using BenchmarkTools # for @btime NewtonRaphson(; autodiff = AutoSparse(AutoForwardDiff(; chunksize = 32)), linsolve = KLUFactorization())); @btime solve(prob_brusselator_2d, - NewtonRaphson(; autodiff = AutoSparse(AutoForwardDiff(; chunksize = 32)), + NewtonRaphson(; autodiff = AutoForwardDiff(; chunksize = 32), linsolve = KrylovJL_GMRES())); nothing # hide ``` @@ -226,7 +226,8 @@ left and right preconditioners, matrices which approximate the inverse of `W = I used in the solution of the ODE. An example of this with using [IncompleteLU.jl](https://github.com/haampie/IncompleteLU.jl) is as follows: -```@example ill_conditioned_nlprob +```julia +# FIXME: On 1.10+ this is broken. Skipping this for now. using IncompleteLU function incompletelu(W, du, u, p, t, newW, Plprev, Prprev, solverdata) diff --git a/lib/SciMLJacobianOperators/LICENSE b/lib/SciMLJacobianOperators/LICENSE new file mode 100644 index 000000000..abb594d1e --- /dev/null +++ b/lib/SciMLJacobianOperators/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2024 SciML + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/lib/SciMLJacobianOperators/Project.toml b/lib/SciMLJacobianOperators/Project.toml new file mode 100644 index 000000000..5c1bc20ff --- /dev/null +++ b/lib/SciMLJacobianOperators/Project.toml @@ -0,0 +1,54 @@ +name = "SciMLJacobianOperators" +uuid = "19f34311-ddf3-4b8b-af20-060888a46c0e" +authors = ["Avik Pal and contributors"] +version = "0.1.0" + +[deps] +ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" +ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471" +ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9" +DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" +EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869" +FastClosures = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" +SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961" + +[compat] +ADTypes = "1.8.1" +Aqua = "0.8.7" +ConcreteStructs = "0.2.3" +ConstructionBase = "1.5" +DifferentiationInterface = "0.5" +Enzyme = "0.12, 0.13" +EnzymeCore = "0.7, 0.8" +ExplicitImports = "1.9.0" +FastClosures = "0.3.2" +FiniteDiff = "2.24.0" +ForwardDiff = "0.10.36" +InteractiveUtils = "<0.0.1, 1" +LinearAlgebra = "1.10" +ReverseDiff = "1.15" +SciMLBase = "2.54.0" +SciMLOperators = "0.3" +Test = "1.10" +TestItemRunner = "1" +Tracker = "0.2.35" +Zygote = "0.6.70" +julia = "1.10" + +[extras] +Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" +ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7" +FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" +ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +TestItemRunner = "f8b46487-2199-4994-9208-9a1283c18c0a" +Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" + +[targets] +test = ["Aqua", "Enzyme", "ExplicitImports", "FiniteDiff", "ForwardDiff", "InteractiveUtils", "ReverseDiff", "Test", "TestItemRunner", "Tracker", "Zygote"] diff --git a/lib/SciMLJacobianOperators/src/SciMLJacobianOperators.jl b/lib/SciMLJacobianOperators/src/SciMLJacobianOperators.jl new file mode 100644 index 000000000..a173dc395 --- /dev/null +++ b/lib/SciMLJacobianOperators/src/SciMLJacobianOperators.jl @@ -0,0 +1,399 @@ +module SciMLJacobianOperators + +using ADTypes: ADTypes, AutoSparse, AutoEnzyme +using ConcreteStructs: @concrete +using ConstructionBase: ConstructionBase +using DifferentiationInterface: DifferentiationInterface +using EnzymeCore: EnzymeCore +using FastClosures: @closure +using LinearAlgebra: LinearAlgebra +using SciMLBase: SciMLBase, AbstractNonlinearProblem, AbstractNonlinearFunction +using SciMLOperators: AbstractSciMLOperator + +const DI = DifferentiationInterface +const True = Val(true) +const False = Val(false) + +abstract type AbstractJacobianOperator{T} <: AbstractSciMLOperator{T} end + +abstract type AbstractMode end + +struct VJP <: AbstractMode end +struct JVP <: AbstractMode end + +flip_mode(::VJP) = JVP() +flip_mode(::JVP) = VJP() + +""" + JacobianOperator{iip, T} <: AbstractJacobianOperator{T} <: AbstractSciMLOperator{T} + +A Jacobian Operator Provides both JVP and VJP without materializing either (if possible). + +### Constructor + +```julia +JacobianOperator(prob::AbstractNonlinearProblem, fu, u; jvp_autodiff = nothing, + vjp_autodiff = nothing, skip_vjp::Val = Val(false), skip_jvp::Val = Val(false)) +``` + +By default, the `JacobianOperator` will compute `JVP`. Use `Base.adjoint` or +`Base.transpose` to switch to `VJP`. + +### Computing the VJP + +Computing the VJP is done according to the following rules: + + - If `f` has a `vjp` method, then we use that. + - If `f` has a `jac` method and no `vjp_autodiff` is provided, then we use `jac * v`. + - If `vjp_autodiff` is provided we using DifferentiationInterface.jl to compute the VJP. + +### Computing the JVP + +Computing the JVP is done according to the following rules: + + - If `f` has a `jvp` method, then we use that. + - If `f` has a `jac` method and no `jvp_autodiff` is provided, then we use `v * jac`. + - If `jvp_autodiff` is provided we using DifferentiationInterface.jl to compute the JVP. + +### Special Case (Number) + +For Number inputs, VJP and JVP are not distinct. Hence, if either `vjp` or `jvp` is +provided, then we use that. If neither is provided, then we use `v * jac` if `jac` is +provided. Finally, we use the respective autodiff methods to compute the derivative +using DifferentiationInterface.jl and multiply by `v`. + +### Methods Provided + +!!! warning + + Currently it is expected that `p` during problem construction is same as `p` during + operator evaluation. This restriction will be lifted in the future. + + - `(op::JacobianOperator)(v, u, p)`: Computes `∂f(u, p)/∂u * v` or `∂f(u, p)/∂uᵀ * v`. + - `(op::JacobianOperator)(res, v, u, p)`: Computes `∂f(u, p)/∂u * v` or `∂f(u, p)/∂uᵀ * v` + and stores the result in `res`. + +See also [`VecJacOperator`](@ref) and [`JacVecOperator`](@ref). +""" +@concrete struct JacobianOperator{iip, T <: Real} <: AbstractJacobianOperator{T} + mode <: AbstractMode + + jvp_op + vjp_op + + size + + output_cache + input_cache +end + +SciMLBase.isinplace(::JacobianOperator{iip}) where {iip} = iip + +function ConstructionBase.constructorof(::Type{<:JacobianOperator{iip, T}}) where {iip, T} + return JacobianOperator{iip, T} +end + +Base.size(J::JacobianOperator) = J.size +Base.size(J::JacobianOperator, d::Integer) = J.size[d] + +for op in (:adjoint, :transpose) + @eval function Base.$(op)(operator::JacobianOperator{iip, T}) where {iip, T} + return JacobianOperator{iip, T}( + flip_mode(operator.mode), operator.jvp_op, operator.vjp_op, + reverse(operator.size), operator.input_cache, operator.output_cache) + end +end + +function JacobianOperator(prob::AbstractNonlinearProblem, fu, u; jvp_autodiff = nothing, + vjp_autodiff = nothing, skip_vjp::Val = False, skip_jvp::Val = False) + @assert !(skip_vjp === True && skip_jvp === True) "Cannot skip both vjp and jvp \ + construction." + f = prob.f + iip = SciMLBase.isinplace(prob) + T = promote_type(eltype(u), eltype(fu)) + + vjp_autodiff = set_function_as_const(get_dense_ad(vjp_autodiff)) + vjp_op = prepare_vjp(skip_vjp, prob, f, u, fu; autodiff = vjp_autodiff) + + jvp_autodiff = set_function_as_const(get_dense_ad(jvp_autodiff)) + jvp_op = prepare_jvp(skip_jvp, prob, f, u, fu; autodiff = jvp_autodiff) + + output_cache = fu isa Number ? T(fu) : similar(fu, T) + input_cache = u isa Number ? T(u) : similar(u, T) + + return JacobianOperator{iip, T}( + JVP(), jvp_op, vjp_op, (length(fu), length(u)), output_cache, input_cache) +end + +function (op::JacobianOperator)(v, u, p) + if op.mode isa VJP + if SciMLBase.isinplace(op) + res = zero(op.output_cache) + op.vjp_op(res, v, u, p) + return res + end + return op.vjp_op(v, u, p) + else + if SciMLBase.isinplace(op) + res = zero(op.output_cache) + op.jvp_op(res, v, u, p) + return res + end + return op.jvp_op(v, u, p) + end +end + +function (op::JacobianOperator)(::Number, ::Number, _, __) + error("Inplace Jacobian Operator not possible for scalars.") +end + +function (op::JacobianOperator)(Jv, v, u, p) + if op.mode isa VJP + if SciMLBase.isinplace(op) + op.vjp_op(Jv, v, u, p) + else + copyto!(Jv, op.vjp_op(v, u, p)) + end + else + if SciMLBase.isinplace(op) + op.jvp_op(Jv, v, u, p) + else + copyto!(Jv, op.jvp_op(v, u, p)) + end + end + return Jv +end + +""" + VecJacOperator(args...; autodiff = nothing, kwargs...) + +Constructs a [`JacobianOperator`](@ref) which only provides the VJP using the +`vjp_autodiff = autodiff`. +""" +function VecJacOperator(args...; autodiff = nothing, kwargs...) + return JacobianOperator(args...; kwargs..., skip_jvp = True, vjp_autodiff = autodiff)' +end + +""" + JacVecOperator(args...; autodiff = nothing, kwargs...) + +Constructs a [`JacobianOperator`](@ref) which only provides the JVP using the +`jvp_autodiff = autodiff`. +""" +function JacVecOperator(args...; autodiff = nothing, kwargs...) + return JacobianOperator(args...; kwargs..., skip_vjp = True, jvp_autodiff = autodiff) +end + +""" + StatefulJacobianOperator(jac_op::JacobianOperator, u, p) + +Wrapper over a [`JacobianOperator`](@ref) which stores the input `u` and `p` and defines +`mul!` and `*` for computing VJPs and JVPs. +""" +@concrete struct StatefulJacobianOperator{M <: AbstractMode, T} <: + AbstractJacobianOperator{T} + mode::M + jac_op <: JacobianOperator + u + p + + function StatefulJacobianOperator(jac_op::JacobianOperator, u, p) + return new{ + typeof(jac_op.mode), eltype(jac_op), typeof(jac_op), typeof(u), typeof(p)}( + jac_op.mode, jac_op, u, p) + end +end + +Base.size(J::StatefulJacobianOperator) = size(J.jac_op) +Base.size(J::StatefulJacobianOperator, d::Integer) = size(J.jac_op, d) + +for op in (:adjoint, :transpose) + @eval function Base.$(op)(operator::StatefulJacobianOperator) + return StatefulJacobianOperator($(op)(operator.jac_op), operator.u, operator.p) + end +end + +Base.:*(J::StatefulJacobianOperator, v::AbstractArray) = J.jac_op(v, J.u, J.p) +Base.:*(J::StatefulJacobianOperator, v::Number) = J.jac_op(v, J.u, J.p) + +function LinearAlgebra.mul!( + Jv::AbstractArray, J::StatefulJacobianOperator, v::AbstractArray) + J.jac_op(Jv, v, J.u, J.p) + return Jv +end + +""" + StatefulJacobianNormalFormOperator(vjp_operator, jvp_operator, cache) + +This constructs a Normal Form Jacobian Operator, i.e. it constructs the operator +corresponding to `JᵀJ` where `J` is the Jacobian Operator. This is not meant to be directly +constructed, rather it is constructed with `*` on two [`StatefulJacobianOperator`](@ref)s. +""" +@concrete mutable struct StatefulJacobianNormalFormOperator{T} <: + AbstractJacobianOperator{T} + vjp_operator <: StatefulJacobianOperator{VJP} + jvp_operator <: StatefulJacobianOperator{JVP} + cache +end + +function Base.size(J::StatefulJacobianNormalFormOperator) + return size(J.vjp_operator, 1), size(J.jvp_operator, 2) +end + +function Base.:*(J1::StatefulJacobianOperator{VJP}, J2::StatefulJacobianOperator{JVP}) + cache = J2 * J2.jac_op.input_cache + T = promote_type(eltype(J1), eltype(J2)) + return StatefulJacobianNormalFormOperator{T}(J1, J2, cache) +end + +function LinearAlgebra.mul!(C::StatefulJacobianNormalFormOperator, + A::StatefulJacobianOperator{VJP}, B::StatefulJacobianOperator{JVP}) + C.vjp_operator = A + C.jvp_operator = B + return C +end + +function Base.:*(JᵀJ::StatefulJacobianNormalFormOperator, x::AbstractArray) + return JᵀJ.vjp_operator * (JᵀJ.jvp_operator * x) +end +function Base.:*(JᵀJ::StatefulJacobianNormalFormOperator, x::Number) + return JᵀJ.vjp_operator * (JᵀJ.jvp_operator * x) +end + +function LinearAlgebra.mul!( + JᵀJx::AbstractArray, JᵀJ::StatefulJacobianNormalFormOperator, x::AbstractArray) + LinearAlgebra.mul!(JᵀJ.cache, JᵀJ.jvp_operator, x) + LinearAlgebra.mul!(JᵀJx, JᵀJ.vjp_operator, JᵀJ.cache) + return JᵀJx +end + +# Helper Functions +prepare_vjp(::Val{true}, args...; kwargs...) = nothing + +function prepare_vjp(::Val{false}, prob::AbstractNonlinearProblem, + f::AbstractNonlinearFunction, u::Number, fu::Number; autodiff = nothing) + return prepare_scalar_op(Val(false), prob, f, u, fu; autodiff) +end + +function prepare_vjp(::Val{false}, prob::AbstractNonlinearProblem, + f::AbstractNonlinearFunction, u, fu; autodiff = nothing) + SciMLBase.has_vjp(f) && return f.vjp + + if autodiff === nothing && SciMLBase.has_jac(f) + if SciMLBase.isinplace(f) + jac_cache = similar(u, eltype(fu), length(fu), length(u)) + return @closure (vJ, v, u, p) -> begin + f.jac(jac_cache, u, p) + LinearAlgebra.mul!(vec(vJ), jac_cache', vec(v)) + return + end + return vjp_op, vjp_extras + else + return @closure (v, u, p) -> reshape(f.jac(u, p)' * vec(v), size(u)) + end + end + + @assert autodiff!==nothing "`vjp_autodiff` must be provided if `f` doesn't have \ + analytic `vjp` or `jac`." + # TODO: Once DI supports const params we can use `p` + fₚ = SciMLBase.JacobianWrapper{SciMLBase.isinplace(f)}(f, prob.p) + if SciMLBase.isinplace(f) + @assert DI.check_twoarg(autodiff) "Backend: $(autodiff) doesn't support in-place \ + problems." + fu_cache = copy(fu) + v_fake = copy(fu) + di_extras = DI.prepare_pullback(fₚ, fu_cache, autodiff, u, v_fake) + return @closure (vJ, v, u, p) -> begin + DI.pullback!(fₚ, fu_cache, reshape(vJ, size(u)), autodiff, + u, reshape(v, size(fu_cache)), di_extras) + return + end + else + di_extras = DI.prepare_pullback(fₚ, autodiff, u, fu) + return @closure (v, u, p) -> begin + return DI.pullback(fₚ, autodiff, u, reshape(v, size(fu)), di_extras) + end + end +end + +prepare_jvp(skip::Val{true}, args...; kwargs...) = nothing + +function prepare_jvp(::Val{false}, prob::AbstractNonlinearProblem, + f::AbstractNonlinearFunction, u::Number, fu::Number; autodiff = nothing) + return prepare_scalar_op(Val(false), prob, f, u, fu; autodiff) +end + +function prepare_jvp(::Val{false}, prob::AbstractNonlinearProblem, + f::AbstractNonlinearFunction, u, fu; autodiff = nothing) + SciMLBase.has_jvp(f) && return f.jvp + + if autodiff === nothing && SciMLBase.has_jac(f) + if SciMLBase.isinplace(f) + jac_cache = similar(u, eltype(fu), length(fu), length(u)) + return @closure (Jv, v, u, p) -> begin + f.jac(jac_cache, u, p) + LinearAlgebra.mul!(vec(Jv), jac_cache, vec(v)) + return + end + else + return @closure (v, u, p) -> reshape(f.jac(u, p) * vec(v), size(u)) + end + end + + @assert autodiff!==nothing "`jvp_autodiff` must be provided if `f` doesn't have \ + analytic `vjp` or `jac`." + # TODO: Once DI supports const params we can use `p` + fₚ = SciMLBase.JacobianWrapper{SciMLBase.isinplace(f)}(f, prob.p) + if SciMLBase.isinplace(f) + @assert DI.check_twoarg(autodiff) "Backend: $(autodiff) doesn't support in-place \ + problems." + fu_cache = copy(fu) + di_extras = DI.prepare_pushforward(fₚ, fu_cache, autodiff, u, u) + return @closure (Jv, v, u, p) -> begin + DI.pushforward!( + fₚ, fu_cache, reshape(Jv, size(fu_cache)), + autodiff, u, reshape(v, size(u)), di_extras) + return + end + else + di_extras = DI.prepare_pushforward(fₚ, autodiff, u, u) + return @closure (v, u, p) -> begin + return DI.pushforward(fₚ, autodiff, u, reshape(v, size(u)), di_extras) + end + end +end + +function prepare_scalar_op(::Val{false}, prob::AbstractNonlinearProblem, + f::AbstractNonlinearFunction, u::Number, fu::Number; autodiff = nothing) + SciMLBase.has_vjp(f) && return f.vjp + SciMLBase.has_jvp(f) && return f.jvp + SciMLBase.has_jac(f) && return @closure((v, u, p)->f.jac(u, p) * v) + + @assert autodiff!==nothing "`autodiff` must be provided if `f` doesn't have \ + analytic `vjp` or `jvp` or `jac`." + # TODO: Once DI supports const params we can use `p` + fₚ = Base.Fix2(f, prob.p) + di_extras = DI.prepare_derivative(fₚ, autodiff, u) + return @closure (v, u, p) -> DI.derivative(fₚ, autodiff, u, di_extras) * v +end + +get_dense_ad(::Nothing) = nothing +get_dense_ad(ad) = ad +function get_dense_ad(ad::AutoSparse) + dense_ad = ADTypes.dense_ad(ad) + @warn "Sparse AD backend: $(ad) is being used for VJP/JVP computation. Using the dense \ + backend: $(dense_ad) instead." + return dense_ad +end + +# In our case we know that it is safe to mark the function as const +set_function_as_const(ad) = ad +function set_function_as_const(ad::AutoEnzyme{M, Nothing}) where {M} + return AutoEnzyme(; ad.mode, function_annotation = EnzymeCore.Const) +end + +export JacobianOperator, VecJacOperator, JacVecOperator +export StatefulJacobianOperator +export StatefulJacobianNormalFormOperator + +end diff --git a/lib/SciMLJacobianOperators/test/core_tests.jl b/lib/SciMLJacobianOperators/test/core_tests.jl new file mode 100644 index 000000000..e3b595221 --- /dev/null +++ b/lib/SciMLJacobianOperators/test/core_tests.jl @@ -0,0 +1,274 @@ +@testitem "Scalar Ops" begin + using ADTypes, SciMLBase + using Enzyme, Zygote, ForwardDiff, FiniteDiff, ReverseDiff, Tracker + using SciMLJacobianOperators + + reverse_ADs = [ + AutoEnzyme(), + AutoEnzyme(; mode = Enzyme.Reverse), + AutoZygote(), + AutoReverseDiff(), + AutoTracker(), + AutoFiniteDiff() + ] + + forward_ADs = [ + AutoEnzyme(), + AutoEnzyme(; mode = Enzyme.Forward), + AutoForwardDiff(), + AutoFiniteDiff() + ] + + prob = NonlinearProblem(NonlinearFunction{false}((u, p) -> u^2 - p), 1.0, 2.0) + + analytic_jac(u, p) = 2 * u + analytic_jvp(v, u, p) = 2 * u * v + analytic_vjp(v, u, p) = analytic_jvp(v, u, p) + + @testset "AutoDiff" begin + @testset for jvp_autodiff in forward_ADs, vjp_autodiff in reverse_ADs + jac_op = JacobianOperator(prob, -1.0, 1.0; jvp_autodiff, vjp_autodiff) + + @testset for u in rand(4), v in rand(4) + sop = StatefulJacobianOperator(jac_op, u, prob.p) + @test (sop * v)≈analytic_jvp(v, u, prob.p) atol=1e-5 + @test (sop' * v)≈analytic_vjp(v, u, prob.p) atol=1e-5 + + normal_form_sop = sop' * sop + JᵀJv = normal_form_sop * v + J_analytic = analytic_jac(u, prob.p) + JᵀJv_analytic = J_analytic' * J_analytic * v + @test JᵀJv≈JᵀJv_analytic atol=1e-5 + end + end + end + + prob = NonlinearProblem( + NonlinearFunction{false}((u, p) -> u^2 - p; jvp = analytic_jvp, vjp = analytic_vjp), + 1.0, 2.0) + + @testset "Analytic JVP/VJP" begin + jac_op = JacobianOperator(prob, -1.0, 1.0) + + @testset for u in rand(4), v in rand(4) + sop = StatefulJacobianOperator(jac_op, u, prob.p) + @test (sop * v)≈analytic_jvp(v, u, prob.p) atol=1e-5 + @test (sop' * v)≈analytic_vjp(v, u, prob.p) atol=1e-5 + + normal_form_sop = sop' * sop + JᵀJv = normal_form_sop * v + J_analytic = analytic_jac(u, prob.p) + JᵀJv_analytic = J_analytic' * J_analytic * v + @test JᵀJv≈JᵀJv_analytic atol=1e-5 + end + end + + prob = NonlinearProblem( + NonlinearFunction{false}((u, p) -> u^2 - p; jac = (u, p) -> 2 * u), 1.0, 2.0) + + @testset "Analytic Jacobian" begin + jac_op = JacobianOperator(prob, -1.0, 1.0) + + @testset for u in rand(4), v in rand(4) + sop = StatefulJacobianOperator(jac_op, u, prob.p) + @test (sop * v)≈2 * u * v atol=1e-5 + @test (sop' * v)≈2 * u * v atol=1e-5 + + normal_form_sop = sop' * sop + JᵀJv = normal_form_sop * v + J_analytic = analytic_jac(u, prob.p) + JᵀJv_analytic = J_analytic' * J_analytic * v + @test JᵀJv≈JᵀJv_analytic atol=1e-5 + end + end +end + +@testitem "Inplace Problems" begin + using ADTypes, SciMLBase + using Enzyme, ForwardDiff, FiniteDiff, ReverseDiff + using SciMLJacobianOperators + + reverse_ADs = [ + AutoEnzyme(), + AutoEnzyme(; mode = Enzyme.Reverse), + AutoReverseDiff(), + AutoFiniteDiff() + ] + + forward_ADs = [ + AutoEnzyme(), + AutoEnzyme(; mode = Enzyme.Forward), + AutoForwardDiff(), + AutoFiniteDiff() + ] + + prob = NonlinearProblem( + NonlinearFunction{true}((du, u, p) -> du .= u .^ 2 .- p .+ u[2] * u[1]), [1.0, 3.0], 2.0) + + analytic_jac(u, p) = [2 * u[1]+u[2] u[1]; u[2] 2 * u[2]+u[1]] + analytic_jvp(v, u, p) = analytic_jac(u, p) * v + analytic_vjp(v, u, p) = analytic_jac(u, p)' * v + + analytic_jac!(J, u, p) = (J .= analytic_jac(u, p)) + analytic_jvp!(J, v, u, p) = (J .= analytic_jvp(v, u, p)) + analytic_vjp!(J, v, u, p) = (J .= analytic_vjp(v, u, p)) + + @testset "AutoDiff" begin + @testset for jvp_autodiff in forward_ADs, vjp_autodiff in reverse_ADs + jac_op = JacobianOperator(prob, [2.0, 3.0], prob.u0; jvp_autodiff, vjp_autodiff) + + @testset for u in [rand(2) for _ in 1:4], v in [rand(2) for _ in 1:4] + sop = StatefulJacobianOperator(jac_op, u, prob.p) + @test (sop * v)≈analytic_jvp(v, u, prob.p) atol=1e-5 + @test (sop' * v)≈analytic_vjp(v, u, prob.p) atol=1e-5 + + normal_form_sop = sop' * sop + JᵀJv = normal_form_sop * v + J_analytic = analytic_jac(u, prob.p) + JᵀJv_analytic = J_analytic' * J_analytic * v + @test JᵀJv≈JᵀJv_analytic atol=1e-5 + end + end + end + + prob = NonlinearProblem( + NonlinearFunction{true}((du, u, p) -> du .= u .^ 2 .- p .+ u[2] * u[1]; + jvp = analytic_jvp!, vjp = analytic_vjp!), [1.0, 3.0], 2.0) + + @testset "Analytic JVP/VJP" begin + jac_op = JacobianOperator(prob, [2.0, 3.0], prob.u0) + + @testset for u in [rand(2) for _ in 1:4], v in [rand(2) for _ in 1:4] + sop = StatefulJacobianOperator(jac_op, u, prob.p) + @test (sop * v)≈analytic_jvp(v, u, prob.p) atol=1e-5 + @test (sop' * v)≈analytic_vjp(v, u, prob.p) atol=1e-5 + + normal_form_sop = sop' * sop + JᵀJv = normal_form_sop * v + J_analytic = analytic_jac(u, prob.p) + JᵀJv_analytic = J_analytic' * J_analytic * v + @test JᵀJv≈JᵀJv_analytic atol=1e-5 + end + end + + prob = NonlinearProblem( + NonlinearFunction{true}((du, u, p) -> du .= u .^ 2 .- p .+ u[2] * u[1]; + jac = analytic_jac!), [1.0, 3.0], 2.0) + + @testset "Analytic Jacobian" begin + jac_op = JacobianOperator(prob, [2.0, 3.0], prob.u0) + + @testset for u in [rand(2) for _ in 1:4], v in [rand(2) for _ in 1:4] + sop = StatefulJacobianOperator(jac_op, u, prob.p) + @test (sop * v)≈analytic_jvp(v, u, prob.p) atol=1e-5 + @test (sop' * v)≈analytic_vjp(v, u, prob.p) atol=1e-5 + + normal_form_sop = sop' * sop + JᵀJv = normal_form_sop * v + J_analytic = analytic_jac(u, prob.p) + JᵀJv_analytic = J_analytic' * J_analytic * v + @test JᵀJv≈JᵀJv_analytic atol=1e-5 + end + end +end + +@testitem "Out-of-place Problems" begin + using ADTypes, SciMLBase + using Enzyme, ForwardDiff, FiniteDiff, ReverseDiff, Zygote, Tracker + using SciMLJacobianOperators + + reverse_ADs = [ + AutoEnzyme(), + AutoEnzyme(; mode = Enzyme.Reverse), + AutoZygote(), + AutoTracker(), + AutoReverseDiff(), + AutoFiniteDiff() + ] + + forward_ADs = [ + AutoEnzyme(), + AutoEnzyme(; mode = Enzyme.Forward), + AutoForwardDiff(), + AutoFiniteDiff() + ] + + prob = NonlinearProblem( + NonlinearFunction{false}((u, p) -> u .^ 2 .- p .+ u[2] * u[1]), [1.0, 3.0], 2.0) + + analytic_jac(u, p) = [2 * u[1]+u[2] u[1]; u[2] 2 * u[2]+u[1]] + analytic_jvp(v, u, p) = analytic_jac(u, p) * v + analytic_vjp(v, u, p) = analytic_jac(u, p)' * v + + @testset "AutoDiff" begin + @testset for jvp_autodiff in forward_ADs, vjp_autodiff in reverse_ADs + jac_op = JacobianOperator(prob, [2.0, 3.0], prob.u0; jvp_autodiff, vjp_autodiff) + + @testset for u in [rand(2) for _ in 1:4], v in [rand(2) for _ in 1:4] + sop = StatefulJacobianOperator(jac_op, u, prob.p) + @test (sop * v)≈analytic_jvp(v, u, prob.p) atol=1e-5 + @test (sop' * v)≈analytic_vjp(v, u, prob.p) atol=1e-5 + + normal_form_sop = sop' * sop + JᵀJv = normal_form_sop * v + J_analytic = analytic_jac(u, prob.p) + JᵀJv_analytic = J_analytic' * J_analytic * v + @test JᵀJv≈JᵀJv_analytic atol=1e-5 + end + end + end + + prob = NonlinearProblem( + NonlinearFunction{false}((u, p) -> u .^ 2 .- p .+ u[2] * u[1]; + vjp = analytic_vjp, jvp = analytic_jvp), [1.0, 3.0], 2.0) + + @testset "Analytic JVP/VJP" begin + jac_op = JacobianOperator(prob, [2.0, 3.0], prob.u0) + + @testset for u in [rand(2) for _ in 1:4], v in [rand(2) for _ in 1:4] + sop = StatefulJacobianOperator(jac_op, u, prob.p) + @test (sop * v)≈analytic_jvp(v, u, prob.p) atol=1e-5 + @test (sop' * v)≈analytic_vjp(v, u, prob.p) atol=1e-5 + + normal_form_sop = sop' * sop + JᵀJv = normal_form_sop * v + J_analytic = analytic_jac(u, prob.p) + JᵀJv_analytic = J_analytic' * J_analytic * v + @test JᵀJv≈JᵀJv_analytic atol=1e-5 + end + end + + prob = NonlinearProblem( + NonlinearFunction{false}((u, p) -> u .^ 2 .- p .+ u[2] * u[1]; + jac = analytic_jac), [1.0, 3.0], 2.0) + + @testset "Analytic Jacobian" begin + jac_op = JacobianOperator(prob, [2.0, 3.0], prob.u0) + + @testset for u in [rand(2) for _ in 1:4], v in [rand(2) for _ in 1:4] + sop = StatefulJacobianOperator(jac_op, u, prob.p) + @test (sop * v)≈analytic_jvp(v, u, prob.p) atol=1e-5 + @test (sop' * v)≈analytic_vjp(v, u, prob.p) atol=1e-5 + + normal_form_sop = sop' * sop + JᵀJv = normal_form_sop * v + J_analytic = analytic_jac(u, prob.p) + JᵀJv_analytic = J_analytic' * J_analytic * v + @test JᵀJv≈JᵀJv_analytic atol=1e-5 + end + end +end + +@testitem "Aqua" begin + using SciMLJacobianOperators, Aqua + + Aqua.test_all(SciMLJacobianOperators) +end + +@testitem "Explicit Imports" begin + using SciMLJacobianOperators, ExplicitImports + + @test check_no_implicit_imports(SciMLJacobianOperators) === nothing + @test check_no_stale_explicit_imports(SciMLJacobianOperators) === nothing + @test check_all_qualified_accesses_via_owners(SciMLJacobianOperators) === nothing +end diff --git a/lib/SciMLJacobianOperators/test/runtests.jl b/lib/SciMLJacobianOperators/test/runtests.jl new file mode 100644 index 000000000..6ea6326b0 --- /dev/null +++ b/lib/SciMLJacobianOperators/test/runtests.jl @@ -0,0 +1,5 @@ +using TestItemRunner, InteractiveUtils + +@info sprint(InteractiveUtils.versioninfo) + +@run_package_tests diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 781f6eae9..2e420a761 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -40,13 +40,14 @@ using Preferences: Preferences, @load_preference, @set_preferences! using RecursiveArrayTools: recursivecopy!, recursivefill! using SciMLBase: AbstractNonlinearAlgorithm, JacobianWrapper, AbstractNonlinearProblem, AbstractSciMLOperator, _unwrap_val, has_jac, isinplace, NLStats +using SciMLJacobianOperators: AbstractJacobianOperator, JacobianOperator, VecJacOperator, + JacVecOperator, StatefulJacobianOperator using SparseArrays: AbstractSparseMatrix, SparseMatrixCSC using SparseDiffTools: SparseDiffTools, AbstractSparsityDetection, ApproximateJacobianSparsity, JacPrototypeSparsityDetection, NoSparsityDetection, PrecomputedJacobianColorvec, - SymbolicsSparsityDetection, auto_jacvec, auto_jacvec!, auto_vecjac, - init_jacobian, num_jacvec, num_jacvec!, num_vecjac, num_vecjac!, - sparse_jacobian, sparse_jacobian!, sparse_jacobian_cache + SymbolicsSparsityDetection, init_jacobian, sparse_jacobian, + sparse_jacobian!, sparse_jacobian_cache using StaticArraysCore: StaticArray, SVector, SArray, MArray, Size, SMatrix using SymbolicIndexingInterface: SymbolicIndexingInterface, ParameterIndexingProxy, symbolic_container, parameter_values, state_values, getu, @@ -72,7 +73,6 @@ include("descent/dogleg.jl") include("descent/damped_newton.jl") include("descent/geodesic_acceleration.jl") -include("internal/operators.jl") include("internal/jacobian.jl") include("internal/forward_diff.jl") include("internal/linear_solve.jl") diff --git a/src/abstract_types.jl b/src/abstract_types.jl index 0aa4fd82e..2fe2fd071 100644 --- a/src/abstract_types.jl +++ b/src/abstract_types.jl @@ -277,7 +277,8 @@ Abstract Type for Damping Functions in DampedNewton. ### `__internal_init` specification ```julia -__internal_init(prob::AbstractNonlinearProblem, f::AbstractDampingFunction, initial_damping, +__internal_init( + prob::AbstractNonlinearProblem, f::AbstractDampingFunction, initial_damping, J, fu, u, args...; internal_norm = DEFAULT_NORM, kwargs...) --> AbstractDampingFunctionCache ``` @@ -361,7 +362,8 @@ Abstract Type for all Jacobian Initialization Algorithms used in NonlinearSolve. ### `__internal_init` specification ```julia -__internal_init(prob::AbstractNonlinearProblem, alg::AbstractJacobianInitialization, solver, +__internal_init( + prob::AbstractNonlinearProblem, alg::AbstractJacobianInitialization, solver, f::F, fu, u, p; linsolve = missing, internalnorm::IN = DEFAULT_NORM, kwargs...) ``` @@ -453,7 +455,8 @@ Abstract Type for all Trust Region Methods used in NonlinearSolve.jl. ### `__internal_init` specification ```julia -__internal_init(prob::AbstractNonlinearProblem, alg::AbstractTrustRegionMethod, f::F, fu, u, +__internal_init( + prob::AbstractNonlinearProblem, alg::AbstractTrustRegionMethod, f::F, fu, u, p, args...; internalnorm::IF = DEFAULT_NORM, kwargs...) where {F, IF} --> AbstractTrustRegionMethodCache ``` diff --git a/src/algorithms/broyden.jl b/src/algorithms/broyden.jl index 679897efa..77d2a06c4 100644 --- a/src/algorithms/broyden.jl +++ b/src/algorithms/broyden.jl @@ -28,7 +28,8 @@ search. useful for specific problems, but whether it will work may depend strongly on the problem """ -function Broyden(; max_resets = 100, linesearch = NoLineSearch(), reset_tolerance = nothing, +function Broyden(; + max_resets = 100, linesearch = NoLineSearch(), reset_tolerance = nothing, init_jacobian::Val{IJ} = Val(:identity), autodiff = nothing, alpha = nothing, update_rule::Val{UR} = Val(:good_broyden)) where {IJ, UR} if IJ === :identity diff --git a/src/algorithms/extension_algs.jl b/src/algorithms/extension_algs.jl index ea9ab79ab..a3555a17b 100644 --- a/src/algorithms/extension_algs.jl +++ b/src/algorithms/extension_algs.jl @@ -252,20 +252,22 @@ function NLsolveJL(; method = :trust_region, autodiff = :central, store_trace = end if store_trace !== missing - Base.depwarn("`store_trace` for NLsolveJL has been deprecated and will be removed \ - in v4. Use the `store_trace` keyword argument via the logging API \ - https://docs.sciml.ai/NonlinearSolve/stable/basics/Logging/ \ - instead.", + Base.depwarn( + "`store_trace` for NLsolveJL has been deprecated and will be removed \ + in v4. Use the `store_trace` keyword argument via the logging API \ + https://docs.sciml.ai/NonlinearSolve/stable/basics/Logging/ \ + instead.", :NLsolveJL) else store_trace = false end if extended_trace !== missing - Base.depwarn("`extended_trace` for NLsolveJL has been deprecated and will be \ - removed in v4. Use the `trace_level = TraceAll()` keyword argument \ - via the logging API \ - https://docs.sciml.ai/NonlinearSolve/stable/basics/Logging/ instead.", + Base.depwarn( + "`extended_trace` for NLsolveJL has been deprecated and will be \ + removed in v4. Use the `trace_level = TraceAll()` keyword argument \ + via the logging API \ + https://docs.sciml.ai/NonlinearSolve/stable/basics/Logging/ instead.", :NLsolveJL) else extended_trace = false diff --git a/src/algorithms/levenberg_marquardt.jl b/src/algorithms/levenberg_marquardt.jl index 510226d7d..66554ee2d 100644 --- a/src/algorithms/levenberg_marquardt.jl +++ b/src/algorithms/levenberg_marquardt.jl @@ -53,7 +53,8 @@ function LevenbergMarquardt(; descent = GeodesicAcceleration(descent, finite_diff_step_geodesic, α_geodesic) end trustregion = LevenbergMarquardtTrustRegion(b_uphill) - return GeneralizedFirstOrderAlgorithm(; concrete_jac = true, name = :LevenbergMarquardt, + return GeneralizedFirstOrderAlgorithm(; + concrete_jac = true, name = :LevenbergMarquardt, trustregion, descent, jacobian_ad = autodiff) end diff --git a/src/core/generalized_first_order.jl b/src/core/generalized_first_order.jl index 84e74478f..38c3786f5 100644 --- a/src/core/generalized_first_order.jl +++ b/src/core/generalized_first_order.jl @@ -57,11 +57,13 @@ function GeneralizedFirstOrderAlgorithm{concrete_jac, name}(; max_shrink_times::Int = typemax(Int)) where {concrete_jac, name} forward_ad = ifelse(forward_ad !== nothing, forward_ad, - ifelse(jacobian_ad !== nothing && ADTypes.mode(jacobian_ad) isa ADTypes.ForwardMode, + ifelse( + jacobian_ad !== nothing && ADTypes.mode(jacobian_ad) isa ADTypes.ForwardMode, jacobian_ad, nothing)) reverse_ad = ifelse(reverse_ad !== nothing, reverse_ad, - ifelse(jacobian_ad !== nothing && ADTypes.mode(jacobian_ad) isa ADTypes.ReverseMode, + ifelse( + jacobian_ad !== nothing && ADTypes.mode(jacobian_ad) isa ADTypes.ReverseMode, jacobian_ad, nothing)) if linesearch !== missing && !(linesearch isa AbstractNonlinearSolveLineSearchAlgorithm) diff --git a/src/default.jl b/src/default.jl index b4613a6f7..8a9aac6af 100644 --- a/src/default.jl +++ b/src/default.jl @@ -432,7 +432,8 @@ function FastShortcutNonlinearPolyalg( end else if __is_complex(T) - algs = (Broyden(), Broyden(; init_jacobian = Val(:true_jacobian), autodiff), + algs = ( + Broyden(), Broyden(; init_jacobian = Val(:true_jacobian), autodiff), Klement(; linsolve, precs, autodiff), NewtonRaphson(; concrete_jac, linsolve, precs, autodiff)) else diff --git a/src/descent/geodesic_acceleration.jl b/src/descent/geodesic_acceleration.jl index e2d8612e5..136795057 100644 --- a/src/descent/geodesic_acceleration.jl +++ b/src/descent/geodesic_acceleration.jl @@ -31,7 +31,8 @@ for other methods are not theorectically or experimentally verified. end function Base.show(io::IO, alg::GeodesicAcceleration) - print(io, "GeodesicAcceleration(descent = $(alg.descent), finite_diff_step_geodesic = ", + print( + io, "GeodesicAcceleration(descent = $(alg.descent), finite_diff_step_geodesic = ", "$(alg.finite_diff_step_geodesic), α = $(alg.α))") end @@ -102,7 +103,8 @@ function __internal_init(prob::AbstractNonlinearProblem, alg::GeodesicAccelerati T(alg.finite_diff_step_geodesic), Jv, fu_cache, u_cache, false) end -function __internal_solve!(cache::GeodesicAccelerationCache, J, fu, u, idx::Val{N} = Val(1); +function __internal_solve!( + cache::GeodesicAccelerationCache, J, fu, u, idx::Val{N} = Val(1); skip_solve::Bool = false, kwargs...) where {N} a, v, δu = get_acceleration(cache, idx), get_velocity(cache, idx), get_du(cache, idx) skip_solve && return DescentResult(; δu, extras = (; a, v)) diff --git a/src/globalization/trust_region.jl b/src/globalization/trust_region.jl index 51b54959e..e6e2cba17 100644 --- a/src/globalization/trust_region.jl +++ b/src/globalization/trust_region.jl @@ -386,11 +386,13 @@ function __internal_init( p1, p2, p3, p4 = __get_parameters(T, alg.method) ϵ = T(1e-8) + reverse_ad = get_concrete_reverse_ad(alg.reverse_ad, prob; check_reverse_mode = true) vjp_operator = alg.method isa RUS.__Yuan || alg.method isa RUS.__Bastin ? - VecJacOperator(prob, fu, u; autodiff = alg.reverse_ad) : nothing + VecJacOperator(prob, fu, u; autodiff = reverse_ad) : nothing + forward_ad = get_concrete_forward_ad(alg.forward_ad, prob; check_forward_mode = true) jvp_operator = alg.method isa RUS.__Bastin ? - JacVecOperator(prob, fu, u; autodiff = alg.forward_ad) : nothing + JacVecOperator(prob, fu, u; autodiff = forward_ad) : nothing if alg.method isa RUS.__Yuan Jᵀfu_cache = StatefulJacobianOperator(vjp_operator, u, prob.p) * _vec(fu) diff --git a/src/internal/jacobian.jl b/src/internal/jacobian.jl index b712b93e2..0c6080112 100644 --- a/src/internal/jacobian.jl +++ b/src/internal/jacobian.jl @@ -25,7 +25,8 @@ Construct a cache for the Jacobian of `f` w.r.t. `u`. - `jvp_autodiff`: Automatic Differentiation or Finite Differencing backend for computing the Jacobian-vector product. - `linsolve`: Linear Solver Algorithm used to determine if we need a concrete jacobian - or if possible we can just use a [`NonlinearSolve.JacobianOperator`](@ref) instead. + or if possible we can just use a [`SciMLJacobianOperators.JacobianOperator`](@ref) + instead. """ @concrete mutable struct JacobianCache{iip} <: AbstractNonlinearSolveJacobianCache{iip} J @@ -85,8 +86,7 @@ function JacobianCache(prob, alg, f::F, fu_, u, p; stats, autodiff = nothing, __similar(fu, promote_type(eltype(fu), eltype(u)), length(fu), length(u)) : copy(f.jac_prototype) elseif f.jac_prototype === nothing - zero(init_jacobian( - jac_cache; preserve_immutable = Val(true))) + zero(init_jacobian(jac_cache; preserve_immutable = Val(true))) else f.jac_prototype end @@ -103,10 +103,10 @@ function JacobianCache(prob, alg, f::F, ::Number, u::Number, p; stats, if !(autodiff isa AutoForwardDiff || autodiff isa AutoPolyesterForwardDiff || autodiff isa AutoFiniteDiff) - autodiff = AutoFiniteDiff() # Other cases are not properly supported so we fallback to finite differencing @warn "Scalar AD is supported only for AutoForwardDiff and AutoFiniteDiff. \ Detected $(autodiff). Falling back to AutoFiniteDiff." + autodiff = AutoFiniteDiff() end return JacobianCache{false}( u, f, uf, u, u, p, nothing, alg, stats, autodiff, nothing, nothing) @@ -114,9 +114,9 @@ end @inline (cache::JacobianCache)(u = cache.u) = cache(cache.J, u, cache.p) @inline function (cache::JacobianCache)(::Nothing) - J = cache.J - J isa JacobianOperator && return StatefulJacobianOperator(J, cache.u, cache.p) - return J + cache.J isa JacobianOperator && + return StatefulJacobianOperator(cache.J, cache.u, cache.p) + return cache.J end function (cache::JacobianCache)(J::JacobianOperator, u, p = cache.p) diff --git a/src/internal/operators.jl b/src/internal/operators.jl deleted file mode 100644 index 5bbfcb0bf..000000000 --- a/src/internal/operators.jl +++ /dev/null @@ -1,278 +0,0 @@ -# We want a general form of this in SciMLOperators. However, we use this extensively and we -# can have a custom implementation here till -# https://github.com/SciML/SciMLOperators.jl/issues/223 is resolved. -""" - JacobianOperator{vjp, iip, T} <: AbstractNonlinearSolveOperator{T} - -A Jacobian Operator Provides both JVP and VJP without materializing either (if possible). - -This is an internal operator, and is not guaranteed to have a stable API. It might even be -moved out of NonlinearSolve.jl in the future, without a deprecation cycle. Usage of this -outside NonlinearSolve.jl (by everyone except Avik) is strictly prohibited. - -`T` denotes if the Jacobian is transposed or not. `T = true` means that the Jacobian is -transposed, and `T = false` means that the Jacobian is not transposed. - -### Constructor - -```julia -JacobianOperator( - prob::AbstractNonlinearProblem, fu, u; jvp_autodiff = nothing, vjp_autodiff = nothing, - skip_vjp::Val{NoVJP} = False, skip_jvp::Val{NoJVP} = False) where {NoVJP, NoJVP} -``` - -See also [`NonlinearSolve.VecJacOperator`](@ref) and -[`NonlinearSolve.JacVecOperator`](@ref). -""" -@concrete struct JacobianOperator{vjp, iip, T} <: AbstractNonlinearSolveOperator{T} - jvp_op - vjp_op - - input_cache - output_cache -end - -Base.size(J::JacobianOperator) = prod(size(J.output_cache)), prod(size(J.input_cache)) -function Base.size(J::JacobianOperator, d::Integer) - d == 1 && return prod(size(J.output_cache)) - d == 2 && return prod(size(J.input_cache)) - error("Invalid dimension $d for JacobianOperator") -end - -for op in (:adjoint, :transpose) - @eval function Base.$(op)(operator::JacobianOperator{vjp, iip, T}) where {vjp, iip, T} - return JacobianOperator{!vjp, iip, T}( - operator.jvp_op, operator.vjp_op, operator.output_cache, operator.input_cache) - end -end - -function JacobianOperator(prob::AbstractNonlinearProblem, fu, u; jvp_autodiff = nothing, - vjp_autodiff = nothing, skip_vjp::Val{NoVJP} = False, - skip_jvp::Val{NoJVP} = False) where {NoVJP, NoJVP} - f = prob.f - iip = isinplace(prob) - uf = JacobianWrapper{iip}(f, prob.p) - - vjp_op = if NoVJP - nothing - elseif SciMLBase.has_vjp(f) - f.vjp - elseif u isa Number # Ignore vjp directives - if ForwardDiff.can_dual(typeof(u)) && (vjp_autodiff === nothing || - vjp_autodiff isa AutoForwardDiff || - vjp_autodiff isa AutoPolyesterForwardDiff) - # VJP is same as VJP for scalars - @closure (v, u, p) -> last(__scalar_jacvec(uf, u, v)) - else - @closure (v, u, p) -> FiniteDiff.finite_difference_derivative(uf, u) * v - end - else - vjp_autodiff = __get_nonsparse_ad(get_concrete_reverse_ad( - vjp_autodiff, prob, False)) - if vjp_autodiff isa AutoZygote - iip && error("`AutoZygote` cannot handle inplace problems.") - @closure (v, u, p) -> auto_vecjac(uf, u, v) - elseif vjp_autodiff isa AutoFiniteDiff - if iip - cache1 = zero(fu) - cache2 = zero(fu) - @closure (Jv, v, u, p) -> num_vecjac!(Jv, uf, u, v, cache1, cache2) - else - @closure (v, u, p) -> num_vecjac(uf, __mutable(u), v) - end - else - error("`vjp_autodiff` = `$(typeof(vjp_autodiff))` is not supported in \ - JacobianOperator.") - end - end - - jvp_op = if NoJVP - nothing - elseif SciMLBase.has_jvp(f) - f.jvp - elseif u isa Number # Ignore jvp directives - # Only ForwardDiff if user didn't override - if ForwardDiff.can_dual(typeof(u)) && (jvp_autodiff === nothing || - jvp_autodiff isa AutoForwardDiff || - jvp_autodiff isa AutoPolyesterForwardDiff) - @closure (v, u, p) -> last(__scalar_jacvec(uf, u, v)) - else - @closure (v, u, p) -> FiniteDiff.finite_difference_derivative(uf, u) * v - end - else - jvp_autodiff = __get_nonsparse_ad(get_concrete_forward_ad( - jvp_autodiff, prob, False)) - if jvp_autodiff isa AutoForwardDiff || jvp_autodiff isa AutoPolyesterForwardDiff - if iip - # FIXME: Technically we should propagate the tag but ignoring that for now - cache1 = Dual{typeof(ForwardDiff.Tag(uf, eltype(u))), eltype(u), - 1}.(zero(u), ForwardDiff.Partials.(tuple.(u))) - cache2 = Dual{typeof(ForwardDiff.Tag(uf, eltype(fu))), eltype(fu), - 1}.(zero(fu), ForwardDiff.Partials.(tuple.(fu))) - @closure (Jv, v, u, p) -> auto_jacvec!(Jv, uf, u, v, cache1, cache2) - else - @closure (v, u, p) -> auto_jacvec(uf, u, v) - end - elseif jvp_autodiff isa AutoFiniteDiff - if iip - cache1 = zero(fu) - cache2 = zero(u) - @closure (Jv, v, u, p) -> num_jacvec!(Jv, uf, u, v, cache1, cache2) - else - @closure (v, u, p) -> num_jacvec(uf, u, v) - end - else - error("`jvp_autodiff` = `$(typeof(jvp_autodiff))` is not supported in \ - JacobianOperator.") - end - end - - return JacobianOperator{false, iip, promote_type(eltype(fu), eltype(u))}( - jvp_op, vjp_op, u, fu) -end - -""" - VecJacOperator(args...; autodiff = nothing, kwargs...) - -Constructs a [`JacobianOperator`](@ref) which only provides the VJP using the -`vjp_autodiff = autodiff`. - -This is very similar to `SparseDiffTools.VecJac` but is geared towards -[`NonlinearProblem`](@ref)s. For arguments and keyword arguments see -[`JacobianOperator`](@ref). -""" -function VecJacOperator(args...; autodiff = nothing, kwargs...) - return JacobianOperator(args...; kwargs..., skip_jvp = True, vjp_autodiff = autodiff)' -end - -""" - JacVecOperator(args...; autodiff = nothing, kwargs...) - -Constructs a [`JacobianOperator`](@ref) which only provides the JVP using the -`jvp_autodiff = autodiff`. - -This is very similar to `SparseDiffTools.JacVec` but is geared towards -[`NonlinearProblem`](@ref)s. For arguments and keyword arguments see -[`JacobianOperator`](@ref). -""" -function JacVecOperator(args...; autodiff = nothing, kwargs...) - return JacobianOperator(args...; kwargs..., skip_vjp = True, jvp_autodiff = autodiff) -end - -function (op::JacobianOperator{vjp, iip})(v, u, p) where {vjp, iip} - if vjp - if iip - res = zero(op.output_cache) - op.vjp_op(res, v, u, p) - return res - else - return op.vjp_op(v, u, p) - end - else - if iip - res = zero(op.output_cache) - op.jvp_op(res, v, u, p) - return res - else - return op.jvp_op(v, u, p) - end - end -end - -# Prevent Ambiguity -function (op::JacobianOperator{vjp, iip})(Jv::Number, v::Number, u, p) where {vjp, iip} - error("Inplace Jacobian Operator not possible for scalars.") -end - -function (op::JacobianOperator{vjp, iip})(Jv, v, u, p) where {vjp, iip} - if vjp - if iip - op.vjp_op(Jv, v, u, p) - else - copyto!(Jv, op.vjp_op(v, u, p)) - end - else - if iip - op.jvp_op(Jv, v, u, p) - else - copyto!(Jv, op.jvp_op(v, u, p)) - end - end - return Jv -end - -""" - StatefulJacobianOperator(jac_op::JacobianOperator, u, p) - -Wrapper over a [`JacobianOperator`](@ref) which stores the input `u` and `p` and defines -`mul!` and `*` for computing VJPs and JVPs. -""" -@concrete struct StatefulJacobianOperator{ - vjp, iip, T, J <: JacobianOperator{vjp, iip, T}} <: AbstractNonlinearSolveOperator{T} - jac_op::J - u - p -end - -Base.size(J::StatefulJacobianOperator) = size(J.jac_op) -Base.size(J::StatefulJacobianOperator, d::Integer) = size(J.jac_op, d) - -for op in (:adjoint, :transpose) - @eval function Base.$op(operator::StatefulJacobianOperator) - return StatefulJacobianOperator($(op)(operator.jac_op), operator.u, operator.p) - end -end - -Base.:*(J::StatefulJacobianOperator, v::AbstractArray) = J.jac_op(v, J.u, J.p) -function Base.:*(J_op::StatefulJacobianOperator{vjp, iip, T, J, <:Number}, - v::Number) where {vjp, iip, T, J} - return J_op.jac_op(v, J_op.u, J_op.p) -end - -function LinearAlgebra.mul!( - Jv::AbstractArray, J::StatefulJacobianOperator, v::AbstractArray) - J.jac_op(Jv, v, J.u, J.p) - return Jv -end - -""" - StatefulJacobianNormalFormOperator(vjp_operator, jvp_operator, cache) - -This constructs a Normal Form Jacobian Operator, i.e. it constructs the operator -corresponding to `JᵀJ` where `J` is the Jacobian Operator. This is not meant to be directly -constructed, rather it is constructed with `*` on two [`StatefulJacobianOperator`](@ref)s. -""" -@concrete mutable struct StatefulJacobianNormalFormOperator{T} <: - AbstractNonlinearSolveOperator{T} - vjp_operator - jvp_operator - cache -end - -function Base.size(J::StatefulJacobianNormalFormOperator) - return size(J.vjp_operator, 1), size(J.jvp_operator, 2) -end - -function Base.:*(J1::StatefulJacobianOperator{true}, J2::StatefulJacobianOperator{false}) - cache = J2 * J2.jac_op.input_cache - T = promote_type(eltype(J1), eltype(J2)) - return StatefulJacobianNormalFormOperator{T}(J1, J2, cache) -end - -function LinearAlgebra.mul!(C::StatefulJacobianNormalFormOperator, - A::StatefulJacobianOperator{true}, B::StatefulJacobianOperator{false}) - C.vjp_operator = A - C.jvp_operator = B - return C -end - -function Base.:*(JᵀJ::StatefulJacobianNormalFormOperator, x::AbstractArray) - return JᵀJ.vjp_operator * (JᵀJ.jvp_operator * x) -end - -function LinearAlgebra.mul!( - JᵀJx::AbstractArray, JᵀJ::StatefulJacobianNormalFormOperator, x::AbstractArray) - mul!(JᵀJ.cache, JᵀJ.jvp_operator, x) - mul!(JᵀJx, JᵀJ.vjp_operator, JᵀJ.cache) - return JᵀJx -end diff --git a/src/internal/termination.jl b/src/internal/termination.jl index b64897744..570bae110 100644 --- a/src/internal/termination.jl +++ b/src/internal/termination.jl @@ -8,7 +8,8 @@ function init_termination_cache( AbsSafeBestTerminationMode(Base.Fix2(norm, 2); max_stalled_steps = 32)) end -function init_termination_cache(prob::Union{NonlinearProblem, NonlinearLeastSquaresProblem}, +function init_termination_cache( + prob::Union{NonlinearProblem, NonlinearLeastSquaresProblem}, abstol, reltol, du, u, tc::AbstractNonlinearTerminationMode) tc_ = if hasfield(typeof(tc), :internalnorm) && tc.internalnorm === nothing internalnorm = ifelse( diff --git a/src/utils.jl b/src/utils.jl index 64399a349..1a4ffd2c8 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -48,6 +48,7 @@ end return deepcopy(x) end @inline __maybe_unaliased(x::AbstractNonlinearSolveOperator, alias::Bool) = x +@inline __maybe_unaliased(x::AbstractJacobianOperator, alias::Bool) = x @inline __cond(J::AbstractMatrix) = cond(J) @inline __cond(J::SVector) = __cond(Diagonal(MVector(J))) diff --git a/test/core/23_test_problems_tests.jl b/test/core/23_test_problems_tests.jl index 7305af96f..7648db53f 100644 --- a/test/core/23_test_problems_tests.jl +++ b/test/core/23_test_problems_tests.jl @@ -122,7 +122,7 @@ end test_on_library(problems, dicts, alg_ops, broken_tests, Sys.isapple() ? 1e-3 : 1e-4) end -@testitem "Klement" retries=5 setup=[RobustnessTesting] tags=[:core] begin +@testitem "Klement" setup=[RobustnessTesting] tags=[:core] begin alg_ops = (Klement(), Klement(; init_jacobian = Val(:true_jacobian_diagonal))) broken_tests = Dict(alg => Int[] for alg in alg_ops) diff --git a/test/core/rootfind_tests.jl b/test/core/rootfind_tests.jl index 880b34ff1..bca02a44b 100644 --- a/test/core/rootfind_tests.jl +++ b/test/core/rootfind_tests.jl @@ -52,7 +52,7 @@ end # --- NewtonRaphson tests --- -@testitem "NewtonRaphson" setup=[CoreRootfindTesting] tags=[:core] timeout=3600 begin +@testitem "NewtonRaphson" setup=[CoreRootfindTesting] tags=[:core] begin @testset "LineSearch: $(_nameof(lsmethod)) LineSearch AD: $(_nameof(ad))" for lsmethod in ( Static(), StrongWolfe(), BackTracking(), HagerZhang(), MoreThuente()), ad in (AutoForwardDiff(), AutoZygote(), AutoFiniteDiff()) @@ -118,14 +118,14 @@ end # --- TrustRegion tests --- -@testitem "TrustRegion" setup=[CoreRootfindTesting] tags=[:core] skip=:(Sys.isapple()) timeout=3600 begin +@testitem "TrustRegion" setup=[CoreRootfindTesting] tags=[:core] begin radius_update_schemes = [RadiusUpdateSchemes.Simple, RadiusUpdateSchemes.NocedalWright, RadiusUpdateSchemes.NLsolve, RadiusUpdateSchemes.Hei, RadiusUpdateSchemes.Yuan, RadiusUpdateSchemes.Fan, RadiusUpdateSchemes.Bastin] u0s = ([1.0, 1.0], @SVector[1.0, 1.0], 1.0) linear_solvers = [nothing, LUFactorization(), KrylovJL_GMRES(), \] - @testset "[OOP] u0: $(typeof(u0)) radius_update_scheme: $(radius_update_scheme) linear_solver: $(linsolve)" for u0 in u0s, + @testset "[OOP] u0: $(typeof(u0)) $(radius_update_scheme) $(_nameof(linsolve))" for u0 in u0s, radius_update_scheme in radius_update_schemes, linsolve in linear_solvers @@ -141,7 +141,7 @@ end @test (@ballocated solve!($cache)) < 200 end - @testset "[IIP] u0: $(typeof(u0)) radius_update_scheme: $(radius_update_scheme) linear_solver: $(linsolve)" for u0 in ([ + @testset "[IIP] u0: $(typeof(u0)) $(radius_update_scheme) $(_nameof(linsolve))" for u0 in ([ 1.0, 1.0],), radius_update_scheme in radius_update_schemes, linsolve in linear_solvers @@ -162,7 +162,7 @@ end @test nlprob_iterator_interface(quadratic_f, p, Val(false), TrustRegion()) ≈ sqrt.(p) @test nlprob_iterator_interface(quadratic_f!, p, Val(true), TrustRegion()) ≈ sqrt.(p) - @testset "ADType: $(autodiff) u0: $(_nameof(u0)) radius_update_scheme: $(radius_update_scheme)" for autodiff in ( + @testset "$(_nameof(autodiff)) u0: $(_nameof(u0)) $(radius_update_scheme)" for autodiff in ( AutoSparse(AutoForwardDiff()), AutoSparse(AutoFiniteDiff()), AutoZygote(), AutoSparse(AutoZygote()), AutoSparse(AutoEnzyme())), u0 in (1.0, [1.0, 1.0]), @@ -236,7 +236,7 @@ end # --- LevenbergMarquardt tests --- -@testitem "LevenbergMarquardt" setup=[CoreRootfindTesting] tags=[:core] timeout=3600 begin +@testitem "LevenbergMarquardt" setup=[CoreRootfindTesting] tags=[:core] begin u0s = ([1.0, 1.0], @SVector[1.0, 1.0], 1.0) @testset "[OOP] u0: $(typeof(u0))" for u0 in u0s sol = benchmark_nlsolve_oop(quadratic_f, u0; solver = LevenbergMarquardt()) @@ -322,7 +322,7 @@ end # --- DFSane tests --- -@testitem "DFSane" setup=[CoreRootfindTesting] tags=[:core] timeout=3600 begin +@testitem "DFSane" setup=[CoreRootfindTesting] tags=[:core] begin u0s = ([1.0, 1.0], @SVector[1.0, 1.0], 1.0) @testset "[OOP] u0: $(typeof(u0))" for u0 in u0s @@ -393,7 +393,7 @@ end # --- PseudoTransient tests --- -@testitem "PseudoTransient" setup=[CoreRootfindTesting] tags=[:core] timeout=3600 begin +@testitem "PseudoTransient" setup=[CoreRootfindTesting] tags=[:core] begin # These are tests for NewtonRaphson so we should set alpha_initial to be high so that we # converge quickly @testset "PT: alpha_initial = 10.0 PT AD: $(ad)" for ad in ( @@ -462,7 +462,7 @@ end # --- Broyden tests --- -@testitem "Broyden" setup=[CoreRootfindTesting] tags=[:core] skip=:(Sys.isapple()) timeout=3600 begin +@testitem "Broyden" setup=[CoreRootfindTesting] tags=[:core] begin @testset "LineSearch: $(_nameof(lsmethod)) LineSearch AD: $(_nameof(ad)) Init Jacobian: $(init_jacobian) Update Rule: $(update_rule)" for lsmethod in ( Static(), StrongWolfe(), BackTracking(), HagerZhang(), MoreThuente(), LiFukushimaLineSearch()), @@ -512,7 +512,7 @@ end # --- Klement tests --- -@testitem "Klement" setup=[CoreRootfindTesting] tags=[:core] skip=:(Sys.isapple()) timeout=3600 begin +@testitem "Klement" setup=[CoreRootfindTesting] tags=[:core] begin @testset "LineSearch: $(_nameof(lsmethod)) LineSearch AD: $(_nameof(ad)) Init Jacobian: $(init_jacobian)" for lsmethod in ( Static(), StrongWolfe(), BackTracking(), HagerZhang(), MoreThuente()), ad in (AutoForwardDiff(), AutoZygote(), AutoFiniteDiff()), @@ -561,7 +561,7 @@ end # --- LimitedMemoryBroyden tests --- -@testitem "LimitedMemoryBroyden" setup=[CoreRootfindTesting] tags=[:core] skip=:(Sys.isapple()) timeout=3600 begin +@testitem "LimitedMemoryBroyden" setup=[CoreRootfindTesting] tags=[:core] begin @testset "LineSearch: $(_nameof(lsmethod)) LineSearch AD: $(_nameof(ad))" for lsmethod in ( Static(), StrongWolfe(), BackTracking(), HagerZhang(), MoreThuente(), LiFukushimaLineSearch()), @@ -611,7 +611,7 @@ end end # Miscellaneous Tests -@testitem "Custom JVP" setup=[CoreRootfindTesting] tags=[:core] timeout=3600 begin +@testitem "Custom JVP" setup=[CoreRootfindTesting] tags=[:core] begin function F(u::Vector{Float64}, p::Vector{Float64}) Δ = Tridiagonal(-ones(99), 2 * ones(100), -ones(99)) return u + 0.1 * u .* Δ * u - p diff --git a/test/gpu/core_tests.jl b/test/gpu/core_tests.jl index 79a20cc97..d6d38ee1a 100644 --- a/test/gpu/core_tests.jl +++ b/test/gpu/core_tests.jl @@ -1,37 +1,39 @@ -@testitem "CUDA Tests" tags=[:cuda] skip=:(using CUDA; !CUDA.functional()) begin +@testitem "CUDA Tests" tags=[:cuda] begin using CUDA, NonlinearSolve, LinearSolve, StableRNGs - CUDA.allowscalar(false) + if CUDA.functional() + CUDA.allowscalar(false) - A = cu(rand(StableRNG(0), 4, 4)) - u0 = cu(rand(StableRNG(0), 4)) - b = cu(rand(StableRNG(0), 4)) + A = cu(rand(StableRNG(0), 4, 4)) + u0 = cu(rand(StableRNG(0), 4)) + b = cu(rand(StableRNG(0), 4)) - linear_f(du, u, p) = (du .= A * u .+ b) + linear_f(du, u, p) = (du .= A * u .+ b) - prob = NonlinearProblem(linear_f, u0) + prob = NonlinearProblem(linear_f, u0) - SOLVERS = (NewtonRaphson(), LevenbergMarquardt(; linsolve = QRFactorization()), - LevenbergMarquardt(; linsolve = KrylovJL_GMRES()), PseudoTransient(), - Klement(), Broyden(; linesearch = LiFukushimaLineSearch()), - LimitedMemoryBroyden(; threshold = 2, linesearch = LiFukushimaLineSearch()), - DFSane(), TrustRegion(; linsolve = QRFactorization()), - TrustRegion(; linsolve = KrylovJL_GMRES(), concrete_jac = true), # Needed if Zygote not loaded - nothing) + SOLVERS = (NewtonRaphson(), LevenbergMarquardt(; linsolve = QRFactorization()), + LevenbergMarquardt(; linsolve = KrylovJL_GMRES()), PseudoTransient(), + Klement(), Broyden(; linesearch = LiFukushimaLineSearch()), + LimitedMemoryBroyden(; threshold = 2, linesearch = LiFukushimaLineSearch()), + DFSane(), TrustRegion(; linsolve = QRFactorization()), + TrustRegion(; linsolve = KrylovJL_GMRES(), concrete_jac = true), # Needed if Zygote not loaded + nothing) - @testset "[IIP] GPU Solvers" begin - for alg in SOLVERS - @test_nowarn sol = solve(prob, alg; abstol = 1.0f-5, reltol = 1.0f-5) + @testset "[IIP] GPU Solvers" begin + for alg in SOLVERS + @test_nowarn sol = solve(prob, alg; abstol = 1.0f-5, reltol = 1.0f-5) + end end - end - linear_f(u, p) = A * u .+ b + linear_f(u, p) = A * u .+ b - prob = NonlinearProblem{false}(linear_f, u0) + prob = NonlinearProblem{false}(linear_f, u0) - @testset "[OOP] GPU Solvers" begin - for alg in SOLVERS - @test_nowarn sol = solve(prob, alg; abstol = 1.0f-5, reltol = 1.0f-5) + @testset "[OOP] GPU Solvers" begin + for alg in SOLVERS + @test_nowarn sol = solve(prob, alg; abstol = 1.0f-5, reltol = 1.0f-5) + end end end end diff --git a/test/misc/qa_tests.jl b/test/misc/qa_tests.jl index 7e8bf9ce6..5c60e5339 100644 --- a/test/misc/qa_tests.jl +++ b/test/misc/qa_tests.jl @@ -24,8 +24,6 @@ end @test check_no_implicit_imports(NonlinearSolve; skip = (NonlinearSolve, Base, Core, SimpleNonlinearSolve, SciMLBase)) === nothing - @test check_no_stale_explicit_imports(NonlinearSolve) === nothing - @test check_all_qualified_accesses_via_owners(NonlinearSolve) === nothing end diff --git a/test/runtests.jl b/test/runtests.jl index 421a96f7b..74676f3ad 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,10 +1,17 @@ -using ReTestItems, CUDA +using ReTestItems, NonlinearSolve, Hwloc, InteractiveUtils -const GROUP = get(ENV, "GROUP", CUDA.functional() ? "All" : "All") +@info sprint(InteractiveUtils.versioninfo) -if GROUP == "All" - ReTestItems.runtests(@__DIR__) -else - tags = [Symbol(lowercase(GROUP))] - ReTestItems.runtests(@__DIR__; tags) -end +const GROUP = lowercase(get(ENV, "GROUP", "All")) + +const RETESTITEMS_NWORKERS = parse( + Int, get(ENV, "RETESTITEMS_NWORKERS", string(min(Hwloc.num_physical_cores(), 4)))) +const RETESTITEMS_NWORKER_THREADS = parse(Int, + get(ENV, "RETESTITEMS_NWORKER_THREADS", + string(max(Hwloc.num_virtual_cores() ÷ RETESTITEMS_NWORKERS, 1)))) + +@info "Running tests for group: $(GROUP) with $(RETESTITEMS_NWORKERS) workers" + +ReTestItems.runtests(NonlinearSolve; tags = (GROUP == "all" ? nothing : [Symbol(GROUP)]), + nworkers = RETESTITEMS_NWORKERS, nworker_threads = RETESTITEMS_NWORKER_THREADS, + testitem_timeout = 3600, retries = 4)