Add a pellet-aware time step calculator. - #2313
Conversation
New time step calculator named pellet_aware_time_step_calculator, created for the coupling between TORAX and HPI2-NN, that aligns time steps with pellet injection (trigger) times and ablation windows, resolving each ablation window as a single step. Generic for every pellet source with trigger_times/frequency and ablation_time in their runtime parameters. Also Added documentation for this new time step calculator and the HPI2-NN pellet source. Co-authored-by: alexpanera <paneralex@hotmail.com>
| PELLET_AWARE = 'pellet_aware' | ||
|
|
||
| @enum.unique | ||
| class BaseTimeStepCalculatorType(enum.Enum): |
There was a problem hiding this comment.
base_calculator is now a nested time step calculator config object instead of a string enum, so this enum is gone.
| self, | ||
| base_calculator_type: str = 'chi', | ||
| trigger_tolerance: float = 1e-8, | ||
| pellet_source_name: str = 'pellet', |
There was a problem hiding this comment.
should not be needed. There is only ever a single pellet source name. Can hard-code it internally to 'pellet'
| pellet_source_name: For the 'pellet_aware' calculator, the name of the pellet | ||
| source to align time steps with. Defaults to 'pellet'. |
There was a problem hiding this comment.
should not be needed. There is only ever a single pellet source name. Can hard-code it internally to 'pellet'
| dtype = dt_standard.dtype | ||
|
|
||
| t = sim_state.t | ||
| pellet_params = runtime_params.sources.get(self._pellet_source_name) |
There was a problem hiding this comment.
just pellet_params = runtime_params.sources.get('pellet') and remove pellet_source_name from everywhere
| CHI = 'chi' | ||
| FIXED = 'fixed' | ||
| FROM_PREVIOUS_DT = 'from_previous_dt' | ||
| PELLET_AWARE = 'pellet_aware' |
There was a problem hiding this comment.
There is currently no validation that if the TimeStepCalculator is PELLET_AWARE , that an appropriate pellet model is actually being used.
Recommend making a new Pydantic validator on ToraxConfig . Here is an agent recommended function for that
@pydantic.model_validator(mode='after')
def _check_pellet_aware_time_step_calculator_compatibility(self) -> typing_extensions.Self:
if self.time_step_calculator.calculator_type == time_step_calculator_pydantic_model.TimeStepCalculatorType.PELLET_AWARE:
if self.sources.pellet is None:
raise ValueError(
"time_step_calculator.calculator_type='pellet_aware' requires a pellet "
"source to be configured under sources.pellet."
)
has_trigger_times = hasattr(self.sources.pellet, 'trigger_times')
has_frequency = hasattr(self.sources.pellet, 'frequency')
has_ablation_time = hasattr(self.sources.pellet, 'ablation_time')
if not ((has_trigger_times or has_frequency) and has_ablation_time):
raise ValueError(
f"time_step_calculator.calculator_type='pellet_aware' requires a pellet "
f"source supporting discrete pellet injection ('trigger_times' or 'frequency', "
f"and 'ablation_time'), e.g. 'hpi2_nn'. "
f"Got pellet model_name='{self.sources.pellet.model_name}'."
)
return self
| ) | ||
| else: | ||
| in_ablation = jnp.logical_and( | ||
| jnp.logical_and(positive_frequency, after_start), |
There was a problem hiding this comment.
this shouldn't be needed. We should validate positive_frequency on the Pydantic level, and as mentioned above, we should only be in this code pathway if the pellet injector is "ON" as seen by the new attribute.
| delta_to_next_period, | ||
| ) | ||
| dt_trigger = jnp.where( | ||
| jnp.logical_and(positive_frequency, after_start), |
| if self._dt_after_pellet is not None: | ||
| dt_after_pellet = jnp.asarray(self._dt_after_pellet, dtype=dtype) | ||
| in_post_pellet_freq = jnp.logical_and( | ||
| jnp.logical_and(positive_frequency, after_start), |
| also speed up simulations where a small dt is required for a short section, | ||
| but then a larger dt is be appropriate for the remainder of the simulation. | ||
|
|
||
| * ``'pellet_aware'`` |
There was a problem hiding this comment.
list here and explain in more detail all the new attributes of the pellet_aware time_step_calculator
|
@tamaranorman could you also take a look at reviewing this one? |
|
|
||
|
|
||
| class PelletAwareTimeStepCalculator(time_step_calculator.TimeStepCalculator): | ||
| """TimeStepCalculator that resolves pellet trigger and ablation windows. |
There was a problem hiding this comment.
Is this something that could be made generic to many other discrete events to allow supporting other things moving forward
There was a problem hiding this comment.
For this PR I have kept it only for pellets.
| do not skip over these events. | ||
|
|
||
| Arguments: | ||
| base_calculator_type: The type of the base time step calculator to use for |
There was a problem hiding this comment.
Should this be instead base_calculator to make sure that any parameter are created if a calculator has these - then it can be something other than chi or fixed potentially
…culator config object. Remove the pellet_source_name option and read the single 'pellet' source directly.
…exactly one of trigger_times/frequency and detecting the trigger locally ('at_trigger = jnp.any(jnp.abs(t - triggers) <= tol') from the pellet source own trigger_tolerance. Vectorize the trigger scan, and drop the redundant array casts, the preventive early returns (now guaranteed by the ToraxConfig validator), and the duplicate exact_t_final limit already applied by the base next_dt.
Assume a strictly positive frequency (validated by the pellet source config) instead of silently substituting 1.0. Replace the frequency_t_start "after start" gate with a generic injection_enabled flag read from the pellet source (defaulting to on), so the injector can be toggled on and off during the run while frequency_t_start only sets the phase. Because floating point prevents a step from landing exactly on a boundary, an injection_enabled that switches on exactly at a pellet time can miss it, so it must be turned on a small margin before the intended pellet time (documented in the code). Also drop the redundant + tol phase offset.
…c build/validation, and the ToraxConfig compatibility validator
jcitrin
left a comment
jcitrin
left a comment
There was a problem hiding this comment.
Thanks for this feature! Any additional minor changes we can take care of internally since I am aware that you are at the end of your project.
jcitrin
added
the
copybara:import-manual
3 participants
New time step calculator named pellet_aware_time_step_calculator, created for the coupling between TORAX and HPI2-NN, that aligns time steps with pellet injection (trigger) times and ablation windows, resolving each ablation window as a single step. Generic for every pellet source with trigger_times/frequency and ablation_time in their runtime parameters. Also Added documentation for this new time step calculator and the HPI2-NN pellet source.