OptimalTransportPlan.h

Go to the documentation of this file.

#ifndef OPTIMAL_TRANSPORT_PLAN_H
#define OPTIMAL_TRANSPORT_PLAN_H
 
#include <memory>
#include <string>
#include <vector>
#include <map>
#include <filesystem>
#include <fstream>
#include <algorithm>
#include <cmath>
#include <stdexcept>
 
#include <deal.II/base/point.h>
#include <deal.II/base/quadrature_lib.h>
#include <deal.II/base/utilities.h>
#include <deal.II/base/function.h>
#include <deal.II/base/parameter_acceptor.h>
#include <deal.II/lac/vector.h>
#include <deal.II/grid/tria.h>
#include <deal.II/fe/fe_values.h>
#include <deal.II/numerics/rtree.h>
#include <deal.II/numerics/vector_tools.h>
#include <deal.II/numerics/data_out.h>
 
#include "SemiDiscreteOT/solvers/Distance.h"
#include "SemiDiscreteOT/utils/utils.h"
 
 
using namespace dealii;
 
namespace OptimalTransportPlanSpace {
 
// Forward declarations
template <int spacedim> class MapApproximationStrategy;
 
template <int spacedim>
class OptimalTransportPlan : public ParameterAcceptor {
public:
    OptimalTransportPlan(const std::string& strategy_name = "modal");
 
    void set_distance_function(const std::function<double(const Point<spacedim>&, const Point<spacedim>&)>& dist)
    {
        distance_function = dist;
    }
 
    void set_source_measure(const std::vector<Point<spacedim>>& points,
                          const std::vector<double>& density);
 
    void set_target_measure(const std::vector<Point<spacedim>>& points,
                          const std::vector<double>& density);
 
    void set_potential(const Vector<double>& potential,
                      const double regularization_param = 0.0);
 
    void set_truncation_radius(double radius);
 
    void compute_map();
 
    void save_map(const std::string& output_dir) const;
 
    void set_strategy(const std::string& strategy_name);
 
    static std::vector<std::string> get_available_strategies();
 
private:
    // Data members
    std::vector<Point<spacedim>> source_points;
    std::vector<double> source_density;
    std::vector<Point<spacedim>> target_points;
    std::vector<double> target_density;
    Vector<double> transport_potential;
    double epsilon;
    double truncation_radius = -1.0;  // Negative means no truncation
 
    // Distance function
    std::function<double(const Point<spacedim>&, const Point<spacedim>&)> distance_function;
 
    // Strategy pattern implementation
    std::unique_ptr<MapApproximationStrategy<spacedim>> strategy;
 
    // Factory method to create strategies
    static std::unique_ptr<MapApproximationStrategy<spacedim>> 
    create_strategy(const std::string& name);
};
 
template <int spacedim>
class MapApproximationStrategy {
public:
    virtual ~MapApproximationStrategy() = default;
 
    virtual void compute_map(
        const std::function<double(const Point<spacedim>&, const Point<spacedim>&)> distance_function,
        const std::vector<Point<spacedim>>& source_points,
        const std::vector<double>& source_density,
        const std::vector<Point<spacedim>>& target_points,
        const std::vector<double>& target_density,
        const Vector<double>& potential,
        const double regularization_param,
        const double truncation_radius) = 0;
 
    virtual void save_results(const std::string& output_dir) const = 0;
 
protected:
    std::vector<Point<spacedim>> source_points; 
    std::vector<Point<spacedim>> mapped_points; 
    std::vector<double> transport_density; 
};
 
 template <int spacedim>
 class ModalStrategy : public MapApproximationStrategy<spacedim> {
 public:
     void compute_map(
        const std::function<double(const Point<spacedim>&, const Point<spacedim>&)> distance_function,
        const std::vector<Point<spacedim>>& source_points,
        const std::vector<double>& source_density,
        const std::vector<Point<spacedim>>& target_points,
        const std::vector<double>& target_density,
        const Vector<double>& potential,
        const double regularization_param,
        const double truncation_radius) override;
 
     void save_results(const std::string& output_dir) const override;
 };
 
template <int spacedim>
class BarycentricStrategy : public MapApproximationStrategy<spacedim> {
public:
    void compute_map(
        const std::function<double(const Point<spacedim>&, const Point<spacedim>&)> distance_function,
        const std::vector<Point<spacedim>>& source_points,
        const std::vector<double>& source_density,
        const std::vector<Point<spacedim>>& target_points,
        const std::vector<double>& target_density,
        const Vector<double>& potential,
        const double regularization_param,
        const double truncation_radius) override;
 
    void save_results(const std::string& output_dir) const override;
};
 
} // namespace OptimalTransportPlanSpace
 
#endif // OPTIMAL_TRANSPORT_PLAN_H