diff --git a/docs/output.rst b/docs/output.rst index fba4f5269..6be3eb450 100644 --- a/docs/output.rst +++ b/docs/output.rst @@ -107,6 +107,13 @@ only has references to the child datasets. as a dictionary. We use this field to store the input config as a json string under the ``config`` key. +For a stitched restart, ``config`` contains the configuration of the resumed +run. The ``previous_config`` attribute contains the configuration saved in the +file from which it was restarted, also as a JSON string. On repeated restarts, +this is the immediately preceding configuration, not the full configuration +history. If the previous file has no ``config`` attribute, ``previous_config`` +is omitted. Unstitched outputs contain only the current ``config``. + To retrieve the input config, see :ref:`output_examples` below. Child datatrees diff --git a/torax/_src/output_tools/output.py b/torax/_src/output_tools/output.py index bbaee7e4b..020fb7aa9 100644 --- a/torax/_src/output_tools/output.py +++ b/torax/_src/output_tools/output.py @@ -117,7 +117,8 @@ def stitch_state_files( datatree: The xr.DataTree to stitch to the end of the previous state file. Returns: - A xr.DataTree containing the stitched dataset. + A xr.DataTree containing the stitched dataset, with the current config and + the previous file's config (if present) saved as root attributes. """ previous_datatree = load_state_file(file_restart.filename) # pyrefly: ignore[bad-argument-type] np.testing.assert_array_equal( @@ -133,7 +134,17 @@ def stitch_state_files( # we are uncertain if file_restart.time is the exact time of the # first time step in this sim output (it takes the nearest time). previous_datatree = previous_datatree.sel(time=slice(None, datatree.time[0])) - return concat_datatrees(previous_datatree, datatree) + stitched = concat_datatrees(previous_datatree, datatree) + # Concatenation keeps the first dataset's attrs. Save the resumed config and + # retain the immediately preceding config separately. + stitched.attrs[output_keys.CONFIG] = datatree.attrs[output_keys.CONFIG] + if output_keys.CONFIG in previous_datatree.attrs: + stitched.attrs[output_keys.PREVIOUS_CONFIG] = previous_datatree.attrs[ + output_keys.CONFIG + ] + else: + stitched.attrs.pop(output_keys.PREVIOUS_CONFIG, None) + return stitched class StateHistory: @@ -403,6 +414,8 @@ def simulation_output_to_xr(self) -> xr.DataTree: - rho_face_norm: The normalized toroidal coordinate on the face grid. - rho_cell_norm: The normalized toroidal coordinate on the cell grid. - config: The ToraxConfig used to run the simulation serialized to JSON. + - previous_config: For stitched restarts, the config saved in the + previous file, if available, serialized to JSON. The child datasets contain the following variables: - numerics: Contains data variables for numeric quantities to do with the simulation. diff --git a/torax/_src/output_tools/output_keys.py b/torax/_src/output_tools/output_keys.py index c5352f04b..f0d7640fd 100644 --- a/torax/_src/output_tools/output_keys.py +++ b/torax/_src/output_tools/output_keys.py @@ -1090,6 +1090,11 @@ def T_fast_ion_key(source_key: str) -> str: # pylint: disable=invalid-name CONFIG: Final[OutputKey] = OutputKey( "config", units=Units.NOT_APPLICABLE, grid_type=GridType.NOT_APPLICABLE ) +PREVIOUS_CONFIG: Final[OutputKey] = OutputKey( + "previous_config", + units=Units.NOT_APPLICABLE, + grid_type=GridType.NOT_APPLICABLE, +) # --------------------------------------------------------------------------- # Geometry: scalar quantities. diff --git a/torax/_src/output_tools/tests/output_stitch_test.py b/torax/_src/output_tools/tests/output_stitch_test.py new file mode 100644 index 000000000..49e60f604 --- /dev/null +++ b/torax/_src/output_tools/tests/output_stitch_test.py @@ -0,0 +1,159 @@ +# Copyright 2024 DeepMind Technologies Limited +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os + +from absl.testing import absltest +from absl.testing import parameterized +import numpy as np +from torax._src.output_tools import output +from torax._src.output_tools import output_keys +from torax._src.torax_pydantic import file_restart +import xarray as xr + + +class StitchStateFilesTest(parameterized.TestCase): + """Tests metadata and overlap handling when stitching restart output.""" + + @staticmethod + def _tree( + times: list[float], values: list[int], attrs: dict[str, str] + ) -> xr.DataTree: + root = xr.Dataset( + coords={ + output_keys.TIME: times, + output_keys.RHO_CELL_NORM: ( + output_keys.RHO_CELL_NORM, + [0.25, 0.75], + ), + }, + attrs=attrs, + ) + profiles = xr.Dataset( + {"value": ((output_keys.TIME,), values)}, + coords={output_keys.TIME: times}, + ) + return xr.DataTree( + dataset=root, + children={output_keys.PROFILES: xr.DataTree(dataset=profiles)}, + ) + + @parameterized.named_parameters( + dict( + testcase_name="normal_prior", + prior_attrs={output_keys.CONFIG: '{"name": "A"}'}, + expected_previous_config='{"name": "A"}', + ), + dict( + testcase_name="legacy_prior", + prior_attrs={}, + expected_previous_config=None, + ), + dict( + testcase_name="prior_has_older_previous_config", + prior_attrs={ + output_keys.CONFIG: '{"name": "B"}', + "previous_config": '{"name": "A"}', + }, + expected_previous_config='{"name": "B"}', + ), + dict( + testcase_name="legacy_prior_has_stale_previous_config", + prior_attrs={"previous_config": '{"name": "A"}'}, + expected_previous_config=None, + ), + ) + def test_stitch_preserves_immediate_previous_config( + self, + prior_attrs: dict[str, str], + expected_previous_config: str | None, + ): + prior = self._tree([0.0, 1.0, 2.0, 3.0], [10, 11, 12, 13], prior_attrs) + current_attrs = {output_keys.CONFIG: '{"name": "C"}'} + current = self._tree([2.0, 3.0, 4.0], [20, 21, 22], current_attrs) + prior_before = prior.copy(deep=True) + current_before = current.copy(deep=True) + prior_path = os.path.join(self.create_tempdir().full_path, "prior.nc") + prior.to_netcdf(prior_path) + + result = output.stitch_state_files( + file_restart.FileRestartConfig( + filename=prior_path, time=2.0, do_restart=True, stitch=True + ), + current, + ) + + self.assertEqual(result.attrs[output_keys.CONFIG], current_attrs["config"]) + if expected_previous_config is None: + self.assertNotIn("previous_config", result.attrs) + else: + self.assertEqual( + result.attrs["previous_config"], expected_previous_config + ) + np.testing.assert_array_equal( + result.children[output_keys.PROFILES].dataset["value"].values, + [10, 11, 12, 21, 22], + ) + np.testing.assert_array_equal(result.time.values, [0, 1, 2, 3, 4]) + xr.testing.assert_identical( + output.load_state_file(prior_path), prior_before + ) + xr.testing.assert_identical(current, current_before) + + result_path = os.path.join(self.create_tempdir().full_path, "result.nc") + result.to_netcdf(result_path) + xr.testing.assert_identical(output.load_state_file(result_path), result) + + def test_stitch_chain_keeps_only_immediate_previous_config(self): + tempdir = self.create_tempdir().full_path + prior_path = os.path.join(tempdir, "prior.nc") + first_path = os.path.join(tempdir, "first.nc") + self._tree( + [0.0, 1.0, 2.0, 3.0], + [10, 11, 12, 13], + {output_keys.CONFIG: '{"name": "A"}'}, + ).to_netcdf(prior_path) + + first = output.stitch_state_files( + file_restart.FileRestartConfig( + filename=prior_path, time=2.0, do_restart=True, stitch=True + ), + self._tree( + [2.0, 3.0, 4.0], + [20, 21, 22], + {output_keys.CONFIG: '{"name": "B"}'}, + ), + ) + first.to_netcdf(first_path) + result = output.stitch_state_files( + file_restart.FileRestartConfig( + filename=first_path, time=4.0, do_restart=True, stitch=True + ), + self._tree( + [4.0, 5.0], + [30, 31], + {output_keys.CONFIG: '{"name": "C"}'}, + ), + ) + + self.assertEqual(result.attrs[output_keys.CONFIG], '{"name": "C"}') + self.assertEqual(result.attrs["previous_config"], '{"name": "B"}') + np.testing.assert_array_equal( + result.children[output_keys.PROFILES].dataset["value"].values, + [10, 11, 12, 21, 22, 31], + ) + + +if __name__ == "__main__": + absltest.main()