Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 9 additions & 11 deletions ode/impl/KokkosODE_BDF_impl.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -95,32 +95,30 @@ struct BDF_system_wrapper {
}
};

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

bool compute_jac = true;
double t, dt, c = 0;

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

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

// std::cout << "f = psi + d - c * f = " << psi(0) << " + " << d(0) << " - "
// << c << " * " << f(0) << std::endl;

// rhs = higher order terms + y_{n+1}^i - y_n - dt*f
// Corrector equation of the NDF step
// 0 = psi + (y - y_predict) - c * f(t+dt, y)
for (int eqIdx = 0; eqIdx < neqs; ++eqIdx) {
f(eqIdx) = psi(eqIdx) + d(eqIdx) - c * f(eqIdx);
f(eqIdx) = psi(eqIdx) + (y(eqIdx) - y_predict(eqIdx)) - c * f(eqIdx);
}
}

Expand Down Expand Up @@ -301,7 +299,7 @@ KOKKOS_FUNCTION void BDFStep(ode_type& ode, scalar_type& t, scalar_type& dt, sca
gamma(4) = 2.08333333;
gamma(5) = 2.28333333;

BDF_system_wrapper2 sys(ode, psi, update, t, dt);
BDF_system_wrapper2 sys(ode, psi, y_predict, t, dt);
const newton_params param(
max_newton_iters, atol,
Kokkos::max(10 * KokkosKernels::ArithTraits<scalar_type>::eps() / rtol, Kokkos::min(0.03, Kokkos::sqrt(rtol))));
Expand Down
120 changes: 120 additions & 0 deletions ode/unit_test/Test_ODE_BDF.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -114,6 +114,31 @@ struct StiffChemistry {
}
};

// Simple quadratic ODE
// used to check the root of the
// corrector equation solved by the
// adaptive BDF (NDF) time stepper.
//
// Equation: y'(t) = y(t)**2
// Jacobian: df/dy = 2*y
struct Quadratic {
static constexpr int neqs = 1;

Quadratic() {}

template <class vec_type1, class vec_type2>
KOKKOS_FUNCTION void evaluate_function(const double /*t*/, const double /*dt*/, const vec_type1& y,
const vec_type2& f) const {
f(0) = y(0) * y(0);
}

template <class vec_type, class mat_type>
KOKKOS_FUNCTION void evaluate_jacobian(const double /*t*/, const double /*dt*/, const vec_type& y,
const mat_type& jac) const {
jac(0, 0) = 2.0 * y(0);
}
}; // Quadratic

template <class ode_type, KokkosODE::Experimental::BDF_type bdf_type, class vec_type, class mv_type, class mat_type,
class scalar_type>
struct BDFSolve_wrapper {
Expand Down Expand Up @@ -174,6 +199,42 @@ struct BDF_Solve_wrapper {
}
};

template <class ode_type, class vec_type, class mat_type, class status_view_type, class scalar_type>
struct BDFCorrectorSolve_wrapper {
ode_type my_ode;
scalar_type t, dt, c;
vec_type y, y_predict, psi, rhs, update, scale;
mat_type jac, tmp;
KokkosODE::Experimental::Newton_params params;
status_view_type status;

BDFCorrectorSolve_wrapper(const ode_type& my_ode_, const scalar_type t_, const scalar_type dt_, const scalar_type c_,
const vec_type& y_, const vec_type& y_predict_, const vec_type& psi_, const vec_type& rhs_,
const vec_type& update_, const vec_type& scale_, const mat_type& jac_, const mat_type& tmp_,
const KokkosODE::Experimental::Newton_params& params_, const status_view_type& status_)
: my_ode(my_ode_),
t(t_),
dt(dt_),
c(c_),
y(y_),
y_predict(y_predict_),
psi(psi_),
rhs(rhs_),
update(update_),
scale(scale_),
jac(jac_),
tmp(tmp_),
params(params_),
status(status_) {}

KOKKOS_FUNCTION
void operator()(const int /*idx*/) const {
KokkosODE::Impl::BDF_system_wrapper2 sys(my_ode, psi, y_predict, t, dt);
sys.c = c;
status(0) = KokkosODE::Experimental::Newton::Solve(sys, params, jac, tmp, y, rhs, update, scale);
}
};

template <class device_type, class scalar_type>
void test_BDF_Logistic() {
using execution_space = typename device_type::execution_space;
Expand Down Expand Up @@ -389,6 +450,64 @@ void test_BDF_StiffChemistry() {
Kokkos::parallel_for(myPolicy, solve_wrapper);
}

// Solve the corrector equation of the adaptive BDF (NDF) time stepper,
// 0 = psi + (y - y_predict) - c * f(t+dt, y),
// through BDF_system_wrapper2 the same way BDFStep does: the Newton
// iteration starts from y = y_predict with a zeroed update vector.
// For the quadratic ODE f(y) = y**2 with psi = 7, y_predict = 3 and
// c = dt = 1/2 the equation reads 0 = y**2 - 2*y - 8 and Newton
// converges to its root y = 4 in a handful of iterations (c = dt makes
// the modified Newton Jacobian I - dt*df/dy exact here).
//
// This guards against a regression where the correction y - y_predict
// was read from the Newton solver's update vector: that vector only
// holds the last Newton step, which vanishes as the iteration
// converges, so the corrector root degenerates to the root of
// 0 = psi - c * f(t+dt, y), i.e. y = sqrt(14) ~ 3.742, and the solve
// either converges to that wrong root or is flagged as divergent.
// The bug only shows from the third Newton iteration on -- after a
// single step the update still equals the full correction -- which is
// why y_predict is placed far enough from the root to require several
// iterations.
template <class device_type, class scalar_type>
void test_BDF_corrector_equation() {
using execution_space = typename device_type::execution_space;
using newton_solver_status = KokkosODE::Experimental::newton_solver_status;
using vec_type = Kokkos::View<scalar_type*, execution_space>;
using mat_type = Kokkos::View<scalar_type**, execution_space>;

Quadratic mySys{};

const scalar_type t = 0.0, dt = 0.5, c = 0.5;

vec_type y("solution", mySys.neqs), y_predict("predictor", mySys.neqs), psi("higher order terms", mySys.neqs);
vec_type rhs("rhs", mySys.neqs), update("update", mySys.neqs), scale("scaling factors", mySys.neqs);
mat_type jac("jacobian", mySys.neqs, mySys.neqs), tmp("temp mem", mySys.neqs, mySys.neqs + 4);
Kokkos::View<newton_solver_status*, execution_space> status("newton status", 1);

Kokkos::deep_copy(psi, 7.0);
Kokkos::deep_copy(y_predict, 3.0);
Kokkos::deep_copy(y, 3.0);
Kokkos::deep_copy(update, 0.0);
Kokkos::deep_copy(scale, 1.0);

const KokkosODE::Experimental::Newton_params params(50, 1e-12, 1e-8);

Kokkos::RangePolicy<execution_space> my_policy(0, 1);
BDFCorrectorSolve_wrapper solve_wrapper(mySys, t, dt, c, y, y_predict, psi, rhs, update, scale, jac, tmp, params,
status);
Kokkos::parallel_for(my_policy, solve_wrapper);
Kokkos::fence();

auto status_h = Kokkos::create_mirror_view(status);
auto y_h = Kokkos::create_mirror_view(y);
Kokkos::deep_copy(status_h, status);
Kokkos::deep_copy(y_h, y);

EXPECT_TRUE(status_h(0) == newton_solver_status::NLS_SUCCESS);
EXPECT_NEAR_KK_REL(y_h(0), static_cast<scalar_type>(4.0), static_cast<scalar_type>(1.0e-5));
} // test_BDF_corrector_equation

// template <class ode_type, KokkosODE::Experimental::BDF_type bdf_type,
// class vec_type, class mv_type, class mat_type, class scalar_type>
// struct BDFSolve_parallel {
Expand Down Expand Up @@ -720,6 +839,7 @@ void test_BDF_adaptive_stiff() {
TEST_F(TestCategory, BDF_Logistic_serial) { ::Test::test_BDF_Logistic<TestDevice, double>(); }
TEST_F(TestCategory, BDF_LotkaVolterra_serial) { ::Test::test_BDF_LotkaVolterra<TestDevice, double>(); }
TEST_F(TestCategory, BDF_StiffChemistry_serial) { ::Test::test_BDF_StiffChemistry<TestDevice, double>(); }
TEST_F(TestCategory, BDF_corrector_equation) { ::Test::test_BDF_corrector_equation<TestDevice, double>(); }
// TEST_F(TestCategory, BDF_parallel_serial) {
// ::Test::test_BDF_parallel<TestDevice, double>();
// }
Expand Down
Loading