diff --git a/t/hej_test.hh b/t/hej_test.hh
index 74f00b0..a5e9d5a 100644
--- a/t/hej_test.hh
+++ b/t/hej_test.hh
@@ -1,97 +1,105 @@
 /**
  *  \authors   The HEJ collaboration (see AUTHORS for details)
  *  \date      2019-2020
  *  \copyright GPLv2 or later
  */
 #pragma once
 
 #include <array>
 #include <string>
 #include <vector>
 
 #include "HEJ/Event.hh"
 #include "HEJ/exceptions.hh"
 
 namespace HEJ {
   struct Particle;
 }
 
 //! throw error if condition not fulfilled
 #define ASSERT(x) if(!(x)) { \
     throw std::logic_error("Assertion '" #x "' failed."); \
   }
 
 //! throw error if prop is different between ev1 and ev2
 #define ASSERT_PROPERTY(ev1,ev2,prop) ASSERT((ev1).prop == (ev2).prop)
 
 //! throw error if condition not fulfilled
 #define ASSERT_THROW(x, exception) try { \
       x; \
       std::cerr << "'" #x "' did not throw an exception.\n"; \
       throw; \
     } catch(exception const &){} \
     catch (...) { \
       std::cerr << "Unexpected exception thrown for '" #x "'.\n"; \
       throw; \
     }
 
 /** @brief get specific Phase Space Points for njets with boson at pos_boson
  *
  * if pos_boson = -1 (or not implemented) -> no boson
  *
  * njet==7 is special: has less jets, i.e. multiple parton in one jet,
  *                     all partons are massive (4 GeV) -> can be boson/decay
  *                     pos_boson < 0 to select process (see list for details)
  */
 HEJ::Event::EventData get_process(int njet, int pos_boson);
 
 //! select process from string input (see also get_process)
 //!
 //! overwrite_boson to force a specific boson position, indepentent from input
 //! (useful for njet == 7)
 HEJ::Event::EventData parse_configuration(
     std::array<std::string,2> const & in, std::vector<std::string> const & out,
     int overwrite_boson = 0
 );
 
+/** repair moment in event:
+ *  - adjust energies so that the ougoing particles are on-shell
+ *  - update incoming four-momenta to ensure momentum conservation
+ *
+ *  *does not adjust decays*
+ */
+void repair_momenta(HEJ::Event::EventData & ev);
+
 /**
  * @brief Generate a rapidity ordered event (with decays)
  * @param in             Incoming particle flavours
  * @param out            Outgoing particle flavours
  * @param reconstruct    Should call reconstruct_intermediate?
  * @param decays         Any decays associated with particles from out
  * @param shuffle        Perform shuffle on outgoing vector?
  * @returns              EventData for in -> out with rapidity ordered final state
  *
  * Generates EventData for the process in -> out. The outgoing particles
  * are generated in rapidity order. The incoming flavours are ordered:
  * {backward, forward}.
  *
  * Decays are provided by decays, if specificed, and via reconstruct_intermediate
  * if requested.
  *
  * Note: Specified decays are handled before reconstruct_intermediate is called.
  */
 
 HEJ::Event::EventData rapidity_order_ps(
     std::array<std::string,2> const & in,
     std::vector<std::string> const & out,
     bool reconstruct = false,
     std::unordered_map<size_t, std::vector<std::string> > decays = {},
     bool shuffle = true
   );
 
 //! shuffle particles around
 void shuffle_particles(HEJ::Event::EventData & ev);
 
 //! Helper function to couple quarksfor flavour-changing bosons
 bool couple_quark(std::string const & boson, std::string & quark);
 
 //! Decay kinematics for 1->2
 std::vector<HEJ::Particle> decay_kinematics( HEJ::Particle const & parent );
 
 //! Decay W boson to lepton & neutrino
 std::vector<HEJ::Particle> decay_W( HEJ::Particle const & parent );
 
 //! Decay Z to electron-positron
 std::vector<HEJ::Particle> decay_Z( HEJ::Particle const & parent );
diff --git a/t/test_decay.cc b/t/test_decay.cc
index eac5021..0fbe3f6 100644
--- a/t/test_decay.cc
+++ b/t/test_decay.cc
@@ -1,441 +1,444 @@
 /**
  *  \authors   The HEJ collaboration (see AUTHORS for details)
  *  \date      2019-2020
  *  \copyright GPLv2 or later
  *
  *  \brief Test classification for (invalid) W decays
  */
 #include "hej_test.hh"
 
 #include <array>
 #include <cstdlib>
 #include <initializer_list>
 #include <iostream>
 #include <memory>
 #include <string>
 #include <unordered_map>
 #include <utility>
 
 #include "fastjet/JetDefinition.hh"
 #include "fastjet/PseudoJet.hh"
 
 #include "HEJ/Event.hh"
 #include "HEJ/event_types.hh"
 #include "HEJ/Particle.hh"
 #include "HEJ/PDG_codes.hh"
 
 namespace {
   const fastjet::JetDefinition JET_DEF{fastjet::JetAlgorithm::antikt_algorithm, 0.4};
   const double MIN_JET_PT{30.};
 
   HEJ::Event::EventData new_event() {
     HEJ::Event::EventData ev;
     ev.outgoing.push_back({HEJ::ParticleID::gluon, { -11,  -96,  -76,  123}, {}});
     ev.outgoing.push_back({HEJ::ParticleID::gluon, { -15,  -70,  -22,   75}, {}});
     ev.outgoing.push_back({HEJ::ParticleID::gluon, {  68,   93,  -20,  117}, {}});
     ev.outgoing.push_back({HEJ::ParticleID::gluon, { -12,   95,   56,  111}, {}});
     ev.outgoing.push_back({HEJ::ParticleID::gluon, { -30,  -22,   25,   45}, {}});
     ev.incoming[0] = {HEJ::ParticleID::gluon, {   0,    0, -254,  254}, {}};
     ev.incoming[1] = {HEJ::ParticleID::gluon, {   0,    0,  217,  217}, {}};
     return ev;
   }
 
   bool test_event(HEJ::Event::EventData data, bool const valid
   ){
     using namespace HEJ::event_type;
     EventType const expected{ valid?FKL:bad_final_state };
     shuffle_particles(data);
     auto const ev = std::move(data).cluster(JET_DEF, MIN_JET_PT);
     if(ev.type() != expected){
       std::cerr << "Event does not match expectation, expected "
         << name(expected) << "\n" << ev << std::endl;
       return false;
     }
     return true;
   }
 
   // Check basic FKL event
   bool check_fkl() {
     auto ev = new_event();
+    repair_momenta(ev);
     return test_event(ev, true);
   }
 
   // Check W decays
   bool check_W_decay() {
     using namespace HEJ::pid;
 
     auto ev = new_event();
     // W position shouldn't matter
     for(auto const W_type: {Wp, Wm}){
       for(std::size_t w_pos = 1; w_pos<ev.outgoing.size()-1; ++w_pos){
         ev = new_event();
         ev.outgoing[w_pos].type = W_type;
         ev.outgoing.back().type = (W_type==Wp)?d:u;
         ev.incoming.back().type = (W_type==Wp)?u:d;
+        repair_momenta(ev);
 
         // no decay
         if(!test_event(ev, false))
           return false;
 
         // working decay (Wp -> nu_e,e+ OR Wp -> e-,nu_e_bar)
         ev.decays[w_pos] = decay_W(ev.outgoing[w_pos]);
         if(!test_event(ev, true))
           return false;
 
         // swapped W+ <-> W-
         ev.decays[w_pos].at(0).type = static_cast<ParticleID>(
                                         -ev.decays[w_pos].at(0).type );
         ev.decays[w_pos].at(1).type = static_cast<ParticleID>(
                                         -ev.decays[w_pos].at(1).type );
         if(!test_event(ev, false))
           return false;
         ev.decays[w_pos].at(0).type = static_cast<ParticleID>(
                                         -ev.decays[w_pos].at(0).type );
         ev.decays[w_pos].at(1).type = static_cast<ParticleID>(
                                         -ev.decays[w_pos].at(1).type );
 
         // replace e -> mu (normal)
         ev.decays[w_pos].at(0).type = static_cast<ParticleID>(
                                         ev.decays[w_pos].at(0).type+2 );
         if(!test_event(ev, false))
           return false;
         ev.decays[w_pos].at(0).type = static_cast<ParticleID>(
                                         ev.decays[w_pos].at(0).type-2 );
 
         // replace e -> mu (anti)
         ev.decays[w_pos].at(1).type = static_cast<ParticleID>(
                                         ev.decays[w_pos].at(1).type-2 );
         if(!test_event(ev, false))
           return false;
 
         // all mu
         ev.decays[w_pos].at(0).type = static_cast<ParticleID>(
                                         ev.decays[w_pos].at(0).type+2 );
         if(!test_event(ev, true))
           return false;
         ev.decays[w_pos].at(0).type = static_cast<ParticleID>(
                                         ev.decays[w_pos].at(0).type-2 );
         ev.decays[w_pos].at(1).type = static_cast<ParticleID>(
                                         ev.decays[w_pos].at(1).type+2 );
 
         // partonic
         ev.decays[w_pos].at(0).type = static_cast<ParticleID>(
                                         ev.decays[w_pos].at(0).type-10 );
         ev.decays[w_pos].at(1).type = static_cast<ParticleID>(
                                         ev.decays[w_pos].at(1).type+10 );
         if(!test_event(ev, false))
           return false;
         ev.decays[w_pos].at(0).type = static_cast<ParticleID>(
                                         ev.decays[w_pos].at(0).type+10 );
         ev.decays[w_pos].at(1).type = static_cast<ParticleID>(
                                         ev.decays[w_pos].at(1).type-10 );
 
         // double check that we undid all changes
         if(!test_event(ev, true))
           return false;
 
         // 1->3 decay
         ev.decays[w_pos].emplace_back(
             HEJ::Particle{photon, fastjet::PseudoJet(0,0,0,0), {}}
         );
         if(!test_event(ev, false))
           return false;
         ev.decays[w_pos].pop_back();
 
         // invalid secondary decay
         ev.decays[0] = decay_W(ev.outgoing[0]);
         ev.decays[0].at(0).type = ev.outgoing[0].type;
         ev.decays[0].at(1).type = gluon;
         if(!test_event(ev, false))
           return false;
       }
     }
     return true;
   }
 
   // Check Z decays
   bool check_Z_decay() {
     using namespace HEJ::pid;
     auto ev = new_event();
     for(size_t z_pos = 1; z_pos<ev.outgoing.size()-1; ++z_pos){
       auto ev = new_event();
       ev.outgoing[z_pos].type = Z_photon_mix;
       ev.outgoing.back().type = u;
       ev.incoming.back().type = u;
+      repair_momenta(ev);
 
       // no decay
       if(!test_event(ev, false))
         return false;
 
       // working decay (Z -> e-,e+)
       ev.decays[z_pos] = decay_Z(ev.outgoing[z_pos]);
       if(!test_event(ev, true))
         return false;
 
       // replace e- -> mu-
       ev.decays[z_pos].at(0).type = static_cast<ParticleID>(
                                       ev.decays[z_pos].at(0).type+2 );
       if(!test_event(ev, false))
         return false;
       ev.decays[z_pos].at(0).type = static_cast<ParticleID>(
                                       ev.decays[z_pos].at(0).type-2 );
 
       // replace e+ -> mu+
       ev.decays[z_pos].at(1).type = static_cast<ParticleID>(
                                       ev.decays[z_pos].at(1).type-2 );
       if(!test_event(ev, false))
         return false;
 
       // all mu
       ev.decays[z_pos].at(0).type = static_cast<ParticleID>(
                                       ev.decays[z_pos].at(0).type+2 );
       if(!test_event(ev, true))
         return false;
       ev.decays[z_pos].at(0).type = static_cast<ParticleID>(
                                       ev.decays[z_pos].at(0).type-2 );
       ev.decays[z_pos].at(1).type = static_cast<ParticleID>(
                                       ev.decays[z_pos].at(1).type+2 );
 
       // replace e- -> nu_e
       ev.decays[z_pos].at(0).type = static_cast<ParticleID>(
                                       ev.decays[z_pos].at(0).type+1 );
       if(!test_event(ev, false))
         return false;
       ev.decays[z_pos].at(0).type = static_cast<ParticleID>(
                                       ev.decays[z_pos].at(0).type-1 );
 
       // replace e+ -> nu_e_bar
       ev.decays[z_pos].at(1).type = static_cast<ParticleID>(
                                       ev.decays[z_pos].at(1).type-1 );
       if(!test_event(ev, false))
         return false;
 
       // neutrino-antineutrino
       ev.decays[z_pos].at(0).type = static_cast<ParticleID>(
                                       ev.decays[z_pos].at(0).type+1 );
       if(!test_event(ev, false))
         return false;
       ev.decays[z_pos].at(0).type = static_cast<ParticleID>(
                                       ev.decays[z_pos].at(0).type-1 );
       ev.decays[z_pos].at(1).type = static_cast<ParticleID>(
                                       ev.decays[z_pos].at(1).type+1 );
 
       // partonic
       ev.decays[z_pos].at(0).type = static_cast<ParticleID>(
                                       ev.decays[z_pos].at(0).type-10 );
       ev.decays[z_pos].at(1).type = static_cast<ParticleID>(
                                       ev.decays[z_pos].at(1).type+10 );
       if(!test_event(ev, false))
         return false;
       ev.decays[z_pos].at(0).type = static_cast<ParticleID>(
                                       ev.decays[z_pos].at(0).type+10 );
       ev.decays[z_pos].at(1).type = static_cast<ParticleID>(
                                       ev.decays[z_pos].at(1).type-10 );
 
       // double check that we undid all changes
       if(!test_event(ev, true))
         return false;
 
       // 1->3 decay
       ev.decays[z_pos].emplace_back(
           HEJ::Particle{photon, fastjet::PseudoJet(0,0,0,0), {}}
       );
       if(!test_event(ev, false))
         return false;
       ev.decays[z_pos].pop_back();
 
       // invalid secondary decay
       ev.decays[0] = decay_Z(ev.outgoing[0]);
       ev.decays[0].at(0).type = ev.outgoing[0].type;
       ev.decays[0].at(1).type = gluon;
       if(!test_event(ev, false))
         return false;
     }
     return true;
   }
 
   // Check 2 bosons
   bool check_2_boson() {
     const std::vector<std::string> bosons {"Wp", "Wm", "Z_photon_mix", "h"};
 
     // Candidate decays
     std::unordered_map<std::string, std::vector<std::string> > test_decay;
     test_decay["Wp"] = {"nu_e","e+"};
     test_decay["Wm"] = {"e-", "nu_e_bar"};
     test_decay["Z_photon_mix"]  = {"e-", "e+"};
     test_decay["h"]  = {"gamma", "gamma"};
 
     size_t njet = 6;
     const std::vector<std::string> all_quarks{"-4","-1","1","2","3","4"};
     for(std::string const & q1 : all_quarks){
       for(std::string const & q2 : all_quarks){
         for(std::string const & b1 : bosons){
           for(std::string const & b2 : bosons){
             std::array<std::string,2> in {q1, q2};
             std::vector<std::string>  out(njet, "g");
             out.front() = q1;
             out.back()  = q2;
 
             if( !couple_quark(b1, out.front()) ||
                 !couple_quark(b2, out.back() )
               ) { continue; } // skip incompatible
 
             std::unordered_map< size_t, std::vector<std::string> > decays;
 
             // Boson 1
             out.emplace_back(b1);
             decays.emplace( out.size()-1, test_decay[b1] );
             // Boson 2
             out.emplace_back(b2);
             decays.emplace( out.size()-1, test_decay[b2] );
 
             auto ev = rapidity_order_ps(in, out, false, decays)
                       .cluster(JET_DEF, MIN_JET_PT);
 
             if(ev.type() == HEJ::event_type::bad_final_state) {
               std::cerr << "Event was expected to not be a bad final-state\n"
                         << ev << std::endl;
               return false;
             }
           }
         }
       }
     }
     return true;
   }
 
   // Check reconstruction
   bool check_reconstruction() {
 
     // reconstruct_intermediate doesn't support reconstruct Higgs!
     const std::vector<std::string> bosons {"Wp", "Wm", "Z_photon_mix"};
 
     // Candidate decays
     std::unordered_map<std::string, std::vector<std::string> > test_decay;
     test_decay["Wp"] = {"nu_e","e+"};
     test_decay["Wm"] = {"e-", "nu_e_bar"};
     test_decay["Z_photon_mix"]  = {"e-", "e+"};
 
     size_t njet = 2;
     const std::vector<std::string> all_quarks{"-4","-1","1","2","3","4"};
 
     // Single-boson Reconstruction
     for(std::string const & q1 : all_quarks){
       for(std::string const & q2 : all_quarks){
         for(std::string const & boson : bosons){
           std::array<std::string,2> in {q1, q2};
           std::vector<std::string>  out(njet, "g");
           out.front() = q1;
           out.back()  = q2;
 
           // skip incompatible
           if( !couple_quark(boson, out.front()) ) { continue; }
 
           // Add decay products to outgoing
           out.insert(out.end(), test_decay[boson].begin(), test_decay[boson].end());
 
           // Generate event, requesting reconstruction
           auto ev = rapidity_order_ps(in, out, true)
                     .cluster(JET_DEF, MIN_JET_PT);
 
           // Find and compare boson types
           HEJ::ParticleID boson_pid = HEJ::to_ParticleID(boson);
           size_t boson_idx = ev.decays().begin() -> first;
           auto res_boson = ev.outgoing().at(boson_idx);
 
           if(res_boson.type != boson_pid) {
             std::cerr << "Boson " <<  name(boson_pid)
                       << " reconstructed to incorrect type " << name(res_boson.type)
                       << "\n" << ev << std::endl;
             return false;
           }
         }
       }
     }
 
     // Two-boson Reconstruction
     /**
      * Note:
      * Event::reconstruct_intermediate() only supports W-boson reconstructions
      * at 2-boson.
      */
     const std::vector< std::vector<std::string> > two_bosons = {
       {"Wp", "Wm"}, {"Wp", "Wp"}, {"Wm", "Wm"}
     };
 
     for(auto const & bosons : two_bosons){
       for(std::string const & q1 : all_quarks){
         for(std::string const & q2 : all_quarks){
           std::string b1 = bosons[0];
           std::string b2 = bosons[1];
 
           std::array<std::string,2> in {q1, q2};
           std::vector<std::string>  out(njet, "g");
           out.front() = q1;
           out.back()  = q2;
 
           // skip incompatible
           if( !couple_quark(b1, out.front()) ||
               !couple_quark(b2, out.back() )
             ) { continue; } // skip incompatible
 
           // Add decay products to outgoing
           out.insert(out.end(), test_decay[b1].begin(), test_decay[b1].end());
           out.insert(out.end(), test_decay[b2].begin(), test_decay[b2].end());
 
           auto ev = rapidity_order_ps(in, out, true)
                     .cluster(JET_DEF, MIN_JET_PT);
 
           // Expecting 2 decays
           if(ev.decays().size() != bosons.size()) {
               std::cerr << "Expect 2 decays\n" << ev << std::endl;
               return false;
           }
 
           // Expected bosons
           std::vector<HEJ::ParticleID> expected_bosons = {
             HEJ::to_ParticleID(b1),
             HEJ::to_ParticleID(b2)
           };
           std::sort(begin(expected_bosons), end(expected_bosons));
 
           // Reconstructed bosons
           auto decays = ev.decays();
           std::vector<HEJ::ParticleID> res_bosons;
           std::transform(
             cbegin(decays), cend(decays), std::back_inserter(res_bosons),
             [&ev] (auto const & decay) -> HEJ::ParticleID {
               return ev.outgoing().at(decay.first).type;
             }
           );
           std::sort(begin(res_bosons), end(res_bosons));
 
           // Matching?
           if(expected_bosons != res_bosons) {
             std::cerr << "Reconstructed bosons did not match expectation ( ";
             std::copy(
               expected_bosons.begin(), expected_bosons.end(),
               std::ostream_iterator<int>(std::cerr, "  ")
             );
             std::cerr << ")\n" << ev << std::endl;
             return false;
           }
 
         }
       }
     }
 
     return true;
   }
 } // namespace anonymous
 
 
 int main() {
   if(!check_fkl()) return EXIT_FAILURE;
   if(!check_W_decay()) return EXIT_FAILURE;
   if(!check_Z_decay()) return EXIT_FAILURE;
   if(!check_2_boson()) return EXIT_FAILURE;
   if(!check_reconstruction()) return EXIT_FAILURE;
 
   return EXIT_SUCCESS;
 }