Merge pull request #2 from firedrakeproject/JHopeCollins/burgers_wc4dvar

 Weak constraint 4DVar for 1D  advection and Burgers equations

.gitignore

@@ -160,3 +160,14 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+# VTK files
+*.vtu
+*.pvd
+
+# figures
+*.pdf
+*.png
+
+# output
+**/output/*

advection/advection_forward.py

@@ -0,0 +1,46 @@
+import firedrake as fd
+from firedrake.output import VTKFile
+from firedrake.petsc import PETSc
+
+from advection_utils import *
+
+import numpy as np
+np.set_printoptions(legacy='1.25')
+
+Print = PETSc.Sys.Print
+
+nx = 50
+cfl = 1.6
+nt = 40
+
+nprint = 5
+
+u = 1.
+dx = 1./nx
+dt = cfl*dx/u
+
+mesh = fd.PeriodicUnitIntervalMesh(nx)
+
+V = fd.FunctionSpace(mesh, "DG", 1)
+
+qn, qn1, stepper = timestepper(mesh, V, dt, u)
+
+# initial conditions
+x, = fd.SpatialCoordinate(mesh)
+ic = fd.Function(V).interpolate(analytic_solution(mesh, u, t=0))
+qn.assign(ic)
+
+file = VTKFile('output/advection.pvd', comm=mesh.comm)
+file.write(qn, t=0)
+
+t = 0.
+for i in range(nt):
+    qn1.assign(qn)
+    stepper.solve()
+    t += dt
+    qn.assign(qn1)
+
+    file.write(qn, t=t)
+
+    if (i+1) % nprint == 0:
+        Print(f"Timestep {str(i+1).rjust(3)} | t = {str(round(t, 4)).ljust(5)} | norm(qn) = {str(round(fd.norm(qn), 5)).ljust(5)}")

advection/advection_sc4dvar_aaorf.py

@@ -0,0 +1,54 @@
+import firedrake as fd
+from firedrake.__future__ import interpolate
+from firedrake.adjoint import (continue_annotation, pause_annotation,
+                               stop_annotating, Control, taylor_test,
+                               ReducedFunctional, minimize)
+from firedrake.adjoint import FourDVarReducedFunctional
+from advection_utils import *
+
+control = background.copy(deepcopy=True)
+
+continue_annotation()
+
+# create 4DVar reduced functional and record
+# background and initial observation functionals
+
+Jhat = FourDVarReducedFunctional(
+    Control(control),
+    background_iprod=norm2(B),
+    observation_iprod=norm2(R),
+    observation_err=observation_error(0),
+    weak_constraint=False)
+
+nstep = 0
+# record observation stages
+with Jhat.recording_stages(nstages=len(targets)-1) as stages:
+    # loop over stages
+    for stage, ctx in stages:
+        # start forward model
+        qn.assign(stage.control)
+
+        # propogate
+        for _ in range(observation_freq):
+            qn1.assign(qn)
+            stepper.solve()
+            qn.assign(qn1)
+            nstep += 1
+
+        # take observation
+        obs_err = observation_error(stage.observation_index)
+        stage.set_observation(qn, obs_err,
+                              observation_iprod=norm2(R))
+
+pause_annotation()
+
+print(f"{taylor_test(Jhat, control, values[0]) = }")
+
+options = {'disp': True, 'ftol': 1e-2}
+derivative_options = {'riesz_representation': None}
+
+opt = minimize(Jhat, options=options, method="L-BFGS-B",
+               derivative_options=derivative_options)
+
+print(f"{fd.errornorm(targets[0], control) = }")
+print(f"{fd.errornorm(targets[0], opt) = }")

advection/advection_sc4dvar_pyadj.py

@@ -0,0 +1,46 @@
+import firedrake as fd
+from firedrake.__future__ import interpolate
+from firedrake.adjoint import (continue_annotation, pause_annotation,
+                               stop_annotating, Control, taylor_test,
+                               ReducedFunctional, minimize)
+from advection_utils import *
+
+# record observation stages
+control = background.copy(deepcopy=True)
+
+continue_annotation()
+
+# background functional
+J = norm2(B)(control - background)
+
+# initial observation functional
+J += norm2(R)(observation_error(0)(control))
+
+nstep = 0
+qn.assign(control)
+
+for i in range(1, len(targets)):
+
+    for _ in range(observation_freq):
+        qn1.assign(qn)
+        stepper.solve()
+        qn.assign(qn1)
+        nstep += 1
+
+    # observation functional
+    J += norm2(R)(observation_error(i)(qn))
+
+pause_annotation()
+
+Jhat = ReducedFunctional(J, Control(control))
+
+print(f"{taylor_test(Jhat, control, values[0]) = }")
+
+options = {'disp': True, 'ftol': 1e-2}
+derivative_options = {'riesz_representation': 'l2'}
+
+opt = minimize(Jhat, options=options, method="L-BFGS-B",
+               derivative_options=derivative_options)
+
+print(f"{fd.errornorm(targets[0], control) = }")
+print(f"{fd.errornorm(targets[0], opt) = }")

advection/advection_utils.py

@@ -0,0 +1,128 @@
+import firedrake as fd
+from firedrake.__future__ import interpolate
+import numpy as np
+from sys import exit
+
+np.set_printoptions(legacy='1.25', precision=6)
+
+
+def norm2(w):
+    def n2(x):
+        return fd.assemble(fd.inner(x, fd.Constant(w)*x)*fd.dx)
+    return n2
+
+
+def timestepper(mesh, V, dt, u):
+    qn = fd.Function(V, name="qn")
+    qn1 = fd.Function(V, name="qn1")
+
+    def mass(q, phi):
+        return fd.inner(q, phi)*fd.dx
+
+    def tendency(q, phi):
+        uc = fd.as_vector([fd.Constant(u)])
+        n = fd.FacetNormal(mesh)
+        un = fd.Constant(0.5)*(fd.dot(uc, n) + abs(fd.dot(uc, n)))
+        return (- q*fd.div(phi*uc)*fd.dx
+                + fd.jump(phi)*fd.jump(un*q)*fd.dS)
+
+    # midpoint rule
+    q = fd.TrialFunction(V)
+    phi = fd.TestFunction(V)
+
+    qh = fd.Constant(0.5)*(q + qn)
+    eqn = mass(q - qn, phi) + fd.Constant(dt)*tendency(qh, phi)
+
+    stepper = fd.LinearVariationalSolver(
+        fd.LinearVariationalProblem(
+            fd.lhs(eqn), fd.rhs(eqn), qn1,
+            constant_jacobian=True))
+
+    return qn, qn1, stepper
+
+
+def analytic_solution(mesh, u, t, mag=1.0, phase=0.0):
+    x, = fd.SpatialCoordinate(mesh)
+    return mag*fd.sin(2*fd.pi*((x + phase) - u*t))
+
+B = 1
+R = 1
+Q = 1
+
+ensemble = fd.Ensemble(fd.COMM_WORLD, 1)
+
+nx = 20
+cfl = 1.6
+
+u = 1.
+dx = 1./nx
+dt = cfl*dx/u
+
+mesh = fd.PeriodicUnitIntervalMesh(nx, comm=ensemble.comm)
+
+V = fd.FunctionSpace(mesh, "DG", 1)
+
+qn, qn1, stepper = timestepper(mesh, V, dt, u)
+
+# initial conditions
+x, = fd.SpatialCoordinate(mesh)
+ic = fd.Function(V).interpolate(analytic_solution(mesh, u, t=0))
+qn.assign(ic)
+
+# observation operator
+observation_freq = 4
+observation_n = 3
+
+# we have an extra observation at the initial time
+observation_times = [i*observation_freq*dt
+                     for i in range(observation_n+1)]
+
+observation_locations = [
+    [x] for x in [0.13, 0.18, 0.34, 0.36, 0.49, 0.61, 0.72, 0.99]
+]
+
+observation_mesh = fd.VertexOnlyMesh(mesh, observation_locations)
+Vobs = fd.FunctionSpace(observation_mesh, "DG", 0)
+
+
+def H(x):
+    return fd.assemble(interpolate(x, Vobs))
+
+
+# ground truth
+targets = [
+    fd.Function(V).interpolate(analytic_solution(mesh, u, t))
+    for t in observation_times
+]
+
+# take observations
+y = [H(x) for x in targets]
+
+
+def observation_error(i):
+    def obs_err(x):
+        err = fd.Function(Vobs)
+        err.assign(H(x) - y[i])
+        return err
+    return obs_err
+
+
+prior_mag = 0.9
+prior_phase = 0.1
+
+# background
+background = fd.Function(V).interpolate(
+    analytic_solution(mesh, u, 0, mag=prior_mag, phase=prior_phase))
+
+# other values to evaluate reduced functional at
+values = [
+    fd.Function(V).interpolate(
+        analytic_solution(mesh, u, t+0.1, mag=1.1, phase=-0.2))
+    for t in observation_times
+]
+
+values1 = [
+    fd.Function(V).interpolate(
+        analytic_solution(mesh, u, t+0.3, mag=0.8, phase=0.3))
+    for t in observation_times
+]

advection/advection_wc4dvar_aaorf.py

@@ -0,0 +1,62 @@
+import firedrake as fd
+from firedrake.__future__ import interpolate
+from firedrake.adjoint import (continue_annotation, pause_annotation, stop_annotating,
+                               Control, taylor_test, minimize)
+from firedrake.adjoint import FourDVarReducedFunctional
+from advection_utils import *
+
+control = fd.EnsembleFunction(
+    ensemble, [V for _ in range(len(targets))])
+
+for x in control.subfunctions:
+    x.assign(background)
+
+continue_annotation()
+
+# create 4DVar reduced functional and record
+# background and initial observation functionals
+
+Jhat = FourDVarReducedFunctional(
+    Control(control),
+    background_iprod=norm2(B),
+    observation_iprod=norm2(R),
+    observation_err=observation_error(0),
+    weak_constraint=True)
+
+nstep = 0
+# record observation stages
+with Jhat.recording_stages() as stages:
+    # loop over stages
+    for stage, ctx in stages:
+        # start forward model
+        qn.assign(stage.control)
+
+        # propogate
+        for _ in range(observation_freq):
+            qn1.assign(qn)
+            stepper.solve()
+            qn.assign(qn1)
+            nstep += 1
+
+        # take observation
+        obs_err = observation_error(stage.observation_index)
+        stage.set_observation(qn, obs_err,
+                              observation_iprod=norm2(R),
+                              forward_model_iprod=norm2(Q))
+
+pause_annotation()
+
+# the perturbation values need to be held in the
+# same type as the control i.e. and EnsembleFunction
+vals = control.copy()
+for v0, v1 in zip(vals.subfunctions, values):
+    v0.assign(v1)
+
+print(f"{Jhat(control) = }")
+print(f"{taylor_test(Jhat, control, vals) = }")
+
+options = {'disp': True, 'ftol': 1e-2}
+derivative_options = {'riesz_representation': 'l2'}
+
+opt = minimize(Jhat, options=options, method="L-BFGS-B",
+               derivative_options=derivative_options)