test_shoc_standard.py

import itertools
import json
import logging
import pathlib

import numpy
import pandas
import pytest
import shapely
import xarray
from matplotlib.figure import Figure
from numpy.testing import assert_equal
from shapely.geometry.polygon import orient

from emsarray import plot
from emsarray.conventions import DimensionGrid, get_dataset_convention
from emsarray.conventions.arakawa_c import (
    ArakawaCGridKind, c_mask_from_centres
)
from emsarray.conventions.shoc import ShocStandard
from emsarray.operations import geometry
from tests.helpers.array import mask_from_strings
from tests.helpers.datasets import (
    DiagonalShocGrid, ShocGridGenerator, ShocLayerGenerator
)
from tests.helpers.memory import track_peak_memory_usage

logger = logging.getLogger(__name__)


def make_dataset(
    *,
    k_size: int = 5,
    j_size: int,
    i_size: int,
    time_size: int = 4,
    grid_type: type[ShocGridGenerator] = DiagonalShocGrid,
    corner_size: int = 0,
) -> xarray.Dataset:
    """
    Make a dummy SHOC standard dataset of a particular size.
    It will have a sheared grid of points located near (0, 0),
    with increasing j moving north east, and increasing i moving south east.

    For a grid with j_size=5, i_size=9, the centre longitude-latitude will be:

    * index (0, 0): (0.1, 0.9),
    * index (0, 1): (0.2, 0.8),
    * index (0, 8): (0.9, 0.1)
    * index (1, 0): (0.2, 0.10)
    * index (1, 8): (0.10, 0.2)
    * index (4, 0): (0.5, 0.13)

    Three data variables will be made: botz, eta, and temp.
    Each will be filled with random data.
    There will be a one-cell border around the edge of the dataset where all
    data variables will be nans.

    The (i=0, j=j_size) corner will not have any coordinates in a 4x4 box.
    The (i=i_size, j=j_size) corner will have coordinates,
    but data variables will be masked off
    """
    coordinate_centre_mask = numpy.full((j_size, i_size), True)
    # Cut a chunk out of the corner where the coordinates will not be defined.
    if corner_size > 1:
        coordinate_centre_mask[-(corner_size - 1):, :+(corner_size - 1)] = False

    # SHOC files have a 1-cell border around the outside where the cells have
    # coordinates, but no data.
    wet_centre_mask = numpy.full((j_size, i_size), True)
    if corner_size > 0:
        wet_centre_mask[-corner_size:, :+corner_size] = False
        wet_centre_mask[-corner_size:, -corner_size:] = False
    wet_centre_mask[:+1, :] = False
    wet_centre_mask[-1:, :] = False
    wet_centre_mask[:, :+1] = False
    wet_centre_mask[:, -1:] = False
    wet_mask = c_mask_from_centres(wet_centre_mask, {
        ArakawaCGridKind.face: ('j_centre', 'i_centre'),
        ArakawaCGridKind.back: ('j_back', 'i_back'),
        ArakawaCGridKind.left: ('j_left', 'i_left'),
        ArakawaCGridKind.node: ('j_node', 'i_node'),
    })

    # These DataArrays are the long/lats of the grid corners. The centres are
    # derived from these by averaging the surrounding four corners.
    grid = grid_type(j=j_size, i=i_size, face_mask=coordinate_centre_mask)
    layers = ShocLayerGenerator(k=k_size)

    t = xarray.DataArray(
        # Note: Using pandas.date_range() directly here will lead to strange
        # behaviours, where the `record` dimension becomes a data variable with
        # a datetime64 dtype. Using a list of datetimes instead seems to avoid
        # this, resulting in record simply being a dimension.
        data=list(pandas.date_range("2021-11-11", periods=time_size)),
        dims=["record"],
        attrs={
            "long_name": "Time",
            "standard_name": "time",
            "coordinate_type": "time",
        },
    )
    # Note: xarray will reformat this in to 1990-01-01T00:00:00+10:00, which
    # EMS fails to parse. There is no way around this using xarray natively,
    # you have to adjust it with nctool after saving it.
    t.encoding["units"] = "days since 1990-01-01 00:00:00 +10"
    # This can not be represented as an int, so explicitly set the dtype.
    # xarray >=2023.09 warns when encoding variables without a dtype
    # that can not be represented exactly.
    t.encoding["dtype"] = "float32"

    botz = xarray.DataArray(
        data=numpy.random.random((j_size, i_size)) * 10 + 50,
        dims=wet_mask["face_mask"].dims,
        attrs={
            "units": "metre",
            "long_name": "Z coordinate at sea-bed at cell centre",
            "standard_name": "depth",
            "positive": "down",
            "outside": "9999",
            "missing_value": -99.,
        }
    ).where(wet_mask.data_vars["face_mask"])
    botz.values[1, 1] = -99.

    eta = xarray.DataArray(
        data=numpy.random.normal(0, 0.2, (time_size, j_size, i_size)),
        dims=["record", *wet_mask["face_mask"].dims],
        attrs={
            "units": "metre",
            "long_name": "Surface elevation",
            "standard_name": "sea_surface_height_above_geoid",
        }
    ).where(wet_mask.data_vars["face_mask"])
    temp = xarray.DataArray(
        data=numpy.random.normal(12, 0.5, (time_size, k_size, j_size, i_size)),
        dims=["record", "k_centre", *wet_mask["face_mask"].dims],
        attrs={
            "units": "degrees C",
            "long_name": "Temperature",
        },
    ).where(wet_mask.data_vars["face_mask"])

    u1 = xarray.DataArray(
        data=numpy.random.normal(0, 2, (time_size, k_size, j_size, i_size + 1)),
        dims=["record", "k_centre", *wet_mask.data_vars["left_mask"].dims],
        attrs={
            "units": "metre second-1",
            "long_name": "I component of current at left face",
        }
    )
    u2 = xarray.DataArray(
        data=numpy.random.normal(0, 2, (time_size, k_size, j_size + 1, i_size)),
        dims=["record", "k_centre", *wet_mask.data_vars["back_mask"].dims],
        attrs={
            "units": "metre per second",
            "long_name": "I component of current at back face",
        }
    )
    xx, yy = numpy.meshgrid(
        numpy.linspace(-5, 5, i_size),
        numpy.linspace(-5, 5, j_size),
    )
    uav = xarray.DataArray(
        data=numpy.sin(xx) + numpy.sin(yy),
        dims=wet_mask.data_vars["face_mask"].dims,
        name="uav",
        attrs={"units": "metres per second", "long_name": "Eastward component of current"}
    )
    vav = xarray.DataArray(
        data=numpy.cos(xx) + numpy.cos(yy),
        dims=wet_mask.data_vars["face_mask"].dims,
        name="vav",
        attrs={"units": "metres per second", "long_name": "Eastward component of current"}
    )
    flag = xarray.DataArray(
        data=numpy.random.randint(0, 256, (time_size, k_size, j_size + 1, i_size + 1)),
        dims=["record", "k_centre", *wet_mask.data_vars["node_mask"].dims],
        attrs={"long_name": "SHOC masking flags"},
    )

    dataset = xarray.Dataset(
        data_vars={
            **layers.standard_vars,
            **grid.standard_vars,
            "botz": botz,
            "t": t,
            "eta": eta,
            "temp": temp,
            "u1": u1,
            "u2": u2,
            "uav": uav,
            "vav": vav,
            "flag": flag,
        },
        attrs={
            "title": "Example SHOC dataset",
            "ems_version": "v1.2.3 fake",
            "Conventions": "CMR/Timeseries/SHOC",
            "nce1": j_size,
            "nce2": i_size,
            "nfe1": j_size + 1,
            "nfe2": i_size + 1,
            "gridtype": "NUMERICAL",
        },
    )
    dataset.encoding["unlimited_dims"] = {"record"}
    return dataset


def test_make_dataset():
    dataset = make_dataset(j_size=5, i_size=9, corner_size=2)

    # Check that this is recognised as a ShocStandard dataset
    assert get_dataset_convention(dataset) is ShocStandard

    # Check that the correct convention is used
    assert isinstance(dataset.ems, ShocStandard)

    # Check the coordinate generation worked.
    x_centre = dataset["x_centre"]
    y_centre = dataset["y_centre"]
    assert x_centre[0, 0] == pytest.approx(0.1)
    assert x_centre[0, 1] == pytest.approx(0.2)
    assert x_centre[0, 8] == pytest.approx(0.9)
    assert x_centre[1, 0] == pytest.approx(0.2)
    assert y_centre[0, 0] == pytest.approx(0.9)
    assert y_centre[0, 1] == pytest.approx(0.8)
    assert y_centre[0, 8] == pytest.approx(0.1)
    assert y_centre[1, 0] == pytest.approx(1.0)

    # Check the coordinate generation worked.
    x_grid = dataset["x_grid"]
    y_grid = dataset["y_grid"]
    assert x_grid[0, 0] == pytest.approx(0.0)
    assert x_grid[0, 1] == pytest.approx(0.1)
    assert x_grid[0, 9] == pytest.approx(0.9)
    assert x_grid[1, 0] == pytest.approx(0.1)
    assert y_grid[0, 0] == pytest.approx(0.9)
    assert y_grid[0, 1] == pytest.approx(0.8)
    assert y_grid[0, 8] == pytest.approx(0.1)
    assert y_grid[0, 9] == pytest.approx(0.0)
    assert y_grid[1, 0] == pytest.approx(1.0)


def test_varnames():
    dataset = make_dataset(j_size=10, i_size=10)
    assert dataset.ems.depth_coordinate.name == 'z_centre'
    assert {c.name for c in dataset.ems.depth_coordinates} == {'z_centre', 'z_grid'}
    assert dataset.ems.time_coordinate.name == 't'


def test_grids():
    dataset = make_dataset(j_size=10, i_size=20)
    grids = dataset.ems.grids
    assert grids.keys() == {
        ArakawaCGridKind.face,
        ArakawaCGridKind.back,
        ArakawaCGridKind.left,
        ArakawaCGridKind.node,
    }


def test_get_grid() -> None:
    dataset = make_dataset(j_size=10, i_size=20)
    convention: ShocStandard = dataset.ems

    assert convention.get_grid(dataset['temp']) is convention.grids['face']
    assert convention.get_grid(dataset['u1']) is convention.grids['left']
    assert convention.get_grid(dataset['u2']) is convention.grids['back']
    assert convention.get_grid(dataset['x_grid']) is convention.grids['node']


def test_face_grid() -> None:
    j_size, i_size = 10, 20
    dataset = make_dataset(j_size=j_size, i_size=i_size)
    face_grid: DimensionGrid = dataset.ems.grids['face']
    assert face_grid.shape == (j_size, i_size)
    assert face_grid.size == j_size * i_size
    assert isinstance(face_grid.geometry, numpy.ndarray)
    assert face_grid.geometry_type is shapely.Polygon

    polygons = face_grid.geometry
    polygon_grid = face_grid.wind(xarray.DataArray(polygons)).values

    # Check some specific polygons
    actual = polygon_grid[1, 1]
    expected = orient(shapely.Polygon([
        (0.2, 2.0), (0.3, 2.1), (0.4, 2.0), (0.3, 1.9), (0.2, 2.0),
    ]))
    assert actual.equals_exact(expected, 1e-6)

    actual = polygon_grid[-2, -6]
    expected = orient(shapely.Polygon([(2.2, 1.4), (2.3, 1.5), (2.4, 1.4), (2.3, 1.3), (2.2, 1.4)]))
    assert actual.equals_exact(expected, 1e-6)


def test_left_grid() -> None:
    j_size, i_size = 10, 20
    dataset = make_dataset(j_size=j_size, i_size=i_size)
    left_grid: DimensionGrid = dataset.ems.grids['left']
    assert left_grid.shape == (j_size, i_size + 1)
    assert left_grid.size == j_size * (i_size + 1)
    assert isinstance(left_grid.geometry, numpy.ndarray)
    assert left_grid.geometry_type is shapely.LineString

    left_edges = left_grid.geometry
    left_edges_grid = left_grid.wind(xarray.DataArray(left_edges)).values

    # Check some specific polygons
    actual = left_edges_grid[1, 1]
    expected = shapely.LineString([(0.2, 2.0), (0.3, 2.1)])
    assert actual.equals_exact(expected, 1e-6)

    actual = left_edges_grid[-2, -6]
    expected = shapely.LineString([(2.3, 1.3), (2.4, 1.4)])
    assert actual.equals_exact(expected, 1e-6)


def test_back_grid() -> None:
    j_size, i_size = 11, 7
    dataset = make_dataset(j_size=j_size, i_size=i_size)
    back_grid: DimensionGrid = dataset.ems.grids['back']
    assert back_grid.shape == (j_size + 1, i_size)
    assert back_grid.size == (j_size + 1) * i_size
    assert isinstance(back_grid.geometry, numpy.ndarray)
    assert back_grid.geometry_type is shapely.LineString

    back_edges = back_grid.geometry
    back_edges_grid = back_grid.wind(xarray.DataArray(back_edges)).values

    # Check some specific polygons
    actual = back_edges_grid[1, 1]
    expected = shapely.LineString([(0.2, 0.7), (0.3, 0.6)])
    assert actual.equals_exact(expected, 1e-6)

    actual = back_edges_grid[-2, -6]
    expected = shapely.LineString([(1.1, 1.6), (1.2, 1.5)])
    assert actual.equals_exact(expected, 1e-6)


def test_node_grid() -> None:
    j_size, i_size = 11, 7
    dataset = make_dataset(j_size=j_size, i_size=i_size)
    node_grid: DimensionGrid = dataset.ems.grids['node']
    assert node_grid.shape == (j_size + 1, i_size + 1)
    assert node_grid.size == (j_size + 1) * (i_size + 1)
    assert isinstance(node_grid.geometry, numpy.ndarray)
    assert node_grid.geometry_type is shapely.Point


def test_face_centres():
    # SHOC standard face centres are taken directly from the coordinates,
    # not calculated from polygon centres.
    dataset = make_dataset(j_size=10, i_size=20, corner_size=3)
    convention: ShocStandard = dataset.ems

    face_centres = convention.default_grid.centroid
    lons = dataset['x_centre'].values
    lats = dataset['y_centre'].values
    for j in range(dataset.sizes['j_centre']):
        for i in range(dataset.sizes['i_centre']):
            lon = lons[j, i]
            lat = lats[j, i]
            linear_index = convention.ravel_index((ArakawaCGridKind.face, j, i))
            point = face_centres[linear_index]
            if point is None:
                assert numpy.isnan(lon)
                assert numpy.isnan(lat)
            else:
                numpy.testing.assert_equal([point.x, point.y], [lon, lat])


def test_make_geojson_geometry():
    dataset = make_dataset(j_size=10, i_size=10, corner_size=3)
    out = json.dumps(geometry.to_geojson(dataset))
    assert isinstance(out, str)


def test_ravel() -> None:
    dataset = make_dataset(j_size=5, i_size=7)
    convention: ShocStandard = dataset.ems

    for ravelled, (j, i) in enumerate(itertools.product(range(5), range(7))):
        index = (ArakawaCGridKind.face, j, i)
        assert convention.ravel_index(index) == ravelled
        assert convention.wind_index(ravelled) == index
        assert convention.wind_index(ravelled, grid_kind=ArakawaCGridKind.face) == index


def test_ravel_left():
    dataset = make_dataset(j_size=5, i_size=7)
    convention: ShocStandard = dataset.ems

    for ravelled, (j, i) in enumerate(itertools.product(range(5), range(8))):
        index = (ArakawaCGridKind.left, j, i)
        assert convention.ravel_index(index) == ravelled
        assert convention.wind_index(ravelled, grid_kind=ArakawaCGridKind.left) == index


def test_ravel_back():
    dataset = make_dataset(j_size=5, i_size=7)
    convention: ShocStandard = dataset.ems

    for ravelled, (j, i) in enumerate(itertools.product(range(6), range(7))):
        index = (ArakawaCGridKind.back, j, i)
        assert convention.ravel_index(index) == ravelled
        assert convention.wind_index(ravelled, grid_kind=ArakawaCGridKind.back) == index


def test_ravel_grid():
    dataset = make_dataset(j_size=5, i_size=7)
    convention: ShocStandard = dataset.ems

    for ravelled, (j, i) in enumerate(itertools.product(range(6), range(8))):
        index = (ArakawaCGridKind.node, j, i)
        assert convention.ravel_index(index) == ravelled
        assert convention.wind_index(ravelled, grid_kind=ArakawaCGridKind.node) == index


def test_grid_kinds():
    dataset = make_dataset(j_size=3, i_size=5)
    convention: ShocStandard = dataset.ems

    assert convention.grid_kinds == frozenset({
        ArakawaCGridKind.face,
        ArakawaCGridKind.left,
        ArakawaCGridKind.back,
        ArakawaCGridKind.node,
    })

    assert convention.default_grid_kind == ArakawaCGridKind.face


@pytest.mark.parametrize(
    ['index', 'selector'],
    (
        [(ArakawaCGridKind.face, 3, 4), xarray.Dataset({
            'j_centre': ((), 3),
            'i_centre': ((), 4),
        })],
        [(ArakawaCGridKind.left, 5, 6), xarray.Dataset({
            'j_left': ((), 5),
            'i_left': ((), 6),
        })],
        [(ArakawaCGridKind.back, 7, 8), xarray.Dataset({
            'j_back': ((), 7),
            'i_back': ((), 8),
        })],
        [(ArakawaCGridKind.node, 9, 10), xarray.Dataset({
            'j_node': ((), 9),
            'i_node': ((), 10),
        })],
    ),
)
def test_selector_for_index(index: tuple[ArakawaCGridKind, int, int], selector: dict):
    dataset = make_dataset(j_size=5, i_size=7)
    convention: ShocStandard = dataset.ems
    assert selector == convention.selector_for_index(index)


# These select_index tests are not specifically about SHOC,
# they are more about how select_index behaves with multiple grid kinds.
def test_select_index_face() -> None:
    dataset = make_dataset(time_size=4, k_size=5, j_size=5, i_size=9)
    convention: ShocStandard = dataset.ems
    face = convention.select_index((ArakawaCGridKind.face, 3, 4))

    assert set(face.variables.keys()) == {
        'botz', 'eta', 'temp', 'uav', 'vav',
    }
    assert face.sizes == {'record': 4, 'k_centre': 5}


def test_select_index_face_keep_geometry() -> None:
    dataset = make_dataset(time_size=4, k_size=5, j_size=5, i_size=9)
    convention: ShocStandard = dataset.ems
    face = convention.select_index((ArakawaCGridKind.face, 3, 4), drop_geometry=False)

    assert set(face.variables.keys()) == {
        'botz', 'eta', 'temp', 'uav', 'vav',
        'x_centre', 'y_centre',
    }
    assert face.sizes == {'record': 4, 'k_centre': 5}
    assert face['x_centre'].values == dataset['x_centre'].values[3, 4]
    assert face['y_centre'].values == dataset['y_centre'].values[3, 4]


def test_select_index_edge() -> None:
    dataset = make_dataset(time_size=4, k_size=5, j_size=5, i_size=9)
    convention: ShocStandard = dataset.ems

    left = convention.select_index((ArakawaCGridKind.left, 3, 4))
    assert set(left.data_vars.keys()) == {
        # This is the only data variable we expect to see,
        # as it is the only one defined on left edges.
        'u1',
    }
    assert left.sizes == {'record': 4, 'k_centre': 5}

    back = convention.select_index((ArakawaCGridKind.back, 3, 4))
    assert set(back.data_vars.keys()) == {
        # This is the only data variable we expect to see,
        # as it is the only one defined on back edges.
        'u2',
    }
    assert back.sizes == {'record': 4, 'k_centre': 5}


def test_select_index_edge_keep_geometry() -> None:
    dataset = make_dataset(time_size=4, k_size=5, j_size=5, i_size=9)
    convention: ShocStandard = dataset.ems

    left = convention.select_index((ArakawaCGridKind.left, 3, 4), drop_geometry=False)
    assert set(left.variables.keys()) == {
        'u1', 'x_left', 'y_left'
    }
    assert left.sizes == {'record': 4, 'k_centre': 5}

    back = convention.select_index((ArakawaCGridKind.back, 3, 4), drop_geometry=False)
    assert set(back.variables.keys()) == {
        'u2', 'x_back', 'y_back'
    }
    assert back.sizes == {'record': 4, 'k_centre': 5}


def test_select_index_grid() -> None:
    dataset = make_dataset(time_size=4, k_size=5, j_size=5, i_size=9)
    convention: ShocStandard = dataset.ems

    node = convention.select_index((ArakawaCGridKind.node, 3, 4))
    assert set(node.data_vars.keys()) == {
        'flag',
    }
    assert node.sizes == {'record': 4, 'k_centre': 5}


def test_select_index_grid_keep_geometry() -> None:
    dataset = make_dataset(time_size=4, k_size=5, j_size=5, i_size=9)
    convention: ShocStandard = dataset.ems

    node = convention.select_index((ArakawaCGridKind.node, 3, 4), drop_geometry=False)
    assert set(node.data_vars.keys()) == {
        'flag', 'x_grid', 'y_grid'
    }
    assert node.sizes == {'record': 4, 'k_centre': 5}


def test_drop_geometry(datasets: pathlib.Path):
    dataset = xarray.open_dataset(datasets / 'shoc_standard.nc')

    dropped = dataset.ems.drop_geometry()
    assert set(dropped.dims) == {'face_i', 'face_j'}
    for topology in [dataset.ems.face, dataset.ems.back, dataset.ems.left, dataset.ems.node]:
        assert topology.longitude_name in dataset.variables
        assert topology.longitude_name in dataset.variables
        assert topology.longitude_name not in dropped.variables
        assert topology.longitude_name not in dropped.variables


def test_values():
    dataset = make_dataset(j_size=10, i_size=20, corner_size=5)
    eta = dataset.data_vars["eta"].isel(record=0)
    grid = dataset.ems.get_grid(eta)
    values = grid.ravel(eta)

    # There should be one value per cell polygon
    assert grid.size == len(values)

    # The values should be in a specific order
    assert numpy.allclose(values, eta.values.ravel(), equal_nan=True)


@pytest.mark.matplotlib
def test_make_artist_face_scalar(tmp_path: pathlib.Path) -> None:
    dataset = make_dataset(j_size=10, i_size=20, corner_size=5)
    surface_temp = dataset.data_vars["temp"].isel(record=0, k_centre=-1)

    figure = Figure()
    axes = figure.add_subplot(projection=dataset.ems.data_crs)
    artist = dataset.ems.make_artist(axes, surface_temp, cmap='Oranges')
    axes.autoscale()

    # Check the right kind of artist was made
    assert isinstance(artist, plot.artists.PolygonScalarCollection)
    # It should have made a colorbar also
    assert len(figure.axes) == 2
    # The artist should have been added to the axes
    assert artist in axes.get_children()
    # kwargs should be passed through to the artist
    assert artist.get_cmap().name == 'Oranges'

    figure.savefig(tmp_path / 'face_scalar.png')


@pytest.mark.matplotlib
def test_make_artist_face_vector(tmp_path: pathlib.Path) -> None:
    dataset = make_dataset(j_size=10, i_size=20, corner_size=5)

    figure = Figure()
    axes = figure.add_subplot(projection=dataset.ems.data_crs)
    artist = dataset.ems.make_artist(
        axes, ('uav', 'vav'),
        scale=20)
    axes.autoscale()

    # Check the right kind of artist was made
    assert isinstance(artist, plot.artists.PolygonVectorQuiver)
    # Only one axes this time, vectors don't get a colorbar
    assert len(figure.axes) == 1
    # The artist should have been added to the axes
    assert artist in axes.get_children()
    # kwargs should be passed through to the artist
    assert artist.scale == 20

    figure.savefig(tmp_path / 'face_vector.png')


@pytest.mark.matplotlib
def test_make_artist_node_scalar(tmp_path: pathlib.Path) -> None:
    dataset = make_dataset(j_size=10, i_size=20, corner_size=5)
    latest_surface = dataset.isel(record=0, k_centre=-1)

    figure = Figure()
    axes = figure.add_subplot(projection=dataset.ems.data_crs)
    artist = dataset.ems.make_artist(
        axes, latest_surface['flag'],
        cmap='Blues')
    axes.autoscale()

    # Check the right kind of artist was made
    assert isinstance(artist, plot.artists.NodeTriMesh)
    # It should have made a colorbar also
    assert len(figure.axes) == 2
    # The artist should have been added to the axes
    assert artist in axes.get_children()
    # kwargs should be passed through to the artist
    assert artist.get_cmap().name == 'Blues'

    figure.savefig(tmp_path / 'node.png')


@pytest.mark.matplotlib(mock_coast=True)
def test_plot_geometry(tmp_path: pathlib.Path) -> None:
    # Not much to test here, mostly that it doesn't throw an error
    dataset = make_dataset(j_size=10, i_size=20, corner_size=5)

    figure = Figure()
    dataset.ems.plot_on_figure(figure)

    assert len(figure.axes) == 1

    figure.savefig(tmp_path / 'geometry.png')


def test_make_clip_mask():
    dataset = make_dataset(j_size=10, i_size=8)
    convention: ShocStandard = dataset.ems

    # The dataset will have cells with centres from 0-.5 longitude, 0-.7 latitude
    clip_geometry = shapely.Polygon([
        (.74, .84), (.86, .84), (.86, .96), (.74, .96), (.74, .84),
    ])

    mask = convention.make_clip_mask(clip_geometry)

    assert mask.data_vars.keys() \
        == {'face_mask', 'left_mask', 'back_mask', 'node_mask'}

    assert_equal(
        mask.data_vars['face_mask'].values,
        mask_from_strings([
            "00000000",
            "00000000",
            "00000000",
            "00010000",
            "00111000",
            "00010000",
            "00000000",
            "00000000",
            "00000000",
            "00000000",
        ])
    )
    assert_equal(
        mask.data_vars['left_mask'].values,
        mask_from_strings([
            "000000000",
            "000000000",
            "000000000",
            "000110000",
            "001111000",
            "000110000",
            "000000000",
            "000000000",
            "000000000",
            "000000000",
        ]),
    )
    assert_equal(
        mask.data_vars['back_mask'].values,
        mask_from_strings([
            "00000000",
            "00000000",
            "00000000",
            "00010000",
            "00111000",
            "00111000",
            "00010000",
            "00000000",
            "00000000",
            "00000000",
            "00000000",
        ]),
    )
    assert_equal(
        mask.data_vars['node_mask'].values,
        mask_from_strings([
            "000000000",
            "000000000",
            "000000000",
            "000110000",
            "001111000",
            "001111000",
            "000110000",
            "000000000",
            "000000000",
            "000000000",
            "000000000",
        ]),
    )

    # Test adding a buffer also
    mask = convention.make_clip_mask(clip_geometry, buffer=1)
    assert_equal(
        mask.data_vars['face_mask'].values,
        mask_from_strings([
            "00000000",
            "00000000",
            "00111000",
            "01111100",
            "01111100",
            "01111100",
            "00111000",
            "00000000",
            "00000000",
            "00000000",
        ]),
    )
    assert_equal(
        mask.data_vars['node_mask'].values,
        mask_from_strings([
            "000000000",
            "000000000",
            "001111000",
            "011111100",
            "011111100",
            "011111100",
            "011111100",
            "001111000",
            "000000000",
            "000000000",
            "000000000",
        ]),
    )


def test_apply_clip_mask(tmp_path):
    dataset = make_dataset(j_size=10, i_size=8)
    convention: ShocStandard = dataset.ems

    # Clip it!
    clip_geometry = shapely.Polygon([
        (.74, .84), (.86, .84), (.86, .96), (.74, .96), (.74, .84),
    ])
    mask = convention.make_clip_mask(clip_geometry)
    clipped = dataset.ems.apply_clip_mask(mask, tmp_path)

    assert isinstance(clipped.ems, ShocStandard)

    # Check that the variable and dimension keys were preserved
    assert set(dataset.data_vars.keys()) == set(clipped.data_vars.keys())
    assert set(dataset.coords.keys()) == set(clipped.coords.keys())
    assert set(dataset.dims) == set(clipped.dims)

    # Check that the new topology seems reasonable
    assert clipped.ems.face.longitude.shape == (3, 3)
    assert clipped.ems.face.latitude.shape == (3, 3)
    assert clipped.ems.node.longitude.shape == (4, 4)
    assert clipped.ems.node.latitude.shape == (4, 4)

    # Check that the data were preserved, beyond being clipped
    def clip_values(values: numpy.ndarray) -> numpy.ndarray:
        values = values[..., 3:6, 2:5].copy()
        values[..., 0, 0] = numpy.nan
        values[..., 0, -1] = numpy.nan
        values[..., -1, -1] = numpy.nan
        values[..., -1, 0] = numpy.nan
        return values

    assert_equal(clipped.data_vars['botz'].values, clip_values(dataset.data_vars['botz'].values))
    assert_equal(clipped.data_vars['eta'].values, clip_values(dataset.data_vars['eta'].values))
    assert_equal(clipped.data_vars['temp'].values, clip_values(dataset.data_vars['temp'].values))

    # Check that the new geometry matches the relevant polygons in the old geometry
    face_grid = convention.grids['face']
    original_polygons = face_grid.geometry.reshape(10, 8)[3:6, 2:5].ravel()

    clipped_face_grid = clipped.ems.grids['face']
    clipped_polygons = clipped_face_grid.geometry
    assert clipped_face_grid.size == 3 * 3
    assert clipped_polygons[0] is None
    assert clipped_polygons[1].equals_exact(original_polygons[1], 1e-6)
    assert clipped_polygons[2] is None
    assert clipped_polygons[3].equals_exact(original_polygons[3], 1e-6)
    assert clipped_polygons[4].equals_exact(original_polygons[4], 1e-6)
    assert clipped_polygons[5].equals_exact(original_polygons[5], 1e-6)
    assert clipped_polygons[6] is None
    assert clipped_polygons[7].equals_exact(original_polygons[7], 1e-6)
    assert clipped_polygons[8] is None


@pytest.mark.memory_usage
def test_make_polygons_memory_usage():
    j_size, i_size = 1000, 2000
    dataset = make_dataset(j_size=j_size, i_size=i_size)
    face_grid = dataset.ems.grids['face']

    with track_peak_memory_usage() as tracker:
        assert len(face_grid.geometry) == j_size * i_size

    logger.info("current memory usage: %d, peak memory usage: %d", tracker.current, tracker.peak)

    target = 134_500_000
    assert tracker.peak < target, "Peak memory allocation is too large"
    assert tracker.peak > target * 0.9, "Peak memory allocation is suspiciously small - did you improve things?"