Closed ddundo closed 5 months ago
After separating out #179, this is now again ready for review. Sorry for the messy commit history - it was a pain to rebase, so I did it semi-manually with git checkout main --
.
I think this should be easy to review now. Most of the changes are due to copy-pasting tape-related methods from mesh_seq.py
to adjoint.py
.
This looks good, but did we actually end up agreeing that yielding in this way was a good idea? There were concerns that it would make it difficult to couple Goalie to other codes whose solvers don't use such an approach.
I thought @stephankramer's concerns about coupling were related to #179 and not this. But yes, probably best to discuss together :) and the main point of this PR is to avoid memory issues (i.e. avoid taping) so maybe there are better options altogether
Thanks both for reviewing this! I made the requested changes :)
@acse-ej321 mentioned today that this didn't work when using Runge-Kutta, since each computed stage should be written on tape. The changes here shouldn't hinder taping - I just paused it inside MeshSeq._solve_forward
. So I assume you need to do pyadjoint.continue_annotation()
in your script?
@acse-ej321 mentioned today that this didn't work when using Runge-Kutta, since each computed stage should be written on tape.
Goalie itself doesn't work for general RK methods. The way it's currently designed, it will only work for theta methods.
Goalie itself doesn't work for general RK methods. The way it's currently designed, it will only work for theta methods.
I may have misunderstood then - I thought it was working on the main branch
Oh okay. I may have misremembered. Will look into it.
I may have misunderstood then - I thought it was working on the main branch
@ddundo which RK methods did you say worked for you?
Okay I've refreshed my memory. Goalie basically takes Pyadjoint apart and stiches it back together with a few assumptions. We assume solutions can be found in the adj_sol
attribute of a SolveBlock
and that the solution at the previous timestep is a dependency of the SolveBlock
. However, what if you have a multi-stage method with tendencies, etc.? In that case surely the SolveBlock
will depend on a tendency rather than a lagged solution field. We don't currently account for this. I'd be happy to be proven wrong, of course.
I realise this discussion is tangential to the PR here. Is it good to merge? If so, maybe open an issue for RK and continue the discussion there?
Thanks for looking into this! But sorry if I wasn't clear - it was @acse-ej321 who said she was using RK - so she'll be able to answer this :)
I realise this discussion is tangential to the PR here. Is it good to merge? If so, maybe open an issue for RK and continue the discussion there?
Well if I understood correctly, @acse-ej321 couldn't get RK to work with yielding. So I think I should wait to hear from her before merging
@jwallwork23 Apologies for the confusion. I likely mixed two topics up in the last meeting. I had asked @ddundo about multi-stage solvers in the previous meeting as a separate topic. I also mentioned that I could not run my current workflow between Goalie and Thetis with the #182 branch. The current workflow is built on Thetis writing to tape for the forward solve. In a recent conversation, @stephankramer mentioned there might need to be some additional functionality added to Thetis for yield
to integrate well with this PR.
Thanks for clarifying @acse-ej321!
I'm still a bit confused if this should then be addressed on the Goalie side (in this PR) :) but it sounds like a simple solution is to add a kwarg to MeshSeq.solve_forward
to prevent calling pyadjoint.pause_annotation()
?
I'll go ahead and merge this. Please open an issue and feel free to assign me if needed :)
This PR partially addresses #145:
get_solver
functions should nowyield
at every timestep instead of returning the solution dictionary at the end of the time integration. This allows us to avoid annotating to tape insolve_forward
(we still can if we want to in the future for some reason). So everything tape-related has been moved toadjoint.py
. Yielding also gives us a lot of flexibility, so I introduced a private methodMeshSeq._solve_forward
which is now called fromMeshSeq.solve_forward
andMeshSeq.get_checkpoints
, which previously had separate implementations.