Merge pull request #20 from impy-project/refactor/api_changes

Refactor/api changes

impy/__init__.py

@@ -29,4 +29,9 @@
 
 pdata = PYTHIAParticleData(
     cache_file=open(
-        join(base_path, impy_config["pdata_cachefile"]), 'wb'))
+        join(base_path, impy_config["pdata_cachefile"]), 'wb'))
+
+
+import impy.models as models
+import impy.kinematics as kinematics
+import impy.constants as constants

impy/common.py

@@ -18,6 +18,28 @@
 from impy.util import info
 
 
+class EventData:
+    """Data structure to keep views of filtered data"""
+    def __init__(self, npart, p_ids, status, charge,
+                 px, py, pz, en, m, vx, vy, vz, vt,
+                 pem_arr, vt_arr):
+        self.npart = npart
+        self.p_ids = p_ids
+        self.status = status
+        self.charge = charge
+        self.px = px
+        self.py = py
+        self.pz = pz
+        self.en = en
+        self.m = m
+        self.vx = vx
+        self.vy = vy
+        self.vz = vz
+        self.vt = vt
+        self._pem_arr = pem_arr
+        self._vt_arr = vt_arr
+
+
 class MCEvent(object, six.with_metaclass(ABCMeta)):
     """The basis of interaction between user and all the event generators.
 
@@ -47,8 +69,10 @@ class MCEvent(object, six.with_metaclass(ABCMeta)):
         vt (np.array)      : temporal order of vertex in ps, fs?
     """
     __sliced_params__ = [
-        'p_ids', 'status', 'px', 'py', 'pz', 'en', 'm', 'vx', 'vy', 'vz', 'vt',
-        'pem_arr', 'vt_arr'
+        'p_ids', 'status', 'charge',
+        'px', 'py', 'pz', 'en', 'm',
+        'vx', 'vy', 'vz', 'vt',
+        '_pem_arr', '_vt_arr'
     ]
 
     def __init__(self, lib, event_kinematics, event_frame, nevent, npart,
@@ -74,26 +98,81 @@ def __init__(self, lib, event_kinematics, event_frame, nevent, npart,
         self.vt = vt
 
         # Full arrays of kinematical vectors
-        self.pem_arr = pem_arr  # (px, py, pz, E, m)
-        self.vt_arr = vt_arr  # (vx, vy, vz, t)
+        self._pem_arr = pem_arr  # (px, py, pz, E, m)
+        self._vt_arr = vt_arr  # (vx, vy, vz, t)
 
         # Initialize current selection to all entries up to npart
         if impy_config['pre_slice']:
             info(10, 'Pre-slice enabled.')
             self.selection = slice(None, self.npart)
+            self.charge = self._charge_init
             self._apply_slicing()
         else: 
             info(10, 'Pre-slice disabled.')
             self.selection = slice(None, None)
+            self.charge = self._charge_init
 
         # The default slice only cuts limits the view to the array to
         # to the current number of entries
         self._is_filtered = False
+        self._views_cache = []
 
         # Apply boosts into frame required by user
         self.kin.apply_boost(self, event_frame, impy_config["user_frame"])
         self.event_frame = impy_config["user_frame"]
 
+
+    def filter_final(self):
+        self.selection = (self.status == 1)
+        self._cache_current_view()
+        self._apply_slicing()
+        return self
+
+    def filter_charged(self):
+        self.selection = (self.charge != 0)
+        self._cache_current_view()
+        self._apply_slicing()
+        return self
+
+    def _cache_current_view(self):
+        """Cache the current view of data"""
+        self._views_cache.append(EventData(
+            self._npart, self.p_ids, self.status, self.charge,
+            self.px, self.py, self.pz, self.en, self.m,
+            self.vx, self.vy, self.vz, self.vt,
+            self._pem_arr, self._vt_arr
+        ))
+
+    def remove_last_filter(self):
+        """Remove the last applied filter"""
+        # if list empty then return
+        if not self._views_cache:
+            return
+        last = self._views_cache.pop()
+        self.npart = last.npart
+        self.p_ids = last.p_ids
+        self.status = last.status
+        self.charge = last.charge
+        self.px = last.px
+        self.py = last.py
+        self.pz = last.pz
+        self.en = last.en
+        self.m = last.m
+        self.vx = last.vx
+        self.vy = last.vy
+        self.vz = last.vz
+        self.vt = last.vt
+        self._pem_arr = last._pem_arr
+        self._vt_arr = last._vt_arr
+
+
+    def remove_all_filters(self):
+        """Remove all filters (return to original data)"""
+        del self._views_cache[1:]
+        self.remove_last_filter()
+        self._is_filtered = False
+
+
     def _apply_slicing(self):
         """Slices/copies the all varaibles according to filter criteria"""
         info(30, 'Slicing attributes.')
@@ -126,7 +205,7 @@ def filter_final_state_charged(self):
         pass
 
     @abstractproperty
-    def charge(self):
+    def _charge_init(self):
         """Electrical charge"""
         pass
 
@@ -335,6 +414,9 @@ def __init__(self, interaction_model_def, settings_dict=dict(), **kwargs):
 
         # FORTRAN LUN that keeps logfile handle
         self.output_lun = None
+
+        # Number of generated events so far
+        self.nevents = 0
 
     def __enter__(self):
         """TEMP: It would be good to actually use the with construct to
@@ -373,7 +455,7 @@ def init_generator(self, event_kinematics, seed='random'):
         
         The maximal energy and particle masses from the event_kinematics
         object define the maximal range, i.e. the energy requested in subsequent
-        `set_event_kinematics` calls should not exceed the one provided here.
+        `_set_event_kinematics` calls should not exceed the one provided here.
 
         Args:
             event_kinematics (object): maximal energy and masses for subsequent runs
@@ -392,7 +474,7 @@ def generate_event(self):
         pass
 
     @abstractmethod
-    def set_event_kinematics(self, evtkin):
+    def _set_event_kinematics(self, evtkin):
         """Set new combination of energy, momentum, projectile
         and target combination for next event.
 
@@ -404,6 +486,14 @@ def set_event_kinematics(self, evtkin):
         """
         pass
 
+    @property
+    def event_kinematics(self):
+        return self._curr_event_kin
+
+    @event_kinematics.setter
+    def event_kinematics(self, evtkin):
+        self._set_event_kinematics(evtkin)
+
     @abstractmethod
     def set_stable(self, pdgid, stable=True):
         """Prevent decay of unstable particles
@@ -455,11 +545,11 @@ def sigma_inel_air(self, **kwargs):
                 k = EventKinematics(ecm=prev_kin.ecm,
                     p1pdg=prev_kin.p1pdg, 
                     nuc2_prop=(iat, int(iat/2)))
-            self.set_event_kinematics(k)
+            self._set_event_kinematics(k)
             cs += f * self.sigma_inel(**kwargs)
 
         # Restore settings
-        self.set_event_kinematics(prev_kin)
+        self._set_event_kinematics(prev_kin)
 
         return cs
 
@@ -535,6 +625,31 @@ def _define_default_fs_particles(self):
         if impy_config['pi0_stable']:
             self.set_stable(111)
 
+    def __call__(self, nevents):
+        """Generator function (in python sence) 
+        which launches the underlying event generator 
+        and returns its the result (event) as MCEvent object
+        """
+        retry_on_rejection = impy_config['retry_on_rejection']
+        # Initialize counters to prevent infinite loops in rejections
+        ntrials = 0
+        nremaining = nevents
+        while nremaining > 0:
+            if self.generate_event() == 0:
+                self.nevents += 1
+                yield self._event_class(self.lib, self._curr_event_kin,
+                                        self._output_frame)
+                nremaining -= 1
+                ntrials += 1
+            elif retry_on_rejection:
+                info(10, 'Rejection occured. Retrying..')
+                ntrials += 1
+                continue
+            elif ntrials > 2 * nevents:
+                raise Exception('Things run bad. Check your input.')
+            else:
+                info(0, 'Rejection occured')
+
     def event_generator(self, event_kinematics, nevents):
         """This is some kind of equivalent to Hans'
         generator concept.
@@ -544,7 +659,7 @@ def event_generator(self, event_kinematics, nevents):
         initialization issue something has to keep track of ownership
         and history. And classes seem to just this.
         """
-        self.set_event_kinematics(event_kinematics)
+        self._set_event_kinematics(event_kinematics)
         retry_on_rejection = impy_config['retry_on_rejection']
         # Initialize counters to prevent infinite loops in rejections
         ntrials = 0
@@ -560,6 +675,6 @@ def event_generator(self, event_kinematics, nevents):
                 ntrials += 1
                 continue
             elif ntrials > 2 * nevents:
-                raise Exception('Things run bad. Check your input.')
+                raise RuntimeError('Things run bad. Check your input.')
             else:
                 info(0, 'Rejection occured')

impy/constants.py

@@ -1,10 +1,10 @@
 # This dependency might be overkill for just reading a few
 # variables. Should be changed at some point.
-import scipy.constants as spc
+from scipy.constants import speed_of_light as _c_SI
 
-c = 1e2 * spc.c
-cm2sec = 1e-2 / spc.c
-sec2cm = spc.c * 1e2
+c = 1e2 * _c_SI
+cm2sec = 1e-2 / _c_SI
+sec2cm = _c_SI * 1e2
 eV = 1e-9
 keV = 1e-6
 MeV = 1e-3

impy/models/__init__.py

@@ -0,0 +1,10 @@
+from impy.models.sophia import Sophia20
+from impy.models.sibyll import Sibyll21, Sibyll23, Sibyll23c, Sibyll23c00, Sibyll23c01,\
+    Sibyll23c02, Sibyll23c03, Sibyll23c04, Sibyll23d
+from impy.models.dpmjetIII import DpmjetIII191, DpmjetIII192, DpmjetIII306
+from impy.models.epos import EposLHC
+from impy.models.phojet import Phojet112, Phojet191
+from impy.models.urqmd import UrQMD34
+from impy.models.qgsjet import QGSJet01c, QGSJetII03, QGSJetII04
+from impy.models.pythia6 import Pyphia6
+from impy.models.pythia8 import Pyphia8

impy/models/dpmjetII.py

@@ -24,6 +24,8 @@ def __init__(self, lib, event_config):
 
         if 'charge_info' in event_config and event_config['charge_info']:
             self.charge = lib.extevt.idch[sel]
+            # This is a patch to allow working with all other generators
+            self._charge_init = self.charge
 
         self.p_ids = evt.idhkk[sel]
         self.pt2 = evt.phkk[0, sel]**2 + evt.phkk[1, sel]**2
@@ -52,7 +54,7 @@ def get_sigma_inel(self):
         k = self.evkin
         return self.lib.siinel(self.lib.mcihad(k.p1pdg), 1, k.ecm)
 
-    def set_event_kinematics(self, event_kinematics):
+    def _set_event_kinematics(self, event_kinematics):
         k = event_kinematics
 
         self.dpmevt_tup = k.elab, k.A1, k.Z1, k.A2, k.Z2, \
@@ -62,7 +64,7 @@ def set_event_kinematics(self, event_kinematics):
 
     def init_generator(self, config):
         # Comprise DPMJET input from the event kinematics object
-        self.set_event_kinematics(self.evkin)
+        self._set_event_kinematics(self.evkin)
         k = self.evkin
 
         try:

impy/models/dpmjetIII.py

@@ -58,7 +58,7 @@ def children(self):
         return self.lib.dtevt1.jdahkk
 
     @property
-    def charge(self):
+    def _charge_init(self):
         return self.lib.dtpart.iich[self.lib.dtevt2.idbam[self.selection] - 1]
 
     # Nuclear collision parameters
@@ -127,7 +127,7 @@ def _dpmjet_tup(self):
              (k.A1, k.Z1, k.A2, k.Z2, self.lib.idt_icihad(k.p1pdg), k.elab))
         return (k.A1, k.Z1, k.A2, k.Z2, self.lib.idt_icihad(k.p1pdg), k.elab)
 
-    def set_event_kinematics(self, event_kinematics):
+    def _set_event_kinematics(self, event_kinematics):
         """Set new combination of energy, momentum, projectile
         and target combination for next event."""
         info(5, 'Setting event kinematics')
@@ -141,7 +141,7 @@ def set_event_kinematics(self, event_kinematics):
 
         # AF: No idea yet, but apparently this functionality was around?!
         # if hasattr(k, 'beam') and hasattr(self.lib, 'init'):
-        #     print(self.class_name + "::set_event_kinematics():" +
+        #     print(self.class_name + "::_set_event_kinematics():" +
         #           "setting beam params", k.beam)
         #     self.lib.dt_setbm(k.A1, k.Z1, k.A2, k.Z2, k.beam[0], k.beam[1])
         #     print 'OK'
@@ -185,7 +185,7 @@ def init_generator(self, event_kinematics, seed='random', logfname=None):
         self._max_A1 = event_kinematics.A1
         self._max_A2 = event_kinematics.A2
 
-        self.set_event_kinematics(event_kinematics)
+        self._set_event_kinematics(event_kinematics)
         k = self._curr_event_kin
         dpm_conf = impy_config['dpmjetIII']
 
@@ -272,3 +272,25 @@ def generate_event(self):
         reject = self.lib.dt_kkinc(*self._dpmjet_tup(), kkmat=-1)
         self.lib.dtevno.nevent += 1
         return reject
+
+class DpmjetIII191(DpmjetIIIRun):
+    def __init__(self, event_kinematics, seed="random", logfname=None):
+        from impy.definitions import interaction_model_by_tag as models_dict
+        interaction_model_def = models_dict["DPMJETIII191"]       
+        super().__init__(interaction_model_def)
+        self.init_generator(event_kinematics, seed, logfname)
+
+class DpmjetIII192(DpmjetIIIRun):
+    def __init__(self, event_kinematics, seed="random", logfname=None):
+        from impy.definitions import interaction_model_by_tag as models_dict
+        interaction_model_def = models_dict["DPMJETIII192"]       
+        super().__init__(interaction_model_def)
+        self.init_generator(event_kinematics, seed, logfname)
+
+class DpmjetIII306(DpmjetIIIRun):
+    def __init__(self, event_kinematics, seed="random", logfname=None):
+        from impy.definitions import interaction_model_by_tag as models_dict
+        interaction_model_def = models_dict["DPMJETIII306"]       
+        super().__init__(interaction_model_def)
+        self.init_generator(event_kinematics, seed, logfname)                
+