Skip to content

Commit

Permalink
Placed cube_planet_free in fakecomp.py module to avoid circular import
Browse files Browse the repository at this point in the history
  • Loading branch information
@VChristiaens
VChristiaens committed Nov 23, 2024
1 parent fa82d1a commit 9b578d5
Show file tree
Hide file tree
Showing 3 changed files with 88 additions and 90 deletions.
2 changes: 1 addition & 1 deletion vip_hci/fm/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@
"""
from .fakecomp import *
from .fakedisk import *
from .utils_negfc import *
from .negfc_fmerit import *
from .negfc_mcmc import *
from .negfc_nested import *
Expand All @@ -28,4 +27,5 @@
from .negfd_simplex import *
from .scattered_light_disk import *
from .utils_mcmc import *
from .utils_negfc import *
from .utils_negfd import *
85 changes: 84 additions & 1 deletion vip_hci/fm/fakecomp.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,8 @@
'normalize_psf',
'cube_inject_companions',
'generate_cube_copies_with_injections',
'frame_inject_companion']
'frame_inject_companion',
'cube_planet_free']

from multiprocessing import cpu_count
import numpy as np
Expand Down Expand Up @@ -810,3 +811,85 @@ def psf_norm_2d(psf, fwhm, threshold, mask_core, full_output, verbose):
msg += " ``vip_hci.preproc.cube_collapse`` first, or loop on the "
msg += "temporal axis."
raise ValueError(msg.format(array.shape[0]))


def cube_planet_free(planet_parameter, cube, angs, psfn, imlib='vip-fft',
interpolation='lanczos4', transmission=None):
"""Return a cube in which we have injected negative fake companion at the\
position/flux given by planet_parameter.
Parameters
----------
planet_parameter: numpy.array or list or tuple
The (r, theta, flux) for all known companions. For a 3D cube, this
parameter must have a shape (n_pl,3) or (3,) -- the latter case assumes
a single planet in the data. For a 4d cube r, theta and flux
must all be 1d arrays with length equal to cube.shape[0]; i.e.
planet_parameter should have shape: (n_pl,3,n_ch).
cube: numpy ndarray
The cube of fits images expressed as a numpy.array.
angs: numpy ndarray
The parallactic angle fits image expressed as a numpy.array.
psfn: 2d or 3d numpy ndarray
The normalized psf expressed as a numpy ndarray. Can be 3d for a 4d
(spectral+ADI) input cube.
imlib : str, optional
See the documentation of the ``vip_hci.preproc.frame_rotate`` function.
interpolation : str, optional
See the documentation of the ``vip_hci.preproc.frame_rotate`` function.
transmission: numpy array, optional
Radial transmission of the coronagraph, if any. Array with either
2 x n_rad, 1+n_ch x n_rad columns. The first column should contain the
radial separation in pixels, while the other column(s) are the
corresponding off-axis transmission (between 0 and 1), for either all,
or each spectral channel (only relevant for a 4D input cube).
Returns
-------
cpf : numpy.array
The cube with negative companions injected at the position given in
planet_parameter.
"""
cpf = np.zeros_like(cube)

# unify planet_parameter format
planet_parameter = np.array(planet_parameter)
cond1 = cube.ndim == 3 and planet_parameter.ndim < 2
cond2 = cube.ndim == 4 and planet_parameter.ndim < 3
if cond1 or cond2:
planet_parameter = planet_parameter[np.newaxis, :]

if cube.ndim == 4:
if planet_parameter.shape[2] != cube.shape[0]:
raise TypeError("Input planet parameter with wrong dimensions.")

for i in range(planet_parameter.shape[0]):
if i == 0:
cube_temp = cube
else:
cube_temp = cpf

if cube.ndim == 4:
for j in range(cube.shape[0]):
flevel = -planet_parameter[i, 2, j]
r = planet_parameter[i, 0, j]
theta = planet_parameter[i, 1, j]
cpf[j] = cube_inject_companions(cube_temp[j], psfn[j], angs,
flevel=flevel,
rad_dists=[r],
n_branches=1,
theta=theta,
imlib=imlib,
interpolation=interpolation,
verbose=False,
transmission=transmission)
else:
cpf = cube_inject_companions(cube_temp, psfn, angs, n_branches=1,
flevel=-planet_parameter[i, 2],
rad_dists=[planet_parameter[i, 0]],
theta=planet_parameter[i, 1],
imlib=imlib, verbose=False,
interpolation=interpolation,
transmission=transmission)
return cpf
91 changes: 3 additions & 88 deletions vip_hci/fm/utils_negfc.py
View file
Original file line number Diff line number Diff line change
@@ -1,97 +1,12 @@
#! /usr/bin/env python
"""
Module with post-processing related functions called from within the NEGFC
algorithm.
Module with utility function called from within the NEGFC algorithm.
"""

__author__ = 'Carlos Alberto Gomez Gonzalez'
__all__ = ['cube_planet_free',
'find_nearest']
__author__ = 'Valentin Christiaens'
__all__ = ['find_nearest']

import numpy as np
from ..fm import cube_inject_companions


def cube_planet_free(planet_parameter, cube, angs, psfn, imlib='vip-fft',
interpolation='lanczos4', transmission=None):
"""Return a cube in which we have injected negative fake companion at the\
position/flux given by planet_parameter.
Parameters
----------
planet_parameter: numpy.array or list or tuple
The (r, theta, flux) for all known companions. For a 3D cube, this
parameter must have a shape (n_pl,3) or (3,) -- the latter case assumes
a single planet in the data. For a 4d cube r, theta and flux
must all be 1d arrays with length equal to cube.shape[0]; i.e.
planet_parameter should have shape: (n_pl,3,n_ch).
cube: numpy ndarray
The cube of fits images expressed as a numpy.array.
angs: numpy ndarray
The parallactic angle fits image expressed as a numpy.array.
psfn: 2d or 3d numpy ndarray
The normalized psf expressed as a numpy ndarray. Can be 3d for a 4d
(spectral+ADI) input cube.
imlib : str, optional
See the documentation of the ``vip_hci.preproc.frame_rotate`` function.
interpolation : str, optional
See the documentation of the ``vip_hci.preproc.frame_rotate`` function.
transmission: numpy array, optional
Radial transmission of the coronagraph, if any. Array with either
2 x n_rad, 1+n_ch x n_rad columns. The first column should contain the
radial separation in pixels, while the other column(s) are the
corresponding off-axis transmission (between 0 and 1), for either all,
or each spectral channel (only relevant for a 4D input cube).
Returns
-------
cpf : numpy.array
The cube with negative companions injected at the position given in
planet_parameter.
"""
cpf = np.zeros_like(cube)

# unify planet_parameter format
planet_parameter = np.array(planet_parameter)
cond1 = cube.ndim == 3 and planet_parameter.ndim < 2
cond2 = cube.ndim == 4 and planet_parameter.ndim < 3
if cond1 or cond2:
planet_parameter = planet_parameter[np.newaxis, :]

if cube.ndim == 4:
if planet_parameter.shape[2] != cube.shape[0]:
raise TypeError("Input planet parameter with wrong dimensions.")

for i in range(planet_parameter.shape[0]):
if i == 0:
cube_temp = cube
else:
cube_temp = cpf

if cube.ndim == 4:
for j in range(cube.shape[0]):
flevel = -planet_parameter[i, 2, j]
r = planet_parameter[i, 0, j]
theta = planet_parameter[i, 1, j]
cpf[j] = cube_inject_companions(cube_temp[j], psfn[j], angs,
flevel=flevel,
rad_dists=[r],
n_branches=1,
theta=theta,
imlib=imlib,
interpolation=interpolation,
verbose=False,
transmission=transmission)
else:
cpf = cube_inject_companions(cube_temp, psfn, angs, n_branches=1,
flevel=-planet_parameter[i, 2],
rad_dists=[planet_parameter[i, 0]],
theta=planet_parameter[i, 1],
imlib=imlib, verbose=False,
interpolation=interpolation,
transmission=transmission)
return cpf


def find_nearest(array, value, output='index', constraint=None, n=1):
Expand Down

0 comments on commit 9b578d5

Please sign in to comment.