Add simple radiation model (Prad as a fraction of P_tot_in) #572

theo-brown · 2024-11-26T20:52:57Z

Add a simplified radiation sink model, P_rad = α*(P_aux + P_alpha + P_ohmic)

This is based on the functionality in JETTO:

Todo:

Calculate prad as a fraction of ptot
Add to outputs and plotting
Unit tests

theo-brown · 2024-11-27T09:32:21Z

@Nush395 if you've got time, could you lend a hand with this? I have a couple of questions about implicit sources and sources that link back.

At the moment I'm finding that my implementation is nonzero when I inspect it during runtime but has no effect on the state. In the output, it is 0.

Nush395 · 2024-11-27T10:55:21Z

@Nush395 if you've got time, could you lend a hand with this? I have a couple of questions about implicit sources and sources that link back.

At the moment I'm finding that my implementation is nonzero when I inspect it during runtime but has no effect on the state. In the output, it is 0.

Yes for sure! Happy to answer questions, feel free to drop them on this PR or we can go over on a call if that's easier :)

Nush395 · 2024-11-27T11:06:22Z

@Nush395 if you've got time, could you lend a hand with this? I have a couple of questions about implicit sources and sources that link back.

At the moment I'm finding that my implementation is nonzero when I inspect it during runtime but has no effect on the state. In the output, it is 0.

Just took a quick look at the change. Just a sanity check, have you run the prad config with

'radiation_heat_sink': {'mode': 'model_based'}

?
(As the default mode that you inherit from the RuntimeParams is set to ZERO.)

theo-brown · 2024-11-27T11:13:17Z

Ah, thanks for spotting that. That would be why it wasn't being used! I'll do that and then see where I get to. Sorry for not spotting the obvious thing!

Nush395 · 2024-11-27T11:22:06Z

Ah, thanks for spotting that. That would be why it wasn't being used! I'll do that and then see where I get to. Sorry for not spotting the obvious thing!

No worries at all, I've run into that gotcha myself!

theo-brown · 2024-11-27T12:04:39Z

As of 69f5ada, unit tests fail because I've added a term to post_processed_outputs that is not present in all of the current .nc files. Hence, all test data will need to be regenerated. Is there a speedy way of doing this?

Nush395 · 2024-11-27T12:13:16Z

As of 69f5ada, unit tests fail because I've added a term to post_processed_outputs that is not present in all of the current .nc files. Hence, all test data will need to be regenerated. Is there a speedy way of doing this?

There is, there is a run_and_save_all_benchmarks.py script that we use. See the bottom of https://torax.readthedocs.io/en/latest/contribution_tips.html#contribution-tips that @jcitrin recently added to. (Will also make this a bit more visible, if you have suggestions on where this should be please let us know!)

theo-brown · 2024-11-27T14:37:08Z

@Nush395 it now makes sense to ask my questions!

Should source tests be written on a per-source basis, or use inheritance? See e.g. BootstrapCurrentSourceTest (which is a SourceTestCase whose test_extraction_of_relevant_profile_from_output is identical to that of other psi sources), vs FusionHeatSourceTest (which is an IonElSourceTestCase to avoid reimplementing such functions for all Te + Ti sources).
Where is the OhmicHeatSource tested? As far as I can tell it's the only other source that has links_back=True and I was hunting around to try and copy its test pattern.

Required regenerating some integration tests with new Prad output

Use source_builder in all tests, except from BootstrapCurrentSource and QeiSource

theo-brown · 2024-11-27T15:53:31Z

For the time being I've gone for only inheriting from SourceTestCase, which means:

Inherit
- test_expected_mesh_states
- test_runtime_params_builds_dynamic_params
Reimplement
- test_extraction_of_relevant_profile_from_output
- test_invalid_source_types_raise_errors
- test_source_value

I couldn't quite figure out why test_extraction_of_relevant_profile_from_output and test_invalid_source_types_raise_errors weren't implemented in the parent class - it seems to me that they should be functions of the source properties only? Regardless, this seemed to be the pattern implemented by most other sources, so I went with that.

Nush395 · 2024-11-27T20:34:14Z

@Nush395 it now makes sense to ask my questions!

Hey Theo, sorry for the slow response on this.

Should source tests be written on a per-source basis, or use inheritance? See e.g. BootstrapCurrentSourceTest (which is a SourceTestCase whose test_extraction_of_relevant_profile_from_output is identical to that of other psi sources), vs FusionHeatSourceTest (which is an IonElSourceTestCase to avoid reimplementing such functions for all Te + Ti sources).

In the general case I'd advise inheriting from SourceTestCase as that has some basic test cases relevant to all sources in general. If your source only affects a single profile there is also the SingleProfileSourceTestCase https://github.com/google-deepmind/torax/blob/main/torax/sources/tests/test_lib.py#L113 that has more specialised tests for sources that only affect a single source (as yours does).

Where is the OhmicHeatSource tested? As far as I can tell it's the only other source that has links_back=True and I was hunting around to try and copy its test pattern.

That's a good question. Taking a quick look now it doesn't seem there is one which is something we can rectify. Thanks for raising this.

Nush395 · 2024-11-27T20:39:14Z

I couldn't quite figure out why test_extraction_of_relevant_profile_from_output and test_invalid_source_types_raise_errors weren't implemented in the parent class - it seems to me that they should be functions of the source properties only? Regardless, this seemed to be the pattern implemented by most other sources, so I went with that.

For test_invalid_source_types_raise_errors that is implemented in the SingleProfileSourceTestCase and you are right for the test_extraction_of_relevant_profile_from_output this is likely a legacy thing that could be tidied up and pushed into the parent class now. I can TAL at doing that and doing some test tidying up.

theo-brown · 2024-11-28T08:06:57Z

I can TAL at doing that and doing some test tidying up.

I'm more than happy to tackle this in a separate PR if you want to have a look at the ohmic source? Unless you'd rather do both!

Nush395 · 2024-11-28T08:53:38Z

I can TAL at doing that and doing some test tidying up.

I'm more than happy to tackle this in a separate PR if you want to have a look at the ohmic source? Unless you'd rather do both!

That would be great if you have time, thanks!

jcitrin

Looks really great, thanks!

Appreciate also the extensive testing + extensions to the tests.

jcitrin · 2024-11-28T17:14:01Z

torax/examples/prad.py

Not sure there's much point in making an example config out of a single feature?

Our standard pattern is to make a new integration test in tests/tests_data with new features for regression and integration testing.

Suggest to remove it here, make a new integration test (to add to tests/sim.py also) based on one of the standard cases like test_iterhybrid_predictor_corrector

jcitrin · 2024-11-29T10:48:47Z

torax/sources/tests/radiation_heat_sink.py

a bit of a nit:

radiation_heat_sink isn't specific enough, it could refer to cyclotron, bremsstrahlung, e.g.

How about impurity_radiation_heat_sink ?

Was also thinking about line_radiation_heat_sink, but some radiation models like Mavrin polynomials include both line radiation + cremsstrahlung, so maybe ok to specify just "impurity"

jcitrin · 2024-11-29T15:39:27Z

torax/sources/radiation_heat_sink.py

+) -> jax.Array:
+    """Model function for radiation heat sink.
+
+    In this model, a fixed % of the total power input to the temp_el equation is lost."""


nit: confusing wording. Maybe "A fixed % of the total power input is lost as a sink of the temp_el equation"

jcitrin · 2024-11-29T16:58:26Z

torax/sources/radiation_heat_sink.py

+
+    Qtot_in = jnp.sum(jnp.stack(source_profiles), axis=0)
+    Ptot_in = math_utils.cell_integration(Qtot_in * geo.vpr, geo)
+    Vtot = math_utils.cell_integration(geo.vpr, geo)


can just use geo.volume_face[-1] . It will differ by a tiny amount (like 0.01% or so, depending on the grid and geometry resolution) from this cell_integration, since geo.volume_face comes directly from the vpr integration on the original geometry grid, as opposed to the interpolation on the TORAX cell grid.

jcitrin · 2024-11-29T17:12:25Z

torax/sources/radiation_heat_sink.py

+    # Manually remove sources that will not be summed
+    sources_to_sum = source_models.temp_el_sources | source_models.temp_ion_sources
+    sources_to_sum.pop(SOURCE_NAME, None)
+    sources_to_sum.pop(bremsstrahlung_heat_sink.SOURCE_NAME, None)


Can easily imagine a bug where we add more heat sinks and forget to update this function.

Can we guard against it somehow? Maybe check all remaining keys if they contain the substring 'sink' and then raise an error. This depends on having the discipline to call all new sinks "something_heat_sink", but it's better than nothing.

jcitrin · 2024-11-29T17:20:56Z

torax/tests/test_data/test_changing_config_before_BASE_171788.nc

jcitrin · 2024-11-29T17:21:01Z

torax/tests/test_data/test_changing_config_before_LOCAL_171788.nc

jcitrin · 2024-11-29T17:21:08Z

torax/tests/test_data/test_changing_config_before_REMOTE_171788.nc

jcitrin · 2024-11-29T17:21:20Z

torax/tests/test_data/test_implicit.nc

how come this needed to be replaced?

jcitrin · 2024-11-29T17:33:43Z

torax/sources/radiation_heat_sink.py

+    # and summing over *both* ion and electron heating
+
+    def get_temp_el_profile(source_name: str, source: source_lib.Source) -> jax.Array:
+        # TODO: Currently this recomputes the profile for each source, which is inefficient


agree. We will prioritize changing the pattern such that this won't happen.

theo-brown changed the title ~~Add simple radiation model (Prad as a fraction of Ptot)~~ Add simple radiation model (Prad as a fraction of P_tot_in) Nov 26, 2024

theo-brown mentioned this pull request Nov 27, 2024

Feature parity for STEP simulations #420

Open

theo-brown added 4 commits November 27, 2024 15:24

Add radiation_heat_sink, proportional to Qtot

a005f07

Add Prad to post_processing and plots

ff0c854

Required regenerating some integration tests with new Prad output

Support sources with links_back=True

cc16cce

Use source_builder in all tests, except from BootstrapCurrentSource and QeiSource

Add unit test for RadiationHeatSink

7a0936a

theo-brown force-pushed the theo-brown/add-prad branch from 3cd6cad to 7a0936a Compare November 27, 2024 15:35

Merge branch 'main' into theo-brown/add-prad

601cde4

theo-brown marked this pull request as ready for review November 27, 2024 15:41

theo-brown requested a review from jcitrin November 27, 2024 15:48

theo-brown force-pushed the theo-brown/add-prad branch from f737ad0 to 601cde4 Compare November 27, 2024 15:54

Sum over both ion and el heat sources

65e5836

Switch to radially constant fraction of Pin

c6710b6

jcitrin reviewed Nov 29, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add simple radiation model (Prad as a fraction of P_tot_in) #572

Add simple radiation model (Prad as a fraction of P_tot_in) #572

theo-brown commented Nov 26, 2024 •

edited

Loading

theo-brown commented Nov 27, 2024

Nush395 commented Nov 27, 2024

Nush395 commented Nov 27, 2024 •

edited

Loading

theo-brown commented Nov 27, 2024

Nush395 commented Nov 27, 2024

theo-brown commented Nov 27, 2024 •

edited

Loading

Nush395 commented Nov 27, 2024

theo-brown commented Nov 27, 2024

theo-brown commented Nov 27, 2024

Nush395 commented Nov 27, 2024 •

edited

Loading

Nush395 commented Nov 27, 2024

theo-brown commented Nov 28, 2024

Nush395 commented Nov 28, 2024

jcitrin left a comment

jcitrin Nov 28, 2024

jcitrin Nov 29, 2024

jcitrin Nov 29, 2024

jcitrin Nov 29, 2024

jcitrin Nov 29, 2024

jcitrin Nov 29, 2024

jcitrin Nov 29, 2024

jcitrin Nov 29, 2024

jcitrin Nov 29, 2024

jcitrin Nov 29, 2024

Add simple radiation model (Prad as a fraction of P_tot_in) #572

Are you sure you want to change the base?

Add simple radiation model (Prad as a fraction of P_tot_in) #572

Conversation

theo-brown commented Nov 26, 2024 • edited Loading

theo-brown commented Nov 27, 2024

Nush395 commented Nov 27, 2024

Nush395 commented Nov 27, 2024 • edited Loading

theo-brown commented Nov 27, 2024

Nush395 commented Nov 27, 2024

theo-brown commented Nov 27, 2024 • edited Loading

Nush395 commented Nov 27, 2024

theo-brown commented Nov 27, 2024

theo-brown commented Nov 27, 2024

Nush395 commented Nov 27, 2024 • edited Loading

Nush395 commented Nov 27, 2024

theo-brown commented Nov 28, 2024

Nush395 commented Nov 28, 2024

jcitrin left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

theo-brown commented Nov 26, 2024 •

edited

Loading

Nush395 commented Nov 27, 2024 •

edited

Loading

theo-brown commented Nov 27, 2024 •

edited

Loading

Nush395 commented Nov 27, 2024 •

edited

Loading