Add script to interpolate velocity data to mesh

demos/helheim/preprocess.py

@@ -0,0 +1,110 @@
+import subprocess
+import numpy as np
+import geojson
+import rasterio
+import firedrake
+import firedrake.adjoint
+from firedrake import assemble, Constant, inner, grad, dx
+import icepack
+import data
+
+# Fetch the glacier outline, generate a mesh, and create a function space
+outline_filename = icepack.datasets.fetch_outline("helheim")
+with open(outline_filename, "r") as outline_file:
+    outline = geojson.load(outline_file)
+
+geometry = icepack.meshing.collection_to_geo(outline)
+with open("helheim.geo", "w") as geometry_file:
+    geometry_file.write(geometry.get_code())
+
+command = "gmsh -2 -v 0 -o helheim.msh helheim.geo"
+subprocess.run(command.split())
+
+mesh = firedrake.Mesh("helheim.msh")
+V = firedrake.VectorFunctionSpace(mesh, "CG", 1)
+
+# Read all the raw velocity data
+coords = np.array(list(geojson.utils.coords(outline)))
+delta = 2.5e3
+xmin, xmax = coords[:, 0].min() - delta, coords[:, 0].max() + delta
+ymin, ymax = coords[:, 1].min() - delta, coords[:, 1].max() + delta
+measures_filenames = data.fetch_measures_velocity()
+
+velocity_data = {}
+for key in ["vx", "vy", "ex", "ey"]:
+    filename = [f for f in measures_filenames if key in f][0]
+    with rasterio.open(filename, "r") as source:
+        window = rasterio.windows.from_bounds(
+            left=xmin, right=xmax, bottom=ymin, top=ymax, transform=source.transform
+        ).round_lengths().round_offsets()
+        transform = source.window_transform(window)
+        velocity_data[key] = source.read(indexes=1, window=window)
+
+# Find the points that are inside the domain and create a point cloud
+indices = np.array(
+    [
+        (i, j)
+        for i in range(window.width)
+        for j in range(window.height)
+        if (
+            mesh.locate_cell(transform * (i, j), tolerance=1e-8) and
+            velocity_data["ex"][j, i] >= 0.0
+        )
+    ]
+)
+xs = np.array([transform * idx for idx in indices])
+point_set = firedrake.VertexOnlyMesh(
+    mesh, xs, missing_points_behaviour="error"
+)
+
+Δ = firedrake.FunctionSpace(point_set, "DG", 0)
+
+# Set up and solve a data assimilation problem
+u = firedrake.Function(V)
+N = Constant(len(indices))
+area = assemble(Constant(1) * dx(mesh))
+
+def loss_functional(u):
+    u_int = firedrake.interpolate(u[0], Δ)
+    v_int = firedrake.interpolate(u[1], Δ)
+    δu = u_int - u_o
+    δv = v_int - v_o
+    return 0.5 / N * ((δu / σ_x)**2 + (δv / σ_y)**2) * dx
+
+def regularization(u):
+    Ω = Constant(area)
+    α = Constant(80.0)   # TODO: tune this
+    return 0.5 * α**2 / Ω * inner(grad(u), grad(u)) * dx
+
+u_o = firedrake.Function(Δ)
+v_o = firedrake.Function(Δ)
+σ_x = firedrake.Function(Δ)
+σ_y = firedrake.Function(Δ)
+
+u_o.dat.data[:] = velocity_data["vx"][indices[:, 1], indices[:, 0]]
+v_o.dat.data[:] = velocity_data["vy"][indices[:, 1], indices[:, 0]]
+σ_x.dat.data[:] = velocity_data["ex"][indices[:, 1], indices[:, 0]]
+σ_y.dat.data[:] = velocity_data["ey"][indices[:, 1], indices[:, 0]]
+
+firedrake.adjoint.continue_annotation()
+problem = icepack.statistics.StatisticsProblem(
+    simulation=lambda u: u.copy(deepcopy=True),
+    loss_functional=loss_functional,
+    regularization=regularization,
+    controls=u,
+)
+
+estimator = icepack.statistics.MaximumProbabilityEstimator(
+    problem,
+    gradient_tolerance=5e-5,
+    step_tolerance=1e-8,
+    max_iterations=50,
+)
+
+u = estimator.solve()
+firedrake.adjoint.pause_annotation()
+
+# Save the results to disk
+with firedrake.CheckpointFile("helheim.h5", "w") as chk:
+    chk.save_mesh(mesh)
+    chk.save_function(u, name="velocity")