Add cubic and RBF interpolation for adjoint inversion

examples/headland_inversion/Makefile

@@ -3,13 +3,16 @@ all: invert plot
 CASE        ?= Uniform # Regions IndependentPointsScheme GradientReg HessianReg
 STATIONS    = stationA stationB stationC stationD stationE stationF stationG
 PARALLEL    = 1
+IP_INTERPOLATION ?= linear # linear cubic rbf
+IP_RBF_KERNEL ?= thin_plate_spline
+IP_RBF_SMOOTHING ?= 0.0
 
 invert:
-	mpiexec --use-hwthread-cpus -np $(PARALLEL) python inverse_problem.py --case $(CASE) --no-consistency-test --no-taylor-test
+	mpiexec --use-hwthread-cpus -np $(PARALLEL) python inverse_problem.py --case $(CASE) --no-consistency-test --no-taylor-test --ip-interpolation $(IP_INTERPOLATION) --ip-rbf-kernel $(IP_RBF_KERNEL) --ip-rbf-smoothing $(IP_RBF_SMOOTHING)
 
 plot:
 	for station in $(STATIONS); do \
-		python3 plot_velocity_progress.py -s $$station --case $(CASE); \
+		python3 plot_velocity_progress.py -s $$station --case $(CASE) --ip-interpolation $(IP_INTERPOLATION); \
 	done; \
 
 clean:

examples/headland_inversion/README.md

@@ -145,11 +145,20 @@ The independent point scheme approach works in the same way as the region-based
 function which tells us how the Manning field changes with respect to our input independent point values. The only thing 
 we need to do is change the masks we generate. 
 
-Instead of masks with 0/1 values, we have masks which describe the contribution of each point to the rest of the 
-domain. Note that this will only work for linear interpolation, as we cannot generate static coefficients for non-linear
-mappings (RBF, quadratic, cubic etc.). In those cases, we would need to 'annotate' the interpolation functions for 
-`pyadjoint` to track the gradient through. We can generate a mapping in the same way to force a smooth surface, but it 
-would not be true RBF/quadratic/cubic interpolation.
+Instead of masks with 0/1 values, we have a mapping which describes the contribution of each point to the rest of the 
+domain. Linear interpolation uses fixed Firedrake basis fields. Cubic and radial-basis interpolation use Thetis'
+independent-point external operator, which evaluates the SciPy interpolation and provides the adjoint vector-Jacobian
+product required by `pyadjoint`.
+
+```sh
+source ~/firedrake/bin/activate
+make invert CASE=IndependentPointsScheme IP_INTERPOLATION=linear
+make invert CASE=IndependentPointsScheme IP_INTERPOLATION=cubic
+make invert CASE=IndependentPointsScheme IP_INTERPOLATION=rbf
+```
+
+The RBF mode uses SciPy's `RBFInterpolator`. The default kernel is `thin_plate_spline` with zero smoothing; these can be
+changed through `IP_RBF_KERNEL` and `IP_RBF_SMOOTHING`.
 
 ## Post-processing
 

examples/headland_inversion/inverse_problem.py

@@ -3,7 +3,6 @@
 from firedrake import *
 from firedrake.adjoint import *
 from model_config import construct_solver
-from scipy.interpolate import LinearNDInterpolator, NearestNDInterpolator
 import h5py
 import numpy as np
 import argparse
@@ -27,10 +26,21 @@
                     help='Skip consistency test')
 parser.add_argument('--no-taylor-test', action='store_true',
                     help='Skip Taylor test')
+parser.add_argument('--ip-interpolation',
+                    help='Interpolation method for the independent point scheme',
+                    choices=['linear', 'cubic', 'rbf'],
+                    default='linear')
+parser.add_argument('--ip-rbf-kernel',
+                    help='SciPy RBFInterpolator kernel for the independent point scheme',
+                    default='thin_plate_spline')
+parser.add_argument('--ip-rbf-smoothing', type=float,
+                    help='SciPy RBFInterpolator smoothing parameter for the independent point scheme',
+                    default=0.0)
 args = parser.parse_args()
 do_consistency_test = not args.no_consistency_test
 do_taylor_test = not args.no_taylor_test
 no_exports = os.getenv('THETIS_REGRESSION_TEST') is not None
+ip_interpolation = args.ip_interpolation
 
 case_to_output_dir = {
     'Uniform': 'uniform_friction',
@@ -44,6 +54,8 @@
 pwd = os.path.abspath(os.path.dirname(__file__))
 output_dir_forward = os.path.join(pwd, 'outputs', 'outputs_forward')
 output_dir_invert = os.path.join(pwd, 'outputs', 'outputs_inverse', case_to_output_dir[selected_case])
+if selected_case == 'IndependentPointsScheme' and ip_interpolation != 'linear':
+    output_dir_invert += f'_{ip_interpolation}'
 
 continue_annotation()
 
@@ -71,7 +83,7 @@
 
 local_N = coordinates.shape[0]
 N = mesh2d.comm.allreduce(local_N, MPI.SUM)  # allreduce sums the local numbers to get the total number of coordinates
-masks, M, m_true = None, 0, []
+masks, point_mapping, M, m_true = None, None, 0, []
 
 # Create a FunctionSpace on the mesh (corresponds to Manning)
 V = get_functionspace(mesh2d, 'CG', 1)
@@ -104,6 +116,7 @@
     for m_, mask_ in zip(m_true, masks):
         manning_2d += m_ * mask_
 elif selected_case == 'IndependentPointsScheme':
+    print_output(f'Using {ip_interpolation} independent point interpolation')
     # Define our values for n
     points = np.array([
         [0, 0], [0, 0.5], [0, 1], [0.5, 0], [1, 0], [1, 0.5], [1, 1], [0.5, 0.5],
@@ -116,24 +129,11 @@
     m_true = [domain_constant(0.03 - 0.0005 * i, mesh2d) for i in range(len(points))]
     M = len(m_true)
 
-    linear_interpolator = LinearNDInterpolator(points, np.eye(len(points)))
-    nearest_interpolator = NearestNDInterpolator(points, np.eye(len(points)))
-
-    # Apply the interpolators to the mesh coordinates
-    linear_coefficients = linear_interpolator(coordinates)
-    nan_mask = np.isnan(linear_coefficients).any(axis=1)
-    linear_coefficients[nan_mask] = nearest_interpolator(coordinates[nan_mask])
-
-    # Create Function objects to store the coefficients
-    masks = [Function(V) for _ in range(len(points))]
-
-    # Assign the linear coefficients to the masks
-    for i, mask in enumerate(masks):
-        mask.dat.data[:] = linear_coefficients[:, i]
-
-    manning_2d.assign(0)
-    for m_, mask_ in zip(m_true, masks):
-        manning_2d += m_ * mask_
+    point_mapping = inversion_tools.IndependentPointControlMapping(
+        V, points, method=ip_interpolation,
+        rbf_kernel=args.ip_rbf_kernel, rbf_smoothing=args.ip_rbf_smoothing)
+    masks = point_mapping.masks
+    point_mapping.assign(manning_2d, m_true)
 elif selected_case in ('GradientReg', 'HessianReg'):
     pass
 else:
@@ -245,12 +245,16 @@
     manning_const = domain_constant(0.03, mesh2d, name='Manning')
     manning_2d.project(manning_const)
     inv_manager.add_control(manning_const)
-elif selected_case in ('Regions', 'IndependentPointsScheme'):
+elif selected_case == 'Regions':
     m_values = [domain_constant(0.03 - 0.0005 * i, mesh2d) for i in range(M)]
     manning_2d.assign(0)
     for m_, mask_ in zip(m_values, masks):
         manning_2d += m_ * mask_
     inv_manager.add_control(m_values, mappings=masks)
+elif selected_case == 'IndependentPointsScheme':
+    m_values = [domain_constant(0.03 - 0.0005 * i, mesh2d) for i in range(M)]
+    point_mapping.assign(manning_2d, m_values)
+    inv_manager.add_control(m_values, mappings=point_mapping)
 else:
     manning_2d.assign(0.04)
     inv_manager.add_control(manning_2d)
@@ -295,14 +299,18 @@
         manning_2d.assign(domain_constant(inv_manager.m_list[-1], mesh2d))
         if not no_exports:
             VTKFile(os.path.join(options.output_directory, 'manning_optimised.pvd')).write(manning_2d)
-    elif selected_case == 'Regions' or selected_case == 'IndependentPointsScheme':
+    elif selected_case == 'Regions':
         P1_2d = get_functionspace(mesh2d, 'CG', 1)
         manning_2d = Function(P1_2d, name='manning2d')
         manning_2d.assign(0)
         for m_, mask_ in zip(inv_manager.m_list, masks):
             manning_2d += m_ * mask_
         if not no_exports:
             VTKFile(os.path.join(options.output_directory, 'manning_optimised.pvd')).write(manning_2d)
+    elif selected_case == 'IndependentPointsScheme':
+        manning_2d = point_mapping.project(inv_manager.m_list, name='manning2d')
+        if not no_exports:
+            VTKFile(os.path.join(options.output_directory, 'manning_optimised.pvd')).write(manning_2d)
     else:
         name = cc.name()
         oc.rename(name)

examples/headland_inversion/plot_velocity_progress.py

@@ -20,6 +20,10 @@
     choices=['Uniform', 'Regions', 'IndependentPointsScheme', 'GradientReg', 'HessianReg'],
     default=['IndependentPointsScheme'],
 )
+parser.add_argument('--ip-interpolation',
+                    help='Interpolation method for the independent point scheme',
+                    choices=['linear', 'cubic', 'rbf'],
+                    default='linear')
 args = parser.parse_args()
 station_names = sorted(args.station)
 station_str = '-'.join(station_names)
@@ -35,6 +39,8 @@
 selected_case = args.case[0]
 output_dir_forward = os.path.join('outputs', 'outputs_forward')
 output_dir_invert = os.path.join('outputs', 'outputs_inverse', case_to_output_dir[selected_case])
+if selected_case == 'IndependentPointsScheme' and args.ip_interpolation != 'linear':
+    output_dir_invert += f'_{args.ip_interpolation}'
 
 # --- Loop through stations ---
 for sta in station_names:

test_adjoint/test_independent_point_mapping.py

@@ -0,0 +1,54 @@
+import numpy
+import pytest
+
+from firedrake import *
+from firedrake.adjoint import *
+from pyadjoint import get_working_tape
+
+from thetis.inversion_tools import IndependentPointControlMapping
+
+
+def make_real_controls(mesh, values, name):
+    R = FunctionSpace(mesh, 'R', 0)
+    controls = []
+    for i, value in enumerate(values):
+        control = Function(R, name=f'{name}_{i:02d}')
+        control.assign(float(value))
+        controls.append(control)
+    return controls
+
+
+@pytest.mark.parametrize('method', ['linear', 'cubic', 'rbf'])
+def test_independent_point_mapping_gradient(method):
+    get_working_tape().clear_tape()
+    continue_annotation()
+
+    mesh = UnitSquareMesh(8, 8)
+    V = FunctionSpace(mesh, 'CG', 1)
+    x, y = SpatialCoordinate(mesh)
+    target = Function(V).interpolate(0.022 + 0.004*sin(pi*x)*cos(pi*y) + 0.002*x)
+
+    points = numpy.array([
+        [0.0, 0.0],
+        [1.0, 0.0],
+        [0.0, 1.0],
+        [1.0, 1.0],
+        [0.2, 0.25],
+        [0.75, 0.2],
+        [0.25, 0.75],
+        [0.8, 0.8],
+        [0.5, 0.5],
+    ])
+    values = numpy.array([0.020, 0.024, 0.023, 0.021, 0.026, 0.019, 0.025, 0.022, 0.027])
+    direction_values = numpy.array([0.30, -0.20, 0.15, -0.10, 0.25, -0.35, 0.40, -0.30, 0.20])
+
+    controls = make_real_controls(mesh, values, f'{method}_control')
+    directions = make_real_controls(mesh, direction_values, f'{method}_direction')
+    mapping = IndependentPointControlMapping(V, points, method=method)
+
+    field = Function(V, name=f'{method}_field')
+    mapping.assign(field, controls)
+    J = assemble(0.5*(field - target)**2*dx)
+    reduced_functional = ReducedFunctional(J, [Control(c) for c in controls])
+
+    assert taylor_test(reduced_functional, controls, directions) > 1.9