Hot start functionality added for GSSHA functions

Author: xaviernogueira

src/xarray_data_accessor/data_converters/to_gssha.py

@@ -1,6 +1,5 @@
 import warnings
 import logging
-import pyproj
 import pandas as pd
 import numpy as np
 import xarray as xr
@@ -28,21 +27,27 @@
     Dict,
     Union,
     Optional,
+    Literal,
     TypedDict,
 )
 
+OPEN_MODES: Dict[bool, Literal['a', 'w']] = {
+    True: 'a',
+    False: 'w',
+}
+
 
 class EventIntervals(TypedDict):
     name: str
     start: datetime
     end: datetime
 
 
-@ DataConversionFactory.register
+@DataConversionFactory.register
 class ConvertToGSSHA(DataConverterBase):
     """Converts xarray datasets to GSSHA input files."""
 
-    @ staticmethod
+    @staticmethod
     def _get_file_path(
         file_dir: Optional[Union[str, Path]] = None,
         file_name: Optional[str] = None,
@@ -63,7 +68,7 @@ def _get_file_path(
         # make sure the file name is valid
         if not file_name:
             file_name = 'gssha_input'
-            logging.warn(
+            logging.warning(
                 f'No file name was provided! Using default file name {file_name}.',
             )
         if not isinstance(file_name, str):
@@ -86,16 +91,17 @@ def _get_file_path(
         # return the file path
         return Path(file_dir / f'{file_name}{file_suffix}')
 
-    @ staticmethod
+    @staticmethod
     def _write_ascii_file(
         text_content: str,
         file_path: Path,
+        hot_start: Optional[bool] = False,
     ) -> None:
         """Writes the text content to the file path."""
-        # write the text content to the file path
+
         with open(
             file_path,
-            'w',
+            OPEN_MODES[hot_start],
             encoding='ascii',
         ) as file:
             file.write(text_content)
@@ -117,7 +123,7 @@ def _write_ascii_file(
                     f'Something went wrong - File {file_path} is not a valid ASCII file.',
                 )
 
-    @ staticmethod
+    @staticmethod
     def _write_precip_coords(
         easting: np.ndarray,
         northing: np.ndarray,
@@ -254,6 +260,7 @@ def make_gssha_precipitation_input(
         file_dir: Optional[Union[str, Path]] = None,
         file_name: Optional[str] = None,
         file_suffix: Optional[str] = None,
+        hot_start: Optional[bool] = False,
     ) -> Path:
         """Creates a GSSHA precipitation input file from an xarray dataset.
 
@@ -269,6 +276,8 @@ def make_gssha_precipitation_input(
             file_dir: The directory to save the file to.
             file_name: The name of the file to save.
             file_suffix: The file suffix to use.
+            hot_start: If true data is appended to the end of the file.
+                Otherwise, the file is overwritten.
 
         Returns:
             The path of the output precipitation ASCII input file.
@@ -356,11 +365,12 @@ def make_gssha_precipitation_input(
         cls._write_ascii_file(
             text_content=ascii_text,
             file_path=file_path,
+            hot_start=hot_start,
         )
         logging.info(f'Precipitation ASCII file saved @ {file_path}.')
         return file_path
 
-    @ classmethod
+    @classmethod
     def make_gssha_grass_ascii(
         cls,
         xarray_dataset: xr.Dataset,
@@ -477,7 +487,7 @@ def make_gssha_grass_ascii(
         )
         return file_paths
 
-    @ classmethod
+    @classmethod
     def make_gssha_hmet_wes(
         cls,
         xarray_dataset: xr.Dataset,
@@ -487,6 +497,7 @@ def make_gssha_hmet_wes(
         file_dir: Optional[Union[str, Path]] = None,
         file_name: Optional[str] = None,
         file_suffix: Optional[str] = None,
+        hot_start: Optional[bool] = False,
         how: Optional[HMETAggregationFunctions] = None,
         xy_coords: Optional[Tuple[str, str]] = None,
     ) -> Path:
@@ -505,6 +516,8 @@ def make_gssha_hmet_wes(
                 NOTE: The file name is automatically generated.
             file_name: The name of the file to save.
             file_suffix: The file suffix to use.
+            hot_start: If true data is appended to the end of the file.
+                Otherwise, the file is overwritten.
             how: The method to use to aggregate the data at each time step.
                 Options include: 'mean', 'median', 'min', 'max', 'sum'.
             xy_coords: The x and y coordinate names to use for aggregation.
@@ -586,12 +599,13 @@ def make_gssha_hmet_wes(
         cls._write_ascii_file(
             text_content=data_str,
             file_path=file_path,
+            hot_start=hot_start,
         )
 
         logging.info(f'HMET WES ASCII file saved @ {file_path}.')
         return file_path
 
-    @ classmethod
+    @classmethod
     def get_conversion_functions(
         cls,
     ) -> Dict[str, DataConverterBase.ConversionFunctionType]:

testing/test_5_gssha.py

@@ -24,6 +24,12 @@ def test_dataset(test_dir) -> xr.Dataset:
     return ds
 
 
+def count_lines(filename: Path):
+    with open(filename, 'r') as file:
+        lines = file.readlines()
+    return len(lines)
+
+
 def test_factory() -> None:
     """Make sure the function was correctly registered."""
     assert 'ConvertToGSSHA' in xda.DataConversionFactory.get_converter_classes().keys()
@@ -43,11 +49,24 @@ def test_precipitation_input(test_dataset) -> None:
     out_path = xda.DataConversionFunctions.make_gssha_precipitation_input(
         test_dataset,
         precipitation_variable='2m_temperature',
+        precipitation_type='GAGE',
         output_epsg=26915,
     )
 
     assert out_path.exists()
     assert out_path.suffix == '.gag'
+    l1 = count_lines(out_path)
+
+    # test the hot start
+    out_path = xda.DataConversionFunctions.make_gssha_precipitation_input(
+        test_dataset,
+        precipitation_variable='2m_temperature',
+        precipitation_type='GAGE',
+        output_epsg=26915,
+        hot_start=True,
+    )
+    l2 = count_lines(out_path)
+    assert l1 < l2
     out_path.unlink()