Improve summary module

pyPRMS/summary/OutputCSV.py

@@ -1,19 +1,20 @@
 import pandas as pd   # type: ignore
 
 from pathlib import Path
-from typing import Optional, Union
 
 from ..base.console import get_console_instance
 
+__all__ = ['OutputCSV']
+
 con = None
 
 
-class OutputCSV(object):
+class OutputCSV:
     """Class for working with PRMS CSV output files.
     """
 
-    def __init__(self, filename: Union[str, Path],
-                 verbose: Optional[bool] = False):
+    def __init__(self, filename: str | Path,
+                 verbose: bool = False):
         """Initialize the OutputCSV object.
 
         :param filename: Name of the PRMS CSV output file
@@ -36,48 +37,70 @@ def __init__(self, filename: Union[str, Path],
         self.__basin_vars = []
         self.__col_var = {}
 
+        if not self.__filename.exists():
+            raise FileNotFoundError(f'CSV output file not found: {self.__filename}')
+
         self._read_csv_header()
         self._read_csv_ascii()
 
+    def __repr__(self) -> str:
+        return f'OutputCSV(filename={self.__filename})'
+
     @property
-    def basin_vars(self):
-        """Returns the basin variables from the CSV output file."""
+    def basin_vars(self) -> list[str]:
+        """Returns the basin variables from the CSV output file.
+
+        :returns: List of basin variable names
+        """
         return self.__basin_vars
 
     @property
-    def data(self):
+    def data(self) -> pd.DataFrame:
+        """Returns the model output data as a pandas DataFrame.
+
+        :returns: DataFrame with time index and output variables as columns
+        """
         return self.__data
 
     @property
-    def pois(self):
+    def pois(self) -> list[str]:
+        """Returns the in-order list of points-of-interest (POI) identifiers.
+
+        :returns: List of POI identifiers
+        """
         return self.__pois
 
     @property
-    def poi_segments(self):
+    def poi_segments(self) -> dict[str, int]:
+        """Returns mapping of POI identifiers to their segment indices.
+
+        :returns: Dictionary mapping POI ID to zero-based segment index
+        """
         return self.__poi_segments
 
     @property
-    def variables(self):
+    def variables(self) -> list[str]:
+        """Returns a sorted list of all variable names in the CSV output file.
+
+        :returns: Sorted list of variable names
+        """
         return sorted(list(self.__col_var.values()))
 
     def _read_csv_header(self):
         """Read the headers from a PRMS CSV model output file"""
 
-        fhdl = open(self.__filename, 'r')
-
-        # First row contains field names
-        # Second row is a a mix of field names (for the date) and data types
-        hdr1 = fhdl.readline().strip()
-        hdr2 = fhdl.readline().strip()
-        fhdl.close()
+        with open(self.__filename, 'r') as fhdl:
+            # First row contains field names
+            # Second row is a a mix of field names (for the date) and data types
+            hdr1 = fhdl.readline().strip()
+            hdr2 = fhdl.readline().strip()
 
         # Determine the value separator
         # Check for comma first; some files have commas and spaces
         self.sep = ' '
 
-        match ',' in hdr1:
-            case True:
-                self.sep = ','
+        if ',' in hdr1:
+            self.sep = ','
 
         if self.verbose:
             con.print(f'[green]INFO[/]: value separator = {self.sep}')

pyPRMS/summary/OutputVariable.py

@@ -3,12 +3,21 @@
 import xarray as xr
 
 from pathlib import Path
-from typing import List, Optional, Union
 
 from ..constants import NEW_PTYPE_TO_DTYPE
 
+__all__ = ['OutputVariable']
 
-class OutputVariable(object):
+_LOCAL_DIM_DESC = {'nhru': 'Local model Hydrologic Response Unit ID (HRU)',
+                   'nsegment': 'Local model segment ID'}
+
+_GLOBAL_DIMS = dict(nhru=dict(varname='nhm_id',
+                              long_name='NHM Hydrologic Response Unit ID (HRU)'),
+                    nsegment=dict(varname='nhm_seg',
+                                  long_name='NHM segment ID'))
+
+
+class OutputVariable:
     """Container for a single output variable
 
     Each OutputVariable instance contains the model output for a single
@@ -17,7 +26,7 @@ class OutputVariable(object):
     """
 
     def __init__(self, name: str,
-                 filename: Union[str, Path],
+                 filename: str | Path,
                  metadata: dict):
         """Initialize the OutputVariable object.
 
@@ -30,6 +39,9 @@ def __init__(self, name: str,
             filename = Path(filename)
         self.__filename = filename
 
+        if not self.__filename.exists():
+            raise FileNotFoundError(f'Output variable file not found: {self.__filename}')
+
         self.__name = name
         self.__data = None
 
@@ -38,6 +50,9 @@ def __init__(self, name: str,
         else:
             self.metadata = metadata[name]
 
+    def __repr__(self) -> str:
+        return f'OutputVariable(name={self.__name!r}, filename={self.__filename!r})'
+
     @property
     def data(self) -> pd.DataFrame:
         """Returns the source model output as a pandas DataFrame
@@ -58,9 +73,9 @@ def filename(self) -> Path:
 
         return self.__filename
 
-    def to_csv(self, filename: Union[str, Path],
-               columns: Optional[List[int]] = None,
-               sep: Optional[str] = ','):
+    def to_csv(self, filename: str | Path,
+               columns: list[int] | None = None,
+               sep: str = ','):
         """Write the output variable to a CSV file.
 
         :param filename: Name of the output file
@@ -73,7 +88,7 @@ def to_csv(self, filename: Union[str, Path],
 
         self.data.to_csv(filename, sep=sep, index=True, header=True, columns=columns, chunksize=50)
 
-    def to_netcdf(self, filename: Union[str, Path]):
+    def to_netcdf(self, filename: str | Path):
         """Write the output variable to a netCDF file.
 
         :param filename: Name of the netCDF output file
@@ -90,19 +105,32 @@ def to_xarray(self) -> xr.DataArray:
         :returns: xarray DataArray
         """
 
-        local_dim_desc = {'nhru': 'Local model Hydrologic Response Unit ID (HRU)',
-                          'nsegment': 'Local model segment ID'}
+        dim_name = self._resolve_dim_name()
+        da = self._build_data_array(dim_name)
+        self._set_time_encoding(da)
+        self._set_variable_attrs(da)
+        return da
+
+    def _resolve_dim_name(self) -> str:
+        """Map the metadata dimension to a canonical dimension name.
 
-        global_dims = dict(nhru=dict(varname='nhm_id',
-                                     long_name='NHM Hydrologic Response Unit ID (HRU)'),
-                           nsegment=dict(varname='nhm_seg',
-                                         long_name='NHM segment ID'))
+        :returns: Canonical dimension name (e.g. 'nhru', 'nsegment', 'one')
+        """
 
         dim_name = self.metadata['dimensions'][0]
 
         if dim_name in ['nssr', 'ngw']:
             dim_name = 'nhru'
 
+        return dim_name
+
+    def _build_data_array(self, dim_name: str) -> xr.DataArray:
+        """Build the xarray DataArray from the output data.
+
+        :param dim_name: Canonical dimension name
+        :returns: xarray DataArray with coordinates and dimension attributes set
+        """
+
         if dim_name == 'one':
             # Basin variable
             da = self.data.squeeze().to_xarray()
@@ -121,34 +149,47 @@ def to_xarray(self) -> xr.DataArray:
 
             if self.metadata.get('is_global', False):
                 # When is_global is true the file header contains global HRU or segment IDs
-                da[global_dims[dim_name]['varname']] = da[dim_name]
-                da[global_dims[dim_name]['varname']].attrs['long_name'] = global_dims[dim_name]['long_name']
+                da[_GLOBAL_DIMS[dim_name]['varname']] = da[dim_name]
+                da[_GLOBAL_DIMS[dim_name]['varname']].attrs['long_name'] = _GLOBAL_DIMS[dim_name]['long_name']
 
                 # Reset the nhru/nsegment coordinate variable values to 1..N
                 da[dim_name] = np.arange(1, self.data.shape[1]+1, dtype=np.int32)
 
             # Set attributes for local model dimensions
-            da[dim_name].attrs['long_name'] = local_dim_desc[dim_name]
+            da[dim_name].attrs['long_name'] = _LOCAL_DIM_DESC[dim_name]
+
+        return da
+
+    def _set_time_encoding(self, da: xr.DataArray):
+        """Set time coordinate attributes and encoding on the DataArray.
+
+        :param da: DataArray to modify in place
+        """
 
-        # Set the time coordinate variable attributes
         first_time = self.data.index[0]
         da.time.attrs['standard_name'] = 'time'
         da.time.attrs['long_name'] = 'time'
         da.time.encoding['units'] = f'days since {first_time.year}-{first_time.month:02d}-{first_time.day:02d} 00:00:00'
         da.time.encoding['calendar'] = 'standard'
 
-        # Output variable attributes
+    def _set_variable_attrs(self, da: xr.DataArray):
+        """Set output variable attributes and compression encoding.
+
+        :param da: DataArray to modify in place
+        """
+
         da.attrs['long_name'] = self.metadata['description']
         da.attrs['units'] = self.metadata['units']
         da.encoding.update(dict(_FillValue=None,
                                 compression='zlib',
                                 complevel=2,
                                 fletcher32=True))
 
-        return da
-
     def _read_file(self):
         """Read model variable output file.
+
+        Parses the CSV file and stores the result as a pandas DataFrame
+        with a time index and appropriately typed columns.
         """
 
         self.__data = pd.read_csv(self.__filename, sep=',', skipinitialspace=True,