Interpolate hybrid-height data onto a finer grid to make a big cube for performance testing.

expanded_hybrid_height.py

# Script to make bigger hybrid-height data.
# E.G. "$ python expanded_hybrid_height.py 5"

import datetime
import sys

import numpy as np

import iris
from iris.aux_factory import HybridHeightFactory
from iris.cube import Cube
from iris.coords import AuxCoord, DimCoord
from iris.analysis import Linear
import iris.tests.stock as istk


if __name__ == '__main__':
    if len(sys.argv) > 1:
        n_times = int(sys.argv[1])
    else:
        n_times = 1
    print 'Expansion factor : {n} * {n}'.format(n=n_times)

    cube = istk.realistic_4d()

    shape = cube.shape[-2:]
    shape = list(n_times * np.array(shape))
    grid = Cube(np.zeros(shape, dtype=np.int8))

    def co_expand(axis_name):
        co = cube.coord(axis=axis_name)
        x0, x1, nx = co.points[0], co.points[-1], co.shape[0]
        co_expanded = DimCoord(
            np.linspace(x0, x1, nx * n_times),
            standard_name=co.standard_name, units=co.units,
            coord_system=co.coord_system)
        return co_expanded

    co_x = co_expand('x')
    co_y = co_expand('y')
    grid.add_dim_coord(co_y, 0)
    grid.add_dim_coord(co_x, 1)

    co_orog = cube.coord('surface_altitude')
    orog = Cube(co_orog.points, units=co_orog.units)
    orog.add_dim_coord(cube.coord(axis='y'), 0)
    orog.add_dim_coord(cube.coord(axis='x'), 1)

    cube.remove_coord(co_orog)
    cube.remove_aux_factory(cube.aux_factories[0])

    big_cube = cube.regrid(grid, Linear())
    big_orog = orog.regrid(grid, Linear())

    big_co_orog = AuxCoord(big_orog.data, standard_name='surface_altitude', units=co_orog.units)
    ndims = len(big_cube.shape)
    big_cube.add_aux_coord(big_co_orog, (ndims-2, ndims-1))
    big_cube.add_aux_factory(
        HybridHeightFactory(
            delta=big_cube.coord('level_height'),
            sigma=big_cube.coord('sigma'),
            orography=big_co_orog))

    print big_cube
    print 'orog shape = ', big_cube.coord('surface_altitude').shape
    file_path = 'big_hybrid_height_cube.nc'
    print 'save to : ', file_path

    iris.FUTURE.netcdf_no_unlimited = True
    iris.save(big_cube, file_path)