Merge pull request #125 from bwheelz36/bw_dev

some bug fixes in filter by time functionality

Author: bwheelz36

ParticlePhaseSpace/DataLoaders.py

@@ -181,9 +181,7 @@ def _import_data(self):
                     raise Exception(f'failed to calculate momentums from topas data. Possible solution is to increase'
                                     f'the value of momentum_precision_factor, currently set to {momentum_precision_factor: 1.2e}'
                                     f'and failed data has value {relative_difference: 1.2e}')
-            warnings.warn(f'{n_negative_locations: d} entries returned invalid pz values and were set to zero.'
-                          f'\nWe will now check that momentum and energy are consistent to within '
-                          f'{self._energy_consistency_check_cutoff: 1.4f} {self._units.energy.label}')
+            warnings.warn(f'{n_negative_locations: d} entries returned invalid pz values and were set to zero.')
 
         ParticleDir = [-1 if elem else 1 for elem in ParticleDir]
         self.data[self._columns['pz']] = np.multiply(np.sqrt(temp), ParticleDir)

ParticlePhaseSpace/_ParticlePhaseSpace.py

@@ -304,27 +304,32 @@ def particle_positions_hist_2D(self, beam_direction: str = 'z', quantity: str =
             ind = self._PS._ps_data['particle type [pdg_code]'] == particle
             ps_data = self._PS._ps_data.loc[ind]
             if beam_direction == 'x':
-                loop_data = zip(ps_data[self._PS.columns['z']], ps_data[self._PS.columns['y']],
-                                ps_data[self._PS.columns['Ek']],
-                                ps_data['weight'])
-                _xlabel = self._PS.columns['z']
-                _ylabel = self._PS.columns['y']
-            if beam_direction == 'y':
-                loop_data = zip(ps_data[self._PS.columns['x']], ps_data[self._PS.columns['z']],
-                                ps_data[self._PS.columns['Ek']],
-                                ps_data['weight'])
-                _xlabel = self._PS.columns['x']
-                _ylabel = self._PS.columns['z']
-            if beam_direction == 'z':
-                loop_data = zip(ps_data[self._PS.columns['x']], ps_data[self._PS.columns['y']],
-                                ps_data[self._PS.columns['Ek']],
-                                ps_data['weight'])
-                _xlabel = self._PS.columns['x']
-                _ylabel = self._PS.columns['y']
+                x_data = ps_data[self._PS.columns['y']]
+                y_data = ps_data[self._PS.columns['z']]
+                _x_label = self._PS.columns['y']
+                _y_label = self._PS.columns['z']
+            elif beam_direction == 'y':
+                x_data = ps_data[self._PS.columns['x']]
+                y_data = ps_data[self._PS.columns['z']]
+                _x_label = self._PS.columns['x']
+                _y_label = self._PS.columns['z']
+            elif beam_direction == 'z':
+                x_data = ps_data[self._PS.columns['x']]
+                y_data = ps_data[self._PS.columns['y']]
+                _x_label = self._PS.columns['x']
+                _y_label = self._PS.columns['y']
+            else:
+                raise NotImplementedError('beam direction must be one of "x", "y" or "z"')
             if xlim is None:
-                xlim = [ps_data[self._PS.columns['x']].min(), ps_data[self._PS.columns['x']].max()]
+                if beam_direction == 'x':
+                    xlim = [ps_data[self._PS.columns['y']].min(), ps_data[self._PS.columns['y']].max()]
+                elif beam_direction == 'y' or beam_direction == 'z':
+                    xlim = [ps_data[self._PS.columns['x']].min(), ps_data[self._PS.columns['x']].max()]
             if ylim is None:
-                ylim = [ps_data[self._PS.columns['y']].min(), ps_data[self._PS.columns['y']].max()]
+                if beam_direction == 'z':
+                    ylim = [ps_data[self._PS.columns['y']].min(), ps_data[self._PS.columns['y']].max()]
+                elif beam_direction == 'x' or beam_direction == 'y':
+                    ylim = [ps_data[self._PS.columns['z']].min(), ps_data[self._PS.columns['z']].max()]
             if quantity == 'intensity':
                 _title = f"n_particles intensity;\n{particle_cfg.particle_properties[particle]['name']}"
                 _weight = ps_data['weight']
@@ -333,24 +338,23 @@ def particle_positions_hist_2D(self, beam_direction: str = 'z', quantity: str =
                 _weight = np.multiply(ps_data['weight'], ps_data[self._PS.columns['Ek']])
             X = np.linspace(xlim[0], xlim[1], bins)
             Y = np.linspace(ylim[0], ylim[1], bins)
-            h, xedges, yedges = np.histogram2d(ps_data[self._PS.columns['x']],
-                                               ps_data[self._PS.columns['y']],
-                                               bins=[X, Y], weights=_weight, )
+            h, xedges, yedges = np.histogram2d(x_data,
+                                               y_data,
+                                               bins=[X, Y], weights=_weight)
+
             if normalize:
-                h = h * 100 / h.max()
-            # _im1 = axs[0, n_axs].hist2d(ps_data[self._PS.columns['x']], ps_data[self._PS.columns['y']],
-            #                           bins=[X,Y],
-            #                           weights=_weight, norm=LogNorm(vmin=1, vmax=100),
-            #                           cmap='inferno',
-            #                           vmin=vmin, vmax=vmax)[3]
+                try:
+                    h = h * 100 / h.max()
+                except ValueError:
+                    pass
             _im1 = axs[0, n_axs].pcolormesh(xedges, yedges, h.T, cmap='inferno',
                                             norm=_scale, rasterized=False, vmin=vmin, vmax=vmax)
 
             fig.colorbar(_im1, ax=axs[0, n_axs])
 
             axs[0, n_axs].set_title(_title)
-            axs[0, n_axs].set_xlabel(_xlabel, fontsize=_FigureSpecs.LabelFontSize)
-            axs[0, n_axs].set_ylabel(_ylabel, fontsize=_FigureSpecs.LabelFontSize)
+            axs[0, n_axs].set_xlabel(_x_label, fontsize=_FigureSpecs.LabelFontSize)
+            axs[0, n_axs].set_ylabel(_y_label, fontsize=_FigureSpecs.LabelFontSize)
             axs[0, n_axs].set_aspect('equal')
             if grid:
                 axs[0, n_axs].grid()
@@ -1376,7 +1380,7 @@ def assess_density_versus_r(self, Rvals=None, verbose: bool = True, beam_directi
             print(density_data)
         return density_data
 
-    def filter_by_time(self, t_start, t_finish):
+    def filter_by_time(self, t_start, t_finish, in_place: bool=False):
         """
         Generates a new PhaseSpace which only contains particles inside t_start and t_finish (inclusive).
         t_start and t_finish should be specfied in ps.
@@ -1389,16 +1393,9 @@ def filter_by_time(self, t_start, t_finish):
         """
         ind = np.logical_and(self._ps_data[self.columns['time']] >= t_start,
                              self._ps_data[self.columns['time']] <= t_finish)
-        ps_data = self._ps_data[ind]
-        for col_name in ps_data.columns:
-            if not col_name in ps_cfg.get_required_column_names(self._units):
-                ps_data.drop(columns=col_name, inplace=True)
-        # create a new instance of _DataImportersBase based on particle_data
-        ps_data_loader = DataLoaders.Load_PandasData(ps_data, units=self._units)
-        new_instance = PhaseSpace(ps_data_loader)
-        print(f'Original data contains {len(self): d} particles')
-        print(f'Filtered data contains {len(new_instance): d} particles')
-        return new_instance
+        new_PS = self.filter_by_boolean_index(ind, in_place=in_place)
+        return new_PS
+
 
     def filter_by_boolean_index(self, boolean_index, in_place: bool = False, split: bool = False):
         """