update era5 soil

Author: larsbuntemeyer

pyremo/preproc/era5.py

@@ -109,6 +109,7 @@ def get_file_from_template(
     code,
     level_type,
     template=None,
+    **kwargs,
 ):
     """Derive filename from filename template
 
@@ -181,18 +182,22 @@ class ERA5:
     dynamic = ["ta", "hus", "ps", "tos", "sic", "clw", "snd"]
     wind = ["svo", "sd"]
     fx = ["orog", "sftlf"]
-    soil = [
+    soil_vars = [
         "tsl1",
         "tsl2",
         "tsl3",
         "tsl4",
-        "tsn",
-        "src",
-        "skt",
         "swvl1",
         "swvl2",
         "swvl3",
         "swvl4",
+        "tsn",
+        "src",
+        "skt",
+        "tos",
+        "sic",
+        "skt",
+        "snd",
     ]
 
     chunks = {}
@@ -213,32 +218,28 @@ def __init__(
         if params is None:
             params = era5_params
         self.params = params
+        if gridfile is None:
+            gridfile = era5_grid_file
         self.gridfile = gridfile
         if template is None:
             template = dkrz_template
         self.template = template
         self.cdo = Cdo(tempdir=scratch)
 
-    def _get_files(self, date):
+    def _get_files(self, date, variables=None):
+        if variables is None:
+            variables = self.dynamic + self.fx + self.wind
         if self.cat:
             return get_files_from_intake(
                 self.cat,
-                {
-                    k: v
-                    for k, v in self.params.items()
-                    if k in self.dynamic + self.fx + self.wind
-                },
+                {k: v for k, v in self.params.items() if k in variables},
                 date,
             )
         else:
             return get_files_from_template(
                 date=date,
                 template=self.template,
-                params={
-                    k: v
-                    for k, v in self.params.items()
-                    if k in self.dynamic + self.fx + self.wind
-                },
+                params={k: v for k, v in self.params.items() if k in variables},
             )
 
     def _seldate(self, filename, date):
@@ -259,7 +260,7 @@ def _to_regulars(self, filenames, gridtypes):
         return {
             v: self._to_regular(f, gridtype=gridtypes[v], setname=v)
             for v, f in filenames.items()
-            if v in self.dynamic + self.fx
+            if v in self.dynamic + self.fx + self.soil_vars
         }
 
     def _get_gridtypes(self, filenames):
@@ -384,6 +385,26 @@ def gfile(self, date, path=None, expid=None, filename=None):
 
         return filename
 
+    def get_soil(self, date, path=None, expid=None, filename=None):
+        files = self._get_files(date, self.soil_vars + self.fx)
+        gridtypes = self._get_gridtypes(files)
+        seldates = self._seldates(files, date)
+        regulars = self._to_regulars(seldates, gridtypes)
+        merge = (
+            f"-setgrid,{self.gridfile} -merge [ "
+            + " ".join(list(regulars.values()))
+            + " ]"
+        )
+        call = f"cdo {self.options} invertlev -invertlat {merge} {filename}"
+        print(f"execute: {call}")
+        subprocess.run(
+            call.split(),
+            check=True,
+            shell=False,
+        )
+
+        return filename
+
 
 def era5_from_gcloud(time=None, chunks="auto", hfreq=6):
     """

pyremo/preproc/physics.py

@@ -231,18 +231,28 @@ def calculate_zdts(fibem, fibge):
     return (fibem - fibge) * zdtfak
 
 
-def initialize_soil_temperatures(
-    tdge, tswge, tslge, td3ge, td4ge, td5ge, fibem, fibge, iland
+def adapt_soil_temperatures(
+    tdge, tswge, tslge, td3ge, td4ge, td5ge, fibem, fibge, blaem
 ):
+    # return tdge, tswge, tslge, td3ge, td4ge, td5ge, fibem, fibge, blaem
+    iland = (blaem * 10000).astype(int)
     zdts = calculate_zdts(fibem, fibge)
     tslem = xr.where(iland == 0, tswge, tslge - zdts)
+    tslem.name = "TSL"
     tswem = tswge.copy()
+    tswem.name = "TSW"
     tsnem = xr.where(iland == 0, tswge, tslem)
+    tsnem.name = "TSN"
     td3em = xr.where(iland == 0, tswge, tslge - ((tslge - td3ge) * zds3) - zdts)
+    td3em.name = "TD3"
     td4em = xr.where(iland == 0, tswge, td4ge - ((td4ge - td5ge) * zds4) - zdts)
+    td4em.name = "TD4"
     td5em = xr.where(iland == 0, tswge, td5ge - ((td5ge - tdge) * zds5) - zdts)
+    td5em.name = "TD5"
     tdem = td5em.copy()
+    tdem.name = "TD"
     tdclem = td5em.copy()
+    tdclem.name = "TDCL"
     return tslem, tswem, tsnem, td3em, td4em, td5em, tdem, tdclem
 
 
@@ -310,10 +320,8 @@ def bodfld(
     )  # Assuming FAKINF is 1.0 for simplicity
 
     # Initialize temperatures
-    tslem, tswem, tsnem, td3em, td4em, td5em, tdem, tdclem = (
-        initialize_soil_temperatures(
-            tswge, tslge, td3ge, td4ge, td5ge, fibem, fibge, iland
-        )
+    tslem, tswem, tsnem, td3em, td4em, td5em, tdem, tdclem = adapt_soil_temperatures(
+        tswge, tslge, td3ge, td4ge, td5ge, fibem, fibge, iland
     )
 
     # Calculate sea ice and tsi

pyremo/preproc/remap.py

@@ -30,6 +30,12 @@
 
 # from pyremo.core.remo_ds import update_meta_info
 
+# required variables for initial conditions:
+
+# ['T', 'U', 'V', 'PS', 'RF', 'QW', 'QD', 'QDBL', 'PSEH', 'TSW', 'TSI', 'SEAICE', 'DTPB',
+#  'TSL', 'TSN', 'TD3', 'TD4', 'TD5', 'TD', 'TDCL', 'WS', 'WL', 'SN',
+#  'FIB', 'BLA', 'AZ0', 'ALB', 'VGRAT', 'VAROR', 'VLT',
+#  'FOREST', 'FAO', 'WSMX', 'BETA', 'WMINLOK', 'WMAXLOK', 'TSLEH', 'GLAC']
 
 xr.set_options(keep_attrs=True)
 
@@ -668,7 +674,7 @@ def update_soil_temperatures(ds):
     return ds
 
 
-def create_initial_soil(ds, domain_info, surflib):
+def remap_era5_soil(ds, domain_info, surflib):
     """Create initial soil dataset from atmospheric dataset."""
     # rename vertical coordinate of input to avoid conflict with output lev
     ds = ds.copy()
@@ -688,13 +694,16 @@ def create_initial_soil(ds, domain_info, surflib):
     # broadcast 1d global coordinates
     lamgm, phigm = broadcast_coords(ds)
 
+    print("getting matrix")
     # compute remap matrix
     indii, indjj = intersect(lamgm, phigm, lamem, phiem)  # .compute()
 
+    print("remapping")
     # horizontal interpolation
     wsge = interpolate_horizontal(
         wsgm, lamem, phiem, lamgm, phigm, "WS", indii=indii, indjj=indjj
     )
+
     td3ge = interpolate_horizontal(
         ds.tsl1, lamem, phiem, lamgm, phigm, "TD3", indii=indii, indjj=indjj
     )
@@ -707,24 +716,46 @@ def create_initial_soil(ds, domain_info, surflib):
     tdge = interpolate_horizontal(
         ds.tsl4, lamem, phiem, lamgm, phigm, "TD", indii=indii, indjj=indjj
     )
-    snge = interpolate_horizontal(
+    snem = interpolate_horizontal(
         ds.snd, lamem, phiem, lamgm, phigm, "SN", indii=indii, indjj=indjj
     )
-    wlge = interpolate_horizontal(
+    wlem = interpolate_horizontal(
         ds.src, lamem, phiem, lamgm, phigm, "WL", indii=indii, indjj=indjj
     )
+    fibge = interpolate_horizontal(
+        ds.orog, lamem, phiem, lamgm, phigm, "FIB", indii=indii, indjj=indjj
+    )
+    tslge = interpolate_horizontal(
+        ds.skt, lamem, phiem, lamgm, phigm, "TSL", indii=indii, indjj=indjj
+    )
+
     tswge = remap_sst(
         ds.tos,
         lamem,
         phiem,
         lamgm,
         phigm,
-        blagm=xr.where(ds.tos.isel(time=0).isnull(), 1.0, 0.0),
+        blagm=xr.where(
+            ds.tos.isel(time=0).isnull() if "time" in ds else ds.tos.isnull(), 1.0, 0.0
+        ),
         blaem=blaem,
     )
-    soil = xr.merge([wsge, td3ge, td4ge, td5ge, tdge, snge, wlge, tswge])
+
+    wsmxem = surflib.WSMX.squeeze(drop=True)
+    wsem = np.minimum(wsmxem, wsge * wsmxem)
+
+    # Initialize temperatures
+    tslem, tswem, tsnem, td3em, td4em, td5em, tdem, tdclem = (
+        physics.adapt_soil_temperatures(
+            tdge, tswge, tslge, td3ge, td4ge, td5ge, fibem, fibge, blaem
+        )
+    )
+
+    soil = xr.merge(
+        [tslem, tswem, tsnem, td3em, td4em, td5em, tdem, tdclem, wlem, snem, wsem]
+    )
     soil.attrs["history"] = "preprocessing with pyremo = {}".format(pr.__version__)
-    soil.attrs["domain_id"] = domain_info.get("domain_id", "no name")
+    soil.attrs["domain_id"] = domain_info.get("domain_id", "UNKNOWNs")
 
     soil = update_attrs(soil)