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

Author: like2000

monitors/monitors.py

@@ -380,7 +380,9 @@ def _write_data_to_file(self, bunch, arrays_dict):
         for quant in self.quantities_to_store:
             quant_values = getattr(bunch, quant)
             all_quantities[quant] = quant_values
-        all_quantities.update(arrays_dict)
+
+        if arrays_dict is not None:
+            all_quantities.update(arrays_dict)
 
         for quant in all_quantities.keys():
             quant_values = all_quantities[quant]

particles/generators.py

@@ -409,19 +409,23 @@ class CutRFBucket6D(ParticleGenerator):
     the optics resp. TWISS parameters taken from a TransverseMap
     instance. The longitudinal phase space is initialised as a
     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.
+    the function given by is_accepted. The usual choices for
+    is_accepted are either RFBucket.is_in_separatrix or
+    RFBucket.is_accepted . The former strictly follows the separatrix
+    of the bucket (the equihamiltonian with a value of 0), while with
+    the latter, tighter boundaries can be used (equihamiltonian lying
+    inside the RFBucket) as a result of which particles are not
+    initialised too close to the separatrix. This option is usually
+    preferred as it avoids bucket leakage and particle losses which may
+    occur as a consequence of the unmatched initialisation. The
+    RFBucket.is_accepted method must be created first however, by
+    calling the RFBucket.make_is_accepted(margin) method with a certain
+    value for the margin (in % of RFBucket.Hmax, 5% by default).
     '''
     def __init__(self, macroparticlenumber, intensity, charge, mass,
                  circumference, gamma_reference,
                  transverse_map, epsn_x, epsn_y,
-                 sigma_z, sigma_dp, is_accepted, margin=0.05,
+                 sigma_z, sigma_dp, is_accepted,
                  *args, **kwargs):
         '''Uses the transverse_map to extract the optics parameters
         and the rf_bucket to match the longitudinal distribution.
@@ -433,7 +437,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, margin, *args, **kwargs)
+            sigma_z, sigma_dp, is_accepted, *args, **kwargs)
         super(CutRFBucket6D, self).__init__(
             macroparticlenumber, intensity, charge, mass, circumference,
             gamma_reference, HEADTAILcoords.coordinates, *args, **kwargs)
@@ -478,27 +482,29 @@ 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.
+    '''For HEADTAIL style matching into RF bucket.
+    The argument is_accepted takes a function (i.e. reference to a
+    function). The usual choices are RFBucket.is_in_separatrix or
+    RFBucket.is_accepted . The former strictly follows the separatrix
+    of the bucket (the equihamiltonian with a value of 0), while with
+    the latter, tighter boundaries can be used (equihamiltonian lying
+    inside the RFBucket) as a result of which particles are not
+    initialised too close to the separatrix. This option is usually
+    preferred as it avoids bucket leakage and particle losses which may
+    occur as a consequence of the unmatched initialisation.
+    The RFBucket.is_accepted method must be created first however, by
+    calling the RFBucket.make_is_accepted(margin) method with a certain
+    value for the margin (in % of RFBucket.Hmax, 5% by default).
 
     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, margin=0.05, *args, **kwargs):
+                 is_accepted, *args, **kwargs):
 
         self.sigma_z = sigma_z
         self.sigma_dp = sigma_dp
         self.is_accepted = is_accepted
-        self.margin = margin
 
         super(CutRFBucket2D, self).__init__(
             macroparticlenumber, intensity, charge, mass, circumference,
@@ -514,12 +520,12 @@ def distribute(self):
 
     def _redistribute(self, z, dp):
 
-        mask_out = ~self.is_accepted(z, dp, self.margin)
+        mask_out = ~self.is_accepted(z, dp)
         while mask_out.any():
             n_gen = np.sum(mask_out)
             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)
+            mask_out = ~self.is_accepted(z, dp)
             self.prints('Reiterate on non-accepted particles')
 
 

testing/interactive-tests/GeneratorTest.ipynb

@@ -1,7 +1,7 @@
 {
  "metadata": {
   "name": "",
-  "signature": "sha256:dfff2f627b552921077821c4363a152f9af9e974b4c53e6434e2d24a4b077c40"
+  "signature": "sha256:40730948c653f14f34c07a1fed2c8a8d06c79f2381ae5f9084b77273c9fcb345"
  },
  "nbformat": 3,
  "nbformat_minor": 0,
@@ -44,7 +44,7 @@
        "output_type": "stream",
        "stream": "stdout",
        "text": [
-        "PyHEADTAIL v1.0.0\n",
+        "PyHEADTAIL v1.0.4\n",
         "\n",
         "\n"
        ]
@@ -95,7 +95,6 @@
       "        for key in particledata[step].keys():\n",
       "            _[:] = particledata[step][key][:]\n",
       "    \n",
-      "    \n",
       "    bunchdata.close()\n",
       "    slicedata.close()\n",
       "    particledata.close()\n",
@@ -149,7 +148,7 @@
      "language": "python",
      "metadata": {},
      "outputs": [],
-     "prompt_number": 3
+     "prompt_number": 12
     },
     {
      "cell_type": "code",
@@ -176,7 +175,7 @@
      "language": "python",
      "metadata": {},
      "outputs": [],
-     "prompt_number": 4
+     "prompt_number": 9
     },
     {
      "cell_type": "heading",
@@ -211,7 +210,7 @@
      "language": "python",
      "metadata": {},
      "outputs": [],
-     "prompt_number": 5
+     "prompt_number": 10
     },
     {
      "cell_type": "code",
@@ -240,13 +239,17 @@
       "    write_buffer_every=20)\n",
       "particle_monitor = ParticleMonitor(filename=particle_filename, stride=10, parameters_dict={'Q_x': Q_x})\n",
       "\n",
+      "arrays_dict = {}\n",
       "map_ = trans_map\n",
       "for i in xrange(n_turns):\n",
       "    for m_ in map_:\n",
       "        m_.track(bunch)\n",
       "    bunch_monitor.dump(bunch)\n",
       "    slice_monitor.dump(bunch)\n",
-      "    particle_monitor.dump(bunch)\n",
+      "    \n",
+      "    slice_set_pmon = bunch.get_slices(unibin_slicer)\n",
+      "    arrays_dict.update({'slidx': slice_set_pmon.slice_index_of_particle, 'zz': bunch.z})\n",
+      "    particle_monitor.dump(bunch, arrays_dict)\n",
       "\n",
       "read_all_data(bunch_filename, slice_filename, particle_filename)\n",
       "read_n_plot_data(bunch_filename, slice_filename, particle_filename)"
@@ -258,11 +261,11 @@
        "output_type": "stream",
        "stream": "stdout",
        "text": [
-        "0.0580668146353\n"
+        "0.0576953432278\n"
        ]
       }
      ],
-     "prompt_number": 8
+     "prompt_number": 13
     },
     {
      "cell_type": "code",

testing/interactive-tests/MonitorTest.ipynb

@@ -261,14 +261,26 @@ def p_max(self, zc):
         f = self.equihamiltonian(zc)
         return np.amax(f(self.zs))
 
-    def is_in_separatrix(self, z, dp, margin):
-        """
-        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.
-        """
+    def is_in_separatrix(self, z, dp):
+        """ Returns boolean whether the coordinate (z, dp) is located
+        strictly inside the separatrix. """
         return np.logical_and(np.logical_and(self.zleft < z, z < self.zright),
-                              self.hamiltonian(z, dp) > margin * self.Hmax)
+                              self.hamiltonian(z, dp) > 0)
+
+    def make_is_accepted(self, margin):
+        """ Returns the function is_accepted(z, dp) definining the
+        equihamiltonian with a value of margin*self.Hmax . For margin 0,
+        the returned is_accepted(z, dp) function is equivalent to
+        self.is_in_separatrix(z, dp). """
+
+        def is_accepted(z, dp):
+            """ Returns boolean whether the coordinate (z, dp) is
+            located inside the equihamiltonian defined by
+            margin*self.Hmax . """
+            return np.logical_and(
+                np.logical_and(self.zleft < z, z < self.zright),
+                self.hamiltonian(z, dp) > margin * self.Hmax)
+        return is_accepted
 
     def bucket_area(self):
         xmin, xmax = self.zleft, self.zright