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

roms_tools/setup/cdr_forcing.py

@@ -2,19 +2,19 @@ import itertools
 import logging
 from collections import Counter
 from collections.abc import Iterator
-from datetime import datetime
+from datetime import datetime, timedelta
 from pathlib import Path
 from typing import Annotated
 
 import matplotlib.gridspec as gridspec
 import matplotlib.pyplot as plt
 import numpy as np
+import pandas as pd
 import xarray as xr
 from pydantic import (
     BaseModel,
     Field,
     RootModel,
-    conlist,
     model_serializer,
     model_validator,
 )
@@ -40,6 +40,7 @@ from roms_tools.setup.utils import (
     from_yaml,
     gc_dist,
     get_target_coords,
+    get_tracer_metadata_dict,
     to_dict,
     validate_names,
     write_to_yaml,
@@ -103,14 +104,16 @@ class ReleaseSimulationManager(BaseModel):
 class ReleaseCollector(RootModel):
     """Collects and validates multiple releases against each other."""
 
-    root: conlist(
-        Annotated[
-            VolumeRelease | TracerPerturbation, Field(discriminator="release_type")
+    root: Annotated[
+        list[
+            Annotated[
+                VolumeRelease | TracerPerturbation, Field(discriminator="release_type")
+            ]
         ],
-        min_length=1,
-    ) = Field(alias="releases")
+        Field(alias="releases", min_length=1),
+    ]
 
-    _release_type: ReleaseType = None
+    _release_type: ReleaseType | None = None
 
     def __iter__(self) -> Iterator[Release]:
         return iter(self.root)
@@ -126,6 +129,9 @@ class ReleaseCollector(RootModel):
         else:
             raise TypeError(f"Invalid key type: {type(item)}. Must be int or str.")
 
+    def __len__(self):
+        return len(self.root)
+
     @model_validator(mode="before")
     @classmethod
     def unpack_dict(cls, data):
@@ -774,6 +780,62 @@ class CDRForcing(BaseModel):
         fig.subplots_adjust(hspace=0.45)
         fig.suptitle(f"Release distribution for: {release_name}")
 
+    def compute_total_cdr_source(self, dt: float) -> pd.DataFrame:
+        """
+        Compute integrated tracer quantities for all releases and return a DataFrame.
+
+        Parameters
+        ----------
+        dt : float
+            Time step in seconds for reconstructing ROMS time stamps.
+
+        Returns
+        -------
+        pd.DataFrame
+            DataFrame with one row per release and one row of units at the top.
+            Columns 'temp' and 'salt' are excluded from integrated totals.
+        """
+        # Reconstruct ROMS time stamps
+        _, rel_seconds = _reconstruct_roms_time_stamps(
+            self.start_time, self.end_time, dt, self.model_reference_date
+        )
+
+        # Collect accounting results for all releases
+        records = []
+        release_names = []
+        for release in self.releases:
+            result = release._do_accounting(rel_seconds, self.model_reference_date)
+            records.append(result)
+            release_names.append(getattr(release, "name", f"release_{len(records)}"))
+
+        # Build DataFrame: rows = releases, columns = tracer names
+        df = pd.DataFrame(records, index=release_names)
+
+        # Exclude temp and salt from units row and integrated totals
+        integrated_tracers = [col for col in df.columns if col not in ("temp", "salt")]
+
+        # Add a row of units only for integrated tracers
+        tracer_meta = get_tracer_metadata_dict(include_bgc=True, unit_type="integrated")
+        units_row = {
+            col: tracer_meta.get(col, {}).get("units", "") for col in integrated_tracers
+        }
+
+        df_units = pd.DataFrame([units_row], index=["units"])
+
+        # Keep only integrated_tracers columns in df, drop temp and salt
+        df_integrated = df[integrated_tracers]
+
+        # Concatenate units row on top
+        df_final = pd.concat([df_units, df_integrated])
+
+        # Store dt as metadata
+        df_final.attrs["time_step"] = dt
+        df_final.attrs["start_time"] = self.start_time
+        df_final.attrs["end_time"] = self.end_time
+        df_final.attrs["title"] = "Integrated tracer releases"
+
+        return df_final
+
     def save(
         self,
         filepath: str | Path,
@@ -1039,21 +1101,67 @@ def _map_3d_gaussian(
         # Stack 2D distribution at that vertical level
         distribution_3d[{"s_rho": vertical_idx}] = distribution_2d
     else:
-        # 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"})
-
         # Compute vertical Gaussian shape
         exponent = -(((depth - release.depth) / release.vsc) ** 2)
         vertical_profile = np.exp(exponent)
 
         # Apply vertical Gaussian scaling
-        distribution_3d = distribution_2d * vertical_profile * dz
+        distribution_3d = distribution_2d * vertical_profile
 
         # Normalize
         distribution_3d /= release.vsc * np.sqrt(np.pi)
         distribution_3d /= distribution_3d.sum()
 
     return distribution_3d
+
+
+def _reconstruct_roms_time_stamps(
+    start_time: datetime,
+    end_time: datetime,
+    dt: float,
+    model_reference_date: datetime,
+) -> tuple[list[datetime], np.ndarray]:
+    """
+    Reconstruct ROMS time stamps between `start_time` and `end_time` with step `dt`.
+
+    Parameters
+    ----------
+    start_time : datetime
+        Beginning of the time series.
+    end_time : datetime
+        End of the time series (inclusive if it falls exactly on a step).
+    dt : float
+        Time step in seconds (can be fractional if needed).
+    model_reference_date : datetime
+        The reference date for ROMS time (elapsed time will be relative to this).
+
+    Returns
+    -------
+    times : list of datetime
+        Sequence of datetimes from `start_time` to `end_time`.
+    rel_days : np.ndarray
+        Array of elapsed times in **seconds** relative to `model_reference_date`.
+
+    Raises
+    ------
+    ValueError
+        If `end_time` is not after `start_time` or if `dt` is not positive.
+    """
+    if end_time <= start_time:
+        raise ValueError("end_time must be after start_time")
+    if dt <= 0:
+        raise ValueError("dt must be positive")
+
+    # Generate absolute times
+    delta = timedelta(seconds=dt)
+    times: list[datetime] = []
+    t = start_time
+    while t <= end_time:
+        times.append(t)
+        t += delta
+
+    # Convert to relative ROMS time (days since model_reference_date)
+    rel_days = convert_to_relative_days(times, model_reference_date)
+    rel_seconds = rel_days * 3600 * 24
+
+    return times, rel_seconds

roms_tools/setup/cdr_release.py

@@ -5,6 +5,8 @@ from datetime import datetime
 from enum import StrEnum, auto
 from typing import Annotated, Literal
 
+import numpy as np
+import pandas as pd
 from annotated_types import Ge, Le
 from pydantic import (
     BaseModel,
@@ -16,10 +18,17 @@ from pydantic import (
 )
 from pydantic_core.core_schema import ValidationInfo
 
-from roms_tools.setup.utils import get_tracer_defaults, get_tracer_metadata_dict
+from roms_tools.setup.utils import (
+    convert_to_relative_days,
+    get_tracer_defaults,
+    get_tracer_metadata_dict,
+)
 
 NonNegativeFloat = Annotated[float, Ge(0)]
 
+# Show all columns when printing a DataFrame
+pd.set_option("display.max_columns", None)
+
 
 @dataclass
 class ValueArray(ABC):
@@ -272,10 +281,68 @@ class Release(BaseModel):
             if self.times[-1] < end_time:
                 self.times.append(end_time)
 
-    @staticmethod
-    def get_tracer_metadata():
+    @classmethod
+    def get_tracer_metadata(cls):
         return {}
 
+    @classmethod
+    def get_metadata(cls):
+        return pd.DataFrame(cls.get_tracer_metadata())
+
+    def _compute_integrated_tracers(
+        self,
+        roms_time_stamps: np.ndarray,
+        model_reference_date: datetime,
+        tracer_series_dict: dict[str, np.ndarray],
+    ) -> dict[str, float]:
+        """
+        Compute time-integrated tracer quantities over ROMS time steps using a left-hold rule.
+
+        This method performs a left-hold (stepwise constant) integration of tracer fluxes
+        over the intervals defined by the ROMS time stamps. It first interpolates the
+        tracer time series from the release schedule onto the ROMS time stamps, then
+        multiplies the value at the start of each interval by the duration of that interval.
+
+        Parameters
+        ----------
+        roms_time_stamps : np.ndarray
+            1D array of ROMS model time stamps in seconds since `model_reference_date`.
+            Must be strictly increasing and contain at least two entries.
+        model_reference_date : datetime
+            Reference datetime of the ROMS model calendar, used to compute relative times
+            for interpolation.
+        tracer_series_dict : dict[str, np.ndarray]
+            Dictionary mapping tracer names to 1D arrays of tracer flux values at the
+            release schedule times (`self.times`). Each array must have the same length
+            as `self.times`.
+
+        Returns
+        -------
+        dict[str, float]
+            Dictionary mapping each tracer name to its integrated quantity over the
+            ROMS time period. Integration is performed using the left-hold rule,
+            ignoring the last release point because it defines the end of the final interval.
+
+        Raises
+        ------
+        ValueError
+            If `roms_time_stamps` has fewer than two entries, since at least one interval
+            is required for integration.
+        """
+        if len(roms_time_stamps) < 2:
+            raise ValueError("Need at least two ROMS time stamps to define intervals.")
+
+        dt = np.diff(roms_time_stamps)
+        results = {}
+        for tracer, series in tracer_series_dict.items():
+            interp_values = np.interp(
+                roms_time_stamps,
+                convert_to_relative_days(self.times, model_reference_date) * 3600 * 24,
+                series,
+            )
+            results[tracer] = np.sum(interp_values[:-1] * dt)
+        return results
+
 
 class VolumeRelease(Release):
     """Represents a CDR release with volume flux and tracer concentrations.
@@ -389,9 +456,11 @@ class VolumeRelease(Release):
         num_times = len(self.times)
 
         for tracer_concentrations in self.tracer_concentrations.values():
-            tracer_concentrations.check_length(num_times)
+            if isinstance(tracer_concentrations, Concentration):
+                tracer_concentrations.check_length(num_times)
 
-        self.volume_fluxes.check_length(num_times)
+        if isinstance(self.volume_fluxes, Flux):
+            self.volume_fluxes.check_length(num_times)
 
         return self
 
@@ -410,7 +479,52 @@ class VolumeRelease(Release):
     @staticmethod
     def get_tracer_metadata():
         """Returns long names and expected units for the tracer concentrations."""
-        return get_tracer_metadata_dict(include_bgc=True, with_flux_units=False)
+        return get_tracer_metadata_dict(include_bgc=True, unit_type="concentration")
+
+    def _do_accounting(
+        self,
+        roms_time_stamps: np.ndarray,
+        model_reference_date: datetime,
+    ) -> dict[str, float]:
+        """
+        Compute time-integrated tracer quantities over ROMS time steps.
+
+        This method interpolates tracer flux time series from the CDR schedule
+        onto the provided ROMS time stamps (in seconds since model reference date),
+        then applies a "left-hold" rule: the interpolated value at t₀ is applied
+        across the full interval [t₀, t₁).
+
+        Parameters
+        ----------
+        roms_time_stamps : np.ndarray
+            1D array of ROMS time stamps (seconds since `model_reference_date`).
+            Must be strictly increasing.
+        model_reference_date : datetime
+            Reference date of the ROMS model calendar.
+
+        Returns
+        -------
+        dict[str, float]
+            Dictionary mapping tracer names to the total integrated quantity over
+            the entire ROMS time period. Each value is the sum of the interpolated
+            tracer fluxes multiplied by the corresponding ROMS time step durations.
+        """
+        tracer_series_dict = {}
+        volume_array = (
+            np.asarray(self.volume_fluxes.values)
+            if isinstance(self.volume_fluxes, Flux)
+            else np.asarray(self.volume_fluxes)
+        )
+        for tracer, conc in self.tracer_concentrations.items():
+            tracer_array = (
+                np.asarray(conc.values)
+                if isinstance(conc, Concentration)
+                else np.asarray(conc)
+            )
+            tracer_series_dict[tracer] = volume_array * tracer_array
+        return self._compute_integrated_tracers(
+            roms_time_stamps, model_reference_date, tracer_series_dict
+        )
 
     @model_serializer(mode="wrap")
     def _simplified_dump(self, pydantic_serializer) -> dict:
@@ -503,7 +617,8 @@ class TracerPerturbation(Release):
     def _check_tracer_flux_lengths(self):
         num_times = len(self.times)
         for flux in self.tracer_fluxes.values():
-            flux.check_length(num_times)
+            if isinstance(flux, Flux):
+                flux.check_length(num_times)
         return self
 
     def _extend_to_endpoints(self, start_time, end_time):
@@ -520,7 +635,45 @@ class TracerPerturbation(Release):
     @staticmethod
     def get_tracer_metadata():
         """Returns long names and expected units for the tracer fluxes."""
-        return get_tracer_metadata_dict(include_bgc=True, with_flux_units=True)
+        return get_tracer_metadata_dict(include_bgc=True, unit_type="flux")
+
+    def _do_accounting(
+        self,
+        roms_time_stamps: np.ndarray,
+        model_reference_date: datetime,
+    ) -> dict[str, float]:
+        """
+        Compute time-integrated tracer quantities over ROMS time steps.
+
+        This method interpolates tracer flux time series from the CDR schedule
+        onto the provided ROMS time stamps (in days since model reference date),
+        then applies a "left-hold" rule: the interpolated value at t₀ is applied
+        across the full interval [t₀, t₁).
+
+        Parameters
+        ----------
+        roms_time_stamps : np.ndarray
+            1D array of ROMS time stamps (days since `model_reference_date`).
+            Must be strictly increasing.
+        model_reference_date : datetime
+            Reference date of the ROMS model calendar.
+
+        Returns
+        -------
+        dict[str, float]
+            Dictionary mapping tracer names to the total integrated quantity over
+            the entire ROMS time period. Each value is the sum of the interpolated
+            tracer fluxes multiplied by the corresponding ROMS time step durations.
+        """
+        tracer_series_dict = {
+            tracer: np.asarray(flux.values)
+            if isinstance(flux, Flux)
+            else np.asarray(flux)
+            for tracer, flux in self.tracer_fluxes.items()
+        }
+        return self._compute_integrated_tracers(
+            roms_time_stamps, model_reference_date, tracer_series_dict
+        )
 
     @model_serializer(mode="wrap")
     def _simplified_dump(self, pydantic_serializer) -> dict: