diff --git a/pleque/core/equilibrium.py b/pleque/core/equilibrium.py index 4e7b774..16d7235 100644 --- a/pleque/core/equilibrium.py +++ b/pleque/core/equilibrium.py @@ -195,6 +195,12 @@ def __init__(self, logger.debug('Recognizing equilibrium type') self._setup_plasma_boundary() + if 'time_unit' in basedata: + # basedata may be an xr.Dataset, so normalise to a plain str + # (e.g. from a 0-d DataArray) rather than storing an array. + self.time_unit = str(np.asarray(basedata['time_unit']).item()) + else: + self.time_unit = "ms" logger.debug('Generating 1D splines') self._setup_1d_profiles(psi_n, F, FFprime, pressure, pprime) diff --git a/pleque/io/compass.py b/pleque/io/compass.py index 0aae005..bae1dab 100644 --- a/pleque/io/compass.py +++ b/pleque/io/compass.py @@ -167,6 +167,9 @@ def get_ds_from_cudb(shot, time=None, revision=-1, variant='', time_unit='s', fi dst['F0'] = (dst.Bvac * dst.R).mean(dim=['R', 'Z']) dst['shot'] = int(shot) + # CUDB time axis is in seconds; record it so the Equilibrium does not fall + # back to the default 'ms' (which would mis-encode the g-file header time). + dst['time_unit'] = 's' if first_wall is None: try: diff --git a/tests/test_gfile.py b/tests/test_gfile.py index 4dd19ac..c5f3774 100644 --- a/tests/test_gfile.py +++ b/tests/test_gfile.py @@ -54,6 +54,38 @@ def test_from_to_gfile(equilibrium): assert np.isclose(equilibrium.magnetic_axis.Z, eq2.magnetic_axis.Z, atol=1e-5, rtol=1e-4) +def test_time_unit_from_dataset_seconds(): + """ + Regression test: a dataset whose ``time`` axis is in seconds (as produced by + the CUDB / Fiesta reader) must yield ``time_unit == 's'`` so the GEQDSK header + time is encoded correctly. Previously ``time_unit`` defaulted to ``'ms'`` and + ``int(time)`` seconds were written into the header (e.g. 0 ms for t < 1 s). + """ + resource_package = "pleque.resources" + nc_file = pkg_resources.resource_filename(resource_package, "test00.nc") + + basedata = xr.load_dataset(nc_file).load() + # test00.nc carries time in seconds but no explicit time_unit. + basedata["time_unit"] = "s" + + eq = pleque.Equilibrium(basedata, cocos=13) + + # time_unit must be a plain string, not a 0-d DataArray + assert eq.time_unit == "s" + assert isinstance(eq.time_unit, str) + + # The GEQDSK header time is int(time * 1000) ms when time_unit == 's'. + tmp_dir = tempfile.TemporaryDirectory() + file_name = "{}/g0000.0000".format(tmp_dir.name) + eq.to_geqdsk(file_name) + with open(file_name, "r") as f: + header = f.readline() + os.remove(file_name) + + expected_ms = int(float(np.asarray(basedata["time"])) * 1000) + assert "{:d}ms".format(expected_ms) in header + + def test_fiesta_gfile_vs_database(): tmp_dir = tempfile.TemporaryDirectory() file_name = f'{tmp_dir.name}/g0000.0000'