Solve nan and encoding in dtype (fill value) (#173)

* added a main_clock with dimension 'mclock' to initialize context, this returns a xr.DataArray. However, I cannot set xs.variables in the initialize step: 'AttributeError: cannot set attribute'

* fixed black

* changed 'master' to 'main' everywhere except in non-breaking testcases.

* fixed space and undid changes of access-clock

* added master_clock_dim + tests, master_clock_coords no tests

* Update xsimlab/xr_accessor.py

Co-authored-by: Benoit Bovy <[email protected]>

* Update xsimlab/xr_accessor.py

Co-authored-by: Benoit Bovy <[email protected]>

* Update xsimlab/xr_accessor.py

Co-authored-by: Benoit Bovy <[email protected]>

* added master/main_clock_coord access tests, updated whats-new

* removed vscode settings

* Apply suggestions from code review

Co-authored-by: Benoit Bovy <[email protected]>

* removed raise checks from test_update_clocks_master_clock_warning

* removed raises in test_update_clocks_master_warning

* removed redundancy in test_update_master_clock

* updated fill_value from dtype and support for encoding

* added more datatypes, added warning for boolean

* very black

* more black

* booleans set to False

* removed warning

* removed redundancy

* tests and docs

* added environment.yml

* added tests

* black

* solved environment.yml (#174)

Co-authored-by: Benoit Bovy <[email protected]>

Author: feefladder

doc/develop.rst

@@ -183,7 +183,7 @@ files) in ``xarray-simlab/doc``.
 To build the documentation locally, first install requirements (for example here
 in a separate conda environment)::
 
-   $ conda env create -n xarray-simlab_doc -f doc/environment.yml
+   $ conda env create -n xarray-simlab_doc -f ci/requirements/doc.yml
    $ conda activate xarray-simlab_doc
 
 Then build documentation with ``make``::

doc/io_storage.rst

@@ -149,6 +149,86 @@ computing with Dask`_ in xarray's docs).
 
 Advanced usage
 --------------
+Encoding options
+~~~~~~~~~~~~~~~~
+
+It is possible to control via some encoding options how Zarr stores simulation
+data. Those options can be set for variables declared in process classes. See
+the ``encoding`` parameter of :func:`~xsimlab.variable` for all available
+options. Encoding options may also be set or overridden when calling
+:func:`~xarray.Dataset.xsimlab.run`. A often encountered use-case for encoding, 
+is to suppress ``nan`` values in the output dataset:
+
+Default fill values
+~~~~~~~~~~~~~~~~~~~
+
+By default, output variables have a fill value that is set to `np.nan` in output. 
+These values are determined during model runtime from the variable's datatype.
+This only affects the output array: during the model run, the actual values are used. 
+
+.. list-table:: Fill values
+   :header-rows: 1
+
+   * - datatype
+     - fill values
+     - example
+   * - (unsigned) integer
+     - maximum possible value
+     - uint8, 255
+   * - float
+     - np.nan
+     -
+   * - string
+     - empty string
+     - ""
+   * - bool
+     - False
+     - 
+   * - complex values
+     - default for each dtype
+     - 
+
+Especially for boolean datatypes, where the default fill value is ``False``, it is
+desireable to suppress this behaviour. There are several options, with different 
+benefits and drawbacks, as outlined below.  
+
+
+If you know beforehand what the ``fill_value`` should be, this can be set directly in the process class:
+
+.. literalinclude:: scripts/encoding.py
+   :lines: 5-13
+
+.. ipython:: python
+   :suppress:
+
+   from encoding import Foo
+
+The resulting output is set to nan when no ``fill_value`` is specified:
+
+.. ipython:: python
+
+   model = xs.Model({"foo":Foo})
+   in_ds = xs.create_setup(
+      model=model,
+      clocks={"clock": [0, 1]},
+      output_vars={"foo__v_bool": None, "foo__v_bool_nan": None},
+   )
+   #this will result in nan values in output
+   in_ds.xsimlab.run(model=model)
+
+Alternatively, encoding (or decoding) options can be set during model run:
+
+.. ipython:: python
+
+   # set encoding options during model run
+   in_ds.xsimlab.run(model=model, encoding={"foo__v_bool_nan": {"fill_value": None}})
+
+   # set mask_and_scale to false
+   in_ds.xsimlab.run(model=model, decoding={"mask_and_scale": False})
+
+However, using ``mask_and_scale:False`` results in non-serializeable attributes 
+in the output dataset, so the other alternatives are preferable.
+
 
 Dynamically sized arrays
 ~~~~~~~~~~~~~~~~~~~~~~~~
@@ -206,22 +286,4 @@ to the xarray Dataset or DataArray :meth:`~xarray.Dataset.stack`,
 
 .. _io_storage_encoding:
 
-Encoding options
-~~~~~~~~~~~~~~~~
 
-It is possible to control via some encoding options how Zarr stores simulation
-data. Those options can be set for variables declared in process classes. See
-the ``encoding`` parameter of :func:`~xsimlab.variable` for all available
-options. Encoding options may also be set or overridden when calling
-:func:`~xarray.Dataset.xsimlab.run`.
-
-.. warning::
-
-   Zarr uses ``0`` as the default fill value for numeric value types. This may
-   badly affect the results, as array elements with the fill value are replaced
-   by NA in the output xarray Dataset. For variables which accept ``0`` as a
-   possible (non-missing) value, it is highly recommended to explicitly provide
-   another ``fill_value``. Alternatively, it is possible to deactivate this
-   value masking behavior by setting the ``mask_and_scale=False`` option and
-   pass it via the ``decoding`` parameter of
-   :func:`~xarray.Dataset.xsimlab.run`.

doc/scripts/encoding.py

@@ -0,0 +1,13 @@
+import xsimlab as xs
+import numpy as np
+
+
+@xs.process
+class Foo:
+    v_bool_nan = xs.variable(dims="x", intent="out")
+    # suppress nan values by setting an explicit fill value:
+    v_bool = xs.variable(dims="x", intent="out", encoding={"fill_value": None})
+
+    def initialize(self):
+        self.v_bool_nan = [True, False]
+        self.v_bool = [True, False]

doc/whats_new.rst

@@ -9,6 +9,10 @@ v0.6.0 (Unreleased)
    ``main_clock``, ``main_clock_dim`` and ``main_clock_coords`` and all
    occurences of ``master`` to ``main`` in the rest of the codebase. all 
    ``master...`` API hooks are still working, but raise a Futurewarning
+ - Changed default ``fill_value`` in the zarr stores to maximum dtype value 
+   for integer dtypes and ``np.nan`` for floating-point variables.
+
+
 
 v0.5.0 (26 January 2021)
 ------------------------

xsimlab/stores.py

@@ -1,5 +1,6 @@
 from collections.abc import MutableMapping
 from typing import Any, Dict, Optional, Tuple, Union
+import warnings
 
 import numpy as np
 import xarray as xr
@@ -60,8 +61,14 @@ def ensure_no_dataset_conflict(zgroup, znames):
 def default_fill_value_from_dtype(dtype=None):
     if dtype is None:
         return 0
-    if dtype.kind == "f":
+    elif dtype.kind == "f":
         return np.nan
+    elif dtype.kind == "i":
+        return np.iinfo(dtype).max
+    elif dtype.kind == "u":
+        return np.iinfo(dtype).max
+    elif dtype.kind == "U":
+        return ""
     elif dtype.kind in "c":
         return (
             default_fill_value_from_dtype(dtype.type().real.dtype),
@@ -190,7 +197,11 @@ def _create_zarr_dataset(
         value = model.cache[var_key]["value"]
         clock = var_info["clock"]
 
-        dtype = getattr(value, "dtype", np.asarray(value).dtype)
+        if "dtype" in var_info["metadata"]["encoding"]:
+            dtype = np.dtype(var_info["metadata"]["encoding"]["dtype"])
+        else:
+            dtype = getattr(value, "dtype", np.asarray(value).dtype)
+
         shape = list(np.shape(value))
         chunks = list(get_auto_chunks(shape, dtype))
 

xsimlab/tests/test_stores.py

@@ -286,6 +286,44 @@ def _get_v2(self):
         assert ztest.p__v3.chunks == (10,)
         assert ztest.p__v4[0] == {"foo": "bar"}
 
+    def test_fill_values(self):
+        @xs.process
+        class Foo:
+            v_int64 = xs.variable(dims="x", intent="out")
+            v_float64 = xs.variable(dims="x", intent="out")
+            v_uint8 = xs.variable(dims="x", intent="out", encoding={"dtype": np.uint8})
+            v_string = xs.variable(dims="x", intent="out")
+            v_bool = xs.variable(dims="x", intent="out")
+
+            def initialize(self):
+                self.v_int64 = [0, np.iinfo("int64").max]
+                self.v_float64 = [0.0, np.nan]
+                self.v_uint8 = [0, 255]
+                self.v_string = ["hello", ""]
+                self.v_bool = [True, False]
+
+        model = xs.Model({"foo": Foo})
+        in_ds = xs.create_setup(
+            model=model,
+            clocks={"clock": [0, 1]},
+            output_vars={
+                "foo__v_int64": None,
+                "foo__v_float64": None,
+                "foo__v_uint8": None,
+                "foo__v_string": None,
+                "foo__v_bool": None,
+            },
+        )
+        out_ds = in_ds.xsimlab.run(model=model)
+        np.testing.assert_equal(out_ds["foo__v_int64"].data, [0, np.nan])
+        np.testing.assert_equal(out_ds["foo__v_float64"].data, [0.0, np.nan])
+        np.testing.assert_equal(out_ds["foo__v_uint8"].data, [0, np.nan])
+        # np.testing.assert_equal does not work for "object" dtypes, so test each value explicitly:
+        assert out_ds["foo__v_string"].data[0] == "hello"
+        assert np.isnan(out_ds["foo__v_string"].data[1])
+        assert out_ds["foo__v_bool"].data[0] == True
+        assert np.isnan(out_ds["foo__v_bool"].data[1])
+
     def test_open_as_xr_dataset(self, store):
         model = store.model
         model.state[("profile", "u")] = np.array([1.0, 2.0, 3.0])