Simplifications to dome notebook + unit test

* Remove the initial SIA solve. This is unnecessary once you use
  continuation or a composite rheology.
* In the initial momentum solve, use the full function space
  `V * Σ * V * Q` and make the `h` block the identity matrix. This lets
  us reuse the momentum part of the equation and only change the mass
  part of the full problem.

Author: danshapero

notebooks/rainier-monolithic.ipynb

@@ -155,8 +155,7 @@
     "r_h = Constant(5e3)\n",
     "H = Constant(100.0)\n",
     "expr = H * firedrake.max_value(0, 1 - inner(x, x) / r_h**2)\n",
-    "h = firedrake.Function(Q).interpolate(expr)\n",
-    "h_0 = h.copy(deepcopy=True)"
+    "h_init = firedrake.Function(Q).interpolate(expr)"
    ]
   },
   {
@@ -168,7 +167,7 @@
    "source": [
     "fig, ax = plt.subplots()\n",
     "ax.set_aspect(\"equal\")\n",
-    "colors = firedrake.tripcolor(h, axes=ax)\n",
+    "colors = firedrake.tripcolor(h_init, axes=ax)\n",
     "fig.colorbar(colors);"
    ]
   },
@@ -187,7 +186,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "s = firedrake.Function(Q).interpolate(b + h)\n",
+    "s = firedrake.Function(Q).interpolate(b + h_init)\n",
     "a = firedrake.Function(Q).interpolate(smb(s))"
    ]
   },
@@ -227,10 +226,12 @@
   },
   {
    "cell_type": "markdown",
-   "id": "bec1a763-ab01-4f63-ae93-d92f0ce0ef2a",
+   "id": "922938b5-e240-4573-be3d-a6f6302951c2",
    "metadata": {},
    "source": [
-    "We'll use the shallow ice approximation, which diagnoses the velocity at a point in a purely local fashion using the thickness and surface slope, to initialize the ice velocity."
+    "To proceed, we have to make function spaces to represent the membrane and basal stresses and create a variable that stores the velocity and stresses.\n",
+    "Here we use piecewise constant finite elements for the membrane stresses.\n",
+    "In general, if you use $CG(k)$ elements for the velocity, you need to use $DG(k - 1)$ elements for the stresses."
    ]
   },
   {
@@ -242,63 +243,10 @@
    "source": [
     "V = firedrake.VectorFunctionSpace(mesh, cg1)\n",
     "\n",
-    "u = firedrake.Function(V)\n",
-    "v = firedrake.TestFunction(V)\n",
-    "\n",
-    "ρ_I = Constant(ice_density)\n",
-    "g = Constant(gravity)\n",
-    "\n",
-    "n = Constant(glen_flow_law)\n",
-    "\n",
-    "P = ρ_I * g * h\n",
-    "S_n = inner(grad(s), grad(s))**((n - 1) / 2)\n",
-    "u_shear = -2 * A * P ** n / (n + 2) * h * S_n * grad(s)\n",
-    "F = inner(u - u_shear, v) * dx\n",
-    "\n",
-    "solver_params = {\"snes_type\": \"ksponly\", \"ksp_type\": \"gmres\"}\n",
-    "fc_params = {\"quadrature_degree\": 6}\n",
-    "params = {\"solver_parameters\": solver_params, \"form_compiler_parameters\": fc_params}\n",
-    "firedrake.solve(F == 0, u, **params)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "5a32887c-dcd3-4aae-844e-1a6c70b183db",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "fig, ax = plt.subplots()\n",
-    "ax.set_aspect(\"equal\")\n",
-    "colors = firedrake.tripcolor(u, axes=ax)\n",
-    "fig.colorbar(colors);"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "9a900c91-a77f-4d7a-86e3-9c014512e3b9",
-   "metadata": {},
-   "source": [
-    "Now we'll get to the real modeling.\n",
-    "To proceed, we have to make function spaces to represent the membrane and basal stresses and create a variable that stores the velocity and stresses.\n",
-    "Here we use piecewise constant finite elements for the stresses.\n",
-    "In general, if you use $CG(k)$ elements for the velocity, you need to use $DG(k - 1)$ elements for the stresses."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "db935250-689e-4ca7-aee5-c5b0f212c4d4",
-   "metadata": {},
-   "outputs": [],
-   "source": [
     "dg0 = firedrake.FiniteElement(\"DG\", \"triangle\", 0)\n",
     "Σ = firedrake.TensorFunctionSpace(mesh, dg0, symmetry=True)\n",
-    "T = firedrake.VectorFunctionSpace(mesh, cg1)\n",
-    "Z = V * Σ * T\n",
-    "z = firedrake.Function(Z)\n",
-    "\n",
-    "z.sub(0).assign(u);"
+    "Z = V * Σ * V * Q\n",
+    "z = firedrake.Function(Z)"
    ]
   },
   {
@@ -322,6 +270,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "n = Constant(glen_flow_law)\n",
+    "\n",
     "τ_c = Constant(0.1)\n",
     "ε_c = Constant(A * τ_c ** n)"
    ]
@@ -351,6 +301,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "u, M, τ, h = firedrake.split(z)\n",
+    "\n",
     "K = h * A / (n + 2)\n",
     "U_c = Constant(100.0)\n",
     "u_c = K * τ_c ** n + U_c"
@@ -378,16 +330,15 @@
     "    \"sliding_coefficient\": u_c / τ_c,\n",
     "}\n",
     "\n",
-    "u, M, τ = firedrake.split(z)\n",
     "fields = {\n",
     "    \"velocity\": u,\n",
     "    \"membrane_stress\": M,\n",
     "    \"basal_stress\": τ,\n",
     "    \"thickness\": h,\n",
-    "    \"surface\": s,\n",
+    "    \"surface\": b + h,\n",
     "}\n",
     "\n",
-    "v, N, σ = firedrake.TestFunctions(Z)"
+    "v, N, σ, η = firedrake.TestFunctions(Z)"
    ]
   },
   {
@@ -441,14 +392,17 @@
     "from icepack2 import model\n",
     "from icepack2.model.variational import momentum_balance, flow_law, friction_law\n",
     "\n",
-    "F = (\n",
+    "F_momentum = (\n",
     "    momentum_balance(**fields, test_function=v)\n",
     "    + firedrake.replace(flow_law(**fields, **glen_rheology, test_function=N), {h: H})\n",
     "    + α * firedrake.replace(flow_law(**fields, **linear_rheology, test_function=N), {h: H})\n",
     "    + friction_law(**fields, **glen_rheology, test_function=σ)\n",
     "    + α * friction_law(**fields, **linear_rheology, test_function=σ)\n",
     ")\n",
     "\n",
+    "F_mass = (h - h_init) * η * dx\n",
+    "\n",
+    "F = F_momentum + F_mass\n",
     "momentum_problem = firedrake.NonlinearVariationalProblem(F, z, **pparams)\n",
     "momentum_solver = firedrake.NonlinearVariationalSolver(momentum_problem, **sparams)"
    ]
@@ -482,7 +436,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "u_init, M_init, τ_init = z.subfunctions\n",
+    "u_init, M_init, τ_init, h_init = z.subfunctions\n",
     "\n",
     "fig, axes = plt.subplots()\n",
     "axes.set_aspect(\"equal\")\n",
@@ -526,21 +480,6 @@
     "Pack in the initial values of the velocity, membrane stress, and basal stress that we computed before."
    ]
   },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "873be577-b14f-4f6a-bd69-fbe9a6cc2f34",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "W = V * Σ * T * Q\n",
-    "w = firedrake.Function(W)\n",
-    "w.sub(0).assign(u_init)\n",
-    "w.sub(1).assign(M_init)\n",
-    "w.sub(2).assign(τ_init)\n",
-    "w.sub(3).assign(h);"
-   ]
-  },
   {
    "cell_type": "markdown",
    "id": "0573021d-eade-4e73-b8b1-563c56da95aa",
@@ -556,27 +495,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "u, M, τ, h = firedrake.split(w)\n",
-    "v, N, σ, η = firedrake.TestFunctions(W)\n",
-    "\n",
-    "fields = {\n",
-    "    \"velocity\": u,\n",
-    "    \"membrane_stress\": M,\n",
-    "    \"basal_stress\": τ,\n",
-    "    \"thickness\": h,\n",
-    "    \"surface\": b + h,\n",
-    "}\n",
-    "\n",
-    "F_momentum = (\n",
-    "    momentum_balance(**fields, test_function=v)\n",
-    "    + firedrake.replace(flow_law(**fields, **glen_rheology, test_function=N), {h: H})\n",
-    "    + α * firedrake.replace(flow_law(**fields, **linear_rheology, test_function=N), {h: H})\n",
-    "    + friction_law(**fields, **glen_rheology, test_function=σ)\n",
-    "    + α * friction_law(**fields, **linear_rheology, test_function=σ)\n",
-    ")\n",
-    "\n",
     "F_mass = model.mass_balance(thickness=h, velocity=u, accumulation=a, test_function=η)\n",
-    "\n",
     "F = F_momentum + F_mass"
    ]
   },
@@ -591,8 +510,8 @@
     "t = Constant(0.0)\n",
     "dt = Constant(1.0 / 6)\n",
     "\n",
-    "lower = firedrake.Function(W)\n",
-    "upper = firedrake.Function(W)\n",
+    "lower = firedrake.Function(Z)\n",
+    "upper = firedrake.Function(Z)\n",
     "lower.assign(-np.inf)\n",
     "upper.assign(+np.inf)\n",
     "lower.subfunctions[3].assign(0.0)\n",
@@ -603,18 +522,17 @@
     "        \"snes_monitor\": \":rainier-output-vi.log\",\n",
     "        \"snes_type\": \"vinewtonrsls\",\n",
     "        \"snes_max_it\": 200,\n",
-    "        #\"snes_linesearch_type\": \"nleqerr\",\n",
     "        \"ksp_type\": \"gmres\",\n",
     "        \"pc_type\": \"lu\",\n",
     "        \"pc_factor_mat_solver_type\": \"mumps\",\n",
     "    },\n",
-    "    \"form_compiler_parameters\": fc_params,\n",
+    "    \"form_compiler_parameters\": {\"quadrature_degree\": 6},\n",
     "    \"stage_type\": \"value\",\n",
     "    \"basis_type\": \"Bernstein\",\n",
     "    \"bounds\": bounds,\n",
     "}\n",
     "\n",
-    "solver = irksome.TimeStepper(F, tableau, t, dt, w, **bparams)"
+    "solver = irksome.TimeStepper(F, tableau, t, dt, z, **bparams)"
    ]
   },
   {
@@ -632,18 +550,18 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "us = [w.subfunctions[0].copy(deepcopy=True)]\n",
-    "hs = [w.subfunctions[3].copy(deepcopy=True)]\n",
+    "us = [z.subfunctions[0].copy(deepcopy=True)]\n",
+    "hs = [z.subfunctions[3].copy(deepcopy=True)]\n",
     "\n",
     "final_time = 500.0\n",
     "num_steps = int(final_time / float(dt))\n",
     "for step in trange(num_steps):\n",
     "    solver.advance()\n",
-    "    h = w.subfunctions[3]\n",
+    "    h = z.subfunctions[3]\n",
     "    a.interpolate(smb(b + h))\n",
     "\n",
-    "    us.append(w.subfunctions[0].copy(deepcopy=True))\n",
-    "    hs.append(w.subfunctions[3].copy(deepcopy=True))"
+    "    us.append(z.subfunctions[0].copy(deepcopy=True))\n",
+    "    hs.append(z.subfunctions[3].copy(deepcopy=True))"
    ]
   },
   {
@@ -679,7 +597,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "u, M, τ, h = w.subfunctions\n",
+    "u, M, τ, h = z.subfunctions\n",
     "\n",
     "fig, axes = plt.subplots()\n",
     "axes.set_aspect(\"equal\")\n",