Start Helheim demo

demos/helheim/data.py

@@ -0,0 +1,38 @@
+import pooch
+from icepack.datasets import EarthDataDownloader
+
+nsidc_url = "https://n5eil01u.ecs.nsidc.org/"
+earthdata_downloader = EarthDataDownloader()
+
+def fetch_bedmachine():
+    base_url = f"{nsidc_url}/ICEBRIDGE/IDBMG4.005/1993.01.01/"
+    filename = "BedMachineGreenland-v5.nc"
+    sha256 = "f7116b8e9e3840649075dcceb796ce98aaeeb5d279d15db489e6e7668e0d80db"
+    bedmachine = pooch.create(
+        path=pooch.os_cache("icepack"),
+        base_url=base_url,
+        registry={filename: sha256},
+    )
+    return bedmachine.fetch(filename, downloader=earthdata_downloader)
+
+
+def fetch_measures_velocity():
+    base_url = f"{nsidc_url}/MEASURES/NSIDC-0670.001/1995.12.01/"
+    filenames = [
+        f"greenland_vel_mosaic250_{key}_v1.tif" for key in ["vx", "vy", "ex", "ey"]
+    ]
+    sha256s = [
+        "e903891d5ed2c5faaccb60705088917bf1595cbd516223e5cea208b55979d68d",
+        "16b2c1cbd7be2ee3a50219eb2e00604cfeca5fca059e48f02c34570e15f1d8c9",
+        "8cc265a8478c2d429cf94606f3386f102fd1da7ebb4c99d85a422663ed0a5a33",
+        "295b9d4a07bbba65104004adb160f7df9b31eabeb935de9e6b843d39dd88b272",
+    ]
+    measures = pooch.create(
+        path=pooch.os_cache("icepack"),
+        base_url=base_url,
+        registry={filename: sha256 for filename, sha256 in zip(filenames, sha256s)}
+    )
+    return [
+        measures.fetch(filename, downloader=earthdata_downloader)
+        for filename in filenames
+    ]

demos/helheim/helheim.ipynb

@@ -0,0 +1,268 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "dacf0e38-7e0b-4001-b872-8ba34bc80582",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import numpy as np\n",
+    "import matplotlib.pyplot as plt\n",
+    "import geojson\n",
+    "import rasterio\n",
+    "import xarray\n",
+    "import firedrake\n",
+    "from firedrake import assemble, Constant, inner, grad, dx\n",
+    "import icepack\n",
+    "import data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "94715d1b-5027-49c2-b124-08e14c1caea9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "outline_filename = icepack.datasets.fetch_outline(\"helheim\")\n",
+    "with open(outline_filename, \"r\") as outline_file:\n",
+    "    outline = geojson.load(outline_file)\n",
+    "\n",
+    "geometry = icepack.meshing.collection_to_geo(outline)\n",
+    "with open(\"helheim.geo\", \"w\") as geometry_file:\n",
+    "    geometry_file.write(geometry.get_code())\n",
+    "\n",
+    "!gmsh -2 -v 0 -o helheim.msh helheim.geo\n",
+    "\n",
+    "mesh = firedrake.Mesh(\"helheim.msh\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "23b2551f-da54-418a-a99b-63f3d88d919f",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, ax = plt.subplots()\n",
+    "ax.set_aspect(\"equal\")\n",
+    "firedrake.triplot(mesh, axes=ax)\n",
+    "ax.legend();"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ad5c3288-c5a0-4dc3-a082-9a98fb5ce857",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "Q = firedrake.FunctionSpace(mesh, \"CG\", 1)\n",
+    "V = firedrake.VectorFunctionSpace(mesh, \"CG\", 1)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7319d41a-72fd-415e-9ef1-f189853a1de0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "bedmachine_filename = data.fetch_bedmachine()\n",
+    "dataset = xarray.open_dataset(bedmachine_filename)\n",
+    "b = icepack.interpolate(dataset[\"bed\"], Q)\n",
+    "h = icepack.interpolate(dataset[\"thickness\"], Q)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6d63226f-c2b8-4620-ae06-7e43ef7a421a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, ax = plt.subplots()\n",
+    "ax.set_aspect(\"equal\")\n",
+    "colors = firedrake.tripcolor(h, axes=ax)\n",
+    "fig.colorbar(colors);"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ad86fd07-5fe3-4c79-bcd2-b7e144425b2b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "coords = np.array(list(geojson.utils.coords(outline)))\n",
+    "delta = 2.5e3\n",
+    "xmin, xmax = coords[:, 0].min() - delta, coords[:, 0].max() + delta\n",
+    "ymin, ymax = coords[:, 1].min() - delta, coords[:, 1].max() + delta"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "70cbb217-6c51-4899-97a7-f81ea3fa2552",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "measures_filenames = data.fetch_measures_velocity()\n",
+    "velocity_data = {}\n",
+    "for key in [\"vx\", \"vy\", \"ex\", \"ey\"]:\n",
+    "    filename = [f for f in measures_filenames if key in f][0]\n",
+    "    with rasterio.open(filename, \"r\") as source:\n",
+    "        window = rasterio.windows.from_bounds(\n",
+    "            left=xmin, right=xmax, bottom=ymin, top=ymax, transform=source.transform\n",
+    "        ).round_lengths().round_offsets()\n",
+    "        transform = source.window_transform(window)\n",
+    "        velocity_data[key] = source.read(indexes=1, window=window)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c6c6c8a0-2380-4b80-a4a8-49ee8bdf0397",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "indices = np.array(\n",
+    "    [\n",
+    "        (i, j)\n",
+    "        for i in range(window.width)\n",
+    "        for j in range(window.height)\n",
+    "        if (\n",
+    "            mesh.locate_cell(transform * (i, j), tolerance=1e-8) and\n",
+    "            velocity_data[\"ex\"][j, i] >= 0.0\n",
+    "        )\n",
+    "    ]\n",
+    ")\n",
+    "xs = np.array([transform * idx for idx in indices])\n",
+    "point_set = firedrake.VertexOnlyMesh(\n",
+    "    mesh, xs, missing_points_behaviour=\"error\"\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5f5d2b18-d3ad-4151-903e-276eb3c81cca",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, axes = plt.subplots()\n",
+    "axes.set_aspect(\"equal\")\n",
+    "axes.set_xlim((300e3, 310e3))\n",
+    "axes.set_ylim((-2.582e6, -2.572e6))\n",
+    "firedrake.triplot(mesh, axes=axes)\n",
+    "axes.scatter(xs[:, 0], xs[:, 1], marker=\".\");"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "358ad29b-685b-4561-8743-f64907ebd78a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "Δ = firedrake.FunctionSpace(point_set, \"DG\", 0)\n",
+    "\n",
+    "u_o = firedrake.Function(Δ)\n",
+    "v_o = firedrake.Function(Δ)\n",
+    "σ_x = firedrake.Function(Δ)\n",
+    "σ_y = firedrake.Function(Δ)\n",
+    "\n",
+    "u_o.dat.data[:] = velocity_data[\"vx\"][indices[:, 1], indices[:, 0]]\n",
+    "v_o.dat.data[:] = velocity_data[\"vy\"][indices[:, 1], indices[:, 0]]\n",
+    "σ_x.dat.data[:] = velocity_data[\"ex\"][indices[:, 1], indices[:, 0]]\n",
+    "σ_y.dat.data[:] = velocity_data[\"ey\"][indices[:, 1], indices[:, 0]]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c9bb179f-31c8-4021-9138-f8f1c36a78be",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "u = firedrake.Function(V)\n",
+    "N = Constant(len(indices))\n",
+    "area = assemble(Constant(1) * dx(mesh))\n",
+    "\n",
+    "def loss_functional(u):\n",
+    "    u_int = firedrake.interpolate(u[0], Δ)\n",
+    "    v_int = firedrake.interpolate(u[1], Δ)\n",
+    "    δu = u_int - u_o\n",
+    "    δv = v_int - v_o\n",
+    "    return 0.5 / N * ((δu / σ_x)**2 + (δv / σ_y)**2) * dx\n",
+    "\n",
+    "def regularization(u):\n",
+    "    Ω = Constant(area)\n",
+    "    α = Constant(1e1)   # TODO: tune this\n",
+    "    return 0.5 * α**2 / Ω * inner(grad(u), grad(u)) * dx"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "46e8737e-5b26-4a84-921c-e248337eab92",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import firedrake.adjoint\n",
+    "firedrake.adjoint.continue_annotation()\n",
+    "problem = icepack.statistics.StatisticsProblem(\n",
+    "    simulation=lambda u: u.copy(deepcopy=True),\n",
+    "    loss_functional=loss_functional,\n",
+    "    regularization=regularization,\n",
+    "    controls=u,\n",
+    ")\n",
+    "\n",
+    "estimator = icepack.statistics.MaximumProbabilityEstimator(\n",
+    "    problem,\n",
+    "    gradient_tolerance=5e-5,\n",
+    "    step_tolerance=1e-8,\n",
+    "    max_iterations=50,\n",
+    ")\n",
+    "\n",
+    "u = estimator.solve()\n",
+    "firedrake.adjoint.pause_annotation()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6a0ab4c7-4526-4169-b2a7-0b52c1923d4a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "fig, axes = plt.subplots()\n",
+    "axes.set_aspect(\"equal\")\n",
+    "colors = firedrake.tripcolor(u, axes=axes)\n",
+    "fig.colorbar(colors);"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "firedrake",
+   "language": "python",
+   "name": "firedrake"
+  },
+  "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.11.5"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}