Add ND time steppers

gyselax · Oct 20, 2023 · 5dd46a9 · 5dd46a9
1 parent d8bbf64
commit 5dd46a9
Show file tree

Hide file tree

Showing 22 changed files with 1,176 additions and 175 deletions.
diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
@@ -11,3 +11,4 @@ add_subdirectory(geometryRTheta)
 add_subdirectory(geometryXYVxVy)
 add_subdirectory(interpolation)
 add_subdirectory(advection)
+add_subdirectory(timestepper)
diff --git a/src/README.md b/src/README.md
@@ -10,4 +10,5 @@ The `src/` folder contains all the code necessary to build a gyrokinetic semi-La
 <!-- - [paraconfpp](./paraconfpp/README.md) - Paraconf utility functions. -->
 - [quadrature](./quadrature/README.md) - Code describing different quadrature methods.
 <!-- - [speciesinfo](./speciesinfo/README.md) - Code used to describe the different species. -->
+- [timestepper](./timestepper/README.md) - Code used to describe time-stepping methods (e.g. Runge-Kutta methods).
 - [utils](./utils/README.md) - Code describing utility functions.
diff --git a/src/geometryXVx/rhs/CMakeLists.txt b/src/geometryXVx/rhs/CMakeLists.txt
@@ -29,6 +29,7 @@ target_link_libraries("rhs_${GEOMETRY_VARIANT}"
         vcx::quadrature
         vcx::speciesinfo
         vcx::initialization_${GEOMETRY_VARIANT}
+	vcx::timestepper
         vcx::utils_${GEOMETRY_VARIANT}
         vcx::utils
 )

diff --git a/src/geometryXVx/rhs/collisions_inter.cpp b/src/geometryXVx/rhs/collisions_inter.cpp
@@ -3,6 +3,7 @@
 #include <fluid_moments.hpp>
 #include <maxwellianequilibrium.hpp>
 #include <pdi.h>
+#include <rk2.hpp>
 
 #include "collisions_inter.hpp"
 #include "collisions_utils.hpp"
@@ -37,89 +38,71 @@ CollisionsInter::CollisionsInter(IDomainSpXVx const& mesh, double nustar0)
 /**
  * right hand side of the equation \partial f / \partial_t = C_ab
  */
-void CollisionsInter::compute_rhs(
-        DSpanSpVx const rhs,
-        DViewSp const nustar_profile,
-        DViewSpVx const allfdistribu) const
+void CollisionsInter::compute_rhs(DSpanSpXVx const rhs, DViewSpXVx const allfdistribu) const
 {
     IDomainSp const dom_sp(ddc::get_domain<IDimSp>(allfdistribu));
     IDomainVx const gridvx(ddc::get_domain<IDimVx>(allfdistribu));
 
     FluidMoments moments(Quadrature<IDimVx>(trapezoid_quadrature_coefficients(gridvx)));
 
-    //Moments computation
-    DFieldSp density(dom_sp);
-    DFieldSp fluid_velocity(dom_sp);
-    DFieldSp temperature(dom_sp);
-    ddc::for_each(dom_sp, [&](IndexSp const isp) {
-        moments(density(isp), allfdistribu[isp], FluidMoments::s_density);
-        moments(fluid_velocity(isp), allfdistribu[isp], density(isp), FluidMoments::s_velocity);
-        moments(temperature(isp),
-                allfdistribu[isp],
-                density(isp),
-                fluid_velocity(isp),
-                FluidMoments::s_temperature);
-    });
+    ddc::for_each(ddc::get_domain<IDimX>(allfdistribu), [&](IndexX const ix) {
+        //Moments computation
+        DFieldSp density(dom_sp);
+        DFieldSp fluid_velocity(dom_sp);
+        DFieldSp temperature(dom_sp);
+        ddc::for_each(dom_sp, [&](IndexSp const isp) {
+            IndexSpX ispx(isp, ix);
+            moments(density(isp), allfdistribu[ispx], FluidMoments::s_density);
+            moments(fluid_velocity(isp),
+                    allfdistribu[ispx],
+                    density(isp),
+                    FluidMoments::s_velocity);
+            moments(temperature(isp),
+                    allfdistribu[ispx],
+                    density(isp),
+                    fluid_velocity(isp),
+                    FluidMoments::s_temperature);
+        });
 
-    //Collision frequencies, momentum and energy exchange terms
-    DFieldSp collfreq_ab(dom_sp);
-    compute_collfreq_ab(collfreq_ab.span_view(), nustar_profile, density, temperature);
-    DFieldSp momentum_exchange_ab(dom_sp);
-    DFieldSp energy_exchange_ab(dom_sp);
-    compute_momentum_energy_exchange(
-            momentum_exchange_ab.span_view(),
-            energy_exchange_ab.span_view(),
-            collfreq_ab,
-            density,
-            fluid_velocity,
-            temperature);
+        //Collision frequencies, momentum and energy exchange terms
+        DFieldSp collfreq_ab(dom_sp);
+        DFieldSp nustar_profile_copy(dom_sp);
+        ddc::deepcopy(nustar_profile_copy, m_nustar_profile[ix]);
+        compute_collfreq_ab(collfreq_ab.span_view(), nustar_profile_copy, density, temperature);
+        DFieldSp momentum_exchange_ab(dom_sp);
+        DFieldSp energy_exchange_ab(dom_sp);
+        compute_momentum_energy_exchange(
+                momentum_exchange_ab.span_view(),
+                energy_exchange_ab.span_view(),
+                collfreq_ab,
+                density,
+                fluid_velocity,
+                temperature);
 
-    ddc::for_each(dom_sp, [&](IndexSp const isp) {
-        DFieldVx fmaxwellian(gridvx);
-        MaxwellianEquilibrium::compute_maxwellian(
-                fmaxwellian.span_view(),
-                density(isp),
-                temperature(isp),
-                fluid_velocity(isp));
+        ddc::for_each(dom_sp, [&](IndexSp const isp) {
+            DFieldVx fmaxwellian(gridvx);
+            MaxwellianEquilibrium::compute_maxwellian(
+                    fmaxwellian.span_view(),
+                    density(isp),
+                    temperature(isp),
+                    fluid_velocity(isp));
 
-        ddc::for_each(gridvx, [&](IndexVx const ivx) {
-            double const coordv = ddc::coordinate(ivx);
-            double const term_v(coordv - fluid_velocity(isp));
-            rhs(isp, ivx) = (2. * energy_exchange_ab(isp)
-                                     * (0.5 / temperature(isp) * term_v * term_v - 0.5)
-                             + momentum_exchange_ab(isp) * term_v)
-                            * fmaxwellian(ivx) / (density(isp) * temperature(isp));
+            ddc::for_each(gridvx, [&](IndexVx const ivx) {
+                double const coordv = ddc::coordinate(ivx);
+                double const term_v(coordv - fluid_velocity(isp));
+                rhs(isp, ix, ivx) = (2. * energy_exchange_ab(isp)
+                                             * (0.5 / temperature(isp) * term_v * term_v - 0.5)
+                                     + momentum_exchange_ab(isp) * term_v)
+                                    * fmaxwellian(ivx) / (density(isp) * temperature(isp));
+            });
         });
     });
 }
 
 
 DSpanSpXVx CollisionsInter::operator()(DSpanSpXVx allfdistribu, double dt) const
 {
-    IDomainVx const gridvx(ddc::get_domain<IDimVx>(allfdistribu));
-    FluidMoments moments(Quadrature<IDimVx>(trapezoid_quadrature_coefficients(gridvx)));
-    ddc::for_each(ddc::get_domain<IDimX>(allfdistribu), [&](IndexX const ix) {
-        //RK2 first step
-        DFieldSpVx coll_term(ddc::get_domain<IDimSp, IDimVx>(allfdistribu));
-        DFieldSp nustar_profile_copy(ddc::get_domain<IDimSp>(allfdistribu));
-        ddc::deepcopy(nustar_profile_copy, m_nustar_profile[ix]);
-        DFieldSpVx allfdistribu_copy(ddc::get_domain<IDimSp, IDimVx>(allfdistribu));
-        ddc::deepcopy(allfdistribu_copy, allfdistribu[ix]);
-        DFieldSpVx allfdistribu_half(ddc::get_domain<IDimSp, IDimVx>(allfdistribu));
-        compute_rhs(coll_term.span_view(), nustar_profile_copy.span_cview(), allfdistribu_copy);
-        ddc::for_each(ddc::get_domain<IDimSp, IDimVx>(allfdistribu), [&](IndexSpVx const ispvx) {
-            IndexSp const isp(ddc::select<IDimSp>(ispvx));
-            IndexVx const ivx(ddc::select<IDimVx>(ispvx));
-            allfdistribu_half(isp, ivx) = allfdistribu(isp, ix, ivx) + coll_term(ispvx) * dt / 2;
-        });
-
-        //RK2 final step
-        compute_rhs(coll_term.span_view(), nustar_profile_copy.span_cview(), allfdistribu_half);
-        ddc::for_each(ddc::get_domain<IDimSp, IDimVx>(allfdistribu), [&](IndexSpVx const ispvx) {
-            IndexSp const isp(ddc::select<IDimSp>(ispvx));
-            IndexVx const ivx(ddc::select<IDimVx>(ispvx));
-            allfdistribu(isp, ix, ivx) = allfdistribu(isp, ix, ivx) + coll_term(ispvx) * dt;
-        });
-    });
+    RK2<DFieldSpXVx> timestepper(allfdistribu.domain());
+    timestepper.update(allfdistribu, dt, [&](DSpanSpXVx dy, DViewSpXVx y) { compute_rhs(dy, y); });
     return allfdistribu;
 }
diff --git a/src/geometryXVx/rhs/collisions_inter.hpp b/src/geometryXVx/rhs/collisions_inter.hpp
@@ -10,22 +10,48 @@
 #include <quadrature.hpp>
 #include <trapezoid_quadrature.hpp>
 
-
+/**
+ * @brief Class describing the inter-species collision operator
+ */
 class CollisionsInter : public IRightHandSide
 {
 private:
     double m_nustar0;
     DFieldSpX m_nustar_profile;
 
 public:
+    /**
+     * @brief The constructor for the operator.
+     *
+     * @param[in] mesh The domain on which the operator will act.
+     * @param[in] nustar0 The collision coefficient.
+     */
     CollisionsInter(IDomainSpXVx const& mesh, double nustar0);
 
     ~CollisionsInter() = default;
 
+    /**
+     * @brief Update the distribution function for inter-species collision.
+     *
+     * Update the distribution function for both electrons and ions to show how
+     * it is modified following collisions between the various species.
+     * This operator only handles collisions between particles of different
+     * species.
+     *
+     * @param[inout] allfdistribu The distribution function.
+     * @param[in] dt The time step over which the collisions occur.
+     *
+     * @return A span referencing the distribution function passed as argument.
+     */
     DSpanSpXVx operator()(DSpanSpXVx allfdistribu, double dt) const override;
 
+    /**
+     * @brief Get the collision coefficient.
+     *
+     * @return The collision coefficient.
+     */
     double get_nustar0() const;
 
 private:
-    void compute_rhs(DSpanSpVx rhs, DViewSp nustar_profile, DViewSpVx allfdistribu) const;
+    void compute_rhs(DSpanSpXVx rhs, DViewSpXVx allfdistribu) const;
 };
diff --git a/src/geometryXVx/rhs/krook_source_adaptive.cpp b/src/geometryXVx/rhs/krook_source_adaptive.cpp
@@ -6,6 +6,7 @@
 
 #include <maxwellianequilibrium.hpp>
 #include <quadrature.hpp>
+#include <rk2.hpp>
 #include <species_info.hpp>
 #include <trapezoid_quadrature.hpp>
 
@@ -109,35 +110,35 @@ void KrookSourceAdaptive::get_amplitudes(DSpanSp amplitudes, DViewSpVx const all
     amplitudes(ielec()) = m_amplitude * (density_ion - m_density) / (density_electron - m_density);
 }
 
-
-DSpanSpXVx KrookSourceAdaptive::operator()(DSpanSpXVx const allfdistribu, double const dt) const
+void KrookSourceAdaptive::get_derivative(
+        DSpanSpXVx df,
+        DViewSpXVx allfdistribu,
+        DViewSpXVx allfdistribu_start) const
 {
-    ddc::for_each(ddc::get_domain<IDimX>(allfdistribu), [&](IndexX const ix) {
-        // RK2 first half step
-        DFieldSpVx allfdistribu_half(allfdistribu[ix]);
-        DFieldSp amplitudes(ddc::get_domain<IDimSp>(allfdistribu));
-        get_amplitudes(amplitudes.span_view(), allfdistribu_half.span_cview());
-        ddc::for_each(ddc::get_domain<IDimSp, IDimVx>(allfdistribu), [&](IndexSpVx const ispvx) {
-            IndexSp isp(ddc::select<IDimSp>(ispvx));
-            IndexVx ivx(ddc::select<IDimVx>(ispvx));
-            IndexSpXVx ispxvx(isp, ix, ivx);
-            double const df = -m_mask(ix) * amplitudes(isp)
-                              * (allfdistribu(ispxvx) - m_ftarget(ivx)) * dt / 2.0;
-            allfdistribu_half(isp, ivx) = allfdistribu(ispxvx) + df;
-        });
+    IDomainSpVx sp_vx_dom = ddc::get_domain<IDimSp, IDimVx>(allfdistribu);
+    IDomainSp sp_dom = ddc::get_domain<IDimSp>(allfdistribu);
+    IDomainX x_dom = ddc::get_domain<IDimX>(allfdistribu);
 
-        // RK2 final step
-        get_amplitudes(amplitudes.span_view(), allfdistribu_half.span_cview());
-        ddc::for_each(ddc::get_domain<IDimSp, IDimVx>(allfdistribu), [&](IndexSpVx const ispvx) {
+    DFieldSp amplitudes(sp_dom);
+    ddc::for_each(x_dom, [&](IndexX const ix) {
+        DFieldSpVx allfdistribu_slice(allfdistribu[ix]);
+        get_amplitudes(amplitudes.span_view(), allfdistribu_slice);
+        ddc::for_each(sp_vx_dom, [&](IndexSpVx const ispvx) {
             IndexSp isp(ddc::select<IDimSp>(ispvx));
             IndexVx ivx(ddc::select<IDimVx>(ispvx));
             IndexSpXVx ispxvx(isp, ix, ivx);
-
-            double const df
-                    = -m_mask(ix) * amplitudes(isp) * (allfdistribu(ispxvx) - m_ftarget(ivx)) * dt;
-            allfdistribu(ispxvx) = allfdistribu(ispxvx) + df;
+            df(ispxvx)
+                    = -m_mask(ix) * amplitudes(isp) * (allfdistribu_start(ispxvx) - m_ftarget(ivx));
         });
     });
+}
+
+DSpanSpXVx KrookSourceAdaptive::operator()(DSpanSpXVx const allfdistribu, double const dt) const
+{
+    RK2<DFieldSpXVx> timestepper(allfdistribu.domain());
+    timestepper.update(allfdistribu, dt, [&](DSpanSpXVx df, DViewSpXVx f) {
+        get_derivative(df, f, allfdistribu);
+    });
 
     return allfdistribu;
 }
diff --git a/src/geometryXVx/rhs/krook_source_adaptive.hpp b/src/geometryXVx/rhs/krook_source_adaptive.hpp
@@ -35,4 +35,5 @@ class KrookSourceAdaptive : public IRightHandSide
 
 private:
     void get_amplitudes(DSpanSp amplitudes, DViewSpVx allfdistribu) const;
+    void get_derivative(DSpanSpXVx df, DViewSpXVx f, DViewSpXVx f0) const;
 };
diff --git a/src/timestepper/CMakeLists.txt b/src/timestepper/CMakeLists.txt
@@ -0,0 +1,14 @@
+
+
+
+add_library(timestepper INTERFACE)
+target_include_directories(timestepper
+    INTERFACE
+        "$<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}>"
+)
+target_link_libraries(timestepper INTERFACE
+    DDC::DDC
+    sll::splines
+    vcx::utils
+)
+add_library(vcx::timestepper ALIAS timestepper)
diff --git a/src/timestepper/README.md b/src/timestepper/README.md
@@ -0,0 +1,15 @@
+# Time Stepping Methods
+
+Time stepping methods are methods for calculating how a value (or dimensioned values) evolves over time. Such an evolution should be expressed as a system of equations of the form:
+
+$$
+dy/dt = f(t, y)\\
+y(t_0) = y_0
+$$
+
+The implemented methods are:
+- Second order Runge Kutta (RK2)
+- Third order Runge Kutta (RK3)
+- Fourth order Runge Kutta (RK4)
+
+These classes all contain an `update` method which carries out one time step of the algorithm.
diff --git a/src/timestepper/itimestepper.hpp b/src/timestepper/itimestepper.hpp
@@ -0,0 +1,8 @@
+#pragma once
+#include <array>
+
+#include <ddc/ddc.hpp>
+
+class ITimeStepper
+{
+};