Feature/muscle3 actor #1439
Feature/muscle3 actor #1439
Conversation
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
| pytest.importorskip("imas_core") | ||
|
|
||
|
|
||
| def source_for_tests(): |
There was a problem hiding this comment.
micronit: maybe rename to "data_source" and "data_sink" or "data_target" , since in integrated modelling "source" and "sink" have clear meanings for terms in the transport PDEs, and it's confusing to see them used here for equilibrium data sources and data targets
| from torax._src.orchestration.muscle3_utils import get_t_next | ||
| from torax._src.orchestration.muscle3_utils import merge_extra_vars | ||
| from torax._src.orchestration.muscle3_utils import receive_equilibrium | ||
| from torax._src.orchestration.muscle3_utils import receive_ids_through_muscle3 |
There was a problem hiding this comment.
Can remove? It looks like this is only used within muscle3_utils as a helper method for receive_equilibrium. Can then make it private in muscle3_utils
| config: model_config.ToraxConfig, | ||
| port_name: str, | ||
| ) -> Tuple[GeometryProvider, float, Optional[float], ExtraVarCollection]: | ||
| # TODO: receive all inputs [equilibrium, core_profiles, core_profiles] |
| extra_var_col = ExtraVarCollection() | ||
| torax_config_dict = get_geometry_config_dict(config) | ||
| torax_config_dict["geometry_type"] = "imas" | ||
| with DBEntry("imas:memory?path=/", "w") as db: |
There was a problem hiding this comment.
Could use some comments here. As I understand it, this provides a geometry provider with all the time slices existent in the sent equilibrium IDS?
| if instance.is_connected("equilibrium_f_init"): | ||
| step_fn._geometry_provider, t_cur, t_next_outer, extra_var_col = receive_equilibrium( | ||
| instance, | ||
| step_fn._geometry_provider, | ||
| sim_state, | ||
| torax_config, | ||
| "equilibrium_f_init", | ||
| ) | ||
| if first_run or instance.is_connected("equilibrium_f_init"): |
There was a problem hiding this comment.
I think this can be cleaner with respect to the branches "equilibrium_f_init_ exists" vs "not".
There are duplicated calculations. prepare_simulation should only be done when we have the "correct" initial config, which depends on whether equilibrium_f_init is connected or not. By reordering things, I think you can avoid the recalculation of specific prepare_simulation outputs.
| ) | ||
| first_run = False | ||
|
|
||
| equilibrium_data = torax_state_to_imas_equilibrium(sim_state, post_processed_outputs) |
There was a problem hiding this comment.
nit: both here and below, can put this (and merge_extra_vars) within the output_all_timeslices if statement, to avoid unnecessary computation if not needed
| t_next=t_next_inner, | ||
| ) | ||
|
|
||
| # S |
There was a problem hiding this comment.
Some more comments would be beneficial here clarifying the role of equilibrium_init vs equilibrium_s.
For equilibrium_init, I can imagine that in some workflows, then incoming equilibrium would even be the entire time-series.
However, for equilibrium_s, which is inter-step, then it only makes sense for a single equilibrium slice to be included in the provider, which will then be used for both geo_t and geo_t_plus_dt.
If both equilibrium_init and equilibrium_s are present, then for the first time in the loop the role of the first equilibrium_s seems duplicate to equilibrium_init, which in this case should only provide a single initial timeslice
| sim_state, | ||
| torax_config, | ||
| "equilibrium_s", | ||
| ) |
There was a problem hiding this comment.
Another note on equilibrium_s, and the need to modify the internal TORAX API for such "new" cases like this where new equilibria are provided inter_step.
Inside the step function, we have _get_geo_and_dynamic_runtime_params_at_t_plus_dt_and_phibdot . This uses geo_t and geo_t_plus_dt to calculate the phibdot term, which is non-zero is geo_t and geo_t_plus_dt have specific differences (e.g. if the toroidal flux at the LCFS is different). This can be important in rampup and rampdown.
However, in workflows with equilibrium_s, I assume that a single new timeslice is provided to then be used inside step function as geo_t and geo_t_plus_dt. This then means the phibdot=0 which is not accurate in general.
Maybe a TODO can be added here. There needs to be a way to treat this. There are several options that can be discussed.
2 participants
Muscle3 actor for torax
Added:
side: