minor fixes

jurjen93 · Oct 7, 2024 · 6167bd0 · 6167bd0
1 parent b8436bf
commit 6167bd0
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 58 deletions.
diff --git a/ds9_helpers/split_polygon_facets.py b/ds9_helpers/split_polygon_facets.py
@@ -1,7 +1,7 @@
 """
-Split a ds9 region file with multiple facets in it into separate facet region files (these correspond to input h5 files).
-This is helpful if you want to make images of individual facets instead of full facet-imaging in wsclean.
-The script also return a polygon_info.csv containing the center of the polygon and the calibrator source direction with a polygon area
+Split a ds9 region file with multiple facets into separate facet region files.
+You need this if you want to make images of individual facets instead of full facet-imaging in wsclean.
+The script also returns a polygon_info.csv containing the center of the polygon and the calibrator source direction with a polygon area
 and an estimate for how many times you can average based on the measurement sets from a 2.5x2.5 degree wide-field (for long-baselines)
 
 Example:

diff --git a/ms_helpers/remove_flagged_stations.py b/ms_helpers/remove_flagged_stations.py
@@ -20,7 +20,7 @@ def remove_flagged_antennas(msin: str = None, msout: str = None, overwrite: bool
 
     # Set name for output if not given
     if msout is None:
-        msout = f"flagged_{msin}"
+        msout = f"flagged_{msin.split('/')[-1]}"
 
     # Read antenna names from Measurement Set
     with table(f"{msin}::ANTENNA", ack=False) as ants:

diff --git a/subtract/subtract_with_wsclean.py b/subtract/subtract_with_wsclean.py
@@ -1,18 +1,19 @@
-import numpy as np
-import sys
 import os
-import casacore.tables as ct
-import pyregion
-from astropy.io import fits
-from astropy.wcs import WCS
+import random
+import re
+import shutil
+import sys
+from argparse import ArgumentParser
 from glob import glob
-import tables
 from itertools import repeat
-import re
+
+import numpy as np
 import pandas as pd
-from argparse import ArgumentParser
-import random
-import shutil
+import pyregion
+import tables
+from astropy.io import fits
+from astropy.wcs import WCS
+import casacore.tables as ct
 
 
 def add_trailing_zeros(s, digitsize=4):
@@ -100,6 +101,7 @@ def parse_history(ms, hist_item):
     print('WARNING:' + hist_item + ' not found')
     return None
 
+
 def make_utf8(inp):
     """
     Convert input to utf8 instead of bytes
@@ -430,57 +432,56 @@ def subtract_col(self, out_column: str = None):
 
         for ms in self.mslist:
             print('Subtract ' + ms)
-            ts = ct.table(ms, readonly=False, ack=False)
-            colnames = ts.colnames()
-
-            if "MODEL_DATA" not in colnames:
-                sys.exit(
-                    f"ERROR: MODEL_DATA does not exist in {ms}.\nThis is most likely due to a failed predict step.")
+            with ct.table(ms, readonly=False, ack=False) as ts:
+                colnames = ts.colnames()
 
-            if not self.onlyprint:
-                if out_column not in colnames:
-                    # get column description from DATA
-                    desc = ts.getcoldesc('DATA')
-                    # create output column
-                    print('Create ' + out_column)
-                    desc['name'] = out_column
-                    # create template for output column
-                    ts.addcols(desc)
+                if "MODEL_DATA" not in colnames:
+                    sys.exit(
+                        f"ERROR: MODEL_DATA does not exist in {ms}.\nThis is most likely due to a failed predict step.")
 
+                if not self.onlyprint:
+                    if out_column not in colnames:
+                        # get column description from DATA
+                        desc = ts.getcoldesc('DATA')
+                        # create output column
+                        print('Create ' + out_column)
+                        desc['name'] = out_column
+                        # create template for output column
+                        ts.addcols(desc)
+
+                    else:
+                        print(out_column, ' already exists')
+
+                # get number of rows
+                nrows = ts.nrows()
+
+                # make sure every slice has the same size
+                best_slice = get_largest_divider(nrows, 1000)
+
+                if self.inverse:
+                    sign = '+'
                 else:
-                    print(out_column, ' already exists')
+                    sign = '-'
 
-            # get number of rows
-            nrows = ts.nrows()
-
-            # make sure every slice has the same size
-            best_slice = get_largest_divider(nrows, 1000)
-
-            if self.inverse:
-                sign = '+'
-            else:
-                sign = '-'
-
-            if 'SUBTRACT_DATA' in colnames:
-                colmn = 'SUBTRACT_DATA'
-            elif 'CORRECTED_DATA' in colnames:
-                colmn = 'CORRECTED_DATA'
-            else:
-                colmn = 'DATA'
+                if 'SUBTRACT_DATA' in colnames:
+                    colmn = 'SUBTRACT_DATA'
+                elif 'CORRECTED_DATA' in colnames:
+                    colmn = 'CORRECTED_DATA'
+                else:
+                    colmn = 'DATA'
 
-            for c in range(0, nrows, best_slice):
+                for c in range(0, nrows, best_slice):
 
-                if c == 0:
-                    print(f'Output --> {colmn} {sign} MODEL_DATA')
+                    if c == 0:
+                        print(f'Output --> {colmn} {sign} MODEL_DATA')
 
-                if not self.onlyprint:
-                    data = ts.getcol(colmn, startrow=c, nrow=best_slice)
+                    if not self.onlyprint:
+                        data = ts.getcol(colmn, startrow=c, nrow=best_slice)
 
-                if not self.onlyprint:
-                    model = ts.getcol('MODEL_DATA', startrow=c, nrow=best_slice)
-                    ts.putcol(out_column, data - model if not self.inverse else data + model, startrow=c, nrow=best_slice)
-            ts.removecols(['MODEL_DATA'])
-            ts.close()
+                    if not self.onlyprint:
+                        model = ts.getcol('MODEL_DATA', startrow=c, nrow=best_slice)
+                        ts.putcol(out_column, data - model if not self.inverse else data + model, startrow=c, nrow=best_slice)
+                ts.removecols(['MODEL_DATA'])
 
         return self