oem_mpi.h

Go to the documentation of this file.

#endif _ARTS_OEM_MPI_H_

#include "invlib/algebra.h"
#include "invlib/algebra/precision_matrix.h"
#include "invlib/algebra/solvers.h"
#include "invlib/interfaces/arts_wrapper.h"
#include "invlib/map.h"
#include "invlib/mpi/log.h"
#include "invlib/mpi/mpi_matrix.h"
#include "invlib/mpi/mpi_vector.h"
#include "invlib/optimization.h"
#include "invlib/profiling/timer.h"

namespace oem {

using MPIMatrix = invlib::Matrix<invlib::MPIMatrix<invlib::Timer<ArtsMatrix>>>;
using MPICovarianceMatrix = invlib::Matrix<
    invlib::MPIMatrix<invlib::Timer<ArtsCovarianceMatrixWrapper>>>;
using MPIVector = invlib::Vector<invlib::MPIVector<invlib::Timer<ArtsVector>>>;

template <typename ForwardModel>
using OEM_STANDARD_MPI = invlib::MAP<ForwardModel,
                                  OEMMatrix,
                                  MPICovarianceMatrix,
                                  MPICovarianceMatrix,
                                  OEMVector,
                                  Formulation::STANDARD>;

using LM_MPI = invlib::LevenbergMarquardt<Numeric, MPISparse, CG>;

class AgendaWrapperMPI {

 public:
  const unsigned int m, n;

  AgendaWrapperMPI(Workspace *ws_,
                   const Agenda *inversion_iterate_agenda,
                   Index m_,
                   Index n_)
      : ws_(ws),
        local_jacobian_(),
        inversion_iterate_agenda_(inversion_iterate_agenda),
        m(static_cast<unsigned int>(m_)),
        n(static_cast<unsigned int>(n_)) {}

  MPIMatrix Jacobian(const OEMVector &xi, OEMVector &yi) {
    yi.resize(m);
    inversion_iterate_agendaExecute(
        *ws_, yi, local_jacobian_, xi, 1, *inversion_iterate_agenda_);
    // Create MPI vector from local results, use conversion to vector
    // to broadcast local results.
    MPIVector yi_mpi(yi);
    yi = yi_mpi;

    MPIMatrix jacobian = local_jacobian_;
    return jacobian;
  }

  OEMVector evaluate(const OEMVector &xi) {
    Matrix dummy = local_jacobian_;
    OEMVector yi;
    yi.resize(m);
    inversion_iterate_agendaExecute(
        *ws_, yi, dummy, xi, 0, *inversion_iterate_agenda_);

    // Create MPI vector from local results, use conversion to vector
    // to broadcast local results.
    MPIVector yi_mpi = yi;
    yi = yi_mpi;
    return yi;
  }

 private:
  Workspace *ws_;
  OEMMatrix local_jacobian_;
  const Agenda *inversion_iterate_agenda_;
};
}       // namespace oem
#endif  // _ARTS_OEM_MPI_H_