pspipe 0.11__tar.gz → 0.13__tar.gz

{pspipe-0.11 → pspipe-0.13}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pspipe
-Version: 0.11
+Version: 0.13
 Summary: Power spectra pipeline for Cosmic Dawn, Epoch of Reionization radio interferometric experiments
 License-Expression: GPL-3.0-or-later
 License-File: LICENSE
@@ -18,7 +18,7 @@ Provides-Extra: docs
 Provides-Extra: ml-gpr
 Requires-Dist: astropy (>6.0)
 Requires-Dist: click (>=8.0)
-Requires-Dist: libpipe (>=0.8)
+Requires-Dist: libpipe (>=0.9)
 Requires-Dist: myst-parser (>=2) ; extra == "docs"
 Requires-Dist: ps-eor (>=1.0)
 Requires-Dist: ps-eor[ml-gpr] (>=1.0) ; extra == "ml-gpr"

{pspipe-0.11 → pspipe-0.13}/pspipe/settings.py

@@ -5,13 +5,20 @@
 import os
 
 import astropy.units as u
-from libpipe import futils
+from libpipe import futils, worker
 from libpipe.settings import BaseSettings
 
 # Lazy loaded:
 # - ps_eor.datacube in Settings.validate_vis_cube()
 from . import utils
 
+
+def worker_resources(s_worker):
+    '''worker.resources as {category: {'cpus': .., 'mem': ..}}, or None.'''
+    resources = s_worker.get('resources') or {}
+    return {k: dict(v) for k, v in resources.items()} or None
+
+
 TEMPLATE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'templates')
 
 TEMPLATE_NAMES = ['a12', 'nenufar', 'hba', 'alo']
@@ -25,6 +32,8 @@ class Settings(BaseSettings):
         BaseSettings.__init__(self, file, d)
 
     def validate_image(self, val_fct):
+        val_fct(self.image.imager in ('wsclean', 'psimager'), 'imager')
+        val_fct(self.image.imager == 'wsclean' or not self.image.clean_niter, 'clean_niter')
         val_fct(u.Unit(self.image.scale).is_equivalent((u.rad, u.s)), 'scale')
         val_fct(len(self.image.size.split()) == 2
                 and all(k.isdigit() for k in self.image.size.split()), 'size')
@@ -49,9 +58,19 @@ class Settings(BaseSettings):
     def validate_worker(self, val_fct):
         val_fct(isinstance(self.worker.max_concurrent, int), 'max_concurrent')
         val_fct(not self.worker.env_file or futils.is_file(self.worker.env_file), 'env_file')
-        val_fct(self.worker.get('max_time') is None
+        val_fct(not self.worker.get('max_time')
                 or (isinstance(self.worker.max_time, (int, float)) and self.worker.max_time > 0),
                 'max_time')
+        slurm_nodes = self.worker.get('slurm_nodes') or 0
+        val_fct(isinstance(slurm_nodes, int) and slurm_nodes >= 0, 'slurm_nodes')
+        try:
+            resources = worker.resolve_resources(worker_resources(self.worker))
+        except (ValueError, TypeError):
+            resources = None
+        val_fct(resources is not None, 'resources')
+        categories = self.worker.get('categories') or {}
+        val_fct(resources is None or all(c in resources for c in categories.values()),
+                'categories')
 
     def validate_merge_ms(self, val_fct):
         val_fct(isinstance(self.merge_ms.apply_aoflagger, bool), 'apply_aoflagger')

{pspipe-0.11 → pspipe-0.13}/pspipe/tasks.py

@@ -14,7 +14,7 @@ import astropy.io.fits as pf
 import numpy as np
 from libpipe import futils, worker
 
-from . import database, utils
+from . import database, settings, utils
 
 CWD = os.getcwd()
 
@@ -48,6 +48,11 @@ def get_target_host(in_file, worker_settings):
 
 class AbstractSubTask:
 
+    # Task type (worker.categories key, usage summary) and resource category as a job step of an
+    # allocation (worker.slurm_nodes): small, medium, high_mem, cpu or heavy
+    label = None
+    category = 'heavy'
+
     def __init__(self, settings_obj, exec_name):
         self.settings = settings_obj
         self.exec_name = exec_name
@@ -74,23 +79,24 @@ class AbstractTask:
         self.name = name
         self.settings = settings_obj.copy()
         self.settings.worker.update(worker_kargs)
-        s_worker = self.settings.worker
 
         if not self.settings.validate('worker'):
             raise ValueError('Invalid worker settings.')
 
-        self.worker = worker.get_worker_pool(
+        self.worker = self._worker_pool(name)
+
+    def _worker_pool(self, name):
+        s_worker = self.settings.worker
+        # 0 and '' (the defaults) leave max_time unset and slurm_options from the environment
+        return worker.get_worker_pool(
                 name, nodes=s_worker.nodes, max_concurrent=s_worker.max_concurrent,
-                env_file=s_worker.env_file, max_time=s_worker.get('max_time'),
-                slurm_options=s_worker.get('slurm_options'))
+                env_file=s_worker.env_file, max_time=s_worker.get('max_time') or None,
+                slurm_options=s_worker.get('slurm_options') or None,
+                slurm_nodes=s_worker.get('slurm_nodes') or None,
+                resources=settings.worker_resources(s_worker))
 
     def init_for_new_task(self, post_task_name):
-        s_worker = self.settings.worker
-        self.worker = worker.get_worker_pool(post_task_name, nodes=s_worker.nodes,
-                                             max_concurrent=s_worker.max_concurrent,
-                                             env_file=s_worker.env_file,
-                                             max_time=s_worker.get('max_time'),
-                                             slurm_options=s_worker.get('slurm_options'))
+        self.worker = self._worker_pool(post_task_name)
 
     def add_sub_task(self, sub_task, run_on_host=None):
         log_file = sub_task.get_log_file()
@@ -103,8 +109,10 @@ class AbstractTask:
         cmd = sub_task.build_command()
         if self.settings.debug:
             print(cmd)
+        label = sub_task.label or sub_task.exec_name
+        category = (self.settings.worker.get('categories') or {}).get(label, sub_task.category)
         self.worker.add(cmd, output_file=log_file, done_callback=sub_task.done_callback,
-                        run_on_host=run_on_host)
+                        run_on_host=run_on_host, category=category, label=label)
 
     def run(self, obs_ids):
         raise NotImplementedError()
@@ -128,12 +136,18 @@ class AbstractTask:
 
 class ImagingSubTask(AbstractSubTask):
 
-    def __init__(self, settings_obj, obs_id, inp_ms_files, even=False, odd=False, index=0):
-        AbstractSubTask.__init__(self, settings_obj, 'wsclean')
+    label = 'image'
+    category = 'heavy'
+
+    def __init__(self, settings_obj, obs_id, inp_ms_files, even=False, odd=False, index=0,
+                 even_odd=False):
+        AbstractSubTask.__init__(self, settings_obj, settings_obj.image.imager)
         self.inp_ms_files = inp_ms_files
         self.obs_id = obs_id
         self.even = even
         self.odd = odd
+        # psimager only: even and odd time steps in one call
+        self.even_odd = even_odd
         self.index = index
         self.img_rev = database.ImageRevision(self.settings)
         self.obs_meta = database.Database(self.settings).get_obs_id_meta(self.obs_id)
@@ -169,11 +183,12 @@ class ImagingSubTask(AbstractSubTask):
 
         return out_prefix
 
-    def register_files(self, files, lst_idx=None, img_name='image'):
+    def register_files(self, files, lst_idx=None, img_name='image', even=None, odd=None):
         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.add_fits(files, even=self.even if even is None else even,
+                      odd=self.odd if odd is None else odd, img_name=img_name)
         meta.save()
 
     def done_callback(self):
@@ -189,32 +204,42 @@ class ImagingSubTask(AbstractSubTask):
         else:
             total_time = self.obs_meta.total_time
 
-        if img.clean_niter > 0 and img.clean_residual:
-            files = glob.glob(f'{self.get_fits_prefix()}*psf.fits')
-            files.extend(glob.glob(f'{self.get_fits_prefix()}*residual.fits'))
-            img_name = 'residual'
-        else:
-            files = glob.glob(f'{self.get_fits_prefix()}*.fits')
-            img_name = 'image'
-
-        for file in files:
-            pf.setval(file, 'PEINTTIM', value=int(self.obs_meta.int_time))
-            pf.setval(file, 'PETOTTIM', value=int(total_time))
-            pf.setval(file, 'PECHWIDT', value=int(self.obs_meta.chan_width))
-
-        if not img.lst_bins:
-            self.register_files(files, img_name=img_name)
+        if self.even_odd:
+            parts = [('-even', True, False), ('-odd', False, True)]
         else:
-            for i, data_lst_idx in enumerate(self.img_rev.get_lst_bins(self.obs_id,
-                    return_time_index=True)[2]):
-                files = glob.glob(f'{self.get_fits_prefix()}*-t{i:04}-*.fits')
-                self.register_files(files, lst_idx=data_lst_idx, img_name=img_name)
-
-        with contextlib.suppress(OSError):
-            os.remove(self.temp_protect_file)
-
-        with contextlib.suppress(OSError):
-            os.rmdir(self.temp_dir)
+            parts = [('', self.even, self.odd)]
+
+        for part, even, odd in parts:
+            prefix = self.get_fits_prefix() + part
+
+            if img.clean_niter > 0 and img.clean_residual:
+                files = glob.glob(f'{prefix}*psf.fits')
+                files.extend(glob.glob(f'{prefix}*residual.fits'))
+                img_name = 'residual'
+            else:
+                files = glob.glob(f'{prefix}*.fits')
+                img_name = 'image'
+
+            for file in files:
+                pf.setval(file, 'PEINTTIM', value=int(self.obs_meta.int_time))
+                pf.setval(file, 'PETOTTIM', value=int(total_time))
+                pf.setval(file, 'PECHWIDT', value=int(self.obs_meta.chan_width))
+
+            if not img.lst_bins:
+                self.register_files(files, img_name=img_name, even=even, odd=odd)
+            else:
+                for i, data_lst_idx in enumerate(self.img_rev.get_lst_bins(self.obs_id,
+                        return_time_index=True)[2]):
+                    files = glob.glob(f'{prefix}*-t{i:04}-*.fits')
+                    self.register_files(files, lst_idx=data_lst_idx, img_name=img_name,
+                                        even=even, odd=odd)
+
+        if self.temp_dir is not None:
+            with contextlib.suppress(OSError):
+                os.remove(self.temp_protect_file)
+
+            with contextlib.suppress(OSError):
+                os.rmdir(self.temp_dir)
 
     def get_parameters(self):
         args = []
@@ -229,29 +254,33 @@ class ImagingSubTask(AbstractSubTask):
             channels_out = self.obs_meta.n_channels // int(channels_out[5:])
         channels_out = int(channels_out)
 
-        args.append(f"-oversampling {img.oversampling} -kernel-size {img.kernel_size} -nwlayers "
-                    f"{img.nwlayers}")
-        args.append(f"-grid-mode kb -taper-edge {img.taper_edge} -padding {img.padding}")
-        args.append(f"-visibility-weighting-mode {img.visibility_weighting_mode} -reorder "
+        args.append(f"-taper-edge {img.taper_edge} -padding {img.padding}")
+        args.append(f"-visibility-weighting-mode {img.visibility_weighting_mode} "
                     f"-make-psf -no-dirty")
 
-        # If MS directory is not writable, we need to define a temp directory
-        if not all(os.access(ms_file, os.W_OK) for ms_file in self.inp_ms_files):
-            self.temp_dir = os.path.join(self.get_dir(),
-                                         'tmp' + database.even_odd_str(self.even, self.odd))
-        else:
-            self.temp_dir = os.path.join(self.inp_ms_files[0],
-                                         'tmp' + database.even_odd_str(self.even, self.odd))
+        if img.imager == 'wsclean':
+            args.append(f"-oversampling {img.oversampling} -kernel-size {img.kernel_size} "
+                        f"-nwlayers {img.nwlayers} -grid-mode kb -reorder")
+
+            # If MS directory is not writable, we need to define a temp directory
+            if not all(os.access(ms_file, os.W_OK) for ms_file in self.inp_ms_files):
+                self.temp_dir = os.path.join(self.get_dir(),
+                                             'tmp' + database.even_odd_str(self.even, self.odd))
+            else:
+                self.temp_dir = os.path.join(self.inp_ms_files[0],
+                                             'tmp' + database.even_odd_str(self.even, self.odd))
 
-        if not os.path.isdir(self.temp_dir):
-            os.makedirs(self.temp_dir)
+            if not os.path.isdir(self.temp_dir):
+                os.makedirs(self.temp_dir)
 
-        # Make sure we do not remove this directory by accident
-        _, self.temp_protect_file = tempfile.mkstemp(dir=self.temp_dir)
+            # Make sure we do not remove this directory by accident
+            _, self.temp_protect_file = tempfile.mkstemp(dir=self.temp_dir)
 
-        args.append(f'-temp-dir {self.temp_dir}')
+            args.append(f'-temp-dir {self.temp_dir}')
 
-        if self.even:
+        if self.even_odd:
+            args.append('-even-odd')
+        elif self.even:
             args.append('-even-timesteps')
         elif self.odd:
             args.append('-odd-timesteps')
@@ -327,7 +356,11 @@ class ImagingTask(AbstractTask):
 
             for i, ms_files_channel in enumerate(ms_files):
                 file_host = get_target_host(ms_files_channel[0], self.settings.worker)
-                if self.settings.image.split_even_odd:
+                if self.settings.image.split_even_odd and self.settings.image.imager == 'psimager':
+                    t = ImagingSubTask(self.settings, obs_id, ms_files_channel, index=i,
+                                       even_odd=True)
+                    self.add_sub_task(t, run_on_host=file_host)
+                elif self.settings.image.split_even_odd:
                     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)
@@ -346,6 +379,9 @@ class ImagingTask(AbstractTask):
 
 class GenVisCubeSubTask(AbstractSubTask):
 
+    label = 'gen_vis_cube'
+    category = 'small'
+
     def __init__(self, settings_obj, obs_id, stokes, even=False, odd=False, lst_idx=None):
         AbstractSubTask.__init__(self, settings_obj, 'pstool gen_vis_cube')
         self.obs_id = obs_id
@@ -401,6 +437,9 @@ class GenVisCubeSubTask(AbstractSubTask):
 
 class EvenOddToSumDiffSubTask(AbstractSubTask):
 
+    label = 'even_odd_to_sum_diff'
+    category = 'small'
+
     def __init__(self, settings_obj, obs_id, stokes, lst_idx=None):
         AbstractSubTask.__init__(self, settings_obj, 'pstool even_odd_to_sum_diff')
         self.obs_id = obs_id
@@ -501,6 +540,9 @@ class GenVisCubeTask(AbstractTask):
 
 
 class ConcatMSSubTask(AbstractSubTask):
+
+    label = 'concat_ms'
+    category = 'medium'
     """
     Subtask that calls the CLI 'mstool concat <MS...> <OUT.MS>'
     to concatenate MS files in time.
@@ -582,6 +624,9 @@ class ConcatMSTask(AbstractTask):
 
 class MergeMSSubTask(AbstractSubTask):
 
+    label = 'merge_ms'
+    category = 'heavy'
+
     def __init__(self, settings_obj, ms_files, obs_id_out):
         AbstractSubTask.__init__(self, settings_obj, 'DP3')
         self.db = database.Database(self.settings)
@@ -696,6 +741,9 @@ class MergeMSTask(AbstractTask):
 
 class AOFlaggerSubTask(AbstractSubTask):
 
+    label = 'ao_flag'
+    category = 'medium'
+
     def __init__(self, settings_obj, ms_file):
         AbstractSubTask.__init__(self, settings_obj, 'DP3')
         self.ms_file = ms_file
@@ -749,6 +797,9 @@ class AOFlaggerTask(AbstractTask):
 
 class DP3SubTask(AbstractSubTask):
 
+    label = 'dp3'
+    category = 'heavy'
+
     def __init__(self, settings_obj, ms_file):
         AbstractSubTask.__init__(self, settings_obj, 'DP3')
         self.ms_file = ms_file
@@ -796,6 +847,9 @@ class DP3Task(AbstractTask):
 
 class MakePSSubTask(AbstractSubTask):
 
+    label = 'make_ps'
+    category = 'small'
+
     def __init__(self, settings_obj, obs_id, lst_idx=None):
         AbstractSubTask.__init__(self, settings_obj, 'pstool make_ps')
         self.obs_id = obs_id
@@ -859,6 +913,9 @@ class MakePSTask(AbstractTask):
 
 class RunGPRSubTask(AbstractSubTask):
 
+    label = 'run_gpr'
+    category = 'cpu'
+
     def __init__(self, settings_obj, obs_id, lst_idx=None):
         AbstractSubTask.__init__(self, settings_obj, 'pstool run_gpr')
         self.obs_id = obs_id
@@ -929,6 +986,9 @@ class RunGPRTask(AbstractTask):
 
 class RunMLGPRSubTask(AbstractSubTask):
 
+    label = 'run_ml_gpr'
+    category = 'cpu'
+
     def __init__(self, settings_obj, obs_id, lst_idx=None):
         AbstractSubTask.__init__(self, settings_obj, 'pstool run_ml_gpr')
         self.obs_id = obs_id
@@ -1002,6 +1062,9 @@ class RunMLGPRTask(AbstractTask):
 
 class RunMLGPRInjSubTask(AbstractSubTask):
 
+    label = 'run_ml_gpr_inj'
+    category = 'cpu'
+
     def __init__(self, settings_obj, obs_id, var_21_over_noise, x1, x2, lst_idx=None):
         AbstractSubTask.__init__(self, settings_obj, 'pstool run_ml_gpr_inj')
         self.obs_id = obs_id
@@ -1208,6 +1271,9 @@ class CombineTask(AbstractTask):
 
 class RunFlaggerSubTask(AbstractSubTask):
 
+    label = 'run_flagger'
+    category = 'small'
+
     def __init__(self, settings_obj, obs_id):
         AbstractSubTask.__init__(self, settings_obj, 'pstool run_flagger')
         self.vis_rev = database.VisRevision(self.settings)
@@ -1222,6 +1288,9 @@ class RunFlaggerSubTask(AbstractSubTask):
 
 class VisToSphSubTask(AbstractSubTask):
 
+    label = 'vis_to_sph'
+    category = 'small'
+
     def __init__(self, settings_obj, obs_id_out, stokes, h5_input, h5_out, h5_flag):
         AbstractSubTask.__init__(self, settings_obj, 'pstool vis_to_sph')
         self.vis_rev = database.VisRevision(self.settings)
@@ -1296,6 +1365,9 @@ class VisToSphTask(AbstractTask):
 
 class ApplyFlagSubTask(AbstractSubTask):
 
+    label = 'apply_flag'
+    category = 'small'
+
     def __init__(self, settings_obj, ms_file, h5_file):
         AbstractSubTask.__init__(self, settings_obj, 'flagtool restore')
         self.ms_file = ms_file
@@ -1383,6 +1455,9 @@ class RestoreFlagTask(AbstractTask):
 
 class DelayFlaggerSubTask(AbstractSubTask):
 
+    label = 'delay_flagger'
+    category = 'medium'
+
     def __init__(self, settings_obj, ms_file):
         AbstractSubTask.__init__(self, settings_obj, 'flagtool delay_flagger')
         self.ms_file = ms_file
@@ -1440,6 +1515,9 @@ class DelayFlagger(AbstractTask):
 
 class VisFlaggerSubTask(AbstractSubTask):
 
+    label = 'vis_flagger'
+    category = 'medium'
+
     def __init__(self, settings_obj, ms_file):
         AbstractSubTask.__init__(self, settings_obj, 'flagtool vis_flagger')
         self.ms_file = ms_file
@@ -1491,6 +1569,9 @@ class VisFlagger(AbstractTask):
 
 class SSINSSubTask(AbstractSubTask):
 
+    label = 'ssins'
+    category = 'medium'
+
     def __init__(self, settings_obj, ms_file):
         AbstractSubTask.__init__(self, settings_obj, 'flagtool ssins')
         self.ms_file = ms_file

{pspipe-0.11 → pspipe-0.13}/pspipe/templates/default_settings.toml

@@ -10,6 +10,13 @@ max_concurrent = 4
 env_file = ''
 run_on_file_host = false
 run_on_file_host_pattern = ''
+slurm_options = ''
+max_time = 0
+slurm_nodes = 0
+
+[worker.resources]
+
+[worker.categories]
 
 [merge_ms]
 apply_aoflagger = true
@@ -63,6 +70,7 @@ config_file = ''
 
 [image]
 name = 'default'
+imager = 'wsclean'
 umin = 50
 umax = 250
 scale = '0.5amin'

{pspipe-0.11 → pspipe-0.13}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [project]
 name = "pspipe"
-version = "0.11"
+version = "0.13"
 description = "Power spectra pipeline for Cosmic Dawn, Epoch of Reionization radio interferometric experiments"
 readme = "README.md"
 license = "GPL-3.0-or-later"
@@ -20,7 +20,7 @@ dependencies = [
     "tables>=3.2",
     "tabulate>=0.8",
     "ps-eor>=1.0",
-    "libpipe>=0.8",
+    "libpipe>=0.9",
     "python-casacore>=3.0",
 ]