Split get_background into 3d and 4d tasks

src/swell/tasks/get_background.py

@@ -8,11 +8,12 @@
 # --------------------------------------------------------------------------------------------------
 
 
-from swell.tasks.base.task_base import taskBase
+import os
+import importlib
 
 import isodate
-import os
-from r2d2 import fetch
+from swell.tasks.base.task_base import taskBase
+#from r2d2 import fetch
 
 
 # --------------------------------------------------------------------------------------------------
@@ -27,6 +28,56 @@
 
 class GetBackground(taskBase):
 
+    def _fetch_all(self, background_prep, r2d2_dict, model_component):
+
+        # Loop over fc
+        # ------------
+        for fc in r2d2_dict['fetch']['fc']:
+
+            # Reset target file
+            # --------------------
+            file_type = fc['file_type']
+            target_file_template = fc['filename']
+
+            # Loop over background steps
+            # --------------------
+            for bkg_step in background_prep.bkg_steps:
+
+                # Set the datetime format for the output files
+                # --------------------------------------------
+                background_time = background_prep.forecast_start_time +\
+                                  isodate.parse_duration(bkg_step)
+
+                # Set the datetime templating in the target file name
+                # ---------------------------------------------------
+                target_file = background_time.strftime(target_file_template)
+
+                args = {'data': str(background_prep.forecast_start_time),
+                        'target_file': str(target_file),
+                        'model': str(r2d2_model_dict[model_component]),
+                        'file_type': str(file_type),
+                        'fc_date_rendering': 'analysis',
+                        'step': str(bkg_step),
+                        'resolution': str(background_prep.horizontal_resolution),
+                        'type': 'fc',
+                        'experiment': str(background_prep.background_experiment)}
+
+                print(str(args), '\n')
+
+                # fetch(
+                #     date=background_prep.forecast_start_time,
+                #     target_file=target_file,
+                #     model=r2d2_model_dict[model_component],
+                #     file_type=file_type,
+                #     fc_date_rendering='analysis',
+                #     step=bkg_step,
+                #     resolution=horizontal_resolution,
+                #     type='fc',
+                #     experiment=background_experiment)
+                #
+                # # Change permission
+                # os.chmod(target_file, 0o644)
+
     def execute(self):
 
         """Acquires background files for a given experiment and cycle
@@ -37,120 +88,35 @@ def execute(self):
              See the taskBase constructor for more information.
         """
 
-        # Get duration into forecast for first background file
-        # ----------------------------------------------------
-        bkg_steps = []
-
-        # Parse config
-        background_experiment = self.config.background_experiment()
-        background_frequency = self.config.background_frequency(None)
-        forecast_offset = self.config.analysis_forecast_window_offset()
-        horizontal_resolution = self.config.horizontal_resolution()
-        window_length = self.config.window_length()
-        window_offset = self.config.window_offset()
         window_type = self.config.window_type()
 
+        # Delegate a "prep background" object.
+        prep_background_module = importlib.import_module(".prep_background", package=__package__)
+        prep_background_class = getattr(prep_background_module, f"PrepBackground{window_type}") # {window_type}")
+
+        # Initialize background_prep object
+        background_prep = prep_background_class.fromConfig(self.config)
+
         # Get window parameters
-        local_background_time = self.da_window_params.local_background_time(window_offset,
-                                                                            window_type)
+        local_background_time =\
+            self.da_window_params.local_background_time(background_prep.window_offset,
+                                                        window_type)
 
         # Add to jedi config rendering dictionary
         self.jedi_rendering.add_key('local_background_time', local_background_time)
 
-        # Convert to datetime durations
-        # -----------------------------
-        window_length_dur = isodate.parse_duration(window_length)
-        forecast_offset_dur = isodate.parse_duration(forecast_offset)
-
-        # Duration between the start of the forecast that generated the background
-        # and the middle of the current window
-        # -------------------------------------------------------------------------------
-        forecast_duration_for_background = window_length_dur - forecast_offset_dur
-
-        # If the window type is 4D then remove the window offset as first background
-        # occurs at the beginning of the window
-        # -------------------------------------------------------------------------------
-        if window_type == "4D":
-            window_offset_dur = isodate.parse_duration(window_offset)
-            forecast_duration_for_background = forecast_duration_for_background - window_offset_dur
-
-        # Append the list of backgrounds to get with the first background
-        # -----------------------------------------------------------------
-        bkg_steps.append(isodate.duration_isoformat(forecast_duration_for_background))
-
-        # If background is provided though files get all backgrounds
-        # ----------------------------------------------------------
-        if window_type == "4D":
-
-            bkg_freq_dur = isodate.parse_duration(background_frequency)
-
-            # Check for a sensible frequency
-            # ------------------------------
-            if (window_length_dur/bkg_freq_dur) % 2:
-                self.logger.abort('Window length not divisible by background frequency')
-
-            # Loop over window
-            print('self.cycle_time_dto()', self.cycle_time_dto())
-            print('window_offset_dur', window_offset_dur)
-
-            start_date = self.cycle_time_dto() - window_offset_dur
-            final_date = self.cycle_time_dto() + window_offset_dur
-
-            loop_date = start_date + bkg_freq_dur
-
-            while loop_date <= final_date:
-                duration_in = loop_date - start_date + forecast_duration_for_background
-                bkg_steps.append(isodate.duration_isoformat(duration_in))
-                loop_date += bkg_freq_dur
-
-        # Get the forecast start time
-        # ---------------------------
-        forecast_start_time = self.cycle_time_dto() - window_length_dur + forecast_offset_dur
+        # Prepare all background parameters
+        background_prep.prep(cycle_time_dto = self.cycle_time_dto())
 
         # Get name of this model component
-        # --------------------------------
         model_component = self.get_model()
 
-        # Loop over background files in the R2D2 config and fetch
-        # -------------------------------------------------------
-        self.logger.info('Background steps being fetched: '+' '.join(str(e) for e in bkg_steps))
+        self.logger.info('Background steps being fetched: ' +\
+                         ' '.join(str(e) for e in background_prep.bkg_steps))
 
         # Get r2d2 dictionary
         r2d2_dict = self.jedi_rendering.render_interface_model('r2d2')
 
-        # Loop over fc
-        # ------------
-        for fc in r2d2_dict['fetch']['fc']:
-
-            # Reset target file
-            # --------------------
-            file_type = fc['file_type']
-            target_file_template = fc['filename']
-
-            # Loop over background steps
-            # --------------------
-            for bkg_step in bkg_steps:
-
-                # Set the datetime format for the output files
-                # --------------------------------------------
-                background_time = forecast_start_time + isodate.parse_duration(bkg_step)
-
-                # Set the datetime templating in the target file name
-                # ---------------------------------------------------
-                target_file = background_time.strftime(target_file_template)
-
-                fetch(
-                    date=forecast_start_time,
-                    target_file=target_file,
-                    model=r2d2_model_dict[model_component],
-                    file_type=file_type,
-                    fc_date_rendering='analysis',
-                    step=bkg_step,
-                    resolution=horizontal_resolution,
-                    type='fc',
-                    experiment=background_experiment)
-
-                # Change permission
-                os.chmod(target_file, 0o644)
+        # Loop over background files in the R2D2 config and fetch
+        self._fetch_all(background_prep, r2d2_dict, model_component)
 
-# --------------------------------------------------------------------------------------------------

src/swell/tasks/prep_background.py

@@ -0,0 +1,140 @@
+#!/usr/bin/env python
+
+# (C) Copyright 2021- United States Government as represented by the Administrator of the
+# National Aeronautics and Space Administration. All Rights Reserved.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+
+# -------------------------------------------------------------------------------------------------
+
+
+from abc import ABC, abstractmethod
+import isodate
+
+# --------------------------------------------------------------------------------------------------
+# Prep Background superclass
+# --------------------------------------------------------------------------------------------------
+class PrepBackground(ABC):
+
+    def __init__(self, config):
+
+        # Delegate config object
+        self.config = config
+
+        # Parse a few config elements
+        self.background_experiment = config.background_experiment()
+        self.background_frequency = config.background_frequency(None)
+        self.forecast_offset = config.analysis_forecast_window_offset()
+        self.horizontal_resolution = config.horizontal_resolution()
+        self.window_length = config.window_length()
+        self.window_offset = config.window_offset()
+
+        # Duration parameters
+        self.window_length_dur = None
+        self.forecast_offset_dur = None
+        self.forecast_duration_for_background = None
+
+        # Offset parameters
+        self.window_offset_dur = None
+
+        # Forecast start time
+        self.forecast_start_time = None
+
+        # Initialize a list of background time steps
+        self.bkg_steps = []
+
+    @classmethod
+    def fromConfig(cls, config):
+        return cls(config)
+
+    def prep(self, cycle_time_dto):
+
+        # Reset bkg_steps to empty list if necessary
+        self.bkg_steps = []
+
+        self._parse_dur()
+        self._offset()
+        self._add_bkg_steps(cycle_time_dto)
+        self._set_forecast_start_time(cycle_time_dto)
+
+    def _parse_dur(self):
+
+        # Convert to datetime durations
+        self.window_length_dur = isodate.parse_duration(self.window_length)
+        self.forecast_offset_dur = isodate.parse_duration(self.forecast_offset)
+        self.forecast_duration_for_background = self.window_length_dur - self.forecast_offset_dur
+
+    def _set_forecast_start_time(self, cycle_time_dto):
+        self.forecast_start_time = cycle_time_dto - self.window_length_dur + self.forecast_offset_dur
+
+    @abstractmethod
+    def _offset(self):
+        pass
+
+    @abstractmethod
+    def _add_bkg_steps(self, *args):
+        pass
+
+# --------------------------------------------------------------------------------------------------
+# Prep 3D Background subclass
+# --------------------------------------------------------------------------------------------------
+class PrepBackground3D(PrepBackground):
+
+    def _offset(self):
+        return None
+
+    def _add_bkg_steps(self, *args):
+
+        self.bkg_steps.append(isodate.duration_isoformat(self.forecast_duration_for_background))
+
+        return None
+
+# --------------------------------------------------------------------------------------------------
+# Prep 4D Background sublcass
+# --------------------------------------------------------------------------------------------------
+class PrepBackground4D(PrepBackground):
+
+    # -------------------------------------------------------------------------------
+    def _offset(self):
+
+        # If the window type is 4D then remove the window offset as first background
+        # occurs at the beginning of the window
+        self.window_offset_dur = isodate.parse_duration(self.window_offset)
+
+        self.forecast_duration_for_background =\
+            self.forecast_duration_for_background - self.window_offset_dur
+
+        return None
+
+    # -------------------------------------------------------------------------------
+    def _add_bkg_steps(self, cycle_time_dto):
+
+        self.bkg_steps.append(isodate.duration_isoformat(self.forecast_duration_for_background))
+
+        # If background is provided though files get all backgrounds
+        bkg_freq_dur = isodate.parse_duration(self.background_frequency)
+
+        # Check for a sensible frequency
+        if (self.window_length_dur/bkg_freq_dur) % 2:
+            # !!!!!!!!!!!!!!!! no logger in self
+            # self.logger.abort('Window length not divisible by background frequency')
+            raise ValueError("Need to call logger somehow.")
+
+        # Loop over window
+        print('cycle_time_dto', cycle_time_dto)
+        print('window_offset_dur', self.window_offset_dur)
+
+        start_date = cycle_time_dto - self.window_offset_dur
+        final_date = cycle_time_dto + self.window_offset_dur
+
+        loop_date = start_date + bkg_freq_dur
+
+        while loop_date <= final_date:
+
+            duration_in = loop_date - start_date + self.forecast_duration_for_background
+            self.bkg_steps.append(isodate.duration_isoformat(duration_in))
+            loop_date += bkg_freq_dur
+
+        return None
+