[FEAT] set default dates for all functions (#56)

* set default dates for all functions * robust to nan dates * Add profile selection in some functions Co-authored-by: Chiara Monforte <[email protected]>
OceanGlidersCommunity · Nov 1, 2024 · df79144 · df79144
1 parent 75fd0da
commit df79144
Show file tree

Hide file tree

Showing 3 changed files with 63 additions and 46 deletions.
diff --git a/glidertest/tools.py b/glidertest/tools.py
@@ -1,5 +1,6 @@
 import matplotlib.dates as mdates
 import matplotlib.pyplot as plt
+import datetime
 import numpy as np
 import pandas as pd
 import seaborn as sns
@@ -399,19 +400,23 @@ def sunset_sunrise(time, lat, lon):
     return sunrise, sunset
 
 
-def day_night_avg(ds, sel_var='CHLA', start_time='2024-04-18', end_time='2024-04-20'):
+def day_night_avg(ds, sel_var='CHLA', start_time=None, end_time=None, start_prof=None, end_prof=None):
     """
-    This function computes night and day averages for a selected variable over a specific period of time
+    This function computes night and day averages for a selected variable over a specific period of time or a specific series of dives
     Data in divided into day and night using the sunset and sunrise time as described in the above function sunset_sunrise from GliderTools
     Parameters
     ----------
     ds: xarray on OG1 format containing at least time, depth, latitude, longitude and the selected variable. 
         Data should not be gridded.
     sel_var: variable to use to compute the day night averages
     start_time: Start date of the data selection. As missions can be long and can make it hard to visualise NPQ effect,
-                we recommend end selecting small section of few days to few weeks.
+                we recommend selecting small section of few days to a few weeks. Defaults to the central week of the deployment
     end_time: End date of the data selection. As missions can be long and can make it hard to visualise NPQ effect,
-                we recommend selecting small section of few days to few weeks.
+                we recommend selecting small section of few days to a few weeks. Defaults to the central week of the deployment
+    start_prof: Start profile of the data selection. If no profile is specified, the specified time selction will be used or the the central week of the deployment.
+                It is important to have a large enough number of dives to have some day and night data otherwise the function will not run
+    end_prof:  End profile of the data selection. If no profile is specified, the specified time selction will be used or the the central week of the deployment.
+            It is important to have a large enough number of dives to have some day and night data otherwise the function will not run
                 
     Returns
     -------
@@ -429,11 +434,20 @@ def day_night_avg(ds, sel_var='CHLA', start_time='2024-04-18', end_time='2024-04
             day: Actual date for the batch 
 
     """
-    if "TIME" in ds.indexes.keys():
-        pass
-    else:
+    if "TIME" not in ds.indexes.keys():
         ds = ds.set_xindex('TIME')
-    ds_sel = ds.sel(TIME=slice(start_time, end_time))
+
+    if not start_time:
+        start_time = ds.TIME.mean() - np.timedelta64(3, 'D')
+    if not end_time:
+        end_time = ds.TIME.mean() + np.timedelta64(3, 'D')
+
+    if start_prof and end_prof:
+        t1 = ds.TIME.where(ds.PROFILE_NUMBER==start_prof).dropna(dim='N_MEASUREMENTS')[0]
+        t2 = ds.TIME.where(ds.PROFILE_NUMBER==end_prof).dropna(dim='N_MEASUREMENTS')[-1]
+        ds_sel = ds.sel(TIME=slice(t1,t2))
+    else:
+        ds_sel = ds.sel(TIME=slice(start_time, end_time))
     sunrise, sunset = sunset_sunrise(ds_sel.TIME, ds_sel.LATITUDE, ds_sel.LONGITUDE)
 
     # creating batches where one batch is a night and the following day
@@ -458,7 +472,7 @@ def day_night_avg(ds, sel_var='CHLA', start_time='2024-04-18', end_time='2024-04
     return day_av, night_av
 
 
-def plot_daynight_avg(day: pd.DataFrame, night: pd.DataFrame, ax: plt.Axes = None, sel_day='2023-09-09',
+def plot_daynight_avg(day: pd.DataFrame, night: pd.DataFrame, ax: plt.Axes = None, sel_day=None,
                       xlabel='Chlorophyll [mg m-3]', **kw: dict, ) -> tuple({plt.Figure, plt.Axes}):
     """
     This function can be used to plot the day and night averages computed with the day_night_avg function
@@ -468,14 +482,18 @@ def plot_daynight_avg(day: pd.DataFrame, night: pd.DataFrame, ax: plt.Axes = Non
     day: pandas dataframe containing the day averages
     night: pandas dataframe containing the night averages
     ax: axis to plot the data
-    sel_day: selected day to plot
+    sel_day: selected day to plot. Defaults to the median day
     xlabel: label for the x-axis
     
     Returns
     -------
     A line plot comparing the day and night average over depth for the selected day
 
     """
+    if not sel_day:
+        dates = list(day.date.dropna().values) + list(night.date.dropna().values)
+        dates.sort()
+        sel_day = dates[int(len(dates)/2)]
     if ax is None:
         fig, ax = plt.subplots(figsize=(5, 5))
     else:
@@ -492,8 +510,8 @@ def plot_daynight_avg(day: pd.DataFrame, night: pd.DataFrame, ax: plt.Axes = Non
     return fig, ax
 
 
-def plot_section_with_srss(ds: xr.Dataset, sel_var: str, ax: plt.Axes = None, start_time='2023-09-06',
-                           end_time='2023-09-10', ylim=45, **kw: dict, ) -> tuple({plt.Figure, plt.Axes}):
+def plot_section_with_srss(ds: xr.Dataset, sel_var: str, ax: plt.Axes = None, start_time=None,
+                           end_time=None,start_prof=None, end_prof=None, ylim=45, **kw: dict, ) -> tuple({plt.Figure, plt.Axes}):
     """
     This function can be used to plot sections for any variable with the sunrise and sunset plotted over
     
@@ -503,8 +521,10 @@ def plot_section_with_srss(ds: xr.Dataset, sel_var: str, ax: plt.Axes = None, st
         Data should not be gridded.
     sel_var: selected variable to plot
     ax: axis to plot the data
-    start_time: Start date of the data selection. As missions can be long and came make it hard to visualise NPQ effect,
-    end_time: End date of the data selection. As missions can be long and came make it hard to visualise NPQ effect,
+    start_time: Start date of the data selection format 'YYYY-MM-DD'. As missions can be long and came make it hard to visualise NPQ effect. Defaults to mid 4 days
+    end_time: End date of the data selection format 'YYYY-MM-DD'. As missions can be long and came make it hard to visualise NPQ effect. Defaults to mid 4 days
+    start_prof: Start profile of the data selection. If no profile is specified, the specified time selction will be used or the mid 4 days of the deployment
+    end_prof:  End profile of the data selection. If no profile is specified, the specified time selction will be used or the mid 4 days of the deployment
     ylim: specified limit for the maximum y-axis value. The minimum is computed as ylim/30
     
     Returns
@@ -518,7 +538,23 @@ def plot_section_with_srss(ds: xr.Dataset, sel_var: str, ax: plt.Axes = None, st
 
     if "TIME" not in ds.indexes.keys():
         ds = ds.set_xindex('TIME')
-    ds_sel = ds.sel(TIME=slice(start_time, end_time))
+
+    if not start_time:
+        start_time = ds.TIME.mean() - np.timedelta64(2, 'D')
+    if not end_time:
+        end_time = ds.TIME.mean() + np.timedelta64(2, 'D')
+
+    if start_prof and end_prof:
+        t1 = ds.TIME.where(ds.PROFILE_NUMBER==start_prof).dropna(dim='N_MEASUREMENTS')[0]
+        t2 = ds.TIME.where(ds.PROFILE_NUMBER==end_prof).dropna(dim='N_MEASUREMENTS')[-1]
+        ds_sel = ds.sel(TIME=slice(t1,t2))
+    else:
+        ds_sel = ds.sel(TIME=slice(start_time, end_time))
+
+    if len(ds_sel.TIME) == 0:
+        msg = f"supplied limits start_time: {start_time} end_time: {end_time} do not overlap with dataset TIME range {str(ds.TIME.values.min())[:10]} - {str(ds.TIME.values.max())[:10]}"
+        raise ValueError(msg)
+
     sunrise, sunset = sunset_sunrise(ds_sel.TIME, ds_sel.LATITUDE, ds_sel.LONGITUDE)
 
     c = ax.scatter(ds_sel.TIME, ds_sel.DEPTH, c=ds_sel[sel_var], s=10, vmin=np.nanpercentile(ds_sel[sel_var], 0.5),
@@ -821,6 +857,7 @@ def plot_ts_histograms(ds: xr.Dataset, ax: plt.Axes = None, **kw: dict) -> tuple
 
     return fig, ax
 
+
 def calc_DEPTH_Z(ds):
     """
     Calculate the depth (Z position) of the glider using the gsw library to convert pressure to depth.
@@ -1040,5 +1077,3 @@ def plot_vertical_speeds_with_histograms(ds, start_prof=None, end_prof=None):
     plt.show()
 
     return fig, axs
-
-
diff --git a/notebooks/demo.ipynb b/notebooks/demo.ipynb
@@ -277,7 +277,7 @@
     "# Let's visually check a section of chlorphyll and see if we observe any NPQ\n",
     "fig, ax = plt.subplots(1, 1, figsize=(15, 5))\n",
     "\n",
-    "tools.plot_section_with_srss(ds, 'CHLA', ax, start_time = '2023-06-04', end_time = '2023-06-06', ylim=35)"
+    "tools.plot_section_with_srss(ds, 'CHLA', ax, ylim=35)"
    ]
   },
   {
@@ -288,9 +288,9 @@
    "outputs": [],
    "source": [
     "# Compute day and night average for chlorophylla and temeparture\n",
-    "dayT, nightT = tools.day_night_avg(ds, sel_var='TEMP',start_time = '2023-06-03', end_time = '2023-06-07')\n",
-    "dayS, nightS = tools.day_night_avg(ds, sel_var='PSAL',start_time = '2023-06-03', end_time = '2023-06-07')\n",
-    "dayC, nightC = tools.day_night_avg(ds, sel_var='CHLA',start_time = '2023-06-03', end_time = '2023-06-07')"
+    "dayT, nightT = tools.day_night_avg(ds, sel_var='TEMP')\n",
+    "dayS, nightS = tools.day_night_avg(ds, sel_var='PSAL')\n",
+    "dayC, nightC = tools.day_night_avg(ds, sel_var='CHLA')"
    ]
   },
   {
@@ -302,9 +302,9 @@
    "source": [
     "fig, ax = plt.subplots(1, 3, figsize=(15, 5))\n",
     "\n",
-    "tools.plot_daynight_avg( dayT, nightT, ax[0],sel_day='2023-06-04', xlabel='Temperature [C]')\n",
-    "tools.plot_daynight_avg( dayS, nightS, ax[1],sel_day='2023-06-04', xlabel='Salinity [PSU]')\n",
-    "tools.plot_daynight_avg( dayC, nightC, ax[2],sel_day='2023-06-04', xlabel='Chlorophyll [mg m-3]')"
+    "tools.plot_daynight_avg( dayT, nightT, ax[0], xlabel='Temperature [C]')\n",
+    "tools.plot_daynight_avg( dayS, nightS, ax[1], xlabel='Salinity [PSU]')\n",
+    "tools.plot_daynight_avg( dayC, nightC, ax[2], xlabel='Chlorophyll [mg m-3]')"
    ]
   },
   {
@@ -340,16 +340,6 @@
     "### Photosyntetically Active Radiation (PAR)"
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "6865c233-5a34-4a9b-9ffb-6b0375ade115",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "ds_par = fetchers.load_sample_dataset(dataset_name=\"sea055_20220104T1536_delayed.nc\")"
-   ]
-  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -358,7 +348,7 @@
    "outputs": [],
    "source": [
     "fig, ax = plt.subplots(1, 1, figsize=(5, 5))\n",
-    "tools.plot_updown_bias(tools.updown_bias(ds_par, var='DPAR', v_res=1), ax,  xlabel='Irradiance')"
+    "tools.plot_updown_bias(tools.updown_bias(ds, var='DPAR', v_res=1), ax,  xlabel='Irradiance')"
    ]
   },
   {
@@ -433,14 +423,6 @@
    "source": [
     "tools.check_temporal_drift(ds, var='DOXY')"
    ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "0bc6bf18-ab77-4446-a8c4-eaecb9460b59",
-   "metadata": {},
-   "outputs": [],
-   "source": []
   }
  ],
  "metadata": {

diff --git a/tests/test_tools.py b/tests/test_tools.py
@@ -26,10 +26,10 @@ def test_quench_sequence():
     if not "TIME" in ds.indexes.keys():
         ds = ds.set_xindex('TIME')
     fig, ax = plt.subplots()
-    tools.plot_section_with_srss(ds, 'CHLA', ax,start_time = '2023-06-04', end_time = '2023-06-06', ylim=35)
-    dayT, nightT = tools.day_night_avg(ds, sel_var='TEMP',start_time = '2023-06-04', end_time = '2023-06-06')
+    tools.plot_section_with_srss(ds, 'CHLA')
+    dayT, nightT = tools.day_night_avg(ds, sel_var='TEMP')
     fig, ax = plt.subplots()
-    tools.plot_daynight_avg( dayT, nightT,ax,sel_day='2023-06-04', xlabel='Temperature [C]') 
+    tools.plot_daynight_avg(dayT, nightT, ax, xlabel='Temperature [C]') 
 
 def test_temporal_drift():
     ds = fetchers.load_sample_dataset()