From 75632b2bb71169906d2d3d315a4e15613d66f8a8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ruben=20Calj=C3=A9?= <r.calje@artesia-water.nl>
Date: Fri, 15 Mar 2024 14:05:32 +0100
Subject: [PATCH] Add get_time_step_length

---
 nlmod/dims/time.py     | 25 +++++++++++++++++++++++++
 tests/test_016_time.py |  4 ++++
 2 files changed, 29 insertions(+)

diff --git a/nlmod/dims/time.py b/nlmod/dims/time.py
index f766354a..f1b21c17 100644
--- a/nlmod/dims/time.py
+++ b/nlmod/dims/time.py
@@ -408,6 +408,31 @@ def estimate_nstp(
         return nstp_ceiled
 
 
+def get_time_step_length(perlen, nstp, tsmult):
+    """
+    Get the length of the timesteps within a singe stress-period.
+
+    Parameters
+    ----------
+    perlen : float
+        The length of the stress period, in the time unit of the model (generally days).
+    nstp : int
+        The numer of timesteps within the stress period.
+    tsmult : float
+        THe time step multiplier, generally equal or lager than 1.
+
+    Returns
+    -------
+    t : np.ndarray
+        An array with the length of each of the timesteps within the stress period, in
+        the same unit as perlen.
+
+    """
+    t = np.array([tsmult**x for x in range(nstp)])
+    t = t * perlen / t.sum()
+    return t
+
+
 def ds_time_from_model(gwf):
     warnings.warn(
         "this function was renamed to `ds_time_idx_from_model`. "
diff --git a/tests/test_016_time.py b/tests/test_016_time.py
index 1e3339c5..056f7a75 100644
--- a/tests/test_016_time.py
+++ b/tests/test_016_time.py
@@ -28,3 +28,7 @@ def test_ds_time_from_tdis_settings():
 
     elapsed = (tidx.to_numpy() - np.datetime64("2000")) / np.timedelta64(1, "D")
     assert np.allclose(elapsed, [100, 150, 200, 233.33333333, 300.0])
+
+
+def test_get_time_step_length():
+    assert (nlmod.time.get_time_step_length(100, 2, 1.5) == np.array([40, 60])).all()