Add `slice_time` and `slice_index` arguments to IMAS geometry loader #1416

theo-brown · 2025-08-05T12:08:57Z

This enables loading an equilibrium from any timeslice in an equilibrium IDS. Useful if resuming from pre-existing simulations.

Implementation
Testing

theo-brown · 2025-08-06T09:27:57Z

@MateoBell what's the provenance of the ITERhybrid_COCOS17_IDS_ddv4.nc file? I'd like to extend it to have multiple time slices so that I can test what I've implemented in this PR.

torax/_src/geometry/tests/imas_test.py

theo-brown · 2025-08-06T16:39:20Z

Linked to #1391. During investigating the different geometry loaders in ST, I am comparing eqdsk from various timesteps with CHEASE, IMAS, etc. Hence, I need to be able to load different slices from the IMAS equilibrium IDS.

MateoBell · 2025-08-06T17:16:31Z

@MateoBell what's the provenance of the ITERhybrid_COCOS17_IDS_ddv4.nc file? I'd like to extend it to have multiple time slices so that I can test what I've implemented in this PR.

The .nc file was generated with chease converting the EQDSK_ITERhybrid_COCOS02.eqdsk (actually not completely sure which eqdsk of the examples, but it is the one which was used to compare to the CHEASE default geometry file). There is a script to run CHEASE from an eqdsk and save it to an IDS in a database, and then I converted it in netCDF. The script probably requires IMAS core so I can try to generate an IDS with multiple time slices with given eqdsk if needed.

theo-brown · 2025-08-08T12:34:56Z

Thanks! I think for this PR we ideally would have an IDS with different profiles across the time steps we're comparing. This will also pave the way for smoothly interpolating between geometries to have evolving geometry over the run, but this may require rerunning CHEASE (or equivalent fixed-boundary solver) to produce something consistent.

We could, for example, run CHEASE twice using the pprime and FFprime profiles from the beginning and end of the TORAX iterhybrid rampup case as inputs, to give two different geometries for the start and end of the ramp.

What do you think?

MateoBell · 2025-08-11T07:51:59Z

We could, for example, run CHEASE twice using the pprime and FFprime profiles from the beginning and end of the TORAX iterhybrid rampup case as inputs, to give two different geometries for the start and end of the ramp.

Yes fully agree ! I think it is a good idea so we can get different geometries at different times to compare.

theo-brown · 2025-08-11T10:51:54Z

I've never run CHEASE before, if you're able would you be able to generate the two geometries?

Failing CI is just the placeholder test that's waiting for the new data.

MateoBell · 2025-08-14T14:30:28Z

Ok I think I should be able to do it, will try to do it asap !

jcitrin · 2025-08-14T17:29:04Z

Thanks both! Having a multiple time-slice IDS loader and a test for it would be ideal

MateoBell · 2025-08-19T07:04:34Z

I added the multiple time slices IDS, but I just noticed the netCDF file is pretty heavy. I generated it with CHEASE for all time slices of the output simulation from iterhybrid_rampup config. If needed I can reduce the number of time slices to reduce the weight.
Also I modified a bit the tests to make sure the IDS was generated properly and passing it, I fixed it rapidly so you might want to do it more properly.

theo-brown · 2025-08-19T09:49:24Z

Incredible, thanks so much @MateoBell. Your fixes to the tests also seem fine!

In the latest force push I just rebased onto main.

For reference, the full .nc is 56.21 MB (GitHub's recommended limit is 50MB). I think we can definitely reduce the number of timeslices - maybe we should go down to ~10 from ~40?

jcitrin · 2025-08-19T10:23:55Z

Thanks both!

~10 from ~40?

Yes, good idea to not load up the repo with large files

MateoBell · 2025-08-19T11:34:53Z

Okay just replaced it with 11 time slices evenly spaced in time, the file is now ~16MB !

torax/_src/geometry/pydantic_model.py

torax/_src/geometry/imas.py

Nush395

thanks Theo! Left a few comments

torax/_src/geometry/pydantic_model.py

torax/_src/geometry/tests/imas_test.py

torax/_src/geometry/imas.py

torax/data/third_party/geo/README

Nush395 · 2025-09-01T14:31:26Z

torax/_src/geometry/imas.py

+  if slice_time is not None:
+    # Find the closest time in the IDS that is <= slice_time
+    slice_index = (
+        np.searchsorted(equilibrium.time, slice_time, side="right") - 1


is equilibrium.time guaranteed to be sorted?

Unsure. Happy to add a check and then sort if it's not.

torax/_src/geometry/imas.py

jcitrin · 2025-09-10T11:24:05Z

@theo-brown , sorry this one slipped through the cracks.

If it's still relevant can you respond to any comments, sync, and we'll re-review?

theo-brown · 2025-09-10T12:22:00Z

No problem, it also slipped off my radar! I didn't see @Nush395 had reviewed it. Changes made! I wasn't sure whether it's ok to assume time is sorted in the equilibrium IDS, I'm not sure whether this is enforced. Can easily add a check and do a sort-then-search if you think it's required.

Nush395

thanks Theo! Looks great

torax/_src/geometry/standard_geometry.py

torax/_src/imas_tools/input/equilibrium.py

torax/_src/imas_tools/input/tests/equilibrium_test.py

torax/data/third_party/geo/README

Added an IDS with multiple time slices for testing

theo-brown · 2025-09-11T18:39:26Z

Thanks for the review @Nush395!

Changes:

Squashed commits
Added check if time is sorted, and sort if not
Removed "(default 0)"
Fixed typo in README
Switched to assertRaisesRegex to check the error properly

Nush395 · 2025-09-12T08:42:00Z

torax/_src/imas_tools/input/equilibrium.py

+        f"{len(equilibrium.time_slice)} time slices"
+    )
+  IMAS_data = equilibrium.time_slice[slice_index]
+  B_0 = np.abs(IMAS_data.global_quantities.magnetic_axis.b_field_phi)


I just noticed this additional unused B_0, should this be removed? If so I can do so.

Yeah this is definitely wrong, not sure why I added it 🤔 please remove!

theo-brown · 2025-09-12T10:18:40Z

@Nush395 I think something might've gone wrong in the copybara import.

From the merge commit: 2e5f70e

torax/torax/_src/geometry/pydantic_model.py

Line 314 in 2e5f70e

slice_time: float | None = None

From the branch: 5d2093b

torax/torax/_src/geometry/pydantic_model.py

Line 314 in 5d2093b

slice_time: pydantic.NonNegativeFloat | None = None

Nush395 · 2025-09-12T10:35:41Z

@Nush395 I think something might've gone wrong in the copybara import.

From the merge commit: 2e5f70e

torax/torax/_src/geometry/pydantic_model.py

Line 314 in 2e5f70e

slice_time: float | None = None

From the branch: 5d2093b

torax/torax/_src/geometry/pydantic_model.py

Line 314 in 5d2093b

slice_time: pydantic.NonNegativeFloat | None = None

@jcitrin commented in internal review that this slice_time could be negative so I modified this before submitting. Will call out any changes more explicitly in future!

jcitrin: "time can be negative, it's a user definition. There are use-cases where t=0 is some "important reference time" in the middle of a simulation"

theo-brown · 2025-09-12T10:44:11Z

Ok, cool! Thanks for clarifying :)

theo-brown force-pushed the imas-geometry-timeslice branch from 2422ada to 6632990 Compare August 5, 2025 12:09

theo-brown force-pushed the imas-geometry-timeslice branch 2 times, most recently from e939fb9 to 678f7cd Compare August 6, 2025 09:56

theo-brown commented Aug 6, 2025

View reviewed changes

torax/_src/geometry/tests/imas_test.py Outdated Show resolved Hide resolved

theo-brown force-pushed the imas-geometry-timeslice branch from 678f7cd to 40fd9cd Compare August 12, 2025 13:51

theo-brown force-pushed the imas-geometry-timeslice branch from 4b4b771 to 6c9e5bd Compare August 19, 2025 09:44

theo-brown requested a review from jcitrin August 19, 2025 12:18

jcitrin reviewed Aug 20, 2025

View reviewed changes

torax/_src/geometry/pydantic_model.py Outdated Show resolved Hide resolved

torax/_src/geometry/imas.py Outdated Show resolved Hide resolved

theo-brown force-pushed the imas-geometry-timeslice branch from 4234aa1 to 95a7fde Compare August 21, 2025 20:49

theo-brown requested a review from jcitrin August 21, 2025 21:03

theo-brown requested a review from Nush395 September 1, 2025 08:12

Nush395 reviewed Sep 1, 2025

View reviewed changes

theo-brown force-pushed the imas-geometry-timeslice branch from 95a7fde to 3cccd78 Compare September 10, 2025 12:19

theo-brown requested a review from Nush395 September 11, 2025 12:53

Nush395 approved these changes Sep 11, 2025

View reviewed changes

theo-brown force-pushed the imas-geometry-timeslice branch from 3cccd78 to e32aa42 Compare September 11, 2025 18:34

theo-brown force-pushed the imas-geometry-timeslice branch from e32aa42 to 51b30bf Compare September 11, 2025 18:35

Add slice_time and slice_index arguments to IMAS geometry loader

5d2093b

Added an IDS with multiple time slices for testing

theo-brown force-pushed the imas-geometry-timeslice branch from 51b30bf to 5d2093b Compare September 11, 2025 18:38

Nush395 added the copybara:import-manual Set when ready for copybara manual import label Sep 11, 2025

Nush395 reviewed Sep 12, 2025

View reviewed changes

copybara-service bot merged commit 2e5f70e into main Sep 12, 2025
19 checks passed

theo-brown deleted the imas-geometry-timeslice branch September 12, 2025 10:13

theo-brown restored the imas-geometry-timeslice branch September 12, 2025 10:14

theo-brown deleted the imas-geometry-timeslice branch September 12, 2025 10:43

Add slice_time and slice_index arguments to IMAS geometry loader #1416

Add slice_time and slice_index arguments to IMAS geometry loader #1416

Uh oh!

Conversation

theo-brown commented Aug 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

theo-brown commented Aug 6, 2025

Uh oh!

Uh oh!

theo-brown commented Aug 6, 2025

Uh oh!

MateoBell commented Aug 6, 2025

Uh oh!

theo-brown commented Aug 8, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

MateoBell commented Aug 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

theo-brown commented Aug 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

MateoBell commented Aug 14, 2025

Uh oh!

jcitrin commented Aug 14, 2025

Uh oh!

MateoBell commented Aug 19, 2025

Uh oh!

theo-brown commented Aug 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jcitrin commented Aug 19, 2025

Uh oh!

MateoBell commented Aug 19, 2025

Uh oh!

Uh oh!

Uh oh!

Nush395 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Nush395 Sep 1, 2025

Choose a reason for hiding this comment

Uh oh!

theo-brown Sep 10, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

jcitrin commented Sep 10, 2025

Uh oh!

theo-brown commented Sep 10, 2025

Uh oh!

Nush395 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

theo-brown commented Sep 11, 2025

Uh oh!

Nush395 Sep 12, 2025

Choose a reason for hiding this comment

Uh oh!

theo-brown Sep 12, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

theo-brown commented Sep 12, 2025

Uh oh!

Nush395 commented Sep 12, 2025

Uh oh!

Add `slice_time` and `slice_index` arguments to IMAS geometry loader #1416

Add `slice_time` and `slice_index` arguments to IMAS geometry loader #1416

theo-brown commented Aug 5, 2025 •

edited

Loading

theo-brown commented Aug 8, 2025 •

edited

Loading

MateoBell commented Aug 11, 2025 •

edited

Loading

theo-brown commented Aug 11, 2025 •

edited

Loading

theo-brown commented Aug 19, 2025 •

edited

Loading