roms-tools 3.0.0__py3-none-any.whl → 3.1.1__py3-none-any.whl

roms_tools/setup/cdr_forcing.py

@@ -1,17 +1,15 @@
 import itertools
 import logging
 from collections import Counter
+from collections.abc import Iterator
 from datetime import datetime
 from pathlib import Path
-from typing import Annotated, Iterator
+from typing import Annotated
 
-import cartopy.crs as ccrs
-import gcm_filters
 import matplotlib.gridspec as gridspec
 import matplotlib.pyplot as plt
 import numpy as np
 import xarray as xr
-from matplotlib.axes import Axes
 from pydantic import (
     BaseModel,
     Field,
@@ -22,9 +20,14 @@ from pydantic import (
 )
 
 from roms_tools import Grid
-from roms_tools.constants import R_EARTH
-from roms_tools.plot import _get_projection, _plot
-from roms_tools.regrid import LateralRegridFromROMS
+from roms_tools.constants import MAX_DISTINCT_COLORS
+from roms_tools.plot import (
+    assign_category_colors,
+    get_projection,
+    plot,
+    plot_2d_horizontal_field,
+    plot_location,
+)
 from roms_tools.setup.cdr_release import (
     Release,
     ReleaseType,
@@ -32,27 +35,29 @@ from roms_tools.setup.cdr_release import (
     VolumeRelease,
 )
 from roms_tools.setup.utils import (
-    _from_yaml,
-    _to_dict,
-    _write_to_yaml,
     add_tracer_metadata_to_ds,
     convert_to_relative_days,
+    from_yaml,
     gc_dist,
     get_target_coords,
+    to_dict,
+    validate_names,
+    write_to_yaml,
 )
 from roms_tools.utils import (
-    _generate_focused_coordinate_range,
-    _remove_edge_nans,
     normalize_longitude,
     save_datasets,
 )
+from roms_tools.vertical_coordinate import compute_depth_coordinates
 
 INCLUDE_ALL_RELEASE_NAMES = "all"
+MAX_RELEASES_TO_PLOT = 20  # must be <= MAX_DISTINCT_COLORS
 
 
 class ReleaseSimulationManager(BaseModel):
     """Validates and adjusts a single release against a ROMS simulation time window and
-    grid."""
+    grid.
+    """
 
     release: Release
     grid: Grid | None = None
@@ -64,8 +69,8 @@ class ReleaseSimulationManager(BaseModel):
         self,
     ) -> "ReleaseSimulationManager":
         """Ensure the release times are within the [start_time, end_time] simulation
-        window."""
-
+        window.
+        """
         times = self.release.times
         if len(times) == 0:
             return self
@@ -125,7 +130,8 @@ class ReleaseCollector(RootModel):
     @classmethod
     def unpack_dict(cls, data):
         """This helps directly translate a dict of {"releases": [...]} into just the
-        list of releases."""
+        list of releases.
+        """
         if isinstance(data, dict):
             try:
                 return data["releases"]
@@ -162,6 +168,7 @@ class ReleaseCollector(RootModel):
 
     @property
     def release_type(self):
+        """Type of all releases."""
         release_types = set(r.release_type for r in self.root)
         return release_types.pop()
 
@@ -169,14 +176,30 @@ class ReleaseCollector(RootModel):
 class CDRForcingDatasetBuilder:
     """Constructs the xarray `Dataset` to be saved as NetCDF."""
 
-    def __init__(self, releases, model_reference_date, release_type: ReleaseType):
+    def __init__(
+        self,
+        releases: ReleaseCollector,
+        model_reference_date: datetime,
+        release_type: ReleaseType,
+    ):
+        """
+        Initialize the dataset builder.
+
+        Parameters
+        ----------
+        releases : list
+            List of release objects.
+        model_reference_date : datetime
+            Reference date for relative time conversion.
+        release_type : ReleaseType
+            Type of release.
+        """
         self.releases = releases
         self.model_reference_date = model_reference_date
         self.release_type = release_type
 
     def build(self) -> xr.Dataset:
         """Build the CDR forcing dataset."""
-
         all_times = itertools.chain.from_iterable(r.times for r in self.releases)
         unique_times = np.unique(np.array(list(all_times), dtype="datetime64[ns]"))
         unique_rel_times = convert_to_relative_days(
@@ -224,7 +247,6 @@ class CDRForcingDatasetBuilder:
         xr.Dataset
             Initialized dataset with time, location, and release-type-dependent variables.
         """
-
         ds = xr.Dataset()
         ds["time"] = ("time", unique_times)
         ds["cdr_time"] = ("time", unique_rel_times)
@@ -317,7 +339,7 @@ class CDRForcing(BaseModel):
     Parameters
     ----------
     grid : Grid, optional
-        Object representing the grid for spatial context.
+        Object representing the grid.
     start_time : datetime
         Start time of the ROMS model simulation.
     end_time : datetime
@@ -326,18 +348,18 @@ class CDRForcing(BaseModel):
         Reference date for converting absolute times to model-relative time. Defaults to Jan 1, 2000.
     releases : list of Release
         A list of one or more CDR release objects.
-
-    Attributes
-    ----------
-    ds : xr.Dataset
-        The xarray dataset containing release metadata and forcing variables.
     """
 
     grid: Grid | None = None
+    """Object representing the grid."""
     start_time: datetime
+    """Start time of the ROMS model simulation."""
     end_time: datetime
+    """End time of the ROMS model simulation."""
     model_reference_date: datetime = datetime(2000, 1, 1)
+    """The reference date for the ROMS simulation."""
     releases: ReleaseCollector
+    """A list of one or more CDR release objects."""
 
     # this is defined during init and shouldn't be serialized
     _ds: xr.Dataset = None
@@ -365,14 +387,19 @@ class CDRForcing(BaseModel):
 
     @property
     def release_type(self) -> ReleaseType:
+        """Type of the release."""
         return self.releases.release_type
 
     @property
     def ds(self) -> xr.Dataset:
+        """The xarray dataset containing release metadata and forcing variables."""
         return self._ds
 
     def plot_volume_flux(
-        self, start=None, end=None, release_names=INCLUDE_ALL_RELEASE_NAMES
+        self,
+        start: datetime | None = None,
+        end: datetime | None = None,
+        release_names: list[str] | str = INCLUDE_ALL_RELEASE_NAMES,
     ):
         """Plot the volume flux for each specified release within the given time range.
 
@@ -394,7 +421,6 @@ class CDRForcing(BaseModel):
             If `release_names` is not a list of strings or "all".
             If any of the specified release names do not exist in `self.releases`.
         """
-
         if self.release_type != ReleaseType.volume:
             raise ValueError(
                 "plot_volume_flux is only supported when all releases are of type VolumeRelease."
@@ -403,12 +429,7 @@ class CDRForcing(BaseModel):
         start = start or self.start_time
         end = end or self.end_time
 
-        valid_release_names = [r.name for r in self.releases]
-
-        if release_names == INCLUDE_ALL_RELEASE_NAMES:
-            release_names = valid_release_names
-
-        _validate_release_input(release_names, valid_release_names)
+        release_names = _validate_release_names(release_names, self.releases)
 
         data = self.ds["cdr_volume"]
 
@@ -424,9 +445,9 @@ class CDRForcing(BaseModel):
     def plot_tracer_concentration(
         self,
         tracer_name: str,
-        start=None,
-        end=None,
-        release_names=INCLUDE_ALL_RELEASE_NAMES,
+        start: datetime | None = None,
+        end: datetime | None = None,
+        release_names: list[str] | str = INCLUDE_ALL_RELEASE_NAMES,
     ):
         """Plot the concentration of a given tracer for each specified release within
         the given time range.
@@ -460,12 +481,7 @@ class CDRForcing(BaseModel):
         start = start or self.start_time
         end = end or self.end_time
 
-        valid_release_names = [r.name for r in self.releases]
-
-        if release_names == INCLUDE_ALL_RELEASE_NAMES:
-            release_names = valid_release_names
-
-        _validate_release_input(release_names, valid_release_names)
+        release_names = _validate_release_names(release_names, self.releases)
 
         tracer_names = list(self.ds["tracer_name"].values)
         if tracer_name not in tracer_names:
@@ -495,9 +511,9 @@ class CDRForcing(BaseModel):
     def plot_tracer_flux(
         self,
         tracer_name: str,
-        start=None,
-        end=None,
-        release_names=INCLUDE_ALL_RELEASE_NAMES,
+        start: datetime | None = None,
+        end: datetime | None = None,
+        release_names: list[str] | str = INCLUDE_ALL_RELEASE_NAMES,
     ):
         """Plot the flux of a given tracer for each specified release within the given
         time range.
@@ -531,12 +547,7 @@ class CDRForcing(BaseModel):
         start = start or self.start_time
         end = end or self.end_time
 
-        valid_release_names = [r.name for r in self.releases]
-
-        if release_names == INCLUDE_ALL_RELEASE_NAMES:
-            release_names = valid_release_names
-
-        _validate_release_input(release_names, valid_release_names)
+        release_names = _validate_release_names(release_names, self.releases)
 
         tracer_names = list(self.ds["tracer_name"].values)
         if tracer_name not in tracer_names:
@@ -561,7 +572,10 @@ class CDRForcing(BaseModel):
     def _plot_line(self, data, release_names, start, end, title="", ylabel=""):
         """Plots a line graph for the specified releases and time range."""
         valid_release_names = [r.name for r in self.releases]
-        colors = _get_release_colors(valid_release_names)
+        if len(valid_release_names) > MAX_DISTINCT_COLORS:
+            colors = assign_category_colors(release_names)
+        else:
+            colors = assign_category_colors(valid_release_names)
 
         fig, ax = plt.subplots(1, 1, figsize=(7, 4))
         for name in release_names:
@@ -580,7 +594,9 @@ class CDRForcing(BaseModel):
         ax.set(title=title, ylabel=ylabel, xlabel="time")
         ax.set_xlim([start, end])
 
-    def plot_locations(self, release_names="all"):
+    def plot_locations(
+        self, release_names: list[str] | str = INCLUDE_ALL_RELEASE_NAMES
+    ):
         """Plot centers of release locations in top-down view.
 
         Parameters
@@ -597,43 +613,49 @@ class CDRForcing(BaseModel):
             If any of the specified release names do not exist in `self.releases`.
             If `self.grid` is not set.
         """
-
         # Ensure that the grid is provided
         if self.grid is None:
             raise ValueError(
                 "A grid must be provided for plotting. Please pass a valid `Grid` object."
             )
 
-        valid_release_names = [r.name for r in self.releases]
-
-        if release_names == "all":
-            release_names = valid_release_names
-
-        _validate_release_input(release_names, valid_release_names)
+        release_names = _validate_release_names(release_names, self.releases)
 
         lon_deg = self.grid.ds.lon_rho
         lat_deg = self.grid.ds.lat_rho
         if self.grid.straddle:
             lon_deg = xr.where(lon_deg > 180, lon_deg - 360, lon_deg)
-        trans = _get_projection(lon_deg, lat_deg)
+        trans = get_projection(lon_deg, lat_deg)
         fig, ax = plt.subplots(1, 1, figsize=(13, 7), subplot_kw={"projection": trans})
 
         # Plot blue background on map
         field = self.grid.ds.mask_rho
+        field = field.where(field)
         field = field.assign_coords({"lon": lon_deg, "lat": lat_deg})
         vmax = 6
         vmin = 0
         cmap = plt.colormaps.get_cmap("Blues")
+        cmap.set_bad(color="gray")
         kwargs = {"vmax": vmax, "vmin": vmin, "cmap": cmap}
-        _plot(field, kwargs=kwargs, ax=ax, c=None, add_colorbar=False)
+        plot_2d_horizontal_field(field, kwargs=kwargs, ax=ax, add_colorbar=False)
 
         # Plot release locations
-        colors = _get_release_colors(valid_release_names)
-        _plot_location(
-            grid=self.grid,
-            releases=[self.releases[name] for name in release_names],
+        valid_release_names = [r.name for r in self.releases]
+        if len(valid_release_names) > MAX_DISTINCT_COLORS:
+            colors = assign_category_colors(release_names)
+        else:
+            colors = assign_category_colors(valid_release_names)
+        plot_location(
+            grid_ds=self.grid.ds,
+            points={
+                name: {
+                    "lat": self.releases[name].lat,
+                    "lon": self.releases[name].lon,
+                    "color": colors.get(name, "k"),
+                }
+                for name in release_names
+            },
             ax=ax,
-            colors=colors,
         )
 
     def plot_distribution(self, release_name: str, mark_release_center: bool = True):
@@ -663,8 +685,13 @@ class CDRForcing(BaseModel):
                 "A grid must be provided for plotting. Please pass a valid `Grid` object."
             )
 
-        valid_release_names = [r.name for r in self.releases]
-        _validate_release_input(release_name, valid_release_names, list_allowed=False)
+        if not isinstance(release_name, str):
+            raise ValueError(
+                f"Only a single release name (string) is allowed. Got: {release_name!r}"
+            )
+
+        release_name = _validate_release_names([release_name], self.releases)[0]
+
         release = self.releases[release_name]
 
         # Prepare grid coordinates
@@ -676,73 +703,73 @@ class CDRForcing(BaseModel):
         # Setup figure
         fig = plt.figure(figsize=(12, 5.5))
         gs = gridspec.GridSpec(nrows=2, ncols=2, figure=fig)
-        trans = _get_projection(lon_deg, lat_deg)
+        trans = get_projection(lon_deg, lat_deg)
         ax0 = fig.add_subplot(gs[:, 0], projection=trans)
         ax1 = fig.add_subplot(gs[0, 1])
         ax2 = fig.add_subplot(gs[1, 1])
-        cmap = plt.colormaps.get_cmap("RdPu")
-        cmap.set_bad(color="gray")
-        kwargs = {"cmap": cmap}
 
-        # Top down view plot
-        horizontal_field = _map_horizontal_gaussian(self.grid, release)
-        horizontal_field = horizontal_field.assign_coords(
+        # Top down view plot of depth-integrated Gaussian
+        depth_integrated_field = _map_horizontal_gaussian(self.grid, release)
+        depth_integrated_field = depth_integrated_field.assign_coords(
             {"lon": lon_deg, "lat": lat_deg}
         )
-        _plot(
-            horizontal_field.where(self.grid.ds.mask_rho),
-            kwargs=kwargs,
+        cmap = plt.colormaps.get_cmap("RdPu")
+        cmap.set_bad(color="gray")
+        plot_2d_horizontal_field(
+            field=depth_integrated_field.where(self.grid.ds.mask_rho),
             ax=ax0,
-            c=None,
+            kwargs={"cmap": cmap},
             add_colorbar=False,
+            title="Depth-integrated distribution",
         )
         if mark_release_center:
-            _plot_location(
-                grid=self.grid, releases=[release], ax=ax0, include_legend=False
+            plot_location(
+                grid_ds=self.grid.ds,
+                points={
+                    release.name: {
+                        "lat": release.lat,
+                        "lon": release.lon,
+                    }
+                },
+                ax=ax0,
+                include_legend=False,
             )
 
-        # Side view plots
-        kwargs = {
-            "cmap": cmap,
-            "y": "depth",
-            "yincrease": False,
-            "add_colorbar": False,
-        }
-        # Plot along latitude
-        vertical_field = _map_vertical_gaussian(
-            self.grid, release, horizontal_field, orientation="latitude"
-        )
-        more_kwargs = {"x": "lat"}
-        vertical_field.plot(**kwargs, **more_kwargs, ax=ax1)
-        if mark_release_center:
-            ax1.plot(
-                release.lat,
-                release.depth,
-                color="k",
-                marker="x",
-                markersize=8,
-                markeredgewidth=2,
-            )
+        # Spread horizontal Gaussian field into the vertical
+        distribution_3d = _map_3d_gaussian(self.grid, release, depth_integrated_field)
 
+        release_lon = normalize_longitude(release.lon, self.grid.straddle)
+
+        plot(
+            field=distribution_3d,
+            grid_ds=self.grid.ds,
+            lon=release_lon,
+            ax=ax1,
+            cmap_name="RdPu",
+            add_colorbar=False,
+        )
         ax1.set(title=f"Longitude: {release.lon}°E", xlabel="Latitude [°N]")
-        # Plot along longitude
-        vertical_field = _map_vertical_gaussian(
-            self.grid, release, horizontal_field, orientation="longitude"
+
+        plot(
+            field=distribution_3d,
+            grid_ds=self.grid.ds,
+            lat=release.lat,
+            ax=ax2,
+            cmap_name="RdPu",
+            add_colorbar=False,
         )
-        more_kwargs = {"x": "lon"}
-        vertical_field.plot(**kwargs, **more_kwargs, ax=ax2)
-        if mark_release_center:
-            release_lon = normalize_longitude(release.lon, self.grid.straddle)
-            ax2.plot(
-                release_lon,
-                release.depth,
-                color="k",
-                marker="x",
-                markersize=8,
-                markeredgewidth=2,
-            )
         ax2.set(title=f"Latitude: {release.lat}°N", xlabel="Longitude [°E]")
 
+        if mark_release_center:
+            kwargs = {
+                "color": "k",
+                "marker": "x",
+                "markersize": 8,
+                "markeredgewidth": 2,
+            }
+            ax1.plot(release.lat, release.depth, **kwargs)
+            ax2.plot(release_lon, release.depth, **kwargs)
+
         # Adjust layout and title
         fig.subplots_adjust(hspace=0.45)
         fig.suptitle(f"Release distribution for: {release_name}")
@@ -763,7 +790,6 @@ class CDRForcing(BaseModel):
         list[Path]
             A list of `Path` objects for the saved files. Each element in the list corresponds to a file that was saved.
         """
-
         # Ensure filepath is a Path object
         filepath = Path(filepath)
 
@@ -780,7 +806,7 @@ class CDRForcing(BaseModel):
 
     @model_serializer
     def _serialize(self) -> dict:
-        return _to_dict(self)
+        return to_dict(self)
 
     def to_yaml(self, filepath: str | Path) -> None:
         """Export the parameters of the class to a YAML file, including the version of
@@ -791,12 +817,11 @@ class CDRForcing(BaseModel):
         filepath : str | Path
             The path to the YAML file where the parameters will be saved.
         """
-
         forcing_dict = self.model_dump()
         metadata = self.releases[0].get_tracer_metadata()
         forcing_dict["CDRForcing"]["_tracer_metadata"] = metadata
 
-        _write_to_yaml(forcing_dict, filepath)
+        write_to_yaml(forcing_dict, filepath)
 
     @classmethod
     def from_yaml(cls, filepath: str | Path) -> "CDRForcing":
@@ -815,114 +840,45 @@ class CDRForcing(BaseModel):
         filepath = Path(filepath)
 
         grid = Grid.from_yaml(filepath)
-        params = _from_yaml(cls, filepath)
+        params = from_yaml(cls, filepath)
         params.pop("_tracer_metadata", None)
 
         return cls(grid=grid, **params)
 
 
-def _validate_release_input(releases, valid_releases, list_allowed=True):
-    """Validates the input for release names in plotting methods to ensure they are in
-    an acceptable format and exist within the set of valid releases.
-
-    This method ensures that the `releases` parameter is either a single release name (string) or a list
-    of release names (strings), and checks that each release exists in the set of valid releases.
-
-    Parameters
-    ----------
-    releases : str or list of str
-        A single release name as a string, or a list of release names (strings) to validate.
-
-    list_allowed : bool, optional
-        If `True`, a list of release names is allowed. If `False`, only a single release name (string)
-        is allowed. Default is `True`.
-
-    Raises
-    ------
-    ValueError
-        If `releases` is not a string or list of strings, or if any release name is invalid (not in `self.releases`).
-
-    Notes
-    -----
-    This method checks that the `releases` input is in a valid format (either a string or a list of strings),
-    and ensures each release is present in the set of valid releases defined in `self.releases`. Invalid releases
-    are reported in the error message.
-
-    If `list_allowed` is set to `False`, only a single release name (string) will be accepted. Otherwise, a
-    list of release names is also acceptable.
+def _validate_release_names(
+    release_names: list[str] | str, releases: ReleaseCollector
+) -> list[str]:
     """
+    Validate and filter a list of release names.
 
-    # Ensure that a list of releases is only allowed if `list_allowed` is True
-    if not list_allowed and not isinstance(releases, str):
-        raise ValueError(
-            f"Only a single release name (string) is allowed. Got: {releases}"
-        )
-
-    if isinstance(releases, str):
-        releases = [releases]  # Convert to list if a single string is provided
-    elif isinstance(releases, list):
-        if not all(isinstance(r, str) for r in releases):
-            raise ValueError("All elements in `releases` list must be strings.")
-    else:
-        raise ValueError(
-            "`releases` should be a string (single release name) or a list of strings (release names)."
-        )
-
-    # Validate that the specified releases exist in self.releases
-    invalid_releases = [
-        release for release in releases if release not in valid_releases
-    ]
-    if invalid_releases:
-        raise ValueError(f"Invalid releases: {', '.join(invalid_releases)}")
-
-
-def _get_release_colors(valid_releases: list[str]) -> dict[str, tuple]:
-    """Returns a dictionary of colors for the valid releases, based on a consistent
-    colormap.
+    Ensures that each release name exists in `releases` and limits the list
+    to `MAX_RELEASES_TO_PLOT` entries with a warning if truncated.
 
     Parameters
     ----------
-    valid_releases : List[str]
-        List of release names to assign colors to.
+    release_names : list of str or INCLUDE_ALL_RELEASE_NAMES
+        Names of releases to plot, or sentinel to include all.
+    releases : ReleaseCollector
+        Object containing valid release names.
 
     Returns
     -------
-    Dict[str, tuple]
-        A dictionary where the keys are release names and the values are their corresponding colors,
-        assigned based on the order of releases in the valid releases list.
+    list of str
+        Validated and truncated list of release names.
 
     Raises
     ------
     ValueError
-        If the number of valid releases exceeds the available colormap capacity.
-
-    Notes
-    -----
-    The colormap is chosen dynamically based on the number of valid releases:
-
-    - If there are 10 or fewer releases, the "tab10" colormap is used.
-    - If there are more than 10 but fewer than or equal to 20 releases, the "tab20" colormap is used.
-    - For more than 20 releases, the "tab20b" colormap is used.
+        If any names are invalid.
     """
-
-    # Determine the colormap based on the number of releases
-    if len(valid_releases) <= 10:
-        color_map = plt.get_cmap("tab10")
-    elif len(valid_releases) <= 20:
-        color_map = plt.get_cmap("tab20")
-    else:
-        color_map = plt.get_cmap("tab20b")
-
-    # Ensure the number of releases doesn't exceed the available colormap capacity
-    if len(valid_releases) > color_map.N:
-        raise ValueError(
-            f"Too many releases. The selected colormap supports up to {color_map.N} releases."
-        )
-
-    # Create a dictionary of colors based on the release indices
-    colors = {name: color_map(i) for i, name in enumerate(valid_releases)}
-
-    return colors
+    return validate_names(
+        release_names,
+        [r.name for r in releases],
+        INCLUDE_ALL_RELEASE_NAMES,
+        MAX_RELEASES_TO_PLOT,
+        label="release",
+    )
 
 
 def _validate_release_location(grid, release: Release):
@@ -944,13 +900,18 @@ def _validate_release_location(grid, release: Release):
 
         dx = 1 / grid.ds.pm
         dy = 1 / grid.ds.pn
-        max_grid_spacing = np.sqrt(dx**2 + dy**2) / 2
+
+        # Compute the maximum half-diagonal grid spacing across the domain
+        max_grid_spacing = (np.sqrt(dx**2 + dy**2) / 2).max()
+
+        # Apply a 10% safety margin
+        max_grid_spacing *= 1.1
 
         # Compute great-circle distance to all grid points
         dist = gc_dist(grid.ds.lon_rho, grid.ds.lat_rho, lon, release.lat)
         dist_min = dist.min(dim=["eta_rho", "xi_rho"])
 
-        if (dist_min > max_grid_spacing).all():
+        if dist_min > max_grid_spacing:
             raise ValueError(
                 f"Release site '{release.name}' is outside of the grid domain. "
                 "Ensure the provided (lat, lon) falls within the model grid extent."
@@ -1017,98 +978,35 @@ def _map_horizontal_gaussian(grid: Grid, release: Release):
     else:
         lon = xr.where(lon < 0, lon + 360, lon)
     dist = gc_dist(target_coords["lon"], target_coords["lat"], lon, release.lat)
-    dist_min = dist.min(dim=["eta_rho", "xi_rho"])
-
-    # Find the indices of the closest grid cell
-    indices = np.where(dist == dist_min)
-    eta_rho = indices[0][0]
-    xi_rho = indices[1][0]
 
-    # Create a delta function at the center of the Gaussian release
-    delta = xr.zeros_like(grid.ds.mask_rho)
-    delta[eta_rho, xi_rho] = 1
-
-    # Calculate the grid cell area from the inverse grid metrics (pm, pn)
-    area = 1 / (grid.ds.pm * grid.ds.pn)
-
-    # Compute the mean grid spacing (dx) as the average of mean dx and dy
-    dx = (((1 / grid.ds.pm).mean() + (1 / grid.ds.pn).mean()) / 2).item()
-
-    # Extend the domain in both dimensions to mitigate boundary effects during filtering
-    # - Create a mask of ones ignoring land (will be renormalized later)
-    mask = xr.ones_like(grid.ds.mask_rho)
-
-    # Extend the mask along eta_rho by concatenating three copies plus a zeroed edge cell
-    margin_mask = xr.concat(
-        [mask, mask, mask, 0 * mask.isel(eta_rho=-1)], dim="eta_rho"
-    )
-    # Similarly extend along xi_rho
-    margin_mask = xr.concat(
-        [margin_mask, margin_mask, margin_mask, 0 * margin_mask.isel(xi_rho=-1)],
-        dim="xi_rho",
-    )
-
-    # Extend the delta function in the same manner to match the enlarged domain
-    delta_extended = xr.concat(
-        [0 * delta, delta, 0 * delta, 0 * delta.isel(eta_rho=-1)], dim="eta_rho"
-    )
-    delta_extended = xr.concat(
-        [
-            0 * delta_extended,
-            delta_extended,
-            0 * delta_extended,
-            0 * delta_extended.isel(xi_rho=-1),
-        ],
-        dim="xi_rho",
-    )
-
-    # Extend the cell area array to match the enlarged domain
-    area_extended = xr.concat([area, area, area, area.isel(eta_rho=-1)], dim="eta_rho")
-    area_extended = xr.concat(
-        [area_extended, area_extended, area_extended, area_extended.isel(xi_rho=-1)],
-        dim="xi_rho",
-    )
-
-    # Define Gaussian filter parameters:
-    # - filter_scale is computed so that the Gaussian's std dev matches
-    #   the equivalent boxcar filter of width equal to release.hsc / dx
-    #   multiplied by sqrt(12) to convert from boxcar to Gaussian std dev.
-    filter_scale = release.hsc / dx * np.sqrt(12)
-
-    filter = gcm_filters.Filter(
-        filter_scale=filter_scale,
-        dx_min=1,
-        filter_shape=gcm_filters.FilterShape.GAUSSIAN,
-        grid_type=gcm_filters.GridType.REGULAR_WITH_LAND_AREA_WEIGHTED,
-        grid_vars={"wet_mask": margin_mask, "area": area_extended},
-    )
-
-    # Apply the Gaussian filter to the extended delta function over spatial dimensions
-    delta_smooth = filter.apply(delta_extended, dims=["eta_rho", "xi_rho"])
+    if release.hsc == 0:
+        # Find the indices of the closest grid cell
+        indices = dist.argmin(dim=["eta_rho", "xi_rho"])
+        eta_rho = indices["eta_rho"].values
+        xi_rho = indices["xi_rho"].values
 
-    # Remove the extended margins to return to the original domain size
-    delta_smooth = delta_smooth.isel(
-        eta_rho=slice(grid.ds.sizes["eta_rho"], -grid.ds.sizes["eta_rho"] - 1),
-        xi_rho=slice(grid.ds.sizes["xi_rho"], -grid.ds.sizes["eta_rho"] - 1),
-    )
+        # Create a delta function at the center of the Gaussian release
+        distribution_2d = xr.zeros_like(grid.ds.mask_rho)
+        distribution_2d[eta_rho, xi_rho] = 1
 
-    # Mask out land cells and set values to zero there
-    delta_smooth = delta_smooth.where(grid.ds.mask_rho, 0.0)
+    else:
+        frac = np.exp(-((dist / release.hsc) ** 2))
+        distribution_2d = frac.where(frac > 1e-3, 0.0)
 
-    # Renormalize so the integral over ocean points sums to 1
-    integral = delta_smooth.sum(dim=["eta_rho", "xi_rho"])
-    delta_smooth = delta_smooth / integral
+        # Mask out land
+        distribution_2d = distribution_2d.where(grid.ds.mask_rho, 0.0)
 
-    return delta_smooth
+        # Renormalize so the integral over ocean points sums to 1
+        integral = distribution_2d.sum(dim=["eta_rho", "xi_rho"])
+        distribution_2d = distribution_2d / integral
 
+    return distribution_2d
 
-def _map_vertical_gaussian(grid, release, field, orientation="latitude"):
-    """Extract a vertical section from a ROMS grid and apply a Gaussian distribution in
-    depth.
 
-    This function interpolates a horizontally Gaussian field (e.g., from a tracer release)
-    along either a constant latitude or longitude section and distributes it vertically using
-    a Gaussian profile centered around the release depth.
+def _map_3d_gaussian(
+    grid: Grid, release: Release, distribution_2d: xr.DataArray
+) -> xr.DataArray:
+    """Extends 2D Gaussian to 3D Gaussian.
 
     Parameters
     ----------
@@ -1118,170 +1016,44 @@ def _map_vertical_gaussian(grid, release, field, orientation="latitude"):
     release : Release
         Release object containing coordinates (`lat`, `lon`, `depth`) and vertical
         spread (`vsc`) for the Gaussian distribution.
-    field : xr.DataArray
+    distribution_2d : xr.DataArray
         2D horizontal tracer field defined on the ROMS grid.
-    orientation : {"latitude", "longitude"}, default "latitude"
-        Orientation of the extracted vertical section. If "latitude", extracts
-        a section along constant longitude. If "longitude", extracts a section
-        along constant latitude.
 
     Returns
     -------
-    vertical_field : xr.DataArray
-        2D field (depth vs. latitude or longitude) with the vertically-distributed
-        Gaussian mapped along the specified section.
+    xr.DataArray
+        3D tracer distribution (z vs lat/lon) with Gaussian vertical structure.
     """
-    meters_per_degree = 2 * np.pi * R_EARTH / 360
-    release_lon = normalize_longitude(release.lon, grid.straddle)
-
-    if orientation == "longitude":
-        lon_deg = grid.ds.lon_rho
-        if grid.straddle:
-            lon_deg = xr.where(lon_deg > 180, lon_deg - 360, lon_deg)
-        hsc_in_degrees = release.hsc / (meters_per_degree * np.cos(release.lat))
-        lons, _ = _generate_focused_coordinate_range(
-            center=release_lon,
-            sc=hsc_in_degrees,
-            min_val=lon_deg.min().values,
-            max_val=lon_deg.max().values,
-            N=2000,
-        )
-        lons = xr.DataArray(lons, dims=["lon"], attrs={"units": "°E"})
-        lats = [release.lat]
-
-    elif orientation == "latitude":
-        lons = [release_lon]
-        hsc_in_degrees = release.hsc / (
-            meters_per_degree * np.cos(grid.ds.lat_rho.mean().values)
-        )
-        lats, _ = _generate_focused_coordinate_range(
-            center=release.lat,
-            sc=hsc_in_degrees,
-            min_val=grid.ds.lat_rho.min().values,
-            max_val=grid.ds.lat_rho.max().values,
-            N=2000,
-        )
-        lats = xr.DataArray(lats, dims=["lat"], attrs={"units": "°N"})
-    else:
-        raise ValueError(
-            "`section_orientation` must be either 'latitude' or 'longitude'."
-        )
-
-    # Regrid 2D horizontal Gaussian onto desired 1D horizontal section
-    target_coords = {"lat": lats, "lon": lons}
-    lateral_regrid = LateralRegridFromROMS(field, target_coords)
-    field = lateral_regrid.apply(field).squeeze()
-    h = lateral_regrid.apply(grid.ds.h).squeeze()
-
-    # Define depth levels
-    depth_levels, _ = _generate_focused_coordinate_range(
-        center=release.depth,
-        sc=release.vsc,
-        min_val=0.0,
-        max_val=h.max().values,
-        N=2000,
-    )
-    depth_levels = xr.DataArray(
-        np.asarray(depth_levels),
-        dims=["depth"],
-        attrs={"long_name": "Depth", "units": "m"},
+    # Compute depth at rho-points (3D: s_rho, eta_rho, xi_rho)
+    depth = compute_depth_coordinates(
+        grid.ds, zeta=0, depth_type="layer", location="rho"
     )
-    depth_levels = depth_levels.astype(np.float32)
+
+    # Initialize 3D distribution with zeros
+    distribution_3d = xr.zeros_like(depth)
 
     if release.vsc == 0:
-        # Find index of depth_level closest to release.depth
-        closest_idx = np.abs(depth_levels - release.depth).argmin()
-        weights = xr.zeros_like(depth_levels)
-        weights[closest_idx] = 1.0
+        # Find vertical index closest to release depth
+        abs_diff = abs(depth - release.depth)
+        vertical_idx = abs_diff.argmin(dim="s_rho")
+        # Stack 2D distribution at that vertical level
+        distribution_3d[{"s_rho": vertical_idx}] = distribution_2d
     else:
-        # Compute Gaussian weights at each layer center
-        weights = np.exp(-0.5 * ((depth_levels - release.depth) / release.vsc) ** 2)
-        weights /= weights.sum()
-
-    # Redistribute Gaussian mass from under topography to open ocean
-    weights = weights.where(depth_levels < h)
-    weights = weights / weights.sum()
-
-    # Map 1D to 2D Gaussian
-    vertical_field = field * weights
-
-    # Remove NaNs at the edges
-    if orientation == "longitude":
-        vertical_field, _ = _remove_edge_nans(vertical_field, "lon", layer_depth=h)
-    if orientation == "latitude":
-        vertical_field, _ = _remove_edge_nans(vertical_field, "lat", layer_depth=h)
-
-    vertical_field = vertical_field.assign_coords({"depth": depth_levels})
-
-    return vertical_field
-
-
-def _plot_location(
-    grid: Grid,
-    releases: ReleaseCollector,
-    ax: Axes,
-    colors: dict[str, tuple] | None = None,
-    include_legend: bool = True,
-) -> None:
-    """Plot the center location of each release on a top-down map view.
-
-    Each release is represented as a point on the map, with its color
-    determined by the `colors` dictionary.
-
-    Parameters
-    ----------
-    grid : Grid
-        The grid object defining the spatial extent and coordinate system for the plot.
-
-    releases : ReleaseCollector
-        Collection of `Release` objects to plot. Each `Release` must have `.lat`, `.lon`,
-        and `.name` attributes.
-
-    ax : matplotlib.axes.Axes
-        The Matplotlib axis object to plot on.
-
-    colors : dict of str to tuple, optional
-        Optional dictionary mapping release names to RGBA color tuples. If not provided,
-        all releases are plotted in a default color (`"#dd1c77"`).
-
-    include_legend : bool, default True
-        Whether to include a legend showing release names.
-
-    Returns
-    -------
-    None
-    """
-
-    lon_deg = grid.ds.lon_rho
-    lat_deg = grid.ds.lat_rho
-    if grid.straddle:
-        lon_deg = xr.where(lon_deg > 180, lon_deg - 360, lon_deg)
-    trans = _get_projection(lon_deg, lat_deg)
+        # Compute layer thickness
+        depth_interface = compute_depth_coordinates(
+            grid.ds, zeta=0, depth_type="interface", location="rho"
+        )
+        dz = depth_interface.diff("s_w").rename({"s_w": "s_rho"})
 
-    proj = ccrs.PlateCarree()
+        # Compute vertical Gaussian shape
+        exponent = -(((depth - release.depth) / release.vsc) ** 2)
+        vertical_profile = np.exp(exponent)
 
-    for release in releases:
-        # transform coordinates to projected space
-        transformed_lon, transformed_lat = trans.transform_point(
-            release.lon,
-            release.lat,
-            proj,
-        )
+        # Apply vertical Gaussian scaling
+        distribution_3d = distribution_2d * vertical_profile * dz
 
-        if colors is not None:
-            color = colors[release.name]
-        else:
-            color = "k"
-
-        ax.plot(
-            transformed_lon,
-            transformed_lat,
-            marker="x",
-            markersize=8,
-            markeredgewidth=2,
-            label=release.name,
-            color=color,
-        )
+        # Normalize
+        distribution_3d /= release.vsc * np.sqrt(np.pi)
+        distribution_3d /= distribution_3d.sum()
 
-    if include_legend:
-        ax.legend(loc="center left", bbox_to_anchor=(1.1, 0.5))
+    return distribution_3d