started on #72

ParticlePhaseSpace/DataGenerators.py

@@ -0,0 +1,31 @@
+from abc import ABC, abstractmethod
+
+class _DataGenerator(ABC):
+
+
+    @abstractmethod
+    def _check_input_data(self):
+        pass
+
+    @abstractmethod
+    def _generate_data(self):
+        pass
+
+    def _check_generated_data(self):
+        """
+        want to be able to use the DataLoaders method for this since they are the same...
+        maybe make that a function?
+        :return:
+        """
+        pass
+
+class Twiss_data_generator(_DataGenerator):
+
+    def __init__(self, x_twiss, y_twiss, E, input_data):
+        super().__init__(input_data)
+
+    def _import_data(self):
+        pass
+
+    def _check_input_data(self):
+        pass

ParticlePhaseSpace/DataLoaders.py

@@ -333,4 +333,46 @@ def _import_data(self):
 
         self.data[self._columns['px']] = px
         self.data[self._columns['py']] = py
-        self.data[self._columns['pz']] = pz
+        self.data[self._columns['pz']] = pz
+
+
+class Generate_Twiss(_DataLoadersBase):
+    """
+    This class will generate phase space data by samping from a user defined twiss ellipse
+    """
+
+    def __init__(self, x_twiss={'epsilon': .004,
+                                'alpha': -0.3,
+                                'beta': 200},
+                       y_twiss={'epsilon': .004,
+                                'alpha': -0.3,
+                                'beta': 200},
+                       n_particles=10000,
+                       energy_func=None, **kwargs):
+        if energy_func is None:
+            # then use some default
+            self._energy_func = self._default_energy_function
+
+        super().__init__(input_data=None, **kwargs)
+
+    def _default_energy_function(self):
+        pass
+
+    def _get_gamma(self):
+        """
+
+        :return:
+        """
+        pass
+
+    def _check_input_data(self):
+        if not self._particle_type:
+            raise Exception('particle_type must be specified when using the twiss generator')
+
+    def _import_data(self):
+        """
+        note that we actually generate rather than import data here, but will keep the name for compatability
+        with other data loaders
+        :return:
+        """
+        pass

ParticlePhaseSpace/_ParticlePhaseSpace.py

@@ -741,7 +741,8 @@ def plot_transverse_trace_space_hist_2D(self, beam_direction: str = 'z', plot_tw
             _scale = 'log'
         else:
             _scale = None
-        self.calculate_twiss_parameters(beam_direction=beam_direction)
+        if not self.twiss_parameters:
+            self.calculate_twiss_parameters(beam_direction=beam_direction)
         fig, axs = plt.subplots(nrows=len(self._unique_particles), ncols=2, squeeze=False)
         row = 0
         for particle in self._unique_particles: