add assign time start end func

hydroutils/hydro_plot.py

@@ -1,7 +1,7 @@
 """
 Author: Wenyu Ouyang
 Date: 2022-12-02 10:59:30
-LastEditTime: 2025-01-14 09:06:45
+LastEditTime: 2025-01-15 16:55:43
 LastEditors: Wenyu Ouyang
 Description: Some common plots for hydrology
 FilePath: \hydroutils\hydroutils\hydro_plot.py
@@ -1440,3 +1440,57 @@ def plot_rainfall_runoff_xu(
     plt.title(title)
 
     plt.legend(loc="upper left")
+
+
+def plot_rainfall_runoff_chai(
+    t,
+    ps,
+    qs,
+    c_lst="rbkgcmy",
+    title="Observation of Precipitation and Streamflow",
+    alpha_lst=None,
+    p_labels=None,
+    q_labels=None,
+):
+    if alpha_lst is None:
+        alpha_lst = [0.5, 0.5]
+    if p_labels is None:
+        p_labels = ["era5land", "gauge"]
+    if q_labels is None:
+        q_labels = ["observation", "simulation"]
+    fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(20, 8))
+    fig.suptitle(title, fontsize=16)
+    for i, p in enumerate(ps):
+        ax1.bar(
+            t,
+            p,
+            color=c_lst[i],
+            label=p_labels[i],
+            width=0.9,
+            alpha=alpha_lst[i],
+        )
+    ax1.set_xlabel("Time")
+    ax1.set_ylabel("Precipitation (mm/d)", color="b")
+    ax1.invert_yaxis()
+    ax1.tick_params(axis="y", labelcolor="b")
+    ax1.legend()
+
+    for j, q in enumerate(qs):
+        ax2.plot(t, q, color=c_lst[j], label=q_labels[j])
+    ax2.set_xlabel("Time")
+    ax2.set_ylabel("Streamflow (m$^3$/s)", color="r")
+    ax2.tick_params(axis="y", labelcolor="r")
+    ax2.set_xlim(ax1.get_xlim())
+    ax2.text(
+        0.05,
+        0.95,
+        "Streamflow",
+        transform=ax2.transAxes,
+        fontsize=12,
+        verticalalignment="top",
+        bbox=dict(facecolor="white", alpha=0.5),
+    )
+    ax2.legend()
+
+    fig.tight_layout()
+    return fig, ax2

hydroutils/hydro_time.py

@@ -203,3 +203,16 @@ def generate_start0101_time_range(start_time, end_time, freq="8D"):
         current_time = next_time
 
     return pd.to_datetime(all_dates)
+
+
+def assign_time_start_end(time_ranges, assign_way="intersection"):
+    # chose all time start from time_ranges:[[start1, end1], [start2, end2], ...]
+    if assign_way == "intersection":
+        time_start = max(t[0] for t in time_ranges)
+        time_end = min(t[1] for t in time_ranges)
+    elif assign_way == "union":
+        time_start = min(t[0] for t in time_ranges)
+        time_end = max(t[1] for t in time_ranges)
+    else:
+        raise NotImplementedError("We don't support this assign_way yet")
+    return time_start, time_end

tests/test_hydro_time.py

@@ -1,6 +1,7 @@
 import pytest
 import pandas as pd
 from hydroutils.hydro_time import generate_start0101_time_range
+from hydroutils.hydro_time import assign_time_start_end
 
 
 def test_generate_start0101_time_range_basic():
@@ -98,3 +99,37 @@ def test_generate_start0101_time_range_single_day():
 def test_generate_start0101_time_range_invalid_dates():
     with pytest.raises(ValueError):
         generate_start0101_time_range("invalid-date", "2024-01-10", freq="8D")
+
+
+def test_assign_time_start_end_intersection():
+    time_ranges = [["2023-01-01", "2023-12-31"], ["2023-06-01", "2023-12-01"]]
+    expected = ("2023-06-01", "2023-12-01")
+    result = assign_time_start_end(time_ranges, assign_way="intersection")
+    assert result == expected
+
+
+def test_assign_time_start_end_union():
+    time_ranges = [["2023-01-01", "2023-12-31"], ["2023-06-01", "2023-12-01"]]
+    expected = ("2023-01-01", "2023-12-31")
+    result = assign_time_start_end(time_ranges, assign_way="union")
+    assert result == expected
+
+
+def test_assign_time_start_end_single_range():
+    time_ranges = [["2023-01-01", "2023-12-31"]]
+    expected = ("2023-01-01", "2023-12-31")
+    result = assign_time_start_end(time_ranges, assign_way="intersection")
+    assert result == expected
+
+
+def test_assign_time_start_end_no_overlap():
+    time_ranges = [["2023-01-01", "2023-06-01"], ["2023-07-01", "2023-12-01"]]
+    expected = ("2023-07-01", "2023-06-01")
+    result = assign_time_start_end(time_ranges, assign_way="intersection")
+    assert result == expected
+
+
+def test_assign_time_start_end_invalid_assign_way():
+    time_ranges = [["2023-01-01", "2023-12-31"], ["2023-06-01", "2023-12-01"]]
+    with pytest.raises(NotImplementedError):
+        assign_time_start_end(time_ranges, assign_way="invalid")