added checkpointing to gauss adjoint #884

acoh64 · 2023-08-26T04:06:50Z

See: #869. I think this needs some more testing first. It only seems to be working at lower tolerances (1e-4), so I need to figure out why

… lower tolerances (1e-4)

ai-maintainer

AI-Maintainer Review for PR - Added checkpointing to gauss adjoint

Title and Description 👍

The Title and Description are clear and focused

The title and description of the pull request are clear and focused. They effectively communicate the purpose of the changes, which is to add checkpointing to the Gauss adjoint algorithm. The author also acknowledges the need for further testing, particularly at higher tolerances.

Scope of Changes 👍

The changes are narrowly focused

The changes in this pull request are narrowly focused on adding checkpointing to the Gauss adjoint algorithm. The modifications are primarily in the `gauss_adjoint.jl` file, with some changes in the `sensitivity_algorithms.jl` and `sensitivity_interface.jl` files to support the new feature. The author is not trying to resolve multiple issues simultaneously.

Testing ⚠️

Testing details are not provided

The description does not provide specific details about how the author tested the changes. It would be helpful for the author to provide more information about the testing process, such as the test environment, test inputs, and expected outcomes, to ensure that the changes have been thoroughly tested and validated.

Documentation ⚠️

Docstrings are missing for some functions, classes, or methods

The following functions, classes, or methods do not have docstrings:

GaussCheckpointSolution
Gaussfindcursor
Gaussreset_p
GaussAdjoint
setvjp
adjoint_sensitivities

These entities should have docstrings added to describe their behavior, arguments, and return values.

Suggested Changes

Please add docstrings to the GaussCheckpointSolution, Gaussfindcursor, Gaussreset_p, GaussAdjoint, setvjp, and adjoint_sensitivities entities to describe their behavior, arguments, and return values.
Please provide more information about how you tested the changes, including the test environment, test inputs, and expected outcomes.

Reviewed with AI Maintainer

ChrisRackauckas · 2023-08-28T15:06:10Z

https://github.com/SciML/SciMLSensitivity.jl/actions/runs/5982716510/job/16232100431?pr=884#step:6:872

ChrisRackauckas · 2023-08-28T15:06:30Z

@avik-pal would you know where that is from?

avik-pal · 2023-08-28T19:32:46Z

Not really. Core2 tests are passing in other CI https://github.com/SciML/SciMLSensitivity.jl/actions/runs/6001815834/job/16276828862?pr=885. Also I can't find any .t in code, and the stacktrace seems incomplete?

ChrisRackauckas · 2023-08-30T11:08:53Z

It only seems to be working at lower tolerances (1e-4), so I need to figure out why

Do the other ones do this too? We only have low tolerance tests.

acoh64 · 2023-08-30T14:21:32Z

src/sensitivity_interface.jl

I think the issue is here. I removed checkpoints from kwargs because of an error that :checkpoints is an unrecognized kwarg for solve. However, I just realized that I don't think the correct checkpoints are being used now since checkpoints is not being passed to anything

acoh64 · 2023-08-30T14:22:02Z

The other ones checkpoint tests (for interpolating adjoint) work at a tolerance of 1e-9.

acoh64 · 2023-08-30T14:35:09Z

Ah I think I see now

acoh64 · 2023-08-30T18:02:50Z

I fixed errors in the adjoint state that came from not passing the proper tolerances into the checkpointing solves. However, the gradient calculation still only works at tolerances of 1e-4

ChrisRackauckas · 2023-08-30T18:04:03Z

Is the solve dense and not using saveat?

acoh64 · 2023-08-30T18:13:24Z

The checkpoint solves are dense but the adjoint solve is not

codecov · 2023-08-30T18:25:55Z

Codecov Report

Merging #884 (9ca3932) into master (5d03a76) will decrease coverage by 13.47%.
Report is 11 commits behind head on master.
The diff coverage is 0.32%.

❗ Current head 9ca3932 differs from pull request most recent head f1cd78c. Consider uploading reports for the commit f1cd78c to get more accurate results

@@             Coverage Diff             @@
##           master     #884       +/-   ##
===========================================
- Coverage   61.91%   48.45%   -13.47%     
===========================================
  Files          20       20               
  Lines        4385     4658      +273     
===========================================
- Hits         2715     2257      -458     
- Misses       1670     2401      +731

Files	Coverage Δ
src/concrete_solve.jl	`65.88% <100.00%> (-3.50%)`	⬇️
src/interpolating_adjoint.jl	`71.02% <ø> (-4.99%)`	⬇️
src/sensitivity_interface.jl	`83.33% <ø> (-8.34%)`	⬇️
src/derivative_wrappers.jl	`77.97% <0.00%> (-12.23%)`	⬇️
src/sensitivity_algorithms.jl	`70.00% <0.00%> (-7.78%)`	⬇️
src/gauss_adjoint.jl	`0.00% <0.00%> (-65.70%)`	⬇️

... and 6 files with indirect coverage changes

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

acoh64 · 2023-09-13T13:48:08Z

Which solve were you referring to? When I use the dense forward solve with checkpointing, then GaussAdjoint works to 1e-9 tolerance, but this feels like cheating.

The issue is in the forward solves, since the adjoint solution matches perfectly and the callback times line up

acoh64 · 2023-09-14T21:01:00Z

@ChrisRackauckas I figured out the problem: the integrand function is made ahead of time with the solution to the nondense solve. Therefore, it is using the parameter jacobians from the nondense solve in the integral calculation, causing it to be less accurate. I think the best solution is to make a callback which takes as an argument the integrator and a sol, however, I don't think this can be done with the current callback interface. Is there a way around this?

ChrisRackauckas · 2023-09-17T03:08:14Z

Add the GaussIntegrand to the ODEGaussAdjointSensitivityFunction and make it a mutable struct so you can mutate integrand.sol when a new cpsol is taken.

… InterpolatingAdjoint

acoh64 · 2023-09-18T18:31:26Z

Thanks, checkpointing now works for GaussAdjoint!

test/sde_scalar_stratonovich.jl

ChrisRackauckas · 2023-09-19T11:03:16Z

Looks like there's still a few test failures?

acoh64 · 2023-09-20T01:57:13Z

Checkpointing tests pass but I am still working on the SDE tests

frankschae

@glatteis @acoh64's GaussAdjoint will probably also solve the efficiency problem you spotted. I think we can add in the callback

dy2, back = Zygote.pullback(y, p) do u, p
    g(u, p, t) .* dW 
end
out2 = back(λ)

ChrisRackauckas · 2023-09-24T18:16:55Z

test/adjoint.jl

+
+
+
+using Lux, Optimization, Plots, Random


This is not in the right place?

ChrisRackauckas · 2023-09-24T18:17:15Z

test/adjoint.jl

+function dudt(u, p, t)
+    global st
+    #input_val = u_vals[Int(round(t*10)+1)]
+    out, st = nn_model(vcat(u[1], ex[Int(round(10 * 0.1))]), p, st)
+    return out
+end
+
+prob = ODEProblem(dudt, u0, tspan, nothing)
+
+function predict_neuralode(p)
+    _prob = remake(prob, p = p)
+    Array(solve(_prob, Tsit5(), saveat = tsteps, abstol = 1e-8, reltol = 1e-6, sensealg=GaussAdjoint(autojacvec=ZygoteVJP())))
+end
+
+function loss(p)
+    sol = predict_neuralode(p)
+    N = length(sol)
+    return sum(abs2.(y[1:N] .- sol')) / N
+end
+
+adtype = Optimization.AutoZygote()
+optf = Optimization.OptimizationFunction((x, p) -> loss(x), adtype)
+optprob = Optimization.OptimizationProblem(optf, p_model)
+
+tmp1 = Zygote.gradient(loss,ComponentArray(p_model))
+tmp2 = Zygote.gradient(loss,p_model)
+
+res0 = Optimization.solve(optprob, PolyOpt(), maxiters = 100)


Test just the adjoint

acoh64 · 2023-09-29T00:29:43Z

@ChrisRackauckas I removed the SDE stuff so this should be ready to go with checkpointing and nonvector parameter implementations

ChrisRackauckas · 2023-09-29T07:59:01Z

It looks like there's two test failures still seen, one in the docs and one in core 3.

added checkpointing to gauss adjoint, but only seems to be working at…

3cbfeb9

… lower tolerances (1e-4)

ai-maintainer bot reviewed Aug 26, 2023

View reviewed changes

acoh64 mentioned this pull request Aug 26, 2023

Add checkpointing to GaussAdjoint #869

Closed

acoh64 changed the title ~~added checkpointing to gauss adjoint, from https://github.com/SciML/SciMLSensitivity.jl/issues/869~~ added checkpointing to gauss adjoint Aug 26, 2023

acoh64 commented Aug 30, 2023

View reviewed changes

fixed errors in adjoint state solve from tolerances

6d364de

acoh64 and others added 3 commits August 31, 2023 10:43

removed print statements

9825793

Merge branch 'SciML:master' into master

3b48909

added stuff for SDE and RODE support, still need to add tests

5021d5f

acoh64 mentioned this pull request Sep 10, 2023

SDE compatibility for GaussAdjoint #870

Open

added test for GaussAdjoint SDEs, test fails

3d04a21

added tests for SDEs, works for bad tolerances so need to find the bug

2d84859

acoh64 and others added 2 commits September 18, 2023 09:40

Merge branch 'SciML:master' into master

16c5c4c

fixed checkpointing for GaussAdjoint - now works to same tolerance as…

a64b610

… InterpolatingAdjoint

GaussAdjoint runs for SDEs, but it is not accurate

ad973d6

ChrisRackauckas reviewed Sep 19, 2023

View reviewed changes

test/sde_scalar_stratonovich.jl Outdated Show resolved Hide resolved

Update test/sde_scalar_stratonovich.jl

8cfe5e5

Merge branch 'SciML:master' into master

32c99cc

frankschae reviewed Sep 20, 2023

View reviewed changes

nonvector parameter integration

b4f389e

ChrisRackauckas reviewed Sep 24, 2023

View reviewed changes

test/adjoint.jl Outdated

using Lux, Optimization, Plots, Random

Copy link

Member

ChrisRackauckas Sep 24, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is not in the right place?

ChrisRackauckas reviewed Sep 24, 2023

View reviewed changes

acoh64 added 2 commits September 24, 2023 14:46

cleaned up the code, will make a new file for nonnvector parameter tests

9ca3932

removed SDE code, should be ready to go now

f1cd78c

avik-pal linked an issue Oct 3, 2023 that may be closed by this pull request

Compatibility with Functors and non-vector parameters in Gauss Adjoint #868

Closed

Merge branch 'SciML:master' into master

93e87a7

avik-pal mentioned this pull request Oct 16, 2023

Gauss Adjoint Fix Tests #914

Merged

ChrisRackauckas closed this Oct 17, 2023

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

added checkpointing to gauss adjoint #884

added checkpointing to gauss adjoint #884

acoh64 commented Aug 26, 2023 •

edited

Loading

ai-maintainer bot left a comment •

edited

Loading

ChrisRackauckas commented Aug 28, 2023

ChrisRackauckas commented Aug 28, 2023

avik-pal commented Aug 28, 2023

ChrisRackauckas commented Aug 30, 2023

acoh64 Aug 30, 2023

acoh64 commented Aug 30, 2023

acoh64 commented Aug 30, 2023

acoh64 commented Aug 30, 2023

ChrisRackauckas commented Aug 30, 2023

acoh64 commented Aug 30, 2023

codecov bot commented Aug 30, 2023 •

edited

Loading

acoh64 commented Sep 13, 2023

acoh64 commented Sep 14, 2023

ChrisRackauckas commented Sep 17, 2023

acoh64 commented Sep 18, 2023

ChrisRackauckas commented Sep 19, 2023

acoh64 commented Sep 20, 2023

frankschae left a comment

ChrisRackauckas Sep 24, 2023

ChrisRackauckas Sep 24, 2023

acoh64 commented Sep 29, 2023

ChrisRackauckas commented Sep 29, 2023

added checkpointing to gauss adjoint #884

added checkpointing to gauss adjoint #884

Conversation

acoh64 commented Aug 26, 2023 • edited Loading

ai-maintainer bot left a comment • edited Loading

Choose a reason for hiding this comment

AI-Maintainer Review for PR - Added checkpointing to gauss adjoint

Title and Description 👍

Scope of Changes 👍

Testing ⚠️

Documentation ⚠️

Suggested Changes

ChrisRackauckas commented Aug 28, 2023

ChrisRackauckas commented Aug 28, 2023

avik-pal commented Aug 28, 2023

ChrisRackauckas commented Aug 30, 2023

acoh64 Aug 30, 2023

Choose a reason for hiding this comment

acoh64 commented Aug 30, 2023

acoh64 commented Aug 30, 2023

acoh64 commented Aug 30, 2023

ChrisRackauckas commented Aug 30, 2023

acoh64 commented Aug 30, 2023

codecov bot commented Aug 30, 2023 • edited Loading

Codecov Report

acoh64 commented Sep 13, 2023

acoh64 commented Sep 14, 2023

ChrisRackauckas commented Sep 17, 2023

acoh64 commented Sep 18, 2023

ChrisRackauckas commented Sep 19, 2023

acoh64 commented Sep 20, 2023

frankschae left a comment

Choose a reason for hiding this comment

ChrisRackauckas Sep 24, 2023

Choose a reason for hiding this comment

ChrisRackauckas Sep 24, 2023

Choose a reason for hiding this comment

acoh64 commented Sep 29, 2023

ChrisRackauckas commented Sep 29, 2023

acoh64 commented Aug 26, 2023 •

edited

Loading

ai-maintainer bot left a comment •

edited

Loading

codecov bot commented Aug 30, 2023 •

edited

Loading