# Path: atmos/constants.py

# 
# Path: atmos/data.py
# def biggify(small, big, tile=False):
# def collapse(arr, axis=-1):
# def nantrapz(y, x=None, axis=-1):
# def rolling_mean(data, nroll, axis=-1, center=True, **kwargs):
# def gradient(data, vec, axis=-1):
# def pres_units(units):
# def pres_convert(pres, units_in, units_out):
# def precip_units(units):
# def precip_convert(precip, units_in, units_out):
# def get_coord(data, coord_name, return_type='values'):
#     def name_options(nm):
# def subset(data, subset_dict, incl_lower=True, incl_upper=True, search=True,
#            copy=True, squeeze=False):
# def dim_mean(data, dimname, lower=None, upper=None, minfrac=0.5):
#     def one_variable(var, dimname, dimvals, minfrac):
# def ncdisp(filename, verbose=True, decode_cf=False, indent=2, width=None):
# def ncload(filename, verbose=True, unpack=True, missing_name=u'missing_value',
#            offset_name=u'add_offset', scale_name=u'scale_factor',
#            decode_cf=False):
# def load_concat(paths, var_ids=None, concat_dim='TIME', subset_dict=None,
#                 func=None, func_args=None, func_kw=None, squeeze=True, verbose=True):
#     def get_data(path, var_ids, subset_dict, func, func_args, func_kw):
# def save_nc(filename, *args):
# def mean_over_files(files, nms=None):
# def latlon_equal(data1, data2, latname1=None, lonname1=None,
#                  latname2=None, lonname2=None):
# def lon_convention(lon):
# def set_lon(data, lonmax=360, lon=None, lonname=None):
# def interp_latlon(data, lat_out, lon_out, lat_in=None, lon_in=None,
#                   checkbounds=False, masked=False, order=1):
# def mask_oceans(data, lat=None, lon=None, inlands=True, resolution='l',
#                 grid=5):
# def mean_over_geobox(data, lat1, lat2, lon1, lon2, lat=None, lon=None,
#                      area_wtd=True, land_only=False):
# def get_ps_clim(lat, lon, datafile='data/topo/ncep2_ps.nc'):
# def correct_for_topography(data, topo_ps, plev=None, lat=None, lon=None):
# def near_surface(data, pdim=-3, return_inds=False):
# def interp_plevels(data, plev_new, plev_in=None, pdim=-3, kind='linear'):
# def int_pres(data, plev=None, pdim=-3, pmin=0, pmax=1e6):
# def split_timedim(data, n, slowfast=True, timename=None, time0_name='time0',
#                   time0_vals=None, time1_name='time1', time1_vals=None):
#     def time_coord(name, size, vals, coords):
# def splitdays(days):
# def daily_from_subdaily(data, n, method='mean', timename=None, dayname='day',
#                         dayvals=None):
#     def process_one(data, n, method, timename, dayname, dayvals):
# def combine_daily_years(varnames, files, years, yearname='Year',
#                         subset_dict=None):
#     SCALE = 60 * 60 * 24
#     NMAX = 3
#     WAIT = 5


import numpy as np
import matplotlib.pyplot as plt
import xray
import atmos.utils as utils
import atmos.plots as ap
import atmos.data as dat
from atmos.constants import const as constants
from atmos.data import get_coord, divergence_spherical_2d

uq = u * q
vq = v * q

uq_int = dat.int_pres(uq, pdim=-3)
vq_int = dat.int_pres(vq, pdim=-3)

t, k = 0, 4
plt.figure(figsize=(7,8))
plt.subplot(211)
ap.pcolor_latlon(uq[t,k])
plt.subplot(212)
ap.pcolor_latlon(vq[t,k])

plt.figure(figsize=(7,8))
plt.subplot(211)
ap.pcolor_latlon(uq_int[t])
plt.subplot(212)
ap.pcolor_latlon(vq_int[t])

# ----------------------------------------------------------------------
# Moisture flux convergence

# DataArray
mfc = - divergence_spherical_2d(uq_int, vq_int)

# ndarray
mfc2 = - divergence_spherical_2d(uq_int.values, vq_int.values, lat, lon)

# components

# Path: atmos/constants.py

# 
# Path: atmos/data.py
# def biggify(small, big, tile=False):
# def collapse(arr, axis=-1):
# def nantrapz(y, x=None, axis=-1):
# def rolling_mean(data, nroll, axis=-1, center=True, **kwargs):
# def gradient(data, vec, axis=-1):
# def pres_units(units):
# def pres_convert(pres, units_in, units_out):
# def precip_units(units):
# def precip_convert(precip, units_in, units_out):
# def get_coord(data, coord_name, return_type='values'):
#     def name_options(nm):
# def subset(data, subset_dict, incl_lower=True, incl_upper=True, search=True,
#            copy=True, squeeze=False):
# def dim_mean(data, dimname, lower=None, upper=None, minfrac=0.5):
#     def one_variable(var, dimname, dimvals, minfrac):
# def ncdisp(filename, verbose=True, decode_cf=False, indent=2, width=None):
# def ncload(filename, verbose=True, unpack=True, missing_name=u'missing_value',
#            offset_name=u'add_offset', scale_name=u'scale_factor',
#            decode_cf=False):
# def load_concat(paths, var_ids=None, concat_dim='TIME', subset_dict=None,
#                 func=None, func_args=None, func_kw=None, squeeze=True, verbose=True):
#     def get_data(path, var_ids, subset_dict, func, func_args, func_kw):
# def save_nc(filename, *args):
# def mean_over_files(files, nms=None):
# def latlon_equal(data1, data2, latname1=None, lonname1=None,
#                  latname2=None, lonname2=None):
# def lon_convention(lon):
# def set_lon(data, lonmax=360, lon=None, lonname=None):
# def interp_latlon(data, lat_out, lon_out, lat_in=None, lon_in=None,
#                   checkbounds=False, masked=False, order=1):
# def mask_oceans(data, lat=None, lon=None, inlands=True, resolution='l',
#                 grid=5):
# def mean_over_geobox(data, lat1, lat2, lon1, lon2, lat=None, lon=None,
#                      area_wtd=True, land_only=False):
# def get_ps_clim(lat, lon, datafile='data/topo/ncep2_ps.nc'):
# def correct_for_topography(data, topo_ps, plev=None, lat=None, lon=None):
# def near_surface(data, pdim=-3, return_inds=False):
# def interp_plevels(data, plev_new, plev_in=None, pdim=-3, kind='linear'):
# def int_pres(data, plev=None, pdim=-3, pmin=0, pmax=1e6):
# def split_timedim(data, n, slowfast=True, timename=None, time0_name='time0',
#                   time0_vals=None, time1_name='time1', time1_vals=None):
#     def time_coord(name, size, vals, coords):
# def splitdays(days):
# def daily_from_subdaily(data, n, method='mean', timename=None, dayname='day',
#                         dayvals=None):
#     def process_one(data, n, method, timename, dayname, dayvals):
# def combine_daily_years(varnames, files, years, yearname='Year',
#                         subset_dict=None):
#     SCALE = 60 * 60 * 24
#     NMAX = 3
#     WAIT = 5


# Path: testing/testing-data-div_spher_2d.py
import numpy as np
import matplotlib.pyplot as plt
import xray
import atmos.utils as utils
import atmos.plots as ap
import atmos.data as dat
from atmos.constants import const as constants
from atmos.data import get_coord, divergence_spherical_2d


uq = u * q
vq = v * q

uq_int = dat.int_pres(uq, pdim=-3)
vq_int = dat.int_pres(vq, pdim=-3)

t, k = 0, 4
plt.figure(figsize=(7,8))
plt.subplot(211)
ap.pcolor_latlon(uq[t,k])
plt.subplot(212)
ap.pcolor_latlon(vq[t,k])

plt.figure(figsize=(7,8))
plt.subplot(211)
ap.pcolor_latlon(uq_int[t])
plt.subplot(212)
ap.pcolor_latlon(vq_int[t])

# ----------------------------------------------------------------------
# Moisture flux convergence

# DataArray
mfc = - divergence_spherical_2d(uq_int, vq_int)

# ndarray
mfc2 = - divergence_spherical_2d(uq_int.values, vq_int.values, lat, lon)

# components

mfc3, mfc_x, mfc_y = divergence_spherical_2d(uq_int, vq_int, return_comp=True)