EddyId

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Do identification
"""
from datetime import datetime
from netCDF4 import Dataset
from py_eddy_tracker import EddyParser
from py_eddy_tracker.dataset.grid import RegularGridDataset, UnRegularGridDataset
import zarr

def id_parser():
    parser = EddyParser('Eddy Identification')
    parser.add_argument('filename')
    parser.add_argument('datetime')
    parser.add_argument('h')
    parser.add_argument('u', help='If it s None, it will be deduce from h')
    parser.add_argument('v', help='If it s None, it will be deduce from h')
    parser.add_argument('longitude')
    parser.add_argument('latitude')
    parser.add_argument('path_out')
    parser.add_argument('--cut_wavelength', default=500, type=float,
                        help='Wavelength for mesoscale filter in km')
    parser.add_argument('--filter_order', default=3, type=int)
    parser.add_argument('--isoline_step', default=0.002, type=float,
                        help='Step between 2 isoline in m')
    parser.add_argument('--fit_errmax', default=55, type=float,
                        help='Error max accepted to fit circle in percent')
    parser.add_argument('--height_unit', default=None, type=str,
                        help='Force height unit')
    parser.add_argument('--speed_unit', default=None, type=str,
                        help='Force speed unit')
    parser.add_argument('--unregular', action='store_true', help='if grid is unregular')
    parser.add_argument('--zarr',
                        action='store_true',
                        help='Output will be wrote in zarr')
    return parser


if __name__ == '__main__':
    args = id_parser().parse_args()
    grid_class = (UnRegularGridDataset if args.unregular else RegularGridDataset)
    h = grid_class(args.filename, args.longitude, args.latitude)
    date = datetime.strptime(args.datetime, '%Y%m%d')
    if args.u == 'None' and args.v == 'None':
        h.add_uv(args.h)
        u, v = 'u', 'v'
    else:
        u, v = args.u, args.v
    if args.cut_wavelength != 0:
        h.bessel_high_filter(args.h, args.cut_wavelength, order=args.filter_order)
    a, c = h.eddy_identification(args.h, u, v, date, args.isoline_step, pixel_limit=(5, 2000),
                                 shape_error=args.fit_errmax, force_height_unit=args.height_unit,
                                 force_speed_unit=args.speed_unit)
    if args.zarr:
        h = zarr.open(args.path_out + date.strftime('/Anticyclonic_%Y%m%d.zarr'), 'w')
        a.to_zarr(h)
        h = zarr.open(args.path_out + date.strftime('/Cyclonic_%Y%m%d.zarr'), 'w')
        c.to_zarr(h)
    else:
        with Dataset(args.path_out + date.strftime('/Anticyclonic_%Y%m%d.nc'), 'w') as h:
            a.to_netcdf(h)
        with Dataset(args.path_out + date.strftime('/Cyclonic_%Y%m%d.nc'), 'w') as h:
            c.to_netcdf(h)