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

roms_tools/setup/cdr_forcing.py

@@ -1,12 +1,12 @@
 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
@@ -22,9 +22,7 @@ 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.plot import get_projection, plot, plot_2d_horizontal_field
 from roms_tools.setup.cdr_release import (
     Release,
     ReleaseType,
@@ -32,27 +30,27 @@ 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,
+    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"
 
 
 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 +62,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 +123,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 +161,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 +169,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 +240,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 +332,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 +341,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,10 +380,12 @@ 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(
@@ -394,7 +411,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."
@@ -597,7 +613,6 @@ 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(
@@ -615,17 +630,19 @@ class CDRForcing(BaseModel):
         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)
@@ -676,73 +693,65 @@ 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
             )
 
-        # 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 +772,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 +788,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 +799,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,7 +822,7 @@ 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)
@@ -851,7 +858,6 @@ def _validate_release_input(releases, valid_releases, list_allowed=True):
     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.
     """
-
     # 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(
@@ -904,7 +910,6 @@ def _get_release_colors(valid_releases: list[str]) -> dict[str, tuple]:
     - 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.
     """
-
     # Determine the colormap based on the number of releases
     if len(valid_releases) <= 10:
         color_map = plt.get_cmap("tab10")
@@ -944,13 +949,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 +1027,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,102 +1065,47 @@ 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()
+        # 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"})
 
-    # Map 1D to 2D Gaussian
-    vertical_field = field * weights
+        # Compute vertical Gaussian shape
+        exponent = -(((depth - release.depth) / release.vsc) ** 2)
+        vertical_profile = np.exp(exponent)
 
-    # 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)
+        # Apply vertical Gaussian scaling
+        distribution_3d = distribution_2d * vertical_profile * dz
 
-    vertical_field = vertical_field.assign_coords({"depth": depth_levels})
+        # Normalize
+        distribution_3d /= release.vsc * np.sqrt(np.pi)
+        distribution_3d /= distribution_3d.sum()
 
-    return vertical_field
+    return distribution_3d
 
 
 def _plot_location(
@@ -1251,12 +1143,11 @@ def _plot_location(
     -------
     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)
+    trans = get_projection(lon_deg, lat_deg)
 
     proj = ccrs.PlateCarree()