Correlation matrices

advection/advection_sc4dvar_aaorf.py

@@ -15,9 +15,9 @@
 
 Jhat = FourDVarReducedFunctional(
     Control(control),
-    background_iprod=norm2(B),
-    observation_iprod=norm2(R),
-    observation_err=observation_error(0),
+    background_covariance=B,
+    observation_covariance=R,
+    observation_error=observation_error(0),
     weak_constraint=False)
 
 nstep = 0
@@ -38,13 +38,13 @@
         # take observation
         obs_err = observation_error(stage.observation_index)
         stage.set_observation(qn, obs_err,
-                              observation_iprod=norm2(R))
+                              observation_covariance=R)
 
 pause_annotation()
 
 print(f"{taylor_test(Jhat, control, values[0]) = }")
 
-options = {'disp': True, 'ftol': 1e-2}
+options = {'disp': fd.COMM_WORLD.rank == 0, 'ftol': 1e-2}
 derivative_options = {'riesz_representation': None}
 
 opt = minimize(Jhat, options=options, method="L-BFGS-B",

advection/advection_sc4dvar_pyadj.py

@@ -11,10 +11,10 @@
 continue_annotation()
 
 # background functional
-J = norm2(B)(control - background)
+J = covariance_norm(control - background, B)
 
 # initial observation functional
-J += norm2(R)(observation_error(0)(control))
+J += covariance_norm(observation_error(0)(control), R)
 
 nstep = 0
 qn.assign(control)
@@ -28,15 +28,15 @@
         nstep += 1
 
     # observation functional
-    J += norm2(R)(observation_error(i)(qn))
+    J += covariance_norm(observation_error(i)(qn), R)
 
 pause_annotation()
 
 Jhat = ReducedFunctional(J, Control(control))
 
 print(f"{taylor_test(Jhat, control, values[0]) = }")
 
-options = {'disp': True, 'ftol': 1e-2}
+options = {'disp': fd.COMM_WORLD.rank == 0, 'ftol': 1e-2}
 derivative_options = {'riesz_representation': 'l2'}
 
 opt = minimize(Jhat, options=options, method="L-BFGS-B",

advection/advection_utils.py

@@ -3,13 +3,9 @@
 import numpy as np
 from sys import exit
 
-np.set_printoptions(legacy='1.25', precision=6)
-
+from firedrake.adjoint.fourdvar_reduced_functional import covariance_norm
 
-def norm2(w):
-    def n2(x):
-        return fd.assemble(fd.inner(x, fd.Constant(w)*x)*fd.dx)
-    return n2
+np.set_printoptions(legacy='1.25', precision=6)
 
 
 def timestepper(mesh, V, dt, u):
@@ -45,9 +41,9 @@ 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
+B = 10.
+R = 0.1
+Q = 0.2*B
 
 ensemble = fd.Ensemble(fd.COMM_WORLD, 1)
 

advection/advection_wc4dvar_aaorf.py

@@ -5,58 +5,78 @@
 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)
+
+def make_fdvrf():
+    if ensemble.ensemble_size == 1:
+        nlocal_observations = len(targets)
+    elif ensemble.ensemble_size == 3:
+        nlocal_observations = 2 if ensemble.ensemble_rank == 0 else 1
+    else:
+        raise ValueError("Must use either 1 or 3 ensemble ranks")
+    control_space = fd.EnsembleFunctionSpace(
+        [V for _ in range(nlocal_observations)], ensemble)
+    control = fd.EnsembleFunction(control_space)
+
+    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_covariance=B,
+        observation_covariance=R,
+        observation_error=observation_error(0),
+        weak_constraint=True)
+
+    nstep = 0
+    # record observation stages
+    with Jhat.recording_stages(nstep=nstep) 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)
+                ctx.nstep += 1
+
+            # take observation
+            obs_err = observation_error(stage.observation_index)
+            stage.set_observation(qn, obs_err,
+                                  observation_covariance=R,
+                                  forward_model_covariance=Q)
+    pause_annotation()
+
+    return Jhat, control
+
+
+if __name__ == '__main__':
+    Jhat, control = make_fdvrf()
+
+    # 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, method="L-BFGS-B", options=options,
+                   derivative_options=derivative_options)
+    J0 = Jhat(control)
+    Jopt = Jhat(opt)
+    print(f"{J0 = :.3e} | {Jopt = :.3e} | {Jopt/J0 = :.3e}")

advection/advection_wc4dvar_aaorf_tao.py

@@ -0,0 +1,62 @@
+import firedrake as fd
+from firedrake.adjoint import (
+    continue_annotation, pause_annotation, minimize,
+    stop_annotating, Control, taylor_test)
+from firedrake.adjoint import FourDVarReducedFunctional
+from firedrake.petsc import PETSc
+from advection_utils import *
+from fdvar import TAOSolver
+from sys import exit
+
+from advection_wc4dvar_aaorf import make_fdvrf
+Jhat, control = make_fdvrf()
+
+Print = PETSc.Sys.Print
+
+# the perturbation values need to be held in the
+# same type as the control i.e. and EnsembleFunction
+x0 = control.copy()
+
+ksp_params = {
+    'monitor_short': None,
+    # 'converged_rate': None,
+    # 'converged_reason': None,
+}
+
+tao_params = {
+    'tao_view': ':tao_view.log',
+    'tao': {
+        'monitor': None,
+        'converged_reason': None,
+        'gatol': 1e-1,
+        'grtol': 1e-1,
+        'gttol': 1e-1,
+    },
+    'tao_type': 'nls',
+    'tao_nls': {
+        'ksp': ksp_params,
+        'ksp_type': 'gmres',
+        'pc_type': 'lmvm',
+        'ksp_rtol': 1e-1,
+        # 'ksp_type': 'gmres',
+        # 'pc_type': 'python',
+        # 'pc_python_type': 'fdvar.AllAtOnceJacobiPC',
+
+    },
+    # 'tao_cg': {
+    #     'ksp': ksp_params,
+    #     'ksp_rtol': 1e-1,
+    #     'type': 'pr',  # fr-pr-prp-hs-dy
+    # },
+}
+tao = TAOSolver(Jhat, options_prefix="",
+                solver_parameters=tao_params)
+tao.solve()
+
+xopt = Jhat.control.copy()
+J0 = Jhat(x0)
+Jopt = Jhat(xopt)
+
+Print(f"{J0 = :.3e} | {Jopt = :.3e} | {Jopt/J0 = :.3e}")
+
+

advection/advection_wc4dvar_covariancemat.py

@@ -0,0 +1,57 @@
+import firedrake as fd
+from firedrake.petsc import PETSc, OptionsManager
+from firedrake.adjoint import pyadjoint  # noqa: F401
+from firedrake.matrix import ImplicitMatrix
+from pyadjoint.optimization.tao_solver import PETScVecInterface
+from advection_wc4dvar_aaorf import make_fdvrf
+from mpi4py import MPI
+import numpy as np
+from typing import Optional, Collection
+from fdvar.mat import *
+
+
+def CovarianceMat(covariancerf):
+    space = covariancerf.controls[0].control.function_space()
+    comm = space.mesh().comm
+    covariance = covariancerf.covariance
+    if isinstance(covariance, Collection):
+        covariance, power = covariance
+
+    sizes = space.dof_dset.layout_vec.sizes
+    shape = (sizes, sizes)
+    covmat = PETSc.Mat().createConstantDiagonal(
+        shape, covariance, comm=comm)
+
+    covmat.setUp()
+    covmat.assemble()
+    return covmat
+
+
+Jhat, control = make_fdvrf()
+ensemble = Jhat.ensemble
+
+# >>>>> Covariance
+
+# Covariance Mat
+# covrf = Jhat.background_norm
+covrf = Jhat.stages[0].model_norm
+covmat = CovarianceMat(covrf)
+
+# Covariance KSP
+covksp = PETSc.KSP().create(comm=ensemble.comm)
+covksp.setOptionsPrefix('cov_')
+covksp.setOperators(covmat)
+
+covksp.pc.setType(PETSc.PC.Type.JACOBI)
+covksp.setType(PETSc.KSP.Type.PREONLY)
+covksp.setFromOptions()
+covksp.setUp()
+print(PETSc.Options().getAll())
+
+x = covmat.createVecRight()
+b = covmat.createVecLeft()
+
+b.array_w[:] = np.random.random_sample(b.array_w.shape)
+print(f'{b.norm() = }')
+covksp.solve(b, x)
+print(f'{x.norm() = }')

advection/advection_wc4dvar_pyadj.py

@@ -12,10 +12,10 @@
 continue_annotation()
 
 # background functional
-J = norm2(B)(control[0] - background)
+J = covariance_norm(control[0] - background, B)
 
 # initial observation functional
-J += norm2(R)(observation_error(0)(control[0]))
+J += covariance_norm(observation_error(0)(control[0]), R)
 
 nstep = 0
 for i in range(1, len(control)):
@@ -33,10 +33,10 @@
         control[i].assign(qn)
 
     # model error functional
-    J += norm2(Q)(qn - control[i])
+    J += covariance_norm(qn - control[i], Q)
 
     # observation functional
-    J += norm2(R)(observation_error(i)(control[i]))
+    J += covariance_norm(observation_error(i)(control[i]), R)
 
 pause_annotation()