Make front probability for 1thr or 2thr kinds

Compute/front_probability_total.py

@@ -5,6 +5,7 @@ For both S/N zones.
 """
 from os import path
 
+import numpy as np
 import xarray as xr
 
 import lib
@@ -14,39 +15,41 @@ import lib.data.SN_separation
 import lib.data.front_probability
 
 
-def main():
+MASKS = ['low', 'mid', 'hi']
 
-    def add_args(parser):
-        parser.add_argument('-threshold', type=float, default=5.)
-
-    args = lib.get_args(['region', 'days', 'fixes',
-                         'scale', 'number', 'coef'],
-                        add_args)
 
+def main(args):
     sst = lib.data.ostia.get_data(args)
     hi = lib.data.hi.get_data(args)
     ds = xr.merge([lib.fix_time_daily(d) for d in [sst, hi]],
                   join='inner')
 
     ds['HI'] = lib.data.hi.apply_coef(ds, lib.data.hi.get_coef(args))
-    frt = (ds.HI > args['threshold']) * 1.
-    frt = frt.where(ds.sst.notnull())
 
-    tot = frt.mean('time')
+    masks = dict(
+        low=ds.HI < args['thr_lo'],
+        mid=(ds.HI > args['thr_lo']) * (ds.HI < args['thr_hi']),
+        hi=ds.HI > args['thr_hi']
+    )
+    masks = [data.expand_dims(mask=[m]) for m, data in masks.items()]
+    ds['frt'] = xr.concat(masks, 'mask') * 1.
+    ds['frt'] = ds.frt.where(ds.sst.notnull())
+
+    tot = ds.frt.mean('time')
     var_name = 'p_frt'
     tot = tot.to_dataset(name=var_name)
 
-    args['period'] = 'total'
-    outdir = lib.data.front_probability.get_root(args)
-    lib.check_output_dir(outdir)
-
-    year = '2000'
+    ofile = lib.data.front_probability.get_filename(args)
+    lib.check_output_dir(ofile, file=True)
     lib.setup_metadata(tot, args)
-    tot.to_netcdf(path.join(outdir, year, year + '0101.nc'),
-                  encoding={var_name: {'zlib': True}})
+    tot.to_netcdf(ofile, encoding={var_name: {'zlib': True}})
 
     return tot
 
 
 if __name__ == '__main__':
-    tot = main()
+    args = lib.get_args(['region', 'days', 'fixes',
+                         'scale', 'number', 'coef',
+                         'thr_lo', 'thr_hi'])
+    args['kind'] = '2thr'
+    tot = main(args)

lib/box.py

@@ -304,12 +304,12 @@ class Image(Box):
         super().__init__(lon_min, lon_max, lat_min, lat_max)
 
     @property
-    def slice_x(self) -> [int, int]:
+    def slice_x(self):
         """Return indices slice in longitude dimension."""
         return slice(self.ix, self.ix + self.nx)
 
     @property
-    def slice_y(self) -> [int, int]:
+    def slice_y(self):
         """Return indices slice in latitude dimension."""
         return slice(self.iy, self.iy + self.ny)
 

lib/data/front_probability.py

@@ -4,19 +4,26 @@ import lib
 import lib.data
 
 
-ARGS = {'region', 'period'}
+ARGS = {'region', 'kind'}
 grid = '4km_EPSG32662'
-DEFAULTS = dict(period='monthly')
 
 
 def PREGEX(args):
-    return "%(Y)/%(Y)%(m)%(d).nc"
+    pregex = "mask_thr"
+    if args['kind'] == '1thr':
+        pregex += "_%(threshold:fmt=.2f)"
+    elif args['kind'] == '2thr':
+        pregex += "_%(thr_lo:fmt=.2f)_%(thr_hi:fmt=.2f)"
+    else:
+        raise KeyError("kind '{}' not supported.".format(args['kind']))
+
+    pregex += ".nc"
+    return pregex
 
 
 def ROOT(args):
     root = lib.data.get_default_directory(args, 'Frt_proba', time_folder=False)
-    root.append(args['period'])
     return root
 
 
-lib.data.create_dataset(__name__, PREGEX, ROOT, ARGS, defaults=DEFAULTS)
+lib.data.create_dataset(__name__, PREGEX, ROOT, ARGS)