fix kerchunk JSON file filtering

- refactor to deduplicate code
- simplify how dates are passed around for the kerchunk functions
- and this way fix the date format mismatch when filtering the JSON files

Author: honzaflash

docs/whats_new.md

@@ -1,5 +1,9 @@
 # What's New
 
+## Unreleased
+
+* Fix kerchunk JSON file filtering.
+
 ## v0.12.1 (April 8, 2025)
 
 * Correction to `interpolator_filename` handling and improvement in testing.

particle_tracking_manager/config_ocean_model.py

@@ -175,21 +175,7 @@ def create_ocean_model_simulation(
     )
 
 
-def get_file_date_string(name: str, date: datetime) -> str:
-    """Get the file date string for the given ocean model name and date."""
-    if name == "NWGOA":
-        return f"{date.year}-{str(date.month).zfill(2)}-{str(date.day).zfill(2)}"
-    elif name == "CIOFSOP":
-        return f"{date.year}-{str(date.month).zfill(2)}-{str(date.day).zfill(2)}"
-    elif name == "CIOFS":
-        return f"{date.year}_{str(date.timetuple().tm_yday - 1).zfill(4)}"
-    elif name == "CIOFSFRESH":
-        return f"{date.year}_{str(date.timetuple().tm_yday - 1).zfill(4)}"
-    else:
-        raise ValueError(f"get_file_date_string not implemented for {name}.")
-
-
-function_map: dict[str, Callable[[str, str, str], dict[Any, Any]]] = {
+function_map: dict[str, Callable[[datetime, datetime, str], dict[Any, Any]]] = {
     "make_nwgoa_kerchunk": make_nwgoa_kerchunk,
     "make_ciofs_kerchunk": make_ciofs_kerchunk,
 }
@@ -208,9 +194,7 @@ def loc_local(
         start_time = start_sim + timedelta(days=1)
         end_time = end_sim - timedelta(days=1)
 
-    start = get_file_date_string(name, start_time)
-    end = get_file_date_string(name, end_time)
-    return function_map[kerchunk_func_str](start, end, name)
+    return function_map[kerchunk_func_str](start_time, end_time, name)
 
 
 def register_on_the_fly(ds_info: dict, ocean_model: str = "ONTHEFLY") -> None:

particle_tracking_manager/models/opendrift/utils.py

@@ -6,6 +6,7 @@
 # Standard library imports
 from datetime import datetime
 from pathlib import Path
+from typing import Callable, Iterable, List
 
 # Third-party imports
 import fsspec
@@ -96,15 +97,47 @@ def apply_user_input_ocean_model_specific_changes(
     return ds
 
 
-def make_ciofs_kerchunk(start: str, end: str, name: str) -> dict:
+def find_json_files_in_date_range(
+    fs2,
+    make_glob_from_year: Callable[[str], str],
+    start: datetime,
+    end: datetime,
+    filename_date_format: str,
+) -> List[str]:
+    """Find JSON files in a date range and return their paths."""
+
+    # only glob start and end year files, order isn't important
+    if abs(start.year - end.year) > 1:
+        raise ValueError(
+            f"Start ({start.year}) and end ({end.year}) "
+            "dates must be at most 1 year apart."
+        )
+    start_year = start.strftime("%Y")
+    end_year = end.strftime("%Y")
+    json_list = fs2.glob(make_glob_from_year(start_year))
+    if end_year != start_year:
+        json_list += fs2.glob(make_glob_from_year(end_year))
+
+    def filter_paths(start: datetime, end: datetime, paths: Iterable[str]) -> List[str]:
+        if end < start:
+            # if going backward in time, swap start and end
+            end, start = start, end
+        return [
+            pth
+            for pth in paths
+            if start <= datetime.strptime(Path(pth).stem, filename_date_format) <= end
+        ]
+
+    return filter_paths(start, end, json_list)
+
+
+def make_ciofs_kerchunk(start: datetime, end: datetime, name: str) -> dict:
     """_summary_
 
     Parameters
     ----------
-    start, end : str
-        Should be something like "2004_0001" for YYYY_0DDD where DDD is dayofyear
-        to match the files in the directory, which are by year and day of year
-        ("ciofs_fresh" or "ciofs") or "YYYY-MM-DD" for "aws_ciofs"
+    start, end : datetime
+        Start and end time of the simulation.
 
     Returns
     -------
@@ -123,50 +156,21 @@ def make_ciofs_kerchunk(start: str, end: str, name: str) -> dict:
 
     fs2 = fsspec.filesystem("")  # local file system to save final jsons to
 
-    if name in ["CIOFS", "CIOFSFRESH"]:
-
+    if name == "CIOFSOP":
         # base for matching
         def base_str(a_time: str) -> str:
-            return f"{output_dir_single_files}/{a_time}_*.json"
-
-        date_format = "%Y_0%j"
+            return f"{output_dir_single_files}/ciofs_{a_time}-*.json"
 
-    elif name == "CIOFSOP":
+        date_format = "ciofs_%Y-%m-%d"
 
+    else:  # name is "CIOFS" or "CIOFSFRESH"
         # base for matching
         def base_str(a_time: str) -> str:
-            return f"{output_dir_single_files}/ciofs_{a_time}-*.json"
-
-        date_format = "ciofs_%Y-%m-%d"
-    else:
-        raise ValueError(f"Name {name} not recognized")
+            return f"{output_dir_single_files}/{a_time}_*.json"
 
-    # only glob start and end year files, order isn't important
-    json_list = fs2.glob(base_str(start[:4]))
-    if end[:4] != start[:4]:
-        json_list += fs2.glob(base_str(end[:4]))
-
-    # forward in time
-    if end > start:
-        json_list = [
-            j
-            for j in json_list
-            if datetime.strptime(Path(j).stem, date_format).isoformat() >= start
-            and datetime.strptime(Path(j).stem, date_format).isoformat() <= end
-        ]
-    # backward in time
-    elif end < start:
-        json_list = [
-            j
-            for j in json_list
-            if datetime.strptime(Path(j).stem, date_format).isoformat() <= start
-            and datetime.strptime(Path(j).stem, date_format).isoformat() >= end
-        ]
+        date_format = "%Y_0%j"
 
-    if json_list == []:
-        raise ValueError(
-            f"No files found in {output_dir_single_files} for {start} to {end}"
-        )
+    json_list = find_json_files_in_date_range(fs2, base_str, start, end, date_format)
 
     # Multi-file JSONs
     # This code uses the output generated above to create a single ensemble dataset,
@@ -280,13 +284,13 @@ def postprocess(out: dict) -> dict:
     return out
 
 
-def make_nwgoa_kerchunk(start: str, end: str, name: str = "NWGOA") -> dict:
+def make_nwgoa_kerchunk(start: datetime, end: datetime, name: str = "NWGOA") -> dict:
     """_summary_
 
     Parameters
     ----------
-    start, end : str
-        Should be something like "1999-01-02" for YYYY-MM-DD
+    start, end : datetime
+        Start and end time of the simulation.
 
     Returns
     -------
@@ -308,33 +312,7 @@ def base_str(a_time: str) -> str:
 
     date_format = "nwgoa_%Y-%m-%d"
 
-    # only glob start and end year files, order isn't important
-    json_list = fs2.glob(base_str(start[:4]))
-
-    if end[:4] != start[:4]:
-        json_list += fs2.glob(base_str(end[:4]))
-
-    # forward in time
-    if end > start:
-        json_list = [
-            j
-            for j in json_list
-            if datetime.strptime(Path(j).stem, date_format).isoformat() >= start
-            and datetime.strptime(Path(j).stem, date_format).isoformat() <= end
-        ]
-    # backward in time
-    elif end < start:
-        json_list = [
-            j
-            for j in json_list
-            if datetime.strptime(Path(j).stem, date_format).isoformat() <= start
-            and datetime.strptime(Path(j).stem, date_format).isoformat() >= end
-        ]
-
-    if json_list == []:
-        raise ValueError(
-            f"No files found in {output_dir_single_files} for {start} to {end}"
-        )
+    json_list = find_json_files_in_date_range(fs2, base_str, start, end, date_format)
 
     # account for double compression
     # Look at individual variables in the files to see what needs to be changed with