From 1fdd00e4d0eb5f3c024e32ae3b0813239af5cddb Mon Sep 17 00:00:00 2001
From: Patrik Bohlinger <patrikb@met.no>
Date: Thu, 29 Aug 2024 13:18:22 +0200
Subject: [PATCH] added fix such that the use of a maximum lead time for models
 (max_lt) works and can loosely be added as part of kwargs

---
 wavy/collocation_module.py | 21 ++++++++++++---------
 wavy/model_module.py       | 17 ++++++++---------
 2 files changed, 20 insertions(+), 18 deletions(-)

diff --git a/wavy/collocation_module.py b/wavy/collocation_module.py
index 2c38ceb9..1aa5456e 100755
--- a/wavy/collocation_module.py
+++ b/wavy/collocation_module.py
@@ -64,12 +64,14 @@ def collocation_fct(obs_lons, obs_lats, model_lons, model_lats):
     return index_array_2d, distance_array, valid_output_index
 
 
-def get_model_filename(nID, d, leadtime):
+def get_model_filename(nID, d, leadtime, **kwargs):
     mco = mc(nID=nID, sd=d, ed=d, leadtime=leadtime)
-    return mco._make_model_filename_wrapper(parse_date(str(d)), leadtime)
+    return mco._make_model_filename_wrapper(parse_date(str(d)),
+                                            leadtime, **kwargs)
 
 
-def find_valid_fc_dates_for_model_and_leadtime(fc_dates, model, leadtime):
+def find_valid_fc_dates_for_model_and_leadtime(fc_dates, model,
+                                               leadtime, **kwargs):
     '''
     Finds valid dates that are close to desired dates at a precision
     of complete hours
@@ -84,13 +86,13 @@ def find_valid_fc_dates_for_model_and_leadtime(fc_dates, model, leadtime):
     #    pass
     #else:
     fc_dates_new = [d for d in fc_dates_new
-                    if get_model_filename(model, d, leadtime) is not None]
+                    if get_model_filename(model, d, leadtime, **kwargs)
+                    is not None]
     return fc_dates_new
 
 
-def check_if_file_is_valid(fc_date, model, leadtime, max_lt=None):
-    fname = get_model_filename(
-                model, fc_date, leadtime, max_lt=max_lt)
+def check_if_file_is_valid(fc_date, model, leadtime, **kwargs):
+    fname = get_model_filename(model, fc_date, leadtime, **kwargs)
     print('Check if requested file:\n', fname, '\nis available and valid')
     try:
         nc = netCDF4.Dataset(fname, mode='r')
@@ -133,7 +135,7 @@ class collocation_class(qls):
     '''
 
     def __init__(self, oco=None, model=None, poi=None,
-    distlim=None, leadtime=None, max_lt=None, **kwargs):
+    distlim=None, leadtime=None, **kwargs):
         print('# ----- ')
         print(" ### Initializing collocation_class object ###")
         print(" ")
@@ -312,7 +314,8 @@ def _collocate_track(self, **kwargs):
         ndt_datetime = [parse_date(str(d)) for d in ndt]
 
         ndt_valid = find_valid_fc_dates_for_model_and_leadtime(
-                                    ndt, self.model, self.leadtime)
+                                    ndt, self.model, self.leadtime,
+                                    **kwargs)
 
         ndt_valid = np.unique(ndt_valid)
 
diff --git a/wavy/model_module.py b/wavy/model_module.py
index 92187ba4..a67637ac 100755
--- a/wavy/model_module.py
+++ b/wavy/model_module.py
@@ -319,8 +319,7 @@ def _make_model_filename(self, fc_date, leadtime):
                            .replace(">", "\\>")
         return filename
 
-    def _make_model_filename_wrapper(
-    self, fc_date, leadtime, max_lt=None, **kwargs):
+    def _make_model_filename_wrapper(self, fc_date, leadtime, **kwargs):
         """
         Wrapper function of make_model_filename. Organizes various cases.
 
@@ -328,7 +327,6 @@ def _make_model_filename_wrapper(
             model - modelname type(str)
             fc_date - datetime object
             leadtime - integer in hours
-            max_lt - maximum lead time allowed
 
         return:
             filename
@@ -362,10 +360,12 @@ def _make_model_filename_wrapper(
                               + " with extended leadtime")
                         leadtime = (leadtime
                                     + vars(self.cfg)['misc']['init_step'])
-                    if max_lt is not None and leadtime > max_lt:
+                    if (kwargs.get('max_lt') is not None
+                        and leadtime > kwargs.get('max_lt')):
                         print("Leadtime:", leadtime,
                               "is greater as maximum allowed leadtime:",
-                              max_lt)
+                              str(kwargs.get('max_lt')))
+                        break
             else:
                 filename = None
         elif (isinstance(fc_date, list) and isinstance(leadtime, int)):
@@ -380,15 +380,14 @@ def _make_model_filename_wrapper(
 
         return filename
 
-    def _make_list_of_model_filenames(self, fc_dates, lt):
+    def _make_list_of_model_filenames(self, fc_dates, lt, **kwargs):
         """
         return: flst - list of model files to be opened
                 dlst - list of dates to be chosen within each file
         """
-        #fn = make_model_filename_wrapper(datetime(2021,1,1,1),1)
         flst = []
         for d in fc_dates:
-            fn = self._make_model_filename_wrapper(d, lt)
+            fn = self._make_model_filename_wrapper(d, lt, **kwargs)
             if fn is not None:
                 flst.append(fn)
         return flst
@@ -493,7 +492,7 @@ def list_input_files(self, show=False, **kwargs):
                                      self.cfg.misc['date_incr'])
 
             pathlst = self._make_list_of_model_filenames(
-                    fc_dates, self.leadtime)
+                    fc_dates, self.leadtime, **kwargs)
         else:
             # if defined path local
             print(" ## Find and list files ...")