google-deepmind / torax

TORAX: Tokamak transport simulation in JAX
https://torax.readthedocs.io
Other
345 stars 31 forks source link

Add support for prescribed source profiles #221

Closed theo-brown closed 1 month ago

theo-brown commented 3 months ago

Pull request for #212.

Plan is to use TimeInterpolatedArray for prescribing potentially time-evolving source profiles.

Key files are torax/sources/source.py and torax/sources/runtime_params.py.

theo-brown commented 3 months ago

Progress stalled pending #222 .

theo-brown commented 2 months ago

Rebased onto main ef5a49d in fbb0837.

Note that due to changes on main to TimeInterpolated variables, tests for prescribed sources no longer pass. Need to check how the changes should be propagated, which I will do in the coming days. resolved.

araju commented 2 months ago

I cleaned up a couple small things in this PR in our internal system, and need one more internal approval before I can merge this thing.

araju commented 2 months ago

Just so you're aware - this PR is not passing all our internal tests yet when merged with HEAD. I'm not sure why and have been debugging, but this is what's blocking things now. Will try and keep you posted.

theo-brown commented 2 months ago

Is there anything I can do to help? I had thought it had passed the relevant tests (using pytest) in the repo, is it failing on those or is it some internal secret (non-open) test?

araju commented 2 months ago

Is there anything I can do to help? I had thought it had passed the relevant tests (using pytest) in the repo, is it failing on those or is it some internal secret (non-open) test?

I was confused here as well. I had made very minor cosmetic changes on top, so I thought it was something I did.

But actually, I just checked out your fork and your branch and tested, and I'm getting the same errors. My guess is the rebase broke something in your code.

I'm looking into it as well, but if you could check out your branch again and run pytest torax/tests/sim.py you'll see the errors I'm getting as well.

theo-brown commented 2 months ago

Problem 1: are we updating jext/jtot/johm properly?

Given that this particular chain of problems is to do with jext, I thought I should do a sense-check on the simulations produced and have a look at their jext profiles.

This is the simulation from test_prescribed_jext.py, where the jext is smoothly interpolated between two Gaussians. In the last commit I updated the data for this test so that the test passes (it was previously failing to pass because I hadn't updated the data to reflect the simulation that was being run :facepalm:).

t=0s + a bit to avoid initial condition transients image

t=5s image

I'm concerned that I've missed a jext update somewhere, as the jtot and johm profiles seem to keep the lump from the original jext profile rather than it moving around as jext is updated. I'll admit to not really having an understanding of what should be happening physically here. @jcitrin would appreciate your input as to whether this particular sim is producing reasonable results, and potentially also insight into how/where updates to jext/jtot/johm profiles are made.

araju commented 2 months ago

Awesome! Nice catch. That definitely fixed the test you added. You may have already seen this, but 3 other test cases from tests/sim.py are stil not passing:

I didn't get a chance to find the underlying issue yesterday but am looking again today.

I'll let @jcitrin comment on the other questions you had.

theo-brown commented 2 months ago

Problem 2: code in this PR for jext calculation breaks examples where jext is a fixed fraction of the total power

The PR is failing existing integration tests. The tests in question are:

Plotting the simulation result for main (1, solid) vs this PR (2, dashed) gives the following for the test_iterhybrid_rampup scenario:

image

As you can see, this PR's solution has the jext profile (purple, bottom left) varying over time, rather than being fixed. This is because the jext is being recalculated at every timestep, and the jext formula scales the jext profile to be a fixed fraction of the total current.

So: there needs to be a way of setting jext to be a fixed fraction of jtot at t=0, and then freezing that jext profile for all time.

fcasson commented 2 months ago

I'm concerned that I've missed a jext update somewhere, as the jtot and johm profiles seem to keep the lump from the original jext profile rather than it moving around as jext is updated. I'll admit to not really having an understanding of what should be happening physically here

If you add external CD, instantaneously nothing changes in j_tot. Because of faradays law, the J_ohm change should exactly oppose what you added. Then you wait for J_ohm perturbation to decay away for the current drive to be realised. So it's not obviously wrong to me, changes in j_tot take time to show up. j_ohm should change rapidly when you change j_ext, but not j_tot

theo-brown commented 2 months ago

Also note that all of these difficulties are actually to do with jext, rather than with the prescribed source profiles. It could be worthwhile splitting this PR up so that the prescribed profiles bit goes in one PR and the jext bit goes in another.

jcitrin commented 2 months ago

As you can see, this PR's solution has the jext profile (purple, bottom left) varying over time, rather than being fixed. This is because the jext is being recalculated at every timestep, and the jext formula scales the jext profile to be a fixed fraction of the total current. So: there needs to be a way of setting jext to be a fixed fraction of jtot at t=0, and then freezing that jext profile for all time.

@araju and I just looked into this in detail.

What's happening is that you fixed a bug that we didn't realize we had, thanks! :)

The default runtime_params for jext is use_absolute_jext==False, meaning that the jext amplitude is set as a Greenwald fraction according to fext. In the "failing cases", the behaviour in your PR in that jext is time-dependent, is correct, since Ip is changing and thus the Greenwald density is changing. These changes are picked up in your PR in source_models.py due to the get_values() addition

image

Previously, jext would come from core_profiles.currents (for legacy reasons that will be now removed), and it turns out that we neglected to update jext in core_profiles.currents for time-dependent cases, so it was frozen at the initial condition.

jcitrin commented 2 months ago

So we will first fix this bug and do some cleanup, and then rebase your PR on top of that and fix any merge conflicts.

Regarding Problem 1 , yes as Francis says, it takes a long time for the current to diffuse for these parameters, so that johm bump will persist. For a sanity check, you can redo the case with a large resistivity_mult in the numerics dict, e.g. 100, which will decrease the current diffusion time and you should see faster johm and jtot dynamics

araju commented 2 months ago

What we're going to do is fix this bug at head, and then we can merge you PR as well right after. You don't need to do anything on this anymore.

I know you might be on vacation for a bit soon. Just in case, do you mind making sure I have push access to this branch in your forked repo? Just so I can push any changes that may need to happen to resolve conflicts. I should be able to handle conflicts and merge the PR without you, but that is just a precaution.

Thanks again for uncovering this!

theo-brown commented 2 months ago

Amazing! Sounds like it's a good result all round. I've sent @araju an invite to be a collaborator on my fork. Thanks for your time in handling this PR, everyone!

araju commented 1 month ago

Hey Theo, sorry for the delay with all of this. Lost a week or so with a cold, so thanks for bearing with me!

Not sure why the commit we merged separately didnt sit on top of this PR. We're still getting the hand of our internal tools with dealing with the Github PRs haha. But we can close this now that 6aeb322 is submitted.