Completed standardization module and unitttests

o3skim/standardization.py

@@ -28,17 +28,35 @@ vmro3_standard_coordinates = [
 ]
 
 
-def standardize(dataset, tco3_zm=None, vmro3_zm=None):
-    pass  # TODO
+@utils.return_on_failure("Error when loading '{0}'".format(tco3_standard_name),
+                         default=xr.Dataset())
+def standardize_tco3(dataset, variable, coordinates):
+    """Standardizes a tco3 dataset"""
+    dataset = squeeze(dataset)
+    dataset = rename_tco3(dataset, variable, coordinates)
+    dataset = sort(dataset)
+    return dataset
 
 
+@utils.return_on_failure("Error when loading '{0}'".format(vmro3_standard_name),
+                         default=xr.Dataset())
 def standardize_vmro3(dataset, variable, coordinates):
     """Standardizes a vmro3 dataset"""
+    dataset = squeeze(dataset)
     dataset = rename_vmro3(dataset, variable, coordinates)
     dataset = sort(dataset)
     return dataset
 
 
+def rename_tco3(dataset, variable, coordinates):
+    """Renames a tco3 dataset variable and coordinates"""
+    logger.debug("Rename of '{0}' var and coords".format(tco3_standard_name))
+    return dataset.rename({
+        **{variable: tco3_standard_name},
+        **{coordinates[x]: x for x in tco3_standard_coordinates}
+    })
+
+
 def rename_vmro3(dataset, variable, coordinates):
     """Renames a vmro3 dataset variable and coordinates"""
     logger.debug("Rename of '{0}' var and coords".format(vmro3_standard_name))
@@ -48,38 +66,88 @@ def rename_vmro3(dataset, variable, coordinates):
     })
 
 
+def squeeze(dataset):
+    """Squeezes the 1size dimensions on a dataset"""
+    logger.debug("Squeezing coordinates in dataset")
+    return dataset.squeeze(drop=True)
+
+
 def sort(dataset):
     """Sorts a dataset by coordinates"""
     logger.debug("Sorting coordinates in dataset")
     return dataset.sortby(list(dataset.coords))
 
 
-def squeeze(dataset):
-    """Squeezes the 1size dimensions on a dataset"""
-    logger.debug("Squeezing coordinates in dataset")
-    pass  # TODO
+class TestsTCO3(unittest.TestCase):
+
+    data_s = np.mgrid[1:3:3j, 1:3:3j, 1:3:3j][0]
+    data_r = np.mgrid[3:1:3j, 3:1:3j, 1:1:1j, 3:1:3j][0]
+    varname = "tco3"
+    coords = {'lon': 'long', 'lat': 'latd', 'time': 'time'}
+
+    @staticmethod
+    def non_standard_ds():
+        return xr.Dataset(
+            data_vars=dict(
+                tco3=(["long", "latd", "high", "time"], TestsTCO3.data_r)
+            ),
+            coords=dict(
+                long=[180, 0, -180],
+                latd=[90, 0, -90],
+                high=[1],
+                time=pd.date_range("2000-01-03", periods=3, freq='-1d')
+            ),
+            attrs=dict(description="Non standardized dataset")
+        )
+
+    @staticmethod
+    def standard_ds():
+        return xr.Dataset(
+            data_vars=dict(
+                tco3_zm=(["lon", "lat", "time"], TestsTCO3.data_s)
+            ),
+            coords=dict(
+                time=pd.date_range("2000-01-01", periods=3, freq='1d'),
+                lat=[-90, 0, 90],
+                lon=[-180, 0, 180]
+            ),
+            attrs=dict(description="Standardized dataset")
+        )
+
+    def test_standardize(self):
+        standardized_tco3 = standardize_tco3(
+            dataset=TestsTCO3.non_standard_ds(),
+            variable=TestsTCO3.varname,
+            coordinates=TestsTCO3.coords)
+        xr.testing.assert_equal(TestsTCO3.standard_ds(), standardized_tco3)
+
+    def test_fail_returns_empty_dataset(self):
+        empty_dataset = standardize_tco3(
+            dataset=TestsTCO3.non_standard_ds(),
+            variable="badVariable",
+            coordinates=TestsTCO3.coords)
+        xr.testing.assert_equal(xr.Dataset(), empty_dataset)
 
 
 class TestsVMRO3(unittest.TestCase):
 
     data_s = np.mgrid[1:3:3j, 1:3:3j, 1:4:4j, 1:3:3j][0]
     data_r = np.mgrid[3:1:3j, 3:1:3j, 4:1:4j, 3:1:3j][0]
-
     varname = "vmro3"
     coords = {'lon': 'longit', 'lat': 'latitu',
-              'plev': 'level', 'time': 'time'}
+              'plev': 'level', 'time': 't'}
 
     @staticmethod
     def non_standard_ds():
         return xr.Dataset(
             data_vars=dict(
-                vmro3=(["longit", "latitu", "level", "time"], TestsVMRO3.data_r)
+                vmro3=(["longit", "latitu", "level", "t"], TestsVMRO3.data_r)
             ),
             coords=dict(
                 longit=[180, 0, -180],
                 latitu=[90, 0, -90],
                 level=[1000, 100, 10, 1],
-                time=pd.date_range("2000-01-03", periods=3, freq='-1d')
+                t=pd.date_range("2000-01-03", periods=3, freq='-1d')
             ),
             attrs=dict(description="Non standardized dataset")
         )
@@ -105,3 +173,10 @@ class TestsVMRO3(unittest.TestCase):
             variable=TestsVMRO3.varname,
             coordinates=TestsVMRO3.coords)
         xr.testing.assert_equal(TestsVMRO3.standard_ds(), standardized_vmro3)
+
+    def test_fail_returns_empty_dataset(self):
+        empty_dataset = standardize_vmro3(
+            dataset=TestsVMRO3.non_standard_ds(),
+            variable="badVariable",
+            coordinates=TestsVMRO3.coords)
+        xr.testing.assert_equal(xr.Dataset(), empty_dataset)