def __init__(self, var, daxis, dx=None):
# {{{
''' __init__()'''
from pygeode.var import Var
self.daxis = daxis = var.whichaxis(daxis)
assert var.shape[
daxis] > 1, "need at least two values along differentiation axis"
if dx is not None:
if dx.naxes == 1:
assert dx.shape[0] == var.shape[daxis]
self.dx = dx.replace_axes(newaxes=(var.axes[daxis], ))
else:
assert all([dx.hasaxis(a) for a in var.axes])
self.dx = dx
else:
self.dx = var.axes[daxis]
self.var = var
# Construct new variable
if var.name != '':
name = 'd2' + var.name
else:
name = 'd2(UnknownVar)'
Var.__init__(self,
var.axes,
var.dtype,
name=name,
atts=var.atts,
plotatts=var.plotatts)
def __init__ (self, var, slices):
# {{{
from pygeode.var import Var, copy_meta
self.var = var
copy_meta (var, self)
#TODO: remove degenerate dimensions when slicing by integer values
if not hasattr(slices,'__len__'): slices = [slices]
# assert len(slices) == len(var.axes), "expected %i parameters, received %i."%(len(var.axes),len(slices))
slices = list(slices) # make a copy of the list
# Append an implicit Ellipsis at the end (makes the logic a big simpler below)
# if Ellipsis not in slices: slices.append(Ellipsis)
if not any (sl is Ellipsis for sl in slices):
slices.append(Ellipsis)
# Handle Ellipsis argument
# assert slices.count(Ellipsis) == 1, "can't handle more than one Ellipsis argument"
ellipsis_index = [i for i,sl in enumerate(slices) if sl is Ellipsis]
assert len(ellipsis_index) == 1, "can't handle more than one Ellipsis argument"
num_missing = var.naxes - len(slices) + 1
assert num_missing >= 0, "too many slices provided"
# i = slices.index(Ellipsis)
i = ellipsis_index.pop()
slices = slices[:i] + [slice(None)]*num_missing + slices[i+1:]
# Slice the output axes
axes = [a.slice[s] for a,s in zip(var.axes,slices)]
Var.__init__(self, axes, dtype=var.dtype, atts=self.atts, plotatts=self.plotatts)
def __init__ (self, v1, opener, files, faxis, timedict):
# {{{
# v1 - var chunk from the first file
# opener - method for opening a file given a filename
# files - list of filenames
# faxis - starting time of each file, as a time axis
# timedict - dictionary of pre-constructed time axes
from pygeode.var import Var
from pygeode.timeaxis import Time
self.opener = opener
self.files = files
self.faxis = faxis
n = v1.getaxis(Time).name if v1.hasaxis(Time) else None
taxis = timedict[n]
T = type(taxis)
axes = list(v1.axes)
# Replace the existing time axis?
# (look for either a Time axis or a generic axis with the name 'time')
if v1.hasaxis(T): axes[v1.whichaxis(T)] = taxis
elif v1.hasaxis('time'): axes[v1.whichaxis('time')] = taxis
else: axes = [taxis] + axes
Var.__init__ (self, axes, dtype=v1.dtype, name=v1.name, atts=v1.atts, plotatts=v1.plotatts)
def __init__ (self, var, saxis, kernel, fft):
# {{{
''' __init__()'''
import numpy as np
assert len(kernel) <= var.shape[saxis], 'Kernel must not be longer than dimension being smoothed.'
# Construct new variable
self.saxis = saxis
self.var = var
self.kernel = kernel
self.fft = fft
self.klen = len(kernel)
# Normalize and reshape kernel
self.kernel /= np.sum(self.kernel)
self.kernel.shape = [self.klen if i == saxis else 1 for i in range(var.naxes)]
# Determine which convolution function to use
from scipy import signal as sg
tdata = np.ones(len(kernel), 'd')
if self.fft:
try:
sg.fftconvolve(tdata, kernel, 'same', old_behaviour=False)
self._convolve = lambda x, y, z: sg.fftconvolve(x, y, z, old_behaviour=False)
except TypeError:
self._convolve = sg.fftconvolve
else:
try:
sg.convolve(tdata, kernel, 'same', old_behaviour=False)
self._convolve = lambda x, y, z: sg.convolve(x, y, z, old_behaviour=False)
except TypeError:
self._convolve = sg.convolve
Var.__init__(self, var.axes, var.dtype, name=var.name, atts=var.atts, plotatts=var.plotatts)
def __init__ (self, var, inner = -1, naxis=None):
#{{{
from numpy import concatenate
if inner == -1:
if isinstance(var, CompositeVar):
out = var.whichaxis(var.event.__class__)
inner = var.whichaxis(var.offset.__class__)
stride = len(var.axes[inner])
vals = concatenate([e + var.axes[inner].values for e in var.axes[out].values])
else:
raise NotImplementedError("You must specify which axis to flatten")
elif inner == 0: raise NotImplementedError("inner axis must not be the outermost")
else:
out = inner - 1
stride = len(var.axes[inner])
vals = concatenate([i*stride + var.axes[inner].values for i in range(len(var.axes[out]))])
if naxis is None:
if isinstance(var.axes[out], NamedAxis):
naxis = var.axes[out].__class__(vals, var.axes[out].name)
else:
naxis = var.axes[out].__class__(vals)
axes = var.axes[:out] + (naxis,) + var.axes[inner+1:]
self.iout = out
self.iin = inner
self.stride = stride
self.source_var = var
self.name = var.name
Var.__init__(self, axes, var.dtype)
def __init__ (self, var, daxis, dx=None, df='centre'):
# {{{
''' __init__()'''
from pygeode.var import Var
self.daxis = daxis = var.whichaxis(daxis)
assert var.shape[daxis] > 1, "need at least two values along differentiation axis"
if dx is not None:
if dx.naxes == 1:
assert dx.shape[0] == var.shape[daxis]
self.dx = dx.replace_axes(newaxes=(var.axes[daxis],))
else:
assert all([dx.hasaxis(a) for a in var.axes])
self.dx = dx
else:
self.dx = var.axes[daxis]
assert df in ['centre', 'left', 'right']
self.df = df
self.var = var
# Construct new variable
if var.name != '':
name = 'd' + var.name
else:
name = 'd(UnknownVar)'
Var.__init__(self, var.axes, var.dtype, name=name, atts=var.atts, plotatts=var.plotatts)
def __init__(self, var, iaxis, lags, reverse=False):
# {{{
import numpy as np
from pygeode import Var
from pygeode.timeaxis import Time
self.iaxis = var.whichaxis(iaxis)
taxis = var.axes[self.iaxis]
assert isinstance(taxis, Time), 'must specify a Time axis'
delt = (taxis.values[1] - taxis.values[0])
if reverse:
delt = -delt
lags = lags[::-1]
self.lags = np.array(lags).astype('i')
lag = Lag(values=delt * self.lags,
units=taxis.units,
startdate={'day': 0})
axes = var.axes[:self.iaxis + 1] + (lag, ) + var.axes[self.iaxis + 1:]
self.var = var
Var.__init__(self,
axes,
dtype=var.dtype,
name=var.name,
atts=var.atts,
plotatts=var.plotatts)
def __init__(self, f, varid):
# {{{
from pygeode.var import Var
from ctypes import c_int, byref, create_string_buffer
from warnings import warn
self._f = f
self._varid = varid
assert f.fileid.value != -1
name = create_string_buffer(NC_MAX_NAME + 1)
vtype = c_int()
ndims = c_int()
dimids = (c_int * NC_MAX_VAR_DIMS)()
natts = c_int()
ier = lib.nc_inq_var(f.fileid, varid, name, byref(vtype), byref(ndims),
dimids, byref(natts))
assert ier == 0
self._vtype = vtype = vtype.value
dtype = numpy_type[vtype]
self._dimids = dimids = [dimids[j] for j in range(ndims.value)]
name = str(name.value.decode())
# Load attributes
atts = get_attributes(f.fileid, varid)
# Axes (just generic dimensions right now, until some filter is applied)
axes = [makedim(f, d) for d in dimids]
Var.__init__(self, axes, name=name, dtype=dtype, atts=atts)
def __init__(self, f, varid):
# {{{
from pygeode.var import Var
from ctypes import c_int, byref, create_string_buffer
from warnings import warn
self._f = f
self._varid = varid
assert f.fileid.value != -1
name = create_string_buffer(NC_MAX_NAME+1)
vtype = c_int()
ndims = c_int()
dimids = (c_int * NC_MAX_VAR_DIMS)()
natts = c_int()
ier = lib.nc_inq_var (f.fileid, varid, name, byref(vtype), byref(ndims), dimids, byref(natts))
assert ier == 0
self._vtype = vtype = vtype.value
dtype = numpy_type[vtype]
self._dimids = dimids = [dimids[j] for j in range(ndims.value)]
name = str(name.value.decode())
# Load attributes
atts = get_attributes (f.fileid, varid)
# Axes (just generic dimensions right now, until some filter is applied)
axes = [makedim(f,d) for d in dimids]
Var.__init__(self, axes, name=name, dtype=dtype, atts=atts)
def __init__ (self, var, slices):
# {{{
from pygeode.var import Var, copy_meta
self.var = var
copy_meta (var, self)
#TODO: remove degenerate dimensions when slicing by integer values
if not hasattr(slices,'__len__'): slices = [slices]
# assert len(slices) == len(var.axes), "expected %i parameters, received %i."%(len(var.axes),len(slices))
slices = list(slices) # make a copy of the list
# Append an implicit Ellipsis at the end (makes the logic a big simpler below)
# if Ellipsis not in slices: slices.append(Ellipsis)
if not any (sl is Ellipsis for sl in slices):
slices.append(Ellipsis)
# Handle Ellipsis argument
# assert slices.count(Ellipsis) == 1, "can't handle more than one Ellipsis argument"
ellipsis_index = [i for i,sl in enumerate(slices) if sl is Ellipsis]
assert len(ellipsis_index) == 1, "can't handle more than one Ellipsis argument"
num_missing = var.naxes - len(slices) + 1
assert num_missing >= 0, "too many slices provided"
# i = slices.index(Ellipsis)
i = ellipsis_index.pop()
slices = slices[:i] + [slice(None)]*num_missing + slices[i+1:]
# Slice the output axes
axes = [a.slice[s] for a,s in zip(var.axes,slices)]
Var.__init__(self, axes, dtype=var.dtype, atts=self.atts, plotatts=self.plotatts)
def __init__ (self, var, saxis, kernel, fft):
# {{{
''' __init__()'''
import numpy as np
assert len(kernel) <= var.shape[saxis], 'Kernel must not be longer than dimension being smoothed.'
# Construct new variable
self.saxis = saxis
self.var = var
self.kernel = kernel
self.fft = fft
self.klen = len(kernel)
# Normalize and reshape kernel
self.kernel /= np.sum(self.kernel)
self.kernel.shape = [self.klen if i == saxis else 1 for i in range(var.naxes)]
# Determine which convolution function to use
from scipy import signal as sg
tdata = np.ones(len(kernel), 'd')
if self.fft:
try:
sg.fftconvolve(tdata, kernel, 'same', old_behaviour=False)
self._convolve = lambda x, y, z: sg.fftconvolve(x, y, z, old_behaviour=False)
except TypeError:
self._convolve = sg.fftconvolve
else:
try:
sg.convolve(tdata, kernel, 'same', old_behaviour=False)
self._convolve = lambda x, y, z: sg.convolve(x, y, z, old_behaviour=False)
except TypeError:
self._convolve = sg.convolve
Var.__init__(self, var.axes, var.dtype, name=var.name, atts=var.atts, plotatts=var.plotatts)
def __init__ (self, var, axis, n):
# {{{
''' __init__()'''
from pygeode.var import Var
df = 'right' # Hard-coded to match numpy behaviour.
# May be extended in the future?
self.daxis = daxis = var.whichaxis(axis)
assert var.shape[daxis] > n, "need at least %d value(s) along difference axis"%n
self.n = n
self.df = df
self.var = var
# Construct new variable
if var.name != '':
name = 'd' + var.name
else:
name = 'd(UnknownVar)'
axes = list(var.axes)
if df == 'left':
axes[daxis] = axes[daxis].slice[n:]
else:
axes[daxis] = axes[daxis].slice[:-n]
Var.__init__(self, axes, var.dtype, name=name, atts=var.atts, plotatts=var.plotatts)
def __init__(self, var, axisdict={}, ignore_mismatch=False, newaxes=None, keep_old_name=True, **kwargs):
from pygeode.var import Var, copy_meta
from inspect import isclass
axisdict = dict(axisdict, **kwargs)
if newaxes is None:
newaxes = list(var.axes)
else:
assert len(newaxes) == var.naxes, "wrong number of axes provided"
for a,newa in axisdict.items():
if not var.hasaxis(a) and ignore_mismatch: continue
i = var.whichaxis(a)
olda = var.axes[i]
# Keep the old axis name?
name = olda.name if keep_old_name else newa.name
# Convert axis class to axis object, using the existing values?
if isclass(newa):
# Cram in any 'auxiliary' attributes, in case they're needed by the new class.
# (Needed if, say, converting from StandardTime to ModelTime365)
# Note: even if these attributes aren't pick up by the new init,
# they'll get stored in the 'auxatts' field and stay there as benign,
# unused values. Ideally, if we knew ahead of time what attributes are
# needed, we could pass only *those* attributes to the new class...
newa = newa(olda.values, name=name, **olda.auxatts)
# Use this new axis
newaxes[i] = newa
for a1, a2 in zip(newaxes, var.axes): assert len(a1) == len(a2)
self.var = var
Var.__init__(self, newaxes, dtype=var.dtype)
copy_meta (var, self)
def __init__ (self, var, indices):
# {{{
from pygeode.var import Var
import numpy as np
from pygeode.tools import combine_axes, common_dtype
# Are we given a list of variables to work on in parallel?
if isinstance(var,(tuple,list)):
axes = combine_axes(var)
dtype = common_dtype(var)
else:
axes = var.axes
dtype = var.dtype
# if not isinstance(indices,(list,tuple)): indices = [indices]
indices = np.sort([var.whichaxis(i) for i in indices])
assert len(indices) > 0, "no reduction axes specified"
N = [len(axes[i]) for i in indices]
# Check for degenerate reductions (ill-defined)
for i,n in enumerate(N):
if n == 0: raise ValueError("Can't do a reduction over axis '%s' - length is 0."%axes[i].name)
N = int(np.product(N))
self.N = N # number of values to reduce over
self.var = var
self.indices = indices
self.in_axes = axes
# Remove the reduction axis from the output variable
axes = [a for i,a in enumerate(axes) if i not in indices]
Var.__init__(self, axes, dtype=dtype, name=var.name, atts=var.atts, plotatts=var.plotatts)
def __init__ (self, var, indices):
# {{{
from pygeode.var import Var
import numpy as np
from pygeode.tools import combine_axes, common_dtype
# Are we given a list of variables to work on in parallel?
if isinstance(var,(tuple,list)):
axes = combine_axes(var)
dtype = common_dtype(var)
else:
axes = var.axes
dtype = var.dtype
# if not isinstance(indices,(list,tuple)): indices = [indices]
indices = np.sort([var.whichaxis(i) for i in indices])
assert len(indices) > 0, "no reduction axes specified"
N = [len(axes[i]) for i in indices]
# Check for degenerate reductions (ill-defined)
for i,n in enumerate(N):
if n == 0: raise ValueError("Can't do a reduction over axis '%s' - length is 0."%axes[i].name)
N = int(np.product(N))
self.N = N # number of values to reduce over
self.var = var
self.indices = indices
self.in_axes = axes
# Remove the reduction axis from the output variable
axes = [a for i,a in enumerate(axes) if i not in indices]
Var.__init__(self, axes, dtype=dtype, name=var.name, atts=var.atts, plotatts=var.plotatts)
def __init__(self, var, axisdict={}, ignore_mismatch=False, newaxes=None, keep_old_name=True, **kwargs):
from pygeode.var import Var, copy_meta
from inspect import isclass
axisdict = dict(axisdict, **kwargs)
if newaxes is None:
newaxes = list(var.axes)
else:
assert len(newaxes) == var.naxes, "wrong number of axes provided"
for a,newa in axisdict.items():
if not var.hasaxis(a) and ignore_mismatch: continue
i = var.whichaxis(a)
olda = var.axes[i]
# Keep the old axis name?
name = olda.name if keep_old_name else newa.name
# Convert axis class to axis object, using the existing values?
if isclass(newa):
# Cram in any 'auxiliary' attributes, in case they're needed by the new class.
# (Needed if, say, converting from StandardTime to ModelTime365)
# Note: even if these attributes aren't pick up by the new init,
# they'll get stored in the 'auxatts' field and stay there as benign,
# unused values. Ideally, if we knew ahead of time what attributes are
# needed, we could pass only *those* attributes to the new class...
newa = newa(olda.values, name=name, **olda.auxatts)
# Use this new axis
newaxes[i] = newa
for a1, a2 in zip(newaxes, var.axes): assert len(a1) == len(a2)
self.var = var
Var.__init__(self, newaxes, dtype=var.dtype)
copy_meta (var, self)
def clim_detrend(var, yrlen, itime=-1, sig=False):
# {{{
''' clim_detrend() - returns detrended time series with a daily trend.'''
from pygeode.timeaxis import Time
from . import stats
from numpy import arange
if itime == -1: itime = var.whichaxis(Time)
tlen = var.shape[itime]
vary = composite(var, itime, list(range(0, tlen, yrlen)), yrlen)
yrs = vary.axes[itime]
yrs.values = arange(len(yrs)).astype(yrs.dtype)
print('Computing regression')
from pygeode.progress import PBar
m, b, p = stats.regress(yrs, vary, pbar=PBar())
varz = flatten(vary - (m * yrs + b), itime + 1)
varz.axes = var.axes
# Since the axes have been modified after initialization, redo the init to get
# shortcuts to the axes names
Var.__init__(varz, varz.axes, varz.dtype)
if var.name != '':
varz.name = var.name + "'"
if sig:
return m, b, varz, p
else:
return m, b, varz
def __init__(self,
var,
newaxes,
name=None,
fillvalue=None,
scale=None,
offset=None,
atts={},
plotatts={}):
from pygeode.var import Var, copy_meta
atts = atts.copy()
plotatts = plotatts.copy()
assert len(newaxes) == len(var.axes)
for a1, a2 in zip(newaxes, var.axes):
assert len(a1) == len(a2)
self.var = var
dtype = var.dtype
if fillvalue is not None or scale is not None or offset is not None:
dtype = 'float32'
self.fillvalue = fillvalue
self.scale = scale
self.offset = offset
Var.__init__(self, newaxes, dtype=dtype)
copy_meta(var, self)
self.atts = atts
self.plotatts = plotatts
if name is not None: self.name = name
def __init__ (self, v1, opener, files, faxis, timedict):
# {{{
# v1 - var chunk from the first file
# opener - method for opening a file given a filename
# files - list of filenames
# faxis - starting time of each file, as a time axis
# timedict - dictionary of pre-constructed time axes
from pygeode.var import Var
from pygeode.timeaxis import Time
self.opener = opener
self.files = files
self.faxis = faxis
n = v1.getaxis(Time).name if v1.hasaxis(Time) else None
taxis = timedict[n]
T = type(taxis)
axes = list(v1.axes)
# Replace the existing time axis?
# (look for either a Time axis or a generic axis with the name 'time')
if v1.hasaxis(T): axes[v1.whichaxis(T)] = taxis
elif v1.hasaxis('time'): axes[v1.whichaxis('time')] = taxis
else: axes = [taxis] + axes
Var.__init__ (self, axes, dtype=v1.dtype, name=v1.name, atts=v1.atts, plotatts=v1.plotatts)
def clim_detrend(var, yrlen, itime = -1, sig=False):
# {{{
''' clim_detrend() - returns detrended time series with a daily trend.'''
from pygeode.timeaxis import Time
from . import stats
from numpy import arange
if itime == -1: itime = var.whichaxis(Time)
tlen = var.shape[itime]
vary = composite(var, itime, list(range(0, tlen, yrlen)), yrlen)
yrs = vary.axes[itime]
yrs.values=arange(len(yrs)).astype(yrs.dtype)
print('Computing regression')
from pygeode.progress import PBar
m, b, p = stats.regress(yrs, vary, pbar=PBar())
varz = flatten(vary - (m*yrs + b), itime + 1)
varz.axes = var.axes
# Since the axes have been modified after initialization, redo the init to get
# shortcuts to the axes names
Var.__init__(varz, varz.axes, varz.dtype)
if var.name != '':
varz.name = var.name+"'"
if sig:
return m, b, varz, p
else:
return m, b, varz
def __init__(self, u, v, perix=True, name=None, atts=None, plotatts=None):
'''
relativeVorticity(u, v, perix=True, name=None, atts=None, plotatts=None)
'''
from pygeode.axis import Lat, Lon
from pygeode.atmdyn.constants import Re
## check axes
# order: latitude and longitude have to be the last two
assert u.whichaxis(Lat)==u.naxes-2 and u.whichaxis(Lon)==u.naxes-1, 'Unsupported axes order.'
# homogeneity
assert u.axes==v.axes, 'Axes are incompatible!'
# handle meta data
if name is None: name = 'ze'
# parameter (set defaults, if not present)
defatts = {'name': 'zeta', 'units': r'$s^{-1}$', 'long_name': 'Relative Vorticity (vertical)',
'standard_name': 'Relative Vorticity', 'Re': Re, 'perix':perix}
if atts: defatts.update(atts)
# plot parameter
defplot = variablePlotatts['ze']
if plotatts: defplot.update(plotatts)
# make proper Var-instance
Var.__init__(self, axes=u.axes, dtype=u.dtype, name=name, values=None, atts=defatts, plotatts=defplot)
# save variables and parameters
self.perix = perix
self.u = u; self.v = v
self.lon = u.getaxis(Lon); self.lat = u.getaxis(Lat)
self.Re = self.atts['Re']
def __init__(self, T, p, name=None, atts=None, plotatts=None):
'''
Theta(T, p, name=None, atts=None, plotatts=None)
'''
from pygeode.atmdyn.constants import p0, kappa
# input checks
if not T.axes==p.axes:
# need to transpose p to fix axes order between T and p
iaxes = range(len(T.axes))
for i in range(len(T.axes)):
assert p.hasaxis(T.axes[i]), 'Axes of T and p are incompatible!'
iaxes[p.whichaxis(T.axes[i])] = i # order of axes in p
p = p.transpose(*iaxes)
# handle meta data
if name is None: name = 'th'
# parameter (set defaults, if not present
defatts = {'name': 'th', 'units': 'K', 'long_name': 'potential temperature',
'standard_name': r'$\theta$', 'p0': p0, 'kappa': kappa}
if atts: defatts.update(atts)
# plot parameter
defplot = {'plottitle': 'theta', 'plotunits': 'K'}
if plotatts: defplot.update(plotatts)
# make proper Var-instance
Var.__init__(self, axes=T.axes, dtype=T.dtype, name=name, values=None, atts=defatts, plotatts=plotatts)
# save T & p and parameters
self.T = T; self.p = p
self.p0 = self.atts['p0']; self.kappa = self.atts['kappa']
def __init__ (self, var, axis, n):
# {{{
'''__init__()'''
from pygeode.var import Var
df = 'right' # Hard-coded to match numpy behaviour.
# May be extended in the future?
self.daxis = daxis = var.whichaxis(axis)
assert var.shape[daxis] > n, "need at least %d value(s) along difference axis"%n
self.n = n
self.df = df
self.var = var
# Construct new variable
if var.name != '':
name = 'd' + var.name
else:
name = 'd(UnknownVar)'
axes = list(var.axes)
if df == 'left':
axes[daxis] = axes[daxis].slice[n:]
else:
axes[daxis] = axes[daxis].slice[:-n]
Var.__init__(self, axes, var.dtype, name=name, atts=var.atts, plotatts=var.plotatts)
def __init__(self, var, *iaxes, **kwargs):
from pygeode.var import Var, copy_meta
# Get the axes to be squeezed
if len(iaxes) == 1 and isinstance(iaxes[0], (list, tuple)):
iaxes = iaxes[0]
if len(iaxes) == 0:
iaxes = [i for i, a in enumerate(var.axes) if len(a) == 1]
# Only remove degenerate axes
iaxes = [var.whichaxis(a) for a in iaxes]
iaxes = [i for i in iaxes if len(var.axes[i]) == 1]
# Slice the var along some axes (passed by keyword argument)?
if len(kwargs) > 0:
for k, v in kwargs.items():
assert var.hasaxis(k), "unknown axis '%s'" % k
a = var.whichaxis(k)
if a not in iaxes: iaxes.append(a)
assert isinstance(
v, (int, float)
), "expected a numerical value for keyword '%s' - received %s instead" % (
k, type(v))
var = var(
**kwargs) # Do the slicing first, before doing this wrapper
self.var = var
Var.__init__(self,
[a for i, a in enumerate(var.axes) if i not in iaxes],
var.dtype)
copy_meta(var, self)
def __init__(self, var, saxis, maxharm):
# {{{
''' __init__()'''
# Construct new variable
self.saxis = saxis
self.var = var
self.maxharm = maxharm
Var.__init__(self, var.axes, var.dtype, name=var.name, atts=var.atts, plotatts=var.plotatts)
def __init__(self, var, pos, *newaxes):
from pygeode.var import Var, copy_meta
self.var = var
self.pos = pos
self.newaxes = newaxes
axes = var.axes[:pos] + tuple(newaxes) + var.axes[pos:]
Var.__init__(self, axes, dtype=var.dtype)
copy_meta(var, self)
def __init__(self, var, saxis, maxharm):
# {{{
''' __init__()'''
# Construct new variable
self.saxis = saxis
self.var = var
self.maxharm = maxharm
Var.__init__(self, var.axes, var.dtype, name=var.name, atts=var.atts, plotatts=var.plotatts)
def __init__(self, sd, axes):
self.sd = sd
self.name = sd.name
# attributes
atts = get_attributes (sd.sds_id, sd.natts)
# if len(atts) > 0: self.atts = atts
# else: atts = None
Var.__init__(self, axes, numpy_type[sd.type], atts=atts)
def __init__(self, sd, axes):
self.sd = sd
self.name = sd.name
# attributes
atts = get_attributes(sd.sds_id, sd.natts)
# if len(atts) > 0: self.atts = atts
# else: atts = None
Var.__init__(self, axes, numpy_type[sd.type], atts=atts)
def __init__(self, var, pos, *newaxes): from pygeode.var import Var, copy_meta self.var = var self.pos = pos self.newaxes = newaxes axes = var.axes[:pos] + tuple(newaxes) + var.axes[pos:] Var.__init__(self, axes, dtype=var.dtype) copy_meta (var, self)
def __init__(self, *coefs):
from pygeode.axis import Coef
from pygeode.var import Var
assert len(set(c.shape for c in coefs)) == 1
#TODO: more checks?
axes = list(coefs[0].axes)
axes.append(Coef(len(coefs)))
self.coefs = coefs
Var.__init__(self, axes, coefs[0].dtype)
def __init__ (self, name, arr): from pygeode.var import Var from pygeode.axis import NamedAxis self._arr = arr # Extract axes and metadata. # Convert unicode strings to str for compatibility with PyGeode. axes = [NamedAxis(n,str(d)) for n,d in zip(arr.shape,arr.dims)] atts = _fix_atts(arr.attrs) Var.__init__(self, axes, name=str(name), dtype=arr.dtype, atts=atts)
def __init__ (self, invar, inaxis, outaxis, inx, outx, interp_type, \
d_below, d_above, omit_nonmonotonic):
# {{{
from pygeode.var import Var
from pygeode.axis import Axis
# Check the types of the input parameters
assert isinstance(invar,Var)
inaxis = invar.getaxis(inaxis)
assert isinstance(outaxis,Axis)
if inx is None: inx = inaxis
if outx is None: outx = outaxis
assert isinstance(inx,Var)
assert isinstance(outx,Var)
# Validate the input axis
assert invar.hasaxis(inaxis)
# We need the interpolation axis to be the fastest-varying axis
# (For the C code to work properly)
iaxis = invar.whichaxis(inaxis)
order = list(range(0,iaxis)) + list(range(iaxis+1,invar.naxes)) + [iaxis]
invar = invar.transpose (*order)
del iaxis, order
# Generate the output axes
outaxes = list(invar.axes)
outaxes[-1] = outaxis
# Validate the coordinate fields
# assert all(invar.hasaxis(a) for a in inx.axes)
for a in inx.axes:
if not invar.hasaxis(a):
if invar.hasaxis(a.__class__):
raise ValueError("mismatching '%s' axis between invar and inx"%a.name)
else:
raise TypeError("invar does not have '%s' axis specified by inx"%a.name)
# assert all(invar.hasaxis(a) or a is outaxis for a in outx.axes)
for a in outx.axes:
if a is outaxis: continue
assert invar.hasaxis(a), "invar doesn't have axis '%s'"%repr(a)
assert inx.hasaxis(inaxis)
assert outx.hasaxis(outaxis)
self.invar = invar
self.inx = inx
self.outx = outx
self.interp_type = interp_type
self.d_below = d_below
self.d_above = d_above
self.omit_nonmonotonic = omit_nonmonotonic
Var.__init__ (self, outaxes, name=invar.name, dtype='d', atts=invar.atts, plotatts=invar.plotatts)
def __init__ (self, invar, iaxis, reverse=False): from pygeode.var import Var self.var = invar iaxis = invar.whichaxis(iaxis) axes = list(invar.axes) oldaxis = axes[iaxis] newaxis = oldaxis.sorted(reverse=reverse) axes[iaxis] = newaxis Var.__init__(self, axes, dtype=invar.dtype, name=invar.name, atts=invar.atts, plotatts=invar.plotatts)
def __init__ (self, *coefs): from pygeode.axis import Coef from pygeode.var import Var assert len(set(c.shape for c in coefs)) == 1 #TODO: more checks? axes = list(coefs[0].axes) axes.append(Coef(len(coefs))) self.coefs = coefs Var.__init__(self, axes, coefs[0].dtype)
def __init__(self, name, arr):
from pygeode.var import Var
from pygeode.axis import NamedAxis
self._arr = arr
# Extract axes and metadata.
# Convert unicode strings to str for compatibility with PyGeode.
axes = [NamedAxis(n, str(d)) for n, d in zip(arr.shape, arr.dims)]
atts = _fix_atts(arr.attrs)
Var.__init__(self, axes, name=str(name), dtype=arr.dtype, atts=atts)
def __init__ (self, var, iaxis, ievents, evlen, evoff=0, saxes=None, sindices=None):
#{{{
# Replace the time axis with a reference time (of an event?), and the offset
# from that event.
import numpy as np
ievents = np.array(ievents)
n = ievents.shape[0]
caxis = var.axes[iaxis]
# Event offsets can either be specified per event or as a single offset
if hasattr(evoff, '__len__'):
evoff = np.array(evoff)
mevoff = evoff.max()
assert evoff.shape == ievents.shape, "The number of event offsets provided does not match the number of events."
else:
mevoff = evoff
evoff = np.ones(n, 'i') * mevoff
# Event lengths can either be specified per event or as a single length
if hasattr(evlen, '__len__'):
evlen = np.array(evlen)
mevlen = (evlen - evoff + mevoff).max()
assert evlen.shape == ievents.shape, "The number of event lengths provided does not match the number of events."
else:
mevlen = evlen
evlen = np.ones(n, 'i') * mevlen
# Construct event and offset axes
from pygeode.timeaxis import Time, Yearless
ev = Event(np.arange(n)+1, indices=ievents)
if isinstance(caxis, Time):
units = caxis.units
delta = caxis.delta()
off = Yearless(values=delta*np.arange(-mevoff, mevlen-mevoff), units=units, startdate={'day':0})
else:
off = Offset(np.arange(-mevoff, mevlen-mevoff))
axes = var.axes[:iaxis] + (ev, off) + var.axes[iaxis+1:]
# Build var object
self.var = var
self.iaxis = iaxis
self.evlens = evlen
self.mevlen = mevlen
self.evoffs = evoff
self.mevoff = mevoff
for i, (iev, el, eo) in enumerate(zip(ievents, evlen, evoff)):
if iev - eo < 0:
self.evoffs[i] += iev - eo
self.evlens[i] -= iev - eo
if iev - eo + el >= len(caxis):
self.evlens[i] = len(caxis) - (iev - eo)
#assert iev - eo >= 0 and iev - eo + el < len(caxis), \
#'Event %d (i: %d) is not fully defined' % (np.where(ievents==iev)[0][0], iev)
Var.__init__(self, axes, dtype=var.dtype, name=var.name, atts=var.atts, plotatts=var.plotatts)
def __init__(self, var, iaxis, dx=None, v0=None, order=1, type='trapz'):
# {{{
''' __init__()'''
from pygeode.var import Var
self.iaxis = var.whichaxis(iaxis)
if var.shape[iaxis] < 2:
raise ValueError(
'At least two values are needed to integrate %s along axis %s'
% (var.name, var.axis[iaxis]))
self.var = var
if v0 is not None:
# Confirm the initial values are consistently shaped
if isinstance(v0, Var):
if v0.hasaxis(var.axes[self.iaxis]):
raise ValueError(
'v0\n %s \n must not share integration axis with \n\n %s'
% (v0, var))
if not all([a in var.axes for a in v0.axes]):
raise ValueError(
'v0\n %s \n\n axes must all match those of \n\n %s' %
(v0, var))
self.v0 = v0
else:
self.v0 = 0.
if dx is not None: # Optionally one can specify a coordinate system
if dx.naxes == 1:
assert dx.shape[0] == var.shape[iaxis]
self.dx = dx.replace_axes(newaxes=(var.axes[iaxis], ))
else: # Must be mappable to integrand, with a matching integration axis
for a in var.axes:
assert dx.hasaxis(a)
self.dx = dx
else:
self.dx = var.axes[iaxis]
self.order = order
if type not in ['trapz', 'rectr', 'rectl']:
raise ValueError(
"type (%s) must be one of 'trapz', 'rectr', 'rectl'." % type)
self.type = type
# Construct new variable
if var.name != '':
name = 'i' + var.name
else:
name = 'i(UnknownVar)'
Var.__init__(self, var.axes, dtype='d', name=name)
def __init__ (self, var, c): from pygeode.axis import Coef from pygeode.var import Var self.var = var self.c = c self.ci = ci = var.whichaxis(Coef) self.caxis = var.axes[ci] axes = list(var.axes) axes = axes[:ci] + axes[ci+1:] Var.__init__(self, axes, dtype=var.dtype) if var.name != '': self.name = var.name + "_coef%i"%c
def __init__(self, var, fill):
from numpy import float32, float64
from pygeode.var import Var, copy_meta
self.var = var
self._fill = fill
# We need floating-point values to have a nan
if var.dtype not in (float32, float64):
dtype = float32
else: dtype = var.dtype
Var.__init__(self, var.axes, dtype)
copy_meta (var, self)
def __init__(self, var, order):
from pygeode.var import Var, copy_meta
self.var = var
outaxes = list(var.axes)
for iaxis, o in order.items():
reverse = {1:False, 0:None, -1:True}[o]
outaxes[iaxis] = var.axes[iaxis].sorted(reverse=reverse)
Var.__init__(self, outaxes, dtype=var.dtype)
copy_meta (var, self)
def __init__(self, var, fill):
from numpy import float32, float64
from pygeode.var import Var, copy_meta
self.var = var
self._fill = fill
# We need floating-point values to have a nan
if var.dtype not in (float32, float64):
dtype = float32
else: dtype = var.dtype
Var.__init__(self, var.axes, dtype)
copy_meta (var, self)
def __init__(self, var, c):
from pygeode.axis import Coef
from pygeode.var import Var
self.var = var
self.c = c
self.ci = ci = var.whichaxis(Coef)
self.caxis = var.axes[ci]
axes = list(var.axes)
axes = axes[:ci] + axes[ci + 1:]
Var.__init__(self, axes, dtype=var.dtype)
if var.name != '': self.name = var.name + "_coef%i" % c
def __init__(self, var, order):
from pygeode.var import Var, copy_meta
self.var = var
outaxes = list(var.axes)
for iaxis, o in order.items():
reverse = {1: False, 0: None, -1: True}[o]
outaxes[iaxis] = var.axes[iaxis].sorted(reverse=reverse)
Var.__init__(self, outaxes, dtype=var.dtype)
copy_meta(var, self)
def test_underscores_in_axes(): from pygeode.axis import ZAxis from pygeode.var import Var # Axis names without underscores works class HeightwrtGround(ZAxis): pass ax = HeightwrtGround(values=[1.5]) x = Var(axes=[ax], name='blah', values=[123.45]) y = x(heightwrtground=0.0) # Axis names with underscores fails class Height_wrt_Ground(ZAxis): pass ax = Height_wrt_Ground(values=[1.5]) x = Var(axes=[ax], name='blah', values=[123.45]) y = x(height_wrt_ground=0.0)
def __init__(self, u, v, th, rho, w=None, z=None, perix=True, name=None, atts=None, plotatts=None):
'''
PotentialVorticity(u, v, th, rho, w=None, z=None, name=None, atts=None, plotatts=None)
'''
from pygeode.axis import Lat, Lon, Height, ZAxis, TAxis
from pygeode.atmdyn.constants import Re, Omega
## check axes
# order
assert u.whichaxis(TAxis)==0 and u.whichaxis(Lat)==2 and u.whichaxis(Lon)==3, 'Unsupported axes order.'
# homogeneity
assert u.axes==v.axes, 'Axes are incompatible!'
assert th.axes==rho.axes, 'Axes are incompatible!'
assert u.axes==th.axes, 'Axes are incompatible!'
if w: assert u.axes==w.axes, 'Axes are incompatible!' # should have same axes as u & v
if z: # z is a field, e.g. geopotential on hybrid axis
zField = True
assert u.whichaxis(ZAxis)==1, 'Position of vertical axis is not supported.'
if not z.axes==th.axes: # sort z's axes if necessary
order = []
for ax in th.axes:
order.append(ax.name)
z = z.transpose(*order)
# assert z.axes==th.axes, 'Axes are incompatible!' # should have same axes as th & rho
else: # just use height axis as z-field
zField = False
u.whichaxis(Height)==1, 'Position of vertical axis is not supported.'
z = u.getaxis(Height) # expand to field later
# handle meta data
if th.name=='th':
if name is None: name = 'PV'
# parameter (set defaults, if not present
defatts = {'name': 'PV', 'units': r'$K m^2 (s kg)^{-1}$', 'long_name': 'Ertel Potential Vorticity',
'standard_name': 'isentropic PV', 'Re': Re, 'Omega': Omega, 'perix':perix}
elif th.name=='s':
if name is None: name = 'PVs'
# parameter (set defaults, if not present
defatts = {'name': 'PVs', 'units': r'$J m^2 (K s)^{-1} kg^{-2}$', 'long_name': 'Entropy Potential Vorticity',
'standard_name': 'Entropy PV', 'Re': Re, 'Omega': Omega, 'perix':perix}
if atts: defatts.update(atts)
# plot parameter
defplot = variablePlotatts[name]
if plotatts: defplot.update(plotatts)
# make proper Var-instance
Var.__init__(self, axes=th.axes, dtype=th.dtype, name=name, values=None, atts=defatts, plotatts=defplot)
# save variables and parameters
self.perix = perix
self.zField = zField
self.u = u; self.v = v; self.w = w
self.th = th; self.rho = rho; self.z = z
self.lon = u.getaxis(Lon); self.lat = u.getaxis(Lat)
self.Re = self.atts['Re']; self.Omega = self.atts['Omega']
def __init__(self, varlist):
from pygeode.var import Var, combine_meta
self.varlist = varlist
# assume the vars have already been checked for consistency
var0 = varlist[0]
axes = list(var0.axes)
# axes = [Ensemble(len(varlist))] + axes
axes = [make_ensemble(len(varlist))] + axes
Var.__init__(self, axes, dtype=var0.dtype)
# copy_meta (var0, self)
# self.atts = common_dict(var.atts for var in varlist)
# self.plotatts = common_dict(var.plotatts for var in varlist)
combine_meta(varlist, self)
self.name = varlist[0].name
def __init__(self, varlist):
from pygeode.var import Var, combine_meta
self.varlist = varlist
# assume the vars have already been checked for consistency
var0 = varlist[0]
axes = list(var0.axes)
# axes = [Ensemble(len(varlist))] + axes
axes = [make_ensemble(len(varlist))] + axes
Var.__init__(self, axes, dtype=var0.dtype)
# copy_meta (var0, self)
# self.atts = common_dict(var.atts for var in varlist)
# self.plotatts = common_dict(var.plotatts for var in varlist)
combine_meta (varlist, self)
self.name = varlist[0].name
def __init__(self, invar, iaxis, reverse=False):
from pygeode.var import Var
self.var = invar
iaxis = invar.whichaxis(iaxis)
axes = list(invar.axes)
oldaxis = axes[iaxis]
newaxis = oldaxis.sorted(reverse=reverse)
axes[iaxis] = newaxis
Var.__init__(self,
axes,
dtype=invar.dtype,
name=invar.name,
atts=invar.atts,
plotatts=invar.plotatts)
def __init__(self, T, ps, phis):
from numpy import diff
from pygeode.axis import Hybrid
from pygeode.varoperations import sorted
# precondition input and store internally
assert T.hasaxis(
Hybrid
), 'this function only computes geopotential from hybrid coordinates'
# T: make vertical axis varying the slowest (default is 2nd slowest, after time)
ietaT = T.whichaxis(Hybrid)
inOrder = [ietaT] + range(0, ietaT) + range(ietaT + 1, T.naxes)
self.T = T.transpose(*inOrder)
# surface fields
self.ps = ps # surface pressure
self.phis = phis
# get vertical coefficients for hybrid coordinate
self.A = T.axes[ietaT].auxarrays['A']
self.B = T.axes[ietaT].auxarrays['B']
# construct output axes: make eta varying the fastest, to prevent break-up in loop-over routine
outAxes = T.axes[0:ietaT] + T.axes[ietaT + 1:] + (T.axes[ietaT], )
# ensure eta axis is ordered properly
if not all(diff(self.T.eta.values) > 0):
self.T = sorted(self.T, eta=1)
from warnings import warn
warn(
'The vertical axis (eta) was not in the expected order - the sorted-fct. has been applied.'
)
# attributes
atts = {}
atts['name'] = 'z'
atts['long_name'] = 'geopotential height'
atts['standard_name'] = 'geopotential'
atts['units'] = 'm'
atts['g0'] = g0
atts['Rd'] = Rd
# plot attributes (defaults)
plotatts = variablePlotatts['z']
# make proper Var-instance
Var.__init__(self,
axes=outAxes,
dtype=self.T.dtype,
name='z',
values=None,
atts=atts,
plotatts=plotatts)
self.ieta = self.whichaxis(Hybrid)
# make sure axes are assigned properly and eta is the innermost axis
assert self.naxes == T.naxes
assert self.ieta == T.naxes - 1
def test_issue061(): from pygeode.tutorial import t1 from pygeode.var import Var # Construct 3 variables, with some attributes that are identical and some # that are different. x = Var(t1.axes, atts=dict(standard_name='x',units='deg C'), dtype=float) y = Var(t1.axes, atts=dict(standard_name='y',units='deg C'), dtype=float) z = Var(t1.axes, atts=dict(standard_name='z',units='deg C'), dtype=float) # Apply a ufunc operator on these variables, to get a new variable w = x*y*z # Make sure the standard_name is removed. assert w.atts == dict(units='deg C')
def finalize (var, time, space, eof, eig, pc, variance, weight, out):
from pygeode.var import Var
import numpy as np
# Keep the name
name = var.name
num = eof.shape[0]
# Conform to the proper shape
eof = eof.reshape((num,)+space.shape)
pc = pc.reshape((num,)+time.shape)
# Use a consistent sign convention
# (the first element of the eof is non-negative)
sign = [-1 if eof.reshape(num,-1)[i,0] < 0 else 1 for i in range(num)]
for i,s in enumerate(sign):
eof[i,...] *= s
pc[i,...] *= s
# Copy the data into fresh array
# (that way, if these are view on much larger arrays, the rest of
# the unused data can be garbage collected)
eof = np.array(eof)
pc = np.array(pc)
eig = np.array(eig)
# EOF axis
orderaxis = order(num,name="order")
eof = Var((orderaxis,)+space.axes, values=eof)
eig = Var([orderaxis], values=eig)
pc = Var((orderaxis,)+time.axes, values = pc)
# Undo normalization by area weight?
eof = remove_weights(eof, weight=weight).load()
eof.name = name + "_EOF"
pc.name = name + "_timeseries"
eig.name = name + "_eigenvalues"
# Other things
# Fraction of total variance
frac = ((eig**2) / variance).load()
frac.name = name + "_fraction_of_variance"
# Eigenvalues of the covariance matrix
eig2 = (eig**2).load()
eig2.name = name + "_eigenvalues"
# Gather up all possible outputs
outdict = dict(eof=eof, eig=eig, eig2=eig2, pc=pc, var=variance, frac=frac)
out = whichout(out)
out = [outdict[o] for o in out]
return out
def __init__(self, var, index):
# {{{
from pygeode.var import Var
import numpy as np
axes = var.axes
index = var.whichaxis(index)
self.N = len(axes[index])
self.var = var
self.indices = [index]
self.in_axes = axes
# Remove the reduction axis from the output variable
axes = [a for i,a in enumerate(axes) if i != index]
Var.__init__(self, axes, dtype='i', name=var.name, atts=var.atts, plotatts=var.plotatts)
def __init__(self, var, index):
# {{{
from pygeode.var import Var
import numpy as np
axes = var.axes
index = var.whichaxis(index)
self.N = len(axes[index])
self.var = var
self.indices = [index]
self.in_axes = axes
# Remove the reduction axis from the output variable
axes = [a for i,a in enumerate(axes) if i != index]
Var.__init__(self, axes, dtype='i', name=var.name, atts=var.atts, plotatts=var.plotatts)
def __init__(self, var, newaxes, name=None, fillvalue=None, scale=None, offset=None, atts={}, plotatts={}):
from pygeode.var import Var, copy_meta
atts = atts.copy()
plotatts = plotatts.copy()
assert len(newaxes) == len(var.axes)
for a1, a2 in zip(newaxes, var.axes): assert len(a1) == len(a2)
self.var = var
dtype = var.dtype
if fillvalue is not None or scale is not None or offset is not None: dtype = 'float32'
self.fillvalue = fillvalue
self.scale = scale
self.offset = offset
Var.__init__(self, newaxes, dtype=dtype)
copy_meta(var, self)
self.atts = atts
self.plotatts = plotatts
if name is not None: self.name = name
def __init__ (self, var):
from pygeode.var import Var
from pygeode.timeaxis import CalendarTime
self.ti = ti = var.whichaxis(CalendarTime)
intime = var.axes[ti]
outtime = self.get_outtime(intime)
# Fudge the input axis? (endow it with extra information that it normally wouldn't have?)
new_intime = self.get_intime(intime)
if new_intime is not intime: var = var.replace_axes({CalendarTime:new_intime})
if var.name != '': self.name = var.name + self.name_suffix1 + self.name_suffix2
self.var = var
axes = list(var.axes)
axes[ti] = outtime
Var.__init__ (self, axes+list(self.extra_dims), dtype = var.dtype)
def clim_anoms(var, yrlen, itime=-1):
# {{{
''' clim_anoms() - quick and dirty implementation;
returns climatology and anomalies of given variable.'''
from pygeode.timeaxis import Time
if itime == -1: itime = var.whichaxis(Time)
tlen = (var.shape[itime] // yrlen) * yrlen
vary = composite(var, itime, list(range(0, tlen, yrlen)), yrlen)
varc = vary.mean(itime).load()
varz = flatten(vary - varc, itime + 1)
varz.axes = var.axes
# Since the axes have been modified after initialization, redo the init to get
# shortcuts to the axes names
Var.__init__(varz, varz.axes, varz.dtype)
if var.name != '':
varz.name = var.name + '_anom'
return varc, varz
def clim_anoms(var, yrlen, itime = -1):
# {{{
''' clim_anoms() - quick and dirty implementation;
returns climatology and anomalies of given variable.'''
from pygeode.timeaxis import Time
if itime == -1: itime = var.whichaxis(Time)
tlen = (var.shape[itime] // yrlen) * yrlen
vary = composite(var, itime, list(range(0, tlen, yrlen)), yrlen)
varc = vary.mean(itime).load()
varz = flatten(vary - varc, itime + 1)
varz.axes = var.axes
# Since the axes have been modified after initialization, redo the init to get
# shortcuts to the axes names
Var.__init__(varz, varz.axes, varz.dtype)
if var.name != '':
varz.name = var.name+'_anom'
return varc, varz
def __init__ (self, var, saxis, kernel):
# {{{
''' __init__()'''
import numpy as np
assert len(kernel) <= var.shape[saxis], 'Kernel must not be longer than dimension being smoothed.'
# Construct new variable
self.saxis = saxis
self.var = var
self.kernel = kernel
self.klen = len(kernel)
# Normalize and reshape kernel
self.kernel /= np.sum(self.kernel)
self.kernel.shape = [self.klen if i == saxis else 1 for i in range(var.naxes)]
Var.__init__(self, var.axes, var.dtype, name=var.name, atts=var.atts, plotatts=var.plotatts)
def decode_string_var (var):
from pygeode.var import Var
name = var.name
if name.endswith('_name'):
name = name[:-5]
data = [''.encode('ascii').join(s) for s in var.get()]
data = [str(s.decode()) for s in data]
return Var(axes=var.axes[:-1], values=data, name=name)
