Supports transforming coordinates for GSSHA output

Author: xaviernogueira

src/xarray_data_accessor/data_converters/to_gssha.py

@@ -1,5 +1,6 @@
 import warnings
 import logging
+import pyproj
 import pandas as pd
 import numpy as np
 import xarray as xr
@@ -118,8 +119,10 @@ def _write_ascii_file(
 
     @ staticmethod
     def _write_precip_coords(
-        easting: pd.Series,
-        northing: pd.Series,
+        easting: np.ndarray,
+        northing: np.ndarray,
+        input_epsg: Optional[int] = None,
+        output_epsg: Optional[int] = None,
     ) -> str:
         """Writes the coordinate lines of a precipitation ASCII file.
 
@@ -136,8 +139,17 @@ def _write_precip_coords(
                 COORD 204555.0  4751268.0 "Center of precipitation pixel #1"
                 COORD 205642.0  4750491.0 "Center of precipitation pixel #2"
         """
+        # reproject if necessary
+        if output_epsg:
+            easting, northing = utility_functions._convert_xy_coordinates(
+                x=easting,
+                y=northing,
+                input_epsg=input_epsg,
+                output_epsg=output_epsg,
+            )
+
         # zip the coordinates
-        coordinates = zip(easting.to_list(), northing.to_list())
+        coordinates = zip(easting.tolist(), northing.tolist())
 
         # get the number of "gages"
         num_gages = len(easting)
@@ -238,6 +250,7 @@ def make_gssha_precipitation_input(
         precipitation_variable: str,
         precipitation_type: Optional[PrecipitationType] = None,
         event_intervals: Optional[List[EventIntervals]] = None,
+        output_epsg: Optional[int] = None,
         file_dir: Optional[Union[str, Path]] = None,
         file_name: Optional[str] = None,
         file_suffix: Optional[str] = None,
@@ -296,8 +309,10 @@ def make_gssha_precipitation_input(
         )
         ts1: datetime = data_df['time'].unique()[0]
         coordinates_header: str = cls._write_precip_coords(
-            easting=data_df.loc[data_df.time == ts1, x_dim],
-            northing=data_df.loc[data_df.time == ts1, y_dim],
+            easting=data_df.loc[data_df.time == ts1, x_dim].to_numpy(),
+            northing=data_df.loc[data_df.time == ts1, y_dim].to_numpy(),
+            input_epsg=xarray_dataset.attrs.get('EPSG', None),
+            output_epsg=output_epsg,
         )
 
         # get events
@@ -353,6 +368,7 @@ def make_gssha_grass_ascii(
         hmet_variable: Optional[str] = None,
         start_time: Optional[TimeInput] = None,
         end_time: Optional[TimeInput] = None,
+        output_epsg: Optional[int] = None,
         file_dir: Optional[Union[str, Path]] = None,
         file_name: Optional[str] = None,
         file_suffix: Optional[str] = None,
@@ -402,11 +418,20 @@ def make_gssha_grass_ascii(
         x_dim = xarray_dataset.attrs['x_dim']
         y_dim = xarray_dataset.attrs['y_dim']
 
+        # get easting and northing coordinate arrays (reproject if necessary)
+        easting, northing = utility_functions._convert_xy_coordinates(
+            x=xarray_dataset[x_dim].values,
+            y=xarray_dataset[y_dim].values,
+            input_epsg=xarray_dataset.attrs.get('EPSG', None),
+            output_epsg=output_epsg,
+        )
+
+        # make bounding box for header
         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),
+            north=np.max(northing),
+            south=np.min(northing),
+            east=np.max(easting),
+            west=np.min(easting),
         )
 
         for direction in ['north', 'south', 'east', 'west']:
@@ -482,6 +507,8 @@ def make_gssha_hmet_wes(
             file_suffix: The file suffix to use.
             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.
+                Note that this must be in the same units as the xarray dataset.
 
         Returns:
             The path of the output WES ASCII file.

src/xarray_data_accessor/utility_functions.py

@@ -221,12 +221,40 @@ def _coords_in_bbox(
     return False
 
 
+def _convert_xy_coordinates(
+    x: np.ndarray,
+    y: np.ndarray,
+    input_epsg: Optional[int] = None,
+    output_epsg: Optional[int] = None,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Transforms coordinates from one EPSG to another."""
+
+    if not input_epsg and output_epsg:
+        raise ValueError(
+            'An input EPSG must be provided if an output EPSG is provided! '
+            'This is pulled from xarray_dataset.attrs["EPSG"]. '
+            'Please set this attribute to a valid integer EPSG code.',
+        )
+    if output_epsg and output_epsg != input_epsg:
+        # Create a pyproj Transformer for the coordinate conversion
+        transformer = pyproj.Transformer.from_crs(
+            input_epsg,
+            output_epsg,
+            always_xy=True,
+        )
+
+        # Perform the coordinate conversion
+        x, y = transformer.transform(x, y)
+    return x, y
+
+
 def _convert_bbox(
     bbox: BoundingBoxDict,
     known_epsg: int,
 ) -> BoundingBoxDict:
     """Converts bbox coordinates to a different EPSG."""
 
+    # TODO: consider deprecating this method
     # EPSG:4326 bbox list
     bbox_list_in = [
         bbox['west'],

testing/test_5_gssha.py

@@ -43,6 +43,7 @@ def test_precipitation_input(test_dataset) -> None:
     out_path = xda.DataConversionFunctions.make_gssha_precipitation_input(
         test_dataset,
         precipitation_variable='2m_temperature',
+        output_epsg=26915,
     )
 
     assert out_path.exists()