Merge pull request #75 from NVlabs/pm/sst-offset

Bugfix: added missing SST offset changes in the data loaders

Author: nbren12

src/cbottle/datasets/amip_sst_loader.py

@@ -102,7 +102,8 @@ class AmipSSTLoader:
     units = "Kelvin"
     path = config.AMIP_MID_MONTH_SST
 
-    def __init__(self, target_grid=None):
+    def __init__(self, target_grid=None, sst_offset: float = 0.0):
+        self.sst_offset = sst_offset
         self.ensure_downloaded()
 
         self.ds = xarray.open_dataset(
@@ -140,7 +141,10 @@ def ensure_downloaded(cls):
 
     async def sel_time(self, times):
         data = self.interp(times)
-        return {("tosbcs", -1): self.regrid(data)}
+        regridded = self.regrid(data)
+        if self.sst_offset != 0.0:
+            regridded = regridded + self.sst_offset
+        return {("tosbcs", -1): regridded}
 
     def interp(self, time: datetime.datetime):
         """Linearly interpolate between the available points"""
@@ -168,7 +172,8 @@ def ensure_downloaded(cls):
         if _is_file(cls.path):
             _download_aimip_sst(Path(cls.path))
 
-    def __init__(self, target_grid=None):
+    def __init__(self, target_grid=None, sst_offset: float = 0.0):
+        self.sst_offset = sst_offset
         self.ensure_downloaded()
 
         self.ds = xarray.open_dataset(

src/cbottle/datasets/dataset_3d.py

@@ -486,6 +486,7 @@ def _get_loaders(
     sst_input: bool = True,
     ibtracs_input: bool = False,
     variable_config: VariableConfig = _default_config,
+    sst_offset: float = 0.0,
 ):
     """Get the appropriate loaders for a given dataset.
 
@@ -494,6 +495,7 @@ def _get_loaders(
         sst_input: Whether to include SST input
         ibtracs_input: Whether to include IBTrACS input
         variable_config: Variable configuration for the dataset
+        sst_offset: Uniform SST offset in Kelvin
 
     Returns:
         List of loaders for the specified dataset
@@ -568,12 +570,13 @@ def _get_loaders(
                 HPX_LEVEL, pixel_order=earth2grid.healpix.PixelOrder.NEST
             )
             if variable_config == VARIABLE_CONFIGS["aimip"]:
-                loaders.append(AImip_SSTLoader(grid))
+                loaders.append(AImip_SSTLoader(grid, sst_offset=sst_offset))
                 print("used aimip sst data")
             else:
                 loaders.append(
                     AmipSSTLoader(
                         grid,
+                        sst_offset=sst_offset,
                     )
                 )
         if ibtracs_input:
@@ -587,11 +590,12 @@ def _get_loaders(
             HPX_LEVEL, pixel_order=earth2grid.healpix.PixelOrder.NEST
         )
         if variable_config == VARIABLE_CONFIGS["aimip"]:
-            loaders = [AImip_SSTLoader(grid)]
+            loaders = [AImip_SSTLoader(grid, sst_offset=sst_offset)]
         else:
             loaders = [
                 AmipSSTLoader(
                     grid,
+                    sst_offset=sst_offset,
                 )
             ]
 
@@ -700,13 +704,15 @@ def get_dataset(
     frame_masker: Optional[FrameMasker] = None,
     variable_config: VariableConfig = _default_config,
     map_style: bool = False,
+    sst_offset: float = 0.0,
 ) -> TimeMergedDataset | TimeMergedMapStyle:
     # Get the appropriate loaders for the dataset
     loaders = _get_loaders(
         dataset,
         sst_input=sst_input,
         ibtracs_input=ibtracs_input,
         variable_config=variable_config,
+        sst_offset=sst_offset,
     )
     times = _get_splits(dataset)[split]
     if times.size == 0: