Update to current firedrake (#20)

Make firedrake-ts usable with the current firedrake

* make _TSContext as a subclass of _SNESContext to reduce some code duplications
* deal with the new notion of cofunction in UFL/firedrake
* modify some examples to run
* bump firedrake-ts to version 0.2

Author: JaroslavHron

examples/adjoint.py

@@ -63,7 +63,9 @@ def ts_monitor(ts, steps, time, X):
 with dJdu.dat.vec as vec:
     dJdu_vec = vec
 
-print(f"Norm of dJdu before the adjoint solve: {fd.norm(dJdu)}")
+fdJdu=dJdu.riesz_representation()
+    
+print(f"Norm of dJdu before the adjoint solve: {dJdu_vec.norm()=} {fd.norm(fdJdu)=}")
 
 # setCostGradients accepts two PETSc Vecs
 # J is the objective function
@@ -75,7 +77,9 @@ def ts_monitor(ts, steps, time, X):
 
 solver.adjoint_solve()
 
-print(f"Norm of dJdu after the adjoint solve: {fd.norm(dJdu)}")
+fdJdu=dJdu.riesz_representation()
+print(f"Norm of dJdu after the adjoint solve: {dJdu_vec.norm()=} {fd.norm(fdJdu)=}")
+
 
 adj_out = fd.File("result/adj.pvd")
-adj_out.write(dJdu, time=0)
+adj_out.write(fdJdu, time=0)

examples/cahn-hilliard.py

@@ -5,7 +5,7 @@
 
 import firedrake_ts
 from firedrake import *
-from pyop2 import op2
+import numpy as np
 
 # Model parameters
 lmbda = 1.0e-02  # surface parameter
@@ -26,18 +26,6 @@
 c, mu = split(u)
 c_t, mu_t = split(u_t)
 
-# Create intial conditions and interpolate
-init_code = "A[0] = 0.63 + 0.02*(0.5 - (double)random()/RAND_MAX);"
-user_code = """int __rank;
-MPI_Comm_rank(MPI_COMM_WORLD, &__rank);
-srandom(2 + __rank);"""
-par_loop(
-    init_code,
-    direct,
-    {"A": (u[0], WRITE)},
-    kernel_kwargs={"headers": ["#include <stdlib.h>"],
-                   "user_code": user_code}
-)
 
 # Compute the chemical potential df/dc
 c = variable(c)
@@ -49,7 +37,13 @@
 F1 = mu * v * dx - dfdc * v * dx - lmbda * dot(grad(c), grad(v)) * dx
 F = F0 + F1
 
-pc = "fieldsplit"
+rng = np.random.default_rng(11)
+c , mu = u.subfunctions
+with c.dat.vec as v:
+    v[:]=0.63 + 0.2*(0.5-rng.random(v.size))
+
+    
+pc = "lu" #"fieldsplit"
 ksp = "lgmres"
 inner_ksp = "preonly"
 maxit = 1
@@ -77,9 +71,9 @@
 
 params["snes_monitor"] = None
 params["ts_monitor"] = None
-params["ts_view"] = None
+#params["ts_view"] = None
 
-problem = firedrake_ts.DAEProblem(F, u, u_t, (0.0, 50 * dt))
+problem = firedrake_ts.DAEProblem(F, u, u_t, (0.0, 2 * dt))
 solver = firedrake_ts.DAESolver(problem, solver_parameters=params)
 
 if pc in ["fieldsplit", "ilu"]:

examples/heat-explicit.py

@@ -1,28 +1,35 @@
 from firedrake import *
 import firedrake_ts
+from firedrake.__future__ import interpolate
 
-mesh = UnitIntervalMesh(11)
+mesh = UnitIntervalMesh(10)
 V = FunctionSpace(mesh, "P", 1)
 
 u = Function(V)
 u_t = Function(V)
 v = TestFunction(V)
 
-# F(u̇, u, t) = G(u, t)
+# F(u_t, u, t) = G(u, t)
 F = inner(u_t, v) * dx
 G = -(inner(grad(u), grad(v)) * dx - 1.0 * v * dx)
 
-bc = DirichletBC(V, 0.0, "on_boundary")
+bc1 = DirichletBC(V, 1.0, 1)
+bc2 = DirichletBC(V, 0.0, 2)
+bcs=[bc1, bc2]
 
 x = SpatialCoordinate(mesh)
-bump = conditional(lt(abs(x[0] - 0.5), 0.1), 1.0, 0.0)
-u.interpolate(bump)
+bump = conditional(lt(x[0], 0.5), 1.0, 0.0)
+assemble(interpolate(bump, u), tensor=u)
+print(f'{u.dat.data}=')
 
-problem = firedrake_ts.DAEProblem(F, u, u_t, (0.0, 1.0), bcs=bc, G=G)
-solver = firedrake_ts.DAESolver(problem)
+
+def monitor(ts, step, t, x):
+    print(f'{solver.ts.getTime()=}')
+    print(f'{u.dat.data}=')
+
+problem = firedrake_ts.DAEProblem(F, u, u_t, (0.0, 0.4), bcs=bcs, G=G)
+
+solver = firedrake_ts.DAESolver(problem, options_prefix='', monitor_callback=monitor)
 
 solver.solve()
 
-print(solver.ts.getTime())
-# solver.ts.view()
-print(u.dat.data)

examples/heat.py

@@ -1,5 +1,6 @@
 from firedrake import *
 import firedrake_ts
+from firedrake.__future__ import interpolate
 
 mesh = UnitIntervalMesh(10)
 V = FunctionSpace(mesh, "P", 1)
@@ -9,16 +10,37 @@
 v = TestFunction(V)
 F = inner(u_t, v) * dx + inner(grad(u), grad(v)) * dx - 1.0 * v * dx
 
-bc = DirichletBC(V, 0.0, "on_boundary")
+bc_val=Constant(0.0)
+bc = DirichletBC(V, bc_val, "on_boundary")
 
 x = SpatialCoordinate(mesh)
-# gaussian = exp(-30*(x[0]-0.5)**2)
-bump = conditional(lt(abs(x[0] - 0.5), 0.1), 1.0, 0.0)
-u.interpolate(bump)
+bump = conditional(lt(x[0], 0.5), 1.0, 0.0)
+assemble(interpolate(bump, u), tensor=u)
+print(f'{u.dat.data}=')
 
-problem = firedrake_ts.DAEProblem(F, u, u_t, (0.0, 1.0), bcs=bc)
-solver = firedrake_ts.DAESolver(problem)
 
+p0=NonlinearVariationalProblem(F, u, bcs=bc)
+s0=NonlinearVariationalSolver(p0)
+s0.solve()
+print(f'{u.dat.data}=')
+
+bc_val.assign(1.0)
+s0.solve()
+print(f'{u.dat.data}=')
+
+
+def monitor(ts, step, t, x):
+    print(f'{solver.ts.getTime()=} {problem.time=} {bc_val=}')
+    print(f'{u.dat.data}=')
+
+problem = firedrake_ts.DAEProblem(F, u, u_t, (0.0, 1.0), bcs=bc, time=bc_val)
+solver = firedrake_ts.DAESolver(problem, options_prefix='', monitor_callback=monitor)
+
+bc_val.assign(2.0)
+solver.solve()
+print(id(bc_val), id(problem.time))
+
+bc_val.assign(4.0)
 solver.solve()
 
 print(solver.ts.getTime())