HMET ASCII works

Author: xaviernogueira

.gitignore

@@ -13,6 +13,7 @@ examples/example_data/*.tif.ovr
 testing/test_data/*.csv
 testing/test_data/*.parquet
 testing/test_data/*.xlsx
+*.asc
 dev_sandbox/**
 lake_erie_data/elements_data/**
 lake_erie_data/nodes_data/**

src/xarray_data_accessor/data_converters/to_gssha.py

@@ -1,60 +1,46 @@
 import warnings
+import logging
 import pandas as pd
 import numpy as np
 import xarray as xr
 from xarray_data_accessor import utility_functions
 from xarray_data_accessor.shared_types import (
     BoundingBoxDict,
+    TimeInput,
 )
 from xarray_data_accessor.data_converters.base import (
     DataConverterBase,
 )
 from xarray_data_accessor.data_converters.factory import (
     DataConversionFactory,
 )
+from xarray_data_accessor.info.gssha import (
+    PrecipitationType,
+    HMETVariables,
+)
 from datetime import datetime
 from pathlib import Path
 from typing import (
     List,
-    Dict,
     Tuple,
+    Dict,
     Union,
-    Literal,
     Optional,
     TypedDict,
 )
 
-PrecipitationType = Literal[
-    'GAGES',
-    'RADAR',
-    'RATES',
-    'ACCUM',
-]
-
-HMETVariables = Literal[
-    'Cloud Cover',
-    'Atmospheric Pressure',
-    'Relative Humidity',
-    'Temperature',
-    'Wind Speed',
-    'Direct Radiation',
-    'Global Radiation',
-]
-
 
 class EventIntervals(TypedDict):
-    """Stores information about a single event interval."""
-
     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,
@@ -98,7 +84,7 @@ 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,
@@ -129,14 +115,14 @@ def _write_ascii_file(
                     f'Something went wrong - File {file_path} is not a valid ASCII file.'
                 )
 
-    @staticmethod
+    @ staticmethod
     def _write_precip_coords(
         easting: pd.Series,
         northing: pd.Series,
     ) -> str:
         """Writes the coordinate lines of a precipitation ASCII file.
 
-        NOTE: This function assumes that the easting and northing coordinates 
+        NOTE: This function assumes that the easting and northing coordinates
         are in the same projection and correspond to the same time step!
 
         Arguments:
@@ -146,7 +132,7 @@ def _write_precip_coords(
         Returns: A string of the precipitation coordinates in ASCII format.
             Output format:
                 NRGAG 2
-                COORD 204555.0  4751268.0 "Center of precipitation pixel #1" 
+                COORD 204555.0  4751268.0 "Center of precipitation pixel #1"
                 COORD 205642.0  4750491.0 "Center of precipitation pixel #2"
         """
         # zip the coordinates
@@ -180,7 +166,7 @@ def make_gssha_precipitation_input(
             xarray_dataset: The xarray dataset to convert.
             precipitation_variable: The name of the precipitation variable.
             precipitation_type: The type of precipitation (i.e., GAGE, RADAR,...).
-            event_intervals: A list of event intervals. Each interval must be 
+            event_intervals: A list of event intervals. Each interval must be
                 a dict of form {event_name: str, start: datetime, end: datetime}.
             file_dir: The directory to save the file to.
             file_name: The name of the file to save.
@@ -269,15 +255,17 @@ def make_gssha_precipitation_input(
             text_content=ascii_text,
             file_path=file_path,
         )
+        logging.info(f'Precipitation ASCII file saved @ {file_path}.')
         return file_path
 
-    @classmethod
+    @ classmethod
     def make_grass_ascii_input(
         cls,
         xarray_dataset: xr.Dataset,
         variable: str,
-        hmet_variable: Optional[HMETVariables] = None,
-        file_name_code: Optional[str] = None,
+        hmet_variable: Optional[str] = None,
+        start_time: Optional[TimeInput] = None,
+        end_time: Optional[TimeInput] = None,
         file_dir: Optional[Union[str, Path]] = None,
         file_suffix: Optional[str] = None,
     ) -> Path:
@@ -289,57 +277,99 @@ def make_grass_ascii_input(
         Arguments:
             xarray_dataset: The xarray dataset to convert.
             variable: The name of the variable to convert.
+            hmet_variable: The name of the GSSHA HMET variable.
+                See: https://www.gsshawiki.com/Continuous:Hydrometeorological_Data
+            start_time: The start time to trim the data to.
+            end_time: The end time to trim the data to.
             file_dir: The directory to save the file to.
-            file_name: The name of the file to save.
+                NOTE: The file name is automatically generated.
             file_suffix: The file suffix to use.
 
         Returns:
-            The path of the output GRASS ASCII input file.
+            The path of the directory containing all GRASS ASCII input files.
         """
-        # get a file path
-        file_path: Path = cls._get_file_path(
-            file_dir=file_dir,
-            file_name=file_name,
-            file_suffix=file_suffix,
-        )
 
-        # TODO: make this work over time https://www.gsshawiki.com/Continuous:Hydrometeorological_Data
-        # make header
-        ascii_text: str = ''
-        bbox: BoundingBoxDict = utility_functions._bbox_from_raster(
-            raster=xarray_dataset,
+        # make sure we are using an appropriate variable name
+        if variable not in xarray_dataset:
+            raise KeyError(
+                f'Variable {variable} not found in xarray dataset!',
+            )
+        if hmet_variable not in HMETVariables.keys():
+            raise KeyError(
+                f'Variable {hmet_variable} not found in HMETVariables! '
+                f'Please specify a valid HMET variable from list: {HMETVariables.keys()}',
+            )
+
+        file_name: str = HMETVariables[hmet_variable].ascii_file_name
+
+        # trim to time range if necessary
+        if start_time or end_time:
+            start_dt = utility_functions._get_datetime(start_time)
+            end_dt = utility_functions._get_datetime(end_time)
+            xarray_dataset = xarray_dataset.time.where(
+                (xarray_dataset.time.dt >=
+                 start_dt and xarray_dataset.time.dt <= end_dt),
+                drop=True,
+            ).copy()
+
+        # make GRASS ASCII header with bounds
+        grass_header: str = ''
+
+        x_dim = xarray_dataset.attrs['x_dim']
+        y_dim = xarray_dataset.attrs['y_dim']
+
+        bbox = BoundingBoxDict(
+            north=np.max(xarray_dataset[y_dim].values),
+            south=np.min(xarray_dataset[y_dim].values),
+            east=np.max(xarray_dataset[x_dim].values),
+            west=np.min(xarray_dataset[x_dim].values),
         )
+
         for direction in ['north', 'south', 'east', 'west']:
-            ascii_text += f'{direction}: {bbox[direction]}\n'
-        ascii_text += f'rows: {len(xarray_dataset.dims[xarray_dataset.attrs["y_dim"]].values)}\n'
-        ascii_text += f'rows: {len(xarray_dataset.dims[xarray_dataset.attrs["x_dim"]].values)}\n'
-
-        # get the data as a string
-        data_str = (
-            np.array2string(
-                xarray_dataset[variable].values,
-                formatter={'float': lambda x: str(x)},
-                separator=' ',
+            grass_header += f'{direction}: {bbox[direction]}\n'
+
+        grass_header += f'rows: {len(xarray_dataset[y_dim].values)}\n'
+        grass_header += f'cols: {len(xarray_dataset[x_dim].values)}\n'
+
+        # iterate over time steps
+        for time in xarray_dataset.time.values:
+            data_str = (
+                np.array2string(
+                    xarray_dataset[variable].sel(time=time).values,
+                    max_line_width=100000000,
+                    formatter={'float': lambda x: str(x)},
+                    separator=' ',
+                )
+                .replace('  ', '')
+                .replace('[', '')
+                .replace(']', '')
+                .replace('\n ', '\n')
             )
-            .replace('  ', '')
-            .replace('[', '')
-            .replace(']', '')
-        )
 
-        ascii_text += data_str
+            ascii_text = grass_header + data_str
 
-        # write the ASCII file
-        cls._write_ascii_file(
-            text_content=ascii_text,
-            file_path=file_path,
-        )
-        return file_path
+            # get a file path (YYYYMMDDHH_TYPE.asc format)
+            timestamp: str = pd.to_datetime(time).strftime('%Y%m%d%H')
+            file_path: Path = cls._get_file_path(
+                file_dir=file_dir,
+                file_name=f'{timestamp}_{file_name}',
+                file_suffix=file_suffix,
+            )
 
-    @classmethod
+            # write the ASCII file
+            cls._write_ascii_file(
+                text_content=ascii_text,
+                file_path=file_path,
+            )
+        logging.info(f'GRASS ASCII files saved to {file_dir}.')
+        return file_path.parent
+
+    @ classmethod
     def get_conversion_functions(
         cls,
     ) -> Dict[str, DataConverterBase.ConversionFunctionType]:
         """Returns a dictionary of conversion functions."""
         return {
             cls.make_gssha_precipitation_input.__name__: cls.make_gssha_precipitation_input,
+            cls.make_grass_ascii_input.__name__: cls.make_grass_ascii_input,
         }