Various fixes

PyFHD/beam_setup/beam.py

@@ -320,6 +320,6 @@ def create_psf(obs: dict, pyfhd_config: dict, logger: Logger) -> dict | File:
         sys.exit(1)
     raise ValueError(
         f"Unknown beam file type {pyfhd_config['beam_file_path'].suffix}. "
-        "Please use a .sav, .h5, .hdf5"
+        "Please use a .sav, .h5, .hdf5 "
         "If you meant for PyFHD to do the beam forming, please set the beam_file_path to None (~ in YAML)."
     )

PyFHD/beam_setup/beam_utils.py

@@ -156,7 +156,11 @@ def beam_image(
     """
 
     psf_dim = psf["dim"]
-    freq_norm = psf["fnorm"]
+    if "freq_norm" in psf:
+        freq_norm = psf["freq_norm"]
+    elif "fnorm" in psf:
+        # handling for older files or imports from IDL FHD
+        freq_norm = psf["fnorm"]
     pix_horizon = psf["pix_horizon"]
     group_id = psf["id"][pol_i, 0, :]
     if "beam_gaussian_params" in psf:
@@ -170,10 +174,12 @@ def beam_image(
         freq_norm = freq_norm[:]
         pix_horizon = pix_horizon[0]
     dimension = elements = obs["dimension"]
-    xl = dimension / 2 - psf_dim / 2 + 1
-    xh = dimension / 2 - psf_dim / 2 + psf_dim
-    yl = elements / 2 - psf_dim / 2 + 1
-    yh = elements / 2 - psf_dim / 2 + psf_dim
+    # these should all be integers b/c dimensions are usually even numbers.
+    # but they have to be cast to ints to be used in slicing.
+    xl = int(dimension / 2 - psf_dim / 2 + 1)
+    xh = int(dimension / 2 - psf_dim / 2 + psf_dim)
+    yl = int(elements / 2 - psf_dim / 2 + 1)
+    yh = int(elements / 2 - psf_dim / 2 + psf_dim)
 
     group_n, _, ri_id = histogram(group_id, min=0)
     gi_use = np.nonzero(group_n)

PyFHD/calibration/calibration_utils.py

@@ -49,7 +49,7 @@ def vis_extract_autocorr(
     if autocorr_i.size > 0:
         auto_tile_i = obs["baseline_info"]["tile_a"][autocorr_i] - 1
         # As auto_tile_i is used for indexing we need to make it an integer array
-        auto_tile_i = auto_tile_i.astype(np.integer)
+        auto_tile_i = auto_tile_i.astype(int)
         auto_tile_i_single = np.unique(auto_tile_i)
         # expect it as a list of 2D arrays, so there might be trouble
         if not pyfhd_config["cal_time_average"]:
@@ -1557,6 +1557,10 @@ def vis_baseline_hist(
                 wh_noflag = np.where(np.abs(model_vals) > 0)[0]
                 if wh_noflag.size > 0:
                     inds = inds[wh_noflag]
+                    # inds changed so we need to update model_vals
+                    # otherwise it has a different shape than vis_cal_use below
+                    # causing errors
+                    model_vals = (vis_model_arr[pol_i]).flatten()[inds]
                 else:
                     continue
                 # if Keyword_Set(calibration_visibilities_subtract) THEN BEGIN

PyFHD/data_setup/uvfits.py

@@ -88,17 +88,33 @@ def extract_header(
     pyfhd_header["frequency_array"] = (
         np.arange(pyfhd_header["n_freq"]) - freq_ref_i
     ) * pyfhd_header["freq_res"] + pyfhd_header["freq_ref"]
+    # the following is guaranteed from the uvfits memo (AIPS memo 117), logic
+    # stolen from pyuvdata. The uvfits memo is available in the pyuvdata repo
+    # under docs/resources and on the NRAO website.
+    if params_header["naxis"] == 7:
+        ra_axis = 6
+        dec_axis = 7
+    else:
+        ra_axis = 5
+        dec_axis = 6
+    # obsra/obsdec/ra/dec are not standard uvfits keywords
     try:
         pyfhd_header["obsra"] = params_header["obsra"]
     except KeyError:
         logger.warning("OBSRA not found in UVFITS file")
-        pyfhd_header["obsra"] = params_header["ra"]
+        if "ra" in params_header:
+            pyfhd_header["obsra"] = params_header["ra"]
+        else:
+            pyfhd_header["obsra"] = params_header[f"CRVAL{ra_axis}"]
 
     try:
         pyfhd_header["obsdec"] = params_header["obsdec"]
     except KeyError:
         logger.warning("OBSDEC not found in UVFITS file")
-        pyfhd_header["obsdec"] = params_header["dec"]
+        if "dec" in params_header:
+            pyfhd_header["obsdec"] = params_header["dec"]
+        else:
+            pyfhd_header["obsdec"] = params_header[f"CRVAL{dec_axis}"]
     # Put in locations of instrument from FITS file or from Astropy site data
     # If you want to see the list of current site names using EarthLocation.get_site_names()
     # If you want to use PyFHD with HERA in the future
@@ -113,6 +129,8 @@ def extract_header(
         # Can also do MWA or Murchison Widefield Array
         location = EarthLocation("mwa")
 
+    # These are all non-standard uvfits keywords. This information is stored in
+    # the antenna table. See pyuvdata for the right way to do this.
     try:
         pyfhd_header["lon"] = params_header["lon"]
     except KeyError:

PyFHD/gridding/gridding_utils.py

@@ -300,7 +300,7 @@ def dirty_image_generate(
     logger: Logger,
     uniform_filter_uv: NDArray[np.float64] | None = None,
     mask: NDArray[np.integer] | None = None,
-    baseline_threshold: int | float = 0,
+    baseline_threshold: int | float | None = None,
     normalization: float | NDArray[np.float64] | None = None,
     resize: int | None = None,
     width_smooth: int | float | None = None,
@@ -330,7 +330,7 @@ def dirty_image_generate(
     mask : NDArray[np.integer] | None, optional
         A 2D {u,v} mask to apply before image creation, by default None
     baseline_threshold : int | float, optional
-        The maximum baseline length to include in units of pixels, by default 0
+        The maximum baseline length to include in units of pixels, default None
     normalization : float | NDArray[np.float64] | None, optional
         A value by which to normalize the image by, by default None
     resize : int | None, optional

PyFHD/plotting/image.py

@@ -176,7 +176,7 @@ def quick_image(
         count_missing = len(wh_missing[0])
         if count_missing > 0:
             image[wh_missing] = np.nan
-            missing_color = 0
+        missing_color = 0
     else:
         count_missing = 0
         wh_missing = None

PyFHD/pyfhd.py

@@ -134,8 +134,12 @@ def main():
 
     pyfhd_successful = False
     try:
-        checkpoint_name = pyfhd_config["description"]
-        if pyfhd_config["description"] is None or pyfhd_config["description"] == "":
+        if (
+            pyfhd_config["description"] is not None
+            and pyfhd_config["description"] != ""
+        ):
+            checkpoint_name = pyfhd_config["description"] + "_" + pyfhd_config["obs_id"]
+        else:
             checkpoint_name = pyfhd_config["obs_id"]
 
         if pyfhd_config["save_checkpoints"]:
@@ -144,9 +148,49 @@ def main():
             )
             pyfhd_config["checkpoint_dir"].mkdir(exist_ok=True)
 
+            obs_checkpoint_file = Path(
+                pyfhd_config["checkpoint_dir"],
+                f"{checkpoint_name}_obs_checkpoint.h5",
+            )
+            if pyfhd_config["obs_checkpoint"] and not obs_checkpoint_file.exists():
+                logger.warning(
+                    "obs_checkpoint is set but obs checkpoint file does not exist. Recalculating obs."
+                )
+                pyfhd_config["obs_checkpoint"] = False
+
+            cal_checkpoint_file = Path(
+                pyfhd_config["checkpoint_dir"],
+                f"{checkpoint_name}_calibrate_checkpoint.h5",
+            )
+            if (
+                pyfhd_config["calibrate_checkpoint"]
+                and not cal_checkpoint_file.exists()
+            ):
+                logger.warning(
+                    "calibrate_checkpoint is set but cal checkpoint file does not exist. Recalculating cal."
+                )
+                pyfhd_config["calibrate_checkpoint"] = False
+
+            grid_checkpoint_file = Path(
+                pyfhd_config["checkpoint_dir"],
+                f"{checkpoint_name}_gridding_checkpoint.h5",
+            )
+            if (
+                pyfhd_config["gridding_checkpoint"]
+                and not grid_checkpoint_file.exists()
+            ):
+                logger.warning(
+                    "gridding_checkpoint is set but grid checkpoint file does not exist. Recalculating grid."
+                )
+                pyfhd_config["gridding_checkpoint"] = False
+        else:
+            pyfhd_config["obs_checkpoint"] = False
+            pyfhd_config["calibrate_checkpoint"] = False
+            pyfhd_config["gridding_checkpoint"] = False
+
         if (
-            pyfhd_config["obs_checkpoint"] is None
-            and pyfhd_config["calibrate_checkpoint"] is None
+            not pyfhd_config["obs_checkpoint"]
+            and not pyfhd_config["calibrate_checkpoint"]
         ):
             header_start = time.time()
             # Get the header
@@ -217,10 +261,7 @@ def main():
                     "vis_weights": vis_weights,
                 }
                 save(
-                    Path(
-                        pyfhd_config["checkpoint_dir"],
-                        f"{checkpoint_name}_obs_checkpoint.h5",
-                    ),
+                    obs_checkpoint_file,
                     checkpoint,
                     "obs_checkpoint",
                     logger=logger,
@@ -231,11 +272,8 @@ def main():
                 )
         else:
             # Load the checkpoint and initialize the required variables
-            if (
-                pyfhd_config["obs_checkpoint"]
-                and Path(pyfhd_config["obs_checkpoint"]).exists()
-            ):
-                obs_checkpoint = load(pyfhd_config["obs_checkpoint"], logger=logger)
+            if pyfhd_config["obs_checkpoint"]:
+                obs_checkpoint = load(obs_checkpoint_file, logger=logger)
                 obs = obs_checkpoint["obs"]
                 params = obs_checkpoint["params"]
                 vis_arr = obs_checkpoint["vis_arr"]
@@ -247,11 +285,8 @@ def main():
 
         # If the calibration checkpoint exists, load it now before loading in the beam
         # to get the observation metadata and visibility parameters
-        if (
-            pyfhd_config["calibrate_checkpoint"] is not None
-            and Path(pyfhd_config["calibrate_checkpoint"]).exists()
-        ):
-            cal_checkpoint = load(pyfhd_config["calibrate_checkpoint"], logger=logger)
+        if pyfhd_config["calibrate_checkpoint"]:
+            cal_checkpoint = load(cal_checkpoint_file, logger=logger)
             obs = cal_checkpoint["obs"]
             params = cal_checkpoint["params"]
             vis_arr = cal_checkpoint["vis_arr"]
@@ -272,7 +307,10 @@ def main():
         )
 
         # Check if the calibrate checkpoint has been used, if not run the calibration steps
-        if pyfhd_config["calibrate_checkpoint"] is None:
+        if (
+            not pyfhd_config["calibrate_checkpoint"]
+            and not pyfhd_config["gridding_checkpoint"]
+        ):
             if pyfhd_config["deproject_w_term"] is not None:
                 w_term_start = time.time()
                 vis_arr = simple_deproject_w_term(
@@ -394,10 +432,7 @@ def main():
                         "cal": cal,
                     }
                     save(
-                        Path(
-                            pyfhd_config["checkpoint_dir"],
-                            f"{checkpoint_name}_calibrate_checkpoint.h5",
-                        ),
+                        cal_checkpoint_file,
                         checkpoint,
                         "calibrate_checkpoint",
                         logger=logger,
@@ -461,18 +496,15 @@ def main():
             pyfhd_successful = True
             sys.exit(0)
 
-        if (
+        if "image_info" not in psf or (
             psf["image_info"]["image_power_beam_arr"] is not None
             and psf["image_info"]["image_power_beam_arr"].shape == 1
         ):
             # Turn off beam_per_baseline if image_power_beam_arr is
             # only one value
             pyfhd_config["beam_per_baseline"] = False
 
-        if (
-            pyfhd_config["recalculate_grid"]
-            and pyfhd_config["gridding_checkpoint"] is None
-        ):
+        if pyfhd_config["recalculate_grid"] or not pyfhd_config["gridding_checkpoint"]:
             grid_start = time.time()
             image_uv = np.empty(
                 (obs["n_pol"], obs["elements"], obs["dimension"]), dtype=np.complex128
@@ -535,6 +567,8 @@ def main():
                 if vis_model_arr is not None:
                     model_uv = crosspol_reformat(model_uv)
             if pyfhd_config["gridding_plots"]:
+                # TODO: move this after the checkpointing so an error in plotting
+                # doesn't require rerunning gridding.
                 logger.info(
                     f"Plotting the continuum gridding outputs into {pyfhd_config['output_dir']/'plots'/'gridding'}"
                 )
@@ -557,10 +591,7 @@ def main():
                 if vis_model_arr is not None:
                     checkpoint["model_uv"] = model_uv
                 save(
-                    Path(
-                        pyfhd_config["checkpoint_dir"],
-                        f"{checkpoint_name}_gridding_checkpoint.h5",
-                    ),
+                    grid_checkpoint_file,
                     checkpoint,
                     "gridding_checkpoint",
                     logger=logger,
@@ -572,20 +603,17 @@ def main():
             grid_end = time.time()
             _print_time_diff(grid_start, grid_end, "Visibilities gridded", logger)
         else:
-            if pyfhd_config["gridding_checkpoint"]:
-                grid_checkpoint = load(
-                    pyfhd_config["gridding_checkpoint"], logger=logger
-                )
-                image_uv = grid_checkpoint["image_uv"]
-                weights_uv = grid_checkpoint["weights_uv"]
-                variance_uv = grid_checkpoint["variance_uv"]
-                uniform_filter_uv = grid_checkpoint["uniform_filter_uv"]
-                if "model_uv" in grid_checkpoint:
-                    model_uv = grid_checkpoint["model_uv"]
-                del grid_checkpoint
-                logger.info(
-                    f"Checkpoint Loaded: The Gridded UV Planes loaded from {Path(pyfhd_config['output_dir'], 'gridding_checkpoint.h5')}"
-                )
+            grid_checkpoint = load(grid_checkpoint_file, logger=logger)
+            image_uv = grid_checkpoint["image_uv"]
+            weights_uv = grid_checkpoint["weights_uv"]
+            variance_uv = grid_checkpoint["variance_uv"]
+            uniform_filter_uv = grid_checkpoint["uniform_filter_uv"]
+            if "model_uv" in grid_checkpoint:
+                model_uv = grid_checkpoint["model_uv"]
+            del grid_checkpoint
+            logger.info(
+                f"Checkpoint Loaded: The Gridded UV Planes loaded from {Path(pyfhd_config['output_dir'], 'gridding_checkpoint.h5')}"
+            )
 
         # Call quickview to save the all the variables if set in the config. Also create dirty images and save
         # FITS files with the dirty images on a per polarization basis