Merge pull request #266 from fnalacceleratormodeling/265-egsternneed-…

…accelerating-phase-separate-from-reference-energy

header for new method to adjust the particle coordinates to match a n…

src/synergia/bunch/bunch.h

@@ -17,6 +17,8 @@
 #include "synergia/utils/hdf5_file.h"
 #include "synergia/utils/logger.h"
 
+#include "synergia/libFF/ff_adjust_bunch_ref_coords.h"
+
 #include <cereal/types/array.hpp>
 #include <cereal/types/map.hpp>
 
@@ -485,6 +487,14 @@ class bunch_t {
  return std::make_pair(boundary, boundary_param);
  }
 
+ // adjust_bunch_particle_reference_energy
+ // This method adjusts the particle coordinates to use
+ void
+ adjust_bunch_particles_reference_energy(double newE)
+ {
+ FF_adjust_bunch_ref_coords::apply(newE, *this);
+ }
+
  // bucket index
  void
  set_bucket_index(int index)

src/synergia/bunch/bunch_pywrap.cc

@@ -205,6 +205,11 @@ PYBIND11_MODULE(bunch, m)
  &Bunch::get_longitudinal_boundary,
  "Get the longitudinal boundary of the bunch")
 
+ .def("adjust_bunch_particles_reference_energy",
+ &Bunch::adjust_bunch_particles_reference_energy,
+ "Set the reference energy of the bunch and adjust the normalization of the particles to match",
+ "newE"_a)
+
  .def(
  "add_diagnostics",
  [](Bunch& self, std::shared_ptr<Diagnostics> const& diag) {

src/synergia/bunch/tests/CMakeLists.txt

@@ -14,4 +14,5 @@ add_mpi_test(test_bunch_particles 1)
 if(BUILD_PYTHON_BINDINGS)
  add_py_test(test_bunch.py)
  add_py_test(test_bunch_reference_particles.py)
-endif()
+ add_py_test(test_adjust_particles_ref_energy.py)
+ endif()

src/synergia/bunch/tests/test_adjust_particles_ref_energy.py

@@ -0,0 +1,69 @@
+#!/usr/bin/env python
+import sys, os
+import numpy as np
+import synergia
+import pytest
+
+macroparticles=16
+realparticles=4.0e10
+
+KE0 = 0.8 # starting kinetic energy
+mp = synergia.foundation.pconstants.mp
+
+# prop_fixture is a Bunch
+@pytest.fixture
+def bunch_fixture():
+ ref_part = synergia.foundation.Reference_particle(1, mp, mp+KE0)
+ sim = synergia.simulation.Bunch_simulator.create_single_bunch_simulator(ref_part, macroparticles, realparticles)
+ bunch = sim.get_bunch()
+ bunch.checkout_particles()
+ lp = bunch.get_particles_numpy()
+ lp[:, 0:6] = 0.0
+ np = lp.shape[0]
+ for i in range(np):
+ lp[i, 1] = i * 1.0e-4
+ lp[i, 3] = (np-i) * 1.0e-4
+ bunch.checkin_particles()
+ return bunch
+
+
+def test_adjust_ref_energy1(bunch_fixture):
+
+ refpart = bunch_fixture.get_design_reference_particle()
+ energy = refpart.get_total_energy()
+ old_p = refpart.get_momentum()
+
+ # add some energy to each particle by changing dp/p.
+ dE = 0.050 # 50 MeV
+ # what is that in dp/p?
+ new_energy = energy + dE
+ print('new energy: ', new_energy)
+ new_p = np.sqrt(new_energy**2 - mp**2)
+ dpop = new_p/old_p - 1.0
+ print('dp/p corresponding to new energy: ', dpop)
+
+ # set the dp/p of all the particles
+ bunch_fixture.checkout_particles()
+ lp = bunch_fixture.get_particles_numpy()
+ lp[:, 5] = dpop
+ bunch_fixture.checkin_particles()
+
+ # adjust the reference energy
+ bunch_fixture.adjust_bunch_particles_reference_energy(new_energy)
+
+ # check that dp/p is set correctly
+ bunch_fixture.checkout_particles()
+ lp = bunch_fixture.get_particles_numpy()
+ print('particles dp/p: ', lp[:, 5])
+ for i in range(lp.shape[0]):
+ assert lp[i, 5] == pytest.approx(0.0, abs=1.0e-12)
+ for i in range(lp.shape[0]):
+ assert lp[i, 1] == pytest.approx(i * 1.0e-4 * old_p/new_p)
+ assert lp[i, 3] == pytest.approx((lp.shape[0] - i) * 1.0e-4 * old_p/new_p)
+
+def main():
+ bf = bunch_fixture()
+ test_adjust_ref_energy1(bf)
+
+if __name__ == "__main__":
+ main()

src/synergia/libFF/ff_adjust_bunch_ref_coords.h

@@ -0,0 +1,153 @@
+#ifndef FF_REF_COORDS_H
+#define FF_REF_COORDS_H
+
+#include "synergia/foundation/physical_constants.h"
+#include "synergia/libFF/ff_algorithm.h"
+#include "synergia/libFF/ff_patterned_propagator.h"
+#include "synergia/utils/simple_timer.h"
+
+namespace ref_energy_impl {
+ struct CoordParams {
+ double mass; // mass
+ double old_E; // old energy
+ double new_E; // new energy
+ double old_p; // old momentum
+ double new_p; // new momentum
+ };
+
+ template <class BunchT>
+ struct AdjustCoords {
+ using gsv_t = typename BunchT::gsv_t;
+
+ typename BunchT::bp_t::parts_t p;
+ typename BunchT::bp_t::const_masks_t m;
+ const CoordParams cp;
+
+ KOKKOS_INLINE_FUNCTION
+ void
+ operator()(const int idx) const
+ {
+ int i = idx * gsv_t::size();
+ int mask = 0;
+ for (int x = i; x < i + gsv_t::size(); ++x)
+ mask |= m(x);
+
+ if (mask) {
+ gsv_t p1(&p(i, 1));
+ gsv_t p3(&p(i, 3));
+ gsv_t p4(&p(i, 4));
+ gsv_t p5(&p(i, 5));
+
+ FF_algorithm::adjust_ref_unit(p1,
+ p3,
+ p4,
+ p5,
+ cp.mass,
+ cp.old_p,
+ cp.new_p,
+ cp.old_E,
+ cp.new_E);
+
+ p1.store(&p(i, 1));
+ p3.store(&p(i, 3));
+ p4.store(&p(i, 4));
+ p5.store(&p(i, 5));
+ }
+ }
+ };
+
+ template <class BunchT>
+ struct PropAdjust {
+ using gsv_t = typename BunchT::gsv_t;
+
+ typename BunchT::bp_t::parts_t p;
+ typename BunchT::bp_t::const_masks_t m;
+ const CoordParams cp;
+
+ KOKKOS_INLINE_FUNCTION
+ void
+ operator()(const int idx) const
+ {
+ int i = idx * gsv_t::size();
+ int mask = 0;
+ for (int x = i; x < i + gsv_t::size(); ++x)
+ mask |= m(x);
+
+ if (mask) {
+ gsv_t p0(&p(i, 0));
+ gsv_t p1(&p(i, 1));
+ gsv_t p2(&p(i, 2));
+ gsv_t p3(&p(i, 3));
+ gsv_t p4(&p(i, 4));
+ gsv_t p5(&p(i, 5));
+
+ FF_algorithm::adjust_ref_unit(p1,
+ p3,
+ p4,
+ p5,
+ cp.mass,
+ cp.old_p,
+ cp.new_p,
+ cp.old_E,
+ cp.new_E);
+
+ p0.store(&p(i, 0));
+ p1.store(&p(i, 1));
+ p2.store(&p(i, 2));
+ p3.store(&p(i, 3));
+ p4.store(&p(i, 4));
+ p5.store(&p(i, 5));
+ }
+ }
+ };
+
+}
+
+namespace FF_adjust_bunch_ref_coords {
+ template <class BunchT>
+ void
+ apply(double new_E, BunchT& bunch)
+ {
+ using namespace ref_energy_impl;
+
+ scoped_simple_timer timer("libFF_ref_energy");
+
+ CoordParams cp;
+
+ // actually we don't care about the slice at all. Maybe
+ // we can eliminate the slice argument
+
+ Reference_particle& ref_l = bunch.get_design_reference_particle();
+ Reference_particle& ref_b = bunch.get_reference_particle();
+
+ // The bunch particles momentum is with respect to the bunch reference
+ // particle
+ cp.old_p = ref_l.get_momentum();
+ cp.old_E = ref_l.get_total_energy();
+ cp.mass = bunch.get_mass();
+
+ cp.new_E = new_E;
+ cp.new_p = sqrt(new_E*new_E - cp.mass*cp.mass);
+
+ // bunch particles
+ auto apply = [&](ParticleGroup pg) {
+ if (!bunch.get_local_num(pg)) return;
+
+ auto parts = bunch.get_local_particles(pg);
+ auto masks = bunch.get_local_particle_masks(pg);
+
+ using exec = typename BunchT::exec_space;
+ auto range = Kokkos::RangePolicy<exec>(0, bunch.size_in_gsv(pg));
+
+ AdjustCoords<BunchT> adjuster{parts, masks, cp};
+ Kokkos::parallel_for(range, adjuster);
+ };
+ apply(ParticleGroup::regular);
+ apply(ParticleGroup::spectator);
+
+ // now set the bunch energy to match
+ ref_b.set_total_energy(cp.new_E);
+ }
+}
+
+#endif // FF_REF_COORDS_H

src/synergia/libFF/ff_algorithm.h

@@ -1070,6 +1070,31 @@ namespace FF_algorithm {
  dpop = sqrt((E - T(m)) * (E + T(m))) / T(new_ref_p) - T(1.0);
  }
 
+ // adjust particle coordinates to use new reference energy
+ template <typename T>
+ KOKKOS_INLINE_FUNCTION void
+ adjust_ref_unit(T& px,
+ T& py,
+ T const& cdt,
+ T& dpop,
+ double m,
+ double old_pref,
+ double new_pref,
+ double old_E,
+ double new_E)
+ {
+
+ T p = T(old_pref) * (dpop + T(1.0));
+ T E = sqrt(p * p + T(m * m));
+
+
+ px = px * T(old_pref / new_pref);
+ py = py * T(old_pref / new_pref);
+
+ dpop = sqrt((E - T(m)) * (E + T(m))) / T(new_pref) - T(1.0);
+
+ }
+
  KOKKOS_INLINE_FUNCTION
  double
  factorial(int n)