rename aao_rf -> 4dvar_rf

firedrakeproject · Dec 23, 2024 · 93fdc83 · 93fdc83
1 parent c22531f
commit 93fdc83
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 14 deletions.
diff --git a/firedrake/adjoint/__init__.py b/firedrake/adjoint/__init__.py
@@ -38,7 +38,7 @@
 from firedrake.adjoint.ufl_constraints import UFLInequalityConstraint, \
     UFLEqualityConstraint  # noqa F401
 from firedrake.adjoint.ensemble_reduced_functional import EnsembleReducedFunctional  # noqa F401
-from firedrake.adjoint.all_at_once_reduced_functional import AllAtOnceReducedFunctional  # noqa F401
+from firedrake.adjoint.fourdvar_reduced_functional import FourDVarReducedFunctional  # noqa F401
 import numpy_adjoint  # noqa F401
 import firedrake.ufl_expr
 import types

diff --git a/...adjoint/all_at_once_reduced_functional.py → ...ke/adjoint/fourdvar_reduced_functional.py b/...adjoint/all_at_once_reduced_functional.py → ...ke/adjoint/fourdvar_reduced_functional.py
@@ -9,7 +9,7 @@
 from contextlib import contextmanager
 from mpi4py import MPI
 
-__all__ = ['AllAtOnceReducedFunctional']
+__all__ = ['FourDVarReducedFunctional']
 
 
 # @set_working_tape()  # ends up using old_tape = None because evaluates when imported - need separate decorator
@@ -54,7 +54,7 @@ def sc_passthrough(func):
     is instantiated then passes args/kwargs through to the
     corresponding strong_reduced_functional method.
 
-    If using weak constraint, returns the AllAtOnceReducedFunctional
+    If using weak constraint, returns the FourDVarReducedFunctional
     method definition.
     """
     @wraps(func)
@@ -93,7 +93,7 @@ def _intermediate_options(final_options):
     }
 
 
-class AllAtOnceReducedFunctional(ReducedFunctional):
+class FourDVarReducedFunctional(ReducedFunctional):
     """ReducedFunctional for 4DVar data assimilation.
 
     Creates either the strong constraint or weak constraint system
@@ -586,7 +586,7 @@ def recording_stages(self, sequential=True, nstages=None, **stage_kwargs):
 
 class ObservationStageSequence:
     def __init__(self, controls: Control,
-                 aaorf: AllAtOnceReducedFunctional,
+                 aaorf: FourDVarReducedFunctional,
                  global_index: int,
                  observation_index: int,
                  stage_kwargs: dict = None,
@@ -663,12 +663,12 @@ class StrongObservationStage:
     ----------
 
     aaorf
-        The strong constraint AllAtOnceReducedFunctional.
+        The strong constraint FourDVarReducedFunctional.
 
     """
 
     def __init__(self, control: OverloadedType,
-                 aaorf: AllAtOnceReducedFunctional,
+                 aaorf: FourDVarReducedFunctional,
                  index: Optional[int] = None,
                  observation_index: Optional[int] = None):
         self.aaorf = aaorf

diff --git a/tests/firedrake/regression/test_4dvar_reduced_functional.py b/tests/firedrake/regression/test_4dvar_reduced_functional.py
@@ -4,7 +4,7 @@
 from firedrake.adjoint import (
     continue_annotation, pause_annotation, stop_annotating,
     set_working_tape, get_working_tape, Control, taylor_test,
-    ReducedFunctional, AllAtOnceReducedFunctional)
+    ReducedFunctional, FourDVarReducedFunctional)
 
 
 @pytest.fixture(autouse=True)
@@ -207,7 +207,7 @@ def strong_fdvar_pyadjoint(V):
 
 
 def strong_fdvar_firedrake(V):
-    """Build an AllAtOnceReducedFunctional for the strong constraint 4DVar system"""
+    """Build an FourDVarReducedFunctional for the strong constraint 4DVar system"""
     qn, qn1, stepper = timestepper(V)
 
     # prior data
@@ -223,7 +223,7 @@ def strong_fdvar_firedrake(V):
     # create 4DVar reduced functional and record
     # background and initial observation functionals
 
-    Jhat = AllAtOnceReducedFunctional(
+    Jhat = FourDVarReducedFunctional(
         Control(control),
         background_iprod=prodB,
         observation_iprod=prodR,
@@ -243,9 +243,8 @@ def strong_fdvar_firedrake(V):
                 stepper.solve()
                 qn.assign(qn1)
 
-            obs_index = stage.index + 1
-
             # take observation
+            obs_index = stage.observation_index
             stage.set_observation(qn, obs_errors(obs_index),
                                   observation_iprod=prodR)
 
@@ -312,7 +311,7 @@ def weak_fdvar_pyadjoint(V):
 
 
 def weak_fdvar_firedrake(V, ensemble):
-    """Build an AllAtOnceReducedFunctional for the weak constraint 4DVar system"""
+    """Build an FourDVarReducedFunctional for the weak constraint 4DVar system"""
     qn, qn1, stepper = timestepper(V)
 
     # One control for each observation time
@@ -338,7 +337,7 @@ def weak_fdvar_firedrake(V, ensemble):
     # create 4DVar reduced functional and record
     # background and initial observation functionals
 
-    Jhat = AllAtOnceReducedFunctional(
+    Jhat = FourDVarReducedFunctional(
         Control(control),
         background_iprod=prodB,
         observation_iprod=prodR,