pspipe 0.7.2__tar.gz → 0.8__tar.gz

{pspipe-0.7.2 → pspipe-0.8}/PKG-INFO

@@ -1,7 +1,8 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: pspipe
-Version: 0.7.2
+Version: 0.8
 Summary: Power spectra pipeline for Cosmic Dawn, Epoch of Reionization radio interferometric experiments
+License-File: LICENSE
 Author: Florent Mertens
 Author-email: florent.mertens@gmail.com
 Requires-Python: >=3.10,<4.0
@@ -10,10 +11,11 @@ Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
 Provides-Extra: ml-gpr
 Requires-Dist: astropy (>6.0)
 Requires-Dist: click (>=8.0)
-Requires-Dist: libpipe (>=0.1.5)
+Requires-Dist: libpipe (>=0.6)
 Requires-Dist: ps-eor (>=0.8)
 Requires-Dist: ps-eor[ml-gpr] (>=0.8) ; extra == "ml-gpr"
 Requires-Dist: python-casacore (>=3.0)

{pspipe-0.7.2 → pspipe-0.8}/pspipe/database.py

@@ -3,6 +3,7 @@
 # Author: F. Mertens
 
 import os
+import hashlib
 import glob
 import json
 import fnmatch
@@ -137,6 +138,41 @@ class ImageRevision(object):
             self.settings.image.name = img_rev_name
         self.name = self.settings.image.name
 
+    def get_distributed_data_root(self, obs_id, persist=True):
+        """
+        Return the root directory used for FITS products for this obs_id.
+
+        - Reads from global meta (img_list.json) if present.
+        - Otherwise computes deterministic root from settings.distributed_data_dirs
+          (or falls back to settings.data_dir).
+        - If persist=True, writes the chosen root into global meta.
+        """
+        meta = self.get_meta(obs_id, lst_idx=None)
+
+        if meta is not None:
+            root = meta.get('data_dir_root')
+            if root:
+                return root
+
+        dirs = self.settings.get('distributed_data_dirs')
+        if dirs:
+            dirs = list(dirs)
+            h = int(hashlib.sha1(obs_id.encode('utf-8')).hexdigest(), 16)
+            root = dirs[h % len(dirs)]
+        else:
+            root = self.settings.get_path('data_dir')
+
+        if not persist:
+            return root
+
+        # Persist in global meta
+        if meta is None:
+            meta = self.new_obs_id(obs_id, lst_idx=None)  # create empty global meta
+
+        meta['data_dir_root'] = root
+        meta.save()
+        return root
+
     def create(self):
         os.makedirs(self.get_path(), exist_ok=True)
         self.settings.duplicate(self.get_settings_path())
@@ -185,6 +221,10 @@ class ImageRevision(object):
     def get_image_path(self, obs_id):
         return os.path.join(self.get_path(), 'images', obs_id)
 
+    def get_image_data_path(self, obs_id, persist=True):
+        root = self.get_distributed_data_root(obs_id, persist=persist)
+        return os.path.join(root, self.settings.image.name, 'images', obs_id)
+
     def get_meta_path(self, obs_id, lst_idx=None):
         lst_str = f'_t{int(lst_idx):03}' if lst_idx is not None else ''
         return os.path.join(self.get_image_path(obs_id), f'img_list{lst_str}.json')
@@ -205,8 +245,9 @@ class ImageRevision(object):
     def has_obs_id(self, obs_id):
         return obs_id in self.list_obs_ids()
 
-    def new_obs_id(self, obs_id, lst_idx=None):
-        img_meta = ImageRevisionMeta.new(self.get_meta_path(obs_id, lst_idx=lst_idx), obs_id, self.settings.image.name)
+    def new_obs_id(self, obs_id, lst_idx=None, data_dir_root=None):
+        img_meta = ImageRevisionMeta.new(self.get_meta_path(obs_id, lst_idx=lst_idx), obs_id, self.settings.image.name,
+                                         data_dir_root=data_dir_root)
         os.makedirs(os.path.dirname(self.get_meta_path(obs_id, lst_idx=lst_idx)), exist_ok=True)
         img_meta.save()
 
@@ -381,8 +422,11 @@ class ObsIDMeta(AbstractJsonMeta):
 class ImageRevisionMeta(AbstractJsonMeta):
 
     @staticmethod
-    def new(file, obs_id, img_rev):
-        return ImageRevisionMeta(file, {'obs_id': obs_id, 'img_rev': img_rev})
+    def new(file, obs_id, img_rev, data_dir_root=None):
+        d = {'obs_id': obs_id, 'img_rev': img_rev}
+        if data_dir_root:
+            d['data_dir_root'] = data_dir_root
+        return ImageRevisionMeta(file, d)
 
     def get_fits(self, stokes_or_psf, even=False, odd=False):
         assert (even != odd) or (even is False)

pspipe-0.8/pspipe/obsdata.py

@@ -0,0 +1,475 @@
+# Direct access to the data
+#
+# Author: F. Mertens
+
+import os
+
+import numpy as np
+
+from ps_eor import pspec, datacube, flagger
+
+from . import database, utils
+
+# Lazy loaded:
+# from ps_eor import ml_gpr
+
+
+class ObsData(object):
+
+    def __init__(self, vis_rev, obs_id, lst_idx=None):
+        self.vis_rev = vis_rev
+        self.obs_id = obs_id
+        self.lst_idx = lst_idx
+        self.settings = self.vis_rev.settings
+        self.flag = None
+        self.uvrange = None
+        self.nfreqs = None
+        self.data = {}
+        self.all_stokes = {}
+        for stokes in database.all_stokes:
+            self.all_stokes[stokes.lower()] = stokes
+            self.all_stokes[stokes.lower() + '_even'] = 'even_%s' % stokes
+            self.all_stokes[stokes.lower() + '_odd'] = 'odd_%s' % stokes
+            self.all_stokes[stokes.lower() + '_dt'] = 'dt_%s' % stokes
+        self._ps_builder = None
+
+    def __getattr__(self, key):
+        assert key in self.all_stokes.keys()
+        return self.get_stokes(self.all_stokes[key])
+
+    def _invalidate_cache(self):
+        self.data = {}
+
+    @property
+    def ps_builder(self):
+        """Power-spectrum builder, loaded on first use."""
+        if self._ps_builder is None:
+            self._ps_builder = pspec.PowerSpectraBuilder(
+                self.settings.power_spectra.get_path('ps_config'),
+                self.settings.power_spectra.get_path('eor_bin_list'))
+        return self._ps_builder
+
+    @ps_builder.setter
+    def ps_builder(self, builder):
+        self._ps_builder = builder
+
+    def _get_first_stokes(self, candidates):
+        for stokes in candidates:
+            if self.has_stokes(stokes):
+                return stokes, self.get_stokes(stokes)
+        return None, None
+
+    def get_stokes_i_like(self):
+        return self._get_first_stokes(['I', 'XX', 'YY'])[1]
+
+    def get_stokes_v_like(self):
+        return self._get_first_stokes(
+            ['V', 'XY', 'YX', 'XX', 'YY', 'I'])[1]
+
+    @property
+    def mjd_mid(self):
+        """Midpoint of the observation in MJD."""
+        cube = self.get_stokes_i_like()
+        if cube is None:
+            raise ValueError(f'No Stokes-I-like data available for {self.obs_id}')
+        return cube.meta.obs_mjd + cube.meta.total_time / 172800.
+
+    @property
+    def lst(self):
+        """Mean local sidereal time of the observation in hours."""
+        return utils.get_lst(self.mjd_mid, self.settings.longitude)
+
+    def set_flag(self, flag):
+        self.flag = flag
+        self._invalidate_cache()
+
+    def filter_uvrange(self, umin, umax):
+        self.uvrange = [umin, umax]
+        self._invalidate_cache()
+
+    def average_freqs(self, nfreqs):
+        self.nfreqs = nfreqs
+        self._invalidate_cache()
+
+    def do_flag(self, flagger_file=None, verbose_flagging=True,
+                plot_flagging=True, flagger_runner=None):
+        if flagger_runner is None:
+            if flagger_file is None:
+                flagger_file = self.settings.power_spectra.get_path('flagger')
+            flagger_runner = flagger.FlaggerRunner.load(flagger_file)
+
+        flagger_runner.verbose = verbose_flagging
+        i_stokes, i_cube = self._get_first_stokes(['I', 'XX', 'YY'])
+        filtered_i, _ = flagger_runner.run(
+            i_cube, self.get_stokes_v_like())
+        if plot_flagging:
+            flagger_runner.plot()
+
+        self.flag = flagger_runner.flag
+        self._invalidate_cache()
+        self.data[i_stokes] = filtered_i
+
+    def get_h5_file(self, stokes):
+        return self.vis_rev.get_h5_file(self.obs_id, stokes, lst_idx=self.lst_idx)
+
+    def get_stokes(self, stokes):
+        assert self.has_stokes(stokes)
+
+        if stokes not in self.data:
+            path = self.get_h5_file(stokes)
+            cube = datacube.CartDataCube.load(path)
+            if self.uvrange is not None:
+                cube.filter_uvrange(self.uvrange[0], self.uvrange[1])
+            if self.nfreqs is not None:
+                cube = cube.average_freqs(self.nfreqs)
+            if self.flag is not None:
+                cube = self.flag.apply(cube)
+            self.data[stokes] = cube
+        return self.data[stokes]
+
+    def has_stokes(self, stokes):
+        return stokes in self.data or os.path.exists(self.get_h5_file(stokes))
+
+    def get_ml_gpr_path(self, eor_bin, umin=50, umax=250, name='', sampler_method='mcmc'):
+        return os.path.join(os.path.dirname(self.get_h5_file('I')), name, f'eor{eor_bin}_u{umin}-{umax}')
+
+    def get_ml_gpr_res(self, eor_bin, umin=50, umax=250, name='', sampler_method='mcmc'):
+        from ps_eor import ml_gpr
+
+        path = self.get_ml_gpr_path(eor_bin, umin=umin, umax=umax, name=name, sampler_method=sampler_method)
+        return ml_gpr.MLGPRResult.load(path, 'ml_gpr_res_I', sampler_method=sampler_method)
+
+    def get_ml_gpr_inj_path(self, eor_bin, var_21_over_noise, x1, x2, umin=50, umax=250, name=''):
+        d = f'inj_eor{eor_bin}_u{umin}-{umax}_a{var_21_over_noise:.3f}_x1{x1:.3f}_x2{x2:.3f}'
+        return os.path.join(os.path.dirname(self.get_h5_file('I')), name, d)
+
+    def get_ml_gpr_inj_res(self, eor_bin, var_21_over_noise, x1, x2, umin=50, umax=250, name='', sampler_method='mcmc'):
+        from ps_eor import ml_gpr
+
+        path = self.get_ml_gpr_inj_path(eor_bin, var_21_over_noise=var_21_over_noise, x1=x1, x2=x2, umin=umin, umax=umax, name=name)
+        data_inj = datacube.CartDataCube.load(os.path.join(path, 'c_injected.h5'))
+        return ml_gpr.MLGPRInjResult.load(path, 'ml_gpr_res_I', data_inj, sampler_method=sampler_method)
+
+    def get_ps_gen(self, **kargs):
+        return self.ps_builder.get(self.get_stokes_i_like(), **kargs)
+
+
+class ObsDataGroup(object):
+    """A group of observations exposed through a combined ObsData-like view."""
+
+    def __init__(self, observations, name, umin, umax, weighting_mode='full',
+                 observation_weights=None, index=None, lst_min=None, lst_max=None):
+        self.observations = tuple(observations)
+        if not self.observations:
+            raise ValueError('observations must not be empty')
+
+        self.name = name
+        self.obs_id = name
+        self.index = index
+        self.lst_min = lst_min
+        self.lst_max = lst_max
+        self.umin = umin
+        self.umax = umax
+        self.weighting_mode = weighting_mode
+        self._stokes_cache = {}
+
+        obs_ids = [obs.obs_id for obs in self.observations]
+        if len(set(obs_ids)) != len(obs_ids):
+            raise ValueError('observations must have unique obs_id values')
+
+        if observation_weights is None:
+            self.observation_weights = dict.fromkeys(obs_ids, 1.)
+        else:
+            try:
+                self.observation_weights = {
+                    obs_id: float(weight)
+                    for obs_id, weight in dict(observation_weights).items()}
+            except (TypeError, ValueError):
+                raise ValueError(
+                    'observation_weights must contain real scalar values')
+            missing = set(obs_ids) - set(self.observation_weights)
+            unknown = set(self.observation_weights) - set(obs_ids)
+            if missing or unknown:
+                raise ValueError(
+                    'observation_weights must contain exactly the group obs_ids')
+        if any(
+                not np.isfinite(weight) or weight < 0
+                for weight in self.observation_weights.values()):
+            raise ValueError(
+                'observation_weights must be finite and non-negative')
+
+        self.all_stokes = dict(self.observations[0].all_stokes)
+
+    def __len__(self):
+        return len(self.observations)
+
+    def __getattr__(self, key):
+        all_stokes = self.__dict__.get('all_stokes', {})
+        if key not in all_stokes:
+            raise AttributeError(
+                f'{type(self).__name__} has no attribute {key!r}')
+        return self.get_stokes(all_stokes[key])
+
+    @property
+    def obs_ids(self):
+        return np.array([obs.obs_id for obs in self.observations])
+
+    @property
+    def lst(self):
+        if self.lst_min is not None and self.lst_max is not None:
+            return .5 * (self.lst_min + self.lst_max)
+        return np.mean(self.lsts)
+
+    @property
+    def lsts(self):
+        return np.array([obs.lst for obs in self.observations])
+
+    @property
+    def mjd_mid(self):
+        return np.array([obs.mjd_mid for obs in self.observations])
+
+    @property
+    def total_time(self):
+        return sum(
+            obs.get_stokes_i_like().meta.total_time
+            for obs in self.observations)
+
+    def clear_cache(self):
+        """Discard combined Stokes cubes cached by this group."""
+        self._stokes_cache = {}
+
+    def has_stokes(self, stokes):
+        """Return whether every observation provides ``stokes``."""
+        return all(obs.has_stokes(stokes) for obs in self.observations)
+
+    def get_observation_cubes(self, stokes):
+        """Return the individual cubes for ``stokes`` in group order."""
+        if not self.has_stokes(stokes):
+            raise ValueError(
+                f'Stokes {stokes!r} is not available in every observation')
+        return tuple(obs.get_stokes(stokes) for obs in self.observations)
+
+    def get_stokes(self, stokes):
+        """Return the lazily combined cube for ``stokes``."""
+        if stokes not in self._stokes_cache:
+            combiner = datacube.DataCubeCombiner(
+                self.umin, self.umax, weighting_mode=self.weighting_mode)
+
+            for obs, cube in zip(
+                    self.observations,
+                    self.get_observation_cubes(stokes)):
+                cube = cube.copy()
+                if cube.weights is None:
+                    raise ValueError(
+                        f'Stokes {stokes!r} for {obs.obs_id} has no weights')
+                cube.weights.uv_scale *= self.observation_weights[obs.obs_id]
+                combiner.add(cube, obs.obs_id)
+
+            self._stokes_cache[stokes] = combiner.get()
+
+        return self._stokes_cache[stokes]
+
+    def get_stokes_i_like(self):
+        for stokes in ['I', 'XX', 'YY']:
+            if self.has_stokes(stokes):
+                return self.get_stokes(stokes)
+
+    def get_stokes_v_like(self):
+        for stokes in ['V', 'XY', 'YX', 'XX', 'YY', 'I']:
+            if self.has_stokes(stokes):
+                return self.get_stokes(stokes)
+
+    def get_ps_gen(self, **kargs):
+        return self.observations[0].ps_builder.get(
+            self.get_stokes_i_like(), **kargs)
+
+
+class ObsDataCollection(object):
+    """Observation groups indexed by contiguous LST bins."""
+
+    def __init__(self, lst_bins, groups):
+        self.lst_bins = np.asarray(lst_bins, dtype=float)
+        if self.lst_bins.ndim != 1 or self.lst_bins.size < 2:
+            raise ValueError('lst_bins must be a 1D array with at least two edges')
+        if not np.all(np.diff(self.lst_bins) > 0):
+            raise ValueError('lst_bins must be strictly increasing')
+
+        self.groups = tuple(groups)
+        indices = [group.index for group in self.groups]
+        if any(index is None for index in indices):
+            raise ValueError('collection groups must have a bin index')
+        if len(set(indices)) != len(indices):
+            raise ValueError('collection groups must have unique bin indices')
+        if any(
+                index < 0 or index >= len(self.lst_bins) - 1
+                for index in indices):
+            raise ValueError('collection group index outside lst_bins')
+
+    def __iter__(self):
+        return iter(self.groups)
+
+    def __len__(self):
+        return len(self.groups)
+
+    @property
+    def observations(self):
+        return tuple(
+            obs for group in self.groups for obs in group.observations)
+
+    @property
+    def obs_ids(self):
+        return np.array([obs.obs_id for obs in self.observations])
+
+    def get_bin(self, index):
+        """Return the populated group for bin ``index``, or ``None``."""
+        for group in self.groups:
+            if group.index == index:
+                return group
+        return None
+
+    def select_lst(self, lst):
+        """Return the populated group containing ``lst``, or ``None``."""
+        index = np.searchsorted(self.lst_bins, lst, side='right') - 1
+        if index < 0 or index >= len(self.lst_bins) - 1:
+            return None
+        return self.get_bin(index)
+
+    @classmethod
+    def from_revision(cls, vis_rev, obs_ids, lst_bins, *, umin, umax,
+                      required_stokes=('I', 'V', 'dt_I'),
+                      excluded_obs_ids=None, flagger_parset=None,
+                      weighting_mode='full', observation_weights=None,
+                      verbose=True):
+        """Load observations from a revision and group them by their LST."""
+        lst_bins = np.asarray(lst_bins, dtype=float)
+        if lst_bins.ndim != 1 or lst_bins.size < 2:
+            raise ValueError('lst_bins must be a 1D array with at least two edges')
+        if not np.all(np.diff(lst_bins) > 0):
+            raise ValueError('lst_bins must be strictly increasing')
+
+        excluded_obs_ids = set(
+            () if excluded_obs_ids is None else excluded_obs_ids)
+        grouped = [[] for _ in range(len(lst_bins) - 1)]
+        flagger_runner = None
+
+        for obs_id in obs_ids:
+            if obs_id in excluded_obs_ids:
+                continue
+
+            obs = vis_rev.get_data(obs_id)
+            missing = [
+                stokes for stokes in required_stokes
+                if not obs.has_stokes(stokes)]
+            if missing:
+                if verbose:
+                    print(f'Warning: {obs_id} missing Stokes {missing}')
+                continue
+
+            obs.filter_uvrange(umin, umax)
+            if flagger_parset is not None:
+                if flagger_runner is None:
+                    flagger_runner = flagger.FlaggerRunner.load(
+                        flagger_parset)
+                obs.do_flag(
+                    flagger_parset, False, False,
+                    flagger_runner=flagger_runner)
+
+            index = np.searchsorted(lst_bins, obs.lst, side='right') - 1
+            if 0 <= index < len(grouped):
+                grouped[index].append(obs)
+
+        groups = []
+        for index, observations in enumerate(grouped):
+            if not observations:
+                continue
+
+            weights = None
+            if observation_weights is not None:
+                weights = {
+                    obs.obs_id: observation_weights[obs.obs_id]
+                    for obs in observations}
+
+            groups.append(ObsDataGroup(
+                observations,
+                f'lst_{.5 * (lst_bins[index] + lst_bins[index + 1]):.2f}',
+                umin,
+                umax,
+                weighting_mode=weighting_mode,
+                observation_weights=weights,
+                index=index,
+                lst_min=lst_bins[index],
+                lst_max=lst_bins[index + 1],
+            ))
+
+        return cls(lst_bins, groups)
+
+    @classmethod
+    def from_revision_interleaved(
+            cls, vis_rev, obs_ids, lst_bins, *, umin, umax,
+            required_stokes=('I', 'V', 'dt_I'),
+            excluded_obs_ids=None, flagger_parset=None,
+            weighting_mode='full', observation_weights=None,
+            verbose=True):
+        """Load observations into groups of alternating uniform LST bins.
+
+        The first group contains bins 0, 2, 4, ... and the second contains
+        bins 1, 3, 5, .... Bin parity follows ``lst_bins`` even when a bin is
+        empty. The result is suitable for constructing two disjoint datasets
+        with similar LST coverage.
+
+        Returns:
+            ``(even_group, odd_group)`` as two :class:`ObsDataGroup` objects.
+        """
+        lst_bins = np.asarray(lst_bins, dtype=float)
+        if lst_bins.ndim != 1 or lst_bins.size < 3:
+            raise ValueError(
+                'lst_bins must be a 1D array with at least three edges')
+        widths = np.diff(lst_bins)
+        if not np.all(widths > 0):
+            raise ValueError('lst_bins must be strictly increasing')
+        if not np.allclose(widths, widths[0]):
+            raise ValueError(
+                'interleaved LST groups require uniformly spaced bins')
+
+        collection = cls.from_revision(
+            vis_rev,
+            obs_ids,
+            lst_bins,
+            umin=umin,
+            umax=umax,
+            required_stokes=required_stokes,
+            excluded_obs_ids=excluded_obs_ids,
+            flagger_parset=flagger_parset,
+            weighting_mode=weighting_mode,
+            observation_weights=observation_weights,
+            verbose=verbose,
+        )
+
+        interleaved = []
+        for parity, name in enumerate(('lst_even', 'lst_odd')):
+            source_groups = [
+                group for group in collection.groups
+                if group.index % 2 == parity]
+            observations = [
+                obs
+                for group in source_groups
+                for obs in group.observations]
+            if not observations:
+                raise ValueError(
+                    f'no observations found in {name.replace("_", " ")}')
+
+            weights = {
+                obs.obs_id: group.observation_weights[obs.obs_id]
+                for group in source_groups
+                for obs in group.observations}
+            interleaved.append(ObsDataGroup(
+                observations,
+                name,
+                umin,
+                umax,
+                weighting_mode=weighting_mode,
+                observation_weights=weights,
+                index=parity,
+            ))
+
+        return tuple(interleaved)

{pspipe-0.7.2 → pspipe-0.8}/pspipe/tasks.py

@@ -107,6 +107,8 @@ class AbstractTask(object):
     def execute(self):
         self.worker.execute()
 
+        return self.get_report().ok
+
     def execute_single_local(self, cmd, log_file):
         w = worker.WorkerPool([worker.localhost_shortname], name=self.name,
                               env_source_file=self.settings.worker.env_file)
@@ -115,6 +117,9 @@ class AbstractTask(object):
             print(cmd)
         w.execute()
 
+    def get_report(self):
+        return worker.WorkerPoolReport(self.worker.tasks)
+
 
 class ImagingSubTask(AbstractSubTask):
 
@@ -133,6 +138,9 @@ class ImagingSubTask(AbstractSubTask):
     def get_dir(self):
         return self.img_rev.get_image_path(self.obs_id)
 
+    def get_image_data_dir(self):
+        return self.img_rev.get_image_data_path(self.obs_id, persist=True)
+
     def get_sb(self):
         res = re.search(r'_SB(\d+)', os.path.basename(self.inp_ms_files[0]))
         if res is not None:
@@ -149,10 +157,16 @@ class ImagingSubTask(AbstractSubTask):
         eo_str = database.even_odd_str(self.even, self.odd)
         w_str = utils.alphanum(img.weight)
         s = f'{self.obs_id}{eo_str}-SB{self.get_sb()}-UV{img.umin}_{img.umax}_{w_str}'
-        return os.path.join(self.get_dir(), s)
+
+        out_prefix = os.path.join(self.get_image_data_dir(), s)
+        os.makedirs(os.path.dirname(out_prefix), exist_ok=True)
+
+        return out_prefix
 
     def register_files(self, files, lst_idx=None, img_name='image'):
         meta = self.img_rev.get_meta(self.obs_id, lst_idx=lst_idx)
+        if meta is None:
+            meta = self.img_rev.new_obs_id(self.obs_id, lst_idx=lst_idx)
         meta.add_fits(files, even=self.even, odd=self.odd, img_name=img_name)
         meta.save()
 
@@ -300,26 +314,28 @@ class ImagingTask(AbstractTask):
                 print(f'Error: Obs ID {obs_id} has no MS files registered.')
                 return []
 
-            if not self.settings.image.lst_bins:
-                self.img_rev.new_obs_id(obs_id)
-            else:
+            data_dir_root = self.img_rev.get_distributed_data_root(obs_id, persist=False)
+
+            self.img_rev.new_obs_id(obs_id, data_dir_root=data_dir_root)
+            if self.settings.image.lst_bins:
                 for i in self.img_rev.get_lst_bins(obs_id).keys():
                     self.img_rev.new_obs_id(obs_id, lst_idx=i)
 
             for i, ms_files_channel in enumerate(ms_files):
-                target_host = get_target_host(ms_files_channel[0], self.settings.worker)
+                file_host = get_target_host(ms_files_channel[0], self.settings.worker)
                 if self.settings.image.split_even_odd:
-                    self.add_sub_task(ImagingSubTask(self.settings, obs_id, ms_files_channel, even=True, index=i),
-                                      run_on_host=target_host)
-                    self.add_sub_task(ImagingSubTask(self.settings, obs_id, ms_files_channel, odd=True, index=i),
-                                      run_on_host=target_host)
+                    t = ImagingSubTask(self.settings, obs_id, ms_files_channel, even=True, index=i)
+                    self.add_sub_task(t, run_on_host=file_host)
+                    t = ImagingSubTask(self.settings, obs_id, ms_files_channel, odd=True, index=i)
+                    self.add_sub_task(t, run_on_host=file_host)
                 else:
-                    self.add_sub_task(ImagingSubTask(self.settings, obs_id, ms_files_channel, index=i),
-                                      run_on_host=target_host)
+                    t = ImagingSubTask(self.settings, obs_id, ms_files_channel, index=i)
+                    self.add_sub_task(t, run_on_host=file_host)
 
             all_obs_id_out.append(obs_id)
 
-        self.execute()
+        if not self.execute():
+            return []
 
         return all_obs_id_out
 
@@ -443,7 +459,8 @@ class GenVisCubeTask(AbstractTask):
 
             all_obs_id_out.append(obs_id)
 
-        self.execute()
+        if not self.execute():
+            return []
 
         if has_even.intersection(has_odd):
             self.init_for_new_task('EvenOddToSumDiff')
@@ -458,7 +475,8 @@ class GenVisCubeTask(AbstractTask):
                         if meta.has_stokes(stokes, odd=True) and meta.has_stokes(stokes, even=True):
                             self.add_sub_task(EvenOddToSumDiffSubTask(self.settings, obs_id, stokes, lst_idx))
 
-            self.execute()
+            if not self.execute():
+                return []
 
         return all_obs_id_out
 
@@ -535,7 +553,9 @@ class ConcatMSTask(AbstractTask):
             self.add_sub_task(ConcatMSSubTask(self.settings, msfiles_flat, obs_id_out))
             all_obs_id_out.append(obs_id_out)
 
-        self.execute()
+        if not self.execute():
+            return []
+
         return all_obs_id_out
 
 
@@ -645,7 +665,8 @@ class MergeMSTask(AbstractTask):
             self.add_sub_task(MergeMSSubTask(self.settings, ms_files, obs_id_out))
             all_obs_id_out.append(obs_id_out)
 
-        self.execute()
+        if not self.execute():
+            return []
 
         return all_obs_id_out
 
@@ -695,7 +716,8 @@ class AOFlaggerTask(AbstractTask):
                 self.add_sub_task(AOFlaggerSubTask(self.settings, ms_file))
             all_obs_id_out.append(obs_id)
 
-        self.execute()
+        if not self.execute():
+            return []
 
         return all_obs_id_out
 
@@ -741,7 +763,8 @@ class DP3Task(AbstractTask):
                 self.add_sub_task(DP3SubTask(self.settings, ms_file))
             all_obs_id_out.append(obs_id)
 
-        self.execute()
+        if not self.execute():
+            return []
 
         return all_obs_id_out
 
@@ -800,7 +823,8 @@ class MakePSTask(AbstractTask):
             for lst_idx in all_lst_idx:
                 self.add_sub_task(MakePSSubTask(self.settings, obs_id, lst_idx))
 
-        self.execute()
+        if not self.execute():
+            return []
 
         return obs_ids
 
@@ -867,7 +891,8 @@ class RunGPRTask(AbstractTask):
             for lst_idx in all_lst_idx:
                 self.add_sub_task(RunGPRSubTask(self.settings, obs_id, lst_idx))
 
-        self.execute()
+        if not self.execute():
+            return []
 
         return obs_ids
 
@@ -936,7 +961,8 @@ class RunMLGPRTask(AbstractTask):
             for lst_idx in all_lst_idx:
                 self.add_sub_task(RunMLGPRSubTask(self.settings, obs_id, lst_idx))
 
-        self.execute()
+        if not self.execute():
+            return []
 
         return obs_ids
 
@@ -1027,7 +1053,8 @@ class RunMLGPRInjTask(AbstractTask):
 
                 self.add_sub_task(RunMLGPRInjSubTask(self.settings, obs_id, 0, 0, 0, lst_idx))
 
-        self.execute()
+        if not self.execute():
+            return []
 
         return obs_ids
 
@@ -1193,7 +1220,9 @@ class VisToSphTask(AbstractTask):
 
             self.add_sub_task(RunFlaggerSubTask(self.settings, obs_id))
 
-        self.execute()
+        if not self.execute():
+            return []
+
         self.init_for_new_task('VisToSph S2')
         new_obs_ids = []
 
@@ -1213,7 +1242,8 @@ class VisToSphTask(AbstractTask):
 
             new_obs_ids.append(obs_id_out)
 
-        self.execute()
+        if not self.execute():
+            return []
 
         return new_obs_ids
 
@@ -1263,7 +1293,8 @@ class ApplyFlagTask(AbstractTask):
 
                 self.add_sub_task(ApplyFlagSubTask(self.settings, ms_file, h5_file))
 
-        self.execute()
+        if not self.execute():
+            return []
 
         return obs_ids
 
@@ -1298,7 +1329,8 @@ class RestoreFlagTask(AbstractTask):
 
                 self.add_sub_task(ApplyFlagSubTask(self.settings, ms_file, h5_file))
 
-        self.execute()
+        if not self.execute():
+            return []
 
         return obs_ids
 
@@ -1353,7 +1385,8 @@ class DelayFlagger(AbstractTask):
             for ms_file in np.array(ms_files).flatten():
                 self.add_sub_task(DelayFlaggerSubTask(self.settings, ms_file))
 
-        self.execute()
+        if not self.execute():
+            return []
 
         return obs_ids
 
@@ -1402,7 +1435,8 @@ class VisFlagger(AbstractTask):
             for ms_file in np.array(ms_files).flatten():
                 self.add_sub_task(VisFlaggerSubTask(self.settings, ms_file))
 
-        self.execute()
+        if not self.execute():
+            return []
 
         return obs_ids
 
@@ -1461,6 +1495,7 @@ class SSINSTask(AbstractTask):
             for ms_file in np.array(ms_files).flatten():
                 self.add_sub_task(SSINSSubTask(self.settings, ms_file))
 
-        self.execute()
+        if not self.execute():
+            return []
 
         return obs_ids

{pspipe-0.7.2 → pspipe-0.8}/pspipe/templates/default_settings.toml

@@ -1,4 +1,5 @@
 data_dir = '.'
+distributed_data_dirs = []
 default_settings = ''
 debug = false
 longitude = 6.57

{pspipe-0.7.2 → pspipe-0.8}/pspipe/tools/pspipe.py

@@ -13,6 +13,25 @@ from .psdb import load_settings
 tasks_help = '\n'.join([f'  {name.ljust(17)}{desc}' for name, desc in tasks.get_all_tasks_descriptions().items()])
 
 
+def print_report(report):
+    print("\n=== Summary of completed tasks ===")
+    print(f"Successful tasks (rc=0):        {report.n_success_rc0} / {report.n_total}")
+    print(f"Completed with non-zero rc:     {report.n_success_nonzero} / {report.n_total}")
+    print(f"Failed tasks (exception/kill):  {report.n_failed} / {report.n_total}")
+
+    if report.success_nonzero:
+        print("\nTasks completed with non-zero return code:")
+        for t in report.success_nonzero:
+            print(f" - {t.name} | host={report.host_of(t)} | rc={t.returncode}")
+
+    if report.failed:
+        print("\nFailed tasks:")
+        for t in report.failed:
+            print(f" - {t.name} | status={t.status.name} | host={report.host_of(t)} | rc={t.returncode}")
+
+    print("==========================\n")
+
+
 @click.command(epilog='\b\nAvailable tasks:\n' + tasks_help)
 @click.version_option(__version__)
 @click.argument('tasks_str', metavar='TASKS')
@@ -39,12 +58,17 @@ def main(tasks_str, settings_file, obs_ids, config_modifier):
 
     obs_ids_set = database.Database(s).filter_obs_ids(obs_ids)
 
-    if not obs_ids:
+    if not obs_ids_set:
         print(f'Error: No obs_ids matching {obs_ids} in database.')
         sys.exit(1)
 
     for task_name in tasks_str.split(','):
-        obs_ids_set = all_tasks[task_name](s).run(obs_ids_set)
+        t = all_tasks[task_name](s)
+        obs_ids_set = t.run(obs_ids_set)
+        report = t.get_report()
+
         if not obs_ids_set:
-            print('\nError: no obs_id to process.')
+            print(f'Task {task_name} finished with errors')
+            print_report(report)
             sys.exit(1)
+        print_report(report)

{pspipe-0.7.2 → pspipe-0.8}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [project]
 name = "pspipe"
-version = "0.7.2"
+version = "0.8"
 description = "Power spectra pipeline for Cosmic Dawn, Epoch of Reionization radio interferometric experiments"
 readme = "README.md"
 authors = [
@@ -20,7 +20,7 @@ dependencies = [
     "tables>=3.2",
     "tabulate>=0.8",
     "ps-eor>=0.8",
-    "libpipe>=0.1.5",
+    "libpipe>=0.6",
     "python-casacore>=3.0",
 ]
 
@@ -34,3 +34,7 @@ pspipe = "pspipe.tools.pspipe:main"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^8.0"
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+addopts = "-ra"

pspipe-0.7.2/pspipe/obsdata.py

@@ -1,127 +0,0 @@
-# Direct access to the data
-#
-# Author: F. Mertens
-
-import os
-
-from ps_eor import pspec, datacube, flagger
-
-from . import database
-
-# Lazy loaded:
-# from ps_eor import fgfit
-# from ps_eor import ml_gpr
-
-
-class ObsData(object):
-
-    def __init__(self, vis_rev, obs_id, lst_idx=None):
-        self.vis_rev = vis_rev
-        self.obs_id = obs_id
-        self.lst_idx = lst_idx
-        self.settings = self.vis_rev.settings
-        self.flag = None
-        self.uvrange = None
-        self.nfreqs = None
-        self.data = {}
-        self.all_stokes = {}
-        for stokes in database.all_stokes:
-            self.all_stokes[stokes.lower()] = stokes
-            self.all_stokes[stokes.lower() + '_even'] = 'even_%s' % stokes
-            self.all_stokes[stokes.lower() + '_odd'] = 'odd_%s' % stokes
-            self.all_stokes[stokes.lower() + '_dt'] = 'dt_%s' % stokes
-        self.ps_builder = pspec.PowerSpectraBuilder(self.settings.power_spectra.get_path('ps_config'),
-                                                    self.settings.power_spectra.get_path('eor_bin_list'))
-
-    def __getattr__(self, key):
-        assert key in self.all_stokes.keys()
-        return self.get_stokes(self.all_stokes[key])
-
-    def _invalidate_cache(self):
-        self.data = {}
-
-    def get_stokes_i_like(self):
-        for s in ['I', 'XX', 'YY']:
-            if self.has_stokes(s):
-                return self.get_stokes(s)
-
-    def get_stokes_v_like(self):
-        for s in ['V', 'XY', 'YX', 'XX', 'YY', 'I']:
-            if self.has_stokes(s):
-                return self.get_stokes(s)
-
-    def set_flag(self, flag):
-        self.flag = flag
-        self._invalidate_cache()
-
-    def filter_uvrange(self, umin, umax):
-        self.uvrange = [umin, umax]
-        self._invalidate_cache()
-
-    def average_freqs(self, nfreqs):
-        self.nfreqs = nfreqs
-        self._invalidate_cache()
-
-    def do_flag(self, flagger_file=None, verbose_flagging=True, plot_flagging=True):
-        if flagger_file is None:
-            flagger_file = self.settings.power_spectra.get_path('flagger')
-        flagger_runner = flagger.FlaggerRunner.load(flagger_file)
-        flagger_runner.verbose = verbose_flagging
-        flagger_runner.run(self.get_stokes_i_like(), self.get_stokes_v_like())
-        if plot_flagging:
-            flagger_runner.plot()
-        self.set_flag(flagger_runner.flag)
-
-    def get_h5_file(self, stokes):
-        return self.vis_rev.get_h5_file(self.obs_id, stokes, lst_idx=self.lst_idx)
-
-    def get_stokes(self, stokes):
-        assert self.has_stokes(stokes)
-
-        if stokes not in self.data:
-            path = self.get_h5_file(stokes)
-            cube = datacube.CartDataCube.load(path)
-            if self.uvrange is not None:
-                cube.filter_uvrange(self.uvrange[0], self.uvrange[1])
-            if self.nfreqs is not None:
-                cube = cube.average_freqs(self.nfreqs)
-            if self.flag is not None:
-                cube = self.flag.apply(cube)
-            self.data[stokes] = cube
-        return self.data[stokes]
-
-    def has_stokes(self, stokes):
-        return stokes in self.data or os.path.exists(self.get_h5_file(stokes))
-
-    def get_gpr_res(self, eor_bin, umin=50, umax=250, stokes='I', name=''):
-        from ps_eor import fgfit
-
-        path = os.path.join(os.path.dirname(self.get_h5_file(stokes)), name, f'eor{eor_bin}_u{umin}-{umax}')
-        i_res = fgfit.GprForegroundResult.load(path, f'gpr_res_{stokes}')
-        i_noise = datacube.CartDataCube.load(os.path.join(path, 'gpr_res_I.noise.h5'))
-        i_res.noise = i_noise
-
-        return i_res
-
-    def get_ml_gpr_path(self, eor_bin, umin=50, umax=250, name='', sampler_method='mcmc'):
-        return os.path.join(os.path.dirname(self.get_h5_file('I')), name, f'eor{eor_bin}_u{umin}-{umax}')
-
-    def get_ml_gpr_res(self, eor_bin, umin=50, umax=250, name='', sampler_method='mcmc'):
-        from ps_eor import ml_gpr
-
-        path = self.get_ml_gpr_path(eor_bin, umin=umin, umax=umax, name=name, sampler_method=sampler_method)
-        return ml_gpr.MLGPRResult.load(path, 'ml_gpr_res_I', sampler_method=sampler_method)
-
-    def get_ml_gpr_inj_path(self, eor_bin, var_21_over_noise, x1, x2, umin=50, umax=250, name=''):
-        d = f'inj_eor{eor_bin}_u{umin}-{umax}_a{var_21_over_noise:.3f}_x1{x1:.3f}_x2{x2:.3f}'
-        return os.path.join(os.path.dirname(self.get_h5_file('I')), name, d)
-
-    def get_ml_gpr_inj_res(self, eor_bin, var_21_over_noise, x1, x2, umin=50, umax=250, name='', sampler_method='mcmc'):
-        from ps_eor import ml_gpr
-
-        path = self.get_ml_gpr_inj_path(eor_bin, var_21_over_noise=var_21_over_noise, x1=x1, x2=x2, umin=umin, umax=umax, name=name)
-        data_inj = datacube.CartDataCube.load(os.path.join(path, 'c_injected.h5'))
-        return ml_gpr.MLGPRInjResult.load(path, 'ml_gpr_res_I', data_inj, sampler_method=sampler_method)
-
-    def get_ps_gen(self, **kargs):
-        return self.ps_builder.get(self.get_stokes_i_like(), **kargs)

pspipe-0.7.2/pspipe/tests/test_settings.py

@@ -1,32 +0,0 @@
-from pspipe import settings
-
-
-def test_validate(tmp_path):
-    s = settings.Settings.get_defaults()
-
-    f = (tmp_path / 'f')
-    f.touch()
-
-    s.vis_to_sph.pre_flag = str(f)
-    s.combine.pre_flag = str(f)
-    s.power_spectra.eor_bin_list = str(f)
-    s.power_spectra.ps_config = str(f)
-    s.power_spectra.flagger = str(f)
-    s.gpr.config_i = str(f)
-    s.gpr.config_v = str(f)
-
-    for m in s.keys():
-        if hasattr(s, f'validate_{m}'):
-            assert s.validate(m), m
-
-    s.gpr.config_v = 'kks'
-    assert not s.validate('gpr')
-
-    s.vis_cube.win_fct = 'hann'
-    assert s.validate('vis_cube')
-
-    s.vis_cube.win_fct = 'aksks'
-    assert not s.validate('vis_cube')
-
-    s.image.lst_bins = [0, 4, 1]
-    assert s.validate('image')

pspipe-0.7.2/pspipe/tests/test_utils.py

@@ -1,28 +0,0 @@
-from pspipe import utils
-
-
-def test_alphanum():
-    assert utils.alphanum('|sd dd 12.s ') == 'sddd12s'
-
-
-def test_all_same():
-    assert utils.all_same([1, 1, 1])
-    assert not utils.all_same([1, 1, 2])
-    assert not utils.all_same([1, 2, 1])
-    assert utils.all_same(['a', 'a', 'a'])
-    assert not utils.all_same(['a', 'a', 'a', 1])
-
-
-def test_all_in_other():
-    assert utils.all_in_other([1, 2], list(range(10)))
-    assert not utils.all_in_other([1, 2, 10], list(range(10)))
-
-
-def test_is_in_lst_bin():
-    assert utils.is_in_lst_bin(2, 1, 3)
-    assert utils.is_in_lst_bin(2, 2, 3)
-    assert not utils.is_in_lst_bin(2, 3, 3)
-    assert not utils.is_in_lst_bin(3, 2, 3)
-    assert not utils.is_in_lst_bin(3, 21, 3)
-    assert utils.is_in_lst_bin(2, 21, 3)
-    assert utils.is_in_lst_bin(21, 21, 3)