From 4953c8cf5a41be5d93c086f2e899dcc6885026fc Mon Sep 17 00:00:00 2001
From: Olender <olender@usp.br>
Date: Tue, 14 Jul 2026 20:44:44 +0100
Subject: [PATCH 1/7] reorganizing wrappers to deal with circular import with
 less confusion

---
 spyro/io/__init__.py               |  21 +-
 spyro/io/basicio.py                | 314 +--------------------
 spyro/io/parallelism_wrappers.py   | 439 +++++++++++++++++++++++++++++
 spyro/meshing/meshing_functions.py |   4 +-
 spyro/solvers/acoustic_wave.py     |   2 +-
 spyro/solvers/inversion.py         |   3 +-
 spyro/solvers/wave.py              |   2 +-
 spyro/utils/utils.py               | 132 +--------
 8 files changed, 464 insertions(+), 453 deletions(-)
 create mode 100644 spyro/io/parallelism_wrappers.py

diff --git a/spyro/io/__init__.py b/spyro/io/__init__.py
index d355e5a6..fa52eda3 100644
--- a/spyro/io/__init__.py
+++ b/spyro/io/__init__.py
@@ -1,6 +1,5 @@
 from .basicio import (
     write_function_to_grid,
-    is_owner,
     save_shots,
     load_shots,
     read_mesh,
@@ -8,21 +7,24 @@
     project_grid_velocity_data,
     # ensemble_forward_ad,
     # ensemble_forward_elastic_waves,
-    ensemble_gradient,
     # ensemble_gradient_elastic_waves,
     parallel_print,
     saving_source_and_receiver_location_in_csv,
-    switch_serial_shot,
-    ensemble_save,
-    ensemble_load,
-    delete_tmp_files,
-    ensemble_shot_record,
-    ensemble_functional,
     read_segy_velocity_model,
     read_bin_velocity_model,
     write_velocity_model
 )
-from .segy_io import create_segy, create_segy_from_grid
+from .parallelism_wrappers import (
+    delete_tmp_files,
+    ensemble_functional,
+    ensemble_gradient,
+    ensemble_load,
+    ensemble_save,
+    ensemble_shot_record,
+    is_owner,
+    switch_serial_shot,
+)
+from .segy_io import create_segy, create_segy_from_grid, export_scalar_field
 from .model_parameters import Model_parameters
 from .backwards_compatibility_io import Dictionary_conversion
 from . import dictionaryio
@@ -63,4 +65,5 @@
     "read_segy_velocity_model",
     "read_bin_velocity_model",
     "write_velocity_model",
+    "export_scalar_field"
 ]
diff --git a/spyro/io/basicio.py b/spyro/io/basicio.py
index 7ad2a387..7fd59d2a 100644
--- a/spyro/io/basicio.py
+++ b/spyro/io/basicio.py
@@ -10,6 +10,7 @@
 import os
 import warnings
 import segyio
+from .parallelism_wrappers import ensemble_save, ensemble_load
 from ..tools.version_control import is_firedrake_new
 
 
@@ -18,293 +19,6 @@
     fire.interpolate = interpolate
 
 
-def delete_tmp_files(wave):
-    """Delete temporary numpy files associated with a wave object."""
-    str_id = f"*{wave.random_id_string}.npy"
-    for file in glob.glob(str_id):
-        os.remove(file)
-
-
-def _run_for_each_shot(obj, func, *args, **kwargs):
-    """Helper to run a function for each shot in spatial parallelism."""
-    results = []
-    for snum in range(obj.number_of_sources):
-        switch_serial_shot(obj, snum)
-        results.append(func(*args, **kwargs))
-    return results
-
-
-def ensemble_shot_record(func):
-    """Decorator for read and write shots for ensemble parallelism"""
-    def wrapper(*args, **kwargs):
-        obj = args[0]
-        if obj.parallelism_type == "spatial" and obj.number_of_sources > 1:
-            return _run_for_each_shot(obj, func, *args, **kwargs)
-    return wrapper
-
-
-def ensemble_save(func):
-    """Decorator for saving files with parallelism.
-
-    Parameters:
-    -----------
-    func: The wrapped function that performs the actual saving operation.
-    Expected to accept a :class:`Wave` based object as first argument.
-
-    Returns:
-    --------
-    wrapper: A decorator function that wraps the original saving function with
-        parallelism logic.
-
-    Notes:
-    ------
-    Handles saving in different scenarions:
-    - For ensemble parallelism or single source: iterates through propagations in
-      each core and saves when the propagation is owned by the current rank.
-    - For spatial-only parallelism with multiple sources: loads shots from temporary
-      files using the switch_serial_shot method and saves to named output files
-    - Requires first object to have attributes: `comm`, `parallelism_type`, `number_of_sources`,
-      and `shot_ids_per_propagation`.
-    - Temporary files are loaded via :meth:`switch_serial_shot()` when using spatial-only parallelism
-    """
-    def wrapper(*args, **kwargs):
-        obj = args[0]  # Requires first arg to be an instant or subclass of Wave
-        _comm = obj.comm
-        if obj.parallelism_type != "spatial" or obj.number_of_sources == 1:
-            for propagation_id, shot_ids_in_propagation in enumerate(obj.shot_ids_per_propagation):
-                if is_owner(_comm, propagation_id) and _comm.comm.rank == 0:
-                    func(obj, **dict(kwargs, shot_ids=shot_ids_in_propagation))
-        else:
-            # For spatial parallelism: load propagation data from tmp files (no file_name) then save wanted data to named files
-            for snum in range(obj.number_of_sources):
-                switch_serial_shot(obj, snum, file_name=None)  # Load from tmp files
-                if _comm.comm.rank == 0:
-                    func(obj, **dict(kwargs, shot_ids=[snum]))
-    return wrapper
-
-
-def ensemble_load(func):
-    """Decorator for loading shots for ensemble parallelism.
-
-    For spatial parallelism with multiple sources, loads from named files directly.
-
-    Parameters:
-    -----------
-    func: The wrapped function that performs the actual loading operation.
-    Expected to accept a :class:`Wave` based object as first argument.
-
-    Returns:
-    --------
-    wrapper: A decorator function that wraps the original loading function with
-        parallelism logic.
-    """
-    def wrapper(*args, **kwargs):
-        obj = args[0]
-        _comm = obj.comm
-        if obj.parallelism_type != "spatial" or obj.number_of_sources == 1:
-            for propagation_id, shot_ids_in_propagation in enumerate(obj.shot_ids_per_propagation):
-                if is_owner(_comm, propagation_id):
-                    func(obj, **dict(kwargs, shot_ids=shot_ids_in_propagation))
-        else:
-            # For spatial parallelism: load data directly from named files (no switch_serial_shot needed)
-            for snum in range(obj.number_of_sources):
-                func(obj, **dict(kwargs, shot_ids=[snum]))
-    return wrapper
-
-
-def ensemble_propagator(func):
-    """Decorator for forward to distribute shots for ensemble parallelism
-
-    Parameters:
-    -----------
-    func: The wrapped function that performs the actual propagation operation.
-    Expected to accept a :class:`Wave` based object as first argument.
-
-    Returns:
-    --------
-    wrapper: A decorator function that wraps the original propagator function with
-        ensemble parallelism logic.
-    """
-
-    def wrapper(*args, **kwargs):
-        if args[0].parallelism_type != "spatial" or args[0].number_of_sources == 1:
-            shot_ids_per_propagation_list = args[0].shot_ids_per_propagation
-            _comm = args[0].comm
-            for propagation_id, shot_ids_in_propagation in enumerate(shot_ids_per_propagation_list):
-                if is_owner(_comm, propagation_id):
-                    func(*args, **dict(kwargs, source_nums=shot_ids_in_propagation))
-        elif args[0].parallelism_type == "spatial" and args[0].number_of_sources > 1:
-            num = args[0].number_of_sources
-            starting_time = args[0].current_time
-            for snum in range(num):
-                args[0].reset_pressure()
-                args[0].current_time = starting_time
-                func(*args, **dict(kwargs, source_nums=[snum]))
-                save_serial_data(args[0], snum)
-
-    return wrapper
-
-
-def _shot_filename(propagation_id, wave, prefix='tmp', random_str_in_use=True):
-    """
-    Helper to construct filenames for shot/receiver data based on propagation and wave information.
-
-    Parameters:
-    -----------
-    propagation_id (int): The index identifying the current propagation.
-
-    wave (object): A :class:`Wave` object containing shot and communication information. Must have attributes:
-        - shot_ids_per_propagation: A list or dict mapping propagation IDs to shot IDs.
-        - comm: The current MPI communicator.
-    prefix (str, optional): Prefix for the filename. Defaults to 'tmp'.
-    random_str_in_use (bool, optional): If True, includes a random string and communicator rank in
-        the filename, gotten from the Wave object, and uses '.npy' extension.
-        If False, omits these and uses '.dat' extension. Defaults to True.
-
-    Returns:
-    --------
-    str: The constructed filename.
-    """
-    shot_ids = wave.shot_ids_per_propagation[propagation_id]
-    if random_str_in_use:
-        id_str = wave.random_id_string
-        spatialcomm = wave.comm.comm.rank
-        comm__str = f"_comm{spatialcomm}"
-        post_fix = "npy"
-    else:
-        id_str = ""
-        comm__str = ""
-        post_fix = "dat"
-    return f"{prefix}{shot_ids}{comm__str}{id_str}.{post_fix}"
-
-
-def save_serial_data(wave, propagation_id):
-    """
-    Save serial data to numpy files.
-
-    Args:
-        wave (:class:`Wave`): The wave object containing the forward solution.
-        propagation_id (int): The propagation ID.
-
-    Returns:
-        None
-    """
-    if wave.forward_solution:
-        # There are cases where forward_solution is empty, e.g. when running
-        # forward_solve for the true model. In that case, we skip saving the
-        # solution on the entire domain, which is not needed.
-        arrays_list = [obj.dat.data[:] for obj in wave.forward_solution]
-        stacked_arrays = np.stack(arrays_list, axis=0)
-        np.save(_shot_filename(propagation_id, wave, prefix='tmp_shot'), stacked_arrays)
-    np.save(_shot_filename(propagation_id, wave, prefix='tmp_rec'), wave.forward_solution_receivers)
-
-
-def switch_serial_shot(wave, propagation_id, file_name=None, just_for_dat_management=False):
-    """
-    Switches the current serial shot for a given wave to shot identified with propagation ID.
-
-    Args:
-        wave (:class:`Wave`): The wave object.
-        propagation_id (int): The propagation ID.
-
-    Returns:
-        None
-    """
-    if file_name is None:
-        forward_solution_filename = _shot_filename(propagation_id, wave, prefix='tmp_shot')
-        if os.path.exists(forward_solution_filename) or wave.forward_solution:
-            # The adjoint propagator consumes forward_solution with pop(). When
-            # switching to the next shot, reload saved snapshots even if the
-            # in-memory list has been emptied.
-            stacked_shot_arrays = np.load(forward_solution_filename)
-            if not wave.forward_solution:
-                rebuild_empty_forward_solution(wave, len(stacked_shot_arrays))
-            for array_i, array in enumerate(stacked_shot_arrays):
-                wave.forward_solution[array_i].dat.data[:] = array
-        receiver_solution_filename = _shot_filename(propagation_id, wave, prefix='tmp_rec')
-    else:
-        receiver_solution_filename = _shot_filename(propagation_id, wave, prefix=file_name, random_str_in_use=False)
-    wave.forward_solution_receivers = np.load(receiver_solution_filename, allow_pickle=True)
-
-
-def ensemble_functional(func):
-    """Decorator for gradient to distribute shots for ensemble parallelism"""
-
-    def wrapper(*args, **kwargs):
-        comm = args[0].comm
-        if args[0].adjoint_type.name == "AUTOMATED_ADJOINT":
-            # pyadjoint needs the annotated Firedrake object, not a numpy scalar
-            # produced by the ensemble reduction path below.
-            return func(*args, **kwargs)
-        if args[0].parallelism_type != "spatial" or args[0].number_of_sources == 1:
-            J = func(*args, **kwargs)
-            J_total = np.zeros((1))
-            J_total[0] += J
-            J_total = fire.COMM_WORLD.allreduce(J_total, op=MPI.SUM)
-            J_total[0] /= comm.comm.size
-
-        elif args[0].parallelism_type == "spatial" and args[0].number_of_sources > 1:
-            residual_list = args[1]
-            J_total = np.zeros((1))
-
-            for snum in range(args[0].number_of_sources):
-                switch_serial_shot(args[0], snum)
-                current_residual = residual_list[snum]
-                J = func(args[0], current_residual)
-                J_total += J
-            J_total[0] /= comm.comm.size
-
-            comm.comm.barrier()
-
-        return J_total[0]
-
-    return wrapper
-
-
-def ensemble_gradient(func):
-    """Decorator for gradient to distribute shots for ensemble parallelism"""
-
-    def wrapper(*args, **kwargs):
-        comm = args[0].comm
-        if args[0].parallelism_type != "spatial" or args[0].number_of_sources == 1:
-            shot_ids_per_propagation_list = args[0].shot_ids_per_propagation
-            for propagation_id, shot_ids_in_propagation in enumerate(shot_ids_per_propagation_list):
-                if is_owner(comm, propagation_id):
-                    grad = func(*args, **kwargs)
-            grad_total = fire.Function(args[0].function_space)
-
-            comm.comm.barrier()
-            grad_total = comm.allreduce(grad, grad_total)
-            grad_total /= comm.ensemble_comm.size
-
-            return grad_total
-        elif args[0].parallelism_type == "spatial" and args[0].number_of_sources > 1:
-            num = args[0].number_of_sources
-            starting_time = args[0].current_time
-            grad_total = fire.Function(args[0].function_space)
-            misfit_list = kwargs.get("misfit")
-
-            for snum in range(num):
-                switch_serial_shot(args[0], snum)
-                current_misfit = misfit_list[snum]
-                args[0].reset_pressure()
-                args[0].current_time = starting_time
-                grad = func(*args,
-                            **dict(
-                                kwargs,
-                                misfit=current_misfit,
-                            )
-                            )
-                grad_total += grad
-
-            grad_total /= num
-            comm.comm.barrier()
-
-            return grad_total
-
-    return wrapper
-
-
 def write_function_to_grid(function, V, grid_spacing, buffer=False):
     """Interpolate a Firedrake function to a structured grid
 
@@ -405,12 +119,6 @@ def save_shots(Wave_obj, file_name="shots/shot_record_", shot_ids=0):
     return None
 
 
-def rebuild_empty_forward_solution(wave, time_steps):
-    wave.forward_solution = []
-    for i in range(time_steps):
-        wave.forward_solution.append(fire.Function(wave.function_space))
-
-
 @ensemble_load
 def load_shots(Wave_obj, file_name="shots/shot_record_", shot_ids=0):
     """Load a `pickle` to a `numpy.ndarray`.
@@ -439,26 +147,6 @@ def load_shots(Wave_obj, file_name="shots/shot_record_", shot_ids=0):
     return None
 
 
-def is_owner(ens_comm, rank):
-    """Distribute shots between processors in using a modulus operator
-
-    Parameters
-    ----------
-    ens_comm: Firedrake.ensemble_communicator
-        An ensemble communicator
-    rank: int
-        The rank of the core
-
-    Returns
-    -------
-    boolean
-        `True` if `rank` owns this shot
-
-    """
-    owner = ens_comm.ensemble_comm.rank == (rank % ens_comm.ensemble_comm.size)
-    return owner
-
-
 def _check_units(c):
     if min(c.dat.data[:]) > 100.0:
         # data is in m/s but must be in km/s
diff --git a/spyro/io/parallelism_wrappers.py b/spyro/io/parallelism_wrappers.py
new file mode 100644
index 00000000..416f2123
--- /dev/null
+++ b/spyro/io/parallelism_wrappers.py
@@ -0,0 +1,439 @@
+import numpy as np
+import firedrake as fire
+
+def run_in_one_core(func):
+    """Decorator to execute function only on rank 0.
+
+    Ensures the decorated function runs only on the root process (rank 0)
+    of the communicator. Other processes skip execution. Useful for I/O
+    operations and serial tasks in parallel environments.
+
+    Parameters
+    ----------
+    func : callable
+        Function to decorate. The first argument of func must be an object
+        with a `comm` attribute containing an MPI communicator.
+
+    Returns
+    -------
+    callable
+        Wrapped function that executes only on rank 0.
+
+    Notes
+    -----
+    The function checks for two types of communicators:
+    - Ensemble communicator: Runs only if both `ensemble_comm.rank` == 0
+      and comm.rank == 0.
+    - Regular communicator: Runs only if comm.rank == 0.
+    - If comm is None, the function runs normally without restrictions.
+
+    The function does not broadcast results to other processes.
+
+    See Also
+    --------
+    run_in_one_core_and_broadcast : Similar decorator that also broadcasts results.
+
+    Examples
+    --------
+    >>> @run_in_one_core
+    ... def save_file(obj, filename):
+    ...     # Only rank 0 writes the file
+    ...     with open(filename, 'w') as f:
+    ...         f.write(str(obj.data))
+    """
+
+    def wrapper(*args, **kwargs):
+        comm = args[0].comm
+        if comm is None:
+            return func(*args, **kwargs)
+        else:
+            if getattr(comm, "ensemble_comm", None) is not None:
+                if comm.ensemble_comm.rank == 0 and comm.comm.rank == 0:
+                    return func(*args, **kwargs)
+            elif getattr(comm, "rank", None) is not None:
+                if comm.rank == 0:
+                    return func(*args, **kwargs)
+
+    return wrapper
+
+
+def run_in_one_core_and_broadcast(func):
+    """Decorator to execute function on rank 0 and broadcast result.
+
+    Ensures the decorated function runs only on the root process (rank 0)
+    and broadcasts the return value to all other processes. Useful for
+    file reading and other operations that should be performed once but
+    shared across all processes.
+
+    Parameters
+    ----------
+    func : callable
+        Function to decorate. The first argument of func must be an object
+        with a 'comm' attribute containing an MPI communicator.
+
+    Returns
+    -------
+    callable
+        Wrapped function that executes on rank 0 and broadcasts the result
+        to all processes.
+
+    Notes
+    -----
+    The function handles two types of communicators:
+    - Ensemble communicator: Executes on rank (0,0), broadcasts within
+      ensemble, then within spatial communicator.
+    - Regular communicator: Executes on rank 0, broadcasts to all.
+    - If comm is None, the function runs normally without MPI operations.
+
+    All processes receive the same return value from the broadcast.
+
+    See Also
+    --------
+    run_in_one_core : Similar decorator without broadcasting.
+
+    Examples
+    --------
+    >>> @run_in_one_core_and_broadcast
+    ... def load_config(obj, filename):
+    ...     # Only rank 0 reads the file, result shared with all
+    ...     with open(filename, 'r') as f:
+    ...         return json.load(f)
+    """
+
+    def wrapper(*args, **kwargs):
+        comm = args[0].comm
+        if comm is None:
+            return func(*args, **kwargs)
+        else:
+            result = None
+            if getattr(comm, "ensemble_comm", None) is not None:
+                # Handle ensemble communicator
+                if comm.ensemble_comm.rank == 0 and comm.comm.rank == 0:
+                    result = func(*args, **kwargs)
+                # Broadcast within ensemble
+                result = comm.ensemble_comm.bcast(result, root=0)
+                # Broadcast within spatial communicator
+                result = comm.comm.bcast(result, root=0)
+            elif getattr(comm, "rank", None) is not None:
+                # Handle regular communicator
+                if comm.rank == 0:
+                    result = func(*args, **kwargs)
+                result = comm.bcast(result, root=0)
+            return result
+
+    return wrapper
+
+
+def _run_for_each_shot(obj, func, *args, **kwargs):
+    """Helper to run a function for each shot in spatial parallelism."""
+    results = []
+    for snum in range(obj.number_of_sources):
+        switch_serial_shot(obj, snum)
+        results.append(func(*args, **kwargs))
+    return results
+
+
+def ensemble_shot_record(func):
+    """Decorator for read and write shots for ensemble parallelism"""
+    def wrapper(*args, **kwargs):
+        obj = args[0]
+        if obj.parallelism_type == "spatial" and obj.number_of_sources > 1:
+            return _run_for_each_shot(obj, func, *args, **kwargs)
+    return wrapper
+
+
+def switch_serial_shot(wave, propagation_id, file_name=None, just_for_dat_management=False):
+    """
+    Switches the current serial shot for a given wave to shot identified with propagation ID.
+
+    Args:
+        wave (:class:`Wave`): The wave object.
+        propagation_id (int): The propagation ID.
+
+    Returns:
+        None
+    """
+    if file_name is None:
+        forward_solution_filename = _shot_filename(propagation_id, wave, prefix='tmp_shot')
+        if os.path.exists(forward_solution_filename) or wave.forward_solution:
+            # The adjoint propagator consumes forward_solution with pop(). When
+            # switching to the next shot, reload saved snapshots even if the
+            # in-memory list has been emptied.
+            stacked_shot_arrays = np.load(forward_solution_filename)
+            if not wave.forward_solution:
+                rebuild_empty_forward_solution(wave, len(stacked_shot_arrays))
+            for array_i, array in enumerate(stacked_shot_arrays):
+                wave.forward_solution[array_i].dat.data[:] = array
+        receiver_solution_filename = _shot_filename(propagation_id, wave, prefix='tmp_rec')
+    else:
+        receiver_solution_filename = _shot_filename(propagation_id, wave, prefix=file_name, random_str_in_use=False)
+    wave.forward_solution_receivers = np.load(receiver_solution_filename, allow_pickle=True)
+
+
+def _shot_filename(propagation_id, wave, prefix='tmp', random_str_in_use=True):
+    """
+    Helper to construct filenames for shot/receiver data based on propagation and wave information.
+
+    Parameters:
+    -----------
+    propagation_id (int): The index identifying the current propagation.
+
+    wave (object): A :class:`Wave` object containing shot and communication information. Must have attributes:
+        - shot_ids_per_propagation: A list or dict mapping propagation IDs to shot IDs.
+        - comm: The current MPI communicator.
+    prefix (str, optional): Prefix for the filename. Defaults to 'tmp'.
+    random_str_in_use (bool, optional): If True, includes a random string and communicator rank in
+        the filename, gotten from the Wave object, and uses '.npy' extension.
+        If False, omits these and uses '.dat' extension. Defaults to True.
+
+    Returns:
+    --------
+    str: The constructed filename.
+    """
+    shot_ids = wave.shot_ids_per_propagation[propagation_id]
+    if random_str_in_use:
+        id_str = wave.random_id_string
+        spatialcomm = wave.comm.comm.rank
+        comm__str = f"_comm{spatialcomm}"
+        post_fix = "npy"
+    else:
+        id_str = ""
+        comm__str = ""
+        post_fix = "dat"
+    return f"{prefix}{shot_ids}{comm__str}{id_str}.{post_fix}"
+
+
+def rebuild_empty_forward_solution(wave, time_steps):
+    wave.forward_solution = []
+    for i in range(time_steps):
+        wave.forward_solution.append(fire.Function(wave.function_space))
+
+
+def ensemble_save(func):
+    """Decorator for saving files with parallelism.
+
+    Parameters:
+    -----------
+    func: The wrapped function that performs the actual saving operation.
+    Expected to accept a :class:`Wave` based object as first argument.
+
+    Returns:
+    --------
+    wrapper: A decorator function that wraps the original saving function with
+        parallelism logic.
+
+    Notes:
+    ------
+    Handles saving in different scenarions:
+    - For ensemble parallelism or single source: iterates through propagations in
+      each core and saves when the propagation is owned by the current rank.
+    - For spatial-only parallelism with multiple sources: loads shots from temporary
+      files using the switch_serial_shot method and saves to named output files
+    - Requires first object to have attributes: `comm`, `parallelism_type`, `number_of_sources`,
+      and `shot_ids_per_propagation`.
+    - Temporary files are loaded via :meth:`switch_serial_shot()` when using spatial-only parallelism
+    """
+    def wrapper(*args, **kwargs):
+        obj = args[0]  # Requires first arg to be an instant or subclass of Wave
+        _comm = obj.comm
+        if obj.parallelism_type != "spatial" or obj.number_of_sources == 1:
+            for propagation_id, shot_ids_in_propagation in enumerate(obj.shot_ids_per_propagation):
+                if is_owner(_comm, propagation_id) and _comm.comm.rank == 0:
+                    func(obj, **dict(kwargs, shot_ids=shot_ids_in_propagation))
+        else:
+            # For spatial parallelism: load propagation data from tmp files (no file_name) then save wanted data to named files
+            for snum in range(obj.number_of_sources):
+                switch_serial_shot(obj, snum, file_name=None)  # Load from tmp files
+                if _comm.comm.rank == 0:
+                    func(obj, **dict(kwargs, shot_ids=[snum]))
+    return wrapper
+
+
+def ensemble_load(func):
+    """Decorator for loading shots for ensemble parallelism.
+
+    For spatial parallelism with multiple sources, loads from named files directly.
+
+    Parameters:
+    -----------
+    func: The wrapped function that performs the actual loading operation.
+    Expected to accept a :class:`Wave` based object as first argument.
+
+    Returns:
+    --------
+    wrapper: A decorator function that wraps the original loading function with
+        parallelism logic.
+    """
+    def wrapper(*args, **kwargs):
+        obj = args[0]
+        _comm = obj.comm
+        if obj.parallelism_type != "spatial" or obj.number_of_sources == 1:
+            for propagation_id, shot_ids_in_propagation in enumerate(obj.shot_ids_per_propagation):
+                if is_owner(_comm, propagation_id):
+                    func(obj, **dict(kwargs, shot_ids=shot_ids_in_propagation))
+        else:
+            # For spatial parallelism: load data directly from named files (no switch_serial_shot needed)
+            for snum in range(obj.number_of_sources):
+                func(obj, **dict(kwargs, shot_ids=[snum]))
+    return wrapper
+
+
+def is_owner(ens_comm, rank):
+    """Distribute shots between processors in using a modulus operator
+
+    Parameters
+    ----------
+    ens_comm: Firedrake.ensemble_communicator
+        An ensemble communicator
+    rank: int
+        The rank of the core
+
+    Returns
+    -------
+    boolean
+        `True` if `rank` owns this shot
+
+    """
+    owner = ens_comm.ensemble_comm.rank == (rank % ens_comm.ensemble_comm.size)
+    return owner
+
+
+def delete_tmp_files(wave):
+    """Delete temporary numpy files associated with a wave object."""
+    str_id = f"*{wave.random_id_string}.npy"
+    for file in glob.glob(str_id):
+        os.remove(file)
+
+
+def ensemble_propagator(func):
+    """Decorator for forward to distribute shots for ensemble parallelism
+
+    Parameters:
+    -----------
+    func: The wrapped function that performs the actual propagation operation.
+    Expected to accept a :class:`Wave` based object as first argument.
+
+    Returns:
+    --------
+    wrapper: A decorator function that wraps the original propagator function with
+        ensemble parallelism logic.
+    """
+
+    def wrapper(*args, **kwargs):
+        if args[0].parallelism_type != "spatial" or args[0].number_of_sources == 1:
+            shot_ids_per_propagation_list = args[0].shot_ids_per_propagation
+            _comm = args[0].comm
+            for propagation_id, shot_ids_in_propagation in enumerate(shot_ids_per_propagation_list):
+                if is_owner(_comm, propagation_id):
+                    func(*args, **dict(kwargs, source_nums=shot_ids_in_propagation))
+        elif args[0].parallelism_type == "spatial" and args[0].number_of_sources > 1:
+            num = args[0].number_of_sources
+            starting_time = args[0].current_time
+            for snum in range(num):
+                args[0].reset_pressure()
+                args[0].current_time = starting_time
+                func(*args, **dict(kwargs, source_nums=[snum]))
+                save_serial_data(args[0], snum)
+
+    return wrapper
+
+
+def save_serial_data(wave, propagation_id):
+    """
+    Save serial data to numpy files.
+
+    Args:
+        wave (:class:`Wave`): The wave object containing the forward solution.
+        propagation_id (int): The propagation ID.
+
+    Returns:
+        None
+    """
+    if wave.forward_solution:
+        # There are cases where forward_solution is empty, e.g. when running
+        # forward_solve for the true model. In that case, we skip saving the
+        # solution on the entire domain, which is not needed.
+        arrays_list = [obj.dat.data[:] for obj in wave.forward_solution]
+        stacked_arrays = np.stack(arrays_list, axis=0)
+        np.save(_shot_filename(propagation_id, wave, prefix='tmp_shot'), stacked_arrays)
+    np.save(_shot_filename(propagation_id, wave, prefix='tmp_rec'), wave.forward_solution_receivers)
+
+
+def ensemble_functional(func):
+    """Decorator for gradient to distribute shots for ensemble parallelism"""
+
+    def wrapper(*args, **kwargs):
+        comm = args[0].comm
+        if args[0].adjoint_type.name == "AUTOMATED_ADJOINT":
+            # pyadjoint needs the annotated Firedrake object, not a numpy scalar
+            # produced by the ensemble reduction path below.
+            return func(*args, **kwargs)
+        if args[0].parallelism_type != "spatial" or args[0].number_of_sources == 1:
+            J = func(*args, **kwargs)
+            J_total = np.zeros((1))
+            J_total[0] += J
+            J_total = fire.COMM_WORLD.allreduce(J_total, op=MPI.SUM)
+            J_total[0] /= comm.comm.size
+
+        elif args[0].parallelism_type == "spatial" and args[0].number_of_sources > 1:
+            residual_list = args[1]
+            J_total = np.zeros((1))
+
+            for snum in range(args[0].number_of_sources):
+                switch_serial_shot(args[0], snum)
+                current_residual = residual_list[snum]
+                J = func(args[0], current_residual)
+                J_total += J
+            J_total[0] /= comm.comm.size
+
+            comm.comm.barrier()
+
+        return J_total[0]
+
+    return wrapper
+
+
+def ensemble_gradient(func):
+    """Decorator for gradient to distribute shots for ensemble parallelism"""
+
+    def wrapper(*args, **kwargs):
+        comm = args[0].comm
+        if args[0].parallelism_type != "spatial" or args[0].number_of_sources == 1:
+            shot_ids_per_propagation_list = args[0].shot_ids_per_propagation
+            for propagation_id, shot_ids_in_propagation in enumerate(shot_ids_per_propagation_list):
+                if is_owner(comm, propagation_id):
+                    grad = func(*args, **kwargs)
+            grad_total = fire.Function(args[0].function_space)
+
+            comm.comm.barrier()
+            grad_total = comm.allreduce(grad, grad_total)
+            grad_total /= comm.ensemble_comm.size
+
+            return grad_total
+        elif args[0].parallelism_type == "spatial" and args[0].number_of_sources > 1:
+            num = args[0].number_of_sources
+            starting_time = args[0].current_time
+            grad_total = fire.Function(args[0].function_space)
+            misfit_list = kwargs.get("misfit")
+
+            for snum in range(num):
+                switch_serial_shot(args[0], snum)
+                current_misfit = misfit_list[snum]
+                args[0].reset_pressure()
+                args[0].current_time = starting_time
+                grad = func(*args,
+                            **dict(
+                                kwargs,
+                                misfit=current_misfit,
+                            )
+                            )
+                grad_total += grad
+
+            grad_total /= num
+            comm.comm.barrier()
+
+            return grad_total
+
+    return wrapper
+
+
diff --git a/spyro/meshing/meshing_functions.py b/spyro/meshing/meshing_functions.py
index 14fb13cc..26ef6059 100644
--- a/spyro/meshing/meshing_functions.py
+++ b/spyro/meshing/meshing_functions.py
@@ -1,7 +1,7 @@
 import numpy as np
 from firedrake import Mesh as FireMeshReader
-from ..io import parallel_print
-from ..io import create_segy_from_grid
+from ..io.basicio import parallel_print
+from ..io.segy_io import create_segy_from_grid
 from .meshing_gmsh2d import build_gmsh_geometry_and_groups, apply_structured_winslow_smoothing2d
 from .meshing_utils import create_sizing_function, calculate_edge_length
 from .firedrake_based_wrappers import rectangle_mesh, periodic_rectangle_mesh, box_mesh
diff --git a/spyro/solvers/acoustic_wave.py b/spyro/solvers/acoustic_wave.py
index 4d8f286e..6e28f02c 100644
--- a/spyro/solvers/acoustic_wave.py
+++ b/spyro/solvers/acoustic_wave.py
@@ -2,7 +2,7 @@
 
 from .wave import Wave
 from pyadjoint import Tape, AdjFloat
-from ..io.basicio import ensemble_gradient
+from ..io.parallelism_wrappers import ensemble_gradient
 from ..io import interpolate
 from .acoustic_solver_construction_no_pml import (
     construct_solver_or_matrix_no_pml,
diff --git a/spyro/solvers/inversion.py b/spyro/solvers/inversion.py
index 814df630..d2bafa9c 100644
--- a/spyro/solvers/inversion.py
+++ b/spyro/solvers/inversion.py
@@ -13,8 +13,9 @@
 from ..utils import Gradient_mask_for_pml, Mask
 from ..utils.typing import WaveType
 from ..plots import plot_model as spyro_plot_model
-from ..io.basicio import parallel_print, switch_serial_shot
+from ..io.basicio import parallel_print
 from ..io.basicio import load_shots, save_shots
+from ..io.parallelism_wrappers import switch_serial_shot
 from ..io import create_segy
 from ..utils import run_in_one_core
 
diff --git a/spyro/solvers/wave.py b/spyro/solvers/wave.py
index e8f42e87..ab76c2f2 100644
--- a/spyro/solvers/wave.py
+++ b/spyro/solvers/wave.py
@@ -9,7 +9,7 @@
 from ..domains.space import check_function_space_type, create_function_space
 from ..io import Model_parameters
 from ..io import material_properties_io
-from ..io.basicio import ensemble_propagator
+from ..io.parallelism_wrappers import ensemble_propagator
 from ..io import parallel_print
 from ..io.field_logger import FieldLogger
 from ..receivers.Receivers import Receivers
diff --git a/spyro/utils/utils.py b/spyro/utils/utils.py
index 229c59c9..cb3f2475 100644
--- a/spyro/utils/utils.py
+++ b/spyro/utils/utils.py
@@ -6,9 +6,12 @@
 from scipy.signal import butter, filtfilt
 import warnings
 
-from ..io import ensemble_functional
-from ..io import parallel_print
-from ..io import write_velocity_model
+from ..io.basicio import parallel_print, write_velocity_model
+from ..io.parallelism_wrappers import (
+    ensemble_functional,
+    run_in_one_core,
+    run_in_one_core_and_broadcast,
+)
 from ..domains.space import create_function_space
 from .typing import FunctionalEvaluationMode, FunctionalType
 
@@ -584,129 +587,6 @@ def __init__(self, Wave_obj):
         super().__init__(boundaries, Wave_obj)
 
 
-def run_in_one_core(func):
-    """Decorator to execute function only on rank 0.
-
-    Ensures the decorated function runs only on the root process (rank 0)
-    of the communicator. Other processes skip execution. Useful for I/O
-    operations and serial tasks in parallel environments.
-
-    Parameters
-    ----------
-    func : callable
-        Function to decorate. The first argument of func must be an object
-        with a `comm` attribute containing an MPI communicator.
-
-    Returns
-    -------
-    callable
-        Wrapped function that executes only on rank 0.
-
-    Notes
-    -----
-    The function checks for two types of communicators:
-    - Ensemble communicator: Runs only if both `ensemble_comm.rank` == 0
-      and comm.rank == 0.
-    - Regular communicator: Runs only if comm.rank == 0.
-    - If comm is None, the function runs normally without restrictions.
-
-    The function does not broadcast results to other processes.
-
-    See Also
-    --------
-    run_in_one_core_and_broadcast : Similar decorator that also broadcasts results.
-
-    Examples
-    --------
-    >>> @run_in_one_core
-    ... def save_file(obj, filename):
-    ...     # Only rank 0 writes the file
-    ...     with open(filename, 'w') as f:
-    ...         f.write(str(obj.data))
-    """
-
-    def wrapper(*args, **kwargs):
-        comm = args[0].comm
-        if comm is None:
-            return func(*args, **kwargs)
-        else:
-            if getattr(comm, "ensemble_comm", None) is not None:
-                if comm.ensemble_comm.rank == 0 and comm.comm.rank == 0:
-                    return func(*args, **kwargs)
-            elif getattr(comm, "rank", None) is not None:
-                if comm.rank == 0:
-                    return func(*args, **kwargs)
-
-    return wrapper
-
-
-def run_in_one_core_and_broadcast(func):
-    """Decorator to execute function on rank 0 and broadcast result.
-
-    Ensures the decorated function runs only on the root process (rank 0)
-    and broadcasts the return value to all other processes. Useful for
-    file reading and other operations that should be performed once but
-    shared across all processes.
-
-    Parameters
-    ----------
-    func : callable
-        Function to decorate. The first argument of func must be an object
-        with a 'comm' attribute containing an MPI communicator.
-
-    Returns
-    -------
-    callable
-        Wrapped function that executes on rank 0 and broadcasts the result
-        to all processes.
-
-    Notes
-    -----
-    The function handles two types of communicators:
-    - Ensemble communicator: Executes on rank (0,0), broadcasts within
-      ensemble, then within spatial communicator.
-    - Regular communicator: Executes on rank 0, broadcasts to all.
-    - If comm is None, the function runs normally without MPI operations.
-
-    All processes receive the same return value from the broadcast.
-
-    See Also
-    --------
-    run_in_one_core : Similar decorator without broadcasting.
-
-    Examples
-    --------
-    >>> @run_in_one_core_and_broadcast
-    ... def load_config(obj, filename):
-    ...     # Only rank 0 reads the file, result shared with all
-    ...     with open(filename, 'r') as f:
-    ...         return json.load(f)
-    """
-
-    def wrapper(*args, **kwargs):
-        comm = args[0].comm
-        if comm is None:
-            return func(*args, **kwargs)
-        else:
-            result = None
-            if getattr(comm, "ensemble_comm", None) is not None:
-                # Handle ensemble communicator
-                if comm.ensemble_comm.rank == 0 and comm.comm.rank == 0:
-                    result = func(*args, **kwargs)
-                # Broadcast within ensemble
-                result = comm.ensemble_comm.bcast(result, root=0)
-                # Broadcast within spatial communicator
-                result = comm.comm.bcast(result, root=0)
-            elif getattr(comm, "rank", None) is not None:
-                # Handle regular communicator
-                if comm.rank == 0:
-                    result = func(*args, **kwargs)
-                result = comm.bcast(result, root=0)
-            return result
-
-    return wrapper
-
-
 @run_in_one_core_and_broadcast
 def write_hdf5_velocity_model(obj_with_comm, segy_filename):
     """Convert a SEG-Y velocity model to HDF5 format.

From c8fcab2aca873b0f9a1c0327cb52a93821ff5daf Mon Sep 17 00:00:00 2001
From: Olender <olender@usp.br>
Date: Tue, 14 Jul 2026 20:51:41 +0100
Subject: [PATCH 2/7] debugging

---
 spyro/io/basicio.py              | 2 --
 spyro/io/parallelism_wrappers.py | 2 ++
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/spyro/io/basicio.py b/spyro/io/basicio.py
index 7fd59d2a..f18cf219 100644
--- a/spyro/io/basicio.py
+++ b/spyro/io/basicio.py
@@ -1,12 +1,10 @@
 from __future__ import with_statement
 
 import pickle
-from mpi4py import MPI
 import firedrake as fire
 import h5py
 import numpy as np
 from scipy.interpolate import griddata
-import glob
 import os
 import warnings
 import segyio
diff --git a/spyro/io/parallelism_wrappers.py b/spyro/io/parallelism_wrappers.py
index 416f2123..8b549074 100644
--- a/spyro/io/parallelism_wrappers.py
+++ b/spyro/io/parallelism_wrappers.py
@@ -1,3 +1,5 @@
+from mpi4py import MPI
+import glob
 import numpy as np
 import firedrake as fire
 

From 596afb9265f50664a3960f6657dd147cc8161600 Mon Sep 17 00:00:00 2001
From: Olender <olender@usp.br>
Date: Tue, 14 Jul 2026 21:40:37 +0100
Subject: [PATCH 3/7] adding wrapper for comm as kwarg. Duplicate code to
 reduce later

---
 spyro/io/parallelism_wrappers.py | 57 ++++++++++++++++++++++++++++++++
 1 file changed, 57 insertions(+)

diff --git a/spyro/io/parallelism_wrappers.py b/spyro/io/parallelism_wrappers.py
index 8b549074..50c6993e 100644
--- a/spyro/io/parallelism_wrappers.py
+++ b/spyro/io/parallelism_wrappers.py
@@ -1,3 +1,4 @@
+import os
 from mpi4py import MPI
 import glob
 import numpy as np
@@ -59,6 +60,62 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
+def run_in_one_core_kwarg_comm(func):
+    """Decorator to execute function only on rank 0.
+
+    Ensures the decorated function runs only on the root process (rank 0)
+    of the communicator. Other processes skip execution. Useful for I/O
+    operations and serial tasks in parallel environments.
+
+    Parameters
+    ----------
+    func : callable
+        Function to decorate. The wrapped function must receive a `comm`
+        keyword argument containing an MPI communicator.
+
+    Returns
+    -------
+    callable
+        Wrapped function that executes only on rank 0.
+
+    Notes
+    -----
+    The function checks for two types of communicators:
+    - Ensemble communicator: Runs only if both `ensemble_comm.rank` == 0
+      and comm.rank == 0.
+    - Regular communicator: Runs only if comm.rank == 0.
+    - If comm is None, the function runs normally without restrictions.
+
+    The function does not broadcast results to other processes.
+
+    See Also
+    --------
+    run_in_one_core_and_broadcast : Similar decorator that also broadcasts results.
+
+    Examples
+    --------
+    >>> @run_in_one_core
+    ... def save_file(obj, filename):
+    ...     # Only rank 0 writes the file
+    ...     with open(filename, 'w') as f:
+    ...         f.write(str(obj.data))
+    """
+
+    def wrapper(*args, **kwargs):
+        comm = kwargs.get("comm")
+        if comm is None:
+            return func(*args, **kwargs)
+        else:
+            if getattr(comm, "ensemble_comm", None) is not None:
+                if comm.ensemble_comm.rank == 0 and comm.comm.rank == 0:
+                    return func(*args, **kwargs)
+            elif getattr(comm, "rank", None) is not None:
+                if comm.rank == 0:
+                    return func(*args, **kwargs)
+
+    return wrapper
+
+
 def run_in_one_core_and_broadcast(func):
     """Decorator to execute function on rank 0 and broadcast result.
 

From 10e36ef10d149c9f4924c1c32cd4c8e4364a4599 Mon Sep 17 00:00:00 2001
From: Olender <olender@usp.br>
Date: Tue, 14 Jul 2026 21:41:35 +0100
Subject: [PATCH 4/7] added export scalar field function

---
 spyro/io/segy_io.py | 44 +++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 43 insertions(+), 1 deletion(-)

diff --git a/spyro/io/segy_io.py b/spyro/io/segy_io.py
index c8f983c0..f6e0189b 100644
--- a/spyro/io/segy_io.py
+++ b/spyro/io/segy_io.py
@@ -1,7 +1,8 @@
 import numpy as np
+from pathlib import Path
 import segyio
 from io import BytesIO
-from ..io import write_function_to_grid
+from .parallelism_wrappers import run_in_one_core_kwarg_comm
 import matplotlib.pyplot as plt
 
 
@@ -139,6 +140,47 @@ def create_segy(function, V, grid_spacing, filename):
     -------
     None
     """
+    from ..io import write_function_to_grid  # Avoiding circular import
     velocity_grid_data = write_function_to_grid(function, V, grid_spacing, buffer=True)
 
     return create_segy_from_grid(velocity_grid_data, filename)
+
+
+@run_in_one_core_kwarg_comm
+def export_scalar_field(function, grid_spacing, output_filename, comm=None):
+    """Export a scalar field to SEG-Y or PNG.
+
+    Parameters
+    ----------
+    function : firedrake.Function
+        Scalar field to export.
+    grid_spacing : float
+        Spacing of output grid points used when sampling ``function``.
+    output_filename : str
+        Destination filename. Supported extensions are ``.segy``/``.sgy`` and
+        ``.png``.
+
+    Returns
+    -------
+    None
+        Writes output files to disk. For ``.png`` output, an intermediate
+        SEG-Y file is created alongside the PNG.
+
+    Notes
+    -----
+    - ``.segy`` and ``.sgy`` outputs are written directly from the sampled
+      field.
+    - ``.png`` output is generated by first writing a temporary SEG-Y file
+      with the same basename and then converting it to PNG.
+    """
+    # if output_filename.lower().endswith((".png", ".segy", ".sgy")):
+    V = function.function_space()
+    if output_filename.lower().endswith((".segy", ".sgy")):
+        create_segy(function, V, grid_spacing, output_filename)
+    elif output_filename.lower().endswith((".png")):
+        segy_filename = Path(output_filename).with_suffix(".segy")
+        create_segy(function, V, grid_spacing, segy_filename)
+        segy_to_png(segy_filename, output_file=output_filename)
+    else:
+        valid_extensions = [".segy", ".sgy", ".png"]
+        raise ValueError(f"Invalid extension, please use {valid_extensions}")

From a33ed345780363d0a126f9128506894e4925697c Mon Sep 17 00:00:00 2001
From: Olender <olender@usp.br>
Date: Wed, 15 Jul 2026 14:20:09 +0100
Subject: [PATCH 5/7] linting

---
 spyro/io/parallelism_wrappers.py | 3 +--
 spyro/utils/utils.py             | 1 -
 2 files changed, 1 insertion(+), 3 deletions(-)

diff --git a/spyro/io/parallelism_wrappers.py b/spyro/io/parallelism_wrappers.py
index 50c6993e..da5b559f 100644
--- a/spyro/io/parallelism_wrappers.py
+++ b/spyro/io/parallelism_wrappers.py
@@ -4,6 +4,7 @@
 import numpy as np
 import firedrake as fire
 
+
 def run_in_one_core(func):
     """Decorator to execute function only on rank 0.
 
@@ -494,5 +495,3 @@ def wrapper(*args, **kwargs):
             return grad_total
 
     return wrapper
-
-
diff --git a/spyro/utils/utils.py b/spyro/utils/utils.py
index cb3f2475..e9b2edd2 100644
--- a/spyro/utils/utils.py
+++ b/spyro/utils/utils.py
@@ -9,7 +9,6 @@
 from ..io.basicio import parallel_print, write_velocity_model
 from ..io.parallelism_wrappers import (
     ensemble_functional,
-    run_in_one_core,
     run_in_one_core_and_broadcast,
 )
 from ..domains.space import create_function_space

From 571cc2e37bc6b2dc7fd7e0d18a9723b46ed8ff50 Mon Sep 17 00:00:00 2001
From: Olender <olender@usp.br>
Date: Thu, 16 Jul 2026 16:52:51 +0100
Subject: [PATCH 6/7] fixing imports

---
 spyro/meshing/gmsh_based_methods.py | 2 +-
 spyro/solvers/inversion.py          | 2 +-
 spyro/utils/__init__.py             | 2 --
 3 files changed, 2 insertions(+), 4 deletions(-)

diff --git a/spyro/meshing/gmsh_based_methods.py b/spyro/meshing/gmsh_based_methods.py
index e4581671..b700f773 100644
--- a/spyro/meshing/gmsh_based_methods.py
+++ b/spyro/meshing/gmsh_based_methods.py
@@ -1,6 +1,6 @@
 import numpy as np
 from scipy.interpolate import RegularGridInterpolator
-from ..utils import run_in_one_core
+from ..io.parallelism_wrappers import run_in_one_core
 from .meshing_utils import check_gmsh, vp_to_sizing
 
 try:
diff --git a/spyro/solvers/inversion.py b/spyro/solvers/inversion.py
index d2bafa9c..740ba371 100644
--- a/spyro/solvers/inversion.py
+++ b/spyro/solvers/inversion.py
@@ -17,7 +17,7 @@
 from ..io.basicio import load_shots, save_shots
 from ..io.parallelism_wrappers import switch_serial_shot
 from ..io import create_segy
-from ..utils import run_in_one_core
+from ..io.parallelism_wrappers import run_in_one_core
 
 
 try:
diff --git a/spyro/utils/__init__.py b/spyro/utils/__init__.py
index 2f4a9c04..e3c3c66c 100644
--- a/spyro/utils/__init__.py
+++ b/spyro/utils/__init__.py
@@ -4,7 +4,6 @@
     compute_functional,
     Mask,
     Gradient_mask_for_pml,
-    run_in_one_core,
     write_hdf5_velocity_model,
     get_real_shot_record,
 )
@@ -21,7 +20,6 @@
     "Mask",
     "Gradient_mask_for_pml",
     "velocity_to_grid",
-    "run_in_one_core",
     "change_scalar_field_resolution",
     "write_hdf5_velocity_model",
     "get_real_shot_record",

From ed59a96657a467eaedfd9b75ad98080f101d8681 Mon Sep 17 00:00:00 2001
From: Olender <olender@usp.br>
Date: Thu, 16 Jul 2026 16:58:34 +0100
Subject: [PATCH 7/7] adding a test that will becoma a demo

---
 .github/workflows/python-tests.yml            |   7 +-
 .../test_simple_fwi_with_segy_output.py       | 230 ++++++++++++++++++
 2 files changed, 236 insertions(+), 1 deletion(-)
 create mode 100644 tests/parallel/test_simple_fwi_with_segy_output.py

diff --git a/.github/workflows/python-tests.yml b/.github/workflows/python-tests.yml
index a9285c23..ae62987e 100644
--- a/.github/workflows/python-tests.yml
+++ b/.github/workflows/python-tests.yml
@@ -148,6 +148,7 @@ jobs:
           mpiexec -n 6 pytest tests/parallel/test_fwi.py
           mpiexec -n 3 pytest tests/parallel/test_supershot_grad.py
           mpiexec -n 2 pytest tests/parallel/test_gradient_serialshots.py
+          mpiexec -n 6 pytest tests/parallel/test_simple_fwi_with_segy_output.py
 
     - name: Covering parallel 3D forward test
       continue-on-error: true
@@ -180,7 +181,11 @@ jobs:
     - name: Covering spatially parallelized shots in serial
       continue-on-error: true
       run: |
-          mpiexec -n 2 pytest --cov-report=xml --cov-append --cov=spyro tests/parallel/test_test_gradient_serialshots.py
+          mpiexec -n 2 pytest --cov-report=xml --cov-append --cov=spyro tests/parallel/test_gradient_serialshots.py
+    - name: Covering source parallelized fwi with segy output
+      continue-on-error: true
+      run: |
+          mpiexec -n 2 pytest --cov-report=xml --cov-append --cov=spyro tests/parallel/test_simple_fwi_with_segy_output.py
     - name: Uploading coverage to Codecov
       run: |
         export CODECOV_TOKEN="6cd21147-54f7-4b77-94ad-4b138053401d" && curl -s https://codecov.io/bash | bash
diff --git a/tests/parallel/test_simple_fwi_with_segy_output.py b/tests/parallel/test_simple_fwi_with_segy_output.py
new file mode 100644
index 00000000..140cdf47
--- /dev/null
+++ b/tests/parallel/test_simple_fwi_with_segy_output.py
@@ -0,0 +1,230 @@
+"""Demo script for running a small full waveform inversion example.
+
+The script builds a synthetic "true" model, generates a shot record, and then
+runs a simple inversion loop against that record.
+"""
+# from mpi4py.MPI import COMM_WORLD
+# import debugpy
+# debugpy.listen(3000 + COMM_WORLD.rank)
+# debugpy.wait_for_client()
+
+from copy import deepcopy
+import numpy as np
+import firedrake as fire
+import spyro
+import pytest
+
+
+def run_forward_real_model(input_dictionary, case="camembert", shot_filename="shots/shot_record_", dt=None):
+    """Generate and save a synthetic shot record for the chosen demo case.
+
+    Parameters
+    ----------
+    input_dictionary : dict
+        Configuration dictionary used to build the forward-modeling object.
+    case : str, optional
+        Demo model to generate. Currently only ``"camembert"`` is supported.
+    output_filename : str, optional
+        Base filename used when saving the generated shot record with NumPy.
+
+    Returns
+    -------
+    None
+        The generated shot record is written to disk.
+    """
+    if dt is not None:
+        original_dt = deepcopy(input_dictionary["time_axis"]["dt"])
+        input_dictionary["time_axis"]["dt"] = dt
+
+    fwi_obj = spyro.FullWaveformInversion(dictionary=input_dictionary)
+
+    fwi_obj.set_real_mesh(input_mesh_parameters={"edge_length": 0.05})
+
+    supported_cases = ["camembert"]
+    if case == "camembert":
+        # Builds the true velocity model based on a conditional
+    
+        center_z = -1.0
+        center_x = 1.0
+        mesh_z = fwi_obj.wave.mesh_z
+        mesh_x = fwi_obj.wave.mesh_x
+        cond = fire.conditional((mesh_z-center_z)**2 + (mesh_x-center_x)**2 < .2**2, 3.0, 2.5)
+    # elif case == "layers":
+    #     # not yet done here
+    #     # Works for any number of horizontal layers and velocity values
+    #     z_switch = [-1.0]  # List of floats representing z value where vp changes
+    #     layer_vps = [2.5, 3.0]  # List of vp values
+    #     cond = multiple_layer_velocity_model(fwi_obj, z_switch, layer_vps)
+    elif case not in supported_cases:
+        return ValueError(f"Case of {case} not part of supported cases: {supported_cases}")
+    else:
+        return ValueError(f"Case of {case} only partially implemented.")
+
+    fwi_obj.set_real_velocity_model(conditional=cond, output=True, dg_velocity_model=False)
+    fwi_obj.generate_real_shot_record(
+        plot_model=True,
+        model_filename="True_experiment.png",
+        shot_filename=shot_filename,
+        abc_points=[(-0.5, 0.5), (-1.5, 0.5), (-1.5, 1.5), (-0.5, 1.5)]
+    )
+    if dt is not None:
+        fwi_obj.wave.dt = original_dt
+
+    return fwi_obj
+
+
+def multiple_layer_velocity_model(fwi_obj, z_switch, layers):
+    """
+    Sets the heterogeneous velocity model to be split into horizontal layers.
+    Each layer's velocity value is defined by the corresponding value in the
+    layers list. The layers are separated by the values in the z_switch list.
+
+    Parameters
+    ----------
+    z_switch : list of floats
+        List of z values that separate the layers.
+    layers : list of floats
+        List of velocity values for each layer.
+    """
+    if len(z_switch) != (len(layers) - 1):
+        raise ValueError(
+            "Float list of z_switch has to have length exactly one less \
+                            than list of layer values"
+        )
+    if len(z_switch) == 0:
+        raise ValueError("Float list of z_switch cannot be empty")
+    for i in range(len(z_switch)):
+        if i == 0:
+            cond = fire.conditional(
+                fwi_obj.wave.mesh_z > z_switch[i], layers[i], layers[i + 1]
+            )
+        else:
+            cond = fire.conditional(
+                fwi_obj.wave.mesh_z > z_switch[i], cond, layers[i + 1]
+            )
+
+    return cond
+
+
+def test_camembert_fwi_with_output(load_real_shot=False):
+    """Run the demo inversion workflow.
+
+    Parameters
+    ----------
+    load_real_shot : bool, optional
+        If ``True``, load the saved shot record from disk. If ``False``,
+        generate a fresh synthetic shot record first.
+
+    Returns
+    -------
+    None
+        The inversion is run for its side effects.
+    """
+    final_time = 0.9
+    real_shot_record_dt = 0.0005
+    simulation_dt = 0.001
+
+    dictionary = {}
+    dictionary["options"] = {
+        "cell_type": "T",  # simplexes such as triangles or tetrahedra (T) or quadrilaterals (Q)
+        "variant": "lumped",  # lumped, equispaced or DG, default is lumped
+        "degree": 4,  # p order
+        "dimension": 2,  # dimension
+    }
+    dictionary["parallelism"] = {
+        "type": "automatic",  # options: automatic (same number of cores for evey processor) or spatial
+    }
+    dictionary["mesh"] = {
+        "length_z": 2.0,  # depth in km - always positive
+        "length_x": 2.0,  # width in km - always positive
+        "length_y": 0.0,  # thickness in km - always positive
+        "mesh_file": None,
+        "mesh_type": "firedrake_mesh",
+    }
+    dictionary["acquisition"] = {
+        "source_type": "ricker",
+        "source_locations": spyro.create_transect((-0.55, 0.7), (-0.55, 1.3), 6),
+        # "source_locations": [(-1.1, 1.5)],
+        "frequency": 5.0,
+        "delay": 0.2,
+        "delay_type": "time",
+        "receiver_locations": spyro.create_transect((-1.45, 0.7), (-1.45, 1.3), 200),
+        "use_vertex_only_mesh": True,
+    }
+    dictionary["time_axis"] = {
+        "initial_time": 0.0,  # Initial time for event
+        "final_time": final_time,  # Final time for event
+        "dt": simulation_dt,  # timestep size
+        "amplitude": 1,  # the Ricker has an amplitude of 1.
+        "output_frequency": 100,  # how frequently to output solution to pvds
+        "gradient_sampling_frequency": 1,  # how frequently to save solution to RAM
+    }
+    dictionary["visualization"] = {
+        "forward_output": False,
+        "forward_output_filename": "results/forward_output.pvd",
+        "fwi_velocity_model_output": False,
+        "velocity_model_filename": None,
+        "gradient_output": False,
+        "gradient_filename": "results/Gradient.pvd",
+        "adjoint_output": False,
+        "adjoint_filename": None,
+        "debug_output": False,
+    }
+    dictionary["inversion"] = {
+        "perform_fwi": True,  # switch to true to make a FWI
+        "initial_guess_model_file": None,
+        "shot_record_file": None,
+    }
+
+    shots_filenames="shots/shot_record_"
+
+    # Setting up to run synthetic real problem
+    if load_real_shot is False:
+        fwi_obj = run_forward_real_model(
+            dictionary,
+            dt=real_shot_record_dt,
+            case="camembert",
+            shot_filename=shots_filenames,
+        )
+
+    else:
+        dictionary["time_axis"]["dt"] = simulation_dt
+        dictionary["inversion"]["real_shot_record_file"] = shots_filenames
+        fwi_obj = spyro.FullWaveformInversion(dictionary=dictionary)
+    # Since the shot record is using a different timestep than our guess model we have to interpolate it
+
+    # for rec_id in [3]:
+    #     spyro.plots.plot_receiver_response(fwi_obj.wave.real_shot_record[:, rec_id], final_time, hold=True, name=f"unchanged{rec_id}")
+
+    fwi_obj.real_shot_record = spyro.io.time_io.interpolate_time_series(
+        fwi_obj.real_shot_record,
+        simulation_dt,
+        0.0,
+        final_time,
+    )
+
+    # for rec_id in [3]:
+    #     spyro.plots.plot_receiver_response(fwi_obj.wave.real_shot_record[:, rec_id], final_time, filename="changed_receiver.png", hold=True, name=f"changed{rec_id}", linestyle="--")
+    # Setting up initial guess problem
+    fwi_obj.set_guess_mesh(input_mesh_parameters={"edge_length": 0.2})
+    fwi_obj.set_guess_velocity_model(constant=2.5)
+
+    # This is deprecated, in more complex cases we mark zero gradient boundaries directly on the mesh
+    # which is a lot more efficient
+    mask_boundaries = {
+        "z_min": -1.3,
+        "z_max": -0.7,
+        "x_min": 0.7,
+        "x_max": 1.3,
+    }
+    fwi_obj.set_gradient_mask(boundaries=mask_boundaries)
+
+    fwi_obj.run_fwi(vmin=2.5, vmax=3.0, maxiter=2)
+    export_grid_spacing = 0.01
+    spyro.io.export_scalar_field(fwi_obj.wave.c, export_grid_spacing, "camembert.png", comm=fwi_obj.wave.comm)
+
+    print("END", flush=True)
+
+
+if __name__ == "__main__":
+    test_camembert_fwi_with_output(load_real_shot=False)