Support different sources for trapped particle fraction#2283
Open
theo-brown wants to merge 5 commits into
Open
Support different sources for trapped particle fraction#2283theo-brown wants to merge 5 commits into
theo-brown wants to merge 5 commits into
Conversation
(1/4) Adds support for different methods of computing f_trap. Defaults to Sauter analytic approximation.
(2/4) Adds support for loading the trapped fraction from CHEASE and IMAS, where it is already computed.
(3/4) Adds support for computing the trapped fraction from the full bounce-averaged integral over the 2D equilibrium, for EQDSK and IMAS.
(4/4) Documents the SAUTER/FILE/EXACT options and switches the STEP flattop example to use the exact bounce-averaged integral.
theo-brown
force-pushed
the
trapped-particle-integral
branch
from
July 13, 2026 12:44
36f8124 to
e8232c1
Compare
jcitrin
requested changes
Jul 21, 2026
jcitrin
left a comment
jcitrin
left a comment
Collaborator
There was a problem hiding this comment.
Great stuff. Beyond the comments in the review:
- Please solve merge conflicts
- The PR is very large and contains several independent refactors/features. For example, the refactor of the sauter formula, and the modification of its calculation in GeometryIntermediates instead of in the neoclassical calculation functions (which is what is leading to the nc file updates), is its own PR. Then beyond that, could consider splitting the rest into several smaller PRs (e.g. on the new config plumbing, then implementation of the FILE cases, then the EXACT formulas and implementation, etc.)
- Should there be a new sim test or two with EXACT and FILE cases?
Comment on lines
+39
to
+44
| self.assertTrue(np.all(trapped_fraction >= 0.0)) | ||
| self.assertTrue(np.all(trapped_fraction <= 1.0)) | ||
| self.assertGreater( | ||
| np.mean(np.diff(trapped_fraction) >= -1e-6), | ||
| 0.8, | ||
| ) |
Comment on lines
+97
to
+108
| # No trapped particles on the magnetic axis, where B is uniform. | ||
| self.assertAlmostEqual(float(trapped_fraction[0]), 0.0) | ||
| # The trapped particle fraction is a fraction, so must lie in [0, 1]. | ||
| self.assertTrue(np.all(trapped_fraction >= 0.0)) | ||
| self.assertTrue(np.all(trapped_fraction <= 1.0)) | ||
| # Trapped fraction increases with normalized radius over most of the | ||
| # profile (small deviations from strict monotonicity are possible near | ||
| # the edge for diverted geometries, due to the X-point). | ||
| self.assertGreater( | ||
| np.mean(np.diff(trapped_fraction) >= -1e-6), | ||
| 0.8, | ||
| ) |
Comment on lines
+89
to
+98
| self.assertIsNotNone(trapped_fraction) | ||
| self.assertTrue(np.all(trapped_fraction >= 0.0)) | ||
| self.assertTrue(np.all(trapped_fraction <= 1.0)) | ||
| # Trapped fraction increases with normalized radius over most of the | ||
| # profile (small deviations from strict monotonicity are possible near | ||
| # the edge for diverted geometries). | ||
| self.assertGreater( | ||
| np.mean(np.diff(trapped_fraction) >= -1e-6), | ||
| 0.8, | ||
| ) |
Comment on lines
+135
to
+140
| self.assertTrue(np.all(trapped_fraction >= 0.0)) | ||
| self.assertTrue(np.all(trapped_fraction <= 1.0)) | ||
| self.assertGreater( | ||
| np.mean(np.diff(trapped_fraction) >= -1e-6), | ||
| 0.8, | ||
| ) |
Comment on lines
+328
to
+341
| case base.TrappedFractionSource.FILE: | ||
| if not IMAS_data.profiles_1d.trapped_fraction: | ||
| raise ValueError( | ||
| "trapped_fraction_source=FILE requires the equilibrium IDS to" | ||
| " populate profiles_1d.trapped_fraction, but this IDS does" | ||
| " not. Use trapped_fraction_source=EXACT to compute it directly" | ||
| " from the 2D equilibrium instead, or SAUTER for the analytic" | ||
| " approximation." | ||
| ) | ||
| trapped_fraction = np.asarray(IMAS_data.profiles_1d.trapped_fraction) | ||
| case base.TrappedFractionSource.EXACT: | ||
| exact_trapped_fraction = _calculate_exact_trapped_fraction( | ||
| IMAS_data, np.asarray(IMAS_data.profiles_1d.gm5) | ||
| ) |
Collaborator
There was a problem hiding this comment.
It would be good to have a test that tests whether the FILE and EXACT calculations agree, for an IDS that has both available
Comment on lines
+452
to
+455
| trapped_fraction = formulas.calculate_sauter_trapped_fraction( | ||
| epsilon=LY['epsilon'], | ||
| delta=0.5 * (LY['deltau'] + LY['deltal']), | ||
| ) |
Collaborator
There was a problem hiding this comment.
can future proof with a match case, or just have a TODO to extend when MEQ team provides FILE values
| _IMAS_RECTANGULAR_GRID_TYPE = 1 | ||
|
|
||
|
|
||
| def _calculate_exact_trapped_fraction( |
Collaborator
There was a problem hiding this comment.
This is not unit tested.
Also, playing Devil's Advocate. Is this needed in TORAX? If trapped fraction using the same formula is part of the equilibrium IDS, then could it just be the responsibility of the data owner to make sure it's provided? In other words, is there a use-case of having both EXACT and FILE for IMAS?
Comment on lines
+354
to
+361
| exact_is_unreliable = ( | ||
| np.isnan(exact_trapped_fraction) | ||
| | (exact_trapped_fraction < 0.0) | ||
| | (exact_trapped_fraction > 1.0) | ||
| ) | ||
| trapped_fraction = np.where( | ||
| exact_is_unreliable, sauter_trapped_fraction, exact_trapped_fraction | ||
| ) |
Collaborator
There was a problem hiding this comment.
same comment as in EQDSK regarding the silent overwriting by Sauter
| The effective trapped particle fraction of this flux surface. | ||
| """ | ||
| B_max = B.max() | ||
| lam = np.linspace(0.0, 1.0, 101) / B_max |
Collaborator
There was a problem hiding this comment.
why 101 and not some other number? Also, this could be a module private constant variable
| ) | ||
|
|
||
|
|
||
| def calculate_bounce_averaged_trapped_fraction( |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The Sauter model for trapped particle fraction is not valid across all aspect ratios.
This PR adds additional options for f_trap:
As this is done once at geometry construction time, it is a one-time cost.