Loading IMAS geometry from equilibrium ids object is incompatible with saving config to output

[theo-brown]

To reproduce:

import copy

from torax._src import path_utils
from torax._src.imas_tools.input import loader
from torax.tests.test_data import test_iterhybrid_predictor_corrector

# Load IDSs
path = (
    path_utils.torax_path()
    / "data"
    / "imas_data"
    / "ITERhybrid_COCOS17_IDS_ddv4.nc"
)
equilibrium_ids = loader.load_imas_data(str(path), "equilibrium")

CONFIG = copy.deepcopy(test_iterhybrid_predictor_corrector.CONFIG)

CONFIG["geometry"] = {
    "geometry_type": "imas",
    "equilibrium_object": equilibrium_ids,
    "imas_filepath": None,
    "Ip_from_parameters": True,
}

File ".../torax/torax/_src/orchestration/run_simulation.py", line 147, in run_simulation
    state_history.simulation_output_to_xr(),
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^
  File .../torax/torax/_src/output_tools/output.py", line 428, in simulation_output_to_xr
    attrs={CONFIG: self.torax_config.model_dump_json()},
                   ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^
  File ".../torax/venv/lib/python3.13/site-packages/pydantic/main.py", line 528, in model_dump_json
    return self.__pydantic_serializer__.to_json(
           ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^
        self,
        ^^^^^
    ...<13 lines>...
        serialize_as_any=serialize_as_any,
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    ).decode()
    ^
pydantic_core._pydantic_core.PydanticSerializationError: Unable to serialize unknown type: <class 'imas.ids_toplevel.IDSToplevel'>

[Nush395]

hey Theo, unfortuntately this config option is not currently compatible with serialisation, see https://source.corp.google.com/piper///depot/google3/third_party/py/torax/_src/geometry/pydantic_model.py;l=295.

The issue with that is that we haven't defined serialisation logic for an IDS, this will likely be revamped in the nearish future when we plan to look again at the geometry configs. For now are you able to use a custom run loop or a different config option?

Also if there is a clearer way to expose the info in the docstring lmk!

[theo-brown]

@Nush395 makes sense! It would be handy if things didn't fall over completely when a config with this options is run with the main run loop. Two suggestions for short-term fixes:

• Raise an exception if trying to save simulation outputs to xarray with equilibrium IDS in the config, and then handle the exception gracefully.
• Add a serialization branch where the equilibrium IDS' ids_properties tree is saved as a dict

It would also be nice to be warned from the simulation main app if you're loading a config that has an equilibrium IDS object, so that you can cancel the run before wasting time :)

Happy to implement any of these, lmk which you think is best.

[Nush395]

Hmm, I would have said for a short term fix we could raise an error if trying to use the run_simulation loop with an invalid config (in this case populating the equilibrium IDS is not consistent with that workflow.)

Add a serialization branch where the equilibrium IDS' ids_properties tree is saved as a dict

but this is a better long term fix for us to think about. My concern/question mark is we would need to define serialisation/deserialisation logic for the IDSTopLevel object which does look to be defined already [here] (but requires access to imas_core which hasn't been open sourced yet). Are you suggesting rolling our own serialisation/deserialisation logic for this object?

[theo-brown]

It was sort of a halfway house suggestion - basically to just save the provenance information for the IDS, rather than trying to save the full data. The reason being that it is useful to still generate an output file for these cases, even if it's not sufficient to fully reproduce the simulation.

[Nush395]

Ah I see, so like a dummy node saved for that particular field? It's a bit not ideal as it introduces an edge case to our current workflow of being able to load a config back from the stored config which is used by quite a few people but there's also an edge case in running run_simulation atm so maybe it's trading one for the other.

Just to check is it desirable to run the standard run_simulation loop with the equilibrium_IDS field populated? (rather than the imas_filepath or imas_uri fields). If so I think adding some serialisation logic that ignores this field if populated sounds reasonable and we can have a TODO to populate properly once imas_core is open sourced.

[theo-brown]

I'm not super committed either way tbh. Given my use case is not super urgent (it was motivated by wanting to tidy up https://github.com/google-deepmind/torax/blob/main/torax/examples/step_flattop_bgb.py) I'm happy to just add a warning to logging and a todo for future.

[Nush395]

Adding a warning to logging definitely sounds like a positive improvement, are you happy to make that change? (Happy to do so if you have other bits to do!)