diff --git a/doxygen/CMakeLists.txt b/doxygen/CMakeLists.txt
index 406117070..202682737 100644
--- a/doxygen/CMakeLists.txt
+++ b/doxygen/CMakeLists.txt
@@ -70,6 +70,7 @@ doxygen_add_docs(asgard_doxygen
                  src/pde/spherical_diffusion.cpp
                  src/pde/sinwav.cpp
                  src/pde/two_stream.cpp
+                 src/pde/two_stream_md.cpp
                  src/pde/vplb.cpp
                  src/pde/sod_shock.cpp
                  src/pde/burgers.cpp
diff --git a/src/asgard_discretization.cpp b/src/asgard_discretization.cpp
index 1ab400fa4..c8d0134f2 100644
--- a/src/asgard_discretization.cpp
+++ b/src/asgard_discretization.cpp
@@ -414,7 +414,7 @@ std::vector<precision> discretization_manager<precision>::get_moment(moment_id i
       else
         throw std::runtime_error("an electric moment was requested but no poisson solver is set");
     }, get_poisson());
-    terms.moms.solve_poisson(terms.grid, terms.conn, terms.hier, terms.interp, terms.kwork, true);
+    terms.moms.solve_poisson(terms.grid, terms.conn, terms.hier, terms.interp, terms.kwork);
     return terms.moms.get_cached_raw(id);
   }
   std::vector<precision> result;
diff --git a/src/asgard_gpu_tensors_cuda.cpp b/src/asgard_gpu_tensors_cuda.cpp
index 6dae9ae2b..394035a88 100644
--- a/src/asgard_gpu_tensors_cuda.cpp
+++ b/src/asgard_gpu_tensors_cuda.cpp
@@ -473,7 +473,7 @@ __global__ void kernel_moment_expand(int pos_block, int vel_block,
     } else
       for (int c = 0; c < num_cycles; c++) {
         P const src = pos_data[i * pos_block + (threadIdx.x + c * cycle_size) / vel_block];
-        vals[j * pos_block * vel_block + threadIdx.x + (threadIdx.x + c * cycle_size)] = src;
+        vals[j * pos_block * vel_block + (threadIdx.x + c * cycle_size)] = src;
       }
 
     blk += gridDim.x * blockDim.y;
diff --git a/src/asgard_gpu_tensors_rocm.cpp b/src/asgard_gpu_tensors_rocm.cpp
index 239ad0ef6..c8ec3a251 100644
--- a/src/asgard_gpu_tensors_rocm.cpp
+++ b/src/asgard_gpu_tensors_rocm.cpp
@@ -473,7 +473,7 @@ __global__ void kernel_moment_expand(int pos_block, int vel_block,
     } else
       for (int c = 0; c < num_cycles; c++) {
         P const src = pos_data[i * pos_block + (threadIdx.x + c * cycle_size) / vel_block];
-        vals[j * pos_block * vel_block + threadIdx.x + (threadIdx.x + c * cycle_size)] = src;
+        vals[j * pos_block * vel_block + (threadIdx.x + c * cycle_size)] = src;
       }
 
     blk += gridDim.x * blockDim.y;
diff --git a/src/asgard_io.cpp b/src/asgard_io.cpp
index 7e2136112..4c53d9c60 100644
--- a/src/asgard_io.cpp
+++ b/src/asgard_io.cpp
@@ -141,8 +141,7 @@ void h5manager<P>::write(prog_opts const &options, pde_domain<P> const &domain,
     if (moms.has_poisson()) {
       moment_id m0 = moms.find_id(moment::zero(moms.num_vel()));
       moms.cache_moment(m0, grid, state);
-      constexpr bool raw_on_cpu = true;
-      moms.solve_poisson(terms.grid, terms.conn, terms.hier, terms.interp, terms.kwork, raw_on_cpu);
+      moms.solve_poisson(terms.grid, terms.conn, terms.hier, terms.interp, terms.kwork);
     }
     for (int i : iindexof(moms.num_moments()))
     {
@@ -152,10 +151,7 @@ void h5manager<P>::write(prog_opts const &options, pde_domain<P> const &domain,
       H5Easy::dump(file, auxstr.name,
                    std::string("__moment_") + moms.get_by_id(mid).to_string());
       if (moms.needs_poisson(mid)) {
-        if (domain.num_pos() == 1)
-          write_vector(auxstr.data, moms.get_cached_level(mid)); // only the level is computed for poisson_1d
-        else
-          write_vector(auxstr.data, moms.get_cached_raw(mid));
+        write_vector(auxstr.data, moms.get_cached_raw(mid, terms.hier));
       } else {
         std::vector<P> vals;
         moms.mcompute(grid, mid, state, vals);
diff --git a/src/asgard_moment_manager.cpp b/src/asgard_moment_manager.cpp
index bb0d97622..9097af6c8 100644
--- a/src/asgard_moment_manager.cpp
+++ b/src/asgard_moment_manager.cpp
@@ -559,17 +559,19 @@ void moment_manager<P>::cache_moments(
     interpolation_manager<P> const &interp, kronmult::workspace<P> &work) const
 {
   // Compute moments on hierarchical grid
+  bool has_poisson = false;
   if (group == group_id::all()) { // do all moments
     tools::time_event performance_("cache all moments");
     for (auto mid : raw_moments_) {
-      if (mid == moment_id::unset() or needs_poisson(mid)) continue;
-      mcompute(grid, mid, state, raw_vals.get(mid));
+      if (mid == moment_id::unset()) continue;
+      if (needs_poisson(mid))
+        has_poisson = true;
+      else
+        mcompute(grid, mid, state, raw_vals.get(mid));
       full_level.get(mid).resize(0);
     }
-    solve_poisson(grid, conn, hier, interp, work);
   } else {
     tools::time_event performance_("cache moments (" + std::to_string(group()) + ")");
-    bool has_poisson = false;
     for (auto mid = first_in(group, raw_moments_);
          *mid != moment_id::unset(); mid++) {
       if (needs_poisson(*mid))
@@ -578,8 +580,8 @@ void moment_manager<P>::cache_moments(
         mcompute(grid, *mid, state, raw_vals.get(*mid));
       full_level.get(*mid).resize(0);
     }
-    if (has_poisson) solve_poisson(grid, conn, hier, interp, work);
   }
+  if (has_poisson) solve_poisson(grid, conn, hier, interp, work);
 
   // Poisson moments are already interpolated if they were computed on GPU
   auto skip_poisson = [&](moment_id mid) -> bool {
@@ -596,13 +598,13 @@ void moment_manager<P>::cache_moments(
   if (group == group_id::all()) {
     for (auto mid : interp_moments_) {
       if (skip_poisson(mid) or mid == moment_id::unset()) continue;
-      make_nodal(mid, interp, work, interp.it1);
+      make_nodal(mid, interp, hier, work, interp.it1);
     }
   } else {
     for (auto mid = first_in(group, interp_moments_);
          *mid != moment_id::unset(); mid++) {
       if (skip_poisson(*mid)) continue;
-      make_nodal(*mid, interp, work, interp.it1);
+      make_nodal(*mid, interp, hier, work, interp.it1);
     }
   }
   interp.it1.resize(num_entries);
@@ -641,99 +643,109 @@ void moment_manager<P>::complete_level(hierarchy_manipulator<P> const &hier,
 template<typename P>
 void moment_manager<P>::solve_poisson(sparse_grid const &grid, connection_patterns const &conn,
                                       hierarchy_manipulator<P> const &hier, interpolation_manager<P> const &interp,
-                                      kronmult::workspace<P> &work, bool raw_on_cpu) const
+                                      kronmult::workspace<P> &work) const
 {
   std::visit([&](auto &p) {
-      if constexpr (std::is_same_v<std::decay_t<decltype(p)>, poisson_1d<P>>) {
-        p.solve_periodic(get_cached_level(p.moment0(), hier), full_level.get(p.moment_electric()));
-      } else if constexpr (std::is_same_v<std::decay_t<decltype(p)>, poisson_md<P>>) {
-        #ifdef ASGARD_USE_GPU
-        compute->set_device(gpu::device{0});
-        int const num_entries = pos_grid.num_dof();
-        std::array<gpu::vector<P>, max_num_gpus> &work1 = interp.gpu_it1;
-        assert(work1[0].size() >= num_entries);
-        work1[0].copy_from_host(num_entries, raw_vals.get(p.moment0()).data());
-        for (moment_id const &mid : p.moments_electric()) {
-          if (mid == moment_id::unset()) continue;
-          gpu::vector<P> &res = gpu_interps[0][mid];
-          res.resize(full_block * grid.num_indexes());
-        }
-        solve_poisson_gpu(conn, hier, interp, work, raw_on_cpu);
-        auto find_mom = [this](moment_id mid) -> const mom_on_gpu &
-        {
-          for (auto const &mom : this->gpu_moments[0]) {
-            if (mom.mid == mid)
-              return mom;
-          }
-          throw std::runtime_error("could not find the moment");
-        };
-        for (moment_id const &mid : p.moments_electric()) {
-          if (mid == moment_id::unset()) continue;
-          mom_on_gpu const &mom = find_mom(mid);
-          if (mom.interp_on_cpu()) continue; // solve_poisson_gpu will already copy if interp_on_cpu() is true
-          gpu::vector<P> &res = gpu_interps[0][mid];
-          res.copy_to_host(interps.get(mid));
-        }
-        #else
-        std::ignore = interp;
-        std::ignore = grid;
-        std::ignore = raw_on_cpu;
-        update_position_grid_dsort();
-        p.solve_periodic(raw_vals.get(p.moment0()), raw_vals, pos_grid, conn, work);
-        #endif
-      }
-    }, poisson_solver);
+    if constexpr (std::is_same_v<std::decay_t<decltype(p)>, poisson_1d<P>>) {
+      p.solve_periodic(get_cached_level(p.moment0(), hier), full_level.get(p.moment_electric()));
+    } else if constexpr (std::is_same_v<std::decay_t<decltype(p)>, poisson_md<P>>) {
+      update_position_grid_dsort();
+
+      #ifdef ASGARD_USE_GPU
+      compute->set_device(gpu::device{0});
+
+      // transfer density to GPU
+      int const num_entries = pos_grid.num_dof();
+      interp.gpu_it1[0] = raw_vals[p.moment0(0)];
+
+      // setup interpolation function
+      auto interp_func = [&](gpu::vector<P> const &efield, moment_id mid) -> void
+      {
+        efield.copy_to_host(this->raw_vals[mid]);
+        this->full_level.get(mid).resize(0);
+        interp.pos2nodal(dev, this->pos_grid, efield.data(), this->wav_scale, w2.vec.data(), work);
+        gpu::vector<P> &res = this->gpu_interps[dev.id][mid];
+        res.resize(this->full_block * grid.num_indexes());
+        moment_expand(pdof, this->pos_grid.num_dims(), num_vel_, reduce_ij[dev.id], w2.vec, res);
+        res.copy_to_host(interps[mid]);
+      };
+
+      // solve
+      p.solve_periodic(w1.vec, pos_grid, conn, interp_func, work);
+      #else
+      std::ignore = interp;
+      std::ignore = grid;
+      p.solve_periodic(raw_vals.get(p.moment0()), raw_vals, pos_grid, conn, work);
+      #endif
+    }
+  }, poisson_solver);
 }
 
 #ifdef ASGARD_USE_GPU
 template<typename P>
-void moment_manager<P>::solve_poisson_gpu(connection_patterns const &conn, hierarchy_manipulator<P> const &hier,
-                                          interpolation_manager<P> const &interp, kronmult::workspace<P> &work, bool raw_on_cpu) const
+template<bool always_interp>
+void moment_manager<P>::solve_poisson_gpu(gpu::device dev, sparse_grid const &grid, connection_patterns const &conn, hierarchy_manipulator<P> const &hier,
+                                          interpolation_manager<P> const &interp, kronmult::workspace<P> &work, bool result_to_cpu) const
 {
+  // Setup
+  int64_t const num_entries = pos_grid.num_dof();
+  std::array<gpu::vector<P>, max_num_gpus> &work1 = interp.gpu_it1;
+  std::array<gpu::vector<P>, max_num_gpus> &work2 = interp.gpu_it2;
+  assert(work1[dev.id].size() >= num_entries);
+  assert(work2[dev.id].size() >= num_entries);
+  gpu::wrap_array<P> w1(work1[dev.id].data(), num_entries);
+  gpu::wrap_array<P> w2(work2[dev.id].data(), num_entries);
+  
   std::visit([&](auto &p) {
-      if constexpr (std::is_same_v<std::decay_t<decltype(p)>, poisson_1d<P>>) {
-        // This is fast enough on CPU
-        int64_t const num_entries = pos_grid.num_dof();
-        std::array<gpu::vector<P>, max_num_gpus> &work1 = interp.gpu_it1;
-        assert(work1[0].size() >= num_entries);
-        gpu::wrap_array<P> w1(work1[0].data(), num_entries);
-        w1.vec.copy_to_host(raw_vals[p.moment0()]);
-        p.solve_periodic(get_cached_level(p.moment0(), hier), full_level.get(p.moment_electric()));
-      } else if constexpr (std::is_same_v<std::decay_t<decltype(p)>, poisson_md<P>>) {
-        // Setup
-        update_position_grid_dsort();
-        int64_t const num_entries = pos_grid.num_dof();
-        std::array<gpu::vector<P>, max_num_gpus> &work1 = interp.gpu_it1;
-        std::array<gpu::vector<P>, max_num_gpus> &work2 = interp.gpu_it2;
-        assert(work1[0].size() >= num_entries);
-        assert(work2[0].size() >= num_entries);
-        gpu::wrap_array<P> w1(work1[0].data(), num_entries);
-        gpu::wrap_array<P> w2(work2[0].data(), num_entries);
-        auto find_mom = [this](moment_id mid) -> const mom_on_gpu &
-        {
-          for (auto const &mom : this->gpu_moments[0]) {
-            if (mom.mid == mid)
-              return mom;
-          }
-          throw std::runtime_error("could not find the moment");
-        };
-        auto interp_func = [&](gpu::vector<P> const &efield, moment_id mid) -> void
-        {
-          if (raw_on_cpu) {
-            efield.copy_to_host(this->raw_vals[mid]);
-            this->full_level.get(mid).resize(0);
-          }
-          interp.pos2nodal(gpu::device{0}, this->pos_grid, efield.data(), this->wav_scale, w2.vec.data(), work);
-          gpu::vector<P> &res = this->gpu_interps[0][mid];
-          moment_expand(pdof, this->pos_grid.num_dims(), num_vel_, reduce_ij[0], w2.vec, res);
-          mom_on_gpu const &mom = find_mom(mid);
-          if (mom.interp_on_cpu()) res.copy_to_host(interps[mom.mid]);
-        };
-        // Solve
-        p.solve_periodic(w1.vec, pos_grid, conn, interp_func, work);
+    if constexpr (std::is_same_v<std::decay_t<decltype(p)>, poisson_1d<P>>) {
+      // transfer density to CPU
+      moment_id const m0 = p.moment0();
+      moment_id const melectric = p.moment_electric();
+      w1.vec.copy_to_host(raw_vals[m0]);
+
+      // solve
+      p.solve_periodic(get_cached_level(m0, hier), full_level.get(melectric));
+
+      // interpolate
+      if constexpr (always_interp) {
+        cache_raw_from_level(melectric, hier);
+        w1.vec.copy_from_host(raw_vals[melectric].size(), raw_vals[melectric].data());
+        interp.pos2nodal(dev, pos_grid, w1.vec.data(), wav_scale, w2.vec.data(), work);
+        gpu::vector<P> &res = gpu_interps[dev.id][melectric];
+        res.resize(full_block * grid.num_indexes());
+        moment_expand(pdof, pos_grid.num_dims(), num_vel_, reduce_ij[dev.id], w2.vec, res);
+        if (result_to_cpu) res.copy_to_host(interps[melectric]);
       }
-    },
-    poisson_solver);
+    } else if constexpr (std::is_same_v<std::decay_t<decltype(p)>, poisson_md<P>>) {
+      // setup
+      update_position_grid_dsort();
+      auto find_mom = [this](moment_id mid) -> const mom_on_gpu &
+      {
+        for (auto const &mom : this->gpu_moments[0]) {
+          if (mom.mid == mid)
+            return mom;
+        }
+        throw std::runtime_error("could not find the moment");
+      };
+      auto interp_func = [&](gpu::vector<P> const &efield, moment_id mid) -> void
+      {
+        mom_on_gpu const &mom = find_mom(mid);
+        if (result_to_cpu or mom.raw_on_cpu()) {
+          efield.copy_to_host(this->raw_vals[mid]);
+          this->full_level.get(mid).resize(0);
+        }
+        interp.pos2nodal(dev, this->pos_grid, efield.data(), this->wav_scale, w2.vec.data(), work);
+        gpu::vector<P> &res = this->gpu_interps[dev.id][mid];
+        res.resize(this->full_block * grid.num_indexes());
+        moment_expand(pdof, this->pos_grid.num_dims(), num_vel_, reduce_ij[dev.id], w2.vec, res);
+        if (result_to_cpu or mom.interp_on_cpu())
+          res.copy_to_host(interps[mid]);
+      };
+
+      // solve
+      p.solve_periodic(w1.vec, pos_grid, conn, interp_func, work);
+    }
+  }, poisson_solver);
 }
 #endif // ASGARD_USE_GPU
 
@@ -751,12 +763,14 @@ void moment_manager<P>::update_position_grid_dsort() const
 template<typename P>
 void moment_manager<P>::make_nodal(
     moment_id id, interpolation_manager<P> const &interp,
+    hierarchy_manipulator<P> const &hier,
     kronmult::workspace<P> &kwork, std::vector<P> &workspace) const
 {
-  assert(not (num_pos_ == 1 and needs_poisson(id))); // poisson_1d solves the moment for the full level
-                                                     // so there is no raw moment to interpolate
   update_position_grid_dsort();
 
+  if (num_pos_ == 1 and needs_poisson(id))
+    cache_raw_from_level(id, hier);
+    
   interp.pos2nodal(pos_grid, raw_vals[id].data(), wav_scale, workspace, kwork);
 
   interps[id].resize(pntr.back() * full_block);
@@ -799,7 +813,7 @@ void moment_manager<P>::compute_interps(
     solve_poisson(grid, conn, hier, interp, work);
   }
   for (auto const &id : ids)
-    make_nodal(id, interp, work, interp.it1);
+    make_nodal(id, interp, hier, work, interp.it1);
   interp.it1.resize(num_entries);
 }
 
@@ -954,12 +968,7 @@ void moment_manager<P>::prepare_pos_grid_gpu(group_id group, sparse_grid const &
 
   if (not has_interp.empty()) {
     if (group == group_id::all() or has_interp[group()]) {
-      if (dsort_generation != pos_grid.generation()) {
-        pos_grid.dsort_  = dimension_sort(pos_grid.iset_);
-        bool constexpr skip_indexes = true; // already loaded in reduce_grid()
-        pos_grid.gpu_sync<skip_indexes>();
-        dsort_generation = pos_grid.generation();
-      }
+      update_position_grid_dsort();
     }
   }
 }
@@ -1012,13 +1021,13 @@ void moment_manager<P>::compute_moments(
 
       moment const mom = mlist[im->mid]; // using this to get the necessary powers
 
+      if (mom.is_electric()) continue; // skip electric field moments
+
       gpu::vector<P> &res = gpu_interps[g][im->mid];
       if (not im->skip_interp()) {
         res.resize(full_block * grid.num_indexes());
       }
 
-      if (mom.is_electric()) continue; // skip electric field moments
-
       std::array<P const *, max_mom_dims> itg =
         {gpu_integ[g][0].data() + mom.pows[0] * integ[0].stride(), nullptr, nullptr};
       for (int i = 1; i < num_vel_; i++)
@@ -1048,10 +1057,8 @@ void moment_manager<P>::compute_moments(
       }
 
       // solve poisson equation while w1 holds density
-      if (mom.is_zero()) {
-        assert(g == 0); // poisson uses gpu 0 for solve
-        solve_poisson_gpu(conn, hier, interp, kwork);
-      }
+      if (mom.is_zero())
+        solve_poisson_gpu(gpu::device{g}, grid, conn, hier, interp, kwork);
 
       if (im->skip_interp()) continue;
 
@@ -1083,6 +1090,7 @@ void moment_manager<P>::compute_moments(
   // which will also allow to avoid the spread when doing term-apply
 
   prepare_pos_grid_gpu(group_id::all(), grid);
+  update_position_grid_dsort(); // interpolation is always done in this overload of compute_moments()
 
   int const pos_block = pos_grid.block_size();
 
@@ -1108,11 +1116,11 @@ void moment_manager<P>::compute_moments(
   {
     moment const mom = mlist[im]; // using this to get the necessary powers
 
+    if (mom.is_electric()) continue; // skip electric field moments
+
     gpu::vector<P> &res = gpu_interps[0][im];
     res.resize(full_block * grid.num_indexes());
 
-    if (mom.is_electric()) continue; // skip electric field moments
-
     std::array<P const *, max_mom_dims> itg =
       {gpu_integ[0][0].data() + mom.pows[0] * integ[0].stride(), nullptr, nullptr};
     for (int i = 1; i < num_vel_; i++)
@@ -1142,11 +1150,15 @@ void moment_manager<P>::compute_moments(
 
     moment_expand(pdof, pos_grid.num_dims(), num_vel_, reduce_ij[0], w2.vec, res);
 
-    if (result_to_cpu) res.copy_to_host(interps[im]);
+    if (result_to_cpu) {
+      res.copy_to_host(interps[im]);
+      full_level.get(im).resize(0);
+    }
 
     // solve poisson equation while w1 holds density
+    constexpr bool always_interp = true;
     if (mom.is_zero() and has_electric)
-      solve_poisson_gpu(conn, hier, interp, kwork);
+      solve_poisson_gpu<always_interp>(gpu::device{0}, grid, conn, hier, interp, kwork, result_to_cpu);
   }
 }
 #endif
diff --git a/src/asgard_moment_manager.hpp b/src/asgard_moment_manager.hpp
index a0b5f4dd1..4aa72313b 100644
--- a/src/asgard_moment_manager.hpp
+++ b/src/asgard_moment_manager.hpp
@@ -88,6 +88,15 @@ class moment_manager {
   //! fill the vector to a full 1d level, only for position 1d
   void complete_level(hierarchy_manipulator<P> const &hier, std::vector<P> const &raw,
                       std::vector<P> &vals) const;
+  //! cache the raw moment from the level, only for position 1d and electric field moments
+  void cache_raw_from_level(moment_id id, hierarchy_manipulator<P> const &hier) const {
+    assert(num_pos_ == 1 and needs_poisson(id)); // this should only be called for 1D electric field moments
+    int const level = pos_grid.current_level(0);
+    size_t const num_cells = pos_grid.num_indexes();
+    constexpr int full_level_pdof = 1;
+    std::vector<P> padded_level = pad_pdof(num_cells, full_level_pdof, pdof, full_level[id]);
+    hier.transform(level, padded_level, raw_vals[id]);
+  }
   //! cache a number of ids listed as the first n entries of a container ids, where ids[i] is moment_id
   template<typename vec_type>
   void cache_levels(int num, hierarchy_manipulator<P> const &hier, vec_type const &ids) const {
@@ -109,19 +118,20 @@ class moment_manager {
   sparse_grid const &get_position_grid() const { return pos_grid; }
   //! return the cached raw moment defined on the position grid with moment id mid
   std::vector<P> const &get_cached_raw(moment_id mid) const {
-    rassert(not (num_pos_ == 1 and needs_poisson(mid)), "The electric field moment is only computed for the full level in 1D");
-    return raw_vals.get(mid);
+    rassert(not (num_pos_ == 1 and needs_poisson(mid)), 
+      "The electric field moment is only computed for the full level in 1D, use get_cached_raw(mid, hier) to convert from the level to the raw moment instead");
+    return raw_vals[mid];
   }
-
-  //! solves the poisson equation and caches them as electric moments
-  void solve_poisson(sparse_grid const &grid, connection_patterns const &conn, hierarchy_manipulator<P> const &hier,
-                     interpolation_manager<P> const &interp, kronmult::workspace<P> &work) const {
-    constexpr bool raw_on_cpu = false;
-    solve_poisson(grid, conn, hier, interp, work, raw_on_cpu);
+  //! return the cached raw moment defined on the position grid with moment id mid, overload for 1d electric moments
+  std::vector<P> const &get_cached_raw(moment_id mid, hierarchy_manipulator<P> const &hier) const {
+    if (num_pos_ == 1 and needs_poisson(mid))
+      cache_raw_from_level(mid, hier);
+    return raw_vals[mid];
   }
+
   //! solves the poisson equation and caches them as electric moments
   void solve_poisson(sparse_grid const &grid, connection_patterns const &conn, hierarchy_manipulator<P> const &hier,
-                     interpolation_manager<P> const &interp, kronmult::workspace<P> &work, bool raw_on_cpu) const;
+                     interpolation_manager<P> const &interp, kronmult::workspace<P> &work) const;
   //! compute the specified interpolated moments
   void compute_interps(std::vector<moment_id> const &ids, sparse_grid const &grid,
                        std::vector<P> const &state, interpolation_manager<P> const &interp,
@@ -175,14 +185,16 @@ class moment_manager {
                        hierarchy_manipulator<P> const &hier, kronmult::workspace<P> &kwork,
                        gpu::vector<P> const &state, bool result_to_cpu) const;
   //! solves the poisson equation and caches them as electric moments, on gpu
-  void solve_poisson_gpu(connection_patterns const &conn, hierarchy_manipulator<P> const &hier,
+  template<bool always_interp = false>
+  void solve_poisson_gpu(gpu::device dev, sparse_grid const &grid, connection_patterns const &conn, hierarchy_manipulator<P> const &hier,
                          interpolation_manager<P> const &interp, kronmult::workspace<P> &work) const {
-    constexpr bool raw_on_cpu = false;
-    solve_poisson_gpu(conn, hier, interp, work, raw_on_cpu);
+    constexpr bool result_to_cpu = false;
+    solve_poisson_gpu<always_interp>(dev, grid, conn, hier, interp, work, result_to_cpu);
   }
   //! solves the poisson equation and caches them as electric moments, on gpu
-  void solve_poisson_gpu(connection_patterns const &conn, hierarchy_manipulator<P> const &hier,
-                         interpolation_manager<P> const &interp, kronmult::workspace<P> &work, bool raw_on_cpu) const;
+  template<bool always_interp = false>
+  void solve_poisson_gpu(gpu::device dev, sparse_grid const &grid, connection_patterns const &conn, hierarchy_manipulator<P> const &hier,
+                         interpolation_manager<P> const &interp, kronmult::workspace<P> &work, bool result_to_cpu) const;
   #endif
   /*!
    * \brief Defines moments that should be used as raw or interpolation
@@ -249,6 +261,7 @@ class moment_manager {
    * \brief computes the nodal values of the moment
    */
   void make_nodal(moment_id id, interpolation_manager<P> const &interp,
+                  hierarchy_manipulator<P> const &hier,
                   kronmult::workspace<P> &work, std::vector<P> &workspace) const;
 
   #ifdef ASGARD_USE_GPU
diff --git a/src/asgard_pde.cpp b/src/asgard_pde.cpp
index f3abe3b69..a0ef39918 100644
--- a/src/asgard_pde.cpp
+++ b/src/asgard_pde.cpp
@@ -583,7 +583,7 @@ void pde_scheme<P>:: update_deps(term_md<P> &tmd) {
       case term_dependence::electric_field_only:
         rassert(1 <= domain_.num_vel() and domain_.num_vel() <= 3,
                 "electric field dependence requires moments which in turn require 1 - 3 velocity dimensions");
-        t1d.mids_ = {this->register_moment(moment::zero(domain_.num_vel())), this->register_electric_moment(dimension_id(d), domain_.num_pos())};
+        t1d.mids_ = {this->register_moment(moment::zero(domain_.num_vel())), this->register_electric_moment(dimension_id(d))};
         break;
       case term_dependence::moment_divided_by_density:
         rassert(1 <= domain_.num_vel() and domain_.num_vel() <= 3,
diff --git a/src/asgard_pde.hpp b/src/asgard_pde.hpp
index 2cb6e2689..893ac5935 100644
--- a/src/asgard_pde.hpp
+++ b/src/asgard_pde.hpp
@@ -1821,8 +1821,9 @@ class pde_scheme
   }
   //! register an electric field moment in the specified dimension and obtain the moment id
   //! i.e. register_electric_moment(dimension_id(0)) corresponds to E_x
-  moment_id register_electric_moment(dimension_id dim, int num_pos_dims) { // this function is not really needed as it just calls the other register function
-    moment mom = moment::electric(dim, num_pos_dims);
+  moment_id register_electric_moment(dimension_id dim) {
+    assert(domain_.num_pos() > 0);
+    moment mom = moment::electric(dim, domain_.num_pos());
     return this->register_moment(mom);
   }
   //! return true if the pde contains an electric field moment
diff --git a/src/asgard_term_build.cpp b/src/asgard_term_build.cpp
index d7b0ae95c..2e9072be8 100644
--- a/src/asgard_term_build.cpp
+++ b/src/asgard_term_build.cpp
@@ -486,7 +486,9 @@ term_manager<P>::term_manager(prog_opts const &options, pde_domain<P> const &dom
         if (not resources.owns(tentry.rec)) continue;
 
         term_md<P> const &tmd = tentry.tmd;
-        has_poisson = has_poisson or tmd.is_electric(moms.moments());
+        has_poisson |= tmd.is_electric(moms.moments());
+        for (moment_id const mid : pde.ref_moments_)
+          has_poisson |= (resources.is_leader() and moms.get_by_id(mid).is_electric());
         if (tentry.is_separable()) { // only separable terms can have 1D moment deps
           for (int d : iindexof(num_dims)) {
             auto const &mids = tmd.dim(d).mids_;
diff --git a/src/asgard_transformations.cpp b/src/asgard_transformations.cpp
index ac3da458c..69f51999c 100644
--- a/src/asgard_transformations.cpp
+++ b/src/asgard_transformations.cpp
@@ -9,6 +9,22 @@
 namespace asgard
 {
 
+template<typename P>
+std::vector<P> pad_pdof(size_t num_cells, int src_pdof, int dest_pdof, std::vector<P> const &src)
+{
+  if (src_pdof == dest_pdof) return src;
+
+  assert(src_pdof > 0);
+  assert(dest_pdof > src_pdof);
+  assert(src.size() == num_cells * src_pdof);
+
+  size_t const dest_size = num_cells * dest_pdof;
+  std::vector<P> dest(dest_size, P{0});
+  for (size_t i = 0; i < num_cells; i++)
+    std::copy_n(src.data() + i * src_pdof, src_pdof, dest.data() + i * dest_pdof);
+  return dest;
+}
+
 template<typename P>
 legendre_basis<P>::legendre_basis(int degree) : pdof(degree + 1) {
 
@@ -1403,6 +1419,8 @@ void hierarchy_manipulator<P>::setup_projection_matrices()
       prec const *trans, int level, prec src[], prec dest[]) const; \
 
 #ifdef ASGARD_ENABLE_DOUBLE
+template std::vector<double> pad_pdof(size_t num_cells, int src_pdof, int dest_pdof, std::vector<double> const &src);
+
 template struct legendre_basis<double>;
 template class hierarchy_manipulator<double>;
 
@@ -1429,6 +1447,8 @@ instantiate_multi(double, -1);
 #endif
 
 #ifdef ASGARD_ENABLE_FLOAT
+template std::vector<float> pad_pdof(size_t num_cells, int src_pdof, int dest_pdof, std::vector<float> const &src);
+
 template struct legendre_basis<float>;
 template class hierarchy_manipulator<float>;
 
diff --git a/src/asgard_transformations.hpp b/src/asgard_transformations.hpp
index 4702d662c..53d5290da 100644
--- a/src/asgard_transformations.hpp
+++ b/src/asgard_transformations.hpp
@@ -5,6 +5,22 @@
 
 namespace asgard
 {
+
+/*!
+ * \internal
+ * \brief Returns a new vector with \p dest_pdof DG DOFs padded with zeros
+ * 
+ * \p src is expected to come from a \p src_pdof DOF representation
+ * and the returned vector will have \p dest_pdof DG DOFs where the higher DOF coefficients
+ * are set to zero.
+ * 
+ * \p dest_pdof must be greater than \p src_pdof
+ *
+ * \endinternal
+ */
+template<typename P>
+std::vector<P> pad_pdof(size_t num_cells, int src_pdof, int dest_pdof, std::vector<P> const &src);
+
 /*!
  * \internal
  * \brief Shorthand for array of diagonal mass matrices
diff --git a/src/pde/landau_damping.cpp b/src/pde/landau_damping.cpp
index 536b65911..7803d5377 100644
--- a/src/pde/landau_damping.cpp
+++ b/src/pde/landau_damping.cpp
@@ -138,8 +138,8 @@ asgard::pde_scheme<P> make_landau(asgard::prog_opts options) {
         y[i] = std::min(P{0}, x[i]);
     };
 
-  asgard::moment_id melectric_x = pde.register_electric_moment(asgard::dimension_id(0), 2);
-  asgard::moment_id melectric_y = pde.register_electric_moment(asgard::dimension_id(1), 2);
+  asgard::moment_id melectric_x = pde.register_electric_moment(asgard::dimension_id(0));
+  asgard::moment_id melectric_y = pde.register_electric_moment(asgard::dimension_id(1));
 
   auto md_positive_x = [=](P /* time */, asgard::vector2d<P> const& /* nodes */,
                     asgard::momentset<P> const &moments, std::vector<P> const &field,
@@ -457,8 +457,8 @@ void test_damping(std::string const &opt_str) {
   // the pde needs only the zeroth moment and computes that internally
   // we are using the other moments to check energy conservation properties
   auto pde = make_landau(options);
-  moment_id const melectric_x = pde.register_electric_moment(asgard::dimension_id(0), 2);
-  moment_id const melectric_y = pde.register_electric_moment(asgard::dimension_id(1), 2);
+  moment_id const melectric_x = pde.register_electric_moment(asgard::dimension_id(0));
+  moment_id const melectric_y = pde.register_electric_moment(asgard::dimension_id(1));
   discretization_manager disc(std::move(pde), verbosity_level::quiet);
 
   int64_t const nt = disc.remaining_steps();
diff --git a/src/pde/two_stream.cpp b/src/pde/two_stream.cpp
index ac29a041f..f9ae90fc4 100644
--- a/src/pde/two_stream.cpp
+++ b/src/pde/two_stream.cpp
@@ -152,6 +152,19 @@ asgard::pde_scheme<P> make_two_stream(asgard::prog_opts options) {
 
   pde.add_initial(asgard::separable_func<P>({ic_x, ic_v}));
 
+  asgard::moment_id melectric = pde.register_electric_moment(asgard::dimension_id(0));
+  auto weight = [=](P /* time */, asgard::vector2d<P> const& /* nodes */,
+                 asgard::momentset<P> const &moments, std::vector<P> const &field,
+                 std::vector<P> &vals)
+    {
+      std::vector<P> const &ef = moments[melectric];
+#pragma omp parallel for
+      for (size_t i = 0; i < vals.size(); i++)
+        vals[i] = field[i] * ef[i] * ef[i];
+    };
+
+  pde.set_adapt_weight(weight, {melectric});
+
   return pde;
 
 #ifndef __ASGARD_DOXYGEN_SKIP
@@ -245,7 +258,7 @@ void test_energy(std::string const &opt_str) {
   moment_id const m0 = pde.register_moment({0});
   moment_id const m1 = pde.register_moment({1}); // needed for verification, but not running
   moment_id const m2 = pde.register_moment({2});
-  moment_id const melectric_x = pde.register_electric_moment(asgard::dimension_id(0), 1);
+  moment_id const melectric_x = pde.register_electric_moment(asgard::dimension_id(0));
   discretization_manager disc(std::move(pde), verbosity_level::quiet);
 
   P E0 = 0; // initial total energy (potential + kinetic), will initialize on first iteration
@@ -278,6 +291,7 @@ void test_energy(std::string const &opt_str) {
     if (disc.current_step() == 1) // first time-step
       E0 = 0.5 * (Ep + Ek);
 
+    // std::cout << std::abs(0.5 * (Ep + Ek) - E0) << std::endl;
     tcheckless(i, std::abs(0.5 * (Ep + Ek) - E0), 1.E-6);
 
     std::vector<P> mom0 = disc.get_moment(m0);
diff --git a/src/pde/two_stream_md.cpp b/src/pde/two_stream_md.cpp
index 6f8183876..4d6b6b791 100644
--- a/src/pde/two_stream_md.cpp
+++ b/src/pde/two_stream_md.cpp
@@ -4,39 +4,38 @@
 
 /*!
  * \internal
- * \file two_stream.cpp
- * \brief Two stream instability example
+ * \file two_stream_md.cpp
+ * \brief Two stream instability in 2X2V or 3X3V example
  * \author The ASGarD Team
- * \ingroup asgard_examples_two_stream
+ * \ingroup asgard_examples_two_stream_md
  *
  * \endinternal
  */
 
 /*!
  * \ingroup asgard_examples
- * \addtogroup asgard_examples_two_stream Example: Two stream instability
+ * \addtogroup asgard_examples_two_stream_md Example: Two stream instability in 2X2V or 3X3V
  *
- * \par Two stream instability
+ * \par Two stream instability (Multi-D)
  * Solves the Vlasov-Poisson equation in a common example
  * often called the two stream instability problem
- * \f[ \frac{\partial}{\partial t} f(x, v) + v \cdot \nabla_x f(x, v, t) + E(x, t) \nabla_v \cdot f(x, v, t) = 0 \f]
+ * \f[ \frac{\partial}{\partial t} f(x, v,t ) + v \cdot \nabla_x f(x, v, t) + E(x, t) \cdot \nabla_v f(x, v, t) = 0 \f]
  * where the electric field term depends on the Poisson equation
- * \f[ E(x,t) = -\nabla_x \Phi(x, t), \qquad - \nabla_x \cdot \nabla_x \Phi(x, t) = \int_v f(x, v, t) dv \f]
+ * \f[ E(x, t) = -\nabla_x \Phi(x, t), \qquad - \nabla_x \cdot \nabla_x \Phi(x, t) = \int_v f(x, v, t) dv \f]
  * The equation represents the evolution of a charged particle field under the effects
  * of self-induced electric field.
  * The right-hand integral represents the density of the particles and creates
  * non-linear coupling between the fields.
  *
  * \par
- * The focus of this example is the coupling with the electric field and Poisson
- * solver.
+ * The focus of this example is to show how to set up a Poisson
+ * solver for multiple dimensions. Internally the Poisson solver in multiple
+ * dimensions (2X or 3X) works differently from the Poisson solver in the 1X1V case.
  *
- * \par
- * <i>This is still work-in-progress, the documentation needs more work.</i>
  */
 
 /*!
- * \ingroup asgard_examples_two_stream
+ * \ingroup asgard_examples_two_stream_md
  * \brief The ratio of circumference to diameter of a circle
  */
 double constexpr PI = asgard::PI;
@@ -47,16 +46,33 @@ void self_test();
 #endif
 
 #ifdef ASGARD_USE_GPU
+/*!
+ * \ingroup asgard_examples_two_stream_md
+ * \brief The GPU kernel for taking the positive component of the elctric field
+ * 
+ * This kernel computes \f[ E(x, t) \cdot \nabla_v f(x, v, t) \f] for the x
+ * locations where E(x, t) is positive.
+ *
+ * \snippet two_stream_md.cpp two_stream_md make
+ */
 template<typename P>
 __global__ void interp_positive_kernel(int64_t num, P const* field, P const* mom, P* out) {
   int64_t i = blockIdx.x * blockDim.x + threadIdx.x;
   while (i < num) {
-    // Avoiding std::max to prevent __device__ compilation header conflicts
     out[i] = field[i] * ((mom[i] > 0.0) ? mom[i] : 0.0); 
     i += blockDim.x * gridDim.x;
   }
 }
 
+/*!
+ * \ingroup asgard_examples_two_stream_md
+ * \brief The GPU kernel for taking the negative component of the elctric field
+ * 
+ * This kernel computes \f[ E(x, t) \cdot \nabla_v f(x, v, t) \f] for the x
+ * locations where E(x, t) is negative.
+ *
+ * \snippet two_stream_md.cpp two_stream_md make
+ */
 template<typename P>
 __global__ void interp_negative_kernel(int64_t num, P const* field, P const* mom, P* out) {
   int64_t i = blockIdx.x * blockDim.x + threadIdx.x;
@@ -66,6 +82,14 @@ __global__ void interp_negative_kernel(int64_t num, P const* field, P const* mom
   }
 }
 
+/*!
+ * \ingroup asgard_examples_two_stream_md
+ * \brief The GPU kernel for computing the weight used for adapting the grid
+ * 
+ * This kernel computes \f[ ||E(x, t)||^2 f(x, v, t) \f] where \f[ x \in R^2 \f]
+ *
+ * \snippet two_stream_md.cpp two_stream_md make
+ */
 template<typename P>
 __global__ void weight_kernel_2d(int64_t num, P const *field, P const *ex, P const *ey, P* out) {
   int64_t i = blockIdx.x * blockDim.x + threadIdx.x;
@@ -75,6 +99,14 @@ __global__ void weight_kernel_2d(int64_t num, P const *field, P const *ex, P con
   }
 }
 
+/*!
+ * \ingroup asgard_examples_two_stream_md
+ * \brief The GPU kernel for computing the weight used for adapting the grid
+ * 
+ * This kernel computes \f[ ||E(x, t)||^2 f(x, v, t) \f] where \f[ x \in R^3 \f]
+ *
+ * \snippet two_stream_md.cpp two_stream_md make
+ */
 template<typename P>
 __global__ void weight_kernel_3d(int64_t num, P const *field, P const *ex, P const *ey, P const *ez, P* out) {
   int64_t i = blockIdx.x * blockDim.x + threadIdx.x;
@@ -86,25 +118,27 @@ __global__ void weight_kernel_3d(int64_t num, P const *field, P const *ex, P con
 #endif
 
 /*!
- * \ingroup asgard_examples_two_stream
+ * \ingroup asgard_examples_two_stream_md
  * \brief Make single two-stream PDE
  *
- * Constructs the pde description for the given umber of dimensions
+ * Constructs the pde description for the given number of dimensions
  * and options.
  *
  * \tparam P is either double or float, the asgard::default_precision will select
  *           first double, if unavailable, will go for float
  *
+ * \param pos_dims is the number of position dimensions (can be 2 or 3)
+ * 
  * \param options is the set of options
  *
  * \returns the asgard::pde_scheme definition
  *
- * \snippet two_stream.cpp two_stream make
+ * \snippet two_stream_md.cpp two_stream_md make
  */
 template<typename P = asgard::default_precision>
 asgard::pde_scheme<P> make_two_stream(int const pos_dims, asgard::prog_opts options) {
 #ifndef __ASGARD_DOXYGEN_SKIP
-//! [two_stream make]
+//! [two_stream_md make]
 #endif
 
   options.title = std::to_string(pos_dims) + "X" + std::to_string(pos_dims) + "V Two Stream Instability";
@@ -145,12 +179,11 @@ asgard::pde_scheme<P> make_two_stream(int const pos_dims, asgard::prog_opts opti
   asgard::pde_scheme<P> pde(options, domain);
 
   // set up moments
-  asgard::moment_id melectric_x = pde.register_electric_moment(asgard::dimension_id(0), pos_dims);
-  asgard::moment_id melectric_y = pde.register_electric_moment(asgard::dimension_id(1), pos_dims);
+  asgard::moment_id melectric_x = pde.register_electric_moment(asgard::dimension_id(0));
+  asgard::moment_id melectric_y = pde.register_electric_moment(asgard::dimension_id(1));
   asgard::moment_id melectric_z;
   if (pos_dims == 3)
-    melectric_z = pde.register_electric_moment(asgard::dimension_id(2), pos_dims);
-
+    melectric_z = pde.register_electric_moment(asgard::dimension_id(2));
   std::vector<asgard::moment_id> mids{melectric_x, melectric_y, melectric_z};
   std::vector<asgard::term_1d<P>> vterms(2 * pos_dims, asgard::term_identity{});
 
@@ -172,7 +205,7 @@ asgard::pde_scheme<P> make_two_stream(int const pos_dims, asgard::prog_opts opti
         y[i] = std::min(P{0}, x[i]);
     };
 
-  // v * df/dx
+  // v * grad_x(f)
   for (int d : asgard::iindexof(pos_dims)) {
     vterms[d] = asgard::term_div<P>(1, asgard::flux_type::upwind, asgard::boundary_type::periodic);
     vterms[pos_dims + d] = asgard::term_volume<P>(positive);
@@ -184,6 +217,8 @@ asgard::pde_scheme<P> make_two_stream(int const pos_dims, asgard::prog_opts opti
     vterms[pos_dims + d] = asgard::term_identity{};
   }
 
+  // set up the function used for determining when to adapt the grid
+  // computes ||E(x, t)||^2 * f(x, v, t)
 #ifdef ASGARD_USE_GPU
   std::function<void(int64_t, P, P const[], asgard::momentset_gpu<P> const&, P const[], P[])> weight;
   if (pos_dims == 2) {
@@ -237,6 +272,7 @@ asgard::pde_scheme<P> make_two_stream(int const pos_dims, asgard::prog_opts opti
   }
 #endif
 
+  // E * grad_v(f)
   for (int d : asgard::iindexof(pos_dims)) {
     asgard::moment_id mid = mids[d];
 
@@ -300,63 +336,72 @@ asgard::pde_scheme<P> make_two_stream(int const pos_dims, asgard::prog_opts opti
 #endif
   }
 
-  // initial conditions in x and v
+  // initial conditions
+  // the initial condition which creates a uniform density with a cosine perturbation
   auto ic_perturbed_pos = [](std::vector<P> const &x, P /* time */, std::vector<P> &fx) ->
     void {
       for (size_t i = 0; i < x.size(); i++)
         fx[i] = 1.0 - 0.5 * std::cos(0.5 * x[i]);
     };
-
-  auto ic_uniform_pos = [](std::vector<P> const &y, P /* time */, std::vector<P> &fy) ->
+  // the initial condition which creates a uniform density
+  auto ic_uniform_pos = [](std::vector<P> const &x, P /* time */, std::vector<P> &fx) ->
     void {
-      for (size_t i = 0; i < y.size(); i++)
-        fy[i] = 1.0;
+      for (size_t i = 0; i < x.size(); i++)
+        fx[i] = 1.0;
     };
-
-  auto ic_twostream_vel = [](std::vector<P> const &vx, P /* time */, std::vector<P> &fv) ->
+  // the initial condition which creates counterstreaming beams of particles
+  auto ic_twostream_vel = [](std::vector<P> const &v, P /* time */, std::vector<P> &fv) ->
     void {
       P const c = P{2} / std::sqrt(PI);
 
-      for (size_t i = 0; i < vx.size(); i++)
-        fv[i] = c * vx[i] * vx[i] * std::exp(-vx[i] * vx[i]);
+      for (size_t i = 0; i < v.size(); i++)
+        fv[i] = c * v[i] * v[i] * std::exp(-v[i] * v[i]);
     };
-
-  auto ic_maxwellian_vel = [](std::vector<P> const &vy, P /* time */, std::vector<P> &fv) ->
+  // the Maxwellian initial condition
+  auto ic_maxwellian_vel = [](std::vector<P> const &v, P /* time */, std::vector<P> &fv) ->
     void {
       P const c = P{1} / std::sqrt(PI);
 
-      for (size_t i = 0; i < vy.size(); i++)
-        fv[i] = c * std::exp(-vy[i] * vy[i]);
+      for (size_t i = 0; i < v.size(); i++)
+        fv[i] = c * std::exp(-v[i] * v[i]);
     };
 
   if (pos_dims == 2) {
+    // the perturbation is applied in x with a uniform distribution in y
+    // there are two counter streaming beams for vx and a maxwellian is used for vy
     pde.add_initial(asgard::separable_func<P>({ic_perturbed_pos, ic_uniform_pos, ic_twostream_vel, ic_maxwellian_vel}));
+
+    // use the adaptive weight function defined earlier
     pde.set_adapt_weight(weight, {melectric_x, melectric_y});
-  } else {
+  } else { // pos_dims = 3
+    // the perturbation is applied in x with a uniform distribution in y and z
+    // there are two counter streaming beams for vx and a maxwellian is used for vy and vz
     pde.add_initial(asgard::separable_func<P>({ic_perturbed_pos, ic_uniform_pos, ic_uniform_pos,
                                                ic_twostream_vel, ic_maxwellian_vel, ic_maxwellian_vel}));
+
+    // use the adaptive weight function defined earlier
     pde.set_adapt_weight(weight, {melectric_x, melectric_y, melectric_z});
   }
 
   return pde;
 
 #ifndef __ASGARD_DOXYGEN_SKIP
-//! [two_stream make]
+//! [two_stream_md make]
 #endif
 }
 
 /*!
- * \ingroup asgard_examples_two_stream
- * \brief main() for the diffusion example
+ * \ingroup asgard_examples_two_stream_md
+ * \brief main() for the two-stream example
  *
  * The main() processes the command line arguments and calls make_two_stream().
  *
- * \snippet two_stream.cpp two_stream main
+ * \snippet two_stream_md.cpp two_stream_md main
  */
 int main(int argc, char** argv)
 {
 #ifndef __ASGARD_DOXYGEN_SKIP
-//! [two_stream main]
+//! [two_stream_md main]
 #endif
 
   // if MPI is enabled, call MPI_Init(), otherwise do nothing
@@ -430,7 +475,7 @@ R"help(<< additional options for this file >>
   return 0;
 
 #ifndef __ASGARD_DOXYGEN_SKIP
-//! [two_stream main]
+//! [two_stream_md main]
 #endif
 };
 
@@ -455,8 +500,8 @@ void test_energy(std::string const &opt_str) {
   // the pde needs only the zeroth moment and computes that internally
   // we are using the other moments to check energy conservation properties
   auto pde = make_two_stream(2, options);
-  moment_id const melectric_x = pde.register_electric_moment(asgard::dimension_id(0), 2);
-  moment_id const melectric_y = pde.register_electric_moment(asgard::dimension_id(1), 2);
+  moment_id const melectric_x = pde.register_electric_moment(asgard::dimension_id(0));
+  moment_id const melectric_y = pde.register_electric_moment(asgard::dimension_id(1));
 
   // needed for verification but not for running
   moment_id const rho = pde.register_moment({0, 0});
@@ -480,17 +525,26 @@ void test_energy(std::string const &opt_str) {
     if (not disc.has_poisson()) // in MPI context, do error checking only on Poisson-ranks
       continue;
 
+    // the area of a cell is needed to rescale the kinetic energy moments
     P area = disc.domain().length(0) * disc.domain().length(1);
 
+    // get the electric field vectors in x and y
     auto efieldx = disc.get_moment(melectric_x);
     auto efieldy = disc.get_moment(melectric_y);
 
+    // compute the electric field energy
     P Ep = 0;
     for (auto ex : efieldx) Ep += ex * ex;
     for (auto ey : efieldy) Ep += ey * ey;
 
+    // get the kinetic energy moments
     std::vector<P> momke0 = disc.get_moment(ke0);
     std::vector<P> momke1 = disc.get_moment(ke1);
+
+    // the total kinetic energy is just the first coefficient
+    // because this is the average accross the domain.
+    // we need to rescale by the area because the wavelet basis
+    // is scaled from (-1, 1)
     P Ek = (momke0[0] + momke1[0]) * std::sqrt(area);
 
     if (disc.current_step() == 1) // first time-step
@@ -499,15 +553,19 @@ void test_energy(std::string const &opt_str) {
     // std::cout << "Total energy error: " << std::abs(0.5 * (Ep + Ek) - E0) << "\n";
     tcheckless(i, std::abs(0.5 * (Ep + Ek) - E0), 7.E-6);
 
+    // get the density and velocity moments
     std::vector<P> mom0 = disc.get_moment(rho);
     std::vector<P> momp0 = disc.get_moment(p0);
     std::vector<P> momp1 = disc.get_moment(p1);
 
-    // integral of moment 0 by moment 1, by delta_ij orthogonality of the basis
+    // integral of moment (0, 0) by moment (1, 0), by delta_ij orthogonality of the basis
     // just sum up the product of the coefficients
     P mv0 = 0;
     for (size_t j = 0; j < mom0.size(); j++)
       mv0 += mom0[j] * momp0[j];
+
+    // integral of moment (0, 0) by moment (0, 1), by delta_ij orthogonality of the basis
+    // just sum up the product of the coefficients
     P mv1 = 0;
     for (size_t j = 0; j < mom0.size(); j++)
       mv1 += mom0[j] * momp1[j];
diff --git a/src/pde/vplb.cpp b/src/pde/vplb.cpp
index c133a1190..deeaba79f 100644
--- a/src/pde/vplb.cpp
+++ b/src/pde/vplb.cpp
@@ -390,7 +390,7 @@ void test_energy(int const vdims, std::string const &opt_str) {
   // we are using the other moments to check energy conservation properties
   auto pde = make_vplb<P>(vdims, options);
   moment_id const m2 = pde.register_moment({2});
-  moment_id const melectric_x = pde.register_electric_moment(asgard::dimension_id(0), 1);
+  moment_id const melectric_x = pde.register_electric_moment(asgard::dimension_id(0));
   discretization_manager disc(std::move(pde), verbosity_level::quiet);
 
   double E0 = 0; // initial total energy (potential + kinetic), will initialize on first iteration