make_cat: write the final catalogue once per save stage, node-local

src/shapepipe/modules/make_cat_package/make_cat.py

@@ -138,8 +138,9 @@ def save_sextractor_data(final_cat_file, sexcat_path, remove_vignet=True):
 
     final_cat_file.save_as_fits(data, ext_name="RESULTS")
     final_cat_file.open()
-    final_cat_file.add_col("TILE_ID", tile_id_array)
-    final_cat_file.add_col("TILE_UNIQUE_ID", unique_id)
+    final_cat_file.add_cols(
+        {"TILE_ID": tile_id_array, "TILE_UNIQUE_ID": unique_id}
+    )
 
     sexcat_file.close()
 
@@ -204,10 +205,11 @@ def save_mask_ext_data(final_cat_file, band_paths, w_log):
     ra = np.copy(final_cat_file.get_data()["XWIN_WORLD"])
     dec = np.copy(final_cat_file.get_data()["YWIN_WORLD"])
 
+    mask_cols = {}
     for band, path in band_paths.items():
         w_log.info(f"Query external mask for band {band}: {path}")
-        values = mask_query.query_map(path, ra, dec)
-        final_cat_file.add_col(f"MASK_{band}", values)
+        mask_cols[f"MASK_{band}"] = mask_query.query_map(path, ra, dec)
+    final_cat_file.add_cols(mask_cols)
 
     final_cat_file.close()
 
@@ -279,9 +281,7 @@ def process(
             )
 
         if err_msg is None:
-
-            for key in self._output_dict.keys():
-                self._final_cat_file.add_col(key, self._output_dict[key])
+            self._final_cat_file.add_cols(self._output_dict)
 
         self._final_cat_file.close()
 

src/shapepipe/modules/make_cat_runner.py

@@ -7,6 +7,7 @@
 """
 
 import os
+import shutil
 
 from shapepipe.modules.make_cat_package import make_cat
 from shapepipe.modules.module_decorator import module_runner
@@ -62,9 +63,23 @@ def make_cat_runner(
     else:
         n_epoch_slots = None
 
+    # The catalogue is built in WORK_DIR when set (e.g. node-local disk:
+    # each save stage rewrites the whole file) and moved to the run's output
+    # directory once complete.
+    if config.has_option(module_config_sec, "WORK_DIR"):
+        work_dir = config.getexpanded(module_config_sec, "WORK_DIR")
+        os.makedirs(work_dir, exist_ok=True)
+    else:
+        work_dir = run_dirs["output"]
+    work_path = make_cat.get_output_name(work_dir, file_number_string)
+    # save_as_fits appends to an existing file: a work file left by an
+    # earlier attempt must not survive into this one.
+    if os.path.exists(work_path):
+        os.remove(work_path)
+
     # Set final output file
     final_cat_file = make_cat.prepare_final_cat_file(
-        run_dirs["output"],
+        work_dir,
         file_number_string,
     )
 
@@ -84,12 +99,7 @@ def make_cat_runner(
 
         # If error message: delete (incomplete) output file and raise error
         if err_msg is not None:
-            os.remove(
-                make_cat.get_output_name(
-                    run_dirs["output"],
-                    file_number_string,
-                )
-            )
+            os.remove(work_path)
             #raise ValueError(err_msg)
             w_log.info(err_msg)
 
@@ -108,4 +118,10 @@ def make_cat_runner(
         w_log.info("Save external mask data")
         make_cat.save_mask_ext_data(final_cat_file, band_paths, w_log)
 
+    if work_dir != run_dirs["output"] and os.path.exists(work_path):
+        shutil.move(
+            work_path,
+            make_cat.get_output_name(run_dirs["output"], file_number_string),
+        )
+
     return None, None

src/shapepipe/pipeline/file_io.py

@@ -1354,7 +1354,7 @@ def add_col(
     ):
         """Add Column.
 
-        Add a Column to the catalogue.
+        Add one column to the catalogue; see :meth:`add_cols`.
 
         Parameters
         ----------
@@ -1365,7 +1365,46 @@ def add_col(
         hdu_no : int
             HDU index
         ext_name : str, optional
-            Change the name of the extansion
+            Change the name of the extension
+        new_cat : bool, optional
+            If true will save the changes into a new catalogue
+        new_cat_inst : io.FITSCatalogue
+            New catalogue object
+
+        """
+        self.add_cols(
+            {col_name: col_data},
+            hdu_no=hdu_no,
+            ext_name=ext_name,
+            new_cat=new_cat,
+            new_cat_inst=new_cat_inst,
+        )
+
+    def add_cols(
+        self,
+        columns,
+        hdu_no=None,
+        ext_name=None,
+        new_cat=False,
+        new_cat_inst=None,
+    ):
+        """Add Columns.
+
+        Append columns to a table HDU, rebuilding the HDU and writing the
+        file once for all of them. The result is identical to one
+        :meth:`add_col` call per column in the order of ``columns``, but the
+        cost is one write of the file instead of one per column, which
+        dominates when many columns are added to a large catalogue.
+
+        Parameters
+        ----------
+        columns : dict
+            Column name to column data (``numpy.ndarray``); columns are
+            appended in the dict's order
+        hdu_no : int
+            HDU index
+        ext_name : str, optional
+            Change the name of the extension
         new_cat : bool, optional
             If true will save the changes into a new catalogue
         new_cat_inst : io.FITSCatalogue
@@ -1388,86 +1427,80 @@ def add_col(
                 open_mode_needed=FITSCatalogue.OpenMode.ReadWrite,
             )
 
-        if type(col_data) != np.ndarray:
-            raise TypeError("col_data must be a numpy.ndarray")
+        for col_data in columns.values():
+            if type(col_data) != np.ndarray:
+                raise TypeError("col_data must be a numpy.ndarray")
+
+        if not columns:
+            return
 
         if hdu_no is None:
             hdu_no = self.hdu_no
         if ext_name is None:
             ext_name = self._cat_data[hdu_no].name
 
-        n_of_hdu = len(self._cat_data)
-        old_hdu_prev = []
-        for i in range(0, hdu_no):
-            old_hdu_prev.append(self._cat_data[i])
-        old_hdu_next = []
-        for i in range(hdu_no + 1, n_of_hdu):
-            old_hdu_next.append(self._cat_data[i])
+        col_list = self._cat_data[hdu_no].data.columns + fits.ColDefs(
+            [
+                self._make_fits_col(col_name, col_data)
+                for col_name, col_data in columns.items()
+            ]
+        )
 
-        new_fits = fits.HDUList(old_hdu_prev)
+        new_fits = fits.HDUList(
+            self._cat_data[:hdu_no]
+            + [fits.BinTableHDU.from_columns(col_list, name=ext_name)]
+            + self._cat_data[hdu_no + 1 :]
+        )
+        new_fits.writeto(output_path, overwrite=True)
 
-        col_list = self._cat_data[hdu_no].data.columns
+        if not new_cat:
+            self._cat_data.close()
+            del self._cat_data
+            self._cat_data = fits.open(
+                self.fullpath,
+                mode=self.open_mode,
+                memmap=self.use_memmap,
+            )
+
+    def _make_fits_col(self, col_name, col_data):
+        """Make FITS Column.
+
+        Build the ``astropy.io.fits.Column`` that :meth:`add_cols` appends
+        for one array: the FITS type from :meth:`_get_fits_col_type`, a
+        repeat count and ``TDIM`` for multi-dimensional arrays, and a width
+        set by the longest entry for strings.
+
+        Parameters
+        ----------
+        col_name : str
+            Column name
+        col_data : numpy.ndarray
+            Column data
 
+        Returns
+        -------
+        astropy.io.fits.Column
+            The column
+
+        """
         data_type = self._get_fits_col_type(col_data)
         data_shape = col_data.shape[1:]
-        dim = str(tuple(data_shape))
+        dim = None
         mem_size = 1
         if len(data_shape) != 0:
             for k in data_shape:
                 mem_size *= k
-            data_format = f"{mem_size}{data_type}"
-            new_col = fits.ColDefs(
-                [
-                    fits.Column(
-                        name=col_name,
-                        format=data_format,
-                        array=col_data,
-                        dim=dim,
-                    )
-                ]
-            )
-            col_list += new_col
+            dim = str(tuple(data_shape))
         elif data_type == "A":
             mem_size *= len(max(col_data, key=len))
-            data_format = f"{mem_size}{data_type}"
-            new_col = fits.ColDefs(
-                [
-                    fits.Column(
-                        name=col_name,
-                        format=data_format,
-                        array=col_data,
-                        dim=str((mem_size,)),
-                    )
-                ]
-            )
-            col_list += new_col
-        else:
-            data_format = f"{mem_size}{data_type}"
-            new_col = fits.ColDefs(
-                [
-                    fits.Column(
-                        name=col_name,
-                        format=data_format,
-                        array=col_data,
-                    )
-                ]
-            )
-            col_list += new_col
-
-        new_fits.append(fits.BinTableHDU.from_columns(col_list, name=ext_name))
+            dim = str((mem_size,))
 
-        new_fits += fits.HDUList(old_hdu_next)
-
-        new_fits.writeto(output_path, overwrite=True)
-
-        if not new_cat:
-            self._cat_data.close()
-            del self._cat_data
-            self._cat_data = fits.open(
-                self.fullpath,
-                mode=self.open_mode,
-                memmap=self.use_memmap,
-            )
+        return fits.Column(
+            name=col_name,
+            format=f"{mem_size}{data_type}",
+            array=col_data,
+            dim=dim,
+        )
 
     def remove_col(self, col_index):
         """Remove Column.