pspipe 0.3__tar.gz → 0.5__tar.gz

{pspipe-0.3 → pspipe-0.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pspipe
-Version: 0.3
+Version: 0.5
 Summary: Power spectra pipeline for Cosmic Dawn, Epoch of Reionization radio interferometric experiments
 Home-page: https://gitlab.com/flomertens/pspipe/
 Author: "Florent Mertens"
@@ -12,14 +12,15 @@ Classifier: Programming Language :: Python :: 3.8
 Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
-Requires-Dist: astropy (>=4.0,<5.0); python_version >= "3.7" and python_version < "3.8"
-Requires-Dist: astropy (>=5.0,<6.0); python_version >= "3.8"
+Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: astropy (>=4.0,<5.0) ; python_version >= "3.7" and python_version < "3.8"
+Requires-Dist: astropy (>=5.0,<6.0) ; python_version >= "3.8"
 Requires-Dist: matplotlib (>=3.0,<4.0)
 Requires-Dist: nenucal (>=0.1.11)
 Requires-Dist: ps-eor (>=0.8)
 Requires-Dist: python-casacore (>=3.0,<4.0)
-Requires-Dist: scipy (>=1.6,<2.0); python_version >= "3.7" and python_version < "3.8"
-Requires-Dist: scipy (>=1.8,<2.0); python_version >= "3.8"
+Requires-Dist: scipy (>=1.6,<2.0) ; python_version >= "3.7" and python_version < "3.8"
+Requires-Dist: scipy (>=1.8,<2.0) ; python_version >= "3.8"
 Requires-Dist: tables (>=3.2,<4.0)
 Requires-Dist: tabulate (>=0.8)
 Requires-Dist: toml (>=0.10)

{pspipe-0.3 → pspipe-0.5}/pspipe/obsdata.py

@@ -33,7 +33,7 @@ class ObsData(object):
         assert key in self.all_stokes.keys()
         return self.get_stokes(self.all_stokes[key])
 
-    def __invalid_cache(self):
+    def _invalidate_cache(self):
         self.data = {}
 
     def get_stokes_i_like(self):
@@ -48,15 +48,15 @@ class ObsData(object):
 
     def set_flag(self, flag):
         self.flag = flag
-        self.__invalid_cache()
+        self._invalidate_cache()
 
     def filter_uvrange(self, umin, umax):
         self.uvrange = [umin, umax]
-        self.__invalid_cache()
+        self._invalidate_cache()
 
     def average_freqs(self, nfreqs):
         self.nfreqs = nfreqs
-        self.__invalid_cache()
+        self._invalidate_cache()
 
     def do_flag(self, flagger_file=None, verbose_flagging=True, plot_flagging=True):
         if flagger_file is None:
@@ -97,9 +97,21 @@ class ObsData(object):
 
         return i_res
 
-    def get_ml_gpr_res(self, eor_bin, umin=50, umax=250, name=''):
-        path = os.path.join(os.path.dirname(self.get_h5_file('I')), name, f'eor{eor_bin}_u{umin}-{umax}')
-        return ml_gpr.MLGPRResult.load(path, 'ml_gpr_res_I')
+    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'):
+        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'):
+        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.3 → pspipe-0.5}/pspipe/tasks.py

@@ -6,6 +6,8 @@ import os
 import re
 import glob
 import tempfile
+import datetime
+import itertools
 
 import numpy as np
 import astropy.io.fits as pf
@@ -87,6 +89,8 @@ class AbstractTask(object):
 
     def add_sub_task(self, sub_task, run_on_host=None):
         log_file = sub_task.get_log_file()
+        if log_file:
+            log_file = log_file.rstrip('/')
 
         if log_file is not None:
             if not os.path.exists(os.path.dirname(log_file)):
@@ -137,7 +141,8 @@ class ImagingSubTask(AbstractSubTask):
             return f'{self.index:04d}'
 
     def get_log_file(self):
-        return os.path.join(self.get_dir(), 'logs', os.path.basename(self.get_fits_prefix()) + '.log')
+        date_str = datetime.datetime.now().strftime('%Y-%m-%dT%H:%M:%S')
+        return os.path.join(self.get_dir(), 'logs', os.path.basename(self.get_fits_prefix()) + f'_{date_str}.log')
 
     def get_fits_prefix(self):
         img = self.settings.image
@@ -212,8 +217,8 @@ class ImagingSubTask(AbstractSubTask):
             channels_out = self.obs_meta.n_channels // int(channels_out[5:])
         channels_out = int(channels_out)
 
-        args.append("-oversampling 4095 -kernel-size 15 -nwlayers 1000 -grid-mode kb -taper-edge 100")
-        args.append("-padding 2 -visibility-weighting-mode unit -reorder -make-psf -no-dirty")
+        args.append(f"-oversampling 4095 -kernel-size 15 -nwlayers 1000 -grid-mode kb -taper-edge {img.taper_edge}")
+        args.append(f"-padding {img.padding} -visibility-weighting-mode unit -reorder -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):
@@ -346,8 +351,9 @@ class GenVisCubeSubTask(AbstractSubTask):
         return self.vis_rev.get_h5_file(self.obs_id, self.stokes.upper(), self.even, self.odd, lst_idx=self.lst_idx)
 
     def get_log_file(self):
+        date_str = datetime.datetime.now().strftime('%Y-%m-%dT%H:%M:%S')
         return os.path.join(self.vis_rev.get_h5_path(self.obs_id, self.lst_idx), 'logs',
-                            os.path.basename(self.get_out_file().replace('.h5', '.log')))
+                            os.path.basename(self.get_out_file().replace('.h5', f'_{date_str}.log')))
 
     def done_callback(self):
         meta = self.vis_rev.get_meta(self.obs_id, self.lst_idx)
@@ -459,7 +465,7 @@ class GenVisCubeTask(AbstractTask):
 class MergeMSSubTask(AbstractSubTask):
 
     def __init__(self, settings_obj, ms_files, obs_id_out):
-        AbstractSubTask.__init__(self, settings_obj, 'DPPP')
+        AbstractSubTask.__init__(self, settings_obj, 'DP3')
         self.db = database.Database(self.settings)
         self.ms_files = ms_files
         self.obs_id_out = obs_id_out
@@ -471,7 +477,8 @@ class MergeMSSubTask(AbstractSubTask):
         return os.path.join(self.get_dir(), f'{self.obs_id_out}.MS')
 
     def get_log_file(self):
-        return os.path.join(self.get_out_file().replace('.MS', '.log'))
+        date_str = datetime.datetime.now().strftime('%Y-%m-%dT%H:%M:%S')
+        return os.path.join(self.get_out_file().replace('.MS', f'_{date_str}.log'))
 
     def make_link(self):
         inp_ms_dir = os.path.join(self.get_dir(), 'input_ms')
@@ -499,7 +506,7 @@ class MergeMSSubTask(AbstractSubTask):
         merge_ms = self.settings.merge_ms
 
         steps = ['filter']
-        args = [f'msin={inp_ms_dir}/*.MS', f'msin.datacolumn={merge_ms.data_col}']
+        args = [f'msin={inp_ms_dir}/*.MS', f'msin.datacolumn={merge_ms.data_col}', f'numthreads={merge_ms.numthreads}']
         args.append(f'msout={self.get_out_file()} msout.overwrite=true')
         args.append(f'filter.type=filter filter.blrange=[{merge_ms.blmin},{merge_ms.blmax}]')
 
@@ -511,7 +518,7 @@ class MergeMSSubTask(AbstractSubTask):
 
         if merge_ms.apply_aoflagger:
             steps.append('flagger')
-            args.append(f'flagger.type=aoflagger flagger.memoryperc=30 flagger.strategy="{merge_ms.aoflagger_strategy}"')
+            args.append(f'flagger.type=aoflagger flagger.memoryperc={merge_ms.memoryperc} flagger.strategy="{merge_ms.aoflagger_strategy}"')
 
         if merge_ms.flag_baselines:
             steps.append('preflag')
@@ -566,6 +573,102 @@ class MergeMSTask(AbstractTask):
         return all_obs_id_out
 
 
+class AOFlaggerSubTask(AbstractSubTask):
+
+    def __init__(self, settings_obj, ms_file):
+        AbstractSubTask.__init__(self, settings_obj, 'DP3')
+        self.ms_file = ms_file
+
+    def get_log_file(self):
+        date_str = datetime.datetime.now().strftime('%Y-%m-%dT%H:%M:%S')
+        return os.path.join(self.ms_file.replace('.MS', f'_aoflag_{date_str}.log'))
+
+    def get_parameters(self):
+        ao_flag = self.settings.ao_flag
+        args = [f'msin={self.ms_file}', 'msout=.', f'msin.datacolumn={ao_flag.data_col}', 'steps=[flagger]', f'numthreads={ao_flag.numthreads}']
+        args.append(f'flagger.type=aoflagger flagger.memoryperc={ao_flag.memoryperc} flagger.strategy="{ao_flag.aoflagger_strategy}"')
+
+        return args
+
+
+class AOFlaggerTask(AbstractTask):
+
+    name = 'ao_flag'
+    desc = 'Run AO flagger on the Measurement Set(s)'
+
+    def __init__(self, settings_obj):
+        AbstractTask.__init__(self, 'AOFlagger', settings_obj, max_concurrent=1)
+        self.db = database.Database(self.settings)
+
+    def run(self, obs_ids):
+        all_obs_id_out = []
+
+        for obs_id in obs_ids:
+            if not self.db.has_obs_id(obs_id):
+                print(f'Error: Obs ID {obs_id} does not exist.')
+                return []
+
+            ms_files = self.db.get_obs_id_meta(obs_id).ms_files
+
+            if not ms_files:
+                print(f'Error: Obs ID {obs_id} has no MS files registered.')
+                return []
+
+            for ms_file in np.array(ms_files).flatten():
+                self.add_sub_task(AOFlaggerSubTask(self.settings, ms_file))
+            all_obs_id_out.append(obs_id)
+
+        self.execute()
+
+        return all_obs_id_out
+
+
+class DP3SubTask(AbstractSubTask):
+
+    def __init__(self, settings_obj, ms_file):
+        AbstractSubTask.__init__(self, settings_obj, 'DP3')
+        self.ms_file = ms_file
+
+    def get_log_file(self):
+        date_str = datetime.datetime.now().strftime('%Y-%m-%dT%H:%M:%S')
+        return os.path.join(self.ms_file.replace('.MS', f'_dp3_{date_str}.log'))
+
+    def get_parameters(self):
+        return [self.settings.dp3.dp3_file, f'msin={self.ms_file}']
+
+
+class DP3Task(AbstractTask):
+
+    name = 'dp3'
+    desc = 'Run DP3 on the Measurement Set(s)'
+
+    def __init__(self, settings_obj):
+        AbstractTask.__init__(self, 'DP3', settings_obj, max_concurrent=1)
+        self.db = database.Database(self.settings)
+
+    def run(self, obs_ids):
+        all_obs_id_out = []
+
+        for obs_id in obs_ids:
+            if not self.db.has_obs_id(obs_id):
+                print(f'Error: Obs ID {obs_id} does not exist.')
+                return []
+
+            ms_files = self.db.get_obs_id_meta(obs_id).ms_files
+
+            if not ms_files:
+                print(f'Error: Obs ID {obs_id} has no MS files registered.')
+                return []
+
+            for ms_file in np.array(ms_files).flatten():
+                self.add_sub_task(DP3SubTask(self.settings, ms_file))
+            all_obs_id_out.append(obs_id)
+
+        self.execute()
+
+        return all_obs_id_out
+
+
 class MakePSSubTask(AbstractSubTask):
 
     def __init__(self, settings_obj, obs_id, lst_idx=None):
@@ -575,7 +678,8 @@ class MakePSSubTask(AbstractSubTask):
         self.vis_rev = database.VisRevision(self.settings)
 
     def get_log_file(self):
-        return os.path.join(self.vis_rev.get_h5_path(self.obs_id, self.lst_idx), 'logs', 'make_ps.log')
+        date_str = datetime.datetime.now().strftime('%Y-%m-%dT%H:%M:%S')
+        return os.path.join(self.vis_rev.get_h5_path(self.obs_id, self.lst_idx), 'logs', f'make_ps_{date_str}.log')
 
     def get_parameters(self):
         h5_i = self.vis_rev.get_h5_file(self.obs_id, 'I', lst_idx=self.lst_idx)
@@ -633,7 +737,8 @@ class RunGPRSubTask(AbstractSubTask):
         self.vis_rev = database.VisRevision(self.settings)
 
     def get_log_file(self):
-        return os.path.join(self.vis_rev.get_path(), self.obs_id, 'logs', 'run_gpr.log')
+        date_str = datetime.datetime.now().strftime('%Y-%m-%dT%H:%M:%S')
+        return os.path.join(self.vis_rev.get_path(), self.obs_id, 'logs', f'run_gpr_{date_str}.log')
 
     def get_parameters(self):
         h5_i = self.vis_rev.get_h5_file(self.obs_id, 'I', lst_idx=self.lst_idx)
@@ -699,13 +804,16 @@ class RunMLGPRSubTask(AbstractSubTask):
         self.vis_rev = database.VisRevision(self.settings)
 
     def get_log_file(self):
-        return os.path.join(self.vis_rev.get_path(), self.obs_id, 'logs', 'run_ml_gpr.log')
+        date_str = datetime.datetime.now().strftime('%Y-%m-%dT%H:%M:%S')
+        return os.path.join(self.vis_rev.get_path(), self.obs_id, 'logs', f'run_ml_gpr_{date_str}.log')
 
     def get_parameters(self):
-        h5_i = self.vis_rev.get_h5_file(self.obs_id, 'I', lst_idx=self.lst_idx)
-        h5_v = self.vis_rev.get_h5_file(self.obs_id, 'V', lst_idx=self.lst_idx)
         psc = self.settings.power_spectra
         gpr = self.settings.ml_gpr
+
+        h5_i = self.vis_rev.get_h5_file(self.obs_id, gpr.stokes_i, lst_idx=self.lst_idx)
+        h5_v = self.vis_rev.get_h5_file(self.obs_id, gpr.stokes_v, lst_idx=self.lst_idx)
+
         out_dir = os.path.join(self.vis_rev.get_h5_path(self.obs_id, self.lst_idx), gpr.name)
         plot_out_dir = os.path.join(self.vis_rev.get_result_path(self.obs_id, self.lst_idx), gpr.name)
 
@@ -714,7 +822,7 @@ class RunMLGPRSubTask(AbstractSubTask):
         args.append(f'-c {psc.get_path("ps_config")} -o {out_dir} -po {plot_out_dir}')
 
         if gpr.use_v_dt_as_noise:
-            h5_v_dt = self.vis_rev.get_h5_file(self.obs_id, 'dt_V', lst_idx=self.lst_idx)
+            h5_v_dt = self.vis_rev.get_h5_file(self.obs_id, gpr.stokes_dt, lst_idx=self.lst_idx)
             args.append(f'--file_noise {h5_v_dt}')
 
         if not gpr.plot_results:
@@ -756,6 +864,97 @@ class RunMLGPRTask(AbstractTask):
         return obs_ids
 
 
+class RunMLGPRInjSubTask(AbstractSubTask):
+
+    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
+        self.lst_idx = lst_idx
+        self.var_21_over_noise = var_21_over_noise
+        self.x1 = x1
+        self.x2 = x2
+        self.vis_rev = database.VisRevision(self.settings)
+
+    def get_log_file(self):
+        date_str = datetime.datetime.now().strftime('%Y-%m-%dT%H:%M:%S')
+        return os.path.join(self.vis_rev.get_path(), self.obs_id, 'logs', f'run_ml_gpr_inj_{date_str}.log')
+
+    def get_parameters(self):
+        h5_i = self.vis_rev.get_h5_file(self.obs_id, 'I', lst_idx=self.lst_idx)
+        h5_v = self.vis_rev.get_h5_file(self.obs_id, 'V', lst_idx=self.lst_idx)
+        psc = self.settings.power_spectra
+        gpr = self.settings.ml_gpr
+        out_dir = os.path.join(self.vis_rev.get_h5_path(self.obs_id, self.lst_idx), gpr.name)
+
+        args = [h5_i, h5_v, gpr.get_path("config")]
+        args.append(f'-f {psc.get_path("flagger")} -e {psc.get_path("eor_bin_list")}')
+        args.append(f'-c {psc.get_path("ps_config")} -o {out_dir}')
+
+        if gpr.use_v_dt_as_noise:
+            h5_v_dt = self.vis_rev.get_h5_file(self.obs_id, 'dt_V', lst_idx=self.lst_idx)
+            args.append(f'--file_noise {h5_v_dt}')
+
+        args.append(f'-a {self.var_21_over_noise} --x1 {self.x1} --x2 {self.x2}')
+
+        if not gpr.plot_results:
+            args.append('--no_plot')
+
+        return args
+
+
+class RunMLGPRInjTask(AbstractTask):
+
+    name = 'run_ml_gpr_inj'
+    desc = 'Run ML-GPR injection and produce power-spectra'
+
+    def __init__(self, settings_obj):
+        AbstractTask.__init__(self, 'ML-GPR', settings_obj)
+        self.vis_rev = database.VisRevision(self.settings)
+        self.img_rev = database.ImageRevision(self.settings)
+
+    def run(self, obs_ids):
+        if not (self.settings.validate('power_spectra') and self.settings.validate('gpr')):
+            return []
+
+        if not self.vis_rev.exist():
+            print(f'Error: visibility revision does not exist.')
+            return []
+
+        for obs_id in obs_ids:
+            if not self.vis_rev.has_obs_id(obs_id):
+                print(f'Error: visibility revision for obs ID {obs_id} not available')
+                return []
+
+            all_lst_idx = self.img_rev.get_lst_bins(obs_id).keys() if self.settings.image.lst_bins else [None]
+
+            for lst_idx in all_lst_idx:
+                var = self.settings.ml_gpr_inj.var_21_over_noise
+                x1s = self.settings.ml_gpr_inj.x1
+                x2s = self.settings.ml_gpr_inj.x2
+                x1_x2_file = self.settings.ml_gpr_inj.x1_x2_file
+
+                try:
+                    if x1_x2_file:
+                        x1_x2s = np.loadtxt(x1_x2_file, delimiter=',').tolist()
+                    else:
+                        x1_x2s = list(itertools.product(x1s, x2s))
+                except IOError as e:
+                    print(f"Error loading file '{x1_x2_file}': {e}")
+                    return []
+                except ValueError as e:
+                    print(f"Error parsing file '{x1_x2_file}': {e}")
+                    return []
+
+                for v, (x1, x2) in itertools.product(var, x1_x2s):
+                    self.add_sub_task(RunMLGPRInjSubTask(self.settings, obs_id, v, x1, x2, lst_idx))
+
+                self.add_sub_task(RunMLGPRInjSubTask(self.settings, obs_id, 0, 0, 0, lst_idx))
+
+        self.execute()
+
+        return obs_ids
+
+
 class CombineTask(AbstractTask):
 
     name = 'combine'
@@ -842,7 +1041,8 @@ class CombineTask(AbstractTask):
         umax = self.settings.vis_cube.umax
         c = self.settings.combine
 
-        out_log = os.path.join(self.get_path(), f'C{obs_id_out_prefix}.log')
+        date_str = datetime.datetime.now().strftime('%Y-%m-%dT%H:%M:%S')
+        out_log = os.path.join(self.get_path(), f'C{obs_id_out_prefix}_{date_str}.log')
 
         if is_sph:
             cmd = (f'pstool combine_sph {self.get_combine_file()} -o {output_template} '
@@ -948,6 +1148,10 @@ class ApplyFlagSubTask(AbstractSubTask):
         self.ms_file = ms_file
         self.h5_file = h5_file
 
+    def get_log_file(self):
+        date_str = datetime.datetime.now().strftime('%Y-%m-%dT%H:%M:%S')
+        return os.path.join(self.ms_file.replace('.MS', f'apply_flags_{date_str}.log'))
+
     def get_parameters(self):
         return [self.ms_file, self.h5_file]
 
@@ -1027,12 +1231,13 @@ class DelayFlaggerSubTask(AbstractSubTask):
     def __init__(self, settings_obj, ms_file):
         AbstractSubTask.__init__(self, settings_obj, 'flagtool delay_flagger')
         self.ms_file = ms_file
+        self.date_str = datetime.datetime.now().strftime('%Y-%m-%dT%H:%M:%S')
 
     def get_log_file(self):
-        return os.path.join(self.ms_file.replace('.MS', '_delay_flagger.log'))
+        return os.path.join(self.ms_file.replace('.MS', f'_delay_flagger_{self.date_str}.log'))
 
     def get_path(self):
-        d = os.path.join(os.path.dirname(self.ms_file), 'flagger')
+        d = os.path.join(os.path.dirname(self.ms_file), f'flagger_{self.date_str}')
         futils.mkdir(d)
         return d
 
@@ -1076,6 +1281,55 @@ class DelayFlagger(AbstractTask):
         return obs_ids
 
 
+class VisFlaggerSubTask(AbstractSubTask):
+
+    def __init__(self, settings_obj, ms_file):
+        AbstractSubTask.__init__(self, settings_obj, 'flagtool vis_flagger')
+        self.ms_file = ms_file
+        self.date_str = datetime.datetime.now().strftime('%Y-%m-%dT%H:%M:%S')
+
+    def get_log_file(self):
+        return os.path.join(self.ms_file.replace('.MS', f'{self.date_str}_vis_flagger.log'))
+
+    def get_path(self):
+        d = os.path.join(os.path.dirname(self.ms_file), f'vis_flagger_{self.date_str}')
+        futils.mkdir(d)
+        return d
+
+    def get_parameters(self):
+        return [self.ms_file, self.settings.vis_flagger.config_file, f'--plot_dir={self.get_path()}',
+                '--backup', f'--backup_file={H5_FLAG_BACKUP}']
+
+
+class VisFlagger(AbstractTask):
+
+    name = 'vis_flagger'
+    desc = 'Apply vis flagger on MS'
+
+    def __init__(self, settings_obj):
+        AbstractTask.__init__(self, 'VisFlagger', settings_obj)
+        self.db = database.Database(self.settings)
+
+    def run(self, obs_ids):
+        for obs_id in obs_ids:
+            if not self.db.has_obs_id(obs_id):
+                print(f'Error: Obs ID {obs_id} does not exist.')
+                return []
+
+            ms_files = self.db.get_obs_id_meta(obs_id).ms_files
+
+            if not ms_files:
+                print(f'Error: Obs ID {obs_id} has no MS files registered.')
+                return []
+
+            for ms_file in np.array(ms_files).flatten():
+                self.add_sub_task(VisFlaggerSubTask(self.settings, ms_file))
+
+        self.execute()
+
+        return obs_ids
+
+
 class SSINSSubTask(AbstractSubTask):
 
     def __init__(self, settings_obj, ms_file):

pspipe-0.5/pspipe/templates/config/eor_bins_hba.parset

@@ -0,0 +1,4 @@
+# name,fmhz_start,fmhz_end,fmhz_fg_start,fmhz_fg_end
+1,122.0,134.3,122.0,134.3
+2,134.3,146.8,134.3,146.8
+3,146.8,159.3,146.8,159.3

{pspipe-0.3 → pspipe-0.5}/pspipe/templates/default_settings.toml

@@ -21,6 +21,18 @@ flag_baselines = ''
 avg_timestep = 1
 time_start_index = 0
 time_end_index = 0
+numthreads = 0
+memoryperc = 30
+
+[ao_flag]
+data_col = 'DATA'
+numthreads = 0
+memoryperc = 30
+aoflagger_strategy = ''
+
+[dp3]
+numthreads = 0
+dp3_file = ''
 
 [ssins]
 apply_ssins = true
@@ -45,6 +57,9 @@ n_times = 20
 n_sigma_i = 6
 n_sigma_v = 6
 
+[vis_flagger]
+config_file = ''
+
 [image]
 name = 'default'
 umin = 50
@@ -64,6 +79,8 @@ clean_auto_threshold = 5
 clean_mgain = 0.6
 clean_residual = true
 clean_spectra_fit_nterms = 4
+taper_edge = 100
+padding = 2
 wsclean_args = {}
 
 [vis_cube]
@@ -100,4 +117,13 @@ use_v_dt_as_noise = false
 name = ''
 config = ''
 plot_results = true
+stokes_i = 'I'
+stokes_v = 'V'
+stokes_dt = 'dt_v'
 use_v_dt_as_noise = false
+
+[ml_gpr_inj]
+var_21_over_noise = [1]
+x1 = [0]
+x2 = [0]
+x1_x2_file = ''

{pspipe-0.3 → pspipe-0.5}/pspipe/tools/psdb.py

@@ -396,5 +396,59 @@ def remove_obs(settings_file, obs_ids):
         print('No changes made.')
 
 
+def concatenate_datacubes(cubes):
+    from ps_eor import datacube
+
+    key_fct = lambda a: a.freqs[0]
+    cubes = sorted(cubes, key=key_fct)
+
+    assert len(cubes) > 1
+    assert np.alltrue([isinstance(cube, datacube.DataCube) for cube in cubes])
+    assert np.alltrue([np.allclose(cubes[0].uu, cube.uu) for cube in cubes[1:]])
+    assert np.alltrue([np.allclose(cubes[0].vv, cube.vv) for cube in cubes[1:]])
+
+    data = np.concatenate([c.data for c in cubes])
+    weights_data = np.concatenate([c.weights.data for c in cubes])
+    freqs = np.concatenate([c.freqs for c in cubes])
+    weights = datacube.CartWeightCube(weights_data, cubes[0].uu, cubes[0].vv, freqs, cubes[0].meta)
+
+    return datacube.CartDataCube(data, cubes[0].uu, cubes[0].vv, freqs, cubes[0].meta, weights=weights)
+
+
+@main.command('concat_vis_cube')
+@click.argument('settings_file', type=t_file)
+@click.argument('obs_ids', type=str)
+@click.argument('new_obs_id', type=str)
+@click.option('--fmin', type=float, help='Min frequency in MHz', default=None)
+@click.option('--fmax', type=float, help='Max frequency in MHz', default=None)
+def concat_vis_cube(settings_file, obs_ids, new_obs_id, fmin, fmax):
+    ''' Remove an Obs ID from the given pspipe database
+    '''
+    rev = database.VisRevision(load_settings(settings_file))
+    db = database.Database(load_settings(settings_file))
+    obs_id_set = list(db.filter_obs_ids(obs_ids))
+
+    if not obs_id_set:
+        print(f'Error: No obs_ids matching {obs_ids} in database.')
+        return
+
+    meta = db.get_obs_id_meta(obs_id_set[0])
+
+    add(db, new_obs_id, None, int_time=meta.int_time, total_time=meta.total_time, chan_width=meta.chan_width)
+    data1 = rev.get_data(obs_id_set[0])
+
+    all_stokes = [s for s in data1.all_stokes.values() if data1.has_stokes(s)]
+
+    for s in all_stokes:
+        data_cube = concatenate_datacubes([rev.get_data(obs_id).get_stokes(s) for obs_id in obs_id_set])
+        if fmin is not None and fmax is not None:
+            data_cube = data_cube.get_slice(fmin, fmax)
+        h5_filename = rev.get_h5_file(new_obs_id, s)
+        os.makedirs(os.path.dirname(h5_filename),  exist_ok=True)
+        data_cube.save(h5_filename)
+
+    print(f'Combined data for {new_obs_id} successfully added.')
+
+
 if __name__ == '__main__':
     main()

{pspipe-0.3 → pspipe-0.5}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pspipe"
-version = "0.3"
+version = "0.5"
 description = "Power spectra pipeline for Cosmic Dawn, Epoch of Reionization radio interferometric experiments"
 authors = ["\"Florent Mertens\" <\"florent.mertens@gmail.com\">"]
 repository = "https://gitlab.com/flomertens/pspipe/"

pspipe-0.3/pspipe/templates/config/eor_bins_hba.parset

@@ -1,5 +0,0 @@
-# name,fmhz_start,fmhz_end,fmhz_fg_start,fmhz_fg_end
-1,122,134,122,134
-2,134,146,134,146
-3,146,158,146,158
-

pspipe-0.3/setup.py

@@ -1,47 +0,0 @@
-# -*- coding: utf-8 -*-
-from setuptools import setup
-
-packages = \
-['pspipe', 'pspipe.tests', 'pspipe.tools']
-
-package_data = \
-{'': ['*'], 'pspipe': ['templates/*', 'templates/config/*']}
-
-install_requires = \
-['matplotlib>=3.0,<4.0',
- 'nenucal>=0.1.11',
- 'ps-eor>=0.8',
- 'python-casacore>=3.0,<4.0',
- 'tables>=3.2,<4.0',
- 'tabulate>=0.8',
- 'toml>=0.10']
-
-extras_require = \
-{':python_version >= "3.7" and python_version < "3.8"': ['astropy>=4.0,<5.0',
-                                                         'scipy>=1.6,<2.0'],
- ':python_version >= "3.8"': ['astropy>=5.0,<6.0', 'scipy>=1.8,<2.0']}
-
-entry_points = \
-{'console_scripts': ['psdb = pspipe.tools.psdb:main',
-                     'pspipe = pspipe.tools.pspipe:main']}
-
-setup_kwargs = {
-    'name': 'pspipe',
-    'version': '0.3',
-    'description': 'Power spectra pipeline for Cosmic Dawn, Epoch of Reionization radio interferometric experiments',
-    'long_description': 'Power Spectra Generation pipeline for CD/EoR experiments\n========================================================\n\npspipe is an analysis pipeline which aims to make generation of high precision power-spectra for Cosmic Dawn and Epoch of Reionization experiments as easy and reproducible as possible. pspipe is the default analysis pipeline of the [Lofar-EoR](http://www.lofar.org/astronomy/eor-ksp/epoch-reionization.html) and [NenuFAR](https://nenufar.obs-nancay.fr/en/homepage-en/) CD projects. It performs all tasks from calibrated data to final the power-spectra, including:\n\n- Post-calibration flagging (using [AOflagger](https://sourceforge.net/p/aoflagger/wiki/Home/) and [SSINS](https://arxiv.org/abs/1906.01093))\n- Visibility gridding (using [WSClean](https://sourceforge.net/p/wsclean/wiki/Home/))\n- Foregrounds modeling and removal (using [GPR](https://doi.org/10.1093/mnras/sty1207), [GMCA](https://arxiv.org/abs/1209.4769), polynomial fitting)\n- Optimal power-spectra generation (using [ps_eor](https://gitlab.com/flomertens/ps_eor))\n\nThanks to its multi-nodes concurrent processing, pspipe is also very fast at crunching hundreds of hour of calibrated data.\n\nInstallation\n------------\n\npspipe can be installed via pip:\n\n    $ pip install pspipe\n\nand requires Python 3.7.0 or higher. For certain tasks, you will also need [DPPP](https://github.com/lofar-astron/DP3/tree/master/DPPP) and [WSClean](https://sourceforge.net/p/wsclean/wiki/Home/)\n\nDocumentation\n--------------\n\nPlease check the [wiki page](https://gitlab.com/flomertens/pspipe/-/wikis/home).\n',
-    'author': '"Florent Mertens"',
-    'author_email': '"florent.mertens@gmail.com"',
-    'maintainer': 'None',
-    'maintainer_email': 'None',
-    'url': 'https://gitlab.com/flomertens/pspipe/',
-    'packages': packages,
-    'package_data': package_data,
-    'install_requires': install_requires,
-    'extras_require': extras_require,
-    'entry_points': entry_points,
-    'python_requires': '>=3.7,<4.0',
-}
-
-
-setup(**setup_kwargs)