Skip to content

Commit

Permalink
Use preprocessing in notebook
Browse files Browse the repository at this point in the history
  • Loading branch information
@danshapero
danshapero committed Oct 2, 2023
1 parent 2cafea4 commit 5e81727
Showing 1 changed file with 157 additions and 105 deletions.
262 changes: 157 additions & 105 deletions demos/helheim/helheim.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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()"
]
},
{
Expand All @@ -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."
]
},
{
Expand All @@ -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",
Expand All @@ -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()"
]
}
],
Expand Down

0 comments on commit 5e81727

Please sign in to comment.