diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 9eb146557..b74c820cc 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -14,12 +14,12 @@ steps: Pkg.Registry.update(); # Install packages present in subdirectories dev_pks = Pkg.PackageSpec[]; - for path in ("lib/SciMLJacobianOperators", "lib/BracketingNonlinearSolve", "lib/NonlinearSolveBase", "lib/SimpleNonlinearSolve") + for path in ("lib/SciMLJacobianOperators", "lib/BracketingNonlinearSolve", "lib/NonlinearSolveBase", "lib/SimpleNonlinearSolve", "lib/NonlinearSolveFirstOrder", "lib/NonlinearSolveSpectralMethods", "lib/NonlinearSolveQuasiNewton") push!(dev_pks, Pkg.PackageSpec(; path)); end Pkg.develop(dev_pks); Pkg.instantiate(); - Pkg.test(; coverage=true)' + Pkg.test(; coverage="user")' agents: queue: "juliagpu" cuda: "*" @@ -42,12 +42,12 @@ steps: Pkg.Registry.update(); # Install packages present in subdirectories dev_pks = Pkg.PackageSpec[]; - for path in ("lib/NonlinearSolveBase", "lib/BracketingNonlinearSolve") + for path in ("lib/NonlinearSolveBase", "lib/BracketingNonlinearSolve", "lib/SciMLJacobianOperators") push!(dev_pks, Pkg.PackageSpec(; path)) end Pkg.develop(dev_pks); Pkg.instantiate(); - Pkg.test(; coverage=true)' + Pkg.test(; coverage="user")' agents: queue: "juliagpu" cuda: "*" diff --git a/.github/workflows/CI_BracketingNonlinearSolve.yml b/.github/workflows/CI_BracketingNonlinearSolve.yml index 6d78e5612..956fa33a2 100644 --- a/.github/workflows/CI_BracketingNonlinearSolve.yml +++ b/.github/workflows/CI_BracketingNonlinearSolve.yml @@ -6,8 +6,9 @@ on: - master paths: - "lib/BracketingNonlinearSolve/**" - - "lib/NonlinearSolveBase/**" - ".github/workflows/CI_BracketingNonlinearSolve.yml" + - "lib/NonlinearSolveBase/**" + - "lib/SciMLJacobianOperators/**" push: branches: - master @@ -25,7 +26,7 @@ jobs: fail-fast: false matrix: version: - - "min" + - "1.10" - "1" os: - ubuntu-latest @@ -52,16 +53,54 @@ jobs: Pkg.Registry.update() # Install packages present in subdirectories dev_pks = Pkg.PackageSpec[] - for path in ("lib/NonlinearSolveBase",) + for path in ("lib/NonlinearSolveBase", "lib/SciMLJacobianOperators") + push!(dev_pks, Pkg.PackageSpec(; path)) + end + Pkg.develop(dev_pks) + Pkg.instantiate() + Pkg.test(; coverage="user") + shell: julia --color=yes --code-coverage=user --depwarn=yes --project=lib/BracketingNonlinearSolve {0} + - uses: julia-actions/julia-processcoverage@v1 + with: + directories: lib/BracketingNonlinearSolve/src,lib/BracketingNonlinearSolve/ext,lib/NonlinearSolveBase/src,lib/NonlinearSolveBase/ext,lib/SciMLJacobianOperators/src + - uses: codecov/codecov-action@v4 + with: + file: lcov.info + token: ${{ secrets.CODECOV_TOKEN }} + verbose: true + fail_ci_if_error: true + + downgrade: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + version: + - "1.10" + steps: + - uses: actions/checkout@v4 + - uses: julia-actions/setup-julia@v2 + with: + version: ${{ matrix.version }} + - uses: julia-actions/julia-downgrade-compat@v1 + with: + skip: NonlinearSolveBase, SciMLJacobianOperators + - 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/NonlinearSolveBase", "lib/SciMLJacobianOperators") push!(dev_pks, Pkg.PackageSpec(; path)) end Pkg.develop(dev_pks) Pkg.instantiate() - Pkg.test(; coverage=true) + Pkg.test(; coverage="user") shell: julia --color=yes --code-coverage=user --depwarn=yes --project=lib/BracketingNonlinearSolve {0} - uses: julia-actions/julia-processcoverage@v1 with: - directories: lib/BracketingNonlinearSolve/src,lib/BracketingNonlinearSolve/ext + directories: lib/BracketingNonlinearSolve/src,lib/BracketingNonlinearSolve/ext,lib/NonlinearSolveBase/src,lib/NonlinearSolveBase/ext,lib/SciMLJacobianOperators/src - uses: codecov/codecov-action@v4 with: file: lcov.info diff --git a/.github/workflows/CI_NonlinearSolve.yml b/.github/workflows/CI_NonlinearSolve.yml index 843a04b54..f685a7e7a 100644 --- a/.github/workflows/CI_NonlinearSolve.yml +++ b/.github/workflows/CI_NonlinearSolve.yml @@ -14,6 +14,9 @@ on: - "lib/BracketingNonlinearSolve/**" - "lib/NonlinearSolveBase/**" - "lib/SimpleNonlinearSolve/**" + - "lib/NonlinearSolveFirstOrder/**" + - "lib/NonlinearSolveSpectralMethods/**" + - "lib/NonlinearSolveQuasiNewton/**" push: branches: - master @@ -31,12 +34,11 @@ jobs: fail-fast: false matrix: group: - - Core - - Downstream - - Misc - - Wrappers + - core + - downstream + - wrappers version: - - "min" + - "1.10" - "1" os: - ubuntu-latest @@ -63,18 +65,62 @@ jobs: Pkg.Registry.update() # Install packages present in subdirectories dev_pks = Pkg.PackageSpec[] - for path in ("lib/SciMLJacobianOperators", "lib/BracketingNonlinearSolve", "lib/NonlinearSolveBase", "lib/SimpleNonlinearSolve") + for path in ("lib/SciMLJacobianOperators", "lib/BracketingNonlinearSolve", "lib/NonlinearSolveBase", "lib/SimpleNonlinearSolve", "lib/NonlinearSolveFirstOrder", "lib/NonlinearSolveSpectralMethods", "lib/NonlinearSolveQuasiNewton") push!(dev_pks, Pkg.PackageSpec(; path)) end Pkg.develop(dev_pks) Pkg.instantiate() - Pkg.test(; coverage=true) + Pkg.test(; coverage="user") 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 + directories: src,ext,lib/SciMLJacobianOperators/src,lib/BracketingNonlinearSolve/src,lib/NonlinearSolveBase/src,lib/NonlinearSolveBase/ext,lib/SimpleNonlinearSolve/src,lib/NonlinearSolveFirstOrder/src,lib/NonlinearSolveSpectralMethods/src,lib/NonlinearSolveQuasiNewton/src + - uses: codecov/codecov-action@v4 + with: + file: lcov.info + token: ${{ secrets.CODECOV_TOKEN }} + verbose: true + fail_ci_if_error: true + + downgrade: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + version: + - "1.10" + group: + - core + - downstream + - wrappers + steps: + - uses: actions/checkout@v4 + - uses: julia-actions/setup-julia@v2 + with: + version: ${{ matrix.version }} + - uses: julia-actions/julia-downgrade-compat@v1 + with: + skip: SciMLJacobianOperators, BracketingNonlinearSolve, NonlinearSolveBase, SimpleNonlinearSolve, NonlinearSolveFirstOrder, NonlinearSolveSpectralMethods, NonlinearSolveQuasiNewton + - 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", "lib/BracketingNonlinearSolve", "lib/NonlinearSolveBase", "lib/SimpleNonlinearSolve", "lib/NonlinearSolveFirstOrder", "lib/NonlinearSolveSpectralMethods", "lib/NonlinearSolveQuasiNewton") + push!(dev_pks, Pkg.PackageSpec(; path)) + end + Pkg.develop(dev_pks) + Pkg.instantiate() + Pkg.test(; coverage="user") + 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,lib/BracketingNonlinearSolve/src,lib/NonlinearSolveBase/src,lib/NonlinearSolveBase/ext,lib/SimpleNonlinearSolve/src,lib/NonlinearSolveFirstOrder/src,lib/NonlinearSolveSpectralMethods/src,lib/NonlinearSolveQuasiNewton/src - uses: codecov/codecov-action@v4 with: file: lcov.info diff --git a/.github/workflows/CI_NonlinearSolveBase.yml b/.github/workflows/CI_NonlinearSolveBase.yml index f3878acef..8b303ae2e 100644 --- a/.github/workflows/CI_NonlinearSolveBase.yml +++ b/.github/workflows/CI_NonlinearSolveBase.yml @@ -7,6 +7,7 @@ on: paths: - "lib/NonlinearSolveBase/**" - ".github/workflows/CI_NonlinearSolveBase.yml" + - "lib/SciMLJacobianOperators/**" push: branches: - master @@ -24,7 +25,7 @@ jobs: fail-fast: false matrix: version: - - "min" + - "1.10" - "1" os: - ubuntu-latest @@ -49,12 +50,56 @@ jobs: 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) + Pkg.test(; coverage="user") shell: julia --color=yes --code-coverage=user --depwarn=yes --project=lib/NonlinearSolveBase {0} - uses: julia-actions/julia-processcoverage@v1 with: - directories: lib/NonlinearSolveBase/src,lib/NonlinearSolveBase/ext + directories: lib/NonlinearSolveBase/src,lib/NonlinearSolveBase/ext,lib/SciMLJacobianOperators/src + - uses: codecov/codecov-action@v4 + with: + file: lcov.info + token: ${{ secrets.CODECOV_TOKEN }} + verbose: true + fail_ci_if_error: true + + downgrade: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + version: + - "1.10" + steps: + - uses: actions/checkout@v4 + - uses: julia-actions/setup-julia@v2 + with: + version: ${{ matrix.version }} + - uses: julia-actions/julia-downgrade-compat@v1 + with: + skip: SciMLJacobianOperators + - 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="user") + shell: julia --color=yes --code-coverage=user --depwarn=yes --project=lib/NonlinearSolveBase {0} + - uses: julia-actions/julia-processcoverage@v1 + with: + directories: lib/NonlinearSolveBase/src,lib/NonlinearSolveBase/ext,lib/SciMLJacobianOperators/src - uses: codecov/codecov-action@v4 with: file: lcov.info diff --git a/.github/workflows/CI_NonlinearSolveFirstOrder.yml b/.github/workflows/CI_NonlinearSolveFirstOrder.yml new file mode 100644 index 000000000..8f68f398b --- /dev/null +++ b/.github/workflows/CI_NonlinearSolveFirstOrder.yml @@ -0,0 +1,109 @@ +name: CI (NonlinearSolveFirstOrder) + +on: + pull_request: + branches: + - master + paths: + - "lib/NonlinearSolveFirstOrder/**" + - ".github/workflows/CI_NonlinearSolveFirstOrder.yml" + - "lib/NonlinearSolveBase/**" + - "lib/SciMLJacobianOperators/**" + 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: + version: + - "1.10" + - "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", "lib/NonlinearSolveBase") + push!(dev_pks, Pkg.PackageSpec(; path)) + end + Pkg.develop(dev_pks) + Pkg.instantiate() + Pkg.test(; coverage="user") + shell: julia --color=yes --code-coverage=user --depwarn=yes --project=lib/NonlinearSolveFirstOrder {0} + - uses: julia-actions/julia-processcoverage@v1 + with: + directories: lib/NonlinearSolveFirstOrder/src,lib/NonlinearSolveBase/src,lib/NonlinearSolveBase/ext,lib/SciMLJacobianOperators/src + - uses: codecov/codecov-action@v4 + with: + file: lcov.info + token: ${{ secrets.CODECOV_TOKEN }} + verbose: true + fail_ci_if_error: true + + downgrade: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + version: + - "1.10" + steps: + - uses: actions/checkout@v4 + - uses: julia-actions/setup-julia@v2 + with: + version: ${{ matrix.version }} + - uses: julia-actions/julia-downgrade-compat@v1 + with: + skip: NonlinearSolveBase, SciMLJacobianOperators + - 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", "lib/NonlinearSolveBase") + push!(dev_pks, Pkg.PackageSpec(; path)) + end + Pkg.develop(dev_pks) + Pkg.instantiate() + Pkg.test(; coverage="user") + shell: julia --color=yes --code-coverage=user --depwarn=yes --project=lib/NonlinearSolveFirstOrder {0} + - uses: julia-actions/julia-processcoverage@v1 + with: + directories: lib/NonlinearSolveFirstOrder/src,lib/NonlinearSolveBase/src,lib/NonlinearSolveBase/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_NonlinearSolveQuasiNewton.yml b/.github/workflows/CI_NonlinearSolveQuasiNewton.yml new file mode 100644 index 000000000..3c0904739 --- /dev/null +++ b/.github/workflows/CI_NonlinearSolveQuasiNewton.yml @@ -0,0 +1,109 @@ +name: CI (NonlinearSolveQuasiNewton) + +on: + pull_request: + branches: + - master + paths: + - "lib/NonlinearSolveQuasiNewton/**" + - ".github/workflows/CI_NonlinearSolveQuasiNewton.yml" + - "lib/NonlinearSolveBase/**" + - "lib/SciMLJacobianOperators/**" + 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: + version: + - "1.10" + - "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", "lib/NonlinearSolveBase") + push!(dev_pks, Pkg.PackageSpec(; path)) + end + Pkg.develop(dev_pks) + Pkg.instantiate() + Pkg.test(; coverage="user") + shell: julia --color=yes --code-coverage=user --depwarn=yes --project=lib/NonlinearSolveQuasiNewton {0} + - uses: julia-actions/julia-processcoverage@v1 + with: + directories: lib/NonlinearSolveQuasiNewton/src,lib/NonlinearSolveBase/src,lib/NonlinearSolveBase/ext,lib/SciMLJacobianOperators/src + - uses: codecov/codecov-action@v4 + with: + file: lcov.info + token: ${{ secrets.CODECOV_TOKEN }} + verbose: true + fail_ci_if_error: true + + downgrade: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + version: + - "1.10" + steps: + - uses: actions/checkout@v4 + - uses: julia-actions/setup-julia@v2 + with: + version: ${{ matrix.version }} + - uses: julia-actions/julia-downgrade-compat@v1 + with: + skip: NonlinearSolveBase, SciMLJacobianOperators + - 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", "lib/NonlinearSolveBase") + push!(dev_pks, Pkg.PackageSpec(; path)) + end + Pkg.develop(dev_pks) + Pkg.instantiate() + Pkg.test(; coverage="user") + shell: julia --color=yes --code-coverage=user --depwarn=yes --project=lib/NonlinearSolveQuasiNewton {0} + - uses: julia-actions/julia-processcoverage@v1 + with: + directories: lib/NonlinearSolveQuasiNewton/src,lib/NonlinearSolveBase/src,lib/NonlinearSolveBase/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_NonlinearSolveSpectralMethods.yml b/.github/workflows/CI_NonlinearSolveSpectralMethods.yml new file mode 100644 index 000000000..f39420efa --- /dev/null +++ b/.github/workflows/CI_NonlinearSolveSpectralMethods.yml @@ -0,0 +1,109 @@ +name: CI (NonlinearSolveSpectralMethods) + +on: + pull_request: + branches: + - master + paths: + - "lib/NonlinearSolveSpectralMethods/**" + - ".github/workflows/CI_NonlinearSolveSpectralMethods.yml" + - "lib/NonlinearSolveBase/**" + - "lib/SciMLJacobianOperators/**" + 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: + version: + - "1.10" + - "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", "lib/NonlinearSolveBase") + push!(dev_pks, Pkg.PackageSpec(; path)) + end + Pkg.develop(dev_pks) + Pkg.instantiate() + Pkg.test(; coverage="user") + shell: julia --color=yes --code-coverage=user --depwarn=yes --project=lib/NonlinearSolveSpectralMethods {0} + - uses: julia-actions/julia-processcoverage@v1 + with: + directories: lib/NonlinearSolveSpectralMethods/src,lib/NonlinearSolveBase/src,lib/NonlinearSolveBase/ext,lib/SciMLJacobianOperators/src + - uses: codecov/codecov-action@v4 + with: + file: lcov.info + token: ${{ secrets.CODECOV_TOKEN }} + verbose: true + fail_ci_if_error: true + + downgrade: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + version: + - "1.10" + steps: + - uses: actions/checkout@v4 + - uses: julia-actions/setup-julia@v2 + with: + version: ${{ matrix.version }} + - uses: julia-actions/julia-downgrade-compat@v1 + with: + skip: NonlinearSolveBase, SciMLJacobianOperators + - 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", "lib/NonlinearSolveBase") + push!(dev_pks, Pkg.PackageSpec(; path)) + end + Pkg.develop(dev_pks) + Pkg.instantiate() + Pkg.test(; coverage="user") + shell: julia --color=yes --code-coverage=user --depwarn=yes --project=lib/NonlinearSolveSpectralMethods {0} + - uses: julia-actions/julia-processcoverage@v1 + with: + directories: lib/NonlinearSolveSpectralMethods/src,lib/NonlinearSolveBase/src,lib/NonlinearSolveBase/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_SciMLJacobianOperators.yml b/.github/workflows/CI_SciMLJacobianOperators.yml index 92daf23e9..4d1f6780d 100644 --- a/.github/workflows/CI_SciMLJacobianOperators.yml +++ b/.github/workflows/CI_SciMLJacobianOperators.yml @@ -24,7 +24,7 @@ jobs: fail-fast: false matrix: version: - - "min" + - "1.10" - "1" os: - ubuntu-latest @@ -50,7 +50,37 @@ jobs: import Pkg Pkg.Registry.update() Pkg.instantiate() - Pkg.test(; coverage=true) + Pkg.test(; coverage="user") + shell: julia --color=yes --code-coverage=user --depwarn=yes --project=lib/SciMLJacobianOperators {0} + - uses: julia-actions/julia-processcoverage@v1 + with: + directories: lib/SciMLJacobianOperators/src + - uses: codecov/codecov-action@v4 + with: + file: lcov.info + token: ${{ secrets.CODECOV_TOKEN }} + verbose: true + fail_ci_if_error: true + + downgrade: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + version: + - "1.10" + steps: + - uses: actions/checkout@v4 + - uses: julia-actions/setup-julia@v2 + with: + version: ${{ matrix.version }} + - uses: julia-actions/julia-downgrade-compat@v1 + - name: "Install Dependencies and Run Tests" + run: | + import Pkg + Pkg.Registry.update() + Pkg.instantiate() + Pkg.test(; coverage="user") shell: julia --color=yes --code-coverage=user --depwarn=yes --project=lib/SciMLJacobianOperators {0} - uses: julia-actions/julia-processcoverage@v1 with: diff --git a/.github/workflows/CI_SimpleNonlinearSolve.yml b/.github/workflows/CI_SimpleNonlinearSolve.yml index 11c3ef7c2..0212232a8 100644 --- a/.github/workflows/CI_SimpleNonlinearSolve.yml +++ b/.github/workflows/CI_SimpleNonlinearSolve.yml @@ -6,9 +6,10 @@ on: - master paths: - "lib/SimpleNonlinearSolve/**" + - ".github/workflows/CI_SimpleNonlinearSolve.yml" - "lib/BracketingNonlinearSolve/**" - "lib/NonlinearSolveBase/**" - - ".github/workflows/CI_SimpleNonlinearSolve.yml" + - "lib/SciMLJacobianOperators/**" push: branches: - master @@ -26,7 +27,7 @@ jobs: fail-fast: false matrix: version: - - "min" + - "1.10" - "1" os: - ubuntu-latest @@ -57,18 +58,62 @@ jobs: Pkg.Registry.update() # Install packages present in subdirectories dev_pks = Pkg.PackageSpec[] - for path in ("lib/NonlinearSolveBase", "lib/BracketingNonlinearSolve") + for path in ("lib/NonlinearSolveBase", "lib/BracketingNonlinearSolve", "lib/SciMLJacobianOperators") + push!(dev_pks, Pkg.PackageSpec(; path)) + end + Pkg.develop(dev_pks) + Pkg.instantiate() + Pkg.test(; coverage="user") + shell: julia --color=yes --code-coverage=user --depwarn=yes --project=lib/SimpleNonlinearSolve {0} + env: + GROUP: ${{ matrix.group }} + - uses: julia-actions/julia-processcoverage@v1 + with: + directories: lib/SimpleNonlinearSolve/src,lib/SimpleNonlinearSolve/ext,lib/SciMLJacobianOperators/src,lib/BracketingNonlinearSolve/src,lib/BracketingNonlinearSolve/ext,lib/NonlinearSolveBase/src,lib/NonlinearSolveBase/ext + - uses: codecov/codecov-action@v4 + with: + file: lcov.info + token: ${{ secrets.CODECOV_TOKEN }} + verbose: true + fail_ci_if_error: true + + downgrade: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + version: + - "1.10" + group: + - core + - adjoint + - alloc_check + steps: + - uses: actions/checkout@v4 + - uses: julia-actions/setup-julia@v2 + with: + version: ${{ matrix.version }} + - uses: julia-actions/julia-downgrade-compat@v1 + with: + skip: NonlinearSolveBase, BracketingNonlinearSolve, SciMLJacobianOperators + - 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/NonlinearSolveBase", "lib/BracketingNonlinearSolve", "lib/SciMLJacobianOperators") push!(dev_pks, Pkg.PackageSpec(; path)) end Pkg.develop(dev_pks) Pkg.instantiate() - Pkg.test(; coverage=true) + Pkg.test(; coverage="user") shell: julia --color=yes --code-coverage=user --depwarn=yes --project=lib/SimpleNonlinearSolve {0} env: GROUP: ${{ matrix.group }} - uses: julia-actions/julia-processcoverage@v1 with: - directories: lib/SimpleNonlinearSolve/src,lib/SimpleNonlinearSolve/ext + directories: lib/SimpleNonlinearSolve/src,lib/SimpleNonlinearSolve/ext,lib/SciMLJacobianOperators/src,lib/BracketingNonlinearSolve/src,lib/BracketingNonlinearSolve/ext,lib/NonlinearSolveBase/src,lib/NonlinearSolveBase/ext - uses: codecov/codecov-action@v4 with: file: lcov.info diff --git a/.github/workflows/Documentation.yml b/.github/workflows/Documentation.yml index 9dc416799..af44cd769 100644 --- a/.github/workflows/Documentation.yml +++ b/.github/workflows/Documentation.yml @@ -22,7 +22,7 @@ jobs: Pkg.Registry.update() # Install packages present in subdirectories dev_pks = Pkg.PackageSpec[] - for path in ("lib/SciMLJacobianOperators", ".", "lib/SimpleNonlinearSolve", "lib/NonlinearSolveBase", "lib/BracketingNonlinearSolve") + for path in ("lib/SciMLJacobianOperators", ".", "lib/SimpleNonlinearSolve", "lib/NonlinearSolveBase", "lib/BracketingNonlinearSolve", "lib/NonlinearSolveFirstOrder", "lib/NonlinearSolveQuasiNewton", "lib/NonlinearSolveSpectralMethods") push!(dev_pks, Pkg.PackageSpec(; path)) end Pkg.develop(dev_pks) diff --git a/.github/workflows/Downgrade.yml b/.github/workflows/Downgrade.yml deleted file mode 100644 index 172457df8..000000000 --- a/.github/workflows/Downgrade.yml +++ /dev/null @@ -1,55 +0,0 @@ -name: Downgrade -on: - pull_request: - branches: - - master - paths-ignore: - - "docs/**" - push: - branches: - - master - paths-ignore: - - "docs/**" -jobs: - test: - runs-on: ubuntu-latest - strategy: - fail-fast: false - matrix: - version: - - "1.10" - group: - - Core - - Downstream - - Misc - - Wrappers - steps: - - uses: actions/checkout@v4 - - uses: julia-actions/setup-julia@v2 - with: - version: ${{ matrix.version }} - - uses: julia-actions/julia-downgrade-compat@v1 - - 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", "lib/BracketingNonlinearSolve", "lib/NonlinearSolveBase", "lib/SimpleNonlinearSolve") - 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/Downstream.yml b/.github/workflows/Downstream.yml index 9e54283aa..dc40d46f2 100644 --- a/.github/workflows/Downstream.yml +++ b/.github/workflows/Downstream.yml @@ -32,7 +32,18 @@ jobs: with: version: ${{ matrix.julia-version }} arch: x64 - - uses: julia-actions/julia-buildpkg@latest + - name: "Install Dependencies" + run: | + import Pkg + Pkg.Registry.update() + # Install packages present in subdirectories + dev_pks = Pkg.PackageSpec[] + for path in ("lib/SciMLJacobianOperators", "lib/BracketingNonlinearSolve", "lib/NonlinearSolveBase", "lib/SimpleNonlinearSolve", "lib/NonlinearSolveFirstOrder", "lib/NonlinearSolveSpectralMethods", "lib/NonlinearSolveQuasiNewton") + push!(dev_pks, Pkg.PackageSpec(; path)) + end + Pkg.develop(dev_pks) + Pkg.instantiate() + shell: julia --color=yes --code-coverage=user --depwarn=yes --project=. {0} - name: Clone Downstream uses: actions/checkout@v4 with: @@ -46,7 +57,7 @@ jobs: # force it to use this PR's version of the package Pkg.develop(PackageSpec(path=".")) # resolver may fail with main deps Pkg.update() - Pkg.test(coverage=true) # resolver may fail with test time deps + Pkg.test(coverage="user") # resolver may fail with test time deps catch err err isa Pkg.Resolve.ResolverError || rethrow() # If we can't resolve that means this is incompatible by SemVer and this is fine diff --git a/.github/workflows/Invalidations.yml b/.github/workflows/Invalidations.yml index 66c86a362..e221a7c75 100644 --- a/.github/workflows/Invalidations.yml +++ b/.github/workflows/Invalidations.yml @@ -20,14 +20,36 @@ jobs: with: version: '1' - uses: actions/checkout@v4 - - uses: julia-actions/julia-buildpkg@v1 + - name: "Install Dependencies" + run: | + import Pkg + Pkg.Registry.update() + # Install packages present in subdirectories + dev_pks = Pkg.PackageSpec[] + for path in ("lib/SciMLJacobianOperators", "lib/BracketingNonlinearSolve", "lib/NonlinearSolveBase", "lib/SimpleNonlinearSolve", "lib/NonlinearSolveFirstOrder", "lib/NonlinearSolveSpectralMethods", "lib/NonlinearSolveQuasiNewton") + push!(dev_pks, Pkg.PackageSpec(; path)) + end + Pkg.develop(dev_pks) + Pkg.instantiate() + shell: julia --color=yes --code-coverage=user --depwarn=yes --project=. {0} - uses: julia-actions/julia-invalidations@v1 id: invs_pr - uses: actions/checkout@v4 with: ref: ${{ github.event.repository.default_branch }} - - uses: julia-actions/julia-buildpkg@v1 + - name: "Install Dependencies" + run: | + import Pkg + Pkg.Registry.update() + # Install packages present in subdirectories + dev_pks = Pkg.PackageSpec[] + for path in ("lib/SciMLJacobianOperators", "lib/BracketingNonlinearSolve", "lib/NonlinearSolveBase", "lib/SimpleNonlinearSolve", "lib/NonlinearSolveFirstOrder", "lib/NonlinearSolveSpectralMethods", "lib/NonlinearSolveQuasiNewton") + push!(dev_pks, Pkg.PackageSpec(; path)) + end + Pkg.develop(dev_pks) + Pkg.instantiate() + shell: julia --color=yes --code-coverage=user --depwarn=yes --project=. {0} - uses: julia-actions/julia-invalidations@v1 id: invs_default diff --git a/Project.toml b/Project.toml index 50a7a8d4f..845611499 100644 --- a/Project.toml +++ b/Project.toml @@ -6,6 +6,7 @@ version = "4.0.0" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9" +BracketingNonlinearSolve = "70df07ce-3d50-431d-a3e7-ca6ddb60ac1e" CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2" ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471" DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e" @@ -13,29 +14,24 @@ DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" FastClosures = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a" FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" -LazyArrays = "5078a376-72f3-5289-bfd5-ec5146d43c02" LineSearch = "87fe0de2-c867-4266-b59a-2f0a94fc965b" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae" -MaybeInplace = "bb5d69b7-63fc-4a16-80bd-7e42200c7bdb" NonlinearSolveBase = "be0214bd-f91f-a760-ac4e-3421ce2b2da0" +NonlinearSolveFirstOrder = "5959db7a-ea39-4486-b5fe-2dd0bf03d60d" +NonlinearSolveQuasiNewton = "9a2c21bd-3a47-402d-9113-8faf9a0ee114" +NonlinearSolveSpectralMethods = "26075421-4e9a-44e1-8bd1-420ed7ad02b2" PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" Preferences = "21216c6a-2e73-6563-6e65-726566657250" -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" -SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961" SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35" StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5" -TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f" [weakdeps] -BandedMatrices = "aae01518-5342-5314-be14-df237901396f" FastLevenbergMarquardt = "7a0df574-e128-4d35-8cbd-3d84502bf7ce" FixedPointAcceleration = "817d07cb-a79a-5c30-9a31-890123675176" LeastSquaresOptim = "0fc2ff8b-aaa3-5acd-a817-1944a5e08891" @@ -50,7 +46,6 @@ SpeedMapping = "f1835b91-879b-4a3f-a438-e4baacf14412" Sundials = "c3572dad-4567-51f8-b174-8c6c989267f4" [extensions] -NonlinearSolveBandedMatricesExt = "BandedMatrices" NonlinearSolveFastLevenbergMarquardtExt = "FastLevenbergMarquardt" NonlinearSolveFixedPointAccelerationExt = "FixedPointAcceleration" NonlinearSolveLeastSquaresOptimExt = "LeastSquaresOptim" @@ -68,11 +63,12 @@ Aqua = "0.8" ArrayInterface = "7.16" BandedMatrices = "1.5" BenchmarkTools = "1.4" +BracketingNonlinearSolve = "1" CUDA = "5.5" CommonSolve = "0.2.4" ConcreteStructs = "0.2.3" -DiffEqBase = "6.155.3" -DifferentiationInterface = "0.6.16" +DiffEqBase = "6.158.3" +DifferentiationInterface = "0.6.18" Enzyme = "0.13.11" ExplicitImports = "1.5" FastClosures = "0.3.2" @@ -82,34 +78,31 @@ FixedPointAcceleration = "0.3" ForwardDiff = "0.10.36" Hwloc = "3" InteractiveUtils = "<0.0.1, 1" -LazyArrays = "1.8.2, 2" LeastSquaresOptim = "0.8.5" LineSearch = "0.1.4" LineSearches = "7.3" LinearAlgebra = "1.10" -LinearSolve = "2.35" +LinearSolve = "2.36.1" MINPACK = "1.2" MPI = "0.20.22" -MaybeInplace = "0.1.4" NLSolvers = "0.5" NLsolve = "4.5" NaNMath = "1" NonlinearProblemLibrary = "0.1.2" NonlinearSolveBase = "1" +NonlinearSolveFirstOrder = "1" +NonlinearSolveQuasiNewton = "1" +NonlinearSolveSpectralMethods = "1" OrdinaryDiffEqTsit5 = "1.1.0" -PETSc = "0.2" +PETSc = "0.3" Pkg = "1.10" PrecompileTools = "1.2" Preferences = "1.4" -Printf = "1.10" Random = "1.10" ReTestItems = "1.24" -RecursiveArrayTools = "3.27" Reexport = "1.2" SIAMFANLEquations = "1.0.1" -SciMLBase = "2.54.0" -SciMLJacobianOperators = "0.1" -SciMLOperators = "0.3.10" +SciMLBase = "2.58" SimpleNonlinearSolve = "2" SparseArrays = "1.10" SparseConnectivityTracer = "0.6.5" @@ -121,7 +114,6 @@ StaticArraysCore = "1.4" Sundials = "4.23.1" SymbolicIndexingInterface = "0.3.31" Test = "1.10" -TimerOutputs = "0.5.23" Zygote = "0.6.69" julia = "1.10" diff --git a/common/common_nlls_testing.jl b/common/common_nlls_testing.jl new file mode 100644 index 000000000..86c33e8ac --- /dev/null +++ b/common/common_nlls_testing.jl @@ -0,0 +1,50 @@ +using NonlinearSolveBase, SciMLBase, StableRNGs, ForwardDiff, Random, LinearAlgebra + +true_function(x, θ) = @. θ[1] * exp(θ[2] * x) * cos(θ[3] * x + θ[4]) +true_function(y, x, θ) = (@. y = θ[1] * exp(θ[2] * x) * cos(θ[3] * x + θ[4])) + +θ_true = [1.0, 0.1, 2.0, 0.5] +x = [-1.0, -0.5, 0.0, 0.5, 1.0] + +const y_target = true_function(x, θ_true) + +function loss_function(θ, p) + ŷ = true_function(p, θ) + return ŷ .- y_target +end + +function loss_function(resid, θ, p) + true_function(resid, p, θ) + resid .= resid .- y_target + return resid +end + +θ_init = θ_true .+ randn!(StableRNG(0), similar(θ_true)) * 0.1 + +function vjp(v, θ, p) + resid = zeros(length(p)) + J = ForwardDiff.jacobian((resid, θ) -> loss_function(resid, θ, p), resid, θ) + return vec(v' * J) +end + +function vjp!(Jv, v, θ, p) + resid = zeros(length(p)) + J = ForwardDiff.jacobian((resid, θ) -> loss_function(resid, θ, p), resid, θ) + mul!(vec(Jv), transpose(J), v) + return nothing +end + +prob_oop = NonlinearLeastSquaresProblem{false}(loss_function, θ_init, x) +prob_iip = NonlinearLeastSquaresProblem{true}( + NonlinearFunction(loss_function; resid_prototype = zero(y_target)), θ_init, x +) +prob_oop_vjp = NonlinearLeastSquaresProblem( + NonlinearFunction{false}(loss_function; vjp), θ_init, x +) +prob_iip_vjp = NonlinearLeastSquaresProblem( + NonlinearFunction{true}(loss_function; resid_prototype = zero(y_target), vjp = vjp!), + θ_init, x +) + +export prob_oop, prob_iip, prob_oop_vjp, prob_iip_vjp +export true_function, θ_true, x, y_target, loss_function, θ_init diff --git a/common/common_rootfind_testing.jl b/common/common_rootfind_testing.jl new file mode 100644 index 000000000..824d35528 --- /dev/null +++ b/common/common_rootfind_testing.jl @@ -0,0 +1,54 @@ +using NonlinearSolveBase, SciMLBase + +const TERMINATION_CONDITIONS = [ + NormTerminationMode(Base.Fix1(maximum, abs)), + RelTerminationMode(), + RelNormTerminationMode(Base.Fix1(maximum, abs)), + RelNormSafeTerminationMode(Base.Fix1(maximum, abs)), + RelNormSafeBestTerminationMode(Base.Fix1(maximum, abs)), + AbsTerminationMode(), + AbsNormTerminationMode(Base.Fix1(maximum, abs)), + AbsNormSafeTerminationMode(Base.Fix1(maximum, abs)), + AbsNormSafeBestTerminationMode(Base.Fix1(maximum, abs)) +] + +quadratic_f(u, p) = u .* u .- p +quadratic_f!(du, u, p) = (du .= u .* u .- p) +quadratic_f2(u, p) = @. p[1] * u * u - p[2] + +function newton_fails(u, p) + return 0.010000000000000002 .+ + 10.000000000000002 ./ (1 .+ + (0.21640425613334457 .+ + 216.40425613334457 ./ (1 .+ + (0.21640425613334457 .+ + 216.40425613334457 ./ (1 .+ 0.0006250000000000001(u .^ 2.0))) .^ 2.0)) .^ + 2.0) .- 0.0011552453009332421u .- p +end + +function solve_oop(f, u0, p = 2.0; solver, kwargs...) + prob = NonlinearProblem{false}(f, u0, p) + return solve(prob, solver; abstol = 1e-9, kwargs...) +end + +function solve_iip(f, u0, p = 2.0; solver, kwargs...) + prob = NonlinearProblem{true}(f, u0, p) + return solve(prob, solver; abstol = 1e-9, kwargs...) +end + +function nlprob_iterator_interface(f, p_range, isinplace, solver) + probN = NonlinearProblem{isinplace}(f, isinplace ? [0.5] : 0.5, p_range[begin]) + cache = init(probN, solver; maxiters = 100, abstol = 1e-10) + sols = zeros(length(p_range)) + for (i, p) in enumerate(p_range) + reinit!(cache, isinplace ? [cache.u[1]] : cache.u; p = p) + sol = solve!(cache) + sols[i] = isinplace ? sol.u[1] : sol.u + end + return sols +end + +export TERMINATION_CONDITIONS +export quadratic_f, quadratic_f!, quadratic_f2, newton_fails +export solve_oop, solve_iip +export nlprob_iterator_interface diff --git a/common/nlls_problem_workloads.jl b/common/nlls_problem_workloads.jl new file mode 100644 index 000000000..0d7c0f7ef --- /dev/null +++ b/common/nlls_problem_workloads.jl @@ -0,0 +1,26 @@ +using SciMLBase: NonlinearLeastSquaresProblem, NonlinearFunction, NoSpecialize + +nonlinear_functions = ( + (NonlinearFunction{false, NoSpecialize}((u, p) -> (u .^ 2 .- p)[1:1]), [0.1, 0.0]), + ( + NonlinearFunction{false, NoSpecialize}((u, p) -> vcat(u .* u .- p, u .* u .- p)), + [0.1, 0.1] + ), + ( + NonlinearFunction{true, NoSpecialize}( + (du, u, p) -> du[1] = u[1] * u[1] - p, resid_prototype = zeros(1) + ), + [0.1, 0.0] + ), + ( + NonlinearFunction{true, NoSpecialize}( + (du, u, p) -> du .= vcat(u .* u .- p, u .* u .- p), resid_prototype = zeros(4) + ), + [0.1, 0.1] + ) +) + +nlls_problems = NonlinearLeastSquaresProblem[] +for (fn, u0) in nonlinear_functions + push!(nlls_problems, NonlinearLeastSquaresProblem(fn, u0, 2.0)) +end diff --git a/common/nonlinear_problem_workloads.jl b/common/nonlinear_problem_workloads.jl new file mode 100644 index 000000000..43d0de29a --- /dev/null +++ b/common/nonlinear_problem_workloads.jl @@ -0,0 +1,12 @@ +using SciMLBase: NonlinearProblem, NonlinearFunction, NoSpecialize + +nonlinear_functions = ( + (NonlinearFunction{false, NoSpecialize}((u, p) -> u .* u .- p), 0.1), + (NonlinearFunction{false, NoSpecialize}((u, p) -> u .* u .- p), [0.1]), + (NonlinearFunction{true, NoSpecialize}((du, u, p) -> du .= u .* u .- p), [0.1]) +) + +nonlinear_problems = NonlinearProblem[] +for (fn, u0) in nonlinear_functions + push!(nonlinear_problems, NonlinearProblem(fn, u0, 2.0)) +end diff --git a/docs/Project.toml b/docs/Project.toml index 9ed14da4e..c6c28820c 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -14,6 +14,9 @@ InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae" NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec" NonlinearSolveBase = "be0214bd-f91f-a760-ac4e-3421ce2b2da0" +NonlinearSolveFirstOrder = "5959db7a-ea39-4486-b5fe-2dd0bf03d60d" +NonlinearSolveQuasiNewton = "9a2c21bd-3a47-402d-9113-8faf9a0ee114" +NonlinearSolveSpectralMethods = "26075421-4e9a-44e1-8bd1-420ed7ad02b2" OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a" PETSc = "ace2c81b-2b5f-4b1e-a30d-d662738edfe0" Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" @@ -32,7 +35,7 @@ AlgebraicMultigrid = "0.5, 0.6" ArrayInterface = "6, 7" BenchmarkTools = "1" BracketingNonlinearSolve = "1" -DiffEqBase = "6.158" +DiffEqBase = "6.158.3" DifferentiationInterface = "0.6.16" Documenter = "1" DocumenterCitations = "1" @@ -42,11 +45,14 @@ InteractiveUtils = "<0.0.1, 1" LinearSolve = "2" NonlinearSolve = "4" NonlinearSolveBase = "1" +NonlinearSolveFirstOrder = "1" +NonlinearSolveQuasiNewton = "1" +NonlinearSolveSpectralMethods = "1" OrdinaryDiffEqTsit5 = "1.1.0" -PETSc = "0.2" +PETSc = "0.3" Plots = "1" Random = "1.10" -SciMLBase = "2.4" +SciMLBase = "2.58" SciMLJacobianOperators = "0.1" SimpleNonlinearSolve = "2" SparseConnectivityTracer = "0.6.5" diff --git a/docs/make.jl b/docs/make.jl index bdc2b7fe9..29b1535dd 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -1,13 +1,23 @@ using Documenter, DocumenterCitations, DocumenterInterLinks -using NonlinearSolve, SimpleNonlinearSolve, Sundials, SteadyStateDiffEq, SciMLBase, - BracketingNonlinearSolve, NonlinearSolveBase -using SciMLJacobianOperators import DiffEqBase -cp(joinpath(@__DIR__, "Manifest.toml"), - joinpath(@__DIR__, "src/assets/Manifest.toml"), force = true) -cp(joinpath(@__DIR__, "Project.toml"), - joinpath(@__DIR__, "src/assets/Project.toml"), force = true) +using Sundials +using NonlinearSolveBase, SciMLBase, DiffEqBase +using SimpleNonlinearSolve, BracketingNonlinearSolve +using NonlinearSolveFirstOrder, NonlinearSolveQuasiNewton, NonlinearSolveSpectralMethods +using SciMLJacobianOperators +using NonlinearSolve, SteadyStateDiffEq + +cp( + joinpath(@__DIR__, "Manifest.toml"), + joinpath(@__DIR__, "src/assets/Manifest.toml"); + force = true +) +cp( + joinpath(@__DIR__, "Project.toml"), + joinpath(@__DIR__, "src/assets/Project.toml"); + force = true +) include("pages.jl") @@ -20,20 +30,29 @@ interlinks = InterLinks( makedocs(; sitename = "NonlinearSolve.jl", - authors = "Chris Rackauckas", - modules = [NonlinearSolve, SimpleNonlinearSolve, SteadyStateDiffEq, DiffEqBase, - Sundials, NonlinearSolveBase, SciMLBase, SciMLJacobianOperators, - BracketingNonlinearSolve], + authors = "SciML", + modules = [ + NonlinearSolveBase, SciMLBase, DiffEqBase, + SimpleNonlinearSolve, BracketingNonlinearSolve, + NonlinearSolveFirstOrder, NonlinearSolveQuasiNewton, NonlinearSolveSpectralMethods, + Sundials, + SciMLJacobianOperators, + NonlinearSolve, SteadyStateDiffEq + ], clean = true, doctest = false, linkcheck = true, - linkcheck_ignore = ["https://twitter.com/ChrisRackauckas/status/1544743542094020615", - "https://link.springer.com/article/10.1007/s40096-020-00339-4"], + linkcheck_ignore = [ + "https://twitter.com/ChrisRackauckas/status/1544743542094020615", + "https://link.springer.com/article/10.1007/s40096-020-00339-4" + ], checkdocs = :exports, warnonly = [:missing_docs], plugins = [bib, interlinks], - format = Documenter.HTML(assets = ["assets/favicon.ico", "assets/citations.css"], - canonical = "https://docs.sciml.ai/NonlinearSolve/stable/"), + format = Documenter.HTML( + assets = ["assets/favicon.ico", "assets/citations.css"], + canonical = "https://docs.sciml.ai/NonlinearSolve/stable/" + ), pages ) diff --git a/docs/src/basics/autodiff.md b/docs/src/basics/autodiff.md index d2d01e00b..5cdffa75b 100644 --- a/docs/src/basics/autodiff.md +++ b/docs/src/basics/autodiff.md @@ -3,7 +3,7 @@ !!! note We support all backends supported by DifferentiationInterface.jl. Please refer to - the [backends page](https://gdalle.github.io/DifferentiationInterface.jl/DifferentiationInterface/stable/explanation/backends/) + the [backends page](https://juliadiff.org/DifferentiationInterface.jl/DifferentiationInterface/stable/explanation/backends/) for more information. ## Summary of Finite Differencing Backends diff --git a/docs/src/basics/faq.md b/docs/src/basics/faq.md index 265cee264..4d428250a 100644 --- a/docs/src/basics/faq.md +++ b/docs/src/basics/faq.md @@ -102,7 +102,8 @@ It is hard to say why your code is not fast. Take a look at the there is type instability. If you are using the defaults for the autodiff and your problem is not a scalar or using -static arrays, ForwardDiff will create type unstable code. See this simple example: +static arrays, ForwardDiff will create type unstable code and lead to dynamic dispatch +internally. See this simple example: ```@example type_unstable using NonlinearSolve, InteractiveUtils @@ -136,14 +137,17 @@ prob = NonlinearProblem(f, [1.0, 2.0], 2.0) nothing # hide ``` -Oh no! This is type unstable. This is because ForwardDiff.jl will chunk the jacobian -computation and the type of this chunksize can't be statically inferred. To fix this, we -directly specify the chunksize: +Ah it is still type stable. But internally since the chunksize is not statically inferred, +it will be dynamic and lead to dynamic dispatch. To fix this, we directly specify the +chunksize: ```@example type_unstable -@code_warntype solve(prob, +@code_warntype solve( + prob, NewtonRaphson(; - autodiff = AutoForwardDiff(; chunksize = NonlinearSolve.pickchunksize(prob.u0)))) + autodiff = AutoForwardDiff(; chunksize = NonlinearSolve.pickchunksize(prob.u0)) + ) +) nothing # hide ``` diff --git a/docs/src/devdocs/algorithm_helpers.md b/docs/src/devdocs/algorithm_helpers.md index c945b6003..2bd038792 100644 --- a/docs/src/devdocs/algorithm_helpers.md +++ b/docs/src/devdocs/algorithm_helpers.md @@ -3,8 +3,8 @@ ## Pseudo Transient Method ```@docs -NonlinearSolve.SwitchedEvolutionRelaxation -NonlinearSolve.SwitchedEvolutionRelaxationCache +NonlinearSolveFirstOrder.SwitchedEvolutionRelaxation +NonlinearSolveFirstOrder.SwitchedEvolutionRelaxationCache ``` ## Approximate Jacobian Methods @@ -12,54 +12,54 @@ NonlinearSolve.SwitchedEvolutionRelaxationCache ### Initialization ```@docs -NonlinearSolve.IdentityInitialization -NonlinearSolve.TrueJacobianInitialization -NonlinearSolve.BroydenLowRankInitialization +NonlinearSolveQuasiNewton.IdentityInitialization +NonlinearSolveQuasiNewton.TrueJacobianInitialization +NonlinearSolveQuasiNewton.BroydenLowRankInitialization ``` ### Jacobian Structure ```@docs -NonlinearSolve.FullStructure -NonlinearSolve.DiagonalStructure +NonlinearSolveQuasiNewton.FullStructure +NonlinearSolveQuasiNewton.DiagonalStructure ``` ### Jacobian Caches ```@docs -NonlinearSolve.InitializedApproximateJacobianCache +NonlinearSolveQuasiNewton.InitializedApproximateJacobianCache ``` ### Reset Methods ```@docs -NonlinearSolve.NoChangeInStateReset -NonlinearSolve.IllConditionedJacobianReset +NonlinearSolveQuasiNewton.NoChangeInStateReset +NonlinearSolveQuasiNewton.IllConditionedJacobianReset ``` ### Update Rules ```@docs -NonlinearSolve.GoodBroydenUpdateRule -NonlinearSolve.BadBroydenUpdateRule -NonlinearSolve.KlementUpdateRule +NonlinearSolveQuasiNewton.GoodBroydenUpdateRule +NonlinearSolveQuasiNewton.BadBroydenUpdateRule +NonlinearSolveQuasiNewton.KlementUpdateRule ``` ## Levenberg Marquardt Method ```@docs -NonlinearSolve.LevenbergMarquardtTrustRegion +NonlinearSolveFirstOrder.LevenbergMarquardtTrustRegion ``` ## Trust Region Method ```@docs -NonlinearSolve.GenericTrustRegionScheme +NonlinearSolveFirstOrder.GenericTrustRegionScheme ``` ## Miscellaneous ```@docs -NonlinearSolve.callback_into_cache! -NonlinearSolve.concrete_jac +NonlinearSolveBase.callback_into_cache! +NonlinearSolveBase.concrete_jac ``` diff --git a/docs/src/devdocs/internal_interfaces.md b/docs/src/devdocs/internal_interfaces.md index 53e20d754..f4474e1aa 100644 --- a/docs/src/devdocs/internal_interfaces.md +++ b/docs/src/devdocs/internal_interfaces.md @@ -3,50 +3,43 @@ ## Solvers ```@docs -NonlinearSolve.AbstractNonlinearSolveAlgorithm -NonlinearSolve.AbstractNonlinearSolveExtensionAlgorithm -NonlinearSolve.AbstractNonlinearSolveCache +NonlinearSolveBase.AbstractNonlinearSolveAlgorithm +NonlinearSolveBase.AbstractNonlinearSolveCache ``` -## Descent Algorithms +## Descent Directions ```@docs -NonlinearSolve.AbstractDescentAlgorithm -NonlinearSolve.AbstractDescentCache +NonlinearSolveBase.AbstractDescentDirection +NonlinearSolveBase.AbstractDescentCache ``` -## Descent Results +### Descent Results ```@docs -NonlinearSolve.DescentResult +NonlinearSolveBase.DescentResult ``` ## Approximate Jacobian ```@docs -NonlinearSolve.AbstractApproximateJacobianStructure -NonlinearSolve.AbstractJacobianInitialization -NonlinearSolve.AbstractApproximateJacobianUpdateRule -NonlinearSolve.AbstractApproximateJacobianUpdateRuleCache -NonlinearSolve.AbstractResetCondition +NonlinearSolveBase.AbstractApproximateJacobianStructure +NonlinearSolveBase.AbstractJacobianInitialization +NonlinearSolveBase.AbstractApproximateJacobianUpdateRule +NonlinearSolveBase.AbstractApproximateJacobianUpdateRuleCache +NonlinearSolveBase.AbstractResetCondition ``` ## Damping Algorithms ```@docs -NonlinearSolve.AbstractDampingFunction -NonlinearSolve.AbstractDampingFunctionCache +NonlinearSolveBase.AbstractDampingFunction +NonlinearSolveBase.AbstractDampingFunctionCache ``` ## Trust Region ```@docs -NonlinearSolve.AbstractTrustRegionMethod -NonlinearSolve.AbstractTrustRegionMethodCache -``` - -## Tracing - -```@docs -NonlinearSolve.AbstractNonlinearSolveTraceLevel +NonlinearSolveBase.AbstractTrustRegionMethod +NonlinearSolveBase.AbstractTrustRegionMethodCache ``` diff --git a/docs/src/devdocs/jacobian.md b/docs/src/devdocs/jacobian.md index 082478f7c..3e54966f4 100644 --- a/docs/src/devdocs/jacobian.md +++ b/docs/src/devdocs/jacobian.md @@ -1,6 +1,5 @@ # Jacobian Wrappers ```@docs -NonlinearSolve.AbstractNonlinearSolveJacobianCache -NonlinearSolve.JacobianCache +NonlinearSolveBase.construct_jacobian_cache ``` diff --git a/docs/src/devdocs/linear_solve.md b/docs/src/devdocs/linear_solve.md index 88fa87440..dadd7dea1 100644 --- a/docs/src/devdocs/linear_solve.md +++ b/docs/src/devdocs/linear_solve.md @@ -1,6 +1,6 @@ # Linear Solve ```@docs -NonlinearSolve.AbstractLinearSolverCache -NonlinearSolve.LinearSolverCache +NonlinearSolveBase.AbstractLinearSolverCache +NonlinearSolveBase.construct_linear_solver ``` diff --git a/docs/src/devdocs/operators.md b/docs/src/devdocs/operators.md index 15d00093a..249eda29d 100644 --- a/docs/src/devdocs/operators.md +++ b/docs/src/devdocs/operators.md @@ -1,13 +1,7 @@ # Custom SciML Operators -## Abstract Operators - -```@docs -NonlinearSolve.AbstractNonlinearSolveOperator -``` - ## Low-Rank Jacobian Operators ```@docs -NonlinearSolve.BroydenLowRankJacobian +NonlinearSolveQuasiNewton.BroydenLowRankJacobian ``` diff --git a/docs/src/native/diagnostics.md b/docs/src/native/diagnostics.md index 35f11552f..1c0571c87 100644 --- a/docs/src/native/diagnostics.md +++ b/docs/src/native/diagnostics.md @@ -5,9 +5,9 @@ These functions are not exported since the names have a potential for conflict. ```@docs -NonlinearSolve.enable_timer_outputs -NonlinearSolve.disable_timer_outputs -NonlinearSolve.@static_timeit +NonlinearSolveBase.enable_timer_outputs +NonlinearSolveBase.disable_timer_outputs +NonlinearSolveBase.@static_timeit ``` ## Tracing API @@ -17,6 +17,3 @@ TraceAll TraceWithJacobianConditionNumber TraceMinimal ``` - -For details about the arguments refer to the documentation of -[`NonlinearSolve.AbstractNonlinearSolveTraceLevel`](@ref). diff --git a/docs/src/native/solvers.md b/docs/src/native/solvers.md index 5653f1fab..a3aa900e0 100644 --- a/docs/src/native/solvers.md +++ b/docs/src/native/solvers.md @@ -81,7 +81,7 @@ All of the previously mentioned solvers are wrappers around the following solver are meant for advanced users and allow building custom solvers. ```@docs -ApproximateJacobianSolveAlgorithm +QuasiNewtonAlgorithm GeneralizedFirstOrderAlgorithm GeneralizedDFSane ``` diff --git a/docs/src/release_notes.md b/docs/src/release_notes.md index 1dc3d9433..bb0b5dec9 100644 --- a/docs/src/release_notes.md +++ b/docs/src/release_notes.md @@ -4,6 +4,7 @@ ### Breaking Changes in `NonlinearSolve.jl` v4 + - `ApproximateJacobianSolveAlgorithm` has been renamed to `QuasiNewtonAlgorithm`. - See [common breaking changes](@ref common-breaking-changes-v4v2) below. ### Breaking Changes in `SimpleNonlinearSolve.jl` v2 diff --git a/ext/NonlinearSolveBandedMatricesExt.jl b/ext/NonlinearSolveBandedMatricesExt.jl deleted file mode 100644 index b79df3578..000000000 --- a/ext/NonlinearSolveBandedMatricesExt.jl +++ /dev/null @@ -1,11 +0,0 @@ -module NonlinearSolveBandedMatricesExt - -using BandedMatrices: BandedMatrix -using LinearAlgebra: Diagonal -using NonlinearSolve: NonlinearSolve -using SparseArrays: sparse - -# This is used if we vcat a Banded Jacobian with a Diagonal Matrix in Levenberg -@inline NonlinearSolve._vcat(B::BandedMatrix, D::Diagonal) = vcat(sparse(B), D) - -end diff --git a/ext/NonlinearSolveFastLevenbergMarquardtExt.jl b/ext/NonlinearSolveFastLevenbergMarquardtExt.jl index e772043b3..a851c3a92 100644 --- a/ext/NonlinearSolveFastLevenbergMarquardtExt.jl +++ b/ext/NonlinearSolveFastLevenbergMarquardtExt.jl @@ -1,72 +1,84 @@ module NonlinearSolveFastLevenbergMarquardtExt -using ArrayInterface: ArrayInterface using FastClosures: @closure + +using ArrayInterface: ArrayInterface using FastLevenbergMarquardt: FastLevenbergMarquardt -using NonlinearSolveBase: NonlinearSolveBase, get_tolerance -using NonlinearSolve: NonlinearSolve, FastLevenbergMarquardtJL -using SciMLBase: SciMLBase, NonlinearLeastSquaresProblem, NonlinearProblem, ReturnCode using StaticArraysCore: SArray -const FastLM = FastLevenbergMarquardt +using NonlinearSolveBase: NonlinearSolveBase +using NonlinearSolve: NonlinearSolve, FastLevenbergMarquardtJL +using SciMLBase: SciMLBase, AbstractNonlinearProblem, ReturnCode -@inline function _fast_lm_solver(::FastLevenbergMarquardtJL{linsolve}, x) where {linsolve} - if linsolve === :cholesky - return FastLM.CholeskySolver(ArrayInterface.undefmatrix(x)) - elseif linsolve === :qr - return FastLM.QRSolver(eltype(x), length(x)) - else - throw(ArgumentError("Unknown FastLevenbergMarquardt Linear Solver: $linsolve")) - end -end -@inline _fast_lm_solver(::FastLevenbergMarquardtJL{linsolve}, ::SArray) where {linsolve} = linsolve +const FastLM = FastLevenbergMarquardt -function SciMLBase.__solve(prob::Union{NonlinearLeastSquaresProblem, NonlinearProblem}, - alg::FastLevenbergMarquardtJL, args...; alias_u0 = false, abstol = nothing, - reltol = nothing, maxiters = 1000, termination_condition = nothing, kwargs...) - NonlinearSolve.__test_termination_condition( - termination_condition, :FastLevenbergMarquardt) +function SciMLBase.__solve( + prob::AbstractNonlinearProblem, alg::FastLevenbergMarquardtJL, args...; + alias_u0 = false, abstol = nothing, reltol = nothing, maxiters = 1000, + termination_condition = nothing, kwargs... +) + NonlinearSolveBase.assert_extension_supported_termination_condition( + termination_condition, alg + ) - fn, u, resid = NonlinearSolve.__construct_extension_f( - prob; alias_u0, can_handle_oop = Val(prob.u0 isa SArray)) + f_wrapped, u, resid = NonlinearSolveBase.construct_extension_function_wrapper( + prob; alias_u0, can_handle_oop = Val(prob.u0 isa SArray) + ) f = if prob.u0 isa SArray - @closure (u, p) -> fn(u) + @closure (u, p) -> f_wrapped(u) else - @closure (du, u, p) -> fn(du, u) + @closure (du, u, p) -> f_wrapped(du, u) end - abstol = get_tolerance(abstol, eltype(u)) - reltol = get_tolerance(reltol, eltype(u)) - _jac_fn = NonlinearSolve.__construct_extension_jac( - prob, alg, u, resid; alg.autodiff, can_handle_oop = Val(prob.u0 isa SArray)) + abstol = NonlinearSolveBase.get_tolerance(abstol, eltype(u)) + reltol = NonlinearSolveBase.get_tolerance(reltol, eltype(u)) + + jac_fn_wrapped = NonlinearSolveBase.construct_extension_jac( + prob, alg, u, resid; alg.autodiff, can_handle_oop = Val(prob.u0 isa SArray) + ) jac_fn = if prob.u0 isa SArray - @closure (u, p) -> _jac_fn(u) + @closure (u, p) -> jac_fn_wrapped(u) else - @closure (J, u, p) -> _jac_fn(J, u) + @closure (J, u, p) -> jac_fn_wrapped(J, u) end - solver_kwargs = (; xtol = reltol, ftol = reltol, gtol = abstol, maxit = maxiters, + solver_kwargs = (; + xtol = reltol, ftol = reltol, gtol = abstol, maxit = maxiters, alg.factor, alg.factoraccept, alg.factorreject, alg.minscale, - alg.maxscale, alg.factorupdate, alg.minfactor, alg.maxfactor) + alg.maxscale, alg.factorupdate, alg.minfactor, alg.maxfactor + ) if prob.u0 isa SArray res, fx, info, iter, nfev, njev = FastLM.lmsolve( - f, jac_fn, prob.u0; solver_kwargs...) + f, jac_fn, prob.u0; solver_kwargs... + ) LM, solver = nothing, nothing else J = prob.f.jac_prototype === nothing ? similar(u, length(resid), length(u)) : zero(prob.f.jac_prototype) - solver = _fast_lm_solver(alg, u) + + solver = if alg.linsolve === :cholesky + FastLM.CholeskySolver(ArrayInterface.undefmatrix(u)) + elseif alg.linsolve === :qr + FastLM.QRSolver(eltype(u), length(u)) + else + throw(ArgumentError("Unknown FastLevenbergMarquardt Linear Solver: \ + $(Meta.quot(alg.linsolve))")) + end + LM = FastLM.LMWorkspace(u, resid, J) res, fx, info, iter, nfev, njev, LM, solver = FastLM.lmsolve!( - f, jac_fn, LM; solver, solver_kwargs...) + f, jac_fn, LM; solver, solver_kwargs... + ) end stats = SciMLBase.NLStats(nfev, njev, -1, -1, iter) retcode = info == -1 ? ReturnCode.MaxIters : ReturnCode.Success - return SciMLBase.build_solution(prob, alg, res, fx; retcode, - original = (res, fx, info, iter, nfev, njev, LM, solver), stats) + return SciMLBase.build_solution( + prob, alg, res, fx; retcode, + original = (res, fx, info, iter, nfev, njev, LM, solver), stats + ) end end diff --git a/ext/NonlinearSolveFixedPointAccelerationExt.jl b/ext/NonlinearSolveFixedPointAccelerationExt.jl index 2d36d7f56..6ff9f240c 100644 --- a/ext/NonlinearSolveFixedPointAccelerationExt.jl +++ b/ext/NonlinearSolveFixedPointAccelerationExt.jl @@ -1,41 +1,51 @@ module NonlinearSolveFixedPointAccelerationExt -using NonlinearSolveBase: NonlinearSolveBase, get_tolerance +using FixedPointAcceleration: FixedPointAcceleration, fixed_point + +using NonlinearSolveBase: NonlinearSolveBase using NonlinearSolve: NonlinearSolve, FixedPointAccelerationJL using SciMLBase: SciMLBase, NonlinearProblem, ReturnCode -using FixedPointAcceleration: FixedPointAcceleration, fixed_point -function SciMLBase.__solve(prob::NonlinearProblem, alg::FixedPointAccelerationJL, args...; +function SciMLBase.__solve( + prob::NonlinearProblem, alg::FixedPointAccelerationJL, args...; abstol = nothing, maxiters = 1000, alias_u0::Bool = false, - show_trace::Val{PrintReports} = Val(false), - termination_condition = nothing, kwargs...) where {PrintReports} - NonlinearSolve.__test_termination_condition( - termination_condition, :FixedPointAccelerationJL) + show_trace::Val = Val(false), termination_condition = nothing, kwargs... +) + NonlinearSolveBase.assert_extension_supported_termination_condition( + termination_condition, alg + ) + + f, u0, resid = NonlinearSolveBase.construct_extension_function_wrapper( + prob; alias_u0, make_fixed_point = Val(true), force_oop = Val(true) + ) - f, u0, resid = NonlinearSolve.__construct_extension_f( - prob; alias_u0, make_fixed_point = Val(true), force_oop = Val(true)) - tol = get_tolerance(abstol, eltype(u0)) + tol = NonlinearSolveBase.get_tolerance(abstol, eltype(u0)) - sol = fixed_point(f, u0; Algorithm = alg.algorithm, MaxIter = maxiters, MaxM = alg.m, + sol = fixed_point( + f, u0; Algorithm = alg.algorithm, MaxIter = maxiters, MaxM = alg.m, ConvergenceMetricThreshold = tol, ExtrapolationPeriod = alg.extrapolation_period, - Dampening = alg.dampening, PrintReports, ReplaceInvalids = alg.replace_invalids, - ConditionNumberThreshold = alg.condition_number_threshold, quiet_errors = true) + Dampening = alg.dampening, PrintReports = show_trace isa Val{true}, + ReplaceInvalids = alg.replace_invalids, + ConditionNumberThreshold = alg.condition_number_threshold, quiet_errors = true + ) if sol.FixedPoint_ === missing u0 = prob.u0 isa Number ? u0[1] : u0 - resid = NonlinearSolve.evaluate_f(prob, u0) + resid = NonlinearSolveBase.Utils.evaluate_f(prob, u0) res = u0 converged = false else res = prob.u0 isa Number ? first(sol.FixedPoint_) : reshape(sol.FixedPoint_, size(prob.u0)) - resid = NonlinearSolve.evaluate_f(prob, res) + resid = NonlinearSolveBase.Utils.evaluate_f(prob, res) converged = maximum(abs, resid) ≤ tol end - return SciMLBase.build_solution(prob, alg, res, resid; original = sol, + return SciMLBase.build_solution( + prob, alg, res, resid; original = sol, retcode = converged ? ReturnCode.Success : ReturnCode.Failure, - stats = SciMLBase.NLStats(sol.Iterations_, 0, 0, 0, sol.Iterations_)) + stats = SciMLBase.NLStats(sol.Iterations_, 0, 0, 0, sol.Iterations_) + ) end end diff --git a/ext/NonlinearSolveLeastSquaresOptimExt.jl b/ext/NonlinearSolveLeastSquaresOptimExt.jl index 20dac092d..834627606 100644 --- a/ext/NonlinearSolveLeastSquaresOptimExt.jl +++ b/ext/NonlinearSolveLeastSquaresOptimExt.jl @@ -1,79 +1,70 @@ module NonlinearSolveLeastSquaresOptimExt -using ConcreteStructs: @concrete using LeastSquaresOptim: LeastSquaresOptim -using NonlinearSolveBase: NonlinearSolveBase, get_tolerance -using NonlinearSolve: NonlinearSolve, LeastSquaresOptimJL, TraceMinimal -using SciMLBase: SciMLBase, NonlinearLeastSquaresProblem, NonlinearProblem, ReturnCode -const LSO = LeastSquaresOptim - -@inline function _lso_solver(::LeastSquaresOptimJL{alg, ls}) where {alg, ls} - linsolve = ls === :qr ? LSO.QR() : - (ls === :cholesky ? LSO.Cholesky() : (ls === :lsmr ? LSO.LSMR() : nothing)) - if alg === :lm - return LSO.LevenbergMarquardt(linsolve) - elseif alg === :dogleg - return LSO.Dogleg(linsolve) - else - throw(ArgumentError("Unknown LeastSquaresOptim Algorithm: $alg")) - end -end - -@concrete struct LeastSquaresOptimJLCache - prob - alg - allocated_prob - kwargs -end - -function Base.show(io::IO, cache::LeastSquaresOptimJLCache) - print(io, "LeastSquaresOptimJLCache()") -end +using NonlinearSolveBase: NonlinearSolveBase, TraceMinimal +using NonlinearSolve: NonlinearSolve, LeastSquaresOptimJL +using SciMLBase: SciMLBase, AbstractNonlinearProblem, ReturnCode -function SciMLBase.reinit!(cache::LeastSquaresOptimJLCache, args...; kwargs...) - error("Reinitialization not supported for LeastSquaresOptimJL.") -end +const LSO = LeastSquaresOptim -function SciMLBase.__init(prob::Union{NonlinearLeastSquaresProblem, NonlinearProblem}, - alg::LeastSquaresOptimJL, args...; alias_u0 = false, abstol = nothing, - show_trace::Val{ShT} = Val(false), trace_level = TraceMinimal(), - reltol = nothing, store_trace::Val{StT} = Val(false), maxiters = 1000, - termination_condition = nothing, kwargs...) where {ShT, StT} - NonlinearSolve.__test_termination_condition(termination_condition, :LeastSquaresOptim) +function SciMLBase.__solve( + prob::AbstractNonlinearProblem, alg::LeastSquaresOptimJL, args...; + alias_u0 = false, abstol = nothing, reltol = nothing, maxiters = 1000, + trace_level = TraceMinimal(), termination_condition = nothing, + show_trace::Val = Val(false), store_trace::Val = Val(false), kwargs... +) + NonlinearSolveBase.assert_extension_supported_termination_condition( + termination_condition, alg + ) - f!, u, resid = NonlinearSolve.__construct_extension_f(prob; alias_u0) - abstol = get_tolerance(abstol, eltype(u)) - reltol = get_tolerance(reltol, eltype(u)) + f!, u, resid = NonlinearSolveBase.construct_extension_function_wrapper(prob; alias_u0) + abstol = NonlinearSolveBase.get_tolerance(abstol, eltype(u)) + reltol = NonlinearSolveBase.get_tolerance(reltol, eltype(u)) if prob.f.jac === nothing && alg.autodiff isa Symbol - lsoprob = LSO.LeastSquaresProblem(; x = u, f!, y = resid, alg.autodiff, - J = prob.f.jac_prototype, output_length = length(resid)) + lsoprob = LSO.LeastSquaresProblem(; + x = u, f!, y = resid, alg.autodiff, J = prob.f.jac_prototype, + output_length = length(resid) + ) else - g! = NonlinearSolve.__construct_extension_jac(prob, alg, u, resid; alg.autodiff) + g! = NonlinearSolveBase.construct_extension_jac(prob, alg, u, resid; alg.autodiff) lsoprob = LSO.LeastSquaresProblem(; x = u, f!, y = resid, g!, J = prob.f.jac_prototype, - output_length = length(resid)) + output_length = length(resid) + ) end - allocated_prob = LSO.LeastSquaresProblemAllocated(lsoprob, _lso_solver(alg)) + linsolve = alg.linsolve === :qr ? LSO.QR() : + (alg.linsolve === :cholesky ? LSO.Cholesky() : + (alg.linsolve === :lsmr ? LSO.LSMR() : nothing)) - return LeastSquaresOptimJLCache(prob, - alg, - allocated_prob, - (; x_tol = reltol, f_tol = abstol, g_tol = abstol, iterations = maxiters, - show_trace = ShT, store_trace = StT, show_every = trace_level.print_frequency)) -end + lso_solver = if alg.alg === :lm + LSO.LevenbergMarquardt(linsolve) + elseif alg.alg === :dogleg + LSO.Dogleg(linsolve) + else + throw(ArgumentError("Unknown LeastSquaresOptim Algorithm: $(Meta.quot(alg.alg))")) + end + + allocated_prob = LSO.LeastSquaresProblemAllocated(lsoprob, lso_solver) + res = LSO.optimize!( + allocated_prob; + x_tol = reltol, f_tol = abstol, g_tol = abstol, iterations = maxiters, + show_trace = show_trace isa Val{true}, store_trace = store_trace isa Val{true}, + show_every = trace_level.print_frequency + ) -function SciMLBase.solve!(cache::LeastSquaresOptimJLCache) - res = LSO.optimize!(cache.allocated_prob; cache.kwargs...) - maxiters = cache.kwargs[:iterations] retcode = res.x_converged || res.f_converged || res.g_converged ? ReturnCode.Success : (res.iterations ≥ maxiters ? ReturnCode.MaxIters : ReturnCode.ConvergenceFailure) stats = SciMLBase.NLStats(res.f_calls, res.g_calls, -1, -1, res.iterations) + + f!(resid, res.minimizer) + return SciMLBase.build_solution( - cache.prob, cache.alg, res.minimizer, res.ssr / 2; retcode, original = res, stats) + prob, alg, res.minimizer, resid; retcode, original = res, stats + ) end end diff --git a/ext/NonlinearSolveMINPACKExt.jl b/ext/NonlinearSolveMINPACKExt.jl index 88adf5753..fc0f919a1 100644 --- a/ext/NonlinearSolveMINPACKExt.jl +++ b/ext/NonlinearSolveMINPACKExt.jl @@ -1,52 +1,67 @@ module NonlinearSolveMINPACKExt +using FastClosures: @closure using MINPACK: MINPACK -using NonlinearSolveBase: NonlinearSolveBase, get_tolerance + +using NonlinearSolveBase: NonlinearSolveBase using NonlinearSolve: NonlinearSolve, CMINPACK using SciMLBase: SciMLBase, NonlinearLeastSquaresProblem, NonlinearProblem, ReturnCode -using FastClosures: @closure function SciMLBase.__solve( - prob::Union{NonlinearLeastSquaresProblem, NonlinearProblem}, alg::CMINPACK, - args...; abstol = nothing, maxiters = 1000, alias_u0::Bool = false, - show_trace::Val{ShT} = Val(false), store_trace::Val{StT} = Val(false), - termination_condition = nothing, kwargs...) where {ShT, StT} - NonlinearSolve.__test_termination_condition(termination_condition, :CMINPACK) - - _f!, u0, resid = NonlinearSolve.__construct_extension_f(prob; alias_u0) - f! = @closure (du, u) -> (_f!(du, u); Cint(0)) + prob::Union{NonlinearLeastSquaresProblem, NonlinearProblem}, alg::CMINPACK, args...; + abstol = nothing, maxiters = 1000, alias_u0::Bool = false, + show_trace::Val = Val(false), store_trace::Val = Val(false), + termination_condition = nothing, kwargs... +) + NonlinearSolveBase.assert_extension_supported_termination_condition( + termination_condition, alg + ) + + f_wrapped!, u0, resid = NonlinearSolveBase.construct_extension_function_wrapper( + prob; alias_u0 + ) + resid_size = size(resid) + f! = @closure (du, u) -> (f_wrapped!(du, u); Cint(0)) m = length(resid) - method = ifelse(alg.method === :auto, - ifelse(prob isa NonlinearLeastSquaresProblem, :lm, :hybr), alg.method) + method = ifelse( + alg.method === :auto, + ifelse(prob isa NonlinearLeastSquaresProblem, :lm, :hybr), alg.method + ) - show_trace = ShT - tracing = StT - tol = get_tolerance(abstol, eltype(u0)) + show_trace = show_trace isa Val{true} + tracing = store_trace isa Val{true} + tol = NonlinearSolveBase.get_tolerance(abstol, eltype(u0)) if alg.autodiff === missing && prob.f.jac === nothing original = MINPACK.fsolve( - f!, u0, m; tol, show_trace, tracing, method, iterations = maxiters) + f!, u0, m; tol, show_trace, tracing, method, iterations = maxiters + ) else autodiff = alg.autodiff === missing ? nothing : alg.autodiff - _jac! = NonlinearSolve.__construct_extension_jac(prob, alg, u0, resid; autodiff) - jac! = @closure (J, u) -> (_jac!(J, u); Cint(0)) + jac_wrapped! = NonlinearSolveBase.construct_extension_jac( + prob, alg, u0, resid; autodiff + ) + jac! = @closure (J, u) -> (jac_wrapped!(J, u); Cint(0)) original = MINPACK.fsolve( - f!, jac!, u0, m; tol, show_trace, tracing, method, iterations = maxiters) + f!, jac!, u0, m; tol, show_trace, tracing, method, iterations = maxiters + ) end u = original.x - resid_ = original.f - objective = maximum(abs, resid_) + resid = original.f + objective = maximum(abs, resid) retcode = ifelse(objective ≤ tol, ReturnCode.Success, ReturnCode.Failure) # These are only meaningful if `store_trace = Val(true)` - stats = SciMLBase.NLStats(original.trace.f_calls, original.trace.g_calls, - original.trace.g_calls, original.trace.g_calls, -1) + stats = SciMLBase.NLStats( + original.trace.f_calls, original.trace.g_calls, + original.trace.g_calls, original.trace.g_calls, -1 + ) - u_ = prob.u0 isa Number ? original.x[1] : reshape(original.x, size(prob.u0)) - resid_ = prob.u0 isa Number ? resid_[1] : reshape(resid_, size(resid)) - return SciMLBase.build_solution(prob, alg, u_, resid_; retcode, original, stats) + u = prob.u0 isa Number ? original.x[1] : reshape(original.x, size(prob.u0)) + resid = prob.u0 isa Number ? resid[1] : reshape(resid, resid_size) + return SciMLBase.build_solution(prob, alg, u, resid; retcode, original, stats) end end diff --git a/ext/NonlinearSolveNLSolversExt.jl b/ext/NonlinearSolveNLSolversExt.jl index a9f41c87c..95aa0c8c3 100644 --- a/ext/NonlinearSolveNLSolversExt.jl +++ b/ext/NonlinearSolveNLSolversExt.jl @@ -1,74 +1,63 @@ module NonlinearSolveNLSolversExt using ADTypes: ADTypes, AutoFiniteDiff, AutoForwardDiff, AutoPolyesterForwardDiff +using DifferentiationInterface: DifferentiationInterface, Constant using FastClosures: @closure -using FiniteDiff: FiniteDiff -using ForwardDiff: ForwardDiff -using LinearAlgebra: norm + using NLSolvers: NLSolvers, NEqOptions, NEqProblem -using NonlinearSolveBase: NonlinearSolveBase, get_tolerance + +using NonlinearSolveBase: NonlinearSolveBase using NonlinearSolve: NonlinearSolve, NLSolversJL using SciMLBase: SciMLBase, NonlinearProblem, ReturnCode -function SciMLBase.__solve(prob::NonlinearProblem, alg::NLSolversJL, args...; - abstol = nothing, reltol = nothing, maxiters = 1000, - alias_u0::Bool = false, termination_condition = nothing, kwargs...) - NonlinearSolve.__test_termination_condition(termination_condition, :NLSolversJL) +const DI = DifferentiationInterface + +function SciMLBase.__solve( + prob::NonlinearProblem, alg::NLSolversJL, args...; + abstol = nothing, reltol = nothing, maxiters = 1000, alias_u0::Bool = false, + termination_condition = nothing, kwargs... +) + NonlinearSolveBase.assert_extension_supported_termination_condition( + termination_condition, alg + ) - abstol = get_tolerance(abstol, eltype(prob.u0)) - reltol = get_tolerance(reltol, eltype(prob.u0)) + abstol = NonlinearSolveBase.get_tolerance(abstol, eltype(prob.u0)) + reltol = NonlinearSolveBase.get_tolerance(reltol, eltype(prob.u0)) options = NEqOptions(; maxiter = maxiters, f_abstol = abstol, f_reltol = reltol) if prob.u0 isa Number - f_scalar = @closure x -> prob.f(x, prob.p) - - if alg.autodiff === nothing - if ForwardDiff.can_dual(typeof(prob.u0)) - autodiff_concrete = :forwarddiff - else - autodiff_concrete = :finitediff - end - else - if alg.autodiff isa AutoForwardDiff || alg.autodiff isa AutoPolyesterForwardDiff - autodiff_concrete = :forwarddiff - elseif alg.autodiff isa AutoFiniteDiff - autodiff_concrete = :finitediff - else - error("Only ForwardDiff or FiniteDiff autodiff is supported.") - end - end + f_scalar = Base.Fix2(prob.f, prob.p) + autodiff = NonlinearSolveBase.select_jacobian_autodiff(prob, alg.autodiff) + + prep = DifferentiationInterface.prepare_derivative( + prob.f, autodiff, prob.u0, Constant(prob.p) + ) - if autodiff_concrete === :forwarddiff - fj_scalar = @closure (Jx, x) -> begin - T = typeof(ForwardDiff.Tag(prob.f, eltype(x))) - x_dual = ForwardDiff.Dual{T}(x, one(x)) - y = prob.f(x_dual, prob.p) - return ForwardDiff.value(y), ForwardDiff.extract_derivative(T, y) - end - else - fj_scalar = @closure (Jx, x) -> begin - _f = Base.Fix2(prob.f, prob.p) - return _f(x), FiniteDiff.finite_difference_derivative(_f, x) - end + fj_scalar = @closure (Jx, x) -> begin + return DifferentiationInterface.value_and_derivative( + prob.f, prep, autodiff, x, Constant(prob.p) + ) end prob_obj = NLSolvers.ScalarObjective(; f = f_scalar, fg = fj_scalar) prob_nlsolver = NEqProblem(prob_obj; inplace = false) res = NLSolvers.solve(prob_nlsolver, prob.u0, alg.method, options) - retcode = ifelse(norm(res.info.best_residual, Inf) ≤ abstol, - ReturnCode.Success, ReturnCode.MaxIters) + retcode = ifelse( + maximum(abs, res.info.best_residual) ≤ abstol, + ReturnCode.Success, ReturnCode.MaxIters + ) stats = SciMLBase.NLStats(-1, -1, -1, -1, res.info.iter) return SciMLBase.build_solution( prob, alg, res.info.solution, res.info.best_residual; - retcode, original = res, stats) + retcode, original = res, stats + ) end - f!, u0, resid = NonlinearSolve.__construct_extension_f(prob; alias_u0) - - jac! = NonlinearSolve.__construct_extension_jac(prob, alg, u0, resid; alg.autodiff) + f!, u0, resid = NonlinearSolveBase.construct_extension_function_wrapper(prob; alias_u0) + jac! = NonlinearSolveBase.construct_extension_jac(prob, alg, u0, resid; alg.autodiff) FJ_vector! = @closure (Fx, Jx, x) -> begin f!(Fx, x) @@ -82,11 +71,15 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::NLSolversJL, args...; res = NLSolvers.solve(prob_nlsolver, u0, alg.method, options) retcode = ifelse( - norm(res.info.best_residual, Inf) ≤ abstol, ReturnCode.Success, ReturnCode.MaxIters) + maximum(abs, res.info.best_residual) ≤ abstol, + ReturnCode.Success, ReturnCode.MaxIters + ) stats = SciMLBase.NLStats(-1, -1, -1, -1, res.info.iter) - return SciMLBase.build_solution(prob, alg, res.info.solution, res.info.best_residual; - retcode, original = res, stats) + return SciMLBase.build_solution( + prob, alg, res.info.solution, res.info.best_residual; + retcode, original = res, stats + ) end end diff --git a/ext/NonlinearSolveNLsolveExt.jl b/ext/NonlinearSolveNLsolveExt.jl index 73d98c062..1b0beb993 100644 --- a/ext/NonlinearSolveNLsolveExt.jl +++ b/ext/NonlinearSolveNLsolveExt.jl @@ -1,44 +1,51 @@ module NonlinearSolveNLsolveExt using LineSearches: Static -using NonlinearSolveBase: NonlinearSolveBase, get_tolerance -using NonlinearSolve: NonlinearSolve, NLsolveJL, TraceMinimal using NLsolve: NLsolve, OnceDifferentiable, nlsolve + +using NonlinearSolveBase: NonlinearSolveBase, Utils, TraceMinimal +using NonlinearSolve: NonlinearSolve, NLsolveJL using SciMLBase: SciMLBase, NonlinearProblem, ReturnCode function SciMLBase.__solve( - prob::NonlinearProblem, alg::NLsolveJL, args...; abstol = nothing, - maxiters = 1000, alias_u0::Bool = false, termination_condition = nothing, - store_trace::Val{StT} = Val(false), show_trace::Val{ShT} = Val(false), - trace_level = TraceMinimal(), kwargs...) where {StT, ShT} - NonlinearSolve.__test_termination_condition(termination_condition, :NLsolveJL) + prob::NonlinearProblem, alg::NLsolveJL, args...; + abstol = nothing, maxiters = 1000, alias_u0::Bool = false, + termination_condition = nothing, trace_level = TraceMinimal(), + store_trace::Val = Val(false), show_trace::Val = Val(false), kwargs... +) + NonlinearSolveBase.assert_extension_supported_termination_condition( + termination_condition, alg + ) - f!, u0, resid = NonlinearSolve.__construct_extension_f(prob; alias_u0) + f!, u0, resid = NonlinearSolveBase.construct_extension_function_wrapper(prob; alias_u0) if prob.f.jac === nothing && alg.autodiff isa Symbol df = OnceDifferentiable(f!, u0, resid; alg.autodiff) else autodiff = alg.autodiff isa Symbol ? nothing : alg.autodiff - jac! = NonlinearSolve.__construct_extension_jac(prob, alg, u0, resid; autodiff) + jac! = NonlinearSolveBase.construct_extension_jac(prob, alg, u0, resid; autodiff) if prob.f.jac_prototype === nothing J = similar( - u0, promote_type(eltype(u0), eltype(resid)), length(u0), length(resid)) + u0, promote_type(eltype(u0), eltype(resid)), length(u0), length(resid) + ) else J = zero(prob.f.jac_prototype) end - df = OnceDifferentiable(f!, jac!, vec(u0), vec(resid), J) + df = OnceDifferentiable(f!, jac!, Utils.safe_vec(u0), Utils.safe_vec(resid), J) end - abstol = get_tolerance(abstol, eltype(u0)) - show_trace = ShT - store_trace = StT - extended_trace = !(trace_level isa TraceMinimal) + abstol = NonlinearSolveBase.get_tolerance(abstol, eltype(u0)) + show_trace = show_trace isa Val{true} + store_trace = store_trace isa Val{true} + extended_trace = !(trace_level.trace_mode isa Val{:minimal}) linesearch = alg.linesearch === missing ? Static() : alg.linesearch - original = nlsolve(df, vec(u0); ftol = abstol, iterations = maxiters, alg.method, - store_trace, extended_trace, linesearch, alg.linsolve, alg.factor, - alg.autoscale, alg.m, alg.beta, show_trace) + original = nlsolve( + df, vec(u0); + ftol = abstol, iterations = maxiters, alg.method, store_trace, extended_trace, + linesearch, alg.linsolve, alg.factor, alg.autoscale, alg.m, alg.beta, show_trace + ) f!(vec(resid), original.zero) u = prob.u0 isa Number ? original.zero[1] : reshape(original.zero, size(prob.u0)) @@ -46,8 +53,10 @@ function SciMLBase.__solve( retcode = original.x_converged || original.f_converged ? ReturnCode.Success : ReturnCode.Failure - stats = SciMLBase.NLStats(original.f_calls, original.g_calls, original.g_calls, - original.g_calls, original.iterations) + stats = SciMLBase.NLStats( + original.f_calls, original.g_calls, original.g_calls, + original.g_calls, original.iterations + ) return SciMLBase.build_solution(prob, alg, u, resid; retcode, original, stats) end diff --git a/ext/NonlinearSolvePETScExt.jl b/ext/NonlinearSolvePETScExt.jl index 9146c0b61..7b36cecce 100644 --- a/ext/NonlinearSolvePETScExt.jl +++ b/ext/NonlinearSolvePETScExt.jl @@ -1,23 +1,31 @@ module NonlinearSolvePETScExt using FastClosures: @closure + using MPI: MPI -using NonlinearSolveBase: NonlinearSolveBase, get_tolerance -using NonlinearSolve: NonlinearSolve, PETScSNES using PETSc: PETSc + +using NonlinearSolveBase: NonlinearSolveBase +using NonlinearSolve: NonlinearSolve, PETScSNES using SciMLBase: SciMLBase, NonlinearProblem, ReturnCode + using SparseArrays: AbstractSparseMatrix function SciMLBase.__solve( - prob::NonlinearProblem, alg::PETScSNES, args...; abstol = nothing, reltol = nothing, + prob::NonlinearProblem, alg::PETScSNES, args...; + abstol = nothing, reltol = nothing, maxiters = 1000, alias_u0::Bool = false, termination_condition = nothing, - show_trace::Val{ShT} = Val(false), kwargs...) where {ShT} + show_trace::Val = Val(false), kwargs... +) # XXX: https://petsc.org/release/manualpages/SNES/SNESSetConvergenceTest/ - termination_condition === nothing || - error("`PETScSNES` does not support termination conditions!") - - _f!, u0, resid = NonlinearSolve.__construct_extension_f(prob; alias_u0) - T = eltype(prob.u0) + NonlinearSolveBase.assert_extension_supported_termination_condition( + termination_condition, alg; abs_norm_supported = false + ) + + f_wrapped!, u0, resid = NonlinearSolveBase.construct_extension_function_wrapper( + prob; alias_u0 + ) + T = eltype(u0) @assert T ∈ PETSc.scalar_types if alg.petsclib === missing @@ -35,8 +43,8 @@ function SciMLBase.__solve( end PETSc.initialized(petsclib) || PETSc.initialize(petsclib) - abstol = get_tolerance(abstol, T) - reltol = get_tolerance(reltol, T) + abstol = NonlinearSolveBase.get_tolerance(abstol, T) + reltol = NonlinearSolveBase.get_tolerance(reltol, T) nf = Ref{Int}(0) @@ -44,77 +52,82 @@ function SciMLBase.__solve( nf[] += 1 fx = cfx isa Ptr{Nothing} ? PETSc.unsafe_localarray(T, cfx; read = false) : cfx x = cx isa Ptr{Nothing} ? PETSc.unsafe_localarray(T, cx; write = false) : cx - _f!(fx, x) + f_wrapped!(fx, x) Base.finalize(fx) Base.finalize(x) return end - snes = PETSc.SNES{T}(petsclib, + snes = PETSc.SNES{T}( + petsclib, alg.mpi_comm === missing ? MPI.COMM_SELF : alg.mpi_comm; - alg.snes_options..., snes_monitor = ShT, snes_rtol = reltol, - snes_atol = abstol, snes_max_it = maxiters) + alg.snes_options..., snes_monitor = show_trace isa Val{true}, snes_rtol = reltol, + snes_atol = abstol, snes_max_it = maxiters + ) PETSc.setfunction!(snes, f!, PETSc.VecSeq(zero(u0))) - if alg.autodiff === missing && prob.f.jac === nothing - _jac! = nothing - njac = Ref{Int}(-1) - else + njac = Ref{Int}(-1) + if alg.autodiff !== missing || prob.f.jac !== nothing autodiff = alg.autodiff === missing ? nothing : alg.autodiff if prob.u0 isa Number - _jac! = NonlinearSolve.__construct_extension_jac( - prob, alg, prob.u0, prob.u0; autodiff) + jac! = NonlinearSolveBase.construct_extension_jac( + prob, alg, prob.u0, prob.u0; autodiff + ) J_init = zeros(T, 1, 1) else - _jac!, J_init = NonlinearSolve.__construct_extension_jac( - prob, alg, u0, resid; autodiff, initial_jacobian = Val(true)) + jac!, J_init = NonlinearSolveBase.construct_extension_jac( + prob, alg, u0, resid; autodiff, initial_jacobian = Val(true) + ) end njac = Ref{Int}(0) if J_init isa AbstractSparseMatrix PJ = PETSc.MatSeqAIJ(J_init) - jac! = @closure (cx, J, _, user_ctx) -> begin + jac_fn! = @closure (cx, J, _, user_ctx) -> begin njac[] += 1 x = cx isa Ptr{Nothing} ? PETSc.unsafe_localarray(T, cx; write = false) : cx if J isa PETSc.AbstractMat - _jac!(user_ctx.jacobian, x) + jac!(user_ctx.jacobian, x) copyto!(J, user_ctx.jacobian) PETSc.assemble(J) else - _jac!(J, x) + jac!(J, x) end Base.finalize(x) return end - PETSc.setjacobian!(snes, jac!, PJ, PJ) + PETSc.setjacobian!(snes, jac_fn!, PJ, PJ) snes.user_ctx = (; jacobian = J_init) else PJ = PETSc.MatSeqDense(J_init) - jac! = @closure (cx, J, _, user_ctx) -> begin + jac_fn! = @closure (cx, J, _, user_ctx) -> begin njac[] += 1 x = cx isa Ptr{Nothing} ? PETSc.unsafe_localarray(T, cx; write = false) : cx - _jac!(J, x) + jac!(J, x) Base.finalize(x) J isa PETSc.AbstractMat && PETSc.assemble(J) return end - PETSc.setjacobian!(snes, jac!, PJ, PJ) + PETSc.setjacobian!(snes, jac_fn!, PJ, PJ) end end res = PETSc.solve!(u0, snes) - _f!(resid, res) - u_ = prob.u0 isa Number ? res[1] : res - resid_ = prob.u0 isa Number ? resid[1] : resid + f_wrapped!(resid, res) + u_res = prob.u0 isa Number ? res[1] : res + resid_res = prob.u0 isa Number ? resid[1] : resid objective = maximum(abs, resid) # XXX: Return Code from PETSc retcode = ifelse(objective ≤ abstol, ReturnCode.Success, ReturnCode.Failure) - return SciMLBase.build_solution(prob, alg, u_, resid_; retcode, original = snes, - stats = SciMLBase.NLStats(nf[], njac[], -1, -1, -1)) + return SciMLBase.build_solution( + prob, alg, u_res, resid_res; + retcode, original = snes, + stats = SciMLBase.NLStats(nf[], njac[], -1, -1, -1) + ) end end diff --git a/ext/NonlinearSolveSIAMFANLEquationsExt.jl b/ext/NonlinearSolveSIAMFANLEquationsExt.jl index 2468064fb..bfd2bd72d 100644 --- a/ext/NonlinearSolveSIAMFANLEquationsExt.jl +++ b/ext/NonlinearSolveSIAMFANLEquationsExt.jl @@ -1,13 +1,14 @@ module NonlinearSolveSIAMFANLEquationsExt using FastClosures: @closure -using NonlinearSolveBase: NonlinearSolveBase, get_tolerance -using NonlinearSolve: NonlinearSolve, SIAMFANLEquationsJL -using SciMLBase: SciMLBase, NonlinearProblem, ReturnCode using SIAMFANLEquations: SIAMFANLEquations, aasol, nsol, nsoli, nsolsc, ptcsol, ptcsoli, ptcsolsc, secant -@inline function __siam_fanl_equations_retcode_mapping(sol) +using NonlinearSolveBase: NonlinearSolveBase +using NonlinearSolve: NonlinearSolve, SIAMFANLEquationsJL +using SciMLBase: SciMLBase, NonlinearProblem, ReturnCode + +function siamfanlequations_retcode_mapping(sol) if sol.errcode == 0 return ReturnCode.Success elseif sol.errcode == 10 @@ -16,102 +17,129 @@ using SIAMFANLEquations: SIAMFANLEquations, aasol, nsol, nsoli, nsolsc, ptcsol, return ReturnCode.Failure elseif sol.errcode == -1 return ReturnCode.Default + else + error("Unknown SIAMFANLEquations return code: $(sol.errcode)") end end -@inline function __zeros_like(x, args...) +function zeros_like(x, args...) z = similar(x, args...) - fill!(z, zero(eltype(x))) + fill!(z, false) return z end # pseudo transient continuation has a fixed cost per iteration, iteration statistics are # not interesting here. -@inline function __siam_fanl_equations_stats_mapping(method, sol) +function siamfanlequations_stats_mapping(method, sol) ((method === :pseudotransient) || (method === :anderson)) && return nothing return SciMLBase.NLStats( - sum(sol.stats.ifun), sum(sol.stats.ijac), 0, 0, sum(sol.stats.iarm)) + sum(sol.stats.ifun), sum(sol.stats.ijac), 0, 0, sum(sol.stats.iarm) + ) end -function SciMLBase.__solve(prob::NonlinearProblem, alg::SIAMFANLEquationsJL, args...; - abstol = nothing, reltol = nothing, alias_u0::Bool = false, - maxiters = 1000, termination_condition = nothing, - show_trace::Val{ShT} = Val(false), kwargs...) where {ShT} - NonlinearSolve.__test_termination_condition(termination_condition, :SIAMFANLEquationsJL) +function SciMLBase.__solve( + prob::NonlinearProblem, alg::SIAMFANLEquationsJL, args...; + abstol = nothing, reltol = nothing, alias_u0::Bool = false, maxiters = 1000, + termination_condition = nothing, show_trace = Val(false), kwargs... +) + NonlinearSolveBase.assert_extension_supported_termination_condition( + termination_condition, alg + ) (; method, delta, linsolve, m, beta) = alg T = eltype(prob.u0) - atol = get_tolerance(abstol, T) - rtol = get_tolerance(reltol, T) + atol = NonlinearSolveBase.get_tolerance(abstol, T) + rtol = NonlinearSolveBase.get_tolerance(reltol, T) + + printerr = show_trace isa Val{true} if prob.u0 isa Number - f = @closure u -> prob.f(u, prob.p) + f = Base.Fix2(prob.f, prob.p) if method == :newton - sol = nsolsc(f, prob.u0; maxit = maxiters, atol, rtol, printerr = ShT) + sol = nsolsc(f, prob.u0; maxit = maxiters, atol, rtol, printerr) elseif method == :pseudotransient sol = ptcsolsc( - f, prob.u0; delta0 = delta, maxit = maxiters, atol, rtol, printerr = ShT) + f, prob.u0; delta0 = delta, maxit = maxiters, atol, rtol, printerr + ) elseif method == :secant - sol = secant(f, prob.u0; maxit = maxiters, atol, rtol, printerr = ShT) + sol = secant(f, prob.u0; maxit = maxiters, atol, rtol, printerr) elseif method == :anderson - f_aa, u, _ = NonlinearSolve.__construct_extension_f( - prob; alias_u0, make_fixed_point = Val(true)) - sol = aasol(f_aa, u, m, __zeros_like(u, 1, 2 * m + 4); - maxit = maxiters, atol, rtol, beta) + f_aa, u, _ = NonlinearSolveBase.construct_extension_function_wrapper( + prob; alias_u0, make_fixed_point = Val(true) + ) + sol = aasol( + f_aa, u, m, zeros_like(u, 1, 2 * m + 4); + maxit = maxiters, atol, rtol, beta + ) end else - f, u, resid = NonlinearSolve.__construct_extension_f( - prob; alias_u0, make_fixed_point = Val(method == :anderson)) + f, u, resid = NonlinearSolveBase.construct_extension_function_wrapper( + prob; alias_u0, make_fixed_point = Val(method == :anderson) + ) N = length(u) - FS = __zeros_like(u, N) + FS = zeros_like(u, N) # Jacobian Free Newton Krylov if linsolve !== nothing # Allocate ahead for Krylov basis - JVS = linsolve == :gmres ? __zeros_like(u, N, 3) : __zeros_like(u, N) + JVS = linsolve == :gmres ? zeros_like(u, N, 3) : zeros_like(u, N) linsolve_alg = String(linsolve) if method == :newton - sol = nsoli(f, u, FS, JVS; lsolver = linsolve_alg, - maxit = maxiters, atol, rtol, printerr = ShT) + sol = nsoli( + f, u, FS, JVS; lsolver = linsolve_alg, + maxit = maxiters, atol, rtol, printerr + ) elseif method == :pseudotransient - sol = ptcsoli(f, u, FS, JVS; lsolver = linsolve_alg, - maxit = maxiters, atol, rtol, printerr = ShT) + sol = ptcsoli( + f, u, FS, JVS; lsolver = linsolve_alg, + maxit = maxiters, atol, rtol, printerr + ) end else if prob.f.jac === nothing && alg.autodiff === missing - FPS = __zeros_like(u, N, N) + FPS = zeros_like(u, N, N) if method == :newton - sol = nsol(f, u, FS, FPS; sham = 1, atol, rtol, - maxit = maxiters, printerr = ShT) + sol = nsol( + f, u, FS, FPS; sham = 1, atol, rtol, maxit = maxiters, printerr + ) elseif method == :pseudotransient - sol = ptcsol(f, u, FS, FPS; atol, rtol, maxit = maxiters, - delta0 = delta, printerr = ShT) + sol = ptcsol( + f, u, FS, FPS; + atol, rtol, maxit = maxiters, delta0 = delta, printerr + ) elseif method == :anderson sol = aasol( - f, u, m, zeros(T, N, 2 * m + 4); atol, rtol, maxit = maxiters, beta) + f, u, m, zeros(T, N, 2 * m + 4); + atol, rtol, maxit = maxiters, beta + ) end else autodiff = alg.autodiff === missing ? nothing : alg.autodiff FPS = prob.f.jac_prototype !== nothing ? zero(prob.f.jac_prototype) : - __zeros_like(u, N, N) - jac = NonlinearSolve.__construct_extension_jac( - prob, alg, u, resid; autodiff) + zeros_like(u, N, N) + jac = NonlinearSolveBase.construct_extension_jac( + prob, alg, u, resid; autodiff + ) AJ! = @closure (J, u, x) -> jac(J, x) if method == :newton - sol = nsol(f, u, FS, FPS, AJ!; sham = 1, atol, - rtol, maxit = maxiters, printerr = ShT) + sol = nsol( + f, u, FS, FPS, AJ!; sham = 1, atol, + rtol, maxit = maxiters, printerr + ) elseif method == :pseudotransient - sol = ptcsol(f, u, FS, FPS, AJ!; atol, rtol, maxit = maxiters, - delta0 = delta, printerr = ShT) + sol = ptcsol( + f, u, FS, FPS, AJ!; atol, rtol, maxit = maxiters, + delta0 = delta, printerr + ) end end end end - retcode = __siam_fanl_equations_retcode_mapping(sol) - stats = __siam_fanl_equations_stats_mapping(method, sol) + retcode = siamfanlequations_retcode_mapping(sol) + stats = siamfanlequations_stats_mapping(method, sol) res = prob.u0 isa Number && method === :anderson ? sol.solution[1] : sol.solution - resid = NonlinearSolve.evaluate_f(prob, res) + resid = NonlinearSolveBase.Utils.evaluate_f(prob, res) return SciMLBase.build_solution(prob, alg, res, resid; retcode, stats, original = sol) end diff --git a/ext/NonlinearSolveSpeedMappingExt.jl b/ext/NonlinearSolveSpeedMappingExt.jl index ff9b4683b..c0f39607b 100644 --- a/ext/NonlinearSolveSpeedMappingExt.jl +++ b/ext/NonlinearSolveSpeedMappingExt.jl @@ -1,30 +1,41 @@ module NonlinearSolveSpeedMappingExt -using NonlinearSolveBase: NonlinearSolveBase, get_tolerance +using SpeedMapping: speedmapping + +using NonlinearSolveBase: NonlinearSolveBase using NonlinearSolve: NonlinearSolve, SpeedMappingJL using SciMLBase: SciMLBase, NonlinearProblem, ReturnCode -using SpeedMapping: speedmapping -function SciMLBase.__solve(prob::NonlinearProblem, alg::SpeedMappingJL, args...; +function SciMLBase.__solve( + prob::NonlinearProblem, alg::SpeedMappingJL, args...; abstol = nothing, maxiters = 1000, alias_u0::Bool = false, - maxtime = nothing, store_trace::Val{store_info} = Val(false), - termination_condition = nothing, kwargs...) where {store_info} - NonlinearSolve.__test_termination_condition(termination_condition, :SpeedMappingJL) + maxtime = nothing, store_trace::Val = Val(false), + termination_condition = nothing, kwargs... +) + NonlinearSolveBase.assert_extension_supported_termination_condition( + termination_condition, alg + ) - m!, u, resid = NonlinearSolve.__construct_extension_f( - prob; alias_u0, make_fixed_point = Val(true)) - tol = get_tolerance(abstol, eltype(u)) + m!, u, resid = NonlinearSolveBase.construct_extension_function_wrapper( + prob; alias_u0, make_fixed_point = Val(true) + ) + tol = NonlinearSolveBase.get_tolerance(abstol, eltype(u)) time_limit = ifelse(maxtime === nothing, 1000, maxtime) - sol = speedmapping(u; m!, tol, Lp = Inf, maps_limit = maxiters, alg.orders, - alg.check_obj, store_info, alg.σ_min, alg.stabilize, time_limit) + sol = speedmapping( + u; m!, tol, Lp = Inf, maps_limit = maxiters, alg.orders, + alg.check_obj, store_info = store_trace isa Val{true}, alg.σ_min, alg.stabilize, + time_limit + ) res = prob.u0 isa Number ? first(sol.minimizer) : sol.minimizer - resid = NonlinearSolve.evaluate_f(prob, res) + resid = NonlinearSolveBase.Utils.evaluate_f(prob, res) - return SciMLBase.build_solution(prob, alg, res, resid; original = sol, - retcode = sol.converged ? ReturnCode.Success : ReturnCode.Failure, - stats = SciMLBase.NLStats(sol.maps, 0, 0, 0, sol.maps)) + return SciMLBase.build_solution( + prob, alg, res, resid; + original = sol, stats = SciMLBase.NLStats(sol.maps, 0, 0, 0, sol.maps), + retcode = ifelse(sol.converged, ReturnCode.Success, ReturnCode.Failure) + ) end end diff --git a/ext/NonlinearSolveSundialsExt.jl b/ext/NonlinearSolveSundialsExt.jl index edea3a49b..303da2167 100644 --- a/ext/NonlinearSolveSundialsExt.jl +++ b/ext/NonlinearSolveSundialsExt.jl @@ -1,16 +1,28 @@ module NonlinearSolveSundialsExt +using Sundials: KINSOL + +using CommonSolve: CommonSolve using NonlinearSolveBase: NonlinearSolveBase, nonlinearsolve_forwarddiff_solve, nonlinearsolve_dual_solution -using NonlinearSolve: DualNonlinearProblem +using NonlinearSolve: NonlinearSolve, DualNonlinearProblem using SciMLBase: SciMLBase -using Sundials: KINSOL function SciMLBase.__solve(prob::DualNonlinearProblem, alg::KINSOL, args...; kwargs...) sol, partials = nonlinearsolve_forwarddiff_solve(prob, alg, args...; kwargs...) dual_soln = nonlinearsolve_dual_solution(sol.u, partials, prob.p) return SciMLBase.build_solution( - prob, alg, dual_soln, sol.resid; sol.retcode, sol.stats, sol.original) + prob, alg, dual_soln, sol.resid; sol.retcode, sol.stats, sol.original + ) +end + +function SciMLBase.__init(prob::DualNonlinearProblem, alg::KINSOL, args...; kwargs...) + p = NonlinearSolveBase.nodual_value(prob.p) + newprob = SciMLBase.remake(prob; u0 = NonlinearSolveBase.nodual_value(prob.u0), p) + cache = CommonSolve.init(newprob, alg, args...; kwargs...) + return NonlinearSolveForwardDiffCache( + cache, newprob, alg, prob.p, p, ForwardDiff.partials(prob.p) + ) end end diff --git a/lib/BracketingNonlinearSolve/Project.toml b/lib/BracketingNonlinearSolve/Project.toml index 6fc241d7d..16cade168 100644 --- a/lib/BracketingNonlinearSolve/Project.toml +++ b/lib/BracketingNonlinearSolve/Project.toml @@ -1,13 +1,14 @@ name = "BracketingNonlinearSolve" uuid = "70df07ce-3d50-431d-a3e7-ca6ddb60ac1e" authors = ["Avik Pal and contributors"] -version = "1.0.0" +version = "1.1.0" [deps] CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2" ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471" NonlinearSolveBase = "be0214bd-f91f-a760-ac4e-3421ce2b2da0" PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" +Reexport = "189a3867-3050-52da-a836-e630ba90ab69" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" [weakdeps] @@ -23,9 +24,10 @@ ConcreteStructs = "0.2.3" ExplicitImports = "1.10.1" ForwardDiff = "0.10.36" InteractiveUtils = "<0.0.1, 1" -NonlinearSolveBase = "1" +NonlinearSolveBase = "1.1" PrecompileTools = "1.2" -SciMLBase = "2.50" +Reexport = "1.2" +SciMLBase = "2.58" Test = "1.10" TestItemRunner = "1" julia = "1.10" diff --git a/lib/BracketingNonlinearSolve/ext/BracketingNonlinearSolveForwardDiffExt.jl b/lib/BracketingNonlinearSolve/ext/BracketingNonlinearSolveForwardDiffExt.jl index b41a88451..09616b5a2 100644 --- a/lib/BracketingNonlinearSolve/ext/BracketingNonlinearSolveForwardDiffExt.jl +++ b/lib/BracketingNonlinearSolve/ext/BracketingNonlinearSolveForwardDiffExt.jl @@ -7,19 +7,23 @@ using SciMLBase: SciMLBase, IntervalNonlinearProblem using BracketingNonlinearSolve: Bisection, Brent, Alefeld, Falsi, ITP, Ridder +const DualIntervalNonlinearProblem{T, V, P} = IntervalNonlinearProblem{ + uType, iip, <:Union{<:Dual{T, V, P}, <:AbstractArray{<:Dual{T, V, P}}} +} where {uType, iip} + for algT in (Bisection, Brent, Alefeld, Falsi, ITP, Ridder) @eval function CommonSolve.solve( - prob::IntervalNonlinearProblem{ - uType, iip, <:Union{<:Dual{T, V, P}, <:AbstractArray{<:Dual{T, V, P}}}}, - alg::$(algT), - args...; - kwargs...) where {uType, iip, T, V, P} + prob::DualIntervalNonlinearProblem{T, V, P}, alg::$(algT), args...; + kwargs... + ) where {T, V, P} sol, partials = nonlinearsolve_forwarddiff_solve(prob, alg, args...; kwargs...) dual_soln = nonlinearsolve_dual_solution(sol.u, partials, prob.p) return SciMLBase.build_solution( - prob, alg, dual_soln, sol.resid; sol.retcode, sol.stats, - sol.original, left = Dual{T, V, P}(sol.left, partials), - right = Dual{T, V, P}(sol.right, partials)) + prob, alg, dual_soln, sol.resid; + sol.retcode, sol.stats, sol.original, + left = Dual{T, V, P}(sol.left, partials), + right = Dual{T, V, P}(sol.right, partials) + ) end end diff --git a/lib/BracketingNonlinearSolve/src/BracketingNonlinearSolve.jl b/lib/BracketingNonlinearSolve/src/BracketingNonlinearSolve.jl index 62cf7a7b7..ae95eecba 100644 --- a/lib/BracketingNonlinearSolve/src/BracketingNonlinearSolve.jl +++ b/lib/BracketingNonlinearSolve/src/BracketingNonlinearSolve.jl @@ -1,14 +1,14 @@ module BracketingNonlinearSolve using ConcreteStructs: @concrete +using PrecompileTools: @compile_workload, @setup_workload +using Reexport: @reexport using CommonSolve: CommonSolve, solve -using NonlinearSolveBase: NonlinearSolveBase -using SciMLBase: SciMLBase, AbstractNonlinearAlgorithm, IntervalNonlinearProblem, ReturnCode - -using PrecompileTools: @compile_workload, @setup_workload +using NonlinearSolveBase: NonlinearSolveBase, AbstractNonlinearSolveAlgorithm +using SciMLBase: SciMLBase, IntervalNonlinearProblem, ReturnCode -abstract type AbstractBracketingAlgorithm <: AbstractNonlinearAlgorithm end +abstract type AbstractBracketingAlgorithm <: AbstractNonlinearSolveAlgorithm end include("common.jl") @@ -30,19 +30,19 @@ end @setup_workload begin for T in (Float32, Float64) prob_brack = IntervalNonlinearProblem{false}( - (u, p) -> u^2 - p, T.((0.0, 2.0)), T(2)) + (u, p) -> u^2 - p, T.((0.0, 2.0)), T(2) + ) algs = (Alefeld(), Bisection(), Brent(), Falsi(), ITP(), Ridder()) @compile_workload begin - for alg in algs - CommonSolve.solve(prob_brack, alg; abstol = 1e-6) + @sync for alg in algs + Threads.@spawn CommonSolve.solve(prob_brack, alg; abstol = 1e-6) end end end end -export IntervalNonlinearProblem -export solve +@reexport using SciMLBase, NonlinearSolveBase export Alefeld, Bisection, Brent, Falsi, ITP, Ridder diff --git a/lib/BracketingNonlinearSolve/src/alefeld.jl b/lib/BracketingNonlinearSolve/src/alefeld.jl index a669900dd..6880f8c95 100644 --- a/lib/BracketingNonlinearSolve/src/alefeld.jl +++ b/lib/BracketingNonlinearSolve/src/alefeld.jl @@ -8,8 +8,10 @@ algorithm 4.1 because, in certain sense, the second algorithm(4.2) is an optimal """ struct Alefeld <: AbstractBracketingAlgorithm end -function CommonSolve.solve(prob::IntervalNonlinearProblem, alg::Alefeld, args...; - maxiters = 1000, abstol = nothing, kwargs...) +function CommonSolve.solve( + prob::IntervalNonlinearProblem, alg::Alefeld, args...; + maxiters = 1000, abstol = nothing, kwargs... +) f = Base.Fix2(prob.f, prob.p) a, b = prob.tspan c = a - (b - a) / (f(b) - f(a)) * f(a) @@ -17,12 +19,14 @@ function CommonSolve.solve(prob::IntervalNonlinearProblem, alg::Alefeld, args... fc = f(c) if a == c || b == c return SciMLBase.build_solution( - prob, alg, c, fc; retcode = ReturnCode.FloatingPointLimit, left = a, right = b) + prob, alg, c, fc; retcode = ReturnCode.FloatingPointLimit, left = a, right = b + ) end if iszero(fc) return SciMLBase.build_solution( - prob, alg, c, fc; retcode = ReturnCode.Success, left = a, right = b) + prob, alg, c, fc; retcode = ReturnCode.Success, left = a, right = b + ) end a, b, d = Impl.bracket(f, a, b, c) @@ -68,13 +72,15 @@ function CommonSolve.solve(prob::IntervalNonlinearProblem, alg::Alefeld, args... if ā == c || b̄ == c return SciMLBase.build_solution( - prob, alg, c, fc; retcode = ReturnCode.FloatingPointLimit, - left = ā, right = b̄) + prob, alg, c, fc; + retcode = ReturnCode.FloatingPointLimit, left = ā, right = b̄ + ) end if iszero(fc) return SciMLBase.build_solution( - prob, alg, c, fc; retcode = ReturnCode.Success, left = ā, right = b̄) + prob, alg, c, fc; retcode = ReturnCode.Success, left = ā, right = b̄ + ) end ā, b̄, d̄ = Impl.bracket(f, ā, b̄, c) @@ -89,13 +95,15 @@ function CommonSolve.solve(prob::IntervalNonlinearProblem, alg::Alefeld, args... if ā == c || b̄ == c return SciMLBase.build_solution( - prob, alg, c, fc; retcode = ReturnCode.FloatingPointLimit, - left = ā, right = b̄) + prob, alg, c, fc; + retcode = ReturnCode.FloatingPointLimit, left = ā, right = b̄ + ) end if iszero(fc) return SciMLBase.build_solution( - prob, alg, c, fc; retcode = ReturnCode.Success, left = ā, right = b̄) + prob, alg, c, fc; retcode = ReturnCode.Success, left = ā, right = b̄ + ) end ā, b̄, d = Impl.bracket(f, ā, b̄, c) @@ -110,12 +118,14 @@ function CommonSolve.solve(prob::IntervalNonlinearProblem, alg::Alefeld, args... if ā == c || b̄ == c return SciMLBase.build_solution( - prob, alg, c, fc; retcode = ReturnCode.FloatingPointLimit, - left = ā, right = b̄) + prob, alg, c, fc; + retcode = ReturnCode.FloatingPointLimit, left = ā, right = b̄ + ) end if iszero(fc) return SciMLBase.build_solution( - prob, alg, c, fc; retcode = ReturnCode.Success, left = ā, right = b̄) + prob, alg, c, fc; retcode = ReturnCode.Success, left = ā, right = b̄ + ) end a, b, d = Impl.bracket(f, ā, b̄, c) end @@ -131,5 +141,6 @@ function CommonSolve.solve(prob::IntervalNonlinearProblem, alg::Alefeld, args... # Return solution when run out of max iteration return SciMLBase.build_solution( - prob, alg, c, fc; retcode = ReturnCode.MaxIters, left = a, right = b) + prob, alg, c, fc; retcode = ReturnCode.MaxIters, left = a, right = b + ) end diff --git a/lib/BracketingNonlinearSolve/src/bisection.jl b/lib/BracketingNonlinearSolve/src/bisection.jl index e51416145..91c17a775 100644 --- a/lib/BracketingNonlinearSolve/src/bisection.jl +++ b/lib/BracketingNonlinearSolve/src/bisection.jl @@ -19,8 +19,10 @@ A common bisection method. exact_right::Bool = false end -function CommonSolve.solve(prob::IntervalNonlinearProblem, alg::Bisection, - args...; maxiters = 1000, abstol = nothing, verbose::Bool = true, kwargs...) +function CommonSolve.solve( + prob::IntervalNonlinearProblem, alg::Bisection, args...; + maxiters = 1000, abstol = nothing, verbose::Bool = true, kwargs... +) @assert !SciMLBase.isinplace(prob) "`Bisection` only supports out-of-place problems." f = Base.Fix2(prob.f, prob.p) @@ -32,12 +34,14 @@ function CommonSolve.solve(prob::IntervalNonlinearProblem, alg::Bisection, if iszero(fl) return SciMLBase.build_solution( - prob, alg, left, fl; retcode = ReturnCode.ExactSolutionLeft, left, right) + prob, alg, left, fl; retcode = ReturnCode.ExactSolutionLeft, left, right + ) end if iszero(fr) return SciMLBase.build_solution( - prob, alg, right, fr; retcode = ReturnCode.ExactSolutionRight, left, right) + prob, alg, right, fr; retcode = ReturnCode.ExactSolutionRight, left, right + ) end if sign(fl) == sign(fr) @@ -45,7 +49,8 @@ function CommonSolve.solve(prob::IntervalNonlinearProblem, alg::Bisection, @warn "The interval is not an enclosing interval, opposite signs at the \ boundaries are required." return SciMLBase.build_solution( - prob, alg, left, fl; retcode = ReturnCode.InitialFailure, left, right) + prob, alg, left, fl; retcode = ReturnCode.InitialFailure, left, right + ) end i = 1 @@ -54,13 +59,15 @@ function CommonSolve.solve(prob::IntervalNonlinearProblem, alg::Bisection, if mid == left || mid == right return SciMLBase.build_solution( - prob, alg, left, fl; retcode = ReturnCode.FloatingPointLimit, left, right) + prob, alg, left, fl; retcode = ReturnCode.FloatingPointLimit, left, right + ) end fm = f(mid) if abs((right - left) / 2) < abstol return SciMLBase.build_solution( - prob, alg, mid, fm; retcode = ReturnCode.Success, left, right) + prob, alg, mid, fm; retcode = ReturnCode.Success, left, right + ) end if iszero(fm) @@ -80,10 +87,12 @@ function CommonSolve.solve(prob::IntervalNonlinearProblem, alg::Bisection, end sol, i, left, right, fl, fr = Impl.bisection( - left, right, fl, fr, f, abstol, maxiters - i, prob, alg) + left, right, fl, fr, f, abstol, maxiters - i, prob, alg + ) sol !== nothing && return sol return SciMLBase.build_solution( - prob, alg, left, fl; retcode = ReturnCode.MaxIters, left, right) + prob, alg, left, fl; retcode = ReturnCode.MaxIters, left, right + ) end diff --git a/lib/BracketingNonlinearSolve/src/brent.jl b/lib/BracketingNonlinearSolve/src/brent.jl index fb3740e98..7baebc90c 100644 --- a/lib/BracketingNonlinearSolve/src/brent.jl +++ b/lib/BracketingNonlinearSolve/src/brent.jl @@ -5,8 +5,10 @@ Left non-allocating Brent method. """ struct Brent <: AbstractBracketingAlgorithm end -function CommonSolve.solve(prob::IntervalNonlinearProblem, alg::Brent, args...; - maxiters = 1000, abstol = nothing, verbose::Bool = true, kwargs...) +function CommonSolve.solve( + prob::IntervalNonlinearProblem, alg::Brent, args...; + maxiters = 1000, abstol = nothing, verbose::Bool = true, kwargs... +) @assert !SciMLBase.isinplace(prob) "`Brent` only supports out-of-place problems." f = Base.Fix2(prob.f, prob.p) @@ -15,16 +17,19 @@ function CommonSolve.solve(prob::IntervalNonlinearProblem, alg::Brent, args...; ϵ = eps(convert(typeof(fl), 1)) abstol = NonlinearSolveBase.get_tolerance( - left, abstol, promote_type(eltype(left), eltype(right))) + left, abstol, promote_type(eltype(left), eltype(right)) + ) if iszero(fl) return SciMLBase.build_solution( - prob, alg, left, fl; retcode = ReturnCode.ExactSolutionLeft, left, right) + prob, alg, left, fl; retcode = ReturnCode.ExactSolutionLeft, left, right + ) end if iszero(fr) return SciMLBase.build_solution( - prob, alg, right, fr; retcode = ReturnCode.ExactSolutionRight, left, right) + prob, alg, right, fr; retcode = ReturnCode.ExactSolutionRight, left, right + ) end if sign(fl) == sign(fr) @@ -32,7 +37,8 @@ function CommonSolve.solve(prob::IntervalNonlinearProblem, alg::Brent, args...; @warn "The interval is not an enclosing interval, opposite signs at the \ boundaries are required." return SciMLBase.build_solution( - prob, alg, left, fl; retcode = ReturnCode.InitialFailure, left, right) + prob, alg, left, fl; retcode = ReturnCode.InitialFailure, left, right + ) end if abs(fl) < abs(fr) @@ -67,8 +73,10 @@ function CommonSolve.solve(prob::IntervalNonlinearProblem, alg::Brent, args...; # Bisection method s = (left + right) / 2 if s == left || s == right - return SciMLBase.build_solution(prob, alg, left, fl; - retcode = ReturnCode.FloatingPointLimit, left, right) + return SciMLBase.build_solution( + prob, alg, left, fl; + retcode = ReturnCode.FloatingPointLimit, left, right + ) end cond = true else @@ -78,7 +86,8 @@ function CommonSolve.solve(prob::IntervalNonlinearProblem, alg::Brent, args...; fs = f(s) if abs((right - left) / 2) < abstol return SciMLBase.build_solution( - prob, alg, s, fs; retcode = ReturnCode.Success, left, right) + prob, alg, s, fs; retcode = ReturnCode.Success, left, right + ) end if iszero(fs) @@ -110,10 +119,12 @@ function CommonSolve.solve(prob::IntervalNonlinearProblem, alg::Brent, args...; end sol, i, left, right, fl, fr = Impl.bisection( - left, right, fl, fr, f, abstol, maxiters - i, prob, alg) + left, right, fl, fr, f, abstol, maxiters - i, prob, alg + ) sol !== nothing && return sol return SciMLBase.build_solution( - prob, alg, left, fl; retcode = ReturnCode.MaxIters, left, right) + prob, alg, left, fl; retcode = ReturnCode.MaxIters, left, right + ) end diff --git a/lib/BracketingNonlinearSolve/src/common.jl b/lib/BracketingNonlinearSolve/src/common.jl index 65239ea63..0127869d8 100644 --- a/lib/BracketingNonlinearSolve/src/common.jl +++ b/lib/BracketingNonlinearSolve/src/common.jl @@ -10,14 +10,16 @@ function bisection(left, right, fl, fr, f::F, abstol, maxiters, prob, alg) where if mid == left || mid == right sol = SciMLBase.build_solution( - prob, alg, left, fl; left, right, retcode = ReturnCode.FloatingPointLimit) + prob, alg, left, fl; left, right, retcode = ReturnCode.FloatingPointLimit + ) break end fm = f(mid) if abs((right - left) / 2) < abstol sol = SciMLBase.build_solution( - prob, alg, mid, fm; left, right, retcode = ReturnCode.Success) + prob, alg, mid, fm; left, right, retcode = ReturnCode.Success + ) break end diff --git a/lib/BracketingNonlinearSolve/src/falsi.jl b/lib/BracketingNonlinearSolve/src/falsi.jl index f56155ef7..3074a5eb4 100644 --- a/lib/BracketingNonlinearSolve/src/falsi.jl +++ b/lib/BracketingNonlinearSolve/src/falsi.jl @@ -5,8 +5,10 @@ A non-allocating regula falsi method. """ struct Falsi <: AbstractBracketingAlgorithm end -function CommonSolve.solve(prob::IntervalNonlinearProblem, alg::Falsi, args...; - maxiters = 1000, abstol = nothing, verbose::Bool = true, kwargs...) +function CommonSolve.solve( + prob::IntervalNonlinearProblem, alg::Falsi, args...; + maxiters = 1000, abstol = nothing, verbose::Bool = true, kwargs... +) @assert !SciMLBase.isinplace(prob) "`False` only supports out-of-place problems." f = Base.Fix2(prob.f, prob.p) @@ -19,12 +21,14 @@ function CommonSolve.solve(prob::IntervalNonlinearProblem, alg::Falsi, args...; if iszero(fl) return SciMLBase.build_solution( - prob, alg, left, fl; retcode = ReturnCode.ExactSolutionLeft, left, right) + prob, alg, left, fl; retcode = ReturnCode.ExactSolutionLeft, left, right + ) end if iszero(fr) return SciMLBase.build_solution( - prob, alg, right, fr; retcode = ReturnCode.ExactSolutionRight, left, right) + prob, alg, right, fr; retcode = ReturnCode.ExactSolutionRight, left, right + ) end if sign(fl) == sign(fr) @@ -32,14 +36,16 @@ function CommonSolve.solve(prob::IntervalNonlinearProblem, alg::Falsi, args...; @warn "The interval is not an enclosing interval, opposite signs at the \ boundaries are required." return SciMLBase.build_solution( - prob, alg, left, fl; retcode = ReturnCode.InitialFailure, left, right) + prob, alg, left, fl; retcode = ReturnCode.InitialFailure, left, right + ) end i = 1 while i ≤ maxiters if Impl.nextfloat_tdir(left, l, r) == right return SciMLBase.build_solution( - prob, alg, left, fl; left, right, retcode = ReturnCode.FloatingPointLimit) + prob, alg, left, fl; left, right, retcode = ReturnCode.FloatingPointLimit + ) end mid = (fr * left - fl * right) / (fr - fl) @@ -52,7 +58,8 @@ function CommonSolve.solve(prob::IntervalNonlinearProblem, alg::Falsi, args...; fm = f(mid) if abs((right - left) / 2) < abstol return SciMLBase.build_solution( - prob, alg, mid, fm; left, right, retcode = ReturnCode.Success) + prob, alg, mid, fm; left, right, retcode = ReturnCode.Success + ) end if abs(fm) < abstol @@ -70,10 +77,12 @@ function CommonSolve.solve(prob::IntervalNonlinearProblem, alg::Falsi, args...; end sol, i, left, right, fl, fr = Impl.bisection( - left, right, fl, fr, f, abstol, maxiters - i, prob, alg) + left, right, fl, fr, f, abstol, maxiters - i, prob, alg + ) sol !== nothing && return sol return SciMLBase.build_solution( - prob, alg, left, fl; retcode = ReturnCode.MaxIters, left, right) + prob, alg, left, fl; retcode = ReturnCode.MaxIters, left, right + ) end diff --git a/lib/BracketingNonlinearSolve/src/itp.jl b/lib/BracketingNonlinearSolve/src/itp.jl index 821047a5a..50443e2e9 100644 --- a/lib/BracketingNonlinearSolve/src/itp.jl +++ b/lib/BracketingNonlinearSolve/src/itp.jl @@ -56,8 +56,10 @@ function ITP(; scaled_k1::Real = 0.2, k2::Real = 2, n0::Int = 10) return ITP(scaled_k1, k2, n0) end -function CommonSolve.solve(prob::IntervalNonlinearProblem, alg::ITP, args...; - maxiters = 1000, abstol = nothing, verbose::Bool = true, kwargs...) +function CommonSolve.solve( + prob::IntervalNonlinearProblem, alg::ITP, args...; + maxiters = 1000, abstol = nothing, verbose::Bool = true, kwargs... +) @assert !SciMLBase.isinplace(prob) "`ITP` only supports out-of-place problems." f = Base.Fix2(prob.f, prob.p) @@ -65,16 +67,19 @@ function CommonSolve.solve(prob::IntervalNonlinearProblem, alg::ITP, args...; fl, fr = f(left), f(right) abstol = NonlinearSolveBase.get_tolerance( - left, abstol, promote_type(eltype(left), eltype(right))) + left, abstol, promote_type(eltype(left), eltype(right)) + ) if iszero(fl) return SciMLBase.build_solution( - prob, alg, left, fl; retcode = ReturnCode.ExactSolutionLeft, left, right) + prob, alg, left, fl; retcode = ReturnCode.ExactSolutionLeft, left, right + ) end if iszero(fr) return SciMLBase.build_solution( - prob, alg, right, fr; retcode = ReturnCode.ExactSolutionRight, left, right) + prob, alg, right, fr; retcode = ReturnCode.ExactSolutionRight, left, right + ) end if sign(fl) == sign(fr) @@ -82,7 +87,8 @@ function CommonSolve.solve(prob::IntervalNonlinearProblem, alg::ITP, args...; @warn "The interval is not an enclosing interval, opposite signs at the \ boundaries are required." return SciMLBase.build_solution( - prob, alg, left, fl; retcode = ReturnCode.InitialFailure, left, right) + prob, alg, left, fl; retcode = ReturnCode.InitialFailure, left, right + ) end ϵ = abstol @@ -115,7 +121,8 @@ function CommonSolve.solve(prob::IntervalNonlinearProblem, alg::ITP, args...; if abs((left - right) / 2) < ϵ return SciMLBase.build_solution( - prob, alg, xt, f(xt); retcode = ReturnCode.Success, left, right) + prob, alg, xt, f(xt); retcode = ReturnCode.Success, left, right + ) end # update @@ -131,7 +138,8 @@ function CommonSolve.solve(prob::IntervalNonlinearProblem, alg::ITP, args...; left, fl = xp, yp else return SciMLBase.build_solution( - prob, alg, xp, yps; retcode = ReturnCode.Success, left, right) + prob, alg, xp, yps; retcode = ReturnCode.Success, left, right + ) end i += 1 @@ -140,10 +148,12 @@ function CommonSolve.solve(prob::IntervalNonlinearProblem, alg::ITP, args...; if Impl.nextfloat_tdir(left, prob.tspan...) == right return SciMLBase.build_solution( - prob, alg, right, fr; retcode = ReturnCode.FloatingPointLimit, left, right) + prob, alg, right, fr; retcode = ReturnCode.FloatingPointLimit, left, right + ) end end return SciMLBase.build_solution( - prob, alg, left, fl; retcode = ReturnCode.MaxIters, left, right) + prob, alg, left, fl; retcode = ReturnCode.MaxIters, left, right + ) end diff --git a/lib/BracketingNonlinearSolve/src/ridder.jl b/lib/BracketingNonlinearSolve/src/ridder.jl index d988c9dc5..9192897c5 100644 --- a/lib/BracketingNonlinearSolve/src/ridder.jl +++ b/lib/BracketingNonlinearSolve/src/ridder.jl @@ -5,8 +5,10 @@ A non-allocating ridder method. """ struct Ridder <: AbstractBracketingAlgorithm end -function CommonSolve.solve(prob::IntervalNonlinearProblem, alg::Ridder, args...; - maxiters = 1000, abstol = nothing, verbose::Bool = true, kwargs...) +function CommonSolve.solve( + prob::IntervalNonlinearProblem, alg::Ridder, args...; + maxiters = 1000, abstol = nothing, verbose::Bool = true, kwargs... +) @assert !SciMLBase.isinplace(prob) "`Ridder` only supports out-of-place problems." f = Base.Fix2(prob.f, prob.p) @@ -14,16 +16,19 @@ function CommonSolve.solve(prob::IntervalNonlinearProblem, alg::Ridder, args...; fl, fr = f(left), f(right) abstol = NonlinearSolveBase.get_tolerance( - left, abstol, promote_type(eltype(left), eltype(right))) + left, abstol, promote_type(eltype(left), eltype(right)) + ) if iszero(fl) return SciMLBase.build_solution( - prob, alg, left, fl; retcode = ReturnCode.ExactSolutionLeft, left, right) + prob, alg, left, fl; retcode = ReturnCode.ExactSolutionLeft, left, right + ) end if iszero(fr) return SciMLBase.build_solution( - prob, alg, right, fr; retcode = ReturnCode.ExactSolutionRight, left, right) + prob, alg, right, fr; retcode = ReturnCode.ExactSolutionRight, left, right + ) end if sign(fl) == sign(fr) @@ -31,7 +36,8 @@ function CommonSolve.solve(prob::IntervalNonlinearProblem, alg::Ridder, args...; @warn "The interval is not an enclosing interval, opposite signs at the \ boundaries are required." return SciMLBase.build_solution( - prob, alg, left, fl; retcode = ReturnCode.InitialFailure, left, right) + prob, alg, left, fl; retcode = ReturnCode.InitialFailure, left, right + ) end xo = oftype(left, Inf) @@ -41,14 +47,16 @@ function CommonSolve.solve(prob::IntervalNonlinearProblem, alg::Ridder, args...; if mid == left || mid == right return SciMLBase.build_solution( - prob, alg, left, fl; retcode = ReturnCode.FloatingPointLimit, left, right) + prob, alg, left, fl; retcode = ReturnCode.FloatingPointLimit, left, right + ) end fm = f(mid) s = sqrt(fm^2 - fl * fr) if iszero(s) return SciMLBase.build_solution( - prob, alg, left, fl; retcode = ReturnCode.Failure, left, right) + prob, alg, left, fl; retcode = ReturnCode.Failure, left, right + ) end x = mid + (mid - left) * sign(fl - fm) * fm / s @@ -56,7 +64,8 @@ function CommonSolve.solve(prob::IntervalNonlinearProblem, alg::Ridder, args...; xo = x if abs((right - left) / 2) < abstol return SciMLBase.build_solution( - prob, alg, mid, fm; retcode = ReturnCode.Success, left, right) + prob, alg, mid, fm; retcode = ReturnCode.Success, left, right + ) end if iszero(fx) @@ -82,10 +91,12 @@ function CommonSolve.solve(prob::IntervalNonlinearProblem, alg::Ridder, args...; end sol, i, left, right, fl, fr = Impl.bisection( - left, right, fl, fr, f, abstol, maxiters - i, prob, alg) + left, right, fl, fr, f, abstol, maxiters - i, prob, alg + ) sol !== nothing && return sol return SciMLBase.build_solution( - prob, alg, left, fl; retcode = ReturnCode.MaxIters, left, right) + prob, alg, left, fl; retcode = ReturnCode.MaxIters, left, right + ) end diff --git a/lib/BracketingNonlinearSolve/test/qa_tests.jl b/lib/BracketingNonlinearSolve/test/qa_tests.jl index c01c493f4..dc78ce75e 100644 --- a/lib/BracketingNonlinearSolve/test/qa_tests.jl +++ b/lib/BracketingNonlinearSolve/test/qa_tests.jl @@ -1,7 +1,12 @@ @testitem "Aqua" tags=[:core] begin using Aqua, BracketingNonlinearSolve - Aqua.test_all(BracketingNonlinearSolve; piracies = false, ambiguities = false) + Aqua.test_all( + BracketingNonlinearSolve; + piracies = false, ambiguities = false, stale_deps = false, deps_compat = false + ) + Aqua.test_stale_deps(BracketingNonlinearSolve; ignore = [:SciMLJacobianOperators]) + Aqua.test_deps_compat(BracketingNonlinearSolve; ignore = [:SciMLJacobianOperators]) Aqua.test_piracies(BracketingNonlinearSolve; treat_as_own = [IntervalNonlinearProblem]) Aqua.test_ambiguities(BracketingNonlinearSolve; recursive = false) end diff --git a/lib/NonlinearSolveBase/Project.toml b/lib/NonlinearSolveBase/Project.toml index 70c0f97d4..4abb81cde 100644 --- a/lib/NonlinearSolveBase/Project.toml +++ b/lib/NonlinearSolveBase/Project.toml @@ -1,10 +1,11 @@ name = "NonlinearSolveBase" uuid = "be0214bd-f91f-a760-ac4e-3421ce2b2da0" authors = ["Avik Pal and contributors"] -version = "1.0.0" +version = "1.1.0" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" +Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9" CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2" Compat = "34da2185-b29b-5c13-b0c7-acf172513d20" @@ -15,28 +16,45 @@ FastClosures = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a" FunctionProperties = "f62d2435-5019-4c03-9749-2d4c77af0cbc" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a" +MaybeInplace = "bb5d69b7-63fc-4a16-80bd-7e42200c7bdb" +Preferences = "21216c6a-2e73-6563-6e65-726566657250" +Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" +SciMLJacobianOperators = "19f34311-ddf3-4b8b-af20-060888a46c0e" +SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961" StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" +SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5" +TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f" [weakdeps] +BandedMatrices = "aae01518-5342-5314-be14-df237901396f" DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +LineSearch = "87fe0de2-c867-4266-b59a-2f0a94fc965b" +LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" +SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35" [extensions] +NonlinearSolveBaseBandedMatricesExt = "BandedMatrices" NonlinearSolveBaseDiffEqBaseExt = "DiffEqBase" NonlinearSolveBaseForwardDiffExt = "ForwardDiff" +NonlinearSolveBaseLineSearchExt = "LineSearch" +NonlinearSolveBaseLinearSolveExt = "LinearSolve" NonlinearSolveBaseSparseArraysExt = "SparseArrays" +NonlinearSolveBaseSparseMatrixColoringsExt = "SparseMatrixColorings" [compat] ADTypes = "1.9" +Adapt = "4.1.0" Aqua = "0.8.7" ArrayInterface = "7.9" +BandedMatrices = "1.5" CommonSolve = "0.2.4" Compat = "4.15" ConcreteStructs = "0.2.3" -DiffEqBase = "6.149" +DiffEqBase = "6.158.3" DifferentiationInterface = "0.6.16" EnzymeCore = "0.8" ExplicitImports = "1.10.1" @@ -44,23 +62,35 @@ FastClosures = "0.3" ForwardDiff = "0.10.36" FunctionProperties = "0.1.2" InteractiveUtils = "<0.0.1, 1" +LineSearch = "0.1.4" LinearAlgebra = "1.10" +LinearSolve = "2.36.1" Markdown = "1.10" +MaybeInplace = "0.1.4" +Preferences = "1.4" +Printf = "1.10" RecursiveArrayTools = "3" -SciMLBase = "2.50" +SciMLBase = "2.58" +SciMLJacobianOperators = "0.1.1" +SciMLOperators = "0.3.10" SparseArrays = "1.10" +SparseMatrixColorings = "0.4.5" StaticArraysCore = "1.4" +SymbolicIndexingInterface = "0.3.31" Test = "1.10" +TimerOutputs = "0.5.23" julia = "1.10" [extras] Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" +BandedMatrices = "aae01518-5342-5314-be14-df237901396f" DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e" ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Aqua", "DiffEqBase", "ExplicitImports", "ForwardDiff", "InteractiveUtils", "SparseArrays", "Test"] +test = ["Aqua", "BandedMatrices", "DiffEqBase", "ExplicitImports", "ForwardDiff", "InteractiveUtils", "LinearAlgebra", "SparseArrays", "Test"] diff --git a/lib/NonlinearSolveBase/ext/NonlinearSolveBaseBandedMatricesExt.jl b/lib/NonlinearSolveBase/ext/NonlinearSolveBaseBandedMatricesExt.jl new file mode 100644 index 000000000..93f01f51f --- /dev/null +++ b/lib/NonlinearSolveBase/ext/NonlinearSolveBaseBandedMatricesExt.jl @@ -0,0 +1,17 @@ +module NonlinearSolveBaseBandedMatricesExt + +using BandedMatrices: BandedMatrix +using LinearAlgebra: Diagonal + +using NonlinearSolveBase: NonlinearSolveBase, Utils + +# This is used if we vcat a Banded Jacobian with a Diagonal Matrix in Levenberg +@inline function Utils.faster_vcat(B::BandedMatrix, D::Diagonal) + if Utils.is_extension_loaded(Val(:SparseArrays)) + @warn "Load `SparseArrays` for an optimized vcat for BandedMatrices." + return vcat(B, D) + end + return vcat(Utils.make_sparse(B), D) +end + +end diff --git a/lib/NonlinearSolveBase/ext/NonlinearSolveBaseForwardDiffExt.jl b/lib/NonlinearSolveBase/ext/NonlinearSolveBaseForwardDiffExt.jl index c4f1dc901..0b16391c4 100644 --- a/lib/NonlinearSolveBase/ext/NonlinearSolveBaseForwardDiffExt.jl +++ b/lib/NonlinearSolveBase/ext/NonlinearSolveBaseForwardDiffExt.jl @@ -25,7 +25,8 @@ Utils.value(x::AbstractArray{<:Dual}) = Utils.value.(x) function NonlinearSolveBase.nonlinearsolve_forwarddiff_solve( prob::Union{IntervalNonlinearProblem, NonlinearProblem, ImmutableNonlinearProblem}, - alg, args...; kwargs...) + alg, args...; kwargs... +) p = Utils.value(prob.p) if prob isa IntervalNonlinearProblem tspan = Utils.value.(prob.tspan) @@ -55,7 +56,8 @@ function NonlinearSolveBase.nonlinearsolve_forwarddiff_solve( end function NonlinearSolveBase.nonlinearsolve_forwarddiff_solve( - prob::NonlinearLeastSquaresProblem, alg, args...; kwargs...) + prob::NonlinearLeastSquaresProblem, alg, args...; kwargs... +) p = Utils.value(prob.p) newprob = remake(prob; p, u0 = Utils.value(prob.u0)) sol = solve(newprob, alg, args...; kwargs...) @@ -168,13 +170,17 @@ function NonlinearSolveBase.nonlinearsolve_∂f_∂u(prob, f::F, u, p) where {F} return ForwardDiff.jacobian(Base.Fix2(f, p), u) end -function NonlinearSolveBase.nonlinearsolve_dual_solution(u::Number, partials, - ::Union{<:AbstractArray{<:Dual{T, V, P}}, Dual{T, V, P}}) where {T, V, P} +function NonlinearSolveBase.nonlinearsolve_dual_solution( + u::Number, partials, + ::Union{<:AbstractArray{<:Dual{T, V, P}}, Dual{T, V, P}} +) where {T, V, P} return Dual{T, V, P}(u, partials) end -function NonlinearSolveBase.nonlinearsolve_dual_solution(u::AbstractArray, partials, - ::Union{<:AbstractArray{<:Dual{T, V, P}}, Dual{T, V, P}}) where {T, V, P} +function NonlinearSolveBase.nonlinearsolve_dual_solution( + u::AbstractArray, partials, + ::Union{<:AbstractArray{<:Dual{T, V, P}}, Dual{T, V, P}} +) where {T, V, P} return map(((uᵢ, pᵢ),) -> Dual{T, V, P}(uᵢ, pᵢ), zip(u, Utils.restructure(u, partials))) end diff --git a/lib/NonlinearSolveBase/ext/NonlinearSolveBaseLineSearchExt.jl b/lib/NonlinearSolveBase/ext/NonlinearSolveBaseLineSearchExt.jl new file mode 100644 index 000000000..3b705b0dc --- /dev/null +++ b/lib/NonlinearSolveBase/ext/NonlinearSolveBaseLineSearchExt.jl @@ -0,0 +1,18 @@ +module NonlinearSolveBaseLineSearchExt + +using LineSearch: LineSearch, AbstractLineSearchCache +using SciMLBase: SciMLBase + +using NonlinearSolveBase: NonlinearSolveBase, InternalAPI + +function NonlinearSolveBase.callback_into_cache!( + topcache, cache::AbstractLineSearchCache, args... +) + return LineSearch.callback_into_cache!(cache, NonlinearSolveBase.get_fu(topcache)) +end + +function InternalAPI.reinit!(cache::AbstractLineSearchCache; kwargs...) + return SciMLBase.reinit!(cache; kwargs...) +end + +end diff --git a/lib/NonlinearSolveBase/ext/NonlinearSolveBaseLinearSolveExt.jl b/lib/NonlinearSolveBase/ext/NonlinearSolveBaseLinearSolveExt.jl new file mode 100644 index 000000000..d1829804d --- /dev/null +++ b/lib/NonlinearSolveBase/ext/NonlinearSolveBaseLinearSolveExt.jl @@ -0,0 +1,138 @@ +module NonlinearSolveBaseLinearSolveExt + +using ArrayInterface: ArrayInterface + +using CommonSolve: CommonSolve, init, solve! +using LinearSolve: LinearSolve, QRFactorization, SciMLLinearSolveAlgorithm +using SciMLBase: ReturnCode, LinearProblem + +using LinearAlgebra: ColumnNorm + +using NonlinearSolveBase: NonlinearSolveBase, LinearSolveJLCache, LinearSolveResult, Utils + +function (cache::LinearSolveJLCache)(; + A = nothing, b = nothing, linu = nothing, du = nothing, p = nothing, + cachedata = nothing, reuse_A_if_factorization = false, verbose = true, kwargs... +) + cache.stats.nsolve += 1 + + update_A!(cache, A, reuse_A_if_factorization) + b !== nothing && setproperty!(cache.lincache, :b, b) + linu !== nothing && NonlinearSolveBase.set_lincache_u!(cache, linu) + + Plprev = cache.lincache.Pl + Prprev = cache.lincache.Pr + + if cache.precs === nothing + Pl, Pr = nothing, nothing + else + Pl, Pr = cache.precs( + cache.lincache.A, du, linu, p, nothing, + A !== nothing, Plprev, Prprev, cachedata + ) + end + + if Pl !== nothing || Pr !== nothing + Pl, Pr = NonlinearSolveBase.wrap_preconditioners(Pl, Pr, linu) + cache.lincache.Pl = Pl + cache.lincache.Pr = Pr + end + + linres = solve!(cache.lincache) + cache.lincache = linres.cache + # Unfortunately LinearSolve.jl doesn't have the most uniform ReturnCode handling + if linres.retcode === ReturnCode.Failure + structured_mat = ArrayInterface.isstructured(cache.lincache.A) + is_gpuarray = ArrayInterface.device(cache.lincache.A) isa ArrayInterface.GPU + + if !(cache.linsolve isa QRFactorization{ColumnNorm}) && !is_gpuarray && + !structured_mat + if verbose + @warn "Potential Rank Deficient Matrix Detected. Attempting to solve using \ + Pivoted QR Factorization." + end + @assert (A !== nothing)&&(b !== nothing) "This case is not yet supported. \ + Please open an issue at \ + https://github.com/SciML/NonlinearSolve.jl" + if cache.additional_lincache === nothing # First time + linprob = LinearProblem(A, b; u0 = linres.u) + cache.additional_lincache = init( + linprob, QRFactorization(ColumnNorm()); alias_u0 = false, + alias_A = false, alias_b = false, cache.lincache.Pl, cache.lincache.Pr) + else + cache.additional_lincache.A = A + cache.additional_lincache.b = b + cache.additional_lincache.Pl = cache.lincache.Pl + cache.additional_lincache.Pr = cache.lincache.Pr + end + linres = solve!(cache.additional_lincache) + cache.additional_lincache = linres.cache + linres.retcode === ReturnCode.Failure && + return LinearSolveResult(; linres.u, success = false) + return LinearSolveResult(; linres.u) + elseif !(cache.linsolve isa QRFactorization{ColumnNorm}) + if verbose + if structured_mat || is_gpuarray + mat_desc = structured_mat ? "Structured" : "GPU" + @warn "Potential Rank Deficient Matrix Detected. But Matrix is \ + $(mat_desc). Currently, we don't attempt to solve Rank Deficient \ + $(mat_desc) Matrices. Please open an issue at \ + https://github.com/SciML/NonlinearSolve.jl" + end + end + end + return LinearSolveResult(; linres.u, success = false) + end + + return LinearSolveResult(; linres.u) +end + +function NonlinearSolveBase.needs_square_A(linsolve::SciMLLinearSolveAlgorithm, ::Any) + return LinearSolve.needs_square_A(linsolve) +end +function NonlinearSolveBase.needs_concrete_A(linsolve::SciMLLinearSolveAlgorithm) + return LinearSolve.needs_concrete_A(linsolve) +end + +update_A!(cache::LinearSolveJLCache, ::Nothing, reuse) = cache +function update_A!(cache::LinearSolveJLCache, A, reuse) + return update_A!(cache, Utils.safe_getproperty(cache.linsolve, Val(:alg)), A, reuse) +end + +function update_A!(cache::LinearSolveJLCache, alg, A, reuse) + # Not a Factorization Algorithm so don't update `nfactors` + set_lincache_A!(cache.lincache, A) + return cache +end +function update_A!(cache::LinearSolveJLCache, ::LinearSolve.AbstractFactorization, A, reuse) + reuse && return cache + set_lincache_A!(cache.lincache, A) + cache.stats.nfactors += 1 + return cache +end +function update_A!( + cache::LinearSolveJLCache, alg::LinearSolve.DefaultLinearSolver, A, reuse) + if alg == + LinearSolve.DefaultLinearSolver(LinearSolve.DefaultAlgorithmChoice.KrylovJL_GMRES) + # Force a reset of the cache. This is not properly handled in LinearSolve.jl + set_lincache_A!(cache.lincache, A) + return cache + end + reuse && return cache + set_lincache_A!(cache.lincache, A) + cache.stats.nfactors += 1 + return cache +end + +function set_lincache_A!(lincache, new_A) + if !LinearSolve.default_alias_A(lincache.alg, new_A, lincache.b) && + ArrayInterface.can_setindex(lincache.A) + copyto!(lincache.A, new_A) + lincache.A = lincache.A # important!! triggers special code in `setproperty!` + return + end + lincache.A = new_A + return +end + +end diff --git a/lib/NonlinearSolveBase/ext/NonlinearSolveBaseSparseArraysExt.jl b/lib/NonlinearSolveBase/ext/NonlinearSolveBaseSparseArraysExt.jl index be13ebb8f..bc7350d21 100644 --- a/lib/NonlinearSolveBase/ext/NonlinearSolveBaseSparseArraysExt.jl +++ b/lib/NonlinearSolveBase/ext/NonlinearSolveBaseSparseArraysExt.jl @@ -1,10 +1,24 @@ module NonlinearSolveBaseSparseArraysExt -using NonlinearSolveBase: NonlinearSolveBase -using SparseArrays: AbstractSparseMatrixCSC, nonzeros +using SparseArrays: AbstractSparseMatrix, AbstractSparseMatrixCSC, nonzeros, sparse + +using NonlinearSolveBase: NonlinearSolveBase, Utils function NonlinearSolveBase.NAN_CHECK(x::AbstractSparseMatrixCSC) return any(NonlinearSolveBase.NAN_CHECK, nonzeros(x)) end +NonlinearSolveBase.sparse_or_structured_prototype(::AbstractSparseMatrix) = true + +Utils.maybe_symmetric(x::AbstractSparseMatrix) = x + +Utils.make_sparse(x) = sparse(x) + +Utils.condition_number(J::AbstractSparseMatrix) = Utils.condition_number(Matrix(J)) + +function Utils.maybe_pinv!!_workspace(A::AbstractSparseMatrix) + dense_A = Matrix(A) + return dense_A, copy(dense_A) +end + end diff --git a/lib/NonlinearSolveBase/ext/NonlinearSolveBaseSparseMatrixColoringsExt.jl b/lib/NonlinearSolveBase/ext/NonlinearSolveBaseSparseMatrixColoringsExt.jl new file mode 100644 index 000000000..4daf5ea98 --- /dev/null +++ b/lib/NonlinearSolveBase/ext/NonlinearSolveBaseSparseMatrixColoringsExt.jl @@ -0,0 +1,25 @@ +module NonlinearSolveBaseSparseMatrixColoringsExt + +using ADTypes: ADTypes, AbstractADType +using SciMLBase: SciMLBase, NonlinearFunction + +using SparseMatrixColorings: ConstantColoringAlgorithm, GreedyColoringAlgorithm, + LargestFirst + +using NonlinearSolveBase: NonlinearSolveBase, Utils + +Utils.is_extension_loaded(::Val{:SparseMatrixColorings}) = true + +function NonlinearSolveBase.select_fastest_coloring_algorithm( + ::Val{:SparseMatrixColorings}, + prototype, f::NonlinearFunction, ad::AbstractADType) + prototype === nothing && return GreedyColoringAlgorithm(LargestFirst()) + if SciMLBase.has_colorvec(f) + return ConstantColoringAlgorithm{ifelse( + ADTypes.mode(ad) isa ADTypes.ReverseMode, :row, :column)}( + prototype, f.colorvec) + end + return GreedyColoringAlgorithm(LargestFirst()) +end + +end diff --git a/lib/NonlinearSolveBase/src/NonlinearSolveBase.jl b/lib/NonlinearSolveBase/src/NonlinearSolveBase.jl index b07b4b168..a08384677 100644 --- a/lib/NonlinearSolveBase/src/NonlinearSolveBase.jl +++ b/lib/NonlinearSolveBase/src/NonlinearSolveBase.jl @@ -1,44 +1,84 @@ module NonlinearSolveBase -using ADTypes: ADTypes, AbstractADType -using ArrayInterface: ArrayInterface -using CommonSolve: CommonSolve using Compat: @compat using ConcreteStructs: @concrete -using DifferentiationInterface: DifferentiationInterface -using EnzymeCore: EnzymeCore using FastClosures: @closure +using Preferences: @load_preference, @set_preferences! + +using ADTypes: ADTypes, AbstractADType, AutoSparse, NoSparsityDetector, + KnownJacobianSparsityDetector +using Adapt: WrappedArray +using ArrayInterface: ArrayInterface +using DifferentiationInterface: DifferentiationInterface, Constant +using StaticArraysCore: StaticArray, SMatrix, SArray, MArray + +using CommonSolve: CommonSolve, init +using EnzymeCore: EnzymeCore using FunctionProperties: hasbranching -using LinearAlgebra: norm -using Markdown: @doc_str +using MaybeInplace: @bb using RecursiveArrayTools: AbstractVectorOfArray, ArrayPartition using SciMLBase: SciMLBase, ReturnCode, AbstractODEIntegrator, AbstractNonlinearProblem, - NonlinearProblem, NonlinearLeastSquaresProblem, AbstractNonlinearFunction, - @add_kwonly, StandardNonlinearProblem, NullParameters, isinplace, - warn_paramtype -using StaticArraysCore: StaticArray + AbstractNonlinearAlgorithm, AbstractNonlinearFunction, + NonlinearProblem, NonlinearLeastSquaresProblem, StandardNonlinearProblem, + NonlinearFunction, NullParameters, NLStats, LinearProblem +using SciMLJacobianOperators: JacobianOperator, StatefulJacobianOperator +using SciMLOperators: AbstractSciMLOperator, IdentityOperator +using SymbolicIndexingInterface: SymbolicIndexingInterface + +using LinearAlgebra: LinearAlgebra, Diagonal, norm, ldiv!, diagind +using Markdown: @doc_str +using Printf: @printf const DI = DifferentiationInterface +const SII = SymbolicIndexingInterface include("public.jl") include("utils.jl") +include("abstract_types.jl") + include("immutable_problem.jl") include("common_defaults.jl") include("termination_conditions.jl") include("autodiff.jl") +include("jacobian.jl") +include("linear_solve.jl") +include("timer_outputs.jl") +include("tracing.jl") +include("wrappers.jl") + +include("descent/common.jl") +include("descent/newton.jl") +include("descent/steepest.jl") +include("descent/damped_newton.jl") +include("descent/dogleg.jl") +include("descent/geodesic_acceleration.jl") + +include("solve.jl") # Unexported Public API @compat(public, (L2_NORM, Linf_NORM, NAN_CHECK, UNITLESS_ABS2, get_tolerance)) @compat(public, (nonlinearsolve_forwarddiff_solve, nonlinearsolve_dual_solution)) @compat(public, - (select_forward_mode_autodiff, select_reverse_mode_autodiff, - select_jacobian_autodiff)) + (select_forward_mode_autodiff, select_reverse_mode_autodiff, select_jacobian_autodiff)) + +# public for NonlinearSolve.jl and subpackages to use +@compat(public, (InternalAPI, supports_line_search, supports_trust_region, set_du!)) +@compat(public, (construct_linear_solver, needs_square_A, needs_concrete_A)) +@compat(public, (construct_jacobian_cache,)) +@compat(public, + (assert_extension_supported_termination_condition, + construct_extension_function_wrapper, construct_extension_jac)) + +export TraceMinimal, TraceWithJacobianConditionNumber, TraceAll export RelTerminationMode, AbsTerminationMode, NormTerminationMode, RelNormTerminationMode, AbsNormTerminationMode, RelNormSafeTerminationMode, AbsNormSafeTerminationMode, RelNormSafeBestTerminationMode, AbsNormSafeBestTerminationMode +export DescentResult, SteepestDescent, NewtonDescent, DampedNewtonDescent, Dogleg, + GeodesicAcceleration + end diff --git a/lib/NonlinearSolveBase/src/abstract_types.jl b/lib/NonlinearSolveBase/src/abstract_types.jl new file mode 100644 index 000000000..6105f5125 --- /dev/null +++ b/lib/NonlinearSolveBase/src/abstract_types.jl @@ -0,0 +1,587 @@ +module InternalAPI + +using SciMLBase: NLStats + +function init end +function solve! end +function step! end + +function reinit! end +function reinit_self! end + +function reinit!(x::Any; kwargs...) + @debug "`InternalAPI.reinit!` is not implemented for $(typeof(x))." + return +end +function reinit_self!(x::Any; kwargs...) + @debug "`InternalAPI.reinit_self!` is not implemented for $(typeof(x))." + return +end + +function reinit!(stats::NLStats) + stats.nf = 0 + stats.nsteps = 0 + stats.nfactors = 0 + stats.njacs = 0 + stats.nsolve = 0 +end + +end + +abstract type AbstractNonlinearSolveBaseAPI end # Mostly used for pretty-printing + +function Base.show(io::IO, ::MIME"text/plain", alg::AbstractNonlinearSolveBaseAPI) + print(io, Utils.clean_sprint_struct(alg)) + return +end + +""" + AbstractDescentDirection + +Abstract Type for all Descent Directions used in NonlinearSolveBase. Given the Jacobian +`J` and the residual `fu`, these algorithms compute the descent direction `δu`. + +For non-square Jacobian problems, if we need to solve a linear solve problem, we use a +least squares solver by default, unless the provided `linsolve` can't handle non-square +matrices, in which case we use the normal form equations ``JᵀJ δu = Jᵀ fu``. Note that +this factorization is often the faster choice, but it is not as numerically stable as +the least squares solver. + +### `InternalAPI.init` specification + +```julia +InternalAPI.init( + prob::AbstractNonlinearProblem, alg::AbstractDescentDirection, J, fu, u; + pre_inverted::Val = Val(false), linsolve_kwargs = (;), + abstol = nothing, reltol = nothing, alias_J::Bool = true, + shared::Val = Val(1), kwargs... +)::AbstractDescentCache +``` + + - `pre_inverted`: whether or not the Jacobian has been pre_inverted. + - `linsolve_kwargs`: keyword arguments to pass to the linear solver. + - `abstol`: absolute tolerance for the linear solver. + - `reltol`: relative tolerance for the linear solver. + - `alias_J`: whether or not to alias the Jacobian. + - `shared`: Store multiple descent directions in the cache. Allows efficient and + correct reuse of factorizations if needed. + +Some of the algorithms also allow additional keyword arguments. See the documentation for +the specific algorithm for more information. + +### Interface Functions + + - `supports_trust_region(alg)`: whether or not the algorithm supports trust region + methods. Defaults to `false`. + - `supports_line_search(alg)`: whether or not the algorithm supports line search + methods. Defaults to `false`. + +See also [`NewtonDescent`](@ref), [`Dogleg`](@ref), [`SteepestDescent`](@ref), +[`DampedNewtonDescent`](@ref). +""" +abstract type AbstractDescentDirection <: AbstractNonlinearSolveBaseAPI end + +supports_line_search(::AbstractDescentDirection) = false +supports_trust_region(::AbstractDescentDirection) = false + +function get_linear_solver(alg::AbstractDescentDirection) + return Utils.safe_getproperty(alg, Val(:linsolve)) +end + +""" + AbstractDescentCache + +Abstract Type for all Descent Caches. + +### `InternalAPI.solve!` specification + +```julia +InternalAPI.solve!( + cache::AbstractDescentCache, J, fu, u, idx::Val; + skip_solve::Bool = false, new_jacobian::Bool = true, kwargs... +)::DescentResult +``` + + - `J`: Jacobian or Inverse Jacobian (if `pre_inverted = Val(true)`). + - `fu`: residual. + - `u`: current state. + - `idx`: index of the descent problem to solve and return. Defaults to `Val(1)`. + - `skip_solve`: Skip the direction computation and return the previous direction. + Defaults to `false`. This is useful for Trust Region Methods where the previous + direction was rejected and we want to try with a modified trust region. + - `new_jacobian`: Whether the Jacobian has been updated. Defaults to `true`. + - `kwargs`: keyword arguments to pass to the linear solver if there is one. + +#### Returned values + + - `descent_result`: Result in a [`DescentResult`](@ref). + +### Interface Functions + + - `get_du(cache)`: get the descent direction. + - `get_du(cache, ::Val{N})`: get the `N`th descent direction. + - `set_du!(cache, δu)`: set the descent direction. + - `set_du!(cache, δu, ::Val{N})`: set the `N`th descent direction. + - `last_step_accepted(cache)`: whether or not the last step was accepted. Checks if the + cache has a `last_step_accepted` field and returns it if it does, else returns `true`. + - `preinverted_jacobian(cache)`: whether or not the Jacobian has been preinverted. + - `normal_form(cache)`: whether or not the linear solver uses normal form. +""" +abstract type AbstractDescentCache <: AbstractNonlinearSolveBaseAPI end + +SciMLBase.get_du(cache::AbstractDescentCache) = cache.δu +SciMLBase.get_du(cache::AbstractDescentCache, ::Val{1}) = SciMLBase.get_du(cache) +SciMLBase.get_du(cache::AbstractDescentCache, ::Val{N}) where {N} = cache.δus[N - 1] +set_du!(cache::AbstractDescentCache, δu) = (cache.δu = δu) +set_du!(cache::AbstractDescentCache, δu, ::Val{1}) = set_du!(cache, δu) +set_du!(cache::AbstractDescentCache, δu, ::Val{N}) where {N} = (cache.δus[N - 1] = δu) + +function last_step_accepted(cache::AbstractDescentCache) + hasfield(typeof(cache), :last_step_accepted) && return cache.last_step_accepted + return true +end + +for fname in (:preinverted_jacobian, :normal_form) + @eval function $(fname)(alg::AbstractDescentCache) + res = Utils.unwrap_val(Utils.safe_getproperty(alg, Val($(QuoteNode(fname))))) + res === missing && return false + return res + end +end + +""" + AbstractDampingFunction + +Abstract Type for Damping Functions in DampedNewton. + +### `InternalAPI.init` specification + +```julia +InternalAPI.init( + prob::AbstractNonlinearProblem, f::AbstractDampingFunction, initial_damping, + J, fu, u, args...; + internalnorm = L2_NORM, kwargs... +)::AbstractDampingFunctionCache +``` + +Returns a [`NonlinearSolveBase.AbstractDampingFunctionCache`](@ref). +""" +abstract type AbstractDampingFunction <: AbstractNonlinearAlgorithm end + +""" + AbstractDampingFunctionCache + +Abstract Type for the Caches created by AbstractDampingFunctions + +### Interface Functions + + - `requires_normal_form_jacobian(alg)`: whether or not the Jacobian is needed in normal + form. No default. + - `requires_normal_form_rhs(alg)`: whether or not the residual is needed in normal form. + No default. + - `returns_norm_form_damping(alg)`: whether or not the damping function returns the + damping factor in normal form. Defaults to + `requires_normal_form_jacobian(alg) || requires_normal_form_rhs(alg)`. + - `(cache::AbstractDampingFunctionCache)(::Nothing)`: returns the damping factor. The type + of the damping factor returned from `solve!` is guaranteed to be the same as this. + +### `InternalAPI.solve!` specification + +```julia +InternalAPI.solve!( + cache::AbstractDampingFunctionCache, J, fu, u, δu, descent_stats +) +``` + +Returns the damping factor. +""" +abstract type AbstractDampingFunctionCache <: AbstractNonlinearAlgorithm end + +function requires_normal_form_jacobian end +function requires_normal_form_rhs end +function returns_norm_form_damping(f::F) where {F} + return requires_normal_form_jacobian(f) || requires_normal_form_rhs(f) +end + +""" + AbstractNonlinearSolveAlgorithm <: AbstractNonlinearAlgorithm + +Abstract Type for all NonlinearSolveBase Algorithms. + +### Interface Functions + + - `concrete_jac(alg)`: whether or not the algorithm uses a concrete Jacobian. Defaults + to `nothing`. +""" +abstract type AbstractNonlinearSolveAlgorithm <: AbstractNonlinearAlgorithm end + +""" + concrete_jac(alg::AbstractNonlinearSolveAlgorithm)::Bool + +Whether the algorithm uses a concrete Jacobian. +""" +function concrete_jac(alg::AbstractNonlinearSolveAlgorithm) + return concrete_jac(Utils.safe_getproperty(alg, Val(:concrete_jac))) +end +concrete_jac(::Missing) = false +concrete_jac(::Nothing) = false +concrete_jac(v::Bool) = v +concrete_jac(::Val{false}) = false +concrete_jac(::Val{true}) = true + +function Base.show(io::IO, ::MIME"text/plain", alg::AbstractNonlinearSolveAlgorithm) + print(io, Utils.clean_sprint_struct(alg, 0)) + return +end + +function show_nonlinearsolve_algorithm( + io::IO, alg::AbstractNonlinearSolveAlgorithm, name, indent::Int = 0 +) + print(io, name) + print(io, Utils.clean_sprint_struct(alg, indent)) +end + +""" + AbstractNonlinearSolveCache + +Abstract Type for all NonlinearSolveBase Caches. + +### Interface Functions + + - `get_fu(cache)`: get the residual. + + - `get_u(cache)`: get the current state. + - `set_fu!(cache, fu)`: set the residual. + - `has_time_limit(cache)`: whether or not the solver has a maximum time limit. + - `not_terminated(cache)`: whether or not the solver has terminated. + - `SciMLBase.set_u!(cache, u)`: set the current state. + - `SciMLBase.reinit!(cache, u0; kwargs...)`: reinitialize the cache with the initial state + `u0` and any additional keyword arguments. + - `SciMLBase.isinplace(cache)`: whether or not the solver is inplace. + - `CommonSolve.step!(cache; kwargs...)`: See [`CommonSolve.step!`](@ref) for more details. + +Additionally implements `SymbolicIndexingInterface` interface Functions. + +#### Expected Fields in Sub-Types + +For the default interface implementations we expect the following fields to be present in +the cache: + + - `fu`: the residual. + - `u`: the current state. + - `maxiters`: the maximum number of iterations. + - `nsteps`: the number of steps taken. + - `force_stop`: whether or not the solver has been forced to stop. + - `retcode`: the return code. + - `stats`: `NLStats` object. + - `alg`: the algorithm. + - `maxtime`: the maximum time limit for the solver. (Optional) + - `timer`: the timer for the solver. (Optional) + - `total_time`: the total time taken by the solver. (Optional) +""" +abstract type AbstractNonlinearSolveCache <: AbstractNonlinearSolveBaseAPI end + +get_u(cache::AbstractNonlinearSolveCache) = cache.u +get_fu(cache::AbstractNonlinearSolveCache) = cache.fu +set_fu!(cache::AbstractNonlinearSolveCache, fu) = (cache.fu = fu) +SciMLBase.set_u!(cache::AbstractNonlinearSolveCache, u) = (cache.u = u) + +function has_time_limit(cache::AbstractNonlinearSolveCache) + maxtime = Utils.safe_getproperty(cache, Val(:maxtime)) + return maxtime !== missing && maxtime !== nothing +end + +function not_terminated(cache::AbstractNonlinearSolveCache) + return !cache.force_stop && cache.nsteps < cache.maxiters +end + +function SciMLBase.reinit!(cache::AbstractNonlinearSolveCache; kwargs...) + return InternalAPI.reinit!(cache; kwargs...) +end +function SciMLBase.reinit!(cache::AbstractNonlinearSolveCache, u0; kwargs...) + return InternalAPI.reinit!(cache; u0, kwargs...) +end + +SciMLBase.isinplace(cache::AbstractNonlinearSolveCache) = SciMLBase.isinplace(cache.prob) + +## SII Interface +SII.symbolic_container(cache::AbstractNonlinearSolveCache) = cache.prob +SII.parameter_values(cache::AbstractNonlinearSolveCache) = SII.parameter_values(cache.prob) +SII.state_values(cache::AbstractNonlinearSolveCache) = SII.state_values(cache.prob) + +function Base.getproperty(cache::AbstractNonlinearSolveCache, sym::Symbol) + if sym === :ps + !hasfield(typeof(cache), :ps) && return SII.ParameterIndexingProxy(cache) + return getfield(cache, :ps) + end + return getfield(cache, sym) +end + +Base.getindex(cache::AbstractNonlinearSolveCache, sym) = SII.getu(cache, sym)(cache) +function Base.setindex!(cache::AbstractNonlinearSolveCache, val, sym) + return SII.setu(cache, sym)(cache, val) +end + +function Base.show(io::IO, ::MIME"text/plain", cache::AbstractNonlinearSolveCache) + return show_nonlinearsolve_cache(io, cache) +end + +function show_nonlinearsolve_cache(io::IO, cache::AbstractNonlinearSolveCache, indent = 0) + println(io, "$(nameof(typeof(cache)))(") + show_nonlinearsolve_algorithm( + io, + cache.alg, + (" "^(indent + 4)) * "alg = ", + indent + 4 + ) + + ustr = sprint(show, get_u(cache); context = (:compact => true, :limit => true)) + println(io, ",\n" * (" "^(indent + 4)) * "u = $(ustr),") + + residstr = sprint(show, get_fu(cache); context = (:compact => true, :limit => true)) + println(io, (" "^(indent + 4)) * "residual = $(residstr),") + + normstr = sprint( + show, norm(get_fu(cache), Inf); context = (:compact => true, :limit => true) + ) + println(io, (" "^(indent + 4)) * "inf-norm(residual) = $(normstr),") + + println(io, " "^(indent + 4) * "nsteps = ", cache.stats.nsteps, ",") + println(io, " "^(indent + 4) * "retcode = ", cache.retcode) + print(io, " "^(indent) * ")") +end + +""" + AbstractLinearSolverCache + +Abstract Type for all Linear Solvers used in NonlinearSolveBase. Subtypes of these are +meant to be constructured via [`construct_linear_solver`](@ref). +""" +abstract type AbstractLinearSolverCache <: AbstractNonlinearSolveBaseAPI end + +""" + AbstractJacobianCache + +Abstract Type for all Jacobian Caches used in NonlinearSolveBase. Subtypes of these are +meant to be constructured via [`construct_jacobian_cache`](@ref). +""" +abstract type AbstractJacobianCache <: AbstractNonlinearSolveBaseAPI end + +""" + AbstractApproximateJacobianStructure + +Abstract Type for all Approximate Jacobian Structures used in NonlinearSolve.jl. + +### Interface Functions + + - `stores_full_jacobian(alg)`: whether or not the algorithm stores the full Jacobian. + Defaults to `false`. + - `get_full_jacobian(cache, alg, J)`: get the full Jacobian. Defaults to throwing an + error if `stores_full_jacobian(alg)` is `false`. +""" +abstract type AbstractApproximateJacobianStructure <: AbstractNonlinearSolveBaseAPI end + +stores_full_jacobian(::AbstractApproximateJacobianStructure) = false +function get_full_jacobian(cache, alg::AbstractApproximateJacobianStructure, J) + stores_full_jacobian(alg) && return J + error("This algorithm does not store the full Jacobian. Define `get_full_jacobian` for \ + this algorithm.") +end + +""" + AbstractJacobianInitialization + +Abstract Type for all Jacobian Initialization Algorithms used in NonlinearSolveBase. + +### Interface Functions + + - `jacobian_initialized_preinverted(alg)`: whether or not the Jacobian is initialized + preinverted. Defaults to `false`. + +### `InternalAPI.init` specification + +```julia +InternalAPI.init( + prob::AbstractNonlinearProblem, alg::AbstractJacobianInitialization, solver, + f, fu, u, p; + linsolve = missing, internalnorm::IN = L2_NORM, kwargs... +)::AbstractJacobianCache +``` + +All subtypes need to define +`(cache::AbstractJacobianCache)(alg::NewSubType, fu, u)` which reinitializes the Jacobian in +`cache.J`. +""" +abstract type AbstractJacobianInitialization <: AbstractNonlinearSolveBaseAPI end + +jacobian_initialized_preinverted(::AbstractJacobianInitialization) = false + +""" + AbstractApproximateJacobianUpdateRule + +Abstract Type for all Approximate Jacobian Update Rules used in NonlinearSolveBase. + +### Interface Functions + + - `store_inverse_jacobian(alg)`: Return `alg.store_inverse_jacobian` + +### `InternalAPI.init` specification + +```julia +InternalAPI.init( + prob::AbstractNonlinearProblem, alg::AbstractApproximateJacobianUpdateRule, J, fu, u, + du, args...; internalnorm = L2_NORM, kwargs... +)::AbstractApproximateJacobianUpdateRuleCache +``` +""" +abstract type AbstractApproximateJacobianUpdateRule <: AbstractNonlinearSolveBaseAPI end + +function store_inverse_jacobian(rule::AbstractApproximateJacobianUpdateRule) + return rule.store_inverse_jacobian +end + +""" + AbstractApproximateJacobianUpdateRuleCache + +Abstract Type for all Approximate Jacobian Update Rule Caches used in NonlinearSolveBase. + +### Interface Functions + + - `store_inverse_jacobian(cache)`: Return `store_inverse_jacobian(cache.rule)` + +### `InternalAPI.solve!` specification + +```julia +InternalAPI.solve!( + cache::AbstractApproximateJacobianUpdateRuleCache, J, fu, u, du; kwargs... +) --> J / J⁻¹ +``` +""" +abstract type AbstractApproximateJacobianUpdateRuleCache <: AbstractNonlinearSolveBaseAPI end + +function store_inverse_jacobian(cache::AbstractApproximateJacobianUpdateRuleCache) + return store_inverse_jacobian(cache.rule) +end + +""" + AbstractResetCondition + +Condition for resetting the Jacobian in Quasi-Newton's methods. + +### `InternalAPI.init` specification + +```julia +InternalAPI.init( + alg::AbstractResetCondition, J, fu, u, du, args...; kwargs... +)::AbstractResetConditionCache +``` +""" +abstract type AbstractResetCondition <: AbstractNonlinearSolveBaseAPI end + +""" + AbstractResetConditionCache + +Abstract Type for all Reset Condition Caches used in NonlinearSolveBase. + +### `InternalAPI.solve!` specification + +```julia +InternalAPI.solve!( + cache::AbstractResetConditionCache, J, fu, u, du; kwargs... +)::Bool +``` +""" +abstract type AbstractResetConditionCache <: AbstractNonlinearSolveBaseAPI end + +""" + AbstractTrustRegionMethod + +Abstract Type for all Trust Region Methods used in NonlinearSolveBase. + +### `InternalAPI.init` specification + +```julia +InternalAPI.init( + prob::AbstractNonlinearProblem, alg::AbstractTrustRegionMethod, f, fu, u, p, args...; + internalnorm = L2_NORM, kwargs... +)::AbstractTrustRegionMethodCache +``` +""" +abstract type AbstractTrustRegionMethod <: AbstractNonlinearSolveBaseAPI end + +""" + AbstractTrustRegionMethodCache + +Abstract Type for all Trust Region Method Caches used in NonlinearSolveBase. + +### Interface Functions + + - `last_step_accepted(cache)`: whether or not the last step was accepted. Defaults to + `cache.last_step_accepted`. Should if overloaded if the field is not present. + +### `InternalAPI.solve!` specification + +```julia +InternalAPI.solve!( + cache::AbstractTrustRegionMethodCache, J, fu, u, δu, descent_stats; kwargs... +) +``` + +Returns `last_step_accepted`, updated `u_cache` and `fu_cache`. If the last step was +accepted then these values should be copied into the toplevel cache. +""" +abstract type AbstractTrustRegionMethodCache <: AbstractNonlinearSolveBaseAPI end + +last_step_accepted(cache::AbstractTrustRegionMethodCache) = cache.last_step_accepted + +# Additional Interface +""" + callback_into_cache!(cache, internalcache, args...) + +Define custom operations on `internalcache` tightly coupled with the calling `cache`. +`args...` contain the sequence of caches calling into `internalcache`. + +This unfortunately makes code very tightly coupled and not modular. It is recommended to not +use this functionality unless it can't be avoided (like in `LevenbergMarquardt`). +""" +callback_into_cache!(cache, internalcache, args...) = nothing # By default do nothing + +# Helper functions to generate cache callbacks and resetting functions +macro internal_caches(cType, internal_cache_names...) + callback_caches = map(internal_cache_names) do name + return quote + $(callback_into_cache!)( + cache, getproperty(internalcache, $(name)), internalcache, args... + ) + end + end + callbacks_self = map(internal_cache_names) do name + return quote + $(callback_into_cache!)(cache, getproperty(cache, $(name))) + end + end + reinit_caches = map(internal_cache_names) do name + return quote + $(InternalAPI.reinit!)(getproperty(cache, $(name)), args...; kwargs...) + end + end + return esc(quote + function NonlinearSolveBase.callback_into_cache!( + cache, internalcache::$(cType), args... + ) + $(callback_caches...) + return + end + function NonlinearSolveBase.callback_into_cache!(cache::$(cType)) + $(callbacks_self...) + return + end + function NonlinearSolveBase.InternalAPI.reinit!( + cache::$(cType), args...; kwargs... + ) + $(reinit_caches...) + $(InternalAPI.reinit_self!)(cache, args...; kwargs...) + return + end + end) +end diff --git a/src/descent/common.jl b/lib/NonlinearSolveBase/src/descent/common.jl similarity index 74% rename from src/descent/common.jl rename to lib/NonlinearSolveBase/src/descent/common.jl index 3d53573ea..b757a990a 100644 --- a/src/descent/common.jl +++ b/lib/NonlinearSolveBase/src/descent/common.jl @@ -1,6 +1,8 @@ """ - DescentResult(; δu = missing, u = missing, success::Bool = true, - linsolve_success::Bool = true, extras = (;)) + DescentResult(; + δu = missing, u = missing, success::Bool = true, linsolve_success::Bool = true, + extras = (;) + ) Construct a `DescentResult` object. @@ -23,8 +25,10 @@ Construct a `DescentResult` object. extras end -function DescentResult(; δu = missing, u = missing, success::Bool = true, - linsolve_success::Bool = true, extras = (;)) +function DescentResult(; + δu = missing, u = missing, success::Bool = true, linsolve_success::Bool = true, + extras = (;) +) @assert δu !== missing || u !== missing return DescentResult(δu, u, success, linsolve_success, extras) end diff --git a/lib/NonlinearSolveBase/src/descent/damped_newton.jl b/lib/NonlinearSolveBase/src/descent/damped_newton.jl new file mode 100644 index 000000000..47ed56207 --- /dev/null +++ b/lib/NonlinearSolveBase/src/descent/damped_newton.jl @@ -0,0 +1,267 @@ +""" + DampedNewtonDescent(; + linsolve = nothing, precs = nothing, initial_damping, damping_fn + ) + +A Newton descent algorithm with damping. The damping factor is computed using the +`damping_fn` function. The descent direction is computed as ``(JᵀJ + λDᵀD) δu = -fu``. For +non-square Jacobians, we default to solving for `Jδx = -fu` and `√λ⋅D δx = 0` +simultaneously. If the linear solver can't handle non-square matrices, we use the normal +form equations ``(JᵀJ + λDᵀD) δu = Jᵀ fu``. Note that this factorization is often the faster +choice, but it is not as numerically stable as the least squares solver. + +The damping factor returned must be a non-negative number. + +### Keyword Arguments + + - `initial_damping`: the initial damping factor to use + - `damping_fn`: the function to use to compute the damping factor. This must satisfy the + [`NonlinearSolveBase.AbstractDampingFunction`](@ref) interface. +""" +@kwdef @concrete struct DampedNewtonDescent <: AbstractDescentDirection + linsolve = nothing + precs = nothing + initial_damping + damping_fn <: AbstractDampingFunction +end + +supports_line_search(::DampedNewtonDescent) = true +supports_trust_region(::DampedNewtonDescent) = true + +@concrete mutable struct DampedNewtonDescentCache <: AbstractDescentCache + J + δu + δus + lincache + JᵀJ_cache + Jᵀfu_cache + rhs_cache + damping_fn_cache + timer + preinverted_jacobian <: Union{Val{false}, Val{true}} + mode <: Union{Val{:normal_form}, Val{:least_squares}, Val{:simple}} +end + +@internal_caches DampedNewtonDescentCache :lincache :damping_fn_cache + +function InternalAPI.init( + prob::AbstractNonlinearProblem, alg::DampedNewtonDescent, J, fu, u; stats, + pre_inverted::Val = Val(false), linsolve_kwargs = (;), + abstol = nothing, reltol = nothing, + timer = get_timer_output(), + alias_J::Bool = true, shared::Val = Val(1), + kwargs... +) + length(fu) != length(u) && + @assert pre_inverted isa Val{false} "Precomputed Inverse for Non-Square Jacobian doesn't make sense." + + @bb δu = similar(u) + δus = Utils.unwrap_val(shared) ≤ 1 ? nothing : map(2:Utils.unwrap_val(shared)) do i + @bb δu_ = similar(u) + end + + normal_form_damping = returns_norm_form_damping(alg.damping_fn) + normal_form_linsolve = needs_square_A(alg.linsolve, u) + + mode = if u isa Number + :simple + elseif prob isa NonlinearProblem + if normal_form_damping + ifelse(normal_form_linsolve, :normal_form, :least_squares) + else + :simple + end + else + if normal_form_linsolve & !normal_form_damping + throw(ArgumentError("Linear Solver expects Normal Form but returned Damping is \ + not Normal Form. This is not supported.")) + end + if normal_form_damping & !normal_form_linsolve + :least_squares + else + ifelse(!normal_form_damping & !normal_form_linsolve, :simple, :normal_form) + end + end + + if mode === :least_squares + if requires_normal_form_jacobian(alg.damping_fn) + JᵀJ = transpose(J) * J # Needed to compute the damping factor + jac_damp = JᵀJ + else + JᵀJ = nothing + jac_damp = J + end + if requires_normal_form_rhs(alg.damping_fn) + Jᵀfu = transpose(J) * Utils.safe_vec(fu) + rhs_damp = Jᵀfu + else + Jᵀfu = nothing + rhs_damp = fu + end + + damping_fn_cache = InternalAPI.init( + prob, alg.damping_fn, alg.initial_damping, jac_damp, rhs_damp, u, Val(false); + stats, kwargs... + ) + D = damping_fn_cache(nothing) + + D isa Number && (D = D * LinearAlgebra.I) + rhs_cache = vcat(Utils.safe_vec(fu), Utils.safe_vec(u)) + J_cache = Utils.faster_vcat(J, D) + A, b = J_cache, rhs_cache + elseif mode === :simple + damping_fn_cache = InternalAPI.init( + prob, alg.damping_fn, alg.initial_damping, J, fu, u, Val(false); kwargs... + ) + J_cache = Utils.maybe_unaliased(J, alias_J) + D = damping_fn_cache(nothing) + + J_damped = dampen_jacobian!!(J_cache, J, D) + J_cache = J_damped + A, b = J_damped, Utils.safe_vec(fu) + JᵀJ, Jᵀfu, rhs_cache = nothing, nothing, nothing + elseif mode === :normal_form + JᵀJ = transpose(J) * J + Jᵀfu = transpose(J) * Utils.safe_vec(fu) + jac_damp = requires_normal_form_jacobian(alg.damping_fn) ? JᵀJ : J + rhs_damp = requires_normal_form_rhs(alg.damping_fn) ? Jᵀfu : fu + + damping_fn_cache = InternalAPI.init( + prob, alg.damping_fn, alg.initial_damping, jac_damp, rhs_damp, u, Val(true); + stats, kwargs... + ) + D = damping_fn_cache(nothing) + + @bb J_cache = similar(JᵀJ) + @bb @. J_cache = 0 + J_damped = dampen_jacobian!!(J_cache, JᵀJ, D) + A, b = Utils.maybe_symmetric(J_damped), Utils.safe_vec(Jᵀfu) + rhs_cache = nothing + end + + lincache = construct_linear_solver( + alg, alg.linsolve, A, b, Utils.safe_vec(u); + stats, abstol, reltol, linsolve_kwargs... + ) + + return DampedNewtonDescentCache( + J_cache, δu, δus, lincache, JᵀJ, Jᵀfu, rhs_cache, + damping_fn_cache, timer, pre_inverted, Val(mode) + ) +end + +function InternalAPI.solve!( + cache::DampedNewtonDescentCache, J, fu, u, idx::Val = Val(1); + skip_solve::Bool = false, new_jacobian::Bool = true, kwargs... +) + δu = SciMLBase.get_du(cache, idx) + skip_solve && return DescentResult(; δu) + + recompute_A = idx === Val(1) + + @static_timeit cache.timer "dampen" begin + if cache.mode isa Val{:least_squares} + if (J !== nothing || new_jacobian) && recompute_A + preinverted_jacobian(cache) && (J = inv(J)) + if requires_normal_form_jacobian(cache.damping_fn_cache) + @bb cache.JᵀJ_cache = transpose(J) × J + jac_damp = cache.JᵀJ_cache + else + jac_damp = J + end + if requires_normal_form_rhs(cache.damping_fn_cache) + @bb cache.Jᵀfu_cache = transpose(J) × vec(fu) + rhs_damp = cache.Jᵀfu_cache + else + rhs_damp = fu + end + D = InternalAPI.solve!( + cache.damping_fn_cache, jac_damp, rhs_damp, Val(false) + ) + if Utils.can_setindex(cache.J) + copyto!(@view(cache.J[1:size(J, 1), :]), J) + cache.J[(size(J, 1) + 1):end, :] .= sqrt.(D) + else + cache.J = Utils.faster_vcat(J, sqrt.(D)) + end + end + A = cache.J + if Utils.can_setindex(cache.rhs_cache) + cache.rhs_cache[1:length(fu)] .= Utils.safe_vec(fu) + cache.rhs_cache[(length(fu) + 1):end] .= false + else + cache.rhs_cache = vcat(Utils.safe_vec(fu), zero(Utils.safe_vec(u))) + end + b = cache.rhs_cache + elseif cache.mode isa Val{:simple} + if (J !== nothing || new_jacobian) && recompute_A + preinverted_jacobian(cache) && (J = inv(J)) + D = InternalAPI.solve!(cache.damping_fn_cache, J, fu, Val(false)) + cache.J = dampen_jacobian!!(cache.J, J, D) + end + A, b = cache.J, Utils.safe_vec(fu) + elseif cache.mode isa Val{:normal_form} + if (J !== nothing || new_jacobian) && recompute_A + preinverted_jacobian(cache) && (J = inv(J)) + @bb cache.JᵀJ_cache = transpose(J) × J + @bb cache.Jᵀfu_cache = transpose(J) × vec(fu) + D = InternalAPI.solve!( + cache.damping_fn_cache, cache.JᵀJ_cache, cache.Jᵀfu_cache, Val(true) + ) + cache.J = dampen_jacobian!!(cache.J, cache.JᵀJ_cache, D) + A = Utils.maybe_symmetric(cache.J) + elseif !recompute_A + @bb cache.Jᵀfu_cache = transpose(J) × vec(fu) + A = Utils.maybe_symmetric(cache.J) + else + A = nothing + end + b = cache.Jᵀfu_cache + else + error("Unknown Mode: $(cache.mode).") + end + end + + @static_timeit cache.timer "linear solve" begin + linres = cache.lincache(; + A, b, + reuse_A_if_factorization = !new_jacobian && !recompute_A, + kwargs..., + linu = Utils.safe_vec(δu) + ) + δu = Utils.restructure(SciMLBase.get_du(cache, idx), linres.u) + if !linres.success + set_du!(cache, δu, idx) + return DescentResult(; δu, success = false, linsolve_success = false) + end + end + + @bb @. δu *= -1 + set_du!(cache, δu, idx) + return DescentResult(; δu) +end + +dampen_jacobian!!(::Any, J::Union{AbstractSciMLOperator, Number}, D) = J + D +function dampen_jacobian!!(J_cache, J::AbstractMatrix, D::Union{AbstractMatrix, Number}) + ArrayInterface.can_setindex(J_cache) || return J .+ D + J_cache !== J && copyto!(J_cache, J) + if ArrayInterface.fast_scalar_indexing(J_cache) + if D isa Number + @simd ivdep for i in axes(J_cache, 1) + @inbounds J_cache[i, i] += D + end + else + @simd ivdep for i in axes(J_cache, 1) + @inbounds J_cache[i, i] += D[i, i] + end + end + else + idxs = diagind(J_cache) + if D isa Number + J_cache[idxs] .+= D + else + J_cache[idxs] .+= @view(D[idxs]) + end + end + return J_cache +end diff --git a/src/descent/dogleg.jl b/lib/NonlinearSolveBase/src/descent/dogleg.jl similarity index 57% rename from src/descent/dogleg.jl rename to lib/NonlinearSolveBase/src/descent/dogleg.jl index 4c96c98f6..f446e9cf6 100644 --- a/src/descent/dogleg.jl +++ b/lib/NonlinearSolveBase/src/descent/dogleg.jl @@ -1,5 +1,5 @@ """ - Dogleg(; linsolve = nothing, precs = DEFAULT_PRECS) + Dogleg(; linsolve = nothing, precs = nothing) Switch between Newton's method and the steepest descent method depending on the size of the trust region. The trust region is specified via keyword argument `trust_region` to @@ -7,82 +7,93 @@ trust region. The trust region is specified via keyword argument `trust_region` See also [`SteepestDescent`](@ref), [`NewtonDescent`](@ref), [`DampedNewtonDescent`](@ref). """ -@concrete struct Dogleg <: AbstractDescentAlgorithm - newton_descent - steepest_descent -end - -function Base.show(io::IO, d::Dogleg) - print(io, - "Dogleg(newton_descent = $(d.newton_descent), steepest_descent = $(d.steepest_descent))") +@concrete struct Dogleg <: AbstractDescentDirection + newton_descent <: Union{NewtonDescent, DampedNewtonDescent} + steepest_descent <: SteepestDescent end supports_trust_region(::Dogleg) = true get_linear_solver(alg::Dogleg) = get_linear_solver(alg.newton_descent) -function Dogleg(; linsolve = nothing, precs = DEFAULT_PRECS, damping = False, +function Dogleg(; linsolve = nothing, precs = nothing, damping = Val(false), damping_fn = missing, initial_damping = missing, kwargs...) - if damping === False + if !Utils.unwrap_val(damping) return Dogleg(NewtonDescent(; linsolve, precs), SteepestDescent(; linsolve, precs)) end if damping_fn === missing || initial_damping === missing throw(ArgumentError("`damping_fn` and `initial_damping` must be supplied if \ `damping = Val(true)`.")) end - return Dogleg(DampedNewtonDescent(; linsolve, precs, damping_fn, initial_damping), - SteepestDescent(; linsolve, precs)) + return Dogleg( + DampedNewtonDescent(; linsolve, precs, damping_fn, initial_damping), + SteepestDescent(; linsolve, precs) + ) end -@concrete mutable struct DoglegCache{pre_inverted, normalform} <: AbstractDescentCache +@concrete mutable struct DoglegCache <: AbstractDescentCache δu δus - newton_cache - cauchy_cache + newton_cache <: Union{NewtonDescentCache, DampedNewtonDescentCache} + cauchy_cache <: SteepestDescentCache internalnorm - JᵀJ_cache + Jᵀδu_cache δu_cache_1 δu_cache_2 δu_cache_mul + preinverted_jacobian <: Union{Val{false}, Val{true}} + normal_form <: Union{Val{false}, Val{true}} end @internal_caches DoglegCache :newton_cache :cauchy_cache -function __internal_init(prob::AbstractNonlinearProblem, alg::Dogleg, J, fu, u; - pre_inverted::Val{INV} = False, linsolve_kwargs = (;), +function InternalAPI.init( + prob::AbstractNonlinearProblem, alg::Dogleg, J, fu, u; + pre_inverted::Val = Val(false), linsolve_kwargs = (;), abstol = nothing, reltol = nothing, internalnorm::F = L2_NORM, - shared::Val{N} = Val(1), kwargs...) where {F, INV, N} - newton_cache = __internal_init(prob, alg.newton_descent, J, fu, u; pre_inverted, - linsolve_kwargs, abstol, reltol, shared, kwargs...) - cauchy_cache = __internal_init(prob, alg.steepest_descent, J, fu, u; pre_inverted, - linsolve_kwargs, abstol, reltol, shared, kwargs...) + shared::Val = Val(1), kwargs... +) where {F} + newton_cache = InternalAPI.init( + prob, alg.newton_descent, J, fu, u; + pre_inverted, linsolve_kwargs, abstol, reltol, shared, kwargs... + ) + cauchy_cache = InternalAPI.init( + prob, alg.steepest_descent, J, fu, u; + pre_inverted, linsolve_kwargs, abstol, reltol, shared, kwargs... + ) + @bb δu = similar(u) - δus = N ≤ 1 ? nothing : map(2:N) do i + δus = Utils.unwrap_val(shared) ≤ 1 ? nothing : map(2:Utils.unwrap_val(shared)) do i @bb δu_ = similar(u) end @bb δu_cache_1 = similar(u) @bb δu_cache_2 = similar(u) @bb δu_cache_mul = similar(u) - T = promote_type(eltype(u), eltype(fu)) - normal_form = prob isa NonlinearLeastSquaresProblem && - __needs_square_A(alg.newton_descent.linsolve, u) - JᵀJ_cache = !normal_form ? J * _vec(δu) : nothing # TODO: Rename + needs_square_A(alg.newton_descent.linsolve, u) + + Jᵀδu_cache = !normal_form ? J * Utils.safe_vec(δu) : nothing - return DoglegCache{INV, normal_form}(δu, δus, newton_cache, cauchy_cache, internalnorm, - JᵀJ_cache, δu_cache_1, δu_cache_2, δu_cache_mul) + return DoglegCache( + δu, δus, newton_cache, cauchy_cache, internalnorm, Jᵀδu_cache, + δu_cache_1, δu_cache_2, δu_cache_mul, pre_inverted, Val(normal_form) + ) end -# If TrustRegion is not specified, then use a Gauss-Newton step -function __internal_solve!(cache::DoglegCache{INV, NF}, J, fu, u, idx::Val{N} = Val(1); - trust_region = nothing, skip_solve::Bool = false, kwargs...) where {INV, NF, N} +# If trust_region is not specified, then use a Gauss-Newton step +function InternalAPI.solve!( + cache::DoglegCache, J, fu, u, idx::Val = Val(1); + trust_region = nothing, skip_solve::Bool = false, kwargs... +) @assert trust_region!==nothing "Trust Region must be specified for Dogleg. Use \ `NewtonDescent` or `SteepestDescent` if you don't \ want to use a Trust Region." - δu = get_du(cache, idx) + + δu = SciMLBase.get_du(cache, idx) T = promote_type(eltype(u), eltype(fu)) - δu_newton = __internal_solve!( - cache.newton_cache, J, fu, u, idx; skip_solve, kwargs...).δu + δu_newton = InternalAPI.solve!( + cache.newton_cache, J, fu, u, idx; skip_solve, kwargs... + ).δu # Newton's Step within the trust region if cache.internalnorm(δu_newton) ≤ trust_region @@ -91,23 +102,24 @@ function __internal_solve!(cache::DoglegCache{INV, NF}, J, fu, u, idx::Val{N} = return DescentResult(; δu, extras = (; δuJᵀJδu = T(NaN))) end - # Take intersection of steepest descent direction and trust region if Cauchy point lies - # outside of trust region - if NF + # Take intersection of steepest descent direction and trust region if Cauchy point + # lies outside of trust region + if normal_form(cache) δu_cauchy = cache.newton_cache.Jᵀfu_cache JᵀJ = cache.newton_cache.JᵀJ_cache @bb @. δu_cauchy *= -1 l_grad = cache.internalnorm(δu_cauchy) @bb cache.δu_cache_mul = JᵀJ × vec(δu_cauchy) - δuJᵀJδu = __dot(δu_cauchy, cache.δu_cache_mul) + δuJᵀJδu = Utils.safe_dot(cache.δu_cache_mul, cache.δu_cache_mul) else - δu_cauchy = __internal_solve!( - cache.cauchy_cache, J, fu, u, idx; skip_solve, kwargs...).δu - J_ = INV ? inv(J) : J + δu_cauchy = InternalAPI.solve!( + cache.cauchy_cache, J, fu, u, idx; skip_solve, kwargs... + ).δu + J_ = preinverted_jacobian(cache) ? inv(J) : J l_grad = cache.internalnorm(δu_cauchy) - @bb cache.JᵀJ_cache = J × vec(δu_cauchy) # TODO: Rename - δuJᵀJδu = __dot(cache.JᵀJ_cache, cache.JᵀJ_cache) + @bb cache.Jᵀδu_cache = J_ × vec(δu_cauchy) + δuJᵀJδu = Utils.safe_dot(cache.Jᵀδu_cache, cache.Jᵀδu_cache) end d_cauchy = (l_grad^3) / δuJᵀJδu @@ -125,8 +137,8 @@ function __internal_solve!(cache::DoglegCache{INV, NF}, J, fu, u, idx::Val{N} = # trust region @bb @. cache.δu_cache_1 = (d_cauchy / l_grad) * δu_cauchy @bb @. cache.δu_cache_2 = δu_newton - cache.δu_cache_1 - a = dot(cache.δu_cache_2, cache.δu_cache_2) - b = 2 * dot(cache.δu_cache_1, cache.δu_cache_2) + a = Utils.safe_dot(cache.δu_cache_2, cache.δu_cache_2) + b = 2 * Utils.safe_dot(cache.δu_cache_1, cache.δu_cache_2) c = d_cauchy^2 - trust_region^2 aux = max(0, b^2 - 4 * a * c) τ = (-b + sqrt(aux)) / (2 * a) diff --git a/src/descent/geodesic_acceleration.jl b/lib/NonlinearSolveBase/src/descent/geodesic_acceleration.jl similarity index 56% rename from src/descent/geodesic_acceleration.jl rename to lib/NonlinearSolveBase/src/descent/geodesic_acceleration.jl index 8e8a305f0..f7758b1e7 100644 --- a/src/descent/geodesic_acceleration.jl +++ b/lib/NonlinearSolveBase/src/descent/geodesic_acceleration.jl @@ -5,7 +5,7 @@ Uses the `descent` algorithm to compute the velocity and acceleration terms for geodesic acceleration method. The velocity and acceleration terms are then combined to compute the descent direction. -This method in its current form was developed for [`LevenbergMarquardt`](@ref). Performance +This method in its current form was developed for `LevenbergMarquardt`. Performance for other methods are not theorectically or experimentally verified. ### Keyword Arguments @@ -24,18 +24,12 @@ for other methods are not theorectically or experimentally verified. `α_geodesic = 0.1` is an effective choice. Defaults to `0.75`. See Section 3 of [transtrum2012improvements](@citet). """ -@concrete struct GeodesicAcceleration <: AbstractDescentAlgorithm +@concrete struct GeodesicAcceleration <: AbstractDescentDirection descent finite_diff_step_geodesic α end -function Base.show(io::IO, alg::GeodesicAcceleration) - print( - io, "GeodesicAcceleration(descent = $(alg.descent), finite_diff_step_geodesic = ", - "$(alg.finite_diff_step_geodesic), α = $(alg.α))") -end - supports_trust_region(::GeodesicAcceleration) = true get_linear_solver(alg::GeodesicAcceleration) = get_linear_solver(alg.descent) @@ -55,71 +49,77 @@ get_linear_solver(alg::GeodesicAcceleration) = get_linear_solver(alg.descent) last_step_accepted::Bool end -function __reinit_internal!( - cache::GeodesicAccelerationCache, args...; p = cache.p, kwargs...) +function InternalAPI.reinit_self!(cache::GeodesicAccelerationCache; p = cache.p, kwargs...) cache.p = p cache.last_step_accepted = false end @internal_caches GeodesicAccelerationCache :descent_cache -get_velocity(cache::GeodesicAccelerationCache) = get_du(cache.descent_cache, Val(1)) -function set_velocity!(cache::GeodesicAccelerationCache, δv) - set_du!(cache.descent_cache, δv, Val(1)) +function get_velocity(cache::GeodesicAccelerationCache) + return SciMLBase.get_du(cache.descent_cache, Val(1)) end function get_velocity(cache::GeodesicAccelerationCache, ::Val{N}) where {N} - get_du(cache.descent_cache, Val(2N - 1)) -end -function set_velocity!(cache::GeodesicAccelerationCache, δv, ::Val{N}) where {N} - set_du!(cache.descent_cache, δv, Val(2N - 1)) + return SciMLBase.get_du(cache.descent_cache, Val(2N - 1)) end -get_acceleration(cache::GeodesicAccelerationCache) = get_du(cache.descent_cache, Val(2)) -function set_acceleration!(cache::GeodesicAccelerationCache, δa) - set_du!(cache.descent_cache, δa, Val(2)) +function get_acceleration(cache::GeodesicAccelerationCache) + return SciMLBase.get_du(cache.descent_cache, Val(2)) end function get_acceleration(cache::GeodesicAccelerationCache, ::Val{N}) where {N} - get_du(cache.descent_cache, Val(2N)) -end -function set_acceleration!(cache::GeodesicAccelerationCache, δa, ::Val{N}) where {N} - set_du!(cache.descent_cache, δa, Val(2N)) + return SciMLBase.get_du(cache.descent_cache, Val(2N)) end -function __internal_init(prob::AbstractNonlinearProblem, alg::GeodesicAcceleration, J, - fu, u; shared::Val{N} = Val(1), pre_inverted::Val{INV} = False, - linsolve_kwargs = (;), abstol = nothing, reltol = nothing, - internalnorm::F = L2_NORM, kwargs...) where {INV, N, F} +function InternalAPI.init( + prob::AbstractNonlinearProblem, alg::GeodesicAcceleration, J, fu, u; + shared::Val = Val(1), pre_inverted::Val = Val(false), linsolve_kwargs = (;), + abstol = nothing, reltol = nothing, + internalnorm::F = L2_NORM, kwargs... +) where {F} T = promote_type(eltype(u), eltype(fu)) @bb δu = similar(u) - δus = N ≤ 1 ? nothing : map(2:N) do i + δus = Utils.unwrap_val(shared) ≤ 1 ? nothing : map(2:Utils.unwrap_val(shared)) do i @bb δu_ = similar(u) end - descent_cache = __internal_init(prob, alg.descent, J, fu, u; shared = Val(N * 2), - pre_inverted, linsolve_kwargs, abstol, reltol, kwargs...) + descent_cache = InternalAPI.init( + prob, alg.descent, J, fu, u; + shared = Val(2 * Utils.unwrap_val(shared)), pre_inverted, linsolve_kwargs, + abstol, reltol, + kwargs... + ) @bb Jv = similar(fu) @bb fu_cache = copy(fu) @bb u_cache = similar(u) return GeodesicAccelerationCache( δu, δus, descent_cache, prob.f, prob.p, T(alg.α), internalnorm, - T(alg.finite_diff_step_geodesic), Jv, fu_cache, u_cache, false) + 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); - skip_solve::Bool = false, kwargs...) where {N} - a, v, δu = get_acceleration(cache, idx), get_velocity(cache, idx), get_du(cache, idx) +function InternalAPI.solve!( + cache::GeodesicAccelerationCache, J, fu, u, idx::Val = Val(1); + skip_solve::Bool = false, kwargs... +) + a = get_acceleration(cache, idx) + v = get_velocity(cache, idx) + δu = SciMLBase.get_du(cache, idx) skip_solve && return DescentResult(; δu, extras = (; a, v)) - v = __internal_solve!( - cache.descent_cache, J, fu, u, Val(2N - 1); skip_solve, kwargs...).δu + + v = InternalAPI.solve!( + cache.descent_cache, J, fu, u, Val(2 * Utils.unwrap_val(idx) - 1); + skip_solve, kwargs... + ).δu @bb @. cache.u_cache = u + cache.h * v - cache.fu_cache = evaluate_f!!(cache.f, cache.fu_cache, cache.u_cache, cache.p) + cache.fu_cache = Utils.evaluate_f!!(cache.f, cache.fu_cache, cache.u_cache, cache.p) J !== nothing && @bb(cache.Jv=J × vec(v)) - Jv = _restructure(cache.fu_cache, cache.Jv) + Jv = Utils.restructure(cache.fu_cache, cache.Jv) @bb @. cache.fu_cache = (2 / cache.h) * ((cache.fu_cache - fu) / cache.h - Jv) - a = __internal_solve!(cache.descent_cache, J, cache.fu_cache, u, Val(2N); - skip_solve, kwargs..., reuse_A_if_factorization = true).δu + a = InternalAPI.solve!( + cache.descent_cache, J, cache.fu_cache, u, Val(2 * Utils.unwrap_val(idx)); + skip_solve, kwargs..., reuse_A_if_factorization = true + ).δu norm_v = cache.internalnorm(v) norm_a = cache.internalnorm(a) diff --git a/lib/NonlinearSolveBase/src/descent/newton.jl b/lib/NonlinearSolveBase/src/descent/newton.jl new file mode 100644 index 000000000..810661507 --- /dev/null +++ b/lib/NonlinearSolveBase/src/descent/newton.jl @@ -0,0 +1,131 @@ +""" + NewtonDescent(; linsolve = nothing, precs = nothing) + +Compute the descent direction as ``J δu = -fu``. For non-square Jacobian problems, this is +commonly referred to as the Gauss-Newton Descent. + +See also [`Dogleg`](@ref), [`SteepestDescent`](@ref), [`DampedNewtonDescent`](@ref). +""" +@kwdef @concrete struct NewtonDescent <: AbstractDescentDirection + linsolve = nothing + precs = nothing +end + +supports_line_search(::NewtonDescent) = true + +@concrete mutable struct NewtonDescentCache <: AbstractDescentCache + δu + δus + lincache + JᵀJ_cache # For normal form else nothing + Jᵀfu_cache + timer + preinverted_jacobian <: Union{Val{false}, Val{true}} + normal_form <: Union{Val{false}, Val{true}} +end + +@internal_caches NewtonDescentCache :lincache + +function InternalAPI.init( + prob::AbstractNonlinearProblem, alg::NewtonDescent, J, fu, u; stats, + shared = Val(1), pre_inverted::Val = Val(false), linsolve_kwargs = (;), + abstol = nothing, reltol = nothing, + timer = get_timer_output(), kwargs... +) + @bb δu = similar(u) + δus = Utils.unwrap_val(shared) ≤ 1 ? nothing : map(2:Utils.unwrap_val(shared)) do i + @bb δu_ = similar(u) + end + if Utils.unwrap_val(pre_inverted) + lincache = nothing + else + lincache = construct_linear_solver( + alg, alg.linsolve, J, Utils.safe_vec(fu), Utils.safe_vec(u); + stats, abstol, reltol, linsolve_kwargs... + ) + end + return NewtonDescentCache( + δu, δus, lincache, nothing, nothing, timer, pre_inverted, Val(false) + ) +end + +function InternalAPI.init( + prob::NonlinearLeastSquaresProblem, alg::NewtonDescent, J, fu, u; stats, + shared = Val(1), pre_inverted::Val = Val(false), linsolve_kwargs = (;), + abstol = nothing, reltol = nothing, + timer = get_timer_output(), kwargs... +) + length(fu) != length(u) && + @assert !Utils.unwrap_val(pre_inverted) "Precomputed Inverse for Non-Square Jacobian doesn't make sense." + + @bb δu = similar(u) + δus = Utils.unwrap_val(shared) ≤ 1 ? nothing : map(2:N) do i + @bb δu_ = similar(u) + end + + normal_form = needs_square_A(alg.linsolve, u) + if normal_form + JᵀJ = transpose(J) * J + Jᵀfu = transpose(J) * Utils.safe_vec(fu) + A, b = Utils.maybe_symmetric(JᵀJ), Jᵀfu + else + JᵀJ, Jᵀfu = nothing, nothing + A, b = J, Utils.safe_vec(fu) + end + + lincache = construct_linear_solver( + alg, alg.linsolve, A, b, Utils.safe_vec(u); + stats, abstol, reltol, linsolve_kwargs... + ) + + return NewtonDescentCache( + δu, δus, lincache, JᵀJ, Jᵀfu, timer, pre_inverted, Val(normal_form) + ) +end + +function InternalAPI.solve!( + cache::NewtonDescentCache, J, fu, u, idx::Val = Val(1); + skip_solve::Bool = false, new_jacobian::Bool = true, kwargs... +) + δu = SciMLBase.get_du(cache, idx) + skip_solve && return DescentResult(; δu) + + if preinverted_jacobian(cache) && !normal_form(cache) + @assert J!==nothing "`J` must be provided when `preinverted_jacobian = Val(true)`." + @bb δu = J × vec(fu) + else + if normal_form(cache) + @assert !preinverted_jacobian(cache) + if idx === Val(1) + @bb cache.JᵀJ_cache = transpose(J) × J + end + @bb cache.Jᵀfu_cache = transpose(J) × vec(fu) + @static_timeit cache.timer "linear solve" begin + linres = cache.lincache(; + A = Utils.maybe_symmetric(cache.JᵀJ_cache), b = cache.Jᵀfu_cache, + kwargs..., linu = Utils.safe_vec(δu), du = Utils.safe_vec(δu), + reuse_A_if_factorization = !new_jacobian || (idx !== Val(1)) + ) + end + else + @static_timeit cache.timer "linear solve" begin + linres = cache.lincache(; + A = J, b = Utils.safe_vec(fu), + kwargs..., linu = Utils.safe_vec(δu), du = Utils.safe_vec(δu), + reuse_A_if_factorization = !new_jacobian || idx !== Val(1) + ) + end + end + @static_timeit cache.timer "linear solve" begin + δu = Utils.restructure(SciMLBase.get_du(cache, idx), linres.u) + if !linres.success + set_du!(cache, δu, idx) + return DescentResult(; δu, success = false, linsolve_success = false) + end + end + end + + @bb @. δu *= -1 + set_du!(cache, δu, idx) + return DescentResult(; δu) +end diff --git a/lib/NonlinearSolveBase/src/descent/steepest.jl b/lib/NonlinearSolveBase/src/descent/steepest.jl new file mode 100644 index 000000000..deb1a9817 --- /dev/null +++ b/lib/NonlinearSolveBase/src/descent/steepest.jl @@ -0,0 +1,75 @@ +""" + SteepestDescent(; linsolve = nothing, precs = nothing) + +Compute the descent direction as ``δu = -Jᵀfu``. The linear solver and preconditioner are +only used if `J` is provided in the inverted form. + +See also [`Dogleg`](@ref), [`NewtonDescent`](@ref), [`DampedNewtonDescent`](@ref). +""" +@kwdef @concrete struct SteepestDescent <: AbstractDescentDirection + linsolve = nothing + precs = nothing +end + +supports_line_search(::SteepestDescent) = true + +@concrete mutable struct SteepestDescentCache <: AbstractDescentCache + δu + δus + lincache + timer + preinverted_jacobian <: Union{Val{false}, Val{true}} +end + +@internal_caches SteepestDescentCache :lincache + +function InternalAPI.init( + prob::AbstractNonlinearProblem, alg::SteepestDescent, J, fu, u; + stats, shared = Val(1), pre_inverted::Val = Val(false), linsolve_kwargs = (;), + abstol = nothing, reltol = nothing, + timer = get_timer_output(), + kwargs... +) + if Utils.unwrap_val(pre_inverted) + @assert length(fu)==length(u) "Non-Square Jacobian Inverse doesn't make sense." + end + @bb δu = similar(u) + δus = Utils.unwrap_val(shared) ≤ 1 ? nothing : map(2:Utils.unwrap_val(shared)) do i + @bb δu_ = similar(u) + end + if Utils.unwrap_val(pre_inverted) + lincache = construct_linear_solver( + alg, alg.linsolve, transpose(J), Utils.safe_vec(fu), Utils.safe_vec(u); + stats, abstol, reltol, linsolve_kwargs... + ) + else + lincache = nothing + end + return SteepestDescentCache(δu, δus, lincache, timer, pre_inverted) +end + +function InternalAPI.solve!( + cache::SteepestDescentCache, J, fu, u, idx::Val = Val(1); + new_jacobian::Bool = true, kwargs... +) + δu = SciMLBase.get_du(cache, idx) + if Utils.unwrap_val(cache.preinverted_jacobian) + A = J === nothing ? nothing : transpose(J) + linres = cache.lincache(; + A, b = Utils.safe_vec(fu), kwargs..., linu = Utils.safe_vec(δu), + du = Utils.safe_vec(δu), + reuse_A_if_factorization = !new_jacobian || idx !== Val(1) + ) + δu = Utils.restructure(SciMLBase.get_du(cache, idx), linres.u) + if !linres.success + set_du!(cache, δu, idx) + return DescentResult(; δu, success = false, linsolve_success = false) + end + else + @assert J!==nothing "`J` must be provided when `preinverted_jacobian = Val(false)`." + @bb δu = transpose(J) × vec(fu) + end + @bb @. δu *= -1 + set_du!(cache, δu, idx) + return DescentResult(; δu) +end diff --git a/lib/NonlinearSolveBase/src/immutable_problem.jl b/lib/NonlinearSolveBase/src/immutable_problem.jl index 03b44a9ec..2d429845f 100644 --- a/lib/NonlinearSolveBase/src/immutable_problem.jl +++ b/lib/NonlinearSolveBase/src/immutable_problem.jl @@ -6,14 +6,14 @@ struct ImmutableNonlinearProblem{uType, iip, P, F, K, PT} <: problem_type::PT kwargs::K - @add_kwonly function ImmutableNonlinearProblem{iip}( + SciMLBase.@add_kwonly function ImmutableNonlinearProblem{iip}( f::AbstractNonlinearFunction{iip}, u0, p = NullParameters(), problem_type = StandardNonlinearProblem(); kwargs...) where {iip} if haskey(kwargs, :p) error("`p` specified as a keyword argument `p = $(kwargs[:p])` to \ `NonlinearProblem`. This is not supported.") end - warn_paramtype(p) + SciMLBase.warn_paramtype(p) return new{ typeof(u0), iip, typeof(p), typeof(f), typeof(kwargs), typeof(problem_type)}( f, u0, p, problem_type, kwargs) @@ -31,12 +31,11 @@ struct ImmutableNonlinearProblem{uType, iip, P, F, K, PT} <: end """ -Define a nonlinear problem using an instance of -[`AbstractNonlinearFunction`](@ref AbstractNonlinearFunction). +Define a nonlinear problem using an instance of [`AbstractNonlinearFunction`](@ref). """ function ImmutableNonlinearProblem( f::AbstractNonlinearFunction, u0, p = NullParameters(); kwargs...) - return ImmutableNonlinearProblem{isinplace(f)}(f, u0, p; kwargs...) + return ImmutableNonlinearProblem{SciMLBase.isinplace(f)}(f, u0, p; kwargs...) end function ImmutableNonlinearProblem(f, u0, p = NullParameters(); kwargs...) @@ -44,14 +43,14 @@ function ImmutableNonlinearProblem(f, u0, p = NullParameters(); kwargs...) end """ -Define a ImmutableNonlinearProblem problem from SteadyStateProblem +Define a ImmutableNonlinearProblem problem from SteadyStateProblem. """ function ImmutableNonlinearProblem(prob::AbstractNonlinearProblem) - return ImmutableNonlinearProblem{isinplace(prob)}(prob.f, prob.u0, prob.p) + return ImmutableNonlinearProblem{SciMLBase.isinplace(prob)}(prob.f, prob.u0, prob.p) end function Base.convert( ::Type{ImmutableNonlinearProblem}, prob::T) where {T <: NonlinearProblem} - return ImmutableNonlinearProblem{isinplace(prob)}( + return ImmutableNonlinearProblem{SciMLBase.isinplace(prob)}( prob.f, prob.u0, prob.p, prob.problem_type; prob.kwargs...) end diff --git a/src/internal/jacobian.jl b/lib/NonlinearSolveBase/src/jacobian.jl similarity index 55% rename from src/internal/jacobian.jl rename to lib/NonlinearSolveBase/src/jacobian.jl index 70cc7021e..6ab31efd2 100644 --- a/src/internal/jacobian.jl +++ b/lib/NonlinearSolveBase/src/jacobian.jl @@ -1,6 +1,10 @@ + """ - JacobianCache(prob, alg, f::F, fu, u, p; autodiff = nothing, - vjp_autodiff = nothing, jvp_autodiff = nothing, linsolve = missing) where {F} + construct_jacobian_cache( + prob, alg, f, fu, u = prob.u0, p = prob.p; + autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing, + linsolve = missing + ) Construct a cache for the Jacobian of `f` w.r.t. `u`. @@ -25,43 +29,33 @@ 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 `SciMLJacobianOperators.JacobianOperator` instead. + or if possible we can just use a [`SciMLJacobianOperators.JacobianOperator`](@ref) + instead. """ -@concrete mutable struct JacobianCache{iip} <: AbstractNonlinearSolveJacobianCache{iip} - J - f - fu - u - p - stats::NLStats - autodiff - di_extras -end - -function reinit_cache!(cache::JacobianCache{iip}, args...; p = cache.p, - u0 = cache.u, kwargs...) where {iip} - cache.u = u0 - cache.p = p -end - -function JacobianCache(prob, alg, f::F, fu_, u, p; stats, autodiff = nothing, - vjp_autodiff = nothing, jvp_autodiff = nothing, linsolve = missing) where {F} - iip = isinplace(prob) - +function construct_jacobian_cache( + prob, alg, f::NonlinearFunction, fu, u = prob.u0, p = prob.p; stats, + autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing, + linsolve = missing +) has_analytic_jac = SciMLBase.has_jac(f) - linsolve_needs_jac = concrete_jac(alg) === nothing && (linsolve === missing || - (linsolve === nothing || __needs_concrete_A(linsolve))) - alg_wants_jac = concrete_jac(alg) !== nothing && concrete_jac(alg) - needs_jac = linsolve_needs_jac || alg_wants_jac + linsolve_needs_jac = !concrete_jac(alg) && (linsolve === missing || + (linsolve === nothing || needs_concrete_A(linsolve))) + needs_jac = linsolve_needs_jac || concrete_jac(alg) - @bb fu = similar(fu_) + @bb fu_cache = similar(fu) if !has_analytic_jac && needs_jac + if autodiff === nothing + throw(ArgumentError("`autodiff` argument to `construct_jacobian_cache` must be \ + specified and cannot be `nothing`. Use \ + `NonlinearSolveBase.select_jacobian_autodiff` for \ + automatic backend selection.")) + end autodiff = construct_concrete_adtype(f, autodiff) - di_extras = if iip - DI.prepare_jacobian(f, fu, autodiff, u, Constant(prob.p)) + di_extras = if SciMLBase.isinplace(f) + DI.prepare_jacobian(f, fu_cache, autodiff, u, Constant(p)) else - DI.prepare_jacobian(f, autodiff, u, Constant(prob.p)) + DI.prepare_jacobian(f, autodiff, u, Constant(p)) end else di_extras = nothing @@ -75,11 +69,12 @@ function JacobianCache(prob, alg, f::F, fu_, u, p; stats, autodiff = nothing, # which is needed to create the linear solver cache stats.njacs += 1 if has_analytic_jac - __similar( - fu, promote_type(eltype(fu), eltype(u)), length(fu), length(u)) + Utils.safe_similar( + fu, promote_type(eltype(fu), eltype(u)), length(fu), length(u) + ) else - if iip - DI.jacobian(f, fu, di_extras, autodiff, u, Constant(p)) + if SciMLBase.isinplace(f) + DI.jacobian(f, fu_cache, di_extras, autodiff, u, Constant(p)) else DI.jacobian(f, di_extras, autodiff, u, Constant(p)) end @@ -93,65 +88,96 @@ function JacobianCache(prob, alg, f::F, fu_, u, p; stats, autodiff = nothing, end end - return JacobianCache{iip}(J, f, fu, u, p, stats, autodiff, di_extras) + return JacobianCache(J, f, fu, u, p, stats, autodiff, di_extras) end -function JacobianCache(prob, alg, f::F, ::Number, u::Number, p; stats, - autodiff = nothing, kwargs...) where {F} - fu = f(u, p) +function construct_jacobian_cache( + prob, alg, f::NonlinearFunction, fu::Number, u::Number = prob.u0, p = prob.p; stats, + autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing, + linsolve = missing +) if SciMLBase.has_jac(f) || SciMLBase.has_vjp(f) || SciMLBase.has_jvp(f) - return JacobianCache{false}(u, f, fu, u, p, stats, autodiff, nothing) + return JacobianCache(u, f, fu, u, p, stats, autodiff, nothing) + end + if autodiff === nothing + throw(ArgumentError("`autodiff` argument to `construct_jacobian_cache` must be \ + specified and cannot be `nothing`. Use \ + `NonlinearSolveBase.select_jacobian_autodiff` for \ + automatic backend selection.")) end - di_extras = DI.prepare_derivative(f, get_dense_ad(autodiff), u, Constant(prob.p)) - return JacobianCache{false}(u, f, fu, u, p, stats, autodiff, di_extras) + @assert !(autodiff isa AutoSparse) "`autodiff` cannot be `AutoSparse` for scalar \ + nonlinear problems." + di_extras = DI.prepare_derivative(f, autodiff, u, Constant(prob.p)) + return JacobianCache(u, f, fu, u, p, stats, autodiff, di_extras) end -(cache::JacobianCache)(u = cache.u) = cache(cache.J, u, cache.p) -function (cache::JacobianCache)(::Nothing) - cache.J isa JacobianOperator && - return StatefulJacobianOperator(cache.J, cache.u, cache.p) - return cache.J +@concrete mutable struct JacobianCache <: AbstractJacobianCache + J + f <: NonlinearFunction + fu + u + p + stats::NLStats + autodiff + di_extras +end + +function InternalAPI.reinit!(cache::JacobianCache; p = cache.p, u0 = cache.u, kwargs...) + cache.u = u0 + cache.p = p end -# Operator -function (cache::JacobianCache)(J::JacobianOperator, u, p = cache.p) +# Core Computation +(cache::JacobianCache)(u) = cache(cache.J, u, cache.p) +function (cache::JacobianCache{<:JacobianOperator})(::Nothing) + return StatefulJacobianOperator(cache.J, cache.u, cache.p) +end +(cache::JacobianCache)(::Nothing) = cache.J + +## Operator +function (cache::JacobianCache{<:JacobianOperator})(J::JacobianOperator, u, p = cache.p) return StatefulJacobianOperator(J, u, p) end -# Numbers -function (cache::JacobianCache)(::Number, u, p = cache.p) # Scalar + +## Numbers +function (cache::JacobianCache{<:Number})(::Number, u, p = cache.p) cache.stats.njacs += 1 - if SciMLBase.has_jac(cache.f) - return cache.f.jac(u, p) + cache.J = if SciMLBase.has_jac(cache.f) + cache.f.jac(u, p) elseif SciMLBase.has_vjp(cache.f) - return cache.f.vjp(one(u), u, p) + cache.f.vjp(one(u), u, p) elseif SciMLBase.has_jvp(cache.f) - return cache.f.jvp(one(u), u, p) + cache.f.jvp(one(u), u, p) + else + DI.derivative(cache.f, cache.di_extras, cache.autodiff, u, Constant(p)) end - return DI.derivative(cache.f, cache.di_extras, cache.autodiff, u, Constant(p)) + return cache.J end -# Actually Compute the Jacobian -function (cache::JacobianCache{iip})( - J::Union{AbstractMatrix, Nothing}, u, p = cache.p) where {iip} + +## Actually Compute the Jacobian +function (cache::JacobianCache)(J::Union{AbstractMatrix, Nothing}, u, p = cache.p) cache.stats.njacs += 1 - if iip + if SciMLBase.isinplace(cache.f) if SciMLBase.has_jac(cache.f) cache.f.jac(J, u, p) else DI.jacobian!( - cache.f, cache.fu, J, cache.di_extras, cache.autodiff, u, Constant(p)) + cache.f, cache.fu, J, cache.di_extras, cache.autodiff, u, Constant(p) + ) end return J else if SciMLBase.has_jac(cache.f) - return cache.f.jac(u, p) + cache.J = cache.f.jac(u, p) else - return DI.jacobian(cache.f, cache.di_extras, cache.autodiff, u, Constant(p)) + cache.J = DI.jacobian(cache.f, cache.di_extras, cache.autodiff, u, Constant(p)) end + return cache.J end end +# Sparse Automatic Differentiation function construct_concrete_adtype(f::NonlinearFunction, ad::AbstractADType) - @assert !(ad isa AutoSparse) "This shouldn't happen. Open an issue." if f.sparsity === nothing if f.jac_prototype === nothing if SciMLBase.has_colorvec(f) @@ -167,11 +193,12 @@ function construct_concrete_adtype(f::NonlinearFunction, ad::AbstractADType) end return ad end + coloring_algorithm = select_fastest_coloring_algorithm(f.jac_prototype, f, ad) + coloring_algorithm === nothing && return ad return AutoSparse( ad; sparsity_detector = KnownJacobianSparsityDetector(f.jac_prototype), - coloring_algorithm = select_fastest_coloring_algorithm( - f.jac_prototype, f, ad) + coloring_algorithm ) end else @@ -184,11 +211,12 @@ function construct_concrete_adtype(f::NonlinearFunction, ad::AbstractADType) provided. Pass only `jac_prototype`.")) end end + coloring_algorithm = select_fastest_coloring_algorithm(f.sparsity, f, ad) + coloring_algorithm === nothing && return ad return AutoSparse( ad; sparsity_detector = KnownJacobianSparsityDetector(f.sparsity), - coloring_algorithm = select_fastest_coloring_algorithm( - f.sparsity, f, ad) + coloring_algorithm ) end @@ -199,11 +227,9 @@ function construct_concrete_adtype(f::NonlinearFunction, ad::AbstractADType) @warn "`colorvec` is provided but `jac_prototype` is not specified. \ `colorvec` will be ignored." end - return AutoSparse( - ad; - sparsity_detector, - coloring_algorithm = GreedyColoringAlgorithm(LargestFirst()) - ) + coloring_algorithm = select_fastest_coloring_algorithm(nothing, f, ad) + coloring_algorithm === nothing && return ad + return AutoSparse(ad; sparsity_detector, coloring_algorithm) else if sparse_or_structured_prototype(f.jac_prototype) if !(sparsity_detector isa NoSparsityDetector) @@ -213,27 +239,13 @@ function construct_concrete_adtype(f::NonlinearFunction, ad::AbstractADType) end sparsity_detector = KnownJacobianSparsityDetector(f.jac_prototype) end - - return AutoSparse( - ad; - sparsity_detector, - coloring_algorithm = select_fastest_coloring_algorithm( - f.jac_prototype, f, ad) - ) + coloring_algorithm = select_fastest_coloring_algorithm(f.jac_prototype, f, ad) + coloring_algorithm === nothing && return ad + return AutoSparse(ad; sparsity_detector, coloring_algorithm) end end end -function select_fastest_coloring_algorithm( - prototype, f::NonlinearFunction, ad::AbstractADType) - if SciMLBase.has_colorvec(f) - return ConstantColoringAlgorithm{ifelse( - ADTypes.mode(ad) isa ADTypes.ReverseMode, :row, :column)}( - prototype, f.colorvec) - end - return GreedyColoringAlgorithm(LargestFirst()) -end - function construct_concrete_adtype(::NonlinearFunction, ad::AutoSparse) error("Specifying a sparse AD type for Nonlinear Problems was removed in v4. \ Instead use the `sparsity`, `jac_prototype`, and `colorvec` to specify \ @@ -241,10 +253,16 @@ function construct_concrete_adtype(::NonlinearFunction, ad::AutoSparse) detection algorithm and coloring algorithm present in $(ad).") end -get_dense_ad(ad) = ad -get_dense_ad(ad::AutoSparse) = ADTypes.dense_ad(ad) +function select_fastest_coloring_algorithm( + prototype, f::NonlinearFunction, ad::AbstractADType) + if !Utils.is_extension_loaded(Val(:SparseMatrixColorings)) + @warn "`SparseMatrixColorings` must be explicitly imported for sparse automatic \ + differentiation to work. Proceeding with Dense Automatic Differentiation." + return nothing + end + return select_fastest_coloring_algorithm(Val(:SparseMatrixColorings), prototype, f, ad) +end -sparse_or_structured_prototype(::AbstractSparseMatrix) = true function sparse_or_structured_prototype(prototype::AbstractMatrix) return ArrayInterface.isstructured(prototype) end diff --git a/lib/NonlinearSolveBase/src/linear_solve.jl b/lib/NonlinearSolveBase/src/linear_solve.jl new file mode 100644 index 000000000..f3a2338c9 --- /dev/null +++ b/lib/NonlinearSolveBase/src/linear_solve.jl @@ -0,0 +1,142 @@ +@kwdef @concrete struct LinearSolveResult + u + success::Bool = true +end + +@concrete mutable struct LinearSolveJLCache <: AbstractLinearSolverCache + lincache + linsolve + additional_lincache::Any + precs + stats::NLStats +end + +@concrete mutable struct NativeJLLinearSolveCache <: AbstractLinearSolverCache + A + b + stats::NLStats +end + +""" + construct_linear_solver(alg, linsolve, A, b, u; stats, kwargs...) + +Construct a cache for solving linear systems of the form `A * u = b`. Following cases are +handled: + + 1. `A` is Number, then we solve it with `u = b / A` + 2. `A` is `SMatrix`, then we solve it with `u = A \\ b` (using the defaults from base + Julia) (unless a preconditioner is specified) + 3. If `linsolve` is `\\`, then we solve it with directly using `ldiv!(u, A, b)` + 4. In all other cases, we use `alg` to solve the linear system using + [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl) + +### Solving the System + +```julia +(cache::LinearSolverCache)(; + A = nothing, b = nothing, linu = nothing, du = nothing, p = nothing, + weight = nothing, cachedata = nothing, reuse_A_if_factorization = false, kwargs...) +``` + +Returns the solution of the system `u` and stores the updated cache in `cache.lincache`. + +#### Special Handling for Rank-deficient Matrix `A` + +If we detect a failure in the linear solve (mostly due to using an algorithm that doesn't +support rank-deficient matrices), we emit a warning and attempt to solve the problem using +Pivoted QR factorization. This is quite efficient if there are only a few rank-deficient +that originate in the problem. However, if these are quite frequent for the main nonlinear +system, then it is recommended to use a different linear solver that supports rank-deficient +matrices. + +#### Keyword Arguments + + - `reuse_A_if_factorization`: If `true`, then the factorization of `A` is reused if + possible. This is useful when solving the same system with different `b` values. + If the algorithm is an iterative solver, then we reset the internal linear solve cache. + +One distinct feature of this compared to the cache from LinearSolve is that it respects the +aliasing arguments even after cache construction, i.e., if we passed in an `A` that `A` is +not mutated, we do this by copying over `A` to a preconstructed cache. +""" +function construct_linear_solver(alg, linsolve, A, b, u; stats, kwargs...) + no_preconditioner = !hasfield(typeof(alg), :precs) || alg.precs === nothing + + if (A isa Number && b isa Number) || (A isa Diagonal) + return NativeJLLinearSolveCache(A, b, stats) + elseif linsolve isa typeof(\) + !no_preconditioner && + error("Default Julia Backsolve Operator `\\` doesn't support Preconditioners") + return NativeJLLinearSolveCache(A, b, stats) + elseif no_preconditioner && linsolve === nothing + if (A isa SMatrix || A isa WrappedArray{<:Any, <:SMatrix}) + return NativeJLLinearSolveCache(A, b, stats) + end + end + + u_fixed = fix_incompatible_linsolve_arguments(A, b, u) + @bb u_cache = copy(u_fixed) + linprob = LinearProblem(A, b; u0 = u_cache, kwargs...) + + if no_preconditioner + precs, Pl, Pr = nothing, nothing, nothing + else + precs = alg.precs + Pl, Pr = precs(A, nothing, u, ntuple(Returns(nothing), 6)...) + end + Pl, Pr = wrap_preconditioners(Pl, Pr, u) + + # unlias here, we will later use these as caches + lincache = init(linprob, linsolve; alias_A = false, alias_b = false, Pl, Pr) + return LinearSolveJLCache(lincache, linsolve, nothing, precs, stats) +end + +function (cache::NativeJLLinearSolveCache)(; + A = nothing, b = nothing, linu = nothing, kwargs...) + cache.stats.nsolve += 1 + cache.stats.nfactors += 1 + + A === nothing || (cache.A = A) + b === nothing || (cache.b = b) + + if linu !== nothing && ArrayInterface.can_setindex(linu) && + applicable(ldiv!, linu, cache.A, cache.b) && applicable(ldiv!, cache.A, linu) + ldiv!(linu, cache.A, cache.b) + res = linu + else + res = cache.A \ cache.b + end + return LinearSolveResult(; u = res) +end + +fix_incompatible_linsolve_arguments(A, b, u) = u +fix_incompatible_linsolve_arguments(::SArray, ::SArray, u::SArray) = u +function fix_incompatible_linsolve_arguments(A, b, u::SArray) + (Core.Compiler.return_type(\, Tuple{typeof(A), typeof(b)}) <: typeof(u)) && return u + @warn "Solving Linear System A::$(typeof(A)) x::$(typeof(u)) = b::$(typeof(u)) is not \ + properly supported. Converting `x` to a mutable array. Check the return type \ + of the nonlinear function provided for optimal performance." maxlog=1 + return MArray(u) +end + +set_lincache_u!(cache, u) = setproperty!(cache.lincache, :u, u) +function set_lincache_u!(cache, u::SArray) + cache.lincache.u isa MArray && return set_lincache_u!(cache, MArray(u)) + cache.lincache.u = u +end + +function wrap_preconditioners(Pl, Pr, u) + Pl = Pl === nothing ? IdentityOperator(length(u)) : Pl + Pr = Pr === nothing ? IdentityOperator(length(u)) : Pr + return Pl, Pr +end + +# Traits. Core traits are expanded in LinearSolve extension +needs_square_A(::Any, ::Number) = false +needs_square_A(::Nothing, ::Number) = false +needs_square_A(::Nothing, ::Any) = false +needs_square_A(::typeof(\), ::Number) = false +needs_square_A(::typeof(\), ::Any) = false + +needs_concrete_A(::Union{Nothing, Missing}) = false +needs_concrete_A(::typeof(\)) = true diff --git a/lib/NonlinearSolveBase/src/public.jl b/lib/NonlinearSolveBase/src/public.jl index b39aa26d2..2101f4274 100644 --- a/lib/NonlinearSolveBase/src/public.jl +++ b/lib/NonlinearSolveBase/src/public.jl @@ -95,9 +95,11 @@ for norm_type in (:RelNorm, :AbsNorm), safety in (:Safe, :SafeBest) ## Constructor - $($struct_name)(internalnorm; protective_threshold = nothing, + $($struct_name)( + internalnorm; protective_threshold = nothing, patience_steps = 100, patience_objective_multiplier = 3, - min_max_factor = 1.3, max_stalled_steps = nothing) + min_max_factor = 1.3, max_stalled_steps = nothing + ) $($TERM_INTERNALNORM_DOCS). """ diff --git a/lib/NonlinearSolveBase/src/solve.jl b/lib/NonlinearSolveBase/src/solve.jl new file mode 100644 index 000000000..08b60e4db --- /dev/null +++ b/lib/NonlinearSolveBase/src/solve.jl @@ -0,0 +1,108 @@ +function SciMLBase.__solve( + prob::AbstractNonlinearProblem, alg::AbstractNonlinearSolveAlgorithm, args...; + kwargs... +) + cache = SciMLBase.__init(prob, alg, args...; kwargs...) + return CommonSolve.solve!(cache) +end + +function CommonSolve.solve!(cache::AbstractNonlinearSolveCache) + while not_terminated(cache) + CommonSolve.step!(cache) + end + + # The solver might have set a different `retcode` + if cache.retcode == ReturnCode.Default + cache.retcode = ifelse( + cache.nsteps ≥ cache.maxiters, ReturnCode.MaxIters, ReturnCode.Success + ) + end + + update_from_termination_cache!(cache.termination_cache, cache) + + update_trace!( + cache.trace, cache.nsteps, get_u(cache), get_fu(cache), nothing, nothing, nothing; + last = Val(true) + ) + + return SciMLBase.build_solution( + cache.prob, cache.alg, get_u(cache), get_fu(cache); + cache.retcode, cache.stats, cache.trace + ) +end + +""" + step!( + cache::AbstractNonlinearSolveCache; + recompute_jacobian::Union{Nothing, Bool} = nothing + ) + +Performs one step of the nonlinear solver. + +### Keyword Arguments + + - `recompute_jacobian`: allows controlling whether the jacobian is recomputed at the + current step. If `nothing`, then the algorithm determines whether to recompute the + jacobian. If `true` or `false`, then the jacobian is recomputed or not recomputed, + respectively. For algorithms that don't use jacobian information, this keyword is + ignored with a one-time warning. +""" +function CommonSolve.step!(cache::AbstractNonlinearSolveCache, args...; kwargs...) + not_terminated(cache) || return + + has_time_limit(cache) && (time_start = time()) + + res = @static_timeit cache.timer "solve" begin + InternalAPI.step!(cache, args...; kwargs...) + end + + cache.stats.nsteps += 1 + cache.nsteps += 1 + + if has_time_limit(cache) + cache.total_time += time() - time_start + + if !cache.force_stop && cache.retcode == ReturnCode.Default && + cache.total_time ≥ cache.maxtime + cache.retcode = ReturnCode.MaxTime + cache.force_stop = true + end + end + + return res +end + +# Some algorithms don't support creating a cache and doing `solve!`, this unfortunately +# makes it difficult to write generic code that supports caching. For the algorithms that +# don't have a `__init` function defined, we create a "Fake Cache", which just calls +# `__solve` from `solve!` +# Warning: This doesn't implement all the necessary interface functions +@concrete mutable struct NonlinearSolveNoInitCache <: AbstractNonlinearSolveCache + prob + alg + args + kwargs::Any +end + +function SciMLBase.reinit!( + cache::NonlinearSolveNoInitCache, u0 = cache.prob.u0; p = cache.prob.p, kwargs... +) + cache.prob = SciMLBase.remake(cache.prob; u0, p) + cache.kwargs = merge(cache.kwargs, kwargs) + return cache +end + +function Base.show(io::IO, ::MIME"text/plain", cache::NonlinearSolveNoInitCache) + print(io, "NonlinearSolveNoInitCache(alg = $(cache.alg))") +end + +function SciMLBase.__init( + prob::AbstractNonlinearProblem, alg::AbstractNonlinearSolveAlgorithm, args...; + kwargs... +) + return NonlinearSolveNoInitCache(prob, alg, args, kwargs) +end + +function CommonSolve.solve!(cache::NonlinearSolveNoInitCache) + return CommonSolve.solve(cache.prob, cache.alg, cache.args...; cache.kwargs...) +end diff --git a/lib/NonlinearSolveBase/src/termination_conditions.jl b/lib/NonlinearSolveBase/src/termination_conditions.jl index 9f20e46bc..cca9134d1 100644 --- a/lib/NonlinearSolveBase/src/termination_conditions.jl +++ b/lib/NonlinearSolveBase/src/termination_conditions.jl @@ -1,7 +1,9 @@ const RelNormModes = Union{ - RelNormTerminationMode, RelNormSafeTerminationMode, RelNormSafeBestTerminationMode} + RelNormTerminationMode, RelNormSafeTerminationMode, RelNormSafeBestTerminationMode +} const AbsNormModes = Union{ - AbsNormTerminationMode, AbsNormSafeTerminationMode, AbsNormSafeBestTerminationMode} + AbsNormTerminationMode, AbsNormSafeTerminationMode, AbsNormSafeBestTerminationMode +} # Core Implementation @concrete mutable struct NonlinearTerminationModeCache{uType, T} @@ -32,7 +34,8 @@ end function CommonSolve.init( ::AbstractNonlinearProblem, mode::AbstractNonlinearTerminationMode, du, u, - saved_value_prototype...; abstol = nothing, reltol = nothing, kwargs...) + saved_value_prototype...; abstol = nothing, reltol = nothing, kwargs... +) T = promote_type(eltype(du), eltype(u)) abstol = get_tolerance(u, abstol, T) reltol = get_tolerance(u, reltol, T) @@ -77,12 +80,14 @@ function CommonSolve.init( return NonlinearTerminationModeCache( u_unaliased, ReturnCode.Default, abstol, reltol, best_value, mode, initial_objective, objectives_trace, 0, saved_value_prototype, - u0_norm, step_norm_trace, max_stalled_steps, u_diff_cache) + u0_norm, step_norm_trace, max_stalled_steps, u_diff_cache + ) end function SciMLBase.reinit!( cache::NonlinearTerminationModeCache, du, u, saved_value_prototype...; - abstol = cache.abstol, reltol = cache.reltol, kwargs...) + abstol = cache.abstol, reltol = cache.reltol, kwargs... +) T = eltype(cache.abstol) length(saved_value_prototype) != 0 && (cache.saved_values = saved_value_prototype) @@ -113,7 +118,8 @@ end ## This dispatch is needed based on how Terminating Callback works! function (cache::NonlinearTerminationModeCache)( - integrator::AbstractODEIntegrator, abstol::Number, reltol::Number, min_t) + integrator::AbstractODEIntegrator, abstol::Number, reltol::Number, min_t +) if min_t === nothing || integrator.t ≥ min_t return cache(cache.mode, SciMLBase.get_du(integrator), integrator.u, integrator.uprev, abstol, reltol) @@ -125,7 +131,8 @@ function (cache::NonlinearTerminationModeCache)(du, u, uprev, args...) end function (cache::NonlinearTerminationModeCache)( - mode::AbstractNonlinearTerminationMode, du, u, uprev, abstol, reltol, args...) + mode::AbstractNonlinearTerminationMode, du, u, uprev, abstol, reltol, args... +) if check_convergence(mode, du, u, uprev, abstol, reltol) cache.retcode = ReturnCode.Success return true @@ -134,7 +141,8 @@ function (cache::NonlinearTerminationModeCache)( end function (cache::NonlinearTerminationModeCache)( - mode::AbstractSafeNonlinearTerminationMode, du, u, uprev, abstol, reltol, args...) + mode::AbstractSafeNonlinearTerminationMode, du, u, uprev, abstol, reltol, args... +) if mode isa AbsNormSafeTerminationMode || mode isa AbsNormSafeBestTerminationMode objective = Utils.apply_norm(mode.internalnorm, du) criteria = abstol @@ -251,7 +259,8 @@ end # High-Level API with defaults. ## This is mostly for internal usage in NonlinearSolve and SimpleNonlinearSolve function default_termination_mode( - ::Union{ImmutableNonlinearProblem, NonlinearProblem}, ::Val{:simple}) + ::Union{ImmutableNonlinearProblem, NonlinearProblem}, ::Val{:simple} +) return AbsNormTerminationMode(Base.Fix1(maximum, abs)) end function default_termination_mode(::NonlinearLeastSquaresProblem, ::Val{:simple}) @@ -259,7 +268,8 @@ function default_termination_mode(::NonlinearLeastSquaresProblem, ::Val{:simple} end function default_termination_mode( - ::Union{ImmutableNonlinearProblem, NonlinearProblem}, ::Val{:regular}) + ::Union{ImmutableNonlinearProblem, NonlinearProblem}, ::Val{:regular} +) return AbsNormSafeBestTerminationMode(Base.Fix1(maximum, abs); max_stalled_steps = 32) end @@ -268,16 +278,53 @@ function default_termination_mode(::NonlinearLeastSquaresProblem, ::Val{:regular end function init_termination_cache( - prob::AbstractNonlinearProblem, abstol, reltol, du, u, ::Nothing, callee::Val) + prob::AbstractNonlinearProblem, abstol, reltol, du, u, ::Nothing, callee::Val +) return init_termination_cache( prob, abstol, reltol, du, u, default_termination_mode(prob, callee), callee) end function init_termination_cache(prob::AbstractNonlinearProblem, abstol, reltol, du, - u, tc::AbstractNonlinearTerminationMode, ::Val) + u, tc::AbstractNonlinearTerminationMode, ::Val +) T = promote_type(eltype(du), eltype(u)) abstol = get_tolerance(u, abstol, T) reltol = get_tolerance(u, reltol, T) - cache = CommonSolve.init(prob, tc, du, u; abstol, reltol) + cache = init(prob, tc, du, u; abstol, reltol) return abstol, reltol, cache end + +function check_and_update!(cache, fu, u, uprev) + return check_and_update!( + cache.termination_cache, cache, fu, u, uprev, cache.termination_cache.mode + ) +end + +function check_and_update!(tc_cache, cache, fu, u, uprev, mode) + if tc_cache(fu, u, uprev) + cache.retcode = tc_cache.retcode + update_from_termination_cache!(tc_cache, cache, mode, u) + cache.force_stop = true + end +end + +function update_from_termination_cache!(tc_cache, cache, u = get_u(cache)) + return update_from_termination_cache!(tc_cache, cache, tc_cache.mode, u) +end + +function update_from_termination_cache!( + tc_cache, cache, ::AbstractNonlinearTerminationMode, u = get_u(cache) +) + Utils.evaluate_f!(cache, u, cache.p) +end + +function update_from_termination_cache!( + tc_cache, cache, ::AbstractSafeBestNonlinearTerminationMode, u = get_u(cache) +) + if SciMLBase.isinplace(cache) + copyto!(get_u(cache), tc_cache.u) + else + SciMLBase.set_u!(cache, tc_cache.u) + end + Utils.evaluate_f!(cache, get_u(cache), cache.p) +end diff --git a/src/timer_outputs.jl b/lib/NonlinearSolveBase/src/timer_outputs.jl similarity index 83% rename from src/timer_outputs.jl rename to lib/NonlinearSolveBase/src/timer_outputs.jl index 510e1d5ed..623e7ea15 100644 --- a/src/timer_outputs.jl +++ b/lib/NonlinearSolveBase/src/timer_outputs.jl @@ -1,32 +1,9 @@ # Timer Outputs has some overhead, so we only use it if we are debugging -# Even `@static_timeit` has overhead so we write our custom version of that using -# Preferences +# Even `@timeit` has overhead so we write our custom version of that using Preferences const TIMER_OUTPUTS_ENABLED = @load_preference("enable_timer_outputs", false) @static if TIMER_OUTPUTS_ENABLED - using TimerOutputs -end - -""" - enable_timer_outputs() - -Enable `TimerOutput` for all `NonlinearSolve` algorithms. This is useful for debugging -but has some overhead, so it is disabled by default. -""" -function enable_timer_outputs() - @set_preferences!("enable_timer_outputs"=>true) - @info "Timer Outputs Enabled. Restart the Julia session for this to take effect." -end - -""" - disable_timer_outputs() - -Disable `TimerOutput` for all `NonlinearSolve` algorithms. This should be used when -`NonlinearSolve` is being used in performance-critical code. -""" -function disable_timer_outputs() - @set_preferences!("enable_timer_outputs"=>false) - @info "Timer Outputs Disabled. Restart the Julia session for this to take effect." + using TimerOutputs: TimerOutput, timer_expr, reset_timer! end function get_timer_output() @@ -41,11 +18,11 @@ end @static_timeit to name expr Like `TimerOutputs.@timeit_debug` but has zero overhead if `TimerOutputs` is disabled via -[`NonlinearSolve.disable_timer_outputs()`](@ref). +[`NonlinearSolveBase.disable_timer_outputs()`](@ref). """ macro static_timeit(to, name, expr) @static if TIMER_OUTPUTS_ENABLED - return TimerOutputs.timer_expr(__module__, false, to, name, expr) + return timer_expr(__module__, false, to, name, expr) else return esc(expr) end @@ -54,3 +31,25 @@ end @static if !TIMER_OUTPUTS_ENABLED @inline reset_timer!(::Nothing) = nothing end + +""" + enable_timer_outputs() + +Enable `TimerOutput` for all `NonlinearSolve` algorithms. This is useful for debugging +but has some overhead, so it is disabled by default. +""" +function enable_timer_outputs() + @set_preferences!("enable_timer_outputs"=>true) + @info "Timer Outputs Enabled. Restart the Julia session for this to take effect." +end + +""" + disable_timer_outputs() + +Disable `TimerOutput` for all `NonlinearSolve` algorithms. This should be used when +`NonlinearSolve` is being used in performance-critical code. +""" +function disable_timer_outputs() + @set_preferences!("enable_timer_outputs"=>false) + @info "Timer Outputs Disabled. Restart the Julia session for this to take effect." +end diff --git a/lib/NonlinearSolveBase/src/tracing.jl b/lib/NonlinearSolveBase/src/tracing.jl new file mode 100644 index 000000000..d090c7c40 --- /dev/null +++ b/lib/NonlinearSolveBase/src/tracing.jl @@ -0,0 +1,233 @@ +@concrete struct NonlinearSolveTracing + trace_mode <: Union{Val{:minimal}, Val{:condition_number}, Val{:all}} + print_frequency::Int + store_frequency::Int +end + +""" + TraceMinimal(freq) + TraceMinimal(; print_frequency = 1, store_frequency::Int = 1) + +Trace Minimal Information + + 1. Iteration Number + 2. f(u) inf-norm + 3. Step 2-norm + +See also [`TraceWithJacobianConditionNumber`](@ref) and [`TraceAll`](@ref). +""" +function TraceMinimal(; print_frequency = 1, store_frequency::Int = 1) + return NonlinearSolveTracing(Val(:minimal), print_frequency, store_frequency) +end + +""" + TraceWithJacobianConditionNumber(freq) + TraceWithJacobianConditionNumber(; print_frequency = 1, store_frequency::Int = 1) + +[`TraceMinimal`](@ref) + Print the Condition Number of the Jacobian. + +See also [`TraceMinimal`](@ref) and [`TraceAll`](@ref). +""" +function TraceWithJacobianConditionNumber(; print_frequency = 1, store_frequency::Int = 1) + return NonlinearSolveTracing(Val(:condition_number), print_frequency, store_frequency) +end + +""" + TraceAll(freq) + TraceAll(; print_frequency = 1, store_frequency::Int = 1) + +[`TraceWithJacobianConditionNumber`](@ref) + Store the Jacobian, u, f(u), and δu. + +!!! warning + + This is very expensive and makes copyies of the Jacobian, u, f(u), and δu. + +See also [`TraceMinimal`](@ref) and [`TraceWithJacobianConditionNumber`](@ref). +""" +function TraceAll(; print_frequency = 1, store_frequency::Int = 1) + return NonlinearSolveTracing(Val(:all), print_frequency, store_frequency) +end + +for Tr in (:TraceMinimal, :TraceWithJacobianConditionNumber, :TraceAll) + @eval $(Tr)(freq) = $(Tr)(; print_frequency = freq, store_frequency = freq) +end + +# NonlinearSolve Tracing Utilities +@concrete struct NonlinearSolveTraceEntry + iteration::Int + fnorm + stepnorm + condJ + storage + norm_type::Symbol +end + +function Base.getproperty(entry::NonlinearSolveTraceEntry, sym::Symbol) + hasfield(typeof(entry), sym) && return getfield(entry, sym) + return getproperty(entry.storage, sym) +end + +function print_top_level(io::IO, entry::NonlinearSolveTraceEntry) + if entry.condJ === nothing + @printf io "%-8s\t%-20s\t%-20s\n" "----" "-------------" "-----------" + if entry.norm_type === :L2 + @printf io "%-8s\t%-20s\t%-20s\n" "Iter" "f(u) 2-norm" "Step 2-norm" + else + @printf io "%-8s\t%-20s\t%-20s\n" "Iter" "f(u) inf-norm" "Step 2-norm" + end + @printf io "%-8s\t%-20s\t%-20s\n" "----" "-------------" "-----------" + else + @printf io "%-8s\t%-20s\t%-20s\t%-20s\n" "----" "-------------" "-----------" "-------" + if entry.norm_type === :L2 + @printf io "%-8s\t%-20s\t%-20s\t%-20s\n" "Iter" "f(u) 2-norm" "Step 2-norm" "cond(J)" + else + @printf io "%-8s\t%-20s\t%-20s\t%-20s\n" "Iter" "f(u) inf-norm" "Step 2-norm" "cond(J)" + end + @printf io "%-8s\t%-20s\t%-20s\t%-20s\n" "----" "-------------" "-----------" "-------" + end +end + +function Base.show(io::IO, ::MIME"text/plain", entry::NonlinearSolveTraceEntry) + entry.iteration == 0 && print_top_level(io, entry) + if entry.iteration < 0 # Special case for final entry + @printf io "%-8s\t%-20.8e\n" "Final" entry.fnorm + @printf io "%-28s\n" "----------------------" + elseif entry.condJ === nothing + @printf io "%-8d\t%-20.8e\t%-20.8e\n" entry.iteration entry.fnorm entry.stepnorm + else + @printf io "%-8d\t%-20.8e\t%-20.8e\t%-20.8e\n" entry.iteration entry.fnorm entry.stepnorm entry.condJ + end +end + +function NonlinearSolveTraceEntry(prob::AbstractNonlinearProblem, iteration, fu, δu, J, u) + norm_type = ifelse(prob isa NonlinearLeastSquaresProblem, :L2, :Inf) + fnorm = prob isa NonlinearLeastSquaresProblem ? L2_NORM(fu) : Linf_NORM(fu) + condJ = J !== missing ? Utils.condition_number(J) : nothing + storage = if u === missing + nothing + else + (; + u = ArrayInterface.ismutable(u) ? copy(u) : u, + fu = ArrayInterface.ismutable(fu) ? copy(fu) : fu, + δu = ArrayInterface.ismutable(δu) ? copy(δu) : δu, + J = ArrayInterface.ismutable(J) ? copy(J) : J + ) + end + return NonlinearSolveTraceEntry( + iteration, fnorm, L2_NORM(δu), condJ, storage, norm_type + ) +end + +@concrete struct NonlinearSolveTrace + show_trace <: Union{Val{false}, Val{true}} + store_trace <: Union{Val{false}, Val{true}} + history + trace_level <: NonlinearSolveTracing + prob +end + +reset!(trace::NonlinearSolveTrace) = reset!(trace.history) +reset!(::Nothing) = nothing +reset!(history::Vector) = empty!(history) + +function Base.show(io::IO, ::MIME"text/plain", trace::NonlinearSolveTrace) + if trace.history !== nothing + foreach(trace.history) do entry + show(io, MIME"text/plain"(), entry) + end + else + print(io, "Tracing Disabled") + end +end + +function init_nonlinearsolve_trace( + prob, alg, u, fu, J, δu; show_trace::Val = Val(false), + trace_level::NonlinearSolveTracing = TraceMinimal(), store_trace::Val = Val(false), + uses_jac_inverse = Val(false), kwargs... +) + return init_nonlinearsolve_trace( + prob, alg, show_trace, trace_level, store_trace, u, fu, J, δu, uses_jac_inverse + ) +end + +function init_nonlinearsolve_trace( + prob::AbstractNonlinearProblem, alg, show_trace::Val, + trace_level::NonlinearSolveTracing, store_trace::Val, u, fu, J, δu, + uses_jac_inverse::Val +) + if show_trace isa Val{true} + print("\nAlgorithm: ") + str = Utils.clean_sprint_struct(alg, 0) + Base.printstyled(str, "\n\n"; color = :green, bold = true) + end + J = uses_jac_inverse isa Val{true} ? + (trace_level.trace_mode isa Val{:minimal} ? J : LinearAlgebra.pinv(J)) : J + history = init_trace_history(prob, show_trace, trace_level, store_trace, u, fu, J, δu) + return NonlinearSolveTrace(show_trace, store_trace, history, trace_level, prob) +end + +function init_trace_history( + prob::AbstractNonlinearProblem, show_trace::Val, trace_level, + store_trace::Val, u, fu, J, δu +) + store_trace isa Val{false} && show_trace isa Val{false} && return nothing + entry = if trace_level.trace_mode isa Val{:minimal} + NonlinearSolveTraceEntry(prob, 0, fu, δu, missing, missing) + elseif trace_level.trace_mode isa Val{:condition_number} + NonlinearSolveTraceEntry(prob, 0, fu, δu, J, missing) + else + NonlinearSolveTraceEntry(prob, 0, fu, δu, J, u) + end + show_trace isa Val{true} && show(stdout, MIME"text/plain"(), entry) + store_trace isa Val{true} && return NonlinearSolveTraceEntry[entry] + return nothing +end + +function update_trace!( + trace::NonlinearSolveTrace, iter, u, fu, J, δu, α = true; last::Val = Val(false) +) + trace.store_trace isa Val{false} && trace.show_trace isa Val{false} && return nothing + + if last isa Val{true} + norm_type = ifelse(trace.prob isa NonlinearLeastSquaresProblem, :L2, :Inf) + fnorm = trace.prob isa NonlinearLeastSquaresProblem ? L2_NORM(fu) : Linf_NORM(fu) + entry = NonlinearSolveTraceEntry(-1, fnorm, NaN32, nothing, nothing, norm_type) + trace.show_trace isa Val{true} && show(stdout, MIME"text/plain"(), entry) + return trace + end + + show_now = trace.show_trace isa Val{true} && + (mod1(iter, trace.trace_level.print_frequency) == 1) + store_now = trace.store_trace isa Val{true} && + (mod1(iter, trace.trace_level.store_frequency) == 1) + if show_now || store_now + entry = if trace.trace_level.trace_mode isa Val{:minimal} + NonlinearSolveTraceEntry(trace.prob, iter, fu, δu .* α, missing, missing) + else + J = convert(AbstractArray, J) + if trace.trace_level.trace_mode isa Val{:condition_number} + NonlinearSolveTraceEntry(trace.prob, iter, fu, δu .* α, J, missing) + else + NonlinearSolveTraceEntry(trace.prob, iter, fu, δu .* α, J, u) + end + end + show_now && show(stdout, MIME"text/plain"(), entry) + store_now && push!(trace.history, entry) + end + return trace +end + +function update_trace!(cache, α = true; uses_jac_inverse = Val(false)) + trace = Utils.safe_getproperty(cache, Val(:trace)) + trace === missing && return nothing + + J = Utils.safe_getproperty(cache, Val(:J)) + if J === missing + update_trace!( + trace, cache.nsteps + 1, get_u(cache), get_fu(cache), nothing, cache.du, α + ) + else + J = uses_jac_inverse isa Val{true} ? Utils.Pinv(cache.J) : cache.J + update_trace!(trace, cache.nsteps + 1, get_u(cache), get_fu(cache), J, cache.du, α) + end +end diff --git a/lib/NonlinearSolveBase/src/utils.jl b/lib/NonlinearSolveBase/src/utils.jl index 825f63767..99e8e50c5 100644 --- a/lib/NonlinearSolveBase/src/utils.jl +++ b/lib/NonlinearSolveBase/src/utils.jl @@ -1,14 +1,32 @@ module Utils using ArrayInterface: ArrayInterface +using ConcreteStructs: @concrete using FastClosures: @closure -using LinearAlgebra: norm -using RecursiveArrayTools: AbstractVectorOfArray, ArrayPartition +using LinearAlgebra: LinearAlgebra, Diagonal, Symmetric, norm, dot, cond, diagind, pinv +using MaybeInplace: @bb +using RecursiveArrayTools: AbstractVectorOfArray, ArrayPartition, recursivecopy! +using SciMLOperators: AbstractSciMLOperator +using SciMLBase: SciMLBase, AbstractNonlinearProblem, NonlinearFunction +using StaticArraysCore: StaticArray, SArray, SMatrix -using ..NonlinearSolveBase: L2_NORM, Linf_NORM +using ..NonlinearSolveBase: NonlinearSolveBase, L2_NORM, Linf_NORM + +is_extension_loaded(::Val) = false fast_scalar_indexing(xs...) = all(ArrayInterface.fast_scalar_indexing, xs) +@concrete struct Pinv + J +end + +function Base.convert(::Type{AbstractArray}, A::Pinv) + hasmethod(pinv, Tuple{typeof(A.J)}) && return pinv(A.J) + @warn "`pinv` not defined for $(typeof(A.J)). Jacobian will not be inverted when \ + tracing." maxlog=1 + return A.J +end + function nonallocating_isapprox(x::Number, y::Number; atol = false, rtol = atol > 0 ? false : sqrt(eps(promote_type(typeof(x), typeof(y))))) return isapprox(x, y; atol, rtol) @@ -73,6 +91,12 @@ convert_real(::Type{T}, ::Nothing) where {T} = nothing convert_real(::Type{T}, x) where {T} = real(T(x)) restructure(::Number, x::Number) = x +function restructure( + y::T1, x::T2 +) where {T1 <: AbstractSciMLOperator, T2 <: AbstractSciMLOperator} + @assert size(y)==size(x) "cannot restructure operators. ensure their sizes match." + return x +end restructure(y, x) = ArrayInterface.restructure(y, x) function safe_similar(x, args...; kwargs...) @@ -90,4 +114,208 @@ end safe_reshape(x::Number, args...) = x safe_reshape(x, args...) = reshape(x, args...) +@generated function safe_getproperty(s::S, ::Val{X}) where {S, X} + hasfield(S, X) && return :(getproperty(s, $(Meta.quot(X)))) + return :(missing) +end + +@generated function safe_vec(v) + hasmethod(vec, Tuple{typeof(v)}) || return :(vec(v)) + return :(v) +end +safe_vec(v::Number) = v +safe_vec(v::AbstractVector) = v + +safe_dot(x, y) = dot(safe_vec(x), safe_vec(y)) + +unwrap_val(x) = x +unwrap_val(::Val{x}) where {x} = unwrap_val(x) + +is_default_value(::Any, ::Symbol, ::Nothing) = true +is_default_value(::Any, ::Symbol, ::Missing) = true +is_default_value(::Any, ::Symbol, val::Int) = val == typemax(typeof(val)) +is_default_value(::Any, ::Symbol, ::Any) = false +is_default_value(::Any, ::Any, ::Any) = false + +maybe_symmetric(x) = Symmetric(x) +maybe_symmetric(x::Number) = x +## LinearSolve with `nothing` doesn't dispatch correctly here +maybe_symmetric(x::StaticArray) = x # XXX: Can we remove this? +maybe_symmetric(x::AbstractSciMLOperator) = x + +# Define special concatenation for certain Array combinations +faster_vcat(x, y) = vcat(x, y) + +maybe_unaliased(x::Union{Number, SArray}, ::Bool) = x +function maybe_unaliased(x::AbstractArray, alias::Bool) + (alias || !ArrayInterface.can_setindex(typeof(x))) && return x + return copy(x) +end +maybe_unaliased(x::AbstractSciMLOperator, ::Bool) = x + +can_setindex(x) = ArrayInterface.can_setindex(x) +can_setindex(::Number) = false + +function evaluate_f!!(prob::AbstractNonlinearProblem, fu, u, p = prob.p) + return evaluate_f!!(prob.f, fu, u, p) +end +function evaluate_f!!(f::NonlinearFunction, fu, u, p) + if SciMLBase.isinplace(f) + f(fu, u, p) + return fu + end + return f(u, p) +end + +function evaluate_f(prob::AbstractNonlinearProblem, u) + if SciMLBase.isinplace(prob) + fu = prob.f.resid_prototype === nothing ? similar(u) : + similar(prob.f.resid_prototype) + prob.f(fu, u, prob.p) + else + fu = prob.f(u, prob.p) + end + return fu +end + +function evaluate_f!(cache, u, p) + cache.stats.nf += 1 + if SciMLBase.isinplace(cache) + cache.prob.f(NonlinearSolveBase.get_fu(cache), u, p) + else + NonlinearSolveBase.set_fu!(cache, cache.prob.f(u, p)) + end +end + +function make_sparse end + +condition_number(J::AbstractMatrix) = cond(J) +function condition_number(J::AbstractVector) + if !ArrayInterface.can_setindex(J) + J′ = similar(J) + copyto!(J′, J) + J = J′ + end + return cond(Diagonal(J)) +end +condition_number(::Any) = -1 + +# compute `pinv` if `inv` won't work +maybe_pinv!!_workspace(A) = nothing, A + +maybe_pinv!!(workspace, A::Union{Number, AbstractMatrix}) = pinv(A) +function maybe_pinv!!(workspace, A::Diagonal) + D = A.diag + @bb @. D = pinv(D) + return Diagonal(D) +end +maybe_pinv!!(workspace, A::AbstractVector) = maybe_pinv!!(workspace, Diagonal(A)) +function maybe_pinv!!(workspace, A::StridedMatrix) + LinearAlgebra.checksquare(A) + if LinearAlgebra.istriu(A) + issingular = any(iszero, @view(A[diagind(A)])) + A_ = LinearAlgebra.UpperTriangular(A) + !issingular && return LinearAlgebra.triu!(parent(inv(A_))) + elseif LinearAlgebra.istril(A) + A_ = LinearAlgebra.LowerTriangular(A) + issingular = any(iszero, @view(A_[diagind(A_)])) + !issingular && return LinearAlgebra.tril!(parent(inv(A_))) + else + F = LinearAlgebra.lu(A; check = false) + if LinearAlgebra.issuccess(F) + Ai = LinearAlgebra.inv!(F) + return convert(typeof(parent(Ai)), Ai) + end + end + return pinv(A) +end + +function initial_jacobian_scaling_alpha(α, u, fu, ::Any) + return convert(promote_type(eltype(u), eltype(fu)), α) +end +function initial_jacobian_scaling_alpha(::Nothing, u, fu, internalnorm::F) where {F} + fu_norm = internalnorm(fu) + fu_norm < 1e-5 && return initial_jacobian_scaling_alpha(true, u, fu, internalnorm) + return (2 * fu_norm) / max(L2_NORM(u), true) +end + +make_identity!!(::T, α) where {T <: Number} = T(α) +function make_identity!!(A::AbstractVector{T}, α) where {T} + if ArrayInterface.can_setindex(A) + @. A = α + else + A = one.(A) .* α + end + return A +end +function make_identity!!(::SMatrix{S1, S2, T, L}, α) where {S1, S2, T, L} + return SMatrix{S1, S2, T, L}(LinearAlgebra.I * α) +end +function make_identity!!(A::AbstractMatrix{T}, α) where {T} + A = ArrayInterface.can_setindex(A) ? A : similar(A) + fill!(A, false) + if ArrayInterface.fast_scalar_indexing(A) + @simd ivdep for i in axes(A, 1) + @inbounds A[i, i] = α + end + else + A[diagind(A)] .= α + end + return A +end + +function reinit_common!(cache, u0, p, alias_u0::Bool) + if SciMLBase.isinplace(cache) + recursivecopy!(cache.u, u0) + cache.prob.f(cache.fu, cache.u, p) + else + cache.u = maybe_unaliased(u0, alias_u0) + NonlinearSolveBase.set_fu!(cache, cache.prob.f(u0, p)) + end + cache.p = p +end + +function clean_sprint_struct(x) + x isa Symbol && return "$(Meta.quot(x))" + x isa Number && return string(x) + (!Base.isstructtype(typeof(x)) || x isa Val) && return string(x) + + modifiers = String[] + name = nameof(typeof(x)) + for field in fieldnames(typeof(x)) + val = getfield(x, field) + if field === :name && val isa Symbol && val !== :unknown + name = val + continue + end + is_default_value(x, field, val) && continue + push!(modifiers, "$(field) = $(clean_sprint_struct(val))") + end + + return "$(nameof(typeof(x)))($(join(modifiers, ", ")))" +end + +function clean_sprint_struct(x, indent::Int) + x isa Symbol && return "$(Meta.quot(x))" + x isa Number && return string(x) + (!Base.isstructtype(typeof(x)) || x isa Val) && return string(x) + + modifiers = String[] + name = nameof(typeof(x)) + for field in fieldnames(typeof(x)) + val = getfield(x, field) + if field === :name && val isa Symbol && val !== :unknown + name = val + continue + end + is_default_value(x, field, val) && continue + push!(modifiers, "$(field) = $(clean_sprint_struct(val, indent + 4))") + end + spacing = " "^indent * " " + spacing_last = " "^indent + + length(modifiers) == 0 && return "$(nameof(typeof(x)))()" + return "$(name)(\n$(spacing)$(join(modifiers, ",\n$(spacing)"))\n$(spacing_last))" +end + end diff --git a/lib/NonlinearSolveBase/src/wrappers.jl b/lib/NonlinearSolveBase/src/wrappers.jl new file mode 100644 index 000000000..0b96836ad --- /dev/null +++ b/lib/NonlinearSolveBase/src/wrappers.jl @@ -0,0 +1,113 @@ +function assert_extension_supported_termination_condition( + termination_condition, alg; abs_norm_supported = true +) + no_termination_condition = termination_condition === nothing + no_termination_condition && return nothing + abs_norm_supported && termination_condition isa AbsNormTerminationMode && return nothing + throw(AssertionError("`$(nameof(typeof(alg)))` does not support termination conditions!")) +end + +function construct_extension_function_wrapper( + prob::AbstractNonlinearProblem; alias_u0::Bool = false, + can_handle_oop::Val = Val(false), can_handle_scalar::Val = Val(false), + make_fixed_point::Val = Val(false), force_oop::Val = Val(false) +) + if can_handle_oop isa Val{false} && can_handle_scalar isa Val{true} + error("Incorrect Specification: OOP not supported but scalar supported.") + end + + resid = Utils.evaluate_f(prob, prob.u0) + u0 = can_handle_scalar isa Val{true} || !(prob.u0 isa Number) ? + Utils.maybe_unaliased(prob.u0, alias_u0) : [prob.u0] + + fₚ = if make_fixed_point isa Val{true} + if SciMLBase.isinplace(prob) + @closure (du, u) -> begin + prob.f(du, u, prob.p) + du .+= u + return du + end + else + @closure u -> prob.f(u, prob.p) .+ u + end + else + if SciMLBase.isinplace(prob) + @closure (du, u) -> begin + prob.f(du, u, prob.p) + return du + end + else + Base.Fix2(prob.f, prob.p) + end + end + + f_flat_structure = if SciMLBase.isinplace(prob) + u0_size, du_size = size(u0), size(resid) + @closure (du, u) -> begin + fₚ(reshape(du, du_size), reshape(u, u0_size)) + return du + end + else + if prob.u0 isa Number + if can_handle_scalar isa Val{true} + fₚ + elseif can_handle_oop isa Val{true} + @closure u -> [fₚ(first(u))] + else + @closure (du, u) -> begin + du[1] = fₚ(first(u)) + return du + end + end + else + u0_size = size(u0) + if can_handle_oop isa Val{true} + @closure u -> vec(fₚ(reshape(u, u0_size))) + else + @closure (du, u) -> begin + copyto!(du, fₚ(reshape(u, u0_size))) + return du + end + end + end + end + + f_final = if force_oop isa Val{true} && applicable(f_flat_structure, u0, u0) + resid = resid isa Number ? [resid] : Utils.safe_vec(resid) + du = Utils.safe_vec(zero(resid)) + @closure u -> begin + f_flat_structure(du, u) + return du + end + else + f_flat_structure + end + + return f_final, Utils.safe_vec(u0), (resid isa Number ? [resid] : Utils.safe_vec(resid)) +end + +function construct_extension_jac( + prob, alg, u0, fu; + can_handle_oop::Val = Val(false), can_handle_scalar::Val = Val(false), + autodiff = nothing, initial_jacobian = Val(false), kwargs... +) + autodiff = select_jacobian_autodiff(prob, autodiff) + + Jₚ = construct_jacobian_cache( + prob, alg, prob.f, fu, u0, prob.p; + stats = NLStats(0, 0, 0, 0, 0), autodiff, kwargs... + ) + + J_no_scalar = can_handle_scalar isa Val{false} && prob.u0 isa Number ? + @closure(u->[Jₚ(u[1])]) : Jₚ + + J_final = if can_handle_oop isa Val{false} && !SciMLBase.isinplace(prob) + @closure((J, u)->copyto!(J, J_no_scalar(u))) + else + J_no_scalar + end + + initial_jacobian isa Val{false} && return J_final + + return J_final, Jₚ(nothing) +end diff --git a/lib/NonlinearSolveBase/test/runtests.jl b/lib/NonlinearSolveBase/test/runtests.jl index 07c0f14c6..8f4f4d8a1 100644 --- a/lib/NonlinearSolveBase/test/runtests.jl +++ b/lib/NonlinearSolveBase/test/runtests.jl @@ -8,7 +8,10 @@ using InteractiveUtils, Test @testset "Aqua" begin using Aqua, NonlinearSolveBase - Aqua.test_all(NonlinearSolveBase; piracies = false, ambiguities = false) + Aqua.test_all( + NonlinearSolveBase; piracies = false, ambiguities = false, stale_deps = false + ) + Aqua.test_stale_deps(NonlinearSolveBase; ignore = [:TimerOutputs]) Aqua.test_piracies(NonlinearSolveBase) Aqua.test_ambiguities(NonlinearSolveBase; recursive = false) end @@ -21,4 +24,13 @@ using InteractiveUtils, Test @test check_no_stale_explicit_imports(NonlinearSolveBase) === nothing @test check_all_qualified_accesses_via_owners(NonlinearSolveBase) === nothing end + + @testset "Banded Matrix vcat" begin + using BandedMatrices, LinearAlgebra, SparseArrays + + b = BandedMatrix(Ones(5, 5), (1, 1)) + d = Diagonal(ones(5, 5)) + + @test NonlinearSolveBase.Utils.faster_vcat(b, d) == vcat(sparse(b), d) + end end diff --git a/lib/NonlinearSolveFirstOrder/LICENSE b/lib/NonlinearSolveFirstOrder/LICENSE new file mode 100644 index 000000000..46c972b17 --- /dev/null +++ b/lib/NonlinearSolveFirstOrder/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2024 Avik Pal + +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/NonlinearSolveFirstOrder/Project.toml b/lib/NonlinearSolveFirstOrder/Project.toml new file mode 100644 index 000000000..94ab4fdbc --- /dev/null +++ b/lib/NonlinearSolveFirstOrder/Project.toml @@ -0,0 +1,88 @@ +name = "NonlinearSolveFirstOrder" +uuid = "5959db7a-ea39-4486-b5fe-2dd0bf03d60d" +authors = ["Avik Pal and contributors"] +version = "1.0.0" + +[deps] +ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" +ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9" +CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2" +ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471" +DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e" +FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae" +MaybeInplace = "bb5d69b7-63fc-4a16-80bd-7e42200c7bdb" +NonlinearSolveBase = "be0214bd-f91f-a760-ac4e-3421ce2b2da0" +PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" +Reexport = "189a3867-3050-52da-a836-e630ba90ab69" +SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" +SciMLJacobianOperators = "19f34311-ddf3-4b8b-af20-060888a46c0e" +Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46" +StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" + +[compat] +ADTypes = "1.9.0" +Aqua = "0.8" +ArrayInterface = "7.16.0" +BandedMatrices = "1.7.5" +BenchmarkTools = "1.5.0" +CommonSolve = "0.2.4" +ConcreteStructs = "0.2.3" +DiffEqBase = "6.158.3" +Enzyme = "0.13.12" +ExplicitImports = "1.5" +FiniteDiff = "2.24" +ForwardDiff = "0.10.36" +Hwloc = "3" +InteractiveUtils = "<0.0.1, 1" +LineSearch = "0.1.4" +LineSearches = "7.3.0" +LinearAlgebra = "1.10" +LinearSolve = "2.36.1" +MaybeInplace = "0.1.4" +NonlinearProblemLibrary = "0.1.2" +NonlinearSolveBase = "1.1" +Pkg = "1.10" +PrecompileTools = "1.2" +Random = "1.10" +ReTestItems = "1.24" +Reexport = "1" +SciMLBase = "2.58" +SciMLJacobianOperators = "0.1.0" +Setfield = "1.1.1" +SparseArrays = "1.10" +SparseConnectivityTracer = "0.6.8" +SparseMatrixColorings = "0.4.5" +StableRNGs = "1" +StaticArrays = "1.9.8" +StaticArraysCore = "1.4.3" +Test = "1.10" +Zygote = "0.6.72" +julia = "1.10" + +[extras] +Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" +BandedMatrices = "aae01518-5342-5314-be14-df237901396f" +BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" +ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7" +Hwloc = "0e44f5e4-bd66-52a0-8798-143a42290a1d" +InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" +LineSearch = "87fe0de2-c867-4266-b59a-2f0a94fc965b" +LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255" +NonlinearProblemLibrary = "b7050fa9-e91f-4b37-bcee-a89a063da141" +Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823" +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" +SparseConnectivityTracer = "9f842d2f-2579-4b1d-911e-f412cf18a3f5" +SparseMatrixColorings = "0a514795-09f3-496d-8182-132a7b665d35" +StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" +StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" + +[targets] +test = ["Aqua", "BandedMatrices", "BenchmarkTools", "Enzyme", "ExplicitImports", "Hwloc", "InteractiveUtils", "LineSearch", "LineSearches", "NonlinearProblemLibrary", "Pkg", "Random", "ReTestItems", "SparseArrays", "SparseConnectivityTracer", "SparseMatrixColorings", "StableRNGs", "StaticArrays", "Test", "Zygote"] diff --git a/lib/NonlinearSolveFirstOrder/src/NonlinearSolveFirstOrder.jl b/lib/NonlinearSolveFirstOrder/src/NonlinearSolveFirstOrder.jl new file mode 100644 index 000000000..b611b9051 --- /dev/null +++ b/lib/NonlinearSolveFirstOrder/src/NonlinearSolveFirstOrder.jl @@ -0,0 +1,68 @@ +module NonlinearSolveFirstOrder + +using ConcreteStructs: @concrete +using PrecompileTools: @compile_workload, @setup_workload +using Reexport: @reexport +using Setfield: @set! + +using ADTypes: ADTypes +using ArrayInterface: ArrayInterface +using LinearAlgebra: LinearAlgebra, Diagonal, dot, diagind +using StaticArraysCore: SArray + +using CommonSolve: CommonSolve +using DiffEqBase: DiffEqBase # Needed for `init` / `solve` dispatches +using LinearSolve: LinearSolve # Trigger Linear Solve extension in NonlinearSolveBase +using MaybeInplace: @bb +using NonlinearSolveBase: NonlinearSolveBase, AbstractNonlinearSolveAlgorithm, + AbstractNonlinearSolveCache, AbstractDampingFunction, + AbstractDampingFunctionCache, AbstractTrustRegionMethod, + AbstractTrustRegionMethodCache, + Utils, InternalAPI, get_timer_output, @static_timeit, + update_trace!, L2_NORM, + NewtonDescent, DampedNewtonDescent, GeodesicAcceleration, + Dogleg +using SciMLBase: SciMLBase, AbstractNonlinearProblem, NLStats, ReturnCode, NonlinearFunction +using SciMLJacobianOperators: VecJacOperator, JacVecOperator, StatefulJacobianOperator + +using FiniteDiff: FiniteDiff # Default Finite Difference Method +using ForwardDiff: ForwardDiff # Default Forward Mode AD + +include("raphson.jl") +include("gauss_newton.jl") +include("levenberg_marquardt.jl") +include("trust_region.jl") +include("pseudo_transient.jl") + +include("solve.jl") + +@setup_workload begin + include("../../../common/nonlinear_problem_workloads.jl") + include("../../../common/nlls_problem_workloads.jl") + + nlp_algs = [NewtonRaphson(), TrustRegion(), LevenbergMarquardt()] + nlls_algs = [GaussNewton(), TrustRegion(), LevenbergMarquardt()] + + @compile_workload begin + @sync begin + for prob in nonlinear_problems, alg in nlp_algs + Threads.@spawn CommonSolve.solve(prob, alg; abstol = 1e-2, verbose = false) + end + + for prob in nlls_problems, alg in nlls_algs + Threads.@spawn CommonSolve.solve(prob, alg; abstol = 1e-2, verbose = false) + end + end + end +end + +@reexport using SciMLBase, NonlinearSolveBase + +export NewtonRaphson, PseudoTransient +export GaussNewton, LevenbergMarquardt, TrustRegion + +export RadiusUpdateSchemes + +export GeneralizedFirstOrderAlgorithm + +end diff --git a/lib/NonlinearSolveFirstOrder/src/gauss_newton.jl b/lib/NonlinearSolveFirstOrder/src/gauss_newton.jl new file mode 100644 index 000000000..59b90b76e --- /dev/null +++ b/lib/NonlinearSolveFirstOrder/src/gauss_newton.jl @@ -0,0 +1,22 @@ +""" + GaussNewton(; + concrete_jac = nothing, linsolve = nothing, linesearch = missing, precs = nothing, + autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing + ) + +An advanced GaussNewton implementation with support for efficient handling of sparse +matrices via colored automatic differentiation and preconditioned linear solvers. Designed +for large-scale and numerically-difficult nonlinear systems. +""" +function GaussNewton(; + concrete_jac = nothing, linsolve = nothing, linesearch = missing, precs = nothing, + autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing +) + return GeneralizedFirstOrderAlgorithm(; + linesearch, + descent = NewtonDescent(; linsolve, precs), + autodiff, vjp_autodiff, jvp_autodiff, + concrete_jac, + name = :GaussNewton + ) +end diff --git a/lib/NonlinearSolveFirstOrder/src/levenberg_marquardt.jl b/lib/NonlinearSolveFirstOrder/src/levenberg_marquardt.jl new file mode 100644 index 000000000..f9c11ae28 --- /dev/null +++ b/lib/NonlinearSolveFirstOrder/src/levenberg_marquardt.jl @@ -0,0 +1,294 @@ +""" + LevenbergMarquardt(; + linsolve = nothing, precs = nothing, + damping_initial::Real = 1.0, α_geodesic::Real = 0.75, disable_geodesic = Val(false), + damping_increase_factor::Real = 2.0, damping_decrease_factor::Real = 3.0, + finite_diff_step_geodesic = 0.1, b_uphill::Real = 1.0, min_damping_D::Real = 1e-8, + autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing + ) + +An advanced Levenberg-Marquardt implementation with the improvements suggested in +[transtrum2012improvements](@citet). Designed for large-scale and numerically-difficult +nonlinear systems. + +### Keyword Arguments + + - `damping_initial`: the starting value for the damping factor. The damping factor is + inversely proportional to the step size. The damping factor is adjusted during each + iteration. Defaults to `1.0`. See Section 2.1 of [transtrum2012improvements](@citet). + - `damping_increase_factor`: the factor by which the damping is increased if a step is + rejected. Defaults to `2.0`. See Section 2.1 of [transtrum2012improvements](@citet). + - `damping_decrease_factor`: the factor by which the damping is decreased if a step is + accepted. Defaults to `3.0`. See Section 2.1 of [transtrum2012improvements](@citet). + - `min_damping_D`: the minimum value of the damping terms in the diagonal damping matrix + `DᵀD`, where `DᵀD` is given by the largest diagonal entries of `JᵀJ` yet encountered, + where `J` is the Jacobian. It is suggested by [transtrum2012improvements](@citet) to use + a minimum value of the elements in `DᵀD` to prevent the damping from being too small. + Defaults to `1e-8`. + - `disable_geodesic`: Disables Geodesic Acceleration if set to `Val(true)`. It provides + a way to trade-off robustness for speed, though in most situations Geodesic Acceleration + should not be disabled. + +For the remaining arguments, see [`GeodesicAcceleration`](@ref) and +[`NonlinearSolveFirstOrder.LevenbergMarquardtTrustRegion`](@ref) documentations. +""" +function LevenbergMarquardt(; + linsolve = nothing, precs = nothing, + damping_initial::Real = 1.0, α_geodesic::Real = 0.75, disable_geodesic = Val(false), + damping_increase_factor::Real = 2.0, damping_decrease_factor::Real = 3.0, + finite_diff_step_geodesic = 0.1, b_uphill::Real = 1.0, min_damping_D::Real = 1e-8, + autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing +) + descent = DampedNewtonDescent(; + linsolve, + precs, + initial_damping = damping_initial, + damping_fn = LevenbergMarquardtDampingFunction( + damping_increase_factor, damping_decrease_factor, min_damping_D + ) + ) + if disable_geodesic isa Val{false} + descent = GeodesicAcceleration(descent, finite_diff_step_geodesic, α_geodesic) + end + trustregion = LevenbergMarquardtTrustRegion(b_uphill) + return GeneralizedFirstOrderAlgorithm(; + trustregion, + descent, + autodiff, + vjp_autodiff, + jvp_autodiff, + name = :LevenbergMarquardt, + concrete_jac = Val(true) + ) +end + +@concrete struct LevenbergMarquardtDampingFunction <: AbstractDampingFunction + increase_factor + decrease_factor + min_damping +end + +function InternalAPI.init( + prob::AbstractNonlinearProblem, f::LevenbergMarquardtDampingFunction, + initial_damping, J, fu, u, normal_form::Val; kwargs... +) + T = promote_type(eltype(u), eltype(fu)) + DᵀD = init_levenberg_marquardt_diagonal(u, T(f.min_damping)) + if normal_form isa Val{true} + J_diag_cache = nothing + else + @bb J_diag_cache = similar(u) + end + J_damped = T(initial_damping) .* DᵀD + return LevenbergMarquardtDampingCache( + T(f.increase_factor), T(f.decrease_factor), T(f.min_damping), + T(f.increase_factor), T(initial_damping), DᵀD, J_diag_cache, J_damped, f, + T(initial_damping) + ) +end + +@concrete mutable struct LevenbergMarquardtDampingCache <: AbstractDampingFunctionCache + increase_factor + decrease_factor + min_damping + λ_factor + λ + DᵀD + J_diag_cache + J_damped + damping_f + initial_damping +end + +function InternalAPI.reinit!(cache::LevenbergMarquardtDampingCache, args...; kwargs...) + cache.λ = cache.initial_damping + cache.λ_factor = cache.damping_f.increase_factor + if !(cache.DᵀD isa Number) + if ArrayInterface.can_setindex(cache.DᵀD.diag) + cache.DᵀD.diag .= cache.min_damping + else + cache.DᵀD = Diagonal(ones(typeof(cache.DᵀD.diag)) * cache.min_damping) + end + end + cache.J_damped = cache.λ .* cache.DᵀD + return +end + +function NonlinearSolveBase.requires_normal_form_jacobian(::Union{ + LevenbergMarquardtDampingFunction, LevenbergMarquardtDampingCache}) + return false +end +function NonlinearSolveBase.requires_normal_form_rhs(::Union{ + LevenbergMarquardtDampingFunction, LevenbergMarquardtDampingCache}) + return false +end +function NonlinearSolveBase.returns_norm_form_damping(::Union{ + LevenbergMarquardtDampingFunction, LevenbergMarquardtDampingCache}) + return true +end + +(damping::LevenbergMarquardtDampingCache)(::Nothing) = damping.J_damped + +function InternalAPI.solve!( + cache::LevenbergMarquardtDampingCache, J, fu, ::Val{false}; kwargs... +) + if cache.J_diag_cache isa Number + cache.J_diag_cache = abs2(J) + elseif ArrayInterface.can_setindex(cache.J_diag_cache) + sum!(abs2, Utils.safe_vec(cache.J_diag_cache), J') + else + cache.J_diag_cache = dropdims(sum(abs2, J'; dims = 1); dims = 1) + end + cache.DᵀD = update_levenberg_marquardt_diagonal!!( + cache.DᵀD, Utils.safe_vec(cache.J_diag_cache) + ) + @bb @. cache.J_damped = cache.λ * cache.DᵀD + return cache.J_damped +end + +function InternalAPI.solve!( + cache::LevenbergMarquardtDampingCache, JᵀJ, fu, ::Val{true}; kwargs... +) + cache.DᵀD = update_levenberg_marquardt_diagonal!!(cache.DᵀD, JᵀJ) + @bb @. cache.J_damped = cache.λ * cache.DᵀD + return cache.J_damped +end + +function NonlinearSolveBase.callback_into_cache!( + topcache, cache::LevenbergMarquardtDampingCache, args... +) + if NonlinearSolveBase.last_step_accepted(topcache.trustregion_cache) && + NonlinearSolveBase.last_step_accepted(topcache.descent_cache) + cache.λ_factor = 1 / cache.decrease_factor + end + cache.λ *= cache.λ_factor + cache.λ_factor = cache.increase_factor +end + +""" + LevenbergMarquardtTrustRegion(b_uphill) + +Trust Region method for [`LevenbergMarquardt`](@ref). This method is tightly coupled with +the Levenberg-Marquardt method and works by directly updating the damping parameter instead +of specifying a trust region radius. + +### Arguments + + - `b_uphill`: a factor that determines if a step is accepted or rejected. The standard + choice in the Levenberg-Marquardt method is to accept all steps that decrease the cost + and reject all steps that increase the cost. Although this is a natural and safe choice, + it is often not the most efficient. Therefore downhill moves are always accepted, but + uphill moves are only conditionally accepted. To decide whether an uphill move will be + accepted at each iteration ``i``, we compute + ``\\beta_i = \\cos(v_{\\text{new}}, v_{\\text{old}})``, which denotes the cosine angle + between the proposed velocity ``v_{\\text{new}}`` and the velocity of the last accepted + step ``v_{\\text{old}}``. The idea is to accept uphill moves if the angle is small. To + specify, uphill moves are accepted if + ``(1-\\beta_i)^{b_{\\text{uphill}}} C_{i+1} \\le C_i``, where ``C_i`` is the cost at + iteration ``i``. Reasonable choices for `b_uphill` are `1.0` or `2.0`, with + `b_uphill = 2.0` allowing higher uphill moves than `b_uphill = 1.0`. When + `b_uphill = 0.0`, no uphill moves will be accepted. Defaults to `1.0`. See Section 4 of + [transtrum2012improvements](@citet). +""" +@concrete struct LevenbergMarquardtTrustRegion <: AbstractTrustRegionMethod + β_uphill +end + +function InternalAPI.init( + prob::AbstractNonlinearProblem, alg::LevenbergMarquardtTrustRegion, + f::NonlinearFunction, fu, u, p, args...; + stats, internalnorm::F = L2_NORM, kwargs... +) where {F} + T = promote_type(eltype(u), eltype(fu)) + @bb v = copy(u) + @bb u_cache = similar(u) + @bb fu_cache = similar(fu) + return LevenbergMarquardtTrustRegionCache( + f, p, T(Inf), v, T(Inf), internalnorm, T(alg.β_uphill), false, + u_cache, fu_cache, stats + ) +end + +@concrete mutable struct LevenbergMarquardtTrustRegionCache <: + AbstractTrustRegionMethodCache + f + p + loss_old + v_cache + norm_v_old + internalnorm + β_uphill + last_step_accepted::Bool + u_cache + fu_cache + stats::NLStats +end + +function InternalAPI.reinit!( + cache::LevenbergMarquardtTrustRegionCache; p = cache.p, u0 = cache.v_cache, kwargs... +) + cache.p = p + @bb copyto!(cache.v_cache, u0) + cache.loss_old = oftype(cache.loss_old, Inf) + cache.norm_v_old = oftype(cache.norm_v_old, Inf) + cache.last_step_accepted = false +end + +function InternalAPI.solve!( + cache::LevenbergMarquardtTrustRegionCache, J, fu, u, δu, descent_stats +) + # This should be true if Geodesic Acceleration is being used + v = hasfield(typeof(descent_stats), :v) ? descent_stats.v : δu + norm_v = cache.internalnorm(v) + β = dot(v, cache.v_cache) / (norm_v * cache.norm_v_old) + + @bb @. cache.u_cache = u + δu + cache.fu_cache = Utils.evaluate_f!!(cache.f, cache.fu_cache, cache.u_cache, cache.p) + cache.stats.nf += 1 + + loss = cache.internalnorm(cache.fu_cache) + + if (1 - β)^cache.β_uphill * loss ≤ cache.loss_old # Accept Step + cache.last_step_accepted = true + cache.norm_v_old = norm_v + @bb copyto!(cache.v_cache, v) + else + cache.last_step_accepted = false + end + + return cache.last_step_accepted, cache.u_cache, cache.fu_cache +end + +update_levenberg_marquardt_diagonal!!(y::Number, x::Number) = max(y, x) +function update_levenberg_marquardt_diagonal!!(y::Diagonal, x::AbstractVecOrMat) + if ArrayInterface.can_setindex(y.diag) + if ArrayInterface.fast_scalar_indexing(y.diag) + if ndims(x) == 1 + @simd ivdep for i in axes(x, 1) + @inbounds y.diag[i] = max(y.diag[i], x[i]) + end + else + @simd ivdep for i in axes(x, 1) + @inbounds y.diag[i] = max(y.diag[i], x[i, i]) + end + end + else + if ndims(x) == 1 + @. y.diag = max(y.diag, x) + else + y.diag .= max.(y.diag, @view(x[diagind(x)])) + end + end + return y + end + ndims(x) == 1 && return Diagonal(max.(y.diag, x)) + return Diagonal(max.(y.diag, @view(x[diagind(x)]))) +end + +init_levenberg_marquardt_diagonal(u::Number, v) = oftype(u, v) +init_levenberg_marquardt_diagonal(u::SArray, v) = Diagonal(ones(typeof(vec(u))) * v) +function init_levenberg_marquardt_diagonal(u, v) + d = similar(vec(u)) + d .= v + return Diagonal(d) +end diff --git a/src/algorithms/pseudo_transient.jl b/lib/NonlinearSolveFirstOrder/src/pseudo_transient.jl similarity index 62% rename from src/algorithms/pseudo_transient.jl rename to lib/NonlinearSolveFirstOrder/src/pseudo_transient.jl index 1e6b94763..a985cc100 100644 --- a/src/algorithms/pseudo_transient.jl +++ b/lib/NonlinearSolveFirstOrder/src/pseudo_transient.jl @@ -1,7 +1,9 @@ """ - PseudoTransient(; concrete_jac = nothing, linsolve = nothing, - linesearch = NoLineSearch(), precs = DEFAULT_PRECS, autodiff = nothing, - jvp_autodiff = nothing, vjp_autodiff = nothing) + PseudoTransient(; + concrete_jac = nothing, linesearch = missing, alpha_initial = 1e-3, + linsolve = nothing, precs = nothing, + autodiff = nothing, jvp_autodiff = nothing, vjp_autodiff = nothing + ) An implementation of PseudoTransient Method [coffey2003pseudotransient](@cite) that is used to solve steady state problems in an accelerated manner. It uses an adaptive time-stepping @@ -16,13 +18,22 @@ This implementation specifically uses "switched evolution relaxation" you are going to need more iterations to converge but it can be more stable. """ function PseudoTransient(; - concrete_jac = nothing, linsolve = nothing, linesearch = nothing, - precs = DEFAULT_PRECS, alpha_initial = 1e-3, autodiff = nothing, - jvp_autodiff = nothing, vjp_autodiff = nothing) - descent = DampedNewtonDescent(; linsolve, precs, initial_damping = alpha_initial, - damping_fn = SwitchedEvolutionRelaxation()) - return GeneralizedFirstOrderAlgorithm{concrete_jac, :PseudoTransient}(; - linesearch, descent, autodiff, vjp_autodiff, jvp_autodiff) + concrete_jac = nothing, linesearch = missing, alpha_initial = 1e-3, + linsolve = nothing, precs = nothing, + autodiff = nothing, jvp_autodiff = nothing, vjp_autodiff = nothing +) + return GeneralizedFirstOrderAlgorithm(; + linesearch, + descent = DampedNewtonDescent(; + linsolve, precs, initial_damping = alpha_initial, + damping_fn = SwitchedEvolutionRelaxation() + ), + autodiff, + jvp_autodiff, + vjp_autodiff, + concrete_jac, + name = :PseudoTransient + ) end """ @@ -44,27 +55,32 @@ Cache for the [`SwitchedEvolutionRelaxation`](@ref) method. internalnorm end -function requires_normal_form_jacobian(::Union{ +function NonlinearSolveBase.requires_normal_form_jacobian(::Union{ SwitchedEvolutionRelaxation, SwitchedEvolutionRelaxationCache}) return false end -function requires_normal_form_rhs(::Union{ + +function NonlinearSolveBase.requires_normal_form_rhs(::Union{ SwitchedEvolutionRelaxation, SwitchedEvolutionRelaxationCache}) return false end -function __internal_init( - prob::AbstractNonlinearProblem, f::SwitchedEvolutionRelaxation, initial_damping, - J, fu, u, args...; internalnorm::F = L2_NORM, kwargs...) where {F} +function InternalAPI.init( + prob::AbstractNonlinearProblem, f::SwitchedEvolutionRelaxation, + initial_damping, J, fu, u, args...; + internalnorm::F = L2_NORM, kwargs... +) where {F} T = promote_type(eltype(u), eltype(fu)) return SwitchedEvolutionRelaxationCache( - internalnorm(fu), T(1 / initial_damping), internalnorm) + internalnorm(fu), T(inv(initial_damping)), internalnorm + ) end (damping::SwitchedEvolutionRelaxationCache)(::Nothing) = damping.α⁻¹ -function __internal_solve!( - damping::SwitchedEvolutionRelaxationCache, J, fu, args...; kwargs...) +function InternalAPI.solve!( + damping::SwitchedEvolutionRelaxationCache, J, fu, args...; kwargs... +) res_norm = damping.internalnorm(fu) damping.α⁻¹ *= res_norm / damping.res_norm damping.res_norm = res_norm diff --git a/lib/NonlinearSolveFirstOrder/src/raphson.jl b/lib/NonlinearSolveFirstOrder/src/raphson.jl new file mode 100644 index 000000000..5c14f10b7 --- /dev/null +++ b/lib/NonlinearSolveFirstOrder/src/raphson.jl @@ -0,0 +1,22 @@ +""" + NewtonRaphson(; + concrete_jac = nothing, linsolve = nothing, linesearch = missing, precs = nothing, + autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing + ) + +An advanced NewtonRaphson implementation with support for efficient handling of sparse +matrices via colored automatic differentiation and preconditioned linear solvers. Designed +for large-scale and numerically-difficult nonlinear systems. +""" +function NewtonRaphson(; + concrete_jac = nothing, linsolve = nothing, linesearch = missing, precs = nothing, + autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing +) + return GeneralizedFirstOrderAlgorithm(; + linesearch, + descent = NewtonDescent(; linsolve, precs), + autodiff, vjp_autodiff, jvp_autodiff, + concrete_jac, + name = :NewtonRaphson + ) +end diff --git a/lib/NonlinearSolveFirstOrder/src/solve.jl b/lib/NonlinearSolveFirstOrder/src/solve.jl new file mode 100644 index 000000000..c9c8c77a8 --- /dev/null +++ b/lib/NonlinearSolveFirstOrder/src/solve.jl @@ -0,0 +1,324 @@ +""" + GeneralizedFirstOrderAlgorithm(; + descent, linesearch = missing, + trustregion = missing, autodiff = nothing, vjp_autodiff = nothing, + jvp_autodiff = nothing, max_shrink_times::Int = typemax(Int), + concrete_jac = Val(false), name::Symbol = :unknown + ) + +This is a Generalization of First-Order (uses Jacobian) Nonlinear Solve Algorithms. The most +common example of this is Newton-Raphson Method. + +First Order here refers to the order of differentiation, and should not be confused with the +order of convergence. + +### Keyword Arguments + + - `trustregion`: Globalization using a Trust Region Method. This needs to follow the + [`NonlinearSolveBase.AbstractTrustRegionMethod`](@ref) interface. + - `descent`: The descent method to use to compute the step. This needs to follow the + [`NonlinearSolveBase.AbstractDescentDirection`](@ref) interface. + - `max_shrink_times`: The maximum number of times the trust region radius can be shrunk + before the algorithm terminates. +""" +@concrete struct GeneralizedFirstOrderAlgorithm <: AbstractNonlinearSolveAlgorithm + linesearch + trustregion + descent + max_shrink_times::Int + + autodiff + vjp_autodiff + jvp_autodiff + + concrete_jac <: Union{Val{false}, Val{true}} + name::Symbol +end + +function GeneralizedFirstOrderAlgorithm(; + descent, linesearch = missing, trustregion = missing, autodiff = nothing, + vjp_autodiff = nothing, jvp_autodiff = nothing, max_shrink_times::Int = typemax(Int), + concrete_jac = Val(false), name::Symbol = :unknown +) + concrete_jac = concrete_jac isa Bool ? Val(concrete_jac) : + (concrete_jac isa Val ? concrete_jac : Val(concrete_jac !== nothing)) + return GeneralizedFirstOrderAlgorithm( + linesearch, trustregion, descent, max_shrink_times, + autodiff, vjp_autodiff, jvp_autodiff, + concrete_jac, name + ) +end + +@concrete mutable struct GeneralizedFirstOrderAlgorithmCache <: AbstractNonlinearSolveCache + # Basic Requirements + fu + u + u_cache + p + du # Aliased to `get_du(descent_cache)` + J # Aliased to `jac_cache.J` + alg <: GeneralizedFirstOrderAlgorithm + prob <: AbstractNonlinearProblem + globalization <: Union{Val{:LineSearch}, Val{:TrustRegion}, Val{:None}} + + # Internal Caches + jac_cache + descent_cache + linesearch_cache + trustregion_cache + + # Counters + stats::NLStats + nsteps::Int + maxiters::Int + maxtime + max_shrink_times::Int + + # Timer + timer + total_time::Float64 + + # State Affect + make_new_jacobian::Bool + + # Termination & Tracking + termination_cache + trace + retcode::ReturnCode.T + force_stop::Bool + kwargs +end + +function InternalAPI.reinit_self!( + cache::GeneralizedFirstOrderAlgorithmCache, args...; p = cache.p, u0 = cache.u, + alias_u0::Bool = false, maxiters = 1000, maxtime = nothing, kwargs... +) + Utils.reinit_common!(cache, u0, p, alias_u0) + + InternalAPI.reinit!(cache.stats) + cache.nsteps = 0 + cache.maxiters = maxiters + cache.maxtime = maxtime + cache.total_time = 0.0 + cache.force_stop = false + cache.retcode = ReturnCode.Default + cache.make_new_jacobian = true + + NonlinearSolveBase.reset!(cache.trace) + SciMLBase.reinit!( + cache.termination_cache, NonlinearSolveBase.get_fu(cache), + NonlinearSolveBase.get_u(cache); kwargs... + ) + NonlinearSolveBase.reset_timer!(cache.timer) + return +end + +NonlinearSolveBase.@internal_caches(GeneralizedFirstOrderAlgorithmCache, + :jac_cache, :descent_cache, :linesearch_cache, :trustregion_cache) + +function SciMLBase.__init( + prob::AbstractNonlinearProblem, alg::GeneralizedFirstOrderAlgorithm, args...; + stats = NLStats(0, 0, 0, 0, 0), alias_u0 = false, maxiters = 1000, + abstol = nothing, reltol = nothing, maxtime = nothing, + termination_condition = nothing, internalnorm = L2_NORM, + linsolve_kwargs = (;), kwargs... +) + @set! alg.autodiff = NonlinearSolveBase.select_jacobian_autodiff(prob, alg.autodiff) + provided_jvp_autodiff = alg.jvp_autodiff !== nothing + @set! alg.jvp_autodiff = if !provided_jvp_autodiff && alg.autodiff !== nothing && + (ADTypes.mode(alg.autodiff) isa ADTypes.ForwardMode || + ADTypes.mode(alg.autodiff) isa + ADTypes.ForwardOrReverseMode) + NonlinearSolveBase.select_forward_mode_autodiff(prob, alg.autodiff) + else + NonlinearSolveBase.select_forward_mode_autodiff(prob, alg.jvp_autodiff) + end + provided_vjp_autodiff = alg.vjp_autodiff !== nothing + @set! alg.vjp_autodiff = if !provided_vjp_autodiff && alg.autodiff !== nothing && + (ADTypes.mode(alg.autodiff) isa ADTypes.ReverseMode || + ADTypes.mode(alg.autodiff) isa + ADTypes.ForwardOrReverseMode) + NonlinearSolveBase.select_reverse_mode_autodiff(prob, alg.autodiff) + else + NonlinearSolveBase.select_reverse_mode_autodiff(prob, alg.vjp_autodiff) + end + + timer = get_timer_output() + @static_timeit timer "cache construction" begin + u = Utils.maybe_unaliased(prob.u0, alias_u0) + fu = Utils.evaluate_f(prob, u) + @bb u_cache = copy(u) + + linsolve = NonlinearSolveBase.get_linear_solver(alg.descent) + + abstol, reltol, termination_cache = NonlinearSolveBase.init_termination_cache( + prob, abstol, reltol, fu, u, termination_condition, Val(:regular) + ) + linsolve_kwargs = merge((; abstol, reltol), linsolve_kwargs) + + jac_cache = NonlinearSolveBase.construct_jacobian_cache( + prob, alg, prob.f, fu, u, prob.p; + stats, alg.autodiff, linsolve, alg.jvp_autodiff, alg.vjp_autodiff + ) + J = jac_cache(nothing) + + descent_cache = InternalAPI.init( + prob, alg.descent, J, fu, u; stats, abstol, reltol, internalnorm, + linsolve_kwargs, timer + ) + du = SciMLBase.get_du(descent_cache) + + has_linesearch = alg.linesearch !== missing && alg.linesearch !== nothing + has_trustregion = alg.trustregion !== missing && alg.trustregion !== nothing + + if has_trustregion && has_linesearch + error("TrustRegion and LineSearch methods are algorithmically incompatible.") + end + + globalization = Val(:None) + linesearch_cache = nothing + trustregion_cache = nothing + + if has_trustregion + NonlinearSolveBase.supports_trust_region(alg.descent) || + error("Trust Region not supported by $(alg.descent).") + trustregion_cache = InternalAPI.init( + prob, alg.trustregion, prob.f, fu, u, prob.p; + alg.vjp_autodiff, alg.jvp_autodiff, stats, internalnorm, kwargs... + ) + globalization = Val(:TrustRegion) + end + + if has_linesearch + NonlinearSolveBase.supports_line_search(alg.descent) || + error("Line Search not supported by $(alg.descent).") + linesearch_cache = CommonSolve.init( + prob, alg.linesearch, fu, u; stats, internalnorm, + autodiff = ifelse( + provided_jvp_autodiff, alg.jvp_autodiff, alg.vjp_autodiff + ), + kwargs... + ) + globalization = Val(:LineSearch) + end + + trace = NonlinearSolveBase.init_nonlinearsolve_trace( + prob, alg, u, fu, J, du; kwargs... + ) + + return GeneralizedFirstOrderAlgorithmCache( + fu, u, u_cache, prob.p, du, J, alg, prob, globalization, + jac_cache, descent_cache, linesearch_cache, trustregion_cache, + stats, 0, maxiters, maxtime, alg.max_shrink_times, timer, + 0.0, true, termination_cache, trace, ReturnCode.Default, false, kwargs + ) + end +end + +function InternalAPI.step!( + cache::GeneralizedFirstOrderAlgorithmCache; + recompute_jacobian::Union{Nothing, Bool} = nothing +) + @static_timeit cache.timer "jacobian" begin + if (recompute_jacobian === nothing || recompute_jacobian) && cache.make_new_jacobian + J = cache.jac_cache(cache.u) + new_jacobian = true + else + J = cache.jac_cache(nothing) + new_jacobian = false + end + end + + @static_timeit cache.timer "descent" begin + if cache.trustregion_cache !== nothing && + hasfield(typeof(cache.trustregion_cache), :trust_region) + descent_result = InternalAPI.solve!( + cache.descent_cache, J, cache.fu, cache.u; + new_jacobian, cache.trustregion_cache.trust_region, cache.kwargs... + ) + else + descent_result = InternalAPI.solve!( + cache.descent_cache, J, cache.fu, cache.u; new_jacobian, cache.kwargs... + ) + end + end + + if !descent_result.linsolve_success + if new_jacobian + # Jacobian Information is current and linear solve failed terminate the solve + cache.retcode = ReturnCode.InternalLinearSolveFailed + cache.force_stop = true + return + else + # Jacobian Information is not current and linear solve failed, recompute it + if !haskey(cache.kwargs, :verbose) || cache.kwargs[:verbose] + @warn "Linear Solve Failed but Jacobian Information is not current. \ + Retrying with updated Jacobian." + end + # In the 2nd call the `new_jacobian` is guaranteed to be `true`. + cache.make_new_jacobian = true + InternalAPI.step!(cache; recompute_jacobian = true, kwargs...) + return + end + end + + δu, descent_intermediates = descent_result.δu, descent_result.extras + + if descent_result.success + cache.make_new_jacobian = true + if cache.globalization isa Val{:LineSearch} + @static_timeit cache.timer "linesearch" begin + linesearch_sol = CommonSolve.solve!(cache.linesearch_cache, cache.u, δu) + linesearch_failed = !SciMLBase.successful_retcode(linesearch_sol.retcode) + α = linesearch_sol.step_size + end + if linesearch_failed + cache.retcode = ReturnCode.InternalLineSearchFailed + cache.force_stop = true + end + @static_timeit cache.timer "step" begin + @bb axpy!(α, δu, cache.u) + Utils.evaluate_f!(cache, cache.u, cache.p) + end + elseif cache.globalization isa Val{:TrustRegion} + @static_timeit cache.timer "trustregion" begin + tr_accepted, u_new, fu_new = InternalAPI.solve!( + cache.trustregion_cache, J, cache.fu, cache.u, δu, descent_intermediates + ) + if tr_accepted + @bb copyto!(cache.u, u_new) + @bb copyto!(cache.fu, fu_new) + α = true + else + α = false + cache.make_new_jacobian = false + end + if hasfield(typeof(cache.trustregion_cache), :shrink_counter) && + cache.trustregion_cache.shrink_counter > cache.max_shrink_times + cache.retcode = ReturnCode.ShrinkThresholdExceeded + cache.force_stop = true + end + end + elseif cache.globalization isa Val{:None} + @static_timeit cache.timer "step" begin + @bb axpy!(1, δu, cache.u) + Utils.evaluate_f!(cache, cache.u, cache.p) + end + α = true + else + error("Unknown Globalization Strategy: $(cache.globalization). Allowed values \ + are (:LineSearch, :TrustRegion, :None)") + end + NonlinearSolveBase.check_and_update!(cache, cache.fu, cache.u, cache.u_cache) + else + α = false + cache.make_new_jacobian = false + end + + update_trace!(cache, α) + @bb copyto!(cache.u_cache, cache.u) + + NonlinearSolveBase.callback_into_cache!(cache) + + return nothing +end diff --git a/src/globalization/trust_region.jl b/lib/NonlinearSolveFirstOrder/src/trust_region.jl similarity index 64% rename from src/globalization/trust_region.jl rename to lib/NonlinearSolveFirstOrder/src/trust_region.jl index 248f5307c..92d2655ac 100644 --- a/src/globalization/trust_region.jl +++ b/lib/NonlinearSolveFirstOrder/src/trust_region.jl @@ -1,93 +1,45 @@ """ - LevenbergMarquardtTrustRegion(b_uphill) - -Trust Region method for [`LevenbergMarquardt`](@ref). This method is tightly coupled with -the Levenberg-Marquardt method and works by directly updating the damping parameter instead -of specifying a trust region radius. - -### Arguments - - - `b_uphill`: a factor that determines if a step is accepted or rejected. The standard - choice in the Levenberg-Marquardt method is to accept all steps that decrease the cost - and reject all steps that increase the cost. Although this is a natural and safe choice, - it is often not the most efficient. Therefore downhill moves are always accepted, but - uphill moves are only conditionally accepted. To decide whether an uphill move will be - accepted at each iteration ``i``, we compute - ``\\beta_i = \\cos(v_{\\text{new}}, v_{\\text{old}})``, which denotes the cosine angle - between the proposed velocity ``v_{\\text{new}}`` and the velocity of the last accepted - step ``v_{\\text{old}}``. The idea is to accept uphill moves if the angle is small. To - specify, uphill moves are accepted if - ``(1-\\beta_i)^{b_{\\text{uphill}}} C_{i+1} \\le C_i``, where ``C_i`` is the cost at - iteration ``i``. Reasonable choices for `b_uphill` are `1.0` or `2.0`, with - `b_uphill = 2.0` allowing higher uphill moves than `b_uphill = 1.0`. When - `b_uphill = 0.0`, no uphill moves will be accepted. Defaults to `1.0`. See Section 4 of - [transtrum2012improvements](@citet). -""" -@concrete struct LevenbergMarquardtTrustRegion <: AbstractTrustRegionMethod - β_uphill -end - -function Base.show(io::IO, alg::LevenbergMarquardtTrustRegion) - print(io, "LevenbergMarquardtTrustRegion(β_uphill = $(alg.β_uphill))") -end - -@concrete mutable struct LevenbergMarquardtTrustRegionCache <: - AbstractTrustRegionMethodCache - f - p - loss_old - v_cache - norm_v_old - internalnorm - β_uphill - last_step_accepted::Bool - u_cache - fu_cache - stats::NLStats -end + TrustRegion(; + concrete_jac = nothing, linsolve = nothing, precs = nothing, + radius_update_scheme = RadiusUpdateSchemes.Simple, max_trust_radius::Real = 0 // 1, + initial_trust_radius::Real = 0 // 1, step_threshold::Real = 1 // 10000, + shrink_threshold::Real = 1 // 4, expand_threshold::Real = 3 // 4, + shrink_factor::Real = 1 // 4, expand_factor::Real = 2 // 1, + max_shrink_times::Int = 32, + vjp_autodiff = nothing, autodiff = nothing, jvp_autodiff = nothing + ) + +An advanced TrustRegion implementation with support for efficient handling of sparse +matrices via colored automatic differentiation and preconditioned linear solvers. Designed +for large-scale and numerically-difficult nonlinear systems. -function reinit_cache!(cache::LevenbergMarquardtTrustRegionCache, args...; - p = cache.p, u0 = cache.v_cache, kwargs...) - cache.p = p - @bb copyto!(cache.v_cache, u0) - cache.loss_old = oftype(cache.loss_old, Inf) - cache.norm_v_old = oftype(cache.norm_v_old, Inf) - cache.last_step_accepted = false -end - -function __internal_init( - prob::AbstractNonlinearProblem, alg::LevenbergMarquardtTrustRegion, f::F, fu, - u, p, args...; stats, internalnorm::IF = L2_NORM, kwargs...) where {F, IF} - T = promote_type(eltype(u), eltype(fu)) - @bb v = copy(u) - @bb u_cache = similar(u) - @bb fu_cache = similar(fu) - return LevenbergMarquardtTrustRegionCache(f, p, T(Inf), v, T(Inf), internalnorm, - alg.β_uphill, false, u_cache, fu_cache, stats) -end - -function __internal_solve!( - cache::LevenbergMarquardtTrustRegionCache, J, fu, u, δu, descent_stats) - # This should be true if Geodesic Acceleration is being used - v = hasfield(typeof(descent_stats), :v) ? descent_stats.v : δu - norm_v = cache.internalnorm(v) - β = dot(v, cache.v_cache) / (norm_v * cache.norm_v_old) - - @bb @. cache.u_cache = u + δu - cache.fu_cache = evaluate_f!!(cache.f, cache.fu_cache, cache.u_cache, cache.p) - cache.stats.nf += 1 - - loss = cache.internalnorm(cache.fu_cache) +### Keyword Arguments - if (1 - β)^cache.β_uphill * loss ≤ cache.loss_old # Accept Step - cache.last_step_accepted = true - cache.norm_v_old = norm_v - @bb copyto!(cache.v_cache, v) - else - cache.last_step_accepted = false - end + - `radius_update_scheme`: the scheme used to update the trust region radius. Defaults to + `RadiusUpdateSchemes.Simple`. See [`RadiusUpdateSchemes`](@ref) for more details. For a + review on trust region radius update schemes, see [yuan2015recent](@citet). - return cache.last_step_accepted, cache.u_cache, cache.fu_cache +For the remaining arguments, see [`NonlinearSolveFirstOrder.GenericTrustRegionScheme`](@ref) +documentation. +""" +function TrustRegion(; + concrete_jac = nothing, linsolve = nothing, precs = nothing, + radius_update_scheme = RadiusUpdateSchemes.Simple, max_trust_radius::Real = 0 // 1, + initial_trust_radius::Real = 0 // 1, step_threshold::Real = 1 // 10000, + shrink_threshold::Real = 1 // 4, expand_threshold::Real = 3 // 4, + shrink_factor::Real = 1 // 4, expand_factor::Real = 2 // 1, + max_shrink_times::Int = 32, + autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing +) + descent = Dogleg(; linsolve, precs) + trustregion = GenericTrustRegionScheme(; + method = radius_update_scheme, step_threshold, shrink_threshold, expand_threshold, + shrink_factor, expand_factor, initial_trust_radius, max_trust_radius + ) + return GeneralizedFirstOrderAlgorithm(; + trustregion, descent, autodiff, vjp_autodiff, jvp_autodiff, max_shrink_times, + concrete_jac, name = :TrustRegion + ) end # Don't Pollute the namespace @@ -195,10 +147,12 @@ end const RUS = RadiusUpdateSchemes """ - GenericTrustRegionScheme(; method = RadiusUpdateSchemes.Simple, + GenericTrustRegionScheme(; + method = RadiusUpdateSchemes.Simple, max_trust_radius = nothing, initial_trust_radius = nothing, step_threshold = nothing, shrink_threshold = nothing, expand_threshold = nothing, - shrink_factor = nothing, expand_factor = nothing) + shrink_factor = nothing, expand_factor = nothing + ) Trust Region Method that updates and stores the current trust region radius in `trust_region`. For any of the keyword arguments, if the value is `nothing`, then we use @@ -234,9 +188,8 @@ the value used in the respective paper. - `expand_factor`: the factor to expand the trust region radius with if `expand_threshold < r` (with `r` defined in `shrink_threshold`). Defaults to `2.0`. """ -@kwdef @concrete struct GenericTrustRegionScheme{M <: - RadiusUpdateSchemes.AbstractRadiusUpdateScheme} - method::M = RadiusUpdateSchemes.Simple +@kwdef @concrete struct GenericTrustRegionScheme <: AbstractTrustRegionMethod + method <: RUS.AbstractRadiusUpdateScheme = RUS.Simple step_threshold = nothing shrink_threshold = nothing shrink_factor = nothing @@ -246,8 +199,60 @@ the value used in the respective paper. initial_trust_radius = nothing end -function Base.show(io::IO, alg::GenericTrustRegionScheme) - print(io, "GenericTrustRegionScheme(method = $(alg.method))") +function InternalAPI.init( + prob::AbstractNonlinearProblem, alg::GenericTrustRegionScheme, f, fu, u, p, + args...; stats, internalnorm::F = L2_NORM, vjp_autodiff = nothing, + jvp_autodiff = nothing, kwargs... +) where {F} + T = promote_type(eltype(u), eltype(fu)) + u0_norm = internalnorm(u) + fu_norm = internalnorm(fu) + + # Common Setup + mtr = max_trust_radius(alg.max_trust_radius, T, alg.method, u, fu_norm) + itr = initial_trust_radius( + alg.initial_trust_radius, T, alg.method, mtr, u0_norm, fu_norm + ) + stt = step_threshold(alg.step_threshold, T, alg.method) + sht = shrink_threshold(alg.shrink_threshold, T, alg.method) + shf = shrink_factor(alg.shrink_factor, T, alg.method) + et = expand_threshold(alg.expand_threshold, T, alg.method) + ef = expand_factor(alg.expand_factor, T, alg.method) + + # Scheme Specific Setup + p1, p2, p3, p4 = get_parameters(T, alg.method) + ϵ = T(1e-8) + + vjp_operator = alg.method isa RUS.__Yuan || alg.method isa RUS.__Bastin ? + VecJacOperator(prob, fu, u; autodiff = vjp_autodiff) : nothing + + jvp_operator = alg.method isa RUS.__Bastin ? + JacVecOperator(prob, fu, u; autodiff = jvp_autodiff) : nothing + + if alg.method isa RUS.__Yuan + Jᵀfu_cache = StatefulJacobianOperator(vjp_operator, u, prob.p) * Utils.safe_vec(fu) + itr = T(p1 * internalnorm(Jᵀfu_cache)) + elseif u isa Number + Jᵀfu_cache = u + else + @bb Jᵀfu_cache = similar(u) + end + + if alg.method isa RUS.__Bastin + @bb δu_cache = similar(u) + else + δu_cache = nothing + end + + @bb u_cache = similar(u) + @bb fu_cache = similar(fu) + @bb Jδu_cache = similar(fu) + + return GenericTrustRegionSchemeCache( + alg.method, f, p, mtr, itr, itr, stt, sht, et, shf, ef, + p1, p2, p3, p4, ϵ, T(0), vjp_operator, jvp_operator, Jᵀfu_cache, Jδu_cache, + δu_cache, internalnorm, u_cache, fu_cache, false, 0, stats, alg + ) end @concrete mutable struct GenericTrustRegionSchemeCache <: AbstractTrustRegionMethodCache @@ -282,76 +287,72 @@ end alg end -function reinit_cache!( - cache::GenericTrustRegionSchemeCache, args...; u0 = nothing, p = cache.p, kwargs...) - T = eltype(cache.u_cache) +function InternalAPI.reinit!( + cache::GenericTrustRegionSchemeCache; p = cache.p, u0 = nothing, kwargs... +) cache.p = p if u0 !== nothing u0_norm = cache.internalnorm(u0) - cache.trust_region = __initial_trust_radius( - cache.alg.initial_trust_radius, T, cache.alg.method, - cache.max_trust_radius, u0_norm, u0_norm) # FIXME: scheme specific end cache.last_step_accepted = false cache.shrink_counter = 0 end # Defaults -for func in (:__max_trust_radius, :__initial_trust_radius, :__step_threshold, - :__shrink_threshold, :__shrink_factor, :__expand_threshold, :__expand_factor) - @eval begin - @inline function $(func)(val, ::Type{T}, args...) where {T} - val_T = T(val) - iszero(val_T) && return $(func)(nothing, T, args...) - return val_T - end +for func in ( + :max_trust_radius, :initial_trust_radius, :step_threshold, :shrink_threshold, + :shrink_factor, :expand_threshold, :expand_factor +) + @eval function $(func)(val, ::Type{T}, args...) where {T} + iszero(val) && return $(func)(nothing, T, args...) + return T(val) end end -@inline __max_trust_radius(::Nothing, ::Type{T}, method, u, fu_norm) where {T} = T(Inf) -@inline function __max_trust_radius( - ::Nothing, ::Type{T}, ::Union{RUS.__Simple, RUS.__NocedalWright}, +max_trust_radius(::Nothing, ::Type{T}, method, u, fu_norm) where {T} = T(Inf) +function max_trust_radius(::Nothing, ::Type{T}, ::Union{RUS.__Simple, RUS.__NocedalWright}, u, fu_norm) where {T} u_min, u_max = extrema(u) return max(T(fu_norm), u_max - u_min) end -@inline function __initial_trust_radius( - ::Nothing, ::Type{T}, method, max_tr, u0_norm, fu_norm) where {T} +function initial_trust_radius( + ::Nothing, ::Type{T}, method, max_tr, u0_norm, fu_norm +) where {T} method isa RUS.__NLsolve && return T(ifelse(u0_norm > 0, u0_norm, 1)) (method isa RUS.__Hei || method isa RUS.__Bastin) && return T(1) method isa RUS.__Fan && return T((fu_norm^0.99) / 10) return T(max_tr / 11) end -@inline function __step_threshold(::Nothing, ::Type{T}, method) where {T} +function step_threshold(::Nothing, ::Type{T}, method) where {T} method isa RUS.__Hei && return T(0) method isa RUS.__Yuan && return T(1 // 1000) method isa RUS.__Bastin && return T(1 // 20) return T(1 // 10000) end -@inline function __shrink_threshold(::Nothing, ::Type{T}, method) where {T} +function shrink_threshold(::Nothing, ::Type{T}, method) where {T} method isa RUS.__Hei && return T(0) (method isa RUS.__NLsolve || method isa RUS.__Bastin) && return T(1 // 20) return T(1 // 4) end -@inline function __expand_threshold(::Nothing, ::Type{T}, method) where {T} +function expand_threshold(::Nothing, ::Type{T}, method) where {T} method isa RUS.__NLsolve && return T(9 // 10) method isa RUS.__Hei && return T(0) method isa RUS.__Bastin && return T(9 // 10) return T(3 // 4) end -@inline function __shrink_factor(::Nothing, ::Type{T}, method) where {T} +function shrink_factor(::Nothing, ::Type{T}, method) where {T} method isa RUS.__NLsolve && return T(1 // 2) method isa RUS.__Hei && return T(0) method isa RUS.__Bastin && return T(1 // 20) return T(1 // 4) end -@inline function __get_parameters(::Type{T}, method) where {T} +function get_parameters(::Type{T}, method) where {T} method isa RUS.__NLsolve && return (T(1 // 2), T(0), T(0), T(0)) method isa RUS.__Hei && return (T(5), T(1 // 10), T(15 // 100), T(15 // 100)) method isa RUS.__Yuan && return (T(2), T(1 // 6), T(6), T(0)) @@ -360,80 +361,35 @@ end return (T(0), T(0), T(0), T(0)) end -@inline __expand_factor(::Nothing, ::Type{T}, method) where {T} = T(2) - -function __internal_init( - prob::AbstractNonlinearProblem, alg::GenericTrustRegionScheme, f::F, fu, u, - p, args...; stats, internalnorm::IF = L2_NORM, vjp_autodiff = nothing, - jvp_autodiff = nothing, kwargs...) where {F, IF} - T = promote_type(eltype(u), eltype(fu)) - u0_norm = internalnorm(u) - fu_norm = internalnorm(fu) - - # Common Setup - max_trust_radius = __max_trust_radius(alg.max_trust_radius, T, alg.method, u, fu_norm) - initial_trust_radius = __initial_trust_radius( - alg.initial_trust_radius, T, alg.method, max_trust_radius, u0_norm, fu_norm) - step_threshold = __step_threshold(alg.step_threshold, T, alg.method) - shrink_threshold = __shrink_threshold(alg.shrink_threshold, T, alg.method) - expand_threshold = __expand_threshold(alg.expand_threshold, T, alg.method) - shrink_factor = __shrink_factor(alg.shrink_factor, T, alg.method) - expand_factor = __expand_factor(alg.expand_factor, T, alg.method) - - # Scheme Specific Setup - p1, p2, p3, p4 = __get_parameters(T, alg.method) - ϵ = T(1e-8) +expand_factor(::Nothing, ::Type{T}, method) where {T} = T(2) - vjp_operator = alg.method isa RUS.__Yuan || alg.method isa RUS.__Bastin ? - VecJacOperator(prob, fu, u; autodiff = vjp_autodiff) : nothing - - jvp_operator = alg.method isa RUS.__Bastin ? - JacVecOperator(prob, fu, u; autodiff = jvp_autodiff) : nothing - - if alg.method isa RUS.__Yuan - Jᵀfu_cache = StatefulJacobianOperator(vjp_operator, u, prob.p) * _vec(fu) - initial_trust_radius = T(p1 * internalnorm(Jᵀfu_cache)) - else - if u isa Number - Jᵀfu_cache = u - else - @bb Jᵀfu_cache = similar(u) - end - end - - if alg.method isa RUS.__Bastin - @bb δu_cache = similar(u) - else - δu_cache = nothing - end - - @bb u_cache = similar(u) - @bb fu_cache = similar(fu) - @bb Jδu_cache = similar(fu) - - return GenericTrustRegionSchemeCache( - alg.method, f, p, max_trust_radius, initial_trust_radius, initial_trust_radius, - step_threshold, shrink_threshold, expand_threshold, shrink_factor, - expand_factor, p1, p2, p3, p4, ϵ, T(0), vjp_operator, jvp_operator, Jᵀfu_cache, - Jδu_cache, δu_cache, internalnorm, u_cache, fu_cache, false, 0, stats, alg) +function rfunc_adaptive_trust_region( + r::R, c2::R, M::R, γ1::R, γ2::R, β::R +) where {R <: Real} + return ifelse( + r ≥ c2, + (2 * (M - 1 - γ2) * atan(r - c2) + (1 + γ2)) / R(π), + (1 - γ1 - β) * (exp(r - c2) + β / (1 - γ1 - β)) + ) end -function __internal_solve!( - cache::GenericTrustRegionSchemeCache, J, fu, u, δu, descent_stats) +function InternalAPI.solve!( + cache::GenericTrustRegionSchemeCache, J, fu, u, δu, descent_stats +) T = promote_type(eltype(u), eltype(fu)) @bb @. cache.u_cache = u + δu - cache.fu_cache = evaluate_f!!(cache.f, cache.fu_cache, cache.u_cache, cache.p) + cache.fu_cache = Utils.evaluate_f!!(cache.f, cache.fu_cache, cache.u_cache, cache.p) cache.stats.nf += 1 if hasfield(typeof(descent_stats), :δuJᵀJδu) && !isnan(descent_stats.δuJᵀJδu) δuJᵀJδu = descent_stats.δuJᵀJδu else @bb cache.Jδu_cache = J × vec(δu) - δuJᵀJδu = __dot(cache.Jδu_cache, cache.Jδu_cache) + δuJᵀJδu = Utils.safe_dot(cache.Jδu_cache, cache.Jδu_cache) end @bb cache.Jᵀfu_cache = transpose(J) × vec(fu) num = (cache.internalnorm(cache.fu_cache)^2 - cache.internalnorm(fu)^2) / 2 - denom = __dot(δu, cache.Jᵀfu_cache) + δuJᵀJδu / 2 + denom = Utils.safe_dot(δu, cache.Jᵀfu_cache) + δuJᵀJδu / 2 cache.ρ = num / denom if cache.ρ > cache.step_threshold @@ -462,7 +418,8 @@ function __internal_solve!( cache.trust_region = cache.expand_factor * cache.internalnorm(δu) elseif cache.ρ ≥ cache.p1 cache.trust_region = max( - cache.trust_region, cache.expand_factor * cache.internalnorm(δu)) + cache.trust_region, cache.expand_factor * cache.internalnorm(δu) + ) end end elseif cache.method isa RUS.__NocedalWright @@ -477,9 +434,9 @@ function __internal_solve!( end end elseif cache.method isa RUS.__Hei - tr_new = __rfunc( - cache.ρ, cache.shrink_threshold, cache.p1, cache.p3, cache.p4, cache.p2) * - cache.internalnorm(δu) + tr_new = rfunc_adaptive_trust_region( + cache.ρ, cache.shrink_threshold, cache.p1, cache.p3, cache.p4, cache.p2 + ) * cache.internalnorm(δu) if tr_new < cache.trust_region cache.shrink_counter += 1 else @@ -514,15 +471,15 @@ function __internal_solve!( if cache.ρ > cache.step_threshold jvp_op = StatefulJacobianOperator(cache.jvp_operator, cache.u_cache, cache.p) vjp_op = StatefulJacobianOperator(cache.vjp_operator, cache.u_cache, cache.p) - @bb cache.Jδu_cache = jvp_op × vec(δu) + @bb cache.Jδu_cache = jvp_op × vec(cache.δu_cache) @bb cache.Jᵀfu_cache = vjp_op × vec(cache.fu_cache) - denom_1 = dot(_vec(δu), cache.Jᵀfu_cache) + denom_1 = dot(Utils.safe_vec(cache.Jᵀfu_cache), cache.Jᵀfu_cache) @bb cache.Jᵀfu_cache = vjp_op × vec(cache.Jδu_cache) - denom_2 = dot(_vec(δu), cache.Jᵀfu_cache) + denom_2 = dot(Utils.safe_vec(cache.Jᵀfu_cache), cache.Jᵀfu_cache) denom = denom_1 + denom_2 / 2 ρ = num / denom if ρ ≥ cache.expand_threshold - cache.trust_region = cache.p1 * cache.internalnorm(δu) + cache.trust_region = cache.p1 * cache.internalnorm(cache.δu_cache) end cache.shrink_counter = 0 else @@ -535,9 +492,3 @@ function __internal_solve!( return cache.last_step_accepted, cache.u_cache, cache.fu_cache end - -# R-function for adaptive trust region method -function __rfunc(r::R, c2::R, M::R, γ1::R, γ2::R, β::R) where {R <: Real} - return ifelse(r ≥ c2, (2 * (M - 1 - γ2) * atan(r - c2) + (1 + γ2)) / R(π), - (1 - γ1 - β) * (exp(r - c2) + β / (1 - γ1 - β))) -end diff --git a/lib/NonlinearSolveFirstOrder/test/least_squares_tests.jl b/lib/NonlinearSolveFirstOrder/test/least_squares_tests.jl new file mode 100644 index 000000000..bf6ed48c8 --- /dev/null +++ b/lib/NonlinearSolveFirstOrder/test/least_squares_tests.jl @@ -0,0 +1,56 @@ +@testsetup module CoreNLLSTesting + +using NonlinearSolveFirstOrder, LineSearch, ADTypes, LinearSolve +using LineSearches: LineSearches +using ForwardDiff, FiniteDiff, Zygote + +include("../../../common/common_nlls_testing.jl") + +linesearches = [] +for lsmethod in [ + LineSearches.Static(), LineSearches.HagerZhang(), LineSearches.MoreThuente(), + LineSearches.StrongWolfe(), LineSearches.BackTracking() +] + push!(linesearches, LineSearchesJL(; method = lsmethod)) +end +push!(linesearches, BackTracking()) + +solvers = [] +for linsolve in [nothing, LUFactorization(), KrylovJL_GMRES(), KrylovJL_LSMR()] + vjp_autodiffs = linsolve isa KrylovJL ? [nothing, AutoZygote(), AutoFiniteDiff()] : + [nothing] + for linesearch in linesearches, vjp_autodiff in vjp_autodiffs + push!(solvers, GaussNewton(; linsolve, linesearch, vjp_autodiff)) + end +end +append!(solvers, + [ + LevenbergMarquardt(), + LevenbergMarquardt(; linsolve = LUFactorization()), + LevenbergMarquardt(; linsolve = KrylovJL_GMRES()), + LevenbergMarquardt(; linsolve = KrylovJL_LSMR()) + ] +) +for radius_update_scheme in [ + RadiusUpdateSchemes.Simple, RadiusUpdateSchemes.NocedalWright, + RadiusUpdateSchemes.NLsolve, RadiusUpdateSchemes.Hei, + RadiusUpdateSchemes.Yuan, RadiusUpdateSchemes.Fan, RadiusUpdateSchemes.Bastin +] + push!(solvers, TrustRegion(; radius_update_scheme)) +end + +export solvers + +end + +@testitem "General NLLS Solvers" setup=[CoreNLLSTesting] tags=[:core] begin + using LinearAlgebra + + nlls_problems = [prob_oop, prob_iip, prob_oop_vjp, prob_iip_vjp] + + for prob in nlls_problems, solver in solvers + sol = solve(prob, solver; maxiters = 10000, abstol = 1e-6) + @test SciMLBase.successful_retcode(sol) + @test norm(sol.resid, 2) < 1e-6 + end +end diff --git a/test/misc/issues_tests.jl b/lib/NonlinearSolveFirstOrder/test/misc_tests.jl similarity index 89% rename from test/misc/issues_tests.jl rename to lib/NonlinearSolveFirstOrder/test/misc_tests.jl index 24e39943c..40fcb2c55 100644 --- a/test/misc/issues_tests.jl +++ b/lib/NonlinearSolveFirstOrder/test/misc_tests.jl @@ -1,4 +1,4 @@ -@testitem "Issue #451" tags=[:misc] begin +@testitem "Scalar Jacobians: Issue #451" tags=[:core] begin f(u, p) = u^2 - p jac_calls = 0 diff --git a/lib/NonlinearSolveFirstOrder/test/qa_tests.jl b/lib/NonlinearSolveFirstOrder/test/qa_tests.jl new file mode 100644 index 000000000..7bf163723 --- /dev/null +++ b/lib/NonlinearSolveFirstOrder/test/qa_tests.jl @@ -0,0 +1,20 @@ +@testitem "Aqua" tags=[:core] begin + using Aqua, NonlinearSolveFirstOrder + + Aqua.test_all( + NonlinearSolveFirstOrder; + piracies = false, ambiguities = false + ) + Aqua.test_piracies(NonlinearSolveFirstOrder) + Aqua.test_ambiguities(NonlinearSolveFirstOrder; recursive = false) +end + +@testitem "Explicit Imports" tags=[:core] begin + using ExplicitImports, NonlinearSolveFirstOrder + + @test check_no_implicit_imports( + NonlinearSolveFirstOrder; skip = (Base, Core, SciMLBase) + ) === nothing + @test check_no_stale_explicit_imports(NonlinearSolveFirstOrder) === nothing + @test check_all_qualified_accesses_via_owners(NonlinearSolveFirstOrder) === nothing +end diff --git a/lib/NonlinearSolveFirstOrder/test/rootfind_tests.jl b/lib/NonlinearSolveFirstOrder/test/rootfind_tests.jl new file mode 100644 index 000000000..f554e02d6 --- /dev/null +++ b/lib/NonlinearSolveFirstOrder/test/rootfind_tests.jl @@ -0,0 +1,458 @@ +@testsetup module CoreRootfindTesting + +include("../../../common/common_rootfind_testing.jl") + +end + +@testitem "NewtonRaphson" setup=[CoreRootfindTesting] tags=[:core] begin + using ADTypes, LineSearch, LinearAlgebra, Random, LinearSolve + using LineSearches: LineSearches + using BenchmarkTools: @ballocated + using StaticArrays: @SVector + using Zygote, Enzyme, ForwardDiff, FiniteDiff + + u0s = ([1.0, 1.0], @SVector[1.0, 1.0], 1.0) + + preconditioners = [ + u0 -> nothing, + u0 -> ((args...) -> (Diagonal(rand!(similar(u0))), nothing)) + ] + + @testset for ad in (AutoForwardDiff(), AutoZygote(), AutoFiniteDiff(), AutoEnzyme()) + @testset "$(nameof(typeof(linesearch)))" for linesearch in ( + LineSearchesJL(; method = LineSearches.Static(), autodiff = ad), + LineSearchesJL(; method = LineSearches.BackTracking(), autodiff = ad), + LineSearchesJL(; method = LineSearches.MoreThuente(), autodiff = ad), + LineSearchesJL(; method = LineSearches.StrongWolfe(), autodiff = ad), + LineSearchesJL(; method = LineSearches.HagerZhang(), autodiff = ad), + BackTracking(; autodiff = ad) + ) + @testset "[OOP] u0: $(typeof(u0))" for u0 in u0s + solver = NewtonRaphson(; linesearch, autodiff = ad) + sol = solve_oop(quadratic_f, u0; solver) + @test SciMLBase.successful_retcode(sol) + err = maximum(abs, quadratic_f(sol.u, 2.0)) + @test err < 1e-9 + + cache = init( + NonlinearProblem{false}(quadratic_f, u0, 2.0), solver, abstol = 1e-9 + ) + @test (@ballocated solve!($cache)) < 200 + end + + @testset "[IIP] u0: $(typeof(u0))" for u0 in ([1.0, 1.0],) + ad isa AutoZygote && continue + + @testset for preconditioner in preconditioners, + linsolve in (nothing, KrylovJL_GMRES(), \) + + precs = preconditioner(u0) + typeof(linsolve) <: typeof(\) && precs !== nothing && continue + + solver = NewtonRaphson(; linsolve, precs, linesearch, autodiff = ad) + + sol = solve_iip(quadratic_f!, u0; solver) + @test SciMLBase.successful_retcode(sol) + err = maximum(abs, quadratic_f(sol.u, 2.0)) + @test err < 1e-9 + + cache = init( + NonlinearProblem{true}(quadratic_f!, u0, 2.0), solver, abstol = 1e-9 + ) + @test (@ballocated solve!($cache)) ≤ 64 + end + end + end + end +end + +@testitem "NewtonRaphson: Iterator Interface" setup=[CoreRootfindTesting] tags=[:core] begin + p = range(0.01, 2, length = 200) + @test nlprob_iterator_interface(quadratic_f, p, false, NewtonRaphson()) ≈ sqrt.(p) + @test nlprob_iterator_interface(quadratic_f!, p, true, NewtonRaphson()) ≈ sqrt.(p) +end + +@testitem "NewtonRaphson Termination Conditions" setup=[CoreRootfindTesting] tags=[:core] begin + using StaticArrays: @SVector + + @testset "TC: $(nameof(typeof(termination_condition)))" for termination_condition in TERMINATION_CONDITIONS + @testset "u0: $(typeof(u0))" for u0 in ([1.0, 1.0], 1.0, @SVector([1.0, 1.0])) + probN = NonlinearProblem(quadratic_f, u0, 2.0) + sol = solve(probN, NewtonRaphson(); termination_condition) + err = maximum(abs, quadratic_f(sol.u, 2.0)) + @test err < 1e-9 + end + end +end + +@testitem "PseudoTransient" setup=[CoreRootfindTesting] tags=[:core] begin + using ADTypes, Random, LinearSolve, LinearAlgebra + using BenchmarkTools: @ballocated + using StaticArrays: @SVector + using Zygote, Enzyme, ForwardDiff, FiniteDiff + + preconditioners = [ + (u0) -> nothing, + u0 -> ((args...) -> (Diagonal(rand!(similar(u0))), nothing)) + ] + + @testset for ad in (AutoForwardDiff(), AutoZygote(), AutoFiniteDiff(), AutoEnzyme()) + u0s = ([1.0, 1.0], @SVector[1.0, 1.0], 1.0) + + @testset "[OOP] u0: $(typeof(u0))" for u0 in u0s + solver = PseudoTransient(; alpha_initial = 10.0, autodiff = ad) + sol = solve_oop(quadratic_f, u0; solver) + @test SciMLBase.successful_retcode(sol) + err = maximum(abs, quadratic_f(sol.u, 2.0)) + @test err < 1e-9 + + cache = init( + NonlinearProblem{false}(quadratic_f, u0, 2.0), solver, abstol = 1e-9 + ) + @test (@ballocated solve!($cache)) < 200 + end + + @testset "[IIP] u0: $(typeof(u0))" for u0 in ([1.0, 1.0],) + ad isa AutoZygote && continue + + @testset for preconditioner in preconditioners, + linsolve in (nothing, KrylovJL_GMRES(), \) + + precs = preconditioner(u0) + typeof(linsolve) <: typeof(\) && precs !== nothing && continue + + solver = PseudoTransient(; + alpha_initial = 10.0, linsolve, precs, autodiff = ad) + sol = solve_iip(quadratic_f!, u0; solver) + @test SciMLBase.successful_retcode(sol) + err = maximum(abs, quadratic_f(sol.u, 2.0)) + @test err < 1e-9 + + cache = init( + NonlinearProblem{true}(quadratic_f!, u0, 2.0), solver, abstol = 1e-9 + ) + @test (@ballocated solve!($cache)) ≤ 64 + end + end + end +end + +@testitem "PseudoTransient: Iterator Interface" setup=[CoreRootfindTesting] tags=[:core] begin + p = range(0.01, 2, length = 200) + @test nlprob_iterator_interface( + quadratic_f, p, false, PseudoTransient(; alpha_initial = 10.0) + ) ≈ sqrt.(p) + @test nlprob_iterator_interface( + quadratic_f!, p, true, PseudoTransient(; alpha_initial = 10.0) + ) ≈ sqrt.(p) +end + +@testitem "PseudoTransient Termination Conditions" setup=[CoreRootfindTesting] tags=[:core] begin + using StaticArrays: @SVector + + @testset "TC: $(nameof(typeof(termination_condition)))" for termination_condition in TERMINATION_CONDITIONS + @testset "u0: $(typeof(u0))" for u0 in ([1.0, 1.0], 1.0, @SVector([1.0, 1.0])) + probN = NonlinearProblem(quadratic_f, u0, 2.0) + sol = solve(probN, PseudoTransient(); termination_condition) + err = maximum(abs, quadratic_f(sol.u, 2.0)) + @test err < 1e-9 + end + end +end + +@testitem "TrustRegion" setup=[CoreRootfindTesting] tags=[:core] begin + using ADTypes, LinearSolve, LinearAlgebra + using BenchmarkTools: @ballocated + using StaticArrays: @SVector + using Zygote, Enzyme, ForwardDiff, FiniteDiff + + radius_update_schemes = [ + RadiusUpdateSchemes.Simple, RadiusUpdateSchemes.NocedalWright, + RadiusUpdateSchemes.NLsolve, RadiusUpdateSchemes.Hei, + RadiusUpdateSchemes.Yuan, RadiusUpdateSchemes.Fan, RadiusUpdateSchemes.Bastin + ] + + @testset for ad in (AutoForwardDiff(), AutoZygote(), AutoFiniteDiff(), AutoEnzyme()) + @testset for radius_update_scheme in radius_update_schemes, + linsolve in (nothing, LUFactorization(), KrylovJL_GMRES(), \) + + @testset "[OOP] u0: $(typeof(u0))" for u0 in ( + [1.0, 1.0], 1.0, @SVector[1.0, 1.0], 1.0) + abstol = ifelse(linsolve isa KrylovJL, 1e-6, 1e-9) + solver = TrustRegion(; radius_update_scheme, linsolve) + sol = solve_oop(quadratic_f, u0; solver, abstol) + @test SciMLBase.successful_retcode(sol) + err = maximum(abs, quadratic_f(sol.u, 2.0)) + @test err < abstol + + cache = init(NonlinearProblem{false}(quadratic_f, u0, 2.0), solver, abstol) + @test (@ballocated solve!($cache)) < 200 + end + + @testset "[IIP] u0: $(typeof(u0))" for u0 in ([1.0, 1.0],) + ad isa AutoZygote && continue + + abstol = ifelse(linsolve isa KrylovJL, 1e-6, 1e-9) + solver = TrustRegion(; radius_update_scheme, linsolve) + sol = solve_iip(quadratic_f!, u0; solver, abstol) + @test SciMLBase.successful_retcode(sol) + err = maximum(abs, quadratic_f(sol.u, 2.0)) + @test err < abstol + + cache = init(NonlinearProblem{true}(quadratic_f!, u0, 2.0), solver, abstol) + @test (@ballocated solve!($cache)) ≤ 64 + end + end + end +end + +@testitem "TrustRegion: Iterator Interface" setup=[CoreRootfindTesting] tags=[:core] begin + p = range(0.01, 2, length = 200) + @test nlprob_iterator_interface(quadratic_f, p, false, TrustRegion()) ≈ sqrt.(p) + @test nlprob_iterator_interface(quadratic_f!, p, true, TrustRegion()) ≈ sqrt.(p) +end + +@testitem "TrustRegion NewtonRaphson Fails" setup=[CoreRootfindTesting] tags=[:core] begin + u0 = [-10.0, -1.0, 1.0, 2.0, 3.0, 4.0, 10.0] + p = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] + sol = solve_oop(newton_fails, u0, p; solver = TrustRegion()) + @test SciMLBase.successful_retcode(sol) + @test all(abs.(newton_fails(sol.u, p)) .< 1e-9) +end + +@testitem "TrustRegion: Kwargs" setup=[CoreRootfindTesting] tags=[:core] begin + max_trust_radius = [10.0, 100.0, 1000.0] + initial_trust_radius = [10.0, 1.0, 0.1] + step_threshold = [0.0, 0.01, 0.25] + shrink_threshold = [0.25, 0.3, 0.5] + expand_threshold = [0.5, 0.8, 0.9] + shrink_factor = [0.1, 0.3, 0.5] + expand_factor = [1.5, 2.0, 3.0] + max_shrink_times = [10, 20, 30] + + list_of_options = zip( + max_trust_radius, initial_trust_radius, step_threshold, shrink_threshold, + expand_threshold, shrink_factor, expand_factor, max_shrink_times + ) + + for options in list_of_options + alg = TrustRegion(; + max_trust_radius = options[1], initial_trust_radius = options[2], + step_threshold = options[3], shrink_threshold = options[4], + expand_threshold = options[5], shrink_factor = options[6], + expand_factor = options[7], max_shrink_times = options[8] + ) + + sol = solve_oop(quadratic_f, [1.0, 1.0], 2.0; solver = alg) + @test SciMLBase.successful_retcode(sol) + err = maximum(abs, quadratic_f(sol.u, 2.0)) + @test err < 1e-9 + end +end + +@testitem "TrustRegion OOP / IIP Consistency" setup=[CoreRootfindTesting] tags=[:core] begin + maxiterations = [2, 3, 4, 5] + u0 = [1.0, 1.0] + @testset for radius_update_scheme in [ + RadiusUpdateSchemes.Simple, RadiusUpdateSchemes.NocedalWright, + RadiusUpdateSchemes.NLsolve, RadiusUpdateSchemes.Hei, + RadiusUpdateSchemes.Yuan, RadiusUpdateSchemes.Fan, RadiusUpdateSchemes.Bastin + ] + @testset for maxiters in maxiterations + solver = TrustRegion(; radius_update_scheme) + sol_iip = solve_iip(quadratic_f!, u0; solver, maxiters) + sol_oop = solve_oop(quadratic_f, u0; solver, maxiters) + @test sol_iip.u ≈ sol_oop.u + end + end +end + +@testitem "TrustRegion Termination Conditions" setup=[CoreRootfindTesting] tags=[:core] begin + using StaticArrays: @SVector + + @testset "TC: $(nameof(typeof(termination_condition)))" for termination_condition in TERMINATION_CONDITIONS + @testset "u0: $(typeof(u0))" for u0 in ([1.0, 1.0], 1.0, @SVector([1.0, 1.0])) + probN = NonlinearProblem(quadratic_f, u0, 2.0) + sol = solve(probN, TrustRegion(); termination_condition) + err = maximum(abs, quadratic_f(sol.u, 2.0)) + @test err < 1e-9 + end + end +end + +@testitem "LevenbergMarquardt" setup=[CoreRootfindTesting] tags=[:core] begin + using ADTypes, LinearSolve, LinearAlgebra + using BenchmarkTools: @ballocated + using StaticArrays: SVector, @SVector + using Zygote, Enzyme, ForwardDiff, FiniteDiff + + @testset for ad in (AutoForwardDiff(), AutoZygote(), AutoFiniteDiff(), AutoEnzyme()) + solver = LevenbergMarquardt(; autodiff = ad) + + @testset "[OOP] u0: $(typeof(u0))" for u0 in ([1.0, 1.0], 1.0, @SVector([1.0, 1.0])) + if ad isa ADTypes.AutoZygote && u0 isa SVector + # Zygote converts SVector to a Matrix that triggers a bug upstream + @test_broken solve_oop(quadratic_f, u0; solver) + continue + end + + sol = solve_oop(quadratic_f, u0; solver) + @test SciMLBase.successful_retcode(sol) + err = maximum(abs, quadratic_f(sol.u, 2.0)) + @test err < 1e-9 + + cache = init( + NonlinearProblem{false}(quadratic_f, u0, 2.0), + LevenbergMarquardt(), abstol = 1e-9 + ) + @test (@ballocated solve!($cache)) < 200 + end + + @testset "[IIP] u0: $(typeof(u0))" for u0 in ([1.0, 1.0],) + ad isa AutoZygote && continue + + sol = solve_iip(quadratic_f!, u0; solver) + @test SciMLBase.successful_retcode(sol) + err = maximum(abs, quadratic_f(sol.u, 2.0)) + @test err < 1e-9 + + cache = init( + NonlinearProblem{true}(quadratic_f!, u0, 2.0), + LevenbergMarquardt(), abstol = 1e-9 + ) + @test (@ballocated solve!($cache)) ≤ 64 + end + end +end + +@testitem "LevenbergMarquardt NewtonRaphson Fails" setup=[CoreRootfindTesting] tags=[:core] begin + u0 = [-10.0, -1.0, 1.0, 2.0, 3.0, 4.0, 10.0] + p = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] + sol = solve_oop(newton_fails, u0, p; solver = LevenbergMarquardt()) + @test SciMLBase.successful_retcode(sol) + @test all(abs.(newton_fails(sol.u, p)) .< 1e-9) +end + +@testitem "LevenbergMarquardt: Iterator Interface" setup=[CoreRootfindTesting] tags=[:core] begin + p = range(0.01, 2, length = 200) + @test nlprob_iterator_interface(quadratic_f, p, false, LevenbergMarquardt()) ≈ sqrt.(p) + @test nlprob_iterator_interface(quadratic_f!, p, true, LevenbergMarquardt()) ≈ sqrt.(p) +end + +@testitem "LevenbergMarquardt Termination Conditions" setup=[CoreRootfindTesting] tags=[:core] begin + using StaticArrays: @SVector + + @testset "TC: $(nameof(typeof(termination_condition)))" for termination_condition in TERMINATION_CONDITIONS + @testset "u0: $(typeof(u0))" for u0 in ([1.0, 1.0], 1.0, @SVector([1.0, 1.0])) + probN = NonlinearProblem(quadratic_f, u0, 2.0) + sol = solve(probN, LevenbergMarquardt(); termination_condition) + err = maximum(abs, quadratic_f(sol.u, 2.0)) + @test err < 1e-9 + end + end +end + +@testitem "LevenbergMarquardt: Kwargs" setup=[CoreRootfindTesting] tags=[:core] begin + damping_initial = [0.5, 2.0, 5.0] + damping_increase_factor = [1.5, 3.0, 10.0] + damping_decrease_factor = Float64[2, 5, 10.0] + finite_diff_step_geodesic = [0.02, 0.2, 0.3] + α_geodesic = [0.6, 0.8, 0.9] + b_uphill = Float64[0, 1, 2] + min_damping_D = [1e-12, 1e-9, 1e-4] + + list_of_options = zip( + damping_initial, damping_increase_factor, damping_decrease_factor, + finite_diff_step_geodesic, α_geodesic, b_uphill, min_damping_D + ) + for options in list_of_options + alg = LevenbergMarquardt(; + damping_initial = options[1], damping_increase_factor = options[2], + damping_decrease_factor = options[3], + finite_diff_step_geodesic = options[4], α_geodesic = options[5], + b_uphill = options[6], min_damping_D = options[7] + ) + + sol = solve_oop(quadratic_f, [1.0, 1.0], 2.0; solver = alg, maxiters = 10000) + @test SciMLBase.successful_retcode(sol) + err = maximum(abs, quadratic_f(sol.u, 2.0)) + @test err < 1e-9 + end +end + +@testitem "Simple Sparse AutoDiff" setup=[CoreRootfindTesting] tags=[:core] begin + using ADTypes, SparseConnectivityTracer, SparseMatrixColorings + + @testset for ad in (AutoForwardDiff(), AutoFiniteDiff(), AutoZygote(), AutoEnzyme()) + @testset for u0 in ([1.0, 1.0], 1.0) + prob = NonlinearProblem( + NonlinearFunction(quadratic_f; sparsity = TracerSparsityDetector()), u0, 2.0 + ) + + @testset "Newton Raphson" begin + sol = solve(prob, NewtonRaphson(; autodiff = ad)) + err = maximum(abs, quadratic_f(sol.u, 2.0)) + @test err < 1e-9 + end + + @testset "Trust Region" begin + sol = solve(prob, TrustRegion(; autodiff = ad)) + err = maximum(abs, quadratic_f(sol.u, 2.0)) + @test err < 1e-9 + end + end + end +end + +@testitem "Custom JVP" setup=[CoreRootfindTesting] tags=[:core] begin + using LinearAlgebra, LinearSolve, ADTypes + + function F(u::Vector{Float64}, p::Vector{Float64}) + Δ = Tridiagonal(-ones(99), 2 * ones(100), -ones(99)) + return u + 0.1 * u .* Δ * u - p + end + + function F!(du::Vector{Float64}, u::Vector{Float64}, p::Vector{Float64}) + Δ = Tridiagonal(-ones(99), 2 * ones(100), -ones(99)) + du .= u + 0.1 * u .* Δ * u - p + return nothing + end + + function JVP(v::Vector{Float64}, u::Vector{Float64}, p::Vector{Float64}) + Δ = Tridiagonal(-ones(99), 2 * ones(100), -ones(99)) + return v + 0.1 * (u .* Δ * v + v .* Δ * u) + end + + function JVP!( + du::Vector{Float64}, v::Vector{Float64}, u::Vector{Float64}, p::Vector{Float64}) + Δ = Tridiagonal(-ones(99), 2 * ones(100), -ones(99)) + du .= v + 0.1 * (u .* Δ * v + v .* Δ * u) + return nothing + end + + u0 = rand(100) + + prob = NonlinearProblem(NonlinearFunction{false}(F; jvp = JVP), u0, u0) + sol = solve(prob, NewtonRaphson(; linsolve = KrylovJL_GMRES()); abstol = 1e-13) + err = maximum(abs, sol.resid) + @test err < 1e-6 + + sol = solve( + prob, TrustRegion(; linsolve = KrylovJL_GMRES(), vjp_autodiff = AutoFiniteDiff()); + abstol = 1e-13 + ) + err = maximum(abs, sol.resid) + @test err < 1e-6 + + prob = NonlinearProblem(NonlinearFunction{true}(F!; jvp = JVP!), u0, u0) + sol = solve(prob, NewtonRaphson(; linsolve = KrylovJL_GMRES()); abstol = 1e-13) + err = maximum(abs, sol.resid) + @test err < 1e-6 + + sol = solve( + prob, TrustRegion(; linsolve = KrylovJL_GMRES(), vjp_autodiff = AutoFiniteDiff()); + abstol = 1e-13 + ) + err = maximum(abs, sol.resid) + @test err < 1e-6 +end diff --git a/lib/NonlinearSolveFirstOrder/test/runtests.jl b/lib/NonlinearSolveFirstOrder/test/runtests.jl new file mode 100644 index 000000000..d19d33de8 --- /dev/null +++ b/lib/NonlinearSolveFirstOrder/test/runtests.jl @@ -0,0 +1,23 @@ +using ReTestItems, NonlinearSolveFirstOrder, Hwloc, InteractiveUtils, Pkg + +@info sprint(InteractiveUtils.versioninfo) + +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( + NonlinearSolveFirstOrder; tags = (GROUP == "all" ? nothing : [Symbol(GROUP)]), + nworkers = RETESTITEMS_NWORKERS, nworker_threads = RETESTITEMS_NWORKER_THREADS, + testitem_timeout = 3600 +) diff --git a/test/misc/bruss_tests.jl b/lib/NonlinearSolveFirstOrder/test/sparsity_tests.jl similarity index 50% rename from test/misc/bruss_tests.jl rename to lib/NonlinearSolveFirstOrder/test/sparsity_tests.jl index 32a3fff68..3e4df284f 100644 --- a/test/misc/bruss_tests.jl +++ b/lib/NonlinearSolveFirstOrder/test/sparsity_tests.jl @@ -1,5 +1,6 @@ -@testitem "Brusselator 2D" tags=[:misc] begin - using LinearAlgebra, SparseArrays, SparseConnectivityTracer, ADTypes +@testitem "Brusselator 2D" tags=[:core] begin + using LinearAlgebra, SparseArrays, SparseConnectivityTracer, ADTypes, + SparseMatrixColorings const N = 32 const xyd_brusselator = range(0, stop = 1, length = N) @@ -48,7 +49,8 @@ prob_brusselator_2d_sparse = NonlinearProblem( NonlinearFunction(brusselator_2d_loop; sparsity = TracerSparsityDetector()), - u0, p) + u0, p + ) sol = solve(prob_brusselator_2d_sparse, NewtonRaphson(); abstol = 1e-8) @test norm(sol.resid, Inf) < 1e-8 @@ -63,6 +65,75 @@ @test norm(sol.resid, Inf) < 1e-8 sol = solve(prob_brusselator_2d, - NewtonRaphson(autodiff = AutoFiniteDiff()); abstol = 1e-8) + NewtonRaphson(autodiff = AutoFiniteDiff()); abstol = 1e-8 + ) @test norm(sol.resid, Inf) < 1e-8 end + +@testitem "Structured Jacobians" tags=[:core] begin + using SparseConnectivityTracer, BandedMatrices, LinearAlgebra, SparseArrays, + SparseMatrixColorings + + N = 16 + p = rand(N) + u0 = rand(N) + + function f!(du, u, p) + for i in 2:(length(u) - 1) + du[i] = u[i - 1] - 2u[i] + u[i + 1] + p[i] + end + du[1] = -2u[1] + u[2] + p[1] + du[end] = u[end - 1] - 2u[end] + p[end] + return nothing + end + + function f(u, p) + du = similar(u, promote_type(eltype(u), eltype(p))) + f!(du, u, p) + return du + end + + for nlf in (f, f!) + @testset "Dense AD" begin + nlprob = NonlinearProblem(NonlinearFunction(nlf), u0, p) + + cache = init(nlprob, NewtonRaphson(); abstol = 1e-9) + @test cache.jac_cache.J isa Matrix + sol = solve!(cache) + @test SciMLBase.successful_retcode(sol) + end + + @testset "Unstructured Sparse AD" begin + nlprob_autosparse = NonlinearProblem( + NonlinearFunction(nlf; sparsity = TracerSparsityDetector()), + u0, p) + + cache = init(nlprob_autosparse, NewtonRaphson(); abstol = 1e-9) + @test cache.jac_cache.J isa SparseMatrixCSC + sol = solve!(cache) + @test SciMLBase.successful_retcode(sol) + end + + @testset "Structured Sparse AD: Banded Jacobian" begin + jac_prototype = BandedMatrix(-1 => ones(N - 1), 0 => ones(N), 1 => ones(N - 1)) + nlprob_sparse_structured = NonlinearProblem( + NonlinearFunction(nlf; jac_prototype), u0, p) + + cache = init(nlprob_sparse_structured, NewtonRaphson(); abstol = 1e-9) + @test cache.jac_cache.J isa BandedMatrix + sol = solve!(cache) + @test SciMLBase.successful_retcode(sol) + end + + @testset "Structured Sparse AD: Tridiagonal Jacobian" begin + jac_prototype = Tridiagonal(ones(N - 1), ones(N), ones(N - 1)) + nlprob_sparse_structured = NonlinearProblem( + NonlinearFunction(nlf; jac_prototype), u0, p) + + cache = init(nlprob_sparse_structured, NewtonRaphson(); abstol = 1e-9) + @test cache.jac_cache.J isa Tridiagonal + sol = solve!(cache) + @test SciMLBase.successful_retcode(sol) + end + end +end diff --git a/lib/NonlinearSolveQuasiNewton/LICENSE b/lib/NonlinearSolveQuasiNewton/LICENSE new file mode 100644 index 000000000..46c972b17 --- /dev/null +++ b/lib/NonlinearSolveQuasiNewton/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2024 Avik Pal + +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/NonlinearSolveQuasiNewton/Project.toml b/lib/NonlinearSolveQuasiNewton/Project.toml new file mode 100644 index 000000000..2f00863d8 --- /dev/null +++ b/lib/NonlinearSolveQuasiNewton/Project.toml @@ -0,0 +1,76 @@ +name = "NonlinearSolveQuasiNewton" +uuid = "9a2c21bd-3a47-402d-9113-8faf9a0ee114" +authors = ["Avik Pal and contributors"] +version = "1.0.0" + +[deps] +ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9" +CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2" +ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471" +DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae" +MaybeInplace = "bb5d69b7-63fc-4a16-80bd-7e42200c7bdb" +NonlinearSolveBase = "be0214bd-f91f-a760-ac4e-3421ce2b2da0" +PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" +Reexport = "189a3867-3050-52da-a836-e630ba90ab69" +SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" +SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961" +StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" + +[compat] +ADTypes = "1.9.0" +Aqua = "0.8" +ArrayInterface = "7.16.0" +BenchmarkTools = "1.5.0" +CommonSolve = "0.2.4" +ConcreteStructs = "0.2.3" +DiffEqBase = "6.158.3" +Enzyme = "0.13.12" +ExplicitImports = "1.5" +FiniteDiff = "2.24" +ForwardDiff = "0.10.36" +Hwloc = "3" +InteractiveUtils = "<0.0.1, 1" +LineSearch = "0.1.4" +LineSearches = "7.3.0" +LinearAlgebra = "1.10" +LinearSolve = "2.36.1" +MaybeInplace = "0.1.4" +NonlinearProblemLibrary = "0.1.2" +NonlinearSolveBase = "1.1" +Pkg = "1.10" +PrecompileTools = "1.2" +ReTestItems = "1.24" +Reexport = "1" +SciMLBase = "2.58" +SciMLOperators = "0.3.11" +StableRNGs = "1" +StaticArrays = "1.9.8" +StaticArraysCore = "1.4.3" +Test = "1.10" +Zygote = "0.6.72" +julia = "1.10" + +[extras] +ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" +Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" +BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" +ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7" +FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +Hwloc = "0e44f5e4-bd66-52a0-8798-143a42290a1d" +InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" +LineSearch = "87fe0de2-c867-4266-b59a-2f0a94fc965b" +LineSearches = "d3d80556-e9d4-5f37-9878-2ab0fcc64255" +NonlinearProblemLibrary = "b7050fa9-e91f-4b37-bcee-a89a063da141" +Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" +ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823" +StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" +StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" + +[targets] +test = ["ADTypes", "Aqua", "BenchmarkTools", "Enzyme", "ExplicitImports", "FiniteDiff", "ForwardDiff", "Hwloc", "InteractiveUtils", "LineSearch", "LineSearches", "NonlinearProblemLibrary", "Pkg", "ReTestItems", "StableRNGs", "StaticArrays", "Test", "Zygote"] diff --git a/lib/NonlinearSolveQuasiNewton/src/NonlinearSolveQuasiNewton.jl b/lib/NonlinearSolveQuasiNewton/src/NonlinearSolveQuasiNewton.jl new file mode 100644 index 000000000..3b63615e8 --- /dev/null +++ b/lib/NonlinearSolveQuasiNewton/src/NonlinearSolveQuasiNewton.jl @@ -0,0 +1,52 @@ +module NonlinearSolveQuasiNewton + +using ConcreteStructs: @concrete +using PrecompileTools: @compile_workload, @setup_workload +using Reexport: @reexport + +using ArrayInterface: ArrayInterface +using StaticArraysCore: StaticArray, Size, MArray + +using CommonSolve: CommonSolve +using DiffEqBase: DiffEqBase # Needed for `init` / `solve` dispatches +using LinearAlgebra: LinearAlgebra, Diagonal, dot, diag +using LinearSolve: LinearSolve # Trigger Linear Solve extension in NonlinearSolveBase +using MaybeInplace: @bb +using NonlinearSolveBase: NonlinearSolveBase, AbstractNonlinearSolveAlgorithm, + AbstractNonlinearSolveCache, AbstractResetCondition, + AbstractResetConditionCache, AbstractApproximateJacobianStructure, + AbstractJacobianCache, AbstractJacobianInitialization, + AbstractApproximateJacobianUpdateRule, AbstractDescentDirection, + AbstractApproximateJacobianUpdateRuleCache, + Utils, InternalAPI, get_timer_output, @static_timeit, + update_trace!, L2_NORM, NewtonDescent +using SciMLBase: SciMLBase, AbstractNonlinearProblem, NLStats, ReturnCode +using SciMLOperators: AbstractSciMLOperator + +include("reset_conditions.jl") +include("structure.jl") +include("initialization.jl") + +include("broyden.jl") +include("lbroyden.jl") +include("klement.jl") + +include("solve.jl") + +@setup_workload begin + include("../../../common/nonlinear_problem_workloads.jl") + + algs = [Broyden(), Klement()] + + @compile_workload begin + @sync for prob in nonlinear_problems, alg in algs + Threads.@spawn CommonSolve.solve(prob, alg; abstol = 1e-2, verbose = false) + end + end +end + +@reexport using SciMLBase, NonlinearSolveBase + +export Broyden, LimitedMemoryBroyden, Klement, QuasiNewtonAlgorithm + +end diff --git a/lib/NonlinearSolveQuasiNewton/src/broyden.jl b/lib/NonlinearSolveQuasiNewton/src/broyden.jl new file mode 100644 index 000000000..ed487bf94 --- /dev/null +++ b/lib/NonlinearSolveQuasiNewton/src/broyden.jl @@ -0,0 +1,162 @@ +""" + Broyden(; + max_resets::Int = 100, linesearch = nothing, reset_tolerance = nothing, + init_jacobian::Val = Val(:identity), autodiff = nothing, alpha = nothing, + update_rule = Val(:good_broyden) + ) + +An implementation of `Broyden`'s Method [broyden1965class](@cite) with resetting and line +search. + +### Keyword Arguments + + - `max_resets`: the maximum number of resets to perform. Defaults to `100`. + + - `reset_tolerance`: the tolerance for the reset check. Defaults to + `sqrt(eps(real(eltype(u))))`. + - `alpha`: If `init_jacobian` is set to `Val(:identity)`, then the initial Jacobian + inverse is set to be `(αI)⁻¹`. Defaults to `nothing` which implies + `α = max(norm(u), 1) / (2 * norm(fu))`. + - `init_jacobian`: the method to use for initializing the jacobian. Defaults to + `Val(:identity)`. Choices include: + + + `Val(:identity)`: Identity Matrix. + + `Val(:true_jacobian)`: True Jacobian. This is a good choice for differentiable + problems. + - `update_rule`: Update Rule for the Jacobian. Choices are: + + + `Val(:good_broyden)`: Good Broyden's Update Rule + + `Val(:bad_broyden)`: Bad Broyden's Update Rule + + `Val(:diagonal)`: Only update the diagonal of the Jacobian. This algorithm may be + useful for specific problems, but whether it will work may depend strongly on the + problem +""" +function Broyden(; + max_resets::Int = 100, linesearch = nothing, reset_tolerance = nothing, + init_jacobian::Val = Val(:identity), autodiff = nothing, alpha = nothing, + update_rule = Val(:good_broyden) +) + return QuasiNewtonAlgorithm(; + linesearch, + descent = NewtonDescent(), + update_rule = broyden_update_rule(update_rule), + max_resets, + initialization = broyden_init(init_jacobian, update_rule, autodiff, alpha), + reinit_rule = NoChangeInStateReset(; reset_tolerance), + concrete_jac = Val(init_jacobian isa Val{:true_jacobian}), + name = :Broyden + ) +end + +function broyden_init(::Val{:identity}, ::Val{:diagonal}, autodiff, alpha) + return IdentityInitialization(alpha, DiagonalStructure()) +end +function broyden_init(::Val{:identity}, ::Val, autodiff, alpha) + IdentityInitialization(alpha, FullStructure()) +end +function broyden_init(::Val{:true_jacobian}, ::Val, autodiff, alpha) + return TrueJacobianInitialization(FullStructure(), autodiff) +end +function broyden_init(::Val{IJ}, ::Val{UR}, autodiff, alpha) where {IJ, UR} + error("Unknown combination of `init_jacobian = Val($(Meta.quot(IJ)))` and \ + `update_rule = Val($(Meta.quot(UR)))`. Please choose a valid combination.") +end + +broyden_update_rule(::Val{:good_broyden}) = GoodBroydenUpdateRule() +broyden_update_rule(::Val{:bad_broyden}) = BadBroydenUpdateRule() +broyden_update_rule(::Val{:diagonal}) = GoodBroydenUpdateRule() +function broyden_update_rule(::Val{UR}) where {UR} + error("Unknown update rule `update_rule = Val($(Meta.quot(UR)))`. Please choose a \ + valid update rule.") +end + +""" + BadBroydenUpdateRule() + +Broyden Update Rule corresponding to "bad broyden's method" [broyden1965class](@cite). +""" +struct BadBroydenUpdateRule <: AbstractApproximateJacobianUpdateRule end + +""" + GoodBroydenUpdateRule() + +Broyden Update Rule corresponding to "good broyden's method" [broyden1965class](@cite). +""" +struct GoodBroydenUpdateRule <: AbstractApproximateJacobianUpdateRule end + +for rType in (:BadBroydenUpdateRule, :GoodBroydenUpdateRule) + @eval function Base.getproperty(rule::$(rType), sym::Symbol) + sym == :store_inverse_jacobian && return Val(true) + return getfield(rule, sym) + end +end + +function InternalAPI.init( + prob::AbstractNonlinearProblem, + alg::Union{BadBroydenUpdateRule, GoodBroydenUpdateRule}, + J, fu, u, du, args...; internalnorm::F = L2_NORM, kwargs... +) where {F} + @bb J⁻¹dfu = similar(u) + @bb dfu = copy(fu) + if alg isa GoodBroydenUpdateRule || J isa Diagonal + @bb u_cache = similar(u) + else + u_cache = nothing + end + if J isa Diagonal + du_cache = nothing + else + @bb du_cache = similar(du) + end + return BroydenUpdateRuleCache(J⁻¹dfu, dfu, u_cache, du_cache, internalnorm, alg) +end + +@concrete mutable struct BroydenUpdateRuleCache <: + AbstractApproximateJacobianUpdateRuleCache + J⁻¹dfu + dfu + u_cache + du_cache + internalnorm + rule <: Union{BadBroydenUpdateRule, GoodBroydenUpdateRule} +end + +function InternalAPI.solve!( + cache::BroydenUpdateRuleCache, J⁻¹, fu, u, du; kwargs... +) + T = eltype(u) + @bb @. cache.dfu = fu - cache.dfu + @bb cache.J⁻¹dfu = J⁻¹ × vec(cache.dfu) + if cache.rule isa GoodBroydenUpdateRule + @bb cache.u_cache = transpose(J⁻¹) × vec(du) + denom = dot(du, cache.J⁻¹dfu) + rmul = transpose(Utils.safe_vec(cache.u_cache)) + else + denom = cache.internalnorm(cache.dfu)^2 + rmul = transpose(Utils.safe_vec(cache.dfu)) + end + @bb @. cache.du_cache = (du - cache.J⁻¹dfu) / ifelse(iszero(denom), T(1e-5), denom) + @bb J⁻¹ += vec(cache.du_cache) × rmul + @bb copyto!(cache.dfu, fu) + return J⁻¹ +end + +function InternalAPI.solve!( + cache::BroydenUpdateRuleCache, J⁻¹::Diagonal, fu, u, du; kwargs... +) + T = eltype(u) + @bb @. cache.dfu = fu - cache.dfu + J⁻¹_diag = Utils.restructure(cache.dfu, diag(J⁻¹)) + if cache.rule isa GoodBroydenUpdateRule + @bb @. cache.J⁻¹dfu = J⁻¹_diag * cache.dfu * du + denom = sum(cache.J⁻¹dfu) + @bb @. J⁻¹_diag += (du - cache.J⁻¹dfu) * du * J⁻¹_diag / + ifelse(iszero(denom), T(1e-5), denom) + else + denom = cache.internalnorm(cache.dfu)^2 + @bb @. J⁻¹_diag += (du - J⁻¹_diag * cache.dfu) * cache.dfu / + ifelse(iszero(denom), T(1e-5), denom) + end + @bb copyto!(cache.dfu, fu) + return Diagonal(vec(J⁻¹_diag)) +end diff --git a/lib/NonlinearSolveQuasiNewton/src/initialization.jl b/lib/NonlinearSolveQuasiNewton/src/initialization.jl new file mode 100644 index 000000000..f730b541b --- /dev/null +++ b/lib/NonlinearSolveQuasiNewton/src/initialization.jl @@ -0,0 +1,285 @@ +""" + InitializedApproximateJacobianCache( + J, structure, alg, cache, initialized::Bool, internalnorm + ) + +A cache for Approximate Jacobian. + +### Arguments + + - `J`: The current Jacobian. + - `structure`: The structure of the Jacobian. + - `alg`: The initialization algorithm. + - `cache`: The Jacobian cache [`NonlinearSolveBase.construct_jacobian_cache`](@ref) + (if needed). + - `initialized`: A boolean indicating whether the Jacobian has been initialized. + - `internalnorm`: The norm to be used. + +### Interface + +```julia +(cache::InitializedApproximateJacobianCache)(::Nothing) +``` + +Returns the current Jacobian `cache.J` with the proper `structure`. + +```julia +InternalAPI.solve!(cache::InitializedApproximateJacobianCache, fu, u, ::Val{reinit}) +``` + +Solves for the Jacobian `cache.J` and returns it. If `reinit` is `true`, then the Jacobian +is reinitialized. +""" +@concrete mutable struct InitializedApproximateJacobianCache <: AbstractJacobianCache + J + structure + alg + cache + initialized::Bool + internalnorm +end + +function InternalAPI.reinit_self!(cache::InitializedApproximateJacobianCache; kwargs...) + cache.initialized = false +end + +NonlinearSolveBase.@internal_caches InitializedApproximateJacobianCache :cache + +function (cache::InitializedApproximateJacobianCache)(::Nothing) + return NonlinearSolveBase.get_full_jacobian(cache, cache.structure, cache.J) +end + +function InternalAPI.solve!( + cache::InitializedApproximateJacobianCache, fu, u, reinit::Val +) + if reinit isa Val{true} || !cache.initialized + cache(cache.alg, fu, u) + cache.initialized = true + end + if NonlinearSolveBase.stores_full_jacobian(cache.structure) + full_J = cache.J + else + full_J = NonlinearSolveBase.get_full_jacobian(cache, cache.structure, cache.J) + end + return full_J +end + +""" + IdentityInitialization(alpha, structure) + +Initialize the Jacobian to be an Identity Matrix scaled by `alpha` and maintain the +structure as specified by `structure`. +""" +@concrete struct IdentityInitialization <: AbstractJacobianInitialization + alpha + structure +end + +function InternalAPI.init( + prob::AbstractNonlinearProblem, alg::IdentityInitialization, solver, f::F, + fu, u, p; internalnorm::IN = L2_NORM, kwargs... +) where {F, IN} + α = Utils.initial_jacobian_scaling_alpha(alg.alpha, u, fu, internalnorm) + if u isa Number + J = α + else + if alg.structure isa DiagonalStructure + @assert length(u)==length(fu) "Diagonal Jacobian Structure must be square!" + J = one.(Utils.safe_vec(fu)) .* α + else + # A simple trick to get the correct jacobian structure + J = alg.structure(Utils.make_identity!!(vec(fu) * vec(u)', α); alias = true) + end + end + return InitializedApproximateJacobianCache( + J, alg.structure, alg, nothing, true, internalnorm + ) +end + +function (cache::InitializedApproximateJacobianCache)( + alg::IdentityInitialization, fu, u +) + α = Utils.initial_jacobian_scaling_alpha(alg.alpha, u, fu, cache.internalnorm) + cache.J = Utils.make_identity!!(cache.J, α) + return +end + +""" + TrueJacobianInitialization(structure, autodiff) + +Initialize the Jacobian to be the true Jacobian and maintain the structure as specified +by `structure`. `autodiff` is used to compute the true Jacobian and if not specified we +make a selection automatically. +""" +@concrete struct TrueJacobianInitialization <: AbstractJacobianInitialization + structure + autodiff +end + +function InternalAPI.init( + prob::AbstractNonlinearProblem, alg::TrueJacobianInitialization, + solver, f::F, fu, u, p; stats, linsolve = missing, + internalnorm::IN = L2_NORM, kwargs... +) where {F, IN} + autodiff = NonlinearSolveBase.select_jacobian_autodiff(prob, alg.autodiff) + jac_cache = NonlinearSolveBase.construct_jacobian_cache( + prob, solver, prob.f, fu, u, p; stats, autodiff, linsolve + ) + J = alg.structure(jac_cache(nothing)) + return InitializedApproximateJacobianCache( + J, alg.structure, alg, jac_cache, false, internalnorm + ) +end + +function (cache::InitializedApproximateJacobianCache)(::TrueJacobianInitialization, fu, u) + cache.J = cache.structure(cache.J, cache.cache(u)) + return +end + +""" + BroydenLowRankInitialization(alpha, threshold::Val) + +An initialization for `LimitedMemoryBroyden` that uses a low rank approximation of the +Jacobian. The low rank updates to the Jacobian matrix corresponds to what SciPy calls +["simple"](https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.broyden2.html#scipy-optimize-broyden2). +""" +@concrete struct BroydenLowRankInitialization <: AbstractJacobianInitialization + alpha + threshold <: Val +end + +NonlinearSolveBase.jacobian_initialized_preinverted(::BroydenLowRankInitialization) = true + +function InternalAPI.init( + prob::AbstractNonlinearProblem, alg::BroydenLowRankInitialization, + solver, f::F, fu, u, p; + internalnorm::IN = L2_NORM, maxiters = 1000, kwargs... +) where {F, IN} + if u isa Number # Use the standard broyden + return InternalAPI.init( + prob, IdentityInitialization(true, FullStructure()), + solver, f, fu, u, p; internalnorm, maxiters, kwargs... + ) + end + # Pay to cost of slightly more allocations to prevent type-instability for StaticArrays + α = inv(Utils.initial_jacobian_scaling_alpha(alg.alpha, u, fu, internalnorm)) + if u isa StaticArray + J = BroydenLowRankJacobian(fu, u; alg.threshold, alpha = α) + else + threshold = min(Utils.unwrap_val(alg.threshold), maxiters) + if threshold > length(u) + @warn "`threshold` is larger than the size of the state, which may cause \ + numerical instability. Consider reducing `threshold`." + end + J = BroydenLowRankJacobian(fu, u; threshold, alpha = α) + end + return InitializedApproximateJacobianCache( + J, FullStructure(), alg, nothing, true, internalnorm + ) +end + +function (cache::InitializedApproximateJacobianCache)( + alg::BroydenLowRankInitialization, fu, u +) + α = Utils.initial_jacobian_scaling_alpha(alg.alpha, u, fu, cache.internalnorm) + cache.J.idx = 0 + cache.J.alpha = inv(α) + return +end + +""" + BroydenLowRankJacobian{T}(U, Vᵀ, idx, cache, alpha) + +Low Rank Approximation of the Jacobian Matrix. Currently only used for +[`LimitedMemoryBroyden`](@ref). This computes the Jacobian as ``U \\times V^T``. +""" +@concrete mutable struct BroydenLowRankJacobian{T} <: AbstractSciMLOperator{T} + U + Vᵀ + idx::Int + cache + alpha +end + +Utils.maybe_pinv!!(workspace, A::BroydenLowRankJacobian) = A # Already Inverted form + +function get_components(op::BroydenLowRankJacobian) + op.idx ≥ size(op.U, 2) && return op.cache, op.U, transpose(op.Vᵀ) + cache = op.cache === nothing ? op.cache : view(op.cache, 1:(op.idx)) + return cache, view(op.U, :, 1:(op.idx)), transpose(view(op.Vᵀ, :, 1:(op.idx))) +end + +Base.size(op::BroydenLowRankJacobian) = size(op.U, 1), size(op.Vᵀ, 1) + +function Base.transpose(op::BroydenLowRankJacobian{T}) where {T} + return BroydenLowRankJacobian{T}(op.Vᵀ, op.U, op.idx, op.cache, op.alpha) +end +Base.adjoint(op::BroydenLowRankJacobian{<:Real}) = transpose(op) + +# Storing the transpose to ensure contiguous memory on splicing +function BroydenLowRankJacobian( + fu::StaticArray, u::StaticArray; alpha = true, threshold::Val = Val(10) +) + T = promote_type(eltype(u), eltype(fu)) + U = MArray{Tuple{prod(Size(fu)), Utils.unwrap_val(threshold)}, T}(undef) + Vᵀ = MArray{Tuple{prod(Size(u)), Utils.unwrap_val(threshold)}, T}(undef) + return BroydenLowRankJacobian{T}(U, Vᵀ, 0, nothing, T(alpha)) +end + +function BroydenLowRankJacobian(fu, u; threshold::Int = 10, alpha = true) + T = promote_type(eltype(u), eltype(fu)) + U = Utils.safe_similar(fu, T, length(fu), threshold) + Vᵀ = Utils.safe_similar(u, T, length(u), threshold) + cache = Utils.safe_similar(u, T, threshold) + return BroydenLowRankJacobian{T}(U, Vᵀ, 0, cache, T(alpha)) +end + +function Base.:*(J::BroydenLowRankJacobian, x::AbstractVector) + J.idx == 0 && return -x + _, U, Vᵀ = get_components(J) + return U * (Vᵀ * x) .- J.alpha .* x +end + +function LinearAlgebra.mul!(y::AbstractVector, J::BroydenLowRankJacobian, x::AbstractVector) + if J.idx == 0 + @. y = -J.alpha * x + return y + end + cache, U, Vᵀ = get_components(J) + @bb cache = Vᵀ × x + LinearAlgebra.mul!(y, U, cache) + @bb @. y -= J.alpha * x + return y +end + +function Base.:*(x::AbstractVector, J::BroydenLowRankJacobian) + J.idx == 0 && return -x + _, U, Vᵀ = get_components(J) + return Vᵀ' * (U' * x) .- J.alpha .* x +end + +function LinearAlgebra.mul!(y::AbstractVector, x::AbstractVector, J::BroydenLowRankJacobian) + if J.idx == 0 + @. y = -J.alpha * x + return y + end + cache, U, Vᵀ = get_components(J) + @bb cache = transpose(U) × x + LinearAlgebra.mul!(y, transpose(Vᵀ), cache) + @bb @. y -= J.alpha * x + return y +end + +function LinearAlgebra.mul!( + J::BroydenLowRankJacobian, u::AbstractArray, vᵀ::LinearAlgebra.AdjOrTransAbsVec, + α::Bool, β::Bool +) + @assert α & β + idx_update = mod1(J.idx + 1, size(J.U, 2)) + copyto!(@view(J.U[:, idx_update]), Utils.safe_vec(u)) + copyto!(@view(J.Vᵀ[:, idx_update]), Utils.safe_vec(vᵀ)) + J.idx += 1 + return J +end + +ArrayInterface.restructure(::BroydenLowRankJacobian, J::BroydenLowRankJacobian) = J diff --git a/src/algorithms/klement.jl b/lib/NonlinearSolveQuasiNewton/src/klement.jl similarity index 56% rename from src/algorithms/klement.jl rename to lib/NonlinearSolveQuasiNewton/src/klement.jl index be94efc9c..4bd09f45e 100644 --- a/src/algorithms/klement.jl +++ b/lib/NonlinearSolveQuasiNewton/src/klement.jl @@ -1,7 +1,9 @@ """ - Klement(; max_resets = 100, linsolve = nothing, linesearch = nothing, - precs = DEFAULT_PRECS, alpha = nothing, init_jacobian::Val = Val(:identity), - autodiff = nothing) + Klement(; + max_resets = 100, linsolve = nothing, linesearch = nothing, + precs = nothing, alpha = nothing, init_jacobian::Val = Val(:identity), + autodiff = nothing + ) An implementation of `Klement` [klement2014using](@citep) with line search, preconditioning and customizable linear solves. It is recommended to use [`Broyden`](@ref) for most problems @@ -24,16 +26,23 @@ over this. + `Val(:true_jacobian_diagonal)`: Diagonal of True Jacobian. This is a good choice for differentiable problems. """ -function Klement(; max_resets::Int = 100, linsolve = nothing, alpha = nothing, - linesearch = nothing, precs = DEFAULT_PRECS, - autodiff = nothing, init_jacobian::Val = Val(:identity)) - initialization = klement_init(init_jacobian, autodiff, alpha) - CJ = init_jacobian isa Val{:true_jacobian} || - init_jacobian isa Val{:true_jacobian_diagonal} - return ApproximateJacobianSolveAlgorithm{CJ, :Klement}(; - linesearch, descent = NewtonDescent(; linsolve, precs), +function Klement(; + max_resets = 100, linsolve = nothing, linesearch = nothing, + precs = nothing, alpha = nothing, init_jacobian::Val = Val(:identity), + autodiff = nothing +) + concrete_jac = Val(init_jacobian isa Val{:true_jacobian} || + init_jacobian isa Val{:true_jacobian_diagonal}) + return QuasiNewtonAlgorithm(; + linesearch, + descent = NewtonDescent(; linsolve, precs), update_rule = KlementUpdateRule(), - reinit_rule = IllConditionedJacobianReset(), max_resets, initialization) + reinit_rule = IllConditionedJacobianReset(), + max_resets, + initialization = klement_init(init_jacobian, autodiff, alpha), + concrete_jac, + name = :Klement + ) end function klement_init(::Val{:identity}, autodiff, alpha) @@ -50,55 +59,22 @@ function klement_init(::Val{IJ}, autodiff, alpha) where {IJ} `init_jacobian`.") end -# Essentially checks ill conditioned Jacobian -""" - IllConditionedJacobianReset() - -Recommend resetting the Jacobian if the current jacobian is ill-conditioned. This is used -in [`Klement`](@ref). -""" -struct IllConditionedJacobianReset <: AbstractResetCondition end - -@concrete struct IllConditionedJacobianResetCache - condition_number_threshold -end - -function __internal_init(alg::IllConditionedJacobianReset, J, fu, u, du, args...; kwargs...) - condition_number_threshold = if J isa AbstractMatrix - inv(eps(real(eltype(J)))^(1 // 2)) - else - nothing - end - return IllConditionedJacobianResetCache(condition_number_threshold) -end - -function __internal_solve!(cache::IllConditionedJacobianResetCache, J, fu, u, du) - J isa Number && return iszero(J) - J isa Diagonal && return any(iszero, diag(J)) - J isa AbstractMatrix && return cond(J) ≥ cache.condition_number_threshold - J isa AbstractVector && return any(iszero, J) - return false -end - -# Update Rule """ KlementUpdateRule() Update rule for [`Klement`](@ref). """ -@concrete struct KlementUpdateRule <: AbstractApproximateJacobianUpdateRule{false} end +struct KlementUpdateRule <: AbstractApproximateJacobianUpdateRule end -@concrete mutable struct KlementUpdateRuleCache <: - AbstractApproximateJacobianUpdateRuleCache{false} - Jdu - J_cache - J_cache_2 - Jdu_cache - fu_cache +function Base.getproperty(rule::KlementUpdateRule, sym::Symbol) + sym == :store_inverse_jacobian && return Val(false) + return getfield(rule, sym) end -function __internal_init(prob::AbstractNonlinearProblem, alg::KlementUpdateRule, - J, fu, u, du, args...; kwargs...) +function InternalAPI.init( + prob::AbstractNonlinearProblem, alg::KlementUpdateRule, + J, fu, u, du, args...; kwargs... +) @bb Jdu = similar(fu) if J isa Diagonal || J isa Number J_cache, J_cache_2, Jdu_cache = nothing, nothing, nothing @@ -108,29 +84,43 @@ function __internal_init(prob::AbstractNonlinearProblem, alg::KlementUpdateRule, @bb Jdu_cache = similar(Jdu) end @bb fu_cache = copy(fu) - return KlementUpdateRuleCache(Jdu, J_cache, J_cache_2, Jdu_cache, fu_cache) + return KlementUpdateRuleCache(Jdu, J_cache, J_cache_2, Jdu_cache, fu_cache, alg) +end + +@concrete mutable struct KlementUpdateRuleCache <: + AbstractApproximateJacobianUpdateRuleCache + Jdu + J_cache + J_cache_2 + Jdu_cache + fu_cache + rule <: KlementUpdateRule end -function __internal_solve!(cache::KlementUpdateRuleCache, J::Number, fu, u, du) +function InternalAPI.solve!( + cache::KlementUpdateRuleCache, J::Number, fu, u, du; kwargs... +) Jdu = J^2 * du^2 J = J + ((fu - cache.fu_cache - J * du) / ifelse(iszero(Jdu), 1e-5, Jdu)) * du * J^2 cache.fu_cache = fu return J end -function __internal_solve!(cache::KlementUpdateRuleCache, J_::Diagonal, fu, u, du) +function InternalAPI.solve!( + cache::KlementUpdateRuleCache, J::Diagonal, fu, u, du; kwargs... +) T = eltype(u) - J = _restructure(u, diag(J_)) + J = Utils.restructure(u, diag(J)) @bb @. cache.Jdu = (J^2) * (du^2) @bb @. J += ((fu - cache.fu_cache - J * du) / - ifelse(iszero(cache.Jdu), T(1e-5), cache.Jdu)) * - du * - (J^2) + ifelse(iszero(cache.Jdu), T(1e-5), cache.Jdu)) * du * (J^2) @bb copyto!(cache.fu_cache, fu) return Diagonal(vec(J)) end -function __internal_solve!(cache::KlementUpdateRuleCache, J::AbstractMatrix, fu, u, du) +function InternalAPI.solve!( + cache::KlementUpdateRuleCache, J::AbstractMatrix, fu, u, du; kwargs... +) T = eltype(u) @bb @. cache.J_cache = J'^2 @bb @. cache.Jdu = du^2 @@ -138,7 +128,7 @@ function __internal_solve!(cache::KlementUpdateRuleCache, J::AbstractMatrix, fu, @bb cache.Jdu = J × vec(du) @bb @. cache.fu_cache = (fu - cache.fu_cache - cache.Jdu) / ifelse(iszero(cache.Jdu_cache), T(1e-5), cache.Jdu_cache) - @bb cache.J_cache = vec(cache.fu_cache) × transpose(_vec(du)) + @bb cache.J_cache = vec(cache.fu_cache) × transpose(Utils.safe_vec(du)) @bb @. cache.J_cache *= J @bb cache.J_cache_2 = cache.J_cache × J @bb J .+= cache.J_cache_2 diff --git a/lib/NonlinearSolveQuasiNewton/src/lbroyden.jl b/lib/NonlinearSolveQuasiNewton/src/lbroyden.jl new file mode 100644 index 000000000..886905a5b --- /dev/null +++ b/lib/NonlinearSolveQuasiNewton/src/lbroyden.jl @@ -0,0 +1,35 @@ +""" + LimitedMemoryBroyden(; + max_resets::Int = 3, linesearch = nothing, threshold::Val = Val(10), + reset_tolerance = nothing, alpha = nothing + ) + +An implementation of `LimitedMemoryBroyden` [ziani2008autoadaptative](@cite) with resetting +and line search. + +### Keyword Arguments + + - `max_resets`: the maximum number of resets to perform. Defaults to `3`. + - `reset_tolerance`: the tolerance for the reset check. Defaults to + `sqrt(eps(real(eltype(u))))`. + - `threshold`: the number of vectors to store in the low rank approximation. Defaults + to `Val(10)`. + - `alpha`: The initial Jacobian inverse is set to be `(αI)⁻¹`. Defaults to `nothing` + which implies `α = max(norm(u), 1) / (2 * norm(fu))`. +""" +function LimitedMemoryBroyden(; + max_resets::Int = 3, linesearch = nothing, threshold::Union{Val, Int} = Val(10), + reset_tolerance = nothing, alpha = nothing +) + threshold isa Int && (threshold = Val(threshold)) + return QuasiNewtonAlgorithm(; + linesearch, + descent = NewtonDescent(), + update_rule = GoodBroydenUpdateRule(), + max_resets, + initialization = BroydenLowRankInitialization(alpha, threshold), + reinit_rule = NoChangeInStateReset(; reset_tolerance), + name = :LimitedMemoryBroyden, + concrete_jac = Val(false) + ) +end diff --git a/lib/NonlinearSolveQuasiNewton/src/reset_conditions.jl b/lib/NonlinearSolveQuasiNewton/src/reset_conditions.jl new file mode 100644 index 000000000..b8d92237c --- /dev/null +++ b/lib/NonlinearSolveQuasiNewton/src/reset_conditions.jl @@ -0,0 +1,120 @@ +""" + NoChangeInStateReset(; + nsteps::Int = 3, reset_tolerance = nothing, + check_du::Bool = true, check_dfu::Bool = true + ) + +Recommends a reset if the state or the function value has not changed significantly in +`nsteps` steps. This is used in [`Broyden`](@ref). + +### Keyword Arguments + + - `nsteps`: the number of steps to check for no change. Defaults to `3`. + - `reset_tolerance`: the tolerance for the reset check. Defaults to + `eps(real(eltype(u)))^(3 // 4)`. + - `check_du`: whether to check the state. Defaults to `true`. + - `check_dfu`: whether to check the function value. Defaults to `true`. +""" +@kwdef @concrete struct NoChangeInStateReset <: AbstractResetCondition + nsteps::Int = 3 + reset_tolerance = nothing + check_du::Bool = true + check_dfu::Bool = true +end + +function InternalAPI.init( + condition::NoChangeInStateReset, J, fu, u, du, args...; kwargs... +) + if condition.check_dfu + @bb dfu = copy(fu) + else + dfu = fu + end + T = real(eltype(u)) + tol = condition.reset_tolerance === nothing ? eps(T)^(3 // 4) : + T(condition.reset_tolerance) + return NoChangeInStateResetCache(dfu, tol, condition, 0, 0) +end + +@concrete mutable struct NoChangeInStateResetCache <: AbstractResetConditionCache + dfu + reset_tolerance + condition <: NoChangeInStateReset + steps_since_change_du::Int + steps_since_change_dfu::Int +end + +function InternalAPI.reinit!(cache::NoChangeInStateResetCache; u0 = nothing, kwargs...) + if u0 !== nothing && cache.condition.reset_tolerance === nothing + cache.reset_tolerance = eps(real(eltype(u0)))^(3 // 4) + end + cache.steps_since_change_dfu = 0 + cache.steps_since_change_du = 0 +end + +function InternalAPI.solve!(cache::NoChangeInStateResetCache, J, fu, u, du; kwargs...) + cond = ≤(cache.reset_tolerance) ∘ abs + if cache.condition.check_du + if any(cond, du) + cache.steps_since_change_du += 1 + if cache.steps_since_change_du ≥ cache.condition.nsteps + cache.steps_since_change_du = 0 + cache.steps_since_change_dfu = 0 + return true + end + else + cache.steps_since_change_du = 0 + cache.steps_since_change_dfu = 0 + end + end + if cache.condition.check_dfu + @bb @. cache.dfu = fu - cache.dfu + if any(cond, cache.dfu) + cache.steps_since_change_dfu += 1 + if cache.steps_since_change_dfu ≥ cache.condition.nsteps + cache.steps_since_change_dfu = 0 + cache.steps_since_change_du = 0 + @bb copyto!(cache.dfu, fu) + return true + end + else + cache.steps_since_change_dfu = 0 + cache.steps_since_change_du = 0 + end + @bb copyto!(cache.dfu, fu) + end + return false +end + +""" + IllConditionedJacobianReset() + +Recommend resetting the Jacobian if the current jacobian is ill-conditioned. This is used +in [`Klement`](@ref). +""" +struct IllConditionedJacobianReset <: AbstractResetCondition end + +function InternalAPI.init( + condition::IllConditionedJacobianReset, J, fu, u, du, args...; kwargs... +) + condition_number_threshold = J isa AbstractMatrix ? inv(eps(real(eltype(J))))^(1 // 2) : + nothing + return IllConditionedJacobianResetCache(condition_number_threshold) +end + +@concrete struct IllConditionedJacobianResetCache <: AbstractResetConditionCache + condition_number_threshold +end + +# NOTE: we don't need a reinit! since we establish the threshold based on the eltype + +function InternalAPI.solve!( + cache::IllConditionedJacobianResetCache, J, fu, u, du; kwargs... +) + J isa Number && return iszero(J) + J isa Diagonal && return any(iszero, diag(J)) + J isa AbstractVector && return any(iszero, J) + J isa AbstractMatrix && + return Utils.condition_number(J) ≥ cache.condition_number_threshold + return false +end diff --git a/lib/NonlinearSolveQuasiNewton/src/solve.jl b/lib/NonlinearSolveQuasiNewton/src/solve.jl new file mode 100644 index 000000000..c52a425ae --- /dev/null +++ b/lib/NonlinearSolveQuasiNewton/src/solve.jl @@ -0,0 +1,392 @@ +""" + QuasiNewtonAlgorithm(; + linesearch = missing, trustregion = missing, descent, update_rule, reinit_rule, + initialization, max_resets::Int = typemax(Int), name::Symbol = :unknown, + max_shrink_times::Int = typemax(Int), concrete_jac = Val(false) + ) + +Nonlinear Solve Algorithms using an Iterative Approximation of the Jacobian. Most common +examples include [`Broyden`](@ref)'s Method. + +### Keyword Arguments + + - `trustregion`: Globalization using a Trust Region Method. This needs to follow the + [`NonlinearSolveBase.AbstractTrustRegionMethod`](@ref) interface. + - `descent`: The descent method to use to compute the step. This needs to follow the + [`NonlinearSolveBase.AbstractDescentDirection`](@ref) interface. + - `max_shrink_times`: The maximum number of times the trust region radius can be shrunk + before the algorithm terminates. + - `update_rule`: The update rule to use to update the Jacobian. This needs to follow the + [`NonlinearSolveBase.AbstractApproximateJacobianUpdateRule`](@ref) interface. + - `reinit_rule`: The reinitialization rule to use to reinitialize the Jacobian. This + needs to follow the [`NonlinearSolveBase.AbstractResetCondition`](@ref) interface. + - `initialization`: The initialization method to use to initialize the Jacobian. This + needs to follow the [`NonlinearSolveBase.AbstractJacobianInitialization`](@ref) + interface. +""" +@concrete struct QuasiNewtonAlgorithm <: AbstractNonlinearSolveAlgorithm + linesearch + trustregion + descent <: AbstractDescentDirection + update_rule <: AbstractApproximateJacobianUpdateRule + reinit_rule <: AbstractResetCondition + initialization <: AbstractJacobianInitialization + + max_resets::Int + max_shrink_times::Int + + concrete_jac <: Union{Val{false}, Val{true}} + name::Symbol +end + +function QuasiNewtonAlgorithm(; + linesearch = missing, trustregion = missing, descent, update_rule, reinit_rule, + initialization, max_resets::Int = typemax(Int), name::Symbol = :unknown, + max_shrink_times::Int = typemax(Int), concrete_jac = Val(false) +) + return QuasiNewtonAlgorithm( + linesearch, trustregion, descent, update_rule, reinit_rule, initialization, + max_resets, max_shrink_times, concrete_jac, name + ) +end + +@concrete mutable struct QuasiNewtonCache <: AbstractNonlinearSolveCache + # Basic Requirements + fu + u + u_cache + p + du # Aliased to `get_du(descent_cache)` + J # Aliased to `initialization_cache.J` if !inverted_jac + alg <: QuasiNewtonAlgorithm + prob <: AbstractNonlinearProblem + globalization <: Union{Val{:LineSearch}, Val{:TrustRegion}, Val{:None}} + + # Internal Caches + initialization_cache + descent_cache + linesearch_cache + trustregion_cache + update_rule_cache + reinit_rule_cache + + inv_workspace + + # Counters + stats::NLStats + nsteps::Int + nresets::Int + max_resets::Int + maxiters::Int + maxtime + max_shrink_times::Int + steps_since_last_reset::Int + + # Timer + timer + total_time::Float64 + + # Termination & Tracking + termination_cache + trace + retcode::ReturnCode.T + force_stop::Bool + force_reinit::Bool + kwargs +end + +function InternalAPI.reinit_self!( + cache::QuasiNewtonCache, args...; p = cache.p, u0 = cache.u, + alias_u0::Bool = false, maxiters = 1000, maxtime = nothing, kwargs... +) + Utils.reinit_common!(cache, u0, p, alias_u0) + + InternalAPI.reinit!(cache.stats) + cache.nsteps = 0 + cache.nresets = 0 + cache.steps_since_last_reset = 0 + cache.maxiters = maxiters + cache.maxtime = maxtime + cache.total_time = 0.0 + cache.force_stop = false + cache.force_reinit = false + cache.retcode = ReturnCode.Default + + NonlinearSolveBase.reset!(cache.trace) + SciMLBase.reinit!( + cache.termination_cache, NonlinearSolveBase.get_fu(cache), + NonlinearSolveBase.get_u(cache); kwargs... + ) + NonlinearSolveBase.reset_timer!(cache.timer) + return +end + +NonlinearSolveBase.@internal_caches(QuasiNewtonCache, + :initialization_cache, :descent_cache, :linesearch_cache, :trustregion_cache, + :update_rule_cache, :reinit_rule_cache) + +function SciMLBase.__init( + prob::AbstractNonlinearProblem, alg::QuasiNewtonAlgorithm, args...; + stats = NLStats(0, 0, 0, 0, 0), alias_u0 = false, maxtime = nothing, + maxiters = 1000, abstol = nothing, reltol = nothing, + linsolve_kwargs = (;), termination_condition = nothing, + internalnorm::F = L2_NORM, kwargs... +) where {F} + timer = get_timer_output() + @static_timeit timer "cache construction" begin + u = Utils.maybe_unaliased(prob.u0, alias_u0) + fu = Utils.evaluate_f(prob, u) + @bb u_cache = copy(u) + + inverted_jac = NonlinearSolveBase.store_inverse_jacobian(alg.update_rule) + + linsolve = NonlinearSolveBase.get_linear_solver(alg.descent) + + initialization_cache = InternalAPI.init( + prob, alg.initialization, alg, prob.f, fu, u, prob.p; + stats, linsolve, maxiters, internalnorm + ) + + abstol, reltol, termination_cache = NonlinearSolveBase.init_termination_cache( + prob, abstol, reltol, fu, u, termination_condition, Val(:regular) + ) + linsolve_kwargs = merge((; abstol, reltol), linsolve_kwargs) + + J = initialization_cache(nothing) + + inv_workspace, J = Utils.unwrap_val(inverted_jac) ? + Utils.maybe_pinv!!_workspace(J) : (nothing, J) + + descent_cache = InternalAPI.init( + prob, alg.descent, J, fu, u; + stats, abstol, reltol, internalnorm, + linsolve_kwargs, pre_inverted = inverted_jac, timer + ) + du = SciMLBase.get_du(descent_cache) + + reinit_rule_cache = InternalAPI.init(alg.reinit_rule, J, fu, u, du) + + has_linesearch = alg.linesearch !== missing && alg.linesearch !== nothing + has_trustregion = alg.trustregion !== missing && alg.trustregion !== nothing + + if has_trustregion && has_linesearch + error("TrustRegion and LineSearch methods are algorithmically incompatible.") + end + + globalization = Val(:None) + linesearch_cache = nothing + trustregion_cache = nothing + + if has_trustregion + NonlinearSolveBase.supports_trust_region(alg.descent) || + error("Trust Region not supported by $(alg.descent).") + trustregion_cache = InternalAPI.init( + prob, alg.trustregion, fu, u, p; stats, internalnorm, kwargs... + ) + globalization = Val(:TrustRegion) + end + + if has_linesearch + NonlinearSolveBase.supports_line_search(alg.descent) || + error("Line Search not supported by $(alg.descent).") + linesearch_cache = CommonSolve.init( + prob, alg.linesearch, fu, u; stats, internalnorm, kwargs... + ) + globalization = Val(:LineSearch) + end + + update_rule_cache = InternalAPI.init( + prob, alg.update_rule, J, fu, u, du; stats, internalnorm + ) + + trace = NonlinearSolveBase.init_nonlinearsolve_trace( + prob, alg, u, fu, J, du; + uses_jacobian_inverse = inverted_jac, kwargs... + ) + + return QuasiNewtonCache( + fu, u, u_cache, prob.p, du, J, alg, prob, globalization, + initialization_cache, descent_cache, linesearch_cache, + trustregion_cache, update_rule_cache, reinit_rule_cache, + inv_workspace, stats, 0, 0, alg.max_resets, maxiters, maxtime, + alg.max_shrink_times, 0, timer, 0.0, termination_cache, trace, + ReturnCode.Default, false, false, kwargs + ) + end +end + +function InternalAPI.step!( + cache::QuasiNewtonCache; recompute_jacobian::Union{Nothing, Bool} = nothing +) + new_jacobian = true + @static_timeit cache.timer "jacobian init/reinit" begin + if cache.nsteps == 0 # First Step is special ignore kwargs + J_init = InternalAPI.solve!( + cache.initialization_cache, cache.fu, cache.u, Val(false) + ) + if Utils.unwrap_val(NonlinearSolveBase.store_inverse_jacobian(cache.update_rule_cache)) + if NonlinearSolveBase.jacobian_initialized_preinverted( + cache.initialization_cache.alg + ) + cache.J = J_init + else + cache.J = Utils.maybe_pinv!!(cache.inv_workspace, J_init) + end + else + if NonlinearSolveBase.jacobian_initialized_preinverted( + cache.initialization_cache.alg + ) + cache.J = Utils.maybe_pinv!!(cache.inv_workspace, J_init) + else + cache.J = J_init + end + end + J = cache.J + cache.steps_since_last_reset += 1 + else + countable_reinit = false + if cache.force_reinit + reinit, countable_reinit = true, true + cache.force_reinit = false + elseif recompute_jacobian === nothing + # Standard Step + reinit = InternalAPI.solve!( + cache.reinit_rule_cache, cache.J, cache.fu, cache.u, cache.du + ) + reinit && (countable_reinit = true) + elseif recompute_jacobian + reinit = true # Force ReInitialization: Don't count towards resetting + else + new_jacobian = false # Jacobian won't be updated in this step + reinit = false # Override Checks: Unsafe operation + end + + if countable_reinit + cache.nresets += 1 + if cache.nresets ≥ cache.max_resets + cache.retcode = ReturnCode.ConvergenceFailure + cache.force_stop = true + return + end + end + + if reinit + J_init = InternalAPI.solve!( + cache.initialization_cache, cache.fu, cache.u, Val(true) + ) + cache.J = Utils.unwrap_val(NonlinearSolveBase.store_inverse_jacobian(cache.update_rule_cache)) ? + Utils.maybe_pinv!!(cache.inv_workspace, J_init) : J_init + J = cache.J + cache.steps_since_last_reset = 0 + else + J = cache.J + cache.steps_since_last_reset += 1 + end + end + end + + @static_timeit cache.timer "descent" begin + if cache.trustregion_cache !== nothing && + hasfield(typeof(cache.trustregion_cache), :trust_region) + descent_result = InternalAPI.solve!( + cache.descent_cache, J, cache.fu, cache.u; new_jacobian, + cache.trustregion_cache.trust_region, cache.kwargs... + ) + else + descent_result = InternalAPI.solve!( + cache.descent_cache, J, cache.fu, cache.u; new_jacobian, cache.kwargs... + ) + end + end + + if !descent_result.linsolve_success + if new_jacobian && cache.steps_since_last_reset == 0 + # Extremely pathological case. Jacobian was just reset and linear solve + # failed. Should ideally never happen in practice unless true jacobian init + # is used. + cache.retcode = ReturnCode.InternalLinearSolveFailed + cache.force_stop = true + return + else + # Force a reinit because the problem is currently un-solvable + if !haskey(cache.kwargs, :verbose) || cache.kwargs[:verbose] + @warn "Linear Solve Failed but Jacobian Information is not current. \ + Retrying with reinitialized Approximate Jacobian." + end + cache.force_reinit = true + InternalAPI.step!(cache; recompute_jacobian = true) + return + end + end + + δu, descent_intermediates = descent_result.δu, descent_result.extras + + if descent_result.success + if cache.globalization isa Val{:LineSearch} + @static_timeit cache.timer "linesearch" begin + linesearch_sol = CommonSolve.solve!(cache.linesearch_cache, cache.u, δu) + needs_reset = !SciMLBase.successful_retcode(linesearch_sol.retcode) + α = linesearch_sol.step_size + end + if needs_reset && cache.steps_since_last_reset > 5 # Reset after a burn-in period + cache.force_reinit = true + else + @static_timeit cache.timer "step" begin + @bb axpy!(α, δu, cache.u) + Utils.evaluate_f!(cache, cache.u, cache.p) + end + end + elseif cache.globalization isa Val{:TrustRegion} + @static_timeit cache.timer "trustregion" begin + tr_accepted, u_new, fu_new = InternalAPI.solve!( + cache.trustregion_cache, J, cache.fu, cache.u, δu, descent_intermediates + ) + if tr_accepted + @bb copyto!(cache.u, u_new) + @bb copyto!(cache.fu, fu_new) + end + if hasfield(typeof(cache.trustregion_cache), :shrink_counter) && + cache.trustregion_cache.shrink_counter > cache.max_shrink_times + cache.retcode = ReturnCode.ShrinkThresholdExceeded + cache.force_stop = true + end + end + α = true + elseif cache.globalization isa Val{:None} + @static_timeit cache.timer "step" begin + @bb axpy!(1, δu, cache.u) + Utils.evaluate_f!(cache, cache.u, cache.p) + end + α = true + else + error("Unknown Globalization Strategy: $(cache.globalization). Allowed values \ + are (:LineSearch, :TrustRegion, :None)") + end + + NonlinearSolveBase.check_and_update!(cache, cache.fu, cache.u, cache.u_cache) + else + α = false + cache.force_reinit = true + end + + update_trace!( + cache, α; + uses_jac_inverse = NonlinearSolveBase.store_inverse_jacobian(cache.update_rule_cache) + ) + @bb copyto!(cache.u_cache, cache.u) + + if (cache.force_stop || cache.force_reinit || + (recompute_jacobian !== nothing && !recompute_jacobian)) + NonlinearSolveBase.callback_into_cache!(cache) + return nothing + end + + @static_timeit cache.timer "jacobian update" begin + cache.J = InternalAPI.solve!( + cache.update_rule_cache, cache.J, cache.fu, cache.u, δu + ) + NonlinearSolveBase.callback_into_cache!(cache) + end + + return nothing +end diff --git a/lib/NonlinearSolveQuasiNewton/src/structure.jl b/lib/NonlinearSolveQuasiNewton/src/structure.jl new file mode 100644 index 000000000..09438289b --- /dev/null +++ b/lib/NonlinearSolveQuasiNewton/src/structure.jl @@ -0,0 +1,53 @@ +""" + DiagonalStructure() + +Preserves only the Diagonal of the Matrix. +""" +struct DiagonalStructure <: AbstractApproximateJacobianStructure end + +NonlinearSolveBase.get_full_jacobian(cache, ::DiagonalStructure, J::Number) = J +function NonlinearSolveBase.get_full_jacobian(cache, ::DiagonalStructure, J) + return Diagonal(Utils.safe_vec(J)) +end + +function (::DiagonalStructure)(J::AbstractMatrix; alias::Bool = false) + @assert size(J, 1)==size(J, 2) "Diagonal Jacobian Structure must be square!" + return LinearAlgebra.diag(J) +end +(::DiagonalStructure)(J::AbstractVector; alias::Bool = false) = alias ? J : @bb(copy(J)) +(::DiagonalStructure)(J::Number; alias::Bool = false) = J + +(::DiagonalStructure)(::Number, J_new::Number) = J_new +function (::DiagonalStructure)(J::AbstractVector, J_new::AbstractMatrix) + if ArrayInterface.can_setindex(J) + if ArrayInterface.fast_scalar_indexing(J) + @simd ivdep for i in eachindex(J) + @inbounds J[i] = J_new[i, i] + end + else + J .= @view(J_new[diagind(J_new)]) + end + return J + end + return LinearAlgebra.diag(J_new) +end +function (st::DiagonalStructure)(J::AbstractArray, J_new::AbstractMatrix) + return ArrayInterface.restructure(J, st(vec(J), J_new)) +end + +""" + FullStructure() + +Stores the full matrix. +""" +struct FullStructure <: AbstractApproximateJacobianStructure end + +NonlinearSolveBase.stores_full_jacobian(::FullStructure) = true + +(::FullStructure)(J; alias::Bool = false) = alias ? J : @bb(copy(J)) + +function (::FullStructure)(J, J_new) + J === J_new && return J + @bb copyto!(J, J_new) + return J +end diff --git a/lib/NonlinearSolveQuasiNewton/test/core_tests.jl b/lib/NonlinearSolveQuasiNewton/test/core_tests.jl new file mode 100644 index 000000000..d22481b4c --- /dev/null +++ b/lib/NonlinearSolveQuasiNewton/test/core_tests.jl @@ -0,0 +1,228 @@ +@testsetup module CoreRootfindTesting + +include("../../../common/common_rootfind_testing.jl") + +end + +@testitem "Broyden" setup=[CoreRootfindTesting] tags=[:core] begin + using ADTypes, LineSearch + using LineSearches: LineSearches + using BenchmarkTools: @ballocated + using StaticArrays: @SVector + using Zygote, Enzyme, ForwardDiff, FiniteDiff + + u0s = ([1.0, 1.0], @SVector[1.0, 1.0], 1.0) + + @testset for ad in (AutoForwardDiff(), AutoZygote(), AutoFiniteDiff(), AutoEnzyme()) + @testset "$(nameof(typeof(linesearch)))" for linesearch in ( + # LineSearchesJL(; method = LineSearches.Static(), autodiff = ad), + # LineSearchesJL(; method = LineSearches.BackTracking(), autodiff = ad), + # LineSearchesJL(; method = LineSearches.MoreThuente(), autodiff = ad), + # LineSearchesJL(; method = LineSearches.StrongWolfe(), autodiff = ad), + LineSearchesJL(; method = LineSearches.HagerZhang(), autodiff = ad), + BackTracking(; autodiff = ad), + LiFukushimaLineSearch() + ) + @testset for init_jacobian in (Val(:identity), Val(:true_jacobian)), + update_rule in (Val(:good_broyden), Val(:bad_broyden), Val(:diagonal)) + + @testset "[OOP] u0: $(typeof(u0))" for u0 in ( + [1.0, 1.0], @SVector[1.0, 1.0], 1.0 + ) + solver = Broyden(; linesearch, init_jacobian, update_rule) + sol = solve_oop(quadratic_f, u0; solver) + @test SciMLBase.successful_retcode(sol) + err = maximum(abs, quadratic_f(sol.u, 2.0)) + @test err < 1e-9 + + cache = init( + NonlinearProblem{false}(quadratic_f, u0, 2.0), solver, abstol = 1e-9 + ) + @test (@ballocated solve!($cache)) < 200 + end + + @testset "[IIP] u0: $(typeof(u0))" for u0 in ([1.0, 1.0],) + ad isa AutoZygote && continue + + solver = Broyden(; linesearch, init_jacobian, update_rule) + sol = solve_iip(quadratic_f!, u0; solver) + @test SciMLBase.successful_retcode(sol) + err = maximum(abs, quadratic_f(sol.u, 2.0)) + @test err < 1e-9 + + cache = init( + NonlinearProblem{true}(quadratic_f!, u0, 2.0), solver, abstol = 1e-9 + ) + @test (@ballocated solve!($cache)) ≤ 64 + end + end + end + end +end + +@testitem "Broyden: Iterator Interface" setup=[CoreRootfindTesting] tags=[:core] begin + p = range(0.01, 2, length = 200) + @test nlprob_iterator_interface(quadratic_f, p, false, Broyden()) ≈ sqrt.(p) + @test nlprob_iterator_interface(quadratic_f!, p, true, Broyden()) ≈ sqrt.(p) +end + +@testitem "Broyden Termination Conditions" setup=[CoreRootfindTesting] tags=[:core] begin + using StaticArrays: @SVector + + @testset "TC: $(nameof(typeof(termination_condition)))" for termination_condition in TERMINATION_CONDITIONS + @testset "u0: $(typeof(u0))" for u0 in ([1.0, 1.0], 1.0, @SVector([1.0, 1.0])) + probN = NonlinearProblem(quadratic_f, u0, 2.0) + sol = solve(probN, Broyden(); termination_condition) + err = maximum(abs, quadratic_f(sol.u, 2.0)) + @test err < 1e-9 + end + end +end + +@testitem "Klement" setup=[CoreRootfindTesting] tags=[:core] begin + using ADTypes, LineSearch + using LineSearches: LineSearches + using BenchmarkTools: @ballocated + using StaticArrays: @SVector + using Zygote, Enzyme, ForwardDiff, FiniteDiff + + @testset for ad in (AutoForwardDiff(), AutoZygote(), AutoFiniteDiff(), AutoEnzyme()) + @testset "$(nameof(typeof(linesearch)))" for linesearch in ( + # LineSearchesJL(; method = LineSearches.Static(), autodiff = ad), + # LineSearchesJL(; method = LineSearches.BackTracking(), autodiff = ad), + # LineSearchesJL(; method = LineSearches.MoreThuente(), autodiff = ad), + # LineSearchesJL(; method = LineSearches.StrongWolfe(), autodiff = ad), + LineSearchesJL(; method = LineSearches.HagerZhang(), autodiff = ad), + BackTracking(; autodiff = ad), + LiFukushimaLineSearch() + ) + @testset for init_jacobian in ( + Val(:identity), Val(:true_jacobian), Val(:true_jacobian_diagonal)) + @testset "[OOP] u0: $(typeof(u0))" for u0 in ( + [1.0, 1.0], @SVector[1.0, 1.0], 1.0 + ) + solver = Klement(; linesearch, init_jacobian) + sol = solve_oop(quadratic_f, u0; solver) + # XXX: some tests are failing by a margin + # @test SciMLBase.successful_retcode(sol) + err = maximum(abs, quadratic_f(sol.u, 2.0)) + @test err < 1e-9 + + cache = init( + NonlinearProblem{false}(quadratic_f, u0, 2.0), solver, abstol = 1e-9 + ) + @test (@ballocated solve!($cache)) < 200 + end + + @testset "[IIP] u0: $(typeof(u0))" for u0 in ([1.0, 1.0],) + ad isa AutoZygote && continue + + solver = Klement(; linesearch, init_jacobian) + sol = solve_iip(quadratic_f!, u0; solver) + @test SciMLBase.successful_retcode(sol) + err = maximum(abs, quadratic_f(sol.u, 2.0)) + @test err < 1e-9 + + cache = init( + NonlinearProblem{true}(quadratic_f!, u0, 2.0), solver, abstol = 1e-9 + ) + @test (@ballocated solve!($cache)) ≤ 64 + end + end + end + end +end + +@testitem "Klement: Iterator Interface" setup=[CoreRootfindTesting] tags=[:core] begin + p = range(0.01, 2, length = 200) + @test nlprob_iterator_interface(quadratic_f, p, false, Klement()) ≈ sqrt.(p) + @test nlprob_iterator_interface(quadratic_f!, p, true, Klement()) ≈ sqrt.(p) +end + +@testitem "Klement Termination Conditions" setup=[CoreRootfindTesting] tags=[:core] begin + using StaticArrays: @SVector + + @testset "TC: $(nameof(typeof(termination_condition)))" for termination_condition in TERMINATION_CONDITIONS + @testset "u0: $(typeof(u0))" for u0 in ([1.0, 1.0], 1.0, @SVector([1.0, 1.0])) + probN = NonlinearProblem(quadratic_f, u0, 2.0) + sol = solve(probN, Klement(); termination_condition) + err = maximum(abs, quadratic_f(sol.u, 2.0)) + @test err < 1e-9 + end + end +end + +@testitem "LimitedMemoryBroyden" setup=[CoreRootfindTesting] tags=[:core] begin + using ADTypes, LineSearch + using LineSearches: LineSearches + using BenchmarkTools: @ballocated + using StaticArrays: @SVector + using Zygote, Enzyme, ForwardDiff, FiniteDiff + + @testset for ad in (AutoForwardDiff(), AutoZygote(), AutoFiniteDiff(), AutoEnzyme()) + @testset "$(nameof(typeof(linesearch)))" for linesearch in ( + # LineSearchesJL(; method = LineSearches.Static(), autodiff = ad), + # LineSearchesJL(; method = LineSearches.BackTracking(), autodiff = ad), + # LineSearchesJL(; method = LineSearches.MoreThuente(), autodiff = ad), + # LineSearchesJL(; method = LineSearches.StrongWolfe(), autodiff = ad), + LineSearchesJL(; method = LineSearches.HagerZhang(), autodiff = ad), + BackTracking(; autodiff = ad), + LiFukushimaLineSearch() + ) + @testset "[OOP] u0: $(typeof(u0))" for u0 in (ones(32), @SVector(ones(2)), 1.0) + broken = Sys.iswindows() && u0 isa Vector{Float64} && + linesearch isa BackTracking && ad isa AutoFiniteDiff + + solver = LimitedMemoryBroyden(; linesearch) + sol = solve_oop(quadratic_f, u0; solver) + @test SciMLBase.successful_retcode(sol) broken=broken + err = maximum(abs, quadratic_f(sol.u, 2.0)) + @test err<1e-9 broken=broken + + cache = init( + NonlinearProblem{false}(quadratic_f, u0, 2.0), solver, abstol = 1e-9 + ) + @test (@ballocated solve!($cache)) ≤ 400 + end + + @testset "[IIP] u0: $(typeof(u0))" for u0 in (ones(32),) + ad isa AutoZygote && continue + + broken = Sys.iswindows() && u0 isa Vector{Float64} && + linesearch isa BackTracking && ad isa AutoFiniteDiff + + solver = LimitedMemoryBroyden(; linesearch) + sol = solve_iip(quadratic_f!, u0; solver) + @test SciMLBase.successful_retcode(sol) broken=broken + err = maximum(abs, quadratic_f(sol.u, 2.0)) + @test err<1e-9 broken=broken + + cache = init( + NonlinearProblem{true}(quadratic_f!, u0, 2.0), solver, abstol = 1e-9 + ) + @test (@ballocated solve!($cache)) ≤ 64 + end + end + end +end + +@testitem "LimitedMemoryBroyden: Iterator Interface" setup=[CoreRootfindTesting] tags=[:core] begin + p = range(0.01, 2, length = 200) + @test nlprob_iterator_interface(quadratic_f, p, false, LimitedMemoryBroyden()) ≈ + sqrt.(p) + @test nlprob_iterator_interface(quadratic_f!, p, true, LimitedMemoryBroyden()) ≈ + sqrt.(p) +end + +@testitem "LimitedMemoryBroyden Termination Conditions" setup=[CoreRootfindTesting] tags=[:core] begin + using StaticArrays: @SVector + + @testset "TC: $(nameof(typeof(termination_condition)))" for termination_condition in TERMINATION_CONDITIONS + @testset "u0: $(typeof(u0))" for u0 in ([1.0, 1.0], 1.0, @SVector([1.0, 1.0])) + probN = NonlinearProblem(quadratic_f, u0, 2.0) + sol = solve(probN, LimitedMemoryBroyden(); termination_condition) + err = maximum(abs, quadratic_f(sol.u, 2.0)) + @test err < 1e-9 + end + end +end diff --git a/lib/NonlinearSolveQuasiNewton/test/qa_tests.jl b/lib/NonlinearSolveQuasiNewton/test/qa_tests.jl new file mode 100644 index 000000000..badeab3ef --- /dev/null +++ b/lib/NonlinearSolveQuasiNewton/test/qa_tests.jl @@ -0,0 +1,22 @@ +@testitem "Aqua" tags=[:core] begin + using Aqua, NonlinearSolveQuasiNewton + + Aqua.test_all( + NonlinearSolveQuasiNewton; + piracies = false, ambiguities = false, stale_deps = false, deps_compat = false + ) + Aqua.test_stale_deps(NonlinearSolveQuasiNewton; ignore = [:SciMLJacobianOperators]) + Aqua.test_deps_compat(NonlinearSolveQuasiNewton; ignore = [:SciMLJacobianOperators]) + Aqua.test_piracies(NonlinearSolveQuasiNewton) + Aqua.test_ambiguities(NonlinearSolveQuasiNewton; recursive = false) +end + +@testitem "Explicit Imports" tags=[:core] begin + using ExplicitImports, NonlinearSolveQuasiNewton + + @test check_no_implicit_imports( + NonlinearSolveQuasiNewton; skip = (Base, Core, SciMLBase) + ) === nothing + @test check_no_stale_explicit_imports(NonlinearSolveQuasiNewton) === nothing + @test check_all_qualified_accesses_via_owners(NonlinearSolveQuasiNewton) === nothing +end diff --git a/lib/NonlinearSolveQuasiNewton/test/runtests.jl b/lib/NonlinearSolveQuasiNewton/test/runtests.jl new file mode 100644 index 000000000..807441bbc --- /dev/null +++ b/lib/NonlinearSolveQuasiNewton/test/runtests.jl @@ -0,0 +1,23 @@ +using ReTestItems, NonlinearSolveQuasiNewton, Hwloc, InteractiveUtils, Pkg + +@info sprint(InteractiveUtils.versioninfo) + +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( + NonlinearSolveQuasiNewton; tags = (GROUP == "all" ? nothing : [Symbol(GROUP)]), + nworkers = RETESTITEMS_NWORKERS, nworker_threads = RETESTITEMS_NWORKER_THREADS, + testitem_timeout = 3600 +) diff --git a/lib/NonlinearSolveSpectralMethods/LICENSE b/lib/NonlinearSolveSpectralMethods/LICENSE new file mode 100644 index 000000000..46c972b17 --- /dev/null +++ b/lib/NonlinearSolveSpectralMethods/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2024 Avik Pal + +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/NonlinearSolveSpectralMethods/Project.toml b/lib/NonlinearSolveSpectralMethods/Project.toml new file mode 100644 index 000000000..bb9367554 --- /dev/null +++ b/lib/NonlinearSolveSpectralMethods/Project.toml @@ -0,0 +1,54 @@ +name = "NonlinearSolveSpectralMethods" +uuid = "26075421-4e9a-44e1-8bd1-420ed7ad02b2" +authors = ["Avik Pal and contributors"] +version = "1.0.0" + +[deps] +CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2" +ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471" +DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e" +LineSearch = "87fe0de2-c867-4266-b59a-2f0a94fc965b" +MaybeInplace = "bb5d69b7-63fc-4a16-80bd-7e42200c7bdb" +NonlinearSolveBase = "be0214bd-f91f-a760-ac4e-3421ce2b2da0" +PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" +Reexport = "189a3867-3050-52da-a836-e630ba90ab69" +SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" + +[compat] +Aqua = "0.8" +BenchmarkTools = "1.5.0" +CommonSolve = "0.2.4" +ConcreteStructs = "0.2.3" +DiffEqBase = "6.158.3" +ExplicitImports = "1.5" +Hwloc = "3" +InteractiveUtils = "<0.0.1, 1" +LineSearch = "0.1.4" +MaybeInplace = "0.1.4" +NonlinearProblemLibrary = "0.1.2" +NonlinearSolveBase = "1.1" +Pkg = "1.10" +PrecompileTools = "1.2" +ReTestItems = "1.24" +Reexport = "1" +SciMLBase = "2.58" +StableRNGs = "1" +StaticArrays = "1.9.8" +Test = "1.10" +julia = "1.10" + +[extras] +Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" +BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" +ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7" +Hwloc = "0e44f5e4-bd66-52a0-8798-143a42290a1d" +InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" +NonlinearProblemLibrary = "b7050fa9-e91f-4b37-bcee-a89a063da141" +Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" +ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823" +StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" +StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + +[targets] +test = ["Aqua", "BenchmarkTools", "ExplicitImports", "Hwloc", "InteractiveUtils", "NonlinearProblemLibrary", "Pkg", "ReTestItems", "StableRNGs", "StaticArrays", "Test"] diff --git a/lib/NonlinearSolveSpectralMethods/src/NonlinearSolveSpectralMethods.jl b/lib/NonlinearSolveSpectralMethods/src/NonlinearSolveSpectralMethods.jl new file mode 100644 index 000000000..8872fce35 --- /dev/null +++ b/lib/NonlinearSolveSpectralMethods/src/NonlinearSolveSpectralMethods.jl @@ -0,0 +1,36 @@ +module NonlinearSolveSpectralMethods + +using ConcreteStructs: @concrete +using Reexport: @reexport +using PrecompileTools: @compile_workload, @setup_workload + +using CommonSolve: CommonSolve +using DiffEqBase: DiffEqBase # Needed for `init` / `solve` dispatches +using LineSearch: RobustNonMonotoneLineSearch +using MaybeInplace: @bb +using NonlinearSolveBase: NonlinearSolveBase, AbstractNonlinearSolveAlgorithm, + AbstractNonlinearSolveCache, Utils, InternalAPI, get_timer_output, + @static_timeit, update_trace! +using SciMLBase: SciMLBase, AbstractNonlinearProblem, NLStats, ReturnCode + +include("dfsane.jl") + +include("solve.jl") + +@setup_workload begin + include("../../../common/nonlinear_problem_workloads.jl") + + algs = [DFSane()] + + @compile_workload begin + @sync for prob in nonlinear_problems, alg in algs + Threads.@spawn CommonSolve.solve(prob, alg; abstol = 1e-2, verbose = false) + end + end +end + +@reexport using SciMLBase, NonlinearSolveBase + +export GeneralizedDFSane, DFSane + +end diff --git a/lib/NonlinearSolveSpectralMethods/src/dfsane.jl b/lib/NonlinearSolveSpectralMethods/src/dfsane.jl new file mode 100644 index 000000000..108281e7b --- /dev/null +++ b/lib/NonlinearSolveSpectralMethods/src/dfsane.jl @@ -0,0 +1,33 @@ +""" + DFSane(; + sigma_min = 1 // 10^10, sigma_max = 1e10, sigma_1 = 1, M::Int = 10, + gamma = 1 // 10^4, tau_min = 1 // 10, tau_max = 1 // 2, n_exp::Int = 2, + max_inner_iterations::Int = 100, eta_strategy = (fn_1, n, x_n, f_n) -> fn_1 / n^2 + ) + +A low-overhead and allocation-free implementation of the df-sane method for solving +large-scale nonlinear systems of equations. For in depth information about all the +parameters and the algorithm, see [la2006spectral](@citet). + +### Keyword Arguments + + - `sigma_min`: the minimum value of the spectral coefficient `σ` which is related to the + step size in the algorithm. Defaults to `1e-10`. + - `sigma_max`: the maximum value of the spectral coefficient `σₙ` which is related to the + step size in the algorithm. Defaults to `1e10`. + +For other keyword arguments, see [`RobustNonMonotoneLineSearch`](@ref). +""" +function DFSane(; + sigma_min = 1 // 10^10, sigma_max = 1e10, sigma_1 = 1, M::Int = 10, + gamma = 1 // 10^4, tau_min = 1 // 10, tau_max = 1 // 2, n_exp::Int = 2, + max_inner_iterations::Int = 100, eta_strategy::F = (fn_1, n, x_n, f_n) -> fn_1 / n^2 +) where {F} + linesearch = RobustNonMonotoneLineSearch(; + gamma = gamma, sigma_1 = sigma_1, M, tau_min = tau_min, tau_max = tau_max, + n_exp, η_strategy = eta_strategy, maxiters = max_inner_iterations + ) + return GeneralizedDFSane(; + linesearch, sigma_min, sigma_max, sigma_1 = nothing, name = :DFSane + ) +end diff --git a/src/core/spectral_methods.jl b/lib/NonlinearSolveSpectralMethods/src/solve.jl similarity index 51% rename from src/core/spectral_methods.jl rename to lib/NonlinearSolveSpectralMethods/src/solve.jl index 31e988b70..b3a7d216e 100644 --- a/src/core/spectral_methods.jl +++ b/lib/NonlinearSolveSpectralMethods/src/solve.jl @@ -2,7 +2,7 @@ # papers, this seems to be the only one that is widely used. If we have a list of more # papers we can see what is the right level of abstraction to implement here """ - GeneralizedDFSane{name}(linesearch, σ_min, σ_max, σ_1) + GeneralizedDFSane(; linesearch, sigma_min, sigma_max, sigma_1, name::Symbol = :unknown) A generalized version of the DF-SANE algorithm. This algorithm is a Jacobian-Free Spectral Method. @@ -11,35 +11,29 @@ Method. - `linesearch`: Globalization using a Line Search Method. This is not optional currently, but that restriction might be lifted in the future. - - `σ_min`: The minimum spectral parameter allowed. This is used to ensure that the + - `sigma_min`: The minimum spectral parameter allowed. This is used to ensure that the spectral parameter is not too small. - - `σ_max`: The maximum spectral parameter allowed. This is used to ensure that the + - `sigma_max`: The maximum spectral parameter allowed. This is used to ensure that the spectral parameter is not too large. - - `σ_1`: The initial spectral parameter. If this is not provided, then the algorithm - initializes it as `σ_1 = / `. + - `sigma_1`: The initial spectral parameter. If this is not provided, then the algorithm + initializes it as `sigma_1 = / `. """ -@concrete struct GeneralizedDFSane{name} <: AbstractNonlinearSolveAlgorithm{name} +@concrete struct GeneralizedDFSane <: AbstractNonlinearSolveAlgorithm linesearch σ_min σ_max σ_1 -end -function __show_algorithm(io::IO, alg::GeneralizedDFSane, name, indent) - modifiers = String[] - __is_present(alg.linesearch) && push!(modifiers, "linesearch = $(alg.linesearch)") - push!(modifiers, "σ_min = $(alg.σ_min)") - push!(modifiers, "σ_max = $(alg.σ_max)") - push!(modifiers, "σ_1 = $(alg.σ_1)") - spacing = " "^indent * " " - spacing_last = " "^indent - print(io, "$(name)(\n$(spacing)$(join(modifiers, ",\n$(spacing)"))\n$(spacing_last))") + name::Symbol end -concrete_jac(::GeneralizedDFSane) = nothing +function GeneralizedDFSane(; + linesearch, sigma_min, sigma_max, sigma_1, name::Symbol = :unknown +) + return GeneralizedDFSane(linesearch, sigma_min, sigma_max, sigma_1, name) +end -@concrete mutable struct GeneralizedDFSaneCache{iip, timeit} <: - AbstractNonlinearSolveCache{iip, timeit} +@concrete mutable struct GeneralizedDFSaneCache <: AbstractNonlinearSolveCache # Basic Requirements fu fu_cache @@ -47,8 +41,8 @@ concrete_jac(::GeneralizedDFSane) = nothing u_cache p du - alg - prob + alg <: GeneralizedDFSane + prob <: AbstractNonlinearProblem # Parameters σ_n @@ -66,7 +60,7 @@ concrete_jac(::GeneralizedDFSane) = nothing # Timer timer - total_time::Float64 # Simple Counter which works even if TimerOutput is disabled + total_time::Float64 # Termination & Tracking termination_cache @@ -76,23 +70,17 @@ concrete_jac(::GeneralizedDFSane) = nothing kwargs end -function __reinit_internal!( - cache::GeneralizedDFSaneCache{iip}, args...; p = cache.p, u0 = cache.u, - alias_u0::Bool = false, maxiters = 1000, maxtime = nothing, kwargs...) where {iip} - if iip - recursivecopy!(cache.u, u0) - cache.prob.f(cache.fu, cache.u, p) - else - cache.u = __maybe_unaliased(u0, alias_u0) - set_fu!(cache, cache.prob.f(cache.u, p)) - end - cache.p = p +function InternalAPI.reinit_self!( + cache::GeneralizedDFSaneCache, args...; p = cache.p, u0 = cache.u, + alias_u0::Bool = false, maxiters = 1000, maxtime = nothing, kwargs... +) + Utils.reinit_common!(cache, u0, p, alias_u0) if cache.alg.σ_1 === nothing - σ_n = dot(cache.u, cache.u) / dot(cache.u, cache.fu) + σ_n = Utils.safe_dot(cache.u, cache.u) / Utils.safe_dot(cache.u, cache.fu) # Spectral parameter bounds check if !(cache.alg.σ_min ≤ abs(σ_n) ≤ cache.alg.σ_max) - test_norm = dot(cache.fu, cache.fu) + test_norm = NonlinearSolveBase.L2_NORM(cache.fu) σ_n = clamp(inv(test_norm), T(1), T(1e5)) end else @@ -100,61 +88,76 @@ function __reinit_internal!( end cache.σ_n = σ_n - reset_timer!(cache.timer) + NonlinearSolveBase.reset_timer!(cache.timer) cache.total_time = 0.0 - reset!(cache.trace) - reinit!(cache.termination_cache, get_fu(cache), get_u(cache); kwargs...) - __reinit_internal!(cache.stats) + NonlinearSolveBase.reset!(cache.trace) + SciMLBase.reinit!( + cache.termination_cache, NonlinearSolveBase.get_fu(cache), + NonlinearSolveBase.get_u(cache); kwargs... + ) + + InternalAPI.reinit!(cache.stats) cache.nsteps = 0 cache.maxiters = maxiters cache.maxtime = maxtime cache.force_stop = false cache.retcode = ReturnCode.Default + return end -@internal_caches GeneralizedDFSaneCache :linesearch_cache +NonlinearSolveBase.@internal_caches GeneralizedDFSaneCache :linesearch_cache -function SciMLBase.__init(prob::AbstractNonlinearProblem, alg::GeneralizedDFSane, args...; - stats = empty_nlstats(), alias_u0 = false, maxiters = 1000, +function SciMLBase.__init( + prob::AbstractNonlinearProblem, alg::GeneralizedDFSane, args...; + stats = NLStats(0, 0, 0, 0, 0), alias_u0 = false, maxiters = 1000, abstol = nothing, reltol = nothing, termination_condition = nothing, - maxtime = nothing, kwargs...) + maxtime = nothing, kwargs... +) timer = get_timer_output() + @static_timeit timer "cache construction" begin - u = __maybe_unaliased(prob.u0, alias_u0) + u = Utils.maybe_unaliased(prob.u0, alias_u0) T = eltype(u) @bb du = similar(u) @bb u_cache = copy(u) - fu = evaluate_f(prob, u) + fu = Utils.evaluate_f(prob, u) @bb fu_cache = copy(fu) - linesearch_cache = init(prob, alg.linesearch, fu, u; stats, kwargs...) + linesearch_cache = CommonSolve.init(prob, alg.linesearch, fu, u; stats, kwargs...) abstol, reltol, tc_cache = NonlinearSolveBase.init_termination_cache( - prob, abstol, reltol, fu, u_cache, termination_condition, Val(:regular)) - trace = init_nonlinearsolve_trace(prob, alg, u, fu, nothing, du; kwargs...) + prob, abstol, reltol, fu, u_cache, termination_condition, Val(:regular) + ) + trace = NonlinearSolveBase.init_nonlinearsolve_trace( + prob, alg, u, fu, nothing, du; kwargs... + ) if alg.σ_1 === nothing - σ_n = dot(u, u) / dot(u, fu) + σ_n = Utils.safe_dot(u, u) / Utils.safe_dot(u, fu) # Spectral parameter bounds check if !(alg.σ_min ≤ abs(σ_n) ≤ alg.σ_max) - test_norm = dot(fu, fu) + test_norm = NonlinearSolveBase.L2_NORM(fu) σ_n = clamp(inv(test_norm), T(1), T(1e5)) end else σ_n = T(alg.σ_1) end - return GeneralizedDFSaneCache{isinplace(prob), maxtime !== nothing}( - fu, fu_cache, u, u_cache, prob.p, du, alg, prob, σ_n, T(alg.σ_min), - T(alg.σ_max), linesearch_cache, stats, 0, maxiters, maxtime, - timer, 0.0, tc_cache, trace, ReturnCode.Default, false, kwargs) + return GeneralizedDFSaneCache( + fu, fu_cache, u, u_cache, prob.p, du, alg, prob, + σ_n, T(alg.σ_min), T(alg.σ_max), + linesearch_cache, stats, 0, maxiters, maxtime, timer, 0.0, + tc_cache, trace, ReturnCode.Default, false, kwargs + ) end end -function __step!(cache::GeneralizedDFSaneCache{iip}; - recompute_jacobian::Union{Nothing, Bool} = nothing, kwargs...) where {iip} +function InternalAPI.step!( + cache::GeneralizedDFSaneCache; recompute_jacobian::Union{Nothing, Bool} = nothing, + kwargs... +) if recompute_jacobian !== nothing @warn "GeneralizedDFSane is a Jacobian-Free Algorithm. Ignoring \ `recompute_jacobian`" maxlog=1 @@ -165,7 +168,7 @@ function __step!(cache::GeneralizedDFSaneCache{iip}; end @static_timeit cache.timer "linesearch" begin - linesearch_sol = solve!(cache.linesearch_cache, cache.u, cache.du) + linesearch_sol = CommonSolve.solve!(cache.linesearch_cache, cache.u, cache.du) linesearch_failed = !SciMLBase.successful_retcode(linesearch_sol.retcode) α = linesearch_sol.step_size end @@ -178,23 +181,24 @@ function __step!(cache::GeneralizedDFSaneCache{iip}; @static_timeit cache.timer "step" begin @bb axpy!(α, cache.du, cache.u) - evaluate_f!(cache, cache.u, cache.p) + Utils.evaluate_f!(cache, cache.u, cache.p) end update_trace!(cache, α) - check_and_update!(cache, cache.fu, cache.u, cache.u_cache) + + NonlinearSolveBase.check_and_update!(cache, cache.fu, cache.u, cache.u_cache) # Update Spectral Parameter @static_timeit cache.timer "update spectral parameter" begin @bb @. cache.u_cache = cache.u - cache.u_cache @bb @. cache.fu_cache = cache.fu - cache.fu_cache - cache.σ_n = __dot(cache.u_cache, cache.u_cache) / - __dot(cache.u_cache, cache.fu_cache) + cache.σ_n = Utils.safe_dot(cache.u_cache, cache.u_cache) / + Utils.safe_dot(cache.u_cache, cache.fu_cache) # Spectral parameter bounds check if !(cache.σ_min ≤ abs(cache.σ_n) ≤ cache.σ_max) - test_norm = dot(cache.fu, cache.fu) + test_norm = NonlinearSolveBase.L2_NORM(cache.fu) T = eltype(cache.σ_n) cache.σ_n = clamp(inv(test_norm), T(1), T(1e5)) end @@ -204,7 +208,7 @@ function __step!(cache::GeneralizedDFSaneCache{iip}; @bb copyto!(cache.u_cache, cache.u) @bb copyto!(cache.fu_cache, cache.fu) - callback_into_cache!(cache, cache.linesearch_cache) + NonlinearSolveBase.callback_into_cache!(cache, cache.linesearch_cache) return end diff --git a/lib/NonlinearSolveSpectralMethods/test/core_tests.jl b/lib/NonlinearSolveSpectralMethods/test/core_tests.jl new file mode 100644 index 000000000..e127d1e16 --- /dev/null +++ b/lib/NonlinearSolveSpectralMethods/test/core_tests.jl @@ -0,0 +1,87 @@ +@testsetup module CoreRootfindTesting + +include("../../../common/common_rootfind_testing.jl") + +end + +@testitem "DFSane" setup=[CoreRootfindTesting] tags=[:core] begin + using BenchmarkTools: @ballocated + using StaticArrays: @SVector + + u0s = ([1.0, 1.0], @SVector[1.0, 1.0], 1.0) + + @testset "[OOP] u0: $(typeof(u0))" for u0 in u0s + sol = solve_oop(quadratic_f, u0; solver = DFSane()) + @test SciMLBase.successful_retcode(sol) + err = maximum(abs, quadratic_f(sol.u, 2.0)) + @test err < 1e-9 + + cache = init(NonlinearProblem{false}(quadratic_f, u0, 2.0), DFSane(), abstol = 1e-9) + @test (@ballocated solve!($cache)) < 200 + end + + @testset "[IIP] u0: $(typeof(u0))" for u0 in ([1.0, 1.0],) + sol = solve_iip(quadratic_f!, u0; solver = DFSane()) + @test SciMLBase.successful_retcode(sol) + err = maximum(abs, quadratic_f(sol.u, 2.0)) + @test err < 1e-9 + + cache = init(NonlinearProblem{true}(quadratic_f!, u0, 2.0), DFSane(), abstol = 1e-9) + @test (@ballocated solve!($cache)) ≤ 64 + end +end + +@testitem "DFSane Iterator Interface" setup=[CoreRootfindTesting] tags=[:core] begin + p = range(0.01, 2, length = 200) + @test nlprob_iterator_interface(quadratic_f, p, false, DFSane()) ≈ sqrt.(p) + @test nlprob_iterator_interface(quadratic_f!, p, true, DFSane()) ≈ sqrt.(p) +end + +@testitem "DFSane NewtonRaphson Fails" setup=[CoreRootfindTesting] tags=[:core] begin + u0 = [-10.0, -1.0, 1.0, 2.0, 3.0, 4.0, 10.0] + p = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] + sol = solve_oop(newton_fails, u0, p; solver = DFSane()) + @test SciMLBase.successful_retcode(sol) + @test all(abs.(newton_fails(sol.u, p)) .< 1e-9) +end + +@testitem "DFSane: Kwargs" setup=[CoreRootfindTesting] tags=[:core] begin + σ_min = [1e-10, 1e-5, 1e-4] + σ_max = [1e10, 1e5, 1e4] + σ_1 = [1.0, 0.5, 2.0] + M = [10, 1, 100] + γ = [1e-4, 1e-3, 1e-5] + τ_min = [0.1, 0.2, 0.3] + τ_max = [0.5, 0.8, 0.9] + nexp = [2, 1, 2] + η_strategy = [ + (f_1, k, x, F) -> f_1 / k^2, (f_1, k, x, F) -> f_1 / k^3, + (f_1, k, x, F) -> f_1 / k^4 + ] + + list_of_options = zip(σ_min, σ_max, σ_1, M, γ, τ_min, τ_max, nexp, η_strategy) + for options in list_of_options + local probN, sol, alg + alg = DFSane(; + sigma_min = options[1], sigma_max = options[2], sigma_1 = options[3], + M = options[4], gamma = options[5], tau_min = options[6], + tau_max = options[7], n_exp = options[8], eta_strategy = options[9] + ) + + probN = NonlinearProblem{false}(quadratic_f, [1.0, 1.0], 2.0) + sol = solve(probN, alg, abstol = 1e-11) + @test all(abs.(quadratic_f(sol.u, 2.0)) .< 1e-6) + end +end + +@testitem "DFSane Termination Conditions" setup=[CoreRootfindTesting] tags=[:core] begin + using StaticArrays: @SVector + + @testset "TC: $(nameof(typeof(termination_condition)))" for termination_condition in TERMINATION_CONDITIONS + @testset "u0: $(typeof(u0))" for u0 in ([1.0, 1.0], 1.0, @SVector([1.0, 1.0])) + probN = NonlinearProblem(quadratic_f, u0, 2.0) + sol = solve(probN, DFSane(); termination_condition) + @test all(abs.(quadratic_f(sol.u, 2.0)) .< 1e-10) + end + end +end diff --git a/lib/NonlinearSolveSpectralMethods/test/qa_tests.jl b/lib/NonlinearSolveSpectralMethods/test/qa_tests.jl new file mode 100644 index 000000000..f02aec840 --- /dev/null +++ b/lib/NonlinearSolveSpectralMethods/test/qa_tests.jl @@ -0,0 +1,22 @@ +@testitem "Aqua" tags=[:core] begin + using Aqua, NonlinearSolveSpectralMethods + + Aqua.test_all( + NonlinearSolveSpectralMethods; + piracies = false, ambiguities = false, stale_deps = false, deps_compat = false + ) + Aqua.test_stale_deps(NonlinearSolveSpectralMethods; ignore = [:SciMLJacobianOperators]) + Aqua.test_deps_compat(NonlinearSolveSpectralMethods; ignore = [:SciMLJacobianOperators]) + Aqua.test_piracies(NonlinearSolveSpectralMethods) + Aqua.test_ambiguities(NonlinearSolveSpectralMethods; recursive = false) +end + +@testitem "Explicit Imports" tags=[:core] begin + using ExplicitImports, NonlinearSolveSpectralMethods + + @test check_no_implicit_imports( + NonlinearSolveSpectralMethods; skip = (Base, Core, SciMLBase) + ) === nothing + @test check_no_stale_explicit_imports(NonlinearSolveSpectralMethods) === nothing + @test check_all_qualified_accesses_via_owners(NonlinearSolveSpectralMethods) === nothing +end diff --git a/lib/NonlinearSolveSpectralMethods/test/runtests.jl b/lib/NonlinearSolveSpectralMethods/test/runtests.jl new file mode 100644 index 000000000..5256e0e6f --- /dev/null +++ b/lib/NonlinearSolveSpectralMethods/test/runtests.jl @@ -0,0 +1,23 @@ +using ReTestItems, NonlinearSolveSpectralMethods, Hwloc, InteractiveUtils, Pkg + +@info sprint(InteractiveUtils.versioninfo) + +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( + NonlinearSolveSpectralMethods; tags = (GROUP == "all" ? nothing : [Symbol(GROUP)]), + nworkers = RETESTITEMS_NWORKERS, nworker_threads = RETESTITEMS_NWORKER_THREADS, + testitem_timeout = 3600 +) diff --git a/lib/SciMLJacobianOperators/Project.toml b/lib/SciMLJacobianOperators/Project.toml index c889eb199..7be83ae28 100644 --- a/lib/SciMLJacobianOperators/Project.toml +++ b/lib/SciMLJacobianOperators/Project.toml @@ -1,10 +1,11 @@ name = "SciMLJacobianOperators" uuid = "19f34311-ddf3-4b8b-af20-060888a46c0e" authors = ["Avik Pal and contributors"] -version = "0.1.0" +version = "0.1.1" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" +ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9" ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471" ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9" DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" @@ -16,18 +17,19 @@ SciMLOperators = "c0aeaf25-5076-4817-a8d5-81caf7dfa961" [compat] ADTypes = "1.8.1" Aqua = "0.8.7" +ArrayInterface = "7.16" ConcreteStructs = "0.2.3" ConstructionBase = "1.5" DifferentiationInterface = "0.6.16" Enzyme = "0.13.11" ExplicitImports = "1.9.0" FastClosures = "0.3.2" -FiniteDiff = "2.24.0" +FiniteDiff = "2.24" ForwardDiff = "0.10.36" InteractiveUtils = "<0.0.1, 1" LinearAlgebra = "1.10" ReverseDiff = "1.15" -SciMLBase = "2.54.0" +SciMLBase = "2.58" SciMLOperators = "0.3" Test = "1.10" TestItemRunner = "1" diff --git a/lib/SciMLJacobianOperators/src/SciMLJacobianOperators.jl b/lib/SciMLJacobianOperators/src/SciMLJacobianOperators.jl index bb7a60441..a90c983b3 100644 --- a/lib/SciMLJacobianOperators/src/SciMLJacobianOperators.jl +++ b/lib/SciMLJacobianOperators/src/SciMLJacobianOperators.jl @@ -1,6 +1,7 @@ module SciMLJacobianOperators using ADTypes: ADTypes, AutoSparse +using ArrayInterface: ArrayInterface using ConcreteStructs: @concrete using ConstructionBase: ConstructionBase using DifferentiationInterface: DifferentiationInterface, Constant @@ -15,6 +16,14 @@ const False = Val(false) abstract type AbstractJacobianOperator{T} <: AbstractSciMLOperator{T} end +ArrayInterface.can_setindex(::AbstractJacobianOperator) = false +function ArrayInterface.restructure( + y::AbstractJacobianOperator, x::AbstractJacobianOperator +) + @assert size(y)==size(x) "cannot restructure operators. ensure their sizes match." + return x +end + abstract type AbstractMode end struct VJP <: AbstractMode end diff --git a/lib/SimpleNonlinearSolve/Project.toml b/lib/SimpleNonlinearSolve/Project.toml index cfda24544..8492ac67f 100644 --- a/lib/SimpleNonlinearSolve/Project.toml +++ b/lib/SimpleNonlinearSolve/Project.toml @@ -5,7 +5,6 @@ version = "2.0.0" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" -Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697" ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9" BracketingNonlinearSolve = "70df07ce-3d50-431d-a3e7-ca6ddb60ac1e" CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2" @@ -21,6 +20,7 @@ NonlinearSolveBase = "be0214bd-f91f-a760-ac4e-3421ce2b2da0" PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" Reexport = "189a3867-3050-52da-a836-e630ba90ab69" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" +Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46" StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" [weakdeps] @@ -37,33 +37,33 @@ SimpleNonlinearSolveTrackerExt = "Tracker" [compat] ADTypes = "1.2" -Accessors = "0.1" Aqua = "0.8.7" ArrayInterface = "7.16" -BracketingNonlinearSolve = "1" +BracketingNonlinearSolve = "1.1" ChainRulesCore = "1.24" CommonSolve = "0.2.4" ConcreteStructs = "0.2.3" -DiffEqBase = "6.149" +DiffEqBase = "6.158.3" DifferentiationInterface = "0.6.16" Enzyme = "0.13.11" ExplicitImports = "1.9" FastClosures = "0.3.2" -FiniteDiff = "2.24.0" +FiniteDiff = "2.24" ForwardDiff = "0.10.36" InteractiveUtils = "<0.0.1, 1" LineSearch = "0.1.3" LinearAlgebra = "1.10" MaybeInplace = "0.1.4" NonlinearProblemLibrary = "0.1.2" -NonlinearSolveBase = "1" +NonlinearSolveBase = "1.1" Pkg = "1.10" PolyesterForwardDiff = "0.1" PrecompileTools = "1.2" Random = "1.10" Reexport = "1.2" ReverseDiff = "1.15" -SciMLBase = "2.50" +SciMLBase = "2.58" +Setfield = "1.1.1" StaticArrays = "1.9" StaticArraysCore = "1.4.3" Test = "1.10" diff --git a/lib/SimpleNonlinearSolve/ext/SimpleNonlinearSolveChainRulesCoreExt.jl b/lib/SimpleNonlinearSolve/ext/SimpleNonlinearSolveChainRulesCoreExt.jl index f56dee537..50905279a 100644 --- a/lib/SimpleNonlinearSolve/ext/SimpleNonlinearSolveChainRulesCoreExt.jl +++ b/lib/SimpleNonlinearSolve/ext/SimpleNonlinearSolveChainRulesCoreExt.jl @@ -1,21 +1,26 @@ module SimpleNonlinearSolveChainRulesCoreExt using ChainRulesCore: ChainRulesCore, NoTangent + using NonlinearSolveBase: ImmutableNonlinearProblem using SciMLBase: ChainRulesOriginator, NonlinearLeastSquaresProblem using SimpleNonlinearSolve: SimpleNonlinearSolve, simplenonlinearsolve_solve_up, solve_adjoint -function ChainRulesCore.rrule(::typeof(simplenonlinearsolve_solve_up), +function ChainRulesCore.rrule( + ::typeof(simplenonlinearsolve_solve_up), prob::Union{ImmutableNonlinearProblem, NonlinearLeastSquaresProblem}, - sensealg, u0, u0_changed, p, p_changed, alg, args...; kwargs...) + sensealg, u0, u0_changed, p, p_changed, alg, args...; kwargs... +) out, ∇internal = solve_adjoint( - prob, sensealg, u0, p, ChainRulesOriginator(), alg, args...; kwargs...) + prob, sensealg, u0, p, ChainRulesOriginator(), alg, args...; kwargs... + ) function ∇simplenonlinearsolve_solve_up(Δ) ∂f, ∂prob, ∂sensealg, ∂u0, ∂p, _, ∂args... = ∇internal(Δ) return ( - ∂f, ∂prob, ∂sensealg, ∂u0, NoTangent(), ∂p, NoTangent(), NoTangent(), ∂args...) + ∂f, ∂prob, ∂sensealg, ∂u0, NoTangent(), ∂p, NoTangent(), NoTangent(), ∂args... + ) end return out, ∇simplenonlinearsolve_solve_up end diff --git a/lib/SimpleNonlinearSolve/ext/SimpleNonlinearSolveReverseDiffExt.jl b/lib/SimpleNonlinearSolve/ext/SimpleNonlinearSolveReverseDiffExt.jl index 0a407986e..d34d8bac7 100644 --- a/lib/SimpleNonlinearSolve/ext/SimpleNonlinearSolveReverseDiffExt.jl +++ b/lib/SimpleNonlinearSolve/ext/SimpleNonlinearSolveReverseDiffExt.jl @@ -1,10 +1,11 @@ module SimpleNonlinearSolveReverseDiffExt -using ArrayInterface: ArrayInterface using NonlinearSolveBase: ImmutableNonlinearProblem -using ReverseDiff: ReverseDiff, TrackedArray, TrackedReal using SciMLBase: ReverseDiffOriginator, NonlinearLeastSquaresProblem, remake +using ArrayInterface: ArrayInterface +using ReverseDiff: ReverseDiff, TrackedArray, TrackedReal + using SimpleNonlinearSolve: SimpleNonlinearSolve, solve_adjoint import SimpleNonlinearSolve: simplenonlinearsolve_solve_up diff --git a/lib/SimpleNonlinearSolve/ext/SimpleNonlinearSolveTrackerExt.jl b/lib/SimpleNonlinearSolve/ext/SimpleNonlinearSolveTrackerExt.jl index d29c2ac61..d56854316 100644 --- a/lib/SimpleNonlinearSolve/ext/SimpleNonlinearSolveTrackerExt.jl +++ b/lib/SimpleNonlinearSolve/ext/SimpleNonlinearSolveTrackerExt.jl @@ -1,8 +1,9 @@ module SimpleNonlinearSolveTrackerExt -using ArrayInterface: ArrayInterface using NonlinearSolveBase: ImmutableNonlinearProblem using SciMLBase: TrackerOriginator, NonlinearLeastSquaresProblem, remake + +using ArrayInterface: ArrayInterface using Tracker: Tracker, TrackedArray, TrackedReal using SimpleNonlinearSolve: SimpleNonlinearSolve, solve_adjoint diff --git a/lib/SimpleNonlinearSolve/src/SimpleNonlinearSolve.jl b/lib/SimpleNonlinearSolve/src/SimpleNonlinearSolve.jl index cee8f8dd3..3d8258f7e 100644 --- a/lib/SimpleNonlinearSolve/src/SimpleNonlinearSolve.jl +++ b/lib/SimpleNonlinearSolve/src/SimpleNonlinearSolve.jl @@ -1,33 +1,46 @@ module SimpleNonlinearSolve -using Accessors: @reset -using BracketingNonlinearSolve: BracketingNonlinearSolve -using CommonSolve: CommonSolve, solve, init, solve! using ConcreteStructs: @concrete using FastClosures: @closure -using LineSearch: LiFukushimaLineSearch -using LinearAlgebra: LinearAlgebra, dot -using MaybeInplace: @bb, setindex_trait, CannotSetindex, CanSetindex using PrecompileTools: @compile_workload, @setup_workload using Reexport: @reexport -using SciMLBase: SciMLBase, AbstractNonlinearAlgorithm, NonlinearFunction, NonlinearProblem, - NonlinearLeastSquaresProblem, IntervalNonlinearProblem, ReturnCode, remake +using Setfield: @set! + +using BracketingNonlinearSolve: BracketingNonlinearSolve +using CommonSolve: CommonSolve, solve, init, solve! +using LineSearch: LiFukushimaLineSearch +using MaybeInplace: @bb +using NonlinearSolveBase: NonlinearSolveBase, ImmutableNonlinearProblem, L2_NORM, + nonlinearsolve_forwarddiff_solve, nonlinearsolve_dual_solution, + AbstractNonlinearSolveAlgorithm +using SciMLBase: SciMLBase, NonlinearFunction, NonlinearProblem, + NonlinearLeastSquaresProblem, ReturnCode, remake + +using LinearAlgebra: LinearAlgebra, dot + using StaticArraysCore: StaticArray, SArray, SVector, MArray # AD Dependencies using ADTypes: ADTypes, AutoForwardDiff using DifferentiationInterface: DifferentiationInterface using FiniteDiff: FiniteDiff -using ForwardDiff: ForwardDiff - -using NonlinearSolveBase: NonlinearSolveBase, ImmutableNonlinearProblem, L2_NORM, - nonlinearsolve_forwarddiff_solve, nonlinearsolve_dual_solution +using ForwardDiff: ForwardDiff, Dual const DI = DifferentiationInterface -abstract type AbstractSimpleNonlinearSolveAlgorithm <: AbstractNonlinearAlgorithm end +const DualNonlinearProblem = NonlinearProblem{ + <:Union{Number, <:AbstractArray}, iip, + <:Union{<:Dual{T, V, P}, <:AbstractArray{<:Dual{T, V, P}}} +} where {iip, T, V, P} + +const DualNonlinearLeastSquaresProblem = NonlinearLeastSquaresProblem{ + <:Union{Number, <:AbstractArray}, iip, + <:Union{<:Dual{T, V, P}, <:AbstractArray{<:Dual{T, V, P}}} +} where {iip, T, V, P} -const safe_similar = NonlinearSolveBase.Utils.safe_similar +abstract type AbstractSimpleNonlinearSolveAlgorithm <: AbstractNonlinearSolveAlgorithm end + +const NLBUtils = NonlinearSolveBase.Utils is_extension_loaded(::Val) = false @@ -42,61 +55,66 @@ include("raphson.jl") include("trust_region.jl") # By Pass the highlevel checks for NonlinearProblem for Simple Algorithms -function CommonSolve.solve(prob::NonlinearProblem, - alg::AbstractSimpleNonlinearSolveAlgorithm, args...; kwargs...) +function CommonSolve.solve( + prob::NonlinearProblem, alg::AbstractSimpleNonlinearSolveAlgorithm, args...; + kwargs... +) prob = convert(ImmutableNonlinearProblem, prob) return solve(prob, alg, args...; kwargs...) end function CommonSolve.solve( - prob::NonlinearProblem{<:Union{Number, <:AbstractArray}, iip, - <:Union{ - <:ForwardDiff.Dual{T, V, P}, <:AbstractArray{<:ForwardDiff.Dual{T, V, P}}}}, - alg::AbstractSimpleNonlinearSolveAlgorithm, - args...; - kwargs...) where {T, V, P, iip} + prob::DualNonlinearProblem, alg::AbstractSimpleNonlinearSolveAlgorithm, + args...; kwargs... +) if hasfield(typeof(alg), :autodiff) && alg.autodiff === nothing - @reset alg.autodiff = AutoForwardDiff() + @set! alg.autodiff = AutoForwardDiff() end prob = convert(ImmutableNonlinearProblem, prob) sol, partials = nonlinearsolve_forwarddiff_solve(prob, alg, args...; kwargs...) dual_soln = nonlinearsolve_dual_solution(sol.u, partials, prob.p) return SciMLBase.build_solution( - prob, alg, dual_soln, sol.resid; sol.retcode, sol.stats, sol.original) + prob, alg, dual_soln, sol.resid; sol.retcode, sol.stats, sol.original + ) end function CommonSolve.solve( - prob::NonlinearLeastSquaresProblem{<:Union{Number, <:AbstractArray}, iip, - <:Union{ - <:ForwardDiff.Dual{T, V, P}, <:AbstractArray{<:ForwardDiff.Dual{T, V, P}}}}, - alg::AbstractSimpleNonlinearSolveAlgorithm, - args...; - kwargs...) where {T, V, P, iip} + prob::DualNonlinearLeastSquaresProblem, alg::AbstractSimpleNonlinearSolveAlgorithm, + args...; kwargs... +) if hasfield(typeof(alg), :autodiff) && alg.autodiff === nothing - @reset alg.autodiff = AutoForwardDiff() + @set! alg.autodiff = AutoForwardDiff() end sol, partials = nonlinearsolve_forwarddiff_solve(prob, alg, args...; kwargs...) dual_soln = nonlinearsolve_dual_solution(sol.u, partials, prob.p) return SciMLBase.build_solution( - prob, alg, dual_soln, sol.resid; sol.retcode, sol.stats, sol.original) + prob, alg, dual_soln, sol.resid; sol.retcode, sol.stats, sol.original + ) end function CommonSolve.solve( prob::Union{ImmutableNonlinearProblem, NonlinearLeastSquaresProblem}, alg::AbstractSimpleNonlinearSolveAlgorithm, - args...; sensealg = nothing, u0 = nothing, p = nothing, kwargs...) + args...; sensealg = nothing, u0 = nothing, p = nothing, kwargs... +) if sensealg === nothing && haskey(prob.kwargs, :sensealg) sensealg = prob.kwargs[:sensealg] end new_u0 = u0 !== nothing ? u0 : prob.u0 new_p = p !== nothing ? p : prob.p - return simplenonlinearsolve_solve_up(prob, sensealg, new_u0, u0 === nothing, new_p, - p === nothing, alg, args...; prob.kwargs..., kwargs...) + return simplenonlinearsolve_solve_up( + prob, sensealg, + new_u0, u0 === nothing, + new_p, p === nothing, + alg, args...; + prob.kwargs..., kwargs... + ) end function simplenonlinearsolve_solve_up( prob::Union{ImmutableNonlinearProblem, NonlinearLeastSquaresProblem}, sensealg, u0, - u0_changed, p, p_changed, alg, args...; kwargs...) + u0_changed, p, p_changed, alg, args...; kwargs... +) (u0_changed || p_changed) && (prob = remake(prob; u0, p)) return SciMLBase.__solve(prob, alg, args...; kwargs...) end @@ -130,10 +148,8 @@ function solve_adjoint_internal end #!format: on @compile_workload begin - @sync for alg in algs - for prob in (prob_scalar, prob_iip, prob_oop) - Threads.@spawn CommonSolve.solve(prob, alg; abstol = 1e-2) - end + @sync for prob in (prob_scalar, prob_iip, prob_oop), alg in algs + Threads.@spawn CommonSolve.solve(prob, alg; abstol = 1e-2, verbose = false) end end end diff --git a/lib/SimpleNonlinearSolve/src/broyden.jl b/lib/SimpleNonlinearSolve/src/broyden.jl index 6537a4d2d..48a056b7d 100644 --- a/lib/SimpleNonlinearSolve/src/broyden.jl +++ b/lib/SimpleNonlinearSolve/src/broyden.jl @@ -18,17 +18,19 @@ array problems. end function SimpleBroyden(; - linesearch::Union{Bool, Val{true}, Val{false}} = Val(false), alpha = nothing) + linesearch::Union{Bool, Val{true}, Val{false}} = Val(false), alpha = nothing +) linesearch = linesearch isa Bool ? Val(linesearch) : linesearch return SimpleBroyden(linesearch, alpha) end -function SciMLBase.__solve(prob::ImmutableNonlinearProblem, alg::SimpleBroyden, args...; +function SciMLBase.__solve( + prob::ImmutableNonlinearProblem, alg::SimpleBroyden, args...; abstol = nothing, reltol = nothing, maxiters = 1000, - alias_u0 = false, termination_condition = nothing, kwargs...) - x = Utils.maybe_unaliased(prob.u0, alias_u0) - fx = Utils.get_fx(prob, x) - fx = Utils.eval_f(prob, fx, x) + alias_u0 = false, termination_condition = nothing, kwargs... +) + x = NLBUtils.maybe_unaliased(prob.u0, alias_u0) + fx = NLBUtils.evaluate_f(prob, x) T = promote_type(eltype(fx), eltype(x)) iszero(fx) && @@ -54,9 +56,10 @@ function SciMLBase.__solve(prob::ImmutableNonlinearProblem, alg::SimpleBroyden, @bb δJ⁻¹ = copy(J⁻¹) abstol, reltol, tc_cache = NonlinearSolveBase.init_termination_cache( - prob, abstol, reltol, fx, x, termination_condition, Val(:simple)) + prob, abstol, reltol, fx, x, termination_condition, Val(:simple) + ) - if alg.linesearch === Val(true) + if alg.linesearch isa Val{true} ls_alg = LiFukushimaLineSearch(; nan_maxiters = nothing) ls_cache = init(prob, ls_alg, fx, x) else @@ -75,7 +78,7 @@ function SciMLBase.__solve(prob::ImmutableNonlinearProblem, alg::SimpleBroyden, end @bb @. x = xo + α * δx - fx = Utils.eval_f(prob, fx, x) + fx = NLBUtils.evaluate_f!!(prob, fx, x) @bb @. δf = fx - fprev # Termination Checks @@ -88,8 +91,8 @@ function SciMLBase.__solve(prob::ImmutableNonlinearProblem, alg::SimpleBroyden, @bb @. δJ⁻¹n = (δx - J⁻¹δf) / d - δJ⁻¹n_ = Utils.safe_vec(δJ⁻¹n) - xᵀJ⁻¹_ = Utils.safe_vec(xᵀJ⁻¹) + δJ⁻¹n_ = NLBUtils.safe_vec(δJ⁻¹n) + xᵀJ⁻¹_ = NLBUtils.safe_vec(xᵀJ⁻¹) @bb δJ⁻¹ = δJ⁻¹n_ × transpose(xᵀJ⁻¹_) @bb J⁻¹ .+= δJ⁻¹ diff --git a/lib/SimpleNonlinearSolve/src/dfsane.jl b/lib/SimpleNonlinearSolve/src/dfsane.jl index 0d400b0ce..fb371e3c3 100644 --- a/lib/SimpleNonlinearSolve/src/dfsane.jl +++ b/lib/SimpleNonlinearSolve/src/dfsane.jl @@ -1,7 +1,9 @@ """ - SimpleDFSane(; σ_min::Real = 1e-10, σ_max::Real = 1e10, σ_1::Real = 1.0, + SimpleDFSane(; + σ_min::Real = 1e-10, σ_max::Real = 1e10, σ_1::Real = 1.0, M::Union{Int, Val} = Val(10), γ::Real = 1e-4, τ_min::Real = 0.1, τ_max::Real = 0.5, - nexp::Int = 2, η_strategy::Function = (f_1, k, x, F) -> f_1 ./ k^2) + nexp::Int = 2, η_strategy::Function = (f_1, k, x, F) -> f_1 ./ k^2 + ) A low-overhead implementation of the df-sane method for solving large-scale nonlinear systems of equations. For in depth information about all the parameters and the algorithm, @@ -48,20 +50,26 @@ see [la2006spectral](@citet). M <: Val end -# XXX[breaking]: we should change the names to not have unicode -function SimpleDFSane(; σ_min::Real = 1e-10, σ_max::Real = 1e10, σ_1::Real = 1.0, - M::Union{Int, Val} = Val(10), γ::Real = 1e-4, τ_min::Real = 0.1, τ_max::Real = 0.5, - nexp::Int = 2, η_strategy::F = (f_1, k, x, F) -> f_1 ./ k^2) where {F} +function SimpleDFSane(; + sigma_min::Real = 1e-10, sigma_max::Real = 1e10, sigma_1::Real = 1.0, + M::Union{Int, Val} = Val(10), gamma::Real = 1e-4, tau_min::Real = 0.1, + tau_max::Real = 0.5, n_exp::Int = 2, + eta_strategy::F = (fn_1, n, x_n, f_n) -> fn_1 / n^2 +) where {F} M = M isa Int ? Val(M) : M - return SimpleDFSane(σ_min, σ_max, σ_1, γ, τ_min, τ_max, nexp, η_strategy, M) + return SimpleDFSane( + sigma_min, sigma_max, sigma_1, gamma, tau_min, tau_max, n_exp, + eta_strategy, M + ) end -function SciMLBase.__solve(prob::ImmutableNonlinearProblem, alg::SimpleDFSane, args...; +function SciMLBase.__solve( + prob::ImmutableNonlinearProblem, alg::SimpleDFSane, args...; abstol = nothing, reltol = nothing, maxiters = 1000, alias_u0 = false, - termination_condition = nothing, kwargs...) - x = Utils.maybe_unaliased(prob.u0, alias_u0) - fx = Utils.get_fx(prob, x) - fx = Utils.eval_f(prob, fx, x) + termination_condition = nothing, kwargs... +) + x = NLBUtils.maybe_unaliased(prob.u0, alias_u0) + fx = NLBUtils.evaluate_f(prob, x) T = promote_type(eltype(fx), eltype(x)) σ_min = T(alg.σ_min) @@ -74,7 +82,8 @@ function SciMLBase.__solve(prob::ImmutableNonlinearProblem, alg::SimpleDFSane, a τ_max = T(alg.τ_max) abstol, reltol, tc_cache = NonlinearSolveBase.init_termination_cache( - prob, abstol, reltol, fx, x, termination_condition, Val(:simple)) + prob, abstol, reltol, fx, x, termination_condition, Val(:simple) + ) fx_norm = L2_NORM(fx)^nexp α_1 = one(T) @@ -104,7 +113,7 @@ function SciMLBase.__solve(prob::ImmutableNonlinearProblem, alg::SimpleDFSane, a @bb @. x_cache = x + α_p * d - fx = Utils.eval_f(prob, fx, x_cache) + fx = NLBUtils.evaluate_f!!(prob, fx, x_cache) fx_norm_new = L2_NORM(fx)^nexp while k < maxiters @@ -113,7 +122,7 @@ function SciMLBase.__solve(prob::ImmutableNonlinearProblem, alg::SimpleDFSane, a α_tp = α_p^2 * fx_norm / (fx_norm_new + (T(2) * α_p - T(1)) * fx_norm) @bb @. x_cache = x - α_m * d - fx = Utils.eval_f(prob, fx, x_cache) + fx = NLBUtils.evaluate_f!!(prob, fx, x_cache) fx_norm_new = L2_NORM(fx)^nexp (fx_norm_new ≤ (f_bar + η - γ * α_m^2 * fx_norm)) && break @@ -123,7 +132,7 @@ function SciMLBase.__solve(prob::ImmutableNonlinearProblem, alg::SimpleDFSane, a α_m = clamp(α_tm, τ_min * α_m, τ_max * α_m) @bb @. x_cache = x + α_p * d - fx = Utils.eval_f(prob, fx, x_cache) + fx = NLBUtils.evaluate_f!!(prob, fx, x_cache) fx_norm_new = L2_NORM(fx)^nexp k += 1 @@ -146,11 +155,11 @@ function SciMLBase.__solve(prob::ImmutableNonlinearProblem, alg::SimpleDFSane, a fx_norm = fx_norm_new # Store function value - idx = mod1(k, SciMLBase._unwrap_val(alg.M)) + idx = mod1(k, NLBUtils.unwrap_val(alg.M)) if history_f_k isa SVector history_f_k = Base.setindex(history_f_k, fx_norm_new, idx) elseif history_f_k isa NTuple - @reset history_f_k[idx] = fx_norm_new + @set! history_f_k[idx] = fx_norm_new else history_f_k[idx] = fx_norm_new end diff --git a/lib/SimpleNonlinearSolve/src/halley.jl b/lib/SimpleNonlinearSolve/src/halley.jl index 30eb1a821..2d8446d90 100644 --- a/lib/SimpleNonlinearSolve/src/halley.jl +++ b/lib/SimpleNonlinearSolve/src/halley.jl @@ -23,33 +23,37 @@ end function SciMLBase.__solve( prob::ImmutableNonlinearProblem, alg::SimpleHalley, args...; abstol = nothing, reltol = nothing, maxiters = 1000, - alias_u0 = false, termination_condition = nothing, kwargs...) - x = Utils.maybe_unaliased(prob.u0, alias_u0) - fx = Utils.get_fx(prob, x) - fx = Utils.eval_f(prob, fx, x) + alias_u0 = false, termination_condition = nothing, kwargs... +) + x = NLBUtils.maybe_unaliased(prob.u0, alias_u0) + fx = NLBUtils.evaluate_f(prob, x) T = promote_type(eltype(fx), eltype(x)) iszero(fx) && return SciMLBase.build_solution(prob, alg, x, fx; retcode = ReturnCode.Success) abstol, reltol, tc_cache = NonlinearSolveBase.init_termination_cache( - prob, abstol, reltol, fx, x, termination_condition, Val(:simple)) + prob, abstol, reltol, fx, x, termination_condition, Val(:simple) + ) # The way we write the 2nd order derivatives, we know Enzyme won't work there autodiff = alg.autodiff === nothing ? AutoForwardDiff() : alg.autodiff + @set! alg.autodiff = autodiff @bb xo = copy(x) - strait = setindex_trait(x) - - A = strait isa CanSetindex ? safe_similar(x, length(x), length(x)) : x - Aaᵢ = strait isa CanSetindex ? safe_similar(x, length(x)) : x - cᵢ = strait isa CanSetindex ? safe_similar(x) : x + if NLBUtils.can_setindex(x) + A = NLBUtils.safe_similar(x, length(x), length(x)) + Aaᵢ = NLBUtils.safe_similar(x, length(x)) + cᵢ = NLBUtils.safe_similar(x) + else + A, Aaᵢ, cᵢ = x, x, x + end for _ in 1:maxiters fx, J, H = Utils.compute_jacobian_and_hessian(autodiff, prob, fx, x) - strait isa CannotSetindex && (A = J) + NLBUtils.can_setindex(x) || (A = J) # Factorize Once and Reuse J_fact = if J isa Number @@ -57,22 +61,23 @@ function SciMLBase.__solve( else fact = LinearAlgebra.lu(J; check = false) !LinearAlgebra.issuccess(fact) && return SciMLBase.build_solution( - prob, alg, x, fx; retcode = ReturnCode.Unstable) + prob, alg, x, fx; retcode = ReturnCode.Unstable + ) fact end - aᵢ = J_fact \ Utils.safe_vec(fx) - A_ = Utils.safe_vec(A) + aᵢ = J_fact \ NLBUtils.safe_vec(fx) + A_ = NLBUtils.safe_vec(A) @bb A_ = H × aᵢ - A = Utils.restructure(A, A_) + A = NLBUtils.restructure(A, A_) @bb Aaᵢ = A × aᵢ @bb A .*= -1 - bᵢ = J_fact \ Utils.safe_vec(Aaᵢ) + bᵢ = J_fact \ NLBUtils.safe_vec(Aaᵢ) - cᵢ_ = Utils.safe_vec(cᵢ) + cᵢ_ = NLBUtils.safe_vec(cᵢ) @bb @. cᵢ_ = (aᵢ * aᵢ) / (-aᵢ + (T(0.5) * bᵢ)) - cᵢ = Utils.restructure(cᵢ, cᵢ_) + cᵢ = NLBUtils.restructure(cᵢ, cᵢ_) solved, retcode, fx_sol, x_sol = Utils.check_termination(tc_cache, fx, x, xo, prob) solved && return SciMLBase.build_solution(prob, alg, x_sol, fx_sol; retcode) diff --git a/lib/SimpleNonlinearSolve/src/klement.jl b/lib/SimpleNonlinearSolve/src/klement.jl index 31c4cca96..a8fb7705c 100644 --- a/lib/SimpleNonlinearSolve/src/klement.jl +++ b/lib/SimpleNonlinearSolve/src/klement.jl @@ -6,16 +6,18 @@ method is non-allocating on scalar and static array problems. """ struct SimpleKlement <: AbstractSimpleNonlinearSolveAlgorithm end -function SciMLBase.__solve(prob::ImmutableNonlinearProblem, alg::SimpleKlement, args...; +function SciMLBase.__solve( + prob::ImmutableNonlinearProblem, alg::SimpleKlement, args...; abstol = nothing, reltol = nothing, maxiters = 1000, - alias_u0 = false, termination_condition = nothing, kwargs...) - x = Utils.maybe_unaliased(prob.u0, alias_u0) + alias_u0 = false, termination_condition = nothing, kwargs... +) + x = NLBUtils.maybe_unaliased(prob.u0, alias_u0) T = eltype(x) - fx = Utils.get_fx(prob, x) - fx = Utils.eval_f(prob, fx, x) + fx = NLBUtils.evaluate_f(prob, x) abstol, reltol, tc_cache = NonlinearSolveBase.init_termination_cache( - prob, abstol, reltol, fx, x, termination_condition, Val(:simple)) + prob, abstol, reltol, fx, x, termination_condition, Val(:simple) + ) @bb δx = copy(x) @bb fprev = copy(fx) @@ -31,7 +33,7 @@ function SciMLBase.__solve(prob::ImmutableNonlinearProblem, alg::SimpleKlement, @bb @. δx = fprev / J @bb @. x = xo - δx - fx = Utils.eval_f(prob, fx, x) + fx = NLBUtils.evaluate_f!!(prob, fx, x) # Termination Checks solved, retcode, fx_sol, x_sol = Utils.check_termination(tc_cache, fx, x, xo, prob) diff --git a/lib/SimpleNonlinearSolve/src/lbroyden.jl b/lib/SimpleNonlinearSolve/src/lbroyden.jl index d2bd6ef83..a0d33f942 100644 --- a/lib/SimpleNonlinearSolve/src/lbroyden.jl +++ b/lib/SimpleNonlinearSolve/src/lbroyden.jl @@ -1,6 +1,7 @@ """ - SimpleLimitedMemoryBroyden(; threshold::Union{Val, Int} = Val(27), - linesearch = Val(false), alpha = nothing) + SimpleLimitedMemoryBroyden(; + threshold::Union{Val, Int} = Val(27), linesearch = Val(false), alpha = nothing + ) A limited memory implementation of Broyden. This method applies the L-BFGS scheme to Broyden's method. @@ -40,7 +41,8 @@ function SciMLBase.__solve( if termination_condition === nothing || termination_condition isa NonlinearSolveBase.AbsNormTerminationMode return internal_static_solve( - prob, alg, args...; termination_condition, kwargs...) + prob, alg, args...; termination_condition, kwargs... + ) end @warn "Specifying `termination_condition = $(termination_condition)` for \ `SimpleLimitedMemoryBroyden` with `SArray` is not non-allocating. Use \ @@ -53,22 +55,26 @@ end @views function internal_generic_solve( prob::ImmutableNonlinearProblem, alg::SimpleLimitedMemoryBroyden, args...; abstol = nothing, reltol = nothing, maxiters = 1000, - alias_u0 = false, termination_condition = nothing, kwargs...) - x = Utils.maybe_unaliased(prob.u0, alias_u0) - η = min(SciMLBase._unwrap_val(alg.threshold), maxiters) + alias_u0 = false, termination_condition = nothing, kwargs... +) + x = NLBUtils.maybe_unaliased(prob.u0, alias_u0) + η = min(NLBUtils.unwrap_val(alg.threshold), maxiters) # For scalar problems / if the threshold is larger than problem size just use Broyden if x isa Number || length(x) ≤ η - return SciMLBase.__solve(prob, SimpleBroyden(; alg.linesearch), args...; abstol, - reltol, maxiters, termination_condition, kwargs...) + return SciMLBase.__solve( + prob, SimpleBroyden(; alg.linesearch), args...; + abstol, reltol, maxiters, termination_condition, kwargs... + ) end - fx = Utils.get_fx(prob, x) + fx = NLBUtils.evaluate_f(prob, x) U, Vᵀ = init_low_rank_jacobian(x, fx, x isa StaticArray ? alg.threshold : Val(η)) abstol, reltol, tc_cache = NonlinearSolveBase.init_termination_cache( - prob, abstol, reltol, fx, x, termination_condition, Val(:simple)) + prob, abstol, reltol, fx, x, termination_condition, Val(:simple) + ) @bb xo = copy(x) @bb δx = copy(fx) @@ -80,7 +86,7 @@ end Tcache = lbroyden_threshold_cache(x, x isa StaticArray ? alg.threshold : Val(η)) @bb mat_cache = copy(x) - if alg.linesearch === Val(true) + if alg.linesearch isa Val{true} ls_alg = LiFukushimaLineSearch(; nan_maxiters = nothing) ls_cache = init(prob, ls_alg, fx, x) else @@ -96,7 +102,7 @@ end end @bb @. x = xo + α * δx - fx = Utils.eval_f(prob, fx, x) + fx = NLBUtils.evaluate_f!!(prob, fx, x) @bb @. δf = fx - fo # Termination Checks @@ -111,8 +117,8 @@ end d = dot(vᵀ, δf) @bb @. δx = (δx - mvec) / d - selectdim(U, 2, mod1(i, η)) .= Utils.safe_vec(δx) - selectdim(Vᵀ, 1, mod1(i, η)) .= Utils.safe_vec(vᵀ) + selectdim(U, 2, mod1(i, η)) .= NLBUtils.safe_vec(δx) + selectdim(Vᵀ, 1, mod1(i, η)) .= NLBUtils.safe_vec(vᵀ) Uₚ = selectdim(U, 2, 1:min(η, i)) Vᵀₚ = selectdim(Vᵀ, 1, 1:min(η, i)) @@ -130,10 +136,11 @@ end # finicky, so we'll implement it separately from the generic version # Ignore termination_condition. Don't pass things into internal functions function internal_static_solve( - prob::ImmutableNonlinearProblem{<:SArray}, alg::SimpleLimitedMemoryBroyden, - args...; abstol = nothing, maxiters = 1000, kwargs...) + prob::ImmutableNonlinearProblem{<:SArray}, alg::SimpleLimitedMemoryBroyden, args...; + abstol = nothing, maxiters = 1000, kwargs... +) x = prob.u0 - fx = Utils.get_fx(prob, x) + fx = NLBUtils.evaluate_f(prob, x) U, Vᵀ = init_low_rank_jacobian(vec(x), vec(fx), alg.threshold) @@ -165,7 +172,7 @@ function internal_static_solve( xo, fo, δx = res.x, res.fx, res.δx - for i in 1:(maxiters - SciMLBase._unwrap_val(alg.threshold)) + for i in 1:(maxiters - NLBUtils.unwrap_val(alg.threshold)) if ls_cache === nothing α = true else @@ -174,22 +181,22 @@ function internal_static_solve( end x = xo + α * δx - fx = Utils.eval_f(prob, fx, x) + fx = NLBUtils.evaluate_f!!(prob, fx, x) δf = fx - fo maximum(abs, fx) ≤ abstol && return SciMLBase.build_solution(prob, alg, x, fx; retcode = ReturnCode.Success) - vᵀ = Utils.restructure(x, rmatvec!!(U, Vᵀ, vec(δx), init_α)) - mvec = Utils.restructure(x, matvec!!(U, Vᵀ, vec(δf), init_α)) + vᵀ = NLBUtils.restructure(x, rmatvec!!(U, Vᵀ, vec(δx), init_α)) + mvec = NLBUtils.restructure(x, matvec!!(U, Vᵀ, vec(δf), init_α)) d = dot(vᵀ, δf) δx = @. (δx - mvec) / d - U = Base.setindex(U, vec(δx), mod1(i, SciMLBase._unwrap_val(alg.threshold))) - Vᵀ = Base.setindex(Vᵀ, vec(vᵀ), mod1(i, SciMLBase._unwrap_val(alg.threshold))) + U = Base.setindex(U, vec(δx), mod1(i, NLBUtils.unwrap_val(alg.threshold))) + Vᵀ = Base.setindex(Vᵀ, vec(vᵀ), mod1(i, NLBUtils.unwrap_val(alg.threshold))) - δx = -Utils.restructure(fx, matvec!!(U, Vᵀ, vec(fx), init_α)) + δx = -NLBUtils.restructure(fx, matvec!!(U, Vᵀ, vec(fx), init_α)) xo, fo = x, fx end @@ -198,8 +205,8 @@ function internal_static_solve( end @generated function internal_unrolled_lbroyden_initial_iterations( - prob, xo, fo, δx, abstol, U, Vᵀ, ::Val{threshold}, - ls_cache, init_α) where {threshold} + prob, xo, fo, δx, abstol, U, Vᵀ, ::Val{threshold}, ls_cache, init_α +) where {threshold} calls = [] for i in 1:threshold static_idx, static_idx_p1 = Val(i - 1), Val(i) @@ -219,8 +226,8 @@ end Uₚ = first_n_getindex(U, $(static_idx)) Vᵀₚ = first_n_getindex(Vᵀ, $(static_idx)) - vᵀ = Utils.restructure(x, rmatvec!!(Uₚ, Vᵀₚ, vec(δx), init_α)) - mvec = Utils.restructure(x, matvec!!(Uₚ, Vᵀₚ, vec(δf), init_α)) + vᵀ = NLBUtils.restructure(x, rmatvec!!(Uₚ, Vᵀₚ, vec(δx), init_α)) + mvec = NLBUtils.restructure(x, matvec!!(Uₚ, Vᵀₚ, vec(δf), init_α)) d = dot(vᵀ, δf) δx = @. (δx - mvec) / d @@ -230,7 +237,7 @@ end Uₚ = first_n_getindex(U, $(static_idx_p1)) Vᵀₚ = first_n_getindex(Vᵀ, $(static_idx_p1)) - δx = -Utils.restructure(fx, matvec!!(Uₚ, Vᵀₚ, vec(fx), init_α)) + δx = -NLBUtils.restructure(fx, matvec!!(Uₚ, Vᵀₚ, vec(fx), init_α)) x0, fo = x, fx end) @@ -284,7 +291,8 @@ function fast_mapdot(x::SVector{S1}, Y::SVector{S2, <:SVector{S1}}) where {S1, S return map(Base.Fix1(dot, x), Y) end @generated function fast_mapTdot( - x::SVector{S1}, Y::SVector{S1, <:SVector{S2}}) where {S1, S2} + x::SVector{S1}, Y::SVector{S1, <:SVector{S2}} +) where {S1, S2} calls = [] syms = [gensym("m$(i)") for i in 1:S1] for i in 1:S1 @@ -301,22 +309,26 @@ end return :(return SVector{$N, $T}(($(getcalls...)))) end -lbroyden_threshold_cache(x, ::Val{threshold}) where {threshold} = safe_similar(x, threshold) +function lbroyden_threshold_cache(x, ::Val{threshold}) where {threshold} + return NLBUtils.safe_similar(x, threshold) +end function lbroyden_threshold_cache(x::StaticArray, ::Val{threshold}) where {threshold} return zeros(MArray{Tuple{threshold}, eltype(x)}) end lbroyden_threshold_cache(::SArray, ::Val{threshold}) where {threshold} = nothing -function init_low_rank_jacobian(u::StaticArray{S1, T1}, fu::StaticArray{S2, T2}, - ::Val{threshold}) where {S1, S2, T1, T2, threshold} +function init_low_rank_jacobian( + u::StaticArray{S1, T1}, fu::StaticArray{S2, T2}, ::Val{threshold} +) where {S1, S2, T1, T2, threshold} T = promote_type(T1, T2) fuSize, uSize = Size(fu), Size(u) Vᵀ = MArray{Tuple{threshold, prod(uSize)}, T}(undef) U = MArray{Tuple{prod(fuSize), threshold}, T}(undef) return U, Vᵀ end -@generated function init_low_rank_jacobian(u::SVector{Lu, T1}, fu::SVector{Lfu, T2}, - ::Val{threshold}) where {Lu, Lfu, T1, T2, threshold} +@generated function init_low_rank_jacobian( + u::SVector{Lu, T1}, fu::SVector{Lfu, T2}, ::Val{threshold} +) where {Lu, Lfu, T1, T2, threshold} T = promote_type(T1, T2) inner_inits_Vᵀ = [:(zeros(SVector{$Lu, $T})) for i in 1:threshold] inner_inits_U = [:(zeros(SVector{$Lfu, $T})) for i in 1:threshold] @@ -327,7 +339,7 @@ end end end function init_low_rank_jacobian(u, fu, ::Val{threshold}) where {threshold} - Vᵀ = safe_similar(u, threshold, length(u)) - U = safe_similar(u, length(fu), threshold) + Vᵀ = NLBUtils.safe_similar(u, threshold, length(u)) + U = NLBUtils.safe_similar(u, length(fu), threshold) return U, Vᵀ end diff --git a/lib/SimpleNonlinearSolve/src/raphson.jl b/lib/SimpleNonlinearSolve/src/raphson.jl index ebbb5f9f9..34efcbb90 100644 --- a/lib/SimpleNonlinearSolve/src/raphson.jl +++ b/lib/SimpleNonlinearSolve/src/raphson.jl @@ -27,35 +27,37 @@ function SciMLBase.__solve( prob::Union{ImmutableNonlinearProblem, NonlinearLeastSquaresProblem}, alg::SimpleNewtonRaphson, args...; abstol = nothing, reltol = nothing, maxiters = 1000, - alias_u0 = false, termination_condition = nothing, kwargs...) - x = Utils.maybe_unaliased(prob.u0, alias_u0) - fx = Utils.get_fx(prob, x) - fx = Utils.eval_f(prob, fx, x) + alias_u0 = false, termination_condition = nothing, kwargs... +) + x = NLBUtils.maybe_unaliased(prob.u0, alias_u0) + fx = NLBUtils.evaluate_f(prob, x) iszero(fx) && return SciMLBase.build_solution(prob, alg, x, fx; retcode = ReturnCode.Success) abstol, reltol, tc_cache = NonlinearSolveBase.init_termination_cache( - prob, abstol, reltol, fx, x, termination_condition, Val(:simple)) + prob, abstol, reltol, fx, x, termination_condition, Val(:simple) + ) autodiff = SciMLBase.has_jac(prob.f) ? alg.autodiff : NonlinearSolveBase.select_jacobian_autodiff(prob, alg.autodiff) + @set! alg.autodiff = autodiff @bb xo = similar(x) fx_cache = (SciMLBase.isinplace(prob) && !SciMLBase.has_jac(prob.f)) ? - safe_similar(fx) : fx + NLBUtils.safe_similar(fx) : fx jac_cache = Utils.prepare_jacobian(prob, autodiff, fx_cache, x) J = Utils.compute_jacobian!!(nothing, prob, autodiff, fx_cache, x, jac_cache) for _ in 1:maxiters @bb copyto!(xo, x) - δx = Utils.restructure(x, J \ Utils.safe_vec(fx)) + δx = NLBUtils.restructure(x, J \ NLBUtils.safe_vec(fx)) @bb x .-= δx solved, retcode, fx_sol, x_sol = Utils.check_termination(tc_cache, fx, x, xo, prob) solved && return SciMLBase.build_solution(prob, alg, x_sol, fx_sol; retcode) - fx = Utils.eval_f(prob, fx, x) + fx = NLBUtils.evaluate_f!!(prob, fx, x) J = Utils.compute_jacobian!!(J, prob, autodiff, fx_cache, x, jac_cache) end diff --git a/lib/SimpleNonlinearSolve/src/trust_region.jl b/lib/SimpleNonlinearSolve/src/trust_region.jl index 32e7a6219..d9ac54235 100644 --- a/lib/SimpleNonlinearSolve/src/trust_region.jl +++ b/lib/SimpleNonlinearSolve/src/trust_region.jl @@ -1,9 +1,11 @@ """ - SimpleTrustRegion(; autodiff = AutoForwardDiff(), max_trust_radius = 0.0, + SimpleTrustRegion(; + autodiff = AutoForwardDiff(), max_trust_radius = 0.0, initial_trust_radius = 0.0, step_threshold = nothing, shrink_threshold = nothing, expand_threshold = nothing, shrink_factor = 0.25, expand_factor = 2.0, max_shrink_times::Int = 32, - nlsolve_update_rule = Val(false)) + nlsolve_update_rule = Val(false) + ) A low-overhead implementation of a trust-region solver. This method is non-allocating on scalar and static array problems. @@ -18,18 +20,18 @@ scalar and static array problems. - `initial_trust_radius`: the initial trust region radius. Defaults to `max_trust_radius / 11`. - `step_threshold`: the threshold for taking a step. In every iteration, the threshold is - compared with a value `r`, which is the actual reduction in the objective function divided - by the predicted reduction. If `step_threshold > r` the model is not a good approximation, - and the step is rejected. Defaults to `0.1`. For more details, see + compared with a value `r`, which is the actual reduction in the objective function + divided by the predicted reduction. If `step_threshold > r` the model is not a good + approximation, and the step is rejected. Defaults to `0.1`. For more details, see [Rahpeymaii, F.](https://link.springer.com/article/10.1007/s40096-020-00339-4) - `shrink_threshold`: the threshold for shrinking the trust region radius. In every - iteration, the threshold is compared with a value `r` which is the actual reduction in the - objective function divided by the predicted reduction. If `shrink_threshold > r` the trust - region radius is shrunk by `shrink_factor`. Defaults to `0.25`. For more details, see - [Rahpeymaii, F.](https://link.springer.com/article/10.1007/s40096-020-00339-4) + iteration, the threshold is compared with a value `r` which is the actual reduction in + the objective function divided by the predicted reduction. If `shrink_threshold > r` the + trust region radius is shrunk by `shrink_factor`. Defaults to `0.25`. For more details, + see [Rahpeymaii, F.](https://link.springer.com/article/10.1007/s40096-020-00339-4) - `expand_threshold`: the threshold for expanding the trust region radius. If a step is - taken, i.e `step_threshold < r` (with `r` defined in `shrink_threshold`), a check is also - made to see if `expand_threshold < r`. If that is true, the trust region radius is + taken, i.e `step_threshold < r` (with `r` defined in `shrink_threshold`), a check is + also made to see if `expand_threshold < r`. If that is true, the trust region radius is expanded by `expand_factor`. Defaults to `0.75`. - `shrink_factor`: the factor to shrink the trust region radius with if `shrink_threshold > r` (with `r` defined in `shrink_threshold`). Defaults to `0.25`. @@ -55,29 +57,32 @@ scalar and static array problems. nlsolve_update_rule = Val(false) end -function SciMLBase.__solve(prob::ImmutableNonlinearProblem, alg::SimpleTrustRegion, - args...; abstol = nothing, reltol = nothing, maxiters = 1000, - alias_u0 = false, termination_condition = nothing, kwargs...) - x = Utils.maybe_unaliased(prob.u0, alias_u0) +function SciMLBase.__solve( + prob::Union{ImmutableNonlinearProblem, NonlinearLeastSquaresProblem}, + alg::SimpleTrustRegion, args...; + abstol = nothing, reltol = nothing, maxiters = 1000, + alias_u0 = false, termination_condition = nothing, kwargs... +) + x = NLBUtils.maybe_unaliased(prob.u0, alias_u0) T = eltype(x) Δₘₐₓ = T(alg.max_trust_radius) Δ = T(alg.initial_trust_radius) η₁ = T(alg.step_threshold) if alg.shrink_threshold === nothing - η₂ = T(ifelse(SciMLBase._unwrap_val(alg.nlsolve_update_rule), 0.05, 0.25)) + η₂ = T(ifelse(NLBUtils.unwrap_val(alg.nlsolve_update_rule), 0.05, 0.25)) else η₂ = T(alg.shrink_threshold) end if alg.expand_threshold === nothing - η₃ = T(ifelse(SciMLBase._unwrap_val(alg.nlsolve_update_rule), 0.9, 0.75)) + η₃ = T(ifelse(NLBUtils.unwrap_val(alg.nlsolve_update_rule), 0.9, 0.75)) else η₃ = T(alg.expand_threshold) end if alg.shrink_factor === nothing - t₁ = T(ifelse(SciMLBase._unwrap_val(alg.nlsolve_update_rule), 0.5, 0.25)) + t₁ = T(ifelse(NLBUtils.unwrap_val(alg.nlsolve_update_rule), 0.5, 0.25)) else t₁ = T(alg.shrink_factor) end @@ -88,23 +93,23 @@ function SciMLBase.__solve(prob::ImmutableNonlinearProblem, alg::SimpleTrustRegi autodiff = SciMLBase.has_jac(prob.f) ? alg.autodiff : NonlinearSolveBase.select_jacobian_autodiff(prob, alg.autodiff) - fx = Utils.get_fx(prob, x) - fx = Utils.eval_f(prob, fx, x) + fx = NLBUtils.evaluate_f(prob, x) norm_fx = L2_NORM(fx) @bb xo = copy(x) fx_cache = (SciMLBase.isinplace(prob) && !SciMLBase.has_jac(prob.f)) ? - safe_similar(fx) : fx + NLBUtils.safe_similar(fx) : fx jac_cache = Utils.prepare_jacobian(prob, autodiff, fx_cache, x) J = Utils.compute_jacobian!!(nothing, prob, autodiff, fx_cache, x, jac_cache) abstol, reltol, tc_cache = NonlinearSolveBase.init_termination_cache( - prob, abstol, reltol, fx, x, termination_condition, Val(:simple)) + prob, abstol, reltol, fx, x, termination_condition, Val(:simple) + ) # Set default trust region radius if not specified by user. iszero(Δₘₐₓ) && (Δₘₐₓ = max(L2_NORM(fx), maximum(x) - minimum(x))) if iszero(Δ) - if SciMLBase._unwrap_val(alg.nlsolve_update_rule) + if NLBUtils.unwrap_val(alg.nlsolve_update_rule) norm_x = L2_NORM(x) Δ = T(ifelse(norm_x > 0, norm_x, 1)) else @@ -114,7 +119,7 @@ function SciMLBase.__solve(prob::ImmutableNonlinearProblem, alg::SimpleTrustRegi fₖ = 0.5 * norm_fx^2 H = transpose(J) * J - g = Utils.restructure(x, J' * Utils.safe_vec(fx)) + g = NLBUtils.restructure(x, J' * NLBUtils.safe_vec(fx)) shrink_counter = 0 @bb δsd = copy(x) @@ -128,7 +133,7 @@ function SciMLBase.__solve(prob::ImmutableNonlinearProblem, alg::SimpleTrustRegi δ = dogleg_method!!(dogleg_cache, J, fx, g, Δ) @bb @. x = xo + δ - fx = Utils.eval_f(prob, fx, x) + fx = NLBUtils.evaluate_f!!(prob, fx, x) fₖ₊₁ = L2_NORM(fx)^2 / T(2) @@ -149,17 +154,18 @@ function SciMLBase.__solve(prob::ImmutableNonlinearProblem, alg::SimpleTrustRegi if r ≥ η₁ # Termination Checks solved, retcode, fx_sol, x_sol = Utils.check_termination( - tc_cache, fx, x, xo, prob) + tc_cache, fx, x, xo, prob + ) solved && return SciMLBase.build_solution(prob, alg, x_sol, fx_sol; retcode) # Take the step. @bb copyto!(xo, x) J = Utils.compute_jacobian!!(J, prob, autodiff, fx_cache, x, jac_cache) - fx = Utils.eval_f(prob, fx, x) + fx = NLBUtils.evaluate_f!!(prob, fx, x) # Update the trust region radius. - if !SciMLBase._unwrap_val(alg.nlsolve_update_rule) && r > η₃ + if !NLBUtils.unwrap_val(alg.nlsolve_update_rule) && r > η₃ Δ = min(t₂ * Δ, Δₘₐₓ) end fₖ = fₖ₊₁ @@ -168,7 +174,7 @@ function SciMLBase.__solve(prob::ImmutableNonlinearProblem, alg::SimpleTrustRegi @bb g = transpose(J) × vec(fx) end - if SciMLBase._unwrap_val(alg.nlsolve_update_rule) + if NLBUtils.unwrap_val(alg.nlsolve_update_rule) if r > η₃ Δ = t₂ * L2_NORM(δ) elseif r > 0.5 @@ -184,7 +190,7 @@ function dogleg_method!!(cache, J, f, g, Δ) (; δsd, δN_δsd, δN) = cache # Compute the Newton step - @bb δN .= Utils.restructure(δN, J \ Utils.safe_vec(f)) + @bb δN .= NLBUtils.restructure(δN, J \ NLBUtils.safe_vec(f)) @bb δN .*= -1 # Test if the full step is within the trust region (L2_NORM(δN) ≤ Δ) && return δN diff --git a/lib/SimpleNonlinearSolve/src/utils.jl b/lib/SimpleNonlinearSolve/src/utils.jl index bd7368bd7..8c35a324f 100644 --- a/lib/SimpleNonlinearSolve/src/utils.jl +++ b/lib/SimpleNonlinearSolve/src/utils.jl @@ -1,70 +1,31 @@ module Utils using ArrayInterface: ArrayInterface -using ConcreteStructs: @concrete using DifferentiationInterface: DifferentiationInterface, Constant using FastClosures: @closure using LinearAlgebra: LinearAlgebra, I, diagind -using NonlinearSolveBase: NonlinearSolveBase, ImmutableNonlinearProblem, - AbstractNonlinearTerminationMode, +using NonlinearSolveBase: NonlinearSolveBase, AbstractNonlinearTerminationMode, AbstractSafeNonlinearTerminationMode, AbstractSafeBestNonlinearTerminationMode -using SciMLBase: SciMLBase, AbstractNonlinearProblem, NonlinearLeastSquaresProblem, - NonlinearProblem, NonlinearFunction, ReturnCode +using SciMLBase: SciMLBase, ReturnCode using StaticArraysCore: StaticArray, SArray, SMatrix, SVector const DI = DifferentiationInterface - -const safe_similar = NonlinearSolveBase.Utils.safe_similar - -pickchunksize(n::Int) = min(n, 12) - -can_dual(::Type{<:Real}) = true -can_dual(::Type) = false - -maybe_unaliased(x::Union{Number, SArray}, ::Bool) = x -function maybe_unaliased(x::T, alias::Bool) where {T <: AbstractArray} - (alias || !ArrayInterface.can_setindex(T)) && return x - return copy(x) -end - -# NOTE: This doesn't initialize the `f(x)` but just returns a buffer of the same size -function get_fx(prob::NonlinearLeastSquaresProblem, x) - if SciMLBase.isinplace(prob) && prob.f.resid_prototype === nothing - error("Inplace NonlinearLeastSquaresProblem requires a `resid_prototype` to be \ - specified.") - end - return get_fx(prob.f, x, prob.p) -end -function get_fx(prob::Union{ImmutableNonlinearProblem, NonlinearProblem}, x) - return get_fx(prob.f, x, prob.p) -end -function get_fx(f::NonlinearFunction, x, p) - if SciMLBase.isinplace(f) - f.resid_prototype === nothing || return eltype(x).(f.resid_prototype) - return safe_similar(x) - end - return f(x, p) -end - -function eval_f(prob, fx, x) - SciMLBase.isinplace(prob) || return prob.f(x, prob.p) - prob.f(fx, x, prob.p) - return fx -end - -function fixed_parameter_function(prob::AbstractNonlinearProblem) - SciMLBase.isinplace(prob) && return @closure (du, u) -> prob.f(du, u, prob.p) - return Base.Fix2(prob.f, prob.p) -end +const NLBUtils = NonlinearSolveBase.Utils function identity_jacobian(u::Number, fu::Number, α = true) return convert(promote_type(eltype(u), eltype(fu)), α) end function identity_jacobian(u, fu, α = true) - J = safe_similar(u, promote_type(eltype(u), eltype(fu)), length(fu), length(u)) - fill!(J, zero(eltype(J))) - J[diagind(J)] .= eltype(J)(α) + J = NLBUtils.safe_similar(u, promote_type(eltype(u), eltype(fu)), length(fu), length(u)) + fill!(J, false) + if ArrayInterface.fast_scalar_indexing(J) + @simd ivdep for i in axes(J, 1) + @inbounds J[i, i] = α + end + else + J[diagind(J)] .= α + end return J end function identity_jacobian(u::StaticArray, fu, α = true) @@ -97,30 +58,21 @@ function check_termination(cache, fx, x, xo, _, ::AbstractSafeNonlinearTerminati return cache(fx, x, xo), cache.retcode, fx, x end function check_termination( - cache, fx, x, xo, prob, ::AbstractSafeBestNonlinearTerminationMode) + cache, fx, x, xo, prob, ::AbstractSafeBestNonlinearTerminationMode +) if cache(fx, x, xo) x = cache.u - if SciMLBase.isinplace(prob) - prob.f(fx, x, prob.p) - else - fx = prob.f(x, prob.p) - end + fx = NLBUtils.evaluate_f!!(prob, fx, x) return true, cache.retcode, fx, x end return false, ReturnCode.Default, fx, x end -restructure(y, x) = ArrayInterface.restructure(y, x) -restructure(::Number, x::Number) = x - -safe_vec(x::AbstractArray) = vec(x) -safe_vec(x::Number) = x - abstract type AbstractJacobianMode end struct AnalyticJacobian <: AbstractJacobianMode end -@concrete struct DIExtras <: AbstractJacobianMode - prep +struct DIExtras{P} <: AbstractJacobianMode + prep::P end struct DINoPreparation <: AbstractJacobianMode end @@ -161,7 +113,7 @@ end function compute_jacobian!!(J, prob, autodiff, fx, x, ::AnalyticJacobian) if J === nothing if SciMLBase.isinplace(prob.f) - J = safe_similar(fx, length(fx), length(x)) + J = NLBUtils.safe_similar(fx, length(fx), length(x)) prob.f.jac(J, x, prob.p) return J else @@ -214,16 +166,16 @@ end function compute_jacobian_and_hessian(autodiff, prob, fx, x) if SciMLBase.isinplace(prob) jac_fn = @closure (u, p) -> begin - du = safe_similar(fx, promote_type(eltype(fx), eltype(u))) + du = NLBUtils.safe_similar(fx, promote_type(eltype(fx), eltype(u))) return DI.jacobian(prob.f, du, autodiff, u, Constant(p)) end J, H = DI.value_and_jacobian(jac_fn, autodiff, x, Constant(prob.p)) - fx = Utils.eval_f(prob, fx, x) + fx = NLBUtils.evaluate_f!!(prob, fx, x) return fx, J, H else jac_fn = @closure (u, p) -> DI.jacobian(prob.f, autodiff, u, Constant(p)) J, H = DI.value_and_jacobian(jac_fn, autodiff, x, Constant(prob.p)) - fx = Utils.eval_f(prob, fx, x) + fx = NLBUtils.evaluate_f!!(prob, fx, x) return fx, J, H end end diff --git a/lib/SimpleNonlinearSolve/test/core/23_test_problems_tests.jl b/lib/SimpleNonlinearSolve/test/core/23_test_problems_tests.jl deleted file mode 100644 index 6dd85b94b..000000000 --- a/lib/SimpleNonlinearSolve/test/core/23_test_problems_tests.jl +++ /dev/null @@ -1,106 +0,0 @@ -@testsnippet RobustnessTestSnippet begin - using NonlinearProblemLibrary, NonlinearSolveBase, LinearAlgebra, ADTypes - - problems = NonlinearProblemLibrary.problems - dicts = NonlinearProblemLibrary.dicts - - function test_on_library( - problems, dicts, alg_ops, broken_tests, ϵ = 1e-4; skip_tests = nothing) - for (idx, (problem, dict)) in enumerate(zip(problems, dicts)) - x = dict["start"] - res = similar(x) - nlprob = NonlinearProblem(problem, copy(x)) - @testset "$idx: $(dict["title"])" begin - for alg in alg_ops - try - sol = solve(nlprob, alg; - termination_condition = AbsNormTerminationMode(norm)) - problem(res, sol.u, nothing) - - skip = skip_tests !== nothing && idx in skip_tests[alg] - if skip - @test_skip norm(res) ≤ ϵ - continue - end - broken = idx in broken_tests[alg] ? true : false - @test norm(res)≤ϵ broken=broken - catch e - @error e - broken = idx in broken_tests[alg] ? true : false - if broken - @test false broken=true - else - @test 1 == 2 - end - end - end - end - end - end -end - -@testitem "23 Test Problems: SimpleNewtonRaphson" setup=[RobustnessTestSnippet] tags=[:core] begin - alg_ops = (SimpleNewtonRaphson(; autodiff = AutoForwardDiff()),) - - broken_tests = Dict(alg => Int[] for alg in alg_ops) - broken_tests[alg_ops[1]] = [] - - test_on_library(problems, dicts, alg_ops, broken_tests) -end - -@testitem "23 Test Problems: SimpleHalley" setup=[RobustnessTestSnippet] tags=[:core] begin - alg_ops = (SimpleHalley(; autodiff = AutoForwardDiff()),) - - broken_tests = Dict(alg => Int[] for alg in alg_ops) - if Sys.isapple() - broken_tests[alg_ops[1]] = [1, 5, 11, 15, 16, 18] - else - broken_tests[alg_ops[1]] = [1, 5, 15, 16, 18] - end - - test_on_library(problems, dicts, alg_ops, broken_tests) -end - -@testitem "23 Test Problems: SimpleTrustRegion" setup=[RobustnessTestSnippet] tags=[:core] begin - alg_ops = ( - SimpleTrustRegion(; autodiff = AutoForwardDiff()), - SimpleTrustRegion(; nlsolve_update_rule = Val(true), autodiff = AutoForwardDiff()) - ) - - broken_tests = Dict(alg => Int[] for alg in alg_ops) - broken_tests[alg_ops[1]] = [3, 15, 16, 21] - broken_tests[alg_ops[2]] = [15, 16] - - test_on_library(problems, dicts, alg_ops, broken_tests) -end - -@testitem "23 Test Problems: SimpleDFSane" setup=[RobustnessTestSnippet] tags=[:core] begin - alg_ops = (SimpleDFSane(),) - - broken_tests = Dict(alg => Int[] for alg in alg_ops) - if Sys.isapple() - broken_tests[alg_ops[1]] = [1, 2, 3, 5, 6, 21] - else - broken_tests[alg_ops[1]] = [1, 2, 3, 4, 5, 6, 11, 21] - end - - test_on_library(problems, dicts, alg_ops, broken_tests) -end - -@testitem "23 Test Problems: SimpleBroyden" setup=[RobustnessTestSnippet] tags=[:core] begin - alg_ops = (SimpleBroyden(),) - - broken_tests = Dict(alg => Int[] for alg in alg_ops) - broken_tests[alg_ops[1]] = [1, 5, 11] - - test_on_library(problems, dicts, alg_ops, broken_tests) -end - -@testitem "23 Test Problems: SimpleKlement" setup=[RobustnessTestSnippet] tags=[:core] begin - alg_ops = (SimpleKlement(),) - - broken_tests = Dict(alg => Int[] for alg in alg_ops) - broken_tests[alg_ops[1]] = [1, 2, 4, 5, 11, 12, 22] - - test_on_library(problems, dicts, alg_ops, broken_tests) -end diff --git a/lib/SimpleNonlinearSolve/test/core/qa_tests.jl b/lib/SimpleNonlinearSolve/test/core/qa_tests.jl index cef74ac38..071d77d8d 100644 --- a/lib/SimpleNonlinearSolve/test/core/qa_tests.jl +++ b/lib/SimpleNonlinearSolve/test/core/qa_tests.jl @@ -1,10 +1,17 @@ @testitem "Aqua" tags=[:core] begin using Aqua, SimpleNonlinearSolve - Aqua.test_all(SimpleNonlinearSolve; piracies = false, ambiguities = false) + Aqua.test_all( + SimpleNonlinearSolve; + piracies = false, ambiguities = false, stale_deps = false, deps_compat = false + ) + Aqua.test_stale_deps(SimpleNonlinearSolve; ignore = [:SciMLJacobianOperators]) + Aqua.test_deps_compat(SimpleNonlinearSolve; ignore = [:SciMLJacobianOperators]) Aqua.test_piracies(SimpleNonlinearSolve; treat_as_own = [ - NonlinearProblem, NonlinearLeastSquaresProblem, IntervalNonlinearProblem]) + NonlinearProblem, NonlinearLeastSquaresProblem, IntervalNonlinearProblem + ] + ) Aqua.test_ambiguities(SimpleNonlinearSolve; recursive = false) end diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl index 465a776ed..7d6ac7b09 100644 --- a/src/NonlinearSolve.jl +++ b/src/NonlinearSolve.jl @@ -1,207 +1,101 @@ module NonlinearSolve +using ConcreteStructs: @concrete using Reexport: @reexport using PrecompileTools: @compile_workload, @setup_workload +using FastClosures: @closure -using ArrayInterface: ArrayInterface, can_setindex, restructure, fast_scalar_indexing, - ismutable -using CommonSolve: solve, init, solve! -using ConcreteStructs: @concrete +using ADTypes: ADTypes +using ArrayInterface: ArrayInterface +using CommonSolve: CommonSolve, solve, solve! using DiffEqBase: DiffEqBase # Needed for `init` / `solve` dispatches -using FastClosures: @closure -using LazyArrays: LazyArrays, ApplyArray, cache -using LinearAlgebra: LinearAlgebra, ColumnNorm, Diagonal, I, LowerTriangular, Symmetric, - UpperTriangular, axpy!, cond, diag, diagind, dot, issuccess, istril, - istriu, lu, mul!, norm, pinv, tril!, triu! -using LineSearch: LineSearch, AbstractLineSearchCache, LineSearchesJL, NoLineSearch, - RobustNonMonotoneLineSearch, BackTracking, LiFukushimaLineSearch -using LinearSolve: LinearSolve, QRFactorization, needs_concrete_A, AbstractFactorization, - DefaultAlgorithmChoice, DefaultLinearSolver -using MaybeInplace: @bb -using NonlinearSolveBase: NonlinearSolveBase, nonlinearsolve_forwarddiff_solve, - nonlinearsolve_dual_solution, nonlinearsolve_∂f_∂p, - nonlinearsolve_∂f_∂u, L2_NORM, AbsNormTerminationMode, - AbstractNonlinearTerminationMode, - AbstractSafeBestNonlinearTerminationMode, - select_forward_mode_autodiff, select_reverse_mode_autodiff, - select_jacobian_autodiff -using Printf: @printf -using Preferences: Preferences, @load_preference, @set_preferences! -using RecursiveArrayTools: recursivecopy! -using SciMLBase: SciMLBase, AbstractNonlinearAlgorithm, AbstractNonlinearProblem, - _unwrap_val, isinplace, NLStats, NonlinearFunction, - NonlinearLeastSquaresProblem, NonlinearProblem, ReturnCode, get_du, step!, - set_u!, LinearProblem, IdentityOperator -using SciMLOperators: AbstractSciMLOperator +using LinearAlgebra: LinearAlgebra, norm +using LineSearch: BackTracking +using NonlinearSolveBase: NonlinearSolveBase, InternalAPI, AbstractNonlinearSolveAlgorithm, + AbstractNonlinearSolveCache, Utils, L2_NORM, + enable_timer_outputs, disable_timer_outputs + +using Preferences: set_preferences! +using SciMLBase: SciMLBase, NLStats, ReturnCode, AbstractNonlinearProblem, NonlinearProblem, + NonlinearLeastSquaresProblem +using SymbolicIndexingInterface: SymbolicIndexingInterface +using StaticArraysCore: StaticArray + +# Default Algorithm +using NonlinearSolveFirstOrder: NewtonRaphson, TrustRegion, LevenbergMarquardt, GaussNewton, + RUS +using NonlinearSolveQuasiNewton: Broyden, Klement +using SimpleNonlinearSolve: SimpleBroyden, SimpleKlement + +# Default AD Support +using FiniteDiff: FiniteDiff # Default Finite Difference Method +using ForwardDiff: ForwardDiff, Dual # Default Forward Mode AD + +# Sparse AD Support: Implemented via extensions +using SparseArrays: SparseArrays +using SparseMatrixColorings: SparseMatrixColorings + +# Sub-Packages that are re-exported by NonlinearSolve +using BracketingNonlinearSolve: BracketingNonlinearSolve +using LineSearch: LineSearch +using LinearSolve: LinearSolve +using NonlinearSolveFirstOrder: NonlinearSolveFirstOrder, GeneralizedFirstOrderAlgorithm +using NonlinearSolveQuasiNewton: NonlinearSolveQuasiNewton, QuasiNewtonAlgorithm +using NonlinearSolveSpectralMethods: NonlinearSolveSpectralMethods, GeneralizedDFSane using SimpleNonlinearSolve: SimpleNonlinearSolve -using StaticArraysCore: StaticArray, SVector, SArray, MArray, Size, SMatrix -using SymbolicIndexingInterface: SymbolicIndexingInterface, ParameterIndexingProxy, - symbolic_container, parameter_values, state_values, getu, - setu - -# AD Support -using ADTypes: ADTypes, AbstractADType, AutoFiniteDiff, AutoForwardDiff, - AutoPolyesterForwardDiff, AutoZygote, AutoEnzyme, AutoSparse, - NoSparsityDetector, KnownJacobianSparsityDetector -using DifferentiationInterface: DifferentiationInterface, Constant -using FiniteDiff: FiniteDiff -using ForwardDiff: ForwardDiff, Dual -using SciMLJacobianOperators: AbstractJacobianOperator, JacobianOperator, VecJacOperator, - JacVecOperator, StatefulJacobianOperator - -## Sparse AD Support -using SparseArrays: AbstractSparseMatrix, SparseMatrixCSC -using SparseMatrixColorings: ConstantColoringAlgorithm, GreedyColoringAlgorithm, - LargestFirst - -const DI = DifferentiationInterface - -const True = Val(true) -const False = Val(false) - -include("abstract_types.jl") -include("timer_outputs.jl") -include("internal/helpers.jl") - -include("descent/common.jl") -include("descent/newton.jl") -include("descent/steepest.jl") -include("descent/dogleg.jl") -include("descent/damped_newton.jl") -include("descent/geodesic_acceleration.jl") - -include("internal/jacobian.jl") -include("internal/linear_solve.jl") -include("internal/termination.jl") -include("internal/tracing.jl") -include("internal/approximate_initialization.jl") - -include("globalization/line_search.jl") -include("globalization/trust_region.jl") - -include("core/generic.jl") -include("core/approximate_jacobian.jl") -include("core/generalized_first_order.jl") -include("core/spectral_methods.jl") -include("core/noinit.jl") - -include("algorithms/raphson.jl") -include("algorithms/pseudo_transient.jl") -include("algorithms/broyden.jl") -include("algorithms/klement.jl") -include("algorithms/lbroyden.jl") -include("algorithms/dfsane.jl") -include("algorithms/gauss_newton.jl") -include("algorithms/levenberg_marquardt.jl") -include("algorithms/trust_region.jl") -include("algorithms/extension_algs.jl") - -include("utils.jl") + +const SII = SymbolicIndexingInterface + +include("polyalg.jl") +include("extension_algs.jl") + include("default.jl") const ALL_SOLVER_TYPES = [ - Nothing, AbstractNonlinearSolveAlgorithm, GeneralizedDFSane, - GeneralizedFirstOrderAlgorithm, ApproximateJacobianSolveAlgorithm, + Nothing, AbstractNonlinearSolveAlgorithm, + GeneralizedDFSane, GeneralizedFirstOrderAlgorithm, QuasiNewtonAlgorithm, LeastSquaresOptimJL, FastLevenbergMarquardtJL, NLsolveJL, NLSolversJL, SpeedMappingJL, FixedPointAccelerationJL, SIAMFANLEquationsJL, CMINPACK, PETScSNES, - NonlinearSolvePolyAlgorithm{:NLLS, <:Any}, NonlinearSolvePolyAlgorithm{:NLS, <:Any} + NonlinearSolvePolyAlgorithm ] -include("internal/forward_diff.jl") # we need to define after the algorithms +include("forward_diff.jl") @setup_workload begin - nlfuncs = ((NonlinearFunction{false}((u, p) -> u .* u .- p), 0.1), - (NonlinearFunction{true}((du, u, p) -> du .= u .* u .- p), [0.1])) - probs_nls = NonlinearProblem[] - for (fn, u0) in nlfuncs - push!(probs_nls, NonlinearProblem(fn, u0, 2.0)) - end - - nls_algs = ( - NewtonRaphson(), - TrustRegion(), - LevenbergMarquardt(), - Broyden(), - Klement(), - nothing - ) - - probs_nlls = NonlinearLeastSquaresProblem[] - nlfuncs = ( - (NonlinearFunction{false}((u, p) -> (u .^ 2 .- p)[1:1]), [0.1, 0.0]), - (NonlinearFunction{false}((u, p) -> vcat(u .* u .- p, u .* u .- p)), [0.1, 0.1]), - ( - NonlinearFunction{true}( - (du, u, p) -> du[1] = u[1] * u[1] - p, resid_prototype = zeros(1)), - [0.1, 0.0]), - ( - NonlinearFunction{true}((du, u, p) -> du .= vcat(u .* u .- p, u .* u .- p), - resid_prototype = zeros(4)), - [0.1, 0.1] - ) - ) - for (fn, u0) in nlfuncs - push!(probs_nlls, NonlinearLeastSquaresProblem(fn, u0, 2.0)) - end - - nlls_algs = ( - LevenbergMarquardt(), - GaussNewton(), - TrustRegion(), - nothing - ) + include("../common/nonlinear_problem_workloads.jl") + include("../common/nlls_problem_workloads.jl") @compile_workload begin @sync begin - for T in (Float32, Float64), (fn, u0) in nlfuncs - Threads.@spawn NonlinearProblem(fn, T.(u0), T(2)) - end - for (fn, u0) in nlfuncs - Threads.@spawn NonlinearLeastSquaresProblem(fn, u0, 2.0) - end - for prob in probs_nls, alg in nls_algs - Threads.@spawn solve(prob, alg; abstol = 1e-2, verbose = false) + for prob in nonlinear_problems + Threads.@spawn CommonSolve.solve( + prob, nothing; abstol = 1e-2, verbose = false + ) end - for prob in probs_nlls, alg in nlls_algs - Threads.@spawn solve(prob, alg; abstol = 1e-2, verbose = false) + + for prob in nlls_problems + Threads.@spawn CommonSolve.solve( + prob, nothing; abstol = 1e-2, verbose = false + ) end end end end # Rexexports -@reexport using SciMLBase, SimpleNonlinearSolve, NonlinearSolveBase +@reexport using SciMLBase, NonlinearSolveBase, LineSearch, ADTypes +@reexport using NonlinearSolveFirstOrder, NonlinearSolveSpectralMethods, + NonlinearSolveQuasiNewton, SimpleNonlinearSolve, BracketingNonlinearSolve +@reexport using LinearSolve -# Core Algorithms -export NewtonRaphson, PseudoTransient, Klement, Broyden, LimitedMemoryBroyden, DFSane -export GaussNewton, LevenbergMarquardt, TrustRegion -export NonlinearSolvePolyAlgorithm, RobustMultiNewton, FastShortcutNonlinearPolyalg, - FastShortcutNLLSPolyalg +# Poly Algorithms +export NonlinearSolvePolyAlgorithm, + RobustMultiNewton, FastShortcutNonlinearPolyalg, FastShortcutNLLSPolyalg # Extension Algorithms export LeastSquaresOptimJL, FastLevenbergMarquardtJL, NLsolveJL, NLSolversJL, FixedPointAccelerationJL, SpeedMappingJL, SIAMFANLEquationsJL export PETScSNES, CMINPACK -# Advanced Algorithms -- Without Bells and Whistles -export GeneralizedFirstOrderAlgorithm, ApproximateJacobianSolveAlgorithm, GeneralizedDFSane - -# Descent Algorithms -export NewtonDescent, SteepestDescent, Dogleg, DampedNewtonDescent, GeodesicAcceleration - -# Globalization -## Line Search Algorithms -export LineSearch, BackTracking, NoLineSearch, RobustNonMonotoneLineSearch, - LiFukushimaLineSearch, LineSearchesJL -## Trust Region Algorithms -export RadiusUpdateSchemes - -# Tracing Functionality -export TraceAll, TraceMinimal, TraceWithJacobianConditionNumber - -# Reexport ADTypes -export AutoFiniteDiff, AutoForwardDiff, AutoPolyesterForwardDiff, AutoZygote, AutoEnzyme, - AutoSparse - end diff --git a/src/abstract_types.jl b/src/abstract_types.jl deleted file mode 100644 index 255c5e541..000000000 --- a/src/abstract_types.jl +++ /dev/null @@ -1,503 +0,0 @@ -function __internal_init end -function __internal_solve! end - -""" - AbstractDescentAlgorithm - -Given the Jacobian `J` and the residual `fu`, this type of algorithm computes the descent -direction `δu`. - -For non-square Jacobian problems, if we need to solve a linear solve problem, we use a least -squares solver by default, unless the provided `linsolve` can't handle non-square matrices, -in which case we use the normal form equations ``JᵀJ δu = Jᵀ fu``. Note that this -factorization is often the faster choice, but it is not as numerically stable as the least -squares solver. - -### `__internal_init` specification - -```julia -__internal_init(prob::NonlinearProblem{uType, iip}, alg::AbstractDescentAlgorithm, J, - fu, u; pre_inverted::Val{INV} = Val(false), linsolve_kwargs = (;), - abstol = nothing, reltol = nothing, alias_J::Bool = true, - shared::Val{N} = Val(1), kwargs...) where {INV, N, uType, iip} --> AbstractDescentCache - -__internal_init( - prob::NonlinearLeastSquaresProblem{uType, iip}, alg::AbstractDescentAlgorithm, - J, fu, u; pre_inverted::Val{INV} = Val(false), linsolve_kwargs = (;), - abstol = nothing, reltol = nothing, alias_J::Bool = true, - shared::Val{N} = Val(1), kwargs...) where {INV, N, uType, iip} --> AbstractDescentCache -``` - - - `pre_inverted`: whether or not the Jacobian has been pre_inverted. Defaults to `False`. - Note that for most algorithms except `NewtonDescent` setting it to `Val(true)` is - generally a bad idea. - - `linsolve_kwargs`: keyword arguments to pass to the linear solver. Defaults to `(;)`. - - `abstol`: absolute tolerance for the linear solver. Defaults to `nothing`. - - `reltol`: relative tolerance for the linear solver. Defaults to `nothing`. - - `alias_J`: whether or not to alias the Jacobian. Defaults to `true`. - - `shared`: Store multiple descent directions in the cache. Allows efficient and correct - reuse of factorizations if needed, - -Some of the algorithms also allow additional keyword arguments. See the documentation for -the specific algorithm for more information. - -### Interface Functions - - - `supports_trust_region(alg)`: whether or not the algorithm supports trust region - methods. Defaults to `false`. - - `supports_line_search(alg)`: whether or not the algorithm supports line search - methods. Defaults to `false`. - -See also [`NewtonDescent`](@ref), [`Dogleg`](@ref), [`SteepestDescent`](@ref), -[`DampedNewtonDescent`](@ref). -""" -abstract type AbstractDescentAlgorithm end - -supports_trust_region(::AbstractDescentAlgorithm) = false -supports_line_search(::AbstractDescentAlgorithm) = false - -get_linear_solver(alg::AbstractDescentAlgorithm) = __getproperty(alg, Val(:linsolve)) - -""" - AbstractDescentCache - -Abstract Type for all Descent Caches. - -### `__internal_solve!` specification - -```julia -descent_result = __internal_solve!( - cache::AbstractDescentCache, J, fu, u, idx::Val; skip_solve::Bool = false, kwargs...) -``` - - - `J`: Jacobian or Inverse Jacobian (if `pre_inverted = Val(true)`). - - `fu`: residual. - - `u`: current state. - - `idx`: index of the descent problem to solve and return. Defaults to `Val(1)`. - - `skip_solve`: Skip the direction computation and return the previous direction. - Defaults to `false`. This is useful for Trust Region Methods where the previous - direction was rejected and we want to try with a modified trust region. - - `kwargs`: keyword arguments to pass to the linear solver if there is one. - -#### Returned values - - - `descent_result`: Result in a [`DescentResult`](@ref). - -### Interface Functions - - - `get_du(cache)`: get the descent direction. - - `get_du(cache, ::Val{N})`: get the `N`th descent direction. - - `set_du!(cache, δu)`: set the descent direction. - - `set_du!(cache, δu, ::Val{N})`: set the `N`th descent direction. - - `last_step_accepted(cache)`: whether or not the last step was accepted. Checks if the - cache has a `last_step_accepted` field and returns it if it does, else returns `true`. -""" -abstract type AbstractDescentCache end - -SciMLBase.get_du(cache::AbstractDescentCache) = cache.δu -SciMLBase.get_du(cache::AbstractDescentCache, ::Val{1}) = get_du(cache) -SciMLBase.get_du(cache::AbstractDescentCache, ::Val{N}) where {N} = cache.δus[N - 1] -set_du!(cache::AbstractDescentCache, δu) = (cache.δu = δu) -set_du!(cache::AbstractDescentCache, δu, ::Val{1}) = set_du!(cache, δu) -set_du!(cache::AbstractDescentCache, δu, ::Val{N}) where {N} = (cache.δus[N - 1] = δu) - -function last_step_accepted(cache::AbstractDescentCache) - hasfield(typeof(cache), :last_step_accepted) && return cache.last_step_accepted - return true -end - -""" - AbstractNonlinearSolveLineSearchCache - -Abstract Type for all Line Search Caches used in NonlinearSolve.jl. - -### `__internal_solve!` specification - -```julia -__internal_solve!(cache::AbstractNonlinearSolveLineSearchCache, u, du; kwargs...) -``` - -Returns 2 values: - - - `unsuccessful`: If `true` it means that the Line Search Failed. - - `alpha`: The step size. -""" -abstract type AbstractNonlinearSolveLineSearchCache end - -function reinit_cache!( - cache::AbstractNonlinearSolveLineSearchCache, args...; p = cache.p, kwargs...) - cache.p = p -end - -""" - AbstractNonlinearSolveAlgorithm{name} <: AbstractNonlinearAlgorithm - -Abstract Type for all NonlinearSolve.jl Algorithms. `name` can be used to define custom -dispatches by wrapped solvers. - -### Interface Functions - - - `concrete_jac(alg)`: whether or not the algorithm uses a concrete Jacobian. Defaults - to `nothing`. - - `get_name(alg)`: get the name of the algorithm. -""" -abstract type AbstractNonlinearSolveAlgorithm{name} <: AbstractNonlinearAlgorithm end - -""" - concrete_jac(alg::AbstractNonlinearSolveAlgorithm) - -Whether the algorithm uses a concrete Jacobian. Defaults to `nothing` if it is unknown or -not applicable. Else a boolean value is returned. -""" -concrete_jac(::AbstractNonlinearSolveAlgorithm) = nothing - -function Base.show(io::IO, alg::AbstractNonlinearSolveAlgorithm) - __show_algorithm(io, alg, get_name(alg), 0) -end - -get_name(::AbstractNonlinearSolveAlgorithm{name}) where {name} = name - -""" - AbstractNonlinearSolveExtensionAlgorithm <: AbstractNonlinearSolveAlgorithm{:Extension} - -Abstract Type for all NonlinearSolve.jl Extension Algorithms, i.e. wrappers over 3rd party -solvers. -""" -abstract type AbstractNonlinearSolveExtensionAlgorithm <: - AbstractNonlinearSolveAlgorithm{:Extension} end - -""" - AbstractNonlinearSolveCache{iip, timeit} - -Abstract Type for all NonlinearSolve.jl Caches. - -### Interface Functions - - - `get_fu(cache)`: get the residual. - - `get_u(cache)`: get the current state. - - `set_fu!(cache, fu)`: set the residual. - - `set_u!(cache, u)`: set the current state. - - `reinit!(cache, u0; kwargs...)`: reinitialize the cache with the initial state `u0` and - any additional keyword arguments. - - `step!(cache; kwargs...)`: See [`SciMLBase.step!`](@ref) for more details. - - `not_terminated(cache)`: whether or not the solver has terminated. - - `isinplace(cache)`: whether or not the solver is inplace. -""" -abstract type AbstractNonlinearSolveCache{iip, timeit} end - -function SymbolicIndexingInterface.symbolic_container(cache::AbstractNonlinearSolveCache) - return cache.prob -end -function SymbolicIndexingInterface.parameter_values(cache::AbstractNonlinearSolveCache) - return parameter_values(symbolic_container(cache)) -end -function SymbolicIndexingInterface.state_values(cache::AbstractNonlinearSolveCache) - return state_values(symbolic_container(cache)) -end - -function Base.getproperty(cache::AbstractNonlinearSolveCache, sym::Symbol) - sym == :ps && return ParameterIndexingProxy(cache) - return getfield(cache, sym) -end - -function Base.getindex(cache::AbstractNonlinearSolveCache, sym) - return getu(cache, sym)(cache) -end - -function Base.setindex!(cache::AbstractNonlinearSolveCache, val, sym) - return setu(cache, sym)(cache, val) -end - -function Base.show(io::IO, cache::AbstractNonlinearSolveCache) - __show_cache(io, cache, 0) -end - -function __show_cache(io::IO, cache::AbstractNonlinearSolveCache, indent = 0) - println(io, "$(nameof(typeof(cache)))(") - __show_algorithm(io, cache.alg, - (" "^(indent + 4)) * "alg = " * string(get_name(cache.alg)), indent + 4) - - ustr = sprint(show, get_u(cache); context = (:compact => true, :limit => true)) - println(io, ",\n" * (" "^(indent + 4)) * "u = $(ustr),") - - residstr = sprint(show, get_fu(cache); context = (:compact => true, :limit => true)) - println(io, (" "^(indent + 4)) * "residual = $(residstr),") - - normstr = sprint( - show, norm(get_fu(cache), Inf); context = (:compact => true, :limit => true)) - println(io, (" "^(indent + 4)) * "inf-norm(residual) = $(normstr),") - - println(io, " "^(indent + 4) * "nsteps = ", cache.stats.nsteps, ",") - println(io, " "^(indent + 4) * "retcode = ", cache.retcode) - print(io, " "^(indent) * ")") -end - -SciMLBase.isinplace(::AbstractNonlinearSolveCache{iip}) where {iip} = iip - -get_fu(cache::AbstractNonlinearSolveCache) = cache.fu -get_u(cache::AbstractNonlinearSolveCache) = cache.u -set_fu!(cache::AbstractNonlinearSolveCache, fu) = (cache.fu = fu) -SciMLBase.set_u!(cache::AbstractNonlinearSolveCache, u) = (cache.u = u) - -function SciMLBase.reinit!(cache::AbstractNonlinearSolveCache; kwargs...) - return reinit_cache!(cache; kwargs...) -end -function SciMLBase.reinit!(cache::AbstractNonlinearSolveCache, u0; kwargs...) - return reinit_cache!(cache; u0, kwargs...) -end - -""" - AbstractLinearSolverCache <: Function - -Abstract Type for all Linear Solvers used in NonlinearSolve.jl. -""" -abstract type AbstractLinearSolverCache <: Function end - -""" - AbstractDampingFunction - -Abstract Type for Damping Functions in DampedNewton. - -### `__internal_init` specification - -```julia -__internal_init( - prob::AbstractNonlinearProblem, f::AbstractDampingFunction, initial_damping, - J, fu, u, args...; internal_norm = L2_NORM, kwargs...) --> AbstractDampingFunctionCache -``` - -Returns a [`AbstractDampingFunctionCache`](@ref). -""" -abstract type AbstractDampingFunction end - -""" - AbstractDampingFunctionCache - -Abstract Type for the Caches created by AbstractDampingFunctions - -### Interface Functions - - - `requires_normal_form_jacobian(f)`: whether or not the Jacobian is needed in normal - form. No default. - - `requires_normal_form_rhs(f)`: whether or not the residual is needed in normal form. - No default. - - `returns_norm_form_damping(f)`: whether or not the damping function returns the - damping factor in normal form. Defaults to `requires_normal_form_jacobian(f) || requires_normal_form_rhs(f)`. - - `(cache::AbstractDampingFunctionCache)(::Nothing)`: returns the damping factor. The type - of the damping factor returned from `solve!` is guaranteed to be the same as this. - -### `__internal_solve!` specification - -```julia -__internal_solve!(cache::AbstractDampingFunctionCache, J, fu, args...; kwargs...) -``` - -Returns the damping factor. -""" -abstract type AbstractDampingFunctionCache end - -function requires_normal_form_jacobian end -function requires_normal_form_rhs end -function returns_norm_form_damping(f::F) where {F} - return requires_normal_form_jacobian(f) || requires_normal_form_rhs(f) -end - -""" - AbstractNonlinearSolveOperator <: AbstractSciMLOperator - -NonlinearSolve.jl houses a few custom operators. These will eventually be moved out but till -then this serves as the abstract type for them. -""" -abstract type AbstractNonlinearSolveOperator{T} <: AbstractSciMLOperator{T} end - -# Approximate Jacobian Algorithms -""" - AbstractApproximateJacobianStructure - -Abstract Type for all Approximate Jacobian Structures used in NonlinearSolve.jl. - -### Interface Functions - - - `stores_full_jacobian(alg)`: whether or not the algorithm stores the full Jacobian. - Defaults to `false`. - - `get_full_jacobian(cache, alg, J)`: get the full Jacobian. Defaults to throwing an - error if `stores_full_jacobian(alg)` is `false`. -""" -abstract type AbstractApproximateJacobianStructure end - -stores_full_jacobian(::AbstractApproximateJacobianStructure) = false -function get_full_jacobian(cache, alg::AbstractApproximateJacobianStructure, J) - stores_full_jacobian(alg) && return J - error("This algorithm does not store the full Jacobian. Define `get_full_jacobian` for \ - this algorithm.") -end - -""" - AbstractJacobianInitialization - -Abstract Type for all Jacobian Initialization Algorithms used in NonlinearSolve.jl. - -### Interface Functions - - - `jacobian_initialized_preinverted(alg)`: whether or not the Jacobian is initialized - preinverted. Defaults to `false`. - -### `__internal_init` specification - -```julia -__internal_init( - prob::AbstractNonlinearProblem, alg::AbstractJacobianInitialization, solver, - f::F, fu, u, p; linsolve = missing, internalnorm::IN = L2_NORM, kwargs...) -``` - -Returns a [`NonlinearSolve.InitializedApproximateJacobianCache`](@ref). - -All subtypes need to define -`(cache::InitializedApproximateJacobianCache)(alg::NewSubType, fu, u)` which reinitializes -the Jacobian in `cache.J`. -""" -abstract type AbstractJacobianInitialization end - -function Base.show(io::IO, alg::AbstractJacobianInitialization) - modifiers = String[] - hasfield(typeof(alg), :structure) && - push!(modifiers, "structure = $(nameof(typeof(alg.structure)))()") - print(io, "$(nameof(typeof(alg)))($(join(modifiers, ", ")))") - return nothing -end - -jacobian_initialized_preinverted(::AbstractJacobianInitialization) = false - -""" - AbstractApproximateJacobianUpdateRule{INV} - -Abstract Type for all Approximate Jacobian Update Rules used in NonlinearSolve.jl. - -### Interface Functions - - - `store_inverse_jacobian(alg)`: Return `INV` - -### `__internal_init` specification - -```julia -__internal_init( - prob::AbstractNonlinearProblem, alg::AbstractApproximateJacobianUpdateRule, J, fu, u, - du, args...; internalnorm::F = L2_NORM, kwargs...) where {F} --> -AbstractApproximateJacobianUpdateRuleCache{INV} -``` -""" -abstract type AbstractApproximateJacobianUpdateRule{INV} end - -store_inverse_jacobian(::AbstractApproximateJacobianUpdateRule{INV}) where {INV} = INV - -""" - AbstractApproximateJacobianUpdateRuleCache{INV} - -Abstract Type for all Approximate Jacobian Update Rule Caches used in NonlinearSolve.jl. - -### Interface Functions - - - `store_inverse_jacobian(alg)`: Return `INV` - -### `__internal_solve!` specification - -```julia -__internal_solve!( - cache::AbstractApproximateJacobianUpdateRuleCache, J, fu, u, du; kwargs...) --> J / J⁻¹ -``` -""" -abstract type AbstractApproximateJacobianUpdateRuleCache{INV} end - -store_inverse_jacobian(::AbstractApproximateJacobianUpdateRuleCache{INV}) where {INV} = INV - -""" - AbstractResetCondition - -Condition for resetting the Jacobian in Quasi-Newton's methods. - -### `__internal_init` specification - -```julia -__internal_init(alg::AbstractResetCondition, J, fu, u, du, args...; kwargs...) --> -ResetCache -``` - -### `__internal_solve!` specification - -```julia -__internal_solve!(cache::ResetCache, J, fu, u, du) --> Bool -``` -""" -abstract type AbstractResetCondition end - -""" - AbstractTrustRegionMethod - -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, p, args...; - internalnorm::IF = L2_NORM, kwargs...) where {F, IF} --> AbstractTrustRegionMethodCache -``` -""" -abstract type AbstractTrustRegionMethod end - -""" - AbstractTrustRegionMethodCache - -Abstract Type for all Trust Region Method Caches used in NonlinearSolve.jl. - -### Interface Functions - - - `last_step_accepted(cache)`: whether or not the last step was accepted. Defaults to - `cache.last_step_accepted`. Should if overloaded if the field is not present. - -### `__internal_solve!` specification - -```julia -__internal_solve!(cache::AbstractTrustRegionMethodCache, J, fu, u, δu, descent_stats) -``` - -Returns `last_step_accepted`, updated `u_cache` and `fu_cache`. If the last step was -accepted then these values should be copied into the toplevel cache. -""" -abstract type AbstractTrustRegionMethodCache end - -last_step_accepted(cache::AbstractTrustRegionMethodCache) = cache.last_step_accepted - -""" - AbstractNonlinearSolveJacobianCache{iip} <: Function - -Abstract Type for all Jacobian Caches used in NonlinearSolve.jl. -""" -abstract type AbstractNonlinearSolveJacobianCache{iip} <: Function end - -SciMLBase.isinplace(::AbstractNonlinearSolveJacobianCache{iip}) where {iip} = iip - -""" - AbstractNonlinearSolveTraceLevel - -### Common Arguments - - - `freq`: Sets both `print_frequency` and `store_frequency` to `freq`. - -### Common Keyword Arguments - - - `print_frequency`: Print the trace every `print_frequency` iterations if - `show_trace == Val(true)`. - - `store_frequency`: Store the trace every `store_frequency` iterations if - `store_trace == Val(true)`. -""" -abstract type AbstractNonlinearSolveTraceLevel end - -# Default Printing -for aType in (AbstractTrustRegionMethod, AbstractResetCondition, - AbstractApproximateJacobianUpdateRule, AbstractDampingFunction, - AbstractNonlinearSolveExtensionAlgorithm) - @eval function Base.show(io::IO, alg::$(aType)) - print(io, "$(nameof(typeof(alg)))()") - end -end diff --git a/src/algorithms/broyden.jl b/src/algorithms/broyden.jl deleted file mode 100644 index 39962bc6a..000000000 --- a/src/algorithms/broyden.jl +++ /dev/null @@ -1,227 +0,0 @@ -""" - Broyden(; max_resets::Int = 100, linesearch = nothing, reset_tolerance = nothing, - init_jacobian::Val = Val(:identity), autodiff = nothing, alpha = nothing) - -An implementation of `Broyden`'s Method [broyden1965class](@cite) with resetting and line -search. - -### Keyword Arguments - - - `max_resets`: the maximum number of resets to perform. Defaults to `100`. - - - `reset_tolerance`: the tolerance for the reset check. Defaults to - `sqrt(eps(real(eltype(u))))`. - - `alpha`: If `init_jacobian` is set to `Val(:identity)`, then the initial Jacobian - inverse is set to be `(αI)⁻¹`. Defaults to `nothing` which implies - `α = max(norm(u), 1) / (2 * norm(fu))`. - - `init_jacobian`: the method to use for initializing the jacobian. Defaults to - `Val(:identity)`. Choices include: - - + `Val(:identity)`: Identity Matrix. - + `Val(:true_jacobian)`: True Jacobian. This is a good choice for differentiable - problems. - - `update_rule`: Update Rule for the Jacobian. Choices are: - - + `Val(:good_broyden)`: Good Broyden's Update Rule - + `Val(:bad_broyden)`: Bad Broyden's Update Rule - + `Val(:diagonal)`: Only update the diagonal of the Jacobian. This algorithm may be - useful for specific problems, but whether it will work may depend strongly on the - problem -""" -function Broyden(; - max_resets = 100, linesearch = nothing, reset_tolerance = nothing, - init_jacobian = Val(:identity), autodiff = nothing, alpha = nothing, - update_rule = Val(:good_broyden)) - initialization = broyden_init(init_jacobian, update_rule, autodiff, alpha) - update_rule = broyden_update_rule(update_rule) - return ApproximateJacobianSolveAlgorithm{ - init_jacobian isa Val{:true_jacobian}, :Broyden}(; - linesearch, descent = NewtonDescent(), update_rule, max_resets, initialization, - reinit_rule = NoChangeInStateReset(; reset_tolerance)) -end - -function broyden_init(::Val{:identity}, ::Val{:diagonal}, autodiff, alpha) - return IdentityInitialization(alpha, DiagonalStructure()) -end -function broyden_init(::Val{:identity}, ::Val, autodiff, alpha) - IdentityInitialization(alpha, FullStructure()) -end -function broyden_init(::Val{:true_jacobian}, ::Val, autodiff, alpha) - return TrueJacobianInitialization(FullStructure(), autodiff) -end -function broyden_init(::Val{IJ}, ::Val{UR}, autodiff, alpha) where {IJ, UR} - error("Unknown combination of `init_jacobian = Val($(Meta.quot(IJ)))` and \ - `update_rule = Val($(Meta.quot(UR)))`. Please choose a valid combination.") -end - -broyden_update_rule(::Val{:good_broyden}) = GoodBroydenUpdateRule() -broyden_update_rule(::Val{:bad_broyden}) = BadBroydenUpdateRule() -broyden_update_rule(::Val{:diagonal}) = GoodBroydenUpdateRule() -function broyden_update_rule(::Val{UR}) where {UR} - error("Unknown update rule `update_rule = Val($(Meta.quot(UR)))`. Please choose a \ - valid update rule.") -end - -# Checks for no significant change for `nsteps` -""" - NoChangeInStateReset(; nsteps::Int = 3, reset_tolerance = nothing, - check_du::Bool = true, check_dfu::Bool = true) - -Recommends a reset if the state or the function value has not changed significantly in -`nsteps` steps. This is used in [`Broyden`](@ref). - -### Keyword Arguments - - - `nsteps`: the number of steps to check for no change. Defaults to `3`. - - `reset_tolerance`: the tolerance for the reset check. Defaults to - `sqrt(eps(real(eltype(u))))`. - - `check_du`: whether to check the state. Defaults to `true`. - - `check_dfu`: whether to check the function value. Defaults to `true`. -""" -@kwdef @concrete struct NoChangeInStateReset <: AbstractResetCondition - nsteps::Int = 3 - reset_tolerance = nothing - check_du::Bool = true - check_dfu::Bool = true -end - -@concrete mutable struct NoChangeInStateResetCache - dfu - reset_tolerance - check_du - check_dfu - nsteps::Int - steps_since_change_du::Int - steps_since_change_dfu::Int -end - -function reinit_cache!(cache::NoChangeInStateResetCache, args...; kwargs...) - cache.steps_since_change_du = 0 - cache.steps_since_change_dfu = 0 -end - -function __internal_init(alg::NoChangeInStateReset, J, fu, u, du, args...; kwargs...) - if alg.check_dfu - @bb dfu = copy(fu) - else - dfu = fu - end - T = real(eltype(u)) - tol = alg.reset_tolerance === nothing ? eps(T)^(3 // 4) : T(alg.reset_tolerance) - return NoChangeInStateResetCache( - dfu, tol, alg.check_du, alg.check_dfu, alg.nsteps, 0, 0) -end - -function __internal_solve!(cache::NoChangeInStateResetCache, J, fu, u, du) - reset_tolerance = cache.reset_tolerance - if cache.check_du - if any(@closure(x->abs(x) ≤ reset_tolerance), du) - cache.steps_since_change_du += 1 - if cache.steps_since_change_du ≥ cache.nsteps - cache.steps_since_change_du = 0 - cache.steps_since_change_dfu = 0 - return true - end - else - cache.steps_since_change_du = 0 - cache.steps_since_change_dfu = 0 - end - end - if cache.check_dfu - @bb @. cache.dfu = fu - cache.dfu - if any(@closure(x->abs(x) ≤ reset_tolerance), cache.dfu) - cache.steps_since_change_dfu += 1 - if cache.steps_since_change_dfu ≥ cache.nsteps - cache.steps_since_change_dfu = 0 - cache.steps_since_change_du = 0 - @bb copyto!(cache.dfu, fu) - return true - end - else - cache.steps_since_change_dfu = 0 - cache.steps_since_change_du = 0 - end - @bb copyto!(cache.dfu, fu) - end - return false -end - -# Broyden Update Rules -""" - BadBroydenUpdateRule() - -Broyden Update Rule corresponding to "bad broyden's method" [broyden1965class](@cite). -""" -@concrete struct BadBroydenUpdateRule <: AbstractApproximateJacobianUpdateRule{true} end - -""" - GoodBroydenUpdateRule() - -Broyden Update Rule corresponding to "good broyden's method" [broyden1965class](@cite). -""" -@concrete struct GoodBroydenUpdateRule <: AbstractApproximateJacobianUpdateRule{true} end - -@concrete mutable struct BroydenUpdateRuleCache{mode} <: - AbstractApproximateJacobianUpdateRuleCache{true} - J⁻¹dfu - dfu - u_cache - du_cache - internalnorm -end - -function __internal_init(prob::AbstractNonlinearProblem, - alg::Union{GoodBroydenUpdateRule, BadBroydenUpdateRule}, J, fu, u, - du, args...; internalnorm::F = L2_NORM, kwargs...) where {F} - @bb J⁻¹dfu = similar(u) - @bb dfu = copy(fu) - if alg isa GoodBroydenUpdateRule || J isa Diagonal - @bb u_cache = similar(u) - else - u_cache = nothing - end - if J isa Diagonal - du_cache = nothing - else - @bb du_cache = similar(du) - end - mode = alg isa GoodBroydenUpdateRule ? :good : :bad - return BroydenUpdateRuleCache{mode}(J⁻¹dfu, dfu, u_cache, du_cache, internalnorm) -end - -function __internal_solve!(cache::BroydenUpdateRuleCache{mode}, J⁻¹, fu, u, du) where {mode} - T = eltype(u) - @bb @. cache.dfu = fu - cache.dfu - @bb cache.J⁻¹dfu = J⁻¹ × vec(cache.dfu) - if mode === :good - @bb cache.u_cache = transpose(J⁻¹) × vec(du) - denom = dot(du, cache.J⁻¹dfu) - rmul = transpose(_vec(cache.u_cache)) - else - denom = cache.internalnorm(cache.dfu)^2 - rmul = transpose(_vec(cache.dfu)) - end - @bb @. cache.du_cache = (du - cache.J⁻¹dfu) / ifelse(iszero(denom), T(1e-5), denom) - @bb J⁻¹ += vec(cache.du_cache) × rmul - @bb copyto!(cache.dfu, fu) - return J⁻¹ -end - -function __internal_solve!( - cache::BroydenUpdateRuleCache{mode}, J⁻¹::Diagonal, fu, u, du) where {mode} - T = eltype(u) - @bb @. cache.dfu = fu - cache.dfu - J⁻¹_diag = _restructure(cache.dfu, diag(J⁻¹)) - if mode === :good - @bb @. cache.J⁻¹dfu = J⁻¹_diag * cache.dfu * du - denom = sum(cache.J⁻¹dfu) - @bb @. J⁻¹_diag += (du - J⁻¹_diag * cache.dfu) * du * J⁻¹_diag / - ifelse(iszero(denom), T(1e-5), denom) - else - denom = cache.internalnorm(cache.dfu)^2 - @bb @. J⁻¹_diag += (du - J⁻¹_diag * cache.dfu) * cache.dfu / - ifelse(iszero(denom), T(1e-5), denom) - end - @bb copyto!(cache.dfu, fu) - return Diagonal(J⁻¹_diag) -end diff --git a/src/algorithms/dfsane.jl b/src/algorithms/dfsane.jl deleted file mode 100644 index 1ece5f5da..000000000 --- a/src/algorithms/dfsane.jl +++ /dev/null @@ -1,27 +0,0 @@ -# XXX: remove kwargs with unicode -""" - DFSane(; σ_min = 1 // 10^10, σ_max = 1e10, σ_1 = 1, M::Int = 10, γ = 1 // 10^4, - τ_min = 1 // 10, τ_max = 1 // 2, n_exp::Int = 2, max_inner_iterations::Int = 100, - η_strategy = (fn_1, n, x_n, f_n) -> fn_1 / n^2) - -A low-overhead and allocation-free implementation of the df-sane method for solving -large-scale nonlinear systems of equations. For in depth information about all the -parameters and the algorithm, see [la2006spectral](@citet). - -### Keyword Arguments - - - `σ_min`: the minimum value of the spectral coefficient `σₙ` which is related to the step - size in the algorithm. Defaults to `1e-10`. - - `σ_max`: the maximum value of the spectral coefficient `σₙ` which is related to the step - size in the algorithm. Defaults to `1e10`. - -For other keyword arguments, see [`RobustNonMonotoneLineSearch`](@ref). -""" -function DFSane(; σ_min = 1 // 10^10, σ_max = 1e10, σ_1 = 1, M::Int = 10, γ = 1 // 10^4, - τ_min = 1 // 10, τ_max = 1 // 2, n_exp::Int = 2, max_inner_iterations::Int = 100, - η_strategy::ETA = (fn_1, n, x_n, f_n) -> fn_1 / n^2) where {ETA} - linesearch = RobustNonMonotoneLineSearch(; - gamma = γ, sigma_1 = σ_1, M, tau_min = τ_min, tau_max = τ_max, - n_exp, η_strategy, maxiters = max_inner_iterations) - return GeneralizedDFSane{:DFSane}(linesearch, σ_min, σ_max, nothing) -end diff --git a/src/algorithms/gauss_newton.jl b/src/algorithms/gauss_newton.jl deleted file mode 100644 index 0f4ec5cd9..000000000 --- a/src/algorithms/gauss_newton.jl +++ /dev/null @@ -1,15 +0,0 @@ -""" - GaussNewton(; concrete_jac = nothing, linsolve = nothing, precs = DEFAULT_PRECS, - linesearch = nothing, vjp_autodiff = nothing, autodiff = nothing, - jvp_autodiff = nothing) - -An advanced GaussNewton implementation with support for efficient handling of sparse -matrices via colored automatic differentiation and preconditioned linear solvers. Designed -for large-scale and numerically-difficult nonlinear least squares problems. -""" -function GaussNewton(; concrete_jac = nothing, linsolve = nothing, precs = DEFAULT_PRECS, - linesearch = nothing, vjp_autodiff = nothing, autodiff = nothing, - jvp_autodiff = nothing) - return GeneralizedFirstOrderAlgorithm{concrete_jac, :GaussNewton}(; linesearch, - descent = NewtonDescent(; linsolve, precs), autodiff, vjp_autodiff, jvp_autodiff) -end diff --git a/src/algorithms/lbroyden.jl b/src/algorithms/lbroyden.jl deleted file mode 100644 index 89df6ece5..000000000 --- a/src/algorithms/lbroyden.jl +++ /dev/null @@ -1,169 +0,0 @@ -""" - LimitedMemoryBroyden(; max_resets::Int = 3, linesearch = nothing, - threshold::Val = Val(10), reset_tolerance = nothing, alpha = nothing) - -An implementation of `LimitedMemoryBroyden` [ziani2008autoadaptative](@cite) with resetting -and line search. - -### Keyword Arguments - - - `max_resets`: the maximum number of resets to perform. Defaults to `3`. - - `reset_tolerance`: the tolerance for the reset check. Defaults to - `sqrt(eps(real(eltype(u))))`. - - `threshold`: the number of vectors to store in the low rank approximation. Defaults - to `Val(10)`. - - `alpha`: The initial Jacobian inverse is set to be `(αI)⁻¹`. Defaults to `nothing` - which implies `α = max(norm(u), 1) / (2 * norm(fu))`. -""" -function LimitedMemoryBroyden(; max_resets::Int = 3, linesearch = nothing, - threshold::Union{Val, Int} = Val(10), reset_tolerance = nothing, alpha = nothing) - threshold isa Int && (threshold = Val(threshold)) - initialization = BroydenLowRankInitialization{_unwrap_val(threshold)}(alpha, threshold) - return ApproximateJacobianSolveAlgorithm{false, :LimitedMemoryBroyden}(; linesearch, - descent = NewtonDescent(), update_rule = GoodBroydenUpdateRule(), - max_resets, initialization, reinit_rule = NoChangeInStateReset(; reset_tolerance)) -end - -""" - BroydenLowRankInitialization{T}(alpha, threshold::Val{T}) - -An initialization for `LimitedMemoryBroyden` that uses a low rank approximation of the -Jacobian. The low rank updates to the Jacobian matrix corresponds to what SciPy calls -["simple"](https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.broyden2.html#scipy-optimize-broyden2). -""" -@concrete struct BroydenLowRankInitialization{T} <: AbstractJacobianInitialization - alpha - threshold::Val{T} -end - -jacobian_initialized_preinverted(::BroydenLowRankInitialization) = true - -function __internal_init( - prob::AbstractNonlinearProblem, alg::BroydenLowRankInitialization{T}, - solver, f::F, fu, u, p; maxiters = 1000, - internalnorm::IN = L2_NORM, kwargs...) where {T, F, IN} - if u isa Number # Use the standard broyden - return __internal_init(prob, IdentityInitialization(true, FullStructure()), - solver, f, fu, u, p; maxiters, kwargs...) - end - # Pay to cost of slightly more allocations to prevent type-instability for StaticArrays - α = inv(__initial_alpha(alg.alpha, u, fu, internalnorm)) - if u isa StaticArray - J = BroydenLowRankJacobian(fu, u; alg.threshold, alpha = α) - else - threshold = min(_unwrap_val(alg.threshold), maxiters) - J = BroydenLowRankJacobian(fu, u; threshold, alpha = α) - end - return InitializedApproximateJacobianCache( - J, FullStructure(), alg, nothing, true, internalnorm) -end - -function (cache::InitializedApproximateJacobianCache)( - alg::BroydenLowRankInitialization, fu, u) - α = __initial_alpha(alg.alpha, u, fu, cache.internalnorm) - cache.J.idx = 0 - cache.J.alpha = inv(α) - return -end - -""" - BroydenLowRankJacobian{T}(U, Vᵀ, idx, cache, alpha) - -Low Rank Approximation of the Jacobian Matrix. Currently only used for -[`LimitedMemoryBroyden`](@ref). This computes the Jacobian as ``U \\times V^T``. -""" -@concrete mutable struct BroydenLowRankJacobian{T} <: AbstractNonlinearSolveOperator{T} - U - Vᵀ - idx::Int - cache - alpha -end - -__safe_inv!!(workspace, op::BroydenLowRankJacobian) = op # Already Inverted form - -@inline function __get_components(op::BroydenLowRankJacobian) - op.idx ≥ size(op.U, 2) && return op.cache, op.U, transpose(op.Vᵀ) - _cache = op.cache === nothing ? op.cache : view(op.cache, 1:(op.idx)) - return (_cache, view(op.U, :, 1:(op.idx)), transpose(view(op.Vᵀ, :, 1:(op.idx)))) -end - -Base.size(op::BroydenLowRankJacobian) = size(op.U, 1), size(op.Vᵀ, 1) -function Base.size(op::BroydenLowRankJacobian, d::Integer) - return ifelse(d == 1, size(op.U, 1), size(op.Vᵀ, 1)) -end - -for op in (:adjoint, :transpose) - # FIXME: adjoint might be a problem here. Fix if a complex number issue shows up - @eval function Base.$(op)(operator::BroydenLowRankJacobian{T}) where {T} - return BroydenLowRankJacobian{T}( - operator.Vᵀ, operator.U, operator.idx, operator.cache, operator.alpha) - end -end - -# Storing the transpose to ensure contiguous memory on splicing -function BroydenLowRankJacobian( - fu::StaticArray{S2, T2}, u::StaticArray{S1, T1}; alpha = true, - threshold::Val{Th} = Val(10)) where {S1, S2, T1, T2, Th} - T = promote_type(T1, T2) - fuSize, uSize = Size(fu), Size(u) - U = MArray{Tuple{prod(fuSize), Th}, T}(undef) - Vᵀ = MArray{Tuple{prod(uSize), Th}, T}(undef) - return BroydenLowRankJacobian{T}(U, Vᵀ, 0, nothing, T(alpha)) -end - -function BroydenLowRankJacobian(fu, u; threshold::Int = 10, alpha = true) - T = promote_type(eltype(u), eltype(fu)) - U = __similar(fu, T, length(fu), threshold) - Vᵀ = __similar(u, T, length(u), threshold) - cache = __similar(u, T, threshold) - return BroydenLowRankJacobian{T}(U, Vᵀ, 0, cache, T(alpha)) -end - -function Base.:*(J::BroydenLowRankJacobian, x::AbstractVector) - J.idx == 0 && return -x - cache, U, Vᵀ = __get_components(J) - return U * (Vᵀ * x) .- J.alpha .* x -end - -function LinearAlgebra.mul!(y::AbstractVector, J::BroydenLowRankJacobian, x::AbstractVector) - if J.idx == 0 - @. y = -J.alpha * x - return y - end - cache, U, Vᵀ = __get_components(J) - @bb cache = Vᵀ × x - mul!(y, U, cache) - @bb @. y -= J.alpha * x - return y -end - -function Base.:*(x::AbstractVector, J::BroydenLowRankJacobian) - J.idx == 0 && return -x - cache, U, Vᵀ = __get_components(J) - return Vᵀ' * (U' * x) .- J.alpha .* x -end - -function LinearAlgebra.mul!(y::AbstractVector, x::AbstractVector, J::BroydenLowRankJacobian) - if J.idx == 0 - @. y = -J.alpha * x - return y - end - cache, U, Vᵀ = __get_components(J) - @bb cache = transpose(U) × x - mul!(y, transpose(Vᵀ), cache) - @bb @. y -= J.alpha * x - return y -end - -function LinearAlgebra.mul!(J::BroydenLowRankJacobian, u::AbstractArray, - vᵀ::LinearAlgebra.AdjOrTransAbsVec, α::Bool, β::Bool) - @assert α & β - idx_update = mod1(J.idx + 1, size(J.U, 2)) - copyto!(@view(J.U[:, idx_update]), _vec(u)) - copyto!(@view(J.Vᵀ[:, idx_update]), _vec(vᵀ)) - J.idx += 1 - return J -end - -ArrayInterface.restructure(::BroydenLowRankJacobian, J::BroydenLowRankJacobian) = J diff --git a/src/algorithms/levenberg_marquardt.jl b/src/algorithms/levenberg_marquardt.jl deleted file mode 100644 index 01896bd14..000000000 --- a/src/algorithms/levenberg_marquardt.jl +++ /dev/null @@ -1,176 +0,0 @@ -""" - LevenbergMarquardt(; linsolve = nothing, - precs = DEFAULT_PRECS, damping_initial::Real = 1.0, α_geodesic::Real = 0.75, - damping_increase_factor::Real = 2.0, damping_decrease_factor::Real = 3.0, - finite_diff_step_geodesic = 0.1, b_uphill::Real = 1.0, autodiff = nothing, - min_damping_D::Real = 1e-8, disable_geodesic = Val(false), vjp_autodiff = nothing, - jvp_autodiff = nothing) - -An advanced Levenberg-Marquardt implementation with the improvements suggested in -[transtrum2012improvements](@citet). Designed for large-scale and numerically-difficult -nonlinear systems. - -### Keyword Arguments - - - `damping_initial`: the starting value for the damping factor. The damping factor is - inversely proportional to the step size. The damping factor is adjusted during each - iteration. Defaults to `1.0`. See Section 2.1 of [transtrum2012improvements](@citet). - - `damping_increase_factor`: the factor by which the damping is increased if a step is - rejected. Defaults to `2.0`. See Section 2.1 of [transtrum2012improvements](@citet). - - `damping_decrease_factor`: the factor by which the damping is decreased if a step is - accepted. Defaults to `3.0`. See Section 2.1 of [transtrum2012improvements](@citet). - - `min_damping_D`: the minimum value of the damping terms in the diagonal damping matrix - `DᵀD`, where `DᵀD` is given by the largest diagonal entries of `JᵀJ` yet encountered, - where `J` is the Jacobian. It is suggested by [transtrum2012improvements](@citet) to use - a minimum value of the elements in `DᵀD` to prevent the damping from being too small. - Defaults to `1e-8`. - - `disable_geodesic`: Disables Geodesic Acceleration if set to `Val(true)`. It provides - a way to trade-off robustness for speed, though in most situations Geodesic Acceleration - should not be disabled. - -For the remaining arguments, see [`GeodesicAcceleration`](@ref) and -[`NonlinearSolve.LevenbergMarquardtTrustRegion`](@ref) documentations. -""" -function LevenbergMarquardt(; - linsolve = nothing, precs = DEFAULT_PRECS, damping_initial::Real = 1.0, - α_geodesic::Real = 0.75, damping_increase_factor::Real = 2.0, - damping_decrease_factor::Real = 3.0, finite_diff_step_geodesic = 0.1, - b_uphill::Real = 1.0, min_damping_D::Real = 1e-8, disable_geodesic = False, - autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing) - descent = DampedNewtonDescent(; linsolve, precs, initial_damping = damping_initial, - damping_fn = LevenbergMarquardtDampingFunction( - damping_increase_factor, damping_decrease_factor, min_damping_D)) - if disable_geodesic === False - descent = GeodesicAcceleration(descent, finite_diff_step_geodesic, α_geodesic) - end - trustregion = LevenbergMarquardtTrustRegion(b_uphill) - return GeneralizedFirstOrderAlgorithm{true, :LevenbergMarquardt}(; - trustregion, descent, autodiff, vjp_autodiff, jvp_autodiff) -end - -@concrete struct LevenbergMarquardtDampingFunction <: AbstractDampingFunction - increase_factor - decrease_factor - min_damping -end - -@concrete mutable struct LevenbergMarquardtDampingCache <: AbstractDampingFunctionCache - increase_factor - decrease_factor - min_damping - λ_factor - λ - DᵀD - J_diag_cache - J_damped - damping_f - initial_damping -end - -function reinit_cache!(cache::LevenbergMarquardtDampingCache, args...; kwargs...) - cache.λ = cache.initial_damping - cache.λ_factor = cache.damping_f.increase_factor - if !(cache.DᵀD isa Number) - if can_setindex(cache.DᵀD.diag) - cache.DᵀD.diag .= cache.min_damping - else - cache.DᵀD = Diagonal(ones(typeof(cache.DᵀD.diag)) * cache.min_damping) - end - end - cache.J_damped = cache.λ .* cache.DᵀD -end - -function requires_normal_form_jacobian(::Union{ - LevenbergMarquardtDampingFunction, LevenbergMarquardtDampingCache}) - return false -end -function requires_normal_form_rhs(::Union{ - LevenbergMarquardtDampingFunction, LevenbergMarquardtDampingCache}) - return false -end -function returns_norm_form_damping(::Union{ - LevenbergMarquardtDampingFunction, LevenbergMarquardtDampingCache}) - return true -end - -function __internal_init( - prob::AbstractNonlinearProblem, f::LevenbergMarquardtDampingFunction, - initial_damping, J, fu, u, ::Val{NF}; - internalnorm::F = L2_NORM, kwargs...) where {F, NF} - T = promote_type(eltype(u), eltype(fu)) - DᵀD = __init_diagonal(u, T(f.min_damping)) - if NF - J_diag_cache = nothing - else - @bb J_diag_cache = similar(u) - end - J_damped = T(initial_damping) .* DᵀD - return LevenbergMarquardtDampingCache( - T(f.increase_factor), T(f.decrease_factor), T(f.min_damping), T(f.increase_factor), - T(initial_damping), DᵀD, J_diag_cache, J_damped, f, T(initial_damping)) -end - -(damping::LevenbergMarquardtDampingCache)(::Nothing) = damping.J_damped - -function __internal_solve!( - damping::LevenbergMarquardtDampingCache, J, fu, ::Val{false}; kwargs...) - if __can_setindex(damping.J_diag_cache) - sum!(abs2, _vec(damping.J_diag_cache), J') - elseif damping.J_diag_cache isa Number - damping.J_diag_cache = abs2(J) - else - damping.J_diag_cache = dropdims(sum(abs2, J'; dims = 1); dims = 1) - end - damping.DᵀD = __update_LM_diagonal!!(damping.DᵀD, _vec(damping.J_diag_cache)) - @bb @. damping.J_damped = damping.λ * damping.DᵀD - return damping.J_damped -end - -function __internal_solve!( - damping::LevenbergMarquardtDampingCache, JᵀJ, fu, ::Val{true}; kwargs...) - damping.DᵀD = __update_LM_diagonal!!(damping.DᵀD, JᵀJ) - @bb @. damping.J_damped = damping.λ * damping.DᵀD - return damping.J_damped -end - -function callback_into_cache!(topcache, cache::LevenbergMarquardtDampingCache, args...) - if last_step_accepted(topcache.trustregion_cache) && - last_step_accepted(topcache.descent_cache) - cache.λ_factor = 1 / cache.decrease_factor - end - cache.λ *= cache.λ_factor - cache.λ_factor = cache.increase_factor -end - -@inline __update_LM_diagonal!!(y::Number, x::Number) = max(y, x) -@inline function __update_LM_diagonal!!(y::Diagonal, x::AbstractVector) - if __can_setindex(y.diag) - @. y.diag = max(y.diag, x) - return y - else - return Diagonal(max.(y.diag, x)) - end -end -@inline function __update_LM_diagonal!!(y::Diagonal, x::AbstractMatrix) - if __can_setindex(y.diag) - if fast_scalar_indexing(y.diag) - @simd for i in axes(x, 1) - @inbounds y.diag[i] = max(y.diag[i], x[i, i]) - end - return y - else - y .= max.(y.diag, @view(x[diagind(x)])) - return y - end - else - return Diagonal(max.(y.diag, @view(x[diagind(x)]))) - end -end - -@inline __init_diagonal(u::Number, v) = oftype(u, v) -@inline __init_diagonal(u::SArray, v) = Diagonal(ones(typeof(vec(u))) * v) -@inline function __init_diagonal(u, v) - d = similar(vec(u)) - d .= v - return Diagonal(d) -end diff --git a/src/algorithms/raphson.jl b/src/algorithms/raphson.jl deleted file mode 100644 index e5dee4c91..000000000 --- a/src/algorithms/raphson.jl +++ /dev/null @@ -1,15 +0,0 @@ -""" - NewtonRaphson(; concrete_jac = nothing, linsolve = nothing, linesearch = missing, - precs = DEFAULT_PRECS, autodiff = nothing, vjp_autodiff = nothing, - jvp_autodiff = nothing) - -An advanced NewtonRaphson implementation with support for efficient handling of sparse -matrices via colored automatic differentiation and preconditioned linear solvers. Designed -for large-scale and numerically-difficult nonlinear systems. -""" -function NewtonRaphson(; concrete_jac = nothing, linsolve = nothing, linesearch = nothing, - precs = DEFAULT_PRECS, autodiff = nothing, vjp_autodiff = nothing, - jvp_autodiff = nothing) - return GeneralizedFirstOrderAlgorithm{concrete_jac, :NewtonRaphson}(; linesearch, - descent = NewtonDescent(; linsolve, precs), autodiff, vjp_autodiff, jvp_autodiff) -end diff --git a/src/algorithms/trust_region.jl b/src/algorithms/trust_region.jl deleted file mode 100644 index dab6843e8..000000000 --- a/src/algorithms/trust_region.jl +++ /dev/null @@ -1,36 +0,0 @@ -""" - TrustRegion(; concrete_jac = nothing, linsolve = nothing, precs = DEFAULT_PRECS, - radius_update_scheme = RadiusUpdateSchemes.Simple, max_trust_radius::Real = 0 // 1, - initial_trust_radius::Real = 0 // 1, step_threshold::Real = 1 // 10000, - shrink_threshold::Real = 1 // 4, expand_threshold::Real = 3 // 4, - shrink_factor::Real = 1 // 4, expand_factor::Real = 2 // 1, - max_shrink_times::Int = 32, - vjp_autodiff = nothing, autodiff = nothing, jvp_autodiff = nothing) - -An advanced TrustRegion implementation with support for efficient handling of sparse -matrices via colored automatic differentiation and preconditioned linear solvers. Designed -for large-scale and numerically-difficult nonlinear systems. - -### Keyword Arguments - - - `radius_update_scheme`: the scheme used to update the trust region radius. Defaults to - `RadiusUpdateSchemes.Simple`. See [`RadiusUpdateSchemes`](@ref) for more details. For a - review on trust region radius update schemes, see [yuan2015recent](@citet). - -For the remaining arguments, see [`NonlinearSolve.GenericTrustRegionScheme`](@ref) -documentation. -""" -function TrustRegion(; concrete_jac = nothing, linsolve = nothing, precs = DEFAULT_PRECS, - radius_update_scheme = RadiusUpdateSchemes.Simple, max_trust_radius::Real = 0 // 1, - initial_trust_radius::Real = 0 // 1, step_threshold::Real = 1 // 10000, - shrink_threshold::Real = 1 // 4, expand_threshold::Real = 3 // 4, - shrink_factor::Real = 1 // 4, expand_factor::Real = 2 // 1, - max_shrink_times::Int = 32, - autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing) - descent = Dogleg(; linsolve, precs) - trustregion = GenericTrustRegionScheme(; - method = radius_update_scheme, step_threshold, shrink_threshold, expand_threshold, - shrink_factor, expand_factor, initial_trust_radius, max_trust_radius) - return GeneralizedFirstOrderAlgorithm{concrete_jac, :TrustRegion}(; - trustregion, descent, autodiff, vjp_autodiff, jvp_autodiff, max_shrink_times) -end diff --git a/src/core/approximate_jacobian.jl b/src/core/approximate_jacobian.jl deleted file mode 100644 index b8f479a6d..000000000 --- a/src/core/approximate_jacobian.jl +++ /dev/null @@ -1,377 +0,0 @@ -""" - ApproximateJacobianSolveAlgorithm{concrete_jac, name}(; linesearch = missing, - trustregion = missing, descent, update_rule, reinit_rule, initialization, - max_resets::Int = typemax(Int), max_shrink_times::Int = typemax(Int)) - ApproximateJacobianSolveAlgorithm(; concrete_jac = nothing, - name::Symbol = :unknown, kwargs...) - -Nonlinear Solve Algorithms using an Iterative Approximation of the Jacobian. Most common -examples include [`Broyden`](@ref)'s Method. - -### Keyword Arguments - - - `trustregion`: Globalization using a Trust Region Method. This needs to follow the - [`NonlinearSolve.AbstractTrustRegionMethod`](@ref) interface. - - `descent`: The descent method to use to compute the step. This needs to follow the - [`NonlinearSolve.AbstractDescentAlgorithm`](@ref) interface. - - `max_shrink_times`: The maximum number of times the trust region radius can be shrunk - before the algorithm terminates. - - `update_rule`: The update rule to use to update the Jacobian. This needs to follow the - [`NonlinearSolve.AbstractApproximateJacobianUpdateRule`](@ref) interface. - - `reinit_rule`: The reinitialization rule to use to reinitialize the Jacobian. This - needs to follow the [`NonlinearSolve.AbstractResetCondition`](@ref) interface. - - `initialization`: The initialization method to use to initialize the Jacobian. This - needs to follow the [`NonlinearSolve.AbstractJacobianInitialization`](@ref) interface. -""" -@concrete struct ApproximateJacobianSolveAlgorithm{concrete_jac, name} <: - AbstractNonlinearSolveAlgorithm{name} - linesearch - trustregion - descent - update_rule - reinit_rule - max_resets::Int - max_shrink_times::Int - initialization -end - -function __show_algorithm(io::IO, alg::ApproximateJacobianSolveAlgorithm, name, indent) - modifiers = String[] - __is_present(alg.linesearch) && push!(modifiers, "linesearch = $(alg.linesearch)") - __is_present(alg.trustregion) && push!(modifiers, "trustregion = $(alg.trustregion)") - push!(modifiers, "descent = $(alg.descent)") - push!(modifiers, "update_rule = $(alg.update_rule)") - push!(modifiers, "reinit_rule = $(alg.reinit_rule)") - push!(modifiers, "max_resets = $(alg.max_resets)") - push!(modifiers, "initialization = $(alg.initialization)") - store_inverse_jacobian(alg.update_rule) && push!(modifiers, "inverse_jacobian = true") - spacing = " "^indent * " " - spacing_last = " "^indent - print(io, "$(name)(\n$(spacing)$(join(modifiers, ",\n$(spacing)"))\n$(spacing_last))") -end - -function ApproximateJacobianSolveAlgorithm(; - concrete_jac = nothing, name::Symbol = :unknown, kwargs...) - return ApproximateJacobianSolveAlgorithm{concrete_jac, name}(; kwargs...) -end - -function ApproximateJacobianSolveAlgorithm{concrete_jac, name}(; - linesearch = missing, trustregion = missing, descent, update_rule, - reinit_rule, initialization, max_resets::Int = typemax(Int), - max_shrink_times::Int = typemax(Int)) where {concrete_jac, name} - return ApproximateJacobianSolveAlgorithm{concrete_jac, name}( - linesearch, trustregion, descent, update_rule, - reinit_rule, max_resets, max_shrink_times, initialization) -end - -@inline concrete_jac(::ApproximateJacobianSolveAlgorithm{CJ}) where {CJ} = CJ - -@concrete mutable struct ApproximateJacobianSolveCache{INV, GB, iip, timeit} <: - AbstractNonlinearSolveCache{iip, timeit} - # Basic Requirements - fu - u - u_cache - p - du # Aliased to `get_du(descent_cache)` - J # Aliased to `initialization_cache.J` if !INV - alg - prob - - # Internal Caches - initialization_cache - descent_cache - linesearch_cache - trustregion_cache - update_rule_cache - reinit_rule_cache - - inv_workspace - - # Counters - stats::NLStats - nsteps::Int - nresets::Int - max_resets::Int - maxiters::Int - maxtime - max_shrink_times::Int - steps_since_last_reset::Int - - # Timer - timer - total_time::Float64 # Simple Counter which works even if TimerOutput is disabled - - # Termination & Tracking - termination_cache - trace - retcode::ReturnCode.T - force_stop::Bool - force_reinit::Bool - kwargs -end - -store_inverse_jacobian(::ApproximateJacobianSolveCache{INV}) where {INV} = INV - -function __reinit_internal!(cache::ApproximateJacobianSolveCache{INV, GB, iip}, - args...; p = cache.p, u0 = cache.u, alias_u0::Bool = false, - maxiters = 1000, maxtime = nothing, kwargs...) where {INV, GB, iip} - if iip - recursivecopy!(cache.u, u0) - cache.prob.f(cache.fu, cache.u, p) - else - cache.u = __maybe_unaliased(u0, alias_u0) - set_fu!(cache, cache.prob.f(cache.u, p)) - end - cache.p = p - - __reinit_internal!(cache.stats) - cache.nsteps = 0 - cache.nresets = 0 - cache.steps_since_last_reset = 0 - cache.maxiters = maxiters - cache.maxtime = maxtime - cache.total_time = 0.0 - cache.force_stop = false - cache.force_reinit = false - cache.retcode = ReturnCode.Default - - reset!(cache.trace) - reinit!(cache.termination_cache, get_fu(cache), get_u(cache); kwargs...) - reset_timer!(cache.timer) -end - -@internal_caches ApproximateJacobianSolveCache :initialization_cache :descent_cache :linesearch_cache :trustregion_cache :update_rule_cache :reinit_rule_cache - -function SciMLBase.__init( - prob::AbstractNonlinearProblem{uType, iip}, alg::ApproximateJacobianSolveAlgorithm, - args...; stats = empty_nlstats(), alias_u0 = false, maxtime = nothing, - maxiters = 1000, abstol = nothing, reltol = nothing, - linsolve_kwargs = (;), termination_condition = nothing, - internalnorm::F = L2_NORM, kwargs...) where {uType, iip, F} - timer = get_timer_output() - @static_timeit timer "cache construction" begin - (; f, u0, p) = prob - u = __maybe_unaliased(u0, alias_u0) - fu = evaluate_f(prob, u) - @bb u_cache = copy(u) - - INV = store_inverse_jacobian(alg.update_rule) - - linsolve = get_linear_solver(alg.descent) - initialization_cache = __internal_init(prob, alg.initialization, alg, f, fu, u, p; - stats, linsolve, maxiters, internalnorm) - - abstol, reltol, termination_cache = NonlinearSolveBase.init_termination_cache( - prob, abstol, reltol, fu, u, termination_condition, Val(:regular)) - linsolve_kwargs = merge((; abstol, reltol), linsolve_kwargs) - - J = initialization_cache(nothing) - inv_workspace, J = INV ? __safe_inv_workspace(J) : (nothing, J) - descent_cache = __internal_init(prob, alg.descent, J, fu, u; stats, abstol, reltol, - internalnorm, linsolve_kwargs, pre_inverted = Val(INV), timer) - du = get_du(descent_cache) - - reinit_rule_cache = __internal_init(alg.reinit_rule, J, fu, u, du) - - has_linesearch = alg.linesearch !== missing && alg.linesearch !== nothing - has_trustregion = alg.trustregion !== missing && alg.trustregion !== nothing - - if has_trustregion && has_linesearch - error("TrustRegion and LineSearch methods are algorithmically incompatible.") - end - - GB = :None - linesearch_cache = nothing - trustregion_cache = nothing - - if has_trustregion - supports_trust_region(alg.descent) || error("Trust Region not supported by \ - $(alg.descent).") - trustregion_cache = __internal_init( - prob, alg.trustregion, f, fu, u, p; stats, internalnorm, kwargs...) - GB = :TrustRegion - end - - if has_linesearch - supports_line_search(alg.descent) || error("Line Search not supported by \ - $(alg.descent).") - linesearch_cache = init( - prob, alg.linesearch, fu, u; stats, internalnorm, kwargs...) - GB = :LineSearch - end - - update_rule_cache = __internal_init( - prob, alg.update_rule, J, fu, u, du; stats, internalnorm) - - trace = init_nonlinearsolve_trace(prob, alg, u, fu, ApplyArray(__zero, J), du; - uses_jacobian_inverse = Val(INV), kwargs...) - - return ApproximateJacobianSolveCache{INV, GB, iip, maxtime !== nothing}( - fu, u, u_cache, p, du, J, alg, prob, initialization_cache, - descent_cache, linesearch_cache, trustregion_cache, update_rule_cache, - reinit_rule_cache, inv_workspace, stats, 0, 0, alg.max_resets, - maxiters, maxtime, alg.max_shrink_times, 0, timer, 0.0, - termination_cache, trace, ReturnCode.Default, false, false, kwargs) - end -end - -function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip}; - recompute_jacobian::Union{Nothing, Bool} = nothing) where {INV, GB, iip} - new_jacobian = true - @static_timeit cache.timer "jacobian init/reinit" begin - if cache.nsteps == 0 # First Step is special ignore kwargs - J_init = __internal_solve!( - cache.initialization_cache, cache.fu, cache.u, Val(false)) - if INV - if jacobian_initialized_preinverted(cache.initialization_cache.alg) - cache.J = J_init - else - cache.J = __safe_inv!!(cache.inv_workspace, J_init) - end - else - if jacobian_initialized_preinverted(cache.initialization_cache.alg) - cache.J = __safe_inv!!(cache.inv_workspace, J_init) - else - cache.J = J_init - end - end - J = cache.J - cache.steps_since_last_reset += 1 - else - countable_reinit = false - if cache.force_reinit - reinit, countable_reinit = true, true - cache.force_reinit = false - elseif recompute_jacobian === nothing - # Standard Step - reinit = __internal_solve!( - cache.reinit_rule_cache, cache.J, cache.fu, cache.u, cache.du) - reinit && (countable_reinit = true) - elseif recompute_jacobian - reinit = true # Force ReInitialization: Don't count towards resetting - else - new_jacobian = false # Jacobian won't be updated in this step - reinit = false # Override Checks: Unsafe operation - end - - if countable_reinit - cache.nresets += 1 - if cache.nresets ≥ cache.max_resets - cache.retcode = ReturnCode.ConvergenceFailure - cache.force_stop = true - return - end - end - - if reinit - J_init = __internal_solve!( - cache.initialization_cache, cache.fu, cache.u, Val(true)) - cache.J = INV ? __safe_inv!!(cache.inv_workspace, J_init) : J_init - J = cache.J - cache.steps_since_last_reset = 0 - else - J = cache.J - cache.steps_since_last_reset += 1 - end - end - end - - @static_timeit cache.timer "descent" begin - if cache.trustregion_cache !== nothing && - hasfield(typeof(cache.trustregion_cache), :trust_region) - descent_result = __internal_solve!( - cache.descent_cache, J, cache.fu, cache.u; new_jacobian, - trust_region = cache.trustregion_cache.trust_region, cache.kwargs...) - else - descent_result = __internal_solve!( - cache.descent_cache, J, cache.fu, cache.u; new_jacobian, cache.kwargs...) - end - end - - if !descent_result.linsolve_success - if new_jacobian && cache.steps_since_last_reset == 0 - # Extremely pathological case. Jacobian was just reset and linear solve - # failed. Should ideally never happen in practice unless true jacobian init - # is used. - cache.retcode = LinearSolveFailureCode - cache.force_stop = true - return - else - # Force a reinit because the problem is currently un-solvable - if !haskey(cache.kwargs, :verbose) || cache.kwargs[:verbose] - @warn "Linear Solve Failed but Jacobian Information is not current. \ - Retrying with reinitialized Approximate Jacobian." - end - cache.force_reinit = true - __step!(cache; recompute_jacobian = true) - return - end - end - - δu, descent_intermediates = descent_result.δu, descent_result.extras - - if descent_result.success - if GB === :LineSearch - @static_timeit cache.timer "linesearch" begin - linesearch_sol = solve!(cache.linesearch_cache, cache.u, δu) - needs_reset = !SciMLBase.successful_retcode(linesearch_sol.retcode) - α = linesearch_sol.step_size - end - if needs_reset && cache.steps_since_last_reset > 5 # Reset after a burn-in period - cache.force_reinit = true - else - @static_timeit cache.timer "step" begin - @bb axpy!(α, δu, cache.u) - evaluate_f!(cache, cache.u, cache.p) - end - end - elseif GB === :TrustRegion - @static_timeit cache.timer "trustregion" begin - tr_accepted, u_new, fu_new = __internal_solve!( - cache.trustregion_cache, J, cache.fu, - cache.u, δu, descent_intermediates) - if tr_accepted - @bb copyto!(cache.u, u_new) - @bb copyto!(cache.fu, fu_new) - end - if hasfield(typeof(cache.trustregion_cache), :shrink_counter) && - cache.trustregion_cache.shrink_counter > cache.max_shrink_times - cache.retcode = ReturnCode.ShrinkThresholdExceeded - cache.force_stop = true - end - end - α = true - elseif GB === :None - @static_timeit cache.timer "step" begin - @bb axpy!(1, δu, cache.u) - evaluate_f!(cache, cache.u, cache.p) - end - α = true - else - error("Unknown Globalization Strategy: $(GB). Allowed values are (:LineSearch, \ - :TrustRegion, :None)") - end - check_and_update!(cache, cache.fu, cache.u, cache.u_cache) - else - α = false - cache.force_reinit = true - end - - update_trace!(cache, α) - @bb copyto!(cache.u_cache, cache.u) - - if (cache.force_stop || - cache.force_reinit || - (recompute_jacobian !== nothing && !recompute_jacobian)) - callback_into_cache!(cache) - return nothing - end - - @static_timeit cache.timer "jacobian update" begin - cache.J = __internal_solve!(cache.update_rule_cache, cache.J, cache.fu, cache.u, δu) - callback_into_cache!(cache) - end - - return nothing -end diff --git a/src/core/generalized_first_order.jl b/src/core/generalized_first_order.jl deleted file mode 100644 index d5ee3eeab..000000000 --- a/src/core/generalized_first_order.jl +++ /dev/null @@ -1,324 +0,0 @@ -""" - GeneralizedFirstOrderAlgorithm{concrete_jac, name}(; descent, linesearch = missing, - trustregion = missing, autodiff = nothing, vjp_autodiff = nothing, - jvp_autodiff = nothing, max_shrink_times::Int = typemax(Int)) - GeneralizedFirstOrderAlgorithm(; concrete_jac = nothing, name::Symbol = :unknown, - kwargs...) - -This is a Generalization of First-Order (uses Jacobian) Nonlinear Solve Algorithms. The most -common example of this is Newton-Raphson Method. - -First Order here refers to the order of differentiation, and should not be confused with the -order of convergence. - -!!! danger - - `trustregion` and `linesearch` cannot be specified together. - -### Keyword Arguments - - - `trustregion`: Globalization using a Trust Region Method. This needs to follow the - [`NonlinearSolve.AbstractTrustRegionMethod`](@ref) interface. - - `descent`: The descent method to use to compute the step. This needs to follow the - [`NonlinearSolve.AbstractDescentAlgorithm`](@ref) interface. - - `max_shrink_times`: The maximum number of times the trust region radius can be shrunk - before the algorithm terminates. -""" -@concrete struct GeneralizedFirstOrderAlgorithm{concrete_jac, name} <: - AbstractNonlinearSolveAlgorithm{name} - linesearch - trustregion - descent - max_shrink_times::Int - - autodiff - vjp_autodiff - jvp_autodiff -end - -function __show_algorithm(io::IO, alg::GeneralizedFirstOrderAlgorithm, name, indent) - modifiers = String[] - __is_present(alg.linesearch) && push!(modifiers, "linesearch = $(alg.linesearch)") - __is_present(alg.trustregion) && push!(modifiers, "trustregion = $(alg.trustregion)") - push!(modifiers, "descent = $(alg.descent)") - __is_present(alg.autodiff) && push!(modifiers, "autodiff = $(alg.autodiff)") - __is_present(alg.vjp_autodiff) && push!(modifiers, "vjp_autodiff = $(alg.vjp_autodiff)") - __is_present(alg.jvp_autodiff) && push!(modifiers, "jvp_autodiff = $(alg.jvp_autodiff)") - spacing = " "^indent * " " - spacing_last = " "^indent - print(io, "$(name)(\n$(spacing)$(join(modifiers, ",\n$(spacing)"))\n$(spacing_last))") -end - -function GeneralizedFirstOrderAlgorithm(; - concrete_jac = nothing, name::Symbol = :unknown, kwargs...) - return GeneralizedFirstOrderAlgorithm{concrete_jac, name}(; kwargs...) -end - -function GeneralizedFirstOrderAlgorithm{concrete_jac, name}(; - descent, linesearch = missing, trustregion = missing, - autodiff = nothing, jvp_autodiff = nothing, vjp_autodiff = nothing, - max_shrink_times::Int = typemax(Int)) where {concrete_jac, name} - return GeneralizedFirstOrderAlgorithm{concrete_jac, name}( - linesearch, trustregion, descent, max_shrink_times, - autodiff, vjp_autodiff, jvp_autodiff) -end - -concrete_jac(::GeneralizedFirstOrderAlgorithm{CJ}) where {CJ} = CJ - -@concrete mutable struct GeneralizedFirstOrderAlgorithmCache{iip, GB, timeit} <: - AbstractNonlinearSolveCache{iip, timeit} - # Basic Requirements - fu - u - u_cache - p - du # Aliased to `get_du(descent_cache)` - J # Aliased to `jac_cache.J` - alg - prob - - # Internal Caches - jac_cache - descent_cache - linesearch_cache - trustregion_cache - - # Counters - stats::NLStats - nsteps::Int - maxiters::Int - maxtime - max_shrink_times::Int - - # Timer - timer - total_time::Float64 # Simple Counter which works even if TimerOutput is disabled - - # State Affect - make_new_jacobian::Bool - - # Termination & Tracking - termination_cache - trace - retcode::ReturnCode.T - force_stop::Bool - kwargs -end - -SymbolicIndexingInterface.state_values(cache::GeneralizedFirstOrderAlgorithmCache) = cache.u -function SymbolicIndexingInterface.parameter_values(cache::GeneralizedFirstOrderAlgorithmCache) - cache.p -end - -function __reinit_internal!( - cache::GeneralizedFirstOrderAlgorithmCache{iip}, args...; p = cache.p, u0 = cache.u, - alias_u0::Bool = false, maxiters = 1000, maxtime = nothing, kwargs...) where {iip} - if iip - recursivecopy!(cache.u, u0) - cache.prob.f(cache.fu, cache.u, p) - else - cache.u = __maybe_unaliased(u0, alias_u0) - set_fu!(cache, cache.prob.f(cache.u, p)) - end - cache.p = p - - __reinit_internal!(cache.stats) - cache.nsteps = 0 - cache.maxiters = maxiters - cache.maxtime = maxtime - cache.total_time = 0.0 - cache.force_stop = false - cache.retcode = ReturnCode.Default - cache.make_new_jacobian = true - - reset!(cache.trace) - reinit!(cache.termination_cache, get_fu(cache), get_u(cache); kwargs...) - reset_timer!(cache.timer) -end - -@internal_caches GeneralizedFirstOrderAlgorithmCache :jac_cache :descent_cache :linesearch_cache :trustregion_cache - -function SciMLBase.__init( - prob::AbstractNonlinearProblem{uType, iip}, alg::GeneralizedFirstOrderAlgorithm, - args...; stats = empty_nlstats(), alias_u0 = false, maxiters = 1000, - abstol = nothing, reltol = nothing, maxtime = nothing, - termination_condition = nothing, internalnorm = L2_NORM, - linsolve_kwargs = (;), kwargs...) where {uType, iip} - autodiff = select_jacobian_autodiff(prob, alg.autodiff) - jvp_autodiff = if alg.jvp_autodiff === nothing && alg.autodiff !== nothing && - (ADTypes.mode(alg.autodiff) isa ADTypes.ForwardMode || - ADTypes.mode(alg.autodiff) isa ADTypes.ForwardOrReverseMode) - select_forward_mode_autodiff(prob, alg.autodiff) - else - select_forward_mode_autodiff(prob, alg.jvp_autodiff) - end - vjp_autodiff = if alg.vjp_autodiff === nothing && alg.autodiff !== nothing && - (ADTypes.mode(alg.autodiff) isa ADTypes.ReverseMode || - ADTypes.mode(alg.autodiff) isa ADTypes.ForwardOrReverseMode) - select_reverse_mode_autodiff(prob, alg.autodiff) - else - select_reverse_mode_autodiff(prob, alg.vjp_autodiff) - end - - timer = get_timer_output() - @static_timeit timer "cache construction" begin - (; f, u0, p) = prob - u = __maybe_unaliased(u0, alias_u0) - fu = evaluate_f(prob, u) - @bb u_cache = copy(u) - - linsolve = get_linear_solver(alg.descent) - - abstol, reltol, termination_cache = NonlinearSolveBase.init_termination_cache( - prob, abstol, reltol, fu, u, termination_condition, Val(:regular)) - linsolve_kwargs = merge((; abstol, reltol), linsolve_kwargs) - - jac_cache = JacobianCache( - prob, alg, f, fu, u, p; stats, autodiff, linsolve, jvp_autodiff, vjp_autodiff) - J = jac_cache(nothing) - descent_cache = __internal_init(prob, alg.descent, J, fu, u; stats, abstol, - reltol, internalnorm, linsolve_kwargs, timer) - du = get_du(descent_cache) - - has_linesearch = alg.linesearch !== missing && alg.linesearch !== nothing - has_trustregion = alg.trustregion !== missing && alg.trustregion !== nothing - - if has_trustregion && has_linesearch - error("TrustRegion and LineSearch methods are algorithmically incompatible.") - end - - GB = :None - linesearch_cache = nothing - trustregion_cache = nothing - - if has_trustregion - supports_trust_region(alg.descent) || error("Trust Region not supported by \ - $(alg.descent).") - trustregion_cache = __internal_init( - prob, alg.trustregion, f, fu, u, p; stats, internalnorm, kwargs..., - autodiff, jvp_autodiff, vjp_autodiff) - GB = :TrustRegion - end - - if has_linesearch - supports_line_search(alg.descent) || error("Line Search not supported by \ - $(alg.descent).") - linesearch_cache = init( - prob, alg.linesearch, fu, u; stats, autodiff = jvp_autodiff, kwargs...) - GB = :LineSearch - end - - trace = init_nonlinearsolve_trace( - prob, alg, u, fu, ApplyArray(__zero, J), du; kwargs...) - - return GeneralizedFirstOrderAlgorithmCache{iip, GB, maxtime !== nothing}( - fu, u, u_cache, p, du, J, alg, prob, jac_cache, descent_cache, linesearch_cache, - trustregion_cache, stats, 0, maxiters, maxtime, alg.max_shrink_times, timer, - 0.0, true, termination_cache, trace, ReturnCode.Default, false, kwargs) - end -end - -function __step!(cache::GeneralizedFirstOrderAlgorithmCache{iip, GB}; - recompute_jacobian::Union{Nothing, Bool} = nothing, kwargs...) where {iip, GB} - @static_timeit cache.timer "jacobian" begin - if (recompute_jacobian === nothing || recompute_jacobian) && cache.make_new_jacobian - J = cache.jac_cache(cache.u) - new_jacobian = true - else - J = cache.jac_cache(nothing) - new_jacobian = false - end - end - - @static_timeit cache.timer "descent" begin - if cache.trustregion_cache !== nothing && - hasfield(typeof(cache.trustregion_cache), :trust_region) - descent_result = __internal_solve!( - cache.descent_cache, J, cache.fu, cache.u; new_jacobian, - trust_region = cache.trustregion_cache.trust_region, cache.kwargs...) - else - descent_result = __internal_solve!( - cache.descent_cache, J, cache.fu, cache.u; new_jacobian, cache.kwargs...) - end - end - - if !descent_result.linsolve_success - if new_jacobian - # Jacobian Information is current and linear solve failed terminate the solve - cache.retcode = LinearSolveFailureCode - cache.force_stop = true - return - else - # Jacobian Information is not current and linear solve failed, recompute - # Jacobian - if !haskey(cache.kwargs, :verbose) || cache.kwargs[:verbose] - @warn "Linear Solve Failed but Jacobian Information is not current. \ - Retrying with updated Jacobian." - end - # In the 2nd call the `new_jacobian` is guaranteed to be `true`. - cache.make_new_jacobian = true - __step!(cache; recompute_jacobian = true, kwargs...) - return - end - end - - δu, descent_intermediates = descent_result.δu, descent_result.extras - - if descent_result.success - cache.make_new_jacobian = true - if GB === :LineSearch - @static_timeit cache.timer "linesearch" begin - linesearch_sol = solve!(cache.linesearch_cache, cache.u, δu) - linesearch_failed = !SciMLBase.successful_retcode(linesearch_sol.retcode) - α = linesearch_sol.step_size - end - if linesearch_failed - cache.retcode = ReturnCode.InternalLineSearchFailed - cache.force_stop = true - end - @static_timeit cache.timer "step" begin - @bb axpy!(α, δu, cache.u) - evaluate_f!(cache, cache.u, cache.p) - end - elseif GB === :TrustRegion - @static_timeit cache.timer "trustregion" begin - tr_accepted, u_new, fu_new = __internal_solve!( - cache.trustregion_cache, J, cache.fu, - cache.u, δu, descent_intermediates) - if tr_accepted - @bb copyto!(cache.u, u_new) - @bb copyto!(cache.fu, fu_new) - α = true - else - α = false - cache.make_new_jacobian = false - end - if hasfield(typeof(cache.trustregion_cache), :shrink_counter) && - cache.trustregion_cache.shrink_counter > cache.max_shrink_times - cache.retcode = ReturnCode.ShrinkThresholdExceeded - cache.force_stop = true - end - end - elseif GB === :None - @static_timeit cache.timer "step" begin - @bb axpy!(1, δu, cache.u) - evaluate_f!(cache, cache.u, cache.p) - end - α = true - else - error("Unknown Globalization Strategy: $(GB). Allowed values are (:LineSearch, \ - :TrustRegion, :None)") - end - check_and_update!(cache, cache.fu, cache.u, cache.u_cache) - else - α = false - cache.make_new_jacobian = false - end - - update_trace!(cache, α) - @bb copyto!(cache.u_cache, cache.u) - - callback_into_cache!(cache) - - return nothing -end diff --git a/src/core/generic.jl b/src/core/generic.jl deleted file mode 100644 index 44e6e1152..000000000 --- a/src/core/generic.jl +++ /dev/null @@ -1,67 +0,0 @@ -function SciMLBase.__solve(prob::Union{NonlinearProblem, NonlinearLeastSquaresProblem}, - alg::AbstractNonlinearSolveAlgorithm, args...; stats = empty_nlstats(), kwargs...) - cache = SciMLBase.__init(prob, alg, args...; stats, kwargs...) - return solve!(cache) -end - -function not_terminated(cache::AbstractNonlinearSolveCache) - return !cache.force_stop && cache.nsteps < cache.maxiters -end - -function SciMLBase.solve!(cache::AbstractNonlinearSolveCache) - while not_terminated(cache) - step!(cache) - end - - # The solver might have set a different `retcode` - if cache.retcode == ReturnCode.Default - cache.retcode = ifelse( - cache.nsteps ≥ cache.maxiters, ReturnCode.MaxIters, ReturnCode.Success) - end - - update_from_termination_cache!(cache.termination_cache, cache) - - update_trace!(cache.trace, cache.nsteps, get_u(cache), - get_fu(cache), nothing, nothing, nothing; last = True) - - return SciMLBase.build_solution(cache.prob, cache.alg, get_u(cache), get_fu(cache); - cache.retcode, cache.stats, cache.trace) -end - -""" - step!(cache::AbstractNonlinearSolveCache; - recompute_jacobian::Union{Nothing, Bool} = nothing) - -Performs one step of the nonlinear solver. - -### Keyword Arguments - - - `recompute_jacobian`: allows controlling whether the jacobian is recomputed at the - current step. If `nothing`, then the algorithm determines whether to recompute the - jacobian. If `true` or `false`, then the jacobian is recomputed or not recomputed, - respectively. For algorithms that don't use jacobian information, this keyword is - ignored with a one-time warning. -""" -function SciMLBase.step!(cache::AbstractNonlinearSolveCache{iip, timeit}, - args...; kwargs...) where {iip, timeit} - not_terminated(cache) || return - timeit && (time_start = time()) - res = @static_timeit cache.timer "solve" begin - __step!(cache, args...; kwargs...) - end - - cache.stats.nsteps += 1 - cache.nsteps += 1 - - if timeit - cache.total_time += time() - time_start - if !cache.force_stop && - cache.retcode == ReturnCode.Default && - cache.total_time ≥ cache.maxtime - cache.retcode = ReturnCode.MaxTime - cache.force_stop = true - end - end - - return res -end diff --git a/src/core/noinit.jl b/src/core/noinit.jl deleted file mode 100644 index b51c09c23..000000000 --- a/src/core/noinit.jl +++ /dev/null @@ -1,37 +0,0 @@ -# Some algorithms don't support creating a cache and doing `solve!`, this unfortunately -# makes it difficult to write generic code that supports caching. For the algorithms that -# don't have a `__init` function defined, we create a "Fake Cache", which just calls -# `__solve` from `solve!` -@concrete mutable struct NonlinearSolveNoInitCache{iip, timeit} <: - AbstractNonlinearSolveCache{iip, timeit} - prob - alg - args - kwargs::Any -end - -function SciMLBase.reinit!( - cache::NonlinearSolveNoInitCache, u0 = cache.prob.u0; p = cache.prob.p, kwargs...) - prob = remake(cache.prob; u0, p) - cache.prob = prob - cache.kwargs = merge(cache.kwargs, kwargs) - return cache -end - -function Base.show(io::IO, cache::NonlinearSolveNoInitCache) - print(io, "NonlinearSolveNoInitCache(alg = $(cache.alg))") -end - -function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip}, - alg::Union{AbstractNonlinearSolveAlgorithm, - SimpleNonlinearSolve.AbstractSimpleNonlinearSolveAlgorithm}, - args...; - maxtime = nothing, - kwargs...) where {uType, iip} - return NonlinearSolveNoInitCache{iip, maxtime !== nothing}( - prob, alg, args, merge((; maxtime), kwargs)) -end - -function SciMLBase.solve!(cache::NonlinearSolveNoInitCache) - return solve(cache.prob, cache.alg, cache.args...; cache.kwargs...) -end diff --git a/src/default.jl b/src/default.jl index a7cc550d8..6021a98b1 100644 --- a/src/default.jl +++ b/src/default.jl @@ -1,496 +1,3 @@ -# Poly Algorithms -""" - NonlinearSolvePolyAlgorithm(algs, ::Val{pType} = Val(:NLS); - start_index = 1) where {pType} - -A general way to define PolyAlgorithms for `NonlinearProblem` and -`NonlinearLeastSquaresProblem`. This is a container for a tuple of algorithms that will be -tried in order until one succeeds. If none succeed, then the algorithm with the lowest -residual is returned. - -### Arguments - - - `algs`: a tuple of algorithms to try in-order! (If this is not a Tuple, then the - returned algorithm is not type-stable). - - `pType`: the problem type. Defaults to `:NLS` for `NonlinearProblem` and `:NLLS` for - `NonlinearLeastSquaresProblem`. This is used to determine the correct problem type to - dispatch on. - -### Keyword Arguments - - - `start_index`: the index to start at. Defaults to `1`. - -### Example - -```julia -using NonlinearSolve - -alg = NonlinearSolvePolyAlgorithm((NewtonRaphson(), Broyden())) -``` -""" -struct NonlinearSolvePolyAlgorithm{pType, N, A} <: AbstractNonlinearSolveAlgorithm{:PolyAlg} - algs::A - start_index::Int - - function NonlinearSolvePolyAlgorithm( - algs, ::Val{pType} = Val(:NLS); start_index::Int = 1) where {pType} - @assert pType ∈ (:NLS, :NLLS) - @assert 0 < start_index ≤ length(algs) - algs = Tuple(algs) - return new{pType, length(algs), typeof(algs)}(algs, start_index) - end -end - -function Base.show(io::IO, alg::NonlinearSolvePolyAlgorithm{pType, N}) where {pType, N} - problem_kind = ifelse(pType == :NLS, "NonlinearProblem", "NonlinearLeastSquaresProblem") - println(io, "NonlinearSolvePolyAlgorithm for $(problem_kind) with $(N) algorithms") - for i in 1:N - num = " [$(i)]: " - print(io, num) - __show_algorithm(io, alg.algs[i], get_name(alg.algs[i]), length(num)) - i == N || println(io) - end -end - -@concrete mutable struct NonlinearSolvePolyAlgorithmCache{iip, N, timeit} <: - AbstractNonlinearSolveCache{iip, timeit} - caches - alg - best::Int - current::Int - nsteps::Int - stats::NLStats - total_time::Float64 - maxtime - retcode::ReturnCode.T - force_stop::Bool - maxiters::Int - internalnorm - u0 - u0_aliased - alias_u0::Bool -end - -function SymbolicIndexingInterface.symbolic_container(cache::NonlinearSolvePolyAlgorithmCache) - cache.caches[cache.current] -end -SymbolicIndexingInterface.state_values(cache::NonlinearSolvePolyAlgorithmCache) = cache.u0 - -function Base.show( - io::IO, cache::NonlinearSolvePolyAlgorithmCache{pType, N}) where {pType, N} - problem_kind = ifelse(pType == :NLS, "NonlinearProblem", "NonlinearLeastSquaresProblem") - println(io, "NonlinearSolvePolyAlgorithmCache for $(problem_kind) with $(N) algorithms") - best_alg = ifelse(cache.best == -1, "nothing", cache.best) - println(io, "Best algorithm: $(best_alg)") - println(io, "Current algorithm: $(cache.current)") - println(io, "nsteps: $(cache.nsteps)") - println(io, "retcode: $(cache.retcode)") - __show_cache(io, cache.caches[cache.current], 0) -end - -function reinit_cache!(cache::NonlinearSolvePolyAlgorithmCache, args...; kwargs...) - foreach(c -> reinit_cache!(c, args...; kwargs...), cache.caches) - cache.current = cache.alg.start_index - __reinit_internal!(cache.stats) - cache.nsteps = 0 - cache.total_time = 0.0 -end - -for (probType, pType) in ((:NonlinearProblem, :NLS), (:NonlinearLeastSquaresProblem, :NLLS)) - algType = NonlinearSolvePolyAlgorithm{pType} - @eval begin - function SciMLBase.__init( - prob::$probType, alg::$algType{N}, args...; stats = empty_nlstats(), - maxtime = nothing, maxiters = 1000, internalnorm = L2_NORM, - alias_u0 = false, verbose = true, kwargs...) where {N} - if (alias_u0 && !ismutable(prob.u0)) - verbose && @warn "`alias_u0` has been set to `true`, but `u0` is \ - immutable (checked using `ArrayInterface.ismutable`)." - alias_u0 = false # If immutable don't care about aliasing - end - u0 = prob.u0 - if alias_u0 - u0_aliased = copy(u0) - else - u0_aliased = u0 # Irrelevant - end - alias_u0 && (prob = remake(prob; u0 = u0_aliased)) - return NonlinearSolvePolyAlgorithmCache{isinplace(prob), N, maxtime !== nothing}( - map( - solver -> SciMLBase.__init(prob, solver, args...; stats, maxtime, - internalnorm, alias_u0, verbose, kwargs...), - alg.algs), - alg, - -1, - alg.start_index, - 0, - stats, - 0.0, - maxtime, - ReturnCode.Default, - false, - maxiters, - internalnorm, - u0, - u0_aliased, - alias_u0) - end - end -end - -@generated function SciMLBase.solve!(cache::NonlinearSolvePolyAlgorithmCache{ - iip, N}) where {iip, N} - calls = [quote - 1 ≤ cache.current ≤ length(cache.caches) || - error("Current choices shouldn't get here!") - end] - - cache_syms = [gensym("cache") for i in 1:N] - sol_syms = [gensym("sol") for i in 1:N] - u_result_syms = [gensym("u_result") for i in 1:N] - for i in 1:N - push!(calls, - quote - $(cache_syms[i]) = cache.caches[$(i)] - if $(i) == cache.current - cache.alias_u0 && copyto!(cache.u0_aliased, cache.u0) - $(sol_syms[i]) = SciMLBase.solve!($(cache_syms[i])) - if SciMLBase.successful_retcode($(sol_syms[i])) - stats = $(sol_syms[i]).stats - if cache.alias_u0 - copyto!(cache.u0, $(sol_syms[i]).u) - $(u_result_syms[i]) = cache.u0 - else - $(u_result_syms[i]) = $(sol_syms[i]).u - end - fu = get_fu($(cache_syms[i])) - return __build_solution_less_specialize( - $(sol_syms[i]).prob, cache.alg, $(u_result_syms[i]), - fu; retcode = $(sol_syms[i]).retcode, stats, - original = $(sol_syms[i]), trace = $(sol_syms[i]).trace) - elseif cache.alias_u0 - # For safety we need to maintain a copy of the solution - $(u_result_syms[i]) = copy($(sol_syms[i]).u) - end - cache.current = $(i + 1) - end - end) - end - - resids = map(x -> Symbol("$(x)_resid"), cache_syms) - for (sym, resid) in zip(cache_syms, resids) - push!(calls, :($(resid) = @isdefined($(sym)) ? get_fu($(sym)) : nothing)) - end - push!(calls, quote - fus = tuple($(Tuple(resids)...)) - minfu, idx = __findmin(cache.internalnorm, fus) - end) - for i in 1:N - push!(calls, quote - if idx == $(i) - if cache.alias_u0 - u = $(u_result_syms[i]) - else - u = get_u(cache.caches[$i]) - end - end - end) - end - push!(calls, - quote - retcode = cache.caches[idx].retcode - if cache.alias_u0 - copyto!(cache.u0, u) - u = cache.u0 - end - return __build_solution_less_specialize( - cache.caches[idx].prob, cache.alg, u, fus[idx]; - retcode, stats = cache.stats, cache.caches[idx].trace) - end) - - return Expr(:block, calls...) -end - -@generated function __step!( - cache::NonlinearSolvePolyAlgorithmCache{iip, N}, args...; kwargs...) where {iip, N} - calls = [] - cache_syms = [gensym("cache") for i in 1:N] - for i in 1:N - push!(calls, - quote - $(cache_syms[i]) = cache.caches[$(i)] - if $(i) == cache.current - __step!($(cache_syms[i]), args...; kwargs...) - $(cache_syms[i]).nsteps += 1 - if !not_terminated($(cache_syms[i])) - if SciMLBase.successful_retcode($(cache_syms[i]).retcode) - cache.best = $(i) - cache.force_stop = true - cache.retcode = $(cache_syms[i]).retcode - else - cache.current = $(i + 1) - end - end - return - end - end) - end - - push!(calls, quote - if !(1 ≤ cache.current ≤ length(cache.caches)) - minfu, idx = __findmin(cache.internalnorm, cache.caches) - cache.best = idx - cache.retcode = cache.caches[cache.best].retcode - cache.force_stop = true - return - end - end) - - return Expr(:block, calls...) -end - -for (probType, pType) in ((:NonlinearProblem, :NLS), (:NonlinearLeastSquaresProblem, :NLLS)) - algType = NonlinearSolvePolyAlgorithm{pType} - @eval begin - @generated function SciMLBase.__solve( - prob::$probType, alg::$algType{N}, args...; stats = empty_nlstats(), - alias_u0 = false, verbose = true, kwargs...) where {N} - sol_syms = [gensym("sol") for _ in 1:N] - prob_syms = [gensym("prob") for _ in 1:N] - u_result_syms = [gensym("u_result") for _ in 1:N] - calls = [quote - current = alg.start_index - if (alias_u0 && !ismutable(prob.u0)) - verbose && @warn "`alias_u0` has been set to `true`, but `u0` is \ - immutable (checked using `ArrayInterface.ismutable`)." - alias_u0 = false # If immutable don't care about aliasing - end - u0 = prob.u0 - u0_aliased = alias_u0 ? zero(u0) : u0 - end] - for i in 1:N - cur_sol = sol_syms[i] - push!(calls, - quote - if current == $i - if alias_u0 - copyto!(u0_aliased, u0) - $(prob_syms[i]) = remake(prob; u0 = u0_aliased) - else - $(prob_syms[i]) = prob - end - $(cur_sol) = SciMLBase.__solve( - $(prob_syms[i]), alg.algs[$(i)], args...; - stats, alias_u0, verbose, kwargs...) - if SciMLBase.successful_retcode($(cur_sol)) - if alias_u0 - copyto!(u0, $(cur_sol).u) - $(u_result_syms[i]) = u0 - else - $(u_result_syms[i]) = $(cur_sol).u - end - return __build_solution_less_specialize( - prob, alg, $(u_result_syms[i]), $(cur_sol).resid; - $(cur_sol).retcode, $(cur_sol).stats, - original = $(cur_sol), trace = $(cur_sol).trace) - elseif alias_u0 - # For safety we need to maintain a copy of the solution - $(u_result_syms[i]) = copy($(cur_sol).u) - end - current = $(i + 1) - end - end) - end - - resids = map(x -> Symbol("$(x)_resid"), sol_syms) - for (sym, resid) in zip(sol_syms, resids) - push!(calls, :($(resid) = @isdefined($(sym)) ? $(sym).resid : nothing)) - end - - push!(calls, quote - resids = tuple($(Tuple(resids)...)) - minfu, idx = __findmin(L2_NORM, resids) - end) - - for i in 1:N - push!(calls, - quote - if idx == $i - if alias_u0 - copyto!(u0, $(u_result_syms[i])) - $(u_result_syms[i]) = u0 - else - $(u_result_syms[i]) = $(sol_syms[i]).u - end - return __build_solution_less_specialize( - prob, alg, $(u_result_syms[i]), $(sol_syms[i]).resid; - $(sol_syms[i]).retcode, $(sol_syms[i]).stats, - $(sol_syms[i]).trace, original = $(sol_syms[i])) - end - end) - end - push!(calls, :(error("Current choices shouldn't get here!"))) - - return Expr(:block, calls...) - end - end -end - -""" - RobustMultiNewton(::Type{T} = Float64; concrete_jac = nothing, linsolve = nothing, - precs = DEFAULT_PRECS, autodiff = nothing) - -A polyalgorithm focused on robustness. It uses a mixture of Newton methods with different -globalizing techniques (trust region updates, line searches, etc.) in order to find a -method that is able to adequately solve the minimization problem. - -Basically, if this algorithm fails, then "most" good ways of solving your problem fail and -you may need to think about reformulating the model (either there is an issue with the model, -or more precision / more stable linear solver choice is required). - -### Arguments - - - `T`: The eltype of the initial guess. It is only used to check if some of the algorithms - are compatible with the problem type. Defaults to `Float64`. -""" -function RobustMultiNewton(::Type{T} = Float64; concrete_jac = nothing, linsolve = nothing, - precs = DEFAULT_PRECS, autodiff = nothing) where {T} - if __is_complex(T) - # Let's atleast have something here for complex numbers - algs = (NewtonRaphson(; concrete_jac, linsolve, precs, autodiff),) - else - algs = (TrustRegion(; concrete_jac, linsolve, precs, autodiff), - TrustRegion(; concrete_jac, linsolve, precs, autodiff, - radius_update_scheme = RadiusUpdateSchemes.Bastin), - NewtonRaphson(; concrete_jac, linsolve, precs, autodiff), - NewtonRaphson(; concrete_jac, linsolve, precs, - linesearch = BackTracking(), autodiff), - TrustRegion(; concrete_jac, linsolve, precs, - radius_update_scheme = RadiusUpdateSchemes.NLsolve, autodiff), - TrustRegion(; concrete_jac, linsolve, precs, - radius_update_scheme = RadiusUpdateSchemes.Fan, autodiff)) - end - return NonlinearSolvePolyAlgorithm(algs, Val(:NLS)) -end - -""" - FastShortcutNonlinearPolyalg(::Type{T} = Float64; concrete_jac = nothing, - linsolve = nothing, precs = DEFAULT_PRECS, must_use_jacobian::Val = Val(false), - prefer_simplenonlinearsolve::Val{SA} = Val(false), autodiff = nothing, - u0_len::Union{Int, Nothing} = nothing) where {T} - -A polyalgorithm focused on balancing speed and robustness. It first tries less robust methods -for more performance and then tries more robust techniques if the faster ones fail. - -### Arguments - - - `T`: The eltype of the initial guess. It is only used to check if some of the algorithms - are compatible with the problem type. Defaults to `Float64`. - -### Keyword Arguments - - - `u0_len`: The length of the initial guess. If this is `nothing`, then the length of the - initial guess is not checked. If this is an integer and it is less than `25`, we use - jacobian based methods. -""" -function FastShortcutNonlinearPolyalg( - ::Type{T} = Float64; concrete_jac = nothing, linsolve = nothing, - precs = DEFAULT_PRECS, must_use_jacobian::Val{JAC} = Val(false), - prefer_simplenonlinearsolve::Val{SA} = Val(false), - u0_len::Union{Int, Nothing} = nothing, autodiff = nothing) where {T, JAC, SA} - start_index = 1 - if JAC - if __is_complex(T) - algs = (NewtonRaphson(; concrete_jac, linsolve, precs, autodiff),) - else - algs = (NewtonRaphson(; concrete_jac, linsolve, precs, autodiff), - NewtonRaphson(; concrete_jac, linsolve, precs, - linesearch = BackTracking(), autodiff), - TrustRegion(; concrete_jac, linsolve, precs, autodiff), - TrustRegion(; concrete_jac, linsolve, precs, - radius_update_scheme = RadiusUpdateSchemes.Bastin, autodiff)) - end - else - # SimpleNewtonRaphson and SimpleTrustRegion are not robust to singular Jacobians - # and thus are not included in the polyalgorithm - if SA - if __is_complex(T) - algs = (SimpleBroyden(), - Broyden(; init_jacobian = Val(:true_jacobian), autodiff), - SimpleKlement(), - NewtonRaphson(; concrete_jac, linsolve, precs, autodiff)) - else - start_index = u0_len !== nothing ? (u0_len ≤ 25 ? 4 : 1) : 1 - algs = (SimpleBroyden(), - Broyden(; init_jacobian = Val(:true_jacobian), autodiff), - SimpleKlement(), - NewtonRaphson(; concrete_jac, linsolve, precs, autodiff), - NewtonRaphson(; concrete_jac, linsolve, precs, - linesearch = BackTracking(), autodiff), - TrustRegion(; concrete_jac, linsolve, precs, - radius_update_scheme = RadiusUpdateSchemes.Bastin, autodiff)) - end - else - if __is_complex(T) - algs = ( - Broyden(), Broyden(; init_jacobian = Val(:true_jacobian), autodiff), - Klement(; linsolve, precs, autodiff), - NewtonRaphson(; concrete_jac, linsolve, precs, autodiff)) - else - # TODO: This number requires a bit rigorous testing - start_index = u0_len !== nothing ? (u0_len ≤ 25 ? 4 : 1) : 1 - algs = ( - Broyden(; autodiff), - Broyden(; init_jacobian = Val(:true_jacobian), autodiff), - Klement(; linsolve, precs, autodiff), - NewtonRaphson(; concrete_jac, linsolve, precs, autodiff), - NewtonRaphson(; concrete_jac, linsolve, precs, - linesearch = BackTracking(), autodiff), - TrustRegion(; concrete_jac, linsolve, precs, autodiff), - TrustRegion(; concrete_jac, linsolve, precs, - radius_update_scheme = RadiusUpdateSchemes.Bastin, autodiff)) - end - end - end - return NonlinearSolvePolyAlgorithm(algs, Val(:NLS); start_index) -end - -""" - FastShortcutNLLSPolyalg(::Type{T} = Float64; concrete_jac = nothing, linsolve = nothing, - precs = DEFAULT_PRECS, autodiff = nothing, kwargs...) - -A polyalgorithm focused on balancing speed and robustness. It first tries less robust methods -for more performance and then tries more robust techniques if the faster ones fail. - -### Arguments - - - `T`: The eltype of the initial guess. It is only used to check if some of the algorithms - are compatible with the problem type. Defaults to `Float64`. -""" -function FastShortcutNLLSPolyalg( - ::Type{T} = Float64; concrete_jac = nothing, linsolve = nothing, - precs = DEFAULT_PRECS, autodiff = nothing, kwargs...) where {T} - if __is_complex(T) - algs = (GaussNewton(; concrete_jac, linsolve, precs, autodiff, kwargs...), - LevenbergMarquardt(; - linsolve, precs, autodiff, disable_geodesic = Val(true), kwargs...), - LevenbergMarquardt(; linsolve, precs, autodiff, kwargs...)) - else - algs = (GaussNewton(; concrete_jac, linsolve, precs, autodiff, kwargs...), - LevenbergMarquardt(; - linsolve, precs, disable_geodesic = Val(true), autodiff, kwargs...), - TrustRegion(; concrete_jac, linsolve, precs, autodiff, kwargs...), - GaussNewton(; concrete_jac, linsolve, precs, - linesearch = BackTracking(), autodiff, kwargs...), - TrustRegion(; concrete_jac, linsolve, precs, - radius_update_scheme = RadiusUpdateSchemes.Bastin, autodiff, kwargs...), - LevenbergMarquardt(; linsolve, precs, autodiff, kwargs...)) - end - return NonlinearSolvePolyAlgorithm(algs, Val(:NLLS)) -end - -## Defaults - ## TODO: In the long run we want to use an `Assumptions` API like LinearSolve to specify ## the conditioning of the Jacobian and such @@ -503,31 +10,43 @@ end ## the trouble of specifying a custom jacobian function, we should use algorithms that ## can use that! function SciMLBase.__init(prob::NonlinearProblem, ::Nothing, args...; kwargs...) - must_use_jacobian = Val(prob.f.jac !== nothing) - return SciMLBase.__init(prob, + must_use_jacobian = Val(SciMLBase.has_jac(prob.f)) + return SciMLBase.__init( + prob, FastShortcutNonlinearPolyalg( - eltype(prob.u0); must_use_jacobian, u0_len = length(prob.u0)), + eltype(prob.u0); must_use_jacobian, u0_len = length(prob.u0) + ), args...; - kwargs...) + kwargs... + ) end function SciMLBase.__solve(prob::NonlinearProblem, ::Nothing, args...; kwargs...) - must_use_jacobian = Val(prob.f.jac !== nothing) - prefer_simplenonlinearsolve = Val(prob.u0 isa SArray) - return SciMLBase.__solve(prob, - FastShortcutNonlinearPolyalg(eltype(prob.u0); must_use_jacobian, - prefer_simplenonlinearsolve, u0_len = length(prob.u0)), + must_use_jacobian = Val(SciMLBase.has_jac(prob.f)) + prefer_simplenonlinearsolve = Val(prob.u0 isa StaticArray) + return SciMLBase.__solve( + prob, + FastShortcutNonlinearPolyalg( + eltype(prob.u0); + must_use_jacobian, + prefer_simplenonlinearsolve, + u0_len = length(prob.u0) + ), args...; - kwargs...) + kwargs... + ) end function SciMLBase.__init(prob::NonlinearLeastSquaresProblem, ::Nothing, args...; kwargs...) return SciMLBase.__init( - prob, FastShortcutNLLSPolyalg(eltype(prob.u0)), args...; kwargs...) + prob, FastShortcutNLLSPolyalg(eltype(prob.u0)), args...; kwargs... + ) end function SciMLBase.__solve( - prob::NonlinearLeastSquaresProblem, ::Nothing, args...; kwargs...) + prob::NonlinearLeastSquaresProblem, ::Nothing, args...; kwargs... +) return SciMLBase.__solve( - prob, FastShortcutNLLSPolyalg(eltype(prob.u0)), args...; kwargs...) + prob, FastShortcutNLLSPolyalg(eltype(prob.u0)), args...; kwargs... + ) end diff --git a/src/descent/damped_newton.jl b/src/descent/damped_newton.jl deleted file mode 100644 index ba3e1d028..000000000 --- a/src/descent/damped_newton.jl +++ /dev/null @@ -1,256 +0,0 @@ -""" - DampedNewtonDescent(; linsolve = nothing, precs = DEFAULT_PRECS, initial_damping, - damping_fn) - -A Newton descent algorithm with damping. The damping factor is computed using the -`damping_fn` function. The descent direction is computed as ``(JᵀJ + λDᵀD) δu = -fu``. For -non-square Jacobians, we default to solving for `Jδx = -fu` and `√λ⋅D δx = 0` -simultaneously. If the linear solver can't handle non-square matrices, we use the normal -form equations ``(JᵀJ + λDᵀD) δu = Jᵀ fu``. Note that this factorization is often the faster -choice, but it is not as numerically stable as the least squares solver. - -The damping factor returned must be a non-negative number. - -### Keyword Arguments - - - `initial_damping`: The initial damping factor to use - - `damping_fn`: The function to use to compute the damping factor. This must satisfy the - [`NonlinearSolve.AbstractDampingFunction`](@ref) interface. -""" -@kwdef @concrete struct DampedNewtonDescent <: AbstractDescentAlgorithm - linsolve = nothing - precs = DEFAULT_PRECS - initial_damping - damping_fn -end - -function Base.show(io::IO, d::DampedNewtonDescent) - modifiers = String[] - d.linsolve !== nothing && push!(modifiers, "linsolve = $(d.linsolve)") - d.precs !== DEFAULT_PRECS && push!(modifiers, "precs = $(d.precs)") - push!(modifiers, "initial_damping = $(d.initial_damping)") - push!(modifiers, "damping_fn = $(d.damping_fn)") - print(io, "DampedNewtonDescent($(join(modifiers, ", ")))") -end - -supports_line_search(::DampedNewtonDescent) = true -supports_trust_region(::DampedNewtonDescent) = true - -@concrete mutable struct DampedNewtonDescentCache{pre_inverted, mode} <: - AbstractDescentCache - J - δu - δus - lincache - JᵀJ_cache - Jᵀfu_cache - rhs_cache - damping_fn_cache - timer -end - -@internal_caches DampedNewtonDescentCache :lincache :damping_fn_cache - -function __internal_init(prob::AbstractNonlinearProblem, alg::DampedNewtonDescent, J, fu, - u; stats, pre_inverted::Val{INV} = False, linsolve_kwargs = (;), - abstol = nothing, timer = get_timer_output(), reltol = nothing, - alias_J = true, shared::Val{N} = Val(1), kwargs...) where {INV, N} - length(fu) != length(u) && - @assert !INV "Precomputed Inverse for Non-Square Jacobian doesn't make sense." - @bb δu = similar(u) - δus = N ≤ 1 ? nothing : map(2:N) do i - @bb δu_ = similar(u) - end - - normal_form_damping = returns_norm_form_damping(alg.damping_fn) - normal_form_linsolve = __needs_square_A(alg.linsolve, u) - if u isa Number - mode = :simple - elseif prob isa NonlinearProblem - mode = ifelse(!normal_form_damping, :simple, - ifelse(normal_form_linsolve, :normal_form, :least_squares)) - else - if normal_form_linsolve & !normal_form_damping - throw(ArgumentError("Linear Solver expects Normal Form but returned Damping is \ - not Normal Form. This is not supported.")) - end - mode = ifelse(normal_form_damping & !normal_form_linsolve, :least_squares, - ifelse(!normal_form_damping & !normal_form_linsolve, :simple, :normal_form)) - end - - if mode === :least_squares - if requires_normal_form_jacobian(alg.damping_fn) - JᵀJ = transpose(J) * J # Needed to compute the damping factor - jac_damp = JᵀJ - else - JᵀJ = nothing - jac_damp = J - end - if requires_normal_form_rhs(alg.damping_fn) - Jᵀfu = transpose(J) * _vec(fu) - rhs_damp = Jᵀfu - else - Jᵀfu = nothing - rhs_damp = fu - end - damping_fn_cache = __internal_init(prob, alg.damping_fn, alg.initial_damping, - jac_damp, rhs_damp, u, False; stats, kwargs...) - D = damping_fn_cache(nothing) - D isa Number && (D = D * I) - rhs_cache = vcat(_vec(fu), _vec(u)) - J_cache = _vcat(J, D) - A, b = J_cache, rhs_cache - elseif mode === :simple - damping_fn_cache = __internal_init( - prob, alg.damping_fn, alg.initial_damping, J, fu, u, False; kwargs...) - J_cache = __maybe_unaliased(J, alias_J) - D = damping_fn_cache(nothing) - J_damped = __dampen_jacobian!!(J_cache, J, D) - J_cache = J_damped - A, b = J_damped, _vec(fu) - JᵀJ, Jᵀfu, rhs_cache = nothing, nothing, nothing - elseif mode === :normal_form - JᵀJ = transpose(J) * J - Jᵀfu = transpose(J) * _vec(fu) - jac_damp = requires_normal_form_jacobian(alg.damping_fn) ? JᵀJ : J - rhs_damp = requires_normal_form_rhs(alg.damping_fn) ? Jᵀfu : fu - damping_fn_cache = __internal_init(prob, alg.damping_fn, alg.initial_damping, - jac_damp, rhs_damp, u, True; stats, kwargs...) - D = damping_fn_cache(nothing) - @bb J_cache = similar(JᵀJ) - @bb @. J_cache = 0 - J_damped = __dampen_jacobian!!(J_cache, JᵀJ, D) - A, b = __maybe_symmetric(J_damped), _vec(Jᵀfu) - rhs_cache = nothing - end - - lincache = LinearSolverCache( - alg, alg.linsolve, A, b, _vec(u); stats, abstol, reltol, linsolve_kwargs...) - - return DampedNewtonDescentCache{INV, mode}( - J_cache, δu, δus, lincache, JᵀJ, Jᵀfu, rhs_cache, damping_fn_cache, timer) -end - -function __internal_solve!(cache::DampedNewtonDescentCache{INV, mode}, J, fu, - u, idx::Val{N} = Val(1); skip_solve::Bool = false, - new_jacobian::Bool = true, kwargs...) where {INV, N, mode} - δu = get_du(cache, idx) - skip_solve && return DescentResult(; δu) - - recompute_A = idx === Val(1) - - @static_timeit cache.timer "dampen" begin - if mode === :least_squares - if (J !== nothing || new_jacobian) && recompute_A - INV && (J = inv(J)) - if requires_normal_form_jacobian(cache.damping_fn_cache) - @bb cache.JᵀJ_cache = transpose(J) × J - jac_damp = cache.JᵀJ_cache - else - jac_damp = J - end - if requires_normal_form_rhs(cache.damping_fn_cache) - @bb cache.Jᵀfu_cache = transpose(J) × fu - rhs_damp = cache.Jᵀfu_cache - else - rhs_damp = fu - end - D = __internal_solve!(cache.damping_fn_cache, jac_damp, rhs_damp, False) - if __can_setindex(cache.J) - copyto!(@view(cache.J[1:size(J, 1), :]), J) - cache.J[(size(J, 1) + 1):end, :] .= sqrt.(D) - else - cache.J = _vcat(J, sqrt.(D)) - end - end - A = cache.J - if __can_setindex(cache.rhs_cache) - cache.rhs_cache[1:length(fu)] .= _vec(fu) - cache.rhs_cache[(length(fu) + 1):end] .= false - else - cache.rhs_cache = vcat(_vec(fu), zero(_vec(u))) - end - b = cache.rhs_cache - elseif mode === :simple - if (J !== nothing || new_jacobian) && recompute_A - INV && (J = inv(J)) - D = __internal_solve!(cache.damping_fn_cache, J, fu, False) - cache.J = __dampen_jacobian!!(cache.J, J, D) - end - A, b = cache.J, _vec(fu) - elseif mode === :normal_form - if (J !== nothing || new_jacobian) && recompute_A - INV && (J = inv(J)) - @bb cache.JᵀJ_cache = transpose(J) × J - @bb cache.Jᵀfu_cache = transpose(J) × vec(fu) - D = __internal_solve!( - cache.damping_fn_cache, cache.JᵀJ_cache, cache.Jᵀfu_cache, True) - cache.J = __dampen_jacobian!!(cache.J, cache.JᵀJ_cache, D) - A = __maybe_symmetric(cache.J) - elseif !recompute_A - @bb cache.Jᵀfu_cache = transpose(J) × vec(fu) - A = __maybe_symmetric(cache.J) - else - A = nothing - end - b = _vec(cache.Jᵀfu_cache) - else - error("Unknown mode: $(mode)") - end - end - - @static_timeit cache.timer "linear solve" begin - linres = cache.lincache(; - A, b, reuse_A_if_factorization = !new_jacobian && !recompute_A, - kwargs..., linu = _vec(δu)) - δu = _restructure(get_du(cache, idx), linres.u) - if !linres.success - set_du!(cache, δu, idx) - return DescentResult(; δu, success = false, linsolve_success = false) - end - end - - @bb @. δu *= -1 - set_du!(cache, δu, idx) - return DescentResult(; δu) -end - -# Define special concatenation for certain Array combinations -@inline _vcat(x, y) = vcat(x, y) - -# J_cache is allowed to alias J -## Compute ``J + D`` -@inline __dampen_jacobian!!(J_cache, J::AbstractSciMLOperator, D) = J + D -@inline __dampen_jacobian!!(J_cache, J::Number, D) = J + D -@inline function __dampen_jacobian!!(J_cache, J::AbstractMatrix, D::AbstractMatrix) - if __can_setindex(J_cache) - copyto!(J_cache, J) - if fast_scalar_indexing(J_cache) - @simd for i in axes(J_cache, 1) - @inbounds J_cache[i, i] += D[i, i] - end - else - idxs = diagind(J_cache) - J_cache[idxs] .= @view(J[idxs]) .+ @view(D[idxs]) - end - return J_cache - else - return @. J + D - end -end -@inline function __dampen_jacobian!!(J_cache, J::AbstractMatrix, D::Number) - if __can_setindex(J_cache) - copyto!(J_cache, J) - if fast_scalar_indexing(J_cache) - @simd for i in axes(J_cache, 1) - @inbounds J_cache[i, i] += D - end - else - idxs = diagind(J_cache) - J_cache[idxs] .= @view(J[idxs]) .+ D - end - return J_cache - else - return @. J + D - end -end diff --git a/src/descent/newton.jl b/src/descent/newton.jl deleted file mode 100644 index 2fe7abf9a..000000000 --- a/src/descent/newton.jl +++ /dev/null @@ -1,122 +0,0 @@ -""" - NewtonDescent(; linsolve = nothing, precs = DEFAULT_PRECS) - -Compute the descent direction as ``J δu = -fu``. For non-square Jacobian problems, this is -commonly referred to as the Gauss-Newton Descent. - -See also [`Dogleg`](@ref), [`SteepestDescent`](@ref), [`DampedNewtonDescent`](@ref). -""" -@kwdef @concrete struct NewtonDescent <: AbstractDescentAlgorithm - linsolve = nothing - precs = DEFAULT_PRECS -end - -function Base.show(io::IO, d::NewtonDescent) - modifiers = String[] - d.linsolve !== nothing && push!(modifiers, "linsolve = $(d.linsolve)") - d.precs !== DEFAULT_PRECS && push!(modifiers, "precs = $(d.precs)") - print(io, "NewtonDescent($(join(modifiers, ", ")))") -end - -supports_line_search(::NewtonDescent) = true - -@concrete mutable struct NewtonDescentCache{pre_inverted, normalform} <: - AbstractDescentCache - δu - δus - lincache - JᵀJ_cache # For normal form else nothing - Jᵀfu_cache - timer -end - -@internal_caches NewtonDescentCache :lincache - -function __internal_init(prob::NonlinearProblem, alg::NewtonDescent, J, fu, u; stats, - shared::Val{N} = Val(1), pre_inverted::Val{INV} = False, - linsolve_kwargs = (;), abstol = nothing, reltol = nothing, - timer = get_timer_output(), kwargs...) where {INV, N} - @bb δu = similar(u) - δus = N ≤ 1 ? nothing : map(2:N) do i - @bb δu_ = similar(u) - end - INV && return NewtonDescentCache{true, false}(δu, δus, nothing, nothing, nothing, timer) - lincache = LinearSolverCache( - alg, alg.linsolve, J, _vec(fu), _vec(u); stats, abstol, reltol, linsolve_kwargs...) - return NewtonDescentCache{false, false}(δu, δus, lincache, nothing, nothing, timer) -end - -function __internal_init(prob::NonlinearLeastSquaresProblem, alg::NewtonDescent, J, fu, - u; stats, pre_inverted::Val{INV} = False, linsolve_kwargs = (;), - shared::Val{N} = Val(1), abstol = nothing, reltol = nothing, - timer = get_timer_output(), kwargs...) where {INV, N} - length(fu) != length(u) && - @assert !INV "Precomputed Inverse for Non-Square Jacobian doesn't make sense." - - normal_form = __needs_square_A(alg.linsolve, u) - if normal_form - JᵀJ = transpose(J) * J - Jᵀfu = transpose(J) * _vec(fu) - A, b = __maybe_symmetric(JᵀJ), Jᵀfu - else - JᵀJ, Jᵀfu = nothing, nothing - A, b = J, _vec(fu) - end - lincache = LinearSolverCache( - alg, alg.linsolve, A, b, _vec(u); stats, abstol, reltol, linsolve_kwargs...) - @bb δu = similar(u) - δus = N ≤ 1 ? nothing : map(2:N) do i - @bb δu_ = similar(u) - end - return NewtonDescentCache{false, normal_form}(δu, δus, lincache, JᵀJ, Jᵀfu, timer) -end - -function __internal_solve!( - cache::NewtonDescentCache{INV, false}, J, fu, u, idx::Val = Val(1); - skip_solve::Bool = false, new_jacobian::Bool = true, kwargs...) where {INV} - δu = get_du(cache, idx) - skip_solve && return DescentResult(; δu) - if INV - @assert J!==nothing "`J` must be provided when `pre_inverted = Val(true)`." - @bb δu = J × vec(fu) - else - @static_timeit cache.timer "linear solve" begin - linres = cache.lincache(; - A = J, b = _vec(fu), kwargs..., linu = _vec(δu), du = _vec(δu), - reuse_A_if_factorization = !new_jacobian || (idx !== Val(1))) - δu = _restructure(get_du(cache, idx), linres.u) - if !linres.success - set_du!(cache, δu, idx) - return DescentResult(; δu, success = false, linsolve_success = false) - end - end - end - @bb @. δu *= -1 - set_du!(cache, δu, idx) - return DescentResult(; δu) -end - -function __internal_solve!( - cache::NewtonDescentCache{false, true}, J, fu, u, idx::Val = Val(1); - skip_solve::Bool = false, new_jacobian::Bool = true, kwargs...) - δu = get_du(cache, idx) - skip_solve && return DescentResult(; δu) - if idx === Val(1) - @bb cache.JᵀJ_cache = transpose(J) × J - end - @bb cache.Jᵀfu_cache = transpose(J) × vec(fu) - @static_timeit cache.timer "linear solve" begin - linres = cache.lincache(; - A = __maybe_symmetric(cache.JᵀJ_cache), b = cache.Jᵀfu_cache, - kwargs..., linu = _vec(δu), du = _vec(δu), - reuse_A_if_factorization = !new_jacobian || (idx !== Val(1))) - δu = _restructure(get_du(cache, idx), linres.u) - if !linres.success - set_du!(cache, δu, idx) - return DescentResult(; δu, success = false, linsolve_success = false) - end - end - @bb @. δu *= -1 - set_du!(cache, δu, idx) - return DescentResult(; δu) -end diff --git a/src/descent/steepest.jl b/src/descent/steepest.jl deleted file mode 100644 index cc2f4d128..000000000 --- a/src/descent/steepest.jl +++ /dev/null @@ -1,73 +0,0 @@ -""" - SteepestDescent(; linsolve = nothing, precs = DEFAULT_PRECS) - -Compute the descent direction as ``δu = -Jᵀfu``. The linear solver and preconditioner are -only used if `J` is provided in the inverted form. - -See also [`Dogleg`](@ref), [`NewtonDescent`](@ref), [`DampedNewtonDescent`](@ref). -""" -@kwdef @concrete struct SteepestDescent <: AbstractDescentAlgorithm - linsolve = nothing - precs = DEFAULT_PRECS -end - -function Base.show(io::IO, d::SteepestDescent) - modifiers = String[] - d.linsolve !== nothing && push!(modifiers, "linsolve = $(d.linsolve)") - d.precs !== DEFAULT_PRECS && push!(modifiers, "precs = $(d.precs)") - print(io, "SteepestDescent($(join(modifiers, ", ")))") -end - -supports_line_search(::SteepestDescent) = true - -@concrete mutable struct SteepestDescentCache{pre_inverted} <: AbstractDescentCache - δu - δus - lincache - timer -end - -@internal_caches SteepestDescentCache :lincache - -@inline function __internal_init( - prob::AbstractNonlinearProblem, alg::SteepestDescent, J, fu, u; - stats, shared::Val{N} = Val(1), pre_inverted::Val{INV} = False, - linsolve_kwargs = (;), abstol = nothing, reltol = nothing, - timer = get_timer_output(), kwargs...) where {INV, N} - INV && @assert length(fu)==length(u) "Non-Square Jacobian Inverse doesn't make sense." - @bb δu = similar(u) - δus = N ≤ 1 ? nothing : map(2:N) do i - @bb δu_ = similar(u) - end - if INV - lincache = LinearSolverCache(alg, alg.linsolve, transpose(J), _vec(fu), _vec(u); - stats, abstol, reltol, linsolve_kwargs...) - else - lincache = nothing - end - return SteepestDescentCache{INV}(δu, δus, lincache, timer) -end - -function __internal_solve!(cache::SteepestDescentCache{INV}, J, fu, u, idx::Val = Val(1); - new_jacobian::Bool = true, kwargs...) where {INV} - δu = get_du(cache, idx) - if INV - A = J === nothing ? nothing : transpose(J) - @static_timeit cache.timer "linear solve" begin - linres = cache.lincache(; - A, b = _vec(fu), kwargs..., linu = _vec(δu), du = _vec(δu), - reuse_A_if_factorization = !new_jacobian || idx !== Val(1)) - δu = _restructure(get_du(cache, idx), linres.u) - if !linres.success - set_du!(cache, δu, idx) - return DescentResult(; δu, success = false, linsolve_success = false) - end - end - else - @assert J!==nothing "`J` must be provided when `pre_inverted = Val(false)`." - @bb δu = transpose(J) × vec(fu) - end - @bb @. δu *= -1 - set_du!(cache, δu, idx) - return DescentResult(; δu) -end diff --git a/src/algorithms/extension_algs.jl b/src/extension_algs.jl similarity index 79% rename from src/algorithms/extension_algs.jl rename to src/extension_algs.jl index fda61b693..1fad79839 100644 --- a/src/algorithms/extension_algs.jl +++ b/src/extension_algs.jl @@ -22,8 +22,11 @@ for solving `NonlinearLeastSquaresProblem`. This algorithm is only available if `LeastSquaresOptim.jl` is installed and loaded. """ -struct LeastSquaresOptimJL{alg, linsolve} <: AbstractNonlinearSolveExtensionAlgorithm +@concrete struct LeastSquaresOptimJL <: AbstractNonlinearSolveAlgorithm autodiff + alg::Symbol + linsolve <: Union{Symbol, Nothing} + name::Symbol end function LeastSquaresOptimJL(alg = :lm; linsolve = nothing, autodiff = :central) @@ -32,27 +35,24 @@ function LeastSquaresOptimJL(alg = :lm; linsolve = nothing, autodiff = :central) autodiff isa Symbol && @assert autodiff in (:central, :forward) if Base.get_extension(@__MODULE__, :NonlinearSolveLeastSquaresOptimExt) === nothing - error("LeastSquaresOptimJL requires LeastSquaresOptim.jl to be loaded") + error("`LeastSquaresOptimJL` requires `LeastSquaresOptim.jl` to be loaded") end - return LeastSquaresOptimJL{alg, linsolve}(autodiff) + return LeastSquaresOptimJL(autodiff, alg, linsolve, :LeastSquaresOptimJL) end """ - FastLevenbergMarquardtJL(linsolve::Symbol = :cholesky; factor = 1e-6, - factoraccept = 13.0, factorreject = 3.0, factorupdate = :marquardt, + FastLevenbergMarquardtJL( + linsolve::Symbol = :cholesky; + factor = 1e-6, factoraccept = 13.0, factorreject = 3.0, factorupdate = :marquardt, minscale = 1e-12, maxscale = 1e16, minfactor = 1e-28, maxfactor = 1e32, - autodiff = nothing) + autodiff = nothing + ) Wrapper over [FastLevenbergMarquardt.jl](https://github.com/kamesy/FastLevenbergMarquardt.jl) for solving `NonlinearLeastSquaresProblem`. For details about the other keyword arguments see the documentation for `FastLevenbergMarquardt.jl`. -!!! warning - - This is not really the fastest solver. It is called that since the original package - is called "Fast". `LevenbergMarquardt()` is almost always a better choice. - ### Arguments - `linsolve`: Linear solver to use. Can be `:qr` or `:cholesky`. @@ -67,9 +67,9 @@ see the documentation for `FastLevenbergMarquardt.jl`. This algorithm is only available if `FastLevenbergMarquardt.jl` is installed and loaded. """ -@concrete struct FastLevenbergMarquardtJL{linsolve} <: - AbstractNonlinearSolveExtensionAlgorithm +@concrete struct FastLevenbergMarquardtJL <: AbstractNonlinearSolveAlgorithm autodiff + linsolve::Symbol factor factoraccept factorreject @@ -78,21 +78,25 @@ see the documentation for `FastLevenbergMarquardt.jl`. maxscale minfactor maxfactor + name::Symbol end function FastLevenbergMarquardtJL( linsolve::Symbol = :cholesky; factor = 1e-6, factoraccept = 13.0, factorreject = 3.0, factorupdate = :marquardt, minscale = 1e-12, - maxscale = 1e16, minfactor = 1e-28, maxfactor = 1e32, autodiff = nothing) + maxscale = 1e16, minfactor = 1e-28, maxfactor = 1e32, autodiff = nothing +) @assert linsolve in (:qr, :cholesky) @assert factorupdate in (:marquardt, :nielson) if Base.get_extension(@__MODULE__, :NonlinearSolveFastLevenbergMarquardtExt) === nothing - error("FastLevenbergMarquardtJL requires FastLevenbergMarquardt.jl to be loaded") + error("`FastLevenbergMarquardtJL` requires `FastLevenbergMarquardt.jl` to be loaded") end - return FastLevenbergMarquardtJL{linsolve}(autodiff, factor, factoraccept, factorreject, - factorupdate, minscale, maxscale, minfactor, maxfactor) + return FastLevenbergMarquardtJL( + autodiff, linsolve, factor, factoraccept, factorreject, + factorupdate, minscale, maxscale, minfactor, maxfactor, :FastLevenbergMarquardtJL + ) end """ @@ -141,22 +145,25 @@ NonlinearLeastSquaresProblem. This algorithm is only available if `MINPACK.jl` is installed and loaded. """ -@concrete struct CMINPACK <: AbstractNonlinearSolveExtensionAlgorithm +@concrete struct CMINPACK <: AbstractNonlinearSolveAlgorithm method::Symbol autodiff + name::Symbol end function CMINPACK(; method::Symbol = :auto, autodiff = missing) if Base.get_extension(@__MODULE__, :NonlinearSolveMINPACKExt) === nothing - error("CMINPACK requires MINPACK.jl to be loaded") + error("`CMINPACK` requires `MINPACK.jl` to be loaded") end - return CMINPACK(method, autodiff) + return CMINPACK(method, autodiff, :CMINPACK) end """ - NLsolveJL(; method = :trust_region, autodiff = :central, linesearch = Static(), + NLsolveJL(; + method = :trust_region, autodiff = :central, linesearch = Static(), linsolve = (x, A, b) -> copyto!(x, A\\b), factor = one(Float64), autoscale = true, - m = 10, beta = one(Float64)) + m = 10, beta = one(Float64) + ) ### Keyword Arguments @@ -201,7 +208,7 @@ For more information on these arguments, consult the This algorithm is only available if `NLsolve.jl` is installed and loaded. """ -@concrete struct NLsolveJL <: AbstractNonlinearSolveExtensionAlgorithm +@concrete struct NLsolveJL <: AbstractNonlinearSolveAlgorithm method::Symbol autodiff linesearch @@ -210,20 +217,24 @@ For more information on these arguments, consult the autoscale::Bool m::Int beta + name::Symbol end -function NLsolveJL(; method = :trust_region, autodiff = :central, linesearch = missing, - linsolve = (x, A, b) -> copyto!(x, A \ b), factor = 1.0, - autoscale = true, m = 10, beta = one(Float64)) +function NLsolveJL(; + method = :trust_region, autodiff = :central, linesearch = missing, beta = 1.0, + linsolve = (x, A, b) -> copyto!(x, A \ b), factor = 1.0, autoscale = true, m = 10 +) if Base.get_extension(@__MODULE__, :NonlinearSolveNLsolveExt) === nothing - error("NLsolveJL requires NLsolve.jl to be loaded") + error("`NLsolveJL` requires `NLsolve.jl` to be loaded") end if autodiff isa Symbol && autodiff !== :central && autodiff !== :forward error("`autodiff` must be `:central` or `:forward`.") end - return NLsolveJL(method, autodiff, linesearch, linsolve, factor, autoscale, m, beta) + return NLsolveJL( + method, autodiff, linesearch, linsolve, factor, autoscale, m, beta, :NLsolveJL + ) end """ @@ -242,16 +253,17 @@ jacobian function and supply it to the solver. any valid ADTypes.jl autodiff type (conditional on that backend being supported in NonlinearSolve.jl). """ -struct NLSolversJL{M, AD} <: AbstractNonlinearSolveExtensionAlgorithm +struct NLSolversJL{M, AD} <: AbstractNonlinearSolveAlgorithm method::M autodiff::AD + name::Symbol function NLSolversJL(method, autodiff) if Base.get_extension(@__MODULE__, :NonlinearSolveNLSolversExt) === nothing error("NLSolversJL requires NLSolvers.jl to be loaded") end - return new{typeof(method), typeof(autodiff)}(method, autodiff) + return new{typeof(method), typeof(autodiff)}(method, autodiff, :NLSolversJL) end end @@ -259,8 +271,10 @@ NLSolversJL(method; autodiff = nothing) = NLSolversJL(method, autodiff) NLSolversJL(; method, autodiff = nothing) = NLSolversJL(method, autodiff) """ - SpeedMappingJL(; σ_min = 0.0, stabilize::Bool = false, check_obj::Bool = false, - orders::Vector{Int} = [3, 3, 2], time_limit::Real = 1000) + SpeedMappingJL(; + σ_min = 0.0, stabilize::Bool = false, check_obj::Bool = false, + orders::Vector{Int} = [3, 3, 2] + ) Wrapper over [SpeedMapping.jl](https://nicolasl-s.github.io/SpeedMapping.jl/) for solving Fixed Point Problems. We allow using this algorithm to solve root finding problems as well. @@ -277,32 +291,35 @@ Fixed Point Problems. We allow using this algorithm to solve root finding proble means no extrapolation). The two recommended orders are `[3, 2]` and `[3, 3, 2]`, the latter being potentially better for highly non-linear applications (see [lepage2021alternating](@cite)). - - `time_limit`: time limit for the algorithm. !!! note This algorithm is only available if `SpeedMapping.jl` is installed and loaded. """ -@concrete struct SpeedMappingJL <: AbstractNonlinearSolveExtensionAlgorithm +@concrete struct SpeedMappingJL <: AbstractNonlinearSolveAlgorithm σ_min stabilize::Bool check_obj::Bool orders::Vector{Int} + name::Symbol end -function SpeedMappingJL(; σ_min = 0.0, stabilize::Bool = false, check_obj::Bool = false, - orders::Vector{Int} = [3, 3, 2]) +function SpeedMappingJL(; + σ_min = 0.0, stabilize::Bool = false, check_obj::Bool = false, + orders::Vector{Int} = [3, 3, 2] +) if Base.get_extension(@__MODULE__, :NonlinearSolveSpeedMappingExt) === nothing - error("SpeedMappingJL requires SpeedMapping.jl to be loaded") + error("`SpeedMappingJL` requires `SpeedMapping.jl` to be loaded") end - return SpeedMappingJL(σ_min, stabilize, check_obj, orders) + return SpeedMappingJL(σ_min, stabilize, check_obj, orders, :SpeedMappingJL) end """ - FixedPointAccelerationJL(; algorithm = :Anderson, m = missing, - condition_number_threshold = missing, extrapolation_period = missing, - replace_invalids = :NoAction) + FixedPointAccelerationJL(; + algorithm = :Anderson, m = missing, condition_number_threshold = missing, + extrapolation_period = missing, replace_invalids = :NoAction + ) Wrapper over [FixedPointAcceleration.jl](https://s-baumann.github.io/FixedPointAcceleration.jl/) for solving Fixed Point Problems. We allow using this algorithm to solve root finding @@ -326,60 +343,55 @@ problems as well. This algorithm is only available if `FixedPointAcceleration.jl` is installed and loaded. """ -@concrete struct FixedPointAccelerationJL <: AbstractNonlinearSolveExtensionAlgorithm +@concrete struct FixedPointAccelerationJL <: AbstractNonlinearSolveAlgorithm algorithm::Symbol extrapolation_period::Int replace_invalids::Symbol dampening m::Int condition_number_threshold + name::Symbol end function FixedPointAccelerationJL(; algorithm = :Anderson, m = missing, condition_number_threshold = missing, - extrapolation_period = missing, replace_invalids = :NoAction, dampening = 1.0) + extrapolation_period = missing, replace_invalids = :NoAction, dampening = 1.0 +) if Base.get_extension(@__MODULE__, :NonlinearSolveFixedPointAccelerationExt) === nothing - error("FixedPointAccelerationJL requires FixedPointAcceleration.jl to be loaded") + error("`FixedPointAccelerationJL` requires `FixedPointAcceleration.jl` to be loaded") end @assert algorithm in (:Anderson, :MPE, :RRE, :VEA, :SEA, :Simple, :Aitken, :Newton) @assert replace_invalids in (:ReplaceInvalids, :ReplaceVector, :NoAction) if algorithm !== :Anderson - if condition_number_threshold !== missing - error("`condition_number_threshold` is only valid for Anderson acceleration") - end - if m !== missing - error("`m` is only valid for Anderson acceleration") - end + @assert condition_number_threshold===missing "`condition_number_threshold` is only valid for Anderson acceleration" + @assert m===missing "`m` is only valid for Anderson acceleration" end condition_number_threshold === missing && (condition_number_threshold = 1e3) m === missing && (m = 10) if algorithm !== :MPE && algorithm !== :RRE && algorithm !== :VEA && algorithm !== :SEA - if extrapolation_period !== missing - error("`extrapolation_period` is only valid for MPE, RRE, VEA and SEA") - end + @assert extrapolation_period===missing "`extrapolation_period` is only valid for MPE, RRE, VEA and SEA" end if extrapolation_period === missing - if algorithm === :SEA || algorithm === :VEA - extrapolation_period = 6 - else - extrapolation_period = 7 - end + extrapolation_period = algorithm === :SEA || algorithm === :VEA ? 6 : 7 else if (algorithm === :SEA || algorithm === :VEA) && extrapolation_period % 2 != 0 - error("`extrapolation_period` must be multiples of 2 for SEA and VEA") + throw(AssertionError("`extrapolation_period` must be multiples of 2 for SEA and VEA")) end end - return FixedPointAccelerationJL(algorithm, extrapolation_period, replace_invalids, - dampening, m, condition_number_threshold) + return FixedPointAccelerationJL( + algorithm, extrapolation_period, replace_invalids, + dampening, m, condition_number_threshold, :FixedPointAccelerationJL + ) end """ - SIAMFANLEquationsJL(; method = :newton, delta = 1e-3, linsolve = nothing, - autodiff = missing) + SIAMFANLEquationsJL(; + method = :newton, delta = 1e-3, linsolve = nothing, autodiff = missing + ) ### Keyword Arguments @@ -404,22 +416,26 @@ end This algorithm is only available if `SIAMFANLEquations.jl` is installed and loaded. """ -@concrete struct SIAMFANLEquationsJL{L <: Union{Symbol, Nothing}} <: - AbstractNonlinearSolveExtensionAlgorithm +@concrete struct SIAMFANLEquationsJL <: AbstractNonlinearSolveAlgorithm method::Symbol delta - linsolve::L + linsolve <: Union{Symbol, Nothing} m::Int beta autodiff + name::Symbol end -function SIAMFANLEquationsJL(; method = :newton, delta = 1e-3, linsolve = nothing, - m = 0, beta = 1.0, autodiff = missing) +function SIAMFANLEquationsJL(; + method = :newton, delta = 1e-3, linsolve = nothing, m = 0, beta = 1.0, + autodiff = missing +) if Base.get_extension(@__MODULE__, :NonlinearSolveSIAMFANLEquationsExt) === nothing - error("SIAMFANLEquationsJL requires SIAMFANLEquations.jl to be loaded") + error("`SIAMFANLEquationsJL` requires `SIAMFANLEquations.jl` to be loaded") end - return SIAMFANLEquationsJL(method, delta, linsolve, m, beta, autodiff) + return SIAMFANLEquationsJL( + method, delta, linsolve, m, beta, autodiff, :SIAMFANLEquationsJL + ) end """ @@ -459,16 +475,16 @@ These options are forwarded from `solve` to the PETSc SNES solver. If these are This algorithm is only available if `PETSc.jl` is installed and loaded. """ -@concrete struct PETScSNES <: AbstractNonlinearSolveExtensionAlgorithm +@concrete struct PETScSNES <: AbstractNonlinearSolveAlgorithm petsclib mpi_comm - autodiff <: Union{Missing, Nothing, ADTypes.AbstractADType} + autodiff snes_options end function PETScSNES(; petsclib = missing, autodiff = nothing, mpi_comm = missing, kwargs...) if Base.get_extension(@__MODULE__, :NonlinearSolvePETScExt) === nothing - error("PETScSNES requires PETSc.jl to be loaded") + error("`PETScSNES` requires `PETSc.jl` to be loaded") end return PETScSNES(petsclib, mpi_comm, autodiff, kwargs) end diff --git a/src/forward_diff.jl b/src/forward_diff.jl new file mode 100644 index 000000000..410e818c3 --- /dev/null +++ b/src/forward_diff.jl @@ -0,0 +1,98 @@ +const DualNonlinearProblem = NonlinearProblem{ + <:Union{Number, <:AbstractArray}, iip, + <:Union{<:Dual{T, V, P}, <:AbstractArray{<:Dual{T, V, P}}} +} where {iip, T, V, P} +const DualNonlinearLeastSquaresProblem = NonlinearLeastSquaresProblem{ + <:Union{Number, <:AbstractArray}, iip, + <:Union{<:Dual{T, V, P}, <:AbstractArray{<:Dual{T, V, P}}} +} where {iip, T, V, P} +const DualAbstractNonlinearProblem = Union{ + DualNonlinearProblem, DualNonlinearLeastSquaresProblem +} + +for algType in ALL_SOLVER_TYPES + @eval function SciMLBase.__solve( + prob::DualAbstractNonlinearProblem, alg::$(algType), args...; kwargs... + ) + sol, partials = NonlinearSolveBase.nonlinearsolve_forwarddiff_solve( + prob, alg, args...; kwargs... + ) + dual_soln = NonlinearSolveBase.nonlinearsolve_dual_solution(sol.u, partials, prob.p) + return SciMLBase.build_solution( + prob, alg, dual_soln, sol.resid; sol.retcode, sol.stats, sol.original + ) + end +end + +@concrete mutable struct NonlinearSolveForwardDiffCache <: AbstractNonlinearSolveCache + cache + prob + alg + p + values_p + partials_p +end + +function InternalAPI.reinit!( + cache::NonlinearSolveForwardDiffCache, args...; + p = cache.p, u0 = NonlinearSolveBase.get_u(cache.cache), kwargs... +) + inner_cache = InternalAPI.reinit!( + cache.cache; p = nodual_value(p), u0 = nodual_value(u0), kwargs... + ) + cache.cache = inner_cache + cache.p = p + cache.values_p = nodual_value(p) + cache.partials_p = ForwardDiff.partials(p) + return cache +end + +for algType in ALL_SOLVER_TYPES + # XXX: Extend to DualNonlinearLeastSquaresProblem + @eval function SciMLBase.__init( + prob::DualNonlinearProblem, alg::$(algType), args...; kwargs... + ) + p = nodual_value(prob.p) + newprob = SciMLBase.remake(prob; u0 = nodual_value(prob.u0), p) + cache = init(newprob, alg, args...; kwargs...) + return NonlinearSolveForwardDiffCache( + cache, newprob, alg, prob.p, p, ForwardDiff.partials(prob.p) + ) + end +end + +function CommonSolve.solve!(cache::NonlinearSolveForwardDiffCache) + sol = solve!(cache.cache) + prob = cache.prob + + uu = sol.u + Jₚ = NonlinearSolveBase.nonlinearsolve_∂f_∂p(prob, prob.f, uu, cache.values_p) + Jᵤ = NonlinearSolveBase.nonlinearsolve_∂f_∂u(prob, prob.f, uu, cache.values_p) + + z_arr = -Jᵤ \ Jₚ + + sumfun = ((z, p),) -> map(zᵢ -> zᵢ * ForwardDiff.partials(p), z) + if cache.p isa Number + partials = sumfun((z_arr, cache.p)) + else + partials = sum(sumfun, zip(eachcol(z_arr), cache.p)) + end + + dual_soln = NonlinearSolveBase.nonlinearsolve_dual_solution(sol.u, partials, cache.p) + return SciMLBase.build_solution( + prob, cache.alg, dual_soln, sol.resid; sol.retcode, sol.stats, sol.original + ) +end + +nodual_value(x) = x +nodual_value(x::Dual) = ForwardDiff.value(x) +nodual_value(x::AbstractArray{<:Dual}) = map(ForwardDiff.value, x) + +""" + pickchunksize(x) = pickchunksize(length(x)) + pickchunksize(x::Int) + +Determine the chunk size for ForwardDiff and PolyesterForwardDiff based on the input length. +""" +@inline pickchunksize(x) = pickchunksize(length(x)) +@inline pickchunksize(x::Int) = ForwardDiff.pickchunksize(x) diff --git a/src/globalization/line_search.jl b/src/globalization/line_search.jl deleted file mode 100644 index 7549f1f9d..000000000 --- a/src/globalization/line_search.jl +++ /dev/null @@ -1,7 +0,0 @@ -function callback_into_cache!(topcache, cache::AbstractLineSearchCache, args...) - LineSearch.callback_into_cache!(cache, get_fu(topcache)) -end - -function reinit_cache!(cache::AbstractLineSearchCache, args...; kwargs...) - return SciMLBase.reinit!(cache, args...; kwargs...) -end diff --git a/src/internal/approximate_initialization.jl b/src/internal/approximate_initialization.jl deleted file mode 100644 index a8196c367..000000000 --- a/src/internal/approximate_initialization.jl +++ /dev/null @@ -1,280 +0,0 @@ -# Jacobian Structure -""" - DiagonalStructure() - -Preserves only the Diagonal of the Matrix. -""" -struct DiagonalStructure <: AbstractApproximateJacobianStructure end - -get_full_jacobian(cache, ::DiagonalStructure, J::Number) = J -get_full_jacobian(cache, ::DiagonalStructure, J) = Diagonal(_vec(J)) - -function (::DiagonalStructure)(J::AbstractMatrix; alias::Bool = false) - @assert size(J, 1)==size(J, 2) "Diagonal Jacobian Structure must be square!" - return diag(J) -end -(::DiagonalStructure)(J::AbstractVector; alias::Bool = false) = alias ? J : @bb(copy(J)) -(::DiagonalStructure)(J::Number; alias::Bool = false) = J - -(::DiagonalStructure)(::Number, J_new::Number) = J_new -function (::DiagonalStructure)(J::AbstractVector, J_new::AbstractMatrix) - if __can_setindex(J) - if fast_scalar_indexing(J) - @simd for i in eachindex(J) - @inbounds J[i] = J_new[i, i] - end - else - J .= @view(J_new[diagind(J_new)]) - end - return J - end - return diag(J_new) -end -function (st::DiagonalStructure)(J::AbstractArray, J_new::AbstractMatrix) - return _restructure(J, st(vec(J), J_new)) -end - -""" - FullStructure() - -Stores the full matrix. -""" -struct FullStructure <: AbstractApproximateJacobianStructure end - -stores_full_jacobian(::FullStructure) = true - -(::FullStructure)(J; alias::Bool = false) = alias ? J : @bb(copy(J)) - -function (::FullStructure)(J, J_new) - J === J_new && return J - @bb copyto!(J, J_new) - return J -end - -# Initialization Strategies -""" - IdentityInitialization(alpha, structure) - -Initialize the Jacobian to be an Identity Matrix scaled by `alpha` and maintain the -structure as specified by `structure`. -""" -@concrete struct IdentityInitialization <: AbstractJacobianInitialization - alpha - structure -end - -function __internal_init( - prob::AbstractNonlinearProblem, alg::IdentityInitialization, solver, f::F, - fu, u::Number, p; internalnorm::IN = L2_NORM, kwargs...) where {F, IN} - α = __initial_alpha(alg.alpha, u, fu, internalnorm) - return InitializedApproximateJacobianCache( - α, alg.structure, alg, nothing, true, internalnorm) -end -function __internal_init(prob::AbstractNonlinearProblem, alg::IdentityInitialization, - solver, f::F, fu::StaticArray, u::StaticArray, p; - internalnorm::IN = L2_NORM, kwargs...) where {IN, F} - α = __initial_alpha(alg.alpha, u, fu, internalnorm) - if alg.structure isa DiagonalStructure - @assert length(u)==length(fu) "Diagonal Jacobian Structure must be square!" - J = one.(_vec(fu)) .* α - else - T = promote_type(eltype(u), eltype(fu)) - if fu isa SArray - J_ = SArray{Tuple{prod(Size(fu)), prod(Size(u))}, T}(I * α) - else - J_ = MArray{Tuple{prod(Size(fu)), prod(Size(u))}, T}(I * α) - end - J = alg.structure(J_; alias = true) - end - return InitializedApproximateJacobianCache( - J, alg.structure, alg, nothing, true, internalnorm) -end -function __internal_init( - prob::AbstractNonlinearProblem, alg::IdentityInitialization, solver, - f::F, fu, u, p; internalnorm::IN = L2_NORM, kwargs...) where {F, IN} - α = __initial_alpha(alg.alpha, u, fu, internalnorm) - if alg.structure isa DiagonalStructure - @assert length(u)==length(fu) "Diagonal Jacobian Structure must be square!" - J = one.(_vec(fu)) .* α - else - J_ = __similar(fu, promote_type(eltype(fu), eltype(u)), length(fu), length(u)) - J = alg.structure(__make_identity!!(J_, α); alias = true) - end - return InitializedApproximateJacobianCache( - J, alg.structure, alg, nothing, true, internalnorm) -end - -@inline function __initial_alpha(α, u, fu, internalnorm::F) where {F} - return convert(promote_type(eltype(u), eltype(fu)), α) -end -@inline function __initial_alpha(::Nothing, u, fu, internalnorm::F) where {F} - fu_norm = internalnorm(fu) - return ifelse(fu_norm ≥ 1e-5, (2 * fu_norm) / max(norm(u), true), - __initial_alpha(true, u, fu, internalnorm)) -end - -@inline __make_identity!!(A::Number, α) = one(A) * α -@inline __make_identity!!(A::AbstractVector, α) = __can_setindex(A) ? (A .= α) : - (one.(A) .* α) -@inline function __make_identity!!(A::AbstractMatrix{T}, α) where {T} - if A isa SMatrix - Sz = Size(A) - return SArray{Tuple{Sz[1], Sz[2]}, eltype(Sz)}(I * α) - end - @assert __can_setindex(A) "__make_identity!!(::AbstractMatrix) only works on mutable arrays!" - fill!(A, false) - if fast_scalar_indexing(A) - @inbounds for i in axes(A, 1) - A[i, i] = α - end - else - A[diagind(A)] .= α - end - return A -end - -""" - TrueJacobianInitialization(structure, autodiff) - -Initialize the Jacobian to be the true Jacobian and maintain the structure as specified -by `structure`. `autodiff` is used to compute the true Jacobian and if not specified we -make a selection automatically. -""" -@concrete struct TrueJacobianInitialization <: AbstractJacobianInitialization - structure - autodiff -end - -function __internal_init(prob::AbstractNonlinearProblem, alg::TrueJacobianInitialization, - solver, f::F, fu, u, p; stats, linsolve = missing, - internalnorm::IN = L2_NORM, kwargs...) where {F, IN} - autodiff = select_jacobian_autodiff(prob, alg.autodiff) - jac_cache = JacobianCache(prob, solver, prob.f, fu, u, p; stats, autodiff, linsolve) - J = alg.structure(jac_cache(nothing)) - return InitializedApproximateJacobianCache( - J, alg.structure, alg, jac_cache, false, internalnorm) -end - -""" - InitializedApproximateJacobianCache(J, structure, alg, cache, initialized::Bool, - internalnorm) - -A cache for Approximate Jacobian. - -### Arguments - - - `J`: The current Jacobian. - - `structure`: The structure of the Jacobian. - - `alg`: The initialization algorithm. - - `cache`: The Jacobian cache [`NonlinearSolve.JacobianCache`](@ref) (if needed). - - `initialized`: A boolean indicating whether the Jacobian has been initialized. - - `internalnorm`: The norm to be used. - -### Interface - -```julia -(cache::InitializedApproximateJacobianCache)(::Nothing) -``` - -Returns the current Jacobian `cache.J` with the proper `structure`. - -```julia -__internal_solve!(cache::InitializedApproximateJacobianCache, fu, u, ::Val{reinit}) -``` - -Solves for the Jacobian `cache.J` and returns it. If `reinit` is `true`, then the Jacobian -is reinitialized. -""" -@concrete mutable struct InitializedApproximateJacobianCache - J - structure - alg - cache - initialized::Bool - internalnorm -end - -function __reinit_internal!(cache::InitializedApproximateJacobianCache, args...; kwargs...) - cache.initialized = false -end - -@internal_caches InitializedApproximateJacobianCache :cache - -function (cache::InitializedApproximateJacobianCache)(::Nothing) - return get_full_jacobian(cache, cache.structure, cache.J) -end - -function __internal_solve!( - cache::InitializedApproximateJacobianCache, fu, u, ::Val{reinit}) where {reinit} - if reinit || !cache.initialized - cache(cache.alg, fu, u) - cache.initialized = true - end - if stores_full_jacobian(cache.structure) - full_J = cache.J - else - full_J = get_full_jacobian(cache, cache.structure, cache.J) - end - return full_J -end - -function (cache::InitializedApproximateJacobianCache)(alg::IdentityInitialization, fu, u) - α = __initial_alpha(alg.alpha, u, fu, cache.internalnorm) - cache.J = __make_identity!!(cache.J, α) - return -end - -function (cache::InitializedApproximateJacobianCache)( - alg::TrueJacobianInitialization, fu, u) - J_new = cache.cache(u) - cache.J = cache.structure(cache.J, J_new) - return -end - -# Matrix Inversion -@inline __safe_inv_workspace(A) = nothing, A -@inline __safe_inv_workspace(A::ApplyArray) = __safe_inv_workspace(X) -@inline __safe_inv_workspace(A::SparseMatrixCSC) = Matrix(A), Matrix(A) - -@inline __safe_inv!!(workspace, A::Number) = pinv(A) -@inline __safe_inv!!(workspace, A::AbstractMatrix) = pinv(A) -@inline function __safe_inv!!(workspace, A::Diagonal) - D = A.diag - @bb @. D = pinv(D) - return Diagonal(D) -end -@inline function __safe_inv!!(workspace, A::AbstractVector{T}) where {T} - @. A = ifelse(iszero(A), zero(T), one(T) / A) - return A -end -@inline __safe_inv!!(workspace, A::ApplyArray) = __safe_inv!!(workspace, A.f(A.args...)) -@inline function __safe_inv!!(workspace::AbstractMatrix, A::SparseMatrixCSC) - copyto!(workspace, A) - return __safe_inv!!(nothing, workspace) -end -@inline function __safe_inv!!(workspace, A::StridedMatrix{T}) where {T} - LinearAlgebra.checksquare(A) - if istriu(A) - issingular = any(iszero, @view(A[diagind(A)])) - A_ = UpperTriangular(A) - !issingular && return triu!(parent(inv(A_))) - elseif istril(A) - A_ = LowerTriangular(A) - issingular = any(iszero, @view(A_[diagind(A_)])) - !issingular && return tril!(parent(inv(A_))) - else - F = lu(A; check = false) - if issuccess(F) - Ai = LinearAlgebra.inv!(F) - return convert(typeof(parent(Ai)), Ai) - end - end - return pinv(A) -end - -@inline __safe_inv(x) = __safe_inv!!(first(__safe_inv_workspace(x)), x) - -LazyArrays.applied_eltype(::typeof(__safe_inv), x) = eltype(x) -LazyArrays.applied_ndims(::typeof(__safe_inv), x) = ndims(x) -LazyArrays.applied_size(::typeof(__safe_inv), x) = size(x) -LazyArrays.applied_axes(::typeof(__safe_inv), x) = axes(x) diff --git a/src/internal/forward_diff.jl b/src/internal/forward_diff.jl deleted file mode 100644 index 269ebc99a..000000000 --- a/src/internal/forward_diff.jl +++ /dev/null @@ -1,75 +0,0 @@ -const DualNonlinearProblem = NonlinearProblem{<:Union{Number, <:AbstractArray}, iip, - <:Union{<:Dual{T, V, P}, <:AbstractArray{<:Dual{T, V, P}}}} where {iip, T, V, P} -const DualNonlinearLeastSquaresProblem = NonlinearLeastSquaresProblem{ - <:Union{Number, <:AbstractArray}, iip, - <:Union{<:Dual{T, V, P}, <:AbstractArray{<:Dual{T, V, P}}}} where {iip, T, V, P} -const DualAbstractNonlinearProblem = Union{ - DualNonlinearProblem, DualNonlinearLeastSquaresProblem} - -for algType in ALL_SOLVER_TYPES - @eval function SciMLBase.__solve( - prob::DualNonlinearProblem, alg::$(algType), args...; kwargs...) - sol, partials = nonlinearsolve_forwarddiff_solve(prob, alg, args...; kwargs...) - dual_soln = nonlinearsolve_dual_solution(sol.u, partials, prob.p) - return SciMLBase.build_solution( - prob, alg, dual_soln, sol.resid; sol.retcode, sol.stats, sol.original) - end -end - -@concrete mutable struct NonlinearSolveForwardDiffCache - cache - prob - alg - p - values_p - partials_p -end - -@internal_caches NonlinearSolveForwardDiffCache :cache - -function reinit_cache!(cache::NonlinearSolveForwardDiffCache; - p = cache.p, u0 = get_u(cache.cache), kwargs...) - inner_cache = reinit_cache!(cache.cache; p = __value(p), u0 = __value(u0), kwargs...) - cache.cache = inner_cache - cache.p = p - cache.values_p = __value(p) - cache.partials_p = ForwardDiff.partials(p) - return cache -end - -for algType in ALL_SOLVER_TYPES - @eval function SciMLBase.__init( - prob::DualNonlinearProblem, alg::$(algType), args...; kwargs...) - p = __value(prob.p) - newprob = NonlinearProblem(prob.f, __value(prob.u0), p; prob.kwargs...) - cache = init(newprob, alg, args...; kwargs...) - return NonlinearSolveForwardDiffCache( - cache, newprob, alg, prob.p, p, ForwardDiff.partials(prob.p)) - end -end - -function SciMLBase.solve!(cache::NonlinearSolveForwardDiffCache) - sol = solve!(cache.cache) - prob = cache.prob - - uu = sol.u - Jₚ = nonlinearsolve_∂f_∂p(prob, prob.f, uu, cache.values_p) - Jᵤ = nonlinearsolve_∂f_∂u(prob, prob.f, uu, cache.values_p) - - z_arr = -Jᵤ \ Jₚ - - sumfun = ((z, p),) -> map(zᵢ -> zᵢ * ForwardDiff.partials(p), z) - if cache.p isa Number - partials = sumfun((z_arr, cache.p)) - else - partials = sum(sumfun, zip(eachcol(z_arr), cache.p)) - end - - dual_soln = nonlinearsolve_dual_solution(sol.u, partials, cache.p) - return SciMLBase.build_solution( - prob, cache.alg, dual_soln, sol.resid; sol.retcode, sol.stats, sol.original) -end - -@inline __value(x) = x -@inline __value(x::Dual) = ForwardDiff.value(x) -@inline __value(x::AbstractArray{<:Dual}) = map(ForwardDiff.value, x) diff --git a/src/internal/helpers.jl b/src/internal/helpers.jl deleted file mode 100644 index 6a154e0cc..000000000 --- a/src/internal/helpers.jl +++ /dev/null @@ -1,166 +0,0 @@ -# Evaluate the residual function at a given point -function evaluate_f(prob::AbstractNonlinearProblem{uType, iip}, u) where {uType, iip} - (; f, p) = prob - if iip - fu = f.resid_prototype === nothing ? zero(u) : similar(f.resid_prototype) - f(fu, u, p) - else - fu = f(u, p) - end - return fu -end - -function evaluate_f!(cache, u, p) - cache.stats.nf += 1 - if isinplace(cache) - cache.prob.f(get_fu(cache), u, p) - else - set_fu!(cache, cache.prob.f(u, p)) - end -end - -evaluate_f!!(prob::AbstractNonlinearProblem, fu, u, p) = evaluate_f!!(prob.f, fu, u, p) -function evaluate_f!!(f::NonlinearFunction{iip}, fu, u, p) where {iip} - if iip - f(fu, u, p) - return fu - end - return f(u, p) -end - -# Callbacks -""" - callback_into_cache!(cache, internalcache, args...) - -Define custom operations on `internalcache` tightly coupled with the calling `cache`. -`args...` contain the sequence of caches calling into `internalcache`. - -This unfortunately makes code very tightly coupled and not modular. It is recommended to not -use this functionality unless it can't be avoided (like in [`LevenbergMarquardt`](@ref)). -""" -@inline callback_into_cache!(cache, internalcache, args...) = nothing # By default do nothing - -# Extension Algorithm Helpers -function __test_termination_condition(termination_condition, alg) - !(termination_condition isa AbsNormTerminationMode) && - termination_condition !== nothing && - error("`$(alg)` does not support termination conditions!") -end - -function __construct_extension_f(prob::AbstractNonlinearProblem; alias_u0::Bool = false, - can_handle_oop::Val = False, can_handle_scalar::Val = False, - make_fixed_point::Val = False, force_oop::Val = False) - if can_handle_oop === False && can_handle_scalar === True - error("Incorrect Specification: OOP not supported but scalar supported.") - end - - resid = evaluate_f(prob, prob.u0) - u0 = can_handle_scalar === True || !(prob.u0 isa Number) ? - __maybe_unaliased(prob.u0, alias_u0) : [prob.u0] - - fₚ = if make_fixed_point === True - if isinplace(prob) - @closure (du, u) -> (prob.f(du, u, prob.p); du .+= u) - else - @closure u -> prob.f(u, prob.p) .+ u - end - else - if isinplace(prob) - @closure (du, u) -> prob.f(du, u, prob.p) - else - @closure u -> prob.f(u, prob.p) - end - end - - 𝐟 = if isinplace(prob) - u0_size, du_size = size(u0), size(resid) - @closure (du, u) -> (fₚ(reshape(du, du_size), reshape(u, u0_size)); du) - else - if prob.u0 isa Number - if can_handle_scalar === True - fₚ - elseif can_handle_oop === True - @closure u -> [fₚ(first(u))] - else - @closure (du, u) -> (du[1] = fₚ(first(u)); du) - end - else - u0_size = size(u0) - if can_handle_oop === True - @closure u -> vec(fₚ(reshape(u, u0_size))) - else - @closure (du, u) -> (copyto!(du, fₚ(reshape(u, u0_size))); du) - end - end - end - - 𝐅 = if force_oop === True && applicable(𝐟, u0, u0) - _resid = resid isa Number ? [resid] : _vec(resid) - du = _vec(zero(_resid)) - @closure u -> begin - 𝐟(du, u) - return du - end - else - 𝐟 - end - - return 𝐅, _vec(u0), (resid isa Number ? [resid] : _vec(resid)) -end - -function __construct_extension_jac(prob, alg, u0, fu; can_handle_oop::Val = False, - can_handle_scalar::Val = False, autodiff = nothing, initial_jacobian = False, - kwargs...) - autodiff = select_jacobian_autodiff(prob, autodiff) - - Jₚ = JacobianCache( - prob, alg, prob.f, fu, u0, prob.p; stats = empty_nlstats(), autodiff, kwargs...) - - 𝓙 = (can_handle_scalar === False && prob.u0 isa Number) ? @closure(u->[Jₚ(u[1])]) : Jₚ - - 𝐉 = (can_handle_oop === False && !isinplace(prob)) ? - @closure((J, u)->copyto!(J, 𝓙(u))) : 𝓙 - - initial_jacobian === False && return 𝐉 - - return 𝐉, Jₚ(nothing) -end - -function reinit_cache! end -reinit_cache!(cache::Nothing, args...; kwargs...) = nothing -reinit_cache!(cache, args...; kwargs...) = nothing - -function __reinit_internal! end -__reinit_internal!(::Nothing, args...; kwargs...) = nothing -__reinit_internal!(cache, args...; kwargs...) = nothing - -# Auto-generate some of the helper functions -macro internal_caches(cType, internal_cache_names...) - return __internal_caches(cType, internal_cache_names) -end - -function __internal_caches(cType, internal_cache_names::Tuple) - callback_caches = map( - name -> :($(callback_into_cache!)( - cache, getproperty(internalcache, $(name)), internalcache, args...)), - internal_cache_names) - callbacks_self = map( - name -> :($(callback_into_cache!)( - internalcache, getproperty(internalcache, $(name)))), - internal_cache_names) - reinit_caches = map( - name -> :($(reinit_cache!)(getproperty(cache, $(name)), args...; kwargs...)), - internal_cache_names) - return esc(quote - function callback_into_cache!(cache, internalcache::$(cType), args...) - $(callback_caches...) - end - function callback_into_cache!(internalcache::$(cType)) - $(callbacks_self...) - end - function reinit_cache!(cache::$(cType), args...; kwargs...) - $(reinit_caches...) - $(__reinit_internal!)(cache, args...; kwargs...) - end - end) -end diff --git a/src/internal/linear_solve.jl b/src/internal/linear_solve.jl deleted file mode 100644 index 707790ff3..000000000 --- a/src/internal/linear_solve.jl +++ /dev/null @@ -1,245 +0,0 @@ -const LinearSolveFailureCode = isdefined(ReturnCode, :InternalLinearSolveFailure) ? - ReturnCode.InternalLinearSolveFailure : ReturnCode.Failure - -""" - LinearSolverCache(alg, linsolve, A, b, u; stats, kwargs...) - -Construct a cache for solving linear systems of the form `A * u = b`. Following cases are -handled: - - 1. `A` is Number, then we solve it with `u = b / A` - 2. `A` is `SMatrix`, then we solve it with `u = A \\ b` (using the defaults from base - Julia) - 3. `A` is `Diagonal`, then we solve it with `u = b ./ A.diag` - 4. In all other cases, we use `alg` to solve the linear system using - [LinearSolve.jl](https://github.com/SciML/LinearSolve.jl). - -### Solving the System - -```julia -(cache::LinearSolverCache)(; - A = nothing, b = nothing, linu = nothing, du = nothing, p = nothing, - weight = nothing, cachedata = nothing, reuse_A_if_factorization = false, kwargs...) -``` - -Returns the solution of the system `u` and stores the updated cache in `cache.lincache`. - -#### Special Handling for Rank-deficient Matrix `A` - -If we detect a failure in the linear solve (mostly due to using an algorithm that doesn't -support rank-deficient matrices), we emit a warning and attempt to solve the problem using -Pivoted QR factorization. This is quite efficient if there are only a few rank-deficient -that originate in the problem. However, if these are quite frequent for the main nonlinear -system, then it is recommended to use a different linear solver that supports rank-deficient -matrices. - -#### Keyword Arguments - - - `reuse_A_if_factorization`: If `true`, then the factorization of `A` is reused if - possible. This is useful when solving the same system with different `b` values. - If the algorithm is an iterative solver, then we reset the internal linear solve cache. - -One distinct feature of this compared to the cache from LinearSolve is that it respects the -aliasing arguments even after cache construction, i.e., if we passed in an `A` that `A` is -not mutated, we do this by copying over `A` to a preconstructed cache. -""" -@concrete mutable struct LinearSolverCache <: AbstractLinearSolverCache - lincache - linsolve - additional_lincache::Any - A - b - precs - stats::NLStats -end - -@inline __fix_strange_type_combination(A, b, u) = u -@inline function __fix_strange_type_combination(A, b, u::SArray) - A isa SArray && b isa SArray && return u - @warn "Solving Linear System A::$(typeof(A)) x::$(typeof(u)) = b::$(typeof(u)) is not \ - properly supported. Converting `x` to a mutable array. Check the return type \ - of the nonlinear function provided for optimal performance." maxlog=1 - return MArray(u) -end - -@inline __set_lincache_u!(cache, u) = (cache.lincache.u = u) -@inline function __set_lincache_u!(cache, u::SArray) - cache.lincache.u isa MArray && return __set_lincache_u!(cache, MArray(u)) - cache.lincache.u = u -end - -function LinearSolverCache(alg, linsolve, A, b, u; stats, kwargs...) - u_fixed = __fix_strange_type_combination(A, b, u) - - if (A isa Number && b isa Number) || - (linsolve === nothing && A isa SMatrix) || - (A isa Diagonal) || - (linsolve isa typeof(\)) - return LinearSolverCache(nothing, nothing, nothing, A, b, nothing, stats) - end - @bb u_ = copy(u_fixed) - linprob = LinearProblem(A, b; u0 = u_, kwargs...) - - if __hasfield(alg, Val(:precs)) - precs = alg.precs - Pl_, Pr_ = precs(A, nothing, u, ntuple(Returns(nothing), 6)...) - else - precs, Pl_, Pr_ = nothing, nothing, nothing - end - Pl, Pr = __wrapprecs(Pl_, Pr_, u) - - # Unalias here, we will later use these as caches - lincache = init(linprob, linsolve; alias_A = false, alias_b = false, Pl, Pr) - - return LinearSolverCache(lincache, linsolve, nothing, nothing, nothing, precs, stats) -end - -@kwdef @concrete struct LinearSolveResult - u - success::Bool = true -end - -# Direct Linear Solve Case without Caching -function (cache::LinearSolverCache{Nothing})(; - A = nothing, b = nothing, linu = nothing, kwargs...) - cache.stats.nsolve += 1 - cache.stats.nfactors += 1 - A === nothing || (cache.A = A) - b === nothing || (cache.b = b) - if A isa Diagonal - _diag = _restructure(cache.b, cache.A.diag) - @bb @. linu = cache.b / _diag - res = linu - else - res = cache.A \ cache.b - end - return LinearSolveResult(; u = res) -end - -# Use LinearSolve.jl -function (cache::LinearSolverCache)(; - A = nothing, b = nothing, linu = nothing, du = nothing, - p = nothing, weight = nothing, cachedata = nothing, - reuse_A_if_factorization = false, verbose = true, kwargs...) - cache.stats.nsolve += 1 - - __update_A!(cache, A, reuse_A_if_factorization) - b !== nothing && (cache.lincache.b = b) - linu !== nothing && __set_lincache_u!(cache, linu) - - Plprev = cache.lincache.Pl - Prprev = cache.lincache.Pr - - if cache.precs === nothing - _Pl, _Pr = nothing, nothing - else - _Pl, _Pr = cache.precs(cache.lincache.A, du, linu, p, nothing, - A !== nothing, Plprev, Prprev, cachedata) - end - - if (_Pl !== nothing || _Pr !== nothing) - Pl, Pr = __wrapprecs(_Pl, _Pr, linu) - cache.lincache.Pl = Pl - cache.lincache.Pr = Pr - end - - linres = solve!(cache.lincache) - cache.lincache = linres.cache - # Unfortunately LinearSolve.jl doesn't have the most uniform ReturnCode handling - if linres.retcode === ReturnCode.Failure - structured_mat = ArrayInterface.isstructured(cache.lincache.A) - is_gpuarray = ArrayInterface.device(cache.lincache.A) isa ArrayInterface.GPU - if !(cache.linsolve isa QRFactorization{ColumnNorm}) && !is_gpuarray && - !structured_mat - if verbose - @warn "Potential Rank Deficient Matrix Detected. Attempting to solve using \ - Pivoted QR Factorization." - end - @assert (A !== nothing)&&(b !== nothing) "This case is not yet supported. \ - Please open an issue at \ - https://github.com/SciML/NonlinearSolve.jl" - if cache.additional_lincache === nothing # First time - linprob = LinearProblem(A, b; u0 = linres.u) - cache.additional_lincache = init( - linprob, QRFactorization(ColumnNorm()); alias_u0 = false, - alias_A = false, alias_b = false, cache.lincache.Pl, cache.lincache.Pr) - else - cache.additional_lincache.A = A - cache.additional_lincache.b = b - cache.additional_lincache.Pl = cache.lincache.Pl - cache.additional_lincache.Pr = cache.lincache.Pr - end - linres = solve!(cache.additional_lincache) - cache.additional_lincache = linres.cache - linres.retcode === ReturnCode.Failure && - return LinearSolveResult(; u = linres.u, success = false) - return LinearSolveResult(; u = linres.u) - elseif !(cache.linsolve isa QRFactorization{ColumnNorm}) - if verbose - if structured_mat || is_gpuarray - mat_desc = structured_mat ? "Structured" : "GPU" - @warn "Potential Rank Deficient Matrix Detected. But Matrix is \ - $(mat_desc). Currently, we don't attempt to solve Rank Deficient \ - $(mat_desc) Matrices. Please open an issue at \ - https://github.com/SciML/NonlinearSolve.jl" - end - end - end - return LinearSolveResult(; u = linres.u, success = false) - end - - return LinearSolveResult(; u = linres.u) -end - -@inline __update_A!(cache::LinearSolverCache, ::Nothing, reuse) = cache -@inline function __update_A!(cache::LinearSolverCache, A, reuse) - return __update_A!(cache, __getproperty(cache.lincache, Val(:alg)), A, reuse) -end -@inline function __update_A!(cache, alg, A, reuse) - # Not a Factorization Algorithm so don't update `nfactors` - __set_lincache_A(cache.lincache, A) - return cache -end -@inline function __update_A!(cache, ::AbstractFactorization, A, reuse) - reuse && return cache - __set_lincache_A(cache.lincache, A) - cache.stats.nfactors += 1 - return cache -end -@inline function __update_A!(cache, alg::DefaultLinearSolver, A, reuse) - if alg == DefaultLinearSolver(DefaultAlgorithmChoice.KrylovJL_GMRES) - # Force a reset of the cache. This is not properly handled in LinearSolve.jl - __set_lincache_A(cache.lincache, A) - return cache - end - reuse && return cache - __set_lincache_A(cache.lincache, A) - cache.stats.nfactors += 1 - return cache -end - -function __set_lincache_A(lincache, new_A) - if LinearSolve.default_alias_A(lincache.alg, new_A, lincache.b) - lincache.A = new_A - else - if can_setindex(lincache.A) - copyto!(lincache.A, new_A) - lincache.A = lincache.A - else - lincache.A = new_A - end - end -end - -function __wrapprecs(_Pl, _Pr, u) - Pl = _Pl !== nothing ? _Pl : IdentityOperator(length(u)) - Pr = _Pr !== nothing ? _Pr : IdentityOperator(length(u)) - return Pl, Pr -end - -@inline __needs_square_A(_, ::Number) = false -@inline __needs_square_A(::Nothing, ::Number) = false -@inline __needs_square_A(::Nothing, _) = false -@inline __needs_square_A(linsolve, _) = LinearSolve.needs_square_A(linsolve) -@inline __needs_square_A(::typeof(\), _) = false -@inline __needs_square_A(::typeof(\), ::Number) = false # Ambiguity Fix diff --git a/src/internal/termination.jl b/src/internal/termination.jl deleted file mode 100644 index 7728aea69..000000000 --- a/src/internal/termination.jl +++ /dev/null @@ -1,34 +0,0 @@ -function check_and_update!(cache, fu, u, uprev) - return check_and_update!(cache.termination_cache, cache, fu, u, uprev) -end - -function check_and_update!(tc_cache, cache, fu, u, uprev) - return check_and_update!(tc_cache, cache, fu, u, uprev, tc_cache.mode) -end - -function check_and_update!(tc_cache, cache, fu, u, uprev, mode) - if tc_cache(fu, u, uprev) - cache.retcode = tc_cache.retcode - update_from_termination_cache!(tc_cache, cache, mode, u) - cache.force_stop = true - end -end - -function update_from_termination_cache!(tc_cache, cache, u = get_u(cache)) - return update_from_termination_cache!(tc_cache, cache, tc_cache.mode, u) -end - -function update_from_termination_cache!( - tc_cache, cache, ::AbstractNonlinearTerminationMode, u = get_u(cache)) - evaluate_f!(cache, u, cache.p) -end - -function update_from_termination_cache!( - tc_cache, cache, ::AbstractSafeBestNonlinearTerminationMode, u = get_u(cache)) - if isinplace(cache) - copyto!(get_u(cache), tc_cache.u) - else - set_u!(cache, tc_cache.u) - end - evaluate_f!(cache, get_u(cache), cache.p) -end diff --git a/src/internal/tracing.jl b/src/internal/tracing.jl deleted file mode 100644 index fe35f1a9c..000000000 --- a/src/internal/tracing.jl +++ /dev/null @@ -1,217 +0,0 @@ -""" - TraceMinimal(freq) - TraceMinimal(; print_frequency = 1, store_frequency::Int = 1) - -Trace Minimal Information - - 1. Iteration Number - 2. f(u) inf-norm - 3. Step 2-norm - -See also [`TraceWithJacobianConditionNumber`](@ref) and [`TraceAll`](@ref). -""" -@kwdef struct TraceMinimal <: AbstractNonlinearSolveTraceLevel - print_frequency::Int = 1 - store_frequency::Int = 1 -end - -""" - TraceWithJacobianConditionNumber(freq) - TraceWithJacobianConditionNumber(; print_frequency = 1, store_frequency::Int = 1) - -[`TraceMinimal`](@ref) + Print the Condition Number of the Jacobian. - -See also [`TraceMinimal`](@ref) and [`TraceAll`](@ref). -""" -@kwdef struct TraceWithJacobianConditionNumber <: AbstractNonlinearSolveTraceLevel - print_frequency::Int = 1 - store_frequency::Int = 1 -end - -""" - TraceAll(freq) - TraceAll(; print_frequency = 1, store_frequency::Int = 1) - -[`TraceWithJacobianConditionNumber`](@ref) + Store the Jacobian, u, f(u), and δu. - -!!! warning - - This is very expensive and makes copyies of the Jacobian, u, f(u), and δu. - -See also [`TraceMinimal`](@ref) and [`TraceWithJacobianConditionNumber`](@ref). -""" -@kwdef struct TraceAll <: AbstractNonlinearSolveTraceLevel - print_frequency::Int = 1 - store_frequency::Int = 1 -end - -for Tr in (:TraceMinimal, :TraceWithJacobianConditionNumber, :TraceAll) - @eval begin - $(Tr)(freq) = $(Tr)(; print_frequency = freq, store_frequency = freq) - end -end - -# NonlinearSolve Tracing Utilities -@concrete struct NonlinearSolveTraceEntry{nType} - iteration::Int - fnorm - stepnorm - condJ - J - u - fu - δu -end - -function __show_top_level(io::IO, entry::NonlinearSolveTraceEntry{nType}) where {nType} - if entry.condJ === nothing - @printf io "%-8s %-20s %-20s\n" "----" "-------------" "-----------" - if nType === :L2 - @printf io "%-8s %-20s %-20s\n" "Iter" "f(u) 2-norm" "Step 2-norm" - else - @printf io "%-8s %-20s %-20s\n" "Iter" "f(u) inf-norm" "Step 2-norm" - end - @printf io "%-8s %-20s %-20s\n" "----" "-------------" "-----------" - else - @printf io "%-8s %-20s %-20s %-20s\n" "----" "-------------" "-----------" "-------" - if nType === :L2 - @printf io "%-8s %-20s %-20s %-20s\n" "Iter" "f(u) 2-norm" "Step 2-norm" "cond(J)" - else - @printf io "%-8s %-20s %-20s %-20s\n" "Iter" "f(u) inf-norm" "Step 2-norm" "cond(J)" - end - @printf io "%-8s %-20s %-20s %-20s\n" "----" "-------------" "-----------" "-------" - end -end - -function Base.show(io::IO, entry::NonlinearSolveTraceEntry{nType}) where {nType} - entry.iteration == 0 && __show_top_level(io, entry) - if entry.iteration < 0 - # Special case for final entry - @printf io "%-8s %-20.8e\n" "Final" entry.fnorm - @printf io "%-28s\n" "----------------------" - elseif entry.condJ === nothing - @printf io "%-8d %-20.8e %-20.8e\n" entry.iteration entry.fnorm entry.stepnorm - else - @printf io "%-8d %-20.8e %-20.8e %-20.8e\n" entry.iteration entry.fnorm entry.stepnorm entry.condJ - end - return nothing -end - -function NonlinearSolveTraceEntry(prob::AbstractNonlinearProblem, iteration, fu, δu) - nType = ifelse(prob isa NonlinearLeastSquaresProblem, :L2, :Inf) - fnorm = prob isa NonlinearLeastSquaresProblem ? norm(fu, 2) : norm(fu, Inf) - return NonlinearSolveTraceEntry{nType}( - iteration, fnorm, norm(δu, 2), nothing, nothing, nothing, nothing, nothing) -end - -function NonlinearSolveTraceEntry(prob::AbstractNonlinearProblem, iteration, fu, δu, J) - nType = ifelse(prob isa NonlinearLeastSquaresProblem, :L2, :Inf) - fnorm = prob isa NonlinearLeastSquaresProblem ? norm(fu, 2) : norm(fu, Inf) - return NonlinearSolveTraceEntry{nType}( - iteration, fnorm, norm(δu, 2), __cond(J), nothing, nothing, nothing, nothing) -end - -function NonlinearSolveTraceEntry(prob::AbstractNonlinearProblem, iteration, fu, δu, J, u) - nType = ifelse(prob isa NonlinearLeastSquaresProblem, :L2, :Inf) - fnorm = prob isa NonlinearLeastSquaresProblem ? norm(fu, 2) : norm(fu, Inf) - return NonlinearSolveTraceEntry{nType}(iteration, fnorm, norm(δu, 2), __cond(J), - __copy(J), __copy(u), __copy(fu), __copy(δu)) -end - -@concrete struct NonlinearSolveTrace{ - show_trace, store_trace, Tr <: AbstractNonlinearSolveTraceLevel} - history - trace_level::Tr - prob -end - -function reset!(trace::NonlinearSolveTrace) - (trace.history !== nothing && resize!(trace.history, 0)) -end - -function Base.show(io::IO, trace::NonlinearSolveTrace) - if trace.history !== nothing - foreach(entry -> show(io, entry), trace.history) - else - print(io, "Tracing Disabled") - end - return nothing -end - -function init_nonlinearsolve_trace(prob, alg, u, fu, J, δu; show_trace::Val = Val(false), - trace_level::AbstractNonlinearSolveTraceLevel = TraceMinimal(), - store_trace::Val = Val(false), uses_jac_inverse = Val(false), kwargs...) - return init_nonlinearsolve_trace( - prob, alg, show_trace, trace_level, store_trace, u, fu, J, δu, uses_jac_inverse) -end - -function init_nonlinearsolve_trace(prob::AbstractNonlinearProblem, alg, ::Val{show_trace}, - trace_level::AbstractNonlinearSolveTraceLevel, ::Val{store_trace}, u, fu, J, - δu, ::Val{uses_jac_inverse}) where {show_trace, store_trace, uses_jac_inverse} - if show_trace - print("\nAlgorithm: ") - Base.printstyled(alg, "\n\n"; color = :green, bold = true) - end - J_ = uses_jac_inverse ? (trace_level isa TraceMinimal ? J : __safe_inv(J)) : J - history = __init_trace_history( - prob, Val{show_trace}(), trace_level, Val{store_trace}(), u, fu, J_, δu) - return NonlinearSolveTrace{show_trace, store_trace}(history, trace_level, prob) -end - -function __init_trace_history( - prob::AbstractNonlinearProblem, ::Val{show_trace}, trace_level, - ::Val{store_trace}, u, fu, J, δu) where {show_trace, store_trace} - !store_trace && !show_trace && return nothing - entry = __trace_entry(prob, trace_level, 0, u, fu, J, δu) - show_trace && show(entry) - store_trace && return NonlinearSolveTraceEntry[entry] - return nothing -end - -function __trace_entry(prob, ::TraceMinimal, iter, u, fu, J, δu, α = 1) - return NonlinearSolveTraceEntry(prob, iter, fu, δu .* α) -end -function __trace_entry(prob, ::TraceWithJacobianConditionNumber, iter, u, fu, J, δu, α = 1) - return NonlinearSolveTraceEntry(prob, iter, fu, δu .* α, J) -end -function __trace_entry(prob, ::TraceAll, iter, u, fu, J, δu, α = 1) - return NonlinearSolveTraceEntry(prob, iter, fu, δu .* α, J, u) -end - -function update_trace!(trace::NonlinearSolveTrace{ShT, StT}, iter, u, fu, J, δu, - α = 1; last::Val{L} = Val(false)) where {ShT, StT, L} - !StT && !ShT && return nothing - - if L - nType = ifelse(trace.prob isa NonlinearLeastSquaresProblem, :L2, :Inf) - fnorm = trace.prob isa NonlinearLeastSquaresProblem ? norm(fu, 2) : norm(fu, Inf) - entry = NonlinearSolveTraceEntry{nType}( - -1, fnorm, NaN32, nothing, nothing, nothing, nothing, nothing) - ShT && show(entry) - return trace - end - - show_now = ShT && (mod1(iter, trace.trace_level.print_frequency) == 1) - store_now = StT && (mod1(iter, trace.trace_level.store_frequency) == 1) - (show_now || store_now) && - (entry = __trace_entry(trace.prob, trace.trace_level, iter, u, fu, J, δu, α)) - store_now && push!(trace.history, entry) - show_now && show(entry) - return trace -end - -function update_trace!(cache::AbstractNonlinearSolveCache, α = true) - trace = __getproperty(cache, Val(:trace)) - trace === nothing && return nothing - - J = __getproperty(cache, Val(:J)) - if J === nothing - update_trace!( - trace, cache.nsteps + 1, get_u(cache), get_fu(cache), nothing, cache.du, α) - elseif cache isa ApproximateJacobianSolveCache && store_inverse_jacobian(cache) - update_trace!(trace, cache.nsteps + 1, get_u(cache), get_fu(cache), - ApplyArray(__safe_inv, J), cache.du, α) - else - update_trace!(trace, cache.nsteps + 1, get_u(cache), get_fu(cache), J, cache.du, α) - end -end diff --git a/src/polyalg.jl b/src/polyalg.jl new file mode 100644 index 000000000..5cf21d6e1 --- /dev/null +++ b/src/polyalg.jl @@ -0,0 +1,543 @@ +""" + NonlinearSolvePolyAlgorithm(algs; start_index::Int = 1) + +A general way to define PolyAlgorithms for `NonlinearProblem` and +`NonlinearLeastSquaresProblem`. This is a container for a tuple of algorithms that will be +tried in order until one succeeds. If none succeed, then the algorithm with the lowest +residual is returned. + +### Arguments + + - `algs`: a tuple of algorithms to try in-order! (If this is not a Tuple, then the + returned algorithm is not type-stable). + +### Keyword Arguments + + - `start_index`: the index to start at. Defaults to `1`. + +### Example + +```julia +using NonlinearSolve + +alg = NonlinearSolvePolyAlgorithm((NewtonRaphson(), Broyden())) +``` +""" +@concrete struct NonlinearSolvePolyAlgorithm <: AbstractNonlinearSolveAlgorithm + static_length <: Val + algs <: Tuple + start_index::Int +end + +function NonlinearSolvePolyAlgorithm(algs; start_index::Int = 1) + @assert 0 < start_index ≤ length(algs) + algs = Tuple(algs) + return NonlinearSolvePolyAlgorithm(Val(length(algs)), algs, start_index) +end + +@concrete mutable struct NonlinearSolvePolyAlgorithmCache <: AbstractNonlinearSolveCache + static_length <: Val + prob <: AbstractNonlinearProblem + + caches <: Tuple + alg <: NonlinearSolvePolyAlgorithm + + best::Int + current::Int + nsteps::Int + + stats::NLStats + total_time::Float64 + maxtime + + retcode::ReturnCode.T + force_stop::Bool + + maxiters::Int + internalnorm + + u0 + u0_aliased + alias_u0::Bool +end + +function SII.symbolic_container(cache::NonlinearSolvePolyAlgorithmCache) + return cache.caches[cache.current] +end +SII.state_values(cache::NonlinearSolvePolyAlgorithmCache) = cache.u0 + +function Base.show(io::IO, ::MIME"text/plain", cache::NonlinearSolvePolyAlgorithmCache) + println(io, "NonlinearSolvePolyAlgorithmCache with \ + $(Utils.unwrap_val(cache.static_length)) algorithms:") + best_alg = ifelse(cache.best == -1, "nothing", cache.best) + println(io, " Best Algorithm: $(best_alg)") + println( + io, " Current Algorithm: [$(cache.current) / $(Utils.unwrap_val(cache.static_length))]" + ) + println(io, " nsteps: $(cache.nsteps)") + println(io, " retcode: $(cache.retcode)") + print(io, " Current Cache: ") + NonlinearSolveBase.show_nonlinearsolve_cache(io, cache.caches[cache.current], 4) +end + +function InternalAPI.reinit!( + cache::NonlinearSolvePolyAlgorithmCache, args...; p = cache.p, u0 = cache.u0 +) + foreach(cache.caches) do cache + InternalAPI.reinit!(cache, args...; p, u0) + end + cache.current = cache.alg.start_index + InternalAPI.reinit!(cache.stats) + cache.nsteps = 0 + cache.total_time = 0.0 +end + +function SciMLBase.__init( + prob::AbstractNonlinearProblem, alg::NonlinearSolvePolyAlgorithm, args...; + stats = NLStats(0, 0, 0, 0, 0), maxtime = nothing, maxiters = 1000, + internalnorm = L2_NORM, alias_u0 = false, verbose = true, kwargs... +) + if alias_u0 && !ArrayInterface.ismutable(prob.u0) + verbose && @warn "`alias_u0` has been set to `true`, but `u0` is \ + immutable (checked using `ArrayInterface.ismutable`)." + alias_u0 = false # If immutable don't care about aliasing + end + + u0 = prob.u0 + u0_aliased = alias_u0 ? copy(u0) : u0 + alias_u0 && (prob = SciMLBase.remake(prob; u0 = u0_aliased)) + + return NonlinearSolvePolyAlgorithmCache( + alg.static_length, prob, + map(alg.algs) do solver + SciMLBase.__init( + prob, solver, args...; + stats, maxtime, internalnorm, alias_u0, verbose, kwargs... + ) + end, + alg, -1, alg.start_index, 0, stats, 0.0, maxtime, + ReturnCode.Default, false, maxiters, internalnorm, + u0, u0_aliased, alias_u0 + ) +end + +@generated function CommonSolve.solve!(cache::NonlinearSolvePolyAlgorithmCache{Val{N}}) where {N} + calls = [quote + 1 ≤ cache.current ≤ $(N) || error("Current choices shouldn't get here!") + end] + + cache_syms = [gensym("cache") for i in 1:N] + sol_syms = [gensym("sol") for i in 1:N] + u_result_syms = [gensym("u_result") for i in 1:N] + + for i in 1:N + push!(calls, + quote + $(cache_syms[i]) = cache.caches[$(i)] + if $(i) == cache.current + cache.alias_u0 && copyto!(cache.u0_aliased, cache.u0) + $(sol_syms[i]) = CommonSolve.solve!($(cache_syms[i])) + if SciMLBase.successful_retcode($(sol_syms[i])) + stats = $(sol_syms[i]).stats + if cache.alias_u0 + copyto!(cache.u0, $(sol_syms[i]).u) + $(u_result_syms[i]) = cache.u0 + else + $(u_result_syms[i]) = $(sol_syms[i]).u + end + fu = NonlinearSolveBase.get_fu($(cache_syms[i])) + return build_solution_less_specialize( + cache.prob, cache.alg, $(u_result_syms[i]), fu; + retcode = $(sol_syms[i]).retcode, stats, + original = $(sol_syms[i]), trace = $(sol_syms[i]).trace + ) + elseif cache.alias_u0 + # For safety we need to maintain a copy of the solution + $(u_result_syms[i]) = copy($(sol_syms[i]).u) + end + cache.current = $(i + 1) + end + end) + end + + resids = map(Base.Fix2(Symbol, :resid), cache_syms) + for (sym, resid) in zip(cache_syms, resids) + push!(calls, :($(resid) = @isdefined($(sym)) ? $(sym).resid : nothing)) + end + push!(calls, quote + fus = tuple($(Tuple(resids)...)) + minfu, idx = findmin_caches(cache.prob, fus) + end) + for i in 1:N + push!(calls, + quote + if idx == $(i) + u = cache.alias_u0 ? $(u_result_syms[i]) : + NonlinearSolveBase.get_u(cache.caches[$(i)]) + end + end) + end + push!(calls, + quote + retcode = cache.caches[idx].retcode + if cache.alias_u0 + copyto!(cache.u0, u) + u = cache.u0 + end + return build_solution_less_specialize( + cache.prob, cache.alg, u, fus[idx]; + retcode, cache.stats, cache.caches[idx].trace + ) + end) + + return Expr(:block, calls...) +end + +@generated function InternalAPI.step!( + cache::NonlinearSolvePolyAlgorithmCache{Val{N}}, args...; kwargs... +) where {N} + calls = [] + cache_syms = [gensym("cache") for i in 1:N] + for i in 1:N + push!(calls, + quote + $(cache_syms[i]) = cache.caches[$(i)] + if $(i) == cache.current + InternalAPI.step!($(cache_syms[i]), args...; kwargs...) + $(cache_syms[i]).nsteps += 1 + if !NonlinearSolveBase.not_terminated($(cache_syms[i])) + if SciMLBase.successful_retcode($(cache_syms[i]).retcode) + cache.best = $(i) + cache.force_stop = true + cache.retcode = $(cache_syms[i]).retcode + else + cache.current = $(i + 1) + end + end + return + end + end) + end + + push!(calls, quote + if !(1 ≤ cache.current ≤ length(cache.caches)) + minfu, idx = findmin_caches(cache.prob, cache.caches) + cache.best = idx + cache.retcode = cache.caches[idx].retcode + cache.force_stop = true + return + end + end) + + return Expr(:block, calls...) +end + +@generated function SciMLBase.__solve( + prob::AbstractNonlinearProblem, alg::NonlinearSolvePolyAlgorithm{Val{N}}, args...; + stats = NLStats(0, 0, 0, 0, 0), alias_u0 = false, verbose = true, kwargs... +) where {N} + sol_syms = [gensym("sol") for _ in 1:N] + prob_syms = [gensym("prob") for _ in 1:N] + u_result_syms = [gensym("u_result") for _ in 1:N] + calls = [quote + current = alg.start_index + if alias_u0 && !ArrayInterface.ismutable(prob.u0) + verbose && @warn "`alias_u0` has been set to `true`, but `u0` is \ + immutable (checked using `ArrayInterface.ismutable`)." + alias_u0 = false # If immutable don't care about aliasing + end + u0 = prob.u0 + u0_aliased = alias_u0 ? zero(u0) : u0 + end] + for i in 1:N + cur_sol = sol_syms[i] + push!(calls, + quote + if current == $(i) + if alias_u0 + copyto!(u0_aliased, u0) + $(prob_syms[i]) = SciMLBase.remake(prob; u0 = u0_aliased) + else + $(prob_syms[i]) = prob + end + $(cur_sol) = SciMLBase.__solve( + $(prob_syms[i]), alg.algs[$(i)], args...; + stats, alias_u0, verbose, kwargs... + ) + if SciMLBase.successful_retcode($(cur_sol)) + if alias_u0 + copyto!(u0, $(cur_sol).u) + $(u_result_syms[i]) = u0 + else + $(u_result_syms[i]) = $(cur_sol).u + end + return build_solution_less_specialize( + prob, alg, $(u_result_syms[i]), $(cur_sol).resid; + $(cur_sol).retcode, $(cur_sol).stats, + $(cur_sol).trace, original = $(cur_sol) + ) + elseif alias_u0 + # For safety we need to maintain a copy of the solution + $(u_result_syms[i]) = copy($(cur_sol).u) + end + current = $(i + 1) + end + end) + end + + resids = map(Base.Fix2(Symbol, :resid), sol_syms) + for (sym, resid) in zip(sol_syms, resids) + push!(calls, :($(resid) = @isdefined($(sym)) ? $(sym).resid : nothing)) + end + + push!(calls, quote + resids = tuple($(Tuple(resids)...)) + minfu, idx = findmin_resids(prob, resids) + end) + + for i in 1:N + push!(calls, + quote + if idx == $(i) + if alias_u0 + copyto!(u0, $(u_result_syms[i])) + $(u_result_syms[i]) = u0 + else + $(u_result_syms[i]) = $(sol_syms[i]).u + end + return build_solution_less_specialize( + prob, alg, $(u_result_syms[i]), $(sol_syms[i]).resid; + $(sol_syms[i]).retcode, $(sol_syms[i]).stats, + $(sol_syms[i]).trace, original = $(sol_syms[i]) + ) + end + end) + end + push!(calls, :(error("Current choices shouldn't get here!"))) + + return Expr(:block, calls...) +end + +""" + RobustMultiNewton( + ::Type{T} = Float64; + concrete_jac = nothing, + linsolve = nothing, precs = nothing, + autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing + ) + +A polyalgorithm focused on robustness. It uses a mixture of Newton methods with different +globalizing techniques (trust region updates, line searches, etc.) in order to find a +method that is able to adequately solve the minimization problem. + +Basically, if this algorithm fails, then "most" good ways of solving your problem fail and +you may need to think about reformulating the model (either there is an issue with the model, +or more precision / more stable linear solver choice is required). + +### Arguments + + - `T`: The eltype of the initial guess. It is only used to check if some of the algorithms + are compatible with the problem type. Defaults to `Float64`. +""" +function RobustMultiNewton( + ::Type{T} = Float64; + concrete_jac = nothing, + linsolve = nothing, precs = nothing, + autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing +) where {T} + common_kwargs = (; concrete_jac, linsolve, precs, autodiff, vjp_autodiff, jvp_autodiff) + if T <: Complex # Let's atleast have something here for complex numbers + algs = (NewtonRaphson(; common_kwargs...),) + else + algs = ( + TrustRegion(; common_kwargs...), + TrustRegion(; common_kwargs..., radius_update_scheme = RUS.Bastin), + NewtonRaphson(; common_kwargs...), + NewtonRaphson(; common_kwargs..., linesearch = BackTracking()), + TrustRegion(; common_kwargs..., radius_update_scheme = RUS.NLsolve), + TrustRegion(; common_kwargs..., radius_update_scheme = RUS.Fan) + ) + end + return NonlinearSolvePolyAlgorithm(algs) +end + +""" + FastShortcutNonlinearPolyalg( + ::Type{T} = Float64; + concrete_jac = nothing, + linsolve = nothing, precs = nothing, + must_use_jacobian::Val = Val(false), + prefer_simplenonlinearsolve::Val = Val(false), + autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing, + u0_len::Union{Int, Nothing} = nothing + ) where {T} + +A polyalgorithm focused on balancing speed and robustness. It first tries less robust methods +for more performance and then tries more robust techniques if the faster ones fail. + +### Arguments + + - `T`: The eltype of the initial guess. It is only used to check if some of the algorithms + are compatible with the problem type. Defaults to `Float64`. + +### Keyword Arguments + + - `u0_len`: The length of the initial guess. If this is `nothing`, then the length of the + initial guess is not checked. If this is an integer and it is less than `25`, we use + jacobian based methods. +""" +function FastShortcutNonlinearPolyalg( + ::Type{T} = Float64; + concrete_jac = nothing, + linsolve = nothing, precs = nothing, + must_use_jacobian::Val = Val(false), + prefer_simplenonlinearsolve::Val = Val(false), + autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing, + u0_len::Union{Int, Nothing} = nothing +) where {T} + start_index = 1 + common_kwargs = (; concrete_jac, linsolve, precs, autodiff, vjp_autodiff, jvp_autodiff) + if must_use_jacobian isa Val{true} + if T <: Complex + algs = (NewtonRaphson(; common_kwargs...),) + else + algs = ( + NewtonRaphson(; common_kwargs...), + NewtonRaphson(; common_kwargs..., linesearch = BackTracking()), + TrustRegion(; common_kwargs...), + TrustRegion(; common_kwargs..., radius_update_scheme = RUS.Bastin) + ) + end + else + # SimpleNewtonRaphson and SimpleTrustRegion are not robust to singular Jacobians + # and thus are not included in the polyalgorithm + if prefer_simplenonlinearsolve isa Val{true} + if T <: Complex + algs = ( + SimpleBroyden(), + Broyden(; init_jacobian = Val(:true_jacobian), autodiff), + SimpleKlement(), + NewtonRaphson(; common_kwargs...) + ) + else + start_index = u0_len !== nothing ? (u0_len ≤ 25 ? 4 : 1) : 1 + algs = ( + SimpleBroyden(), + Broyden(; init_jacobian = Val(:true_jacobian), autodiff), + SimpleKlement(), + NewtonRaphson(; common_kwargs...), + NewtonRaphson(; common_kwargs..., linesearch = BackTracking()), + TrustRegion(; common_kwargs...), + TrustRegion(; common_kwargs..., radius_update_scheme = RUS.Bastin) + ) + end + else + if T <: Complex + algs = ( + Broyden(; autodiff), + Broyden(; init_jacobian = Val(:true_jacobian), autodiff), + Klement(; linsolve, precs, autodiff), + NewtonRaphson(; common_kwargs...) + ) + else + # TODO: This number requires a bit rigorous testing + start_index = u0_len !== nothing ? (u0_len ≤ 25 ? 4 : 1) : 1 + algs = ( + Broyden(; autodiff), + Broyden(; init_jacobian = Val(:true_jacobian), autodiff), + Klement(; linsolve, precs, autodiff), + NewtonRaphson(; common_kwargs...), + NewtonRaphson(; common_kwargs..., linesearch = BackTracking()), + TrustRegion(; common_kwargs...), + TrustRegion(; common_kwargs..., radius_update_scheme = RUS.Bastin) + ) + end + end + end + return NonlinearSolvePolyAlgorithm(algs; start_index) +end + +""" + FastShortcutNLLSPolyalg( + ::Type{T} = Float64; + concrete_jac = nothing, + linsolve = nothing, precs = nothing, + autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing + ) + +A polyalgorithm focused on balancing speed and robustness. It first tries less robust methods +for more performance and then tries more robust techniques if the faster ones fail. + +### Arguments + + - `T`: The eltype of the initial guess. It is only used to check if some of the algorithms + are compatible with the problem type. Defaults to `Float64`. +""" +function FastShortcutNLLSPolyalg( + ::Type{T} = Float64; + concrete_jac = nothing, + linsolve = nothing, precs = nothing, + autodiff = nothing, vjp_autodiff = nothing, jvp_autodiff = nothing +) where {T} + common_kwargs = (; linsolve, precs, autodiff, vjp_autodiff, jvp_autodiff) + if T <: Complex + algs = ( + GaussNewton(; common_kwargs..., concrete_jac), + LevenbergMarquardt(; common_kwargs..., disable_geodesic = Val(true)), + LevenbergMarquardt(; common_kwargs...) + ) + else + algs = ( + GaussNewton(; common_kwargs..., concrete_jac), + LevenbergMarquardt(; common_kwargs..., disable_geodesic = Val(true)), + TrustRegion(; common_kwargs..., concrete_jac), + GaussNewton(; common_kwargs..., linesearch = BackTracking(), concrete_jac), + TrustRegion(; + common_kwargs..., radius_update_scheme = RUS.Bastin, concrete_jac + ), + LevenbergMarquardt(; common_kwargs...) + ) + end + return NonlinearSolvePolyAlgorithm(algs) +end + +# Original is often determined on runtime information especially for PolyAlgorithms so it +# is best to never specialize on that +function build_solution_less_specialize( + prob::AbstractNonlinearProblem, alg, u, resid; + retcode = ReturnCode.Default, original = nothing, left = nothing, + right = nothing, stats = nothing, trace = nothing, kwargs... +) + return SciMLBase.NonlinearSolution{ + eltype(eltype(u)), ndims(u), typeof(u), typeof(resid), typeof(prob), + typeof(alg), Any, typeof(left), typeof(stats), typeof(trace) + }( + u, resid, prob, alg, retcode, original, left, right, stats, trace + ) +end + +function findmin_caches(prob::AbstractNonlinearProblem, caches) + resids = map(caches) do cache + cache === nothing && return nothing + return NonlinearSolveBase.get_fu(cache) + end + return findmin_resids(prob, resids) +end + +@views function findmin_resids(prob::AbstractNonlinearProblem, caches) + norm_fn = prob isa NonlinearLeastSquaresProblem ? Base.Fix2(norm, 2) : + Base.Fix2(norm, Inf) + idx = findfirst(Base.Fix2(!==, nothing), caches) + # This is an internal function so we assume that inputs are consistent and there is + # atleast one non-`nothing` value + fx_idx = norm_fn(caches[idx]) + idx == length(caches) && return fx_idx, idx + fmin = @closure xᵢ -> begin + xᵢ === nothing && return oftype(fx_idx, Inf) + fx = norm_fn(xᵢ) + return ifelse(isnan(fx), oftype(fx, Inf), fx) + end + x_min, x_min_idx = findmin(fmin, caches[(idx + 1):length(caches)]) + x_min < fx_idx && return x_min, x_min_idx + idx + return fx_idx, idx +end diff --git a/src/utils.jl b/src/utils.jl deleted file mode 100644 index 069fde86e..000000000 --- a/src/utils.jl +++ /dev/null @@ -1,152 +0,0 @@ -# Defaults -@inline DEFAULT_PRECS(W, du, u, p, t, newW, Plprev, Prprev, cachedata) = nothing, nothing - -# Helper Functions -@inline __hasfield(::T, ::Val{field}) where {T, field} = hasfield(T, field) - -@generated function __getproperty(s::S, ::Val{X}) where {S, X} - hasfield(S, X) && return :(s.$X) - return :(missing) -end - -@inline __needs_concrete_A(::Nothing) = false -@inline __needs_concrete_A(::typeof(\)) = true -@inline __needs_concrete_A(linsolve) = needs_concrete_A(linsolve) - -@inline __maybe_mutable(x, ::AutoSparse{<:AutoEnzyme}) = __mutable(x) # TODO: remove? -@inline __maybe_mutable(x, _) = x - -@inline @generated function _vec(v) - hasmethod(vec, Tuple{typeof(v)}) || return :(vec(v)) - return :(v) -end -@inline _vec(v::Number) = v -@inline _vec(v::AbstractVector) = v - -@inline _restructure(y, x) = restructure(y, x) -@inline _restructure(y::Number, x::Number) = x - -@inline __maybe_unaliased(x::Union{Number, SArray}, ::Bool) = x -@inline function __maybe_unaliased(x::AbstractArray, alias::Bool) - # Spend time coping iff we will mutate the array - (alias || !__can_setindex(typeof(x))) && return x - 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))) -@inline __cond(J::AbstractVector) = __cond(Diagonal(J)) -@inline __cond(J::ApplyArray) = __cond(J.f(J.args...)) -@inline __cond(J::SparseMatrixCSC) = __cond(Matrix(J)) -@inline __cond(J) = -1 # Covers cases where `J` is a Operator, nothing, etc. - -@inline __copy(x::AbstractArray) = copy(x) -@inline __copy(x::Number) = x -@inline __copy(x) = x - -# LazyArrays for tracing -__zero(x::AbstractArray) = zero(x) -__zero(x) = x -LazyArrays.applied_eltype(::typeof(__zero), x) = eltype(x) -LazyArrays.applied_ndims(::typeof(__zero), x) = ndims(x) -LazyArrays.applied_size(::typeof(__zero), x) = size(x) -LazyArrays.applied_axes(::typeof(__zero), x) = axes(x) - -# Use Symmetric Matrices if known to be efficient -@inline __maybe_symmetric(x) = Symmetric(x) -@inline __maybe_symmetric(x::Number) = x -## LinearSolve with `nothing` doesn't dispatch correctly here -@inline __maybe_symmetric(x::StaticArray) = x -@inline __maybe_symmetric(x::AbstractSparseMatrix) = x -@inline __maybe_symmetric(x::AbstractSciMLOperator) = x - -# SparseAD --> NonSparseAD -@inline __get_nonsparse_ad(backend::AutoSparse) = ADTypes.dense_ad(backend) -@inline __get_nonsparse_ad(ad) = ad - -# Simple Checks -@inline __is_present(::Nothing) = false -@inline __is_present(::Missing) = false -@inline __is_present(::Any) = true -@inline __is_present(::NoLineSearch) = false - -@inline __is_complex(::Type{ComplexF64}) = true -@inline __is_complex(::Type{ComplexF32}) = true -@inline __is_complex(::Type{Complex}) = true -@inline __is_complex(::Type{T}) where {T} = false - -@inline __findmin_caches(f::F, caches) where {F} = __findmin(f ∘ get_fu, caches) -# FIXME: L2_NORM makes an Array of NaNs not a NaN (atleast according to `isnan`) -@generated function __findmin(f::F, x) where {F} - # JET shows dynamic dispatch if this is not written as a generated function - F === typeof(L2_NORM) && return :(return __findmin_impl(Base.Fix1(maximum, abs), x)) - return :(return __findmin_impl(f, x)) -end -@inline @views function __findmin_impl(f::F, x) where {F} - idx = findfirst(Base.Fix2(!==, nothing), x) - # This is an internal function so we assume that inputs are consistent and there is - # atleast one non-`nothing` value - fx_idx = f(x[idx]) - idx == length(x) && return fx_idx, idx - fmin = @closure xᵢ -> begin - xᵢ === nothing && return oftype(fx_idx, Inf) - fx = f(xᵢ) - return ifelse(isnan(fx), oftype(fx, Inf), fx) - end - x_min, x_min_idx = findmin(fmin, x[(idx + 1):length(x)]) - x_min < fx_idx && return x_min, x_min_idx + idx - return fx_idx, idx -end - -@inline __can_setindex(x) = can_setindex(x) -@inline __can_setindex(::Number) = false - -@inline function __mutable(x) - __can_setindex(x) && return x - y = similar(x) - copyto!(y, x) - return y -end -@inline __mutable(x::SArray) = MArray(x) - -@inline __dot(x, y) = dot(_vec(x), _vec(y)) - -""" - pickchunksize(x) = pickchunksize(length(x)) - pickchunksize(x::Int) - -Determine the chunk size for ForwardDiff and PolyesterForwardDiff based on the input length. -""" -@inline pickchunksize(x) = pickchunksize(length(x)) -@inline pickchunksize(x::Int) = ForwardDiff.pickchunksize(x) - -# Original is often determined on runtime information especially for PolyAlgorithms so it -# is best to never specialize on that -function __build_solution_less_specialize(prob::AbstractNonlinearProblem, alg, u, resid; - retcode = ReturnCode.Default, original = nothing, left = nothing, - right = nothing, stats = nothing, trace = nothing, kwargs...) - T = eltype(eltype(u)) - N = ndims(u) - - return SciMLBase.NonlinearSolution{ - T, N, typeof(u), typeof(resid), typeof(prob), typeof(alg), - Any, typeof(left), typeof(stats), typeof(trace)}( - u, resid, prob, alg, retcode, original, left, right, stats, trace) -end - -@inline empty_nlstats() = NLStats(0, 0, 0, 0, 0) -function __reinit_internal!(stats::NLStats) - stats.nf = 0 - stats.nsteps = 0 - stats.nfactors = 0 - stats.njacs = 0 - stats.nsolve = 0 -end - -function __similar(x, args...; kwargs...) - y = similar(x, args...; kwargs...) - fill!(y, false) - return y -end diff --git a/test/core/23_test_problems_tests.jl b/test/23_test_problems_tests.jl similarity index 51% rename from test/core/23_test_problems_tests.jl rename to test/23_test_problems_tests.jl index 7648db53f..8fa4c47b6 100644 --- a/test/core/23_test_problems_tests.jl +++ b/test/23_test_problems_tests.jl @@ -10,27 +10,25 @@ function test_on_library( x = dict["start"] res = similar(x) nlprob = NonlinearProblem(problem, copy(x)) - @testset "$idx: $(dict["title"])" begin - for alg in alg_ops - try - sol = solve(nlprob, alg; maxiters = 10000) - problem(res, sol.u, nothing) - - skip = skip_tests !== nothing && idx in skip_tests[alg] - if skip - @test_skip norm(res, Inf) ≤ ϵ - continue - end - broken = idx in broken_tests[alg] ? true : false - @test norm(res, Inf)≤ϵ broken=broken - catch err - @error err - broken = idx in broken_tests[alg] ? true : false - if broken - @test false broken=true - else - @test 1 == 2 - end + @testset "$idx: $(dict["title"]) | alg #$(alg_id)" for (alg_id, alg) in enumerate(alg_ops) + try + sol = solve(nlprob, alg; maxiters = 10000) + problem(res, sol.u, nothing) + + skip = skip_tests !== nothing && idx in skip_tests[alg] + if skip + @test_skip norm(res, Inf) ≤ ϵ + continue + end + broken = idx in broken_tests[alg] ? true : false + @test norm(res, Inf)≤ϵ broken=broken + catch err + @error err + broken = idx in broken_tests[alg] ? true : false + if broken + @test false broken=true + else + @test 1 == 2 end end end @@ -40,7 +38,7 @@ end export test_on_library, problems, dicts end -@testitem "PolyAlgorithms" setup=[RobustnessTesting] tags=[:core] begin +@testitem "23 Test Problems: PolyAlgorithms" setup=[RobustnessTesting] tags=[:core] begin alg_ops = (RobustMultiNewton(), FastShortcutNonlinearPolyalg()) broken_tests = Dict(alg => Int[] for alg in alg_ops) @@ -50,8 +48,11 @@ end test_on_library(problems, dicts, alg_ops, broken_tests) end -@testitem "NewtonRaphson" setup=[RobustnessTesting] tags=[:core] begin - alg_ops = (NewtonRaphson(),) +@testitem "23 Test Problems: NewtonRaphson" setup=[RobustnessTesting] tags=[:core] begin + alg_ops = ( + NewtonRaphson(), + SimpleNewtonRaphson() + ) broken_tests = Dict(alg => Int[] for alg in alg_ops) broken_tests[alg_ops[1]] = [1] @@ -59,13 +60,26 @@ end test_on_library(problems, dicts, alg_ops, broken_tests) end -@testitem "TrustRegion" setup=[RobustnessTesting] tags=[:core] begin - alg_ops = (TrustRegion(; radius_update_scheme = RadiusUpdateSchemes.Simple), +@testitem "23 Test Problems: Halley" setup=[RobustnessTesting] tags=[:core] begin + alg_ops = (SimpleHalley(; autodiff = AutoForwardDiff()),) + + broken_tests = Dict(alg => Int[] for alg in alg_ops) + broken_tests[alg_ops[1]] = [1, 5, 15, 16, 18] + + test_on_library(problems, dicts, alg_ops, broken_tests) +end + +@testitem "23 Test Problems: TrustRegion" setup=[RobustnessTesting] tags=[:core] begin + alg_ops = ( + TrustRegion(; radius_update_scheme = RadiusUpdateSchemes.Simple), TrustRegion(; radius_update_scheme = RadiusUpdateSchemes.Fan), TrustRegion(; radius_update_scheme = RadiusUpdateSchemes.Hei), TrustRegion(; radius_update_scheme = RadiusUpdateSchemes.Yuan), TrustRegion(; radius_update_scheme = RadiusUpdateSchemes.Bastin), - TrustRegion(; radius_update_scheme = RadiusUpdateSchemes.NLsolve)) + TrustRegion(; radius_update_scheme = RadiusUpdateSchemes.NLsolve), + SimpleTrustRegion(), + SimpleTrustRegion(; nlsolve_update_rule = Val(true)) + ) broken_tests = Dict(alg => Int[] for alg in alg_ops) broken_tests[alg_ops[1]] = [11, 21] @@ -74,15 +88,20 @@ end broken_tests[alg_ops[4]] = [8, 11, 21] broken_tests[alg_ops[5]] = [21] broken_tests[alg_ops[6]] = [11, 21] + broken_tests[alg_ops[7]] = [3, 15, 16, 21] + broken_tests[alg_ops[8]] = [15, 16] test_on_library(problems, dicts, alg_ops, broken_tests) end -@testitem "LevenbergMarquardt" setup=[RobustnessTesting] tags=[:core] begin +@testitem "23 Test Problems: LevenbergMarquardt" setup=[RobustnessTesting] tags=[:core] begin using LinearSolve - alg_ops = (LevenbergMarquardt(), LevenbergMarquardt(; α_geodesic = 0.1), - LevenbergMarquardt(; linsolve = CholeskyFactorization())) + alg_ops = ( + LevenbergMarquardt(), + LevenbergMarquardt(; α_geodesic = 0.1), + LevenbergMarquardt(; linsolve = CholeskyFactorization()) + ) broken_tests = Dict(alg => Int[] for alg in alg_ops) broken_tests[alg_ops[1]] = [11, 21] @@ -92,47 +111,63 @@ end test_on_library(problems, dicts, alg_ops, broken_tests) end -@testitem "DFSane" setup=[RobustnessTesting] tags=[:core] begin - alg_ops = (DFSane(),) +@testitem "23 Test Problems: DFSane" setup=[RobustnessTesting] tags=[:core] begin + alg_ops = ( + DFSane(), + SimpleDFSane() + ) broken_tests = Dict(alg => Int[] for alg in alg_ops) broken_tests[alg_ops[1]] = [1, 2, 3, 5, 21] + if Sys.isapple() + broken_tests[alg_ops[2]] = [1, 2, 3, 5, 6, 21] + else + broken_tests[alg_ops[2]] = [1, 2, 3, 5, 6, 11, 21] + end test_on_library(problems, dicts, alg_ops, broken_tests) end -@testitem "Broyden" setup=[RobustnessTesting] tags=[:core] begin - alg_ops = (Broyden(), Broyden(; init_jacobian = Val(:true_jacobian)), +@testitem "23 Test Problems: Broyden" setup=[RobustnessTesting] tags=[:core] retries=3 begin + alg_ops = ( + Broyden(), + Broyden(; init_jacobian = Val(:true_jacobian)), Broyden(; update_rule = Val(:bad_broyden)), - Broyden(; init_jacobian = Val(:true_jacobian), update_rule = Val(:bad_broyden))) + Broyden(; init_jacobian = Val(:true_jacobian), update_rule = Val(:bad_broyden)), + SimpleBroyden() + ) broken_tests = Dict(alg => Int[] for alg in alg_ops) + broken_tests[alg_ops[2]] = [1, 5, 8, 11, 18] + broken_tests[alg_ops[4]] = [5, 6, 8, 11] if Sys.isapple() broken_tests[alg_ops[1]] = [1, 5, 11] - broken_tests[alg_ops[2]] = [1, 5, 8, 11, 18] broken_tests[alg_ops[3]] = [1, 5, 6, 9, 11] - broken_tests[alg_ops[4]] = [5, 6, 8, 11] else broken_tests[alg_ops[1]] = [1, 5, 11, 15] - broken_tests[alg_ops[2]] = [1, 5, 8, 11, 18] - broken_tests[alg_ops[3]] = [1, 5, 9, 11] - broken_tests[alg_ops[4]] = [5, 6, 8, 11] + broken_tests[alg_ops[3]] = [1, 5, 9, 11, 16] end + broken_tests[alg_ops[5]] = [1, 5, 11] test_on_library(problems, dicts, alg_ops, broken_tests, Sys.isapple() ? 1e-3 : 1e-4) end -@testitem "Klement" setup=[RobustnessTesting] tags=[:core] begin - alg_ops = (Klement(), Klement(; init_jacobian = Val(:true_jacobian_diagonal))) +@testitem "23 Test Problems: Klement" setup=[RobustnessTesting] tags=[:core] begin + alg_ops = ( + Klement(), + Klement(; init_jacobian = Val(:true_jacobian_diagonal)), + SimpleKlement() + ) broken_tests = Dict(alg => Int[] for alg in alg_ops) broken_tests[alg_ops[1]] = [1, 2, 4, 5, 11, 18, 22] broken_tests[alg_ops[2]] = [2, 4, 5, 7, 18, 22] + broken_tests[alg_ops[3]] = [1, 2, 4, 5, 11, 22] test_on_library(problems, dicts, alg_ops, broken_tests) end -@testitem "PseudoTransient" setup=[RobustnessTesting] tags=[:core] begin +@testitem "23 Test Problems: PseudoTransient" setup=[RobustnessTesting] tags=[:core] begin # PT relies on the root being a stable equilibrium for convergence, so it won't work on # most problems alg_ops = (PseudoTransient(),) diff --git a/test/core/nlls_tests.jl b/test/core/nlls_tests.jl deleted file mode 100644 index 040627c00..000000000 --- a/test/core/nlls_tests.jl +++ /dev/null @@ -1,122 +0,0 @@ -@testsetup module CoreNLLSTesting -using Reexport -@reexport using NonlinearSolve, LinearSolve, LinearAlgebra, StableRNGs, Random, ForwardDiff, - Zygote -using LineSearches: LineSearches, Static, HagerZhang, MoreThuente, StrongWolfe - -linesearches = [] -for ls in ( - Static(), HagerZhang(), MoreThuente(), StrongWolfe(), LineSearches.BackTracking()) - push!(linesearches, LineSearchesJL(; method = ls)) -end -push!(linesearches, BackTracking()) - -true_function(x, θ) = @. θ[1] * exp(θ[2] * x) * cos(θ[3] * x + θ[4]) -true_function(y, x, θ) = (@. y = θ[1] * exp(θ[2] * x) * cos(θ[3] * x + θ[4])) - -const θ_true = [1.0, 0.1, 2.0, 0.5] - -const x = [-1.0, -0.5, 0.0, 0.5, 1.0] - -const y_target = true_function(x, θ_true) - -function loss_function(θ, p) - ŷ = true_function(p, θ) - return ŷ .- y_target -end - -function loss_function(resid, θ, p) - true_function(resid, p, θ) - resid .= resid .- y_target - return resid -end - -const θ_init = θ_true .+ randn!(StableRNG(0), similar(θ_true)) * 0.1 - -solvers = [] -for linsolve in [nothing, LUFactorization(), KrylovJL_GMRES(), KrylovJL_LSMR()] - vjp_autodiffs = linsolve isa KrylovJL ? [nothing, AutoZygote(), AutoFiniteDiff()] : - [nothing] - for linesearch in linesearches, vjp_autodiff in vjp_autodiffs - push!(solvers, GaussNewton(; linsolve, linesearch, vjp_autodiff)) - end -end -append!(solvers, - [LevenbergMarquardt(), LevenbergMarquardt(; linsolve = LUFactorization()), - LevenbergMarquardt(; linsolve = KrylovJL_GMRES()), - LevenbergMarquardt(; linsolve = KrylovJL_LSMR()), nothing]) -for radius_update_scheme in [RadiusUpdateSchemes.Simple, RadiusUpdateSchemes.NocedalWright, - RadiusUpdateSchemes.NLsolve, RadiusUpdateSchemes.Hei, - RadiusUpdateSchemes.Yuan, RadiusUpdateSchemes.Fan, RadiusUpdateSchemes.Bastin] - push!(solvers, TrustRegion(; radius_update_scheme)) -end - -export solvers, θ_init, x, y_target, true_function, θ_true, loss_function -end - -@testitem "General NLLS Solvers" setup=[CoreNLLSTesting] tags=[:core] begin - prob_oop = NonlinearLeastSquaresProblem{false}(loss_function, θ_init, x) - prob_iip = NonlinearLeastSquaresProblem( - NonlinearFunction(loss_function; resid_prototype = zero(y_target)), θ_init, x) - - nlls_problems = [prob_oop, prob_iip] - - for prob in nlls_problems, solver in solvers - sol = solve(prob, solver; maxiters = 10000, abstol = 1e-6) - @test SciMLBase.successful_retcode(sol) - @test norm(sol.resid, 2) < 1e-6 - end -end - -@testitem "Custom VJP" setup=[CoreNLLSTesting] tags=[:core] begin - # This is just for testing that we can use vjp provided by the user - function vjp(v, θ, p) - resid = zeros(length(p)) - J = ForwardDiff.jacobian((resid, θ) -> loss_function(resid, θ, p), resid, θ) - return vec(v' * J) - end - - function vjp!(Jv, v, θ, p) - resid = zeros(length(p)) - J = ForwardDiff.jacobian((resid, θ) -> loss_function(resid, θ, p), resid, θ) - mul!(vec(Jv), transpose(J), v) - return nothing - end - - probs = [ - NonlinearLeastSquaresProblem( - NonlinearFunction{true}( - loss_function; resid_prototype = zero(y_target), vjp = vjp!), - θ_init, - x), - NonlinearLeastSquaresProblem( - NonlinearFunction{false}( - loss_function; resid_prototype = zero(y_target), vjp = vjp), - θ_init, - x)] - - for prob in probs, solver in solvers - sol = solve(prob, solver; maxiters = 10000, abstol = 1e-6) - @test SciMLBase.successful_retcode(sol) - @test norm(sol.resid, 2) < 1e-6 - end -end - -@testitem "NLLS Analytic Jacobian" tags=[:core] begin - dataIn = 1:10 - f(x, p) = x[1] * dataIn .^ 2 .+ x[2] * dataIn .+ x[3] - dataOut = f([1, 2, 3], nothing) + 0.1 * randn(10, 1) - - resid(x, p) = f(x, p) - dataOut - jac(x, p) = [dataIn .^ 2 dataIn ones(10, 1)] - x0 = [1, 1, 1] - - prob = NonlinearLeastSquaresProblem(resid, x0) - sol1 = solve(prob) - - nlfunc = NonlinearFunction(resid; jac) - prob = NonlinearLeastSquaresProblem(nlfunc, x0) - sol2 = solve(prob) - - @test sol1.u ≈ sol2.u -end diff --git a/test/core/rootfind_tests.jl b/test/core/rootfind_tests.jl deleted file mode 100644 index 2d1662570..000000000 --- a/test/core/rootfind_tests.jl +++ /dev/null @@ -1,697 +0,0 @@ -@testsetup module CoreRootfindTesting -using Reexport -@reexport using BenchmarkTools, LinearSolve, NonlinearSolve, StaticArrays, Random, - LinearAlgebra, ForwardDiff, Zygote, Enzyme, SparseConnectivityTracer, - NonlinearSolveBase -using LineSearches: LineSearches - -_nameof(x) = applicable(nameof, x) ? nameof(x) : _nameof(typeof(x)) - -quadratic_f(u, p) = u .* u .- p -quadratic_f!(du, u, p) = (du .= u .* u .- p) -quadratic_f2(u, p) = @. p[1] * u * u - p[2] - -function newton_fails(u, p) - return 0.010000000000000002 .+ - 10.000000000000002 ./ (1 .+ - (0.21640425613334457 .+ - 216.40425613334457 ./ (1 .+ - (0.21640425613334457 .+ - 216.40425613334457 ./ (1 .+ 0.0006250000000000001(u .^ 2.0))) .^ 2.0)) .^ - 2.0) .- 0.0011552453009332421u .- p -end - -const TERMINATION_CONDITIONS = [ - NormTerminationMode(Base.Fix1(maximum, abs)), - RelTerminationMode(), - RelNormTerminationMode(Base.Fix1(maximum, abs)), - RelNormSafeTerminationMode(Base.Fix1(maximum, abs)), - RelNormSafeBestTerminationMode(Base.Fix1(maximum, abs)), - AbsTerminationMode(), - AbsNormTerminationMode(Base.Fix1(maximum, abs)), - AbsNormSafeTerminationMode(Base.Fix1(maximum, abs)), - AbsNormSafeBestTerminationMode(Base.Fix1(maximum, abs)) -] - -function benchmark_nlsolve_oop(f, u0, p = 2.0; solver, kwargs...) - prob = NonlinearProblem{false}(f, u0, p) - return solve(prob, solver; abstol = 1e-9, kwargs...) -end - -function benchmark_nlsolve_iip(f, u0, p = 2.0; solver, kwargs...) - prob = NonlinearProblem{true}(f, u0, p) - return solve(prob, solver; abstol = 1e-9, kwargs...) -end - -function nlprob_iterator_interface(f, p_range, ::Val{iip}, solver) where {iip} - probN = NonlinearProblem{iip}(f, iip ? [0.5] : 0.5, p_range[begin]) - cache = init(probN, solver; maxiters = 100, abstol = 1e-10) - sols = zeros(length(p_range)) - for (i, p) in enumerate(p_range) - reinit!(cache, iip ? [cache.u[1]] : cache.u; p = p) - sol = solve!(cache) - sols[i] = iip ? sol.u[1] : sol.u - end - return sols -end - -for alg in (:Static, :StrongWolfe, :BackTracking, :MoreThuente, :HagerZhang) - algname = Symbol(:LineSearches, alg) - @eval function $(algname)(args...; autodiff = nothing, initial_alpha = true, kwargs...) - return LineSearch.LineSearchesJL(; - method = LineSearches.$(alg)(args...; kwargs...), autodiff, initial_alpha) - end -end - -export nlprob_iterator_interface, benchmark_nlsolve_oop, benchmark_nlsolve_iip, - TERMINATION_CONDITIONS, _nameof, newton_fails, quadratic_f, quadratic_f! -export LineSearchesStatic, LineSearchesStrongWolfe, LineSearchesBackTracking, - LineSearchesMoreThuente, LineSearchesHagerZhang - -end - -# --- NewtonRaphson tests --- - -@testitem "NewtonRaphson" setup=[CoreRootfindTesting] tags=[:core] begin - @testset "LineSearch: $(_nameof(linesearch)) LineSearch AD: $(_nameof(ad))" for ad in ( - AutoForwardDiff(), AutoZygote(), AutoFiniteDiff() - ), - linesearch in ( - LineSearchesStatic(; autodiff = ad), LineSearchesStrongWolfe(; autodiff = ad), - LineSearchesBackTracking(; autodiff = ad), BackTracking(; autodiff = ad), - LineSearchesHagerZhang(; autodiff = ad), - LineSearchesMoreThuente(; autodiff = ad) - ) - - u0s = ([1.0, 1.0], @SVector[1.0, 1.0], 1.0) - - @testset "[OOP] u0: $(typeof(u0))" for u0 in u0s - solver = NewtonRaphson(; linesearch) - sol = benchmark_nlsolve_oop(quadratic_f, u0; solver) - @test SciMLBase.successful_retcode(sol) - @test all(abs.(sol.u .* sol.u .- 2) .< 1e-9) - - cache = init(NonlinearProblem{false}(quadratic_f, u0, 2.0), - NewtonRaphson(), abstol = 1e-9) - @test (@ballocated solve!($cache)) < 200 - end - - precs = [(u0) -> NonlinearSolve.DEFAULT_PRECS, - u0 -> ((args...) -> (Diagonal(rand!(similar(u0))), nothing))] - - @testset "[IIP] u0: $(typeof(u0)) precs: $(_nameof(prec)) linsolve: $(_nameof(linsolve))" for u0 in ([ - 1.0, 1.0],), - prec in precs, - linsolve in (nothing, KrylovJL_GMRES(), \) - - ad isa AutoZygote && continue - if prec === :Random - prec = (args...) -> (Diagonal(randn!(similar(u0))), nothing) - end - solver = NewtonRaphson(; linsolve, precs = prec(u0), linesearch) - sol = benchmark_nlsolve_iip(quadratic_f!, u0; solver) - @test SciMLBase.successful_retcode(sol) - @test all(abs.(sol.u .* sol.u .- 2) .< 1e-9) - - cache = init(NonlinearProblem{true}(quadratic_f!, u0, 2.0), - NewtonRaphson(; linsolve, precs = prec(u0)), abstol = 1e-9) - @test (@ballocated solve!($cache)) ≤ 64 - end - end - - # Iterator interface - p = range(0.01, 2, length = 200) - @test nlprob_iterator_interface(quadratic_f, p, Val(false), NewtonRaphson()) ≈ sqrt.(p) - @test nlprob_iterator_interface(quadratic_f!, p, Val(true), NewtonRaphson()) ≈ sqrt.(p) - - @testset "Sparsity ADType: $(autodiff) u0: $(_nameof(u0))" for autodiff in ( - AutoForwardDiff(), AutoFiniteDiff(), AutoZygote(), AutoEnzyme()), - u0 in (1.0, [1.0, 1.0]) - - probN = NonlinearProblem( - NonlinearFunction(quadratic_f; sparsity = TracerSparsityDetector()), u0, 2.0) - @test all(solve(probN, NewtonRaphson(; autodiff)).u .≈ sqrt(2.0)) - end - - @testset "Termination condition: $(_nameof(termination_condition)) u0: $(_nameof(u0))" for termination_condition in TERMINATION_CONDITIONS, - u0 in (1.0, [1.0, 1.0]) - - probN = NonlinearProblem(quadratic_f, u0, 2.0) - @test all(solve(probN, NewtonRaphson(); termination_condition).u .≈ sqrt(2.0)) - end -end - -# --- TrustRegion tests --- - -@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) $(_nameof(linsolve))" for u0 in u0s, - radius_update_scheme in radius_update_schemes, - linsolve in linear_solvers - - abstol = ifelse(linsolve isa KrylovJL, 1e-6, 1e-9) - - solver = TrustRegion(; radius_update_scheme, linsolve) - sol = benchmark_nlsolve_oop(quadratic_f, u0; solver, abstol) - @test SciMLBase.successful_retcode(sol) - @test all(abs.(sol.u .* sol.u .- 2) .< abstol) - - cache = init(NonlinearProblem{false}(quadratic_f, u0, 2.0), - TrustRegion(; radius_update_scheme, linsolve); abstol) - @test (@ballocated solve!($cache)) < 200 - end - - @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 - - abstol = ifelse(linsolve isa KrylovJL, 1e-6, 1e-9) - solver = TrustRegion(; radius_update_scheme, linsolve) - sol = benchmark_nlsolve_iip(quadratic_f!, u0; solver, abstol) - @test SciMLBase.successful_retcode(sol) - @test all(abs.(sol.u .* sol.u .- 2) .< abstol) - - cache = init(NonlinearProblem{true}(quadratic_f!, u0, 2.0), - TrustRegion(; radius_update_scheme); abstol) - @test (@ballocated solve!($cache)) ≤ 64 - end - - # Iterator interface - p = range(0.01, 2, length = 200) - @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 "$(_nameof(autodiff)) u0: $(_nameof(u0)) $(radius_update_scheme)" for autodiff in ( - AutoForwardDiff(), AutoFiniteDiff(), AutoZygote(), AutoEnzyme()), - u0 in (1.0, [1.0, 1.0]), - radius_update_scheme in radius_update_schemes - - probN = NonlinearProblem( - NonlinearFunction(quadratic_f; sparsity = TracerSparsityDetector()), u0, 2.0) - @test all(solve(probN, TrustRegion(; autodiff, radius_update_scheme)).u .≈ - sqrt(2.0)) - end - - # Test that `TrustRegion` passes a test that `NewtonRaphson` fails on. - @testset "Newton Raphson Fails: radius_update_scheme: $(radius_update_scheme)" for radius_update_scheme in [ - RadiusUpdateSchemes.Simple, RadiusUpdateSchemes.Fan, RadiusUpdateSchemes.Bastin] - u0 = [-10.0, -1.0, 1.0, 2.0, 3.0, 4.0, 10.0] - p = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] - solver = TrustRegion(; radius_update_scheme) - sol = benchmark_nlsolve_oop(newton_fails, u0, p; solver) - @test SciMLBase.successful_retcode(sol) - @test all(abs.(newton_fails(sol.u, p)) .< 1e-9) - end - - # Test kwargs in `TrustRegion` - @testset "Keyword Arguments" begin - max_trust_radius = [10.0, 100.0, 1000.0] - initial_trust_radius = [10.0, 1.0, 0.1] - step_threshold = [0.0, 0.01, 0.25] - shrink_threshold = [0.25, 0.3, 0.5] - expand_threshold = [0.5, 0.8, 0.9] - shrink_factor = [0.1, 0.3, 0.5] - expand_factor = [1.5, 2.0, 3.0] - max_shrink_times = [10, 20, 30] - - list_of_options = zip( - max_trust_radius, initial_trust_radius, step_threshold, shrink_threshold, - expand_threshold, shrink_factor, expand_factor, max_shrink_times) - for options in list_of_options - local probN, sol, alg - alg = TrustRegion( - max_trust_radius = options[1], initial_trust_radius = options[2], - step_threshold = options[3], shrink_threshold = options[4], - expand_threshold = options[5], shrink_factor = options[6], - expand_factor = options[7], max_shrink_times = options[8]) - - probN = NonlinearProblem{false}(quadratic_f, [1.0, 1.0], 2.0) - sol = solve(probN, alg, abstol = 1e-10) - @test all(abs.(quadratic_f(sol.u, 2.0)) .< 1e-10) - end - end - - # Testing consistency of iip vs oop iterations - @testset "OOP / IIP Consistency" begin - maxiterations = [2, 3, 4, 5] - u0 = [1.0, 1.0] - @testset "radius_update_scheme: $(radius_update_scheme) maxiters: $(maxiters)" for radius_update_scheme in radius_update_schemes, - maxiters in maxiterations - - solver = TrustRegion(; radius_update_scheme) - sol_iip = benchmark_nlsolve_iip(quadratic_f!, u0; solver, maxiters) - sol_oop = benchmark_nlsolve_oop(quadratic_f, u0; solver, maxiters) - @test sol_iip.u ≈ sol_oop.u - end - end - - @testset "Termination condition: $(_nameof(termination_condition)) u0: $(_nameof(u0))" for termination_condition in TERMINATION_CONDITIONS, - u0 in (1.0, [1.0, 1.0]) - - probN = NonlinearProblem(quadratic_f, u0, 2.0) - @test all(solve(probN, TrustRegion(); termination_condition).u .≈ sqrt(2.0)) - end -end - -# --- LevenbergMarquardt tests --- - -@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()) - @test SciMLBase.successful_retcode(sol) - @test all(abs.(sol.u .* sol.u .- 2) .< 1e-9) - - cache = init(NonlinearProblem{false}(quadratic_f, u0, 2.0), - LevenbergMarquardt(), abstol = 1e-9) - @test (@ballocated solve!($cache)) < 200 - end - - @testset "[IIP] u0: $(typeof(u0))" for u0 in ([1.0, 1.0],) - sol = benchmark_nlsolve_iip(quadratic_f!, u0; solver = LevenbergMarquardt()) - @test SciMLBase.successful_retcode(sol) - @test all(abs.(sol.u .* sol.u .- 2) .< 1e-9) - - cache = init(NonlinearProblem{true}(quadratic_f!, u0, 2.0), - LevenbergMarquardt(), abstol = 1e-9) - @test (@ballocated solve!($cache)) ≤ 64 - end - - @testset "ADType: $(autodiff) u0: $(_nameof(u0))" for autodiff in ( - AutoForwardDiff(), AutoFiniteDiff(), AutoZygote(), AutoEnzyme()), - u0 in (1.0, [1.0, 1.0]) - - probN = NonlinearProblem( - NonlinearFunction(quadratic_f; sparsity = TracerSparsityDetector()), u0, 2.0) - @test all(solve( - probN, LevenbergMarquardt(; autodiff); abstol = 1e-9, reltol = 1e-9).u .≈ - sqrt(2.0)) - end - - # Test that `LevenbergMarquardt` passes a test that `NewtonRaphson` fails on. - @testset "Newton Raphson Fails" begin - u0 = [-10.0, -1.0, 1.0, 2.0, 3.0, 4.0, 10.0] - p = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] - sol = benchmark_nlsolve_oop(newton_fails, u0, p; solver = LevenbergMarquardt()) - @test SciMLBase.successful_retcode(sol) - @test all(abs.(newton_fails(sol.u, p)) .< 1e-9) - end - - # Iterator interface - p = range(0.01, 2, length = 200) - @test abs.(nlprob_iterator_interface( - quadratic_f, p, Val(false), LevenbergMarquardt())) ≈ sqrt.(p) - @test abs.(nlprob_iterator_interface( - quadratic_f!, p, Val(true), LevenbergMarquardt())) ≈ sqrt.(p) - - # Test kwargs in `LevenbergMarquardt` - @testset "Keyword Arguments" begin - damping_initial = [0.5, 2.0, 5.0] - damping_increase_factor = [1.5, 3.0, 10.0] - damping_decrease_factor = Float64[2, 5, 10.0] - finite_diff_step_geodesic = [0.02, 0.2, 0.3] - α_geodesic = [0.6, 0.8, 0.9] - b_uphill = Float64[0, 1, 2] - min_damping_D = [1e-12, 1e-9, 1e-4] - - list_of_options = zip( - damping_initial, damping_increase_factor, damping_decrease_factor, - finite_diff_step_geodesic, α_geodesic, b_uphill, min_damping_D) - for options in list_of_options - local probN, sol, alg - alg = LevenbergMarquardt(; - damping_initial = options[1], damping_increase_factor = options[2], - damping_decrease_factor = options[3], - finite_diff_step_geodesic = options[4], α_geodesic = options[5], - b_uphill = options[6], min_damping_D = options[7]) - - probN = NonlinearProblem{false}(quadratic_f, [1.0, 1.0], 2.0) - sol = solve(probN, alg; abstol = 1e-13, maxiters = 10000) - @test all(abs.(quadratic_f(sol.u, 2.0)) .< 1e-10) - end - end - - @testset "Termination condition: $(_nameof(termination_condition)) u0: $(_nameof(u0))" for termination_condition in TERMINATION_CONDITIONS, - u0 in (1.0, [1.0, 1.0]) - - probN = NonlinearProblem(quadratic_f, u0, 2.0) - @test all(solve(probN, LevenbergMarquardt(); termination_condition).u .≈ sqrt(2.0)) - end -end - -# --- DFSane tests --- - -@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 - sol = benchmark_nlsolve_oop(quadratic_f, u0; solver = DFSane()) - @test SciMLBase.successful_retcode(sol) - @test all(abs.(sol.u .* sol.u .- 2) .< 1e-9) - - cache = init(NonlinearProblem{false}(quadratic_f, u0, 2.0), DFSane(), abstol = 1e-9) - @test (@ballocated solve!($cache)) < 200 - end - - @testset "[IIP] u0: $(typeof(u0))" for u0 in ([1.0, 1.0],) - sol = benchmark_nlsolve_iip(quadratic_f!, u0; solver = DFSane()) - @test SciMLBase.successful_retcode(sol) - @test all(abs.(sol.u .* sol.u .- 2) .< 1e-9) - - cache = init(NonlinearProblem{true}(quadratic_f!, u0, 2.0), DFSane(), abstol = 1e-9) - @test (@ballocated solve!($cache)) ≤ 64 - end - - # Iterator interface - p = range(0.01, 2, length = 200) - @test abs.(nlprob_iterator_interface(quadratic_f, p, Val(false), DFSane())) ≈ sqrt.(p) - @test abs.(nlprob_iterator_interface(quadratic_f!, p, Val(true), DFSane())) ≈ sqrt.(p) - - # Test that `DFSane` passes a test that `NewtonRaphson` fails on. - @testset "Newton Raphson Fails" begin - u0 = [-10.0, -1.0, 1.0, 2.0, 3.0, 4.0, 10.0] - p = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] - sol = benchmark_nlsolve_oop(newton_fails, u0, p; solver = DFSane()) - @test SciMLBase.successful_retcode(sol) - @test all(abs.(newton_fails(sol.u, p)) .< 1e-9) - end - - # Test kwargs in `DFSane` - @testset "Keyword Arguments" begin - σ_min = [1e-10, 1e-5, 1e-4] - σ_max = [1e10, 1e5, 1e4] - σ_1 = [1.0, 0.5, 2.0] - M = [10, 1, 100] - γ = [1e-4, 1e-3, 1e-5] - τ_min = [0.1, 0.2, 0.3] - τ_max = [0.5, 0.8, 0.9] - nexp = [2, 1, 2] - η_strategy = [(f_1, k, x, F) -> f_1 / k^2, (f_1, k, x, F) -> f_1 / k^3, - (f_1, k, x, F) -> f_1 / k^4] - - list_of_options = zip(σ_min, σ_max, σ_1, M, γ, τ_min, τ_max, nexp, η_strategy) - for options in list_of_options - local probN, sol, alg - alg = DFSane(σ_min = options[1], σ_max = options[2], σ_1 = options[3], - M = options[4], γ = options[5], τ_min = options[6], - τ_max = options[7], n_exp = options[8], η_strategy = options[9]) - - probN = NonlinearProblem{false}(quadratic_f, [1.0, 1.0], 2.0) - sol = solve(probN, alg, abstol = 1e-11) - @test all(abs.(quadratic_f(sol.u, 2.0)) .< 1e-6) - end - end - - @testset "Termination condition: $(_nameof(termination_condition)) u0: $(_nameof(u0))" for termination_condition in TERMINATION_CONDITIONS, - u0 in (1.0, [1.0, 1.0]) - - probN = NonlinearProblem(quadratic_f, u0, 2.0) - @test all(solve(probN, DFSane(); termination_condition).u .≈ sqrt(2.0)) - end -end - -# --- PseudoTransient tests --- - -@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 ( - AutoFiniteDiff(), AutoZygote()) - u0s = ([1.0, 1.0], @SVector[1.0, 1.0], 1.0) - - @testset "[OOP] u0: $(typeof(u0))" for u0 in u0s - solver = PseudoTransient(; alpha_initial = 10.0) - sol = benchmark_nlsolve_oop(quadratic_f, u0; solver) - # Failing by a margin for some - # @test SciMLBase.successful_retcode(sol) - @test all(abs.(sol.u .* sol.u .- 2) .< 1e-9) - - cache = init(NonlinearProblem{false}(quadratic_f, u0, 2.0), - PseudoTransient(alpha_initial = 10.0), abstol = 1e-9) - @test (@ballocated solve!($cache)) < 200 - end - - precs = [NonlinearSolve.DEFAULT_PRECS, :Random] - - @testset "[IIP] u0: $(typeof(u0)) precs: $(_nameof(prec)) linsolve: $(_nameof(linsolve))" for u0 in ([ - 1.0, 1.0],), - prec in precs, - linsolve in (nothing, KrylovJL_GMRES(), \) - - ad isa AutoZygote && continue - if prec === :Random - prec = (args...) -> (Diagonal(randn!(similar(u0))), nothing) - end - solver = PseudoTransient(; alpha_initial = 10.0, linsolve, precs = prec) - sol = benchmark_nlsolve_iip(quadratic_f!, u0; solver) - @test SciMLBase.successful_retcode(sol) - @test all(abs.(sol.u .* sol.u .- 2) .< 1e-9) - - cache = init( - NonlinearProblem{true}(quadratic_f!, u0, 2.0), solver; abstol = 1e-9) - @test (@ballocated solve!($cache)) ≤ 64 - end - end - - p = range(0.01, 2, length = 200) - @test nlprob_iterator_interface( - quadratic_f, p, Val(false), PseudoTransient(; alpha_initial = 10.0)) ≈ sqrt.(p) - @test nlprob_iterator_interface( - quadratic_f!, p, Val(true), PseudoTransient(; alpha_initial = 10.0)) ≈ sqrt.(p) - - @testset "ADType: $(autodiff) u0: $(_nameof(u0))" for autodiff in ( - AutoForwardDiff(), AutoFiniteDiff(), AutoZygote(), AutoEnzyme()), - u0 in (1.0, [1.0, 1.0]) - - probN = NonlinearProblem( - NonlinearFunction(quadratic_f; sparsity = TracerSparsityDetector()), u0, 2.0) - @test all(solve(probN, PseudoTransient(; alpha_initial = 10.0, autodiff)).u .≈ - sqrt(2.0)) - end - - @testset "Termination condition: $(_nameof(termination_condition)) u0: $(_nameof(u0))" for termination_condition in TERMINATION_CONDITIONS, - u0 in (1.0, [1.0, 1.0]) - - probN = NonlinearProblem(quadratic_f, u0, 2.0) - @test all(solve( - probN, PseudoTransient(; alpha_initial = 10.0); termination_condition).u .≈ - sqrt(2.0)) - end -end - -# --- Broyden tests --- - -@testitem "Broyden" setup=[CoreRootfindTesting] tags=[:core] begin - @testset "LineSearch: $(_nameof(linesearch)) LineSearch AD: $(_nameof(ad)) Init Jacobian: $(init_jacobian) Update Rule: $(update_rule)" for ad in ( - AutoForwardDiff(), AutoZygote(), AutoFiniteDiff() - ), - linesearch in ( - LineSearchesStatic(; autodiff = ad), LineSearchesStrongWolfe(; autodiff = ad), - LineSearchesBackTracking(; autodiff = ad), BackTracking(; autodiff = ad), - LineSearchesHagerZhang(; autodiff = ad), - LineSearchesMoreThuente(; autodiff = ad) - ), - init_jacobian in (Val(:identity), Val(:true_jacobian)), - update_rule in (Val(:good_broyden), Val(:bad_broyden), Val(:diagonal)) - - u0s = ([1.0, 1.0], @SVector[1.0, 1.0], 1.0) - - @testset "[OOP] u0: $(typeof(u0))" for u0 in u0s - solver = Broyden(; linesearch, init_jacobian, update_rule) - sol = benchmark_nlsolve_oop(quadratic_f, u0; solver) - @test SciMLBase.successful_retcode(sol) - @test all(abs.(sol.u .* sol.u .- 2) .< 1e-9) - - cache = init(NonlinearProblem{false}(quadratic_f, u0, 2.0), - Broyden(; linesearch, update_rule, init_jacobian), abstol = 1e-9) - @test (@ballocated solve!($cache)) < 200 - end - - @testset "[IIP] u0: $(typeof(u0))" for u0 in ([1.0, 1.0],) - ad isa AutoZygote && continue - solver = Broyden(; linesearch, init_jacobian, update_rule) - sol = benchmark_nlsolve_iip(quadratic_f!, u0; solver) - @test SciMLBase.successful_retcode(sol) - @test all(abs.(sol.u .* sol.u .- 2) .< 1e-9) - - cache = init(NonlinearProblem{true}(quadratic_f!, u0, 2.0), - Broyden(; linesearch, update_rule, init_jacobian), abstol = 1e-9) - @test (@ballocated solve!($cache)) ≤ 64 - end - end - - # Iterator interface - p = range(0.01, 2, length = 200) - @test nlprob_iterator_interface(quadratic_f, p, Val(false), Broyden()) ≈ sqrt.(p) - @test nlprob_iterator_interface(quadratic_f!, p, Val(true), Broyden()) ≈ sqrt.(p) - - @testset "Termination condition: $(_nameof(termination_condition)) u0: $(_nameof(u0))" for termination_condition in TERMINATION_CONDITIONS, - u0 in (1.0, [1.0, 1.0]) - - probN = NonlinearProblem(quadratic_f, u0, 2.0) - @test all(solve(probN, Broyden(); termination_condition).u .≈ sqrt(2.0)) - end -end - -# --- Klement tests --- - -@testitem "Klement" setup=[CoreRootfindTesting] tags=[:core] begin - @testset "LineSearch: $(_nameof(linesearch)) LineSearch AD: $(_nameof(ad)) Init Jacobian: $(init_jacobian)" for ad in ( - AutoForwardDiff(), AutoZygote(), AutoFiniteDiff() - ), - linesearch in ( - LineSearchesStatic(; autodiff = ad), LineSearchesStrongWolfe(; autodiff = ad), - LineSearchesBackTracking(; autodiff = ad), BackTracking(; autodiff = ad), - LineSearchesHagerZhang(; autodiff = ad), - LineSearchesMoreThuente(; autodiff = ad) - ), - init_jacobian in (Val(:identity), Val(:true_jacobian), Val(:true_jacobian_diagonal)) - - u0s = ([1.0, 1.0], @SVector[1.0, 1.0], 1.0) - - @testset "[OOP] u0: $(typeof(u0))" for u0 in u0s - solver = Klement(; linesearch, init_jacobian) - sol = benchmark_nlsolve_oop(quadratic_f, u0; solver) - # Some are failing by a margin - # @test SciMLBase.successful_retcode(sol) - @test all(abs.(sol.u .* sol.u .- 2) .< 3e-9) - - cache = init(NonlinearProblem{false}(quadratic_f, u0, 2.0), - Klement(; linesearch), abstol = 1e-9) - @test (@ballocated solve!($cache)) < 200 - end - - @testset "[IIP] u0: $(typeof(u0))" for u0 in ([1.0, 1.0],) - ad isa AutoZygote && continue - solver = Klement(; linesearch, init_jacobian) - sol = benchmark_nlsolve_iip(quadratic_f!, u0; solver) - @test SciMLBase.successful_retcode(sol) - @test all(abs.(sol.u .* sol.u .- 2) .< 1e-9) - - cache = init(NonlinearProblem{true}(quadratic_f!, u0, 2.0), - Klement(; linesearch), abstol = 1e-9) - @test (@ballocated solve!($cache)) ≤ 64 - end - end - - # Iterator interface - p = range(0.01, 2, length = 200) - @test nlprob_iterator_interface(quadratic_f, p, Val(false), Klement()) ≈ sqrt.(p) - @test nlprob_iterator_interface(quadratic_f!, p, Val(true), Klement()) ≈ sqrt.(p) - - @testset "Termination condition: $(_nameof(termination_condition)) u0: $(_nameof(u0))" for termination_condition in TERMINATION_CONDITIONS, - u0 in (1.0, [1.0, 1.0]) - - probN = NonlinearProblem(quadratic_f, u0, 2.0) - @test all(solve(probN, Klement(); termination_condition).u .≈ sqrt(2.0)) - end -end - -# --- LimitedMemoryBroyden tests --- - -@testitem "LimitedMemoryBroyden" setup=[CoreRootfindTesting] tags=[:core] begin - @testset "LineSearch: $(_nameof(linesearch)) LineSearch AD: $(_nameof(ad))" for ad in ( - AutoForwardDiff(), AutoZygote(), AutoFiniteDiff() - ), - linesearch in ( - LineSearchesStatic(; autodiff = ad), LineSearchesStrongWolfe(; autodiff = ad), - LineSearchesBackTracking(; autodiff = ad), BackTracking(; autodiff = ad), - LineSearchesHagerZhang(; autodiff = ad), - LineSearchesMoreThuente(; autodiff = ad), LiFukushimaLineSearch() - ) - - u0s = ([1.0, 1.0], @SVector[1.0, 1.0], 1.0) - - @testset "[OOP] u0: $(typeof(u0))" for u0 in u0s - broken = linesearch isa BackTracking && ad isa AutoFiniteDiff && u0 isa Vector - - solver = LimitedMemoryBroyden(; linesearch) - sol = benchmark_nlsolve_oop(quadratic_f, u0; solver) - @test SciMLBase.successful_retcode(sol) broken=broken - @test all(abs.(sol.u .* sol.u .- 2) .< 1e-9) broken=broken - - cache = init(NonlinearProblem{false}(quadratic_f, u0, 2.0), - LimitedMemoryBroyden(; linesearch), abstol = 1e-9) - @test (@ballocated solve!($cache)) < 200 - end - - @testset "[IIP] u0: $(typeof(u0))" for u0 in ([1.0, 1.0],) - broken = linesearch isa BackTracking && ad isa AutoFiniteDiff && u0 isa Vector - ad isa AutoZygote && continue - - solver = LimitedMemoryBroyden(; linesearch) - sol = benchmark_nlsolve_iip(quadratic_f!, u0; solver) - @test SciMLBase.successful_retcode(sol) broken=broken - @test all(abs.(sol.u .* sol.u .- 2) .< 1e-9) broken=broken - - cache = init(NonlinearProblem{true}(quadratic_f!, u0, 2.0), - LimitedMemoryBroyden(; linesearch), abstol = 1e-9) - @test (@ballocated solve!($cache)) ≤ 64 - end - end - - # Iterator interface - p = range(0.01, 2, length = 200) - @test nlprob_iterator_interface( - quadratic_f, p, Val(false), LimitedMemoryBroyden())≈sqrt.(p) atol=1e-2 - @test nlprob_iterator_interface( - quadratic_f!, p, Val(true), LimitedMemoryBroyden())≈sqrt.(p) atol=1e-2 - - @testset "Termination condition: $(_nameof(termination_condition)) u0: $(_nameof(u0))" for termination_condition in TERMINATION_CONDITIONS, - u0 in (1.0, [1.0, 1.0]) - - probN = NonlinearProblem(quadratic_f, u0, 2.0) - @test all(solve(probN, LimitedMemoryBroyden(); termination_condition).u .≈ - sqrt(2.0)) - end -end - -# Miscellaneous Tests -@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 - end - - function F!(du::Vector{Float64}, u::Vector{Float64}, p::Vector{Float64}) - Δ = Tridiagonal(-ones(99), 2 * ones(100), -ones(99)) - du .= u + 0.1 * u .* Δ * u - p - return nothing - end - - function JVP(v::Vector{Float64}, u::Vector{Float64}, p::Vector{Float64}) - Δ = Tridiagonal(-ones(99), 2 * ones(100), -ones(99)) - return v + 0.1 * (u .* Δ * v + v .* Δ * u) - end - - function JVP!( - du::Vector{Float64}, v::Vector{Float64}, u::Vector{Float64}, p::Vector{Float64}) - Δ = Tridiagonal(-ones(99), 2 * ones(100), -ones(99)) - du .= v + 0.1 * (u .* Δ * v + v .* Δ * u) - return nothing - end - - u0 = rand(100) - - prob = NonlinearProblem(NonlinearFunction{false}(F; jvp = JVP), u0, u0) - sol = solve(prob, NewtonRaphson(; linsolve = KrylovJL_GMRES()); abstol = 1e-13) - @test norm(sol.resid, Inf) ≤ 1e-6 - sol = solve( - prob, TrustRegion(; linsolve = KrylovJL_GMRES(), vjp_autodiff = AutoFiniteDiff()); - abstol = 1e-13) - @test norm(sol.resid, Inf) ≤ 1e-6 - - prob = NonlinearProblem(NonlinearFunction{true}(F!; jvp = JVP!), u0, u0) - sol = solve(prob, NewtonRaphson(; linsolve = KrylovJL_GMRES()); abstol = 1e-13) - @test norm(sol.resid, Inf) ≤ 1e-6 - sol = solve( - prob, TrustRegion(; linsolve = KrylovJL_GMRES(), vjp_autodiff = AutoFiniteDiff()); - abstol = 1e-13) - @test norm(sol.resid, Inf) ≤ 1e-6 -end diff --git a/test/core_tests.jl b/test/core_tests.jl new file mode 100644 index 000000000..301b0b389 --- /dev/null +++ b/test/core_tests.jl @@ -0,0 +1,416 @@ +@testitem "NLLS Analytic Jacobian" tags=[:core] begin + dataIn = 1:10 + f(x, p) = x[1] * dataIn .^ 2 .+ x[2] * dataIn .+ x[3] + dataOut = f([1, 2, 3], nothing) + 0.1 * randn(10, 1) + + resid(x, p) = f(x, p) - dataOut + jac(x, p) = [dataIn .^ 2 dataIn ones(10, 1)] + x0 = [1, 1, 1] + + prob = NonlinearLeastSquaresProblem(resid, x0) + sol1 = solve(prob) + + nlfunc = NonlinearFunction(resid; jac) + prob = NonlinearLeastSquaresProblem(nlfunc, x0) + sol2 = solve(prob) + + @test sol1.u ≈ sol2.u +end + +@testitem "Basic PolyAlgorithms" tags=[:core] begin + f(u, p) = u .* u .- 2 + u0 = [1.0, 1.0] + + prob = NonlinearProblem(f, u0) + + polyalgs = ( + RobustMultiNewton(), FastShortcutNonlinearPolyalg(), nothing, missing, + NonlinearSolvePolyAlgorithm((Broyden(), LimitedMemoryBroyden())) + ) + + @testset "Direct Solve" begin + @testset for alg in polyalgs + alg = alg === missing ? () : (alg,) + sol = solve(prob, alg...; abstol = 1e-9) + @test SciMLBase.successful_retcode(sol) + err = maximum(abs, f(sol.u, 2.0)) + @test err < 1e-9 + end + end + + @testset "Caching Interface" begin + @testset for alg in polyalgs + alg = alg === missing ? () : (alg,) + cache = init(prob, alg...; abstol = 1e-9) + solver = solve!(cache) + @test SciMLBase.successful_retcode(solver) + end + end + + @testset "Step Interface" begin + @testset for alg in polyalgs + alg = alg === missing ? () : (alg,) + cache = init(prob, alg...; abstol = 1e-9) + for i in 1:10000 + step!(cache) + cache.force_stop && break + end + @test SciMLBase.successful_retcode(cache.retcode) + end + end +end + +@testitem "PolyAlgorithms Autodiff" tags=[:core] begin + cache = zeros(2) + function f(du, u, p) + cache .= u .* u + du .= cache .- 2 + end + u0 = [1.0, 1.0] + probN = NonlinearProblem{true}(f, u0) + + custom_polyalg = NonlinearSolvePolyAlgorithm(( + Broyden(; autodiff = AutoFiniteDiff()), LimitedMemoryBroyden()) + ) + + # Uses the `__solve` function + @test_throws MethodError solve(probN; abstol = 1e-9) + @test_throws MethodError solve(probN, RobustMultiNewton()) + + sol = solve(probN, RobustMultiNewton(; autodiff = AutoFiniteDiff())) + @test SciMLBase.successful_retcode(sol) + + sol = solve( + probN, FastShortcutNonlinearPolyalg(; autodiff = AutoFiniteDiff()); abstol = 1e-9 + ) + @test SciMLBase.successful_retcode(sol) + + quadratic_f(u::Float64, p) = u^2 - p + + prob = NonlinearProblem(quadratic_f, 2.0, 4.0) + + @test_throws MethodError solve(prob) + @test_throws MethodError solve(prob, RobustMultiNewton()) + + sol = solve(prob, RobustMultiNewton(; autodiff = AutoFiniteDiff())) + @test SciMLBase.successful_retcode(sol) +end + +@testitem "PolyAlgorithm Aliasing" tags=[:core] begin + using NonlinearProblemLibrary + + # Use a problem that the initial solvers cannot solve and cause the initial value to + # diverge. If we don't alias correctly, all the subsequent algorithms will also fail. + prob = NonlinearProblemLibrary.nlprob_23_testcases["Generalized Rosenbrock function"].prob + u0 = copy(prob.u0) + prob = remake(prob; u0 = copy(u0)) + + # If aliasing is not handled properly this will diverge + sol = solve( + prob; abstol = 1e-6, alias_u0 = true, + termination_condition = AbsNormTerminationMode(Base.Fix1(maximum, abs)) + ) + + @test sol.u === prob.u0 + @test SciMLBase.successful_retcode(sol.retcode) + + prob = remake(prob; u0 = copy(u0)) + + cache = init( + prob; abstol = 1e-6, alias_u0 = true, + termination_condition = AbsNormTerminationMode(Base.Fix1(maximum, abs)) + ) + sol = solve!(cache) + + @test sol.u === prob.u0 + @test SciMLBase.successful_retcode(sol.retcode) +end + +@testitem "Ensemble Nonlinear Problems" tags=[:core] begin + prob_func(prob, i, repeat) = remake(prob; u0 = prob.u0[:, i]) + + prob_nls_oop = NonlinearProblem((u, p) -> u .* u .- p, rand(4, 128), 2.0) + prob_nls_iip = NonlinearProblem((du, u, p) -> du .= u .* u .- p, rand(4, 128), 2.0) + prob_nlls_oop = NonlinearLeastSquaresProblem((u, p) -> u .^ 2 .- p, rand(4, 128), 2.0) + prob_nlls_iip = NonlinearLeastSquaresProblem( + NonlinearFunction{true}((du, u, p) -> du .= u .^ 2 .- p; resid_prototype = rand(4)), + rand(4, 128), 2.0 + ) + + for prob in (prob_nls_oop, prob_nls_iip, prob_nlls_oop, prob_nlls_iip) + ensembleprob = EnsembleProblem(prob; prob_func) + + for ensemblealg in (EnsembleThreads(), EnsembleSerial()) + sim = solve(ensembleprob, nothing, ensemblealg; trajectories = size(prob.u0, 2)) + @test all(SciMLBase.successful_retcode, sim.u) + end + end +end + +@testitem "BigFloat Support" tags=[:core] begin + using LinearAlgebra + + fn_iip = NonlinearFunction{true}((du, u, p) -> du .= u .* u .- p) + fn_oop = NonlinearFunction{false}((u, p) -> u .* u .- p) + + u0 = BigFloat[1.0, 1.0, 1.0] + prob_iip_bf = NonlinearProblem{true}(fn_iip, u0, BigFloat(2)) + prob_oop_bf = NonlinearProblem{false}(fn_oop, u0, BigFloat(2)) + + for alg in (NewtonRaphson(), Broyden(), Klement(), DFSane(), TrustRegion()) + sol = solve(prob_oop_bf, alg) + @test norm(sol.resid, Inf) < 1e-6 + @test SciMLBase.successful_retcode(sol.retcode) + + sol = solve(prob_iip_bf, alg) + @test norm(sol.resid, Inf) < 1e-6 + @test SciMLBase.successful_retcode(sol.retcode) + end +end + +@testitem "Singular Exception: Issue #153" tags=[:core] begin + function f(du, u, p) + s1, s1s2, s2 = u + k1, c1, Δt = p + + du[1] = -0.25 * c1 * k1 * s1 * s2 + du[2] = 0.25 * c1 * k1 * s1 * s2 + du[3] = -0.25 * c1 * k1 * s1 * s2 + end + + prob = NonlinearProblem(f, [2.0, 2.0, 2.0], [1.0, 2.0, 2.5]) + sol = solve(prob; abstol = 1e-9) + @test SciMLBase.successful_retcode(sol) +end + +@testitem "Simple Scalar Problem: Issue #187" tags=[:core] begin + using NaNMath + + # https://github.com/SciML/NonlinearSolve.jl/issues/187 + # If we use a General Nonlinear Solver the solution might go out of the domain! + ff_interval(u, p) = 0.5 / 1.5 * NaNMath.log.(u ./ (1.0 .- u)) .- 2.0 * u .+ 1.0 + + uspan = (0.02, 0.1) + prob = IntervalNonlinearProblem(ff_interval, uspan) + sol = solve(prob; abstol = 1e-9) + @test SciMLBase.successful_retcode(sol) + + u0 = 0.06 + p = 2.0 + prob = NonlinearProblem(ff_interval, u0, p) + sol = solve(prob; abstol = 1e-9) + @test SciMLBase.successful_retcode(sol) +end + +# Shooting Problem: Taken from BoundaryValueDiffEq.jl +# Testing for Complex Valued Root Finding. For Complex valued inputs we drop some of the +# algorithms which dont support those. +@testitem "Complex Valued Problems: Single-Shooting" tags=[:core] begin + using OrdinaryDiffEqTsit5 + + function ode_func!(du, u, p, t) + du[1] = u[2] + du[2] = -u[1] + return nothing + end + + function objective_function!(resid, u0, p) + odeprob = ODEProblem{true}(ode_func!, u0, (0.0, 100.0), p) + sol = solve(odeprob, Tsit5(), abstol = 1e-9, reltol = 1e-9, verbose = false) + resid[1] = sol(0.0)[1] + resid[2] = sol(100.0)[1] - 1.0 + return nothing + end + + prob = NonlinearProblem{true}(objective_function!, [0.0, 1.0] .+ 1im) + sol = solve(prob; abstol = 1e-10) + @test SciMLBase.successful_retcode(sol) + # This test is not meant to return success but test that all the default solvers can handle + # complex valued problems + @test_nowarn solve(prob; abstol = 1e-19, maxiters = 10) + @test_nowarn solve( + prob, RobustMultiNewton(eltype(prob.u0)); abstol = 1e-19, maxiters = 10 + ) +end + +@testitem "No AD" tags=[:core] begin + no_ad_fast = FastShortcutNonlinearPolyalg(autodiff = AutoFiniteDiff()) + no_ad_robust = RobustMultiNewton(autodiff = AutoFiniteDiff()) + no_ad_algs = Set([no_ad_fast, no_ad_robust, no_ad_fast.algs..., no_ad_robust.algs...]) + + @testset "Inplace" begin + f_iip = Base.Experimental.@opaque (du, u, p) -> du .= u .* u .- p + u0 = [0.5] + prob = NonlinearProblem(f_iip, u0, 1.0) + for alg in no_ad_algs + sol = solve(prob, alg) + @test isapprox(only(sol.u), 1.0) + @test SciMLBase.successful_retcode(sol.retcode) + end + end + + @testset "Out of Place" begin + f_oop = Base.Experimental.@opaque (u, p) -> u .* u .- p + u0 = [0.5] + prob = NonlinearProblem{false}(f_oop, u0, 1.0) + for alg in no_ad_algs + sol = solve(prob, alg) + @test isapprox(only(sol.u), 1.0) + @test SciMLBase.successful_retcode(sol.retcode) + end + end +end + +@testitem "Infeasible" tags=[:core] begin + using LinearAlgebra, StaticArrays + + # this is infeasible + function f1_infeasible!(out, u, p) + μ = 3.986004415e14 + x = 7000.0e3 + y = -6.970561549987071e-9 + z = -3.784706123246018e-9 + v_x = 8.550491684548064e-12 + u[1] + v_y = 6631.60076191005 + u[2] + v_z = 3600.665431405663 + u[3] + r = @SVector [x, y, z] + v = @SVector [v_x, v_y, v_z] + h = cross(r, v) + ev = cross(v, h) / μ - r / norm(r) + i = acos(h[3] / norm(h)) + e = norm(ev) + a = 1 / (2 / norm(r) - (norm(v)^2 / μ)) + out .= [a - 42.0e6, e - 1e-5, i - 1e-5] + return nothing + end + + function f1_infeasible(u, p) + μ = 3.986004415e14 + x = 7000.0e3 + y = -6.970561549987071e-9 + z = -3.784706123246018e-9 + v_x = 8.550491684548064e-12 + u[1] + v_y = 6631.60076191005 + u[2] + v_z = 3600.665431405663 + u[3] + r = [x, y, z] + v = [v_x, v_y, v_z] + h = cross(r, v) + ev = cross(v, h) / μ - r / norm(r) + i = acos(h[3] / norm(h)) + e = norm(ev) + a = 1 / (2 / norm(r) - (norm(v)^2 / μ)) + return [a - 42.0e6, e - 1e-5, i - 1e-5] + end + + u0 = [0.0, 0.0, 0.0] + prob = NonlinearProblem(f1_infeasible!, u0) + sol = solve(prob) + + @test all(!isnan, sol.u) + @test !SciMLBase.successful_retcode(sol.retcode) + @inferred solve(prob) + + u0 = [0.0, 0.0, 0.0] + prob = NonlinearProblem(f1_infeasible, u0) + sol = solve(prob) + + @test all(!isnan, sol.u) + @test !SciMLBase.successful_retcode(sol.retcode) + @inferred solve(prob) + + u0 = @SVector [0.0, 0.0, 0.0] + prob = NonlinearProblem(f1_infeasible, u0) + + sol = solve(prob) + @test all(!isnan, sol.u) + @test !SciMLBase.successful_retcode(sol.retcode) +end + +@testitem "NoInit Caching" tags=[:core] begin + using LinearAlgebra + + solvers = [ + SimpleNewtonRaphson(), SimpleTrustRegion(), SimpleDFSane() + ] + + prob = NonlinearProblem((u, p) -> u .^ 2 .- p, [0.1, 0.3], 2.0) + + for alg in solvers + cache = init(prob, alg) + sol = solve!(cache) + @test SciMLBase.successful_retcode(sol) + @test norm(sol.resid, Inf) ≤ 1e-6 + + reinit!(cache; p = 5.0) + @test cache.prob.p == 5.0 + sol = solve!(cache) + @test SciMLBase.successful_retcode(sol) + @test norm(sol.resid, Inf) ≤ 1e-6 + @test norm(sol.u .^ 2 .- 5.0, Inf) ≤ 1e-6 + end +end + +@testitem "Out-of-place Matrix Resizing" tags=[:core] begin + using StableRNGs + + ff(u, p) = u .* u .- p + u0 = rand(StableRNG(0), 2, 2) + p = 2.0 + vecprob = NonlinearProblem(ff, vec(u0), p) + prob = NonlinearProblem(ff, u0, p) + + for alg in ( + NewtonRaphson(), TrustRegion(), LevenbergMarquardt(), + PseudoTransient(), RobustMultiNewton(), FastShortcutNonlinearPolyalg(), + Broyden(), Klement(), LimitedMemoryBroyden(; threshold = 2) + ) + @test vec(solve(prob, alg).u) == solve(vecprob, alg).u + end +end + +@testitem "Inplace Matrix Resizing" tags=[:core] begin + using StableRNGs + + fiip(du, u, p) = (du .= u .* u .- p) + u0 = rand(StableRNG(0), 2, 2) + p = 2.0 + vecprob = NonlinearProblem(fiip, vec(u0), p) + prob = NonlinearProblem(fiip, u0, p) + + for alg in ( + NewtonRaphson(), TrustRegion(), LevenbergMarquardt(), + PseudoTransient(), RobustMultiNewton(), FastShortcutNonlinearPolyalg(), + Broyden(), Klement(), LimitedMemoryBroyden(; threshold = 2) + ) + @test vec(solve(prob, alg).u) == solve(vecprob, alg).u + end +end + +@testitem "Singular Systems -- Auto Linear Solve Switching" tags=[:core] begin + using LinearAlgebra + + function f!(du, u, p) + du[1] = 2u[1] - 2 + du[2] = (u[1] - 4u[2])^2 + 0.1 + end + + u0 = [0.0, 0.0] # Singular Jacobian at u0 + + prob = NonlinearProblem(f!, u0) + + sol = solve(prob) # This doesn't have a root so let's just test the switching + @test sol.u≈[1.0, 0.25] atol=1e-3 rtol=1e-3 + + function nlls!(du, u, p) + du[1] = 2u[1] - 2 + du[2] = (u[1] - 4u[2])^2 + 0.1 + du[3] = 0 + end + + u0 = [0.0, 0.0] + + prob = NonlinearProblem(NonlinearFunction(nlls!, resid_prototype = zeros(3)), u0) + + solve(prob) + @test sol.u≈[1.0, 0.25] atol=1e-3 rtol=1e-3 +end diff --git a/test/gpu/core_tests.jl b/test/cuda_tests.jl similarity index 93% rename from test/gpu/core_tests.jl rename to test/cuda_tests.jl index 91d6178a4..bd992e3fc 100644 --- a/test/gpu/core_tests.jl +++ b/test/cuda_tests.jl @@ -15,8 +15,7 @@ SOLVERS = ( NewtonRaphson(), LevenbergMarquardt(; linsolve = QRFactorization()), - # XXX: Fails currently - # LevenbergMarquardt(; linsolve = KrylovJL_GMRES()), + LevenbergMarquardt(; linsolve = KrylovJL_GMRES()), PseudoTransient(), Klement(), Broyden(; linesearch = LiFukushimaLineSearch()), diff --git a/test/core/forward_ad_tests.jl b/test/forward_ad_tests.jl similarity index 57% rename from test/core/forward_ad_tests.jl rename to test/forward_ad_tests.jl index 24711b9d9..dcd1c2e3a 100644 --- a/test/core/forward_ad_tests.jl +++ b/test/forward_ad_tests.jl @@ -29,52 +29,62 @@ function solve_with(::Val{mode}, u, alg) where {mode} return f end -__compatible(::Any, ::Val{:oop}) = true -__compatible(::Any, ::Val{:oop_cache}) = true -__compatible(::Number, ::Val{:iip}) = false -__compatible(::AbstractArray, ::Val{:iip}) = true -__compatible(::StaticArray, ::Val{:iip}) = false -__compatible(::Number, ::Val{:iip_cache}) = false -__compatible(::AbstractArray, ::Val{:iip_cache}) = true -__compatible(::StaticArray, ::Val{:iip_cache}) = false +compatible(::Any, ::Val{:oop}) = true +compatible(::Any, ::Val{:oop_cache}) = true +compatible(::Number, ::Val{:iip}) = false +compatible(::AbstractArray, ::Val{:iip}) = true +compatible(::StaticArray, ::Val{:iip}) = false +compatible(::Number, ::Val{:iip_cache}) = false +compatible(::AbstractArray, ::Val{:iip_cache}) = true +compatible(::StaticArray, ::Val{:iip_cache}) = false -__compatible(::Any, ::Number) = true -__compatible(::Number, ::AbstractArray) = false -__compatible(u::AbstractArray, p::AbstractArray) = size(u) == size(p) +compatible(::Any, ::Number) = true +compatible(::Number, ::AbstractArray) = false +compatible(u::AbstractArray, p::AbstractArray) = size(u) == size(p) -__compatible(u::Number, ::SciMLBase.AbstractNonlinearAlgorithm) = true -__compatible(u::Number, ::Union{CMINPACK, NLsolveJL, KINSOL}) = true -__compatible(u::AbstractArray, ::SciMLBase.AbstractNonlinearAlgorithm) = true -__compatible(u::AbstractArray{T, N}, ::KINSOL) where {T, N} = N == 1 # Need to be fixed upstream -__compatible(u::StaticArray{S, T, N}, ::KINSOL) where {S <: Tuple, T, N} = false -__compatible(u::StaticArray, ::SciMLBase.AbstractNonlinearAlgorithm) = true -__compatible(u::StaticArray, ::Union{CMINPACK, NLsolveJL, KINSOL}) = false -__compatible(u, ::Nothing) = true +compatible(u::Number, ::SciMLBase.AbstractNonlinearAlgorithm) = true +compatible(u::Number, ::Union{CMINPACK, NLsolveJL, KINSOL}) = true +compatible(u::AbstractArray, ::SciMLBase.AbstractNonlinearAlgorithm) = true +compatible(u::AbstractArray{T, N}, ::KINSOL) where {T, N} = N == 1 # Need to be fixed upstream +compatible(u::StaticArray{S, T, N}, ::KINSOL) where {S <: Tuple, T, N} = false +compatible(u::StaticArray, ::SciMLBase.AbstractNonlinearAlgorithm) = true +compatible(u::StaticArray, ::Union{CMINPACK, NLsolveJL, KINSOL}) = false +compatible(u, ::Nothing) = true -__compatible(::Any, ::Any) = true -__compatible(::CMINPACK, ::Val{:iip_cache}) = false -__compatible(::CMINPACK, ::Val{:oop_cache}) = false -__compatible(::NLsolveJL, ::Val{:iip_cache}) = false -__compatible(::NLsolveJL, ::Val{:oop_cache}) = false -__compatible(::KINSOL, ::Val{:iip_cache}) = false -__compatible(::KINSOL, ::Val{:oop_cache}) = false +compatible(::Any, ::Any) = true +compatible(::CMINPACK, ::Val{:iip_cache}) = false +compatible(::CMINPACK, ::Val{:oop_cache}) = false +compatible(::NLsolveJL, ::Val{:iip_cache}) = false +compatible(::NLsolveJL, ::Val{:oop_cache}) = false +compatible(::KINSOL, ::Val{:iip_cache}) = false +compatible(::KINSOL, ::Val{:oop_cache}) = false -export test_f!, test_f, jacobian_f, solve_with, __compatible +export test_f!, test_f, jacobian_f, solve_with, compatible end @testitem "ForwardDiff.jl Integration" setup=[ForwardADTesting] tags=[:core] begin - for alg in (NewtonRaphson(), TrustRegion(), LevenbergMarquardt(), - PseudoTransient(; alpha_initial = 10.0), Broyden(), Klement(), DFSane(), - nothing, NLsolveJL(), CMINPACK(), KINSOL(; globalization_strategy = :LineSearch)) + @testset for alg in ( + NewtonRaphson(), + TrustRegion(), + LevenbergMarquardt(), + PseudoTransient(; alpha_initial = 10.0), + Broyden(), + Klement(), + DFSane(), + nothing, + NLsolveJL(), + CMINPACK(), + KINSOL(; globalization_strategy = :LineSearch) + ) us = (2.0, @SVector[1.0, 1.0], [1.0, 1.0], ones(2, 2), @SArray ones(2, 2)) alg isa CMINPACK && Sys.isapple() && continue @testset "Scalar AD" begin for p in 1.0:0.1:100.0, u0 in us, mode in (:iip, :oop, :iip_cache, :oop_cache) - __compatible(u0, alg) || continue - __compatible(u0, Val(mode)) || continue - __compatible(alg, Val(mode)) || continue + compatible(u0, alg) || continue + compatible(u0, Val(mode)) || continue + compatible(alg, Val(mode)) || continue sol = solve(NonlinearProblem(test_f, u0, p), alg) if SciMLBase.successful_retcode(sol) @@ -94,10 +104,10 @@ end p in ([2.0, 1.0], [2.0 1.0; 3.0 4.0]), mode in (:iip, :oop, :iip_cache, :oop_cache) - __compatible(u0, p) || continue - __compatible(u0, alg) || continue - __compatible(u0, Val(mode)) || continue - __compatible(alg, Val(mode)) || continue + compatible(u0, p) || continue + compatible(u0, alg) || continue + compatible(u0, Val(mode)) || continue + compatible(alg, Val(mode)) || continue sol = solve(NonlinearProblem(test_f, u0, p), alg) if SciMLBase.successful_retcode(sol) diff --git a/test/misc/aliasing_tests.jl b/test/misc/aliasing_tests.jl deleted file mode 100644 index 78b8ec798..000000000 --- a/test/misc/aliasing_tests.jl +++ /dev/null @@ -1,25 +0,0 @@ -@testitem "PolyAlgorithm Aliasing" tags=[:misc] begin - using NonlinearProblemLibrary - - # Use a problem that the initial solvers cannot solve and cause the initial value to - # diverge. If we don't alias correctly, all the subsequent algorithms will also fail. - prob = NonlinearProblemLibrary.nlprob_23_testcases["Generalized Rosenbrock function"].prob - u0 = copy(prob.u0) - prob = remake(prob; u0 = copy(u0)) - - # If aliasing is not handled properly this will diverge - sol = solve(prob; abstol = 1e-6, alias_u0 = true, - termination_condition = AbsNormTerminationMode(Base.Fix1(maximum, abs))) - - @test sol.u === prob.u0 - @test SciMLBase.successful_retcode(sol.retcode) - - prob = remake(prob; u0 = copy(u0)) - - cache = init(prob; abstol = 1e-6, alias_u0 = true, - termination_condition = AbsNormTerminationMode(Base.Fix1(maximum, abs))) - sol = solve!(cache) - - @test sol.u === prob.u0 - @test SciMLBase.successful_retcode(sol.retcode) -end diff --git a/test/misc/banded_matrices_tests.jl b/test/misc/banded_matrices_tests.jl deleted file mode 100644 index a9a2bf4ca..000000000 --- a/test/misc/banded_matrices_tests.jl +++ /dev/null @@ -1,8 +0,0 @@ -@testitem "Banded Matrix vcat" tags=[:misc] begin - using BandedMatrices, LinearAlgebra, SparseArrays - - b = BandedMatrix(Ones(5, 5), (1, 1)) - d = Diagonal(ones(5, 5)) - - @test NonlinearSolve._vcat(b, d) == vcat(sparse(b), d) -end diff --git a/test/misc/ensemble_tests.jl b/test/misc/ensemble_tests.jl deleted file mode 100644 index c034bf6bc..000000000 --- a/test/misc/ensemble_tests.jl +++ /dev/null @@ -1,21 +0,0 @@ -@testitem "Ensemble Nonlinear Problems" tags=[:misc] begin - using NonlinearSolve - - prob_func(prob, i, repeat) = remake(prob; u0 = prob.u0[:, i]) - - prob_nls_oop = NonlinearProblem((u, p) -> u .* u .- p, rand(4, 128), 2.0) - prob_nls_iip = NonlinearProblem((du, u, p) -> du .= u .* u .- p, rand(4, 128), 2.0) - prob_nlls_oop = NonlinearLeastSquaresProblem((u, p) -> u .^ 2 .- p, rand(4, 128), 2.0) - prob_nlls_iip = NonlinearLeastSquaresProblem( - NonlinearFunction{true}((du, u, p) -> du .= u .^ 2 .- p; resid_prototype = rand(4)), - rand(4, 128), 2.0) - - for prob in (prob_nls_oop, prob_nls_iip, prob_nlls_oop, prob_nlls_iip) - ensembleprob = EnsembleProblem(prob; prob_func) - - for ensemblealg in (EnsembleThreads(), EnsembleSerial()) - sim = solve(ensembleprob, nothing, ensemblealg; trajectories = size(prob.u0, 2)) - @test all(SciMLBase.successful_retcode, sim.u) - end - end -end diff --git a/test/misc/exotic_type_tests.jl b/test/misc/exotic_type_tests.jl deleted file mode 100644 index 57a9c28ed..000000000 --- a/test/misc/exotic_type_tests.jl +++ /dev/null @@ -1,27 +0,0 @@ -# File for different types of exotic types -@testsetup module NonlinearSolveExoticTypeTests -using NonlinearSolve - -fn_iip = NonlinearFunction{true}((du, u, p) -> du .= u .* u .- p) -fn_oop = NonlinearFunction{false}((u, p) -> u .* u .- p) - -u0 = BigFloat[1.0, 1.0, 1.0] -prob_iip_bf = NonlinearProblem{true}(fn_iip, u0, BigFloat(2)) -prob_oop_bf = NonlinearProblem{false}(fn_oop, u0, BigFloat(2)) - -export fn_iip, fn_oop, u0, prob_iip_bf, prob_oop_bf -end - -@testitem "BigFloat Support" tags=[:misc] setup=[NonlinearSolveExoticTypeTests] begin - using NonlinearSolve, LinearAlgebra - - for alg in [NewtonRaphson(), Broyden(), Klement(), DFSane(), TrustRegion()] - sol = solve(prob_oop_bf, alg) - @test norm(sol.resid, Inf) < 1e-6 - @test SciMLBase.successful_retcode(sol.retcode) - - sol = solve(prob_iip_bf, alg) - @test norm(sol.resid, Inf) < 1e-6 - @test SciMLBase.successful_retcode(sol.retcode) - end -end diff --git a/test/misc/linsolve_switching_tests.jl b/test/misc/linsolve_switching_tests.jl deleted file mode 100644 index 34d10e294..000000000 --- a/test/misc/linsolve_switching_tests.jl +++ /dev/null @@ -1,28 +0,0 @@ -@testitem "Singular Systems -- Auto Linear Solve Switching" tags=[:misc] begin - using LinearSolve, NonlinearSolve - - function f!(du, u, p) - du[1] = 2u[1] - 2 - du[2] = (u[1] - 4u[2])^2 + 0.1 - end - - u0 = [0.0, 0.0] # Singular Jacobian at u0 - - prob = NonlinearProblem(f!, u0) - - sol = solve(prob) # This doesn't have a root so let's just test the switching - @test sol.u≈[1.0, 0.25] atol=1e-3 rtol=1e-3 - - function nlls!(du, u, p) - du[1] = 2u[1] - 2 - du[2] = (u[1] - 4u[2])^2 + 0.1 - du[3] = 0 - end - - u0 = [0.0, 0.0] - - prob = NonlinearProblem(NonlinearFunction(nlls!, resid_prototype = zeros(3)), u0) - - solve(prob) - @test sol.u≈[1.0, 0.25] atol=1e-3 rtol=1e-3 -end diff --git a/test/misc/matrix_resizing_tests.jl b/test/misc/matrix_resizing_tests.jl deleted file mode 100644 index 2aa8151ce..000000000 --- a/test/misc/matrix_resizing_tests.jl +++ /dev/null @@ -1,31 +0,0 @@ -@testitem "Out-of-place Matrix Resizing" tags=[:misc] begin - using StableRNGs - - ff(u, p) = u .* u .- p - u0 = rand(StableRNG(0), 2, 2) - p = 2.0 - vecprob = NonlinearProblem(ff, vec(u0), p) - prob = NonlinearProblem(ff, u0, p) - - for alg in (NewtonRaphson(), TrustRegion(), LevenbergMarquardt(), - PseudoTransient(), RobustMultiNewton(), FastShortcutNonlinearPolyalg(), - Broyden(), Klement(), LimitedMemoryBroyden(; threshold = 2)) - @test vec(solve(prob, alg).u) == solve(vecprob, alg).u - end -end - -@testitem "Inplace Matrix Resizing" tags=[:misc] begin - using StableRNGs - - fiip(du, u, p) = (du .= u .* u .- p) - u0 = rand(StableRNG(0), 2, 2) - p = 2.0 - vecprob = NonlinearProblem(fiip, vec(u0), p) - prob = NonlinearProblem(fiip, u0, p) - - for alg in (NewtonRaphson(), TrustRegion(), LevenbergMarquardt(), - PseudoTransient(), RobustMultiNewton(), FastShortcutNonlinearPolyalg(), - Broyden(), Klement(), LimitedMemoryBroyden(; threshold = 2)) - @test vec(solve(prob, alg).u) == solve(vecprob, alg).u - end -end diff --git a/test/misc/noinit_caching_tests.jl b/test/misc/noinit_caching_tests.jl deleted file mode 100644 index 4ab4138f8..000000000 --- a/test/misc/noinit_caching_tests.jl +++ /dev/null @@ -1,23 +0,0 @@ -@testitem "NoInit Caching" tags=[:misc] begin - using LinearAlgebra - import NLsolve, NLSolvers - - solvers = [SimpleNewtonRaphson(), SimpleTrustRegion(), SimpleDFSane(), NLsolveJL(), - NLSolversJL(NLSolvers.LineSearch(NLSolvers.Newton(), NLSolvers.Backtracking()))] - - prob = NonlinearProblem((u, p) -> u .^ 2 .- p, [0.1, 0.3], 2.0) - - for alg in solvers - cache = init(prob, alg) - sol = solve!(cache) - @test SciMLBase.successful_retcode(sol) - @test norm(sol.resid, Inf) ≤ 1e-6 - - reinit!(cache; p = 5.0) - @test cache.prob.p == 5.0 - sol = solve!(cache) - @test SciMLBase.successful_retcode(sol) - @test norm(sol.resid, Inf) ≤ 1e-6 - @test norm(sol.u .^ 2 .- 5.0, Inf) ≤ 1e-6 - end -end diff --git a/test/misc/polyalg_tests.jl b/test/misc/polyalg_tests.jl deleted file mode 100644 index a106ace44..000000000 --- a/test/misc/polyalg_tests.jl +++ /dev/null @@ -1,256 +0,0 @@ -@testitem "Basic PolyAlgorithms" tags=[:misc] begin - f(u, p) = u .* u .- 2 - u0 = [1.0, 1.0] - probN = NonlinearProblem{false}(f, u0) - - custom_polyalg = NonlinearSolvePolyAlgorithm((Broyden(), LimitedMemoryBroyden())) - - # Uses the `__solve` function - solver = solve(probN; abstol = 1e-9) - @test SciMLBase.successful_retcode(solver) - solver = solve(probN, RobustMultiNewton(); abstol = 1e-9) - @test SciMLBase.successful_retcode(solver) - solver = solve(probN, FastShortcutNonlinearPolyalg(); abstol = 1e-9) - @test SciMLBase.successful_retcode(solver) - solver = solve(probN, custom_polyalg; abstol = 1e-9) - @test SciMLBase.successful_retcode(solver) - - # Test the caching interface - cache = init(probN; abstol = 1e-9) - solver = solve!(cache) - @test SciMLBase.successful_retcode(solver) - cache = init(probN, RobustMultiNewton(); abstol = 1e-9) - solver = solve!(cache) - @test SciMLBase.successful_retcode(solver) - cache = init(probN, FastShortcutNonlinearPolyalg(); abstol = 1e-9) - solver = solve!(cache) - @test SciMLBase.successful_retcode(solver) - cache = init(probN, custom_polyalg; abstol = 1e-9) - solver = solve!(cache) - @test SciMLBase.successful_retcode(solver) - - # Test the step interface - cache = init(probN; abstol = 1e-9) - for i in 1:10000 - step!(cache) - cache.force_stop && break - end - @test SciMLBase.successful_retcode(cache.retcode) - cache = init(probN, RobustMultiNewton(); abstol = 1e-9) - for i in 1:10000 - step!(cache) - cache.force_stop && break - end - @test SciMLBase.successful_retcode(cache.retcode) - cache = init(probN, FastShortcutNonlinearPolyalg(); abstol = 1e-9) - for i in 1:10000 - step!(cache) - cache.force_stop && break - end - @test SciMLBase.successful_retcode(cache.retcode) - cache = init(probN, custom_polyalg; abstol = 1e-9) - for i in 1:10000 - step!(cache) - cache.force_stop && break - end -end - -@testitem "Testing #153 Singular Exception" tags=[:misc] begin - # https://github.com/SciML/NonlinearSolve.jl/issues/153 - function f(du, u, p) - s1, s1s2, s2 = u - k1, c1, Δt = p - - du[1] = -0.25 * c1 * k1 * s1 * s2 - du[2] = 0.25 * c1 * k1 * s1 * s2 - du[3] = -0.25 * c1 * k1 * s1 * s2 - end - - prob = NonlinearProblem(f, [2.0, 2.0, 2.0], [1.0, 2.0, 2.5]) - sol = solve(prob; abstol = 1e-9) - @test SciMLBase.successful_retcode(sol) -end - -@testitem "PolyAlgorithms Autodiff" tags=[:misc] begin - cache = zeros(2) - function f(du, u, p) - cache .= u .* u - du .= cache .- 2 - end - u0 = [1.0, 1.0] - probN = NonlinearProblem{true}(f, u0) - - custom_polyalg = NonlinearSolvePolyAlgorithm(( - Broyden(; autodiff = AutoFiniteDiff()), LimitedMemoryBroyden())) - - # Uses the `__solve` function - @test_throws MethodError solve(probN; abstol = 1e-9) - @test_throws MethodError solve(probN, RobustMultiNewton(); abstol = 1e-9) - sol = solve(probN, RobustMultiNewton(; autodiff = AutoFiniteDiff()); abstol = 1e-9) - @test SciMLBase.successful_retcode(sol) - sol = solve( - probN, FastShortcutNonlinearPolyalg(; autodiff = AutoFiniteDiff()); abstol = 1e-9) - @test SciMLBase.successful_retcode(sol) - sol = solve(probN, custom_polyalg; abstol = 1e-9) - @test SciMLBase.successful_retcode(sol) - - quadratic_f(u::Float64, p) = u^2 - p - - prob = NonlinearProblem(quadratic_f, 2.0, 4.0) - - @test_throws MethodError solve(prob) - @test_throws MethodError solve(prob, RobustMultiNewton()) - sol = solve(prob, RobustMultiNewton(; autodiff = AutoFiniteDiff())) - @test SciMLBase.successful_retcode(sol) -end - -@testitem "Simple Scalar Problem #187" tags=[:misc] begin - using NaNMath - - # https://github.com/SciML/NonlinearSolve.jl/issues/187 - # If we use a General Nonlinear Solver the solution might go out of the domain! - ff_interval(u, p) = 0.5 / 1.5 * NaNMath.log.(u ./ (1.0 .- u)) .- 2.0 * u .+ 1.0 - - uspan = (0.02, 0.1) - prob = IntervalNonlinearProblem(ff_interval, uspan) - sol = solve(prob; abstol = 1e-9) - @test SciMLBase.successful_retcode(sol) - - u0 = 0.06 - p = 2.0 - prob = NonlinearProblem(ff_interval, u0, p) - sol = solve(prob; abstol = 1e-9) - @test SciMLBase.successful_retcode(sol) -end - -# Shooting Problem: Taken from BoundaryValueDiffEq.jl -# Testing for Complex Valued Root Finding. For Complex valued inputs we drop some of the -# algorithms which dont support those. -@testitem "Complex Valued Problems: Single-Shooting" tags=[:misc] begin - using OrdinaryDiffEqTsit5 - - function ode_func!(du, u, p, t) - du[1] = u[2] - du[2] = -u[1] - return nothing - end - - function objective_function!(resid, u0, p) - odeprob = ODEProblem{true}(ode_func!, u0, (0.0, 100.0), p) - sol = solve(odeprob, Tsit5(), abstol = 1e-9, reltol = 1e-9, verbose = false) - resid[1] = sol(0.0)[1] - resid[2] = sol(100.0)[1] - 1.0 - return nothing - end - - prob = NonlinearProblem{true}(objective_function!, [0.0, 1.0] .+ 1im) - sol = solve(prob; abstol = 1e-10) - @test SciMLBase.successful_retcode(sol) - # This test is not meant to return success but test that all the default solvers can handle - # complex valued problems - @test_nowarn solve(prob; abstol = 1e-19, maxiters = 10) - @test_nowarn solve( - prob, RobustMultiNewton(eltype(prob.u0)); abstol = 1e-19, maxiters = 10) -end - -@testitem "No AD" tags=[:misc] begin - no_ad_fast = FastShortcutNonlinearPolyalg(autodiff = AutoFiniteDiff()) - no_ad_robust = RobustMultiNewton(autodiff = AutoFiniteDiff()) - no_ad_algs = Set([no_ad_fast, no_ad_robust, no_ad_fast.algs..., no_ad_robust.algs...]) - - @testset "Inplace" begin - f_iip = Base.Experimental.@opaque (du, u, p) -> du .= u .* u .- p - u0 = [0.5] - prob = NonlinearProblem(f_iip, u0, 1.0) - for alg in no_ad_algs - sol = solve(prob, alg) - @test isapprox(only(sol.u), 1.0) - @test SciMLBase.successful_retcode(sol.retcode) - end - end - - @testset "Out of Place" begin - f_oop = Base.Experimental.@opaque (u, p) -> u .* u .- p - u0 = [0.5] - prob = NonlinearProblem{false}(f_oop, u0, 1.0) - for alg in no_ad_algs - sol = solve(prob, alg) - @test isapprox(only(sol.u), 1.0) - @test SciMLBase.successful_retcode(sol.retcode) - end - end -end - -@testsetup module InfeasibleFunction -using LinearAlgebra, StaticArrays - -# this is infeasible -function f1_infeasible!(out, u, p) - μ = 3.986004415e14 - x = 7000.0e3 - y = -6.970561549987071e-9 - z = -3.784706123246018e-9 - v_x = 8.550491684548064e-12 + u[1] - v_y = 6631.60076191005 + u[2] - v_z = 3600.665431405663 + u[3] - r = @SVector [x, y, z] - v = @SVector [v_x, v_y, v_z] - h = cross(r, v) - ev = cross(v, h) / μ - r / norm(r) - i = acos(h[3] / norm(h)) - e = norm(ev) - a = 1 / (2 / norm(r) - (norm(v)^2 / μ)) - out .= [a - 42.0e6, e - 1e-5, i - 1e-5] - return nothing -end - -# this is unfeasible -function f1_infeasible(u, p) - μ = 3.986004415e14 - x = 7000.0e3 - y = -6.970561549987071e-9 - z = -3.784706123246018e-9 - v_x = 8.550491684548064e-12 + u[1] - v_y = 6631.60076191005 + u[2] - v_z = 3600.665431405663 + u[3] - r = [x, y, z] - v = [v_x, v_y, v_z] - h = cross(r, v) - ev = cross(v, h) / μ - r / norm(r) - i = acos(h[3] / norm(h)) - e = norm(ev) - a = 1 / (2 / norm(r) - (norm(v)^2 / μ)) - return [a - 42.0e6, e - 1e-5, i - 1e-5] -end - -export f1_infeasible!, f1_infeasible -end - -@testitem "[IIP] Infeasible" setup=[InfeasibleFunction] tags=[:misc] begin - u0 = [0.0, 0.0, 0.0] - prob = NonlinearProblem(f1_infeasible!, u0) - sol = solve(prob) - - @test all(!isnan, sol.u) - @test !SciMLBase.successful_retcode(sol.retcode) - @inferred solve(prob) -end - -@testitem "[OOP] Infeasible" setup=[InfeasibleFunction] tags=[:misc] begin - using LinearAlgebra, StaticArrays - - u0 = [0.0, 0.0, 0.0] - prob = NonlinearProblem(f1_infeasible, u0) - sol = solve(prob) - - @test all(!isnan, sol.u) - @test !SciMLBase.successful_retcode(sol.retcode) - @inferred solve(prob) - - u0 = @SVector [0.0, 0.0, 0.0] - prob = NonlinearProblem(f1_infeasible, u0) - - sol = solve(prob) - @test all(!isnan, sol.u) - @test !SciMLBase.successful_retcode(sol.retcode) -end diff --git a/test/misc/qa_tests.jl b/test/misc/qa_tests.jl deleted file mode 100644 index 6aafd4024..000000000 --- a/test/misc/qa_tests.jl +++ /dev/null @@ -1,28 +0,0 @@ -@testitem "Aqua" tags=[:misc] begin - using NonlinearSolve, SimpleNonlinearSolve, Aqua - - Aqua.find_persistent_tasks_deps(NonlinearSolve) - Aqua.test_ambiguities(NonlinearSolve; recursive = false) - Aqua.test_deps_compat(NonlinearSolve) - Aqua.test_piracies(NonlinearSolve, - treat_as_own = [NonlinearProblem, NonlinearLeastSquaresProblem, - SimpleNonlinearSolve.AbstractSimpleNonlinearSolveAlgorithm]) - Aqua.test_project_extras(NonlinearSolve) - # Timer Outputs needs to be enabled via Preferences - Aqua.test_stale_deps(NonlinearSolve; ignore = [:TimerOutputs]) - Aqua.test_unbound_args(NonlinearSolve) - Aqua.test_undefined_exports(NonlinearSolve) -end - -@testitem "Explicit Imports" tags=[:misc] begin - using NonlinearSolve, ADTypes, SimpleNonlinearSolve, SciMLBase - import BandedMatrices, FastLevenbergMarquardt, FixedPointAcceleration, - LeastSquaresOptim, MINPACK, NLsolve, NLSolvers, SIAMFANLEquations, SpeedMapping - - using ExplicitImports - - @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/misc/structured_jacobian_tests.jl b/test/misc/structured_jacobian_tests.jl deleted file mode 100644 index 41b316911..000000000 --- a/test/misc/structured_jacobian_tests.jl +++ /dev/null @@ -1,67 +0,0 @@ -@testitem "Structured Jacobians" tags=[:misc] begin - using NonlinearSolve, SparseConnectivityTracer, BandedMatrices, LinearAlgebra, - SparseArrays - - N = 16 - p = rand(N) - u0 = rand(N) - - function f!(du, u, p) - for i in 2:(length(u) - 1) - du[i] = u[i - 1] - 2u[i] + u[i + 1] + p[i] - end - du[1] = -2u[1] + u[2] + p[1] - du[end] = u[end - 1] - 2u[end] + p[end] - return nothing - end - - function f(u, p) - du = similar(u, promote_type(eltype(u), eltype(p))) - f!(du, u, p) - return du - end - - for nlf in (f, f!) - @testset "Dense AD" begin - nlprob = NonlinearProblem(NonlinearFunction(nlf), u0, p) - - cache = init(nlprob, NewtonRaphson(); abstol = 1e-9) - @test cache.jac_cache.J isa Matrix - sol = solve!(cache) - @test SciMLBase.successful_retcode(sol) - end - - @testset "Unstructured Sparse AD" begin - nlprob_autosparse = NonlinearProblem( - NonlinearFunction(nlf; sparsity = TracerSparsityDetector()), - u0, p) - - cache = init(nlprob_autosparse, NewtonRaphson(); abstol = 1e-9) - @test cache.jac_cache.J isa SparseMatrixCSC - sol = solve!(cache) - @test SciMLBase.successful_retcode(sol) - end - - @testset "Structured Sparse AD: Banded Jacobian" begin - jac_prototype = BandedMatrix(-1 => ones(N - 1), 0 => ones(N), 1 => ones(N - 1)) - nlprob_sparse_structured = NonlinearProblem( - NonlinearFunction(nlf; jac_prototype), u0, p) - - cache = init(nlprob_sparse_structured, NewtonRaphson(); abstol = 1e-9) - @test cache.jac_cache.J isa BandedMatrix - sol = solve!(cache) - @test SciMLBase.successful_retcode(sol) - end - - @testset "Structured Sparse AD: Tridiagonal Jacobian" begin - jac_prototype = Tridiagonal(ones(N - 1), ones(N), ones(N - 1)) - nlprob_sparse_structured = NonlinearProblem( - NonlinearFunction(nlf; jac_prototype), u0, p) - - cache = init(nlprob_sparse_structured, NewtonRaphson(); abstol = 1e-9) - @test cache.jac_cache.J isa Tridiagonal - sol = solve!(cache) - @test SciMLBase.successful_retcode(sol) - end - end -end diff --git a/test/downstream/mtk_cache_indexing_tests.jl b/test/mtk_cache_indexing_tests.jl similarity index 98% rename from test/downstream/mtk_cache_indexing_tests.jl rename to test/mtk_cache_indexing_tests.jl index d29cb1da7..3109d1612 100644 --- a/test/downstream/mtk_cache_indexing_tests.jl +++ b/test/mtk_cache_indexing_tests.jl @@ -13,7 +13,8 @@ @testset "$integtype" for (alg, integtype) in [ (NewtonRaphson(), NonlinearSolve.GeneralizedFirstOrderAlgorithmCache), (FastShortcutNonlinearPolyalg(), NonlinearSolve.NonlinearSolvePolyAlgorithmCache), - (SimpleNewtonRaphson(), NonlinearSolve.NonlinearSolveNoInitCache)] + (SimpleNewtonRaphson(), NonlinearSolve.NonlinearSolveNoInitCache) + ] nint = init(nlprob, alg) @test nint isa integtype diff --git a/test/qa_tests.jl b/test/qa_tests.jl new file mode 100644 index 000000000..ffac5c7bb --- /dev/null +++ b/test/qa_tests.jl @@ -0,0 +1,24 @@ +@testitem "Aqua" tags=[:misc] begin + using NonlinearSolve, SimpleNonlinearSolve, Aqua + + Aqua.test_all(NonlinearSolve; ambiguities = false, piracies = false) + Aqua.test_ambiguities(NonlinearSolve; recursive = false) + Aqua.test_piracies(NonlinearSolve, + treat_as_own = [ + NonlinearProblem, NonlinearLeastSquaresProblem, + SimpleNonlinearSolve.AbstractSimpleNonlinearSolveAlgorithm + ] + ) +end + +@testitem "Explicit Imports" tags=[:misc] begin + using NonlinearSolve, ADTypes, SimpleNonlinearSolve, SciMLBase + import FastLevenbergMarquardt, FixedPointAcceleration, LeastSquaresOptim, MINPACK, + NLsolve, NLSolvers, SIAMFANLEquations, SpeedMapping + + using ExplicitImports + + @test check_no_implicit_imports(NonlinearSolve; skip = (Base, Core)) === 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 33ca5da7a..59a43c2f9 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -10,13 +10,19 @@ const EXTRA_PKGS = Pkg.PackageSpec[] length(EXTRA_PKGS) ≥ 1 && Pkg.add(EXTRA_PKGS) const RETESTITEMS_NWORKERS = parse( - Int, get(ENV, "RETESTITEMS_NWORKERS", string(min(Hwloc.num_physical_cores(), 4)))) + 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)))) + 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)]), +ReTestItems.runtests( + NonlinearSolve; tags = (GROUP == "all" ? nothing : [Symbol(GROUP)]), nworkers = RETESTITEMS_NWORKERS, nworker_threads = RETESTITEMS_NWORKER_THREADS, - testitem_timeout = 3600) + testitem_timeout = 3600 +) diff --git a/test/wrappers/fixedpoint_tests.jl b/test/wrappers/fixedpoint_tests.jl index a3ec497c7..a04e35069 100644 --- a/test/wrappers/fixedpoint_tests.jl +++ b/test/wrappers/fixedpoint_tests.jl @@ -1,11 +1,6 @@ -@testsetup module WrapperFixedPointImports -using Reexport -@reexport using LinearAlgebra -import SIAMFANLEquations, FixedPointAcceleration, SpeedMapping, NLsolve -end +@testitem "Simple Scalar Problem" tags=[:wrappers] begin + import SpeedMapping, SIAMFANLEquations, NLsolve, FixedPointAcceleration -# Simple Scalar Problem -@testitem "Simple Scalar Problem" setup=[WrapperFixedPointImports] tags=[:wrappers] begin f1(x, p) = cos(x) - x prob = NonlinearProblem(f1, 1.1) @@ -22,7 +17,9 @@ end end # Simple Vector Problem -@testitem "Simple Vector Problem" setup=[WrapperFixedPointImports] tags=[:wrappers] begin +@testitem "Simple Vector Problem" tags=[:wrappers] begin + import SpeedMapping, SIAMFANLEquations, NLsolve, FixedPointAcceleration + f2(x, p) = cos.(x) .- x prob = NonlinearProblem(f2, [1.1, 1.1]) @@ -41,7 +38,10 @@ end # Fixed Point for Power Method # Taken from https://github.com/NicolasL-S/SpeedMapping.jl/blob/95951db8f8a4457093090e18802ad382db1c76da/test/runtests.jl -@testitem "Power Method" setup=[WrapperFixedPointImports] tags=[:wrappers] begin +@testitem "Power Method" tags=[:wrappers] begin + using LinearAlgebra + import SpeedMapping, SIAMFANLEquations, NLsolve, FixedPointAcceleration + C = [1 2 3; 4 5 6; 7 8 9] A = C + C' B = Hermitian(ones(10) * ones(10)' .* im + Diagonal(1:10)) diff --git a/test/wrappers/nlls_tests.jl b/test/wrappers/least_squares_tests.jl similarity index 58% rename from test/wrappers/nlls_tests.jl rename to test/wrappers/least_squares_tests.jl index 53cea758d..99d3e0ef3 100644 --- a/test/wrappers/nlls_tests.jl +++ b/test/wrappers/least_squares_tests.jl @@ -1,52 +1,32 @@ @testsetup module WrapperNLLSSetup -using Reexport -@reexport using LinearAlgebra, StableRNGs, StaticArrays, Random, ForwardDiff, Zygote -import FastLevenbergMarquardt, LeastSquaresOptim, MINPACK -true_function(x, θ) = @. θ[1] * exp(θ[2] * x) * cos(θ[3] * x + θ[4]) -true_function(y, x, θ) = (@. y = θ[1] * exp(θ[2] * x) * cos(θ[3] * x + θ[4])) +include("../../common/common_nlls_testing.jl") -θ_true = [1.0, 0.1, 2.0, 0.5] - -x = [-1.0, -0.5, 0.0, 0.5, 1.0] - -const y_target = true_function(x, θ_true) - -function loss_function(θ, p) - ŷ = true_function(p, θ) - return ŷ .- y_target -end - -function loss_function(resid, θ, p) - true_function(resid, p, θ) - resid .= resid .- y_target - return resid -end - -θ_init = θ_true .+ randn!(StableRNG(0), similar(θ_true)) * 0.1 - -export loss_function, θ_init, y_target, true_function, x, θ_true end @testitem "LeastSquaresOptim.jl" setup=[WrapperNLLSSetup] tags=[:wrappers] begin - prob_oop = NonlinearLeastSquaresProblem{false}(loss_function, θ_init, x) - prob_iip = NonlinearLeastSquaresProblem( - NonlinearFunction(loss_function; resid_prototype = zero(y_target)), θ_init, x) + import LeastSquaresOptim nlls_problems = [prob_oop, prob_iip] - solvers = [LeastSquaresOptimJL(alg; autodiff) - for alg in (:lm, :dogleg), - autodiff in (nothing, AutoForwardDiff(), AutoFiniteDiff(), :central, :forward)] + solvers = [] + for alg in (:lm, :dogleg), + autodiff in (nothing, AutoForwardDiff(), AutoFiniteDiff(), :central, :forward) + + push!(solvers, LeastSquaresOptimJL(alg; autodiff)) + end for prob in nlls_problems, solver in solvers sol = solve(prob, solver; maxiters = 10000, abstol = 1e-8) @test SciMLBase.successful_retcode(sol) - @test norm(sol.resid, Inf) < 1e-6 + @test maximum(abs, sol.resid) < 1e-6 end end @testitem "FastLevenbergMarquardt.jl + CMINPACK: Jacobian Provided" setup=[WrapperNLLSSetup] tags=[:wrappers] begin + import FastLevenbergMarquardt, MINPACK + using ForwardDiff + function jac!(J, θ, p) resid = zeros(length(p)) ForwardDiff.jacobian!(J, (resid, θ) -> loss_function(resid, θ, p), resid, θ) @@ -58,19 +38,24 @@ end probs = [ NonlinearLeastSquaresProblem( NonlinearFunction{true}( - loss_function; resid_prototype = zero(y_target), jac = jac!), - θ_init, - x), + loss_function; resid_prototype = zero(y_target), jac = jac! + ), + θ_init, x + ), NonlinearLeastSquaresProblem( NonlinearFunction{false}( - loss_function; resid_prototype = zero(y_target), jac = jac), - θ_init, - x), + loss_function; resid_prototype = zero(y_target), jac = jac + ), + θ_init, x + ), NonlinearLeastSquaresProblem( - NonlinearFunction{false}(loss_function; jac), θ_init, x)] + NonlinearFunction{false}(loss_function; jac), θ_init, x + ) + ] solvers = Any[FastLevenbergMarquardtJL(linsolve) for linsolve in (:cholesky, :qr)] Sys.isapple() || push!(solvers, CMINPACK()) + for solver in solvers, prob in probs sol = solve(prob, solver; maxiters = 10000, abstol = 1e-8) @test maximum(abs, sol.resid) < 1e-6 @@ -78,28 +63,41 @@ end end @testitem "FastLevenbergMarquardt.jl + CMINPACK: Jacobian Not Provided" setup=[WrapperNLLSSetup] tags=[:wrappers] begin + import FastLevenbergMarquardt, MINPACK + probs = [ NonlinearLeastSquaresProblem( NonlinearFunction{true}(loss_function; resid_prototype = zero(y_target)), - θ_init, x), + θ_init, x + ), NonlinearLeastSquaresProblem( NonlinearFunction{false}(loss_function; resid_prototype = zero(y_target)), - θ_init, x), - NonlinearLeastSquaresProblem(NonlinearFunction{false}(loss_function), θ_init, x)] + θ_init, x + ), + NonlinearLeastSquaresProblem(NonlinearFunction{false}(loss_function), θ_init, x) + ] - solvers = vec(Any[FastLevenbergMarquardtJL(linsolve; autodiff) - for linsolve in (:cholesky, :qr), - autodiff in (nothing, AutoForwardDiff(), AutoFiniteDiff())]) - Sys.isapple() || - append!(solvers, [CMINPACK(; method) for method in (:auto, :lm, :lmdif)]) + solvers = [] + for linsolve in (:cholesky, :qr), + autodiff in (nothing, AutoForwardDiff(), AutoFiniteDiff()) + + push!(solvers, FastLevenbergMarquardtJL(linsolve; autodiff)) + end + if Sys.isapple() + for method in (:auto, :lm, :lmdif) + push!(solvers, CMINPACK(; method)) + end + end for solver in solvers, prob in probs sol = solve(prob, solver; maxiters = 10000, abstol = 1e-8) - @test norm(sol.resid, Inf) < 1e-6 + @test maximum(abs, sol.resid) < 1e-6 end end @testitem "FastLevenbergMarquardt.jl + StaticArrays" setup=[WrapperNLLSSetup] tags=[:wrappers] begin + using StaticArrays, FastLevenbergMarquardt + x_sa = SA[-1.0, -0.5, 0.0, 0.5, 1.0] const y_target_sa = true_function(x_sa, θ_true) @@ -113,5 +111,5 @@ end prob_sa = NonlinearLeastSquaresProblem{false}(loss_function_sa, θ_init_sa, x) sol = solve(prob_sa, FastLevenbergMarquardtJL()) - @test norm(sol.resid, Inf) < 1e-6 + @test maximum(abs, sol.resid) < 1e-6 end diff --git a/test/wrappers/rootfind_tests.jl b/test/wrappers/rootfind_tests.jl index 852368ae3..073d31fe4 100644 --- a/test/wrappers/rootfind_tests.jl +++ b/test/wrappers/rootfind_tests.jl @@ -1,13 +1,6 @@ -@testsetup module WrapperRootfindImports -using Reexport -@reexport using LinearAlgebra -import NLSolvers, NLsolve, SIAMFANLEquations, MINPACK, PETSc +@testitem "Steady State Problems" tags=[:wrappers] begin + import NLSolvers, NLsolve, SIAMFANLEquations, MINPACK, PETSc -export NLSolvers -end - -@testitem "Steady State Problems" setup=[WrapperRootfindImports] tags=[:wrappers] begin - # IIP Tests function f_iip(du, u, p, t) du[1] = 2 - 2u[1] du[2] = u[1] - 4u[2] @@ -30,7 +23,6 @@ end @test maximum(abs, sol.resid) < 1e-6 end - # OOP Tests f_oop(u, p, t) = [2 - 2u[1], u[1] - 4u[2]] u0 = zeros(2) prob_oop = SteadyStateProblem(f_oop, u0) @@ -51,9 +43,10 @@ end end end -# Can lead to segfaults -@testitem "Nonlinear Root Finding Problems" setup=[WrapperRootfindImports] tags=[:wrappers] retries=3 begin - # IIP Tests +@testitem "Nonlinear Root Finding Problems" tags=[:wrappers] begin + using LinearAlgebra + import NLSolvers, NLsolve, SIAMFANLEquations, MINPACK, PETSc + function f_iip(du, u, p) du[1] = 2 - 2u[1] du[2] = u[1] - 4u[2] @@ -77,7 +70,6 @@ end @test maximum(abs, sol.resid) < 1e-6 end - # OOP Tests f_oop(u, p) = [2 - 2u[1], u[1] - 4u[2]] u0 = zeros(2) prob_oop = NonlinearProblem{false}(f_oop, u0) @@ -97,7 +89,6 @@ end @test maximum(abs, sol.resid) < 1e-6 end - # Tolerance Tests f_tol(u, p) = u^2 - 2 prob_tol = NonlinearProblem(f_tol, 1.0) for tol in [1e-1, 1e-3, 1e-6, 1e-10, 1e-15], @@ -156,9 +147,11 @@ end sol = solve(ProbN, NLsolveJL(); abstol = 1e-8) @test maximum(abs, sol.resid) < 1e-6 - sol = solve(ProbN, + sol = solve( + ProbN, NLSolversJL(NLSolvers.LineSearch(NLSolvers.Newton(), NLSolvers.Backtracking())); - abstol = 1e-8) + abstol = 1e-8 + ) @test maximum(abs, sol.resid) < 1e-6 sol = solve(ProbN, SIAMFANLEquationsJL(; method = :newton); abstol = 1e-8) @test maximum(abs, sol.resid) < 1e-6 @@ -170,7 +163,9 @@ end end end -@testitem "PETSc SNES Floating Points" setup=[WrapperRootfindImports] tags=[:wrappers] skip=:(Sys.iswindows()) begin +@testitem "PETSc SNES Floating Points" tags=[:wrappers] skip=:(Sys.iswindows()) begin + import PETSc + f(u, p) = u .* u .- 2 u0 = [1.0, 1.0]