generate_start0101_time_range func

hydroutils/hydro_time.py

@@ -1,7 +1,7 @@
 """
 Author: Wenyu Ouyang
 Date: 2022-12-02 11:03:04
-LastEditTime: 2024-08-15 11:21:53
+LastEditTime: 2024-09-14 13:57:36
 LastEditors: Wenyu Ouyang
 Description: some functions to deal with time
 FilePath: \hydroutils\hydroutils\hydro_time.py
@@ -11,11 +11,11 @@
 import datetime
 from typing import Union
 import numpy as np
+import pandas as pd
 import pytz
 import tzfpy
 
 
-
 def t2str(t_: Union[str, datetime.datetime]):
     if type(t_) is str:
         return datetime.datetime.strptime(t_, "%Y-%m-%d")
@@ -162,3 +162,44 @@ def calculate_utc_offset(lat, lng, date=None):
         if offset is not None:
             return int(offset.total_seconds() / 3600)
     return None
+
+
+def generate_start0101_time_range(start_time, end_time, freq="8D"):
+    """Generate a time range with a flexible start date and each year starting from 01-01.
+
+    Parameters
+    ----------
+    start_time : str or pd.Timestamp
+        The start time for the range (could be any date, string or Timestamp).
+    end_time : str or pd.Timestamp
+        The end time for the range (could be any date, string or Timestamp).
+    freq : str, optional
+        Time frequency for intervals, by default '8D'. Could be '7D', '10D', etc.
+
+    Returns
+    -------
+    pd.DatetimeIndex
+        A time range index with custom intervals and annual reset at 01-01.
+    """
+    all_dates = []
+
+    # Ensure the start and end times are of type pd.Timestamp
+    current_time = pd.Timestamp(start_time)
+    end_time = pd.Timestamp(end_time)
+
+    # Parse the frequency interval correctly
+    interval_days = pd.Timedelta(freq)  # Ensure it's a Timedelta
+
+    while current_time <= end_time:
+        all_dates.append(current_time)
+
+        # Calculate next date with the specified interval
+        next_time = current_time + interval_days
+
+        # If next_time crosses into a new year, reset to 01-01 of the new year
+        if next_time.year > current_time.year:
+            next_time = pd.Timestamp(f"{next_time.year}-01-01")
+
+        current_time = next_time
+
+    return pd.to_datetime(all_dates)

tests/test_hydro_time.py

@@ -0,0 +1,100 @@
+import pytest
+import pandas as pd
+from hydroutils.hydro_time import generate_start0101_time_range
+
+
+def test_generate_start0101_time_range_basic():
+    start_time = "2023-12-25"
+    end_time = "2024-01-10"
+    expected = pd.to_datetime(["2023-12-25", "2024-01-01", "2024-01-09"])
+    result = generate_start0101_time_range(start_time, end_time, freq="8D")
+    pd.testing.assert_index_equal(result, expected)
+
+
+def test_generate_start0101_time_range_multiple_years():
+    start_time = "2022-12-25"
+    end_time = "2024-01-10"
+    expected = pd.to_datetime(
+        [
+            "2022-12-25",
+            "2023-01-01",
+            "2023-01-09",
+            "2023-01-17",
+            "2023-01-25",
+            "2023-02-02",
+            "2023-02-10",
+            "2023-02-18",
+            "2023-02-26",
+            "2023-03-06",
+            "2023-03-14",
+            "2023-03-22",
+            "2023-03-30",
+            "2023-04-07",
+            "2023-04-15",
+            "2023-04-23",
+            "2023-05-01",
+            "2023-05-09",
+            "2023-05-17",
+            "2023-05-25",
+            "2023-06-02",
+            "2023-06-10",
+            "2023-06-18",
+            "2023-06-26",
+            "2023-07-04",
+            "2023-07-12",
+            "2023-07-20",
+            "2023-07-28",
+            "2023-08-05",
+            "2023-08-13",
+            "2023-08-21",
+            "2023-08-29",
+            "2023-09-06",
+            "2023-09-14",
+            "2023-09-22",
+            "2023-09-30",
+            "2023-10-08",
+            "2023-10-16",
+            "2023-10-24",
+            "2023-11-01",
+            "2023-11-09",
+            "2023-11-17",
+            "2023-11-25",
+            "2023-12-03",
+            "2023-12-11",
+            "2023-12-19",
+            "2023-12-27",
+            "2024-01-01",
+            "2024-01-09",
+        ]
+    )
+    result = generate_start0101_time_range(start_time, end_time, freq="8D")
+    pd.testing.assert_index_equal(result, expected)
+
+
+def test_generate_start0101_time_range_different_freq():
+    start_time = "2023-12-25"
+    end_time = "2024-01-10"
+    expected = pd.to_datetime(["2023-12-25", "2024-01-01", "2024-01-08"])
+    result = generate_start0101_time_range(start_time, end_time, freq="7D")
+    pd.testing.assert_index_equal(result, expected)
+
+
+def test_generate_start0101_time_range_no_reset():
+    start_time = "2023-12-25"
+    end_time = "2024-01-05"
+    expected = pd.to_datetime(["2023-12-25", "2024-01-01"])
+    result = generate_start0101_time_range(start_time, end_time, freq="8D")
+    pd.testing.assert_index_equal(result, expected)
+
+
+def test_generate_start0101_time_range_single_day():
+    start_time = "2023-01-01"
+    end_time = "2023-01-01"
+    expected = pd.to_datetime(["2023-01-01"])
+    result = generate_start0101_time_range(start_time, end_time, freq="8D")
+    pd.testing.assert_index_equal(result, expected)
+
+
+def test_generate_start0101_time_range_invalid_dates():
+    with pytest.raises(ValueError):
+        generate_start0101_time_range("invalid-date", "2024-01-10", freq="8D")