Skip to content

Commit

Permalink
ODE - RK: adding time step counter
Browse files Browse the repository at this point in the history
  • Loading branch information
@lucbv
lucbv committed Jun 11, 2024
1 parent 6e3346c commit 497fee7
Show file tree
Hide file tree
Showing 8 changed files with 54 additions and 47 deletions.
2 changes: 1 addition & 1 deletion ode/impl/KokkosODE_BDF_impl.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -154,7 +154,7 @@ struct BDF_system_wrapper2 {
if (compute_jac) {
mySys.evaluate_jacobian(t, dt, y, jac);

// J = I - dt*(dy/dy)
// J = I - dt*(df/dy)
for (int rowIdx = 0; rowIdx < neqs; ++rowIdx) {
for (int colIdx = 0; colIdx < neqs; ++colIdx) {
jac(rowIdx, colIdx) = -dt * jac(rowIdx, colIdx);
Expand Down
16 changes: 4 additions & 12 deletions ode/impl/KokkosODE_RungeKutta_impl.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -80,25 +80,20 @@ KOKKOS_FUNCTION void RKStep(ode_type& ode, const table_type& table,
} // RKStep

template <class ode_type, class table_type, class vec_type, class mv_type,
class scalar_type, bool record_count>
class scalar_type>
KOKKOS_FUNCTION Experimental::ode_solver_status RKSolve(
const ode_type& ode, const table_type& table,
const KokkosODE::Experimental::ODE_params& params,
const scalar_type t_start, const scalar_type t_end, const vec_type& y0,
const vec_type& y, const vec_type& temp, const mv_type& k_vecs, int count) {
const vec_type& y, const vec_type& temp, const mv_type& k_vecs, int* const count) {
constexpr scalar_type error_threshold = 1;
scalar_type error_n;
bool adapt = params.adaptivity;
bool dt_was_reduced;
if (std::is_same_v<table_type, ButcherTableau<0, 0>>) {
adapt = false;
}

if constexpr (record_count) {
count = 0;
} else {
count = -1;
}
*count = 0;

// Set current time and initial time step
scalar_type t_now = t_start;
Expand Down Expand Up @@ -163,10 +158,7 @@ KOKKOS_FUNCTION Experimental::ode_solver_status RKSolve(

// Update time and initial condition for next time step
t_now += dt;
if constexpr(record_count) {
count += 1;
std::cout << "current count: " << count << std::endl;
}
*count += 1;
for (int eqIdx = 0; eqIdx < ode.neqs; ++eqIdx) {
y0(eqIdx) = y(eqIdx);
}
Expand Down
2 changes: 1 addition & 1 deletion ode/src/KokkosODE_BDF.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -118,7 +118,7 @@ struct BDF {
KokkosODE::Experimental::ODE_params params(table.order - 1);
for (int stepIdx = 0; stepIdx < init_steps; ++stepIdx) {
KokkosODE::Experimental::RungeKutta<RK_type::RKF45>::Solve(
ode, params, t, t + dt, y0, y, update, kstack);
ode, params, t, t + dt, y0, y, update, kstack, nullptr);

for (int eqIdx = 0; eqIdx < ode.neqs; ++eqIdx) {
y_vecs(eqIdx, stepIdx + 1) = y(eqIdx);
Expand Down
10 changes: 4 additions & 6 deletions ode/src/KokkosODE_RungeKutta.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -90,7 +90,7 @@ struct RK_Tableau_helper<RK_type::RKDP> {
///
/// \tparam RK_type an RK_type enum value used to specify
/// which Runge Kutta method is to be used.
template <RK_type T, bool record_count = false>
template <RK_type T>
struct RungeKutta {
using table_type = typename RK_Tableau_helper<T>::table_type;

Expand Down Expand Up @@ -129,12 +129,10 @@ struct RungeKutta {
KOKKOS_FUNCTION static ode_solver_status Solve(
const ode_type& ode, const KokkosODE::Experimental::ODE_params& params,
const scalar_type t_start, const scalar_type t_end, const vec_type& y0,
const vec_type& y, const vec_type& temp, const mv_type& k_vecs, int count = -1) {
const vec_type& y, const vec_type& temp, const mv_type& k_vecs, int* const count) {
table_type table;
return KokkosODE::Impl::RKSolve<ode_type, table_type, vec_type, mv_type,
scalar_type, record_count>(ode, table, params,
t_start, t_end, y0, y,
temp, k_vecs, count);
return KokkosODE::Impl::RKSolve(ode, table, params, t_start, t_end, y0, y,
temp, k_vecs, count);
}
};

Expand Down
32 changes: 19 additions & 13 deletions ode/unit_test/Test_ODE_RK.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -86,7 +86,7 @@ struct solution_wrapper {

template <class ode_type, KokkosODE::Experimental::RK_type rk_type,
class vec_type, class mv_type, class scalar_type,
bool record_count = false>
class count_type>
struct RKSolve_wrapper {
using ode_params = KokkosODE::Experimental::ODE_params;

Expand All @@ -96,12 +96,13 @@ struct RKSolve_wrapper {
int max_steps;
vec_type y_old, y_new, tmp;
mv_type kstack;
int count;
count_type count;

RKSolve_wrapper(const ode_type& my_ode_, const ode_params& params_,
const scalar_type tstart_, const scalar_type tend_,
const vec_type& y_old_, const vec_type& y_new_,
const vec_type& tmp_, const mv_type& kstack_)
const vec_type& tmp_, const mv_type& kstack_,
const count_type& count_)
: my_ode(my_ode_),
params(params_),
tstart(tstart_),
Expand All @@ -110,13 +111,12 @@ struct RKSolve_wrapper {
y_new(y_new_),
tmp(tmp_),
kstack(kstack_),
count(0) {}
count(count_) {}

KOKKOS_FUNCTION
void operator()(const int /*idx*/) const {
std::cout << "counting step? " << record_count << std::endl;
KokkosODE::Experimental::RungeKutta<rk_type, record_count>::Solve(
my_ode, params, tstart, tend, y_old, y_new, tmp, kstack, count);
KokkosODE::Experimental::RungeKutta<rk_type>::Solve(
my_ode, params, tstart, tend, y_old, y_new, tmp, kstack, count.data());
}
};

Expand All @@ -131,14 +131,16 @@ void test_method(const std::string label, ode_type& my_ode,
typename vec_type::HostMirror y_ref_h) {
using execution_space = typename vec_type::execution_space;
using solver_type = KokkosODE::Experimental::RungeKutta<rk_type>;
using count_type = Kokkos::View<int*, execution_space>;

KokkosODE::Experimental::ODE_params params(num_steps);
vec_type tmp("tmp vector", my_ode.neqs);
mv_type kstack("k stack", solver_type::num_stages(), my_ode.neqs);
count_type count("time step count", 1);

Kokkos::RangePolicy<execution_space> my_policy(0, 1);
RKSolve_wrapper<ode_type, rk_type, vec_type, mv_type, scalar_type>
solve_wrapper(my_ode, params, tstart, tend, y_old, y_new, tmp, kstack);
RKSolve_wrapper<ode_type, rk_type, vec_type, mv_type, scalar_type, count_type>
solve_wrapper(my_ode, params, tstart, tend, y_old, y_new, tmp, kstack, count);
Kokkos::parallel_for(my_policy, solve_wrapper);

auto y_new_h = Kokkos::create_mirror_view(y_new);
Expand Down Expand Up @@ -325,7 +327,9 @@ void test_rate(ode_type& my_ode, const scalar_type& tstart,
typename vec_type::HostMirror& error) {
using execution_space = typename vec_type::execution_space;
using solver_type = KokkosODE::Experimental::RungeKutta<rk_type>;
using count_type = Kokkos::View<int*, execution_space>;

count_type count("time step count", 1);
vec_type tmp("tmp vector", my_ode.neqs);
mv_type kstack("k stack", solver_type::num_stages(), my_ode.neqs);

Expand All @@ -338,8 +342,8 @@ void test_rate(ode_type& my_ode, const scalar_type& tstart,
KokkosODE::Experimental::ODE_params params(num_steps(idx));
Kokkos::deep_copy(y_old, y_old_h);
Kokkos::deep_copy(y_new, y_old_h);
RKSolve_wrapper<ode_type, rk_type, vec_type, mv_type, scalar_type>
solve_wrapper(my_ode, params, tstart, tend, y_old, y_new, tmp, kstack);
RKSolve_wrapper<ode_type, rk_type, vec_type, mv_type, scalar_type, count_type>
solve_wrapper(my_ode, params, tstart, tend, y_old, y_new, tmp, kstack, count);
Kokkos::parallel_for(my_policy, solve_wrapper);

Kokkos::deep_copy(y_new_h, y_new);
Expand Down Expand Up @@ -475,9 +479,11 @@ void test_adaptivity() {
using RK_type = KokkosODE::Experimental::RK_type;
using vec_type = Kokkos::View<double*, Device>;
using mv_type = Kokkos::View<double**, Device>;
using count_type = Kokkos::View<int*, execution_space>;

duho my_oscillator(1, 1, 4);
const int neqs = my_oscillator.neqs;
count_type count("time step count", 1);

vec_type y("solution", neqs), f("function", neqs);
auto y_h = Kokkos::create_mirror(y);
Expand Down Expand Up @@ -526,9 +532,9 @@ void test_adaptivity() {
minStepSize);
Kokkos::deep_copy(y_old, y_old_h);
Kokkos::deep_copy(y_new, y_old_h);
RKSolve_wrapper<duho, RK_type::RKF45, vec_type, mv_type, double>
RKSolve_wrapper<duho, RK_type::RKF45, vec_type, mv_type, double, count_type>
solve_wrapper(my_oscillator, params, tstart, tend, y_old, y_new, tmp,
kstack);
kstack, count);
Kokkos::parallel_for(my_policy, solve_wrapper);

auto y_new_h = Kokkos::create_mirror(y_new);
Expand Down
11 changes: 7 additions & 4 deletions ode/unit_test/Test_ODE_RK_chem.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -96,6 +96,7 @@ void test_chem() {
using mv_type = Kokkos::View<double**, Device>;
using RK_type = KokkosODE::Experimental::RK_type;
using solver_type = KokkosODE::Experimental::RungeKutta<RK_type::RKCK>;
using count_type = Kokkos::View<int*, execution_space>;

{
chem_model_1 chem_model;
Expand All @@ -105,6 +106,7 @@ void test_chem() {
KokkosODE::Experimental::ODE_params params(num_steps);
vec_type tmp("tmp vector", neqs);
mv_type kstack("k stack", solver_type::num_stages(), neqs);
count_type count("time steps count", 1);

// Set initial conditions
vec_type y_new("solution", neqs);
Expand All @@ -116,9 +118,9 @@ void test_chem() {
Kokkos::deep_copy(y_new, y_old_h);

Kokkos::RangePolicy<execution_space> my_policy(0, 1);
RKSolve_wrapper<chem_model_1, RK_type::RKCK, vec_type, mv_type, double>
RKSolve_wrapper<chem_model_1, RK_type::RKCK, vec_type, mv_type, double, count_type>
solve_wrapper(chem_model, params, chem_model.tstart, chem_model.tend,
y_old, y_new, tmp, kstack);
y_old, y_new, tmp, kstack, count);
Kokkos::parallel_for(my_policy, solve_wrapper);

auto y_new_h = Kokkos::create_mirror(y_new);
Expand Down Expand Up @@ -146,6 +148,7 @@ void test_chem() {
KokkosODE::Experimental::ODE_params params(num_steps);
vec_type tmp("tmp vector", neqs);
mv_type kstack("k stack", solver_type::num_stages(), neqs);
count_type count("time steps count", 1);

// Set initial conditions
vec_type y_new("solution", neqs);
Expand All @@ -162,9 +165,9 @@ void test_chem() {
Kokkos::deep_copy(y_new, y_old_h);

Kokkos::RangePolicy<execution_space> my_policy(0, 1);
RKSolve_wrapper<chem_model_2, RK_type::RKCK, vec_type, mv_type, double>
RKSolve_wrapper<chem_model_2, RK_type::RKCK, vec_type, mv_type, double, count_type>
solve_wrapper(chem_model, params, chem_model.tstart, chem_model.tend,
y_old, y_new, tmp, kstack);
y_old, y_new, tmp, kstack, count);
Kokkos::parallel_for(my_policy, solve_wrapper);

auto y_new_h = Kokkos::create_mirror(y_new);
Expand Down
8 changes: 5 additions & 3 deletions ode/unit_test/Test_ODE_RK_counts.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ void test_RK_count() {
using RK_type = KokkosODE::Experimental::RK_type;
using vec_type = Kokkos::View<double*, Device>;
using mv_type = Kokkos::View<double**, Device>;
using count_type = Kokkos::View<int*, execution_space>;

TestProblem::DegreeOnePoly myODE{};
constexpr int neqs = myODE.neqs;
Expand All @@ -39,6 +40,7 @@ void test_RK_count() {
// double dt = (tend - tstart) / num_steps;
vec_type y("solution", neqs), f("function", neqs);
vec_type y_new("y new", neqs), y_old("y old", neqs);
count_type count("time step count", 1);

auto y_h = Kokkos::create_mirror(y);
y_h(0) = myODE.expected_val(tstart, 0);
Expand Down Expand Up @@ -67,16 +69,16 @@ void test_RK_count() {
minStepSize);
Kokkos::deep_copy(y_old, y_old_h);
Kokkos::deep_copy(y_new, y_old_h);
RKSolve_wrapper<TestProblem::DegreeOnePoly, RK_type::RKF45, vec_type, mv_type, double, true>
RKSolve_wrapper<TestProblem::DegreeOnePoly, RK_type::RKF45, vec_type, mv_type, double, count_type>
solve_wrapper(myODE, params, tstart, tend, y_old, y_new, tmp,
kstack);
kstack, count);
Kokkos::parallel_for(my_policy, solve_wrapper);

auto y_new_h = Kokkos::create_mirror(y_new);
Kokkos::deep_copy(y_new_h, y_new);

std::cout << "y_ref=" << y_ref_h(0) << ", y_new=" << y_new_h(0) << std::endl;
std::cout << "time steps taken: " << solve_wrapper.count << std::endl;
std::cout << "time steps taken: " << count(0) << std::endl;

} // test_RK_count

Expand Down
20 changes: 13 additions & 7 deletions perf_test/ode/KokkosODE_RK.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -97,7 +97,7 @@ struct chem_model_2 {
};

template <class ode_type, class table_type, class vec_type, class mv_type,
class scalar_type>
class scalar_type, class count_type>
struct RKSolve_wrapper {
using ode_params = KokkosODE::Experimental::ODE_params;

Expand All @@ -108,12 +108,13 @@ struct RKSolve_wrapper {
scalar_type tstart, tend;
vec_type y_old, y_new, tmp;
mv_type kstack;
count_type count;

RKSolve_wrapper(const ode_type& my_ode_, const table_type& table_,
const ode_params& params_, const scalar_type tstart_,
const scalar_type tend_, const vec_type& y_old_,
const vec_type& y_new_, const vec_type& tmp_,
const mv_type& kstack_)
const mv_type& kstack_, const count_type& count_)
: my_ode(my_ode_),
table(table_),
params(params_),
Expand All @@ -122,7 +123,8 @@ struct RKSolve_wrapper {
y_old(y_old_),
y_new(y_new_),
tmp(tmp_),
kstack(kstack_) {}
kstack(kstack_),
count(count_) {}

KOKKOS_FUNCTION
void operator()(const int idx) const {
Expand All @@ -134,12 +136,13 @@ struct RKSolve_wrapper {
auto local_tmp = Kokkos::subview(tmp, Kokkos::pair(2 * idx, 2 * idx + 1));
auto local_kstack = Kokkos::subview(kstack, Kokkos::ALL(),
Kokkos::pair(2 * idx, 2 * idx + 1));
auto local_count = Kokkos::subview(count, idx, Kokkos::ALL());

// Run Runge-Kutta time integrator
// This should be replaced by a call to the public interface!
KokkosODE::Impl::RKSolve<ode_type, table_type, vec_type, mv_type, double, false>(
KokkosODE::Impl::RKSolve(
my_ode, table, params, tstart, tend, local_y_old, local_y_new,
local_tmp, local_kstack, -1);
local_tmp, local_kstack, local_count.data());
}
};

Expand All @@ -165,6 +168,7 @@ void run_ode_chem(benchmark::State& state, const rk_input_parameters& inputs) {
using mv_type = Kokkos::View<double**, execution_space>;
using table_type = KokkosODE::Impl::ButcherTableau<4, 5, 1>;
using ode_params = KokkosODE::Experimental::ODE_params;
using count_type = Kokkos::View<int**, execution_space>;

const int num_odes = inputs.num_odes;
const int model = inputs.model;
Expand All @@ -175,6 +179,7 @@ void run_ode_chem(benchmark::State& state, const rk_input_parameters& inputs) {
const int neqs = chem_model.neqs;
const int num_steps = 15000;
const double dt = 0.1;
count_type count("time steps count", num_odes, 1);

table_type table;
ode_params params(num_steps);
Expand All @@ -193,7 +198,7 @@ void run_ode_chem(benchmark::State& state, const rk_input_parameters& inputs) {
Kokkos::RangePolicy<execution_space> my_policy(0, num_odes);
RKSolve_wrapper solve_wrapper(chem_model, table, params,
chem_model.tstart, chem_model.tend, y_old,
y_new, tmp, kstack);
y_new, tmp, kstack, count);

Kokkos::Timer time;
time.reset();
Expand Down Expand Up @@ -227,6 +232,7 @@ void run_ode_chem(benchmark::State& state, const rk_input_parameters& inputs) {
const int neqs = chem_model.neqs;
const int num_steps = 15000;
const double dt = 0.1;
count_type count("time steps count", num_odes, 1);

table_type table;
ode_params params(num_steps);
Expand All @@ -250,7 +256,7 @@ void run_ode_chem(benchmark::State& state, const rk_input_parameters& inputs) {
Kokkos::RangePolicy<execution_space> my_policy(0, num_odes);
RKSolve_wrapper solve_wrapper(chem_model, table, params,
chem_model.tstart, chem_model.tend, y_old,
y_new, tmp, kstack);
y_new, tmp, kstack, count);

Kokkos::Timer time;
time.reset();
Expand Down

0 comments on commit 497fee7

Please sign in to comment.