Skip to content

Commit

Permalink
Update for new deps (#80)
Browse files Browse the repository at this point in the history
* Version bumps

* Bump verions

* Fix beamcon API

* Flag the rms!!

* Fix prefect verions

* Bump versions

* Bump racs-tools

* API fixes

* Bump racs-tools
  • Loading branch information
AlecThomson authored Dec 8, 2024
1 parent b8e447f commit 90e6042
Show file tree
Hide file tree
Showing 4 changed files with 63 additions and 37 deletions.
65 changes: 45 additions & 20 deletions arrakis/imager.py
Original file line number Diff line number Diff line change
Expand Up @@ -580,14 +580,31 @@ def make_cube(
logger = get_run_logger()

logger.info(f"Creating cube for {pol=} {image_set.ms=}")
image_list = image_set.image_lists[pol]
image_list = [Path(i) for i in image_set.image_lists[pol]]

image_type = "restored" if aux_mode is None else aux_mode

# Create a cube name
old_name = image_list[0]
out_dir = os.path.dirname(old_name)
old_base = os.path.basename(old_name)
new_base = old_base
b_idx = new_base.find("beam") + len("beam") + 2
sub = new_base[b_idx:]
new_base = new_base.replace(sub, ".conv.fits")
new_base = new_base.replace("image", f"image.{image_type}.{pol.lower()}")
new_name = Path(out_dir) / new_base

# First combine images into cubes
hdu_list, freqs = combine_fits(file_list=image_list, create_blanks=True)
new_header = hdu_list[0].header
data_cube = hdu_list[0].data
_ = combine_fits(
file_list=image_list,
out_cube=new_name,
create_blanks=True,
overwrite=True,
)
with fits.open(new_name, mode="denywrite", memmap=True) as hdu_list:
new_header = hdu_list[0].header
data_cube = hdu_list[0].data

# Add pol angle to header
new_header["INSTRUMENT_RECEPTOR_ANGLE"] = (
Expand All @@ -610,16 +627,6 @@ def make_cube(
# Calculate rms noise
rmss_arr = mad_std(data_cube, axis=(1, 2, 3), ignore_nan=True)

# Create a cube name
old_name = image_list[0]
out_dir = os.path.dirname(old_name)
old_base = os.path.basename(old_name)
new_base = old_base
b_idx = new_base.find("beam") + len("beam") + 2
sub = new_base[b_idx:]
new_base = new_base.replace(sub, ".conv.fits")
new_base = new_base.replace("image", f"image.{image_type}.{pol.lower()}")
new_name = os.path.join(out_dir, new_base)
# Deserialise beam
with open(common_beam_pkl, "rb") as f:
common_beam = pickle.load(f)
Expand All @@ -633,9 +640,9 @@ def make_cube(
# 0 1234.5
# 1 6789.0
# etc.
new_w_name = new_name.replace(
f"image.{image_type}", f"weights.{image_type}"
).replace(".fits", ".txt")
new_w_name = Path(
new_name.as_posix().replace(f"image.{image_type}", f"weights.{image_type}")
).with_suffix(".txt")
data = dict(
Channel=np.arange(len(rmss_arr)),
Weight=1 / rmss_arr**2, # Want inverse variance
Expand Down Expand Up @@ -673,8 +680,26 @@ def get_beam(image_set: ImageSet, cutoff: Optional[float]) -> Path:

# Create a unique hash for the beam log filename
image_hash = hashlib.md5("".join(image_list).encode()).hexdigest()
try:
common_beam, _ = beamcon_2D.get_common_beam(files=image_list, cutoff=cutoff)
except Exception as e:
import traceback
import sys

tbe = traceback.TracebackException.from_exception(e)
logger.error(f"Local {''.join(tbe.format())}")
f = sys.exc_info()[2].tb_frame
f = f.f_back
while f is not None:
tbe.stack.append(
traceback.FrameSummary(
f.f_code.co_filename, f.f_lineno, f.f_code.co_name
)
)
f = f.f_back

common_beam, _ = beamcon_2D.getmaxbeam(files=image_list, cutoff=cutoff)
logger.error(f"Full {''.join(tbe.format())}")
raise e

logger.info(f"The common beam is: {common_beam=}")

Expand Down Expand Up @@ -740,8 +765,8 @@ def smooth_imageset(
for img in pol_images:
logger.info(f"Smoothing {img}")
last_result = executor.submit(
beamcon_2D.worker,
file=img,
beamcon_2D.beamcon_2d_on_fits,
file=Path(img),
outdir=None,
new_beam=common_beam,
conv_mode="robust",
Expand Down
13 changes: 5 additions & 8 deletions arrakis/linmos.py
Original file line number Diff line number Diff line change
Expand Up @@ -123,21 +123,18 @@ def smooth_images(
"""
smooth_dict: Dict[str, ImagePaths] = {}
for stoke, image_list in image_dict.items():
infiles: List[str] = []
infiles: List[Path] = []
for im in image_list.images:
if im.suffix == ".fits":
infiles.append(im.resolve().as_posix())
datadict = beamcon_3D.main(
infile=[im.resolve().as_posix() for im in image_list.images],
infiles.append(im.resolve())
_, _, output_files = beamcon_3D.smooth_fits_cube(
infiles_list=infiles,
uselogs=False,
mode="total",
conv_mode="robust",
suffix="cres",
)
smooth_files: List[Path] = []
for key, val in datadict.items():
smooth_files.append(Path(val["outfile"]))
smooth_dict[stoke] = ImagePaths(smooth_files, image_list.weights)
smooth_dict[stoke] = ImagePaths(output_files, image_list.weights)

return smooth_dict

Expand Down
9 changes: 7 additions & 2 deletions arrakis/rmsynth_oncuts.py
Original file line number Diff line number Diff line change
Expand Up @@ -359,10 +359,15 @@ def sigma_clip_spectra(
for spectrum in stokes_spectra:
filtered_data = spectrum.data.copy()
filtered_data[filter_idx] = np.nan
filtered_rms = spectrum.rms.copy()
# Set the RMS to NaN where the data is NaN!
filtered_rms[filter_idx] = np.nan
filtered_bkg = spectrum.bkg.copy()
filtered_bkg[filter_idx] = np.nan
filtered_spectrum = Spectrum(
data=filtered_data,
rms=spectrum.rms,
bkg=spectrum.bkg,
rms=filtered_rms,
bkg=filtered_bkg,
filename=spectrum.filename,
header=spectrum.header,
)
Expand Down
13 changes: 6 additions & 7 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ dask = "*"
distributed = {git="https://github.com/AlecThomson/distributed", branch="drainclose"}
dask_jobqueue = {version=">=0.8.3", optional=true}
dask_mpi = "*"
FRion = ">=1.1.3"
FRion = ">=1.1.4"
h5py = "*"
ipython = "*"
matplotlib = ">=3.8"
Expand All @@ -46,9 +46,9 @@ pymongo = "*"
pymultinest = "*"
pytest = "*"
python_casacore = "*"
RACS-tools = ">=3.0.5"
RACS-tools = ">=4.1.2"
radio_beam = "*"
RMextract = {git = "https://github.com/lofar-astron/RMextract", optional=true}
RMextract = ">=0.5.1"
schwimmbad = "*"
scipy = "*"
spectral_cube = ">=0.6.3"
Expand All @@ -57,14 +57,13 @@ tqdm = "*"
vorbin = "*"
graphviz = "*"
bokeh = "<3"
prefect = ">=2"
prefect-dask = "*"
prefect = {version=">=2,<3", extras=["dask"]}
RMTable = ">=1.2.1"
RM-Tools = ">=1.4.2"
PolSpectra = ">=1.1.0"
setuptools = "*"
fixms = ">=0.2.6"
fitscube = ">=0.3"
fixms = ">=0.3.3"
fitscube = ">=0.5.4"
psycopg2-binary = "*"
sqlalchemy = "*"
scikit-image = ">=0.23"
Expand Down

0 comments on commit 90e6042

Please sign in to comment.