Parallel ensemble I/O for State and Increment

src/soca/Geometry/soca_geom_mod.F90

@@ -17,7 +17,7 @@ module soca_geom_mod
 
 ! mom6 / fms modules
 use fms_mod, only : fms_init, fms_end
-use soca_io_mod, only : soca_io_writer, soca_io_reader
+use soca_io_mod, only : soca_io_writer, soca_io_reader, soca_io_config_from_yaml
 use MOM, only : MOM_control_struct, initialize_MOM, MOM_end, get_MOM_state_elements
 use MOM_restart, only :MOM_restart_CS ! NOTE remove this when updating MOM6
 use MOM_domains, only : MOM_domain_type, MOM_domains_init, MOM_infra_init, MOM_infra_end
@@ -181,6 +181,10 @@ subroutine soca_geom_init(self, f_conf, f_comm, gen)
   ! MPI communicator
   self%f_comm = f_comm
 
+  ! Resolve the soca_io dispatch knobs (parallel/sequential, scatter/strided,
+  ! async mpi) from geometry.io once; values persist module-level for the run.
+  call soca_io_config_from_yaml(f_conf)
+
   ! use MOM6 to setup domain decomposition
   call mpp_init(localcomm=f_comm%communicator())
   call fms_init()

src/soca/Increment/Increment.cc

@@ -19,6 +19,7 @@
 #include "soca/Increment/IncrementFortran.h"
 #include "soca/State/State.h"
 
+#include "eckit/config/LocalConfiguration.h"
 #include "eckit/exception/Exceptions.h"
 
 #include "oops/base/GeometryData.h"
@@ -373,6 +374,55 @@ namespace soca {
 
   // -----------------------------------------------------------------------------
 
+  void Increment::writeEnsemble(const std::vector<const Increment*> & increments,
+                                const std::vector<eckit::LocalConfiguration> & configs) {
+    Log::trace() << "soca::Increment::writeEnsemble starting (" << increments.size()
+                 << " members)" << std::endl;
+    ASSERT(increments.size() == configs.size());
+    if (increments.empty()) return;
+
+    const size_t n = increments.size();
+    std::vector<F90flds> keys(n);
+    std::vector<const eckit::Configuration*> confPtrs(n);
+    std::vector<const util::DateTime*> dtPtrs(n);
+    for (size_t i = 0; i < n; ++i) {
+      keys[i]     = increments[i]->keyFlds_;
+      confPtrs[i] = &configs[i];
+      dtPtrs[i]   = &increments[i]->time_;
+    }
+    const int nm = static_cast<int>(n);
+    soca_increment_write_ensemble_f90(nm, keys.data(), confPtrs.data(), dtPtrs.data());
+    Log::trace() << "soca::Increment::writeEnsemble done" << std::endl;
+  }
+
+  // -----------------------------------------------------------------------------
+
+  void Increment::readEnsemble(const std::vector<Increment*> & increments,
+                               const std::vector<eckit::LocalConfiguration> & configs) {
+    Log::trace() << "soca::Increment::readEnsemble starting (" << increments.size()
+                 << " members)" << std::endl;
+    ASSERT(increments.size() == configs.size());
+    if (increments.empty()) return;
+
+    const size_t n = increments.size();
+    std::vector<F90flds> keys(n);
+    std::vector<const eckit::Configuration*> confPtrs(n);
+    std::vector<util::DateTime*> dtPtrs(n);
+    for (size_t i = 0; i < n; ++i) {
+      keys[i]     = increments[i]->keyFlds_;
+      confPtrs[i] = &configs[i];
+      dtPtrs[i]   = &increments[i]->time_;
+    }
+    const int nm = static_cast<int>(n);
+    soca_increment_read_ensemble_f90(nm, keys.data(), confPtrs.data(), dtPtrs.data());
+
+    for (auto* x : increments) x->fieldSet_.set_dirty();
+
+    Log::trace() << "soca::Increment::readEnsemble done" << std::endl;
+  }
+
+  // -----------------------------------------------------------------------------
+
   void Increment::horiz_scales(const eckit::Configuration & config) {
     soca_increment_horiz_scales_f90(toFortran(), &config);
     Log::trace() << "Horiz decorrelation length scales computed." << std::endl;

src/soca/Increment/Increment.h

@@ -94,6 +94,18 @@ namespace soca {
       /// I/O and diagnostics
       void read(const eckit::Configuration &);
       void write(const eckit::Configuration &) const;
+
+      /// Bulk parallel write across ensemble members. Mirrors State::writeEnsemble
+      /// for ensemble-of-increments output (e.g. LETKF posterior ensemble
+      /// increments). Each member is gathered to a strided writer PE and
+      /// written concurrently.
+      static void writeEnsemble(const std::vector<const Increment*> &,
+                                const std::vector<eckit::LocalConfiguration> &);
+
+      /// Bulk read across ensemble members. Mirrors State::readEnsemble.
+      static void readEnsemble(const std::vector<Increment*> &,
+                               const std::vector<eckit::LocalConfiguration> &);
+
       void horiz_scales(const eckit::Configuration &);
       void vert_scales(const double &);
       std::vector<double> rmsByLevel(const std::string &) const;

src/soca/Increment/IncrementFortran.h

@@ -46,6 +46,14 @@ namespace soca {
     void soca_increment_horiz_scales_f90(F90flds &,
                                          const eckit::Configuration * const &);
     void soca_increment_vert_scales_f90(F90flds &, const double &);
+    void soca_increment_write_ensemble_f90(const int &,
+                                           const F90flds * const,
+                                           const eckit::Configuration * const *,
+                                           const util::DateTime * const *);
+    void soca_increment_read_ensemble_f90(const int &,
+                                          const F90flds * const,
+                                          const eckit::Configuration * const *,
+                                          util::DateTime * const *);
   }
 }  // namespace soca
 #endif  // SOCA_INCREMENT_INCREMENTFORTRAN_H_

src/soca/Increment/soca_increment.interface.F90

@@ -14,6 +14,8 @@ module soca_increment_mod_c
 use oops_variables_mod, only : oops_variables
 
 ! soca modules
+use soca_fields_mod, only: soca_fields, soca_fields_ptr_t, &
+                           soca_fields_write_ensemble, soca_fields_read_ensemble
 use soca_geom_mod_c, only: soca_geom_registry
 use soca_geom_mod, only: soca_geom
 use soca_increment_mod, only : soca_increment
@@ -197,5 +199,64 @@ subroutine soca_increment_vert_scales_c(c_key_self, c_vert) bind(c,name='soca_in
 
 end subroutine soca_increment_vert_scales_c
 
+
+! ------------------------------------------------------------------------------
+!> C++ interface: bulk write multiple increments in one ensemble pass. Each
+!! member is gathered to a rotated writer PE and written concurrently with the
+!! others.
+subroutine soca_increment_write_ensemble_c(c_n, c_keys, c_confs, c_dts) &
+    bind(c, name='soca_increment_write_ensemble_f90')
+  integer(c_int),       intent(in) :: c_n
+  integer(c_int),       intent(in) :: c_keys(*)
+  type(c_ptr),          intent(in) :: c_confs(*)
+  type(c_ptr),          intent(in) :: c_dts(*)
+
+  type(soca_fields_ptr_t), allocatable :: fptrs(:)
+  type(c_ptr),             allocatable :: confs_local(:)
+  type(datetime),          allocatable :: fdates(:)
+  type(soca_increment), pointer :: fld
+  integer :: i
+
+  if (c_n <= 0) return
+  allocate(fptrs(c_n), confs_local(c_n), fdates(c_n))
+  do i = 1, c_n
+    call soca_increment_registry%get(c_keys(i), fld)
+    fptrs(i)%p => fld
+    confs_local(i) = c_confs(i)
+    call c_f_datetime(c_dts(i), fdates(i))
+  end do
+
+  call soca_fields_write_ensemble(fptrs, confs_local, fdates)
+end subroutine soca_increment_write_ensemble_c
+
+
+! ------------------------------------------------------------------------------
+!> C++ interface: bulk read multiple increments in one ensemble pass.
+subroutine soca_increment_read_ensemble_c(c_n, c_keys, c_confs, c_dts) &
+    bind(c, name='soca_increment_read_ensemble_f90')
+  integer(c_int),       intent(in) :: c_n
+  integer(c_int),       intent(in) :: c_keys(*)
+  type(c_ptr),          intent(in) :: c_confs(*)
+  type(c_ptr),       intent(inout) :: c_dts(*)
+
+  type(soca_fields_ptr_t), allocatable :: fptrs(:)
+  type(c_ptr),             allocatable :: confs_local(:)
+  type(datetime),          allocatable :: fdates(:)
+  type(soca_increment), pointer :: fld
+  integer :: i
+
+  if (c_n <= 0) return
+  allocate(fptrs(c_n), confs_local(c_n), fdates(c_n))
+  do i = 1, c_n
+    call soca_increment_registry%get(c_keys(i), fld)
+    fptrs(i)%p => fld
+    confs_local(i) = c_confs(i)
+    call c_f_datetime(c_dts(i), fdates(i))
+  end do
+
+  call soca_fields_read_ensemble(fptrs, confs_local, fdates)
+end subroutine soca_increment_read_ensemble_c
+
+
 ! ------------------------------------------------------------------------------
 end module

src/soca/State/State.cc

@@ -289,6 +289,59 @@ namespace soca {
 
   // -----------------------------------------------------------------------------
 
+  void State::writeEnsemble(const std::vector<const State*> & states,
+                            const std::vector<eckit::LocalConfiguration> & configs) {
+    Log::trace() << "soca::State::writeEnsemble starting (" << states.size()
+                 << " members)" << std::endl;
+    ASSERT(states.size() == configs.size());
+    if (states.empty()) return;
+
+    const size_t n = states.size();
+    std::vector<F90flds> keys(n);
+    std::vector<const eckit::Configuration*> confPtrs(n);
+    std::vector<const util::DateTime*> dtPtrs(n);
+    for (size_t i = 0; i < n; ++i) {
+      keys[i]     = states[i]->keyFlds_;
+      confPtrs[i] = &configs[i];
+      dtPtrs[i]   = &states[i]->time_;
+    }
+    const int nm = static_cast<int>(n);
+    soca_state_write_ensemble_f90(nm, keys.data(), confPtrs.data(), dtPtrs.data());
+    Log::trace() << "soca::State::writeEnsemble done" << std::endl;
+  }
+
+  // -----------------------------------------------------------------------------
+
+  void State::readEnsemble(const std::vector<State*> & states,
+                           const std::vector<eckit::LocalConfiguration> & configs) {
+    Log::trace() << "soca::State::readEnsemble starting (" << states.size()
+                 << " members)" << std::endl;
+    ASSERT(states.size() == configs.size());
+    if (states.empty()) return;
+
+    const size_t n = states.size();
+    std::vector<F90flds> keys(n);
+    std::vector<const eckit::Configuration*> confPtrs(n);
+    std::vector<util::DateTime*> dtPtrs(n);
+    for (size_t i = 0; i < n; ++i) {
+      keys[i]     = states[i]->keyFlds_;
+      confPtrs[i] = &configs[i];
+      dtPtrs[i]   = &states[i]->time_;
+    }
+    const int nm = static_cast<int>(n);
+    soca_state_read_ensemble_f90(nm, keys.data(), confPtrs.data(), dtPtrs.data());
+
+    // Fortran-side soca_fields_read_finalize marks each atlas::Field dirty,
+    // but the C++ atlas::FieldSet wrapper keeps its own dirty-bit cache that
+    // doesn't auto-update from the underlying fields. Mirror the writeEnsemble
+    // pattern and re-mark here so downstream FieldSet consumers refresh.
+    for (auto* x : states) x->fieldSet_.set_dirty();
+
+    Log::trace() << "soca::State::readEnsemble done" << std::endl;
+  }
+
+  // -----------------------------------------------------------------------------
+
   void State::updateFields(const oops::Variables & vars) {
     // remove fields from the fieldset that are no longer in vars
     atlas::FieldSet orig = util::shareFields(fieldSet_);

src/soca/State/State.h

@@ -83,6 +83,21 @@ namespace soca {
       void read(const eckit::Configuration &);
       void write(const eckit::Configuration &) const;
 
+      /// Bulk parallel write across ensemble members. Each member is gathered
+      /// to a strided writer PE (assigned by soca_io_ensemble_root_pe in
+      /// soca_io_mod) and written in phase 2 with no MPI, so per-member
+      /// netCDF writes happen concurrently across writer PEs. See
+      /// soca_io_writers_commit_ensemble for the staging.
+      static void writeEnsemble(const std::vector<const State*> &,
+                                const std::vector<eckit::LocalConfiguration> &);
+
+      /// Bulk read across ensemble members via soca_io_readers_commit_ensemble.
+      /// Members are read in parallel (rotated reader PEs), honoring the
+      /// geometry.io.'ensemble read' (scatter vs strided) and 'async mpi'
+      /// toggles. See soca_io_readers_commit_ensemble for the staging.
+      static void readEnsemble(const std::vector<State*> &,
+                               const std::vector<eckit::LocalConfiguration> &);
+
       int & toFortran() {return keyFlds_;}
       const int & toFortran() const {return keyFlds_;}
 

src/soca/State/StateFortran.h

@@ -39,6 +39,14 @@ namespace soca {
                                  const eckit::Configuration * const &,
                                  util::DateTime * const *);
     void soca_state_update_fields_f90(const F90flds &, const oops::Variables &);
+    void soca_state_write_ensemble_f90(const int &,
+                                       const F90flds * const,
+                                       const eckit::Configuration * const *,
+                                       const util::DateTime * const *);
+    void soca_state_read_ensemble_f90(const int &,
+                                      const F90flds * const,
+                                      const eckit::Configuration * const *,
+                                      util::DateTime * const *);
   }
 }  // namespace soca
 #endif  // SOCA_STATE_STATEFORTRAN_H_

src/soca/State/soca_state.interface.F90

@@ -14,7 +14,8 @@ module soca_state_mod_c
 use oops_variables_mod, only: oops_variables
 
 ! soca modules
-use soca_fields_mod, only: soca_field
+use soca_fields_mod, only: soca_field, soca_fields_ptr_t, &
+                           soca_fields_write_ensemble, soca_fields_read_ensemble
 use soca_geom_mod_c, only: soca_geom_registry
 use soca_geom_mod, only: soca_geom
 use soca_increment_mod, only: soca_increment
@@ -167,4 +168,62 @@ subroutine soca_state_update_fields_c(c_key_self, c_vars) &
 
 end subroutine soca_state_update_fields_c
 
+
+! ------------------------------------------------------------------------------
+!> C++ interface: bulk write multiple states in one ensemble pass. Each member
+!! is gathered to a rotated writer PE and written concurrently with the others.
+subroutine soca_state_write_ensemble_c(c_n, c_keys, c_confs, c_dts) &
+    bind(c, name='soca_state_write_ensemble_f90')
+  integer(c_int),       intent(in) :: c_n
+  integer(c_int),       intent(in) :: c_keys(*)
+  type(c_ptr),          intent(in) :: c_confs(*)
+  type(c_ptr),          intent(in) :: c_dts(*)
+
+  type(soca_fields_ptr_t), allocatable :: fptrs(:)
+  type(c_ptr),             allocatable :: confs_local(:)
+  type(datetime),          allocatable :: fdates(:)
+  type(soca_state), pointer :: fld
+  integer :: i
+
+  if (c_n <= 0) return
+  allocate(fptrs(c_n), confs_local(c_n), fdates(c_n))
+  do i = 1, c_n
+    call soca_state_registry%get(c_keys(i), fld)
+    fptrs(i)%p => fld
+    confs_local(i) = c_confs(i)
+    call c_f_datetime(c_dts(i), fdates(i))
+  end do
+
+  call soca_fields_write_ensemble(fptrs, confs_local, fdates)
+end subroutine soca_state_write_ensemble_c
+
+
+! ------------------------------------------------------------------------------
+!> C++ interface: bulk read multiple states in one ensemble pass.
+subroutine soca_state_read_ensemble_c(c_n, c_keys, c_confs, c_dts) &
+    bind(c, name='soca_state_read_ensemble_f90')
+  integer(c_int),       intent(in) :: c_n
+  integer(c_int),       intent(in) :: c_keys(*)
+  type(c_ptr),          intent(in) :: c_confs(*)
+  type(c_ptr),       intent(inout) :: c_dts(*)
+
+  type(soca_fields_ptr_t), allocatable :: fptrs(:)
+  type(c_ptr),             allocatable :: confs_local(:)
+  type(datetime),          allocatable :: fdates(:)
+  type(soca_state), pointer :: fld
+  integer :: i
+
+  if (c_n <= 0) return
+  allocate(fptrs(c_n), confs_local(c_n), fdates(c_n))
+  do i = 1, c_n
+    call soca_state_registry%get(c_keys(i), fld)
+    fptrs(i)%p => fld
+    confs_local(i) = c_confs(i)
+    call c_f_datetime(c_dts(i), fdates(i))
+  end do
+
+  call soca_fields_read_ensemble(fptrs, confs_local, fdates)
+end subroutine soca_state_read_ensemble_c
+
+
 end module soca_state_mod_c