remask_cube.py

import iris
import gridfill
import numpy as np
from scipy import interpolate


def remask_cube(cube, new_mask, interp_method='nearest'):
    ma_cube = cube.copy()
    data = cube.data
    if np.ma.is_masked(data):
        x = np.arange(data.shape[1])
        y = np.arange(data.shape[0])
        xx, yy = np.meshgrid(x, y)
        non_ma_x = xx[~data.mask]
        non_ma_y = yy[~data.mask]
        non_ma_data = data[~data.mask]
        if method in ('linear', 'nearest', 'cubic'):
            data = interpolate.griddata((non_ma_x, non_ma_y),
                                        non_ma_data.ravel(),
                                        (xx, yy),
                                        method=method)
        elif method == 'poisson':
            # use gridfill package
            data = gridfill.fill_cube(ma_cube, 0.01).data
        else:
            raise ValueError("Unknown interpolation method '{}'")
    ma_cube.data = np.ma.masked_where(new_mask, data)
    return ma_cube