diff --git a/.github/workflows/pytest.yaml b/.github/workflows/pytest.yaml index f36b1ce..fd2d71a 100644 --- a/.github/workflows/pytest.yaml +++ b/.github/workflows/pytest.yaml @@ -11,7 +11,7 @@ jobs: runs-on: ubuntu-latest steps: - uses: actions/checkout@v2 - - name: Set up Python 3.9 + - name: Set up Python 3.12 uses: actions/setup-python@v2 with: python-version: ${{ matrix.python-version }} diff --git a/.gitignore b/.gitignore index 4b7fd4e..fe64310 100644 --- a/.gitignore +++ b/.gitignore @@ -16,9 +16,12 @@ ignore/ - - +dev_plan.md +*.drawio DRUID.egg-info build backup notepad.ipynb +temp +_* +*.lock \ No newline at end of file diff --git a/DRUID/__init__.py b/DRUID/__init__.py index 0e72604..4e241be 100644 --- a/DRUID/__init__.py +++ b/DRUID/__init__.py @@ -1,2 +1 @@ from .main import sf -from .src import * \ No newline at end of file diff --git a/DRUID/main.py b/DRUID/main.py index 503b3bc..a20bc17 100644 --- a/DRUID/main.py +++ b/DRUID/main.py @@ -1,969 +1,632 @@ -""" -File: main.py -Author: Rhys Shaw -Date: 23/12/2023 -Version: 0.0 -Description: Main file for DRUID -""" +version = "1.0" -version = "0.0-test" import setproctitle setproctitle.setproctitle("DRUID") +import numpy as np +import astropy.io.fits +from astropy.table import Table +import os +import ast +import time +import polars as pl +import polars.selectors as cs +from functools import partial +from multiprocessing import get_context +from multiprocessing import shared_memory +from rich.progress import Progress +import multiprocessing +from scipy.ndimage import gaussian_filter + from .src import utils -from .src import homology_new as homology +from .src.utils import ( + TITLE, + LINK, + GOLD, + RESET, + BOLD, + NOTICE, + ERROR, + WARNING, + CODEBLOCK, + GREEN, + BLACK, +) +from .src import homology from .src import background -from matplotlib import colors from .src import source -import numpy as np -from skimage import measure -from tqdm import tqdm -import matplotlib.pyplot as plt -from astropy.io import fits -from astropy.wcs import WCS -import astropy -import pandas as pd -from multiprocessing import Pool -import time -import os -from scipy import ndimage -import logging - -DRUID_MESSAGE = """ - - -############################################# - -_______ _______ _________ ______ -( __ \ ( ____ )|\ /|\__ __/( __ \ -| ( \ )| ( )|| ) ( | ) ( | ( \ ) -| | ) || (____)|| | | | | | | | ) | -| | | || __)| | | | | | | | | | -| | ) || (\ ( | | | | | | | | ) | -| (__/ )| ) \ \__| (___) |___) (___| (__/ ) -(______/ |/ \__/(_______)\_______/(______/ - - -############################################# +from .src import properties -Detector of astRonomical soUrces in optIcal and raDio images +# Prevent Polars from thread oversubscription during multiprocessing +os.environ["POLARS_MAX_THREADS"] = "1" -Version: {} +DRUID_MESSAGE = rf""" +{TITLE} + _____ _____ _ _ _____ _____ + | __ \| __ \| | | |_ _| __ \ + | | | | |__) | | | | | | | | | | + | | | | _ /| | | | | | | | | | + | |__| | | \ \| |__| |_| |_| |__| | + |_____/|_| \_\\____/|_____|_____/ + +{RESET} + +{BOLD}Detector of astRonomical soUrces in optIcal and raDio images{RESET} + +{GOLD}Version{RESET}: {version} For more information see: -https://github.com/RhysAlfShaw/DRUID - """.format(version) +{LINK}https://github.com/RhysAlfShaw/DRUID{RESET} +""" +global_image = None +global_smoothed_image = None +global_background_map = None +global_background_rms_map = None + +# Keep shared memory objects alive in the worker +shm_img = None +shm_smooth = None +shm_bg = None +shm_rms = None + + +def _worker_init( + shm_img_name, + img_shape, + img_dtype, + shm_smooth_name, + smooth_shape, + smooth_dtype, + shm_bg_name, + bg_shape, + bg_dtype, + shm_rms_name, + rms_shape, + rms_dtype, +): + """ + Initializer for multiprocessing pool. + Attaches to shared memory blocks created by the main process. + """ + global global_image, global_smoothed_image, global_background_map, global_background_rms_map + global shm_img, shm_smooth, shm_bg, shm_rms + + from multiprocessing import shared_memory + import numpy as np + + shm_img = shared_memory.SharedMemory(name=shm_img_name) + global_image = np.ndarray(shape=img_shape, dtype=img_dtype, buffer=shm_img.buf) + + shm_smooth = shared_memory.SharedMemory(name=shm_smooth_name) + global_smoothed_image = np.ndarray( + shape=smooth_shape, dtype=smooth_dtype, buffer=shm_smooth.buf + ) + + shm_bg = shared_memory.SharedMemory(name=shm_bg_name) + global_background_map = np.ndarray( + shape=bg_shape, dtype=bg_dtype, buffer=shm_bg.buf + ) + + shm_rms = shared_memory.SharedMemory(name=shm_rms_name) + global_background_rms_map = np.ndarray( + shape=rms_shape, dtype=rms_dtype, buffer=shm_rms.buf + ) + + +def _worker( + island_info, + analysis_threshold, + lifetime_limit, + lifetime_limit_fraction, + mode=None, + BMAJ=None, + BMIN=None, + EFFRON=None, + EFFGAIN=None, + EXPTIME=None, +) -> pl.DataFrame: + """ + Worker function to compute homology for a single source island. + """ + bbox, position = island_info + min_row, min_col, max_row, max_col = bbox + + # Cutouts extraction + raw_image_cutout = global_image[min_row:max_row, min_col:max_col] + smoothed_image_cutout = global_smoothed_image[min_row:max_row, min_col:max_col] + bg_cutout = global_background_map[min_row:max_row, min_col:max_col] + bg_rms_cutout = global_background_rms_map[min_row:max_row, min_col:max_col] + + local_threshold = bg_cutout + (analysis_threshold * bg_rms_cutout) + + # Island mask derived strictly from the smoothed cutout + island_mask = smoothed_image_cutout > local_threshold + smoothed_cutout_masked = np.where(island_mask, smoothed_image_cutout, 0) + + # Topology computed on smoothed data + cat = homology.compute_homology( + smoothed_cutout_masked, + analysis_threshold=analysis_threshold * np.mean(bg_rms_cutout), + lifetime_limit=lifetime_limit, + lifetime_limit_fraction=lifetime_limit_fraction, + ) + + if cat is not None and not cat.is_empty(): + # Properties utilize BOTH raw and smoothed arrays + cat = properties.calculate_properties( + cat, + raw_image=raw_image_cutout, + smoothed_image=smoothed_cutout_masked, + background=bg_cutout, + background_rms=bg_rms_cutout, + position=position, + analysis_threshold=analysis_threshold, + mode=mode, + BMAJ=BMAJ, + BMIN=BMIN, + EFFRON=EFFRON, + EFFGAIN=EFFGAIN, + EXPTIME=EXPTIME, + ) + + cat = cat.with_columns( + [ + pl.lit(position[0]).alias("Island_Y"), + pl.lit(position[1]).alias("Island_X"), + ] + ) + + return cat -class sf: +class sf: def __init__( self, - image: np.ndarray = None, - image_path: str = None, + image: str | np.ndarray = None, mode: str = None, - pb_path: str = None, - cutup: bool = False, - cutup_size: int = 500, - cutup_buff: int = None, - output: bool = True, - area_limit: int = 5, - smooth_sigma=1, - nproc: int = 1, - GPU: bool = False, + verbose: bool = True, + area_limit: int = 0, + max_area_limit: int = 10000, + smooth_sigma: float = 0, + num_threads: int = 1, + chunksize: int = 10, header: astropy.io.fits.header.Header = None, - Xoff: int = None, - Yoff: int = None, - debug_mode=False, - remove_edge=True, - ) -> None: - """Initialise the DRUID, here general parameters can be set.. - - Args: - image (np.ndarray, optional): _description_. Defaults to None. - image_path (str, optional): _description_. Defaults to None. - mode (str, optional): _description_. Defaults to None. - pb_path (str, optional): _description_. Defaults to None. - cutup (bool, optional): _description_. Defaults to False. - cutup_size (int, optional): _description_. Defaults to 500. - output (bool, optional): _description_. Defaults to True. - area_limit (int, optional): _description_. Defaults to 5. - smooth_sigma (int, optional): _description_. Defaults to 1. - nproc (int, optional): _description_. Defaults to 1. - GPU (bool, optional): _description_. Defaults to False. - header (astropy.io.fits.header.Header, optional): _description_. Defaults to None. - Xoff (int, optional): _description_. Defaults to None. - Yoff (int, optional): _description_. Defaults to None. - - Raises: - ValueError: _description_ - """ - # start up message! - print(DRUID_MESSAGE) - - if debug_mode: - logging.basicConfig(format="%(asctime)s - %(message)s", level=logging.DEBUG) - logging.debug("Debug mode enabled") - else: - logging.basicConfig(format="%(asctime)s - %(message)s", level=logging.INFO) - - self.cutup = cutup - self.output = output - self.image_path = image_path - self.area_limit = area_limit - self.smooth_sigma = smooth_sigma - self.GPU = GPU - self.Xoff = Xoff - self.Yoff = Yoff - self.cutup_buff = cutup_buff - self.remove_edge = remove_edge - - if self.GPU: - - # Lets try importing the GPU stuff, if it fails then we can just use the CPU. - - try: - import cupy as cp - from cupyx.scipy.ndimage import label as cp_label + working_directory: str = "./druid-working-dir", + cache: bool = False, + output_arg: str = "", + no_message: bool = False, + ): + error_msg = f""" + {ERROR}===================================================================={RESET} + {BOLD}DRUID MULTIPROCESSING {ERROR}ERROR!{RESET} - # num_gpus = cp.cuda.runtime.getDeviceCount() - # print(f'Found {num_gpus} GPUs') + It looks like you are running DRUID with `num_threads > 1` without + protecting your execution code. - # if num_gpus > 0: - # #print('GPUs are avalible, GPU functions will now be avalible.') - # GPU_AVALIBLE = True - # else: - # print('No GPUs avalible, using CPU') - # GPU_AVALIBLE = False - except: + Because DRUID uses Python's robust multiprocessing, you must wrap your + top-level code in the `if __name__ == '__main__':` block. - # print('Could not import cupy. DRUID GPU functions will not be avalible') - GPU_AVALIBLE = False + {BOLD}Please update your script to look like this:{RESET} - self.nproc = nproc + {BLACK}from DRUID import sf - self.header = header + def main(): + findmysource = sf(num_threads={num_threads}, ...) + findmysource.set_background(...) + findmysource.phsf(...) - if self.image_path is None: - self.image = image + if __name__ == "__main__": + main() + {ERROR}===================================================================={RESET} + """ - else: - self.image, self.header = utils.open_image(self.image_path) + if multiprocessing.current_process().name != "MainProcess": + raise RuntimeError(error_msg) + if num_threads > 1 and multiprocessing.current_process().name == "MainProcess": + try: + import __main__ - if mode == "Radio": - pass - # self.image = np.pad( - # self.image, ((1, 1), (1, 1)), mode="constant", constant_values=0 - # ) + if not hasattr(__main__, "__file__") or not os.path.exists( + __main__.__file__ + ): + raise RuntimeError( + f"{error_msg} (Cannot verify script safety in interactive environments)" + ) - if self.smooth_sigma != 0: - self.image_smooth = utils.smoothing(self.image, self.smooth_sigma) - logging.info("Image smoothed with sigma = {}".format(self.smooth_sigma)) - else: - self.image_smooth = self.image - self.mode = mode + with open(__main__.__file__, "r", encoding="utf-8") as f: + source_code = f.read() - if self.mode not in ["Radio", "optical", "other"]: - raise ValueError("Mode must be either radio, optical or other.") + tree = ast.parse(source_code) - self.pb_path = pb_path + is_protected = False + for node in tree.body: + if isinstance(node, ast.If): + if isinstance(node.test, ast.Compare): + left = node.test.left + if isinstance(left, ast.Name) and left.id == "__name__": + is_protected = True + break - if self.pb_path is not None: - self.pb_image, self.pb_header = utils.open_image(self.pb_path) + if not is_protected: + raise RuntimeError(error_msg) - if self.cutup: - self.cutup_size = cutup_size - if self.cutup_buff is not None: - self.cutouts_smooth, self.coords = utils.cut_image_buff( - self.image_smooth, cutup_size, buffer_size=self.cutup_buff + except Exception as e: + raise RuntimeError( + f"Failed to validate safe multiprocessing execution: {e}" ) - self.cutouts, self.coords = utils.cut_image_buff( - self.image, cutup_size, buffer_size=self.cutup_buff - ) - else: - self.cutouts_smooth, self.coords = utils.cut_image( - cutup_size, self.image_smooth - ) - self.cutouts, self.coords = utils.cut_image(cutup_size, self.image) - if self.pb_path is not None: - self.pb_cutouts, self.pb_coords = utils.cut_image( - cutup_size, self.pb_image - ) + self.no_message = no_message + if not self.no_message: + print(DRUID_MESSAGE) - else: - self.pb_cutouts = None - self.pb_coords = None - else: - self.cutup_size = None - self.cutouts = None - self.cutouts_smooth = None - self.coords = None - self.pb_cutouts = None - self.pb_coords = None - - def phsf(self, lifetime_limit: float = 0, lifetime_limit_fraction: float = 2): - """Performs the persistent homology source finding algorithm. - - Args: - - lifetime_limit (float): The lifetime limit for the persistent homology algorithm. - - Returns: - - None. The catalogue is stored in the self.catalogue attribute. + self.mode = mode + self.verbose = verbose + self.area_limit = area_limit + self.max_area_limit = max_area_limit + self.smooth_sigma = smooth_sigma + self.num_threads = num_threads + self.chunksize = chunksize + self.header = header + self.cache = cache + self.output_arg = output_arg + self.smoothed_image = None - """ + if image is None: + raise ValueError( + f"{ERROR}No image provided. Please provide a file path or a NumPy array.{RESET}" + ) - if self.cutup == True: + if isinstance(image, str): + try: + self.image, self.header = utils.get_image_from_path(image) + except Exception as e: + raise ValueError( + f"{ERROR}Could not load image from path{RESET}: {image}" + ) from e + elif isinstance(image, np.ndarray): + self.image = image + self.header = header + else: + raise TypeError( + f"{ERROR}Image must be a file path (str) or a NumPy array (np.ndarray).{RESET}" + ) + self.working_directory = working_directory + if self.working_directory: + if not os.path.exists(working_directory): + os.makedirs(working_directory) - catalogue_list = [] - IDoffset = 0 - for i, cutout in tqdm( - enumerate(self.cutouts_smooth), - total=len(self.cutouts_smooth), - desc="Processing Cutouts", - disable=not self.output, - ): + self.BMAJ, self.BMIN = None, None + self.EFFRON, self.EFFGAIN, self.EXPTIME = None, None, None + if self.mode == "radio" and self.header: + try: + self.BMAJ = self.header.get("BMAJ") + self.BMIN = self.header.get("BMIN") + except KeyError: print( - "Computing for Cutout number :{}/{}".format( - i + 1, len(self.cutouts_smooth) - ) + f"{WARNING}Warning: Could not find BMAJ or BMIN in header.{RESET}" ) - - catalogue = homology.compute_ph_components( - cutout, - self.local_bg, - analysis_threshold_val=self.analysis_threshold_val, - lifetime_limit=lifetime_limit, - output=self.output, - bg_map=self.bg_map, - area_limit=self.area_limit, - GPU=self.GPU, - lifetime_limit_fraction=lifetime_limit_fraction, - mean_bg=self.mean_bg, - IDoffset=IDoffset, - box_size=self.box_size, - detection_threshold=self.sigma, - Cutout_X_offset=self.coords[i][1], - Cutout_Y_offset=self.coords[i][0], - ) - if len(catalogue) == 0: - continue - - IDoffset += len(catalogue) - - catalogue["Y0_cutout"] = self.coords[i][0] - catalogue["X0_cutout"] = self.coords[i][1] - catalogue["x1"] = catalogue["x1"] + self.coords[i][0] - catalogue["x2"] = catalogue["x2"] + self.coords[i][0] - catalogue["y1"] = catalogue["y1"] + self.coords[i][1] - catalogue["y2"] = catalogue["y2"] + self.coords[i][1] - catalogue["bbox1"] = catalogue["bbox1"] + self.coords[i][0] - catalogue["bbox2"] = catalogue["bbox2"] + self.coords[i][1] - catalogue["bbox3"] = catalogue["bbox3"] + self.coords[i][0] - catalogue["bbox4"] = catalogue["bbox4"] + self.coords[i][1] - catalogue["distance_from_center"] = ( - (catalogue["x1"] - cutout.shape[0] / 2) ** 2 - + (catalogue["y1"] - cutout.shape[1] / 2) ** 2 - ) ** 0.5 - catalogue["cutup_number"] = i - catalogue_list.append(catalogue) - self.catalogue = pd.concat(catalogue_list) - # print(self.catalogue) - # remove duplicates and keep the one closest to its cutout centre. - # print('before duplicated removal :',len(self.catalogue)) - # self.catalogue = utils.remove_duplicates(self.catalogue) - # drop any with edge_flag == 1 - # set edge flag False to 0 - self.catalogue["edge_flag"] = self.catalogue["edge_flag"].astype(int) - - if self.remove_edge: - self.catalogue = self.catalogue[self.catalogue.edge_flag != 1] - self.catalogue = self.catalogue.sort_values( - by=["distance_from_center"], ascending=True - ) - self.catalogue = self.catalogue.drop_duplicates( - subset=["x1", "y1", "Birth"], keep="first" + elif self.mode == "optical" and self.header: + try: + self.EFFRON = self.header.get("EFFRON") + self.EFFGAIN = self.header.get("EFFGAIN") + self.EXPTIME = self.header.get("EXPTIME") + except KeyError: + print( + f"{WARNING}Warning: Could not find EFFRON, EFFGAIN, or EXPTIME.{RESET}" ) - else: - - for i, row in catalogue.iterrows(): - if row.edge_flag == 1: - if row.bbox1 == 0: - row.bbox1 = 1 - if row.bbox2 == 0: - row.bbox2 = 1 - if row.bbox3 == self.image.shape[0]: - row.bbox3 = self.image.shape[0] - 1 - if row.bbox4 == self.image.shape[1]: - row.bbox4 = self.image.shape[1] - 1 - - self.catalogue.at[i, "bbox1"] = row.bbox1 - self.catalogue.at[i, "bbox2"] = row.bbox2 - self.catalogue.at[i, "bbox3"] = row.bbox3 - self.catalogue.at[i, "bbox4"] = row.bbox4 - else: - IDoffset = 0 - catalogue = homology.compute_ph_components( - self.image_smooth, - self.local_bg, - analysis_threshold_val=self.analysis_threshold_val, - lifetime_limit=lifetime_limit, - output=self.output, - bg_map=self.bg_map, - area_limit=self.area_limit, - GPU=self.GPU, - lifetime_limit_fraction=lifetime_limit_fraction, - mean_bg=self.mean_bg, - IDoffset=IDoffset, - box_size=self.cutup_size, - detection_threshold=self.sigma, + def phsf(self, lifetime_limit: float = 0.0, lifetime_limit_fraction: float = 1.0): + if ( + getattr(self, "background_map", None) is None + or getattr(self, "background_rms_map", None) is None + ): + raise ValueError( + f"{ERROR}Background maps must be set before running source finding.{RESET}" ) - self.catalogue = catalogue - - self.catalogue["Y0_cutout"] = 0 - self.catalogue["X0_cutout"] = 0 - self.catalogue["edge_flag"] = self.catalogue["edge_flag"].astype(int) - - if self.remove_edge: - self.catalogue = self.catalogue[self.catalogue.edge_flag != 1] - else: - for i, row in catalogue.iterrows(): - if row.edge_flag == 1: - if row.bbox1 == 0: - row.bbox1 = 1 - if row.bbox2 == 0: - row.bbox2 = 1 - if row.bbox3 == self.image_smooth.shape[0]: - row.bbox3 = self.image_smooth.shape[0] - 1 - if row.bbox4 == self.image_smooth.shape[1]: - row.bbox4 = self.image_smooth.shape[1] - 1 - - self.catalogue.at[i, "bbox1"] = row.bbox1 - self.catalogue.at[i, "bbox2"] = row.bbox2 - self.catalogue.at[i, "bbox3"] = row.bbox3 - self.catalogue.at[i, "bbox4"] = row.bbox4 - - # print(self.catalogue) - - self.catalogue = self.catalogue.sort_values(by=["lifetime"], ascending=False) - # print('after duplicate removal :',len(self.catalogue)) - - # do enclosed_i evaluation with the bounding box to ensure we dont use the whole image. - # plt.figure(figsize=(20,20)) - # plt.imshow(self.image,cmap='gray_r',norm=colors.LogNorm(clip=True)) - # plt.scatter(self.catalogue.y1,self.catalogue.x1,c='r',marker='x') - # # # # plot the bounding boxes - # for i, row in self.catalogue.iterrows(): - # ymin = row.bbox1 - # ymax = row.bbox3 - # xmin = row.bbox2 - # xmax = row.bbox4 - # plt.plot([xmin,xmax,xmax,xmin,xmin],[ymin,ymin,ymax,ymax,ymin],c='r') - # plt.savefig('test.png') - # time.sleep(3) - enclosed_i_list = [] t0 = time.time() - for i in tqdm( - range(0, len(self.catalogue)), - total=len(self.catalogue), - desc="Calculating enclosed_i", - disable=not self.output, - ): - row = self.catalogue.iloc[i] - x1 = row.x1 - row.bbox1 + 1 - y1 = row.y1 - row.bbox2 + 1 - Birth = row.Birth - Death = row.Death - # is this a new row? - # if row.new_row == 1: - - img = self.image_smooth[ - int(row.bbox1) - 1 : int(row.bbox3) + 1, - int(row.bbox2) - 1 : int(row.bbox4) + 1, - ] - # reduce the cat to just the sources in the bounding box. - - cat = self.catalogue[self.catalogue["x1"] > int(row.bbox1)] - cat = cat[cat["x1"] < int(row.bbox3)] - cat = cat[cat["y1"] > int(row.bbox2)] - cat = cat[cat["y1"] < int(row.bbox4)] - cat["x1"] = cat["x1"] - int(row.bbox1) + 1 - cat["y1"] = cat["y1"] - int(row.bbox2) + 1 - - # plt.imshow(img,cmap='gray_r',norm=colors.LogNorm(clip=True,vmin=1E-13,vmax=1E-9)) - # plt.scatter(cat.y1,cat.x1,c='r',marker='x') - # plt.savefig('test.png') - # plt.close() - # # sleep - # time.sleep(2) - - if self.GPU == True: - import cupy as cp - - # this is not the best way to deal with this. We should crop the gpu version of the image. - img_gpu = cp.asarray(img, dtype=cp.float64) - enclosed_i = homology.make_point_enclosure_assoc_GPU( - 0, x1, y1, Birth, Death, cat, img_gpu - ) - enclosed_i_list.append(enclosed_i) - else: - # print(self.catalogue) - enclosed_i = homology.make_point_enclosure_assoc_CPU( - 0, x1, y1, Birth, Death, cat, img - ) - enclosed_i_list.append(enclosed_i) - - # print('enclosed_i calculated! t='+str(time.time()-t0)+' s') - self.catalogue["enclosed_i"] = enclosed_i_list - # print(self.catalogue) - # print("Enclosed_i calculated! TESTESTEST") - # correct for first destruction - # t0_parent_tag = time.time() - # self.catalogue = homology.parent_tag_func_vectorized_new(self.catalogue) - # t1_parent_tag = time.time() - # print("Time to assign parent tags: ", t1_parent_tag - t0_parent_tag) - # print(self.catalogue) - # print(len(self.catalogue)) - t0_correct_firs = time.time() - # print('BEfore',len(self.catalogue)) - self.catalogue = homology.correct_first_destruction( - self.catalogue, output=not self.output - ) - t1_correct_firs = time.time() - print("Time to correct first destruction: ", t1_correct_firs - t0_correct_firs) - # print('Time to correct first destruction: ',t1_correct_firs-t0_correct_firs) - # print('After',len(self.catalogue)) - # parent tag - # print("Assigning parent tags..") - t0_parent_tag = time.time() - self.catalogue = homology.parent_tag_func_optimized(self.catalogue) - t1_parent_tag = time.time() - - print("Time to assign parent tags: ", t1_parent_tag - t0_parent_tag) - - t0_classify = time.time() - self.catalogue["Class"] = self.catalogue.apply(homology.classify_single, axis=1) - t1_classify = time.time() - print("Time to classify sources: ", t1_classify - t0_classify) - - def set_background( - self, - detection_threshold, - analysis_threshold, - set_bg=None, - bg_map_bool=False, - box_size=None, - mode="mad_std", - smooth=True, - ): - self.sigma = detection_threshold - self.analysis_threshold = analysis_threshold - self.bg_map = bg_map_bool - self.box_size = box_size - self.bgmode = mode - self.set_bg = set_bg - self.bg_map_bool = bg_map_bool - # mode should be MAD_Std, RMS or other. - - if smooth == True: - img = self.image_smooth + # Apply structural smoothing before thresholding + if self.smooth_sigma > 0: + if self.verbose: + print( + f"{NOTICE}Applying Gaussian smoothing with sigma={self.smooth_sigma}...{RESET}" + ) + self.smoothed_image = gaussian_filter(self.image, sigma=self.smooth_sigma) else: - img = self.image - - if mode == "Radio": - # old verion was called radio. - bgmode = "mad_std" - - # need to account dor the cutputs if not usinh bg_map. + self.smoothed_image = self.image - # bg_map and cutup are require only the same code. + if self.verbose: + print(f"{NOTICE}Thresholding to find source islands...{RESET}") - # bg_map and no cuput is the same as bg_map and cutup. - - # no cutup and no bg_map is just one estimation for the whole image. - - if self.cutup == True: + t0 = time.time() + source_islands = source.create_source_islands( + self.smoothed_image, + self.background_map, + self.background_rms_map, + detection_threshold=self.detection_threshold, + analysis_threshold=self.analysis_threshold, + area_limit=self.area_limit, + max_area_limit=self.max_area_limit, + verbose=self.verbose, + ) + t1 = time.time() - # we want to do the bg_map but for the cutout dims. as the box size is in pixels. - self.bg_map = True - bg_map_bool = True - # which is smaller the cutout size or the box size? - if box_size is None: - box_size = self.cutup_size - else: - if box_size > self.cutup_size: - box_size = self.cutup_size - else: - pass - - self.box_size = box_size - - if bg_map_bool == True: - # print('Creating a background map. Inputed Box size = ',box_size) - # these will be returned as arrays like a map. - std, mean_bg = background.calculate_background_map( - img, box_size, mode=self.bgmode + if self.verbose: + print(f"{NOTICE}Thresholding took {t1 - t0:.2f} seconds.{RESET}") + print( + f"{NOTICE}Found {len(source_islands['bboxes'])} source islands.{RESET}" ) - # print('Background map created.') - # print('Mean Background across cutouts: ', np.nanmean(std)) - # print('Median of bg distribution: ', np.nanmean(mean_bg)) - else: - # print('Not creating a background map.') - std, mean_bg = background.calculate_background(img, mode=self.bgmode) - # print('Background set to: ',std) - # print('Background mean set to: ',mean_bg) - - if set_bg is not None: - # set bg should be a tuple of (std,mean_bg) - # print('User has set the background.') - std = set_bg[0] - mean_bg = set_bg[1] - - self.local_bg = std * self.sigma - self.analysis_threshold_val = std * self.analysis_threshold - self.mean_bg = mean_bg - - def set_background_old( - self, - detection_threshold: float, - analysis_threshold, - set_bg: float = None, - bg_map: bool = None, - box_size: int = 10, - mode: str = "Radio", - ): - """Sets the background for the source finding algorithm. - - Args: - detection_threshold (int): _description_ - analysis_threshold (int): _description_ - set_bg (float, optional): _description_. Defaults to None. - bg_map (bool, optional): _description_. Defaults to None. - box_size (int, optional): _description_. Defaults to 10. - mode (str, optional): _description_. Defaults to 'Radio'. - """ + iterable_islands = list( + zip(source_islands["bboxes"], source_islands["positions"]) + ) + iterable_islands.sort( + key=lambda item: (item[0][2] - item[0][0]) * (item[0][3] - item[0][1]), + reverse=True, + ) - self.bg_map = bg_map - self.sigma = detection_threshold - self.analysis_threshold = analysis_threshold - if mode == "Radio": - if self.cutup: - - # loop though each cutout and calculate the local background. - - if bg_map is not None: - - # users wants to use background map so lets make it - local_bg_list = [] - analysis_threshold_list = [] - mean_bg_list = [] - for i, cutout in enumerate(self.cutouts): - local_bg_map, mean_bg = background.radio_background_map( - cutout, box_size - ) - analysis_threshold_list.append( - local_bg_map * self.analysis_threshold - ) - local_bg_list.append(local_bg_map * self.sigma) - mean_bg_list.append(mean_bg) - else: - local_bg_list = [] - analysis_threshold_list = [] - mean_bg_list = [] - for cutout in self.cutouts: - local_bg, mean_bg = background.radio_background(cutout) - analysis_threshold_list.append( - local_bg * self.analysis_threshold - ) - local_bg_list.append(local_bg * self.sigma) - mean_bg_list.append(mean_bg) - local_bg = local_bg_list - analysis_threshold = analysis_threshold_list - mean_bg = mean_bg_list + if not iterable_islands: + if self.verbose: + print( + "{WARNING}Warning: No source islands found. Returning empty catalog.{RESET}" + ) + self.catalog = pl.DataFrame() + return - else: + t0 = time.time() - # Radio background is calculated using the median absolute deviation of the total image. - if bg_map is not None: - local_bg_o, mean_bg = background.radio_background_map( - self.image, box_size - ) - local_bg = local_bg_o * self.sigma - analysis_threshold = local_bg_o * self.analysis_threshold - else: - local_bg_o, mean_bg = background.radio_background(self.image) - local_bg = local_bg_o * self.sigma - analysis_threshold = local_bg_o * self.analysis_threshold - - if mode == "Optical": - # Optical background is calculated using a random sample of pixels - mean_bg, std_bg = background.optical_background(nsamples=1000) - local_bg = mean_bg + self.sigma * std_bg - analysis_threshold = mean_bg + std_bg * self.analysis_threshold - - if mode == "other": - # If the user has a custom background function, they can pass it in here. - local_bg = set_bg * self.sigma - analysis_threshold = local_bg * self.analysis_threshold - # print('Background set to: ',local_bg) - # print('Analysis threshold set to: ',analysis_threshold) - - self.analysis_threshold_val = analysis_threshold - self.local_bg = local_bg - self.mean_bg = mean_bg - # print(self.mean_bg) - - # if bg_map: - # print('Using bg_map for analysis.') - # else: - # if self.cutup: - - # print('Mean Background across cutouts: ', np.nanmean(self.local_bg)) - # print('Median of bg distribution: ', np.nanmean(self.mean_bg)) - - # else: - # print('Background set to: ',self.local_bg) - - def source_characterising(self, use_gpu: bool = False): - """Source Characterising function. This function takes the catalogue and the image and calculates the source properties. - - Args: - use_gpu (bool, optional): Option to use the GPU True to use and False to not, - requires cupy module and a avalible GPU. Defaults to False. - """ - self.set_background( - detection_threshold=self.sigma, + worker_func = partial( + _worker, analysis_threshold=self.analysis_threshold, - set_bg=self.set_bg, - bg_map_bool=self.bg_map_bool, - box_size=self.box_size, - mode=self.bgmode, - smooth=False, - ) - - self.catalogue, self.polygons = source.measure_source_properties( - use_gpu=use_gpu, - catalogue=self.catalogue, - cutout=self.image, - smooth_cutout=self.image_smooth, - background_map=self.local_bg, - output=self.output, - cutupts=self.cutouts, + lifetime_limit=lifetime_limit, + lifetime_limit_fraction=lifetime_limit_fraction, mode=self.mode, - header=self.header, - sigma=self.sigma, + BMAJ=self.BMAJ, + BMIN=self.BMIN, + EFFRON=self.EFFRON, + EFFGAIN=self.EFFGAIN, + EXPTIME=self.EXPTIME, ) - if self.Xoff is not None: - # correct for the poistion of the cutout. when using cutout from a larger image. - self.catalogue["Xc"] = self.catalogue["Xc"] + self.Xoff - self.catalogue["bbox1"] = self.catalogue["bbox1"] + self.Xoff - self.catalogue["bbox3"] = self.catalogue["bbox3"] + self.Xoff - - if self.Yoff is not None: - # correct for the poistion of the cutout. when using cutout from a larger image. - self.catalogue["Yc"] = self.catalogue["Yc"] + self.Yoff - self.catalogue["bbox2"] = self.catalogue["bbox2"] + self.Yoff - self.catalogue["bbox4"] = self.catalogue["bbox4"] + self.Yoff - - if self.header is not None: - # try: - # print('Converting Xc and Yc to RA and DEC') - print(self.catalogue["Xc"], self.catalogue["Yc"]) - Ra, Dec = utils.xy_to_RaDec( - self.catalogue["Xc"], self.catalogue["Yc"], self.header, mode=self.mode - ) - self.catalogue["RA"] = Ra - self.catalogue["DEC"] = Dec - self.catalogue["RA"] = self.catalogue["RA"].astype(float) - self.catalogue["DEC"] = self.catalogue["DEC"].astype(float) - # except: - # pass - - self._set_types_of_dataframe() - - if self.mode == "optical": - - def ABmag(flux): - return -2.5 * np.log10(flux) - - def RONoise(EFFRON, EFFGAIN, EXPTIME, Area): - return np.sqrt(Area) * (EFFRON / EFFGAIN) * EXPTIME - - def SkyNoise(sky): - return np.sqrt(sky) - - def SourceNoise(Flux): - return np.sqrt(Flux) - - def Flux_err(EFFRON, EFFGAIN, EXPTIME, Area, sky, Flux): - return np.sqrt( - RONoise(EFFRON, EFFGAIN, EXPTIME, Area) ** 2 - + SkyNoise(sky) - + SourceNoise(Flux) - ) - - def NOISE(row, local_ng): - return np.sum( - np.random.normal(row["mean_bg"], local_ng, int(row["Area"])) + results = [] + if self.num_threads > 1: + if self.verbose: + print( + f"{NOTICE}Processing in parallel with {self.num_threads} threads.{RESET}" ) + optimal_chunksize = self.chunksize - EFFGAIN = utils.get_EFFGAIN(self.header) - EXPTIME = utils.get_EXPTIME(self.header) - EFFRON = utils.get_EFFRON(self.header) - # print('EFFGAIN: ',EFFGAIN) - # print('EXPTIME: ',EXPTIME) - # print('EFFRON: ',EFFRON) - # print(self.catalogue['Noise']) - self.catalogue["Flux_total_new"] = ( - self.catalogue["Flux_total"] * EFFGAIN * EXPTIME - - self.catalogue["Noise"] * EFFGAIN * EXPTIME - - RONoise(EFFRON, EFFGAIN, EXPTIME, self.catalogue["Area"]) + # Shared memory allocations + shm_img = shared_memory.SharedMemory(create=True, size=self.image.nbytes) + shm_smooth = shared_memory.SharedMemory( + create=True, size=self.smoothed_image.nbytes ) - # print('Flux_total_new: ',self.catalogue['Flux_total_new']) - self.catalogue["Flux_total_err"] = Flux_err( - EFFRON, - EFFGAIN, - EXPTIME, - self.catalogue["Area"], - self.catalogue["Noise"] * EFFGAIN * EXPTIME, - self.catalogue["Flux_total_new"], + shm_bg = shared_memory.SharedMemory( + create=True, size=self.background_map.nbytes ) - # print('Flux_total_err: ',self.catalogue['Flux_total_err']) - self.catalogue["SNR"] = ( - self.catalogue["Flux_total_new"] - / self.catalogue["Flux_total_err"] - / self.catalogue["Area"] + shm_rms = shared_memory.SharedMemory( + create=True, size=self.background_rms_map.nbytes ) - # print('SNR: ',self.catalogue['SNR']) - self.catalogue["MAG_err"] = ( - 1 / self.catalogue["SNR"] - ) # use the approximate error for the magnitude. - self.catalogue["Flux_total_new"] = self.catalogue["Flux_total_new"] / ( - EFFGAIN * EXPTIME - ) - self.catalogue["MAG_flux"] = ABmag(self.catalogue["Flux_total_new"]) - def create_polygons(self, use_gpu=False): - """ - Creates Polygons/contours best when you just want segmentations and not source charateristics. - """ + np.ndarray(self.image.shape, dtype=self.image.dtype, buffer=shm_img.buf)[ + : + ] = self.image[:] + np.ndarray( + self.smoothed_image.shape, + dtype=self.smoothed_image.dtype, + buffer=shm_smooth.buf, + )[:] = self.smoothed_image[:] + np.ndarray( + self.background_map.shape, + dtype=self.background_map.dtype, + buffer=shm_bg.buf, + )[:] = self.background_map[:] + np.ndarray( + self.background_rms_map.shape, + dtype=self.background_rms_map.dtype, + buffer=shm_rms.buf, + )[:] = self.background_rms_map[:] + + with get_context("spawn").Pool( + self.num_threads, + initializer=_worker_init, + initargs=( + shm_img.name, + self.image.shape, + self.image.dtype, + shm_smooth.name, + self.smoothed_image.shape, + self.smoothed_image.dtype, + shm_bg.name, + self.background_map.shape, + self.background_map.dtype, + shm_rms.name, + self.background_rms_map.shape, + self.background_rms_map.dtype, + ), + ) as p: + with Progress(disable=not self.verbose) as progress: + task = progress.add_task( + "[magenta]:mage: Computing...", total=len(iterable_islands) + ) - self.catalogue = source.create_polygons( - use_gpu=use_gpu, - catalogue=self.catalogue, - cutout=self.image, - output=self.output, - cutupts=self.cutouts, - ) + results = [] + for result in p.imap_unordered( + worker_func, iterable_islands, chunksize=optimal_chunksize + ): + results.append(result) + progress.advance(task) # Update the progress bar incrementally + + # Flush memory + shm_img.close() + shm_img.unlink() + shm_smooth.close() + shm_smooth.unlink() + shm_bg.close() + shm_bg.unlink() + shm_rms.close() + shm_rms.unlink() + else: + global global_image, global_smoothed_image, global_background_map, global_background_rms_map + global_image = self.image + global_smoothed_image = self.smoothed_image + global_background_map = self.background_map + global_background_rms_map = self.background_rms_map + + # Single-threaded Rich progress bar implementation + with Progress(disable=not self.verbose) as progress: + task = progress.add_task( + f"[magenta]:mage: Computing...", total=len(iterable_islands) + ) - def temp_create_polygon_workaround(self): - """ - This is a temporary workaround for the polygon creation. It creates a polygon in the bounding box of the source. This is not ideal but it is much faster than the original method and allows for the use of the catalogue for other purposes. - """ - polygons = [] - for index, row in tqdm( - self.catalogue.iterrows(), - total=len(self.catalogue), - desc="Creating polygons", - ): - contour = utils._get_polygons_CPU( - x1=row.x1, - y1=row.y1, - birth=row.Birth, - death=row.Death, - image=self.image, + for island in iterable_islands: + results.append(worker_func(island)) + progress.advance(task) + + results = [res for res in results if res is not None and not res.is_empty()] + if results: + self.catalog = utils.combine_polars_catalogs(results) + # calculate the ra and dec columns if the header is available + # add island offsets to the centroid and contour coordinates + self.catalog = self.catalog.with_columns( + [ + (pl.col("centroid_x") + pl.col("Island_X")).alias("centroid_x"), + (pl.col("centroid_y") + pl.col("Island_Y")).alias("centroid_y"), + ] ) - polygons.append(contour) - - self.polygons = polygons - - def create_polygons_fast(self): - - # since we have a bounding box, we can just create a polygon in the bounding box. - - polygons = [] - for index, row in tqdm( - self.catalogue.iterrows(), - total=len(self.catalogue), - desc="Creating polygons", - ): - contour = utils._get_polygons_in_bbox( - row.bbox2 - 2, - row.bbox4 + 2, - row.bbox1 - 2, - row.bbox3 + 2, - row.x1, - row.y1, - row.Birth, - row.Death, + # add island offsets to the contour coordinates + self.catalog = ( + self.catalog.with_row_index("__row_id") + .explode("contour") + .with_columns( + pl.concat_list( + [ + pl.col("contour").list.get(0) + pl.col("Island_X"), + pl.col("contour").list.get(1) + pl.col("Island_Y"), + ] + ).alias("contour") + ) + .group_by("__row_id", maintain_order=True) + .agg( + pl.all().exclude("contour").first(), + pl.col("contour"), + ) + .drop("__row_id") ) - polygons.append(contour) - - self.polygons = polygons - - def create_polygons_gpu(self): - """ - Recommended when using GPU acceleration. and not using the bounding box to simplify the polygon creation. - """ - polygons = [] - self.image_gpu = cp.asarray(self.image, dtype=cp.float64) - - for index, row in self.catalogue.iterrows(): - t0 = time.time() - contour = utils._get_polygons_gpu(row.x1, row.y1, row.Birth, row.Death) - t1 = time.time() - # print('Time to create polygon: ',t1-t0) - polygons.append(contour) - self.polygons = polygons - - def _set_types_of_dataframe(self): - """ - Sets the catalogue to the correct data types. This is important to allow for writing data. - Otherwise the datatypes will remain object which will try to be pickled. - - """ - self.catalogue["ID"] = self.catalogue["ID"].astype(int) - self.catalogue["Birth"] = self.catalogue["Birth"].astype(float) - self.catalogue["Death"] = self.catalogue["Death"].astype(float) - self.catalogue["x1"] = self.catalogue["x1"].astype(float) - self.catalogue["y1"] = self.catalogue["y1"].astype(float) - self.catalogue["x2"] = self.catalogue["x2"].astype(float) - self.catalogue["y2"] = self.catalogue["y2"].astype(float) - self.catalogue["Flux_total"] = self.catalogue["Flux_total"].astype(float) - self.catalogue["Flux_peak"] = self.catalogue["Flux_peak"].astype(float) - self.catalogue["Area"] = self.catalogue["Area"].astype(float) - self.catalogue["Xc"] = self.catalogue["Xc"].astype(float) - self.catalogue["Yc"] = self.catalogue["Yc"].astype(float) - self.catalogue["bbox1"] = self.catalogue["bbox1"].astype(float) - self.catalogue["bbox2"] = self.catalogue["bbox2"].astype(float) - self.catalogue["bbox3"] = self.catalogue["bbox3"].astype(float) - self.catalogue["bbox4"] = self.catalogue["bbox4"].astype(float) - self.catalogue["Maj"] = self.catalogue["Maj"].astype(float) - self.catalogue["Min"] = self.catalogue["Min"].astype(float) - self.catalogue["Pa"] = self.catalogue["Pa"].astype(float) - self.catalogue["parent_tag"] = self.catalogue["parent_tag"].astype(float) - self.catalogue["Class"] = self.catalogue["Class"].astype(float) - if self.cutup: - self.catalogue["Y0_cutout"] = self.catalogue["Y0_cutout"].astype(float) - self.catalogue["X0_cutout"] = self.catalogue["X0_cutout"].astype(float) - self.catalogue["SNR"] = self.catalogue["SNR"].astype(float) - self.catalogue["Noise"] = self.catalogue["Noise"].astype(float) - - def plot_sources(self, cmap, figsize=(10, 10), norm="linear", save_path=None): - """Plots the source polygons on the image. - - Args: - cmap (str): matplotlib cmap to use, e.g. 'gray'. See https://matplotlib.org/stable/tutorials/colors/colormaps.html for more info. - figsize (tuple, optional): Desired figure size. Defaults to (10,10). - norm (str, optional): _description_. Defaults to 'linear'. - save_path (str, optional): Save path if you desire to save the figure. Defaults to None. - - """ - plt.figure(figsize=figsize) - plt.imshow(self.image, cmap=cmap, origin="lower", norm=norm) - # plt.scatter(self.catalogue['Xc'],self.catalogue['Yc'],s=10,c='r') - - for i, poly in enumerate(self.polygons): - if poly is not None: - plt.plot(poly[:, 1], poly[:, 0]) - if save_path is not None: - plt.savefig(save_path) - plt.show() - - def save_catalogue(self, save_path, filetype=None, overwrite=False): - """Save Catalogue to a file. - - Args: - save_path (str): Desired path to save the catalogue. - filetype (str, optional): Specify the file type or include the approprate file extention. Defaults to None. - overwrite (bool, optional): Overwrite the save file if the name is the same. Defaults to False. - - """ - - # get the extension from the save_path - # print('Saving Catalogue to file: ',save_path) - fileextention = save_path.split(".")[-1] - - if filetype is None: - filetype = fileextention - - if filetype == "csv": - self.catalogue.to_csv(save_path, index=False, overwrite=overwrite) - # print('Catalogue saved to: ',save_path) - - if filetype == "fits": - from astropy.table import Table - - # print('Saving to fits with astropy') - enclosed_i = self.catalogue["enclosed_i"] - - for i in range(len(enclosed_i)): - for j in range(len(enclosed_i[i])): - enclosed_i[i][j] = int(enclosed_i[i][j]) - if len(enclosed_i[i]) == 0: - enclosed_i[i] = [0] - self.catalogue["enclosed_i"] = enclosed_i - # print(self.catalogue) - t = Table.from_pandas(self.catalogue) - t.write(save_path, overwrite=overwrite) + if self.header is not None: + self.catalog = utils.calculate_radec(self.catalog, self.header) - if filetype == "hdf": - self.catalogue.to_hdf(save_path, key="catalogue", mode="w") - # print('Catalogue saved to: ',save_path) - - if filetype == ("txt" or "ascii"): - self.catalogue.to_csv(save_path, index=False, overwrite=overwrite) - # print('Catalogue saved to: ',save_path) - - def open_catalogue(self, file_path, filetype=None): - from astropy.table import Table + else: + print( + f"{WARNING}Warning{RESET}: No FITS header provided. RA and Dec columns will not be calculated." + ) + self.catalog = self.catalog.with_columns( + pl.lit(None).alias("ra"), pl.lit(None).alias("dec") + ) + desired_order = [ + "ID", + "ra", + "dec", + "centroid_x", + "centroid_y", + "flux", + "flux_peak", + "flux_err", + "bg", + "snr", + "maj", + "min", + "pa", + "area", + "contour", + "lifetime", + "birth", + "death", + "x1", + "y1", + "x2", + "y2", + "encloses", + "new_row", + "parent_tag", + "class", + "lifetimeFrac", + "bbox_min_y", + "bbox_min_x", + "bbox_max_y", + "bbox_max_x", + "Island_X", + "Island_Y", + ] + self.catalog = self.catalog.select( + *desired_order, cs.all().exclude(desired_order) + ) - self.catalogue = Table.read(file_path) + # save catalog to working directory + catalog_file = os.path.join( + self.working_directory, + f"druid_source_catalog_{self.output_arg}", + ) + self.catalog.write_parquet(f"{catalog_file}.parquet") - for i in range(len(self.catalogue)): - self.catalogue["contour"][i] = np.array( - self.catalogue["contour"][i] - ).reshape(-1, 2) + print(f"{NOTICE}Catalog saved to {catalog_file}.parquet{RESET}") + else: + self.catalog = pl.DataFrame() + + t1 = time.time() + if self.verbose: + print(f"{NOTICE}Homology computation took {t1 - t0:.2f} seconds.{RESET}") + print(f"{GREEN}---------------CATALOG SUMMARY---------------------{RESET}") + print(f"Total sources detected: {self.catalog.height}") + print( + f"Number of large sources (area > {self.max_area_limit}): {self.catalog.filter(pl.col('area') > self.max_area_limit).height}" + ) + print("Average Background: ", self.background_map.mean()) + print("Average Background RMS: ", self.background_rms_map.mean()) + print(f"{GREEN}---------------------------------------------------{RESET}") - self.catalogue = self.catalogue.to_pandas() + def set_background( + self, + method: str = "rms", + detection_threshold: int = 5, + analysis_threshold: int = 3, + box_size: tuple = (50, 50), + filter_size: tuple = (3, 3), + kernel_size: int = 3, + ): + if self.verbose: + print(f"{NOTICE}Calculating background map and RMS map...{RESET}") + t0 = time.time() + self.detection_threshold = detection_threshold + self.analysis_threshold = analysis_threshold - def save_polygons_to_ds9(self, filename): - """ - Saves the polygons to a ds9 region file. - """ + bg_file = os.path.join(self.working_directory or "", "background_map.npy") + rms_file = os.path.join(self.working_directory or "", "background_rms_map.npy") - with open(filename, "w") as f: - f.write("# Region file format: DS9 version 4.1\n") - f.write( - 'global color=green dashlist=8 3 width=1 font="helvetica 10 normal roman" select=1 highlite=1 dash=0 fixed=0 edit=1 move=1 delete=1 include=1 source=1\n' + if self.cache and os.path.exists(bg_file) and os.path.exists(rms_file): + if self.verbose: + print(f"{NOTICE}Background maps exist. Loading from disk.{RESET}") + self.background_map = np.load(bg_file) + self.background_rms_map = np.load(rms_file) + else: + self.background_map, self.background_rms_map = ( + background.calculate_background_maps( + self.image, + bg_estimator=method, + box_size=box_size, + filter_size=filter_size, + nsigma=detection_threshold, + kernel_size=kernel_size, + ) ) - for polygon in self.polygons: - f.write("polygon(") - for i, point in enumerate(polygon): - f.write( - "{:.2f},{:.2f}".format(point[1], point[0]) - ) # note this transformation as the index in some CARTA inmages start at -1. - if i < len(polygon) - 1: - f.write(",") - f.write(")\n") - - def save_polygons_to_hdf5(self, filename): - """ - Saves the polygons to a hdf5 file. - """ - import h5py + if self.cache: + np.save(bg_file, self.background_map) + np.save(rms_file, self.background_rms_map) - hf = h5py.File(filename, "w") - for i in range(len(self.catalogue)): - key = self.catalogue["ID"][i] - hf.create_dataset(str(key), data=self.catalogue["contour"][i]) - hf.close() + t1 = time.time() + if self.verbose: + print(f"{NOTICE}Background calculation took {t1 - t0:.2f} seconds.{RESET}") diff --git a/DRUID/src/background.py b/DRUID/src/background.py index 0b80660..9669ca6 100644 --- a/DRUID/src/background.py +++ b/DRUID/src/background.py @@ -1,167 +1,78 @@ -""" - -File: background.py -Author: Rhys Shaw -Date: 23/12/2023 -Version: v1.0 -Description: Functions for calculating the background of an image. - -""" - import numpy as np -from astropy.stats import mad_std, sigma_clipped_stats -import matplotlib.pyplot as plt -from photutils.background import SExtractorBackground - - -def radio_background_map(cutout: np.ndarray, box_size: int): - """ - - This function takes an image and a box size and calculates the radio_background() for each box to create a map of local background. - - """ - - step_size = box_size // 2 - - # initialize the map - map_shape = (cutout.shape[0] // step_size, cutout.shape[1] // step_size) - bg_map = np.zeros(map_shape) - mean_bg_map = np.zeros(map_shape) - - # box - box = np.ones((box_size, box_size)) - - for i in range(0, cutout.shape[0], step_size): - for j in range(0, cutout.shape[1], step_size): - # get the box - box_image = cutout[i : i + box_size, j : j + box_size] - # calculate the radio background - local_bg = mad_std(box_image, ignore_nan=True) - mean_bg = np.nanmedian(box_image) - # set the value in the map - bg_map[i // step_size, j // step_size] = local_bg - mean_bg_map[i // step_size, j // step_size] = mean_bg - - # now upsample the map to the original image size - bg_map = np.repeat(bg_map, step_size, axis=0) - bg_map = np.repeat(bg_map, step_size, axis=1) - - mean_bg_map = np.repeat(mean_bg_map, step_size, axis=0) - mean_bg_map = np.repeat(mean_bg_map, step_size, axis=1) - # shift the map to the correct position - - return bg_map, mean_bg_map - - -def calculate_background_map(image, box_size, mode="mad_std"): - - image_height, image_width = len(image), len(image[0]) - box_sum = 0 - box_mean_bg = np.zeros((image_height // box_size + 1, image_width // box_size + 1)) - box_std_bg = np.zeros((image_height // box_size + 1, image_width // box_size + 1)) - for i in range(image_height // box_size + 1): - for j in range(image_width // box_size + 1): - xmin = i * box_size - ymin = j * box_size - subarray = image[xmin : xmin + box_size, ymin : ymin + box_size] - box_mean_bg[i, j], box_std_bg[i, j] = calculate_background( - subarray, mode=mode - ) - - return box_mean_bg, box_std_bg - - -def get_bg_value_from_result_image(original_location_in_full_image, box_size, bg_map): - i, j = original_location_in_full_image - x = i // box_size - y = j // box_size - if x >= bg_map.shape[0]: - x = bg_map.shape[0] - 1 - if y >= bg_map.shape[1]: - y = bg_map.shape[1] - 1 - - result_value = bg_map[x, y] - return result_value - - -def calculate_background(image, mode="mad_std"): - """ - This function calculates the background of an image. - - Args: - image (np.ndarray): The image. - mode (str, optional): Can choose from mad_std or rms. - - Returns: - background (float): The background of the image. - """ - - if mode == "mad_std" or mode == "rms": - - local_bg, mean_bg = radio_background(image, metric=mode) - - elif mode == "SEX" or mode == "sigma_clip": - - local_bg, mean_bg = get_optical_background_estimate(image, mode) - +from astropy.io import fits +from astropy.stats import sigma_clipped_stats + +from photutils.background import ( + Background2D, + MedianBackground, + StdBackgroundRMS, + MADStdBackgroundRMS, + BiweightLocationBackground, + BiweightScaleBackgroundRMS, + MMMBackground, + MeanBackground, + ModeEstimatorBackground, + SExtractorBackground, + BackgroundBase, +) +from photutils.segmentation import detect_sources + + +def make_source_mask(data, nsigma=3.0, kernel_size=3): + mean, median, std = sigma_clipped_stats(data, sigma=nsigma) + threshold = median + nsigma * std + segm = detect_sources(data, threshold, npixels=kernel_size**2) + if segm is None: + return np.zeros(data.shape, dtype=bool) + return segm.data > 0 + + +def calculate_background_maps( + image, + bg_estimator="median", + box_size=(50, 50), + filter_size=(3, 3), + nsigma=3.0, + kernel_size=3, +): + if isinstance(image, str): + with fits.open(image) as hdul: + data = hdul[0].data + elif isinstance(image, np.ndarray): + data = image else: - # some other method can be added here. - raise ValueError("mode not recognised. Please use mad_std or rms") - - return local_bg, mean_bg - - -def get_optical_background_estimate(image, mode): - """Returns the background estimates from the image. using the SExtractorBackground or Sigma Clipping. - - Args: - image (_type_): _description_ - mode (_type_): _description_ - - Returns: - _type_: _description_ - """ - if mode == "sigma_clip": - mean, median, std = sigma_clipped_stats(image, sigma=3, maxiters=5) - - return std, median - - if mode == "SEX": - bkg = SExtractorBackground() - bkg_sigma = bkg.sigma_clip(image).std() - bkg_meadian = np.median(bkg.sigma_clip(image)) - return bkg_sigma, bkg_meadian - - -def radio_background(image: np.ndarray, metric: str = "mad_std"): - """ - - Returns the local background of an image. - - Args: - image (np.ndarray): The image. - metric (str, optional): Can Choose from mad_std or rms. - - Raises: - ValueError: If metric is not recognised. - - Returns: - - local_bg (float): The local background of the image./ - - """ - - if metric == "mad_std": - - local_bg = mad_std(image, ignore_nan=True) - - elif metric == "rms": # rmse (tends to rms when mean is 0.) - predicted = np.nanmean(image) - local_bg = np.sqrt(np.nanmean((image - predicted) ** 2)) - + raise TypeError("Image must be a path or a numpy array.") + + mask = make_source_mask(data, nsigma=nsigma, kernel_size=kernel_size) + + available_estimators = { + "median": MedianBackground, + "std": StdBackgroundRMS, + "mad_std": MADStdBackgroundRMS, + "rms": StdBackgroundRMS, + "biweightlocation": BiweightLocationBackground, + "biweightscale": BiweightScaleBackgroundRMS, + "mm": MMMBackground, + "mean": MeanBackground, + "mode": ModeEstimatorBackground, + "sex": SExtractorBackground, + } + + if isinstance(bg_estimator, str): + bkg_estimator = available_estimators.get( + bg_estimator.lower(), MedianBackground + )() + elif isinstance(bg_estimator, BackgroundBase): + bkg_estimator = bg_estimator else: - raise ValueError("metric not recognised. Please use mad_std or rms") + bkg_estimator = MedianBackground() - mean_bg = np.nanmedian(image) + bkg = Background2D( + data, + box_size, + filter_size=filter_size, + mask=mask, + bkg_estimator=bkg_estimator, + ) - return local_bg, mean_bg + return bkg.background, bkg.background_rms diff --git a/DRUID/src/homology.py b/DRUID/src/homology.py new file mode 100644 index 0000000..864bed7 --- /dev/null +++ b/DRUID/src/homology.py @@ -0,0 +1,272 @@ +""" +Author: Rhys Shaw +Date: 11-06-2025 +""" + +import cripser +import numpy as np +import polars as pl +from scipy.ndimage import label as scipy_label +from skimage import measure + + +def get_enclosing_mask_CPU(x, y, mask): + """ + Returns the connected components inside the mask starting from the point (x, y). + """ + labeled_mask, _ = scipy_label(mask) + if 0 <= x < mask.shape[1] and 0 <= y < mask.shape[0]: + label_at_pixel = labeled_mask[y, x] + if label_at_pixel != 0: + return labeled_mask == label_at_pixel + return None + + +def _get_polygons_CPU(x1, y1, birth, death, image: np.ndarray): + """ + Returns the polygon of the enclosed area of the point (x,y) in the mask. + """ + image_padded = np.pad(image, pad_width=1, mode="constant", constant_values=0) + mask = (image_padded <= birth) & (image_padded > death) + enclosed_mask = get_enclosing_mask_CPU(int(y1) + 1, int(x1) + 1, mask) + + # Return empty list instead of [0] to maintain consistent Polars schema + if enclosed_mask is None: + return [] + + # 0.5 is mathematically correct to find the boundary of a boolean (0/1) mask + contours = measure.find_contours(enclosed_mask, 0.5) + if not contours: + return [] + + contour = contours[0] + # Shift coordinates back due to padding + contour[:, 0] -= 1 + contour[:, 1] -= 1 + + # skimage returns (row, col). Convert to standard (x, y) for plotting + contour_xy = np.column_stack((contour[:, 1], contour[:, 0])) + + return contour_xy.tolist() + + +def get_mask_CPU(x1, y1, Birth, Death, img): + mask = (img <= Birth) & (img > Death) + return get_enclosing_mask_CPU(int(y1), int(x1), mask) + + +def make_point_enclosure_assoc_CPU(x1, y1, Birth, Death, x1_arr, y1_arr, ids, img): + mask = get_mask_CPU(x1, y1, Birth, Death, img) + if mask is None: + return [] + + # FIX: cripser returns x1 as the row (axis 0) and y1 as the col (axis 1). + # Therefore we index the numpy array using mask[x1, y1]. + valid_coords = mask[x1_arr.astype(int), y1_arr.astype(int)] + return ids[valid_coords].tolist() + + +def correct_first_destruction_pl(df: pl.DataFrame) -> pl.DataFrame: + if "new_row" not in df.columns: + df = df.with_columns(pl.lit(0, dtype=pl.Int8).alias("new_row")) + + islands_to_split = df.filter(pl.col("encloses").list.len() > 1) + if islands_to_split.is_empty(): + return df + + new_rows_base = islands_to_split.join( + df.select(["ID", "death"]), + left_on=pl.col("encloses").list.get(0), + right_on="ID", + how="inner", + suffix="_parent", + ) + + if new_rows_base.is_empty(): + return df + + max_id = df["ID"].max() + num_new_rows = len(new_rows_base) + new_ids = pl.int_range( + start=max_id + 1, + end=max_id + num_new_rows + 1, + dtype=df.schema["ID"], + eager=True, + ) + + new_rows = ( + new_rows_base.with_columns( + ID=new_ids, + death=pl.col("death_parent"), + parent_tag=pl.col("ID_parent"), + new_row=pl.lit(1, dtype=pl.Int8), + encloses=pl.lit(None, dtype=df.schema["encloses"]), + ) + .drop(["ID_parent", "death_parent"]) + .select(df.columns) + ) + + return pl.concat([df, new_rows], how="vertical") + + +def parent_tag_func_pl(df: pl.DataFrame) -> pl.DataFrame: + parents = df.filter(pl.col("encloses").list.len() > 1).select( + pl.col("ID").alias("parent_id"), pl.col("encloses") + ) + + mapping = ( + parents.explode("encloses") + .rename({"encloses": "child_id"}) + .filter(pl.col("child_id") != pl.col("parent_id")) + ) + + df_with_parent_info = df.join( + mapping, left_on="ID", right_on="child_id", how="left" + ) + + return df_with_parent_info.with_columns( + parent_tag=pl.when(pl.col("parent_id").is_not_null()) + .then(pl.col("parent_id")) + .otherwise(pl.col("ID")) + ).drop("parent_id") + + +def assign_ph_class(df: pl.DataFrame) -> pl.DataFrame: + # assign a class based on parent_tag, new_row, and encloses length + + is_new_row = pl.col("new_row") != 0 + has_children = pl.col("encloses").list.len() > 1 + has_parent = pl.col("parent_tag") != pl.col("ID") + + # 2. Chain the logic + df = df.with_columns( + pl.when(is_new_row) + .then(1) + .when(~has_children & ~has_parent) + .then(0) + .when(~has_children & has_parent) + .then(2) + .when(has_children & ~has_parent) + .then(4) + .otherwise(3) + .alias("class") # Replace with your desired column name + ) + + return df + + +def compute_homology( + img: np.ndarray, + analysis_threshold: float, + lifetime_limit: float = 0.0, + lifetime_limit_fraction: float = 1.0, + area_size_threshold: int = 2, +) -> pl.DataFrame: + + pd_data = cripser.computePH(-img, maxdim=0) + columns = ["dim", "birth", "death", "x1", "y1", "z1", "x2", "y2", "z2"] + polar_df = pl.DataFrame(pd_data, schema=columns).drop(["dim", "z1", "z2"]) + + polar_df = polar_df.with_columns( + [(-pl.col("birth")).alias("birth"), (-pl.col("death")).alias("death")] + ) + + polar_df = polar_df.with_columns( + pl.when(pl.col("death") < analysis_threshold) + .then(pl.lit(analysis_threshold)) + .otherwise(pl.col("death")) + .alias("death") + ) + + polar_df = polar_df.with_columns( + (abs(pl.col("death") - pl.col("birth"))).alias("lifetime"), + (pl.col("birth") / pl.col("death")).alias("lifetimeFrac"), + ) + + polar_df = polar_df.filter( + (pl.col("lifetimeFrac") > lifetime_limit_fraction) + & (pl.col("lifetime") > lifetime_limit) + ) + + polar_df = polar_df.filter(pl.col("lifetime") > analysis_threshold) + + if polar_df.is_empty(): + return None + + polar_df = polar_df.with_columns( + pl.when(pl.col("lifetime") == pl.col("lifetime").max()) + .then(pl.lit(0)) + .otherwise(pl.col("death")) + .alias("death") + ) + + # Fast NumPy extraction to avoid iter_rows bottleneck + births = polar_df["birth"].to_numpy() + deaths = polar_df["death"].to_numpy() + x1s = polar_df["x1"].to_numpy() + y1s = polar_df["y1"].to_numpy() + + areas, min_ys, min_xs, max_ys, max_xs = [], [], [], [], [] + + for b, d, x, y in zip(births, deaths, x1s, y1s): + mask = get_mask_CPU(x, y, b, d, img) + if mask is not None: + rows, cols = np.where(mask) + areas.append(mask.sum()) + min_ys.append(rows.min()) + min_xs.append(cols.min()) + max_ys.append(rows.max()) + max_xs.append(cols.max()) + else: + areas.append(0) + min_ys.append(np.nan) + min_xs.append(np.nan) + max_ys.append(np.nan) + max_xs.append(np.nan) + + polar_df = polar_df.with_columns( + [ + pl.Series("area", areas), + pl.Series("bbox_min_y", min_ys), + pl.Series("bbox_min_x", min_xs), + pl.Series("bbox_max_y", max_ys), + pl.Series("bbox_max_x", max_xs), + ] + ) + + polar_df = polar_df.filter(pl.col("area") > area_size_threshold) + + if polar_df.is_empty(): + return None + + polar_df = polar_df.with_columns(pl.Series("ID", range(len(polar_df)))) + + # ---> FIX: Re-extract arrays from the FILTERED DataFrame <--- + # This ensures ids, x1s, and y1s all have the exact same length + ids = polar_df["ID"].to_numpy() + filtered_x1s = polar_df["x1"].to_numpy() + filtered_y1s = polar_df["y1"].to_numpy() + + encloses = [ + make_point_enclosure_assoc_CPU(x, y, b, d, filtered_x1s, filtered_y1s, ids, img) + for b, d, x, y in zip( + polar_df["birth"], polar_df["death"], polar_df["x1"], polar_df["y1"] + ) + ] + polar_df = polar_df.with_columns(pl.Series("encloses", encloses)) + + polar_df = correct_first_destruction_pl(polar_df) + polar_df = parent_tag_func_pl(polar_df) + + contours = [ + _get_polygons_CPU(x, y, b, d, img) + for b, d, x, y in zip( + polar_df["birth"], polar_df["death"], polar_df["x1"], polar_df["y1"] + ) + ] + + polar_df = assign_ph_class(polar_df) + + return polar_df.with_columns( + pl.Series("contour", contours, dtype=pl.List(pl.List(pl.Float64))) + ) diff --git a/DRUID/src/homology/__init__.py b/DRUID/src/homology/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/DRUID/src/homology/homology.py b/DRUID/src/homology/homology.py deleted file mode 100644 index a8cc616..0000000 --- a/DRUID/src/homology/homology.py +++ /dev/null @@ -1,561 +0,0 @@ -""" -File: src/homology.py -author: Rhys Shaw -date: 27-12-2023 -Description: This file contains the functions that deal with calculating - persistence diagrams from a given image. -""" - - -import cripser -import numpy as np -from .src import utils -import time -from functools import partial -import pandas -from tqdm import tqdm -from multiprocessing import Pool, freeze_support -import pdb - -from collections import deque - -try: - import cupy as cp - from cupyx.scipy.ndimage import label as cupy_label - GPU_AVAILABLE = True -except: - GPU_AVAILABLE = False - - - - - - -def parent_tag_func_vectorized(df): - """ - - Vectorised implenetation of parent tag function. - - - Args: - - df: pd.Dataframe - data frame for which we calculate the parent tags. - - - Returns: - - df: pd.Dataframe - Pandas data frame with addition parent tag column. - - """ - - enclosed_i_dict = {row['ID']: set(row['enclosed_i']) for idx, row in df.iterrows()} - - def find_parent_tag(row): - #if row.new_row == 0: # only set it if it is not a new row. - - for ID, enclosed_i_set in enclosed_i_dict.items(): - if row.ID in enclosed_i_set: - return ID - return np.nan - - #else: # if it is a new row then we already know the parent tag. - # return row.parent_tags - - return df.apply(find_parent_tag, axis=1) - - - - - - -def classify_single(row): - - """Classifiying the Rows based on orgin. - - Args: - row: pd.series - The row that is being classified. - - Returns: - Class: int - the Class integer that indiceates the class the row belongs too. - - """ - if row.new_row == 0: - if len(row.enclosed_i) == 0: # no children - if np.isnan(row.parent_tag): # no parent - return 0 # no children, no parent. - else: - return 1 # no child has parent. - else: - if np.isnan(row.parent_tag): - return 2 # has children, no parent. - else: - return 3 # has children, has parent. - else: - return 4 # new row has children. - - - - - - -def make_point_enclosure_assoc(row,pd,img): - """Returns a list of the indices of the points that are enclosed by the mask pd point. - - Args: - row (pd.Series): _description_ - pd (pd.DataFrame): _description_ - img (np.ndarray): _description_ - - Returns: - enclosed_list (list): _description_ - - """ - - mask = utils.get_mask_CPU(row,img) - - encloses = [] - for i in range(len(pd)): - point = pd.iloc[i] - # we dont want to include ourselves - if point['ID'] == row['ID']: - continue - if mask[int(point.x1),int(point.y1)]: - encloses.append(point['ID']) - - return encloses - - - - - -def make_point_enclosure_assoc_GPU(Birth,Death,row,pd,img,img_gpu): - """Returns a list of the ID of the points that are enclosed by the mask pd point. - Uses GPU for computation. - - Args: - Birth (float): _description_ - Death (float): _description_ - row (pd.Series): _description_ - pd (pd.DataFrame): _description_ - img (np.ndarray): _description_ - img_gpu (cp.ndarray): _description_ - - Returns: - enclosed_list (list): _description_ - """ - - mask = utils.get_mask_GPU(Birth,Death,row,img_gpu) - #pdb.set_trace() - mask_coords = np.column_stack((pd['x1'], pd['y1'])) - points_inside_mask = mask[mask_coords[:, 0].astype(int), mask_coords[:, 1].astype(int)] - encloses_vectorized = pd.iloc[points_inside_mask]['ID'].tolist() - # remove self from list - #print(row.ID) - #print(encloses_vectorized) - #encloses_vectorized.remove(row.ID) - #print(encloses_vectorized) - #pdb.set_trace() - return encloses_vectorized - - - -def make_point_enclosure_assoc_GPU_second(Birth,Death,row,pd,img,img_gpu): - """Returns a list of the ID of the points that are enclosed by the mask pd point. - Uses GPU for computation. - - Args: - Birth (float): _description_ - Death (float): _description_ - row (pd.Series): _description_ - pd (pd.DataFrame): _description_ - img (np.ndarray): _description_ - img_gpu (cp.ndarray): _description_ - - Returns: - enclosed_list (list): _description_ - """ - - mask = utils.get_mask_GPU(Birth,Death,row,img_gpu) - #pdb.set_trace() - mask_coords = np.column_stack((pd['x1'], pd['y1'])) - points_inside_mask = mask[mask_coords[:, 0].astype(int), mask_coords[:, 1].astype(int)] - encloses_vectorized = pd.iloc[points_inside_mask]['ID'].tolist() - # remove self from list - #print(row.ID) - #print(encloses_vectorized) - encloses_vectorized.remove(row.ID) - #print(encloses_vectorized) - #pdb.set_trace() - return encloses_vectorized - - - - -def correct_first_destruction(pd,output,img=None,img_gpu=None,GPU=False): - """ - Function for correcting for the First destruction of a parent Island. - - Args: - pd (pd.DataFrame): Input catalogue of sources to correct. - output (bool): True if you want interation logginf with tqdm. - - Returns: - pd (pd.DataFrame): The new Catalogue. - - """ - - pd['new_row'] = 0 - - for i in tqdm(range(0,len(pd)),total=len(pd),desc='Correcting first destruction',disable=output): - - row = pd.iloc[i] - #print(row) - enlosed_i = row['enclosed_i'] - - if len(enlosed_i) >= 1: - new_row = row.copy() - - new_row['Death'] = pd.loc[pd['ID'] == enlosed_i[0]]['Death'] - new_row['parent_tag'] = pd.loc[pd['ID'] == enlosed_i[0]]['ID'] - ## this accounts for a bug were the entire series is placed in the death column. - # not sure on the origin of this but the following corrects for it. It only occationally happends so this is not - # computationally expensive. - if type(new_row['Death']) == pandas.core.series.Series: - # get the Death value from the first item in the Series. - new_row['Death'] = new_row['Death'].iloc[0] - - new_row['new_row'] = 1 - new_row['ID'] = pd['ID'].max() + 1 - new_row['enclosed_i'] = [] - - - pd = pandas.concat((pd,new_row.to_frame().T), ignore_index=False) - - return pd - - - - - - - - - -def calculate_area_CPU(row, img): - """Calculates area of source mask (for CPU) - - Args: - row (pd.series): _description_ - img (np.ndarray): Image - - Returns: - _type_: _description_ - """ - mask = utils.get_mask_CPU(row,img) - area = np.sum(mask) - return area - - - - - - - -def calculate_area_GPU(Birth,Death,row, img_gpu): - """Calcualtes are of source mask (for GPU) - - Args: - Birth (float): - Death (float): - row (pd.series): - img_gpu (cp.ndarray): - Returns: - area (float): the calculated area of the source mask. - """ - mask = utils.get_mask_GPU(Birth,Death,row,img_gpu) - # evalute if mask is True on an edge. - edge = utils.check_edge(mask) - if edge: - edge = 1 - - area = np.sum(mask) - return area, edge - - -def process_area(i,pd,img): - # handles worker function - return calculate_area_CPU(pd.iloc[i], img) - - -def process_assoc(i): - # hanldes worker function - return make_point_enclosure_assoc_CPU(pd.iloc[i], pd, img) - - - - - - - - - -def compute_ph_components(img,local_bg,analysis_threshold_val,lifetime_limit=0,output=True,bg_map=False,area_limit=3,nproc=1,GPU=False,lifetime_limit_fraction=2,mean_bg=None,IDoffset=None): - - - global GPU_Option - GPU_Option = GPU - t0_compute_ph = time.time() - pd = cripser.computePH(-img,maxdim=0) - t1_compute_ph = time.time() - print('PH computed! t='+str(t1_compute_ph-t0_compute_ph)+' s') - pd = pandas.DataFrame(pd,columns=['dim','Birth','Death','x1','y1','z1','x2','y2','z2'],index=range(1,len(pd)+1)) - pd.drop(columns=['dim','z1','z2'],inplace=True) - pd['lifetime'] = pd['Death'] - pd['Birth'] - pd['Birth'] = -pd['Birth'] - pd['Death'] = -pd['Death'] - print("mean_bg: ",mean_bg) - pd['mean_bg'] = mean_bg - pd['bg'] = 0 - pd['edge_flag'] = 0 - - if bg_map: - - list_of_index_to_drop = [] - - - for index, row in pd.iterrows(): - # check if local_bg is a map or a value - if row['Birth'] < local_bg[int(row.x1),int(row.y1)]: - list_of_index_to_drop.append(index) - - pd.drop(list_of_index_to_drop,inplace=True) - - - # for each row evaluate if death is below analysis thresholdval map value at its birth point. if its below then set Death to bg map value. - for index, row in pd.iterrows(): - Analy_val = analysis_threshold_val[int(row.y1),int(row.x1)] - if row['Death'] < Analy_val: - row['Death'] = Analy_val - # assign each row the local bg value - row['bg'] = local_bg[int(row.y1),int(row.x1)] - - - else: - - pd = pd[pd['Birth']>local_bg] # maybe this should be at the beginning. - pd['Death'] = np.where(pd['Death'] < analysis_threshold_val, analysis_threshold_val, pd['Death']) - pd['bg'] = local_bg - - pd['lifetime'] = abs(pd['Death'] - pd['Birth']) - - pd['lifetimeFrac'] = pd['Birth']/pd['Death'] - - # fiter by lifetimeFrac - - pd = pd[pd['lifetimeFrac']>lifetime_limit_fraction] - - - print('Persis Diagram computed. Length: ',len(pd)) - - if lifetime_limit > 0: - pd = pd[pd['lifetime'] > lifetime_limit] - - - pd.sort_values(by='lifetime',ascending=False,inplace=True,ignore_index=True) - - pd['ID'] = pd.index + IDoffset - - if len(pd) > 0: - - - - area_list = [] - edge_list = [] - - if nproc == 1: - - if GPU_Option == True: - - if GPU_AVAILABLE == True: - - # convert img to cupy array and define type so it does not have to be converted each time. - img_gpu = cp.asarray(img,dtype=cp.float64) - # Calculate area and enforce area limit Single Process. - print('Calculating area with GPU...') - t0 = time.time() - - for i in tqdm(range(0,len(pd)),total=len(pd),desc='Calculating area',disable=not output): - - row = pd.iloc[i] - Birth = row.Birth - Death = row.Death - area, edge = calculate_area_GPU(Birth,Death,row,img_gpu) - area_list.append(area) - edge_list.append(edge) - #percentage_completed = (i/len(pd))*100 - - #if percentage_completed % 10 == 0: - # print(percentage_completed,'%') - - print('Area calculated! t='+str(time.time()-t0)+' s') - - pd['area'] = area_list - pd['edge_flag'] = edge_list - pd = pd[pd['area'] > area_limit] - - enclosed_i_list = [] - print('Calculating enclosed_i with GPU...') - t0 = time.time() - for i in tqdm(range(0,len(pd)),total=len(pd),desc='Calculating enclosed_i',disable=not output): - row = pd.iloc[i] - Birth = row.Birth - Death = row.Death - enclosed_i = make_point_enclosure_assoc_GPU(Birth,Death,row,pd,img,img_gpu) - enclosed_i_list.append(enclosed_i) - - print('enclosed_i calculated! t='+str(time.time()-t0)+' s') - - pd['enclosed_i'] = enclosed_i_list - - - elif GPU_Option == False: - # 1 Core no GPU. - - # Calculate area and enforce area limit Single Process. - - t0 = time.time() - print("No GPU") - - for i in tqdm(range(0,len(pd)),total=len(pd),desc='Calculating area',disable=not output): - area = calculate_area_CPU(pd.iloc[i],img) - #print(area) - area_list.append(area) - - print('Area calculated! t='+str(time.time()-t0)+' s') - - pd['area'] = area_list - pd = pd[pd['area'] > area_limit] - - - # Parent Associations Single Process - - enclosed_i_list = [] - - t0 = time.time() - for i in tqdm(range(0,len(pd)),total=len(pd),desc='Calculating enclosed_i',disable=not output): - row = pd.iloc[i] - enclosed_i = make_point_enclosure_assoc(row,pd,img) - enclosed_i_list.append(enclosed_i) - - print('enclosed_i calculated! t='+str(time.time()-t0)+' s') - - pd['enclosed_i'] = enclosed_i_list - - - else: - # Multiple CPUs *** currently not working - - print('Calculating area with ',nproc,' processes') - - # Calculate area and enforce area limit Multi Process. - - t0 = time.time() - print("Chunksize: ",len(pd)//nproc) - index_and_args_list = [(i, pd, img) for i in range(len(pd))] - - with Pool(nproc) as p: - area_list = list(p.starmap(process_area, range(len(pd)),chunksize=len(pd)//nproc)) - print('Area calculated! t='+str(time.time()-t0)+' s') - - pd['area'] = area_list - pd = pd[pd['area'] > area_limit] # remove 1 pixel points - - print(len(pd)) - - ## Parent Associations Multi Process. - print('Calculating enclosed_i with ',nproc,' processes') - t0 = time.time() - - print("Chunksize: ",len(pd)//nproc) - - with Pool(nproc) as pool: - enclosed_i_list = list(pool.imap(process_assoc, range(len(pd)),chunksize=len(pd)//nproc)) - - print('enclosed_i calculated! t='+str(time.time()-t0)+' s') - - pd['enclosed_i'] = enclosed_i_list - - - #pd['parent_tag'] = 0 - pd = correct_first_destruction(pd,output=not output,img=img,img_gpu=img_gpu,GPU=GPU) - - # what if we do this all together after the first destruction correction? - - pd['lifetime'] = pd['Birth'] - pd['Death'] - #print(pd) - pd.sort_values(by='lifetime',ascending=False,inplace=True,ignore_index=False) - #print(pd) - #pdb.set_trace() - # corrected for first destruction. points need to have enlosed_i updated. - # update enclosed_i - enclosed_i_list = [] - print('Updating enclosed_i') - if GPU_Option: - enclosed_i_list = [] - t0 = time.time() - for i in tqdm(range(0,len(pd)),total=len(pd),desc='Calculating enclosed_i',disable=not output): - row = pd.iloc[i] - # is this a new row? - #if row.new_row == 1: - Birth = row.Birth - Death = row.Death - enclosed_i = make_point_enclosure_assoc_GPU_second(Birth,Death,row,pd,img,img_gpu) - enclosed_i_list.append(enclosed_i) - #else: - # enclosed_i_list.append(row.enclosed_i) - print('enclosed_i calculated! t='+str(time.time()-t0)+' s') - - pd['enclosed_i'] = enclosed_i_list - - else: - t0 = time.time() - for i in tqdm(range(0,len(pd)),total=len(pd),desc='Calculating enclosed_i',disable=not output): - row = pd.iloc[i] - # is this a new row? - #if row.new_row == 1: - enclosed_i = make_point_enclosure_assoc(row,pd,img) - enclosed_i_list.append(enclosed_i) - #else: - enclosed_i_list.append(row.enclosed_i) - print('enclosed_i calculated! t='+str(time.time()-t0)+' s') - - pd['enclosed_i'] = enclosed_i_list - - - print('Calculating parent_tags... ') - - t0_parent_tag = time.time() - parent_tag_list = parent_tag_func_vectorized(pd) - pd['parent_tag'] = parent_tag_list - t1_parent_tag = time.time() - print('parent_tag calculated! t='+str(t1_parent_tag-t0_parent_tag)+' s') - - #print(pd) - #pd['parent_tag'] = pd.apply(lambda row: parent_tag_func(row,pd), axis=1) - - print('Assigning Class ...') - t0_CLass = time.time() - pd['Class'] = pd.apply(classify_single,axis=1) - t1_Class = time.time() - print('Class assigned! t='+str(t1_Class-t0_CLass)+' s') - # drop the enclosed_i column - #pd.drop(columns=['enclosed_i'],inplace=True) - # distance from center of each image - - pd['distance_from_center'] = ((pd['x1'] - img.shape[0]/2)**2 + (pd['y1'] - img.shape[1]/2)**2)**0.5 - - return pd - - else: - - return pd \ No newline at end of file diff --git a/DRUID/src/homology_new.py b/DRUID/src/homology_new.py deleted file mode 100644 index b28201f..0000000 --- a/DRUID/src/homology_new.py +++ /dev/null @@ -1,602 +0,0 @@ -""" -File: src/homology_new.py -Author: Rhys Shaw -Date: 27/12/2025 -Version: v1.0 -Description: Functions for calculating source properties for sources. - -""" - -import cripser -import numpy as np -from ..src import utils -from ..src import background - -import time -import pandas -from tqdm import tqdm - - -try: - import cupy as cp - from cupyx.scipy.ndimage import label as cupy_label - - GPU_AVAILABLE = True -except: - GPU_AVAILABLE = False - -# used for debugging -# import matplotlib.pyplot as plt -# import pdb - - -def make_point_enclosure_assoc_GPU(id, x1, y1, Birth, Death, pd, img_gpu): - """Returns a list of the ID of the points that are enclosed by the mask pd point. - Uses GPU for computation. - - Args: - Birth (float): _description_ - Death (float): _description_ - row (pd.Series): _description_ - pd (pd.DataFrame): _description_ - img (np.ndarray): _description_ - img_gpu (cp.ndarray): _description_ - - Returns: - enclosed_list (list): _description_ - """ - # print(pd) - mask = utils.get_mask_GPU(Birth, Death, x1, y1, img_gpu).get() - # pdb.set_trace() - mask_coords = np.column_stack((pd["x1"], pd["y1"])) - points_inside_mask = mask[ - mask_coords[:, 0].astype(int), mask_coords[:, 1].astype(int) - ] - encloses_vectorized = pd.iloc[points_inside_mask]["ID"].tolist() - # remove self from list - # print(row.ID) - # print(encloses_vectorized) - # encloses_vectorized.remove(row.ID) - # print(encloses_vectorized) - # pdb.set_trace() - return encloses_vectorized - - -def make_point_enclosure_assoc_CPU(ID, x1, y1, Birth, Death, pd, img): - """Returns a list of the ID of the points that are enclosed by the mask pd point. - Uses GPU for computation. - - Args: - Birth (float): _description_ - Death (float): _description_ - row (pd.Series): _description_ - pd (pd.DataFrame): _description_ - img (np.ndarray): _description_ - img_gpu (cp.ndarray): _description_ - - Returns: - enclosed_list (list): _description_ - """ - - mask = utils.get_mask_CPU(x1, y1, Birth, Death, img) - # print(mask) - # plt.imshow(mask) - # plt.savefig('test.png') - # pdb.set_trace() - mask_coords = np.column_stack((pd["x1"], pd["y1"])) - points_inside_mask = mask[ - mask_coords[:, 0].astype(int), mask_coords[:, 1].astype(int) - ] - encloses_vectorized = pd.iloc[points_inside_mask]["ID"].tolist() - # remove self from list - # print(row.ID) - # print(encloses_vectorized) - # encloses_vectorized.remove(ID) - # print(encloses_vectorized) - # pdb.set_trace() - return encloses_vectorized - - -def classify_single(row): - """Classifiying the Rows based on orgin. - - Args: - row: pd.series - The row that is being classified. - - Returns: - Class: int - the Class integer that indiceates the class the row belongs too. - - """ - if row.new_row == 0: - if len(row.enclosed_i) <= 1: # no children - if row.parent_tag == row.ID: # no parent - return 0 # no children, no parent. - else: - return 2 # no child has parent. - else: - if row.parent_tag == row.ID: # no parent - return 4 # has children, no parent. - else: - return 3 # has children, has parent. - else: - return 1 # new row has children. - - -def parent_tag_func_vectorized(df): - """ - - Vectorised implenetation of parent tag function. - - - Args: - - df: pd.Dataframe - data frame for which we calculate the parent tags. - - - Returns: - - df: pd.Dataframe - Pandas data frame with addition parent tag column. - - """ - - enclosed_i_dict = {row["ID"]: set(row["enclosed_i"]) for idx, row in df.iterrows()} - - def find_parent_tag(row): - # if row.new_row == 0: # only set it if it is not a new row. - - for ID, enclosed_i_set in enclosed_i_dict.items(): - if row.ID in enclosed_i_set: - return ID - - return np.nan - - # else: # if it is a new row then we already know the parent tag. - # return row.parent_tags - - return df.apply(find_parent_tag, axis=1) - - -def parent_tag_func_vectorized_new(df): - """ - Vectorised implementation of parent tag function. - - Args: - df: pd.Dataframe - data frame for which we calculate the parent tags. - - Returns: - df: pd.Dataframe - Pandas data frame with addition parent tag column. - """ - - # Create a dictionary that maps each enclosed ID to its parent ID - def find_parent(row, df): - potential_parents = [] - for idx, other_row in df.iterrows(): - if row["ID"] in other_row["enclosed_i"] and row["ID"] != other_row["ID"]: - potential_parents.append( - (other_row["ID"], len(other_row["enclosed_i"])) - ) - - if potential_parents: - # Sort by length of enclosed_i (descending) and return the ID of the longest - return sorted(potential_parents, key=lambda x: x[1], reverse=True)[0][0] - - # If no parent found, check if this row has the largest enclosed_i - if len(row["enclosed_i"]) == df["enclosed_i"].apply(len).max(): - return row["ID"] # Set its own ID as parent - - return row["ID"] - - # Apply the function to each row in the DataFrame - df["parent_tag"] = df.apply(find_parent, axis=1, args=(df,)) - return df - - -# def parent_tag_func_optimized(df): -# """ -# Optimized version of the parent tag function. - -# Args: -# df: pd.DataFrame - DataFrame for which we calculate the parent tags. - -# Returns: -# df: pd.DataFrame - DataFrame with additional parent_tag column. -# """ -# # Convert enclosed_i lists to sets for faster lookups and map IDs to their indices -# enclosed_sets = df['enclosed_i'].apply(set) -# id_to_index = {id_: idx for idx, id_ in enumerate(df['ID'])} - -# # Create an array of parent indices, initialized to self (i.e., no parent) -# parent_indices = np.arange(len(df)) - -# for row_idx in range(len(df)): -# potential_parents = [id_to_index[id_] for id_ in df['enclosed_i'].iloc[row_idx] if id_ in id_to_index] - -# if potential_parents: -# # Find the parent with the smallest enclosed set -# min_parent = min(potential_parents, key=lambda idx: len(enclosed_sets.iloc[idx])) -# parent_indices[row_idx] = min_parent - -# # Map indices back to IDs -# df['parent_tag'] = df['ID'].iloc[parent_indices].values - -# return df - - -# def parent_tag_func_optimized(df): -# """ -# Corrected and optimized vectorized implementation of parent tag function. - -# Args: -# df: pd.DataFrame - data frame for which we calculate the parent tags. - -# Returns: -# df: pd.DataFrame - Pandas data frame with additional parent tag column. -# """ -# # Create a dictionary mapping IDs to their index in the dataframe -# id_to_index = {id: idx for idx, id in enumerate(df["ID"])} - -# # Convert enclosed_i lists to sets for faster lookup -# enclosed_sets = df["enclosed_i"].apply(set) - -# # Create a matrix of containment relationships -# containment_matrix = enclosed_sets.apply( -# lambda x: [1 if id in x else 0 for id in df["ID"]] -# ) -# containment_matrix = np.array(containment_matrix.tolist()) - -# # Remove self-containment -# np.fill_diagonal(containment_matrix, 0) - -# def find_immediate_parent(row_idx): -# potential_parents = np.where(containment_matrix[:, row_idx] == 1)[0] -# if len(potential_parents) == 0: -# return row_idx # No parent found, return self - -# # Among potential parents, find the one with the smallest enclosed_i set -# parent_sizes = [len(enclosed_sets.iloc[i]) for i in potential_parents] -# immediate_parent_idx = potential_parents[np.argmin(parent_sizes)] -# return immediate_parent_idx - -# # parent_indices = [ -# # find_immediate_parent(i) -# # for i in tqdm(range(len(df)), desc="Finding immediate parents", total=len(df)) -# # ] - -# # rewrite the linline function to a normal function - -# parent_indices = [] -# for i in tqdm(range(len(df)), desc="Finding immediate parents", total=len(df)): -# immediate_parent_idx = find_immediate_parent(i) -# parent_indices.append(immediate_parent_idx) - -# # Map indices back to IDs -# df["parent_tag"] = df["ID"].iloc[parent_indices].values - -# return df - - -def parent_tag_func_optimized(df): - # Create a DataFrame with only IDs that have enclosed_i of length greater than 1 - enclosed_i = df[["ID", "enclosed_i"]] - enclosed_i = enclosed_i[enclosed_i["enclosed_i"].apply(len) > 1] - - # Initialize parent_tag with ID - df["parent_tag"] = df["ID"] - - # Create a mapping from ID to parent ID - mapping = {} - - for idx, row in enclosed_i.iterrows(): - parent_id = row["ID"] - for child_id in row["enclosed_i"]: - if child_id != parent_id: - mapping[child_id] = parent_id - - # Update parent_tag using the mapping - df["parent_tag"] = df["ID"].map(mapping).combine_first(df["parent_tag"]) - - return df - - -def correct_first_destruction(pd, output): - """ - Function for correcting for the First destruction of a parent Island. - - Args: - pd (pd.DataFrame): Input catalogue of sources to correct. - output (bool): True if you want interation logginf with tqdm. - - Returns: - pd (pd.DataFrame): The new Catalogue. - - """ - - pd["new_row"] = 0 - - for i in tqdm( - range(0, len(pd)), - total=len(pd), - desc="Correcting first destruction", - disable=output, - ): - - row = pd.iloc[i] - # print(row) - enlosed_i = row["enclosed_i"] - if len(enlosed_i) > 1: - new_row = row.copy() - # print(i) - new_row["Death"] = pd.loc[pd["ID"] == enlosed_i[1]]["Death"] - new_row["parent_tag"] = pd.loc[pd["ID"] == enlosed_i[1]]["ID"] - - ## this accounts for a bug were the entire series is placed in the death column. - # not sure on the origin of this but the following corrects for it. It only occationally happends so this is not - # computationally expensive - if type(new_row["Death"]) == pandas.core.series.Series: - # get the Death value from the first item in the Series. - # print(new_row['Death']) - new_row["Death"] = new_row["Death"].iloc[0] # - - new_row["new_row"] = 1 - new_row["ID"] = pd["ID"].max() + 1 - new_row["enclosed_i"] = [] - - # concat the new row to the dataframe. - - pd = pandas.concat([pd, new_row.to_frame().T], ignore_index=True) - - # pd = pd.append(new_row,ignore_index=True) - # pd.loc[len(pd)+1] = new_row - - return pd - - -def calculate_area_GPU(Birth, Death, row, img_gpu): - """Calcualtes are of source mask (for GPU) - - Args: - Birth (float): - Death (float): - row (pd.series): - img_gpu (cp.ndarray): - Returns: - area (float): the calculated area of the source mask. - """ - mask = utils.get_mask_GPU(Birth, Death, row.x1, row.y1, img_gpu) - # get bounding box here - # evalute if mask is True on an edge. - bounding_box = utils.bounding_box_gpu(mask) - mask = mask.get() - edge = utils.check_edge(mask) - if edge: - edge = True - - area = np.sum(mask) - return area, edge, bounding_box - - -def calculate_area_CPU(Birth, Death, row, img): - - mask = utils.get_mask_CPU(row.x1, row.y1, Birth, Death, img) - # get bounding box here - bounding_box = utils.bounding_box_cpu(mask) - edge = utils.check_edge(mask) - if edge: - edge = 1 - area = np.sum(mask) - return area, edge, bounding_box - - -def compute_ph_components( - img, - local_bg, - analysis_threshold_val, - lifetime_limit, - output=False, - bg_map=False, - area_limit=3, - GPU=False, - lifetime_limit_fraction=2, - mean_bg=None, - IDoffset=0, - box_size=None, - detection_threshold=None, - Cutout_X_offset=0, - Cutout_Y_offset=0, -): - - global GPU_Option - GPU_Option = GPU - print("Computing PH components...") - t0_compute_ph = time.time() - pd = cripser.computePH(-img, maxdim=0) - t1_compute_ph = time.time() - print("Time to compute PH: {}".format(t1_compute_ph - t0_compute_ph)) - pd = pandas.DataFrame( - pd, - columns=["dim", "Birth", "Death", "x1", "y1", "z1", "x2", "y2", "z2"], - index=range(1, len(pd) + 1), - ) - pd.drop(columns=["dim", "z1", "z2"], inplace=True) - pd["lifetime"] = pd["Death"] - pd["Birth"] - pd["Birth"] = -pd["Birth"] - pd["Death"] = -pd["Death"] - - # print("mean_bg: ",mean_bg) - # get rid of birth less than 0, helps speed up the code alittle. - # mean_bg_temp = np.nanmean(local_bg)/detection_threshold - # print('mean_bg: ',mean_bg_temp) - # get rid of alot of defintly not sources. - # pd = pd[pd['Birth']>mean_bg_temp] - - pd["bg"] = 0 - pd["edge_flag"] = 0 - pd["mean_bg"] = 0 # this is the mean of the background - # assign each row the local bg valuw from the map. - if bg_map: - - # this is a slow function. Can it be improved? - # for each row we need to assign the local bg value. - pd["bg"] = pd["bg"].astype(float) - pd["mean_bg"] = pd["mean_bg"].astype(float) - for index, row in pd.iterrows(): - - pd.loc[index, "bg"] = float( - background.get_bg_value_from_result_image( - (int(row.x1 + Cutout_Y_offset), int(row.y1 + Cutout_X_offset)), - box_size, - local_bg, - ) - ) - pd.loc[index, "mean_bg"] = float( - background.get_bg_value_from_result_image( - (int(row.x1 + Cutout_Y_offset), int(row.y1 + Cutout_X_offset)), - box_size, - mean_bg, - ) - ) - - # evaluate if the death value is below the analysis threshold value at the birth point. - # if it is then set the death value to the analysis threshold value. - - Anal_val = background.get_bg_value_from_result_image( - (int(row.x1 + Cutout_Y_offset), int(row.y1 + Cutout_X_offset)), - box_size, - analysis_threshold_val, - ) - - if row["Death"] < Anal_val: - pd.loc[index, "Death"] = Anal_val - - else: - # no bg map so just assign the local bg value. asn this should be single value. - pd["bg"] = local_bg - pd["mean_bg"] = mean_bg - - for index, row in pd.iterrows(): - if row["Death"] < analysis_threshold_val: - pd.loc[index, "Death"] = analysis_threshold_val - - # print('Before Cull',len(pd)) - # print(pd['bg']) - # print('mean_bg: ',np.mean(pd['bg'])) - pd = pd[pd["Birth"] > pd["bg"]] # maybe this should be at the beginning. - # also make sure Death is less than Birth - # pd = pd[pd['Death'] < pd['Birth']] - # print('After Cull',len(pd)) - # if bg_map: - - # list_of_index_to_drop = [] - - # for index, row in pd.iterrows(): - # # check if local_bg is a map or a value - # if row['Birth'] < local_bg[int(row.x1),int(row.y1)]: - # list_of_index_to_drop.append(index) - - # pd.drop(list_of_index_to_drop,inplace=True) - - # # for each row evaluate if death is below analysis thresholdval map value at its birth point. if its below then set Death to bg map value. - # for index, row in pd.iterrows(): - # Analy_val = analysis_threshold_val[int(row.y1),int(row.x1)] - # if row['Death'] < Analy_val: - # row['Death'] = Analy_val - # # assign each row the local bg value - # row['bg'] = local_bg[int(row.y1),int(row.x1)] - # row['mean_bg'] = mean_bg[int(row.y1),int(row.x1)] - - # else: - # print(local_bg) - # pd = pd[pd['Birth']>local_bg] # maybe this should be at the beginning. - # pd['Death'] = np.where(pd['Death'] < analysis_threshold_val, analysis_threshold_val, pd['Death']) - # pd['bg'] = local_bg - # pd['mean_bg'] = mean_bg - - pd["lifetime"] = abs(pd["Death"] - pd["Birth"]) - - pd["lifetimeFrac"] = pd["Birth"] / pd["Death"] - pd = pd[pd["lifetimeFrac"] > lifetime_limit_fraction] - pd = pd[pd["lifetime"] > lifetime_limit] - pd.sort_values(by="lifetime", ascending=False, inplace=True, ignore_index=True) - - pd["ID"] = pd.index + IDoffset - - # print('PD: ',pd) - # begins here - - if len(pd) > 0: - - area_list = [] - edge_list = [] - bbox1 = [] - bbox2 = [] - bbox3 = [] - bbox4 = [] - - if GPU_Option == True: - if GPU_AVAILABLE == True: - img_gpu = cp.asarray(img, dtype=cp.float64) - # Calculate area and enforce area limit Single Process. - # print('Calculating area with GPU...') - t0 = time.time() - - for i in tqdm( - range(0, len(pd)), - total=len(pd), - desc="Calculating area", - disable=not output, - ): - - row = pd.iloc[i] - Birth = row.Birth - Death = row.Death - area, edge, bbox = calculate_area_GPU(Birth, Death, row, img_gpu) - area_list.append(area) - edge_list.append(edge) - bbox1.append(bbox[0].get()) - bbox2.append(bbox[1].get()) - bbox3.append(bbox[2].get()) - bbox4.append(bbox[3].get()) - - t1 = time.time() - # print('Time to calculate area and inital bbox: {}'.format(t1-t0)) - pd["area"] = area_list - pd["edge_flag"] = edge_list - pd["bbox1"] = bbox1 - pd["bbox2"] = bbox2 - pd["bbox3"] = bbox3 - pd["bbox4"] = bbox4 - pd = pd[pd["area"] > area_limit] - - else: - - # print('Calculating area with CPU...') - t0 = time.time() - for i in tqdm( - range(0, len(pd)), - total=len(pd), - desc="Calculating area", - disable=not output, - ): - - row = pd.iloc[i] - Birth = row.Birth - Death = row.Death - area, edge, bbox = calculate_area_CPU(Birth, Death, row, img) - area_list.append(area) - edge_list.append(edge) - bbox1.append(bbox[0]) - bbox2.append(bbox[1]) - bbox3.append(bbox[2]) - bbox4.append(bbox[3]) - - t1 = time.time() - # print('Time to calculate area and inital bbox: {}'.format(t1-t0)) - pd["area"] = area_list - pd["edge_flag"] = edge_list - pd["bbox1"] = bbox1 - pd["bbox2"] = bbox2 - pd["bbox3"] = bbox3 - pd["bbox4"] = bbox4 - pd = pd[pd["area"] > area_limit] - - return pd diff --git a/DRUID/src/properties.py b/DRUID/src/properties.py new file mode 100644 index 0000000..e017b45 --- /dev/null +++ b/DRUID/src/properties.py @@ -0,0 +1,149 @@ +""" +Author: Rhys Shaw +Date: 08-09-2025 +""" + +import numpy as np +import polars as pl +from skimage import measure +from scipy.ndimage import label as scipy_label + + +def get_enclosing_mask_CPU(x, y, mask): + labeled_mask, _ = scipy_label(mask) + if 0 <= x < mask.shape[1] and 0 <= y < mask.shape[0]: + label_at_pixel = labeled_mask[y, x] + if label_at_pixel != 0: + return labeled_mask == label_at_pixel + return None + + +def calculate_radio_flux_error(background_rms, area, BMAJ, BMIN): + if BMAJ is None or BMIN is None: + return np.nan + + gfactor = 2 * np.sqrt(2 * np.log(2)) + Beam_area = 2 * np.pi * (BMAJ * BMIN) / gfactor + return np.mean(background_rms) * np.sqrt(area / Beam_area) + + +def optical_flux_err(data, bkg, eff_gain): + # https://photutils.readthedocs.io/en/stable/api/photutils.utils.calc_total_error.html + # photoutils implentation of optical flux error calculations. + # f_err = sqrt(obk^2 + (I/g_eff)^2 + (obk/rms_median)^2) + + from photutils.utils import calc_total_error + + return calc_total_error(data, bkg, eff_gain) + + +def calculate_properties( + cat, + raw_image, + smoothed_image, + background, + background_rms, + position, + analysis_threshold, + mode, + BMAJ=None, + BMIN=None, + EFFRON=None, + EFFGAIN=None, + EXPTIME=None, +): + # Vectorized extraction to avoid Polars iter_rows + births = cat["birth"].to_numpy() + deaths = ( + cat["deaths"].to_numpy() if "deaths" in cat.columns else cat["death"].to_numpy() + ) + x1s = cat["x1"].to_numpy() + y1s = cat["y1"].to_numpy() + areas = cat["area"].to_numpy() + + maj, min_ax, pa, centroid_lst, flux, flux_peak, bg, flux_err, snr = ( + [], + [], + [], + [], + [], + [], + [], + [], + [], + ) + + for b, d, x, y, area in zip(births, deaths, x1s, y1s, areas): + # mask boundaries determined by the SMOOTHED image + mask = (smoothed_image <= b) & (smoothed_image > d) + enclosed_mask = get_enclosing_mask_CPU(int(y), int(x), mask) + + if enclosed_mask is None: + maj.append(np.nan) + min_ax.append(np.nan) + pa.append(np.nan) + centroid_lst.append((np.nan, np.nan)) + flux.append(np.nan) + flux_peak.append(np.nan) + bg.append(np.nan) + flux_err.append(np.nan) + snr.append(np.nan) + continue + + enclosed_mask_int = enclosed_mask.astype(int) + + # properties extracted from the RAW image + props = measure.regionprops(enclosed_mask_int, intensity_image=raw_image) + + if props: + p = props[0] + maj.append(p.axis_major_length) + min_ax.append(p.axis_minor_length) + pa.append(p.orientation) + centroid_lst.append(p.centroid) + else: + maj.append(np.nan) + min_ax.append(np.nan) + pa.append(np.nan) + centroid_lst.append((np.nan, np.nan)) + + # Flux summations computed on RAW data + flux_tot = np.nansum(enclosed_mask_int * (raw_image - background)) + flux.append(flux_tot) + flux_peak.append(np.nanmax(enclosed_mask_int * (raw_image - background))) + + bg_mean = np.mean(background * enclosed_mask_int) + bg.append(bg_mean) + + if mode == "radio": + f_err = calculate_radio_flux_error(background_rms, area, BMAJ, BMIN) + flux_err.append(f_err) + if f_err and not np.isnan(f_err): + snr.append(flux_tot / f_err) + else: + snr.append(np.nan) + + elif mode == "optical": + if EFFGAIN is None: + EFFGAIN = 0 # Default to 1 if not provided + f_err = optical_flux_err(raw_image, background, EFFGAIN).mean() + flux_err.append(f_err) + snr.append(flux_tot / f_err if f_err else 0) + else: + flux_err.append(0) + snr.append(0) + + return cat.with_columns( + [ + pl.Series("maj", maj), + pl.Series("min", min_ax), + pl.Series("pa", pa), + pl.Series("centroid_x", [c[1] for c in centroid_lst]), + pl.Series("centroid_y", [c[0] for c in centroid_lst]), + pl.Series("flux_peak", flux_peak), + pl.Series("bg", bg), + pl.Series("flux_err", flux_err), + pl.Series("flux", flux), + pl.Series("snr", snr), + ] + ) diff --git a/DRUID/src/source.py b/DRUID/src/source.py index bfb46a7..4889829 100644 --- a/DRUID/src/source.py +++ b/DRUID/src/source.py @@ -1,752 +1,123 @@ -""" -File: src/source.py -Author: Rhys Shaw -Date: 27/12/2023 -Version: v1.0 -Description: Functions for calculating source properties for sources. - -""" - -from ..src import utils -import pandas as pd import numpy as np -from tqdm import tqdm -from scipy.ndimage import label -import pdb -from scipy.ndimage import binary_dilation - -try: - import cupy as cp - from cupyx.scipy.ndimage import label as cupy_label - -except: - - pass - - -def create_params_df(cutup: bool, params: list): - """ - - Creates a pandas dataframe from the parameters. - - """ - - params = pd.DataFrame( - params, - columns=[ - "ID", - "Birth", - "Death", - "x1", - "y1", - "x2", - "y2", - "Flux_total", - "Flux_total_err", - "Flux_peak", - "Flux_correction_factor", - "Area", - "Xc", - "Yc", - "bbox1", - "bbox2", - "bbox3", - "bbox4", - "Maj", - "Min", - "Pa", - "parent_tag", - "Class", - "SNR", - "Noise", - "X0_cutout", - "Y0_cutout", - "mean_bg", - "bg_rms", - "Edge_flag", - "contour", - "enclosed_i", - ], - ) - - return params - - -def large_mask_red_image_procc_GPU(Birth, Death, x1, y1, image, X0, Y0): - """ - Does all gpu processing for the large mask. - - return the red_image and red_mask and the bounding box. - - """ - - mask = cp.zeros(image.shape, dtype=cp.bool_) - mask = cp.logical_or(mask, cp.logical_and(image <= Birth, image > Death)) - - # mask_enclosed = self.get_enclosing_mask_gpu(y1,x1,mask) - labeled_mask, num_features = cupy_label(mask) - - # Check if the specified pixel is within the mask - if 0 <= y1 < mask.shape[1] and 0 <= x1 < mask.shape[0]: - label_at_pixel = labeled_mask[x1, y1] - # print(x1) - # print(y1) - # print(label_at_pixel) - if label_at_pixel != 0: - # Extract the connected component containing the specified pixel - component_mask = labeled_mask == label_at_pixel - # plt.imshow(component_mask.get()) - # plt.show() - # pdb.set_trace() - non_zero_indices = cp.nonzero(component_mask) - - # Extract minimum and maximum coordinates - xmin = cp.min(non_zero_indices[1]) - ymin = cp.min(non_zero_indices[0]) - xmax = cp.max(non_zero_indices[1]) - ymax = cp.max(non_zero_indices[0]) - - # correct the bounding box for the cutout. - xmin = xmin # + Y0 - xmax = xmax # + Y0 - ymin = ymin # + X0 - ymax = ymax # + X0 - - # images are not being cropped? - - red_image = image[ymin : ymax + 1, xmin : xmax + 1] - red_mask = component_mask[ymin : ymax + 1, xmin : xmax + 1] - - return red_image, red_mask, xmin, xmax, ymin, ymax - - -def large_mask_red_image_procc_CPU(Birth, Death, x1, y1, image, image_smooth): - """ - Does all gpu processing for the large mask. - - return the red_image and red_mask and the bounding box. - - """ - mask = np.zeros(image.shape) - mask = np.logical_or( - mask, np.logical_and(image_smooth <= Birth, image_smooth > Death) - ) - # print(mask.shape) - # mask_enclosed = self.get_enclosing_mask_gpu(y1,x1,mask) - labeled_mask, num_features = label(mask) - - # Check if the specified pixel is within the mask - if 0 <= y1 < mask.shape[1] and 0 <= x1 < mask.shape[0]: - - label_at_pixel = labeled_mask[x1, y1] - - if label_at_pixel != 0: - # Extract the connected component containing the specified pixel - component_mask = labeled_mask == label_at_pixel - # plt.imshow(component_mask.get()) - # plt.show() - # pdb.set_trace() - non_zero_indices = np.nonzero(component_mask) - - # Extract minimum and maximum coordinates - xmin = np.min(non_zero_indices[1]) - ymin = np.min(non_zero_indices[0]) - xmax = np.max(non_zero_indices[1]) - ymax = np.max(non_zero_indices[0]) - - # images are not being cropped? - - red_image = image[ymin : ymax + 1, xmin : xmax + 1] - red_mask = component_mask[ymin : ymax + 1, xmin : xmax + 1] - - return red_image, red_mask, xmin, xmax, ymin, ymax - - -# import cv2 -# import numpy as np - -# def dilate_mask_circular(mask, radius): -# Create a circular structuring element using cv2.getStructuringElement -# circular_kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2*radius+1, 2*radius+1)) - -# Perform dilation using cv2.dilate -# dilated_mask = cv2.dilate(mask.astype(np.uint8), circular_kernel)# - -# return dilated_mask - - -# def curve_of_growth_dilation(mask,image): -# converged = False -# i = 0 -# dilated_mask = np.zeros(image.shape, dtype=bool) -# mask_before = mask -# while converged == False: -# if i == 50: -# #print('max iterations reached') -# return dilated_mask.astype(int) -# # print(i) -# i += 1 - -# dilated_mask = dilate_mask_circular(mask_before,1) - -# # check if the mask has converged -# flux = np.sum(image*dilated_mask) -# flux_old = np.sum(image*mask_before) -# diff = (flux - flux_old) -# #print(diff) -# if diff<0: -# converged = True -# dilated_mask = mask_before - -# else: -# mask_before = dilated_mask -# #print('converged {}'.format(i)) -# return dilated_mask.astype(int) - - -def measure_source_properties( - use_gpu, - catalogue=None, - cutout=None, - smooth_cutout=None, - background_map=None, - output=None, - cutupts=None, - mode="optical", - header=None, - sigma=5, +from skimage.measure import regionprops_table, label +import polars as pl + +from ..src.utils import ( + TITLE, + LINK, + GOLD, + RESET, + BOLD, + NOTICE, + ERROR, + WARNING, + CODEBLOCK, + BLACK, +) + + +def create_source_islands( + image, + background_map, + background_rms_map, + detection_threshold=5, + analysis_threshold=3, + area_limit=2, + max_area_limit=10000, + verbose=True, ): """ - - Characterising the Source Assuming the input image of of the format of a optical astronomical image. - - # needs to work on cutup images. - # work on the whole image. - + Create source islands using optimized vectorization. + Returns bounding boxes instead of full arrays to save IPC overhead. + Separates massive islands for cataloging without computing homology. """ - if header == None: - mode = "other" - - if mode == "Radio": - print("Radio mode selected") - Beam, BMAJ, BMIN, BPA = utils.calculate_beam(header=header) - - if mode == "optical": - print("Optical mode selected, modeling the PSF as a gaussian.") - psf_fwhm_p = utils.get_psf_FWHM(header) - # EFFRON = utils.get_EFFRON(header) - # EFFGAIN = utils.get_EFFGAIN(header) - # EXPTIME = utils.get_EXPTIME(header) - - if use_gpu: - - try: - - import cupy as cp - - except: - - raise ImportError("cupy not installed. GPU acceleration not possible.") - - image = cutout - if use_gpu: - image_gpu = cp.asarray(image, dtype=cp.float64) - smooth_image_gpu = cp.asarray(smooth_cutout, dtype=cp.float64) - - Birth = catalogue["Birth"].to_numpy() - Death = catalogue["Death"].to_numpy() - parent_tag = catalogue["parent_tag"].to_numpy() - Class = catalogue["Class"].to_numpy() - bg = catalogue["bg"].to_numpy() - X0 = catalogue["X0_cutout"].to_numpy() - Y0 = catalogue["Y0_cutout"].to_numpy() - mean_bg = catalogue["mean_bg"].to_numpy() - enclosed_i = catalogue["enclosed_i"].to_numpy() - IDs = catalogue["ID"].to_numpy() - Edge_flags = catalogue["edge_flag"].to_numpy() - bbox1_og = catalogue["bbox1"].to_list() - bbox2_og = catalogue["bbox2"].to_list() - bbox3_og = catalogue["bbox3"].to_list() - bbox4_og = catalogue["bbox4"].to_list() - - x1 = catalogue["x1"].to_numpy() # - 1 #- X0 - y1 = catalogue["y1"].to_numpy() # - 1#- X0 - x2 = catalogue["x2"].to_numpy() # - 1#- Y0 - y2 = catalogue["y2"].to_numpy() # - 1#- X0 - - # map X0 and Y0 to the cutout number - params = [] - polygons = [] - - # import matplotlib.pylab as plt - - for i, source in tqdm( - enumerate(Birth), - total=len(Birth), - desc="Calculating Source Properties..", - disable=not output, - ): - - if use_gpu == True: - - cropped_image_gpu = image_gpu[ - bbox1_og[i] - 1 : bbox3_og[i] + 1, bbox2_og[i] - 1 : bbox4_og[i] + 1 - ] - - cropped_smooth_image_gpu = smooth_image_gpu[ - bbox1_og[i] - 1 : bbox3_og[i] + 1, bbox2_og[i] - 1 : bbox4_og[i] + 1 - ] - - try: - red_image, red_mask, xmin, xmax, ymin, ymax = ( - large_mask_red_image_procc_GPU( - Birth[i], - Death[i], - x1[i] - bbox1_og[i] + 1, - y1[i] - bbox2_og[i] + 1, - cropped_image_gpu, - X0[i], - Y0[i], - ) - ) - except: - - print("Error in GPU processing!") - print("Source ID: ", IDs[i]) - print("x1:", x1[i] - bbox1_og[i] + 1) - print("y1:", y1[i] - bbox2_og[i] + 1) - print("Birth:", Birth[i]) - print("Death:", Death[i]) - # print('Cropped_image',cropped_image_gpu.get()) - - red_mask = red_mask.astype(int) - - red_image = red_image.get() - red_mask = red_mask.get() - xmin = xmin.get() + bbox2_og[i] - xmax = xmax.get() + bbox2_og[i] - ymin = ymin.get() + bbox1_og[i] - ymax = ymax.get() + bbox1_og[i] - - else: - cropped_image = image[ - bbox1_og[i] - 1 : bbox3_og[i] + 1, bbox2_og[i] - 1 : bbox4_og[i] + 1 - ] - cropped_smooth_image = smooth_cutout[ - bbox1_og[i] - 1 : bbox3_og[i] + 1, bbox2_og[i] - 1 : bbox4_og[i] + 1 - ] - try: - red_image, red_mask, xmin, xmax, ymin, ymax = ( - large_mask_red_image_procc_CPU( - Birth[i], - Death[i], - int(x1[i]) - int(bbox1_og[i]) + 1, - int(y1[i]) - int(bbox2_og[i]) + 1, - cropped_image, - cropped_smooth_image, - ) - ) - except: - print("Error in CPU processing!") - print("Source ID: ", IDs[i]) - print("x1:", x1[i] - bbox1_og[i] + 1) - print("y1:", y1[i] - bbox2_og[i] + 1) - print("Birth:", Birth[i]) - print("Death:", Death[i]) - - red_mask = red_mask.astype(int) - - xmin = xmin + bbox2_og[i] - xmax = xmax + bbox2_og[i] - ymin = ymin + bbox1_og[i] - ymax = ymax + bbox1_og[i] - - # DILATION TEST Not Working. - # print(red_mask) - - # #try: - # if Class[i] == 0: - # if xmin > 50 and ymin > 50: - # expanx = 50 - # expany = 50 - # # increase mask and image size by 100 pixels in each direction. - # # this is to ensure that the mask is large enough to capture the entire source. - # # if the source is less than 50px away from the image edge then we padd with the distance to the edge -1 - # red_mask = np.pad(red_mask,pad_width=((expanx,expanx-1),(expany,expany-1)),mode='constant',constant_values=0) - # #,((expanx,expanx),(expany,expany)),mode='constant',constant_values=0) - # #plt.imshow(red_mask) - # #plt.savefig('red_mask.png') - # red_mask = red_mask#[0:-1,0:-1] - # #print('red_mask shape',red_mask.shape) - # red_image = image[ymin-expanx:ymax+expanx,xmin-expany:xmax+expany] - # #c - # plt.figure(figsize=(10,10)) - # plt.imshow(red_image,cmap='gray',origin='lower',vmin=1E-12,vmax=1E-10) - # plt.contour(red_mask) - # plt.savefig('red_image.png') - - # #print('red_image shape',red_image.shape) - # # dilate the mask until the flux converges. - # # check that the shapes are the same - # if red_mask.shape != red_image.shape: - # print('red_mask and red_image shapes are not the same.') - # print('Skipping source...') - # continue - # red_mask = curve_of_growth_dilation(red_mask,red_image) - # plt.figure(figsize=(10,10)) - # plt.imshow(red_image,cmap='gray',origin='lower',vmin=1E-12,vmax=1E-10) - # plt.contour(red_mask) - # plt.savefig('red_image_after.png') - - # pdb.set_trace() - - # #print('red_mask shape',red_mask.shape) - # # reduce the mask and recalculate the boundig box. - # # plt.imshow(red_mask) - # # plt.savefig('red_mask.png') - # xminn,yminn,xmaxn,ymaxn = utils.bounding_box_cpu(red_mask) - # xmax = xmaxn - expany + xmax - # xmin = xminn - expany + xmin - # ymin = yminn - expanx + ymin - # ymax = ymaxn - expanx + ymax - # red_mask = red_mask[yminn:ymaxn,xminn:xmaxn] - # red_image = red_image[yminn:ymaxn,xminn:xmaxn] - - contour = utils._get_polygons_in_bbox( - xmin, xmax, ymin, ymax, x1[i], y1[i], Birth[i], Death[i], red_mask, 0, 0 + if verbose: + print( + f"{NOTICE}Applying analysis threshold and labeling connected components...{RESET}" ) - # TEST # - # from matplotlib.pylab import plt - - # plt.imshow(red_image, cmap="gray", origin="lower") - # plt.imshow(red_mask, alpha=0.5, origin="lower") - # # plot the countour of the red_mask - # plt.plot(contour[:, 1] - xmin, contour[:, 0] - ymin, color="red") - # plt.show() - - # # # # - - source_props = utils.get_region_props(red_mask, image=red_image) - source_props = utils.props_to_dict(source_props[0]) - # print(background_map) - background_map = bg[i] / sigma # background_map - bg_rms = background_map - # print('Mean bg',mean_bg) - mean_bg_s = mean_bg[i] - - background_map = np.random.normal(mean_bg_s, background_map, red_mask.shape) - - red_background_mask = np.where(red_mask == 0, np.nan, red_mask * background_map) - - Noise = np.nansum(red_background_mask) - - # print('Noise',Noise) - peak_coords = np.where(red_image == source_props["max_intensity"]) - - y_peak_loc = peak_coords[0][0] - x_peak_loc = peak_coords[1][0] - # print('x_peak_loc',x_peak_loc) - # print('y_peak_loc',y_peak_loc) - shape = red_image.shape - - Area = source_props["area"] - # print('Class',Class[i]) - SNR = 0 # default value gets overwritten if calculated - if mode == "Radio" or mode == "optical": - shape = red_image.shape - # if shape is smaller than 100 in any direction then we add padding evenly to each side. - if shape[0] < 100: - pad = int((100 - shape[0]) / 2) - red_image = np.pad( - red_image, ((pad, pad), (0, 0)), mode="constant", constant_values=0 - ) - red_background_mask = np.pad( - red_background_mask, - ((pad, pad), (0, 0)), - mode="constant", - constant_values=0, - ) - red_mask = np.pad( - red_mask, ((pad, pad), (0, 0)), mode="constant", constant_values=0 - ) - shape = red_image.shape - y = y_peak_loc + pad - - else: - x = y_peak_loc - y = x_peak_loc - - if shape[1] < 100: - pad = int((100 - shape[1]) / 2) - red_image = np.pad( - red_image, ((0, 0), (pad, pad)), mode="constant", constant_values=0 - ) - red_background_mask = np.pad( - red_background_mask, - ((0, 0), (pad, pad)), - mode="constant", - constant_values=0, - ) - red_mask = np.pad( - red_mask, ((0, 0), (pad, pad)), mode="constant", constant_values=0 - ) - shape = red_image.shape - x = x_peak_loc + pad - - else: - x = y_peak_loc - y = x_peak_loc - - # print(shape) - - if mode == "optical": - - MAJ = psf_fwhm_p / 2.355 # maybe use full expression in future. - MIN = MAJ - BPA = 0 - - Model_Beam = utils.model_beam_func( - source_props["max_intensity"], shape, x, y, MAJ, MIN, BPA - ) - # plt.figure(figsize=(10,10)) - # plt.imshow(Model_Beam) - # plt.scatter(x,y,color='red',marker='x',s=10) - # plot the contour of the red_mask - # plt.contour(red_mask) - # plt.savefig('Model_Beam_test.png') - # pdb.set_trace() - - else: - BMAJ = BMAJ - BMIN = BMIN - - Model_Beam = utils.model_beam_func( - source_props["max_intensity"], shape, x, y, BMAJ / 2, BMIN / 2, BPA - ) - - if mode == "Radio": - Flux_total = ( - np.nansum(red_mask * red_image - red_background_mask) / Beam - ) - else: - - Flux_total = np.nansum(red_mask * red_image - red_background_mask) - - Flux_peak = ( - np.nanmax(red_mask * red_image) - - red_background_mask[y_peak_loc, x_peak_loc] - ) - - Flux_correction_factor = utils.flux_correction_factor(red_mask, Model_Beam) - Flux_total_err = 0 - SNR = 0 - if Area < 100: - # we assume that the flux cannot be corrected - Flux_total = Flux_total * Flux_correction_factor + # Vectorized boolean mask creation + analysis_mask = image > (background_map + analysis_threshold * background_rms_map) + labeled_image = label(analysis_mask, connectivity=2) - # Calcualte Radio Flux error - if mode == "Radio": - # adapted from https://github.com/mhardcastle/radioflux/blob/master/radioflux/radioflux.py - rms = bg[i] / sigma - gfactor = 2 * np.sqrt(2 * np.log(2)) - Beam_area = 2 * np.pi * (BMAJ * BMIN) / gfactor + if verbose: + print(f"{NOTICE}Measuring region properties...{RESET}") - Flux_total_err = rms * np.sqrt(Area / Beam_area) - SNR = Flux_total / Flux_total_err - - elif mode == "optical": - SNR = 0 - - padding = 0 - - else: - - Flux_total = np.nansum(red_mask * red_image - red_background_mask) - Flux_peak = ( - np.nanmax(red_mask * red_image) - - red_background_mask[y_peak_loc, x_peak_loc] - ) - Flux_correction_factor = np.nan - padding = 0 - Flux_total_err = 0 - - # SNR = Flux_total / Flux_total_err - # Noise = np.std(red_background_mask) + # Use regionprops_table for C-level fast property extraction + properties_table = regionprops_table( + labeled_image, + properties=("label", "bbox", "area"), + ) - Xc = source_props["centroid"][1] + xmin - padding - Yc = source_props["centroid"][0] + ymin - padding + # Initialize Polars DataFrame directly from the dictionary of arrays + props_df = pl.DataFrame(properties_table) - Xc = Xc # + X0[i] - Yc = Yc # + Y0[i] - bbox1 = source_props["bbox"][0] + xmin - padding - bbox2 = source_props["bbox"][1] + ymin - padding - bbox3 = source_props["bbox"][2] + xmin - padding - bbox4 = source_props["bbox"][3] + ymin - padding + if verbose: + print(f"{NOTICE}Initial regions found: {props_df.height}{RESET}") + print( + f"{NOTICE}Filtering regions by area ({area_limit} <= area <= {max_area_limit} pixels)...{RESET}" + ) - Maj = source_props["major_axis_length"] - Min = source_props["minor_axis_length"] - Pa = source_props["orientation"] + # ---> FAST POLARS FILTERING <--- + # Standard processing queue + filtered_props_df = props_df.filter( + (pl.col("area") >= area_limit) & (pl.col("area") <= max_area_limit) + ) - # print(Flux_total_err) - if Edge_flags[i] != 1: + # Flagged massive islands + massive_props_df = props_df.filter(pl.col("area") > max_area_limit) - params.append( - [ - IDs[i], - Birth[i], - Death[i], - x1[i], - y1[i], - x2[i], - y2[i], - Flux_total, - Flux_total_err, - Flux_peak, # this should be the peak flux. - Flux_correction_factor, - Area, - Xc, - Yc, - bbox1, - bbox2, - bbox3, - bbox4, - Maj, - Min, - Pa, - parent_tag[i], - Class[i], - SNR, - Noise, - X0[i], - Y0[i], - mean_bg_s, - bg_rms, - Edge_flags[i], - contour, - enclosed_i[i], - ] + if verbose: + if massive_props_df.height > 0: + print( + f"{WARNING}WARNING{RESET}: Flagged {massive_props_df.height} massive region(s) to retain for the final catalog.{RESET}" ) + print( + f"{NOTICE}Regions queued for homology processing: {filtered_props_df.height}{RESET}" + ) - polygons.append(contour) - # except: - # print('Error in optical characteristing!') - # print('Source ID: ',IDs[i]) - # print('Skipping source...') - # continue - # print(params) - # print(len(params[0])) - - return create_params_df(False, params), polygons - - -def create_polygons(use_gpu, catalogue=None, cutout=None, output=None, cutupts=None): - - if use_gpu: - - try: - - import cupy as cp - - except: - - raise ImportError("cupy not installed. GPU acceleration not possible.") - - image = cutout - if use_gpu: - image_gpu = cp.asarray(image, dtype=cp.float64) - - Birth = catalogue["Birth"].to_numpy() - Death = catalogue["Death"].to_numpy() - - X0 = catalogue["X0_cutout"].to_numpy() - Y0 = catalogue["Y0_cutout"].to_numpy() - IDs = catalogue["ID"].to_numpy() - bbox1_og = catalogue["bbox1"].to_list() - bbox2_og = catalogue["bbox2"].to_list() - bbox3_og = catalogue["bbox3"].to_list() - bbox4_og = catalogue["bbox4"].to_list() - - x1 = catalogue["x1"].to_numpy() # - 1 #- X0 - y1 = catalogue["y1"].to_numpy() # - 1#- X0 - polygons = [] - for i, source in tqdm( - enumerate(Birth), - total=len(Birth), - desc="Creating contours..", - disable=not output, - ): - - if use_gpu == True: - - cropped_image_gpu = image_gpu[ - bbox1_og[i] - 1 : bbox3_og[i] + 1, bbox2_og[i] - 1 : bbox4_og[i] + 1 - ] - - try: - red_image, red_mask, xmin, xmax, ymin, ymax = ( - large_mask_red_image_procc_GPU( - Birth[i], - Death[i], - x1[i] - bbox1_og[i] + 1, - y1[i] - bbox2_og[i] + 1, - cropped_image_gpu, - X0[i], - Y0[i], - ) - ) - except: - - print("Error in GPU processing!") - print("Source ID: ", IDs[i]) - print("x1:", x1[i] - bbox1_og[i] + 1) - print("y1:", y1[i] - bbox2_og[i] + 1) - print("Birth:", Birth[i]) - print("Death:", Death[i]) - # print('Cropped_image',cropped_image_gpu.get()) - - red_mask = red_mask.astype(int) + # Extract standard metadata (for the multiprocessing pool) + bboxes = list( + zip( + filtered_props_df["bbox-0"].to_numpy(), + filtered_props_df["bbox-1"].to_numpy(), + filtered_props_df["bbox-2"].to_numpy(), + filtered_props_df["bbox-3"].to_numpy(), + ) + ) - red_image = red_image.get() - red_mask = red_mask.get() - xmin = xmin.get() + bbox2_og[i] - xmax = xmax.get() + bbox2_og[i] - ymin = ymin.get() + bbox1_og[i] - ymax = ymax.get() + bbox1_og[i] + positions = list( + zip( + filtered_props_df["bbox-0"].to_numpy(), + filtered_props_df["bbox-1"].to_numpy(), + ) + ) - else: - cropped_image = image[ - bbox1_og[i] - 1 : bbox3_og[i] + 1, bbox2_og[i] - 1 : bbox4_og[i] + 1 - ] - try: - red_image, red_mask, xmin, xmax, ymin, ymax = ( - large_mask_red_image_procc_CPU( - Birth[i], - Death[i], - int(x1[i]) - int(bbox1_og[i]) + 1, - int(y1[i]) - int(bbox2_og[i]) + 1, - cropped_image, - ) - ) - except: - print("Error in CPU processing!") - print("Source ID: ", IDs[i]) - print("x1:", x1[i] - bbox1_og[i] + 1) - print("y1:", y1[i] - bbox2_og[i] + 1) - print("Birth:", Birth[i]) - print("Death:", Death[i]) + # Extract massive metadata (to bypass pool but append to catalog) + massive_bboxes = list( + zip( + massive_props_df["bbox-0"].to_numpy(), + massive_props_df["bbox-1"].to_numpy(), + massive_props_df["bbox-2"].to_numpy(), + massive_props_df["bbox-3"].to_numpy(), + ) + ) - red_mask = red_mask.astype(int) + massive_positions = list( + zip( + massive_props_df["bbox-0"].to_numpy(), + massive_props_df["bbox-1"].to_numpy(), + ) + ) - xmin = xmin + bbox2_og[i] - xmax = xmax + bbox2_og[i] - ymin = ymin + bbox1_og[i] - ymax = ymax + bbox1_og[i] + source_islands = { + "bboxes": bboxes, + "positions": positions, + "massive_bboxes": massive_bboxes, # <-- New: Saved massive bounding boxes + "massive_positions": massive_positions, # <-- New: Saved massive coordinates + } - contour = utils._get_polygons_in_bbox( - xmin, xmax, ymin, ymax, x1[i], y1[i], Birth[i], Death[i], red_mask, 0, 0 - ) - polygons.append(contour) + if verbose: + print(f"{NOTICE}Source island creation complete.{RESET}") - catalogue["contour"] = polygons - return catalogue + return source_islands diff --git a/DRUID/src/utils.py b/DRUID/src/utils.py index 07816c1..dd7595c 100644 --- a/DRUID/src/utils.py +++ b/DRUID/src/utils.py @@ -1,589 +1,79 @@ -""" -File: utils.py -Author: Rhys Shaw -Date: 23/12/2023 -Version: v1.0 -Description: Utility functions for DRUID - -""" - -from astropy.io import fits +import polars as pl import numpy as np -from scipy.ndimage import gaussian_filter -from skimage import measure -from scipy.ndimage import label -from astropy.wcs import WCS - - -try: - import cupy as cp - from cupyx.scipy.ndimage import label as cupy_label - -except: - - pass - - -def open_image(PATH: str): - """ - - Function to open fits image and return image and header. - - - Args: - PATH (str): Path to the fits image. - - Returns: - image (np.ndarray): The image data. - header (astropy.io.fits.header.Header): The image header. - - - """ - - hdul = fits.open(PATH) - image = hdul[0].data - header = hdul[0].header - hdul.close() - - image_shape = image.shape - - if len(image_shape) == 4: - - image = np.squeeze(image, axis=(0, 1)) - - if len(image_shape) == 3: - - image = np.squeeze(image, axis=(0)) +from astropy.io import fits +RED = "\033[38;5;34m" +TITLE = "\033[38;5;75m" +LINK = "\033[38;5;21m" +GOLD = "\033[38;5;184m" +RESET = "\033[0m" +BOLD = "\033[1m" +NOTICE = f"{BOLD}{GOLD}Info{RESET}: " +ERROR = "\033[38;5;160m" +WARNING = "\033[38;5;220m" +CODEBLOCK = "\033[48;5;231m" +BLACK = "\033[38;5;16m" +GREEN = "\033[38;5;34m" + + +def get_image_from_path(image_path): + with fits.open(image_path) as hdul: + image = hdul[0].data + header = hdul[0].header + + if image.ndim == 3: + image = image[0, :, :] + elif image.ndim == 4: + image = image[0, 0, :, :] return image, header -def get_psf_FWHM(header: fits.header.Header): - """Return the psf FWHM in pixels. - - Args: - header (fits.header.Header): _description_ - - Returns: - _type_: _description_ - """ - psf_FWHM_arcsec = header["PSF_FWHM"] - ## convert to pixels assumes square pixels and uniform pixel scale - arcseconds_per_pixel = abs(header["CD1_1"]) * 3600 * -1 - psf_FWHM_pixels = psf_FWHM_arcsec / arcseconds_per_pixel - - return psf_FWHM_pixels - - -def get_EFFGAIN(header: fits.header.Header): - """Return the effective gain of the image. - - Args: - header (fits.header.Header): _description_ - - Returns: - _type_: _description_ - """ - EFFGAIN = header["EFFGAIN"] - - return EFFGAIN - - -def get_EFFRON(header: fits.header.Header): - """Return the effective read noise of the image. - - Args: - header (fits.header.Header): _description_ - - Returns: - _type_: _description_ - """ - EFFRON = header["EFFRON"] - - return EFFRON - - -def get_EXPTIME(header: fits.header.Header): - """Return the exposure time of the image. - - Args: - header (fits.header.Header): _description_ - - Returns: - _type_: _description_ - """ - EXPTIME = header["EXPTIME"] - - return EXPTIME - - -def cut_image(size: int, image: np.ndarray): - """Cuts an image into smaller images of the specified size (square). - - Args: - size (int): The size of the cutouts. - image (np.ndarray): The image to be cut. - - Returns: - cutouts (list): A list of the cutout images. - coords (list): A list of the coordinates of the cutouts. - - """ - - cutouts = [] - coords = [] - - for i in range(0, image.shape[0], size): - - for j in range(0, image.shape[1], size): - - cutout = image[i : i + size, j : j + size] - - cutouts.append(cutout) - - coords.append([i, j]) - - return cutouts, coords - - -import numpy as np - - -def cut_image_buff(image, cutout_size, buffer_size): - """ - Cuts an image into smaller images of the specified size (square) with optional overlap. - - Args: - image (np.ndarray): The image to be cut. - size (int): The size of the cutouts.s - buffer (int): The size of the buffer for overlap. - - Returns: - cutouts (list): A list of the cutout images. - coords (list): A list of the coordinates of the cutouts. - """ - - # Initialize lists to store cut-up images and their locations - height, width = image.shape - cutup_images = [] - cutup_locations = [] - - # Iterate through the image with the specified cutout size and buffer - for y in range(0, height, cutout_size - buffer_size): - for x in range(0, width, cutout_size - buffer_size): - # coordinates for the current cutout - left = max(0, x - buffer_size) - upper = max(0, y - buffer_size) - right = min(width, x + cutout_size + buffer_size) - lower = min(height, y + cutout_size + buffer_size) - - # Crop the image array to get the cutout - cutout = image[upper:lower, left:right] - - # Append the cutout and its location to the lists - cutup_images.append(cutout) - cutup_locations.append([upper, left]) - - return cutup_images, cutup_locations - - -def remove_duplicates(catalogue): - """ - - Removes duplicate sources from the catalogue. - - Args: - catalogue (pd.DataFrame): The catalogue to be cleaned. - - Returns: - catalogue (pd.DataFrame): The cleaned catalogue. - - """ - # sort by area - catalogue = catalogue.sort_values(by=["distance_from_center"], ascending=False) - catalogue = catalogue.drop_duplicates( - subset=["x1", "y1", "Birth", "Class"], keep="first" - ) - # sort by ID - # sort by ID - # choose the row with the highest area - - return catalogue - - -def smoothing(image: np.ndarray, sigma: float): - """Smooths an image using a gaussian filter. - - Args: - image (np.ndarray): The image to be smoothed. - sigma (float): The sigma value for the gaussian filter. - - Returns: - smoothed_image (np.ndarray): The smoothed image. - - """ - - smoothed_image = gaussian_filter(image, sigma=sigma) - - return smoothed_image - - -def calculate_beam(header: fits.header.Header): - """ - - Calculates the beam of the image in pixels. - - Args: - header (astropy.io.fits.header.Header): The header of the image. - - returns: - beam (float): The beam correction factor. - BMAJ (float): The BMAJ value in pixels. - BMIN (float): The BMIN value in pixels. - BPA (float): The BPA value in pixels. - - """ - - # get beam info from header - - BMAJ = header["BMAJ"] - BMIN = header["BMIN"] - # convert to pixels - arcseconds_per_pixel = header["CDELT1"] * 3600 * -1 - beam_size_arcseconds = header["BMAJ"] * 3600 - BMAJ_oversampled_spacial_width = (BMAJ**2 + beam_size_arcseconds**2) ** 0.5 - BMAJ = BMAJ_oversampled_spacial_width / arcseconds_per_pixel - - beam_size_arcseconds = header["BMIN"] * 3600 - BMIN_oversampled_spacial_width = (BMIN**2 + beam_size_arcseconds**2) ** 0.5 - BMIN = BMIN_oversampled_spacial_width / arcseconds_per_pixel +def combine_polars_catalogs(catalogs: list) -> pl.DataFrame: + if not catalogs: + raise ValueError("No catalogs provided to combine.") - try: - BPA = header["BPA"] + combined_catalog = pl.concat(catalogs) - except KeyError: - BPA = 0 + # Check for 'id' or 'ID' depending on your upstream schema + if "id" in combined_catalog.columns: + combined_catalog = combined_catalog.with_columns( + pl.int_range(1, pl.len() + 1, dtype=pl.Int64).alias("id") + ) + elif "ID" in combined_catalog.columns: + combined_catalog = combined_catalog.with_columns( + pl.int_range(1, pl.len() + 1, dtype=pl.Int64).alias("ID") + ) - return np.pi * (BMAJ) * (BMIN) / (4 * np.log(2)), BMAJ, BMIN, BPA - - -def check_edge(mask): - - return 1 in mask[0, :] or 1 in mask[-1, :] or 1 in mask[:, 0] or 1 in mask[:, -1] - - -def props_to_dict(regionprops): - """ - - Converts the regionprops to a dictionary. - # slow function when all are called. - """ - - dict = { - "area": regionprops.area, - "bbox": regionprops.bbox, - #'bbox_area': regionprops.bbox_area, - "centroid": regionprops.centroid, - #'convex_area': regionprops.convex_area, - "eccentricity": regionprops.eccentricity, - #'equivalent_diameter': regionprops.equivalent_diameter, - #'euler_number': regionprops.euler_number, - #'extent': regionprops.extent, - #'filled_area': regionprops.filled_area, - "major_axis_length": regionprops.major_axis_length, - "minor_axis_length": regionprops.minor_axis_length, - #'moments': regionprops.moments, - #'perimeter': regionprops.perimeter, - #'solidity': regionprops.solidity, - "orientation": regionprops.orientation, - "max_intensity": regionprops.max_intensity, - } - - return dict - - -def bounding_box_cpu(mask): - # Get the indices of elements that are True - rows, cols = np.where(mask) - # Get the minimum and maximum x and y coordinates - min_y, max_y = np.min(rows), np.max(rows) - min_x, max_x = np.min(cols), np.max(cols) - # Return the bounding box as a tuple of tuples - return min_y, min_x, max_y, max_x - - -def bounding_box_gpu(binary_mask_gpu): - """ - Calculate the bounding box of a binary mask using cupy. - - Parameters: - - binary_mask_gpu: cupy array, binary mask. - - Returns: - - tuple: (min_row, min_col, max_row, max_col), representing the bounding box. - """ - - # Get the indices of elements that are True - - rows, cols = cp.where(binary_mask_gpu) - - # Get the minimum and maximum x and y coordinates - min_y, max_y = cp.min(rows), cp.max(rows) - min_x, max_x = cp.min(cols), cp.max(cols) - - # Return the bounding box as a tuple of tuples - return min_y, min_x, max_y, max_x - - -def get_region_props(mask, image): - - region = measure.regionprops(mask, image) - - return region - - -def model_beam_func(peak_flux, shape, x, y, bmaj, bmin, bpa): - model_beam = np.zeros(shape) - model_beam = generate_2d_gaussian( - peak_flux, shape, (x, y), bmaj, bmin, bpa, norm=False - ) - return model_beam - - -def flux_correction_factor(mask, Model_Beam): - - # calculate the correction factor - - model_beam_flux = np.sum(Model_Beam) - masked_beam_flux = np.sum(mask * Model_Beam) - - correction_factor = model_beam_flux / masked_beam_flux - - return correction_factor + return combined_catalog def generate_2d_gaussian(A, shape, center, sigma_x, sigma_y, angle_deg=0, norm=True): - """ - - Generate a 2D elliptical Gaussian distribution on a 2D array. - - Parameters: - - shape (tuple): Shape of the output array (height, width). - center (tuple): Center of the Gaussian distribution (x, y). - sigma_x (float): Standard deviation along the x-axis. - sigma_y (float): Standard deviation along the y-axis. - angle_deg (float): Rotation angle in degrees (default is 0). - - Returns: - - ndarray: 2D array containing the Gaussian distribution. - - """ x, y = np.meshgrid(np.arange(shape[1]), np.arange(shape[0])) x_c, y_c = center angle_rad = np.radians(angle_deg) - # Rotate coordinates - x_rot = (x - x_c) * np.cos(angle_rad) - (y - y_c) * np.sin(angle_rad) y_rot = (x - x_c) * np.sin(angle_rad) + (y - y_c) * np.cos(angle_rad) - # Calculate Gaussian values - gaussian = A * np.exp(-(x_rot**2 / (2 * sigma_x**2) + y_rot**2 / (2 * sigma_y**2))) if norm: return gaussian / (2 * np.pi * sigma_x * sigma_y) - else: - return gaussian - - -def get_enclosing_mask_gpu(x, y, mask): - """Returns the connected components inside the mask starting from the point (x, y) using the GPU. - - - Args: - x (int): x location of the source in the mask. - y (int): y location of the source in the mask. - mask (cp.ndarray): unfiltered mask of source (on GPU memory). + return gaussian - Returns: - component_mask(np.ndarray): Mask for the source in the provided coordinates. - """ - labeled_mask, num_features = cupy_label(mask) - - # Check if the specified pixel is within the mask - if 0 <= x < mask.shape[1] and 0 <= y < mask.shape[0]: - label_at_pixel = labeled_mask[y, x] - - if label_at_pixel != 0: - # Extract the connected component containing the specified pixel - component_mask = labeled_mask == label_at_pixel - return component_mask # still in GPU memory - else: - return None - else: - return None - - -def get_mask_CPU(x1, y1, Birth, Death, img): - """Get mask for a single row uses the CPU - - Args: - row (pd.Series): row we want to get the mask for. - img (np.ndarray): image that the source is in. - - Returns: - mask_enclosed(np.ndarray): Array of the mask - - """ - - mask = np.zeros(img.shape) - mask = np.logical_or(mask, np.logical_and(img <= Birth, img > Death)) - mask_enclosed = get_enclosing_mask_CPU(int(y1), int(x1), mask) - - return mask_enclosed - - -def get_enclosing_mask_CPU(x, y, mask): - """ - Returns the connected components inside the mask starting from the point (x, y). - """ - labeled_mask, num_features = label(mask) - - # Check if the specified pixel is within the mask - if 0 <= x < mask.shape[1] and 0 <= y < mask.shape[0]: - label_at_pixel = labeled_mask[y, x] - - if label_at_pixel != 0: - # Extract the connected component containing the specified pixel - component_mask = labeled_mask == label_at_pixel - return component_mask - else: - # The specified pixel is not part of any connected component - return None - else: - # The specified pixel is outside the mask - return None - - -def get_mask_GPU(Birth, Death, x1, y1, img): - """Gets mask for a single row using the GPU (requires cupy) - - Args: - Birth (float): Birth value of the row of interest. - Death (float): Death value of the row of interest. - row (pd.series): row of interest. - img (cp.ndarray): Image containing the source (in the GPUs memory) - - Returns: - mask_enclosed (np.ndarray): The mask returned in normal memory. - - """ - - mask = cp.zeros(img.shape, dtype=cp.float64) - mask = cp.logical_or(mask, cp.logical_and(img <= Birth, img > Death)) - mask_enclosed = get_enclosing_mask_gpu(int(y1), int(x1), mask) - - return mask_enclosed - - -def _get_polygons_gpu(x1: int, y1: int, birth: float, death: float, image): - """ - Returns the polygon of the enclosed area of the point (x,y) in the mask. - """ - # is the image on the GPU memory? - # if not cp.is_cuda_array(self.image): - # self.image = cp.asarray(self.image, dtype=cp.float64) - - mask = cp.zeros(image.shape) - mask = cp.logical_or(mask, cp.logical_and(image <= birth, image > death)) - mask = get_enclosing_mask_GPU(int(y1), int(x1), mask) - contour = measure.find_contours(mask, 0)[0] - - return contour - - -def _get_polygons_CPU(x1, y1, birth, death, image: np.ndarray): - """ - Returns the polygon of the enclosed area of the point (x,y) in the mask. - """ - - mask = np.zeros(image.shape) - mask = np.logical_or(mask, np.logical_and(image <= birth, image > death)) - mask = get_enclosing_mask_CPU(int(y1), int(x1), mask) - contour = measure.find_contours(mask, 0)[0] - - return contour - - -import matplotlib.pyplot as plt - - -def _get_polygons_in_bbox( - Xmin, Xmax, Ymin, Ymax, x1, y1, birth, death, mask, X0, Y0, pad=1 -): - - mask = np.pad(mask, pad, mode="constant", constant_values=0) - # print(mask) - # try: - contour = measure.find_contours(mask, 0)[0] - # except: - # plt.imshow(mask) - # plt.savefig('mask.png') - contour = measure.find_contours(mask, 0)[0] - # print(contour) - # remove the border - contour[:, 0] -= pad - contour[:, 1] -= pad - - # correct the coordinates to the original image - contour[:, 0] += Ymin + Y0 - contour[:, 1] += Xmin + X0 - - return contour - - -def xy_to_RaDec(x, y, header, mode): - """ - - Convert an X and Y coordinate to RA and Dec using a header file with astropy. +def model_beam_func(peak_flux, shape, x, y, bmaj, bmin, bpa): + return generate_2d_gaussian(peak_flux, shape, (x, y), bmaj, bmin, bpa, norm=False) - Parameters: - x (float): The X coordinate. - y (float): The Y coordinate. - header_file (str): The path to the FITS header file. - stokes (int): The Stokes dimension. - freq (int): The frequency dimension. - Returns: - tuple: A tuple containing the RA and Dec in degrees. +def calculate_radec(catalog: pl.DataFrame, header) -> pl.DataFrame: + # Convert pixel coordinates to RA/Dec using WCS. + from astropy.wcs import WCS - """ wcs = WCS(header) - print("-----------------") - print("Mode :", mode) - if mode == "Radio": - - stokes = 0 # stokes and freq are not used in this function. - freq = 0 # stokes and freq are not used in this function. - ra, dec, _, _ = wcs.all_pix2world(x, y, stokes, freq, 0) - - elif mode == "optical": - print("Optical") - # try: - # image is 2d so no stokes or freq - ra, dec = wcs.all_pix2world(x, y, 0) - - # except: - # try radio - # stokes = 0 - # freq = 0 - # ra, dec, _, _ = wcs.all_pix2world(x, y, stokes, freq, 0) - elif mode == "other": - print("Other") - ra, dec = wcs.all_pix2world(x, y, 0) - - return ra, dec + ra_dec = wcs.all_pix2world( + catalog["centroid_x"].to_numpy(), catalog["centroid_y"].to_numpy(), 0 + ) + ra, dec = ra_dec + catalog = catalog.with_columns(pl.Series("ra", ra), pl.Series("dec", dec)) + return catalog diff --git a/DRUID/tests/test_background.py b/DRUID/tests/test_background.py index 73e5dcc..e9137a9 100644 --- a/DRUID/tests/test_background.py +++ b/DRUID/tests/test_background.py @@ -1,39 +1,40 @@ -from DRUID.src.background import calculate_background +""" +Unit tests for background and RMS map estimation. +""" + import pytest import numpy as np -# test the use of the calculate_background function - -# prehaps use a sample more realistic array with cahracteristics of Radio and Optical data. -test_array = np.array([[1,2,3,4,5], - [6,7,8,9,10], - [11,12,13,14,15], - [16,17,18,19,20], - [21,22,23,24,25]]) -def test_calculate_background_not_valid(): - # test that a parameter not valid will fail - with pytest.raises(ValueError): - calculate_background(np.ones((10,10)), mode='not_valid') - -def test_mad_std_value(): - # test that the mad_std value is correct - # assert has to be approximately equal to the value because of float - - assert calculate_background(test_array, mode='mad_std')[0]== pytest.approx(8.8956, 0.001) - -def test_rms_std_value(): - # test that the rms_std value is correct - # assert has to be approximately equal to the value because of float - assert calculate_background(test_array, mode='rms')[0]== pytest.approx(14.8660, 0.001) - -def test_mean_value(): - # test that the mean value is correct - # assert has to be approximately equal to the value because of float - assert calculate_background(test_array, mode='mad_std')[1]== pytest.approx(13.0, 0.001) - -def test_sigma_clipping_value(): - # test that the sigma_clipping value is correct - # assert has to be approximately equal to the value because of float - assert calculate_background(test_array, mode='sigma_clip')[0]== pytest.approx(7.21110, 0.001) - -def test_SEX_background_value(): - assert calculate_background(test_array, mode='SEX')[0]== pytest.approx(7.21110, 0.001) \ No newline at end of file +from astropy.io import fits +from photutils.background import MedianBackground +from DRUID.src.background import make_source_mask, calculate_background_maps + + +@pytest.fixture +def dummy_fits_file(tmp_path) -> str: + """Fixture to generate a dummy FITS file with a noise floor.""" + rng = np.random.default_rng(42) # Reproducible seed + data = rng.uniform(5, 15, size=(100, 100)) + file_path = tmp_path / "dummy_image.fits" + fits.PrimaryHDU(data).writeto(file_path) + return str(file_path) + + +def test_calculate_background_maps_defaults(dummy_fits_file): + """Test background map generation with default estimators.""" + bg_map, rms_map = calculate_background_maps(dummy_fits_file) + + assert bg_map.shape == (100, 100) + assert rms_map.shape == (100, 100) + assert bg_map.dtype.type == np.float64 + assert np.all(bg_map > 0) + assert np.all(rms_map > 0) + + +def test_calculate_background_maps_array_input(): + """Test background map generation passing a NumPy array directly.""" + rng = np.random.default_rng(42) + data = rng.normal(10, 1, size=(50, 50)) + bg_map, rms_map = calculate_background_maps(data, bg_estimator=MedianBackground()) + + assert bg_map.shape == (50, 50) + np.testing.assert_allclose(np.median(bg_map), 10.0, rtol=0.1) diff --git a/DRUID/tests/test_homology.py b/DRUID/tests/test_homology.py index a5adc1c..8535394 100644 --- a/DRUID/tests/test_homology.py +++ b/DRUID/tests/test_homology.py @@ -1,4 +1,35 @@ +""" +Unit tests for Persistent Homology computation (Cripser & Polars). +""" import pytest +import numpy as np +import polars as pl +from DRUID.src.homology import compute_homology -def test_homology(): - pass \ No newline at end of file +def test_compute_homology_schema(): + """ + Test that compute_homology returns the mathematically expected schema. + """ + image = np.zeros((50, 50)) + # Add a topological maximum (birth) + image[20:30, 20:30] = 10.0 + image[25, 25] = 20.0 + + df = compute_homology( + image, + analysis_threshold=1.0, + lifetime_limit=0.1, + area_size_threshold=2 + ) + + assert isinstance(df, pl.DataFrame) + assert not df.is_empty() + + expected_cols = { + "birth", "death", "x1", "y1", "lifetime", "lifetimeFrac", + "area", "bbox_min_y", "ID", "encloses", "parent_tag", "contour" + } + assert expected_cols.issubset(set(df.columns)) + + # Check that birth is strictly greater than death + assert df.filter(pl.col("birth") <= pl.col("death")).is_empty() \ No newline at end of file diff --git a/DRUID/tests/test_main.py b/DRUID/tests/test_main.py index 79362f8..9504284 100644 --- a/DRUID/tests/test_main.py +++ b/DRUID/tests/test_main.py @@ -1,16 +1,64 @@ -from DRUID.main import sf +""" +Integration tests for the DRUID main pipeline. +""" import pytest import numpy as np +import polars as pl +from scipy.ndimage import gaussian_filter +from DRUID.main import sf, _worker +import DRUID.main as main_module -def test_sf(): - ''' - Test initalisation of the main Class - ''' - arr2d = np.random.rand(100,100) - assert sf(image=arr2d, image_path=None, - mode="Radio",pb_path=None,cutup=False, - cutup_size=None,cutup_buff=None,output=False, - area_limit=5,smooth_sigma=1,nproc=1,GPU=True, - header=None,Xoff=None,Yoff=None) != None +@pytest.fixture +def pipeline_image(): + """Provides a baseline image with a single synthetic source.""" + img = np.random.normal(5, 0.5, size=(100, 100)) + img[45:55, 45:55] += 15.0 # Bright source + return img + +def test_pipeline_sequential_with_smoothing(pipeline_image): + """Test end-to-end pipeline with smooth_sigma > 0.""" + finder = sf(image=pipeline_image, verbose=False, num_threads=1, smooth_sigma=1.0) + finder.set_background(analysis_threshold=3) + finder.phsf() + + assert hasattr(finder, "catalog") + assert isinstance(finder.catalog, pl.DataFrame) + assert not finder.catalog.is_empty() + assert "ID" in finder.catalog.columns + # Ensure smoothed_image was allocated + assert getattr(finder, "smoothed_image", None) is not None + +def test_worker_function(pipeline_image): + """ + Test the inner multiprocessing worker with raw and smoothed global arrays. + """ + bg = np.ones((100, 100)) * 5.0 + rms = np.ones((100, 100)) * 0.5 + smoothed_image = gaussian_filter(pipeline_image, sigma=1.0) + + # Inject into the main module namespace + main_module.global_image = pipeline_image + main_module.global_smoothed_image = smoothed_image + main_module.global_background_map = bg + main_module.global_background_rms_map = rms + + bbox = (40, 40, 60, 60) + position = (40, 40) + island_info = (bbox, position) + + result_cat = _worker( + island_info, + analysis_threshold=3.0, + lifetime_limit=0.0, + lifetime_limit_fraction=1.0, + mode="radio", BMAJ=2.0, BMIN=2.0 + ) + + assert isinstance(result_cat, pl.DataFrame) + assert not result_cat.is_empty() - \ No newline at end of file + # Cleanup namespace + main_module.global_image = None + main_module.global_smoothed_image = None + main_module.global_background_map = None + main_module.global_background_rms_map = None \ No newline at end of file diff --git a/DRUID/tests/test_properties.py b/DRUID/tests/test_properties.py new file mode 100644 index 0000000..5cf2a42 --- /dev/null +++ b/DRUID/tests/test_properties.py @@ -0,0 +1,53 @@ +""" +Unit tests for source property calculations (flux, SNR, geometry). +""" +import pytest +import numpy as np +import polars as pl +from DRUID.src.properties import calculate_properties, calculate_radio_flux_error + +def test_calculate_radio_flux_error(): + """Test flux error analytic derivation.""" + bg_rms = np.array([1.0, 1.0, 1.0]) + area = 100 + bmaj, bmin = 2.0, 2.0 + + err = calculate_radio_flux_error(bg_rms, area, bmaj, bmin) + assert not np.isnan(err) + assert err > 0 + +def test_calculate_properties(): + """ + Test properties calculated on raw image while mask bounds dictate via smoothed image. + """ + raw_image = np.zeros((30, 30)) + smoothed_image = np.zeros((30, 30)) + bg = np.zeros((30, 30)) + rms = np.ones((30, 30)) + + # Simulate a bright sharp core in raw image + raw_image[10:20, 10:20] = 5.0 + + # Simulate a wider, dimmer dispersion in smoothed image + smoothed_image[8:22, 8:22] = 2.0 + + cat = pl.DataFrame({ + "birth": [2.1], # Mask encompasses smoothed_image's 2.0 block + "death": [0.0], + "x1": [15], + "y1": [15], + "area": [196] # (14 x 14 block area) + }) + + result = calculate_properties( + cat, raw_image=raw_image, smoothed_image=smoothed_image, + background=bg, background_rms=rms, position=(0,0), + analysis_threshold=1.0, mode="radio", BMAJ=2.0, BMIN=2.0 + ) + + # Peak flux should pull from the 5.0 raw array, NOT the 2.0 smoothed one + assert result["flux_peak"][0] == 5.0 + + # Total flux is the 10x10 core inside the 14x14 bounds + # (100 pixels * 5.0) + (96 pixels * 0.0) = 500 + np.testing.assert_allclose(result["flux"][0], 500.0, atol=1e-3) \ No newline at end of file diff --git a/DRUID/tests/test_source.py b/DRUID/tests/test_source.py index 2b1d6fb..02eacef 100644 --- a/DRUID/tests/test_source.py +++ b/DRUID/tests/test_source.py @@ -1,4 +1,33 @@ +""" +Unit tests for initial source island thresholding and bounding box generation. +""" import pytest +import numpy as np +from DRUID.src.source import create_source_islands -def test_source(): - pass \ No newline at end of file +def test_create_source_islands(): + """ + Test extraction of standard and massive source islands based on thresholds. + """ + image = np.zeros((100, 100)) + bg = np.ones((100, 100)) + rms = np.ones((100, 100)) * 0.5 + + # Create one standard source (area ~ 25) + image[20:25, 20:25] = 10 + # Create one massive source (area = 400) + image[50:70, 50:70] = 10 + + # threshold = 1 + (3 * 0.5) = 2.5 + islands = create_source_islands( + image, bg, rms, + analysis_threshold=3, area_limit=10, max_area_limit=200, verbose=False + ) + + # Check standard queue + assert len(islands["bboxes"]) == 1 + assert len(islands["positions"]) == 1 + + # Check massive queue + assert len(islands["massive_bboxes"]) == 1 + assert len(islands["massive_positions"]) == 1 \ No newline at end of file diff --git a/DRUID/tests/test_utils.py b/DRUID/tests/test_utils.py index 9db6172..1747cb2 100644 --- a/DRUID/tests/test_utils.py +++ b/DRUID/tests/test_utils.py @@ -1,93 +1,43 @@ -import pytest +""" +Unit tests for DRUID utilities. +""" +import pytest import numpy as np -from astropy.io import fits -from DRUID.src.utils import smoothing, get_region_props, model_beam_func, flux_correction_factor, bounding_box_cpu, open_image, calculate_beam, xy_to_RaDec, generate_2d_gaussian -import pandas - -PATH_test_image_file = 'https://drive.google.com/uc?id=10a6goXcr6wEHX9U5LQ07cCEo2nGQ9QK5' - - - -test_image = np.array([[1,2,3,4,5], - [6,7,8,9,10], - [11,12,13,14,15], - [16,17,18,19,20], - [21,22,23,24,25]]) - -def test_smoothing_works(): - # test that the smoothing function returns the correct shape image that went in. - - assert smoothing(test_image, 3).shape == (5,5) - - -mask = np.array([[0,0,0,0,0], - [0,1,1,1,0], - [0,1,1,1,0], - [0,1,1,1,0], - [0,0,0,0,0]]) -def test_regionprops_function_params(): - # test that the regionprops function returns the correct shape image that went in. - - region = get_region_props(mask, test_image)[0] - - assert region.area == 9 - assert region.centroid == (2.0,2.0) - assert region.max_intensity == 19 - assert region.major_axis_length == pytest.approx(3.2659, 0.0001) - assert region.minor_axis_length == pytest.approx(3.2659, 0.0001) - assert abs(region.orientation) == pytest.approx(0.7853, 0.001) - -def test_model_beam_function(): - test_beam = np.array([[0.00193045, 0.01426423, 0.03877421, 0.03877421, 0.01426423], - [0.01426423, 0.10539922, 0.2865048 , 0.2865048 , 0.10539922], - [0.03877421, 0.2865048 , 0.77880078, 0.77880078, 0.2865048 ], - [0.03877421, 0.2865048 , 0.77880078, 0.77880078, 0.2865048 ], - [0.01426423, 0.10539922, 0.2865048 , 0.2865048 , 0.10539922]]) - model_beam = model_beam_func(1,(5,5),2.5,2.5,1,1,0) - # make sure the model beam and the test beam are the same. - assert np.allclose(model_beam, test_beam, rtol=1e-05, atol=1e-08) - -def test_flux_correction_factor(): - - # test if the correction value returned is correct. - mask = np.array([[0,0,0,0,0], - [0,1,1,1,0], - [0,1,1,1,0], - [0,1,1,1,0], - [0,0,0,0,0]]) - model_beam = model_beam_func(1,(5,5),2.5,2.5,1,1,0)# gaussian beam model of sixe 5x5. - assert flux_correction_factor(mask, model_beam) == pytest.approx(1.38389, 0.0001) +import polars as pl +from DRUID.src.utils import combine_polars_catalogs, generate_2d_gaussian + +def test_combine_polars_catalogs(): + """ + Test the concatenation and dense ranking of catalog IDs. + """ + df1 = pl.DataFrame({"id": [1, 2], "flux": [10.5, 20.1]}) + df2 = pl.DataFrame({"id": [1, 2], "flux": [30.4, 40.2]}) -def test_bouding_box_cpu(): - # test if the bounding box function returns the correct shape image that went in. - mask = np.array([[0,0,0,0,0], - [0,1,1,1,0], - [0,1,1,1,0], - [0,1,1,1,0], - [0,0,0,0,0]]) - bounding_box = bounding_box_cpu(mask) - assert bounding_box == (1,1,3,3) + combined = combine_polars_catalogs([df1, df2]) -def test_open_image(): - image, header = open_image(PATH_test_image_file) - assert image.shape == (256,256) - assert type(header) == fits.header.Header + assert combined.height == 4 + # Ensure dense ranking re-indexes the IDs sequentially + assert combined["id"].to_list() == [1, 2, 3, 4] + +def test_combine_polars_catalogs_empty(): + """Ensure a ValueError is raised for empty catalog lists.""" + with pytest.raises(ValueError, match="No catalogs provided"): + combine_polars_catalogs([]) + +def test_generate_2d_gaussian(): + """ + Test the 2D Gaussian generation for mathematical correctness. + """ + shape = (50, 50) + center = (25, 25) + sigma = 5 -def test_calculate_beam(): - image, header = open_image(PATH_test_image_file) - header['BMAJ'] = 0.0001388888888888889 - header['BMIN'] = 0.0001388888888888889 - header['BPA'] = 0.0 - beam,bmaj,bmin,bpa = calculate_beam(header) - assert beam == pytest.approx(0.28327, 0.0001) - assert bmaj == pytest.approx(0.5000, 0.0001) - assert bmin == pytest.approx(0.5000, 0.0001) - assert bpa == pytest.approx(0.0, 0.0001) + gauss = generate_2d_gaussian( + A=1.0, shape=shape, center=center, + sigma_x=sigma, sigma_y=sigma, norm=False + ) -def test_xy_to_RaDec(): - x,y = 128,128 - image, header = open_image(PATH_test_image_file) - ra, dec = xy_to_RaDec(128,128,header,mode='Radio') - assert ra == pytest.approx(230.68, 0.0001) - assert dec == pytest.approx(54.64416, 0.0001) - + assert gauss.shape == shape + np.testing.assert_allclose(gauss[25, 25], 1.0, atol=1e-5) + # Check symmetric decay + np.testing.assert_allclose(gauss[20, 25], gauss[30, 25], atol=1e-5) \ No newline at end of file diff --git a/Examples/Resolved_Galaxies.ipynb b/Examples/Resolved_Galaxies.ipynb index f3f6837..a166891 100644 --- a/Examples/Resolved_Galaxies.ipynb +++ b/Examples/Resolved_Galaxies.ipynb @@ -14,6 +14,72 @@ "Some of DRUIDs strengths are demonstrated with its ability to handle nested image features. Like the details resovlable in face on galaxies." ] }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Requirement already satisfied: astrocut in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (1.0.1)\n", + "Requirement already satisfied: asdf>=4.1.0 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from astrocut) (4.5.0)\n", + "Requirement already satisfied: astropy>=5.2 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from astrocut) (7.1.0)\n", + "Requirement already satisfied: cachetools>=5.3.2 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from astrocut) (6.2.0)\n", + "Requirement already satisfied: fsspec>=2022.8.2 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from fsspec[http]>=2022.8.2->astrocut) (2025.9.0)\n", + "Requirement already satisfied: s3fs>=2022.8.2 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from astrocut) (2025.9.0)\n", + "Requirement already satisfied: s3path>=0.5.7 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from astrocut) (0.6.4)\n", + "Requirement already satisfied: roman_datamodels>=0.19.0 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from astrocut) (0.27.0)\n", + "Requirement already satisfied: requests>=2.32.3 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from astrocut) (2.32.5)\n", + "Requirement already satisfied: spherical_geometry>=1.3.0 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from astrocut) (1.3.3)\n", + "Requirement already satisfied: gwcs>=0.21.0 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from astrocut) (0.25.2)\n", + "Requirement already satisfied: scipy in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from astrocut) (1.15.3)\n", + "Requirement already satisfied: Pillow in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from astrocut) (11.2.1)\n", + "Requirement already satisfied: asdf-standard>=1.1.0 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from asdf>=4.1.0->astrocut) (1.4.0)\n", + "Requirement already satisfied: asdf-transform-schemas>=0.3 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from asdf>=4.1.0->astrocut) (0.6.0)\n", + "Requirement already satisfied: jmespath>=0.6.2 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from asdf>=4.1.0->astrocut) (1.0.1)\n", + "Requirement already satisfied: numpy>=1.22 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from asdf>=4.1.0->astrocut) (2.2.6)\n", + "Requirement already satisfied: packaging>=19 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from asdf>=4.1.0->astrocut) (25.0)\n", + "Requirement already satisfied: pyyaml>=5.4.1 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from asdf>=4.1.0->astrocut) (6.0.2)\n", + "Requirement already satisfied: semantic_version>=2.8 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from asdf>=4.1.0->astrocut) (2.10.0)\n", + "Requirement already satisfied: attrs>=22.2.0 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from asdf>=4.1.0->astrocut) (25.3.0)\n", + "Requirement already satisfied: pyerfa>=2.0.1.1 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from astropy>=5.2->astrocut) (2.0.1.5)\n", + "Requirement already satisfied: astropy-iers-data>=0.2025.4.28.0.37.27 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from astropy>=5.2->astrocut) (0.2025.5.26.0.37.21)\n", + "Requirement already satisfied: aiohttp!=4.0.0a0,!=4.0.0a1 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from fsspec[http]>=2022.8.2->astrocut) (3.12.15)\n", + "Requirement already satisfied: aiohappyeyeballs>=2.5.0 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]>=2022.8.2->astrocut) (2.6.1)\n", + "Requirement already satisfied: aiosignal>=1.4.0 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]>=2022.8.2->astrocut) (1.4.0)\n", + "Requirement already satisfied: frozenlist>=1.1.1 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]>=2022.8.2->astrocut) (1.7.0)\n", + "Requirement already satisfied: multidict<7.0,>=4.5 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]>=2022.8.2->astrocut) (6.6.4)\n", + "Requirement already satisfied: propcache>=0.2.0 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]>=2022.8.2->astrocut) (0.3.2)\n", + "Requirement already satisfied: yarl<2.0,>=1.17.0 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]>=2022.8.2->astrocut) (1.20.1)\n", + "Requirement already satisfied: idna>=2.0 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from yarl<2.0,>=1.17.0->aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]>=2022.8.2->astrocut) (3.10)\n", + "Requirement already satisfied: typing-extensions>=4.2 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from aiosignal>=1.4.0->aiohttp!=4.0.0a0,!=4.0.0a1->fsspec[http]>=2022.8.2->astrocut) (4.15.0)\n", + "Requirement already satisfied: asdf_wcs_schemas>=0.5.0 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from gwcs>=0.21.0->astrocut) (0.5.0)\n", + "Requirement already satisfied: asdf-astropy>=0.8.0 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from gwcs>=0.21.0->astrocut) (0.8.0)\n", + "Requirement already satisfied: asdf-coordinates-schemas>=0.4 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from asdf-astropy>=0.8.0->gwcs>=0.21.0->astrocut) (0.4.0)\n", + "Requirement already satisfied: charset_normalizer<4,>=2 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from requests>=2.32.3->astrocut) (3.4.3)\n", + "Requirement already satisfied: urllib3<3,>=1.21.1 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from requests>=2.32.3->astrocut) (2.5.0)\n", + "Requirement already satisfied: certifi>=2017.4.17 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from requests>=2.32.3->astrocut) (2025.8.3)\n", + "Requirement already satisfied: lz4>=4.3.0 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from roman_datamodels>=0.19.0->astrocut) (4.4.4)\n", + "Requirement already satisfied: rad>=0.27.0 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from roman_datamodels>=0.19.0->astrocut) (0.27.0)\n", + "Requirement already satisfied: pyarrow>=10.0.1 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from roman_datamodels>=0.19.0->astrocut) (21.0.0)\n", + "Requirement already satisfied: aiobotocore<3.0.0,>=2.5.4 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from s3fs>=2022.8.2->astrocut) (2.24.2)\n", + "Requirement already satisfied: aioitertools<1.0.0,>=0.5.1 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from aiobotocore<3.0.0,>=2.5.4->s3fs>=2022.8.2->astrocut) (0.12.0)\n", + "Requirement already satisfied: botocore<1.40.19,>=1.40.15 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from aiobotocore<3.0.0,>=2.5.4->s3fs>=2022.8.2->astrocut) (1.40.18)\n", + "Requirement already satisfied: python-dateutil<3.0.0,>=2.1 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from aiobotocore<3.0.0,>=2.5.4->s3fs>=2022.8.2->astrocut) (2.9.0.post0)\n", + "Requirement already satisfied: wrapt<2.0.0,>=1.10.10 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from aiobotocore<3.0.0,>=2.5.4->s3fs>=2022.8.2->astrocut) (1.17.3)\n", + "Requirement already satisfied: six>=1.5 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from python-dateutil<3.0.0,>=2.1->aiobotocore<3.0.0,>=2.5.4->s3fs>=2022.8.2->astrocut) (1.17.0)\n", + "Requirement already satisfied: boto3>=1.16.35 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from s3path>=0.5.7->astrocut) (1.40.18)\n", + "Requirement already satisfied: smart-open>=5.1.0 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from s3path>=0.5.7->astrocut) (7.3.1)\n", + "Requirement already satisfied: s3transfer<0.14.0,>=0.13.0 in /Users/rs17612/anaconda3/envs/DRUID/lib/python3.12/site-packages (from boto3>=1.16.35->s3path>=0.5.7->astrocut) (0.13.1)\n" + ] + } + ], + "source": [ + "!pip install astrocut" + ] + }, { "cell_type": "markdown", "metadata": {}, @@ -23,9 +89,29 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 2, "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Filename: (No file associated with this HDUList)\n", + "No. Name Ver Type Cards Dimensions Format\n", + " 0 PRIMARY 1 PrimaryHDU 9 () \n", + " 1 CUTOUT 1 ImageHDU 93 (1000, 1000) float32 \n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "WARNING: FITSFixedWarning: RADECSYS= 'ICRS ' / Coordinate reference frame \n", + "the RADECSYS keyword is deprecated, use RADESYSa. [astropy]\n", + "WARNING: FITSFixedWarning: 'datfix' made the change 'Set DATE-END to '2016-08-30T09:29:55.786' from MJD-END'. [astropy]\n" + ] + } + ], "source": [ "from astrocut import fits_cut\n", "from astropy.coordinates import SkyCoord\n", @@ -54,23 +140,127 @@ "cell_type": "code", "execution_count": null, "metadata": {}, - "outputs": [], + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + " \n", + "\n", + "\u001b[91m#############################################\u001b[0m\n", + "\u001b[92m\n", + "_______ _______ _________ ______ \n", + "( __ \\ ( ____ )|\\ /|\\__ __/( __ \\ \n", + "| ( \\ )| ( )|| ) ( | ) ( | ( \\ )\n", + "| | ) || (____)|| | | | | | | | ) |\n", + "| | | || __)| | | | | | | | | |\n", + "| | ) || (\\ ( | | | | | | | | ) |\n", + "| (__/ )| ) \\ \\__| (___) |___) (___| (__/ )\n", + "(______/ |/ \\__/(_______)\\_______/(______/ \n", + " \n", + "\u001b[0m\n", + "\u001b[91m#############################################\u001b[0m\n", + "\n", + "\u001b[1mDetector of astRonomical soUrces in optIcal and raDio images\u001b[0m\n", + "\n", + "Version: 1.0\n", + "\n", + "For more information see:\n", + "\u001b[94mhttps://github.com/RhysAlfShaw/DRUID\u001b[0m\n", + "\n", + "Calculating background map and RMS map...\n", + "Calculating background map and RMS map from image.\n", + "Background calculation took 0.21 seconds.\n", + "Background map and RMS map calculated.\n", + "Thresholding to find source islands...\n", + "Labeling connected components took 0.00 seconds. Found 1256 components.\n", + "Calculating region properties took 0.01 seconds. Found 1256 properties.\n", + "Cropping components took 0.02 seconds. Found 143 source islands.\n", + "Thresholding took 0.03 seconds.\n", + "Found 143 source islands in the image with area limit 5.\n", + "Processing 143 source islands sequentially.\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "143it [00:08, 16.56it/s]" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Homology computation took 8.64 seconds.\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "\n" + ] + } + ], "source": [ + "# import package from a local directory\n", + "import sys\n", + "sys.path.append('../') # path to the DRUID package\n", "from DRUID import sf\n", "\n", - "findmysources = sf(image=image, # Image Array\n", - " header=cutout_file[0][1].header, # Image Header\n", - " mode='optical', # Mode (optical, radio, other)\n", - " cutup=False, # Cutup the image into smaller pieces?\n", - " smooth_sigma=1.1, # Sigma for the smoothing the image with gaussian filter\n", - " area_limit=3) # Minimum area for a source to be considered as a source\n", - "findmysources.set_background(detection_threshold=2, # Detection threshold for the background\n", - " analysis_threshold=2, # Analysis threshold for the background\n", - " mode='SEX') # Mode (SEX, rms, mad_std, sigma_clip) \n", - "findmysources.phsf(lifetime_limit=findmysources.local_bg, # Lifetime limit for the sources (in absolute units)\n", - " lifetime_limit_fraction=1.001) # Lifetime limit for the sources (as fraction of brith - death)\n", - "findmysources.source_characterising() # Characterise the sources (i.e. measure the properties of the sources and calculate the contours)\n", - "catalogue = findmysources.catalogue # get the catalogue from the sf object\n" + "findmysources = sf(image=image, \n", + " header=cutout_file[0][1].header, \n", + " mode='optical', \n", + " area_limit=5,\n", + " num_threads=1) \n", + "findmysources.set_background(detection_threshold=5) \n", + "findmysources.phsf(lifetime_limit_fraction=1) \n", + "catalogue = findmysources.catalog " + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "shape: (2_255, 20)\n", + "┌────────────┬────────────┬─────┬─────┬───┬────────────┬─────────────────────┬──────────┬──────────┐\n", + "│ birth ┆ death ┆ x1 ┆ y1 ┆ … ┆ parent_tag ┆ contour ┆ Island_X ┆ Island_Y │\n", + "│ --- ┆ --- ┆ --- ┆ --- ┆ ┆ --- ┆ --- ┆ --- ┆ --- │\n", + "│ f64 ┆ f64 ┆ f64 ┆ f64 ┆ ┆ i64 ┆ list[list[f64]] ┆ i32 ┆ i32 │\n", + "╞════════════╪════════════╪═════╪═════╪═══╪════════════╪═════════════════════╪══════════╪══════════╡\n", + "│ 8.7076e-12 ┆ 6.4912e-12 ┆ 1.0 ┆ 2.0 ┆ … ┆ 1 ┆ [[2.0, 2.0], [1.0, ┆ 3 ┆ 996 │\n", + "│ ┆ ┆ ┆ ┆ ┆ ┆ 1.0], … [2.… ┆ ┆ │\n", + "│ 9.3193e-12 ┆ 0.0 ┆ 2.0 ┆ 3.0 ┆ … ┆ 1 ┆ [[5.0, 2.0], [5.0, ┆ 3 ┆ 996 │\n", + "│ ┆ ┆ ┆ ┆ ┆ ┆ 1.0], … [5.… ┆ ┆ │\n", + "│ 9.3193e-12 ┆ 6.4912e-12 ┆ 2.0 ┆ 3.0 ┆ … ┆ 2 ┆ [[5.0, 1.0], [4.0, ┆ 3 ┆ 996 │\n", + "│ ┆ ┆ ┆ ┆ ┆ ┆ 0.0], … [5.… ┆ ┆ │\n", + "│ 2.3841e-11 ┆ 0.0 ┆ 4.0 ┆ 3.0 ┆ … ┆ 0 ┆ [[9.0, 3.0], [8.0, ┆ 23 ┆ 436 │\n", + "│ ┆ ┆ ┆ ┆ ┆ ┆ 2.0], … [9.… ┆ ┆ │\n", + "│ 8.2292e-12 ┆ 6.7574e-12 ┆ 7.0 ┆ 1.0 ┆ … ┆ 1 ┆ [[8.0, 1.0], [7.0, ┆ 37 ┆ 891 │\n", + "│ ┆ ┆ ┆ ┆ ┆ ┆ 0.0], … [8.… ┆ ┆ │\n", + "│ … ┆ … ┆ … ┆ … ┆ … ┆ … ┆ … ┆ … ┆ … │\n", + "│ 7.0367e-12 ┆ 4.6814e-12 ┆ 1.0 ┆ 5.0 ┆ … ┆ 1 ┆ [[2.0, 5.0], [2.0, ┆ 962 ┆ 448 │\n", + "│ ┆ ┆ ┆ ┆ ┆ ┆ 4.0], … [2.… ┆ ┆ │\n", + "│ 1.5760e-11 ┆ 0.0 ┆ 7.0 ┆ 6.0 ┆ … ┆ 2 ┆ [[13.0, 5.0], ┆ 962 ┆ 448 │\n", + "│ ┆ ┆ ┆ ┆ ┆ ┆ [13.0, 4.0], … [… ┆ ┆ │\n", + "│ 1.5760e-11 ┆ 1.2681e-11 ┆ 7.0 ┆ 6.0 ┆ … ┆ 3 ┆ [[9.0, 7.0], [9.0, ┆ 962 ┆ 448 │\n", + "│ ┆ ┆ ┆ ┆ ┆ ┆ 6.0], … [9.… ┆ ┆ │\n", + "│ 1.0088e-11 ┆ 0.0 ┆ 3.0 ┆ 4.0 ┆ … ┆ 0 ┆ [[6.0, 3.0], [6.0, ┆ 971 ┆ 408 │\n", + "│ ┆ ┆ ┆ ┆ ┆ ┆ 2.0], … [6.… ┆ ┆ │\n", + "│ 1.4803e-11 ┆ 0.0 ┆ 3.0 ┆ 2.0 ┆ … ┆ 0 ┆ [[5.0, 4.0], [5.0, ┆ 978 ┆ 823 │\n", + "│ ┆ ┆ ┆ ┆ ┆ ┆ 3.0], … [5.… ┆ ┆ │\n", + "└────────────┴────────────┴─────┴─────┴───┴────────────┴─────────────────────┴──────────┴──────────┘\n" + ] + } + ], + "source": [ + "print(catalogue)" ] }, { @@ -79,16 +269,30 @@ "metadata": {}, "outputs": [], "source": [ + "import numpy as np\n", "from matplotlib.pylab import plt\n", "\n", "plt.figure(figsize=(10,10))\n", "plt.imshow(image, cmap='gray', vmin=0,vmax=1E-9)\n", "for con in catalogue['contour']:\n", - " plt.plot(con[:, 1], con[:, 0])\n", + " contours = catalogue[\"contour\"].to_list()\n", + " Island_X = catalogue[\"Island_X\"].to_list()\n", + " Island_Y = catalogue[\"Island_Y\"].to_list()\n", + " for i, contour in enumerate(contours):\n", + " contour = np.array(contour)\n", + " Island_X_val = Island_X[i]\n", + " Island_Y_val = Island_Y[i]\n", + " plt.plot(\n", + " contour[:, 1] + Island_Y_val,\n", + " contour[:, 0] + Island_X_val,\n", + " # color=\"red\",\n", + " alpha=1,\n", + " linewidth=1,\n", + " )\n", "plt.xlim(300,700)\n", "plt.ylim(300,700)\n", "plt.savefig('images/sample_gal.png')\n", - "#plt.show()" + "plt.show()" ] }, { @@ -115,7 +319,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.9.13" + "version": "3.12.10" } }, "nbformat": 4, diff --git a/README.md b/README.md index b3b6ea4..c20f6e7 100644 --- a/README.md +++ b/README.md @@ -2,77 +2,145 @@ [![Run tests](https://github.com/RhysAlfShaw/DRUID/actions/workflows/pytest.yaml/badge.svg)](https://github.com/RhysAlfShaw/DRUID/actions/workflows/pytest.yaml) [![codecov](https://codecov.io/gh/RhysAlfShaw/DRUID/graph/badge.svg?token=C4KD4C6IXA)](https://codecov.io/gh/RhysAlfShaw/DRUID) -DRUID is a general purpose source finder for optical and radio images written in `python`. That can be appied in a broad range of situations. It comes some GPU acceleration of key compute functions, making processing of large sources faster. +DRUID is a general-purpose source finder for optical and radio images written in Python, applicable to a broad range of scenarios. -DRUID relies on the use of Persistent homology to find sources and nested components from within the image. This information is then processed as described in Shaw et al (in prep). +DRUID relies on the use of persistent homology to find sources and nested components within an image. This information is then processed as described in Shaw et al. (2025). -Currently DRUID uses the `criper` (https://github.com/shizuo-kaji/CubicalRipser_3dim) library to calculate the persistence of homology groups within the 2d data. +Currently, DRUID uses the [`cripser`](https://github.com/shizuo-kaji/CubicalRipser_3dim) library to calculate the persistence of homology groups within 2D data. -# Installation +## Versions -To use DRUID currently the best way is to clone this repository and install it with its dependencies. +This is the newly parallelized version of DRUID, featuring improved background data handling. The version of DRUID used in Shaw et al. (2025) can be found in the [releases](). + +### Notes on the new version +- DRUID's architecture has been significantly updated to allow for highly efficient parallelization. +- Pandas DataFrames have been replaced with Polars, drastically improving compute speed and memory management. +- GPU usage has been removed, as it was incompatible with the new parallel strategy. + +These changes have increased DRUID's speed by roughly 5-60x. This improvement stems from the new processing architecture and the switch to Polars, with further gains driven by parallelization. To get an idea of its new performance, check out the scaling plot below. This benchmark reflects the processing of an optical image with 0.1" resolution over a 0.57 deg² field of view, containing around 100,000 sources. + +![DRUID Performance Scaling](./docs/assets/druid_performance_scaling.png) + +## Installation + +Currently, the best way to use DRUID is to clone this repository and install it along with its dependencies: + +```bash +git clone https://github.com/RhysAlfShaw/DRUID.git +cd DRUID +``` + +### Conda + +Create conda environement: +```bash +conda env create -f environment.yml +``` ```bash pip install . ``` -You can then test if it is working with. +You can then verify the installation by running: ```python from DRUID import sf ``` -## Note for Apple Silicon +### UV -For Apple Silcon Users! cripser does not provide compiled binaries for apple silicon. So you need to compile the library locally. This should be simple with the command: +For a faster install with a single command using uv, simply. ```bash -pip install -U git+https://github.com/shizuo-kaji/CubicalRipser_3dim +uv sync --python 3.12 ``` -Any installation error from here will likely be from the the version of CMAKE or the C compilers you have installed. See https://github.com/shizuo-kaji/CubicalRipser_3dim for further details on required compilers. +uv will automatically detect the requirements and install DRUID. Test as above or with -## Using the GPU functionality +```bash +uv run python -c "from DRUID import sf" +``` -To use the GPU functions that DRUID offers you need to install `cupy` [https://cupy.dev/]. -This is not done in the normal requirements as it requires access to a Nvidia GPU with Cuda. And hence requires cuda tool kit to be installed on you machine. +No errors indicates a successful install. -If you have sucessfully install cupy then you can use `GPU=True`. -# Using DRUID +### Note for Apple Silicon Users +`cripser` does not provide compiled binaries for Apple Silicon, so you will need to compile the library locally. This can typically be done with the following command: -To use DRUID you need to do the following steps. +```bash +pip install -U git+[https://github.com/shizuo-kaji/CubicalRipser_3dim](https://github.com/shizuo-kaji/CubicalRipser_3dim) +``` + +Any installation errors at this stage will likely stem from the version of CMake or the C compilers you have installed. See the [CubicalRipser_3dim repository](https://github.com/shizuo-kaji/CubicalRipser_3dim) for further details on required compilers. + +## Using DRUID + +To run DRUID, follow these steps: -1. Initailise the sf (source finding) object. +1. **Initialize the `sf` (source finding) object:** ```python -findmysource = sf(image=image,image_path=None, mode='optical',area_limit=5,GPU=True, header=header) +findmysource = sf( + image=image, # image, either a 2d np.array, or path to fits file. + mode="optical", + area_limit=5, # Helps remove noise sources. + smooth_sigma=1, # smooth image before ph analysis, fluxes measured on original image. + num_threads=2, # number of threads, as num_threads increases speed gains decrease. + chunksize=20, # chuncking size for multithreading, only provides minor speedup. + max_area_limit = 1E5, # there are size limits on ph analysis this prevent unintentional infinate compute time. + working_directory=".", # where to save outputs and cache results. + cache=False, # Cache/save results to working directory. + ) ``` -2. Define the background. + +2. **Define the background:** ```python -findmysource.set_background(detection_threshold=5,analysis_threshold=2,mode='rms') +findmysource.set_background( + method='mad_std', # background statistic (sex,rms,mad_std...) + detection_threshold=5, # how many sigmas above the background should we call a source. + analysis_threshold=3, # ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ we analyse for a source. + box_size=10 # boz size for creating a background map. + ) ``` -3. Find and Deblend sources with Persistent Homology. + +3. **Find and deblend sources using Persistent Homology:** ```python -findmysources.phsf() +findmysource.phsf(, + lifetime_limit = 0, # float value for this limit + lifetime_limit_fraction=1.2 # fraction based on birth and death. + ) ``` -4. Now we have a list of sources and a hierachy of nested components, we can charaterise them and measure some properties. -```pythons -findmysources.source_charaterising(use_gpu=False) -``` +This function also calculates source properties. -To explore how DRUID can be used Check out the example notebooks where we demonstraight some of DRUIDs functions. (This is incoming, will be based on analysis in Shaw et al in prep) -## Saving the catalogue. -To save the output catalogue with the contours you should use the ```save_catalogue()``` function. As this will properly save the object. To correctly open the catalogue again use ```open_catalogue()``` after initlising the sf class. +## Runnig in parallel. -# Bugs/issues +To prevent issues with pythons multiprocessing functionality. DRUID will not run unless wrapped in a __main__. If you run this inside a jupyter notebook __main__ is not necessary. + +```python +from DRUID import sf + +def main(): + findmysource = sf( + image=image, + mode="optical", + area_limit=5, + num_threads=2, + ) + findmysource.set_background() + findmysource.phsf() + +if __name__ == "__main__": + main() + +``` -Please report any bug or issues using DRUID to this repositories issue page. Thank you. +## Bugs & Issues -# Further application/developement +Please report any bugs or issues you encounter while using DRUID on this repository's [Issues](#) page. Thank you! -If you want to increase the functionality, whether thats adding additional functionality to improving what is already implemented, feel free to submit a pull request or email me (rhys.shaw@bristol.ac.uk) to discuss. +or email me at [rhys.shaw@bristol.ac.uk](mailto:rhys.shaw@bristol.ac.uk). -# Acknowledgements +## Acknowledgements -If you use DRUID for your research please cite: ([Shaw et al.2025](https://doi.org/10.1093/rasti/rzaf006)) \ No newline at end of file +If you use DRUID for your research, please cite: +> [Shaw et al. 2025](https://doi.org/10.1093/rasti/rzaf006) \ No newline at end of file diff --git a/docs/assets/druid_performance_scaling.png b/docs/assets/druid_performance_scaling.png new file mode 100644 index 0000000..58affac Binary files /dev/null and b/docs/assets/druid_performance_scaling.png differ diff --git a/environment.yml b/environment.yml new file mode 100644 index 0000000..bb5be2e --- /dev/null +++ b/environment.yml @@ -0,0 +1,21 @@ +name: DRUID +channels: + - conda-forge + - defaults +dependencies: + - python=3.12 + - pip + - pip: + - numpy + - polars + - pytest + - numpy + - scikit-image + - astropy + - tqdm + - matplotlib + - setproctitle + - scipy + - photutils + - cripser + - bottleneck \ No newline at end of file diff --git a/large_image_test.py b/large_image_test.py new file mode 100644 index 0000000..2caea5b --- /dev/null +++ b/large_image_test.py @@ -0,0 +1,25 @@ +from DRUID import sf + + +def main(): + path = "/Users/rs17612/Documents/Optical_IR_Data/EUCLID/EUC_MER_BGSUB-MOSAIC-DES-Z_TILE102026098-321DAA_20240407T183900.016572Z_00.00.fits" + working_dir = "DRUID/temp_new" + findmysource = sf( + image=path, + mode="optical", + area_limit=15, + num_threads=2, + working_directory=working_dir, + cashe=False, + ) + findmysource.set_background( + detection_threshold=5, + analysis_threshold=3, + box_size=50, + ) + findmysource.phsf(lifetime_limit_fraction=1.4) + print(findmysource.catalog) + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..cb4bc23 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,47 @@ +[build-system] +requires = ["setuptools>=61.0.0", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "DRUID" +version = "1.0" +description = "Astronomical Source finder build with Persistent Homology." +readme = "README.md" +requires-python = ">=3.12" +authors = [ + { name = "Rhys Shaw", email = "rhys.shaw@bristol.ac.uk" } +] +classifiers = [ + "Intended Audience :: Science/Research", + "Programming Language :: Python :: 3.12", + "Topic :: Scientific/Engineering :: Astronomy", + "Topic :: Scientific/Engineering :: Physics", +] + +dependencies = [ + "numpy", + "polars", + "scipy", + "astropy", + "scikit-image", + "photutils", + "cripser", + "bottleneck", + "matplotlib", + "tqdm", + "setproctitle", + "rich>=15.0.0", +] + +[project.urls] +Homepage = "https://github.com/RhysAlfShaw/DRUID" +Repository = "https://github.com/RhysAlfShaw/DRUID.git" + +[project.optional-dependencies] +test = [ + "pytest>=7.0", + "pytest-cov" +] + +[tool.setuptools.packages.find] +include = ["DRUID*"] diff --git a/requirements.txt b/requirements.txt index b36f66b..951f4d1 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,12 +1,11 @@ numpy -pandas -pytest -numpy -scikit-image -astropy -tqdm -matplotlib -setproctitle +polars scipy +astropy +scikit-image photutils cripser +bottleneck +matplotlib +setproctitle +rich>=15.0.0 \ No newline at end of file diff --git a/setup.py b/setup.py deleted file mode 100644 index 8facd69..0000000 --- a/setup.py +++ /dev/null @@ -1,15 +0,0 @@ -from setuptools import setup, find_packages - -with open('requirements.txt') as f: - requiremets = f.read().splitlines() - -setup( - name='DRUID', - version='0.0.0', - author='Rhys Shaw', - author_email='rhys.shaw@bristol.ac.uk', - url='https://github.com/RhysAlfShaw/DRUID', - description='descriptions', - install_requires=requiremets, - packages=find_packages(), -) \ No newline at end of file diff --git a/test.py b/test.py new file mode 100644 index 0000000..9a6ead5 --- /dev/null +++ b/test.py @@ -0,0 +1,97 @@ +from DRUID import sf +import matplotlib.pyplot as plt +import numpy as np + + +def main(): + working_dir = "DRUID/temp" + image_paths = [ + "/Users/rs17612/Documents/Radio_Data/3CRR/3C401", + "/Users/rs17612/Documents/Radio_Data/3CRR/3C295", + "/Users/rs17612/Documents/Radio_Data/3CRR/3C438", + "/Users/rs17612/Documents/Radio_Data/3CRR/3C452", + "/Users/rs17612/Documents/Radio_Data/3CRR/3C76P1", + ] + + n_images = len(image_paths) + fig, axes = plt.subplots(1, n_images, figsize=(5 * n_images, 5)) + i = 1 + if n_images == 1: + axes = [axes] + + for ax, image_path in zip(axes, image_paths): + findmysource = sf( + image=image_path, + mode="radio", + area_limit=15, + num_threads=2, + working_directory=working_dir, + cashe=False, + ) + findmysource.set_background( + detection_threshold=5, + analysis_threshold=3, + box_size=50, + ) + findmysource.phsf(lifetime_limit_fraction=1.4) + + catalog = findmysource.catalog + # save polars catalog to temp folder + if i == 1: + catalog.write_parquet( + f"{working_dir}/catalog_{image_path.split('/')[-1]}.parquet" + ) + # save image as numpy array + np.save( + f"{working_dir}/image_{image_path.split('/')[-1]}.npy", + findmysource.image, + ) + # save background and background rms as numpy array + np.save( + f"{working_dir}/background_{image_path.split('/')[-1]}.npy", + findmysource.background_map, + ) + np.save( + f"{working_dir}/background_rms_{image_path.split('/')[-1]}.npy", + findmysource.background_rms_map, + ) + i += 1 + + print(catalog) + ax.imshow(findmysource.image, cmap="gray", origin="lower") + # Corrected Scatter Plot + ax.scatter( + catalog["y1"] + catalog["Island_X"], # X + X + catalog["x1"] + catalog["Island_Y"], # Y + Y + s=1, + c="red", + label="Source Islands", + ) + + contours = catalog["contour"].to_list() + Island_X = catalog["Island_X"].to_list() + Island_Y = catalog["Island_Y"].to_list() + + # Corrected Contour Plot + for i, contour in enumerate(contours): + contour = np.array(contour) + Island_X_val = Island_X[i] + Island_Y_val = Island_Y[i] + ax.plot( + contour[:, 1] + Island_X_val, # X + X + contour[:, 0] + Island_Y_val, # Y + Y + alpha=1, + linewidth=1, + ) + ax.set_title(f"{image_path.split('/')[-1]}") + ax.set_xlabel("X Pixel") + ax.set_ylabel("Y Pixel") + ax.legend() + + plt.tight_layout() + plt.savefig(f"{working_dir}/all_source_islands_on_images.png") + # plt.show() + + +if __name__ == "__main__": + main()