Merge branch 'master' into allatoncereducedfunctional

.github/workflows/docker_reuse.yml

@@ -31,15 +31,15 @@ jobs:
       - name: Check out the repo
         uses: actions/checkout@v4
       - name: Log in to Docker Hub
-        uses: docker/login-action@v2
+        uses: docker/login-action@v3
         with:
           username: ${{ secrets.DOCKERHUB_USER }}
           password: ${{ secrets.DOCKERHUB_TOKEN }}
       - name: Set up Docker Buildx
         id: buildx
-        uses: docker/setup-buildx-action@v2
+        uses: docker/setup-buildx-action@v3
       - name: Build and push ${{ inputs.target }}
-        uses: docker/build-push-action@v4
+        uses: docker/build-push-action@v6
         with:
           push: true
           no-cache: true

.github/workflows/lint.yml

@@ -14,7 +14,7 @@ jobs:
     steps:
       - uses: actions/checkout@v4
       - name: Setup python
-        uses: actions/setup-python@v4
+        uses: actions/setup-python@v5
         with:
           python-version: 3.8
       - name: Setup flake8 annotations

.github/workflows/zenodo-canary.yml

@@ -17,7 +17,7 @@ jobs:
         with:
           fetch-depth: 1
       - name: Setup Python
-        uses: actions/setup-python@v4
+        uses: actions/setup-python@v5
         with:
           python-version: 3.12
       - name: Install deps

demos/full_waveform_inversion/full_waveform_inversion.py.rst

@@ -99,19 +99,29 @@ The source number is defined with the ``Ensemble.ensemble_comm`` rank::
     source_number = my_ensemble.ensemble_comm.rank
 
 In this example, we consider a two-dimensional square domain with a side length of 1.0 km. The mesh is
-built over the ``my_ensemble.comm`` (spatial) communicator::
-    
-    Lx, Lz = 1.0, 1.0
-    mesh = UnitSquareMesh(80, 80, comm=my_ensemble.comm)
+built over the ``my_ensemble.comm`` (spatial) communicator.
+
+::
+
+    import os
+    if os.getenv("FIREDRAKE_CI_TESTS") == "1": 
+        # Setup for a faster test execution.
+        dt = 0.03  # time step in seconds
+        final_time = 0.6  # final time in seconds
+        nx, ny = 15, 15
+    else:
+        dt = 0.002  # time step in seconds
+        final_time = 1.0  # final time in seconds
+        nx, ny = 80, 80
 
-The basic input for the FWI problem are defined as follows::
+    mesh = UnitSquareMesh(nx, ny, comm=my_ensemble.comm)
+
+The frequency of the Ricker wavelet, the source and receiver locations are defined as follows::
 
     import numpy as np
+    frequency_peak = 7.0  # The dominant frequency of the Ricker wavelet in Hz.
     source_locations = np.linspace((0.3, 0.1), (0.7, 0.1), num_sources)
     receiver_locations = np.linspace((0.2, 0.9), (0.8, 0.9), 20)
-    dt = 0.002  # time step in seconds
-    final_time = 1.0  # final time in seconds
-    frequency_peak = 7.0  # The dominant frequency of the Ricker wavelet in Hz.
 
 Sources and receivers locations are illustrated in the following figure:
 
@@ -251,6 +261,7 @@ To have the step 4, we need first to tape the forward problem. That is done by c
 
     from firedrake.adjoint import *
     continue_annotation()
+    get_working_tape().progress_bar = ProgressBar
 
 **Steps 2-3**: Solve the wave equation and compute the functional::
 

demos/multigrid/geometric_multigrid.py.rst

@@ -191,7 +191,7 @@ bilinear form to the solver ourselves: ::
       "fieldsplit_0_pc_type": "mg",
       "fieldsplit_1_ksp_type": "preonly",
       "fieldsplit_1_pc_type": "python",
-      "fieldsplit_1_pc_python_type": "__main__.Mass",
+      "fieldsplit_1_pc_python_type": "geometric_multigrid.Mass",
       "fieldsplit_1_aux_pc_type": "bjacobi",
       "fieldsplit_1_aux_sub_pc_type": "icc",
   }
@@ -227,7 +227,7 @@ approximations.
         "mg_coarse_fieldsplit_0_pc_type": "lu",
         "mg_coarse_fieldsplit_1_ksp_type": "preonly",
         "mg_coarse_fieldsplit_1_pc_type": "python",
-        "mg_coarse_fieldsplit_1_pc_python_type": "__main__.Mass",
+        "mg_coarse_fieldsplit_1_pc_python_type": "geometric_multigrid.Mass",
         "mg_coarse_fieldsplit_1_aux_pc_type": "cholesky",
         "mg_levels_ksp_type": "richardson",
         "mg_levels_ksp_max_it": 1,
@@ -245,7 +245,7 @@ approximations.
         "mg_levels_fieldsplit_1_ksp_richardson_self_scale": None,
         "mg_levels_fieldsplit_1_ksp_max_it": 3,
         "mg_levels_fieldsplit_1_pc_type": "python",
-        "mg_levels_fieldsplit_1_pc_python_type": "__main__.Mass",
+        "mg_levels_fieldsplit_1_pc_python_type": "geometric_multigrid.Mass",
         "mg_levels_fieldsplit_1_aux_pc_type": "bjacobi",
         "mg_levels_fieldsplit_1_aux_sub_pc_type": "icc",
   }

docs/source/advanced_tut.rst

@@ -23,4 +23,4 @@ element systems.
    A pressure-convection-diffusion preconditioner for the Navier-Stokes equations.</demos/navier_stokes.py>
    Rayleigh-Benard convection.<demos/rayleigh-benard.py>
    Netgen support.<demos/netgen_mesh.py>
-   Full-waveform inversion: Full-waveform inversion: spatial and wave sources parallelism.<demos/full_waveform_inversion.py>
+   Full-waveform inversion: spatial and wave sources parallelism.<demos/full_waveform_inversion.py>

docs/source/download.rst

@@ -107,7 +107,7 @@ packages can be installed into an existing Firedrake installation using
 System requirements
 -------------------
 
-Firedrake requires Python 3.9 to 3.13. The installation script is
+Firedrake requires Python 3.10 to 3.13. The installation script is
 tested by CI on Ubuntu 24.04 LTS. On Ubuntu 22.04 or later, the system
 installed Python 3 is supported. On MacOS, the homebrew_ installed
 Python 3 is supported::
@@ -126,7 +126,7 @@ they have the system dependencies:
 * A Fortran compiler (for PETSc)
 * Blas and Lapack
 * Git, Mercurial
-* Python version 3.9-3.13
+* Python version 3.10-3.13
 * The Python headers
 * autoconf, automake, libtool
 * CMake

docs/source/install-debug.dot

@@ -8,15 +8,15 @@ digraph triage {
     venv_activated [label="venv activated?"];
     install_script_up_to_date [label="Install script\nup to date?"];
     using_anaconda [label="Using\nAnaconda?"];
-    python_version [label="Python <3.9?"];
+    python_version [label="Python <3.10?"];
     using_macos [label="Using\nMacOS?"];
     using_homebrew [label="Using\nHomebrew?"];
     url_error [label="URL Error with SSL\ncertificate failure?"];
     which_python [label="<which python3> points\nat <$(brew --prefix)/bin/python3>?"];
 
     activate_venv [label="Activate the\nvenv first."];
     uninstall_anaconda [label="Deactivate\nAnaconda."];
-    update_python [label="Get Python 3.9-3.13"];
+    update_python [label="Get Python 3.10-3.13"];
     update_install_script [label="Fetch new\ninstall script"];
     get_homebrew [label="Use Homebrew."];
     brew_doctor [label="brew doctor"];

firedrake/adjoint_utils/blocks/solving.py

@@ -2,8 +2,9 @@
 import ufl
 from ufl import replace
 from ufl.formatting.ufl2unicode import ufl2unicode
+from enum import Enum
 
-from pyadjoint import Block, stop_annotating
+from pyadjoint import Block, stop_annotating, get_working_tape
 from pyadjoint.enlisting import Enlist
 import firedrake
 from firedrake.adjoint_utils.checkpointing import maybe_disk_checkpoint
@@ -24,6 +25,12 @@ def extract_subfunction(u, V):
         return u
 
 
+class Solver(Enum):
+    """Enum for solver types."""
+    FORWARD = 0
+    ADJOINT = 1
+
+
 class GenericSolveBlock(Block):
     pop_kwargs_keys = ["adj_cb", "adj_bdy_cb", "adj2_cb", "adj2_bdy_cb",
                        "forward_args", "forward_kwargs", "adj_args",
@@ -206,15 +213,17 @@ def _assemble_and_solve_adj_eq(self, dFdu_adj_form, dJdu, compute_bdy):
 
         adj_sol_bdy = None
         if compute_bdy:
-            adj_sol_bdy = firedrake.Function(
-                self.function_space.dual(),
-                dJdu_copy.dat - firedrake.assemble(
-                    firedrake.action(dFdu_adj_form, adj_sol)
-                ).dat
-            )
+            adj_sol_bdy = self._compute_adj_bdy(
+                adj_sol, adj_sol_bdy, dFdu_adj_form, dJdu_copy)
 
         return adj_sol, adj_sol_bdy
 
+    def _compute_adj_bdy(self, adj_sol, adj_sol_bdy, dFdu_adj_form, dJdu):
+        adj_sol_bdy = firedrake.Function(
+            self.function_space.dual(), dJdu.dat - firedrake.assemble(
+                firedrake.action(dFdu_adj_form, adj_sol)).dat)
+        return adj_sol_bdy
+
     def evaluate_adj_component(self, inputs, adj_inputs, block_variable, idx,
                                prepared=None):
         if not self.linear and self.func == block_variable.output:
@@ -604,12 +613,11 @@ def _init_solver_parameters(self, args, kwargs):
 
 
 class NonlinearVariationalSolveBlock(GenericSolveBlock):
-    def __init__(self, equation, func, bcs, adj_F, adj_cache, problem_J,
+    def __init__(self, equation, func, bcs, adj_cache, problem_J,
                  solver_params, solver_kwargs, **kwargs):
         lhs = equation.lhs
         rhs = equation.rhs
 
-        self.adj_F = adj_F
         self._adj_cache = adj_cache
         self._dFdm_cache = adj_cache.setdefault("dFdm_cache", {})
         self.problem_J = problem_J
@@ -626,15 +634,62 @@ def _init_solver_parameters(self, args, kwargs):
         super()._init_solver_parameters(args, kwargs)
         solve_init_params(self, args, kwargs, varform=True)
 
+    def recompute_component(self, inputs, block_variable, idx, prepared):
+        tape = get_working_tape()
+        if self._ad_solvers["recompute_count"] == tape.recompute_count - 1:
+            # Update how many times the block has been recomputed.
+            self._ad_solvers["recompute_count"] = tape.recompute_count
+            if self._ad_solvers["forward_nlvs"]._problem._constant_jacobian:
+                self._ad_solvers["forward_nlvs"].invalidate_jacobian()
+                self._ad_solvers["update_adjoint"] = True
+        return super().recompute_component(inputs, block_variable, idx, prepared)
+
     def _forward_solve(self, lhs, rhs, func, bcs, **kwargs):
-        self._ad_nlvs_replace_forms()
-        self._ad_nlvs.parameters.update(self.solver_params)
-        self._ad_nlvs.solve()
-        func.assign(self._ad_nlvs._problem.u)
+        self._ad_solver_replace_forms()
+        self._ad_solvers["forward_nlvs"].parameters.update(self.solver_params)
+        self._ad_solvers["forward_nlvs"].solve()
+        func.assign(self._ad_solvers["forward_nlvs"]._problem.u)
         return func
 
-    def _ad_assign_map(self, form):
-        count_map = self._ad_nlvs._problem._ad_count_map
+    def _adjoint_solve(self, dJdu, compute_bdy):
+        dJdu_copy = dJdu.copy()
+        # Homogenize and apply boundary conditions on adj_dFdu and dJdu.
+        bcs = self._homogenize_bcs()
+        for bc in bcs:
+            bc.apply(dJdu)
+
+        if (
+            self._ad_solvers["forward_nlvs"]._problem._constant_jacobian
+            and self._ad_solvers["update_adjoint"]
+        ):
+            # Update left hand side of the adjoint equation.
+            self._ad_solver_replace_forms(Solver.ADJOINT)
+            self._ad_solvers["adjoint_lvs"].invalidate_jacobian()
+            self._ad_solvers["update_adjoint"] = False
+        elif not self._ad_solvers["forward_nlvs"]._problem._constant_jacobian:
+            # Update left hand side of the adjoint equation.
+            self._ad_solver_replace_forms(Solver.ADJOINT)
+
+        # Update the right hand side of the adjoint equation.
+        # problem.F._component[1] is the right hand side of the adjoint.
+        self._ad_solvers["adjoint_lvs"]._problem.F._components[1].assign(dJdu)
+
+        # Solve the adjoint linear variational solver.
+        self._ad_solvers["adjoint_lvs"].solve()
+        u_sol = self._ad_solvers["adjoint_lvs"]._problem.u
+
+        adj_sol_bdy = None
+        if compute_bdy:
+            jac_adj = self._ad_solvers["adjoint_lvs"]._problem.J
+            adj_sol_bdy = self._compute_adj_bdy(
+                u_sol, adj_sol_bdy, jac_adj, dJdu_copy)
+        return u_sol, adj_sol_bdy
+
+    def _ad_assign_map(self, form, solver):
+        if solver == Solver.FORWARD:
+            count_map = self._ad_solvers["forward_nlvs"]._problem._ad_count_map
+        else:
+            count_map = self._ad_solvers["adjoint_lvs"]._problem._ad_count_map
         assign_map = {}
         form_ad_count_map = dict((count_map[coeff], coeff)
                                  for coeff in form.coefficients())
@@ -647,55 +702,46 @@ def _ad_assign_map(self, form):
                 if coeff_count in form_ad_count_map:
                     assign_map[form_ad_count_map[coeff_count]] = \
                         block_variable.saved_output
+
+        if (
+            solver == Solver.ADJOINT
+            and not self._ad_solvers["forward_nlvs"]._problem._constant_jacobian
+        ):
+            block_variable = self.get_outputs()[0]
+            coeff_count = block_variable.output.count()
+            if coeff_count in form_ad_count_map:
+                assign_map[form_ad_count_map[coeff_count]] = \
+                    block_variable.saved_output
         return assign_map
 
-    def _ad_assign_coefficients(self, form):
-        assign_map = self._ad_assign_map(form)
+    def _ad_assign_coefficients(self, form, solver):
+        assign_map = self._ad_assign_map(form, solver)
         for coeff, value in assign_map.items():
             coeff.assign(value)
 
-    def _ad_nlvs_replace_forms(self):
-        problem = self._ad_nlvs._problem
-        self._ad_assign_coefficients(problem.F)
-        self._ad_assign_coefficients(problem.J)
-
-    def _assemble_dFdu_adj(self, dFdu_adj_form, **kwargs):
-        if "dFdu_adj" in self._adj_cache:
-            dFdu = self._adj_cache["dFdu_adj"]
+    def _ad_solver_replace_forms(self, solver=Solver.FORWARD):
+        if solver == Solver.FORWARD:
+            problem = self._ad_solvers["forward_nlvs"]._problem
+            self._ad_assign_coefficients(problem.F, solver)
+            self._ad_assign_coefficients(problem.J, solver)
         else:
-            dFdu = super()._assemble_dFdu_adj(dFdu_adj_form, **kwargs)
-            if self._ad_nlvs._problem._constant_jacobian:
-                self._adj_cache["dFdu_adj"] = dFdu
-        return dFdu
+            self._ad_assign_coefficients(
+                self._ad_solvers["adjoint_lvs"]._problem.J, solver)
 
     def prepare_evaluate_adj(self, inputs, adj_inputs, relevant_dependencies):
-        dJdu = adj_inputs[0]
-
-        F_form = self._create_F_form()
-
-        dFdu_form = self.adj_F
-        dJdu = dJdu.copy()
-
-        # Replace the form coefficients with checkpointed values.
-        replace_map = self._replace_map(dFdu_form)
-        replace_map[self.func] = self.get_outputs()[0].saved_output
-        dFdu_form = replace(dFdu_form, replace_map)
-
         compute_bdy = self._should_compute_boundary_adjoint(
             relevant_dependencies
         )
-        adj_sol, adj_sol_bdy = self._assemble_and_solve_adj_eq(
-            dFdu_form, dJdu, compute_bdy
-        )
+        adj_sol, adj_sol_bdy = self._adjoint_solve(adj_inputs[0], compute_bdy)
         self.adj_state = adj_sol
         if self.adj_cb is not None:
             self.adj_cb(adj_sol)
         if self.adj_bdy_cb is not None and compute_bdy:
             self.adj_bdy_cb(adj_sol_bdy)
 
         r = {}
-        r["form"] = F_form
-        r["adj_sol"] = adj_sol
+        r["form"] = self._create_F_form()
+        r["adj_sol"] = self.adj_state
         r["adj_sol_bdy"] = adj_sol_bdy
         return r