release-script: Merge branch 'release/v1.12.4'

PyHEADTAIL/_version.py

@@ -1 +1 @@
-__version__ = '1.12.3'
+__version__ = '1.12.4'

PyHEADTAIL/cobra_functions/stats.pyx

@@ -11,6 +11,7 @@ cimport libc.math as cmath
 
 cimport cython.boundscheck
 cimport cython.cdivision
+cimport cython.wraparound
 
 
 
@@ -196,6 +197,7 @@ cpdef count_macroparticles_per_slice(int[::1] slice_index_of_particle,
         s_idx = slice_index_of_particle[particles_within_cuts[i]]
         n_macroparticles[s_idx] += 1
 
+
 @cython.boundscheck(False)
 @cython.cdivision(True)
 cpdef sort_particle_indices_by_slice(int[::1] slice_index_of_particle,
@@ -223,6 +225,7 @@ cpdef sort_particle_indices_by_slice(int[::1] slice_index_of_particle,
         particle_indices_by_slice[pos] = p_idx
         pos_ctr[s_idx] += 1
 
+
 @cython.boundscheck(False)
 @cython.cdivision(True)
 cpdef mean_per_slice(int[::1] slice_index_of_particle,
@@ -245,6 +248,7 @@ cpdef mean_per_slice(int[::1] slice_index_of_particle,
         if n_macroparticles[i]:
             mean_u[i] /= n_macroparticles[i]
 
+
 @cython.boundscheck(False)
 @cython.cdivision(True)
 cpdef std_per_slice(int[::1] slice_index_of_particle,
@@ -360,6 +364,7 @@ cpdef emittance_per_slice_old(int[::1] slice_index_of_particle,
 
         epsn_u[i] = cmath.sqrt(u2[i]*up2[i] - uup[i]*uup[i])
 
+
 @cython.boundscheck(False)
 @cython.cdivision(True)
 cpdef emittance_per_slice(int[::1] slice_index_of_particle,
@@ -375,7 +380,7 @@ cpdef emittance_per_slice(int[::1] slice_index_of_particle,
     for each slice.
     """
     cdef unsigned int n_slices = emittance.shape[0]
-    # allocate arrays for covariances 
+    # allocate arrays for covariances
     cdef double[::1] cov_u2 = np.zeros(n_slices, dtype=np.double)
     cdef double[::1] cov_up2 = np.zeros(n_slices, dtype=np.double)
     cdef double[::1] cov_u_up = np.zeros(n_slices, dtype=np.double)
@@ -409,3 +414,84 @@ cpdef emittance_per_slice(int[::1] slice_index_of_particle,
             emittance[i] = cmath.sqrt(_det_beam_matrix(sigma11, sigma12, sigma22))
         else:
             emittance[i] = 0.
+
+@cython.boundscheck(False)
+@cython.cdivision(True)
+@cython.wraparound(False)
+cpdef calc_cell_stats(bunch, double beta_z, double radial_cut,
+                      int n_rings, int n_azim_slices):
+
+    # Prepare arrays to store cell statistics.
+    cdef int[:,::1] n_particles_cell = np.zeros((n_azim_slices, n_rings),
+                                                dtype=np.int32)
+    cdef double[:,::1] mean_x_cell = np.zeros((n_azim_slices, n_rings),
+                                              dtype=np.double)
+    cdef double[:,::1] mean_xp_cell = np.zeros((n_azim_slices, n_rings),
+                                               dtype=np.double)
+    cdef double[:,::1] mean_y_cell = np.zeros((n_azim_slices, n_rings),
+                                              dtype=np.double)
+    cdef double[:,::1] mean_yp_cell = np.zeros((n_azim_slices, n_rings),
+                                               dtype=np.double)
+    cdef double[:,::1] mean_z_cell = np.zeros((n_azim_slices, n_rings),
+                                              dtype=np.double)
+    cdef double[:,::1] mean_dp_cell = np.zeros((n_azim_slices, n_rings),
+                                               dtype=np.double)
+
+    # Declare datatypes of bunch coords.
+    cdef double[::1] x = bunch.x
+    cdef double[::1] xp = bunch.xp
+    cdef double[::1] y = bunch.y
+    cdef double[::1] yp = bunch.yp
+    cdef double[::1] z = bunch.z
+    cdef double[::1] dp = bunch.dp
+    cdef unsigned int n_particles = x.shape[0]
+
+    cdef double ring_width = radial_cut / <double>n_rings
+    cdef double azim_width = 2.*cmath.M_PI / <double>n_azim_slices
+    cdef double beta_z_square = beta_z*beta_z
+
+    cdef double z_i, dp_i, long_action
+    cdef unsigned int p_idx
+    cdef int ring_idx, azim_idx
+    for p_idx in xrange(n_particles):
+        z_i = z[p_idx]
+        dp_i = dp[p_idx]
+
+        # Slice radially.
+        long_action = cmath.sqrt(z_i*z_i + beta_z_square*dp_i*dp_i)
+        ring_idx = <int>cmath.floor(long_action / ring_width)
+        if ring_idx >= n_rings:
+            continue
+
+        # Slice azimuthally: atan 2 returns values in the interval [-pi, pi];
+        # hence, in order to avoid negative indices, we need to add an offset of
+        # +pi before performing the floor division (this consequently just adds
+        # an offset to the slices indices). The one-to-one mapping between
+        # angles and slice indices is retained as desired. In this
+        # interpretation, slice index 0 corresponds to -pi (i.e. starting in 3rd
+        # quadrant) and slice n-1 correspoinds to +pi (i.e. ending in 2nd
+        # quadrant). This needs to be taken into account when interpreting and
+        # plotting cell monitor data - for this, use
+        # theta = np.linspace(-np.pi, np.pi, n_azim_slices)
+        azim_idx = <int>cmath.floor(
+            (cmath.M_PI + cmath.atan2(beta_z*dp_i, z_i)) / azim_width)
+
+        n_particles_cell[azim_idx, ring_idx] += 1
+        mean_x_cell[azim_idx, ring_idx] += x[p_idx]
+        mean_xp_cell[azim_idx, ring_idx] += xp[p_idx]
+        mean_y_cell[azim_idx, ring_idx] += y[p_idx]
+        mean_yp_cell[azim_idx, ring_idx] += yp[p_idx]
+        mean_z_cell[azim_idx, ring_idx] += z_i
+        mean_dp_cell[azim_idx, ring_idx] += dp_i
+
+    for azim_idx in xrange(n_azim_slices):
+        for ring_idx in xrange(n_rings):
+            if n_particles_cell[azim_idx, ring_idx] > 0:
+                mean_x_cell[azim_idx, ring_idx] /= n_particles_cell[azim_idx, ring_idx]
+                mean_xp_cell[azim_idx, ring_idx] /= n_particles_cell[azim_idx, ring_idx]
+                mean_y_cell[azim_idx, ring_idx] /= n_particles_cell[azim_idx, ring_idx]
+                mean_yp_cell[azim_idx, ring_idx] /= n_particles_cell[azim_idx, ring_idx]
+                mean_z_cell[azim_idx, ring_idx] /= n_particles_cell[azim_idx, ring_idx]
+                mean_dp_cell[azim_idx, ring_idx] /= n_particles_cell[azim_idx, ring_idx]
+
+    return n_particles_cell, mean_x_cell, mean_xp_cell, mean_y_cell, mean_yp_cell, mean_z_cell, mean_dp_cell

PyHEADTAIL/gpu/thrust_code.cu

@@ -1,3 +1,4 @@
+#include <thrust/gather.h>
 #include <thrust/sort.h>
 #include <thrust/binary_search.h>
 #include <thrust/device_vector.h>

PyHEADTAIL/monitors/monitors.py

@@ -17,6 +17,10 @@
 from ..gpu import gpu_utils as gpu_utils
 from ..general import decorators as decorators
 
+from ..cobra_functions import stats as cp
+
+# from .. import cobra_functions.stats.calc_cell_stats as calc_cell_stats
+
 
 class Monitor(Printing):
     """ Abstract base class for monitors. A monitor can request
@@ -122,9 +126,9 @@ def _create_file_structure(self, parameters_dict):
                 h5group.create_dataset(stats, shape=(self.n_steps,),
                                        compression='gzip', compression_opts=9)
             h5file.close()
-        except:
-            self.prints('Creation of bunch monitor file ' + self.filename +
-                   'failed. \n')
+        except Exception as err:
+            self.warns('Problem occurred during Bunch monitor creation.')
+            self.warns(err.message)
             raise
 
     @decorators.synchronize_gpu_streams_after
@@ -296,9 +300,9 @@ def _create_file_structure(self, parameters_dict):
                 h5group_slice.create_dataset(stats, shape=(self.slicer.n_slices,
                     self.n_steps), compression='gzip', compression_opts=9)
             h5file.close()
-        except:
-            self.prints('Creation of slice monitor file ' + self.filename +
-                        'failed. \n')
+        except Exception as err:
+            self.warns('Problem occurred during Slice monitor creation.')
+            self.warns(err.message)
             raise
 
     def _write_data_to_buffer(self, bunch):
@@ -398,7 +402,7 @@ def __init__(self, filename, stride=1, parameters_dict=None,
 
           Optionally pass a list called quantities_to_store which
           specifies which members of the bunch will be called/stored.
-                           """
+        """
         quantities_to_store = [ 'x', 'xp', 'y', 'yp', 'z', 'dp', 'id' ]
         self.quantities_to_store = kwargs.pop('quantities_to_store',
                                               quantities_to_store)
@@ -424,9 +428,9 @@ def _create_file_structure(self, parameters_dict):
                 for key in parameters_dict:
                     h5file.attrs[key] = parameters_dict[key]
             h5file.close()
-        except:
-            self.prints('Creation of particle monitor file ' +
-                        self.filename + 'failed. \n')
+        except Exception as err:
+            self.warns('Problem occurred during Particle monitor creation.')
+            self.warns(err.message)
             raise
 
     def _write_data_to_file(self, bunch, arrays_dict):
@@ -465,3 +469,145 @@ def _write_data_to_file(self, bunch, arrays_dict):
             h5group[quant][:] = quant_values[::self.stride]
 
         h5file.close()
+
+
+class CellMonitor(Monitor):
+    """ Class to store cell (z, dp) specific data (for the moment only
+    mean_x, mean_y, mean_z, mean_dp and n_particles_in_cell) to a HDF5
+    file. This monitor uses a buffer (shift register) to reduce the
+    number of writing operations to file. This also helps to avoid IO
+    errors and loss of data when writing to a file that may become
+    temporarily unavailable (e.g. if file is located on network) during
+    the simulation. """
+
+    def __init__(self, filename, n_steps, n_azimuthal_slices, n_radial_slices,
+                 radial_cut, beta_z, parameters_dict=None,
+                 write_buffer_every=512, buffer_size=4096, *args, **kwargs):
+        """ Create an instance of a CellMonitor class. Apart from
+        initializing the HDF5 file, a buffer self.buffer_cell is
+        prepared to buffer the cell-specific data before writing them
+        to file.
+
+          filename:           Path and name of HDF5 file. Without file
+                              extension.
+          n_steps:            Number of entries to be reserved for each
+                              of the quantities in self.stats_to_store.
+          n_azimuthal_slices: Number of pizza slices (azimuthal slicing).
+          n_radial_slices:    Number of rings (radial slicing).
+          radial_cut:         'Radius' of the outermost ring in
+                              longitudinal phase space (using beta_z*dp)
+          parameters_dict:    Metadata for HDF5 file containing main
+                              simulation parameters.
+          write_buffer_every: Number of steps after which buffer
+                              contents are actually written to file.
+          buffer_size:        Number of steps to be buffered. """
+        stats_to_store = [
+            'mean_x', 'mean_xp',
+            'mean_y', 'mean_yp',
+            'mean_z', 'mean_dp',
+            'macroparticlenumber']
+        self.stats_to_store = kwargs.pop('stats_to_store', stats_to_store)
+
+        self.filename = filename
+        self.n_steps = n_steps
+        self.i_steps = 0
+        self.n_azimuthal_slices = n_azimuthal_slices
+        self.n_radial_slices = n_radial_slices
+        self.radial_cut = radial_cut
+        self.beta_z = beta_z
+
+        # Prepare buffer.
+        self.buffer_size = buffer_size
+        self.write_buffer_every = write_buffer_every
+        self.buffer_cell = {}
+
+        for stats in self.stats_to_store:
+            self.buffer_cell[stats] = (
+                np.zeros((self.n_azimuthal_slices, self.n_radial_slices, self.buffer_size)))
+        self._create_file_structure(parameters_dict)
+
+    def dump(self, bunch):
+        """ Evaluate the statistics for the given cells and write the
+        data to the buffer and/or to the HDF5 file. The buffer is used
+        to reduce the number of writing operations to file. This helps
+        to avoid IO errors and loss of data when writing data to a file
+        that may become temporarily unavailable (e.g. if file is on
+        network) during the simulation. Buffer contents are written to
+        file only every self.write_buffer_every steps. """
+        self._write_data_to_buffer(bunch)
+        if ((self.i_steps + 1) % self.write_buffer_every == 0 or
+                (self.i_steps + 1) == self.n_steps):
+            self._write_buffer_to_file()
+        self.i_steps += 1
+
+    def _create_file_structure(self, parameters_dict):
+        """ Initialize HDF5 file and create its basic structure (groups
+        and datasets). One dataset for each of the quantities defined
+        in self.stats_to_store is generated. If specified by
+        the user, write the contents of the parameters_dict as metadata
+        (attributes) to the file. Maximum file compression is
+        activated. """
+        try:
+            h5file = hp.File(self.filename + '.h5', 'w')
+            if parameters_dict:
+                for key in parameters_dict:
+                    h5file.attrs[key] = parameters_dict[key]
+
+            h5file.create_group('Cells')
+            h5group_cells = h5file['Cells']
+            for stats in self.stats_to_store:
+                h5group_cells.create_dataset(
+                    stats, compression='gzip', compression_opts=9,
+                    shape=(self.n_azimuthal_slices, self.n_radial_slices,
+                           self.n_steps))
+            h5file.close()
+        except Exception as err:
+            self.warns('Problem occurred during Cell monitor creation.')
+            self.warns(err.message)
+            raise
+
+    def _write_data_to_buffer(self, bunch):
+        """ Store the data in the self.buffer dictionary before writing
+        them to file. The buffer is implemented as a shift register. The
+        cell-specific data are computed by a cython function. """
+
+        from PyHEADTAIL.cobra_functions.stats import calc_cell_stats
+        n_cl, x_cl, xp_cl, y_cl, yp_cl, z_cl, dp_cl = calc_cell_stats(
+            bunch, self.beta_z, self.radial_cut, self.n_radial_slices, self.n_azimuthal_slices)
+
+        self.buffer_cell['mean_x'][:,:,0] = x_cl[:,:]
+        self.buffer_cell['mean_xp'][:,:,0] = xp_cl[:,:]
+        self.buffer_cell['mean_y'][:,:,0] = y_cl[:,:]
+        self.buffer_cell['mean_yp'][:,:,0] = yp_cl[:,:]
+        self.buffer_cell['mean_z'][:,:,0] = z_cl[:,:]
+        self.buffer_cell['mean_dp'][:,:,0] = dp_cl[:,:]
+        self.buffer_cell['macroparticlenumber'][:,:,0] = n_cl[:,:]
+
+        for stats in self.stats_to_store:
+            self.buffer_cell[stats] = np.roll(
+                self.buffer_cell[stats], shift=-1, axis=2)
+
+    def _write_buffer_to_file(self):
+        """ Write buffer contents to the HDF5 file. The file is opened
+        and closed each time the buffer is written to file to prevent
+        from loss of data in case of a crash. """
+        # Keep track of where to read from buffers and where to store
+        # data in file.
+        n_entries_in_buffer = min(self.i_steps+1, self.buffer_size)
+        low_pos_in_buffer = self.buffer_size - n_entries_in_buffer
+        low_pos_in_file = self.i_steps + 1 - n_entries_in_buffer
+        up_pos_in_file = self.i_steps + 1
+
+        # Try to write data to file. If file is not available, skip this
+        # step and repeat it again after self.write_buffer_every. As
+        # long as self.buffer_size is not exceeded, no data are lost.
+        try:
+            h5file = hp.File(self.filename + '.h5', 'a')
+            h5group_cells = h5file['Cells']
+
+            for stats in self.stats_to_store:
+                h5group_cells[stats][:,:,low_pos_in_file:up_pos_in_file] = \
+                    self.buffer_cell[stats][:,:,low_pos_in_buffer:]
+            h5file.close()
+        except Exception as err:
+            self.warns(err.message)