EddyIdentification

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
===============================================================================
This file is part of py-eddy-tracker.

    py-eddy-tracker is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    py-eddy-tracker is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with py-eddy-tracker.  If not, see <http://www.gnu.org/licenses/>.

Copyright (c) 2014-2015 by Evan Mason
Email: emason@imedea.uib-csic.es
"""

import logging
from yaml import load as yaml_load
from datetime import datetime
import datetime as dt
from scipy.ndimage import gaussian_filter
from matplotlib.dates import date2num, num2julian
from matplotlib.figure import Figure
from os.path import exists
from os import mkdir
from numpy import arange, atleast_1d, round, int, interp, pi, \
    float_, errstate, ma

from py_eddy_tracker import EddyParser
from py_eddy_tracker.property_functions import \
    collection_loop, func_hann2d_fast
from py_eddy_tracker.property_objects import SwirlSpeed
from py_eddy_tracker.grid.aviso import AvisoGrid
from py_eddy_tracker.tracking_objects import IdentificationList
from py_eddy_tracker.grid import browse_dataset_in


def load_config():
    # Command Parser
    parser = EddyParser("Tool to detect eddies. "
                        "To run use EddyIdentification "
                        "eddy_tracker_configuration.yaml'")
    parser.add_argument('yaml_file',
                        help='Yaml file to configure py-eddy-tracker')
    opts = parser.parse_args()

    # Read yaml configuration file
    with open(opts.yaml_file, 'r') as stream:
        config = yaml_load(stream)
    if config is None:
        raise Exception('Yaml file empty')
    logging.info('Launching with yaml file: %s', opts.yaml_file)

    # Create ouput directory if needed
    save_dir = config['PATHS']['SAVE_DIR']
    logging.info('Outputs saved to %(SAVE_DIR)s', config['PATHS'])
    if not exists(save_dir):
        mkdir(save_dir)

    diagnostic_type = config['DIAGNOSTIC_TYPE']
    if diagnostic_type not in ['SLA']:
        raise Exception('Unknown diagnostic %s', diagnostic_type)

    auto = False
    if 'DOMAIN' not in config:
        auto = True
        config['DATE_STR'] = None
        config['DATE_END'] = None
        config['THE_DOMAIN'] = None
        config['lonmin'] = None
        config['lonmax'] = None
        config['latmin'] = None
        config['latmax'] = None
    else:
        config['THE_DOMAIN'] = config['DOMAIN'].get('THE_DOMAIN', 'Auto')

    # It is not recommended to change values given below
    # for 'Global', 'BlackSea' or 'MedSea'...
    if not auto:
        if 'Global' in config['THE_DOMAIN']:
            config['lonmin'] = -100.
            config['lonmax'] = 290.
            config['latmin'] = -80.
            config['latmax'] = 80.

        elif config['THE_DOMAIN'] == 'Regional':
            config['lonmin'] = config['DOMAIN']['LONMIN']
            config['lonmax'] = config['DOMAIN']['LONMAX']
            config['latmin'] = config['DOMAIN']['LATMIN']
            config['latmax'] = config['DOMAIN']['LATMAX']

        config['DATE_STR'] = config['DOMAIN']['DATE_STR']
        config['DATE_END'] = config['DOMAIN']['DATE_END']
        if config['DATE_STR'] > config['DATE_END']:
            raise Exception('DATE_END must be larger than DATE_STR')

    if 'SLA' in diagnostic_type:
        max_sla = config['CONTOUR_PARAMETER'
        ]['CONTOUR_PARAMETER_SLA']['MAX_SLA']
        interval = config['CONTOUR_PARAMETER'
        ]['CONTOUR_PARAMETER_SLA']['INTERVAL']
        config['CONTOUR_PARAMETER'] = arange(-max_sla, max_sla + interval,
                                             interval)

    smoothing = config['SMOOTHING']
    if smoothing:
        if 'SLA' in diagnostic_type:
            config['SMOOTHING_SLA']['ZWL'] = atleast_1d(config['SMOOTHING_SLA']['ZWL'])
            config['SMOOTHING_SLA']['MWL'] = atleast_1d(config['SMOOTHING_SLA']['MWL'])
    return config


if __name__ == '__main__':
    CONFIG = load_config()
    DATASET_LIST = browse_dataset_in(
        data_dir=CONFIG['DATASET']['DATA_DIR'],
        files_model=CONFIG['DATASET']['FILES_MODEL'],
        date_regexp=CONFIG['DATASET']['DATE_REGEXP'],
        date_model=CONFIG['DATASET']['DATE_MODEL'],
        sub_sampling_step=CONFIG['DATASET']['SUBSAMPLING'],
        start_date=CONFIG['DATE_STR'],
        end_date=CONFIG['DATE_END'],
    )
    DIMENSIONS = CONFIG['DATASET'].get('DIMENSIONS', None)
    # End user configuration setup options
    # ----------------------------------------------------------

    # Set up a grid object using first AVISO file in the list
    SLA_GRD = AvisoGrid(DATASET_LIST[0]['filename'], CONFIG['THE_DOMAIN'],
                        CONFIG['lonmin'], CONFIG['lonmax'],
                        CONFIG['latmin'], CONFIG['latmax'],
                        grid_name=CONFIG['DATASET']['VAR_NAME'],
                        lon_name=CONFIG['DATASET']['LON_NAME'],
                        lat_name=CONFIG['DATASET']['LAT_NAME'])

    if 'Gaussian' in CONFIG['SMOOTHING_SLA']['TYPE']:
        # Get parameters for gaussian_filter
        ZRES, MRES = SLA_GRD.gaussian_resolution(
            CONFIG['SMOOTHING_SLA']['ZWL'],
            CONFIG['SMOOTHING_SLA']['MWL'])

    # See Chelton section B2 (0.4 degree radius)
    # These should give 8 and 1000 for 0.25 deg resolution
    CONFIG['PIXMIN'] = round((pi * CONFIG['RADMIN'] ** 2) /
                             SLA_GRD.resolution ** 2)
    CONFIG['PIXMAX'] = round((pi * CONFIG['RADMAX'] ** 2) /
                             SLA_GRD.resolution ** 2)
    logging.info('Pixel range = %(PIXMIN)d-%(PIXMAX)d', CONFIG)

    # Figure to get contours of parameter
    FIG = Figure()
    CONT_AX = FIG.add_subplot(111)

    START_TIME = datetime.now()

    logging.info('Start identification')

    # Loop over grid
    for filename, date in DATASET_LIST:
        date = date.astype(dt.date)
        # Extract grid date
        SLA_GRD.grid_date = date

        # Initialise two eddy objects to hold data
        A_EDDY = IdentificationList('Anticyclonic', SLA_GRD, date, **CONFIG)
        C_EDDY = IdentificationList('Cyclonic', SLA_GRD, date, **CONFIG)

        logging.info('AVISO_FILE : %s', filename)

        sla = SLA_GRD.get_aviso_data(filename, dimensions=DIMENSIONS)
        SLA_GRD.set_mask(sla)
        SLA_GRD.uvmask()

        if CONFIG['SMOOTHING']:
            if 'Gaussian' in CONFIG['SMOOTHING_SLA']['TYPE']:
                logging.info('applying Gaussian high-pass filter')
                # Set landpoints to zero
                m = SLA_GRD.mask
                sla[SLA_GRD.mask] = 0.
                if hasattr(sla, 'data'):
                    m += sla.data == SLA_GRD.fillval
                    sla[sla.data == SLA_GRD.fillval] = 0.
                # High pass filter, see
                # http://stackoverflow.com/questions/6094957/high-pass-filter-
                # for-image-processing-in-python-by-using-scipy-numpy
                v = gaussian_filter(sla, sigma=[MRES, ZRES])
                w = gaussian_filter(float_(~m), sigma=[MRES, ZRES])

                with errstate(invalid='ignore'):
                    sla -= ma.array(v / w, mask=(w == 0) + m)

                logging.info('applying Gaussian high-pass filter -- OK')

            elif 'Hanning' in SMOOTHING_TYPE:
                logging.info('applying %s passes of Hanning filter',
                             SMOOTH_FAC)
                # Do SMOOTH_FAC passes of 2d Hanning filter
                sla = func_hann2d_fast(sla, SMOOTH_FAC)
                logging.info('applying %s passes of Hanning filter -- OK')

            else:
                raise Exception('Filter unknown : %s', SMOOTHING_TYPE)

        # Apply the landmask
        if SLA_GRD.mask is not None:
            sla.mask += SLA_GRD.mask

        # Multiply by 0.01 for m
        SLA_GRD.set_geostrophic_velocity(sla * 0.01)

        # Remove padded boundary
        sla = sla[SLA_GRD.view_unpad]

        # Calculate EKE
        SLA_GRD.get_eke()

        if 'SLA' in CONFIG['DIAGNOSTIC_TYPE']:
            A_EDDY.sla = sla.copy()
            C_EDDY.sla = sla.copy()

        # Get scalar speed
        A_EDDY.uspd = SLA_GRD.uspd
        C_EDDY.uspd = SLA_GRD.uspd

        # Set interpolation coefficients
        SLA_GRD.set_interp_coeffs(sla, SLA_GRD.uspd)
        A_EDDY.sla_coeffs = SLA_GRD.sla_coeffs
        A_EDDY.uspd_coeffs = SLA_GRD.uspd_coeffs
        C_EDDY.sla_coeffs = SLA_GRD.sla_coeffs
        C_EDDY.uspd_coeffs = SLA_GRD.uspd_coeffs

        if 'SLA' in CONFIG['DIAGNOSTIC_TYPE']:
            logging.info('Processing SLA contours for eddies')
            MIN_SLA, MAX_SLA = A_EDDY.sla.min(), A_EDDY.sla.max()
            if (MIN_SLA < A_EDDY.contour_parameter[0]) or (
                        MAX_SLA > A_EDDY.contour_parameter[-1]):
                logging.warning('SLA values [%f,%f] is larger than levels %s',
                                MIN_SLA, MAX_SLA, A_EDDY.contour_parameter)

            # we reduce levels if don't needed
            i_min, i_max = interp(
                (MIN_SLA, MAX_SLA),
                A_EDDY.contour_parameter,
                arange(len(A_EDDY.contour_parameter))
            )
            i_min, i_max = max(0, int(i_min - 2)), int(i_max + 2)
            A_EDDY.contour_parameter = A_EDDY.contour_parameter[i_min:i_max]
            C_EDDY.contour_parameter = C_EDDY.contour_parameter[i_min:i_max]

            CONTOURS = CONT_AX.contour(SLA_GRD.lon, SLA_GRD.lat, A_EDDY.sla,
                                       levels=A_EDDY.contour_parameter)
            # Note that C_CS is in reverse order
            logging.info('Processing SLA contours for eddies -- Ok')

        # Set contour coordinates and indices for calculation of
        # speed-based radius
        A_EDDY.swirl = SwirlSpeed(CONTOURS)
        C_EDDY.swirl = A_EDDY.swirl

        # Now we loop over the CS collection
        if 'SLA' in CONFIG['DIAGNOSTIC_TYPE']:
            logging.info("Anticyclonic research")
            collection_loop(CONTOURS, SLA_GRD, A_EDDY)
            logging.info("Cyclonic research")
            collection_loop(CONTOURS, SLA_GRD, C_EDDY)

        grd_int_date = num2julian(date2num(date)) + 0.5
        A_EDDY.observations.obs['time'] = grd_int_date
        C_EDDY.observations.obs['time'] = grd_int_date
        # clear the current axis
        CONT_AX.cla()
        # Save file
        A_EDDY.write_netcdf(path=CONFIG['PATHS']['SAVE_DIR'])
        C_EDDY.write_netcdf(path=CONFIG['PATHS']['SAVE_DIR'])
        del CONTOURS, A_EDDY, C_EDDY

    # Total running time
    logging.info('Mean duration by loop : %s',
                 (datetime.now() - START_TIME) / len(DATASET_LIST))
    logging.info('Duration : %s', datetime.now() - START_TIME)
    logging.info('Outputs saved to %s', CONFIG['PATHS']['SAVE_DIR'])