Python 3.12 and LST Prod6 MCs multi-telescope support

pyproject.toml

@@ -20,12 +20,12 @@ classifiers = [
     "Operating System :: OS Independent",
 ]
 dependencies = [
-    "astropy==5.1",
+    "astropy>=5.1",
     "numpy>=1.21.0",
-    "pandas==1.4.1",
+    "pandas>=1.4.1",
     "scipy>=1.5.0",
-    "tables==3.7.0",
-    "PyYAML==5.3.1"
+    "tables>=3.7.0",
+    "PyYAML>=5.3.1"
 ]
 
 [project.urls]

src/srcsim/mc.py

@@ -11,7 +11,101 @@ def power_law(e, e0, norm, index):
     return norm * (e/e0).decompose()**index
 
 
-class MCSample:
+class MCBase:
+    @classmethod
+    def _has_key(cls, file_name, key):
+        with tables.open_file(file_name) as table:
+            has_key = key in table
+
+        return has_key
+
+    @classmethod
+    def _choose_first_valid_key(cls, file_name, keys):
+        for key in keys:
+            if cls._has_key(file_name, key):
+                return key
+
+        return None
+
+    @classmethod
+    def get_config_key(cls, file_name):
+        keys = (
+            '/simulation/config',
+            '/simulation/run_config'
+        )
+        key = cls._choose_first_valid_key(file_name, keys)
+        return key
+
+    @classmethod
+    def get_events_key(cls, file_name):
+        keys = (
+            '/events/parameters',
+            '/dl1/event/telescope/parameters/LST_LSTCam',
+            '/dl2/event/telescope/parameters/LST_LSTCam'
+        )
+        key = cls._choose_first_valid_key(file_name, keys)
+        return key
+
+    @classmethod
+    def read_config(cls, file_name):
+        with tables.open_file(file_name) as table:
+            cfg_table = table.root[cls.get_config_key(file_name)]
+
+            columns = {
+                'n_showers': ('num_showers',),
+                'shower_reuse': (),
+                'min_scatter_range': (),
+                'max_scatter_range': (),
+                'energy_range_min': (),
+                'energy_range_max': (),
+                'spectral_index': (),
+                'min_viewcone_radius': (),
+                'max_viewcone_radius': ()
+            }
+
+            data = {}
+
+            for col_name in columns:
+                if col_name in cfg_table.colnames:
+                    data[col_name] = [
+                        v[col_name] for v in cfg_table.iterrows()
+                    ]
+                else:
+                    for alternative in columns[col_name]:
+                        if alternative in cfg_table.colnames:
+                            data[col_name] = [
+                                v[alternative] for v in cfg_table.iterrows()
+                            ]
+                            break
+                if col_name not in data:
+                    raise RuntimeError(f"could not load config key {col_name} from '{file_name}'")
+
+            if 'obs_id' in cfg_table.colnames:
+                data['obs_id'] = [
+                    v['obs_id'] for v in cfg_table.iterrows()
+                ]
+            else:
+                evt_table = table.root[cls.get_events_key(file_name)]
+                row = next(evt_table.iterrows())
+                data['obs_id'] = [
+                    row['obs_id'] for _ in cfg_table.iterrows()
+                ]
+                print(
+                    "WARN: can not find 'obs_id' in the configuration table, "
+                    f"assuming {row['obs_id']} from the first event. "
+                    "Simulation results may be incorrect."
+                )
+
+        return pd.DataFrame(data=data)
+
+    @classmethod
+    def read_data(cls, file_name):
+        data = pd.read_hdf(file_name, cls.get_events_key(file_name))
+
+        return data
+
+
+class MCSample(MCBase):
     def __init__(self, file_name=None, obs_id=None, data_table=None, config_table=None):
         self.units = dict(
             energy = u.TeV,
@@ -34,29 +128,34 @@ def __init__(self, file_name=None, obs_id=None, data_table=None, config_table=No
             self.file_name = file_name
             self.obs_id = obs_id
             self.config_table = self.read_config(file_name, obs_id)
-            self.data_table = pd.read_hdf(file_name, 'dl2/event/telescope/parameters/LST_LSTCam').query(f'obs_id == {obs_id}')
+            self.data_table = self.read_data(file_name, obs_id)
+
+        self.data_table = self._fix_format(self.data_table)
+        self.config_table = self._fix_format(self.config_table)
 
         # Getting the telescope pointing
-        pointing_data = self.data_table[['mc_az_tel', 'mc_alt_tel']].mean()
-        self.tel_pos = SkyCoord(pointing_data['mc_az_tel'], pointing_data['mc_alt_tel'], unit=self.units['angle'], frame='altaz')
+        pointing_data = self.data_table[['pointing_az', 'pointing_alt']].mean()
+        self.tel_pos = SkyCoord(pointing_data['pointing_az'], pointing_data['pointing_alt'], unit=self.units['angle'], frame='altaz')
 
         # Working out the simulation spectrum
-        rmin, rmax = self.config_table[['min_scatter_range', 'max_scatter_range']].iloc[0] * self.units['distance']
+        rmin, rmax = self.config_table[['min_scatter_range', 'max_scatter_range']].iloc[0].values * self.units['distance']
         ground_area = np.pi * (rmax**2 - rmin**2)
-        nevents = self.config_table['num_showers'].iloc[0] * self.config_table['shower_reuse'].iloc[0]
+        nevents = self.config_table['n_showers'].iloc[0] * self.config_table['shower_reuse'].iloc[0]
         emin = self.config_table['energy_range_min'].iloc[0] * self.units['energy']
         emax = self.config_table['energy_range_max'].iloc[0] * self.units['energy']
         index = self.config_table['spectral_index'].iloc[0]
         self.spec_data = self.get_spec_data(nevents, emin, emax, index)
         self.spec_data['norm'] /= ground_area
 
-        cam_x, cam_y = self.data_table[['src_x', 'src_y']].to_numpy().transpose() * self.units['distance'] * self.cam2angle
-        self.evt_coord = SkyCoord(cam_x, cam_y, frame=self.tel_pos.skyoffset_frame())
+        alt, az = self.data_table[['true_alt', 'true_az']].to_numpy().transpose()
+        _unit = 'rad' if 'mc_alt' in self.data_table.columns else 'deg'
+        self.evt_coord = SkyCoord(az, alt, frame='altaz', unit=_unit)
+        self.evt_coord = self.evt_coord.transform_to(self.tel_pos.skyoffset_frame())
 
-        self.evt_energy = self.data_table['mc_energy'].to_numpy() * self.units['energy']
+        self.evt_energy = self.data_table['true_energy'].to_numpy() * self.units['energy']
 
         # Filtering out events with excessive offsets (e.g. due to the simulation numerical accuracy)
-        offset_min, offset_max = self.config_table[['min_viewcone_radius', 'max_viewcone_radius']].iloc[0] * self.units['viewcone']
+        offset_min, offset_max = self.config_table[['min_viewcone_radius', 'max_viewcone_radius']].iloc[0].values * self.units['viewcone']
         evt_offset = self.evt_coord.separation(self.tel_pos)
 
         in_fov = (evt_offset >= offset_min) & (evt_offset <= offset_max)
@@ -77,35 +176,32 @@ def __repr__(self):
         )
 
         return super().__repr__()
-    
+
     @classmethod
     def read_config(cls, file_name, obs_id):
-        with tables.open_file(file_name) as table:
-            cfg_table = table.root['/simulation/run_config']
-            obs_ids = [v['obs_id'] for v in cfg_table.iterrows()]
-
-            columns = (
-                'obs_id',
-                'num_showers',
-                'shower_reuse',
-                'min_scatter_range',
-                'max_scatter_range',
-                'energy_range_min',
-                'energy_range_max',
-                'spectral_index',
-                'min_viewcone_radius',
-                'max_viewcone_radius'
-            )
-
-            obs_idx = obs_ids.index(obs_id)
-
-            data = {}
-
-            for col_name in columns:
-                col_idx = cfg_table.colnames.index(col_name)
-                data[col_name] = (cfg_table[obs_idx][col_idx], )
+        config = super().read_config(file_name)
+        return config.query(f'obs_id == {obs_id}')
+
+    @classmethod
+    def read_data(cls, file_name, obs_id):
+        data = super().read_data(file_name)
+        return data.query(f'obs_id == {obs_id}')
+
+    def _fix_format(self, table):
+        table = table.copy()
+        old2new = dict(
+            mc_energy = 'true_energy',
+            mc_az = 'true_az',
+            mc_alt = 'true_alt',
+            mc_az_tel = 'pointing_az',
+            mc_alt_tel = 'pointing_alt',
+            num_showers = 'n_showers'
+        )
+        for key in old2new:
+            if key in table.columns:
+                table[old2new[key]] = table[key]
 
-            return pd.DataFrame(data=data)
+        return table
 
     def get_spec_data(self, n_events, emin, emax, index=-1):
         e0 = (emin * emax)**0.5
@@ -129,7 +225,7 @@ def dnde(self, energy):
         return power_law(energy, **self.spec_data)
 
     def dndo(self, coord):
-        offset_min, offset_max = self.config_table[['min_viewcone_radius', 'max_viewcone_radius']].iloc[0] * self.units['viewcone']
+        offset_min, offset_max = self.config_table[['min_viewcone_radius', 'max_viewcone_radius']].iloc[0].values * self.units['viewcone']
         sky_area = 2 * np.pi * (np.cos(offset_min) - np.cos(offset_max)) * u.sr
         norm = 1 / sky_area
 
@@ -141,7 +237,7 @@ def dndedo(self, energy, coord):
         return self.dnde(energy) * self.dndo(coord)
 
 
-class MCCollection:
+class MCCollection(MCBase):
     def __init__(self, file_mask=None, samples=None, log=None):
         self.file_mask = file_mask
 
@@ -167,44 +263,15 @@ def __repr__(self):
 
     @classmethod
     def read_obs_ids(cls, file_name):
-        with tables.open_file(file_name, 'r') as data:
-            cfg_table = data.root['/simulation/run_config']
-            obs_ids = [v['obs_id'] for v in cfg_table.iterrows()]
+        obs_ids = cls.read_config(file_name)['obs_id'].values
 
         return obs_ids
 
-    @classmethod
-    def read_config(cls, file_name):
-        with tables.open_file(file_name) as table:
-            cfg_table = table.root['/simulation/run_config']
-
-            columns = (
-                'obs_id',
-                'num_showers',
-                'shower_reuse',
-                'min_scatter_range',
-                'max_scatter_range',
-                'energy_range_min',
-                'energy_range_max',
-                'spectral_index',
-                'min_viewcone_radius',
-                'max_viewcone_radius'
-            )
-
-            data = {}
-
-            for col_name in columns:
-                data[col_name] = [
-                    v[col_name] for v in cfg_table.iterrows()
-                ]
-
-        return pd.DataFrame(data=data)
-
     @classmethod
     def read_file(cls, file_name):
         obs_ids = cls.read_obs_ids(file_name)
 
-        data = pd.read_hdf(file_name, "dl2/event/telescope/parameters/LST_LSTCam")
+        data = cls.read_data(file_name)
         config = cls.read_config(file_name)
 
         samples = tuple(
@@ -228,13 +295,15 @@ def read_files(cls, file_mask):
         return samples
 
     def get_closest(self, target_position):
+        target_position = SkyCoord(target_position.altaz.az, target_position.altaz.alt, frame='altaz')
         tel_pos = SkyCoord([sample.tel_pos for sample in self.samples])
         separation = tel_pos.separation(target_position)
         idx = separation.argmin()
 
         return MCCollection(samples=(self.samples[idx],))
 
     def get_nearby(self, target_position, search_radius):
+        target_position = SkyCoord(target_position.altaz.az, target_position.altaz.alt, frame='altaz')
         samples = tuple(
             filter(
                 lambda sample: sample.tel_pos.separation(target_position) <= search_radius,
@@ -250,8 +319,8 @@ def get_nearby(self, target_position, search_radius):
         return MCCollection(samples=samples)
 
     def get_in_box(self, target_position, max_lon_offset, max_lat_offset):
-        tel_pos = SkyCoord([sample.tel_pos for sample in self.samples])
         target_position = SkyCoord(target_position.altaz.az, target_position.altaz.alt, frame='altaz')
+        tel_pos = SkyCoord([sample.tel_pos for sample in self.samples])
 
         lon_offset, lat_offset = tel_pos.altaz.spherical_offsets_to(target_position.altaz)
         inbox = (np.absolute(lon_offset) <= max_lon_offset) & (np.absolute(lat_offset) <= max_lat_offset)