Allow float conversion in clamp and replace fill values by NaN functions

In order to simplify the API an 'astype' method was also introduced.

Closes #204.

Signed-off-by: Alexis Jeandet <alexis.jeandet@member.fsf.org>

Author: jeandet

speasy/core/data_containers.py

@@ -121,6 +121,10 @@ def ndim(self):
     def dtype(self):
         return self.__values.dtype
 
+    def astype(self, dtype) -> "DataContainer":
+        return DataContainer(values=self.__values.astype(dtype), meta=self.__meta, name=self.__name,
+                             is_time_dependent=self.__is_time_dependent)
+
     @property
     def unit(self) -> str:
         return self.__meta.get('UNITS')

speasy/products/variable.py

@@ -400,6 +400,25 @@ def shape(self):
     def dtype(self):
         return self.__values_container.dtype
 
+    def astype(self, dtype) -> "SpeasyVariable":
+        """Returns a SpeasyVariable with values converted to given dtype
+
+        Parameters
+        ----------
+        dtype : str or np.dtype or type
+            desired dtype
+
+        Returns
+        -------
+        SpeasyVariable
+            SpeasyVariable with values converted to given dtype
+        """
+        return SpeasyVariable(
+            axes=deepcopy(self.__axes),
+            values=self.__values_container.astype(dtype),
+            columns=deepcopy(self.__columns),
+        )
+
     @property
     def name(self) -> str:
         """SpeasyVariable name
@@ -700,14 +719,16 @@ def plot(self, *args, **kwargs):
             values=self.__values_container, columns_names=self.columns, axes=self.axes
         )
 
-    def replace_fillval_by_nan(self, inplace=False) -> "SpeasyVariable":
-        """Replaces fill values by NaN, non float values are automatically converted to float.
+    def replace_fillval_by_nan(self, inplace=False, convert_to_float=False) -> "SpeasyVariable":
+        """Replaces fill values by NaN, non float values are automatically converted to float if convert_to_float is True.
         Fill value is taken from metadata field "FILLVAL"
 
         Parameters
         ----------
         inplace : bool, optional
             Modifies source variable when true else modifies and returns a copy, by default False
+        convert_to_float : bool, optional
+            Automatically converts variable to float if true and needed, by default False.
 
         Returns
         -------
@@ -719,16 +740,25 @@ def replace_fillval_by_nan(self, inplace=False) -> "SpeasyVariable":
         clamp_with_nan: replaces values outside valid range by NaN
         sanitized: removes fill and invalid values
         """
+        # @TODO replace by a match case when Python 3.9 is EOL
         if inplace:
             res = self
+            if convert_to_float and not np.issubdtype(self.dtype, np.floating):
+                res.__values_container = res.__values_container.astype(float)
         else:
-            res = deepcopy(self)
+            if convert_to_float and not np.issubdtype(self.dtype, np.floating):
+                res = self.astype(float)
+            else:
+                res = deepcopy(self)
         if (fill_value := self.fill_value) is not None:
+            if convert_to_float and not np.issubdtype(res.dtype, np.floating):
+                res.__values_container = res.__values_container.astype(float)
             res[res == fill_value] = np.nan
         return res
 
-    def clamp_with_nan(self, inplace=False, valid_min=None, valid_max=None) -> "SpeasyVariable":
-        """Replaces values outside valid range by NaN, valid range is taken from metadata fields "VALIDMIN" and "VALIDMAX"
+    def clamp_with_nan(self, inplace=False, valid_min=None, valid_max=None, convert_to_float=False) -> "SpeasyVariable":
+        """Replaces values outside valid range by NaN, valid range is taken from metadata fields "VALIDMIN" and "VALIDMAX".
+        Automatically converts variable to float if convert_to_float is True and needed.
 
         Parameters
         ----------
@@ -738,6 +768,8 @@ def clamp_with_nan(self, inplace=False, valid_min=None, valid_max=None) -> "Spea
             Optional minimum valid value, takes metadata field "VALIDMIN" if not provided, by default None
         valid_max : Float, optional
             Optional maximum valid value, takes metadata field "VALIDMAX" if not provided, by default None
+        convert_to_float : bool, optional
+            Automatically converts variable to float if true and needed, by default False.
 
         Returns
         -------
@@ -749,10 +781,16 @@ def clamp_with_nan(self, inplace=False, valid_min=None, valid_max=None) -> "Spea
         replace_fillval_by_nan: replaces fill values by NaN
         sanitized: removes fill and invalid values
         """
+        # @TODO replace by a match case when Python 3.9 is EOL
         if inplace:
             res = self
+            if convert_to_float and not np.issubdtype(self.dtype, np.floating):
+                res.__values_container = res.__values_container.astype(float)
         else:
-            res = deepcopy(self)
+            if convert_to_float and not np.issubdtype(self.dtype, np.floating):
+                res = self.astype(float)
+            else:
+                res = deepcopy(self)
         valid_min = valid_min or self.valid_range[0]
         valid_max = valid_max or self.valid_range[1]
         res[np.logical_or(res > valid_max, res < valid_min)] = np.nan

tests/test_speasy_variable.py

@@ -341,6 +341,13 @@ def test_overrides_plot_arguments(self):
         except ImportError:
             self.skipTest("Can't import matplotlib")
 
+    def test_astype(self):
+        var = make_simple_var(1., 10., 1., 10.)
+        self.assertEqual(var.values.dtype, np.float64)
+        var2 = (var * 100).astype(np.int32)
+        self.assertEqual(var2.values.dtype, np.int32)
+        self.assertTrue(np.all(var2.values == (var.values * 100).astype(np.int32)))
+
     def test_replaces_fill_value(self):
         var = make_simple_var(1., 10., 1., 10., meta={"FILLVAL": 50.})
         self.assertEqual(var.fill_value, 50.)
@@ -350,6 +357,17 @@ def test_replaces_fill_value(self):
         var.replace_fillval_by_nan(inplace=True)
         self.assertTrue(np.isnan(var.values[4, 0]))
 
+    def test_replaces_fill_value_non_float(self):
+        var = make_simple_var(1., 10., 1., 10., meta={"FILLVAL": 50.}).astype(np.int32)
+        self.assertEqual(var.fill_value, 50.)
+        with self.assertRaises(ValueError):
+            var.replace_fillval_by_nan(inplace=False)
+        cleaned_copy = var.replace_fillval_by_nan(inplace=False, convert_to_float=True)
+        self.assertTrue(np.isnan(cleaned_copy.values[4, 0]))
+        self.assertFalse(np.isnan(var.values[4, 0]))
+        var.replace_fillval_by_nan(inplace=True, convert_to_float=True)
+        self.assertTrue(np.isnan(var.values[4, 0]))
+
     def test_clamps(self):
         var = make_simple_var(1., 10., 1., 10., meta={"VALIDMIN": 20., "VALIDMAX": 80.})
         clamped_copy = var.clamp_with_nan()
@@ -361,6 +379,19 @@ def test_clamps(self):
         self.assertTrue(np.all(np.isnan(var.values[8:10, 0])))
         self.assertFalse(np.any(np.isnan(var.values[1:8, 0])))
 
+    def test_clamps_non_float(self):
+        var = make_simple_var(1., 10., 1., 10., meta={"VALIDMIN": 20., "VALIDMAX": 80.}).astype(np.int32)
+        with self.assertRaises(ValueError):
+            var.clamp_with_nan(inplace=False)
+        clamped_copy = var.clamp_with_nan(convert_to_float=True)
+        self.assertTrue(np.all(np.isnan(clamped_copy.values[0:1, 0])))
+        self.assertTrue(np.all(np.isnan(clamped_copy.values[8:10, 0])))
+        self.assertFalse(np.any(np.isnan(clamped_copy.values[1:8, 0])))
+        var.clamp_with_nan(inplace=True, convert_to_float=True)
+        self.assertTrue(np.all(np.isnan(var.values[0:1, 0])))
+        self.assertTrue(np.all(np.isnan(var.values[8:10, 0])))
+        self.assertFalse(np.any(np.isnan(var.values[1:8, 0])))
+
     def test_cleans(self):
         var = make_simple_var(1., 10., 1., 10., meta={"FILLVAL": 50., "VALIDMIN": 20., "VALIDMAX": 80.})
         cleaned_copy = var.sanitized()