From 3efff5918d51eeaeced026b9884e5aac7caef63f Mon Sep 17 00:00:00 2001 From: IainHammond Date: Thu, 5 Mar 2026 14:30:21 +0100 Subject: [PATCH 01/46] adding noise is now faster --- heeps/contrast/background.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/heeps/contrast/background.py b/heeps/contrast/background.py index d48b56e..1dda6bf 100644 --- a/heeps/contrast/background.py +++ b/heeps/contrast/background.py @@ -76,8 +76,16 @@ def background(psf_ON, psf_OFF, header=None, mode='RAVC', lam=3.8e-6, dit=0.3, bckg_noise = dit * flux_bckg * thruput * mask_trans psf_ON += bckg_noise # add photon noise ~ N(0, sqrt(psf)) - np.random.seed(seed) - psf_ON += np.random.normal(0, np.sqrt(psf_ON)) + # np.random.seed(seed) + # psf_ON += np.random.normal(0, np.sqrt(psf_ON)) + # Using default_rng with standard_normal is much faster than + # np.random.normal(0, sigma_array) for large cubes, because + # standard_normal draws from N(0,1) using a fast vectorised path + # and we then scale by sqrt(psf_ON) in a single multiply pass, + # avoiding per-element sigma sampling. + rng = np.random.default_rng(seed) + psf_sqrt = np.sqrt(psf_ON) + psf_ON += psf_sqrt * rng.standard_normal(psf_ON.shape) if verbose is True: print(' dit=%s s, thruput=%.4f, mask_trans=%.4f,'%(dit, thruput, mask_trans)) From 1e239b29f56f8e3783dd4edc4ed6c6728b943cb5 Mon Sep 17 00:00:00 2001 From: IainHammond Date: Fri, 15 May 2026 14:32:27 +0200 Subject: [PATCH 02/46] first interation of disc modelling --- heeps/contrast/disc.py | 307 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 307 insertions(+) create mode 100644 heeps/contrast/disc.py diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py new file mode 100644 index 0000000..f38c995 --- /dev/null +++ b/heeps/contrast/disc.py @@ -0,0 +1,307 @@ +""" +HEEPS disc injection and post‑processing utilities. + +This module provides functions to inject a synthetic disc model into HEEPS +PSF data and to perform post‑processing with VIP. +The entry point is :func:`create_disc_sequence`, which creates a mock +observation sequence with a user-provided disc mode. The :func:`post_proc_disc` function is a +wrapper of the relevevant VIP function for PCA post-processing. +""" + +import numpy as np +from typing import Union +import os + +from vip_hci.fits import open_fits, open_header +from vip_hci.preproc import frame_crop, cube_subsample, cube_crop_frames +from vip_hci.fm import cube_inject_fakedisk +from vip_hci.psfsub import pca + +from heeps.contrast.background import background +from heeps.util.psf_template import psf_template +from heeps.util.paralang import paralang + +__all__ = ['create_disc_sequence', 'postproc_disc'] +__author__ = "Iain Hammond" + +def create_disc_sequence( + disc_model : np.ndarray, + on_axis_psf : Union[np.ndarray, None] = None, + off_axis_psf : Union[np.ndarray, None] = None, + transmission : Union[np.ndarray, None] = None, + rdi : bool = True, + rdi_mag : Union[int, float, None] = None, + rdi_duration : Union[int, float, None] = None, + source_xy : Union[tuple, list, None] = None, + extinction : Union[int, float] = 0, + tag : Union[str, None] = None, + starphot : float = 1e11, + imlib : str = "opencv", + **conf +): + """ + Create a mock observation sequence by injecting a synthetic disc model into a HEEPS + PSF and optionally generate a reference differential imaging (RDI) + cube. Existing on- and off-axis PSF cubes can be provided. + + Parameters + ---------- + disc_model : np.ndarray + 2‑D array representing the disc model to be injected. NaNs are replaced with zeros. + on_axis_psf : Union[np.ndarray, None], optional + Pre‑loaded on‑axis PSF cube. If ``None``, the PSF is loaded from the output directory + using the conf parameters. + off_axis_psf : Union[np.ndarray, None], optional + Pre‑loaded off‑axis PSF frame. If ``None``, the PSF is loaded from the output directory. + transmission : Union[np.ndarray, None], optional + Coronagraph transmission. If ``None``, the appropriate transmission is loaded based + on the instrument mode in conf. Uses VIP conventions. + rdi : bool, default=True + Whether to generate an RDI reference cube from a segment of the on‑axis sequence. + rdi_mag : Union[int, float, None], optional + Magnitude to use for the RDI reference. If ``None``, the magnitude of the science target + is used. + rdi_duration : Union[int, float, None], optional + Desired duration (in seconds) of the RDI reference sequence. If ``None``, a default of + 20 % of the on‑axis sequence length is used. + source_xy : Union[tuple, list, None], optional + (x, y) pixel coordinates of the source for which extinction is applied. + extinction : Union[int, float], default=0 + Extinction in magnitudes to apply to the source at ``source_xy``. + tag : Union[str, None], optional + Tag to prepend to output filenames from earlier HEEPS runs. + starphot : float, default=1e11 + Photometric scaling factor for the star. + imlib : str, default="opencv" + Image library to use for image processing. + **conf : dict + Additional conf parameters from HEEPS (e.g., background addition etc.) + + Returns + ------- + psf_ON : numpy.ndarray + Sequence with the fake disc injected into the on‑axis PSF cube. + pa : numpy.ndarray + Array of parallactic angles used for the injection. + psf_RDI : numpy.ndarray, optional + Reference cube for RDI (if ``rdi`` is ``True``). Returned as the third element when + ``rdi`` is enabled; otherwise only ``psf_ON`` and ``pa`` are returned. + """ + tag = "" if tag is None else "%s_"%tag + loadname = os.path.join(conf['dir_output'], '%s%s_PSF_%s_%s.fits'%(tag,'%s', conf['band'], conf['mode'])) + if off_axis_psf is None: + psf_OFF = open_fits(loadname % 'offaxis', verbose=False) + print(f"Using off-axis PSF from {loadname%'offaxis'}") + assert psf_OFF.ndim == 2, "off-axis PSF frame must be 2-dimensional" + + # set nans to zero + disc_model = np.nan_to_num(disc_model, nan=0) + + # Remove any extra dimensions (e.g., MCFOST cubes) + # first, squeeze singleton dimensions + disc_model = np.squeeze(disc_model) + # then collapse any remaining leading dimensions to a 2-D image + while disc_model.ndim > 2: + disc_model = disc_model[0] + + # ensure the model has an odd size + if disc_model.shape[-1] % 2 == 0: + disc_model = disc_model[1:, 1:] # we assume the star flux is on one pixel + + # crop everything to a common size + # the HEEPS PSFs are usually 293px + min_crop = min(psf_OFF.shape[-1], disc_model.shape[-1], conf["ndet"]) + if psf_OFF.shape[-1] > min_crop: + psf_OFF = frame_crop(psf_OFF, min_crop, verbose=False) + if disc_model.shape[-1] > min_crop: + disc_model = frame_crop(disc_model, min_crop, verbose=False) + + # record the stellar flux in the model and set the pixel to 0 + star_val = np.max(disc_model) + star_y, star_x = np.unravel_index(np.argmax(disc_model), disc_model.shape) + disc_model[star_y, star_x] = 0 + + # apply extinction if requested + if extinction > 0 and source_xy is not None: + source_x, source_y = source_xy + extinction_factor = 10 ** (-0.4 * extinction) + disc_model[source_y, source_x] *= extinction_factor + + # load coronagraph transmission + if transmission is None: + if conf["mode"] == "ELT" or conf["mode"] is None: # no coronagraph + transmission = None + elif conf["mode"] == "CVC": + transmission = open_fits(conf["f_vc_trans"], verbose=False) + print(f"Using transmission from {conf['f_vc_trans']}") + else: # TODO + raise NotImplementedError(f"Mode {conf['mode']} not implemented for disc injection.") + + # determine how many frames are in the on-axis sequence (don't open the whole cube yet to save memory) + if on_axis_psf is None: + nframes = open_header(loadname % 'onaxis')['NAXIS3'] + else: + nframes = on_axis_psf.shape[0] + + # generate parallactic angles and inject the fake disc using VIP + pa = paralang(npts=nframes, dec=conf["dec"], lat=conf["lat"], duration=int(nframes * conf["dit"])) + + cube = cube_inject_fakedisk( + disc_model, + pa, + transmission=transmission, + psf=psf_OFF, + normalize_psf=False, + nproc=conf["cpu_count"], + mask_val=0, + edge_blend="interp", + interp_zeros=True, + ker=1, + imlib=imlib, + ) + + # normalise by the stellar flux that we removed earlier + cube /= star_val + + # open and crop sequence if needed + if on_axis_psf is None: + psf_ON = open_fits(loadname % 'onaxis', verbose=False) + print(f"Using on-axis PSF from {loadname % 'onaxis'}") + + assert psf_ON.ndim == 3, "on-axis PSF cube must be 3-dimensional" + if psf_ON.shape[-1] > min_crop: + psf_ON = cube_crop_frames(psf_ON, min_crop, verbose=False) + + # add the disc model + psf_ON += cube + + # background addition + if conf["add_bckg"]: + psf_ON, psf_OFF = background(psf_ON, psf_OFF, verbose=True, **conf) + + _, _, ap_flux = psf_template(psf_OFF) + psf_ON *= starphot / ap_flux + + # RDI handling (at the moment we only support taking a chunk out of the science sequence + if rdi: + if rdi_duration is None: + print("Warning: rdi_duration was not set. Using 20% of on-axis sequence.", flush=True) + rdi_duration = int(0.2 * nframes) * conf["dit"] + n_ref_frames = rdi_duration / conf["dit"] + start_idx = np.random.randint(0, psf_ON.shape[0] - int(n_ref_frames) + 1) # random segment of the on-axis PSF cube to use as the RDI reference, to better match the background conditions + psf_RDI = open_fits(loadname % "onaxis", verbose=False)[start_idx:start_idx + int(n_ref_frames)] + if psf_RDI.shape[-1] > min_crop: + psf_RDI = cube_crop_frames(psf_RDI, min_crop, verbose=False) + + # remove the RDI segment from psf_ON to correctly represent lost integration time on the science target + psf_ON = np.concatenate([psf_ON[:start_idx], psf_ON[start_idx + int(n_ref_frames):]], axis=0) + pa = np.concatenate([pa[:start_idx], pa[start_idx + int(n_ref_frames):]], axis=0) + + psf_RDI_OFF = open_fits(loadname % "offaxis", verbose=False) + if psf_RDI_OFF.shape[-1] > min_crop: + psf_RDI_OFF = frame_crop(psf_RDI_OFF, min_crop, verbose=False) + + if rdi_mag is None: + print("Warning: rdi_mag was not set. Using mag of science target.", flush=True) + else: + conf["mag"] = rdi_mag + + if conf["add_bckg"]: + psf_RDI, psf_RDI_OFF = background(psf_RDI, psf_RDI_OFF, verbose=True, **conf) + + _, _, ap_flux = psf_template(psf_RDI_OFF) + psf_RDI *= starphot / ap_flux + del psf_RDI_OFF + + return psf_ON, pa, psf_RDI + + else: + return psf_ON, pa + + +def postproc_disc( + cube : np.ndarray, + angle_list : np.ndarray, + cube_ref : Union[np.ndarray, None] = None, + subsample : int = 1, + ncomp: Union[tuple, list, int] = 1, + mask_center_px : int = 0, + imlib : str = "vip-fft", + source_xy : Union[tuple, None] = None, + delta_rot : Union[float, int, None] = 1, + fwhm : Union[float, int] = 5, + mask_rdi: np.ndarray = None, + ref_strategy: str = 'RDI', + nproc : int = 1 +): + """ + Perform PCA-based post-processing on a disc‑injected data cube. + + Parameters + ---------- + cube : np.ndarray + 3‑D array (n_frames, ny, nx) containing the science frames. + angle_list : np.ndarray + 1‑D array of parallactic angles (in degrees) associated with each frame in ``cube``. + cube_ref : np.ndarray or None, optional + Reference cube for RDI (reference differential imaging). If ``None``, only ADI is performed. + subsample : int, default=1 + Factor by which to down‑sample the cube (and reference) for speed. ``1`` means no subsampling. + ncomp : int, tuple, or list, default=1 + Number of principal components to use. If an int, a single component is used. If a tuple ``(min, max)`` or a list, the function will compute results for each component in the range. + mask_center_px : int, default=0 + Radius (in pixels) of a circular mask applied to the centre of each frame before PCA. + imlib : str, default="vip-fft" + Image library used by VIP for FFT‑based operations. + source_xy : tuple or None, optional + (x, y) pixel coordinates of a known source; used to mask the source during PCA if provided. + delta_rot : float or int or None, default=1 + Minimum rotation (in FWHM) between frames for a given pixel to be considered independent. + fwhm : float or int, default=5 + Full‑width at half‑maximum of the PSF, used for delta_rot. + mask_rdi : np.ndarray, optional + Optional mask applied to the RDI reference cube. + ref_strategy : str, default='RDI' + Strategy for reference handling. + nproc : int, default=1 + Number of processes to use for parallel computation. + + Returns + ------- + np.ndarray + Array of shape ``(n_ncomp, ny, nx)`` where ``n_ncomp`` is the number of principal components evaluated. Each slice ``res[i]`` contains the PCA‑processed image for the corresponding number of components. + + Notes + ----- + The function currently supports the subset of ``vip_hci.psfsub.pca`` parameters used in the HEEPS disc injection workflow. Additional ``pca`` arguments can be added in the future. + """ + # subsample the cube and the reference if requested, for efficiency purposes + if subsample > 1: + cube, angle_list = cube_subsample(array=cube, n=subsample, parallactic=angle_list) + if cube_ref is not None: + cube_ref = cube_subsample(array=cube_ref, n=subsample) + + # loop over principal components + if isinstance(ncomp, int): + ncomp = [ncomp] + elif isinstance(ncomp, tuple): + ncomp = np.arange(ncomp[0], ncomp[-1]+1) + + res = np.zeros([len(ncomp), cube.shape[-2], cube.shape[-1]]) + + print("Running PCA post-processing", flush=True) + for i, npc in enumerate(ncomp): + res[i] = pca(cube, angle_list=angle_list, cube_ref=cube_ref, ncomp=npc, + mask_center_px=mask_center_px, imlib=imlib, nproc=nproc, + source_xy=source_xy, delta_rot=delta_rot, fwhm=fwhm, mask_rdi=mask_rdi, + ref_strategy=ref_strategy) + return res + + +# script behaviour? i dont like these +# if __name__ == '__main__': +# # python -m HEEPS.heeps.contrast.disc \ +# # --model_path /path/to/model.fits \ +# # --psf_on_path /path/to/onaxis_psf.fits \ +# # --psf_off_path /path/to/offaxis_psf.fits +# pass From aa126047bcac2f330ed97621cf847157eca24113 Mon Sep 17 00:00:00 2001 From: IainHammond Date: Fri, 15 May 2026 14:32:27 +0200 Subject: [PATCH 03/46] first iteration of disc modelling --- heeps/contrast/disc.py | 307 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 307 insertions(+) create mode 100644 heeps/contrast/disc.py diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py new file mode 100644 index 0000000..f38c995 --- /dev/null +++ b/heeps/contrast/disc.py @@ -0,0 +1,307 @@ +""" +HEEPS disc injection and post‑processing utilities. + +This module provides functions to inject a synthetic disc model into HEEPS +PSF data and to perform post‑processing with VIP. +The entry point is :func:`create_disc_sequence`, which creates a mock +observation sequence with a user-provided disc mode. The :func:`post_proc_disc` function is a +wrapper of the relevevant VIP function for PCA post-processing. +""" + +import numpy as np +from typing import Union +import os + +from vip_hci.fits import open_fits, open_header +from vip_hci.preproc import frame_crop, cube_subsample, cube_crop_frames +from vip_hci.fm import cube_inject_fakedisk +from vip_hci.psfsub import pca + +from heeps.contrast.background import background +from heeps.util.psf_template import psf_template +from heeps.util.paralang import paralang + +__all__ = ['create_disc_sequence', 'postproc_disc'] +__author__ = "Iain Hammond" + +def create_disc_sequence( + disc_model : np.ndarray, + on_axis_psf : Union[np.ndarray, None] = None, + off_axis_psf : Union[np.ndarray, None] = None, + transmission : Union[np.ndarray, None] = None, + rdi : bool = True, + rdi_mag : Union[int, float, None] = None, + rdi_duration : Union[int, float, None] = None, + source_xy : Union[tuple, list, None] = None, + extinction : Union[int, float] = 0, + tag : Union[str, None] = None, + starphot : float = 1e11, + imlib : str = "opencv", + **conf +): + """ + Create a mock observation sequence by injecting a synthetic disc model into a HEEPS + PSF and optionally generate a reference differential imaging (RDI) + cube. Existing on- and off-axis PSF cubes can be provided. + + Parameters + ---------- + disc_model : np.ndarray + 2‑D array representing the disc model to be injected. NaNs are replaced with zeros. + on_axis_psf : Union[np.ndarray, None], optional + Pre‑loaded on‑axis PSF cube. If ``None``, the PSF is loaded from the output directory + using the conf parameters. + off_axis_psf : Union[np.ndarray, None], optional + Pre‑loaded off‑axis PSF frame. If ``None``, the PSF is loaded from the output directory. + transmission : Union[np.ndarray, None], optional + Coronagraph transmission. If ``None``, the appropriate transmission is loaded based + on the instrument mode in conf. Uses VIP conventions. + rdi : bool, default=True + Whether to generate an RDI reference cube from a segment of the on‑axis sequence. + rdi_mag : Union[int, float, None], optional + Magnitude to use for the RDI reference. If ``None``, the magnitude of the science target + is used. + rdi_duration : Union[int, float, None], optional + Desired duration (in seconds) of the RDI reference sequence. If ``None``, a default of + 20 % of the on‑axis sequence length is used. + source_xy : Union[tuple, list, None], optional + (x, y) pixel coordinates of the source for which extinction is applied. + extinction : Union[int, float], default=0 + Extinction in magnitudes to apply to the source at ``source_xy``. + tag : Union[str, None], optional + Tag to prepend to output filenames from earlier HEEPS runs. + starphot : float, default=1e11 + Photometric scaling factor for the star. + imlib : str, default="opencv" + Image library to use for image processing. + **conf : dict + Additional conf parameters from HEEPS (e.g., background addition etc.) + + Returns + ------- + psf_ON : numpy.ndarray + Sequence with the fake disc injected into the on‑axis PSF cube. + pa : numpy.ndarray + Array of parallactic angles used for the injection. + psf_RDI : numpy.ndarray, optional + Reference cube for RDI (if ``rdi`` is ``True``). Returned as the third element when + ``rdi`` is enabled; otherwise only ``psf_ON`` and ``pa`` are returned. + """ + tag = "" if tag is None else "%s_"%tag + loadname = os.path.join(conf['dir_output'], '%s%s_PSF_%s_%s.fits'%(tag,'%s', conf['band'], conf['mode'])) + if off_axis_psf is None: + psf_OFF = open_fits(loadname % 'offaxis', verbose=False) + print(f"Using off-axis PSF from {loadname%'offaxis'}") + assert psf_OFF.ndim == 2, "off-axis PSF frame must be 2-dimensional" + + # set nans to zero + disc_model = np.nan_to_num(disc_model, nan=0) + + # Remove any extra dimensions (e.g., MCFOST cubes) + # first, squeeze singleton dimensions + disc_model = np.squeeze(disc_model) + # then collapse any remaining leading dimensions to a 2-D image + while disc_model.ndim > 2: + disc_model = disc_model[0] + + # ensure the model has an odd size + if disc_model.shape[-1] % 2 == 0: + disc_model = disc_model[1:, 1:] # we assume the star flux is on one pixel + + # crop everything to a common size + # the HEEPS PSFs are usually 293px + min_crop = min(psf_OFF.shape[-1], disc_model.shape[-1], conf["ndet"]) + if psf_OFF.shape[-1] > min_crop: + psf_OFF = frame_crop(psf_OFF, min_crop, verbose=False) + if disc_model.shape[-1] > min_crop: + disc_model = frame_crop(disc_model, min_crop, verbose=False) + + # record the stellar flux in the model and set the pixel to 0 + star_val = np.max(disc_model) + star_y, star_x = np.unravel_index(np.argmax(disc_model), disc_model.shape) + disc_model[star_y, star_x] = 0 + + # apply extinction if requested + if extinction > 0 and source_xy is not None: + source_x, source_y = source_xy + extinction_factor = 10 ** (-0.4 * extinction) + disc_model[source_y, source_x] *= extinction_factor + + # load coronagraph transmission + if transmission is None: + if conf["mode"] == "ELT" or conf["mode"] is None: # no coronagraph + transmission = None + elif conf["mode"] == "CVC": + transmission = open_fits(conf["f_vc_trans"], verbose=False) + print(f"Using transmission from {conf['f_vc_trans']}") + else: # TODO + raise NotImplementedError(f"Mode {conf['mode']} not implemented for disc injection.") + + # determine how many frames are in the on-axis sequence (don't open the whole cube yet to save memory) + if on_axis_psf is None: + nframes = open_header(loadname % 'onaxis')['NAXIS3'] + else: + nframes = on_axis_psf.shape[0] + + # generate parallactic angles and inject the fake disc using VIP + pa = paralang(npts=nframes, dec=conf["dec"], lat=conf["lat"], duration=int(nframes * conf["dit"])) + + cube = cube_inject_fakedisk( + disc_model, + pa, + transmission=transmission, + psf=psf_OFF, + normalize_psf=False, + nproc=conf["cpu_count"], + mask_val=0, + edge_blend="interp", + interp_zeros=True, + ker=1, + imlib=imlib, + ) + + # normalise by the stellar flux that we removed earlier + cube /= star_val + + # open and crop sequence if needed + if on_axis_psf is None: + psf_ON = open_fits(loadname % 'onaxis', verbose=False) + print(f"Using on-axis PSF from {loadname % 'onaxis'}") + + assert psf_ON.ndim == 3, "on-axis PSF cube must be 3-dimensional" + if psf_ON.shape[-1] > min_crop: + psf_ON = cube_crop_frames(psf_ON, min_crop, verbose=False) + + # add the disc model + psf_ON += cube + + # background addition + if conf["add_bckg"]: + psf_ON, psf_OFF = background(psf_ON, psf_OFF, verbose=True, **conf) + + _, _, ap_flux = psf_template(psf_OFF) + psf_ON *= starphot / ap_flux + + # RDI handling (at the moment we only support taking a chunk out of the science sequence + if rdi: + if rdi_duration is None: + print("Warning: rdi_duration was not set. Using 20% of on-axis sequence.", flush=True) + rdi_duration = int(0.2 * nframes) * conf["dit"] + n_ref_frames = rdi_duration / conf["dit"] + start_idx = np.random.randint(0, psf_ON.shape[0] - int(n_ref_frames) + 1) # random segment of the on-axis PSF cube to use as the RDI reference, to better match the background conditions + psf_RDI = open_fits(loadname % "onaxis", verbose=False)[start_idx:start_idx + int(n_ref_frames)] + if psf_RDI.shape[-1] > min_crop: + psf_RDI = cube_crop_frames(psf_RDI, min_crop, verbose=False) + + # remove the RDI segment from psf_ON to correctly represent lost integration time on the science target + psf_ON = np.concatenate([psf_ON[:start_idx], psf_ON[start_idx + int(n_ref_frames):]], axis=0) + pa = np.concatenate([pa[:start_idx], pa[start_idx + int(n_ref_frames):]], axis=0) + + psf_RDI_OFF = open_fits(loadname % "offaxis", verbose=False) + if psf_RDI_OFF.shape[-1] > min_crop: + psf_RDI_OFF = frame_crop(psf_RDI_OFF, min_crop, verbose=False) + + if rdi_mag is None: + print("Warning: rdi_mag was not set. Using mag of science target.", flush=True) + else: + conf["mag"] = rdi_mag + + if conf["add_bckg"]: + psf_RDI, psf_RDI_OFF = background(psf_RDI, psf_RDI_OFF, verbose=True, **conf) + + _, _, ap_flux = psf_template(psf_RDI_OFF) + psf_RDI *= starphot / ap_flux + del psf_RDI_OFF + + return psf_ON, pa, psf_RDI + + else: + return psf_ON, pa + + +def postproc_disc( + cube : np.ndarray, + angle_list : np.ndarray, + cube_ref : Union[np.ndarray, None] = None, + subsample : int = 1, + ncomp: Union[tuple, list, int] = 1, + mask_center_px : int = 0, + imlib : str = "vip-fft", + source_xy : Union[tuple, None] = None, + delta_rot : Union[float, int, None] = 1, + fwhm : Union[float, int] = 5, + mask_rdi: np.ndarray = None, + ref_strategy: str = 'RDI', + nproc : int = 1 +): + """ + Perform PCA-based post-processing on a disc‑injected data cube. + + Parameters + ---------- + cube : np.ndarray + 3‑D array (n_frames, ny, nx) containing the science frames. + angle_list : np.ndarray + 1‑D array of parallactic angles (in degrees) associated with each frame in ``cube``. + cube_ref : np.ndarray or None, optional + Reference cube for RDI (reference differential imaging). If ``None``, only ADI is performed. + subsample : int, default=1 + Factor by which to down‑sample the cube (and reference) for speed. ``1`` means no subsampling. + ncomp : int, tuple, or list, default=1 + Number of principal components to use. If an int, a single component is used. If a tuple ``(min, max)`` or a list, the function will compute results for each component in the range. + mask_center_px : int, default=0 + Radius (in pixels) of a circular mask applied to the centre of each frame before PCA. + imlib : str, default="vip-fft" + Image library used by VIP for FFT‑based operations. + source_xy : tuple or None, optional + (x, y) pixel coordinates of a known source; used to mask the source during PCA if provided. + delta_rot : float or int or None, default=1 + Minimum rotation (in FWHM) between frames for a given pixel to be considered independent. + fwhm : float or int, default=5 + Full‑width at half‑maximum of the PSF, used for delta_rot. + mask_rdi : np.ndarray, optional + Optional mask applied to the RDI reference cube. + ref_strategy : str, default='RDI' + Strategy for reference handling. + nproc : int, default=1 + Number of processes to use for parallel computation. + + Returns + ------- + np.ndarray + Array of shape ``(n_ncomp, ny, nx)`` where ``n_ncomp`` is the number of principal components evaluated. Each slice ``res[i]`` contains the PCA‑processed image for the corresponding number of components. + + Notes + ----- + The function currently supports the subset of ``vip_hci.psfsub.pca`` parameters used in the HEEPS disc injection workflow. Additional ``pca`` arguments can be added in the future. + """ + # subsample the cube and the reference if requested, for efficiency purposes + if subsample > 1: + cube, angle_list = cube_subsample(array=cube, n=subsample, parallactic=angle_list) + if cube_ref is not None: + cube_ref = cube_subsample(array=cube_ref, n=subsample) + + # loop over principal components + if isinstance(ncomp, int): + ncomp = [ncomp] + elif isinstance(ncomp, tuple): + ncomp = np.arange(ncomp[0], ncomp[-1]+1) + + res = np.zeros([len(ncomp), cube.shape[-2], cube.shape[-1]]) + + print("Running PCA post-processing", flush=True) + for i, npc in enumerate(ncomp): + res[i] = pca(cube, angle_list=angle_list, cube_ref=cube_ref, ncomp=npc, + mask_center_px=mask_center_px, imlib=imlib, nproc=nproc, + source_xy=source_xy, delta_rot=delta_rot, fwhm=fwhm, mask_rdi=mask_rdi, + ref_strategy=ref_strategy) + return res + + +# script behaviour? i dont like these +# if __name__ == '__main__': +# # python -m HEEPS.heeps.contrast.disc \ +# # --model_path /path/to/model.fits \ +# # --psf_on_path /path/to/onaxis_psf.fits \ +# # --psf_off_path /path/to/offaxis_psf.fits +# pass From a1252e8dbdfe2d282fd063bbb7730c530a637340 Mon Sep 17 00:00:00 2001 From: IainHammond Date: Fri, 15 May 2026 14:43:20 +0200 Subject: [PATCH 04/46] more docstrings --- heeps/contrast/disc.py | 25 ++++++++++++++----------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index f38c995..102c3b1 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -242,38 +242,41 @@ def postproc_disc( cube : np.ndarray 3‑D array (n_frames, ny, nx) containing the science frames. angle_list : np.ndarray - 1‑D array of parallactic angles (in degrees) associated with each frame in ``cube``. + 1‑D array of parallactic angles associated with each frame in ``cube``. cube_ref : np.ndarray or None, optional Reference cube for RDI (reference differential imaging). If ``None``, only ADI is performed. subsample : int, default=1 Factor by which to down‑sample the cube (and reference) for speed. ``1`` means no subsampling. ncomp : int, tuple, or list, default=1 - Number of principal components to use. If an int, a single component is used. If a tuple ``(min, max)`` or a list, the function will compute results for each component in the range. + Number of principal components to use. If an int, a single component is used. If a tuple ``(min, max)`` or a + list, the function will compute results for each component in the range. mask_center_px : int, default=0 - Radius (in pixels) of a circular mask applied to the centre of each frame before PCA. + Radius (in pixels) of a circular mask applied to the centre of each frame. imlib : str, default="vip-fft" - Image library used by VIP for FFT‑based operations. + Image library used by VIP. source_xy : tuple or None, optional - (x, y) pixel coordinates of a known source; used to mask the source during PCA if provided. + (x, y) pixel coordinates of a known source; used for frame rejection in the PCA library. delta_rot : float or int or None, default=1 - Minimum rotation (in FWHM) between frames for a given pixel to be considered independent. + Minimum rotation (in units of FWHM) of source_xy between frames. fwhm : float or int, default=5 Full‑width at half‑maximum of the PSF, used for delta_rot. mask_rdi : np.ndarray, optional - Optional mask applied to the RDI reference cube. + See description in VIP. ref_strategy : str, default='RDI' - Strategy for reference handling. + Strategy for reference handling (RDI or ARDI). nproc : int, default=1 - Number of processes to use for parallel computation. + Number of processors to use for parallel computation. Returns ------- np.ndarray - Array of shape ``(n_ncomp, ny, nx)`` where ``n_ncomp`` is the number of principal components evaluated. Each slice ``res[i]`` contains the PCA‑processed image for the corresponding number of components. + Array of shape ``(n_ncomp, ny, nx)`` where ``n_ncomp`` is the number of principal components evaluated. + Each slice ``res[i]`` contains the PCA‑processed image for the corresponding number of components. Notes ----- - The function currently supports the subset of ``vip_hci.psfsub.pca`` parameters used in the HEEPS disc injection workflow. Additional ``pca`` arguments can be added in the future. + The function currently supports the subset of ``vip_hci.psfsub.pca`` parameters used in the HEEPS disc injection + workflow. Additional ``pca`` arguments can be added in the future. """ # subsample the cube and the reference if requested, for efficiency purposes if subsample > 1: From 73e04c0bcb0f12d96727882084d4f38824410496 Mon Sep 17 00:00:00 2001 From: IainHammond Date: Fri, 15 May 2026 14:46:11 +0200 Subject: [PATCH 05/46] more docstrings --- heeps/contrast/disc.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index 10ff094..f7525ea 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -242,27 +242,41 @@ def postproc_disc( cube : np.ndarray 3‑D array (n_frames, ny, nx) containing the science frames. angle_list : np.ndarray + 1‑D array of parallactic angles associated with each frame in ``cube``. cube_ref : np.ndarray or None, optional Reference cube for RDI (reference differential imaging). If ``None``, only ADI is performed. subsample : int, default=1 Factor by which to down‑sample the cube (and reference) for speed. ``1`` means no subsampling. ncomp : int, tuple, or list, default=1 + Number of principal components to use. If an int, a single component is used. If a tuple ``(min, max)`` or a + list, the function will compute results for each component in the range. mask_center_px : int, default=0 + Radius (in pixels) of a circular mask applied to the centre of each frame. imlib : str, default="vip-fft" + Image library used by VIP for FFT‑based operations. source_xy : tuple or None, optional + (x, y) pixel coordinates of a known source; used for frame rejection in the PCA library. delta_rot : float or int or None, default=1 + Minimum rotation (in FWHM) between frames. fwhm : float or int, default=5 Full‑width at half‑maximum of the PSF, used for delta_rot. mask_rdi : np.ndarray, optional + See description in VIP. ref_strategy : str, default='RDI' + Strategy for reference handling (RDI or ARDI) nproc : int, default=1 + Number of processors to use for parallel computation. Returns ------- np.ndarray + Array of shape ``(n_ncomp, ny, nx)`` where ``n_ncomp`` is the number of principal components evaluated. + Each slice ``res[i]`` contains the PCA‑processed image for the corresponding number of components. Notes ----- + The function currently supports the subset of ``vip_hci.psfsub.pca`` parameters used in the HEEPS disc injection + workflow. Additional ``pca`` arguments can be added in the future. """ # subsample the cube and the reference if requested, for efficiency purposes if subsample > 1: From 270fb526fa505c0b30bf2b5deb0a80ccd59f1dce Mon Sep 17 00:00:00 2001 From: IainHammond Date: Fri, 15 May 2026 14:46:45 +0200 Subject: [PATCH 06/46] oops I left a typo --- heeps/contrast/disc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index f7525ea..ec647f7 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -5,7 +5,7 @@ PSF data and to perform post‑processing with VIP. The entry point is :func:`create_disc_sequence`, which creates a mock observation sequence with a user-provided disc mode. The :func:`post_proc_disc` function is a -wrapper of the relevevant VIP function for PCA post-processing. +wrapper of the relevant VIP function for PCA post-processing. """ import numpy as np From 381c65b2a1f773864e7c4d457446f371c02fd695 Mon Sep 17 00:00:00 2001 From: IainHammond Date: Fri, 15 May 2026 14:49:24 +0200 Subject: [PATCH 07/46] imlib description --- heeps/contrast/disc.py | 12 ++---------- 1 file changed, 2 insertions(+), 10 deletions(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index ec647f7..7fc11c2 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -73,7 +73,8 @@ def create_disc_sequence( starphot : float, default=1e11 Photometric scaling factor for the star. imlib : str, default="opencv" - Image library to use for image processing. + Image library to use for image processing. Use "opencv" to go fast, or "vio-fft" for better flux conservation + at the cost of speed. See VIP documentation for details. **conf : dict Additional conf parameters from HEEPS (e.g., background addition etc.) @@ -299,12 +300,3 @@ def postproc_disc( source_xy=source_xy, delta_rot=delta_rot, fwhm=fwhm, mask_rdi=mask_rdi, ref_strategy=ref_strategy) return res - - -# script behaviour? i dont like these -# if __name__ == '__main__': -# # python -m HEEPS.heeps.contrast.disc \ -# # --model_path /path/to/model.fits \ -# # --psf_on_path /path/to/onaxis_psf.fits \ -# # --psf_off_path /path/to/offaxis_psf.fits -# pass From 0b9aafd8e16a5e25f63e8ea0603519484fc44ad1 Mon Sep 17 00:00:00 2001 From: IainHammond Date: Fri, 15 May 2026 14:50:15 +0200 Subject: [PATCH 08/46] fin --- heeps/contrast/disc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index 7fc11c2..5551538 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -63,7 +63,7 @@ def create_disc_sequence( is used. rdi_duration : Union[int, float, None], optional Desired duration (in seconds) of the RDI reference sequence. If ``None``, a default of - 20 % of the on‑axis sequence length is used. + 20% of the on‑axis sequence length is used. source_xy : Union[tuple, list, None], optional (x, y) pixel coordinates of the source for which extinction is applied. extinction : Union[int, float], default=0 From 67a07b32067de5a8bdf8cc767cc213f6f676e7c1 Mon Sep 17 00:00:00 2001 From: IainHammond Date: Fri, 15 May 2026 14:51:44 +0200 Subject: [PATCH 09/46] post_proc_disc -> postproc_disc --- heeps/contrast/disc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index 5551538..e491e5f 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -4,7 +4,7 @@ This module provides functions to inject a synthetic disc model into HEEPS PSF data and to perform post‑processing with VIP. The entry point is :func:`create_disc_sequence`, which creates a mock -observation sequence with a user-provided disc mode. The :func:`post_proc_disc` function is a +observation sequence with a user-provided disc mode. The :func:`postproc_disc` function is a wrapper of the relevant VIP function for PCA post-processing. """ From 1d11cf53b63078aa68b979b55469c13cd477820d Mon Sep 17 00:00:00 2001 From: IainHammond Date: Fri, 15 May 2026 15:02:49 +0200 Subject: [PATCH 10/46] final docstring update --- heeps/contrast/disc.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index e491e5f..72cf613 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -40,9 +40,10 @@ def create_disc_sequence( **conf ): """ - Create a mock observation sequence by injecting a synthetic disc model into a HEEPS - PSF and optionally generate a reference differential imaging (RDI) - cube. Existing on- and off-axis PSF cubes can be provided. + Create a mock observation sequence by injecting a synthetic disc model into a HEEPS PSF and optionally generate + a reference differential imaging (RDI) cube. Existing on- and off-axis PSF cubes, and coronagraph transmission, + can be provided. If not provided, the function will attempt to load them from the output directory based on the + conf parameters. Parameters ---------- From 0c15ef48bdab179a4baea17e36a73d556ba046fd Mon Sep 17 00:00:00 2001 From: IainHammond Date: Tue, 19 May 2026 11:58:19 +0200 Subject: [PATCH 11/46] description and type hint improvements --- heeps/contrast/disc.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index 72cf613..003e9e3 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -35,7 +35,7 @@ def create_disc_sequence( source_xy : Union[tuple, list, None] = None, extinction : Union[int, float] = 0, tag : Union[str, None] = None, - starphot : float = 1e11, + starphot : Union[float, int] = 1e11, imlib : str = "opencv", **conf ): @@ -74,10 +74,10 @@ def create_disc_sequence( starphot : float, default=1e11 Photometric scaling factor for the star. imlib : str, default="opencv" - Image library to use for image processing. Use "opencv" to go fast, or "vio-fft" for better flux conservation + Image library to use for image processing. Use "opencv" to go fast, or "vip-fft" for better flux conservation at the cost of speed. See VIP documentation for details. - **conf : dict - Additional conf parameters from HEEPS (e.g., background addition etc.) + conf : dict, optional + Configuration dictionary with additional parameters from HEEPS (e.g., background addition etc.). Returns ------- From 1ba498a8222ba6d02341ea32f08b6b66293dec5b Mon Sep 17 00:00:00 2001 From: IainHammond Date: Tue, 19 May 2026 12:04:42 +0200 Subject: [PATCH 12/46] description and type hint improvements --- heeps/contrast/disc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index 003e9e3..08e59c9 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -77,7 +77,7 @@ def create_disc_sequence( Image library to use for image processing. Use "opencv" to go fast, or "vip-fft" for better flux conservation at the cost of speed. See VIP documentation for details. conf : dict, optional - Configuration dictionary with additional parameters from HEEPS (e.g., background addition etc.). + Configuration dictionary with parameters from HEEPS (e.g., background addition etc.). Returns ------- From 7b106e5e2308d9744808534423f94062c2f1e856 Mon Sep 17 00:00:00 2001 From: IainHammond Date: Tue, 19 May 2026 13:45:42 +0200 Subject: [PATCH 13/46] some bug fixes already and support for the star to be on the centre 4 pixels --- heeps/contrast/disc.py | 26 +++++++++++++++++--------- 1 file changed, 17 insertions(+), 9 deletions(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index 08e59c9..238697c 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -13,7 +13,7 @@ import os from vip_hci.fits import open_fits, open_header -from vip_hci.preproc import frame_crop, cube_subsample, cube_crop_frames +from vip_hci.preproc import frame_crop, cube_subsample, cube_crop_frames, frame_shift from vip_hci.fm import cube_inject_fakedisk from vip_hci.psfsub import pca @@ -48,7 +48,8 @@ def create_disc_sequence( Parameters ---------- disc_model : np.ndarray - 2‑D array representing the disc model to be injected. NaNs are replaced with zeros. + 2‑D array representing the disc model to be injected. NaNs are replaced with zeros. Odd-sized models with the + stellar flux on the centre pixel are preferred. on_axis_psf : Union[np.ndarray, None], optional Pre‑loaded on‑axis PSF cube. If ``None``, the PSF is loaded from the output directory using the conf parameters. @@ -94,21 +95,31 @@ def create_disc_sequence( if off_axis_psf is None: psf_OFF = open_fits(loadname % 'offaxis', verbose=False) print(f"Using off-axis PSF from {loadname%'offaxis'}") + else: + psf_OFF = off_axis_psf assert psf_OFF.ndim == 2, "off-axis PSF frame must be 2-dimensional" # set nans to zero disc_model = np.nan_to_num(disc_model, nan=0) - # Remove any extra dimensions (e.g., MCFOST cubes) + # remove any extra dimensions (e.g., MCFOST cubes) # first, squeeze singleton dimensions disc_model = np.squeeze(disc_model) # then collapse any remaining leading dimensions to a 2-D image while disc_model.ndim > 2: disc_model = disc_model[0] + # record the stellar flux in the model for flux normalisation, and set its pixel to 0 (it will be in the on-axis PSF) + # some codes have the star on the centre 4 pixels instead of the centre pixel only so we need to sum the values + mask = disc_model == disc_model.max() + star_val = disc_model[mask].sum() + disc_model[mask] = 0 + # ensure the model has an odd size if disc_model.shape[-1] % 2 == 0: - disc_model = disc_model[1:, 1:] # we assume the star flux is on one pixel + print("Converting input disc model to odd-dimensions.", flush=True) + disc_model = frame_shift(disc_model, shift_y=0.5, shift_x=0.5, imlib=imlib) + disc_model = disc_model[1:, 1:] # crop everything to a common size # the HEEPS PSFs are usually 293px @@ -118,11 +129,6 @@ def create_disc_sequence( if disc_model.shape[-1] > min_crop: disc_model = frame_crop(disc_model, min_crop, verbose=False) - # record the stellar flux in the model and set the pixel to 0 - star_val = np.max(disc_model) - star_y, star_x = np.unravel_index(np.argmax(disc_model), disc_model.shape) - disc_model[star_y, star_x] = 0 - # apply extinction if requested if extinction > 0 and source_xy is not None: source_x, source_y = source_xy @@ -169,6 +175,8 @@ def create_disc_sequence( if on_axis_psf is None: psf_ON = open_fits(loadname % 'onaxis', verbose=False) print(f"Using on-axis PSF from {loadname % 'onaxis'}") + else: + psf_ON = on_axis_psf assert psf_ON.ndim == 3, "on-axis PSF cube must be 3-dimensional" if psf_ON.shape[-1] > min_crop: From 568c1edc6d5693ee540e64b3b56176c79c6c2352 Mon Sep 17 00:00:00 2001 From: IainHammond Date: Tue, 19 May 2026 14:40:31 +0200 Subject: [PATCH 14/46] corrected the oat file path --- heeps/contrast/disc.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index 238697c..3e73800 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -41,7 +41,7 @@ def create_disc_sequence( ): """ Create a mock observation sequence by injecting a synthetic disc model into a HEEPS PSF and optionally generate - a reference differential imaging (RDI) cube. Existing on- and off-axis PSF cubes, and coronagraph transmission, + a reference differential imaging (RDI) cube. Existing on- and off-axis PSF cubes and coronagraph transmission can be provided. If not provided, the function will attempt to load them from the output directory based on the conf parameters. @@ -135,13 +135,14 @@ def create_disc_sequence( extinction_factor = 10 ** (-0.4 * extinction) disc_model[source_y, source_x] *= extinction_factor - # load coronagraph transmission + # load coronagraph transmission depending on the mode and band if transmission is None: if conf["mode"] == "ELT" or conf["mode"] is None: # no coronagraph transmission = None elif conf["mode"] == "CVC": - transmission = open_fits(conf["f_vc_trans"], verbose=False) - print(f"Using transmission from {conf['f_vc_trans']}") + conf['f_oat'] = conf['dir_input'] + 'optics/vc/oat_%s_%s.fits'%(conf['band'], conf['mode']) + transmission = open_fits(conf['f_oat'], verbose=False) + print(f"Using transmission from {conf['f_oat']}") else: # TODO raise NotImplementedError(f"Mode {conf['mode']} not implemented for disc injection.") From d844c8d86f229d5714b5c2a26b0e5c0204837b5b Mon Sep 17 00:00:00 2001 From: IainHammond Date: Tue, 19 May 2026 15:06:26 +0200 Subject: [PATCH 15/46] implemented VC mode transmission handling --- heeps/contrast/disc.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index 3e73800..47b424c 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -139,10 +139,13 @@ def create_disc_sequence( if transmission is None: if conf["mode"] == "ELT" or conf["mode"] is None: # no coronagraph transmission = None - elif conf["mode"] == "CVC": - conf['f_oat'] = conf['dir_input'] + 'optics/vc/oat_%s_%s.fits'%(conf['band'], conf['mode']) - transmission = open_fits(conf['f_oat'], verbose=False) - print(f"Using transmission from {conf['f_oat']}") + elif "VC" in conf["mode"]: + if conf["band"] in ("L", "M"): + conf['f_oat'] = conf['dir_input'] + 'optics/vc/oat_L_%s.fits' % conf['mode'] + elif conf["band"] == "N2": + conf['f_oat'] = conf['dir_input'] + 'optics/vc/oat_N2_CVC.fits' + transmission = open_fits(conf['f_oat'], verbose=False) + print(f"Using transmission from {conf['f_oat']}") else: # TODO raise NotImplementedError(f"Mode {conf['mode']} not implemented for disc injection.") From c43a2db610d22dd68e8221521991ab1ba9af6e0d Mon Sep 17 00:00:00 2001 From: IainHammond Date: Tue, 19 May 2026 16:23:01 +0200 Subject: [PATCH 16/46] fix for RDI and manually provided PSFs --- heeps/contrast/disc.py | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index 47b424c..d39c1b1 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -181,13 +181,14 @@ def create_disc_sequence( print(f"Using on-axis PSF from {loadname % 'onaxis'}") else: psf_ON = on_axis_psf - assert psf_ON.ndim == 3, "on-axis PSF cube must be 3-dimensional" + if psf_ON.shape[-1] > min_crop: psf_ON = cube_crop_frames(psf_ON, min_crop, verbose=False) # add the disc model psf_ON += cube + del cube # background addition if conf["add_bckg"]: @@ -203,7 +204,13 @@ def create_disc_sequence( rdi_duration = int(0.2 * nframes) * conf["dit"] n_ref_frames = rdi_duration / conf["dit"] start_idx = np.random.randint(0, psf_ON.shape[0] - int(n_ref_frames) + 1) # random segment of the on-axis PSF cube to use as the RDI reference, to better match the background conditions - psf_RDI = open_fits(loadname % "onaxis", verbose=False)[start_idx:start_idx + int(n_ref_frames)] + + if on_axis_psf is None: # reload the on-axis sequence + psf_RDI = open_fits(loadname % "onaxis", verbose=False)[start_idx:start_idx + int(n_ref_frames)] + else: + psf_RDI = on_axis_psf[start_idx:start_idx + int(n_ref_frames)] + del on_axis_psf + if psf_RDI.shape[-1] > min_crop: psf_RDI = cube_crop_frames(psf_RDI, min_crop, verbose=False) @@ -211,7 +218,11 @@ def create_disc_sequence( psf_ON = np.concatenate([psf_ON[:start_idx], psf_ON[start_idx + int(n_ref_frames):]], axis=0) pa = np.concatenate([pa[:start_idx], pa[start_idx + int(n_ref_frames):]], axis=0) - psf_RDI_OFF = open_fits(loadname % "offaxis", verbose=False) + if off_axis_psf is None: + psf_RDI_OFF = open_fits(loadname % "offaxis", verbose=False) + else: + psf_RDI_OFF = off_axis_psf + if psf_RDI_OFF.shape[-1] > min_crop: psf_RDI_OFF = frame_crop(psf_RDI_OFF, min_crop, verbose=False) From 7d757fdd77a5b5c5f03cec6e892d93b2fb5d8df7 Mon Sep 17 00:00:00 2001 From: IainHammond Date: Tue, 19 May 2026 20:19:53 +0200 Subject: [PATCH 17/46] adding noise is even faster --- heeps/contrast/background.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/heeps/contrast/background.py b/heeps/contrast/background.py index 1dda6bf..2605bff 100644 --- a/heeps/contrast/background.py +++ b/heeps/contrast/background.py @@ -82,10 +82,11 @@ def background(psf_ON, psf_OFF, header=None, mode='RAVC', lam=3.8e-6, dit=0.3, # np.random.normal(0, sigma_array) for large cubes, because # standard_normal draws from N(0,1) using a fast vectorised path # and we then scale by sqrt(psf_ON) in a single multiply pass, - # avoiding per-element sigma sampling. + # avoiding per-element sigma sampling rng = np.random.default_rng(seed) psf_sqrt = np.sqrt(psf_ON) - psf_ON += psf_sqrt * rng.standard_normal(psf_ON.shape) + psf_sqrt *= rng.standard_normal(psf_ON.shape) + psf_ON += psf_sqrt if verbose is True: print(' dit=%s s, thruput=%.4f, mask_trans=%.4f,'%(dit, thruput, mask_trans)) From c1a911d35b21762151c0804f4b095770d6792bd9 Mon Sep 17 00:00:00 2001 From: IainHammond Date: Tue, 19 May 2026 20:31:27 +0200 Subject: [PATCH 18/46] cleanup --- heeps/contrast/background.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/heeps/contrast/background.py b/heeps/contrast/background.py index 2605bff..b3c321a 100644 --- a/heeps/contrast/background.py +++ b/heeps/contrast/background.py @@ -8,7 +8,7 @@ def background(psf_ON, psf_OFF, header=None, mode='RAVC', lam=3.8e-6, dit=0.3, """ This function applies background and photon noise to intup PSFs (off-axis - and on-axis), incuding transmission, star flux, and components transmittance. + and on-axis), including transmission, star flux, and components transmittance. Args: psf_ON (float ndarray): @@ -64,7 +64,7 @@ def background(psf_ON, psf_OFF, header=None, mode='RAVC', lam=3.8e-6, dit=0.3, psf_ON *= app_single_psf # scopesim-heeps interface - if call_ScopeSim is True: + if call_ScopeSim: from heeps.contrast.sim_heeps import sim_heeps psf_ON, psf_OFF = sim_heeps(psf_ON, psf_OFF, header, **conf) else: @@ -76,19 +76,17 @@ def background(psf_ON, psf_OFF, header=None, mode='RAVC', lam=3.8e-6, dit=0.3, bckg_noise = dit * flux_bckg * thruput * mask_trans psf_ON += bckg_noise # add photon noise ~ N(0, sqrt(psf)) - # np.random.seed(seed) - # psf_ON += np.random.normal(0, np.sqrt(psf_ON)) - # Using default_rng with standard_normal is much faster than + # default_rng with standard_normal is much faster than # np.random.normal(0, sigma_array) for large cubes, because # standard_normal draws from N(0,1) using a fast vectorised path - # and we then scale by sqrt(psf_ON) in a single multiply pass, + # and then scale by sqrt(psf_ON) in a single multiply pass, # avoiding per-element sigma sampling rng = np.random.default_rng(seed) psf_sqrt = np.sqrt(psf_ON) psf_sqrt *= rng.standard_normal(psf_ON.shape) psf_ON += psf_sqrt - if verbose is True: + if verbose: print(' dit=%s s, thruput=%.4f, mask_trans=%.4f,'%(dit, thruput, mask_trans)) print(' mag=%s, star_signal=%.2e, bckg_noise=%.2e'%(mag, star_signal, bckg_noise)) From 92cc7ea3ff00b20b78a6bbf647f5ae1e99e62b9f Mon Sep 17 00:00:00 2001 From: IainHammond Date: Wed, 20 May 2026 21:58:02 +0200 Subject: [PATCH 19/46] only loading in required vip functions --- heeps/util/psf_template.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/heeps/util/psf_template.py b/heeps/util/psf_template.py index 6d61e4e..9a271f7 100644 --- a/heeps/util/psf_template.py +++ b/heeps/util/psf_template.py @@ -1,5 +1,9 @@ import numpy as np -import vip_hci + +from vip_hci.metrics import aperture_flux +from vip_hci.preproc import frame_shift +from vip_hci.var import fit_2dgaussian + def psf_template(psf, center=None, recenter=True, ncrop=4): @@ -10,7 +14,7 @@ def psf_template(psf, center=None, recenter=True, ncrop=4): else: (cx, cy) = center # fit a 2D Gaussian --> output: fwhm, x-y centroid - fit = vip_hci.var.fit_2dgaussian(psf, True, (cx, cy), debug=False, full_output=True) + fit = fit_2dgaussian(psf, True, (cx, cy), debug=False, full_output=True) # derive the FWHM fwhm = np.mean([fit['fwhm_x'], fit['fwhm_y']]) # recenter @@ -19,11 +23,11 @@ def psf_template(psf, center=None, recenter=True, ncrop=4): assert (yerr < 0.5) and (xerr < 0.5), 'centroid (x,y) error = '\ '(%.1e, %.1e), must be < 0.5 pixel'%(xerr, yerr) shiftx, shifty = cx-fit['centroid_x'].values[0], cy-fit['centroid_y'].values[0] - psf = vip_hci.preproc.frame_shift(psf, shifty, shiftx) + psf = frame_shift(psf, shifty, shiftx) # FWHM aperture photometry - ap_flux = vip_hci.metrics.aperture_flux(psf, [cy], [cx], fwhm, verbose=False)[0] + ap_flux = aperture_flux(psf, [cy], [cx], fwhm, verbose=False)[0] # image radius (cropped), default to 4x FWHM rim = round(ncrop*fwhm) psf_crop = psf[cy-rim:cy+rim+1, cx-rim:cx+rim+1] - return psf_crop, fwhm, ap_flux \ No newline at end of file + return psf_crop, fwhm, ap_flux From cceed2f5276e427c34ea4fdf90af782136f4484e Mon Sep 17 00:00:00 2001 From: IainHammond Date: Sat, 23 May 2026 10:01:15 +0200 Subject: [PATCH 20/46] typo --- heeps/contrast/disc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index d39c1b1..5b279ff 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -4,7 +4,7 @@ This module provides functions to inject a synthetic disc model into HEEPS PSF data and to perform post‑processing with VIP. The entry point is :func:`create_disc_sequence`, which creates a mock -observation sequence with a user-provided disc mode. The :func:`postproc_disc` function is a +observation sequence with a user-provided disc model. The :func:`postproc_disc` function is a wrapper of the relevant VIP function for PCA post-processing. """ From 288705ae40162ad84f86f5eab071b61f14dba8e1 Mon Sep 17 00:00:00 2001 From: IainHammond Date: Tue, 21 Jul 2026 15:03:26 +0200 Subject: [PATCH 21/46] new grid update. new rdi function. --- heeps/contrast/disc.py | 308 +++++++++++++++++++++++------------------ 1 file changed, 172 insertions(+), 136 deletions(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index 5b279ff..2d28b36 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -13,15 +13,13 @@ import os from vip_hci.fits import open_fits, open_header -from vip_hci.preproc import frame_crop, cube_subsample, cube_crop_frames, frame_shift +from vip_hci.preproc import frame_crop, cube_crop_frames, frame_shift from vip_hci.fm import cube_inject_fakedisk -from vip_hci.psfsub import pca -from heeps.contrast.background import background from heeps.util.psf_template import psf_template from heeps.util.paralang import paralang -__all__ = ['create_disc_sequence', 'postproc_disc'] +__all__ = ['create_disc_sequence', 'prepare_rdi_sequence'] __author__ = "Iain Hammond" def create_disc_sequence( @@ -29,12 +27,9 @@ def create_disc_sequence( on_axis_psf : Union[np.ndarray, None] = None, off_axis_psf : Union[np.ndarray, None] = None, transmission : Union[np.ndarray, None] = None, - rdi : bool = True, - rdi_mag : Union[int, float, None] = None, - rdi_duration : Union[int, float, None] = None, + seeing_q : Union[int] = 2, source_xy : Union[tuple, list, None] = None, extinction : Union[int, float] = 0, - tag : Union[str, None] = None, starphot : Union[float, int] = 1e11, imlib : str = "opencv", **conf @@ -58,20 +53,12 @@ def create_disc_sequence( transmission : Union[np.ndarray, None], optional Coronagraph transmission. If ``None``, the appropriate transmission is loaded based on the instrument mode in conf. Uses VIP conventions. - rdi : bool, default=True - Whether to generate an RDI reference cube from a segment of the on‑axis sequence. - rdi_mag : Union[int, float, None], optional - Magnitude to use for the RDI reference. If ``None``, the magnitude of the science target - is used. - rdi_duration : Union[int, float, None], optional - Desired duration (in seconds) of the RDI reference sequence. If ``None``, a default of - 20% of the on‑axis sequence length is used. + seeing_q : int, default=2 + Seeing quartile (1-3) to select the PSF. Uses a default seeing of Q2 if not specified. source_xy : Union[tuple, list, None], optional (x, y) pixel coordinates of the source for which extinction is applied. extinction : Union[int, float], default=0 Extinction in magnitudes to apply to the source at ``source_xy``. - tag : Union[str, None], optional - Tag to prepend to output filenames from earlier HEEPS runs. starphot : float, default=1e11 Photometric scaling factor for the star. imlib : str, default="opencv" @@ -86,12 +73,19 @@ def create_disc_sequence( Sequence with the fake disc injected into the on‑axis PSF cube. pa : numpy.ndarray Array of parallactic angles used for the injection. - psf_RDI : numpy.ndarray, optional - Reference cube for RDI (if ``rdi`` is ``True``). Returned as the third element when - ``rdi`` is enabled; otherwise only ``psf_ON`` and ``pa`` are returned. """ - tag = "" if tag is None else "%s_"%tag - loadname = os.path.join(conf['dir_output'], '%s%s_PSF_%s_%s.fits'%(tag,'%s', conf['band'], conf['mode'])) + # if mag is not in the grid, round to the closest available magnitude + conf["mag"] = _check_mag(conf["mag"]) + + # new grid folder structure, assuming your dir_current points to where the folders are saved + grid_dir = os.path.join( + conf['dir_current'], + f"{conf['band']}_{conf['mode']}_s=Q{seeing_q}_mag={conf['mag']}_{conf['duration']}s_{int(conf['dit'] * 100)}ms" + ) + + # convention for the April 2026 grid files + loadname = os.path.join(grid_dir, '%s_PSF_bckg1_%s_%s.fits' % ('%s', conf['band'], conf['mode'])) + if off_axis_psf is None: psf_OFF = open_fits(loadname % 'offaxis', verbose=False) print(f"Using off-axis PSF from {loadname%'offaxis'}") @@ -122,8 +116,8 @@ def create_disc_sequence( disc_model = disc_model[1:, 1:] # crop everything to a common size - # the HEEPS PSFs are usually 293px min_crop = min(psf_OFF.shape[-1], disc_model.shape[-1], conf["ndet"]) + conf["ndet"] = min_crop if psf_OFF.shape[-1] > min_crop: psf_OFF = frame_crop(psf_OFF, min_crop, verbose=False) if disc_model.shape[-1] > min_crop: @@ -134,6 +128,7 @@ def create_disc_sequence( source_x, source_y = source_xy extinction_factor = 10 ** (-0.4 * extinction) disc_model[source_y, source_x] *= extinction_factor + print("Applied extinction of %.2f mag to source at (x, y) = (%d, %d)" % (extinction, source_x, source_y), flush=True) # load coronagraph transmission depending on the mode and band if transmission is None: @@ -149,9 +144,9 @@ def create_disc_sequence( else: # TODO raise NotImplementedError(f"Mode {conf['mode']} not implemented for disc injection.") - # determine how many frames are in the on-axis sequence (don't open the whole cube yet to save memory) + # determine how many frames are in the on-axis sequence if on_axis_psf is None: - nframes = open_header(loadname % 'onaxis')['NAXIS3'] + nframes = open_header(loadname % 'onaxis')['NAXIS3'] # don't open the whole cube yet to save memory else: nframes = on_axis_psf.shape[0] @@ -171,7 +166,7 @@ def create_disc_sequence( ker=1, imlib=imlib, ) - + # write_fits(conf["dir_current"] + "/prepared_model.fits", cube[0]) # normalise by the stellar flux that we removed earlier cube /= star_val @@ -190,137 +185,178 @@ def create_disc_sequence( psf_ON += cube del cube - # background addition - if conf["add_bckg"]: - psf_ON, psf_OFF = background(psf_ON, psf_OFF, verbose=True, **conf) - _, _, ap_flux = psf_template(psf_OFF) psf_ON *= starphot / ap_flux # RDI handling (at the moment we only support taking a chunk out of the science sequence - if rdi: - if rdi_duration is None: - print("Warning: rdi_duration was not set. Using 20% of on-axis sequence.", flush=True) - rdi_duration = int(0.2 * nframes) * conf["dit"] - n_ref_frames = rdi_duration / conf["dit"] - start_idx = np.random.randint(0, psf_ON.shape[0] - int(n_ref_frames) + 1) # random segment of the on-axis PSF cube to use as the RDI reference, to better match the background conditions + # if rdi: + # if rdi_duration is None: + # print("Warning: rdi_duration was not set. Using 20% of on-axis sequence.", flush=True) + # rdi_duration = int(0.2 * nframes) * conf["dit"] + # n_ref_frames = rdi_duration / conf["dit"] + # start_idx = np.random.randint(0, psf_ON.shape[0] - int(n_ref_frames) + 1) # random segment of the on-axis PSF cube to use as the RDI reference, to better match the background conditions + # + # if on_axis_psf is None: # reload the on-axis sequence + # psf_RDI = open_fits(loadname % "onaxis", verbose=False)[start_idx:start_idx + int(n_ref_frames)] + # else: + # psf_RDI = on_axis_psf[start_idx:start_idx + int(n_ref_frames)] + # del on_axis_psf + # + # if psf_RDI.shape[-1] > min_crop: + # psf_RDI = cube_crop_frames(psf_RDI, min_crop, verbose=False) + # + # # remove the RDI segment from psf_ON to correctly represent lost integration time on the science target + # psf_ON = np.concatenate([psf_ON[:start_idx], psf_ON[start_idx + int(n_ref_frames):]], axis=0) + # pa = np.concatenate([pa[:start_idx], pa[start_idx + int(n_ref_frames):]], axis=0) + # + # if off_axis_psf is None: + # psf_RDI_OFF = open_fits(loadname % "offaxis", verbose=False) + # else: + # psf_RDI_OFF = off_axis_psf + # + # if psf_RDI_OFF.shape[-1] > min_crop: + # psf_RDI_OFF = frame_crop(psf_RDI_OFF, min_crop, verbose=False) + # + # if rdi_mag is None: + # print("Warning: rdi_mag was not set. Using mag of science target.", flush=True) + # else: + # conf["mag"] = rdi_mag + # + # if conf["add_bckg"]: + # psf_RDI, psf_RDI_OFF = background(psf_RDI, psf_RDI_OFF, verbose=True, **conf) + # + # _, _, ap_flux = psf_template(psf_RDI_OFF) + # psf_RDI *= starphot / ap_flux + # del psf_RDI_OFF + # + # return psf_ON, pa, psf_RDI + # + # else: + return psf_ON, pa + + +def prepare_rdi_sequence( + on_axis_psf : Union[np.ndarray, None] = None, + off_axis_psf : Union[np.ndarray, None] = None, + rdi_mag : Union[int, float, None] = None, + rdi_duration : Union[int, float, None] = None, + seeing_q : Union[int] = 2, + starphot : Union[float, int] = 1e11, + **conf +): + """ + Prepare a sequence to be used as a reference for RDI (Reference Differential Imaging) from the PSF grid. + The function extracts a segment of the on-axis PSF cube to serve as the RDI reference depending on duration. + + Parameters + ---------- + on_axis_psf : Union[np.ndarray, None], optional + Pre‑loaded on‑axis PSF cube. If ``None``, the PSF is loaded based off conf and rdi_mag. + using the conf parameters. + off_axis_psf : Union[np.ndarray, None], optional + Pre‑loaded off‑axis PSF frame. If ``None``, the PSF is loaded based off conf and rdi_mag. + rdi_mag : Union[int, float, None], optional + Magnitude to use for the RDI reference. If ``None``, magnitude in conf is used. + rdi_duration : Union[int, float, None], optional + Desired duration (in seconds) of the RDI reference sequence. If ``None``, a default of + 20% of the on‑axis sequence length is used. + seeing_q : int, default=2 + Seeing quartile (1-3) to select the PSF. Uses a default seeing of Q2 if not specified. + starphot : float, default=1e11 + Photometric scaling factor for the star. + + Returns + ------- + psf_RDI : numpy.ndarray + Reference cube for RDI. + """ + if rdi_mag is None: + rdi_mag = conf["mag"] + print("Warning: rdi_mag was not set. Using mag in the conf dictionary.", flush=True) - if on_axis_psf is None: # reload the on-axis sequence - psf_RDI = open_fits(loadname % "onaxis", verbose=False)[start_idx:start_idx + int(n_ref_frames)] - else: - psf_RDI = on_axis_psf[start_idx:start_idx + int(n_ref_frames)] - del on_axis_psf + rdi_mag = _check_mag(rdi_mag) - if psf_RDI.shape[-1] > min_crop: - psf_RDI = cube_crop_frames(psf_RDI, min_crop, verbose=False) + grid_dir = os.path.join( + conf['dir_current'], + f"{conf['band']}_{conf['mode']}_s=Q{seeing_q}_mag={rdi_mag}_{conf['duration']}s_{int(conf['dit'] * 100)}ms" + ) - # remove the RDI segment from psf_ON to correctly represent lost integration time on the science target - psf_ON = np.concatenate([psf_ON[:start_idx], psf_ON[start_idx + int(n_ref_frames):]], axis=0) - pa = np.concatenate([pa[:start_idx], pa[start_idx + int(n_ref_frames):]], axis=0) + # convention for the April 2026 grid files, assuming dir_current points to the folder containing the PSF files + loadname = os.path.join(grid_dir, '%s_PSF_bckg1_%s_%s.fits' % ('%s', conf['band'], conf['mode'])) if off_axis_psf is None: - psf_RDI_OFF = open_fits(loadname % "offaxis", verbose=False) + psf_OFF = open_fits(loadname % 'offaxis', verbose=False) + print(f"Using off-axis PSF from {loadname%'offaxis'}") else: - psf_RDI_OFF = off_axis_psf - - if psf_RDI_OFF.shape[-1] > min_crop: - psf_RDI_OFF = frame_crop(psf_RDI_OFF, min_crop, verbose=False) + psf_OFF = off_axis_psf + assert psf_OFF.ndim == 2, "off-axis PSF frame must be 2-dimensional" - if rdi_mag is None: - print("Warning: rdi_mag was not set. Using mag of science target.", flush=True) + if on_axis_psf is None: + psf_ON = open_fits(loadname % 'onaxis', verbose=False) + print(f"Using on-axis PSF from {loadname % 'onaxis'}") else: - conf["mag"] = rdi_mag + psf_ON = on_axis_psf + assert psf_ON.ndim == 3, "on-axis PSF cube must be 3-dimensional" + + if psf_OFF.shape[-1] > conf["ndet"]: + psf_OFF = frame_crop(psf_OFF, conf["ndet"], verbose=False) + + if psf_ON.shape[-1] > conf["ndet"]: + psf_ON = cube_crop_frames(psf_ON, conf["ndet"], verbose=False) + + if rdi_duration is None: + rdi_duration = int(psf_ON.shape[0] * 0.2) * conf["dit"] + print("Warning: rdi_duration was not set. Using 20% of full sequence.", flush=True) + + # convert into number of frames to extract from the on-axis PSF cube + n_ref_frames = rdi_duration / conf["dit"] - if conf["add_bckg"]: - psf_RDI, psf_RDI_OFF = background(psf_RDI, psf_RDI_OFF, verbose=True, **conf) + # check if the requested reference duration is longer than the available on-axis PSF cube + if n_ref_frames > psf_ON.shape[0]: + n_ref_frames = psf_ON.shape[0] + print(f"WARNING! Requested RDI duration ({rdi_duration}s) exceeded available on-axis PSF cube length " + f"({psf_ON.shape[0] * conf['dit']}s). Setting them to be the same.", flush=True) - _, _, ap_flux = psf_template(psf_RDI_OFF) + # random segment of the on-axis PSF cube to use as the RDI reference + start_idx = np.random.randint(0, psf_ON.shape[0] - int(n_ref_frames) + 1) + psf_RDI = psf_ON[start_idx:start_idx + int(n_ref_frames)] + + _, _, ap_flux = psf_template(psf_OFF) psf_RDI *= starphot / ap_flux - del psf_RDI_OFF + del psf_OFF - return psf_ON, pa, psf_RDI + return psf_RDI - else: - return psf_ON, pa - - -def postproc_disc( - cube : np.ndarray, - angle_list : np.ndarray, - cube_ref : Union[np.ndarray, None] = None, - subsample : int = 1, - ncomp: Union[tuple, list, int] = 1, - mask_center_px : int = 0, - imlib : str = "vip-fft", - source_xy : Union[tuple, None] = None, - delta_rot : Union[float, int, None] = 1, - fwhm : Union[float, int] = 5, - mask_rdi: np.ndarray = None, - ref_strategy: str = 'RDI', - nproc : int = 1 -): + +def _check_mag(mag : Union[float, int] ) -> float: """ - Perform PCA-based post-processing on a disc‑injected data cube. + Check if the provided magnitude is in the Liege grid. If not, round to the closest available magnitude. Parameters ---------- - cube : np.ndarray - 3‑D array (n_frames, ny, nx) containing the science frames. - angle_list : np.ndarray - 1‑D array of parallactic angles associated with each frame in ``cube``. - cube_ref : np.ndarray or None, optional - Reference cube for RDI (reference differential imaging). If ``None``, only ADI is performed. - subsample : int, default=1 - Factor by which to down‑sample the cube (and reference) for speed. ``1`` means no subsampling. - ncomp : int, tuple, or list, default=1 - Number of principal components to use. If an int, a single component is used. If a tuple ``(min, max)`` or a - list, the function will compute results for each component in the range. - mask_center_px : int, default=0 - Radius (in pixels) of a circular mask applied to the centre of each frame. - imlib : str, default="vip-fft" - Image library used by VIP for FFT‑based operations. - source_xy : tuple or None, optional - (x, y) pixel coordinates of a known source; used for frame rejection in the PCA library. - delta_rot : float or int or None, default=1 - Minimum rotation (in FWHM) between frames. - fwhm : float or int, default=5 - Full‑width at half‑maximum of the PSF, used for delta_rot. - mask_rdi : np.ndarray, optional - See description in VIP. - ref_strategy : str, default='RDI' - Strategy for reference handling (RDI or ARDI) - nproc : int, default=1 - Number of processors to use for parallel computation. + mag : float, int + The magnitude to check. Returns ------- - np.ndarray - Array of shape ``(n_ncomp, ny, nx)`` where ``n_ncomp`` is the number of principal components evaluated. - Each slice ``res[i]`` contains the PCA‑processed image for the corresponding number of components. - - Notes - ----- - The function currently supports the subset of ``vip_hci.psfsub.pca`` parameters used in the HEEPS disc injection - workflow. Additional ``pca`` arguments can be added in the future. + new_mag : float + The closest available magnitude in the Liege grid. """ - # subsample the cube and the reference if requested, for efficiency purposes - if subsample > 1: - cube, angle_list = cube_subsample(array=cube, n=subsample, parallactic=angle_list) - if cube_ref is not None: - cube_ref = cube_subsample(array=cube_ref, n=subsample) - - # loop over principal components - if isinstance(ncomp, int): - ncomp = [ncomp] - elif isinstance(ncomp, tuple): - ncomp = np.arange(ncomp[0], ncomp[-1]+1) - - res = np.zeros([len(ncomp), cube.shape[-2], cube.shape[-1]]) - - print("Running PCA post-processing", flush=True) - for i, npc in enumerate(ncomp): - res[i] = pca(cube, angle_list=angle_list, cube_ref=cube_ref, ncomp=npc, - mask_center_px=mask_center_px, imlib=imlib, nproc=nproc, - source_xy=source_xy, delta_rot=delta_rot, fwhm=fwhm, mask_rdi=mask_rdi, - ref_strategy=ref_strategy) - return res + mag_og = float(mag) + + allowed = np.arange(-1.5, 9.0 + 0.5, step=0.5) + + if mag not in allowed: + mag = round(mag / 0.5) * 0.5 + + # clamp to valid range + if mag < allowed[0]: + mag = allowed[0] + print(f"Warning: mag={mag_og} is below minimum. Using mag={mag}") + elif mag > allowed[-1]: + mag = allowed[-1] + print(f"Warning: mag={mag_og} is above maximum. Using mag={mag}") + elif mag != mag_og: + print(f"Warning: mag={mag_og} rounded to nearest grid value mag={mag}") + + return mag From 466c80dda358d9788c15773bb75c4fbde126fdc7 Mon Sep 17 00:00:00 2001 From: IainHammond Date: Tue, 21 Jul 2026 15:56:16 +0200 Subject: [PATCH 22/46] add path to PSF grid --- heeps/contrast/disc.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index 2d28b36..3c14c8c 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -24,6 +24,7 @@ def create_disc_sequence( disc_model : np.ndarray, + db_path : str, on_axis_psf : Union[np.ndarray, None] = None, off_axis_psf : Union[np.ndarray, None] = None, transmission : Union[np.ndarray, None] = None, @@ -45,6 +46,8 @@ def create_disc_sequence( disc_model : np.ndarray 2‑D array representing the disc model to be injected. NaNs are replaced with zeros. Odd-sized models with the stellar flux on the centre pixel are preferred. + db_path : str + Path to the PSF grid directory (the folder that contains run sub-directories). on_axis_psf : Union[np.ndarray, None], optional Pre‑loaded on‑axis PSF cube. If ``None``, the PSF is loaded from the output directory using the conf parameters. @@ -79,7 +82,7 @@ def create_disc_sequence( # new grid folder structure, assuming your dir_current points to where the folders are saved grid_dir = os.path.join( - conf['dir_current'], + db_path, f"{conf['band']}_{conf['mode']}_s=Q{seeing_q}_mag={conf['mag']}_{conf['duration']}s_{int(conf['dit'] * 100)}ms" ) @@ -236,6 +239,7 @@ def create_disc_sequence( def prepare_rdi_sequence( + db_path : str, on_axis_psf : Union[np.ndarray, None] = None, off_axis_psf : Union[np.ndarray, None] = None, rdi_mag : Union[int, float, None] = None, @@ -250,6 +254,8 @@ def prepare_rdi_sequence( Parameters ---------- + db_path : str + Path to the PSF grid directory (the folder that contains run sub-directories). on_axis_psf : Union[np.ndarray, None], optional Pre‑loaded on‑axis PSF cube. If ``None``, the PSF is loaded based off conf and rdi_mag. using the conf parameters. @@ -277,7 +283,7 @@ def prepare_rdi_sequence( rdi_mag = _check_mag(rdi_mag) grid_dir = os.path.join( - conf['dir_current'], + db_path, f"{conf['band']}_{conf['mode']}_s=Q{seeing_q}_mag={rdi_mag}_{conf['duration']}s_{int(conf['dit'] * 100)}ms" ) From 203dc72f8e7ef4103b62cda43e55c39d3a7516cd Mon Sep 17 00:00:00 2001 From: IainHammond Date: Tue, 21 Jul 2026 16:12:21 +0200 Subject: [PATCH 23/46] fix dit bug --- heeps/contrast/disc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index 3c14c8c..b0918a4 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -83,7 +83,7 @@ def create_disc_sequence( # new grid folder structure, assuming your dir_current points to where the folders are saved grid_dir = os.path.join( db_path, - f"{conf['band']}_{conf['mode']}_s=Q{seeing_q}_mag={conf['mag']}_{conf['duration']}s_{int(conf['dit'] * 100)}ms" + f"{conf['band']}_{conf['mode']}_s=Q{seeing_q}_mag={conf['mag']}_{conf['duration']}s_{int(conf['dit'] * 1000)}ms" ) # convention for the April 2026 grid files From 5867f33f3f89428fef476e6f1f71970470212dad Mon Sep 17 00:00:00 2001 From: IainHammond Date: Tue, 21 Jul 2026 16:13:56 +0200 Subject: [PATCH 24/46] reduced warning to attention --- heeps/contrast/disc.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index b0918a4..9dd611d 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -358,11 +358,11 @@ def _check_mag(mag : Union[float, int] ) -> float: # clamp to valid range if mag < allowed[0]: mag = allowed[0] - print(f"Warning: mag={mag_og} is below minimum. Using mag={mag}") + print(f"Attention: mag={mag_og} is below minimum. Using mag={mag}") elif mag > allowed[-1]: mag = allowed[-1] - print(f"Warning: mag={mag_og} is above maximum. Using mag={mag}") + print(f"Attention: mag={mag_og} is above maximum. Using mag={mag}") elif mag != mag_og: - print(f"Warning: mag={mag_og} rounded to nearest grid value mag={mag}") + print(f"Attention: mag={mag_og} rounded to nearest grid value mag={mag}") return mag From d788d53ad0b0e7c9fa8d0095bef05861b5ade4eb Mon Sep 17 00:00:00 2001 From: IainHammond Date: Tue, 21 Jul 2026 16:20:49 +0200 Subject: [PATCH 25/46] fixed dit bug --- heeps/contrast/disc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index 9dd611d..91d2b88 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -284,7 +284,7 @@ def prepare_rdi_sequence( grid_dir = os.path.join( db_path, - f"{conf['band']}_{conf['mode']}_s=Q{seeing_q}_mag={rdi_mag}_{conf['duration']}s_{int(conf['dit'] * 100)}ms" + f"{conf['band']}_{conf['mode']}_s=Q{seeing_q}_mag={rdi_mag}_{conf['duration']}s_{int(conf['dit'] * 1000)}ms" ) # convention for the April 2026 grid files, assuming dir_current points to the folder containing the PSF files From 5555515e34589a621373341c1760f373f700ea56 Mon Sep 17 00:00:00 2001 From: IainHammond Date: Wed, 22 Jul 2026 09:37:06 +0200 Subject: [PATCH 26/46] making _check_mag always return a float --- heeps/contrast/disc.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index 91d2b88..79359dc 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -334,7 +334,7 @@ def prepare_rdi_sequence( return psf_RDI -def _check_mag(mag : Union[float, int] ) -> float: +def _check_mag(mag : Union[float, int]) -> float: """ Check if the provided magnitude is in the Liege grid. If not, round to the closest available magnitude. @@ -345,7 +345,7 @@ def _check_mag(mag : Union[float, int] ) -> float: Returns ------- - new_mag : float + mag : float The closest available magnitude in the Liege grid. """ mag_og = float(mag) @@ -365,4 +365,4 @@ def _check_mag(mag : Union[float, int] ) -> float: elif mag != mag_og: print(f"Attention: mag={mag_og} rounded to nearest grid value mag={mag}") - return mag + return float(mag) From 75b6a2a955ba86223582841d20f74dfe9a105004 Mon Sep 17 00:00:00 2001 From: IainHammond Date: Wed, 22 Jul 2026 10:10:17 +0200 Subject: [PATCH 27/46] clean up prepare_rdi_sequence --- heeps/contrast/disc.py | 26 +++++++++++++------------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index 79359dc..2df88bb 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -249,8 +249,8 @@ def prepare_rdi_sequence( **conf ): """ - Prepare a sequence to be used as a reference for RDI (Reference Differential Imaging) from the PSF grid. - The function extracts a segment of the on-axis PSF cube to serve as the RDI reference depending on duration. + Prepare a sequence from the PSF grid to be used as a reference for RDI (Reference Differential Imaging). + The function extracts a random segment of the on-axis PSF cube depending on the specified duration. Parameters ---------- @@ -298,34 +298,34 @@ def prepare_rdi_sequence( assert psf_OFF.ndim == 2, "off-axis PSF frame must be 2-dimensional" if on_axis_psf is None: - psf_ON = open_fits(loadname % 'onaxis', verbose=False) + psf_RDI = open_fits(loadname % 'onaxis', verbose=False) print(f"Using on-axis PSF from {loadname % 'onaxis'}") else: - psf_ON = on_axis_psf - assert psf_ON.ndim == 3, "on-axis PSF cube must be 3-dimensional" + psf_RDI = on_axis_psf + assert psf_RDI.ndim == 3, "on-axis PSF cube must be 3-dimensional" if psf_OFF.shape[-1] > conf["ndet"]: psf_OFF = frame_crop(psf_OFF, conf["ndet"], verbose=False) - if psf_ON.shape[-1] > conf["ndet"]: - psf_ON = cube_crop_frames(psf_ON, conf["ndet"], verbose=False) + if psf_RDI.shape[-1] > conf["ndet"]: + psf_RDI = cube_crop_frames(psf_RDI, conf["ndet"], verbose=False) if rdi_duration is None: - rdi_duration = int(psf_ON.shape[0] * 0.2) * conf["dit"] + rdi_duration = int(psf_RDI.shape[0] * 0.2) * conf["dit"] print("Warning: rdi_duration was not set. Using 20% of full sequence.", flush=True) # convert into number of frames to extract from the on-axis PSF cube n_ref_frames = rdi_duration / conf["dit"] # check if the requested reference duration is longer than the available on-axis PSF cube - if n_ref_frames > psf_ON.shape[0]: - n_ref_frames = psf_ON.shape[0] + if n_ref_frames > psf_RDI.shape[0]: + n_ref_frames = psf_RDI.shape[0] print(f"WARNING! Requested RDI duration ({rdi_duration}s) exceeded available on-axis PSF cube length " - f"({psf_ON.shape[0] * conf['dit']}s). Setting them to be the same.", flush=True) + f"({psf_RDI.shape[0] * conf['dit']}s). Setting them to be the same.", flush=True) # random segment of the on-axis PSF cube to use as the RDI reference - start_idx = np.random.randint(0, psf_ON.shape[0] - int(n_ref_frames) + 1) - psf_RDI = psf_ON[start_idx:start_idx + int(n_ref_frames)] + start_idx = np.random.randint(low=0, high=psf_RDI.shape[0] - int(n_ref_frames) + 1) + psf_RDI = psf_RDI[start_idx:start_idx + int(n_ref_frames)] _, _, ap_flux = psf_template(psf_OFF) psf_RDI *= starphot / ap_flux From 1a725fbe3a988d0a983c3229d4bf835d70070507 Mon Sep 17 00:00:00 2001 From: IainHammond Date: Wed, 22 Jul 2026 10:49:19 +0200 Subject: [PATCH 28/46] check if starphot is None --- heeps/contrast/disc.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index 2df88bb..4495c34 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -188,8 +188,10 @@ def create_disc_sequence( psf_ON += cube del cube - _, _, ap_flux = psf_template(psf_OFF) - psf_ON *= starphot / ap_flux + if starphot is not None: + _, _, ap_flux = psf_template(psf_OFF) + psf_ON *= starphot / ap_flux + del psf_OFF # RDI handling (at the moment we only support taking a chunk out of the science sequence # if rdi: @@ -327,8 +329,9 @@ def prepare_rdi_sequence( start_idx = np.random.randint(low=0, high=psf_RDI.shape[0] - int(n_ref_frames) + 1) psf_RDI = psf_RDI[start_idx:start_idx + int(n_ref_frames)] - _, _, ap_flux = psf_template(psf_OFF) - psf_RDI *= starphot / ap_flux + if starphot is not None: + _, _, ap_flux = psf_template(psf_OFF) + psf_RDI *= starphot / ap_flux del psf_OFF return psf_RDI From 144e87ee0309a9bee4a58a7d71d4a4da7641a57d Mon Sep 17 00:00:00 2001 From: IainHammond Date: Wed, 22 Jul 2026 11:36:05 +0200 Subject: [PATCH 29/46] removing old code --- heeps/contrast/disc.py | 44 ------------------------------------------ 1 file changed, 44 deletions(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index 4495c34..9298038 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -193,50 +193,6 @@ def create_disc_sequence( psf_ON *= starphot / ap_flux del psf_OFF - # RDI handling (at the moment we only support taking a chunk out of the science sequence - # if rdi: - # if rdi_duration is None: - # print("Warning: rdi_duration was not set. Using 20% of on-axis sequence.", flush=True) - # rdi_duration = int(0.2 * nframes) * conf["dit"] - # n_ref_frames = rdi_duration / conf["dit"] - # start_idx = np.random.randint(0, psf_ON.shape[0] - int(n_ref_frames) + 1) # random segment of the on-axis PSF cube to use as the RDI reference, to better match the background conditions - # - # if on_axis_psf is None: # reload the on-axis sequence - # psf_RDI = open_fits(loadname % "onaxis", verbose=False)[start_idx:start_idx + int(n_ref_frames)] - # else: - # psf_RDI = on_axis_psf[start_idx:start_idx + int(n_ref_frames)] - # del on_axis_psf - # - # if psf_RDI.shape[-1] > min_crop: - # psf_RDI = cube_crop_frames(psf_RDI, min_crop, verbose=False) - # - # # remove the RDI segment from psf_ON to correctly represent lost integration time on the science target - # psf_ON = np.concatenate([psf_ON[:start_idx], psf_ON[start_idx + int(n_ref_frames):]], axis=0) - # pa = np.concatenate([pa[:start_idx], pa[start_idx + int(n_ref_frames):]], axis=0) - # - # if off_axis_psf is None: - # psf_RDI_OFF = open_fits(loadname % "offaxis", verbose=False) - # else: - # psf_RDI_OFF = off_axis_psf - # - # if psf_RDI_OFF.shape[-1] > min_crop: - # psf_RDI_OFF = frame_crop(psf_RDI_OFF, min_crop, verbose=False) - # - # if rdi_mag is None: - # print("Warning: rdi_mag was not set. Using mag of science target.", flush=True) - # else: - # conf["mag"] = rdi_mag - # - # if conf["add_bckg"]: - # psf_RDI, psf_RDI_OFF = background(psf_RDI, psf_RDI_OFF, verbose=True, **conf) - # - # _, _, ap_flux = psf_template(psf_RDI_OFF) - # psf_RDI *= starphot / ap_flux - # del psf_RDI_OFF - # - # return psf_ON, pa, psf_RDI - # - # else: return psf_ON, pa From e267f447be2c3d9c88712aab4fc83195b9d84f32 Mon Sep 17 00:00:00 2001 From: IainHammond Date: Wed, 22 Jul 2026 11:52:40 +0200 Subject: [PATCH 30/46] safer handling of the PSF grid path --- heeps/contrast/disc.py | 51 +++++++++++++++++++++++------------------- 1 file changed, 28 insertions(+), 23 deletions(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index 9298038..9677ff3 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -24,7 +24,7 @@ def create_disc_sequence( disc_model : np.ndarray, - db_path : str, + db_path : str = "vortex_psf_grid/", on_axis_psf : Union[np.ndarray, None] = None, off_axis_psf : Union[np.ndarray, None] = None, transmission : Union[np.ndarray, None] = None, @@ -36,10 +36,9 @@ def create_disc_sequence( **conf ): """ - Create a mock observation sequence by injecting a synthetic disc model into a HEEPS PSF and optionally generate - a reference differential imaging (RDI) cube. Existing on- and off-axis PSF cubes and coronagraph transmission - can be provided. If not provided, the function will attempt to load them from the output directory based on the - conf parameters. + Create a mock observation sequence by injecting a synthetic disc model into a HEEPS PSF. Existing on- and off-axis + PSF cubes and coronagraph transmission can be provided. If not provided, the function will attempt to load them + from the PSF grid directory based on the conf parameters. Parameters ---------- @@ -47,7 +46,8 @@ def create_disc_sequence( 2‑D array representing the disc model to be injected. NaNs are replaced with zeros. Odd-sized models with the stellar flux on the centre pixel are preferred. db_path : str - Path to the PSF grid directory (the folder that contains run sub-directories). + Path to the PSF grid directory (the folder that contains run sub-directories). Required if on_axis_psf or + off_axis_psf are not provided. on_axis_psf : Union[np.ndarray, None], optional Pre‑loaded on‑axis PSF cube. If ``None``, the PSF is loaded from the output directory using the conf parameters. @@ -80,21 +80,18 @@ def create_disc_sequence( # if mag is not in the grid, round to the closest available magnitude conf["mag"] = _check_mag(conf["mag"]) - # new grid folder structure, assuming your dir_current points to where the folders are saved - grid_dir = os.path.join( - db_path, - f"{conf['band']}_{conf['mode']}_s=Q{seeing_q}_mag={conf['mag']}_{conf['duration']}s_{int(conf['dit'] * 1000)}ms" - ) - - # convention for the April 2026 grid files - loadname = os.path.join(grid_dir, '%s_PSF_bckg1_%s_%s.fits' % ('%s', conf['band'], conf['mode'])) - - if off_axis_psf is None: - psf_OFF = open_fits(loadname % 'offaxis', verbose=False) - print(f"Using off-axis PSF from {loadname%'offaxis'}") - else: - psf_OFF = off_axis_psf - assert psf_OFF.ndim == 2, "off-axis PSF frame must be 2-dimensional" + # if one or both of the PSFs are not provided, prepare the path for loading them from the PSF grid directory + if on_axis_psf is None or off_axis_psf is None: + # make sure db_path ends with a slash + if not db_path.endswith("/"): + db_path += "/" + + # new grid folder structure and naming convention for the April 2026 grid files + grid_dir = os.path.join( + db_path, + f"{conf['band']}_{conf['mode']}_s=Q{seeing_q}_mag={conf['mag']}_{conf['duration']}s_{int(conf['dit'] * 1000)}ms" + ) + loadname = os.path.join(grid_dir, '%s_PSF_bckg1_%s_%s.fits' % ('%s', conf['band'], conf['mode'])) # set nans to zero disc_model = np.nan_to_num(disc_model, nan=0) @@ -118,6 +115,13 @@ def create_disc_sequence( disc_model = frame_shift(disc_model, shift_y=0.5, shift_x=0.5, imlib=imlib) disc_model = disc_model[1:, 1:] + if off_axis_psf is None: + psf_OFF = open_fits(loadname % 'offaxis', verbose=False) + print(f"Using off-axis PSF from {loadname%'offaxis'}") + else: + psf_OFF = off_axis_psf + assert psf_OFF.ndim == 2, "off-axis PSF frame must be 2-dimensional" + # crop everything to a common size min_crop = min(psf_OFF.shape[-1], disc_model.shape[-1], conf["ndet"]) conf["ndet"] = min_crop @@ -197,7 +201,7 @@ def create_disc_sequence( def prepare_rdi_sequence( - db_path : str, + db_path : str = "vortex_psf_grid/", on_axis_psf : Union[np.ndarray, None] = None, off_axis_psf : Union[np.ndarray, None] = None, rdi_mag : Union[int, float, None] = None, @@ -213,7 +217,8 @@ def prepare_rdi_sequence( Parameters ---------- db_path : str - Path to the PSF grid directory (the folder that contains run sub-directories). + Path to the PSF grid directory (the folder that contains run sub-directories). Required if on_axis_psf or + off_axis_psf are not provided. on_axis_psf : Union[np.ndarray, None], optional Pre‑loaded on‑axis PSF cube. If ``None``, the PSF is loaded based off conf and rdi_mag. using the conf parameters. From f65ac9e9be0f411674b3c0f9683da05bbd39afc5 Mon Sep 17 00:00:00 2001 From: IainHammond Date: Wed, 22 Jul 2026 11:58:32 +0200 Subject: [PATCH 31/46] docstring clarification --- heeps/contrast/disc.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index 9677ff3..a2d7841 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -40,6 +40,8 @@ def create_disc_sequence( PSF cubes and coronagraph transmission can be provided. If not provided, the function will attempt to load them from the PSF grid directory based on the conf parameters. + The code assumes that the disc model has the stellar flux included to convert to units of contrast. + Parameters ---------- disc_model : np.ndarray @@ -85,7 +87,7 @@ def create_disc_sequence( # make sure db_path ends with a slash if not db_path.endswith("/"): db_path += "/" - + # new grid folder structure and naming convention for the April 2026 grid files grid_dir = os.path.join( db_path, From e255a765ed48d9f5b679db0b521943beeecc5789 Mon Sep 17 00:00:00 2001 From: IainHammond Date: Wed, 22 Jul 2026 12:01:02 +0200 Subject: [PATCH 32/46] docstring update --- heeps/contrast/disc.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index a2d7841..3672f99 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -1,11 +1,14 @@ """ -HEEPS disc injection and post‑processing utilities. +HEEPS disc injection utilities. This module provides functions to inject a synthetic disc model into HEEPS -PSF data and to perform post‑processing with VIP. +PSF cubes for later post‑processing with VIP. + The entry point is :func:`create_disc_sequence`, which creates a mock -observation sequence with a user-provided disc model. The :func:`postproc_disc` function is a -wrapper of the relevant VIP function for PCA post-processing. +observation sequence with a user-provided disc model. + +The :func:`prepare_rdi_sequence` function prepares a reference sequence for +RDI (Reference Differential Imaging) also from the PSF grid. """ import numpy as np From c4518341f4782b31672b31640486f9b707d6ddf8 Mon Sep 17 00:00:00 2001 From: IainHammond Date: Wed, 22 Jul 2026 12:04:06 +0200 Subject: [PATCH 33/46] docstring update --- heeps/contrast/disc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index 3672f99..231d2d7 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -73,7 +73,7 @@ def create_disc_sequence( Image library to use for image processing. Use "opencv" to go fast, or "vip-fft" for better flux conservation at the cost of speed. See VIP documentation for details. conf : dict, optional - Configuration dictionary with parameters from HEEPS (e.g., background addition etc.). + Configuration dictionary with parameters from HEEPS (e.g., band, dit, mag, etc.). Returns ------- From 503072ca26bf8f03f9ac0e39ad4afa03547183c0 Mon Sep 17 00:00:00 2001 From: IainHammond Date: Wed, 22 Jul 2026 12:06:06 +0200 Subject: [PATCH 34/46] another cleanup --- heeps/contrast/disc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index 231d2d7..36383f6 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -178,7 +178,7 @@ def create_disc_sequence( ker=1, imlib=imlib, ) - # write_fits(conf["dir_current"] + "/prepared_model.fits", cube[0]) + # normalise by the stellar flux that we removed earlier cube /= star_val From 16738caa46f1805fbf5701c82c5f6782db8bb771 Mon Sep 17 00:00:00 2001 From: IainHammond Date: Wed, 22 Jul 2026 14:13:11 +0200 Subject: [PATCH 35/46] debugging --- heeps/contrast/disc.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index 36383f6..d24d8ad 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -112,6 +112,8 @@ def create_disc_sequence( # some codes have the star on the centre 4 pixels instead of the centre pixel only so we need to sum the values mask = disc_model == disc_model.max() star_val = disc_model[mask].sum() + print("star_val = %.2e" % star_val, flush=True) + print("mask = %s" % str(mask), flush=True) disc_model[mask] = 0 # ensure the model has an odd size From abdab14b76c2a65643430222b60a72dd09fe3ca9 Mon Sep 17 00:00:00 2001 From: IainHammond Date: Wed, 22 Jul 2026 14:31:27 +0200 Subject: [PATCH 36/46] removing debugging --- heeps/contrast/disc.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index d24d8ad..36383f6 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -112,8 +112,6 @@ def create_disc_sequence( # some codes have the star on the centre 4 pixels instead of the centre pixel only so we need to sum the values mask = disc_model == disc_model.max() star_val = disc_model[mask].sum() - print("star_val = %.2e" % star_val, flush=True) - print("mask = %s" % str(mask), flush=True) disc_model[mask] = 0 # ensure the model has an odd size From f847e1f5abb7c964c8449d67d03b2e704c72c47b Mon Sep 17 00:00:00 2001 From: IainHammond Date: Wed, 22 Jul 2026 17:40:12 +0200 Subject: [PATCH 37/46] major restructure --- heeps/contrast/disc.py | 172 +++++++++++++++++++++++++---------------- 1 file changed, 107 insertions(+), 65 deletions(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index 36383f6..00cf153 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -22,26 +22,25 @@ from heeps.util.psf_template import psf_template from heeps.util.paralang import paralang -__all__ = ['create_disc_sequence', 'prepare_rdi_sequence'] +__all__ = ['create_disc_sequence'] __author__ = "Iain Hammond" def create_disc_sequence( disc_model : np.ndarray, db_path : str = "vortex_psf_grid/", - on_axis_psf : Union[np.ndarray, None] = None, - off_axis_psf : Union[np.ndarray, None] = None, - transmission : Union[np.ndarray, None] = None, seeing_q : Union[int] = 2, source_xy : Union[tuple, list, None] = None, extinction : Union[int, float] = 0, + do_rdi : bool = True, + rdi_mag : Union[int, float, None] = None, + rdi_duration : Union[int, float, None] = None, starphot : Union[float, int] = 1e11, imlib : str = "opencv", **conf ): """ - Create a mock observation sequence by injecting a synthetic disc model into a HEEPS PSF. Existing on- and off-axis - PSF cubes and coronagraph transmission can be provided. If not provided, the function will attempt to load them - from the PSF grid directory based on the conf parameters. + Create a mock observation sequence by injecting a synthetic disc model into a HEEPS PSF. The function will attempt + to load from the PSF grid directory based on the conf parameters. The code assumes that the disc model has the stellar flux included to convert to units of contrast. @@ -51,16 +50,7 @@ def create_disc_sequence( 2‑D array representing the disc model to be injected. NaNs are replaced with zeros. Odd-sized models with the stellar flux on the centre pixel are preferred. db_path : str - Path to the PSF grid directory (the folder that contains run sub-directories). Required if on_axis_psf or - off_axis_psf are not provided. - on_axis_psf : Union[np.ndarray, None], optional - Pre‑loaded on‑axis PSF cube. If ``None``, the PSF is loaded from the output directory - using the conf parameters. - off_axis_psf : Union[np.ndarray, None], optional - Pre‑loaded off‑axis PSF frame. If ``None``, the PSF is loaded from the output directory. - transmission : Union[np.ndarray, None], optional - Coronagraph transmission. If ``None``, the appropriate transmission is loaded based - on the instrument mode in conf. Uses VIP conventions. + Path to the PSF grid directory (the folder that contains run sub-directories). seeing_q : int, default=2 Seeing quartile (1-3) to select the PSF. Uses a default seeing of Q2 if not specified. source_xy : Union[tuple, list, None], optional @@ -85,18 +75,26 @@ def create_disc_sequence( # if mag is not in the grid, round to the closest available magnitude conf["mag"] = _check_mag(conf["mag"]) - # if one or both of the PSFs are not provided, prepare the path for loading them from the PSF grid directory - if on_axis_psf is None or off_axis_psf is None: - # make sure db_path ends with a slash - if not db_path.endswith("/"): - db_path += "/" + if rdi_mag is None: + rdi_mag = conf["mag"] + print("Note: rdi_mag was not set. Using mag in the conf dictionary.", flush=True) - # new grid folder structure and naming convention for the April 2026 grid files - grid_dir = os.path.join( - db_path, - f"{conf['band']}_{conf['mode']}_s=Q{seeing_q}_mag={conf['mag']}_{conf['duration']}s_{int(conf['dit'] * 1000)}ms" - ) - loadname = os.path.join(grid_dir, '%s_PSF_bckg1_%s_%s.fits' % ('%s', conf['band'], conf['mode'])) + rdi_mag = _check_mag(rdi_mag) + + # would the RDI reference come from the exact same grid cube as the science sequence? + reuse_science_cube = do_rdi and (rdi_mag == conf["mag"]) + + # prepare the path for loading the PSF grid directory + # make sure db_path ends with a slash + if not db_path.endswith("/"): + db_path += "/" + + # new grid folder structure and naming convention for the April 2026 grid files + grid_dir = os.path.join( + db_path, + f"{conf['band']}_{conf['mode']}_s=Q{seeing_q}_mag={conf['mag']}_{conf['duration']}s_{int(conf['dit'] * 1000)}ms" + ) + loadname = os.path.join(grid_dir, '%s_PSF_bckg1_%s_%s.fits' % ('%s', conf['band'], conf['mode'])) # set nans to zero disc_model = np.nan_to_num(disc_model, nan=0) @@ -120,11 +118,8 @@ def create_disc_sequence( disc_model = frame_shift(disc_model, shift_y=0.5, shift_x=0.5, imlib=imlib) disc_model = disc_model[1:, 1:] - if off_axis_psf is None: - psf_OFF = open_fits(loadname % 'offaxis', verbose=False) - print(f"Using off-axis PSF from {loadname%'offaxis'}") - else: - psf_OFF = off_axis_psf + psf_OFF = open_fits(loadname % 'offaxis', verbose=False) + print(f"Using off-axis PSF from {loadname%'offaxis'}") assert psf_OFF.ndim == 2, "off-axis PSF frame must be 2-dimensional" # crop everything to a common size @@ -143,28 +138,83 @@ def create_disc_sequence( print("Applied extinction of %.2f mag to source at (x, y) = (%d, %d)" % (extinction, source_x, source_y), flush=True) # load coronagraph transmission depending on the mode and band - if transmission is None: - if conf["mode"] == "ELT" or conf["mode"] is None: # no coronagraph - transmission = None - elif "VC" in conf["mode"]: - if conf["band"] in ("L", "M"): - conf['f_oat'] = conf['dir_input'] + 'optics/vc/oat_L_%s.fits' % conf['mode'] - elif conf["band"] == "N2": - conf['f_oat'] = conf['dir_input'] + 'optics/vc/oat_N2_CVC.fits' - transmission = open_fits(conf['f_oat'], verbose=False) - print(f"Using transmission from {conf['f_oat']}") - else: # TODO - raise NotImplementedError(f"Mode {conf['mode']} not implemented for disc injection.") - - # determine how many frames are in the on-axis sequence - if on_axis_psf is None: - nframes = open_header(loadname % 'onaxis')['NAXIS3'] # don't open the whole cube yet to save memory - else: - nframes = on_axis_psf.shape[0] + if conf["mode"] == "ELT" or conf["mode"] is None: # no coronagraph + transmission = None + elif "VC" in conf["mode"]: + if conf["band"] in ("L", "M"): + conf['f_oat'] = conf['dir_input'] + 'optics/vc/oat_L_%s.fits' % conf['mode'] + elif conf["band"] == "N2": + conf['f_oat'] = conf['dir_input'] + 'optics/vc/oat_N2_CVC.fits' + transmission = open_fits(conf['f_oat'], verbose=False) + print(f"Using transmission from {conf['f_oat']}") + else: # TODO + raise NotImplementedError(f"Mode {conf['mode']} not implemented for disc injection.") + + # open and crop sequence + psf_ON = open_fits(loadname % 'onaxis', verbose=False) + print(f"Using on-axis PSF from {loadname % 'onaxis'}") + assert psf_ON.ndim == 3, "on-axis PSF cube must be 3-dimensional" + + if psf_ON.shape[-1] > min_crop: + psf_ON = cube_crop_frames(psf_ON, min_crop, verbose=False) + + nframes = psf_ON.shape[0] # generate parallactic angles and inject the fake disc using VIP pa = paralang(npts=nframes, dec=conf["dec"], lat=conf["lat"], duration=int(nframes * conf["dit"])) + # RDI handling + if do_rdi: + if rdi_duration is None: + rdi_duration = int(nframes * conf["dit"] * 0.2) + print("Note: rdi_duration was not set. Using 20% of the PSF sequence.", flush=True) + + n_ref_frames = int(round(rdi_duration / conf["dit"])) + + if reuse_science_cube: + # always take the RDI block from the end of the sequence, leaving the science + # sequence as one contiguous block from the start + rdi_idx = slice(nframes - n_ref_frames, nframes) + sci_idx = slice(0, nframes - n_ref_frames) + + psf_RDI = psf_ON[rdi_idx].copy() + psf_ON = psf_ON[sci_idx] + psf_OFF_rdi = psf_OFF.copy() + pa = pa[sci_idx] + + print(f"RDI reference (mag={rdi_mag}) reused from the science cube: " + f"{n_ref_frames} frames removed from the end of the science sequence, " + f"no overlap with science frames.", flush=True) + + else: # we load in a different PSF cube for the RDI reference and extract the requested duration + rdi_grid_dir = os.path.join( + db_path, + f"{conf['band']}_{conf['mode']}_s=Q{seeing_q}_mag={rdi_mag}_{conf['duration']}s_{int(conf['dit'] * 1000)}ms" + ) + rdi_loadname = os.path.join(rdi_grid_dir, '%s_PSF_bckg1_%s_%s.fits' % ('%s', conf['band'], conf['mode'])) + + psf_OFF_rdi = open_fits(rdi_loadname % 'offaxis', verbose=False) + print(f"Using RDI off-axis PSF from {rdi_loadname % 'offaxis'}") + assert psf_OFF_rdi.ndim == 2, "RDI off-axis PSF frame must be 2-dimensional" + + psf_RDI = open_fits(rdi_loadname % 'onaxis', verbose=False) + print(f"Using RDI on-axis PSF from {rdi_loadname % 'onaxis'}") + assert psf_RDI.ndim == 3, "RDI on-axis PSF cube must be 3-dimensional" + + if psf_OFF_rdi.shape[-1] > min_crop: + psf_OFF_rdi = frame_crop(psf_OFF_rdi, min_crop, verbose=False) + if psf_RDI.shape[-1] > min_crop: + psf_RDI = cube_crop_frames(psf_RDI, min_crop, verbose=False) + + start_idx = np.random.randint(low=0, high=psf_RDI.shape[0] - n_ref_frames + 1) + psf_RDI = psf_RDI[start_idx:start_idx + n_ref_frames] + + if starphot is not None: + _, _, ap_flux_rdi = psf_template(psf_OFF_rdi) + psf_RDI *= starphot / ap_flux_rdi + del psf_OFF_rdi + + # where the magic happens cube = cube_inject_fakedisk( disc_model, pa, @@ -179,20 +229,9 @@ def create_disc_sequence( imlib=imlib, ) - # normalise by the stellar flux that we removed earlier + # normalise by the stellar flux that we removed earlier (units of contrast) cube /= star_val - # open and crop sequence if needed - if on_axis_psf is None: - psf_ON = open_fits(loadname % 'onaxis', verbose=False) - print(f"Using on-axis PSF from {loadname % 'onaxis'}") - else: - psf_ON = on_axis_psf - assert psf_ON.ndim == 3, "on-axis PSF cube must be 3-dimensional" - - if psf_ON.shape[-1] > min_crop: - psf_ON = cube_crop_frames(psf_ON, min_crop, verbose=False) - # add the disc model psf_ON += cube del cube @@ -202,10 +241,13 @@ def create_disc_sequence( psf_ON *= starphot / ap_flux del psf_OFF - return psf_ON, pa + if do_rdi: + return psf_ON, pa, psf_RDI + else: + return psf_ON, pa -def prepare_rdi_sequence( +def _prepare_rdi_sequence( db_path : str = "vortex_psf_grid/", on_axis_psf : Union[np.ndarray, None] = None, off_axis_psf : Union[np.ndarray, None] = None, From 4cc74996e122704a8f81d65b0c58b5f23680a7dd Mon Sep 17 00:00:00 2001 From: IainHammond Date: Wed, 22 Jul 2026 19:24:15 +0200 Subject: [PATCH 38/46] cleanup --- heeps/contrast/disc.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index 00cf153..a473cf5 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -75,15 +75,10 @@ def create_disc_sequence( # if mag is not in the grid, round to the closest available magnitude conf["mag"] = _check_mag(conf["mag"]) - if rdi_mag is None: + if do_rdi and rdi_mag is None: rdi_mag = conf["mag"] print("Note: rdi_mag was not set. Using mag in the conf dictionary.", flush=True) - rdi_mag = _check_mag(rdi_mag) - - # would the RDI reference come from the exact same grid cube as the science sequence? - reuse_science_cube = do_rdi and (rdi_mag == conf["mag"]) - # prepare the path for loading the PSF grid directory # make sure db_path ends with a slash if not db_path.endswith("/"): @@ -171,7 +166,7 @@ def create_disc_sequence( n_ref_frames = int(round(rdi_duration / conf["dit"])) - if reuse_science_cube: + if rdi_mag == conf["mag"]: # always take the RDI block from the end of the sequence, leaving the science # sequence as one contiguous block from the start rdi_idx = slice(nframes - n_ref_frames, nframes) From 5513159b0c663014aadec04acb440d2b1ffc0d70 Mon Sep 17 00:00:00 2001 From: IainHammond Date: Wed, 22 Jul 2026 19:59:46 +0200 Subject: [PATCH 39/46] cleanup --- heeps/contrast/disc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index a473cf5..498e6a1 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -155,7 +155,7 @@ def create_disc_sequence( nframes = psf_ON.shape[0] - # generate parallactic angles and inject the fake disc using VIP + # generate parallactic angles pa = paralang(npts=nframes, dec=conf["dec"], lat=conf["lat"], duration=int(nframes * conf["dit"])) # RDI handling From 0653a9ddb4069cc5e62a51da6ebba06f0e637d72 Mon Sep 17 00:00:00 2001 From: IainHammond Date: Wed, 22 Jul 2026 20:31:11 +0200 Subject: [PATCH 40/46] code comments --- heeps/contrast/disc.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index 498e6a1..1f16c4b 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -15,7 +15,7 @@ from typing import Union import os -from vip_hci.fits import open_fits, open_header +from vip_hci.fits import open_fits from vip_hci.preproc import frame_crop, cube_crop_frames, frame_shift from vip_hci.fm import cube_inject_fakedisk @@ -85,6 +85,7 @@ def create_disc_sequence( db_path += "/" # new grid folder structure and naming convention for the April 2026 grid files + # for example, vortex_psf_grid/L_CVC_s=Q2_mag=8.5_3600s_100ms grid_dir = os.path.join( db_path, f"{conf['band']}_{conf['mode']}_s=Q{seeing_q}_mag={conf['mag']}_{conf['duration']}s_{int(conf['dit'] * 1000)}ms" From 6adaca00f2c70547b104db091e467451337ae268 Mon Sep 17 00:00:00 2001 From: IainHammond Date: Thu, 23 Jul 2026 09:23:44 +0200 Subject: [PATCH 41/46] update docstring, final cleanup --- heeps/contrast/disc.py | 121 ++++++----------------------------------- 1 file changed, 16 insertions(+), 105 deletions(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index 1f16c4b..59e1d8f 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -7,8 +7,6 @@ The entry point is :func:`create_disc_sequence`, which creates a mock observation sequence with a user-provided disc model. -The :func:`prepare_rdi_sequence` function prepares a reference sequence for -RDI (Reference Differential Imaging) also from the PSF grid. """ import numpy as np @@ -28,13 +26,13 @@ def create_disc_sequence( disc_model : np.ndarray, db_path : str = "vortex_psf_grid/", - seeing_q : Union[int] = 2, + seeing_q : int = 2, source_xy : Union[tuple, list, None] = None, extinction : Union[int, float] = 0, + starphot : Union[float, int, None] = 1e11, do_rdi : bool = True, rdi_mag : Union[int, float, None] = None, rdi_duration : Union[int, float, None] = None, - starphot : Union[float, int] = 1e11, imlib : str = "opencv", **conf ): @@ -42,7 +40,7 @@ def create_disc_sequence( Create a mock observation sequence by injecting a synthetic disc model into a HEEPS PSF. The function will attempt to load from the PSF grid directory based on the conf parameters. - The code assumes that the disc model has the stellar flux included to convert to units of contrast. + The code assumes that the disc model has the star included to convert to units of contrast. Parameters ---------- @@ -59,6 +57,15 @@ def create_disc_sequence( Extinction in magnitudes to apply to the source at ``source_xy``. starphot : float, default=1e11 Photometric scaling factor for the star. + do_rdi : bool, default=True + Whether to prepare a reference cube for RDI (Reference Differential Imaging). If ``True``, a reference cube is + returned as the third output of this function. + rdi_mag : Union[int, float, None], optional + Magnitude to use for the RDI reference. If ``None``, magnitude in conf is used (i.e., the same as the science + target). + rdi_duration : Union[int, float, None], optional + Desired duration (in seconds) of the RDI reference sequence. If ``None``, a default of 20% of the on‑axis + sequence length is used. imlib : str, default="opencv" Image library to use for image processing. Use "opencv" to go fast, or "vip-fft" for better flux conservation at the cost of speed. See VIP documentation for details. @@ -71,6 +78,8 @@ def create_disc_sequence( Sequence with the fake disc injected into the on‑axis PSF cube. pa : numpy.ndarray Array of parallactic angles used for the injection. + psf_RDI : numpy.ndarray, optional + Reference cube for RDI, if ``do_rdi`` is ``True``. Otherwise, this output is not returned. """ # if mag is not in the grid, round to the closest available magnitude conf["mag"] = _check_mag(conf["mag"]) @@ -84,7 +93,7 @@ def create_disc_sequence( if not db_path.endswith("/"): db_path += "/" - # new grid folder structure and naming convention for the April 2026 grid files + # new grid folder structure and naming convention for the April 2026 grid files, # for example, vortex_psf_grid/L_CVC_s=Q2_mag=8.5_3600s_100ms grid_dir = os.path.join( db_path, @@ -243,104 +252,6 @@ def create_disc_sequence( return psf_ON, pa -def _prepare_rdi_sequence( - db_path : str = "vortex_psf_grid/", - on_axis_psf : Union[np.ndarray, None] = None, - off_axis_psf : Union[np.ndarray, None] = None, - rdi_mag : Union[int, float, None] = None, - rdi_duration : Union[int, float, None] = None, - seeing_q : Union[int] = 2, - starphot : Union[float, int] = 1e11, - **conf -): - """ - Prepare a sequence from the PSF grid to be used as a reference for RDI (Reference Differential Imaging). - The function extracts a random segment of the on-axis PSF cube depending on the specified duration. - - Parameters - ---------- - db_path : str - Path to the PSF grid directory (the folder that contains run sub-directories). Required if on_axis_psf or - off_axis_psf are not provided. - on_axis_psf : Union[np.ndarray, None], optional - Pre‑loaded on‑axis PSF cube. If ``None``, the PSF is loaded based off conf and rdi_mag. - using the conf parameters. - off_axis_psf : Union[np.ndarray, None], optional - Pre‑loaded off‑axis PSF frame. If ``None``, the PSF is loaded based off conf and rdi_mag. - rdi_mag : Union[int, float, None], optional - Magnitude to use for the RDI reference. If ``None``, magnitude in conf is used. - rdi_duration : Union[int, float, None], optional - Desired duration (in seconds) of the RDI reference sequence. If ``None``, a default of - 20% of the on‑axis sequence length is used. - seeing_q : int, default=2 - Seeing quartile (1-3) to select the PSF. Uses a default seeing of Q2 if not specified. - starphot : float, default=1e11 - Photometric scaling factor for the star. - - Returns - ------- - psf_RDI : numpy.ndarray - Reference cube for RDI. - """ - if rdi_mag is None: - rdi_mag = conf["mag"] - print("Warning: rdi_mag was not set. Using mag in the conf dictionary.", flush=True) - - rdi_mag = _check_mag(rdi_mag) - - grid_dir = os.path.join( - db_path, - f"{conf['band']}_{conf['mode']}_s=Q{seeing_q}_mag={rdi_mag}_{conf['duration']}s_{int(conf['dit'] * 1000)}ms" - ) - - # convention for the April 2026 grid files, assuming dir_current points to the folder containing the PSF files - loadname = os.path.join(grid_dir, '%s_PSF_bckg1_%s_%s.fits' % ('%s', conf['band'], conf['mode'])) - - if off_axis_psf is None: - psf_OFF = open_fits(loadname % 'offaxis', verbose=False) - print(f"Using off-axis PSF from {loadname%'offaxis'}") - else: - psf_OFF = off_axis_psf - assert psf_OFF.ndim == 2, "off-axis PSF frame must be 2-dimensional" - - if on_axis_psf is None: - psf_RDI = open_fits(loadname % 'onaxis', verbose=False) - print(f"Using on-axis PSF from {loadname % 'onaxis'}") - else: - psf_RDI = on_axis_psf - assert psf_RDI.ndim == 3, "on-axis PSF cube must be 3-dimensional" - - if psf_OFF.shape[-1] > conf["ndet"]: - psf_OFF = frame_crop(psf_OFF, conf["ndet"], verbose=False) - - if psf_RDI.shape[-1] > conf["ndet"]: - psf_RDI = cube_crop_frames(psf_RDI, conf["ndet"], verbose=False) - - if rdi_duration is None: - rdi_duration = int(psf_RDI.shape[0] * 0.2) * conf["dit"] - print("Warning: rdi_duration was not set. Using 20% of full sequence.", flush=True) - - # convert into number of frames to extract from the on-axis PSF cube - n_ref_frames = rdi_duration / conf["dit"] - - # check if the requested reference duration is longer than the available on-axis PSF cube - if n_ref_frames > psf_RDI.shape[0]: - n_ref_frames = psf_RDI.shape[0] - print(f"WARNING! Requested RDI duration ({rdi_duration}s) exceeded available on-axis PSF cube length " - f"({psf_RDI.shape[0] * conf['dit']}s). Setting them to be the same.", flush=True) - - # random segment of the on-axis PSF cube to use as the RDI reference - start_idx = np.random.randint(low=0, high=psf_RDI.shape[0] - int(n_ref_frames) + 1) - psf_RDI = psf_RDI[start_idx:start_idx + int(n_ref_frames)] - - if starphot is not None: - _, _, ap_flux = psf_template(psf_OFF) - psf_RDI *= starphot / ap_flux - del psf_OFF - - return psf_RDI - - def _check_mag(mag : Union[float, int]) -> float: """ Check if the provided magnitude is in the Liege grid. If not, round to the closest available magnitude. @@ -353,7 +264,7 @@ def _check_mag(mag : Union[float, int]) -> float: Returns ------- mag : float - The closest available magnitude in the Liege grid. + The closest available magnitude in the Liege grid as a float. """ mag_og = float(mag) From f032f1143f2e717e666e3f29a351d4667416274c Mon Sep 17 00:00:00 2001 From: IainHammond Date: Thu, 23 Jul 2026 09:24:05 +0200 Subject: [PATCH 42/46] update docstring, final cleanup --- heeps/contrast/disc.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index 59e1d8f..cb094f0 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -48,7 +48,7 @@ def create_disc_sequence( 2‑D array representing the disc model to be injected. NaNs are replaced with zeros. Odd-sized models with the stellar flux on the centre pixel are preferred. db_path : str - Path to the PSF grid directory (the folder that contains run sub-directories). + Path to the PSF grid directory (the folder that contains run subdirectories). seeing_q : int, default=2 Seeing quartile (1-3) to select the PSF. Uses a default seeing of Q2 if not specified. source_xy : Union[tuple, list, None], optional From 37f7b116ce6d93b2024d4a7b160aa9dcfba15ec0 Mon Sep 17 00:00:00 2001 From: IainHammond Date: Thu, 23 Jul 2026 09:45:34 +0200 Subject: [PATCH 43/46] update docstring, final cleanup --- heeps/contrast/disc.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index cb094f0..55801e1 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -254,7 +254,8 @@ def create_disc_sequence( def _check_mag(mag : Union[float, int]) -> float: """ - Check if the provided magnitude is in the Liege grid. If not, round to the closest available magnitude. + Check if the provided magnitude is in the Liege grid. If not, round to the closest available magnitude. A float + must be returned as the grid contains mag as a float in the filenames. Parameters ---------- @@ -264,7 +265,7 @@ def _check_mag(mag : Union[float, int]) -> float: Returns ------- mag : float - The closest available magnitude in the Liege grid as a float. + The closest available magnitude in the Liege grid. """ mag_og = float(mag) From e7ee2631decd404bd1fece06696a96eeb2044d20 Mon Sep 17 00:00:00 2001 From: IainHammond Date: Fri, 24 Jul 2026 13:48:57 +0200 Subject: [PATCH 44/46] support for M, N1, and N2 --- heeps/contrast/disc.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index 55801e1..1d2d050 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -252,7 +252,7 @@ def create_disc_sequence( return psf_ON, pa -def _check_mag(mag : Union[float, int]) -> float: +def _check_mag(mag : Union[float, int], band : str) -> float: """ Check if the provided magnitude is in the Liege grid. If not, round to the closest available magnitude. A float must be returned as the grid contains mag as a float in the filenames. @@ -261,6 +261,8 @@ def _check_mag(mag : Union[float, int]) -> float: ---------- mag : float, int The magnitude to check. + band : str + The band of the observation (e.g., "L", "M", "N1", "N2".) Returns ------- @@ -269,7 +271,16 @@ def _check_mag(mag : Union[float, int]) -> float: """ mag_og = float(mag) - allowed = np.arange(-1.5, 9.0 + 0.5, step=0.5) + if band == "L": + allowed = np.arange(-1.5, 9.0 + 0.5, step=0.5) + elif band == "M": + allowed = np.arange(-1.5, 7.0 + 0.5, step=0.5) + elif band == "N1": + allowed = np.arange(-1.5, 4.0 + 0.5, step=0.5) + elif band == "N2": + allowed = np.arange(-1.5, 3.0 + 0.5, step=0.5) + else: + raise ValueError(f"Band {band} is not supported. Please use 'L', 'M', 'N1', or 'N2'.") if mag not in allowed: mag = round(mag / 0.5) * 0.5 From 4bb20e510d8bf921b69aeb033d59ea1c32a769e5 Mon Sep 17 00:00:00 2001 From: IainHammond Date: Fri, 24 Jul 2026 14:25:58 +0200 Subject: [PATCH 45/46] support for M, N1, and N2 --- heeps/contrast/disc.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index 1d2d050..9c1b487 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -82,7 +82,7 @@ def create_disc_sequence( Reference cube for RDI, if ``do_rdi`` is ``True``. Otherwise, this output is not returned. """ # if mag is not in the grid, round to the closest available magnitude - conf["mag"] = _check_mag(conf["mag"]) + conf["mag"] = _check_mag(conf["mag"], conf["band"]) if do_rdi and rdi_mag is None: rdi_mag = conf["mag"] @@ -288,10 +288,10 @@ def _check_mag(mag : Union[float, int], band : str) -> float: # clamp to valid range if mag < allowed[0]: mag = allowed[0] - print(f"Attention: mag={mag_og} is below minimum. Using mag={mag}") + print(f"Attention: mag={mag_og} is below minimum for {band}. Using mag={mag}") elif mag > allowed[-1]: mag = allowed[-1] - print(f"Attention: mag={mag_og} is above maximum. Using mag={mag}") + print(f"Attention: mag={mag_og} is above maximum for {band}. Using mag={mag}") elif mag != mag_og: print(f"Attention: mag={mag_og} rounded to nearest grid value mag={mag}") From c61eea119d6ce8fc851b1b85f6720cedf9691e19 Mon Sep 17 00:00:00 2001 From: IainHammond Date: Sat, 25 Jul 2026 16:49:51 +0200 Subject: [PATCH 46/46] support for RAVC --- heeps/contrast/disc.py | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py index 9c1b487..79a6186 100644 --- a/heeps/contrast/disc.py +++ b/heeps/contrast/disc.py @@ -82,7 +82,7 @@ def create_disc_sequence( Reference cube for RDI, if ``do_rdi`` is ``True``. Otherwise, this output is not returned. """ # if mag is not in the grid, round to the closest available magnitude - conf["mag"] = _check_mag(conf["mag"], conf["band"]) + conf["mag"] = _check_grid(conf["mag"], conf["band"], conf["mode"]) if do_rdi and rdi_mag is None: rdi_mag = conf["mag"] @@ -252,7 +252,7 @@ def create_disc_sequence( return psf_ON, pa -def _check_mag(mag : Union[float, int], band : str) -> float: +def _check_grid(mag : Union[float, int], band : str, mode : str) -> float: """ Check if the provided magnitude is in the Liege grid. If not, round to the closest available magnitude. A float must be returned as the grid contains mag as a float in the filenames. @@ -263,24 +263,28 @@ def _check_mag(mag : Union[float, int], band : str) -> float: The magnitude to check. band : str The band of the observation (e.g., "L", "M", "N1", "N2".) + mode : str + The mode of the observation (e.g., "CVC", "RAVC".) Returns ------- mag : float - The closest available magnitude in the Liege grid. + The closest available magnitude in the Liege grid for a given band and mode. """ mag_og = float(mag) - if band == "L": + if band == "L" and mode == "CVC": allowed = np.arange(-1.5, 9.0 + 0.5, step=0.5) - elif band == "M": + elif band == "L" and mode == "RAVC": + allowed = np.arange(-1.5, 4.0 + 0.5, step=0.5) + elif band == "M" and mode == "CVC": allowed = np.arange(-1.5, 7.0 + 0.5, step=0.5) - elif band == "N1": + elif band == "N1" and mode == "CVC": allowed = np.arange(-1.5, 4.0 + 0.5, step=0.5) - elif band == "N2": + elif band == "N2" and mode == "CVC": allowed = np.arange(-1.5, 3.0 + 0.5, step=0.5) else: - raise ValueError(f"Band {band} is not supported. Please use 'L', 'M', 'N1', or 'N2'.") + raise ValueError(f"Band {band} and mode {mode} is not supported.") if mag not in allowed: mag = round(mag / 0.5) * 0.5