From 64d306b0f8cd2919b0232565d350fae6af89c5e4 Mon Sep 17 00:00:00 2001
From: Qiusheng Wu <giswqs@gmail.com>
Date: Mon, 6 May 2024 17:24:43 -0400
Subject: [PATCH] Add filter_pace function (#21)

---
 docs/examples/pace.ipynb | 36 +++++++++++++++++++++++++++++++++++
 hypercoast/pace.py       | 41 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 77 insertions(+)

diff --git a/docs/examples/pace.ipynb b/docs/examples/pace.ipynb
index 248d3ea0..1a9c1aa9 100644
--- a/docs/examples/pace.ipynb
+++ b/docs/examples/pace.ipynb
@@ -95,6 +95,42 @@
    "source": [
     "hypercoast.viz_pace(dataset, wavelengths=[500, 510, 520, 530], ncols=2, crs=\"default\")"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Plot a spectral signature."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "latitude = 25.493961\n",
+    "longitude = -91.25617\n",
+    "hypercoast.filter_pace(dataset, latitude, longitude, return_plot=True)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Plot multiple spectral signatures."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "latitude = (25.49, 25.50)\n",
+    "longitude = (-92, -91.055)\n",
+    "hypercoast.filter_pace(dataset, latitude, longitude, return_plot=True)"
+   ]
   }
  ],
  "metadata": {
diff --git a/hypercoast/pace.py b/hypercoast/pace.py
index 00724448..491faa28 100644
--- a/hypercoast/pace.py
+++ b/hypercoast/pace.py
@@ -176,3 +176,44 @@ def viz_pace(
 
         plt.tight_layout()
         plt.show()
+
+
+def filter_pace(dataset, latitude, longitude, drop=True, return_plot=False, **kwargs):
+    """
+    Filters a PACE dataset based on latitude and longitude.
+
+    Args:
+        dataset (xr.Dataset): The PACE dataset to filter.
+        latitude (float or tuple): The latitude to filter by. If a tuple or list, it represents a range.
+        longitude (float or tuple): The longitude to filter by. If a tuple or list, it represents a range.
+        drop (bool, optional): Whether to drop the filtered out data. Defaults to True.
+
+    Returns:
+        xr.DataArray: The filtered PACE data.
+    """
+    if isinstance(latitude, list) or isinstance(latitude, tuple):
+        lat_con = (dataset["latitude"] > latitude[0]) & (
+            dataset["latitude"] < latitude[1]
+        )
+    else:
+        lat_con = dataset["latitude"] == latitude
+
+    if isinstance(longitude, list) or isinstance(longitude, tuple):
+        lon_con = (dataset["longitude"] > longitude[0]) & (
+            dataset["longitude"] < longitude[1]
+        )
+    else:
+        lon_con = dataset["longitude"] == longitude
+
+    da = dataset["Rrs"].where(lat_con & lon_con, drop=drop, **kwargs)
+    da_filtered = da.dropna(dim="latitude", how="all")
+    da_filtered = da_filtered.dropna(dim="longitude", how="all")
+
+    if return_plot:
+        rrs_stack = da_filtered.stack(
+            {"pixel": ["latitude", "longitude"]},
+            create_index=False,
+        )
+        rrs_stack.plot.line(hue="pixel")
+    else:
+        return da_filtered