Add support for cropping domain with padded lat/lon convex hull

.github/workflows/python-package-pip.yml

@@ -29,7 +29,7 @@ jobs:
 
       - name: Install package with pip
         run: |
-          python -m pip install . "zarr${{ matrix.zarr-version }}"
+          python -m pip install ".[latlon-domain-crop]" "zarr${{ matrix.zarr-version }}"
           python -m pip install pytest
 
       - name: Run tests (non-distributed)

CHANGELOG.md

@@ -5,6 +5,14 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [Unreleased](https://github.com/mllam/mllam-data-prep)
+
+[All changes](https://github.com/mllam/mllam-data-prep/compare/HEAD...v0.6.1)
+
+### Added
+
+- add support for cropping a dataset using the convex hull of the lat/lon coordinates of another dataset (can be used for creating boundary data in Limited Area Modelling setups) [\#45](https://github.com/mllam/mllam-data-prep/pull/45), @leifdenby
+
 ## [v0.6.1](https://github.com/mllam/mllam-data-prep/release/tag/v0.6.1)
 
 [All changes](https://github.com/mllam/mllam-data-prep/compare/v0.6.1...v0.6.0)

README.md

@@ -285,6 +285,27 @@ The `output` section defines three things:
 3. `chunking`: the chunk sizes to use when writing the training dataset to zarr. This is optional, but can be used to optimise the performance of the zarr dataset. By default the chunk sizes are set to the size of the dimension, but this can be overridden by setting the chunk size in the configuration file. A common choice is to set the dimension along which you are batching to align with the of each training item (e.g. if you are training a model with time-step roll-out of 10 timesteps, you might choose a chunksize of 10 along the time dimension).
 4. Splitting and calculation of statistics of the output variables, using the `splitting` section. The `output.splitting.splits` attribute defines the individual splits to create (for example `train`, `val` and `test`) and `output.splitting.dim` defines the dimension to split along. The `compute_statistics` can be optionally set for a given split to calculate the statistical properties requested (for example `mean`, `std`) any method available on `xarray.Dataset.{op}` can be used. In addition methods prefixed by `diff_` (so the operational would be listed as `diff_{op}`) to compute a statistic based on difference of consecutive time-steps, e.g. `diff_mean` to compute the `mean` of the difference between consecutive timesteps (these are used for normalisating increments). The `dims` attribute defines the dimensions to calculate the statistics over (for example `grid_index` and `time`).
 
+In addition the `output` section can also contain a configuration for cropping the output dataset using the convex hull of coordinates from a different dataset. This is used for example when creating training datasets from limited area modelling (LAM) setups, where a separate dataset is used for the boundary data. The example above doesn't include this section, but [see below](#cropping-the-output-dataset-using-convex-hull-of-another-dataset) for an example of how to use this feature.
+
+#### Cropping the output dataset using convex hull of another dataset
+
+When creating training datasets for limited area models (LAMs) it is often useful to crop the training dataset to the convex hull of the coordinates of another dataset (for example a dataset containing boundary data). This can be done by adding a `domain_cropping` section to the `output` section of the configuration file. A full example where ERA5 is cropped is given in [example.danra.yaml](example.danra.yaml). The relevant section is reproduced here for completeness:
+
+```yaml
+output:
+  ...
+  domain_cropping:
+    margin_width_degrees: 10
+    interior_dataset_config_path: example.danra.yaml
+```
+
+The `domain_cropping` section has two required attributes, and one optional attribute:
+- `margin_width_degrees`: the width (in degrees) of the margin to add around the convex hull of the coordinates of the interior dataset. This allows you to control how much extra area to include around the convex hull of the interior dataset (e.g. how thick to make the boundary region for LAM boundary forcing datasets).
+- `interior_dataset_config_path`: the path to the configuration file of the interior dataset. This is used to load the interior dataset and extract the coordinates to calculate the convex hull from.
+- `include_interior_points`: optional, defaults to `false`. If set to `true` the points inside the convex hull of the interior dataset will also be included in the cropped output dataset. If set to `false` only the margin region around the convex hull will be included. This can be useful if you not only want to create a boundary forcing dataset from a global simulation for LAM, but also want to include the interior points of the global simulation that overlap with the LAM domain.
+
+Details on how the convex hull cropping is actually done can be found in the jupyter notebook in [docs/domain-cropping.ipynb](docs/domain-cropping.ipynb).
+
 ### The `inputs` section
 
 ```yaml

docs/domain-cropping.ipynb

@@ -0,0 +1,263 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Cropping a domain using the lat/lon convex hull of another domain\n",
+    "\n",
+    "This notebook demonstrates how convex hull based cropping is implemented in `mllam-data-prep`. To actually use this feature see [the details in the README](../README.md#cropping-a-domain-using-the-latlon-convex-hull-of-another-domain).\n",
+    "\n",
+    "This method crops a domain using the lat/lon convex hull of another domain. This is useful when a) you have two datasets on two different overlapping domains where you only want to keep the overlapping part, or b) in the case where you want to run limited-area simulations and need to create a dataset that provides the boundary conditions for the limited-area domain."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import tests.data as testdata\n",
+    "import tempfile\n",
+    "\n",
+    "import xarray as xr\n",
+    "import numpy as np\n",
+    "import matplotlib.pyplot as plt\n",
+    "import cartopy.crs as ccrs\n",
+    "\n",
+    "import mllam_data_prep as mdp\n",
+    "import tests.data as testdata\n",
+    "from mllam_data_prep.ops import cropping    "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "We will start by creating the some synthetic data on two different domains where one sits within the other, mimicking the scenario where you have a high-resolution dataset within a coarser dataset."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tmpdir = tempfile.TemporaryDirectory()\n",
+    "l = 500 * 1.0e3  # length and width of domain in meters\n",
+    "N = 50  # number of grid points in each direction\n",
+    "config_lam = testdata.create_input_datasets_and_config(\n",
+    "    identifier=\"lam\",\n",
+    "    data_categories=[\"state\"],\n",
+    "    tmpdir=tmpdir,\n",
+    "    xlim=[-l/2.0, l/2.0],\n",
+    "    ylim=[-l/2.0, l/2.0],\n",
+    "    nx=N, ny=N,\n",
+    "    add_latlon=True\n",
+    "\n",
+    ")\n",
+    "# make the global domain twice as large as the LAM domain so that the lam\n",
+    "# domain is contained within the global domain, but half the number of grid\n",
+    "# points in each direction\n",
+    "config_global = testdata.create_input_datasets_and_config(\n",
+    "    identifier=\"global\",\n",
+    "    data_categories=[\"state\"],\n",
+    "    tmpdir=tmpdir,\n",
+    "    xlim=[-l, l],\n",
+    "    ylim=[-l, l],\n",
+    "    nx=N//2, ny=N//2,\n",
+    "    add_latlon=True\n",
+    ")\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ds_lam = mdp.create_dataset(config=config_lam)\n",
+    "ds_global = mdp.create_dataset(config=config_global)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ds_lam"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ds_global"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Let's first make a mask that indicates which of the coarser domain grid-points are within the convex hull of the high-resolution grid points coordinates"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "da_interior_mask = cropping.create_convex_hull_mask(ds=ds_global, ds_reference=ds_lam)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "da_interior_mask = da_interior_mask.where(da_interior_mask, drop=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ds_reference = ds_lam\n",
+    "ds = ds_global\n",
+    "\n",
+    "fig, ax = plt.subplots(subplot_kw=dict(projection=ccrs.PlateCarree()))\n",
+    "xr.plot.scatter(ax=ax, ds=ds, x=\"lon\", y=\"lat\", label=\"ds\", transform=ccrs.PlateCarree())\n",
+    "xr.plot.scatter(ax=ax, ds=ds_reference, x=\"lon\", y=\"lat\", label=\"ds_ref\", transform=ccrs.PlateCarree())\n",
+    "xr.plot.scatter(ax=ax, c=da_interior_mask, ds=da_interior_mask.to_dataset(name=\"mask\"), x=\"lon\", y=\"lat\", marker=\"x\", label=\"ds (interior)\")\n",
+    "ax.coastlines()\n",
+    "ax.gridlines(draw_labels=True)\n",
+    "ax.legend()\n",
+    "plt.show()\n",
+    "    "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Next we will for the points from the larger domain and exterior to this convex hull (i.e. points excluded by the mask created above) compute the distance to the nearest point of the inner domain."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "da_dist = cropping.distance_to_convex_hull_boundary(ds=ds_global, ds_reference=ds_lam)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, ax = plt.subplots(subplot_kw=dict(projection=ccrs.PlateCarree()), figsize=(10, 4))\n",
+    "xr.plot.scatter(ax=ax, ds=ds_reference, x=\"lon\", y=\"lat\", label=\"ds_ref\", marker=\".\", transform=ccrs.PlateCarree())\n",
+    "xr.plot.scatter(ax=ax, hue=\"dist\", ds=da_dist.to_dataset(name=\"dist\"), x=\"lon\", y=\"lat\", marker=\"x\", label=\"ds (distance)\", add_colorbar=True)\n",
+    "ax.coastlines()\n",
+    "ax.gridlines(draw_labels=[\"top\", \"left\"])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "And finally we will use this distance to create a plot that includes a margin around the convex hull of the inner domain."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "max_dist = 1.5 # in degrees\n",
+    "da_global_margin_crop = cropping.crop_with_convex_hull(\n",
+    "    ds=ds_global, ds_reference=ds_lam, margin_thickness=max_dist\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, ax = plt.subplots(subplot_kw=dict(projection=ccrs.PlateCarree()), figsize=(10, 4))\n",
+    "xr.plot.scatter(ax=ax, ds=ds_reference, x=\"lon\", y=\"lat\", label=\"ds_ref\", marker=\".\", transform=ccrs.PlateCarree())\n",
+    "xr.plot.scatter(ax=ax, ds=da_global_margin_crop, x=\"lon\", y=\"lat\", marker=\"x\", label=f\"ds (margin crop, {max_dist}deg)\", transform=ccrs.PlateCarree())\n",
+    "ax.coastlines()\n",
+    "ax.legend()\n",
+    "ax.gridlines(draw_labels=[\"top\", \"left\"])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "max_dist = 1.5 # in degrees\n",
+    "ds_global_margin_crop_without_interior = cropping.crop_with_convex_hull(\n",
+    "    ds=ds_global, ds_reference=ds_lam, margin_thickness=max_dist, include_interior_points=False\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, ax = plt.subplots(subplot_kw=dict(projection=ccrs.PlateCarree()), figsize=(10, 4))\n",
+    "xr.plot.scatter(ax=ax, ds=ds_reference, x=\"lon\", y=\"lat\", label=\"ds_ref\", marker=\".\", transform=ccrs.PlateCarree())\n",
+    "xr.plot.scatter(ax=ax, ds=ds_global_margin_crop_without_interior, x=\"lon\", y=\"lat\", marker=\"x\", label=f\"ds (margin crop, {max_dist}deg)\", transform=ccrs.PlateCarree())\n",
+    "ax.coastlines()\n",
+    "ax.legend()\n",
+    "ax.gridlines(draw_labels=[\"top\", \"left\"])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "mllam-data-prep",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.16"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}