Adding ensemble support in forward runs and inversion by Brazovskis · Pull Request #464 · thetisproject/thetis

Brazovskis · 2026-07-29T09:40:14Z

This PR adds ensemble support to Thetis for both forward simulations and adjoint-based optimisation. At the library level, thetis/exporter.py now supports HDF5 and VTK exporting in ensembles, and thetis/inversion_tools.py now supports ensemble-aware reduced functionals. The headland_inversion example has been updated to demonstrate ensemble usage in both forward runs and inversion.

Changes

Added comm= arguments to I/O-related functions creating output directories or writing files
Made thetis/inversion_tools.py ensemble-aware and added support for EnsembleReducedFunctional to handle cost function evaluation across multiple ensemble members
Added ensemble functionality to forward_run.py and inverse_problem.py in the headland_inversion example
Updated related plotting scripts and README in headland_inversion to accommodate ensemble workflow and document usage

Why work with ensembles?

Ensemble forward runs allow for running Thetis models with multiple permutations of input parameters in parallel, improving efficiency when problem no longer benefits from spatial parallelism
Ensemble-based adjoint optimisation allows for optimising an input field across a number of Thetis model instances that may have different input parameters otherwise. E.g., fitting a Manning field to minimise model discrepancies across multiple observation windows

…ional

stephankramer

Oh I see this is still in "draft" mode, so maybe my review is a bit premature. In any case hope the comments are helpful.
This is great and a nice addition to functionality. I would say the main thing still missing is hooking it up in the test suite. This sort of stuff is easily broken by other devs that don't think about the ensemble case.

stephankramer · 2026-07-29T10:38:48Z

 x, y = coordinates[:, 0], coordinates[:, 1]
-lx = mesh2d.comm.allreduce(np.max(x), MPI.MAX)
-ly = mesh2d.comm.allreduce(np.max(y), MPI.MAX)
+lx = mesh2d.comm.allreduce(numpy.max(x), MPI.MAX)


Ah so this indicates forward_run.py is currently broken? So it seems previously it relied on thetis or firedrake leaking np into its namespace, and now you're relying on thetis leaking numpy instead. Neither is good. Should explicitly "import numpy" on top, but I prefer if you "import numpy as np" and revert these changes (avoiding unnecessary changes)

stephankramer · 2026-07-29T10:55:10Z

                f.store(function)
        else:
-            with CheckpointFile(filename, 'w') as f:
+            with CheckpointFile(filename, 'w', comm=function.function_space().mesh().comm) as f:


the line above it uses function.comm - does that not work? Otherwise this line is a bit long, so I would comm = .... on a separate line before and then use comm=comm here

stephankramer · 2026-07-29T10:56:03Z

        else:
-            with CheckpointFile(filename, 'w') as f:
+            with CheckpointFile(filename, 'w', comm=function.function_space().mesh().comm) as f:
                mesh = function.function_space().mesh()


Oh alternatively move this before the with CheckpointFile block and use mesh.comm

stephankramer · 2026-07-29T10:57:03Z

-        super(VTKExporter, self).__init__(filename, outputdir, next_export_ix,
-                                          verbose)
+        self.comm = fs_visu.mesh().comm
+        ExporterBase.__init__(self, filename, outputdir, next_export_ix, verbose, comm=self.comm)


Why are you not using super().__init__(...) - and why did you nuke the doc-string?

stephankramer · 2026-07-29T10:57:15Z

-        super(HDF5Exporter, self).__init__(filename_prefix, outputdir,
-                                           next_export_ix, verbose)
+        mesh_comm = function_space.mesh().comm
+        ExporterBase.__init__(self, filename_prefix, outputdir, next_export_ix, verbose, comm=mesh_comm)


stephankramer · 2026-07-29T11:08:26Z

+        if self.no_exports:
+            return False
+        return True


The logic here seems different than the name suggests which suggest only the root exports. I'm guessing you later realized station export should be collective? But then it would be much clearer if you just revert to if not self.no_exports:

stephankramer · 2026-07-29T13:06:43Z

-        if not self.no_exports:
+        J = self.reduced_functional(self.control_coeff_list)
+        self._update_objective_from_evaluation(J)
+        self.set_initial_state(self.J if self.ensemble is not None else J,


Can this not be self.J in both cases? Note that before J was explicitly converted to float, but now it no longer is (if ensemble is None) - if you use self.J in both cases you do get the conversion. There might be something more subtle going on, but in that case the code needs to be cleared up as it is very hard to understand why we're using J instead of self.J in the non-ensemble case.

stephankramer · 2026-07-29T13:15:54Z

+            rf_cls = EnsembleReducedFunctional if self.ensemble is not None else ReducedFunctional
+            rf_args = [self.J, self.control_list]
+            if self.ensemble is not None:
+                rf_args.append(self.ensemble)
+            self.Jhat = rf_cls(*rf_args, **self.rf_kwargs)


Suggested change

rf_cls = EnsembleReducedFunctional if self.ensemble is not None else ReducedFunctional

rf_args = [self.J, self.control_list]

if self.ensemble is not None:

rf_args.append(self.ensemble)

self.Jhat = rf_cls(*rf_args, **self.rf_kwargs)

if self.ensemble is None:

self.Jhat = ReducedFunctional(self.J, self.control_list, **self.rf_kwargs)

else:

self.Jhat = EnsembleReducedFunctional(self.J, self.control_list, self.ensemble, **self.rf_kwargs)

Much simpler, no?

stephankramer · 2026-07-29T13:18:24Z

+        if self.ensemble is not None:
+            objective = ReducedFunctionalNumPy(self.reduced_functional)


Why is this needed for ensemble parallel? Presumably the conversion happens automatically for non-ensemble - is this something that needs fixing in pyadjoint?

I have not checked if it would work without the ReducedFunctionalNumPy() wrapping, but can give it a go. Was following the procedure given in Steps 4-6 of this Firedrake example:

The default minimize function uses scipy.minimize, and wraps the ReducedFunctional in a ReducedFunctionalNumPy that handles transferring data between Firedrake and numpy data structures. However, because we have a custom ReducedFunctional, we need to do this ourselves

Seems like ReducedFunctionalNumPy is needed. I'm assuming it's because EnsembleReducedFunctional is a class from firedrake/adjoint rather than being a pyadjoint class like regular ReducedFunctional.

stephankramer · 2026-07-29T13:20:51Z

            func_list.append(dc)
-        minconv = taylor_test(self.reduced_functional, self.control_coeff_list, func_list)
+        if self.ensemble:
+            minconv = taylor_test(self.reduced_functional.local_reduced_functional, self.control_coeff_list, func_list)


Does this mean it runs an independent taylor test for each ensemble? Is that what we want? Do we not want to test correctness of the combined functional and gradient?

cpjordan · 2026-07-29T13:38:29Z

We are maybe overloading this example a bit as well now. It would require a bit of extra work but perhaps we could have two inversion scripts? Inversion script 1 keeps the focus on spatial regularisation and we restrict it to Stations A-E which are for velocity only, then inversion script 2 looks at ensemble to inverting for both velocity (as per inversion script 1) and elevation, using both the friction and boundary forcing. That way the elevation gauges actually contribute as they are obviously more sensitive to boundary forcing in such a small domain.

Happy to help if we split this, given I want to look at boundary forcing inversion for North Sea cases.

Brazovskis · 2026-07-29T14:56:43Z

Thank you for taking the time to look at it Stephan, much appreciated! Will work to address the comments.

@cpjordan headland_inversion is probably getting a bit bloated for an example, I agree. Not sure about inversion script 2 introducing both ensemble operations and the use of boundary forcing as a control in that case – it could end up being as overloaded as headland_inversion is now.

- simplifying comm flags wherever possible

…arameter

… consistency_test() compares an instance of reduced_functional to itself. Taylor tests the combined functional and gradient

Brazovskis added 3 commits July 28, 2026 16:40

Update exporter.py to work with ensembles

0f951b7

Update inversion_tools.py to work with ensembles/EnsembleReducedFunct…

5701ed7

…ional

Update headland_inversion example to work with ensemble runs

9f6fec4

stephankramer reviewed Jul 29, 2026

View reviewed changes

Brazovskis added 6 commits August 3, 2026 15:21

- restoring docstrings

2bf6049

- simplifying comm flags wherever possible

Explicitly importing numpy as np

b37b79f

Simplifying comm flags and adding clarifying notes

4289aee

replacing redundant is_station_export_root function with no_exports p…

00b8add

…arameter

Consistent treatment of J across ensemble and non-ensemble runs

0b9f81c

Restoring original behaviour of consistency_test() and taylor_test().…

026cbd9

… consistency_test() compares an instance of reduced_functional to itself. Taylor tests the combined functional and gradient

		if self.ensemble is not None:
		objective = ReducedFunctionalNumPy(self.reduced_functional)

Uh oh!

Conversation

Brazovskis commented Jul 29, 2026

Changes

Why work with ensembles?

Uh oh!

stephankramer left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

cpjordan commented Jul 29, 2026

Uh oh!

Brazovskis commented Jul 29, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants