search

mesh-adaptation / goalie

Goal-oriented error estimation and mesh adaptation for finite element problems solved using Firedrake
Other
2 stars 0 forks source link

Yielding in get_solver #182

Closed ddundo closed 3 months ago

ddundo commented 5 months ago

This PR partially addresses #145: get_solver functions should now yield at every timestep instead of returning the solution dictionary at the end of the time integration. This allows us to avoid annotating to tape in solve_forward (we still can if we want to in the future for some reason). So everything tape-related has been moved to adjoint.py. Yielding also gives us a lot of flexibility, so I introduced a private method MeshSeq._solve_forward which is now called from MeshSeq.solve_forward and MeshSeq.get_checkpoints, which previously had separate implementations.

ddundo commented 4 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.

ddundo commented 4 months ago

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

ddundo commented 3 months ago

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?

jwallwork23 commented 3 months ago

@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.

ddundo commented 3 months ago

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

jwallwork23 commented 3 months ago

Oh okay. I may have misremembered. Will look into it.

jwallwork23 commented 3 months ago

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.

jwallwork23 commented 3 months ago

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?

ddundo commented 3 months ago

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 :)

ddundo commented 3 months ago

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

acse-ej321 commented 3 months ago

@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.

ddundo commented 3 months ago

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()?

ddundo commented 3 months ago

I'll go ahead and merge this. Please open an issue and feel free to assign me if needed :)