From f07c7ee49c19a57e0f54b7369e6f3e688776558b Mon Sep 17 00:00:00 2001
From: Kristen Thyng <kthyng@gmail.com>
Date: Fri, 22 Sep 2023 17:19:33 -0500
Subject: [PATCH 1/2] preprocessing accounts for more

ROMS preprocessing checks for 3D instead of just 4D data variables now to update their coordinates to work better with cf-xarray.
Also coordinates could say "x_rho" and "y_rho" instead of longitudes and latitudes in which case they are changed to say the coordinates.
---
 extract_model/preprocessing.py | 73 +++++++++++++++++++++++++++++++---
 1 file changed, 67 insertions(+), 6 deletions(-)

diff --git a/extract_model/preprocessing.py b/extract_model/preprocessing.py
index 4e842ac..ace1178 100644
--- a/extract_model/preprocessing.py
+++ b/extract_model/preprocessing.py
@@ -63,6 +63,24 @@ def preprocess_roms(
     for dim in dims:
         ds[dim] = (dim, np.arange(ds.sizes[dim]), {"axis": "Y"})
 
+    # add attributes for lon/lat
+    lon_attrs = {
+        "standard_name": "longitude",
+        "units": "degree_east",
+        "field": "longitude",
+    }
+    lat_attrs = {
+        "standard_name": "latitude",
+        "units": "degree_north",
+        "field": "latitude",
+    }
+    coords = [coord for coord in ds.coords if coord.startswith("lon_")]
+    for coord in coords:
+        ds[coord].attrs = lon_attrs
+    coords = [coord for coord in ds.coords if coord.startswith("lat_")]
+    for coord in coords:
+        ds[coord].attrs = lat_attrs
+
     # Fix standard_name for s_rho/s_w
     if "Vtransform" in ds.data_vars and "s_rho" in ds.coords:
         cond1 = (
@@ -177,17 +195,60 @@ def preprocess_roms(
             ds[zname] = order(ds[zname])
 
     # replace s_rho with z_rho, etc, to make z_rho the vertical coord
-    for sname, zname in name_dict.items():
-        for var in ds.data_vars:
-            if ds[var].ndim == 4:
-                if "coordinates" in ds[var].encoding:
-                    coords = ds[var].encoding["coordinates"]
+    for var in ds.data_vars:
+        if ds[var].ndim >= 3:
+            if "coordinates" in ds[var].encoding:
+                coords = ds[var].encoding["coordinates"]
+                # update s's to z's
+                for sname, zname in name_dict.items():
                     if sname in coords:  # replace if present
                         coords = coords.replace(sname, zname)
                     else:  # still add z_rho or z_w
                         if zname in ds[var].coords and ds[zname].shape == ds[var].shape:
                             coords += f" {zname}"
-                    ds[var].encoding["coordinates"] = coords
+                # could have "x_rho" instead of "lon_rho", etc
+                if "x_" in coords:
+                    xcoord = [element for element in coords.split() if "x_" in element][
+                        0
+                    ]
+                    coords = coords.replace(xcoord, xcoord.replace("x", "lon"))
+                # could have "x_rho" instead of "lon_rho", etc
+                if "y_" in coords:
+                    ycoord = [element for element in coords.split() if "y_" in element][
+                        0
+                    ]
+                    coords = coords.replace(ycoord, ycoord.replace("y", "lat"))
+                ds[var].encoding["coordinates"] = coords
+            # same but coordinates not inside encoding. Do same processing
+            # but also move coordinates from attrs to encoding.
+            elif "coordinates" in ds[var].attrs:
+                coords_here = ds[var].attrs["coordinates"]
+                # update s's to z's
+                for sname, zname in name_dict.items():
+                    if sname in coords_here:  # replace if present
+                        coords_here = coords_here.replace(sname, zname)
+                    else:  # still add z_rho or z_w
+                        if zname in ds[var].coords and ds[zname].shape == ds[var].shape:
+                            coords_here += f" {zname}"
+                # could have "x_rho" instead of "lon_rho", etc
+                if "x_" in coords_here:
+                    xcoord = [
+                        element for element in coords_here.split() if "x_" in element
+                    ][0]
+                    coords_here = coords_here.replace(
+                        xcoord, xcoord.replace("x", "lon")
+                    )
+                # could have "x_rho" instead of "lon_rho", etc
+                if "y_" in coords_here:
+                    ycoord = [
+                        element for element in coords_here.split() if "y_" in element
+                    ][0]
+                    coords_here = coords_here.replace(
+                        ycoord, ycoord.replace("y", "lat")
+                    )
+                # move coords to encoding and delete from attrs
+                ds[var].encoding["coordinates"] = coords_here
+                del ds[var].attrs["coordinates"]
 
     # # easier to remove "coordinates" attribute from any variables than add it to all
     # for var in ds.data_vars:

From b9490536196910e2796a9bc6b16fff3eeb9423da Mon Sep 17 00:00:00 2001
From: Kristen Thyng <kthyng@gmail.com>
Date: Fri, 22 Sep 2023 17:22:01 -0500
Subject: [PATCH 2/2] updated whats new

---
 docs/whats_new.rst | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/docs/whats_new.rst b/docs/whats_new.rst
index 4c5e5cb..c5af1ba 100644
--- a/docs/whats_new.rst
+++ b/docs/whats_new.rst
@@ -1,6 +1,10 @@
 :mod:`What's New`
 -----------------
 
+v1.2.1 (September 22, 2023)
+===========================
+* ROMS preprocessing checks for 3D instead of just 4D data variables now to update their coordinates to work better with cf-xarray. Also coordinates could say "x_rho" and "y_rho" instead of longitudes and latitudes in which case they are changed to say the coordinates
+
 v1.2.0 (September 13, 2023)
 ===========================
 * Improvements to interpolation