diff --git a/PyHEADTAIL/_version.py b/PyHEADTAIL/_version.py index 3081afba..aba17a1d 100644 --- a/PyHEADTAIL/_version.py +++ b/PyHEADTAIL/_version.py @@ -1 +1 @@ -__version__ = '1.12.3' +__version__ = '1.12.4' diff --git a/PyHEADTAIL/cobra_functions/stats.pyx b/PyHEADTAIL/cobra_functions/stats.pyx index f4b66a97..a144ce05 100644 --- a/PyHEADTAIL/cobra_functions/stats.pyx +++ b/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 / n_rings + cdef double azim_width = 2.*cmath.M_PI / 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 = 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 = 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 diff --git a/PyHEADTAIL/gpu/thrust_code.cu b/PyHEADTAIL/gpu/thrust_code.cu index cab93e94..b9c46721 100644 --- a/PyHEADTAIL/gpu/thrust_code.cu +++ b/PyHEADTAIL/gpu/thrust_code.cu @@ -1,3 +1,4 @@ +#include #include #include #include diff --git a/PyHEADTAIL/monitors/monitors.py b/PyHEADTAIL/monitors/monitors.py index b993d483..8ba63046 100644 --- a/PyHEADTAIL/monitors/monitors.py +++ b/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) diff --git a/PyHEADTAIL/particles/generators.py b/PyHEADTAIL/particles/generators.py index 7527d4e2..2e5dbd0f 100644 --- a/PyHEADTAIL/particles/generators.py +++ b/PyHEADTAIL/particles/generators.py @@ -375,6 +375,62 @@ def _uniform2D(n_particles): return _uniform2D +''' +Why we have this ll algo in here? We had a better one (Knuth): +http://citeseerx.ist.psu.edu/viewdoc/download?doi=10.1.1.138.8671&rep=rep1&type=pdf + + +import numpy as np + + +def transform_into_kv(bunch, a_x, a_xp, a_y, a_yp): + + # KV distribution + # ================================================== + d = 4 + for i in range(bunch.macroparticlenumber): + u = np.random.normal(size=d) + r = np.sqrt(np.sum(u**2)) + u *= 1./r + + bunch.x[i] = u[0] + bunch.xp[i] = u[1] + bunch.y[i] = u[2] + bunch.yp[i] = u[3] + # ================================================== + + bunch.x *= a_x + bunch.xp *= a_xp + bunch.y *= a_y + bunch.yp *= a_yp + + +def transform_into_waterbag(bunch, a_x, a_xp, a_y, a_yp): + + # waterbag distribution + # ================================================== + d = 4 + for i in range(bunch.macroparticlenumber): + u = np.random.normal(size=d) + r = np.sqrt(np.sum(u**2)) + u *= (np.random.rand(1))**(1./d)/r + + bunch.x[i] = u[0] + bunch.xp[i] = u[1] + bunch.y[i] = u[2] + bunch.yp[i] = u[3] + # ================================================== + + bunch.x *= a_x + bunch.xp *= a_xp + bunch.y *= a_y + bunch.yp *= a_yp + + +Want to get this back in the near future... +''' + + def kv2D(r_u, r_up): '''Closure which generates a Kapchinski-Vladimirski-type uniform distribution in 2D. The extent is determined by the arguments. diff --git a/PyHEADTAIL/testing/unittests/test_monitor.py b/PyHEADTAIL/testing/unittests/test_monitor.py index c78f3459..071796bb 100644 --- a/PyHEADTAIL/testing/unittests/test_monitor.py +++ b/PyHEADTAIL/testing/unittests/test_monitor.py @@ -18,7 +18,9 @@ from scipy.constants import c, e, m_p import h5py as hp -from PyHEADTAIL.monitors.monitors import BunchMonitor, SliceMonitor +from PyHEADTAIL.particles.particles import Particles +from PyHEADTAIL.monitors.monitors import ( + BunchMonitor, SliceMonitor, CellMonitor) class TestMonitor(unittest.TestCase): ''' Test the BunchMonitor/SliceMonitor''' @@ -57,7 +59,7 @@ def test_bunchmonitor(self): def test_slicemonitor(self): ''' - Test whether the slicemonitor works as excpected, use the mock slicer + Test whether the slicemonitor works as expected, use the mock slicer ''' nslices = 3 mock_slicer = self.generate_mock_slicer(nslices) @@ -79,6 +81,25 @@ def test_slicemonitor(self): self.assertTrue(np.allclose(sd['propertyA'][k,j], k + (j+1)*1000), 'Slices part of SliceMonitor wrong') + def test_cellmonitor(self): + ''' + Test whether the cellmonitor works as expected. + ''' + bunch = self.generate_real_bunch() + cell_monitor = CellMonitor( + filename=self.s_fn, n_steps=self.n_turns, + n_azimuthal_slices=4, n_radial_slices=3, + radial_cut=bunch.sigma_dp() * 3, + beta_z=np.abs(0.003 * bunch.circumference / (2*np.pi*0.004)), + write_buffer_every=9, + ) + for i in xrange(self.n_turns): + cell_monitor.dump(bunch) + s = hp.File(self.s_fn + '.h5') + sc = s['Cells'] + + # to be extended + def generate_mock_bunch(self): ''' @@ -103,6 +124,33 @@ def get_slices(self, slicer, **kwargs): return Mock() + + def generate_real_bunch(self): + + #beam parameters + intensity = 1.234e9 + circumference = 111. + gamma = 20.1 + + #simulation parameters + macroparticlenumber = 2048 + particlenumber_per_mp = intensity / macroparticlenumber + + x = np.random.uniform(-1, 1, macroparticlenumber) + y = np.random.uniform(-1, 1, macroparticlenumber) + z = np.random.uniform(-1, 1, macroparticlenumber) + xp = np.random.uniform(-0.5, 0.5, macroparticlenumber) + yp = np.random.uniform(-0.5, 0.5, macroparticlenumber) + dp = np.random.uniform(-0.5, 0.5, macroparticlenumber) + coords_n_momenta_dict = { + 'x': x, 'y': y, 'z': z, + 'xp': xp, 'yp': yp, 'dp': dp + } + return Particles( + macroparticlenumber, particlenumber_per_mp, e, m_p, + circumference, gamma, coords_n_momenta_dict + ) + def generate_mock_slicer(self, nslices): ''' Create a mock slicer to test behaviour''' class Mock():