roms-tools 3.1.2__py3-none-any.whl → 3.3.0__py3-none-any.whl

roms_tools/tests/test_setup/test_cdr_release.py

@@ -1,6 +1,7 @@
-from datetime import datetime
+from datetime import datetime, timedelta
 from unittest import mock
 
+import numpy as np
 import pytest
 from pydantic import ValidationError
 
@@ -377,3 +378,119 @@ class TestTracerPerturbation:
     def test_get_tracer_metadata(self):
         d = VolumeRelease.get_tracer_metadata()
         assert len(d) == NUM_TRACERS
+
+
+class TestReleaseAccounting:
+    def setup_method(self):
+        self.start = datetime(2020, 1, 1)
+        self.end = datetime(2020, 1, 11)  # 10 days later
+
+    def make_release(self, release_type, fluxes=None, concentrations=None, times=None):
+        times = times or [self.start, self.end]
+
+        if release_type == "volume":
+            vr = VolumeRelease(
+                name="test",
+                lat=0.0,
+                lon=0.0,
+                depth=100.0,
+                times=times,
+                volume_fluxes=Flux("volume", fluxes),
+                tracer_concentrations={"DIC": Concentration("DIC", concentrations)},
+            )
+            vr._extend_to_endpoints(self.start, self.end)
+            return vr
+
+        elif release_type == "tracer":
+            tp = TracerPerturbation(
+                name="test",
+                lat=0.0,
+                lon=0.0,
+                depth=100.0,
+                times=times,
+                tracer_fluxes={"DIC": Flux("DIC", fluxes)},
+            )
+            tp._extend_to_endpoints(self.start, self.end)
+            return tp
+        else:
+            raise ValueError(f"Unknown release type {release_type}")
+
+    @pytest.mark.parametrize("release_type", ["volume", "tracer"])
+    def test_constant_flux(self, release_type):
+        """Case 0: constant fluxes."""
+        flux = 2.0
+        conc = 3.0  # only used for VolumeRelease
+        vr = self.make_release(release_type, fluxes=flux, concentrations=conc)
+        roms_stamps = np.array([0.0, 5.0, 10.0])
+        result = vr._do_accounting(roms_stamps, self.start)
+
+        expected = flux * (roms_stamps[-1] - roms_stamps[0])
+        if release_type == "volume":
+            expected *= conc
+
+        assert result["DIC"] == pytest.approx(expected)
+
+    @pytest.mark.parametrize("release_type", ["volume", "tracer"])
+    def test_aligned_releases(self, release_type):
+        """Case 1: release times exactly aligned with ROMS stamps."""
+        # ROMS time step: 5 days (we want to cover the simulation time of 10 days exactly)
+        roms = np.array([0.0, 5.0, 10.0]) * 24 * 3600  # seconds
+        release_times = [self.start + timedelta(seconds=s) for s in roms]
+        fluxes = [1, 3, 5]
+        concs = [1, 2, 1]
+
+        if release_type == "volume":
+            r = self.make_release(
+                release_type, fluxes=fluxes, concentrations=concs, times=release_times
+            )
+        elif release_type == "tracer":
+            r = self.make_release(release_type, fluxes=fluxes, times=release_times)
+
+        # Expected calculation
+        series = (
+            np.array([f * c for f, c in zip(fluxes, concs)])
+            if release_type == "volume"
+            else np.array(fluxes)
+        )
+        dt = np.diff(roms)
+        expected = np.sum(series[:-1] * dt)
+
+        result = r._do_accounting(roms, self.start)
+        assert result["DIC"] == pytest.approx(expected)
+
+    @pytest.mark.parametrize("release_type", ["volume", "tracer"])
+    def test_unaligned_releases(self, release_type):
+        """Case 2: release times unaligned with ROMS stamps."""
+        # ROMS time step: 5 days (we want to cover the simulation time of 10 days exactly)
+        roms = np.array([0.0, 2.5, 5.0, 7.5, 10.0]) * 24 * 3600  # seconds
+        rel_release_times = np.array([0.0, 1.0, 3.0, 6.0, 10.0]) * 24 * 3600  # seconds
+        release_times = [self.start + timedelta(seconds=s) for s in rel_release_times]
+        fluxes = [1, 2, 3, 4, 5]
+        concs = [1, 1.5, 2, 2.5, 1]
+
+        if release_type == "volume":
+            vr = self.make_release(
+                release_type, fluxes=fluxes, concentrations=concs, times=release_times
+            )
+        elif release_type == "tracer":
+            vr = self.make_release(release_type, fluxes=fluxes, times=release_times)
+
+        # Expected calculation
+        series = (
+            np.array([f * c for f, c in zip(fluxes, concs)])
+            if release_type == "volume"
+            else np.array(fluxes)
+        )
+        interp = np.interp(roms, rel_release_times, series)
+        dt = np.diff(roms)
+        expected = np.sum(interp[:-1] * dt)
+
+        result = vr._do_accounting(roms, self.start)
+        assert result["DIC"] == pytest.approx(expected)
+
+    @pytest.mark.parametrize("release_type", ["volume", "tracer"])
+    def test_raises_with_single_roms_timestamp(self, release_type):
+        vr = self.make_release(release_type, fluxes=1.0, concentrations=1.0)
+        roms_stamps = np.array([0.0])
+        with pytest.raises(ValueError, match="at least two ROMS time stamps"):
+            vr._do_accounting(roms_stamps, self.start)