Seasonal means from monthly data

monseasmean.py

import datetime
import xray
import pandas as pd
import numpy as np
from collections import OrderedDict

dpm = {'noleap': [0, 31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31],
       '365_day': [0, 31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31],
       'standard': [0, 31, 28.25, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31],
       'gregorian': [0, 31, 28.25, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31],
       'proleptic_gregorian': [0, 31, 28.25, 31, 30, 31, 30, 31, 31, 30,
                               31, 30, 31],
       'all_leap': [0, 31, 29, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31],
       '366_day': [0, 31, 29, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31],
       '360_day': [0, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30, 30]}

seas = {12: 'DJF', 1: 'DJF', 2: 'DJF',
        3: 'MAM', 4: 'MAM', 5: 'MAM',
        6: 'JJA', 7: 'JJA', 8: 'JJA',
        9: 'SON', 10: 'SON', 11: 'SON'}

seasons = ('DJF', 'MAM', 'JJA', 'SON')


def make_new_monthly_index(old_index):
    """
    Make a new Pandas datetime index with each record coresponding to the
    first day of the month.
    """
    ns = 1e-9
    t0 = old_index.values[0]
    if type(t0) == tuple:
        d1 = t0[1]
    else:
        d1 = datetime.datetime.utcfromtimestamp(t0.astype(int)*ns)
    return pd.date_range(start=d1.strftime("%Y-%m-01"), periods=len(old_index), freq='M')


def monseasmean(data, calendar='standard', start=None, end=None):
    """
    Calculate seasonal means (DJF, MAM, JJA, SON) from an xray Dataset of
    monthly means

    inputs:     data (xray Dataset time dimension)
    returns:    seasonal means (xray Dataset with 4 time values)
    """

    # determine time slice to average over
    time = make_new_monthly_index(data['time'])
    a, b = time.slice_locs(start, end)
    time_inds = np.arange(a, b)
    time = time[a:b][:]
    ntime = b - a

    # create weights for time slice
    seas_inds = OrderedDict()
    weights = np.empty(ntime, dtype=np.float)
    ds_seas = np.chararray(ntime, itemsize=3)
    for i, m in enumerate(time.month):
        weights[i] = dpm[calendar][m]
        ds_seas[i] = seas[m]
    for s in seasons:
        inds = np.nonzero(ds_seas==s).squeeze()
        weights /= weights[inds].sum()
        seas_inds[s] = inds

    variables = OrderedDict()
    data = data.isel(time=time_inds)
    for name, da in data.vars.items():
        if 'time' in da.coords:
            new = []
            if da.ndim == 3:
                da = weights[:, np.newaxis, np.newaxis] * da
            elif da.ndim == 4:
                da = weights[:, np.newaxis, np.newaxis, np.newaxis] * da
            elif da.ndim == 2:
                da = weights[:, np.newaxis] * da
            for i, s in enumerate(seasons):
                new.append(da.isel(time=seas_inds[s]).sum(dim='time', keep_attrs=True))
            variables[name] = xray.concat(new, dim='season')
    final = xray.Dataset(variables=variables, attrs={'history': 'created by monseasmean'})
    final['season'].values = seasons

    return final