def interpolate(self, x_target, x_src, rising=None):
x_target = np.array(x_target)
x_src = np.array(x_src)
fx_src = np.empty(x_src.shape)
index_interp = IndexInterpolator()
extrap_direct = DirectionExtrapolator()
r1 = stratify.interpolate(x_target,
x_src,
fx_src,
rising=rising,
interpolation=index_interp,
extrapolation=extrap_direct)
if rising is not None:
r2 = stratify.interpolate(-1 * x_target,
-1 * x_src,
fx_src,
rising=not rising,
interpolation=index_interp,
extrapolation=extrap_direct)
assert_array_equal(r1, r2)
return r1
def interp_and_extrap(shape,
interp=stratify.INTERPOLATE_LINEAR,
extrap=stratify.EXTRAPOLATE_NEAREST):
z, fz = src_data(shape)
stratify.interpolate(np.linspace(-20, 120, 50),
z,
fz,
interpolation=interp,
extrapolation=extrap)
def test_npts(self):
interpolation = IndexInterpolator()
extrapolation = stratify.EXTRAPOLATE_LINEAR
msg = (r'Linear extrapolation requires at least 2 '
r'source points. Got 1.')
with self.assertRaisesRegexp(ValueError, msg):
stratify.interpolate([1, 3.], [2], [20],
interpolation=interpolation,
extrapolation=extrapolation,
rising=True)
def _vertical_interpolate(cube, src_levels, levels, interpolation,
extrapolation):
"""Perform vertical interpolation."""
# Determine the source levels and axis for vertical interpolation.
z_axis, = cube.coord_dims(cube.coord(axis='z', dim_coords=True))
# Broadcast the 1d source cube vertical coordinate to fully
# describe the spatial extent that will be interpolated.
src_levels_broadcast = broadcast_to_shape(src_levels.points, cube.shape,
cube.coord_dims(src_levels))
# force mask onto data as nan's
if np.ma.is_masked(cube.data):
cube.data[cube.data.mask] = np.nan
# Now perform the actual vertical interpolation.
new_data = stratify.interpolate(levels,
src_levels_broadcast,
cube.data,
axis=z_axis,
interpolation=interpolation,
extrapolation=extrapolation)
# Calculate the mask based on the any NaN values in the interpolated data.
mask = np.isnan(new_data)
if np.any(mask):
# Ensure that the data is masked appropriately.
new_data = np.ma.array(new_data, mask=mask, fill_value=_MDI)
# Construct the resulting cube with the interpolated data.
return _create_cube(cube, new_data, src_levels, levels.astype(float))
def stratify(self, levels, vertical, axis=1):
"""Short summary.
Parameters
----------
levels : type
Description of parameter `levels`.
vertical : type
Description of parameter `vertical`.
axis : type
Description of parameter `axis`.
Returns
-------
type
Description of returned object.
"""
result = stratify.interpolate(levels,
vertical.chunk(),
self.obj.chunk(),
axis=axis)
dims = self.obj.dims
out = xr.DataArray(result, dims=dims)
for i in dims:
if i != 'z':
out[i] = self.obj[i]
out.attrs = self.obj.attrs.copy()
if len(self.obj.coords) > 0:
for i in self.obj.coords:
out.coords[i] = self.obj.coords[i]
return out
def test(self):
interpolation = IndexInterpolator()
extrapolation = Test_custom_extrap_kernel.my_kernel()
r = stratify.interpolate([1, 3.], [1, 2], [10, 20],
interpolation=interpolation,
extrapolation=extrapolation,
rising=True)
assert_array_equal(r, [0, -10])
def test_single_point(self):
# Test that a single input point that falls exactly on the target
# level triggers a shortcut that avoids the expectation of >=2 source
# points.
interpolation = stratify.INTERPOLATE_LINEAR
extrapolation = DirectionExtrapolator()
r = stratify.interpolate([2], [2], [20],
interpolation=interpolation,
extrapolation=extrapolation,
rising=True)
self.assertEqual(r, 20)
def interpolate(self, x_target):
interpolation = stratify.INTERPOLATE_NEAREST
extrapolation = DirectionExtrapolator()
x_src = np.arange(5)
fx_src = 10 * x_src
# Use -2 to test negative number support.
return stratify.interpolate(np.array(x_target) - 2,
x_src - 2,
fx_src,
interpolation=interpolation,
extrapolation=extrapolation)
def interpolate(self, x_target):
interpolation = IndexInterpolator()
extrapolation = stratify.EXTRAPOLATE_NEAREST
x_src = np.arange(5)
fx_src = 10 * x_src
# Use -2 to test negative number support.
return stratify.interpolate(np.array(x_target) - 2,
x_src - 2,
fx_src,
interpolation=interpolation,
extrapolation=extrapolation)
def calculate(cubes):
"""Compute mole fraction of CO2 at surface."""
co2_cube = cubes.extract_cube(
iris.Constraint(name='mole_fraction_of_carbon_dioxide_in_air'))
ps_cube = cubes.extract_cube(
iris.Constraint(name='surface_air_pressure'))
# Fill masked data if necessary (interpolation fails with masked data)
(z_axis,) = co2_cube.coord_dims(co2_cube.coord(axis='Z',
dim_coords=True))
mask = da.ma.getmaskarray(co2_cube.core_data())
if mask.any():
first_unmasked_data = _get_first_unmasked_data(
co2_cube.core_data(), axis=z_axis)
dim_map = [dim for dim in range(co2_cube.ndim) if dim != z_axis]
first_unmasked_data = iris.util.broadcast_to_shape(
first_unmasked_data, co2_cube.shape, dim_map)
co2_cube.data = da.where(mask, first_unmasked_data,
co2_cube.core_data())
# Interpolation (not supported for dask arrays)
air_pressure_coord = co2_cube.coord('air_pressure')
original_levels = iris.util.broadcast_to_shape(
air_pressure_coord.points, co2_cube.shape,
co2_cube.coord_dims(air_pressure_coord))
target_levels = np.expand_dims(ps_cube.data, axis=z_axis)
co2s_data = stratify.interpolate(
target_levels,
original_levels,
co2_cube.data,
axis=z_axis,
interpolation='linear',
extrapolation='linear',
)
co2s_data = np.squeeze(co2s_data, axis=z_axis)
# Construct co2s cube
indices = [slice(None)] * co2_cube.ndim
indices[z_axis] = 0
co2s_cube = co2_cube[tuple(indices)]
co2s_cube.data = co2s_data
if co2s_cube.coords('air_pressure'):
co2s_cube.remove_coord('air_pressure')
ps_coord = iris.coords.AuxCoord(ps_cube.data,
var_name='plev',
standard_name='air_pressure',
long_name='pressure',
units=ps_cube.units)
co2s_cube.add_aux_coord(ps_coord, np.arange(co2s_cube.ndim))
co2s_cube.convert_units('1e-6')
return co2s_cube
def interpolate(self, x_target):
interpolation = IndexInterpolator()
extrapolation = stratify.EXTRAPOLATE_LINEAR
x_src = np.arange(5)
# To spice things up a bit, let's make x_src non-equal distance.
x_src[4] = 9
fx_src = 10 * x_src
# Use -2 to test negative number support.
return stratify.interpolate(np.array(x_target) - 2,
x_src - 2,
fx_src,
interpolation=interpolation,
extrapolation=extrapolation)
def resample_stratify(da, levels, vertical, axis=1):
import stratify
import xarray as xr
result = stratify.interpolate(
levels, vertical.chunk(), da.chunk(), axis=axis)
dims = da.dims
out = xr.DataArray(result, dims=dims)
for i in dims:
if i != 'z':
out[i] = da[i]
out.attrs = da.attrs.copy()
if len(da.coords) > 0:
for i in da.coords:
if i != 'z':
out.coords[i] = da.coords[i]
return out
def main(gmtb_forcing, wrf_output):
wrf = xarray.open_dataset(glob.glob(wrf_output)[0])
wrf = wrf.isel(south_north=0,
west_east=0,
south_north_stag=0,
west_east_stag=0)
wrf.coords['PRES'] = wrf['P'] + wrf['PB']
wrf['THETA'] = wrf['T'] + wrf_constants.T_BASE
w_on_grid = stratify.interpolate(wrf.ZNU, wrf.ZNW, wrf.W, axis=1)
wrf['W_G'] = (wrf.PRES.dims, w_on_grid)
gmtb = xarray.open_dataset(gmtb_forcing, 'forcing')
gmtb.update(xarray.open_dataset(gmtb_forcing))
gmtb.coords['time'] = wrf['XTIME'].values[0] + pandas.to_timedelta(
gmtb['time'].values, 'seconds')
compare_models(gmtb, wrf)
def compare_variable(gmtb, wrf, ax):
vmax = max(gmtb.max(), wrf.max())
vmin = min(gmtb.min(), wrf.min())
if vmin < 0:
vmax = max(-vmin, vmax)
vmin = None
#fig, ax = plt.subplots(1,3, sharey=True)
gmtb.plot.pcolormesh('time', 'levels', ax=ax[0], vmin=vmin, vmax=vmax)
_, t = numpy.meshgrid(range(wrf.shape[1]), wrf.XTIME)
#ax[1].pcolormesh(t, wrf.PRES, wrf)
wrf.coords['time2'] = (wrf.dims, t)
wrf.plot.pcolormesh('time2', 'PRES', ax=ax[1], vmin=vmin, vmax=vmax)
wrf_on_gmtb = stratify.interpolate(gmtb.levels, wrf.PRES, wrf,
axis=1).T[:, ::3]
#(gmtb - wrf_on_gmtb).plot.pcolormesh('time','levels', ax=ax[2])
ax[0].set_ylim([100000, 0])
def find_falling_level(self, wb_int_data: ndarray, orog_data: ndarray,
height_points: ndarray) -> ndarray:
"""
Find the phase change level by finding the level of the wet-bulb
integral data at the required threshold. Wet-bulb integral data
is only available above ground level and there may be an insufficient
number of levels in the input data, in which case the required
threshold may lie outside the Wet-bulb integral data and the value
at that point will be set to np.nan.
Args:
wb_int_data:
Wet bulb integral data on heights
orog_data:
Orographic data
height_points:
heights agl
Returns:
Phase change level data asl.
"""
from stratify import interpolate
# Create cube of heights above sea level for each height in
# the wet bulb integral cube.
asl = wb_int_data.copy()
for i, height in enumerate(height_points):
asl[i, ::] = orog_data + height
# Calculate phase change level above sea level by
# finding the level corresponding to the falling_level_threshold.
# Interpolate returns an array with height indices
# for falling_level_threshold so we take the 0 index
phase_change_level_data = interpolate(np.array(
[self.falling_level_threshold]),
wb_int_data,
asl,
axis=0)[0]
return phase_change_level_data
def vertical_interpolate(infile, outfile, orogfile, vertlevs):
"""
Perform a vertical interpolation of ancil file 'infile', using the level
definition namelist 'vertlevs'
Args:
infile (string): Path to input UM ancil file
outfile (string): Path to output UM ancil file
orogfile (string): Path to UM orography for true level calculations
vertlevs (string): Path to UM vertical namelist file for target levels
"""
ancil = mule.AncilFile.from_file(infile)
def categorise_fields(m):
df = pandas.DataFrame({'field': m.fields})
df['year'] = df['field'].apply(lambda f: f.lbyr)
df['month'] = df['field'].apply(lambda f: f.lbmon)
df['day'] = df['field'].apply(lambda f: f.lbdat)
df['hour'] = df['field'].apply(lambda f: f.lbhr)
df['minute'] = df['field'].apply(lambda f: f.lbmin)
df['second'] = df['field'].apply(lambda f: f.lbsec)
df['stash'] = df['field'].apply(lambda f: f.lbuser4)
df['vertical_type'] = df['field'].apply(lambda f: f.lbvc)
df['level'] = df['field'].apply(lambda f: f.lblev)
df['pseudo'] = df['field'].apply(lambda f: f.lbuser5)
#df['bulev'] = df['field'].apply(lambda f: f.bulev)
df['blev'] = df['field'].apply(lambda f: f.blev)
df['brlev'] = df['field'].apply(lambda f: f.brlev)
#df['bhulev'] = df['field'].apply(lambda f: f.bhulev)
df['bhlev'] = df['field'].apply(lambda f: f.bhlev)
df['bhrlev'] = df['field'].apply(lambda f: f.bhrlev)
return df
# Categorise the 2d slices in the input file
df = categorise_fields(ancil)
# Get the orography
orog_file = mule.AncilFile.from_file(orogfile)
orog = orog_file.fields[0].get_data()
levtype = 'theta'
target_levels = f90nml.read(vertlevs)['VERTLEVS']
if levtype == 'rho':
# Rho levels
eta = numpy.array(target_levels['eta_rho'])
const_lev = target_levels['first_constant_r_rho_level'] - 1
if levtype == 'theta':
# Theta levels
eta = numpy.array(target_levels['eta_theta'])
const_lev = target_levels['first_constant_r_rho_level'] - 1
# True height of the target levels
target_Zsea = target_levels['z_top_of_model'] * eta
target_C = (1 - eta / eta[const_lev])**2
target_C[const_lev:] = 0
target_Z = target_Zsea[:, numpy.newaxis,
numpy.newaxis] + numpy.multiply.outer(
target_C, orog)
ancil_out = ancil.copy()
# Group the 2d slices with the same field and time value together
for name, g in df.groupby(
['year', 'month', 'day', 'hour', 'minute', 'second', 'stash']):
print("%04d%02d%02dT%02d:%02d:%02d STASH %d" % name)
# Stack the slices into a 3d array
cube = numpy.stack(g['field'].apply(lambda f: f.get_data()))
# True height of each position
Zsea = g['blev']
C = g['bhlev']
Z = Zsea[:, numpy.newaxis, numpy.newaxis] + numpy.multiply.outer(
C, orog)
# Interpolate from the source true height to the target true height
new_cube = stratify.interpolate(target_Z,
Z,
cube,
axis=0,
extrapolation='nearest')
for level in range(1, new_cube.shape[0]):
f = g.iloc[0].at['field'].copy()
f.lblev = level + 1
f.blev = target_Zsea[level]
f.brlev = -1073741824
f.bhlev = target_C[level]
f.bhrlev = -1073741824
f.set_data_provider(mule.ArrayDataProvider(new_cube[level, :, :]))
ancil_out.fields.append(f)
ancil_out.to_file(outfile)
def test_zero_gradient(self):
assert_array_almost_equal(
stratify.interpolate([2], [0, 0], [1, 1], extrapolation='linear'),
[1])
def test_zero_gradient(self):
assert_array_equal(
stratify.interpolate([1], [0, 1, 1, 2], [10, 20, 30, 40],
interpolation='linear'), [20])
