diff --git a/heeps/contrast/background.py b/heeps/contrast/background.py index d48b56e..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,10 +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)) + # 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 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)) diff --git a/heeps/contrast/disc.py b/heeps/contrast/disc.py new file mode 100644 index 0000000..79a6186 --- /dev/null +++ b/heeps/contrast/disc.py @@ -0,0 +1,302 @@ +""" +HEEPS disc injection utilities. + +This module provides functions to inject a synthetic disc model into HEEPS +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. + +""" + +import numpy as np +from typing import Union +import os + +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 + +from heeps.util.psf_template import psf_template +from heeps.util.paralang import paralang + +__all__ = ['create_disc_sequence'] +__author__ = "Iain Hammond" + +def create_disc_sequence( + disc_model : np.ndarray, + db_path : str = "vortex_psf_grid/", + 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, + imlib : str = "opencv", + **conf +): + """ + 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 star included to convert to units of contrast. + + Parameters + ---------- + 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 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 + (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``. + 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. + conf : dict, optional + Configuration dictionary with parameters from HEEPS (e.g., band, dit, mag, 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 ``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_grid(conf["mag"], conf["band"], conf["mode"]) + + 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) + + # 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, + # 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" + ) + 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) + + # 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: + 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:] + + 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 + 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: + disc_model = frame_crop(disc_model, min_crop, verbose=False) + + # 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 + 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 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 + 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 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) + 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, + 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 (units of contrast) + cube /= star_val + + # add the disc model + psf_ON += cube + del cube + + if starphot is not None: + _, _, ap_flux = psf_template(psf_OFF) + psf_ON *= starphot / ap_flux + del psf_OFF + + if do_rdi: + return psf_ON, pa, psf_RDI + else: + return psf_ON, pa + + +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. + + Parameters + ---------- + mag : float, int + 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 for a given band and mode. + """ + mag_og = float(mag) + + if band == "L" and mode == "CVC": + allowed = np.arange(-1.5, 9.0 + 0.5, step=0.5) + 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" and mode == "CVC": + allowed = np.arange(-1.5, 4.0 + 0.5, step=0.5) + elif band == "N2" and mode == "CVC": + allowed = np.arange(-1.5, 3.0 + 0.5, step=0.5) + else: + raise ValueError(f"Band {band} and mode {mode} is not supported.") + + 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"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 for {band}. Using mag={mag}") + elif mag != mag_og: + print(f"Attention: mag={mag_og} rounded to nearest grid value mag={mag}") + + return float(mag) 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