def _regrid_weighted_curvilinear_to_rectilinear__perform(
src_cube, regrid_info):
"""
Second (regrid) part of 'regrid_weighted_curvilinear_to_rectilinear'.
Perform the prepared regrid calculation on a single 2d cube.
"""
from iris.cube import Cube
sparse_matrix, sum_weights, rows, grid_cube = regrid_info
# Calculate the numerator of the weighted mean (M, 1).
is_masked = ma.isMaskedArray(src_cube.data)
if not is_masked:
data = src_cube.data
else:
# Use raw data array
data = src_cube.data.data
# Check if there are any masked source points to take account of.
is_masked = np.ma.is_masked(src_cube.data)
if is_masked:
# Zero any masked source points so they add nothing in output sums.
mask = src_cube.data.mask
data[mask] = 0.0
# Calculate a new 'sum_weights' to allow for missing source points.
# N.B. it is more efficient to use the original once-calculated
# sparse matrix, but in this case we can't.
# Hopefully, this post-multiplying by the validities is less costly
# than repeating the whole sparse calculation.
valid_src_cells = ~mask.flat[:]
src_cell_validity_factors = sparse_diags(
np.array(valid_src_cells, dtype=int), 0)
valid_weights = sparse_matrix * src_cell_validity_factors
sum_weights = valid_weights.sum(axis=1).getA()
# Work out where output cells are missing all contributions.
# This allows for where 'rows' contains output cells that have no
# data because of missing input points.
zero_sums = sum_weights == 0.0
# Make sure we can still divide by sum_weights[rows].
sum_weights[zero_sums] = 1.0
# Calculate sum in each target cell, over contributions from each source
# cell.
numerator = sparse_matrix * data.reshape(-1, 1)
# Create a template for the weighted mean result.
weighted_mean = ma.masked_all(numerator.shape, dtype=numerator.dtype)
# Calculate final results in all relevant places.
weighted_mean[rows] = numerator[rows] / sum_weights[rows]
if is_masked:
# Ensure masked points where relevant source cells were all missing.
if np.any(zero_sums):
# Make masked if it wasn't.
weighted_mean = np.ma.asarray(weighted_mean)
# Mask where contributing sums were zero.
weighted_mean[zero_sums] = np.ma.masked
# Construct the final regridded weighted mean cube.
tx = grid_cube.coord(axis="x", dim_coords=True)
ty = grid_cube.coord(axis="y", dim_coords=True)
(tx_dim, ) = grid_cube.coord_dims(tx)
(ty_dim, ) = grid_cube.coord_dims(ty)
dim_coords_and_dims = list(zip((ty.copy(), tx.copy()), (ty_dim, tx_dim)))
cube = Cube(
weighted_mean.reshape(grid_cube.shape),
dim_coords_and_dims=dim_coords_and_dims,
)
cube.metadata = copy.deepcopy(src_cube.metadata)
for coord in src_cube.coords(dimensions=()):
cube.add_aux_coord(coord.copy())
return cube
def solve_hemholtz(f, boundary_func=lambda x,y: x*y, lap_f=lambda x: None, n=20,xint=(0,1), solver='spdiags', scheme='5-point', k=20):
h = (xint[1]-xint[0])/float(n+1)
if scheme == 'deferred':
#Get the second-order approximation
u_bar = np.zeros((n+2, n+2))
u_bar[1:-1,1:-1] = solve_hemholtz(f, boundary_func, lap_f, n, xint,k=k).reshape((n,n))
#Fill in the boundary data
grid_1d = np.linspace(xint[0], xint[1], n+2)
u_bar[0] = boundary_func(xint[0], grid_1d)
u_bar[-1] = boundary_func(xint[1], grid_1d)
u_bar[:,0] = boundary_func(grid_1d, xint[0])
u_bar[:,-1] = boundary_func(grid_1d, xint[1])
#Compute the xxyy derivative
u_xx = (u_bar[:-2]+u_bar[2:]-2*u_bar[1:-1])/h**2
u_xxyy = (u_xx[:,2:]+u_xx[:,:-2]-2*u_xx[:,1:-1])/h**2
#First, construct the diagonals and the RHS
#Here we assume that the domain is square
b = np.zeros((n,n))
yvals = np.linspace(xint[0], xint[1], n+2)[1:-1]
# print h, yvals[1]-yvals[0]
for i in xrange(n):
xval = yvals[i]
b[i] = f(xval, yvals)
if scheme == '9-point':
b[i] += h**2/12. * (lap_f(xval, yvals) - k**2*f(xval, yvals))
elif scheme == 'deferred':
b[i] += h**2/12. * (lap_f(xval, yvals) - k**2*f(xval, yvals) + k**4 * u_bar[i+1,1:-1] - 2*u_xxyy[i] )
if scheme == '5-point' or scheme == 'deferred':
b *= h**2
else: b *= 6*h**2
# print b
if scheme=='5-point' or scheme == 'deferred':
diags = make_5_point_diags(n, h)
diags[0] += h**2 * k**2
#Now we need to make use of the boundary conditions
b[0] -= boundary_func(xint[0], yvals)
b[-1] -= boundary_func(xint[1], yvals)
b[:,0] -= boundary_func(yvals, xint[0])
b[:,-1] -= boundary_func(yvals, xint[1])
elif scheme=='9-point':
diags = make_9_point_diags(n, h)
diags[0] += 6 *((h*k)**2 - 1./12* (h*k)**4)
#This is like the 5-point stencil
b[0] -= 4*boundary_func(xint[0], yvals)
b[-1] -= 4*boundary_func(xint[1], yvals)
b[:,0] -= 4*boundary_func(yvals, xint[0])
b[:,-1] -= 4*boundary_func(yvals, xint[1])
#Here we handle the corners of the stencil
vals = np.linspace(xint[0], xint[1], n+2)
#Take care of effects to the top and bottom boundaries
b[0] -= boundary_func(xint[0], vals[:-2]) + boundary_func(xint[0], vals[2:])
b[-1] -= boundary_func(xint[1], vals[:-2]) + boundary_func(xint[1], vals[2:])
#Handle the sides - be careful not to double count the corners
b[1:,0] -= boundary_func(vals[1:-2], xint[0])
b[:-1,0] -= boundary_func(vals[2:-1], xint[0])
b[1:,-1] -= boundary_func(vals[1:-2], xint[1])
b[:-1,-1] -= boundary_func(vals[2:-1], xint[1])
else:
raise ValueError("Unrecognized scheme {}".format(scheme))
#Next, create the sparse matrix and solve the system
# print([(d, sum(diags[d]!=0)) for d in diags if isinstance(diags[d], np.ndarray)])
offsets = sorted(diags.keys())
mat = sparse_diags([diags[d] for d in offsets], offsets, format='csc')
# print mat.todense()
if solver == 'spdiags':
# print b
# print mat.shape
return spsolve(mat, b.flatten())
elif solver == 'solve_banded':
ab, l_and_u = diags_to_banded(diags)
# print b.flatten().shape, ab.shape
return solve_banded(l_and_u,ab, b.flatten())
else:
return solver(mat, b.flatten())