Merge branch 'master' of https://github.com/like2000/PyHEADTAIL

Author: like2000

_version.py

@@ -1,2 +1,2 @@
-__version__ = "1.0.3"
+__version__ = "1.0.4"
 

monitors/monitors.py

@@ -341,10 +341,10 @@ def __init__(self, filename, stride=1, parameters_dict=None,
 
         self._create_file_structure(parameters_dict)
 
-    def dump(self, bunch):
+    def dump(self, bunch, arrays_dict=None):
         """ Write particle data to file. See docstring of method
         self._write_data_to_file . """
-        self._write_data_to_file(bunch)
+        self._write_data_to_file(bunch, arrays_dict)
         self.i_steps += 1
 
     def _create_file_structure(self, parameters_dict):
@@ -362,11 +362,12 @@ def _create_file_structure(self, parameters_dict):
                         self.filename + 'failed. \n')
             raise
 
-    def _write_data_to_file(self, bunch):
+    def _write_data_to_file(self, bunch, arrays_dict):
         """ Write macroparticle data (x, xp, y, yp, z, dp, id) of a
-        selection of particles to the HDF5 file. The file is opened and
-        closed every time to prevent from loss of data in case of a
-        crash.
+        selection of particles to the HDF5 file. Optionally, data in
+        additional_quantities can also be added if provided in the
+        constructor. The file is opened and closed every time to prevent
+        from loss of data in case of a crash.
         For each simulation step, a new group with name 'Step#..' is
         created. It contains one dataset for each of the quantities
         given in self.quantities_to_store. """
@@ -375,9 +376,16 @@ def _write_data_to_file(self, bunch):
         dims = (bunch.macroparticlenumber // self.stride,)
 
         # resorting_indices = np.argsort(bunch.id)[::self.stride]
+        all_quantities = {}
         for quant in self.quantities_to_store:
             quant_values = getattr(bunch, quant)
+            all_quantities[quant] = quant_values
+        all_quantities.update(arrays_dict)
+
+        for quant in all_quantities.keys():
+            quant_values = all_quantities[quant]
             h5group.create_dataset(quant, shape=dims, compression='gzip',
                 compression_opts=9, dtype=quant_values.dtype)
             h5group[quant][:] = quant_values[::self.stride]
+
         h5file.close()

particles/generators.py

@@ -411,11 +411,17 @@ class CutRFBucket6D(ParticleGenerator):
     bi-gaussian with given sigma_z and sigma_dp which is then cut along
     the funciotn given by is_accepted which is typically the function
     is_in_separatrix of the rfbucket instance of the longitudinal map.
+    To avoid bucket leakage and particle losses as a consequence of the
+    unmatched initialisation, a margin / tolerance can be specified by
+    the argument margin (in % of RFBucket.Hmax, 5% by default), so that
+    particles will not be initialised too close to the separatrix. The
+    boundary will in fact be given by the equihamiltonian with a value
+    of margin*Hmax.
     '''
     def __init__(self, macroparticlenumber, intensity, charge, mass,
                  circumference, gamma_reference,
                  transverse_map, epsn_x, epsn_y,
-                 sigma_z, sigma_dp, is_accepted=None,
+                 sigma_z, sigma_dp, is_accepted, margin=0.05,
                  *args, **kwargs):
         '''Uses the transverse_map to extract the optics parameters
         and the rf_bucket to match the longitudinal distribution.
@@ -427,7 +433,7 @@ def __init__(self, macroparticlenumber, intensity, charge, mass,
         self._rf_bucket_matcher = CutRFBucket2D(
             macroparticlenumber, intensity, charge, mass,
             circumference, gamma_reference,
-            sigma_z, sigma_dp, is_accepted, *args, **kwargs)
+            sigma_z, sigma_dp, is_accepted, margin, *args, **kwargs)
         super(CutRFBucket6D, self).__init__(
             macroparticlenumber, intensity, charge, mass, circumference,
             gamma_reference, HEADTAILcoords.coordinates, *args, **kwargs)
@@ -474,39 +480,46 @@ def distribute(self):
 class CutRFBucket2D(ParticleGenerator):
     '''
     For HEADTAIL style matching into RF bucket.
+    The argument is_accepted takes a function (usually
+    RFBucket.is_in_separatrix) defining the separatrix of the rf bucket.
+    To avoid bucket leakage and particle losses as a consequence of the
+    unmatched initialisation, a margin / tolerance can be specified by
+    the argument margin (in % of RFBucket.Hmax, 5% by default), so that
+    particles will not be initialised too close to the separatrix. The
+    boundary will in fact be given by the equihamiltonian with a value
+    of margin*Hmax.
+
     BY KEVIN: NEEDS TO BE CLEANED UP BY ADRIAN!
     '''
     def __init__(self, macroparticlenumber, intensity, charge, mass,
                  circumference, gamma_reference, sigma_z, sigma_dp,
-                 is_accepted=None, generator_seed=None, *args, **kwargs):
+                 is_accepted, margin=0.05, *args, **kwargs):
 
         self.sigma_z = sigma_z
         self.sigma_dp = sigma_dp
         self.is_accepted = is_accepted
-        self.random_state = RandomState()
-        self.random_state.seed(generator_seed)
+        self.margin = margin
 
         super(CutRFBucket2D, self).__init__(
             macroparticlenumber, intensity, charge, mass, circumference,
             gamma_reference, HEADTAILcoords.longitudinal, *args, **kwargs)
 
     def distribute(self):
 
-        z  = self.sigma_z  * self.random_state.randn(self.macroparticlenumber)
-        dp = self.sigma_dp * self.random_state.randn(self.macroparticlenumber)
-        if self.is_accepted:
-            self._redistribute(z, dp)
+        z = normal(0, self.sigma_z, self.macroparticlenumber)
+        dp = normal(0, self.sigma_dp, self.macroparticlenumber)
+        self._redistribute(z, dp)
 
         return {'z': z, 'dp': dp}
 
     def _redistribute(self, z, dp):
 
-        mask_out = ~self.is_accepted(z, dp)
+        mask_out = ~self.is_accepted(z, dp, self.margin)
         while mask_out.any():
             n_gen = np.sum(mask_out)
-            z[mask_out] = self.sigma_z * self.random_state.randn(n_gen)
-            dp[mask_out] = self.sigma_dp * self.random_state.randn(n_gen)
-            mask_out = ~self.is_accepted(z, dp)
+            z[mask_out] = normal(0, self.sigma_z, n_gen)
+            dp[mask_out] = normal(0, self.sigma_dp, n_gen)
+            mask_out = ~self.is_accepted(z, dp, self.margin)
             self.prints('Reiterate on non-accepted particles')
 
 

testing/interactive-tests/GeneratorTest.ipynb

@@ -261,13 +261,14 @@ def p_max(self, zc):
         f = self.equihamiltonian(zc)
         return np.amax(f(self.zs))
 
-    def is_in_separatrix(self, z, dp):
+    def is_in_separatrix(self, z, dp, margin):
         """
-        Returns boolean whether this coordinate is located
-        strictly inside the separatrix.
+        Returns boolean whether this coordinate is located inside the
+        separatrix. A margin, in % of self.Hmax, can be specified so
+        that the separatrix is not strictly followed.
         """
         return np.logical_and(np.logical_and(self.zleft < z, z < self.zright),
-                              self.hamiltonian(z, dp) > 0.01 * self.Hmax)
+                              self.hamiltonian(z, dp) > margin * self.Hmax)
 
     def bucket_area(self):
         xmin, xmax = self.zleft, self.zright