|
| 1 | +from __future__ import absolute_import |
| 2 | + |
| 3 | +import numpy as np |
| 4 | + |
| 5 | +from ._conservative import conservative_interpolation |
| 6 | + |
| 7 | + |
| 8 | +def interpolate_conservative(z_target, z_src, fz_src, axis=-1): |
| 9 | + """ |
| 10 | + 1d conservative interpolation across multiple dimensions. |
| 11 | +
|
| 12 | + This function provides the ability to perform 1d interpolation on datasets |
| 13 | + with more than one dimension. For instance, this function can be used to |
| 14 | + interpolate a set of vertical levels, even if the interpolation coordinate |
| 15 | + depends upon other dimensions. |
| 16 | +
|
| 17 | + A good use case might be when wanting to interpolate at a specific height |
| 18 | + for height data which also depends on x and y - e.g. extract 1000hPa level |
| 19 | + from a 3d dataset and associated pressure field. In the case of this |
| 20 | + example, pressure would be the `z` coordinate, and the dataset |
| 21 | + (e.g. geopotential height / temperature etc.) would be `f(z)`. |
| 22 | +
|
| 23 | + Parameters |
| 24 | + ---------- |
| 25 | + z_target: :class:`np.ndarray` |
| 26 | + Target coordinate. |
| 27 | + This coordinate defines the levels to interpolate the source data |
| 28 | + ``fz_src`` to. ``z_target`` must have the same dimensionality as the |
| 29 | + source coordinate ``z_src``, and the shape of ``z_target`` must match |
| 30 | + the shape of ``z_src``, although the axis of interpolation may differ |
| 31 | + in dimension size. |
| 32 | + z_src: :class:`np.ndarray` |
| 33 | + Source coordinate. |
| 34 | + This coordinate defines the levels that the source data ``fz_src`` is |
| 35 | + interpolated from. |
| 36 | + fz_src: :class:`np.ndarray` |
| 37 | + The source data; the phenomenon data values to be interpolated from |
| 38 | + ``z_src`` to ``z_target``. |
| 39 | + The data array must be at least ``z_src.ndim``, and its trailing |
| 40 | + dimensions (i.e. those on its right hand side) must be exactly |
| 41 | + the same as the shape of ``z_src``. |
| 42 | + axis: int (default -1) |
| 43 | + The ``fz_src`` axis to perform the interpolation over. |
| 44 | +
|
| 45 | + Returns |
| 46 | + ------- |
| 47 | + : :class:`np.ndarray` |
| 48 | + fz_src interpolated from z_src to z_target. |
| 49 | +
|
| 50 | + Note |
| 51 | + ---- |
| 52 | + - Support for 1D z_target and corresponding ND z_src will be provided in |
| 53 | + future as driven by user requirement. |
| 54 | + - Those cells, where 'nan' values in the source data contribute, a 'nan' |
| 55 | + value is returned. |
| 56 | +
|
| 57 | + """ |
| 58 | + if z_src.ndim != z_target.ndim: |
| 59 | + msg = ('Expecting source and target levels dimensionality to be ' |
| 60 | + 'identical. {} != {}.') |
| 61 | + raise ValueError(msg.format(z_src.ndim, z_target.ndim)) |
| 62 | + |
| 63 | + # Relative axis |
| 64 | + axis = axis % fz_src.ndim |
| 65 | + axis_relative = axis - (fz_src.ndim - (z_target.ndim-1)) |
| 66 | + |
| 67 | + src_shape = list(z_src.shape) |
| 68 | + src_shape.pop(axis_relative) |
| 69 | + tgt_shape = list(z_target.shape) |
| 70 | + tgt_shape.pop(axis_relative) |
| 71 | + |
| 72 | + if src_shape != tgt_shape: |
| 73 | + src_shape = list(z_src.shape) |
| 74 | + src_shape[axis_relative] = '-' |
| 75 | + tgt_shape = list(z_target.shape) |
| 76 | + src_shape[axis_relative] = '-' |
| 77 | + msg = ('Expecting the shape of the source and target levels except ' |
| 78 | + 'the axis of interpolation to be identical. {} != {}') |
| 79 | + raise ValueError(msg.format(tuple(src_shape), tuple(tgt_shape))) |
| 80 | + |
| 81 | + dat_shape = list(fz_src.shape) |
| 82 | + dat_shape = dat_shape[-(z_src.ndim-1):] |
| 83 | + src_shape = list(z_src.shape[:-1]) |
| 84 | + if dat_shape != src_shape: |
| 85 | + dat_shape = list(fz_src.shape) |
| 86 | + dat_shape[:-(z_src.ndim-1)] = '-' |
| 87 | + msg = ('The provided data is not of compatible shape with the ' |
| 88 | + 'provided source bounds. {} != {}') |
| 89 | + raise ValueError(msg.format(tuple(dat_shape), tuple(src_shape))) |
| 90 | + |
| 91 | + if z_src.shape[-1] != 2: |
| 92 | + msg = 'Unexpected source and target bounds shape. shape[-1] != 2' |
| 93 | + raise ValueError(msg) |
| 94 | + |
| 95 | + # Define our source in a consistent way. |
| 96 | + # [broadcasting_dims, axis_interpolation, z_varying] |
| 97 | + |
| 98 | + # src_data |
| 99 | + bdims = list(range(fz_src.ndim - (z_src.ndim-1))) |
| 100 | + data_vdims = [ind for ind in range(fz_src.ndim) if ind not in |
| 101 | + (bdims + [axis])] |
| 102 | + data_transpose = bdims + [axis] + data_vdims |
| 103 | + fz_src_reshaped = np.transpose(fz_src, data_transpose) |
| 104 | + fz_src_orig = list(fz_src_reshaped.shape) |
| 105 | + shape = ( |
| 106 | + int(np.product(fz_src_reshaped.shape[:len(bdims)])), |
| 107 | + fz_src_reshaped.shape[len(bdims)], |
| 108 | + int(np.product(fz_src_reshaped.shape[len(bdims)+1:]))) |
| 109 | + fz_src_reshaped = fz_src_reshaped.reshape(shape) |
| 110 | + |
| 111 | + # Define our src and target bounds in a consistent way. |
| 112 | + # [axis_interpolation, z_varying, 2] |
| 113 | + vdims = list(set(range(z_src.ndim)) - set([axis_relative])) |
| 114 | + z_src_reshaped = np.transpose(z_src, [axis_relative] + vdims) |
| 115 | + z_target_reshaped = np.transpose(z_target, [axis_relative] + vdims) |
| 116 | + |
| 117 | + shape = int(np.product(z_src_reshaped.shape[1:-1])) |
| 118 | + z_src_reshaped = z_src_reshaped.reshape([z_src_reshaped.shape[0], shape, |
| 119 | + z_src_reshaped.shape[-1]]) |
| 120 | + shape = int(np.product(z_target_reshaped.shape[1:-1])) |
| 121 | + z_target_reshaped = z_target_reshaped.reshape( |
| 122 | + [z_target_reshaped.shape[0], shape, z_target_reshaped.shape[-1]]) |
| 123 | + |
| 124 | + result = conservative_interpolation( |
| 125 | + z_src_reshaped, z_target_reshaped, fz_src_reshaped) |
| 126 | + |
| 127 | + # Turn the result into a shape consistent with the source. |
| 128 | + # First reshape, then reverse transpose. |
| 129 | + shape = fz_src_orig |
| 130 | + shape[len(bdims)] = z_target.shape[axis_relative] |
| 131 | + result = result.reshape(shape) |
| 132 | + invert_transpose = [data_transpose.index(ind) for ind in |
| 133 | + list(range(result.ndim))] |
| 134 | + result = result.transpose(invert_transpose) |
| 135 | + return result |
0 commit comments