add time clock to runtime

xsimlab/drivers.py

@@ -209,9 +209,6 @@ def _maybe_transpose(dataset, model, check_dims, batch_dim):
         if xr_var is None:
             continue
 
-        if any([MAIN_CLOCK in d for d in model.cache[var_key]["metadata"]["dims"]]):
-            raise ValueError("Do not pass xs.MAIN_CLOCK into input vars dimensions")
-
         # all valid dimensions in the right order
         dims = [list(d) for d in model.cache[var_key]["metadata"]["dims"]]
         dims += [[dataset.xsimlab.main_clock_dim] + d for d in dims]

xsimlab/tests/test_model.py

@@ -462,17 +462,30 @@ class Foo:
 
         @xs.runtime(args=["main_clock_values", "main_clock_dataarray"])
         def initialize(self, clock_values, clock_array):
-            self.a = clock_values
-            assert all(self.a == [0, 1, 2, 3])
+            self.a = clock_values * 2
+            assert all(self.a == [0, 2, 4, 6])
             self.b = clock_array * 2
             assert clock_array.dims[0] == "clock"
             assert all(clock_array[clock_array.dims[0]].data == [0, 1, 2, 3])
 
+        @xs.runtime(args=["step_delta", "step"])
+        def run_step(self, dt, n):
+            assert self.a[n] == 2 * n
+            self.a[n] += 1
+
     model = xs.Model({"foo": Foo})
     ds_in = xs.create_setup(
-        model=model, clocks={"clock": [0, 1, 2, 3]}, input_vars={}, output_vars={}
+        model=model,
+        clocks={"clock": range(4)},
+        input_vars={},
+        output_vars={"foo__a": None},
     )
-    ds_in.xsimlab.run(model=model)
+    ds_out = ds_in.xsimlab.run(model=model)
+    assert all(ds_out.foo__a.data == [1, 3, 5, 6])
+
+    # TODO: there is still the problem that the first (0) value of the clock is
+    # set to np.nan in output (still works fine in input) Also, getting
+    # time variables as DataArray as output is not working
 
     # test for error when another dim has the same name as xs.MAIN_CLOCK
     @xs.process
@@ -486,27 +499,27 @@ def run_step(self):
             self.a += self.a
 
     model = xs.Model({"foo": DoubleMainClockDim})
-    # with pytest.raises(ValueError,match="")
-    xs.create_setup(
-        model=model,
-        clocks={"clock": [0, 1, 2, 3]},
-        input_vars={},
-        output_vars={},
-    ).xsimlab.run(model)
+    with pytest.raises(ValueError, match=r"Main clock:*"):
+        xs.create_setup(
+            model=model,
+            clocks={"clock": [0, 1, 2, 3]},
+            input_vars={},
+            output_vars={"foo__a": None},
+        ).xsimlab.run(model)
 
     # test for error when trying to put xs.MAIN_CLOCK as a dim in an input var
     @xs.process
     class InputMainClockDim:
-        a = xs.variable(intent="in", dims=xs.MAIN_CLOCK)
-
-    model = xs.Model({"foo": InputMainClockDim})
-    ds_in = xs.create_setup(
-        model=model,
-        clocks={"clock": [0, 1, 2, 3]},
-        input_vars={"foo__a": 5},
-        output_vars={},
-    )
-    with pytest.raises(
-        ValueError, match="Do not pass xs.MAIN_CLOCK into input vars dimensions"
-    ):
-        ds_in.xsimlab.run(model=model)
+        with pytest.raises(
+            ValueError, match="Do not pass xs.MAIN_CLOCK into input vars dimensions"
+        ):
+            a = xs.variable(intent="in", dims=xs.MAIN_CLOCK)
+
+        with pytest.raises(
+            ValueError, match="Do not pass xs.MAIN_CLOCK into input vars dimensions"
+        ):
+            b = xs.variable(intent="in", dims=(xs.MAIN_CLOCK,))
+        with pytest.raises(
+            ValueError, match="Do not pass xs.MAIN_CLOCK into input vars dimensions"
+        ):
+            c = xs.variable(intent="in", dims=["a", ("a", xs.MAIN_CLOCK)])

xsimlab/tests/test_xr_accessor.py

@@ -158,9 +158,9 @@ def test_master_clock_coords_warning(self):
         )
         with pytest.warns(
             FutureWarning,
-            match="master_clock is to be deprecated in favour of main_clock",
+            match="master_clock_coord is to be deprecated in favour of main_clock",
         ):
-            xr.testing.assert_equal(ds.xsimlab.master_clock_coord, ds.mclock)
+            ds.xsimlab.master_clock_coord
 
     def test_clock_sizes(self):
         ds = xr.Dataset(

xsimlab/utils.py

@@ -16,8 +16,12 @@
 
 
 class _MainClockDim:
-    """singleton class to be used as main clock dimension: update on runtime
-    it has all behaviour that dimensions in a `xr.DataArray` normally have.
+    """Singleton class to be used as a placeholder of the main clock
+    dimension.
+
+    It will be replaced by the actual dimension label set during simulation setup
+    (i.e., ``main_clock`` argument).
+
     """
 
     _singleton = None
@@ -29,7 +33,7 @@ def __new__(cls):
         return _MainClockDim._singleton
 
     def __repr__(self):
-        return "MAIN_CLOCK (uninitialized)"
+        return "MAIN_CLOCK (undefined)"
 
 
 MAIN_CLOCK = _MainClockDim()

xsimlab/variable.py

@@ -58,14 +58,17 @@ def _as_dim_tuple(dims):
 
     """
     # MAIN_CLOCK is sentinel and does not have length (or zero), so check explicitly
-    if dims == MAIN_CLOCK:
+    if dims is MAIN_CLOCK:
         dims = [(dims,)]
     elif not len(dims):
         dims = [()]
     elif isinstance(dims, str):
         dims = [(dims,)]
     elif isinstance(dims, list):
-        dims = [tuple([d]) if isinstance(d, str) else tuple(d) for d in dims]
+        dims = [
+            tuple([d]) if (isinstance(d, str) or d is MAIN_CLOCK) else tuple(d)
+            for d in dims
+        ]
     else:
         dims = [dims]
 
@@ -224,6 +227,9 @@ def variable(
     else:
         _init = True
         _repr = True
+        # also check if MAIN_CLOCK is there
+        if any([MAIN_CLOCK in d for d in metadata["dims"]]):
+            raise ValueError("Do not pass xs.MAIN_CLOCK into input vars dimensions")
 
     return attr.attrib(
         metadata=metadata,
@@ -342,7 +348,7 @@ def on_demand(
         Dictionary specifying how to encode this variable's data into a
         serialized format (i.e., as a zarr dataset). Currently used keys
         include 'dtype', 'compressor', 'fill_value', 'order', 'filters'
-        and 'object_codec'. See :func:`zarr.creation.create` for details
+        and 'object_codec'. See :func:`zarr.creation.df` for details
         about these options. Other keys are ignored.
 
     Notes

xsimlab/xr_accessor.py

@@ -212,6 +212,7 @@ def master_clock_coord(self):
             "master_clock_coord is to be deprecated in favour of main_clock",
             FutureWarning,
         )
+        warnings.warn("Poep in je hoofd!")
         return self.main_clock_coord
 
     @property