timestepper.hpp

Go to the documentation of this file.

#include <petsc/private/tsimpl.h>        /*I "petscts.h"  I*/
#include <petscts.h>
#include <petscksp.h>
#include "mastereq.hpp"
#include <assert.h> 
#include <iostream> 
#include "defs.hpp"
#include "output.hpp"
#include "optimtarget.hpp"
#include <deque>
#pragma once
 
class TimeStepper{
  protected:
    Vec x; 
    Vec xadj; 
    Vec xprimal; 
    std::vector<std::vector<Vec>> trajectory_states; 
    std::vector<Vec> final_states; 
    std::vector<Vec> dpdm_states; 
    int mpirank_world; 
    int mpisize_petsc; 
    int mpirank_petsc; 
    PetscInt localsize_u; 
    PetscInt ilow; 
    PetscInt iupp; 
 
    bool eval_leakage; 
    bool eval_energy; 
    bool eval_dpdm; 
    bool eval_weightedcost; 
    double leakage_integral; 
    double energy_integral; 
    double dpdm_integral; 
    double weightedcost_integral; 
    double weightedcost_width; 
 
    int ntime; 
    double total_time; 
    double dt; 
    bool writeTrajectoryDataFiles;  
 
    Vec redgrad; 
    OptimTarget* optim_target; 
    Output* output; 
    MasterEq* mastereq; 
 
  public: 
    TimeStepper(); 
 
    TimeStepper(const Config& config, MasterEq* mastereq_, Output* output_, int ninit_local); 
 
    virtual ~TimeStepper(); 
 
    double getLeakageIntegral(){ return leakage_integral; };
    double getWeightedCostIntegral(){ return weightedcost_integral; };
    double getEnergyIntegral(){ return energy_integral; };
    double getDPDMIntegral(){ return dpdm_integral; };
    Vec getReducedGradient(){ return redgrad; };
    void setWriteTrajectoryDataFiles(bool write){ writeTrajectoryDataFiles = write; };
 
    void setOptimTarget(OptimTarget* optim_target_){ optim_target = optim_target_; };
 
    Vec getFinalState(size_t iinit_local){ return final_states[iinit_local]; }
 
    virtual double getMinTimestepSize() const { return dt; }
 
    virtual Vec solveODE(int initid, int iinit_local, Vec rho_t0);
 
    virtual void solveAdjointODE(int iinit_local, Vec rho_t0_bar, double Jbar_leakage, double Jbar_weightedcost, double Jbar_dpdm, double Jbar_energy);
 
    double evalLeakage(const Vec x);
 
    void evalLeakage_diff(const Vec x, Vec xbar, double Jbar);
 
    double evalWeightedCost(double time, const Vec x);
 
    void evalWeightedCost_diff(double time, const Vec x, Vec xbar, double Jbar);
 
    double evalDpDm(Vec x, Vec xm1, Vec xm2);
 
    void evalDpDm_diff(int n, Vec xbar, double Jbar);
    
    double evalEnergy(double time);
 
    void evalEnergy_diff(double time, double Jbar, Vec redgrad);
 
    virtual void evolveFWD(const double tstart, const double tstop, Vec x) = 0;
 
    virtual void evolveBWD(const double tstart, const double tstop, const Vec x_stop, Vec x_adj, Vec grad, bool compute_gradient);
};
 
class ExplEuler : public TimeStepper {
  protected:
  Vec stage; 
  public:
    ExplEuler(const Config& config, MasterEq* mastereq_, Output* output_, int ninit_local);
 
    ~ExplEuler();
 
    void evolveFWD(const double tstart, const double tstop, Vec x);
 
    void evolveBWD(const double tstart, const double tstop, const Vec x_stop, Vec x_adj, Vec grad, bool compute_gradient);
};
 
class ImplMidpoint : public TimeStepper {
  protected:
  Vec stage, stage_adj; 
  Vec rhs, rhs_adj; 
  KSP ksp; 
  PC  preconditioner; 
  LinearSolverType linsolve_type; 
  int linsolve_maxiter; 
  double linsolve_abstol; 
  double linsolve_reltol; 
  int linsolve_iterstaken_avg; 
  double linsolve_error_avg; 
  int linsolve_counter; 
  Vec tmp, err; 
 
  public:
    ImplMidpoint(const Config& config, MasterEq* mastereq_, Output* output_, int ninit_local);
 
    ~ImplMidpoint();
 
    virtual void evolveFWD(const double tstart, const double tstop, Vec x);
 
    virtual void evolveBWD(const double tstart, const double tstop, const Vec x_stop, Vec x_adj, Vec grad, bool compute_gradient);
 
    int NeumannSolve(Mat A, Vec b, Vec x, double alpha, bool transpose);
};
 
class CompositionalImplMidpoint : public ImplMidpoint {
  protected:
 
  std::vector<double> gamma; 
  std::vector<Vec> x_stage; 
  Vec aux; 
  int order; 
 
  public:
    CompositionalImplMidpoint(const Config& config, MasterEq* mastereq_, Output* output_, int ninit_local, int order_);
 
    ~CompositionalImplMidpoint();
 
    void evolveFWD(const double tstart, const double tstop, Vec x);
 
    void evolveBWD(const double tstart, const double tstop, const Vec x_stop, Vec x_adj, Vec grad, bool compute_gradient);
};
 
 
class PetscTS : public TimeStepper {
  protected:
    TS ts;      
    std::vector<TS> ts_pool; 
    std::vector<Vec> q_pool; 
    Vec redgrad_ts; 
 
    Mat dRHSdp; 
    Mat dIntegralCostdY; 
    Mat dIntegralCostdP; 
    PetscReal dRHSdp_time; 
    double adj_scale_leakage; 
    double adj_scale_weightedcost; 
    double adj_scale_energy; 
    double min_timestep_size; 
 
  public:
    PetscTS(const Config& config, MasterEq* mastereq_, Output* output_, int ninit_local);
    ~PetscTS();
 
    Vec solveODE(int initid, int iinit_local, Vec rho_t0) override;
 
    void solveAdjointODE(int iinit_local, Vec rho_t0_bar, double Jbar_leakage, double Jbar_weightedcost, double Jbar_dpdm, double Jbar_energy) override;
 
    static PetscErrorCode RHSMatrixUpdate(TS ts, PetscReal t, Vec x, Mat A, Mat B, void *ptr);
 
    static PetscErrorCode dRHSdpMatrixUpdate(TS ts, PetscReal t, Vec x, Mat A, void *ptr);
 
    static PetscErrorCode computedRHSdp(Mat A, Vec x, Vec y);
 
    static PetscErrorCode IntegralCosts(TS ts, PetscReal t, Vec x, Vec F, void *ctx);
 
    static PetscErrorCode dIntegralCostdYUpdate(TS ts, PetscReal t, Vec x, Mat A, Mat B, void *ptr);
 
    static PetscErrorCode dIntegralCostdPUpdate(TS ts, PetscReal t, Vec x, Mat A, void *ptr);
 
    // Callback function during TSSolve to evaluate trajectory data at each accepted time step 
    static PetscErrorCode monitorTrajectory(TS ts, PetscInt step, PetscReal time, Vec state, void *ctx);
 
    // THESE ARE NOT USED. Instead, solveODE and solveAdjointODE overwrites the default time-stepping by calling TSSolve. 
    void evolveFWD(const double, const double, Vec) override {};
    void evolveBWD(const double, const double, const Vec, Vec, Vec, bool) override {};
 
    double getMinTimestepSize() const override { return min_timestep_size; }
};