Start script to extrapolate values to larger mesh

demos/larsen/extrapolate.py

@@ -0,0 +1,77 @@
+import argparse
+import subprocess
+import numpy as np
+import geojson
+import xarray
+import firedrake
+from firedrake import assemble, Constant, inner, grad, dx
+import icepack
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--input", default="larsen-initial.h5")
+parser.add_argument("--outline", default="larsen.geojson")
+parser.add_argument("--degree", type=int, default=1)
+parser.add_argument("--output")
+args = parser.parse_args()
+
+# Read in the input data
+with firedrake.CheckpointFile(args.input, "r") as chk:
+    input_mesh = chk.load_mesh()
+    u_input = chk.load_function(input_mesh, name="velocity")
+    θ_input = chk.load_function(input_mesh, name="log_fluidity")
+
+    Q = θ_input.function_space()
+    μ = firedrake.interpolate(Constant(1), Q)
+
+# Create the mesh and some function spaces
+with open(args.outline, "r") as outline_file:
+    outline = geojson.load(outline_file)
+
+geometry = icepack.meshing.collection_to_geo(outline)
+with open("larsen.geo", "w") as geometry_file:
+    geometry_file.write(geometry.get_code())
+
+command = "gmsh -2 -v 0 -o larsen.msh larsen.geo"
+subprocess.run(command.split())
+
+mesh = firedrake.Mesh("larsen.msh")
+
+import matplotlib.pyplot as plt
+fig, axes = plt.subplots()
+axes.set_aspect("equal")
+firedrake.triplot(mesh, axes=axes)
+axes.legend()
+plt.show()
+
+Q = firedrake.FunctionSpace(mesh, "CG", args.degree)
+V = firedrake.VectorFunctionSpace(mesh, "CG", args.degree)
+
+# Project the mask, log-fluidity, and velocity onto the larger mesh. The
+# regions with no data will be extrapolated by zero.
+μ = firedrake.project(μ, Q)
+μ.interpolate(firedrake.max_value(0, firedrake.min_value(1, μ)))
+
+Eθ = firedrake.project(θ_input, Q)
+Eu = firedrake.project(u_input, V)
+
+θ = Eθ.copy(deepcopy=True)
+u = Eu.copy(deepcopy=True)
+
+bc = firedrake.DirichletBC(V, Eu, [4, 6, 7, 8, 9])
+
+α = Constant(8e3)
+J = 0.5 * (μ * inner(u - Eu, u - Eu) + α**2 * inner(grad(u), grad(u))) * dx
+F = firedrake.derivative(J, u)
+firedrake.solve(F == 0, u, bc)
+
+area = assemble(μ * dx)
+print(np.sqrt(assemble(μ * inner(u - Eu, u - Eu) * dx) / area))
+
+import matplotlib.pyplot as plt
+fig, axes = plt.subplots()
+axes.set_aspect("equal")
+colors = firedrake.tripcolor(u, axes=axes)
+fig.colorbar(colors)
+plt.show()
+
+