From c73027bbec1561645696a1d0807ca218cf1a8adb Mon Sep 17 00:00:00 2001 From: Daniel Shapero Date: Sat, 30 Sep 2023 18:06:53 -0700 Subject: [PATCH] Start Helheim demo --- demos/helheim/data.py | 38 +++++ demos/helheim/helheim.ipynb | 268 ++++++++++++++++++++++++++++++++++++ 2 files changed, 306 insertions(+) create mode 100644 demos/helheim/data.py create mode 100644 demos/helheim/helheim.ipynb diff --git a/demos/helheim/data.py b/demos/helheim/data.py new file mode 100644 index 0000000..9fe5d11 --- /dev/null +++ b/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 + ] diff --git a/demos/helheim/helheim.ipynb b/demos/helheim/helheim.ipynb new file mode 100644 index 0000000..7254f97 --- /dev/null +++ b/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 +}