From 5e8172719771679d2df95a379d2f22b7d58f48af Mon Sep 17 00:00:00 2001 From: Daniel Shapero Date: Mon, 2 Oct 2023 14:36:57 -0700 Subject: [PATCH] Use preprocessing in notebook --- demos/helheim/helheim.ipynb | 262 +++++++++++++++++++++--------------- 1 file changed, 157 insertions(+), 105 deletions(-) diff --git a/demos/helheim/helheim.ipynb b/demos/helheim/helheim.ipynb index 7254f97..5aec22b 100644 --- a/demos/helheim/helheim.ipynb +++ b/demos/helheim/helheim.ipynb @@ -9,33 +9,50 @@ "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", + "from firedrake import (\n", + " assemble,\n", + " interpolate,\n", + " Constant,\n", + " inner,\n", + " max_value,\n", + " min_value,\n", + " grad,\n", + " dx,\n", + ")\n", + "import firedrake.adjoint\n", "import icepack\n", + "from icepack.constants import (\n", + " ice_density as ρ_I,\n", + " water_density as ρ_W,\n", + " gravity as g,\n", + " glen_flow_law as n,\n", + " weertman_sliding_law as m,\n", + ")\n", "import data" ] }, + { + "cell_type": "markdown", + "id": "5e4a2ce8-c69e-47af-b326-bd756d7d88ce", + "metadata": {}, + "source": [ + "Load in a mesh and some velocity data." + ] + }, { "cell_type": "code", "execution_count": null, - "id": "94715d1b-5027-49c2-b124-08e14c1caea9", + "id": "12558a95-c80b-4b78-a488-b1e11a218385", "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", + "with firedrake.CheckpointFile(\"helheim.h5\", \"r\") as chk:\n", + " mesh = chk.load_mesh()\n", + " u = chk.load_function(mesh, name=\"velocity\")\n", "\n", - "mesh = firedrake.Mesh(\"helheim.msh\")" + "V = u.function_space()" ] }, { @@ -54,12 +71,22 @@ { "cell_type": "code", "execution_count": null, - "id": "ad5c3288-c5a0-4dc3-a082-9a98fb5ce857", + "id": "d0e49917-2c6f-48df-ad23-c2f30d05a672", "metadata": {}, "outputs": [], "source": [ - "Q = firedrake.FunctionSpace(mesh, \"CG\", 1)\n", - "V = firedrake.VectorFunctionSpace(mesh, \"CG\", 1)" + "fig, ax = plt.subplots()\n", + "ax.set_aspect(\"equal\")\n", + "colors = firedrake.streamplot(u, resolution=1e3, axes=ax)\n", + "fig.colorbar(colors);" + ] + }, + { + "cell_type": "markdown", + "id": "d91ff639-05d5-4f7a-b4f6-736baeb589a4", + "metadata": {}, + "source": [ + "Load the thickness and bed data." ] }, { @@ -69,6 +96,8 @@ "metadata": {}, "outputs": [], "source": [ + "Q = firedrake.FunctionSpace(mesh, \"CG\", 1)\n", + "\n", "bedmachine_filename = data.fetch_bedmachine()\n", "dataset = xarray.open_dataset(bedmachine_filename)\n", "b = icepack.interpolate(dataset[\"bed\"], Q)\n", @@ -91,156 +120,179 @@ { "cell_type": "code", "execution_count": null, - "id": "ad86fd07-5fe3-4c79-bcd2-b7e144425b2b", + "id": "0d6bc503-e8de-4295-b8a0-d6cbf9378447", "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" + "s_obs = icepack.compute_surface(thickness=h, bed=b)\n", + "\n", + "s = s_obs.copy(deepcopy=True)\n", + "α = Constant(2e3)\n", + "J = 0.5 * ((s - s_obs) ** 2 + α**2 * inner(grad(s), grad(s))) * dx\n", + "F = firedrake.derivative(J, s)\n", + "firedrake.solve(F == 0, s)\n", + "\n", + "τ_d = firedrake.project(-ρ_I * g * h * grad(s), V)" ] }, { "cell_type": "code", "execution_count": null, - "id": "70cbb217-6c51-4899-97a7-f81ea3fa2552", + "id": "fa034371-60ed-4586-887c-bde22b4b3014", "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)" + "fig, ax = plt.subplots()\n", + "ax.set_aspect(\"equal\")\n", + "colors = firedrake.tripcolor(τ_d, axes=ax)\n", + "fig.colorbar(colors);" + ] + }, + { + "cell_type": "markdown", + "id": "0cb43788-ec7c-4a7c-91f2-95c78dfccd32", + "metadata": {}, + "source": [ + "Make some more function spaces." ] }, { "cell_type": "code", "execution_count": null, - "id": "c6c6c8a0-2380-4b80-a4a8-49ee8bdf0397", + "id": "21226f5d-4074-44e5-ad26-b2bea599ec6f", "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", - ")" + "Σ = firedrake.TensorFunctionSpace(mesh, \"DG\", 0, symmetry=True)\n", + "T = firedrake.VectorFunctionSpace(mesh, \"CG\", 1)\n", + "Z = V * Σ * T\n", + "z = firedrake.Function(Z)\n", + "z.sub(0).assign(u)" ] }, { "cell_type": "code", "execution_count": null, - "id": "5f5d2b18-d3ad-4151-903e-276eb3c81cca", + "id": "624e48ee-cd15-4fb4-800e-54fce8394c1c", "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=\".\");" + "ε_c = Constant(0.02)\n", + "τ_c = Constant(0.1)\n", + "u_c = Constant(1000.0)" ] }, { "cell_type": "code", "execution_count": null, - "id": "358ad29b-685b-4561-8743-f64907ebd78a", + "id": "d279d3a2-3819-4e8f-85cd-8e03633e24b0", "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", + "from firedrake import sqrt, sym, grad, tr\n", + "ε = sym(grad(u))\n", + "ε_e = sqrt((inner(ε, ε) + tr(ε) ** 2) / 2)\n", + "I = firedrake.Identity(2)\n", + "M = τ_c / ε_c ** (1 / n) / 2 * ε_e ** (1 / n - 1) * (ε + tr(ε) * I)\n", + "#z.sub(1).project(M)\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]]" + "z.sub(2).interpolate(τ_d);" ] }, { "cell_type": "code", "execution_count": null, - "id": "c9bb179f-31c8-4021-9138-f8f1c36a78be", + "id": "fd5e1e00-bced-4bad-adcf-9c8fcfe051f4", "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", + "inflow_ids = (1, 2, 4, 5, 6)\n", + "outflow_ids = (3,)\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" + "bc = firedrake.DirichletBC(Z.sub(0), u, inflow_ids)" ] }, { "cell_type": "code", "execution_count": null, - "id": "46e8737e-5b26-4a84-921c-e248337eab92", + "id": "05965a2c-91f2-4931-9fd6-9e968215007e", "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", + "from dualform import ice_stream\n", "\n", - "u = estimator.solve()\n", - "firedrake.adjoint.pause_annotation()" + "fns = [\n", + " ice_stream.viscous_power,\n", + " ice_stream.friction_power,\n", + " ice_stream.boundary,\n", + " ice_stream.constraint,\n", + "]" ] }, { "cell_type": "code", "execution_count": null, - "id": "6a0ab4c7-4526-4169-b2a7-0b52c1923d4a", + "id": "855c9b65-1acd-4480-b280-2c8f67d84349", "metadata": {}, "outputs": [], "source": [ - "fig, axes = plt.subplots()\n", - "axes.set_aspect(\"equal\")\n", - "colors = firedrake.tripcolor(u, axes=axes)\n", - "fig.colorbar(colors);" + "u, M, τ = firedrake.split(z)\n", + "s = interpolate(max_value(b + h, (1 - ρ_I / ρ_W) * h), Q)\n", + "h_min = Constant(10.0)\n", + "p_I = ρ_I * g * max_value(h_min, h)\n", + "p_W = -ρ_W * g * min_value(0, s - h)\n", + "f = (1 - p_W / p_I) ** m\n", + "\n", + "fields = {\n", + " \"velocity\": u,\n", + " \"membrane_stress\": M,\n", + " \"basal_stress\": τ,\n", + " \"thickness\": h,\n", + " \"surface\": s,\n", + " \"floating\": f,\n", + "}\n", + "\n", + "params = {\n", + " \"viscous_yield_strain\": ε_c,\n", + " \"viscous_yield_stress\": τ_c,\n", + " \"friction_yield_speed\": u_c,\n", + " \"friction_yield_stress\": τ_c,\n", + "}\n", + "\n", + "exponents_1 = {\n", + " \"flow_law_exponent\": 1,\n", + " \"sliding_law_exponent\": 1,\n", + "}\n", + "\n", + "exponents = {\n", + " \"flow_law_exponent\": n,\n", + " \"sliding_law_exponent\": m,\n", + "}\n", + "\n", + "boundary_ids = {\"outflow_ids\": outflow_ids}\n", + "\n", + "L_1 = sum(fn(**fields, **params, **boundary_ids, **exponents_1) for fn in fns)\n", + "F_1 = firedrake.derivative(L_1, z)\n", + "\n", + "L = sum(fn(**fields, **params, **boundary_ids, **exponents) for fn in fns)\n", + "F = firedrake.derivative(L, z)\n", + "\n", + "sparams = {\n", + " \"solver_parameters\": {\n", + " \"snes_monitor\": None,\n", + " \"ksp_monitor\": None,\n", + " \"ksp_type\": \"gmres\",\n", + " \"pc_type\": \"ilu\",\n", + " },\n", + "}\n", + "problem_1 = firedrake.NonlinearVariationalProblem(F_1, z, bc)\n", + "solver_1 = firedrake.NonlinearVariationalSolver(problem_1, **sparams)\n", + "solver_1.solve()\n", + "\n", + "\n", + "problem = firedrake.NonlinearVariationalProblem(F, z, bc)\n", + "solver = firedrake.NonlinearVariationalSolver(problem, **sparams)\n", + "solver.solve()" ] } ],