Skip to content

Commit

Permalink
ODE: first adaptive version of BDF
Browse files Browse the repository at this point in the history 
The current implementation only allows for adaptivity in time,
no tests has been performed yet but the code is compiling correctly.
Next the correctness of the implementation will be tested and
eventually the order adaptivity will be enabled as well to improve
the algorithmic efficiency of the method.
  • Loading branch information
@lucbv
lucbv committed Sep 20, 2023
1 parent 4d54e13 commit eb4a5c0
Show file tree
Hide file tree
Showing 3 changed files with 256 additions and 2 deletions.
179 changes: 179 additions & 0 deletions ode/impl/KokkosODE_BDF_impl.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,8 @@
#include "Kokkos_Core.hpp"

#include "KokkosODE_Newton.hpp"
#include "KokkosBlas2_serial_gemv.hpp"
#include "KokkosBatched_Gemm_Decl.hpp"

namespace KokkosODE {
namespace Impl {
Expand Down Expand Up @@ -119,6 +121,48 @@ struct BDF_system_wrapper {
}
};

template <class system_type, class subview_type>
struct BDF_system_wrapper2 {
const system_type mySys;
const int neqs;
const subview_type psi, d;

double t, dt, c = 0;

KOKKOS_FUNCTION
BDF_system_wrapper2(const system_type& mySys_,
const subview_type& psi_, const subview_type& d_,
const double t_, const double dt_)
: mySys(mySys_),
neqs(mySys_.neqs),
psi(psi_),
d(d_),
t(t_),
dt(dt_) {}

template <class vec_type>
KOKKOS_FUNCTION void residual(const vec_type& y, const vec_type& f) const {
// f = f(t+dt, y)
mySys.evaluate_function(t, dt, y, f);

for (int eqIdx = 0; eqIdx < neqs; ++eqIdx) {
f(eqIdx) = c * f(eqIdx) - psi(eqIdx) - d(eqIdx);
}
}

template <class vec_type, class mat_type>
KOKKOS_FUNCTION void jacobian(const vec_type& y, const mat_type& jac) const {
mySys.evaluate_jacobian(t, dt, y, jac);

for (int rowIdx = 0; rowIdx < neqs; ++rowIdx) {
for (int colIdx = 0; colIdx < neqs; ++colIdx) {
jac(rowIdx, colIdx) = -dt * jac(rowIdx, colIdx);
}
jac(rowIdx, rowIdx) += 1.0;
}
}
};

template <class ode_type, class table_type, class vec_type, class mv_type,
class mat_type, class scalar_type>
KOKKOS_FUNCTION void BDFStep(ode_type& ode, const table_type& table,
Expand All @@ -145,6 +189,141 @@ KOKKOS_FUNCTION void BDFStep(ode_type& ode, const table_type& table,

} // BDFStep

template <class mat_type, class scalar_type>
KOKKOS_FUNCTION void compute_coeffs(const int order,
const scalar_type factor,
const mat_type& coeffs) {
coeffs(0, 0) = 1.0;
for(int colIdx = 0; colIdx < order; ++colIdx) {
coeffs(0, colIdx + 1) = 1.0;
for(int rowIdx = 0; rowIdx < order; ++rowIdx) {
coeffs(rowIdx + 1, colIdx + 1) = ((rowIdx - factor*(colIdx + 1.0)) / (rowIdx + 1.0)) * coeffs(rowIdx, colIdx + 1);
}
}
}

template <class mat_type, class scalar_type>
KOKKOS_FUNCTION void update_D(const int order, const scalar_type factor,
const mat_type& coeffs,
const mat_type& tempD, const mat_type& D) {
auto subD = Kokkos::subview(D, Kokkos::pair<int, int>(0, order + 1), Kokkos::ALL);
auto subTempD = Kokkos::subview(tempD, Kokkos::pair<int, int>(0, order + 1), Kokkos::ALL);

compute_coeffs(order, factor, coeffs);
auto R = Kokkos::subview(coeffs, Kokkos::pair<int, int>(0, order + 1), Kokkos::pair<int, int>(0, order + 1));
KokkosBatched::SerialGemm<KokkosBatched::Trans::Transpose,
KokkosBatched::Trans::NoTranspose,
KokkosBatched::Algo::Gemm::Blocked>::invoke(1.0, R, subD, 0.0, subTempD);

compute_coeffs(order, 1.0, coeffs);
auto U = Kokkos::subview(coeffs, Kokkos::pair<int, int>(0, order + 1), Kokkos::pair<int, int>(0, order + 1));
KokkosBatched::SerialGemm<KokkosBatched::Trans::Transpose,
KokkosBatched::Trans::NoTranspose,
KokkosBatched::Algo::Gemm::Blocked>::invoke(1.0, U, subTempD, 0.0, subD);
}

template <class ode_type, class table_type, class vec_type, class mv_type,
class mat_type, class scalar_type>
KOKKOS_FUNCTION void BDFStep(ode_type& ode, scalar_type t, scalar_type dt,
scalar_type t_end, int& order, int& num_equal_steps,
const int max_newton_iters,
const scalar_type atol, const scalar_type rtol,
const scalar_type min_factor,
const vec_type& y_old, const vec_type& y_new,
const vec_type& rhs, const vec_type& update,
const mat_type& temp, const mat_type& temp2,
const mat_type& jac) {
using newton_params = KokkosODE::Experimental::Newton_params;
// Kokkos::Array<double, 6> kappa{{0., -0.1850, -1/9 , -0.0823000, -0.0415000, 0.}};
// alpha(i) = (1 - kappa(i)) * gamma(i)
// error_const(i) = kappa(i) * gamma(i) + 1 / (i + 1)
Kokkos::Array<double, 6> gamma{{0., 1. , 1.5 , 1.83333333, 2.08333333, 2.28333333}};
Kokkos::Array<double, 6> alpha{{0., 1.185 , 1.66666667, 1.98421667, 2.16979167, 2.28333333}};
Kokkos::Array<double, 6> error_const{{1. , 0.315 , 0.16666667, 0.09911667, 0.11354167, 0.16666667}};

auto D = Kokkos::subview(temp, Kokkos::pair<int, int>(0, 8), Kokkos::ALL());
auto tempD = Kokkos::subview(temp, Kokkos::pair<int, int>(8, 16), Kokkos::ALL());
auto scale = Kokkos::subview(temp, Kokkos::pair<int, int>(16, 17), Kokkos::ALL());
auto y_predict = Kokkos::subview(temp, Kokkos::pair<int, int>(17, 18), Kokkos::ALL());
auto d = Kokkos::subview(temp, Kokkos::pair<int, int>(18, 19), Kokkos::ALL());
auto psi = Kokkos::subview(temp, Kokkos::pair<int, int>(19, 20), Kokkos::ALL());
auto error = Kokkos::subview(temp, Kokkos::pair<int, int>(20, 21), Kokkos::ALL());
auto coeffs = Kokkos::subview(temp2, Kokkos::pair<int, int>(0, order + 1), Kokkos::ALL());

BDF_system_wrapper2 sys(ode, t, dt, psi, d);
const newton_params param(max_newton_iters, atol, rtol);

scalar_type max_step, min_step;
if(dt > max_step) {
update_D(order, max_step / dt, coeffs, tempD, D);
dt = max_step;
num_equal_steps = 0;
} else if(dt < min_step) {
update_D(order, min_step / dt, coeffs, tempD, D);
dt = min_step;
num_equal_steps = 0;
}

// first set y_new = y_old
for (int eqIdx = 0; eqIdx < sys.neqs; ++eqIdx) {
y_new(eqIdx) = y_old(eqIdx);
}

double t_new = 0;
bool step_accepted = false;
while (!step_accepted) {
if(dt < min_step) {return ;}
t_new = t + dt;

if(t_new > t_end) {
t_new = t_end;
update_D(order, (t_new - t) / dt, coeffs, tempD, D);
num_equal_steps = 0;
}
dt = t_new - t;

for(int eqIdx = 0; eqIdx < sys.neqs; ++eqIdx) {
for(int orderIdx = 0; orderIdx < order + 1; ++orderIdx) {
y_predict(eqIdx) += D(orderIdx, eqIdx);
}
scale(eqIdx) = atol + rtol*Kokkos::abs(y_predict(eqIdx));
}
KokkosBlas::Experimental::serial_gemv('T', 1.0 / alpha[order], D, gamma, 0.0, psi);

sys.c = dt / alpha[order];
KokkosODE::Experimental::newton_solver_status newton_status =
KokkosODE::Experimental::Newton::Solve(sys, param, jac, temp, y_new, rhs,
update);
if(newton_status == KokkosODE::Experimental::newton_solver_status::MAX_ITER) {
dt = 0.5*dt;
update_D(order, 0.5, coeffs, tempD, D);
num_equal_steps = 0;

} else {
// Estimate the solution error
scalar_type safety = 0.9*(2*max_newton_iters + 1) / (2*max_newton_iters + param.iters);
scalar_type error_norm = 0;
for(int eqIdx = 0; eqIdx < sys.neqs; ++eqIdx) {
scale(eqIdx) = atol + rtol*Kokkos::abs(y_new(eqIdx));
error(eqIdx) = error_const[order]*d(eqIdx) / scale(eqIdx);
error_norm += error(eqIdx)*error(eqIdx);
}
error_norm = Kokkos::sqrt(error_norm);

// Check error norm and adapt step size or accept step
if(error_norm > 1) {
scalar_type factor = Kokkos::max(min_factor, safety*Kokkos::pow(error_norm, -1.0 / (order + 1)));
dt = factor*dt;
update_D(order, factor, coeffs, tempD, D);
num_equal_steps = 0;
} else {
step_accepted = true;
}
}
} // while(!step_accepted)

} // BDFStep

} // namespace Impl
} // namespace KokkosODE

Expand Down
2 changes: 1 addition & 1 deletion ode/src/KokkosODE_Types.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,7 @@ struct ODE_params {
enum newton_solver_status { NLS_SUCCESS = 0, MAX_ITER = 1, LIN_SOLVE_FAIL = 2 };

struct Newton_params {
int max_iters;
int max_iters, iters = 0;
double abs_tol, rel_tol;

// Constructor that sets basic solver parameters
Expand Down
77 changes: 76 additions & 1 deletion ode/unit_test/Test_ODE_BDF.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,7 @@ struct LotkaVolterra {
// pp. 178–182, Academic Press, London (1966).
//
// Equations: y0' = -0.04*y0 + 10e4*y1*y2
// y1' = 0.04*y0 - 10e4*y1*y2 - 3e-7 * y1**2
// y1' = 0.04*y0 - 10e4*y1*y2 - 3e7 * y1**2
// y2' = 3e-7 * y1**2
struct StiffChemestry {
static constexpr int neqs = 3;
Expand Down Expand Up @@ -508,6 +508,78 @@ void test_BDF_parallel() {
Kokkos::fence();
}

template <class mat_type, class scalar_type>
void compute_coeffs(const int order, const scalar_type factor, const mat_type& coeffs) {
coeffs(0, 0) = 1.0;
for(int colIdx = 0; colIdx < order; ++colIdx) {
coeffs(0, colIdx + 1) = 1.0;
for(int rowIdx = 0; rowIdx < order; ++rowIdx) {
coeffs(rowIdx + 1, colIdx + 1) = ((rowIdx - factor*(colIdx + 1.0)) / (rowIdx + 1.0)) * coeffs(rowIdx, colIdx + 1);
}
}
}

template <class mat_type, class scalar_type>
void update_D(const int order, const scalar_type factor, const mat_type& coeffs, const mat_type& tempD, const mat_type& D) {
auto subD = Kokkos::subview(D, Kokkos::pair<int, int>(0, order + 1), Kokkos::ALL);
auto subTempD = Kokkos::subview(tempD, Kokkos::pair<int, int>(0, order + 1), Kokkos::ALL);

compute_coeffs(order, factor, coeffs);
auto R = Kokkos::subview(coeffs, Kokkos::pair<int, int>(0, order + 1), Kokkos::pair<int, int>(0, order + 1));
KokkosBatched::SerialGemm<KokkosBatched::Trans::Transpose,
KokkosBatched::Trans::NoTranspose,
KokkosBatched::Algo::Gemm::Blocked>::invoke(1.0, R, subD, 0.0, subTempD);

compute_coeffs(order, 1.0, coeffs);
auto U = Kokkos::subview(coeffs, Kokkos::pair<int, int>(0, order + 1), Kokkos::pair<int, int>(0, order + 1));
KokkosBatched::SerialGemm<KokkosBatched::Trans::Transpose,
KokkosBatched::Trans::NoTranspose,
KokkosBatched::Algo::Gemm::Blocked>::invoke(1.0, U, subTempD, 0.0, subD);
}

template <class execution_space, class scalar_type>
void test_Nordsieck() {
StiffChemestry mySys{};

Kokkos::View<scalar_type**, execution_space> R("coeffs", 6, 6), U("coeffs", 6, 6);
Kokkos::View<scalar_type**, execution_space> D("D", 8, mySys.neqs), tempD("tmp", 8, mySys.neqs);
int order = 1;
double factor = 0.8;

constexpr double t_start = 0.0, t_end = 500.0;
int max_steps = 200000;
double dt = (t_end - t_start) / max_steps;

auto y0 = Kokkos::subview(D, 0, Kokkos::ALL());
auto f = Kokkos::subview(D, 1, Kokkos::ALL());
y0(0) = 1.0;
mySys.evaluate_function(0, 0, y0, f);
for(int eqIdx = 0; eqIdx < mySys.neqs; ++eqIdx) {
f(eqIdx) *= dt;
}

compute_coeffs(order, factor, R);
compute_coeffs(order, 1.0, U);

std::cout << "R: " << std::endl;
for(int i = 0; i < order + 1; ++i) {
std::cout << "{ ";
for(int j = 0; j < order + 1; ++j) {
std::cout << R(i,j) << ", ";
}
std::cout << "}" << std::endl;
}

std::cout << "D before update:" << std::endl;
std::cout << " { " << D(0, 0) << ", " << D(0, 1) << ", " << D(0, 2) << " }" << std::endl;
std::cout << " { " << D(1, 0) << ", " << D(1, 1) << ", " << D(1, 2) << " }" << std::endl;
update_D(order, factor, R, tempD, D);

std::cout << "D after update:" << std::endl;
std::cout << " { " << D(0, 0) << ", " << D(0, 1) << ", " << D(0, 2) << " }" << std::endl;
std::cout << " { " << D(1, 0) << ", " << D(1, 1) << ", " << D(1, 2) << " }" << std::endl;
}

} // namespace Test

TEST_F(TestCategory, BDF_Logistic_serial) {
Expand All @@ -522,3 +594,6 @@ TEST_F(TestCategory, BDF_StiffChemestry_serial) {
TEST_F(TestCategory, BDF_parallel_serial) {
::Test::test_BDF_parallel<TestExecSpace, double>();
}
TEST_F(TestCategory, BDF_Nordsieck) {
::Test::test_Nordsieck<TestExecSpace, double>();
}

0 comments on commit eb4a5c0

Please sign in to comment.