roms-tools 1.2.0__py3-none-any.whl → 1.4.0__py3-none-any.whl

roms_tools/setup/mixins.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 from roms_tools.setup.grid import Grid
-from roms_tools.setup.fill import fill_and_interpolate
+from roms_tools.setup.fill import LateralFill
 from roms_tools.setup.utils import (
     extrapolate_deepest_to_bottom,
     interpolate_from_rho_to_u,
@@ -106,18 +106,25 @@ class ROMSToolsMixins:
         # interpolate onto desired grid
         data_vars = {}
 
+        # Set up solver that does lateral fill
+        lateral_fill = LateralFill(
+            data.ds["mask"],
+            [data.dim_names["latitude"], data.dim_names["longitude"]],
+        )
+
         # 2d interpolation
-        fill_dims = [data.dim_names["latitude"], data.dim_names["longitude"]]
         coords = {data.dim_names["latitude"]: lat, data.dim_names["longitude"]: lon}
         for var in vars_2d:
-            mask = xr.where(data.ds[data.var_names[var]].isel(time=0).isnull(), 0, 1)
-
-            data_vars[var] = fill_and_interpolate(
-                data.ds[data.var_names[var]].astype(np.float64),
-                mask,
-                fill_dims=fill_dims,
-                coords=coords,
-                method="linear",
+            # Propagate ocean values into land via lateral fill
+            data.ds[data.var_names[var]] = lateral_fill.apply(
+                data.ds[data.var_names[var]].astype(np.float64)
+            )
+
+            # Regrid
+            data_vars[var] = (
+                data.ds[data.var_names[var]]
+                .interp(coords, method="linear")
+                .drop_vars(list(coords.keys()))
             )
 
         if vars_3d:
@@ -127,25 +134,26 @@ class ROMSToolsMixins:
                 data.dim_names["latitude"]: lat,
                 data.dim_names["longitude"]: lon,
             }
-        # extrapolate deepest value all the way to bottom ("flooding")
+        # extrapolate deepest value all the way to bottom
         for var in vars_3d:
             data.ds[data.var_names[var]] = extrapolate_deepest_to_bottom(
                 data.ds[data.var_names[var]], data.dim_names["depth"]
             )
-            mask = xr.where(data.ds[data.var_names[var]].isel(time=0).isnull(), 0, 1)
+            # Propagate ocean values into land via lateral fill
+            data.ds[data.var_names[var]] = lateral_fill.apply(
+                data.ds[data.var_names[var]].astype(np.float64)
+            )
 
-            # setting fillvalue_interp to None means that we allow extrapolation in the
+            # Regrid
+            # setting the fill value to None means that we allow extrapolation in the
             # interpolation step to avoid NaNs at the surface if the lowest depth in original
             # data is greater than zero
-
-            data_vars[var] = fill_and_interpolate(
-                data.ds[data.var_names[var]].astype(np.float64),
-                mask,
-                fill_dims=fill_dims,
-                coords=coords,
-                method="linear",
-                fillvalue_interp=None,
+            data_vars[var] = (
+                data.ds[data.var_names[var]]
+                .interp(coords, method="linear", kwargs={"fill_value": None})
+                .drop_vars(list(coords.keys()))
             )
+
             if data.dim_names["time"] != "time":
                 data_vars[var] = data_vars[var].rename({data.dim_names["time"]: "time"})
 
@@ -207,13 +215,6 @@ class ROMSToolsMixins:
             data_vars[vname] = v_rot
 
         if "s_rho" in data_vars[uname].dims and "s_rho" in data_vars[vname].dims:
-            # 3D masks for ROMS domain
-            umask = self.grid.ds.mask_u.expand_dims({"s_rho": data_vars[uname].s_rho})
-            vmask = self.grid.ds.mask_v.expand_dims({"s_rho": data_vars[vname].s_rho})
-
-            data_vars[uname] = data_vars[uname] * umask
-            data_vars[vname] = data_vars[vname] * vmask
-
             # Compute barotropic velocity
             dz = -self.grid.ds["interface_depth_rho"].diff(dim="s_w")
             dz = dz.rename({"s_w": "s_rho"})
@@ -228,169 +229,3 @@ class ROMSToolsMixins:
             ).transpose("time", "eta_v", "xi_rho")
 
         return data_vars
-
-    def get_variable_metadata(self):
-        """
-        Retrieves metadata for commonly used variables in the dataset.
-
-        This method returns a dictionary containing the metadata for various variables, including long names
-        and units for each variable.
-
-        Returns
-        -------
-        dict of str: dict
-            Dictionary where keys are variable names and values are dictionaries with "long_name" and "units" keys.
-
-        """
-
-        d = {
-            "ssh_Re": {"long_name": "Tidal elevation, real part", "units": "m"},
-            "ssh_Im": {"long_name": "Tidal elevation, complex part", "units": "m"},
-            "pot_Re": {"long_name": "Tidal potential, real part", "units": "m"},
-            "pot_Im": {"long_name": "Tidal potential, complex part", "units": "m"},
-            "u_Re": {
-                "long_name": "Tidal velocity in x-direction, real part",
-                "units": "m/s",
-            },
-            "u_Im": {
-                "long_name": "Tidal velocity in x-direction, complex part",
-                "units": "m/s",
-            },
-            "v_Re": {
-                "long_name": "Tidal velocity in y-direction, real part",
-                "units": "m/s",
-            },
-            "v_Im": {
-                "long_name": "Tidal velocity in y-direction, complex part",
-                "units": "m/s",
-            },
-            "uwnd": {"long_name": "10 meter wind in x-direction", "units": "m/s"},
-            "vwnd": {"long_name": "10 meter wind in y-direction", "units": "m/s"},
-            "swrad": {
-                "long_name": "downward short-wave (solar) radiation",
-                "units": "W/m^2",
-            },
-            "lwrad": {
-                "long_name": "downward long-wave (thermal) radiation",
-                "units": "W/m^2",
-            },
-            "Tair": {"long_name": "air temperature at 2m", "units": "degrees Celsius"},
-            "qair": {"long_name": "absolute humidity at 2m", "units": "kg/kg"},
-            "rain": {"long_name": "total precipitation", "units": "cm/day"},
-            "temp": {"long_name": "potential temperature", "units": "degrees Celsius"},
-            "salt": {"long_name": "salinity", "units": "PSU"},
-            "zeta": {"long_name": "sea surface height", "units": "m"},
-            "u": {"long_name": "u-flux component", "units": "m/s"},
-            "v": {"long_name": "v-flux component", "units": "m/s"},
-            "w": {"long_name": "w-flux component", "units": "m/s"},
-            "ubar": {
-                "long_name": "vertically integrated u-flux component",
-                "units": "m/s",
-            },
-            "vbar": {
-                "long_name": "vertically integrated v-flux component",
-                "units": "m/s",
-            },
-            "PO4": {"long_name": "dissolved inorganic phosphate", "units": "mmol/m^3"},
-            "NO3": {"long_name": "dissolved inorganic nitrate", "units": "mmol/m^3"},
-            "SiO3": {"long_name": "dissolved inorganic silicate", "units": "mmol/m^3"},
-            "NH4": {"long_name": "dissolved ammonia", "units": "mmol/m^3"},
-            "Fe": {"long_name": "dissolved inorganic iron", "units": "mmol/m^3"},
-            "Lig": {"long_name": "iron binding ligand", "units": "mmol/m^3"},
-            "O2": {"long_name": "dissolved oxygen", "units": "mmol/m^3"},
-            "DIC": {"long_name": "dissolved inorganic carbon", "units": "mmol/m^3"},
-            "DIC_ALT_CO2": {
-                "long_name": "dissolved inorganic carbon, alternative CO2",
-                "units": "mmol/m^3",
-            },
-            "ALK": {"long_name": "alkalinity", "units": "meq/m^3"},
-            "ALK_ALT_CO2": {
-                "long_name": "alkalinity, alternative CO2",
-                "units": "meq/m^3",
-            },
-            "DOC": {"long_name": "dissolved organic carbon", "units": "mmol/m^3"},
-            "DON": {"long_name": "dissolved organic nitrogen", "units": "mmol/m^3"},
-            "DOP": {"long_name": "dissolved organic phosphorus", "units": "mmol/m^3"},
-            "DOCr": {
-                "long_name": "refractory dissolved organic carbon",
-                "units": "mmol/m^3",
-            },
-            "DONr": {
-                "long_name": "refractory dissolved organic nitrogen",
-                "units": "mmol/m^3",
-            },
-            "DOPr": {
-                "long_name": "refractory dissolved organic phosphorus",
-                "units": "mmol/m^3",
-            },
-            "zooC": {"long_name": "zooplankton carbon", "units": "mmol/m^3"},
-            "spChl": {
-                "long_name": "small phytoplankton chlorophyll",
-                "units": "mg/m^3",
-            },
-            "spC": {"long_name": "small phytoplankton carbon", "units": "mmol/m^3"},
-            "spP": {
-                "long_name": "small phytoplankton phosphorous",
-                "units": "mmol/m^3",
-            },
-            "spFe": {"long_name": "small phytoplankton iron", "units": "mmol/m^3"},
-            "spCaCO3": {"long_name": "small phytoplankton CaCO3", "units": "mmol/m^3"},
-            "diatChl": {"long_name": "diatom chloropyll", "units": "mg/m^3"},
-            "diatC": {"long_name": "diatom carbon", "units": "mmol/m^3"},
-            "diatP": {"long_name": "diatom phosphorus", "units": "mmol/m^3"},
-            "diatFe": {"long_name": "diatom iron", "units": "mmol/m^3"},
-            "diatSi": {"long_name": "diatom silicate", "units": "mmol/m^3"},
-            "diazChl": {"long_name": "diazotroph chloropyll", "units": "mg/m^3"},
-            "diazC": {"long_name": "diazotroph carbon", "units": "mmol/m^3"},
-            "diazP": {"long_name": "diazotroph phosphorus", "units": "mmol/m^3"},
-            "diazFe": {"long_name": "diazotroph iron", "units": "mmol/m^3"},
-            "pco2_air": {"long_name": "atmospheric pCO2", "units": "ppmv"},
-            "pco2_air_alt": {
-                "long_name": "atmospheric pCO2, alternative CO2",
-                "units": "ppmv",
-            },
-            "iron": {"long_name": "iron decomposition", "units": "nmol/cm^2/s"},
-            "dust": {"long_name": "dust decomposition", "units": "kg/m^2/s"},
-            "nox": {"long_name": "NOx decomposition", "units": "kg/m^2/s"},
-            "nhy": {"long_name": "NHy decomposition", "units": "kg/m^2/s"},
-        }
-        return d
-
-    def get_boundary_info(self):
-
-        """
-        This method provides information about the boundary points for the rho, u, and v
-        variables on the grid, specifying the indices for the south, east, north, and west
-        boundaries.
-
-        Returns
-        -------
-        dict
-            A dictionary where keys are variable types ("rho", "u", "v"), and values
-            are nested dictionaries mapping directions ("south", "east", "north", "west")
-            to the corresponding boundary coordinates.
-        """
-
-        # Boundary coordinates
-        bdry_coords = {
-            "rho": {
-                "south": {"eta_rho": 0},
-                "east": {"xi_rho": -1},
-                "north": {"eta_rho": -1},
-                "west": {"xi_rho": 0},
-            },
-            "u": {
-                "south": {"eta_rho": 0},
-                "east": {"xi_u": -1},
-                "north": {"eta_rho": -1},
-                "west": {"xi_u": 0},
-            },
-            "v": {
-                "south": {"eta_v": 0},
-                "east": {"xi_rho": -1},
-                "north": {"eta_v": -1},
-                "west": {"xi_rho": 0},
-            },
-        }
-
-        return bdry_coords

roms_tools/setup/plot.py

@@ -71,7 +71,14 @@ def _plot(
 
     if field is not None:
         p = ax.pcolormesh(lon_deg, lat_deg, field, transform=proj, **kwargs)
-        plt.colorbar(p, label=f"{field.long_name} [{field.units}]")
+        if hasattr(field, "long_name"):
+            label = f"{field.long_name} [{field.units}]"
+        elif hasattr(field, "Long_name"):
+            # this is the case for matlab generated grids
+            label = f"{field.Long_name} [{field.units}]"
+        else:
+            label = ""
+        plt.colorbar(p, label=label)
 
     if depth_contours:
         cs = ax.contour(lon_deg, lat_deg, field.layer_depth, transform=proj, colors="k")