Add min_lead to similarity

unseen/similarity.py

@@ -35,6 +35,7 @@ def similarity_tests(
     fcst,
     obs,
     var,
+    min_lead=None,
     lead_dim="lead_time",
     init_dim="init_date",
     ensemble_dim="ensemble",
@@ -51,6 +52,8 @@ def similarity_tests(
         Observational/comparison dataset with a time dimension.
     var : str
         Variable from the datasets to process.
+    min_lead : int, optional
+        Minimum lead time
     init_dim: str, default 'init_date'
         Name of the initial date dimension in fcst
     lead_dim: str, default 'lead_time'
@@ -73,6 +76,9 @@ def similarity_tests(
     that the two samples are from the same population.
     """
 
+    if min_lead is not None:
+        fcst = fcst.where(fcst[lead_dim] >= min_lead)
+
     if isinstance(fcst, xr.DataArray):
         fcst = fcst.to_dataset()
     if isinstance(obs, xr.DataArray):
@@ -181,6 +187,12 @@ def _parse_command_line():
         default="lead_time",
         help="Name of lead time dimension",
     )
+    parser.add_argument(
+        "--min_lead",
+        type=int,
+        default=None,
+        help="Minimum lead time",
+    )
     parser.add_argument(
         "--by_lead",
         action="store_true",
@@ -211,6 +223,7 @@ def _main():
         ds_fcst,
         ds_obs,
         args.var,
+        min_lead=args.min_lead,
         init_dim=args.init_dim,
         lead_dim=args.lead_dim,
         ensemble_dim=args.ensemble_dim,

unseen/stability.py

@@ -86,24 +86,29 @@ def plot_dist_by_time(ax, sample_da, metric, start_years):
     ax.legend()
 
 
-def return_curve(data, method):
+def return_curve(data, method, params=[]):
     """Return x and y data for a return period curve.
 
     Parameters
     ----------
     data : xarray DataArray
     method : {'gev', 'empirical'}
         Fit a GEV or not to data
+    params : list, default None
+        shape, location and scale parameters (calculated if None)
     """
 
-    if method == "gev":
+    if method == "empirical":
+        return_values = np.sort(data, axis=None)[::-1]
+        return_periods = len(data) / np.arange(1.0, len(data) + 1.0)
+    else:
         return_periods = np.logspace(0, 4, num=10000)
         probabilities = 1.0 / return_periods
-        shape, loc, scale = indices.fit_gev(data, generate_estimates=True)
+        if params:
+            shape, loc, scale = params
+        else:
+            shape, loc, scale = indices.fit_gev(data, generate_estimates=True)
         return_values = gev.isf(probabilities, shape, loc, scale)
-    elif method == "empirical":
-        return_values = np.sort(data, axis=None)[::-1]
-        return_periods = len(data) / np.arange(1.0, len(data) + 1.0)
 
     return return_periods, return_values
 
@@ -187,6 +192,7 @@ def plot_return_by_lead(ax, sample_da, metric, method, uncertainty=False, lead_d
     ax.set_xlabel("return period (years)")
     ax.set_ylabel(sample_da.attrs["units"])
     ax.legend()
+    ax.set_ylim((50, None))
 
 
 def plot_return_by_time(ax, sample_da, metric, start_years, method, uncertainty=False):
@@ -245,6 +251,7 @@ def plot_return_by_time(ax, sample_da, metric, start_years, method, uncertainty=
     ax.set_xscale("log")
     ax.set_xlabel("return period (years)")
     ax.set_ylabel(sample_da.attrs["units"])
+    ax.set_ylim((50, None))
     ax.legend()