roms-tools 2.2.0__py3-none-any.whl → 2.3.0__py3-none-any.whl

roms_tools/setup/initial_conditions.py

@@ -3,10 +3,14 @@ import numpy as np
 import importlib.metadata
 from dataclasses import dataclass, field
 from typing import Dict, Union, List, Optional
-from roms_tools.setup.grid import Grid
+import matplotlib.pyplot as plt
+from pathlib import Path
 from datetime import datetime
+from roms_tools import Grid
+from roms_tools.regrid import LateralRegrid, VerticalRegrid
+from roms_tools.plot import _plot, _section_plot, _profile_plot, _line_plot
+from roms_tools.utils import transpose_dimensions
 from roms_tools.setup.datasets import GLORYSDataset, CESMBGCDataset
-from roms_tools.setup.vertical_coordinate import compute_depth
 from roms_tools.setup.utils import (
     nan_check,
     substitute_nans_by_fillvalue,
@@ -15,16 +19,9 @@ from roms_tools.setup.utils import (
     get_target_coords,
     rotate_velocities,
     compute_barotropic_velocity,
-    transpose_dimensions,
-    interpolate_from_rho_to_u,
-    interpolate_from_rho_to_v,
     _to_yaml,
     _from_yaml,
 )
-from roms_tools.setup.regrid import LateralRegrid, VerticalRegrid
-from roms_tools.setup.plot import _plot, _section_plot, _profile_plot, _line_plot
-import matplotlib.pyplot as plt
-from pathlib import Path
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -396,54 +393,7 @@ class InitialConditions:
             - f"{type}_depth_v": Depth coordinates at v points (if applicable).
         """
 
-        layer_vars = []
-        for location in ["rho"] + additional_locations:
-            layer_vars.append(f"{type}_depth_{location}")
-
-        if all(layer_var in self.grid.ds for layer_var in layer_vars):
-            # Vertical coordinate data already exists
-            pass
-
-        elif f"{type}_depth_rho" in self.grid.ds:
-            depth = self.grid.ds[f"{type}_depth_rho"]
-
-            if "u" in additional_locations or "v" in additional_locations:
-                # interpolation
-                if "u" in additional_locations:
-                    depth_u = interpolate_from_rho_to_u(depth)
-                    depth_u.attrs["long_name"] = f"{type} depth at u-points"
-                    depth_u.attrs["units"] = "m"
-                    self.grid.ds[f"{type}_depth_u"] = depth_u
-                if "v" in additional_locations:
-                    depth_v = interpolate_from_rho_to_v(depth)
-                    depth_v.attrs["long_name"] = f"{type} depth at v-points"
-                    depth_v.attrs["units"] = "m"
-                    self.grid.ds[f"{type}_depth_v"] = depth_v
-        else:
-            h = self.grid.ds["h"]
-            if type == "layer":
-                depth = compute_depth(
-                    0, h, self.grid.hc, self.grid.ds.Cs_r, self.grid.ds.sigma_r
-                )
-            else:
-                depth = compute_depth(
-                    0, h, self.grid.hc, self.grid.ds.Cs_w, self.grid.ds.sigma_w
-                )
-
-            depth.attrs["long_name"] = f"{type} depth at rho-points"
-            depth.attrs["units"] = "m"
-            self.grid.ds[f"{type}_depth_rho"] = depth
-
-            if "u" in additional_locations or "v" in additional_locations:
-                # interpolation
-                depth_u = interpolate_from_rho_to_u(depth)
-                depth_u.attrs["long_name"] = f"{type} depth at u-points"
-                depth_u.attrs["units"] = "m"
-                depth_v = interpolate_from_rho_to_v(depth)
-                depth_v.attrs["long_name"] = f"{type} depth at v-points"
-                depth_v.attrs["units"] = "m"
-                self.grid.ds[f"{type}_depth_u"] = depth_u
-                self.grid.ds[f"{type}_depth_v"] = depth_v
+        self.grid.compute_depth_coordinates(type, additional_locations)
 
     def _write_into_dataset(self, processed_fields, d_meta):
 
@@ -656,18 +606,24 @@ class InitialConditions:
             If the field specified by `var_name` is 2D and both `eta` and `xi` are specified.
         """
 
-        if len(self.ds[var_name].squeeze().dims) == 3 and not any(
-            [s is not None, eta is not None, xi is not None]
-        ):
+        field = self.ds[var_name].squeeze()
+
+        if len(field.dims) == 3:
+            if not any([s is not None, eta is not None, xi is not None]):
+                raise ValueError(
+                    "Invalid input: For 3D fields, you must specify at least one of the dimensions 's', 'eta', or 'xi'."
+                )
+            if all([s is not None, eta is not None, xi is not None]):
+                raise ValueError(
+                    "Ambiguous input: For 3D fields, specify at most two of 's', 'eta', or 'xi'. Specifying all three is not allowed."
+                )
+
+        if len(field.dims) == 2 and all([eta is not None, xi is not None]):
             raise ValueError(
-                "For 3D fields, at least one of s, eta, or xi must be specified."
+                "Conflicting input: For 2D fields, specify only one dimension, either 'eta' or 'xi', not both."
             )
 
-        if len(self.ds[var_name].squeeze().dims) == 2 and all(
-            [eta is not None, xi is not None]
-        ):
-            raise ValueError("For 2D fields, specify either eta or xi, not both.")
-
+        # Load the data
         if self.use_dask:
             from dask.diagnostics import ProgressBar
 
@@ -675,54 +631,73 @@ class InitialConditions:
                 self.ds[var_name].load()
 
         field = self.ds[var_name].squeeze()
-        if s is not None:
-            layer_contours = False
 
+        # Get correct mask and horizontal coordinates
         if all(dim in field.dims for dim in ["eta_rho", "xi_rho"]):
-            if layer_contours:
-                if "interface_depth_rho" in self.grid.ds:
-                    interface_depth = self.grid.ds.interface_depth_rho
-                else:
-                    self.get_vertical_coordinates(
-                        type="interface", additional_locations=[]
-                    )
-            layer_depth = self.grid.ds.layer_depth_rho
-            mask = self.grid.ds.mask_rho
-            field = field.assign_coords(
-                {"lon": self.grid.ds.lon_rho, "lat": self.grid.ds.lat_rho}
-            )
-
+            loc = "rho"
         elif all(dim in field.dims for dim in ["eta_rho", "xi_u"]):
-            if layer_contours:
-                if "interface_depth_u" in self.grid.ds:
-                    interface_depth = self.grid.ds.interface_depth_u
-                else:
-                    self.get_vertical_coordinates(
-                        type="interface", additional_locations=["u", "v"]
-                    )
-            layer_depth = self.grid.ds.layer_depth_u
-            mask = self.grid.ds.mask_u
-            field = field.assign_coords(
-                {"lon": self.grid.ds.lon_u, "lat": self.grid.ds.lat_u}
-            )
+            loc = "u"
 
         elif all(dim in field.dims for dim in ["eta_v", "xi_rho"]):
-            if layer_contours:
-                if "interface_depth_v" in self.grid.ds:
-                    interface_depth = self.grid.ds.interface_depth_v
+            loc = "v"
+        else:
+            ValueError("provided field does not have two horizontal dimension")
+
+        mask = self.grid.ds[f"mask_{loc}"]
+        lat_deg = self.grid.ds[f"lat_{loc}"]
+        lon_deg = self.grid.ds[f"lon_{loc}"]
+
+        if self.grid.straddle:
+            lon_deg = xr.where(lon_deg > 180, lon_deg - 360, lon_deg)
+
+        field = field.assign_coords({"lon": lon_deg, "lat": lat_deg})
+
+        # Retrieve depth coordinates
+        if s is not None:
+            layer_contours = False
+        # Note that `layer_depth_{loc}` has already been computed during `__post_init__`.
+        layer_depth = self.grid.ds[f"layer_depth_{loc}"]
+        if layer_contours:
+            if f"interface_depth_{loc}" not in self.grid.ds:
+                if loc == "rho":
+                    self.get_vertical_coordinates(
+                        type="interface", additional_locations=[]
+                    )
                 else:
                     self.get_vertical_coordinates(
                         type="interface", additional_locations=["u", "v"]
                     )
-            layer_depth = self.grid.ds.layer_depth_v
-            mask = self.grid.ds.mask_v
-            field = field.assign_coords(
-                {"lon": self.grid.ds.lon_v, "lat": self.grid.ds.lat_v}
-            )
+            interface_depth = self.grid.ds[f"interface_depth_{loc}"]
         else:
-            ValueError("provided field does not have two horizontal dimension")
+            interface_depth = None
+
+        # Slice the field as desired
+        def _slice_and_assign(
+            field,
+            mask,
+            layer_depth,
+            interface_depth,
+            title,
+            dim_name,
+            dim_values,
+            idx,
+            layer_contours=False,
+        ):
+            if dim_name in field.dims:
+                title = title + f", {dim_name} = {dim_values[idx].item()}"
+                field = field.isel(**{dim_name: idx})
+                mask = mask.isel(**{dim_name: idx})
+                layer_depth = layer_depth.isel(**{dim_name: idx})
+                if layer_contours:
+                    interface_depth = interface_depth.isel(**{dim_name: idx})
+                if "s_rho" in field.dims:
+                    field = field.assign_coords({"layer_depth": layer_depth})
+            else:
+                raise ValueError(
+                    f"None of the expected dimensions ({dim_name}) found in field."
+                )
+            return field, mask, layer_depth, interface_depth, title
 
-        # slice the field as desired
         title = field.long_name
         if s is not None:
             title = title + f", s_rho = {field.s_rho[s].item()}"
@@ -733,49 +708,32 @@ class InitialConditions:
             depth_contours = False
 
         if eta is not None:
-            if "eta_rho" in field.dims:
-                title = title + f", eta_rho = {field.eta_rho[eta].item()}"
-                field = field.isel(eta_rho=eta)
-                layer_depth = layer_depth.isel(eta_rho=eta)
-                if layer_contours:
-                    interface_depth = interface_depth.isel(eta_rho=eta)
-                if "s_rho" in field.dims:
-                    field = field.assign_coords({"layer_depth": layer_depth})
-            elif "eta_v" in field.dims:
-                title = title + f", eta_v = {field.eta_v[eta].item()}"
-                field = field.isel(eta_v=eta)
-                layer_depth = layer_depth.isel(eta_v=eta)
-                if layer_contours:
-                    interface_depth = interface_depth.isel(eta_v=eta)
-                if "s_rho" in field.dims:
-                    field = field.assign_coords({"layer_depth": layer_depth})
-            else:
-                raise ValueError(
-                    f"None of the expected dimensions (eta_rho, eta_v) found in ds[{var_name}]."
-                )
+            field, mask, layer_depth, interface_depth, title = _slice_and_assign(
+                field,
+                mask,
+                layer_depth,
+                interface_depth,
+                title,
+                "eta_rho" if "eta_rho" in field.dims else "eta_v",
+                field.eta_rho if "eta_rho" in field.dims else field.eta_v,
+                eta,
+                layer_contours,
+            )
+
         if xi is not None:
-            if "xi_rho" in field.dims:
-                title = title + f", xi_rho = {field.xi_rho[xi].item()}"
-                field = field.isel(xi_rho=xi)
-                layer_depth = layer_depth.isel(xi_rho=xi)
-                if layer_contours:
-                    interface_depth = interface_depth.isel(xi_rho=xi)
-                if "s_rho" in field.dims:
-                    field = field.assign_coords({"layer_depth": layer_depth})
-            elif "xi_u" in field.dims:
-                title = title + f", xi_u = {field.xi_u[xi].item()}"
-                field = field.isel(xi_u=xi)
-                layer_depth = layer_depth.isel(xi_u=xi)
-                if layer_contours:
-                    interface_depth = interface_depth.isel(xi_u=xi)
-                if "s_rho" in field.dims:
-                    field = field.assign_coords({"layer_depth": layer_depth})
-            else:
-                raise ValueError(
-                    f"None of the expected dimensions (xi_rho, xi_u) found in ds[{var_name}]."
-                )
+            field, mask, layer_depth, interface_depth, title = _slice_and_assign(
+                field,
+                mask,
+                layer_depth,
+                interface_depth,
+                title,
+                "xi_rho" if "xi_rho" in field.dims else "xi_u",
+                field.xi_rho if "xi_rho" in field.dims else field.xi_u,
+                xi,
+                layer_contours,
+            )
 
-        # chose colorbar
+        # Choose colorbar
         if var_name in ["u", "v", "w", "ubar", "vbar", "zeta"]:
             vmax = max(field.max().values, -field.min().values)
             vmin = -vmax
@@ -792,9 +750,7 @@ class InitialConditions:
 
         if eta is None and xi is None:
             _plot(
-                self.grid.ds,
                 field=field.where(mask),
-                straddle=self.grid.straddle,
                 depth_contours=depth_contours,
                 title=title,
                 kwargs=kwargs,
@@ -813,7 +769,7 @@ class InitialConditions:
 
             if len(field.dims) == 2:
                 _section_plot(
-                    field,
+                    field.where(mask),
                     interface_depth=interface_depth,
                     title=title,
                     kwargs=kwargs,
@@ -821,9 +777,9 @@ class InitialConditions:
                 )
             else:
                 if "s_rho" in field.dims:
-                    _profile_plot(field, title=title, ax=ax)
+                    _profile_plot(field.where(mask), title=title, ax=ax)
                 else:
-                    _line_plot(field, title=title, ax=ax)
+                    _line_plot(field.where(mask), title=title, ax=ax)
 
     def save(
         self, filepath: Union[str, Path], np_eta: int = None, np_xi: int = None

roms_tools/setup/nesting.py

@@ -4,7 +4,10 @@ from scipy.interpolate import griddata
 from dataclasses import dataclass, field
 from typing import Dict, Union
 from pathlib import Path
-from roms_tools.setup.grid import Grid
+import logging
+from scipy.interpolate import interp1d
+from roms_tools import Grid
+from roms_tools.plot import _plot_nesting
 from roms_tools.setup.utils import (
     interpolate_from_rho_to_u,
     interpolate_from_rho_to_v,
@@ -14,9 +17,6 @@ from roms_tools.setup.utils import (
     _to_yaml,
     _from_yaml,
 )
-from roms_tools.setup.plot import _plot_nesting
-import logging
-from scipy.interpolate import interp1d
 
 
 @dataclass(frozen=True, kw_only=True)

roms_tools/setup/river_forcing.py

@@ -2,12 +2,14 @@ import xarray as xr
 import numpy as np
 import logging
 from dataclasses import dataclass, field
-from roms_tools.setup.grid import Grid
+import cartopy.crs as ccrs
 from datetime import datetime
 from typing import Dict, Union, List
-from roms_tools.setup.datasets import DaiRiverDataset
 from pathlib import Path
 import matplotlib.pyplot as plt
+from roms_tools import Grid
+from roms_tools.plot import _get_projection, _add_field_to_ax
+from roms_tools.setup.datasets import DaiRiverDataset
 from roms_tools.setup.utils import (
     get_target_coords,
     gc_dist,
@@ -18,8 +20,6 @@ from roms_tools.setup.utils import (
     _from_yaml,
     get_variable_metadata,
 )
-from roms_tools.setup.plot import _get_projection, _add_field_to_ax
-import cartopy.crs as ccrs
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -416,16 +416,13 @@ class RiverForcing:
         """Plots the original and updated river locations on a map projection."""
 
         field = self.grid.ds.mask_rho
-        field = field.assign_coords(
-            {"lon": self.grid.ds.lon_rho, "lat": self.grid.ds.lat_rho}
-        )
         vmax = 3
         vmin = 0
         cmap = plt.colormaps.get_cmap("Blues")
         kwargs = {"vmax": vmax, "vmin": vmin, "cmap": cmap}
 
-        lon_deg = field.lon
-        lat_deg = field.lat
+        lon_deg = self.grid.ds.lon_rho
+        lat_deg = self.grid.ds.lat_rho
 
         # check if North or South pole are in domain
         if lat_deg.max().values > 89 or lat_deg.min().values < -89:
@@ -435,6 +432,7 @@ class RiverForcing:
 
         if self.grid.straddle:
             lon_deg = xr.where(lon_deg > 180, lon_deg - 360, lon_deg)
+        field = field.assign_coords({"lon": lon_deg, "lat": lat_deg})
 
         trans = _get_projection(lon_deg, lat_deg)
 

roms_tools/setup/surface_forcing.py

@@ -1,11 +1,14 @@
 import xarray as xr
 import importlib.metadata
 from dataclasses import dataclass, field
-from roms_tools.setup.grid import Grid
 from datetime import datetime
 import numpy as np
+import matplotlib.pyplot as plt
+from pathlib import Path
 from typing import Dict, Union, List
-from roms_tools.setup.regrid import LateralRegrid
+from roms_tools import Grid
+from roms_tools.regrid import LateralRegrid
+from roms_tools.plot import _plot
 from roms_tools.setup.datasets import (
     ERA5Dataset,
     ERA5Correction,
@@ -24,9 +27,6 @@ from roms_tools.setup.utils import (
     _to_yaml,
     _from_yaml,
 )
-from roms_tools.setup.plot import _plot
-import matplotlib.pyplot as plt
-from pathlib import Path
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -434,21 +434,23 @@ class SurfaceForcing:
             raise ValueError(f"Variable '{var_name}' is not found in dataset.")
 
         field = self.ds[var_name].isel(time=time)
+
         if self.use_dask:
             from dask.diagnostics import ProgressBar
 
             with ProgressBar():
                 field = field.load()
 
-        title = field.long_name
-
         field = field.where(self.target_coords["mask"])
 
-        field = field.assign_coords(
-            {"lon": self.target_coords["lon"], "lat": self.target_coords["lat"]}
-        )
+        lon_deg = self.target_coords["lon"]
+        lat_deg = self.target_coords["lat"]
+        if self.grid.straddle:
+            lon_deg = xr.where(lon_deg > 180, lon_deg - 360, lon_deg)
+        field = field.assign_coords({"lon": lon_deg, "lat": lat_deg})
+
+        title = field.long_name
 
-        # choose colorbar
         if var_name in ["uwnd", "vwnd"]:
             vmax = max(field.max().values, -field.min().values)
             vmin = -vmax
@@ -465,12 +467,10 @@ class SurfaceForcing:
         kwargs = {"vmax": vmax, "vmin": vmin, "cmap": cmap}
 
         _plot(
-            self.grid.ds,
             field=field,
-            straddle=self.grid.straddle,
             title=title,
-            kwargs=kwargs,
             c="g",
+            kwargs=kwargs,
         )
 
     def save(

roms_tools/setup/tides.py

@@ -3,9 +3,12 @@ import xarray as xr
 import numpy as np
 from typing import Dict, Union, List
 import importlib.metadata
+import matplotlib.pyplot as plt
+from pathlib import Path
 from dataclasses import dataclass, field
-from roms_tools.setup.grid import Grid
-from roms_tools.setup.plot import _plot
+from roms_tools import Grid
+from roms_tools.plot import _plot
+from roms_tools.regrid import LateralRegrid
 from roms_tools.setup.datasets import TPXODataset
 from roms_tools.setup.utils import (
     nan_check,
@@ -20,9 +23,6 @@ from roms_tools.setup.utils import (
     _to_yaml,
     _from_yaml,
 )
-from roms_tools.setup.regrid import LateralRegrid
-import matplotlib.pyplot as plt
-from pathlib import Path
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -319,6 +319,8 @@ class TidalForcing:
         >>> tidal_forcing.plot("ssh_Re", nc=0)
         """
 
+        if var_name not in self.ds:
+            raise ValueError(f"Variable '{var_name}' is not found in dataset.")
         field = self.ds[var_name].isel(ntides=ntides)
 
         if self.use_dask:
@@ -328,25 +330,28 @@ class TidalForcing:
                 field = field.load()
 
         if all(dim in field.dims for dim in ["eta_rho", "xi_rho"]):
-            field = field.where(self.grid.ds.mask_rho)
-            field = field.assign_coords(
-                {"lon": self.grid.ds.lon_rho, "lat": self.grid.ds.lat_rho}
-            )
+            lon_deg = self.grid.ds["lon_rho"]
+            lat_deg = self.grid.ds["lat_rho"]
+            mask = self.grid.ds["mask_rho"]
 
         elif all(dim in field.dims for dim in ["eta_rho", "xi_u"]):
-            field = field.where(self.grid.ds.mask_u)
-            field = field.assign_coords(
-                {"lon": self.grid.ds.lon_u, "lat": self.grid.ds.lat_u}
-            )
+            lon_deg = self.grid.ds["lon_u"]
+            lat_deg = self.grid.ds["lat_u"]
+            mask = self.grid.ds["mask_u"]
 
         elif all(dim in field.dims for dim in ["eta_v", "xi_rho"]):
-            field = field.where(self.grid.ds.mask_v)
-            field = field.assign_coords(
-                {"lon": self.grid.ds.lon_v, "lat": self.grid.ds.lat_v}
-            )
+            lon_deg = self.grid.ds["lon_v"]
+            lat_deg = self.grid.ds["lat_v"]
+            mask = self.grid.ds["mask_v"]
+
         else:
             ValueError("provided field does not have two horizontal dimension")
 
+        field = field.where(mask)
+        if self.grid.straddle:
+            lon_deg = xr.where(lon_deg > 180, lon_deg - 360, lon_deg)
+        field = field.assign_coords({"lon": lon_deg, "lat": lat_deg})
+
         title = "%s, ntides = %i" % (field.long_name, self.ds[var_name].ntides[ntides])
 
         vmax = max(field.max(), -field.min())
@@ -357,12 +362,10 @@ class TidalForcing:
         kwargs = {"vmax": vmax, "vmin": vmin, "cmap": cmap}
 
         _plot(
-            self.grid.ds,
             field=field,
-            straddle=self.grid.straddle,
+            title=title,
             c="g",
             kwargs=kwargs,
-            title=title,
         )
 
     def save(
@@ -464,7 +467,7 @@ class TidalForcing:
             grid=grid,
             **tidal_forcing_params,
             use_dask=use_dask,
-            bypass_validation=bypass_validation
+            bypass_validation=bypass_validation,
         )
 
     def _correct_tides(self, data):

roms_tools/setup/topography.py

@@ -6,8 +6,8 @@ import gcm_filters
 from roms_tools.setup.utils import handle_boundaries
 import warnings
 from itertools import count
+from roms_tools.regrid import LateralRegrid
 from roms_tools.setup.datasets import ETOPO5Dataset, SRTM15Dataset
-from roms_tools.setup.regrid import LateralRegrid
 
 
 def _add_topography(