def rotateGrid(data):
##
## RELATE GRID TO ROTATED POLE
##
import iris.analysis.cartography as ircrt
### LAM Configuration from suite u-bg610
dlat = 0.015714
dlon = 0.016334
frst_lat = -5.6112
frst_lon = 353.0345
pole_lat = 37.5000
pole_lon = 177.5000
lat = np.arange((centlat-(gry*float(0.5)*dlat)),(centlat+(gry*float(0.5)*dlat)),dlat)
lon = np.arange((centlon-(grx*float(0.5)*dlon)),(centlon+(grx*float(0.5)*dlon)),dlon)
rot_lon, rot_lat = ircrt.rotate_pole(lon, lat, frst_lon, frst_lat)
# Print to check conversion
print ('******')
print ('Test of rotated coordinate grid: ')
print ('Lon = ', lon[0])
print ('Lat = ', lat[0])
print ('Rotated lon coord = ', rot_lon[0])
print ('Rotated lat coord = ', rot_lat[0])
print (' ')
# ******
# lat/lon vertices of nest (output): 85.960219407715 80.41973098346767 49.567255645848604 -27.55740381723681
# ******
return lon, lat
def rdf(name, trajectories, pole_lon, pole_lat):
""" Get the values of variable at the start of trajectories
Args:
name (str): Name of variable to calculate at the start of trajectories
trajectories (list): Length n list of trajectory objects
pole_lon (float): Rotated pole longitude
pole_lat (float): Rotated pole latitude
Returns:
rlons (list): Length n list of rotated longitude at end of trajectories
rlons (list): Length n list of rotated latitude at end of trajectories
values (list): Length n list of variable at start of trajectories
"""
# Initialise output lists
rlons = []
rlats = []
values = []
# Loop over all trajectories
for n, trajectory in enumerate(trajectories):
if n % 10000 == 0:
print(n)
# Extract the position at the last trajectory point
time = trajectory.variable('time')[-1]
lon = trajectory.variable('lon')[-1]
lat = trajectory.variable('lat')[-1]
p = trajectory.variable('p')[-1]
pv = trajectory.variable('PV')[-1]
# Extract the data from the forecast
if n == 1:
forecast.set_time(time)
cube = convert.calc(name, forecast.cubelist)
# Get the starting position
lon = trajectory.variable('lon')[0]
lat = trajectory.variable('lat')[0]
rlon, rlat = rotate_pole(np.array(lon), np.array(lat),
pole_lon, pole_lat)
# Save the values for plotting
rlons.append(float(rlon + 360))
rlats.append(float(rlat))
values.append(float(pv))
return rlons, rlats, values, cube
def in_domain(trajectory, pole_lon, pole_lat, domain):
"""Checks whether the trajectory is in the given domain
"""
for n in xrange(len(trajectory)):
lon = trajectory.variable('lon')[n]
lat = trajectory.variable('lat')[n]
rlon, rlat = rotate_pole(np.array(lon), np.array(lat),
pole_lon, pole_lat)
rlon = rlon + 360
if outside_bounds(rlon, rlat, domain):
trajectory.data = trajectory.data[0:(n + 1)]
return
return
def test_multidim(self):
# Testing with >2D data to demonstrate correct operation over
# additional non-XY dimensions (including data masking), which is
# handled by the PointInCell wrapper class.
# Define a simple target grid first, in plain latlon coordinates.
plain_latlon_cs = GeogCS(EARTH_RADIUS)
grid_x_coord = DimCoord(points=[15.0, 25.0, 35.0],
bounds=[[10.0, 20.0], [20.0, 30.0],
[30.0, 40.0]],
standard_name='longitude',
units='degrees',
coord_system=plain_latlon_cs)
grid_y_coord = DimCoord(points=[-30.0, -50.0],
bounds=[[-20.0, -40.0], [-40.0, -60.0]],
standard_name='latitude',
units='degrees',
coord_system=plain_latlon_cs)
grid_cube = Cube(np.zeros((2, 3)))
grid_cube.add_dim_coord(grid_y_coord, 0)
grid_cube.add_dim_coord(grid_x_coord, 1)
# Define some key points in true-lat/lon thta have known positions
# First 3x2 points in the centre of each output cell.
x_centres, y_centres = np.meshgrid(grid_x_coord.points,
grid_y_coord.points)
# An extra point also falling in cell 1, 1
x_in11, y_in11 = 26.3, -48.2
# An extra point completely outside the target grid
x_out, y_out = 70.0, -40.0
# Define a rotated coord system for the source data
pole_lon, pole_lat = -125.3, 53.4
src_cs = RotatedGeogCS(grid_north_pole_latitude=pole_lat,
grid_north_pole_longitude=pole_lon,
ellipsoid=plain_latlon_cs)
# Concatenate all the testpoints in a flat array, and find the rotated
# equivalents.
xx = list(x_centres.flat[:]) + [x_in11, x_out]
yy = list(y_centres.flat[:]) + [y_in11, y_out]
xx, yy = rotate_pole(lons=np.array(xx),
lats=np.array(yy),
pole_lon=pole_lon,
pole_lat=pole_lat)
# Define handy index numbers for all these.
i00, i01, i02, i10, i11, i12, i_in, i_out = range(8)
# Build test data in the shape Z,YX = (3, 8)
data = [[1, 2, 3, 11, 12, 13, 7, 99], [1, 2, 3, 11, 12, 13, 7, 99],
[7, 6, 5, 51, 52, 53, 12, 1]]
mask = [[0, 0, 0, 0, 0, 0, 0, 0], [0, 1, 0, 0, 0, 0, 1, 0],
[0, 0, 0, 0, 0, 0, 0, 0]]
src_data = np.ma.array(data, mask=mask, dtype=float)
# Make the source cube.
src_cube = Cube(src_data)
src_x = AuxCoord(xx,
standard_name='grid_longitude',
units='degrees',
coord_system=src_cs)
src_y = AuxCoord(yy,
standard_name='grid_latitude',
units='degrees',
coord_system=src_cs)
src_z = DimCoord(np.arange(3), long_name='z')
src_cube.add_dim_coord(src_z, 0)
src_cube.add_aux_coord(src_x, 1)
src_cube.add_aux_coord(src_y, 1)
# Add in some extra metadata, to ensure it gets copied over.
src_cube.add_aux_coord(DimCoord([0], long_name='extra_scalar_coord'))
src_cube.attributes['extra_attr'] = 12.3
# Define what the expected answers should be, shaped (3, 2, 3).
expected_result = [
[[1.0, 2.0, 3.0], [11.0, 0.5 * (12 + 7), 13.0]],
[[1.0, -999, 3.0], [11.0, 12.0, 13.0]],
[[7.0, 6.0, 5.0], [51.0, 0.5 * (52 + 12), 53.0]],
]
expected_result = np.ma.masked_less(expected_result, 0)
# Perform the calculation with the regridder.
regridder = Regridder(src_cube, grid_cube)
# Check all is as expected.
result = regridder(src_cube)
self.assertEqual(result.coord('z'), src_cube.coord('z'))
self.assertEqual(result.coord('extra_scalar_coord'),
src_cube.coord('extra_scalar_coord'))
self.assertEqual(result.coord('longitude'),
grid_cube.coord('longitude'))
self.assertEqual(result.coord('latitude'), grid_cube.coord('latitude'))
self.assertMaskedArrayAlmostEqual(result.data, expected_result)
def test_multidim(self):
# Testing with >2D data to demonstrate correct operation over
# additional non-XY dimensions (including data masking), which is
# handled by the PointInCell wrapper class.
# Define a simple target grid first, in plain latlon coordinates.
plain_latlon_cs = GeogCS(EARTH_RADIUS)
grid_x_coord = DimCoord(points=[15.0, 25.0, 35.0],
bounds=[[10.0, 20.0],
[20.0, 30.0],
[30.0, 40.0]],
standard_name='longitude',
units='degrees',
coord_system=plain_latlon_cs)
grid_y_coord = DimCoord(points=[-30.0, -50.0],
bounds=[[-20.0, -40.0], [-40.0, -60.0]],
standard_name='latitude',
units='degrees',
coord_system=plain_latlon_cs)
grid_cube = Cube(np.zeros((2, 3)))
grid_cube.add_dim_coord(grid_y_coord, 0)
grid_cube.add_dim_coord(grid_x_coord, 1)
# Define some key points in true-lat/lon thta have known positions
# First 3x2 points in the centre of each output cell.
x_centres, y_centres = np.meshgrid(grid_x_coord.points,
grid_y_coord.points)
# An extra point also falling in cell 1, 1
x_in11, y_in11 = 26.3, -48.2
# An extra point completely outside the target grid
x_out, y_out = 70.0, -40.0
# Define a rotated coord system for the source data
pole_lon, pole_lat = -125.3, 53.4
src_cs = RotatedGeogCS(grid_north_pole_latitude=pole_lat,
grid_north_pole_longitude=pole_lon,
ellipsoid=plain_latlon_cs)
# Concatenate all the testpoints in a flat array, and find the rotated
# equivalents.
xx = list(x_centres.flat[:]) + [x_in11, x_out]
yy = list(y_centres.flat[:]) + [y_in11, y_out]
xx, yy = rotate_pole(lons=np.array(xx),
lats=np.array(yy),
pole_lon=pole_lon,
pole_lat=pole_lat)
# Define handy index numbers for all these.
i00, i01, i02, i10, i11, i12, i_in, i_out = range(8)
# Build test data in the shape Z,YX = (3, 8)
data = [[1, 2, 3, 11, 12, 13, 7, 99],
[1, 2, 3, 11, 12, 13, 7, 99],
[7, 6, 5, 51, 52, 53, 12, 1]]
mask = [[0, 0, 0, 0, 0, 0, 0, 0],
[0, 1, 0, 0, 0, 0, 1, 0],
[0, 0, 0, 0, 0, 0, 0, 0]]
src_data = np.ma.array(data, mask=mask, dtype=float)
# Make the source cube.
src_cube = Cube(src_data)
src_x = AuxCoord(xx,
standard_name='grid_longitude',
units='degrees',
coord_system=src_cs)
src_y = AuxCoord(yy,
standard_name='grid_latitude',
units='degrees',
coord_system=src_cs)
src_z = DimCoord(np.arange(3), long_name='z')
src_cube.add_dim_coord(src_z, 0)
src_cube.add_aux_coord(src_x, 1)
src_cube.add_aux_coord(src_y, 1)
# Add in some extra metadata, to ensure it gets copied over.
src_cube.add_aux_coord(DimCoord([0], long_name='extra_scalar_coord'))
src_cube.attributes['extra_attr'] = 12.3
# Define what the expected answers should be, shaped (3, 2, 3).
expected_result = [
[[1.0, 2.0, 3.0],
[11.0, 0.5 * (12 + 7), 13.0]],
[[1.0, -999, 3.0],
[11.0, 12.0, 13.0]],
[[7.0, 6.0, 5.0],
[51.0, 0.5 * (52 + 12), 53.0]],
]
expected_result = np.ma.masked_less(expected_result, 0)
# Perform the calculation with the regridder.
regridder = Regridder(src_cube, grid_cube)
# Check all is as expected.
result = regridder(src_cube)
self.assertEqual(result.coord('z'), src_cube.coord('z'))
self.assertEqual(result.coord('extra_scalar_coord'),
src_cube.coord('extra_scalar_coord'))
self.assertEqual(result.coord('longitude'),
grid_cube.coord('longitude'))
self.assertEqual(result.coord('latitude'),
grid_cube.coord('latitude'))
self.assertMaskedArrayAlmostEqual(result.data, expected_result)
def extract_section(cube, waypoints, n_sample_points=30):
"""
Extract a section from the cube given lat-lon waypoints. This works
on any cube with x and y dimension coordinates. If the cube also contains
Z or T etc dimension coordinates then these are retained in the extracted
section.
:param cube: Iris cube
:param waypoints: List of dictionaries, whose keys are 'latitude' and
'longitude' and whose values are the lat/lon points
that should be included in the section.
:param n_sample_points: Number of sample points to include in section
:returns: An iris cube which no longer has X and Y as dimension coordinates
but instead has an 'i_sample_point' coordinate.
If latitude and longitude were not in the original cube, they
are added to to the new section cube.
Load sample data - a gridded surface-only, time-varying cube:
>>> import config
>>> sample_datadir = config.SAMPLE_DATADIR+'gridded_cube_list/'
>>> cube = iris.load_cube(sample_datadir+'aqum_oper_1days.nc',
... 'mass_concentration_of_ozone_in_air')
>>> print(cube) # doctest: +NORMALIZE_WHITESPACE
mass_concentration_of_ozone_in_air / (ug/m3) (time: 25; grid_latitude: 182; \
grid_longitude: 146)
Dimension coordinates:
time x - -
grid_latitude - x -
grid_longitude - - x
Auxiliary coordinates:
forecast_period x - -
surface_altitude - x x
Derived coordinates:
altitude - x x
Scalar coordinates:
atmosphere_hybrid_height_coordinate: 20.000338 m, \
bound=(0.0, 49.998882) m
forecast_day: 1.0 Days
model_level_number: 1
sigma: 0.9977165, bound=(1.0, 0.99429625)
Attributes:
Conventions: CF-1.5
STASH: m01s34i001
label: aqum_oper
short_name: O3
source: Data from Met Office Unified Model
Cell methods:
mean: time (1 hour)
Set up waypoints list:
>>> waypoints = [
... {'latitude': 50.8, 'longitude': -1.8},
... {'latitude': 51.2, 'longitude': -1.2},
... {'latitude': 51.4, 'longitude': -0.9}]
Extract section along these waypoints:
>>> section = extract_section(cube, waypoints)
>>> print(section)
mass_concentration_of_ozone_in_air / (ug/m3) (time: 25; i_sample_point: 30)
Dimension coordinates:
time x -
i_sample_point - x
Auxiliary coordinates:
forecast_period x -
grid_latitude - x
grid_longitude - x
latitude - x
longitude - x
surface_altitude - x
Derived coordinates:
altitude - x
Scalar coordinates:
atmosphere_hybrid_height_coordinate: 20.000338 m, \
bound=(0.0, 49.998882) m
forecast_day: 1.0 Days
model_level_number: 1
sigma: 0.9977165, bound=(1.0, 0.99429625)
Attributes:
Conventions: CF-1.5
STASH: m01s34i001
label: aqum_oper
short_name: O3
source: Data from Met Office Unified Model
Cell methods:
mean: time (1 hour)
"""
xcoord, ycoord = guess_coord_names(cube, ['X', 'Y'])
if None in [xcoord, ycoord]:
raise ValueError('Can not calculate section for this cube')
lat_lon_coord_system = iris.coord_systems.GeogCS(
semi_major_axis=iris.fileformats.pp.EARTH_RADIUS)
lons = np.array([waypoint['longitude'] for waypoint in waypoints])
lats = np.array([waypoint['latitude'] for waypoint in waypoints])
#Calculate waypoints in cube coordinates
if xcoord == 'grid_longitude' and ycoord == 'grid_latitude':
#Rotated pole
pole_lon = cube.coord(xcoord).coord_system.grid_north_pole_longitude
pole_lat = cube.coord(ycoord).coord_system.grid_north_pole_latitude
#Perform rotation
rot_lons, rot_lats = rotate_pole(lons, lats, pole_lon, pole_lat)
#Put back into waypoints format
cube_waypoints = [{
xcoord: lon,
ycoord: lat
} for lat, lon in zip(rot_lats, rot_lons)]
elif xcoord == 'projection_x_coordinate' and \
ycoord == 'projection_y_coordinate':
#Other coordinate system (note this may work for x/ycoords other than
#those considered here
ll_crs = lat_lon_coord_system.as_cartopy_crs()
cube_crs = cube.coord(xcoord).coord_system.as_cartopy_crs()
#Convert to lat/lon points
cube_lonlats = ll_crs.transform_points(cube_crs, lons, lats)
cube_lons = cube_lonlats[:, 0]
cube_lats = cube_lonlats[:, 1]
#Put back into waypoints format
cube_waypoints = [{
xcoord: lon,
ycoord: lat
} for lat, lon in zip(cube_lats, cube_lons)]
elif xcoord == 'longitude' and ycoord == 'latitude':
cube_waypoints = waypoints
else:
raise ValueError('Unable to convert cube x/y points to lat/lon')
#Specify the trajectory (straight-line) by giving the
# start and end of the line and specifying how many points there
# should be on the line (in this case 30)
sample_points = trajectory.Trajectory(cube_waypoints,
sample_count=n_sample_points)
#Extract the data from the cube at the sample points
#(Note this only works with iris vn2.2+)
section = sample_points.interpolate(cube)
#Rename new index coordinate
section.coord('index').rename('i_sample_point')
#Also add latitude and longitude coords to section
#if not in original cube
if xcoord != 'longitude' and ycoord != 'latitude':
cube_lons_samplepts = np.array(
[d[xcoord] for d in sample_points.sampled_points])
cube_lats_samplepts = np.array(
[d[ycoord] for d in sample_points.sampled_points])
if xcoord == 'grid_longitude' and ycoord == 'grid_latitude':
section_lons, section_lats = unrotate_pole(cube_lons_samplepts,
cube_lats_samplepts,
pole_lon, pole_lat)
elif xcoord == 'projection_x_coordinate' and \
ycoord == 'projection_y_coordinate':
#Other coordinate system (note this may work for x/ycoords other than
#those considered here
lonlats = cube_crs.transform_points(ll_crs, cube_lons_samplepts,
cube_lats_samplepts)
section_lons = lonlats[:, 0]
section_lats = lonlats[:, 1]
#Set up lat/lon coords and add as aux coords
lon_coord = iris.coords.AuxCoord(section_lons,
standard_name='longitude',
units='degrees',
coord_system=lat_lon_coord_system)
lat_coord = iris.coords.AuxCoord(section_lats,
standard_name='latitude',
units='degrees',
coord_system=lat_lon_coord_system)
section.add_aux_coord(lon_coord, section.coord_dims(xcoord))
section.add_aux_coord(lat_coord, section.coord_dims(ycoord))
return section
def extract_vert_section(cube, pnts):
"""
Extract vertical slice of a cube
Assume cube's dims: (..., y, x)
"""
sample_count = pnts['sample_count']
zcoord = cube.coord(axis='z', dim_coords=True)
zspan = int(zcoord.points[0]), int(zcoord.points[-1])
if 'x' in pnts and 'y' in pnts:
# Grid indices (no interpolation needed)
if len(pnts['x']) == 2 and len(pnts['y']) == 2:
if ((pnts['x'] == pnts['x'][0]).all() and
not (pnts['y'] == pnts['y'][0]).all()):
# X const, Y varies
xslice = pnts['x'][0]
if pnts['y'][0] < pnts['y'][1]:
stride = 1
else:
stride = -1
yslice = slice(*pnts['y'], stride)
sect_info = dict(v=zcoord.name(),
h='y_axis_ind',
label='z{}-{}y{}-{}x{}'.format(*zspan,
*pnts['y'],
pnts['x'][0]))
elif ((pnts['y'] == pnts['y'][0]).all() and
not (pnts['x'] == pnts['x'][0]).all()):
# Y const, X varies
yslice = pnts['y'][0]
if pnts['x'][0] < pnts['x'][1]:
stride = 1
else:
stride = -1
xslice = slice(*pnts['x'], stride)
sect_info = dict(v=zcoord.name(),
h='x_axis_ind',
label='z{}-{}y{}x{}-{}'.format(*zspan,
pnts['y'][0],
*pnts['x']))
sect = cube[..., yslice, xslice]
else:
_msg = (f'Only 2 point pairs is allowed for grid indices option'
',\n{pnts} are given')
raise NotYetImplementedError(_msg)
elif 'lon' in pnts and 'lat' in pnts:
# Lon / lat coordinates
cs = cube.coord_system()
xcoord = cube.coord(axis='x', dim_coords=True)
ycoord = cube.coord(axis='y', dim_coords=True)
xp, yp = xcoord.points, ycoord.points
if (xp > 180).any():
xp = xp - 360
if isinstance(cs, (iris.coord_systems.RotatedGeogCS, )):
# Rotated coordinate system
# Use iris interpolation along a trajectory
nplat = cs.grid_north_pole_latitude
nplon = cs.grid_north_pole_longitude
waypoints = []
for ilon, ilat in zip(pnts['lon'], pnts['lat']):
waypoints.append({'longitude': ilon, 'latitude': ilat})
if sample_count == 'auto':
_x, _y = [], []
for ilon, ilat in zip(pnts['lon'], pnts['lat']):
_ilon, _ilat = rotate_pole(np.asarray(ilon),
np.asarray(ilat),
nplon, nplat)
_x.append(np.argmin(abs(xp - _ilon[0])))
_y.append(np.argmin(abs(yp - _ilat[0])))
sample_count = int(((np.diff(_x)**2 +
np.diff(_y)**2)**0.5).sum())
else:
assert isinstance(sample_count, int)
traj = trajectory.Trajectory(waypoints, sample_count=sample_count)
lon = np.array([d['longitude'] for d in traj.sampled_points])
lat = np.array([d['latitude'] for d in traj.sampled_points])
rlon, rlat = iris.analysis.cartography.rotate_pole(lon, lat,
nplon, nplat)
sampled_points = [(xcoord.name(), rlon),
(ycoord.name(), rlat)]
sect = trajectory.interpolate(cube, sampled_points)
_add_real_lonlat(sect, lon, lat)
latspan = pnts['lat'][0], pnts['lat'][-1]
lonspan = pnts['lon'][0], pnts['lon'][-1]
sect_info = dict(v=zcoord.name(),
h='index',
label='z{}-{}lat{}-{}lon{}-{}'.format(*zspan,
*latspan,
*lonspan))
elif isinstance(cs, (None, iris.coord_systems.GeogCS)):
# Rectangular grid
if len(pnts['lon']) == 2 and len(pnts['lat']) == 2:
if ((pnts['lon'] == pnts['lon'][0]).all() and
not (pnts['lat'] == pnts['lat'][0]).all()):
# X const, Y varies
xslice = np.argmin(abs(xp - pnts['lon'][0]))
yslice = slice(*[np.argmin(abs(yp - i))
for i in pnts['lat']])
elif ((pnts['lat'] == pnts['lat'][0]).all() and
not (pnts['lon'] == pnts['lon'][0]).all()):
# Y const, X varies
yslice = np.argmin(abs(yp - pnts['lat'][0]))
xslice = slice(*[np.argmin(abs(xp - i))
for i in pnts['lon']])
sect = cube[..., yslice, xslice]
sect_info = {} # TODO
else:
# Diagonal
raise NotYetImplementedError()
else:
raise NotYetImplementedError(f"Can't deal with "
"{cube.coord_system()}")
sect.attributes['sect_info'] = sect_info
sect.attributes['sect_info']['pnts'] = pnts
return sect
def calc_distances(lons, lats, lon, lat):
""" Calculate distances between all lats/lons to specified lat/lon """
lons_rot, lats_rot = cart.rotate_pole(lons, lats, lon, lat)
dists = 111. * (90 - lats_rot)
return dists
