diff --git a/pleque/core/equilibrium.py b/pleque/core/equilibrium.py index ff2e0fe..2252efc 100644 --- a/pleque/core/equilibrium.py +++ b/pleque/core/equilibrium.py @@ -157,7 +157,9 @@ def __init__(self, self.time = -1 if 'time_unit' in basedata: - self.time_unit = basedata['time_unit'] + # 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" diff --git a/pleque/io/compass.py b/pleque/io/compass.py index 2966efc..86fe94f 100644 --- a/pleque/io/compass.py +++ b/pleque/io/compass.py @@ -168,6 +168,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 73d6fdc..c9ddf7a 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 = '{}/g0000.0000'.format(tmp_dir.name)