Efficient FFI loading

src/psfmachine/machine.py

@@ -59,6 +59,7 @@ def __init__(
  rmax=16,
  cut_r=6,
  sparse_dist_lim=40,
+ sources_flux_column="phot_g_mean_flux",
  ):
  """
  Parameters
@@ -104,6 +105,9 @@ def __init__(
  Radial distance used to include pixels around sources when creating delta
  arrays (dra, ddec, r, and phi) as sparse matrices for efficiency.
  Default is 40" (recommended for kepler). (arcseconds)
+ sources_flux_column : str
+ Column name in `sources` table to be used as flux estimate. For Kepler data
+ gaia.phot_g_mean_flux is recommended, for TESS use gaia.phot_rp_mean_flux.
 
  Attributes
  ----------
@@ -172,6 +176,8 @@ def __init__(
  self.quiet = False
  self.contaminant_mag_limit = None
 
+ self.source_flux_estimates = np.copy(self.sources[sources_flux_column].values)
+
  if time_mask is None:
  self.time_mask = np.ones(len(time), bool)
  else:
@@ -231,7 +237,7 @@ def _create_delta_arrays(self, centroid_offset=[0, 0]):
  self.r = np.hypot(self.dra, self.ddec).to("arcsec")
  self.phi = np.arctan2(self.ddec, self.dra)
 
- def _create_delta_sparse_arrays(self, dist_lim=40, centroid_offset=[0, 0]):
+ def _create_delta_sparse_arrays(self, centroid_offset=[0, 0]):
  """
  Creates dra, ddec, r and phi arrays as sparse arrays to be used for dense data,
  e.g. Kepler FFIs or cluster fields. Assuming that there is no flux information
@@ -243,8 +249,6 @@ def _create_delta_sparse_arrays(self, dist_lim=40, centroid_offset=[0, 0]):
 
  Parameters
  ----------
- dist_lim : float
- Distance limit (in arcsecds) at which pixels are keep.
  centroid_offset : list
  Centroid offset for [ra, dec] to be included in dra and ddec computation.
  Default is [0, 0].
@@ -254,7 +258,7 @@ def _create_delta_sparse_arrays(self, dist_lim=40, centroid_offset=[0, 0]):
  if (centroid_offset[0] == centroid_offset[1] == 0) or (
  np.maximum(*np.abs(centroid_offset)) > 4 / 3600
  ):
- # iterate over sources to only keep pixels within dist_lim
+ # iterate over sources to only keep pixels within self.sparse_dist_lim
  # this is inefficient, could be done in a tiled manner? only for squared data
  dra, ddec, sparse_mask = [], [], []
  for i in tqdm(
@@ -328,17 +332,21 @@ def _get_source_mask(
  correct_centroid_offset=True,
  plot=False,
  ):
- """Find the pixel mask that identifies pixels with contributions from ANY NUMBER of Sources
+ """Find the pixel mask that identifies pixels with contributions from ANY
+ NUMBER of Sources
 
- Fits a simple polynomial model to the log of the pixel flux values, in radial dimension and source flux,
- to find the optimum circular apertures for every source.
+ Fits a simple polynomial model to the log of the pixel flux values, in radial
+ dimension and source flux, to find the optimum circular apertures for every
+ source.
 
  Parameters
  ----------
  upper_radius_limit: float
- The radius limit at which we assume there is no flux from a source of any brightness (arcsec)
+ The radius limit at which we assume there is no flux from a source of
+ any brightness (arcsec)
  lower_radius_limit: float
- The radius limit at which we assume there is flux from a source of any brightness (arcsec)
+ The radius limit at which we assume there is flux from a source of any
+ brightness (arcsec)
  upper_flux_limit: float
  The flux at which we assume as source is saturated
  lower_flux_limit: float
@@ -349,12 +357,6 @@ def _get_source_mask(
  plot: bool
  Whether to show diagnostic plot. Default is False
  """
-
- if not hasattr(self, "source_flux_estimates"):
- # gaia estimate flux values per pixel to be used as flux priors
- self.source_flux_estimates = np.copy(
- np.asarray(self.sources.phot_g_mean_flux)
- )
  # We will use the radius a lot, this is for readibility
  # don't do if sparse array
  if isinstance(self.r, u.quantity.Quantity):
@@ -365,14 +367,15 @@ def _get_source_mask(
  # The average flux, which we assume is a good estimate of the whole stack of images
  max_flux = np.nanmax(self.flux[self.time_mask], axis=0)
 
- # Mask out sources that are above the flux limit, and pixels above the radius limit
+ # Mask out sources that are above the flux limit, and pixels above the
+ # radius limit. This is a good estimation only when using phot_g_mean_flux
  source_rad = 0.5 * np.log10(self.source_flux_estimates) ** 1.5 + 3
  # temp_mask for the sparse array case should also be a sparse matrix. Then it is
  # applied to r, max_flux, and, source_flux_estimates to be used later.
  # Numpy array case:
  if not isinstance(r, sparse.csr_matrix):
  # First we make a guess that each source has exactly the gaia flux
- source_flux_estimates = np.asarray(self.sources.phot_g_mean_flux)[
+ source_flux_estimates = np.asarray(self.source_flux_estimates)[
  :, None
  ] * np.ones((self.nsources, self.npixels))
  temp_mask = (r < source_rad[:, None]) & (
@@ -528,7 +531,7 @@ def _get_source_mask(
  if plot:
  k = np.isfinite(f_temp_mask)
  # Make a nice diagnostic plot
- fig, ax = plt.subplots(1, 2, figsize=(8, 3), facecolor="white")
+ fig, ax = plt.subplots(1, 2, figsize=(11, 3), facecolor="white")
 
  ax[0].scatter(r_temp_mask[k], f_temp_mask[k], s=0.4, c="k", label="Data")
  ax[0].scatter(r_temp_mask[k], A[k].dot(w), c="r", s=0.4, label="Model")
@@ -573,8 +576,7 @@ def _get_uncontaminated_pixel_mask(self):
  # we flag sources fainter than mag_limit as non-contaminant
  if isinstance(self.contaminant_mag_limit, (float, int)):
  aux = self.source_mask.multiply(
- self.sources.phot_g_mean_mag.values[:, None]
- < self.contaminant_mag_limit
+ self.source_flux_estimates[:, None] < self.contaminant_mag_limit
  )
  aux.eliminate_zeros()
  self.uncontaminated_source_mask = aux.multiply(
@@ -957,16 +959,10 @@ def build_shape_model(
  Keyword arguments to be passed to `_get_source_mask()`
  """
 
- # gaia estimate flux values per pixel to be used as flux priors
- self.source_flux_estimates = np.copy(np.asarray(self.sources.phot_g_mean_flux))
-
  # Mask of shape nsources x number of pixels, one where flux from a
  # source exists
- if not hasattr(self, "source_mask"):
- self._get_source_mask(**kwargs)
- # Mask of shape npixels (maybe by nt) where not saturated, not faint,
- # not contaminated etc
- self._get_uncontaminated_pixel_mask()
+ # if not hasattr(self, "source_mask"):
+ self._get_source_mask(**kwargs)
 
  # for iter in range(niters):
  flux_estimates = self.source_flux_estimates[:, None]
@@ -1346,24 +1342,23 @@ def get_psf_metrics(self, npoints_per_pixel=10):
  )
  self.perturbed_std = np.nanstd(perturbed_lcs, axis=0)
 
- def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
+ def plot_shape_model(self, frame_index="mean", bin_data=False):
  """
  Diagnostic plot of shape model.
 
  Parameters
  ----------
- radius : float
- Radius (in arcseconds) limit to be shown in the figure.
  frame_index : string or int
  The frame index used to plot the shape model, if "mean" then use the
  mean value across time
+ bin_data : bool
+ Bin or not the pixel data in a 2D historgram, default is False.
 
  Returns
  -------
  fig : matplotlib.Figure
  Figure.
  """
-
  if frame_index == "mean":
  mean_f = np.log10(
  self.uncontaminated_source_mask.astype(float)
@@ -1386,6 +1381,9 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
  dx = dx.data * u.deg.to(u.arcsecond)
  dy = dy.data * u.deg.to(u.arcsecond)
 
+ radius = np.maximum(np.abs(dx).max(), np.abs(dy).max()) * 1.1
+ vmin, vmax = np.nanpercentile(mean_f, [5, 90])
+
  if bin_data:
  nbins = 30 if mean_f.shape[0] <= 5e3 else 90
  _, dx, dy, mean_f, _ = threshold_bin(
@@ -1394,7 +1392,14 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
 
  fig, ax = plt.subplots(3, 2, figsize=(9, 10.5), constrained_layout=True)
  im = ax[0, 0].scatter(
- dx, dy, c=mean_f, cmap="viridis", vmin=-3, vmax=-1, s=3, rasterized=True
+ dx,
+ dy,
+ c=mean_f,
+ cmap="viridis",
+ vmin=vmin,
+ vmax=vmax,
+ s=3,
+ rasterized=True,
  )
  ax[0, 0].set(
  ylabel=r'$\delta y$ ["]',
@@ -1418,7 +1423,7 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
 
  phi, r = np.arctan2(dy, dx), np.hypot(dx, dy)
  im = ax[0, 1].scatter(
- phi, r, c=mean_f, cmap="viridis", vmin=-3, vmax=-1, s=3, rasterized=True
+ phi, r, c=mean_f, cmap="viridis", vmin=vmin, vmax=vmax, s=3, rasterized=True
  )
  ax[0, 1].set(
  ylabel='$r$ ["]',
@@ -1448,8 +1453,8 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
  r,
  c=model,
  cmap="viridis",
- vmin=-3,
- vmax=-1,
+ vmin=vmin,
+ vmax=vmax,
  s=3,
  rasterized=True,
  )
@@ -1464,8 +1469,8 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
  dy,
  c=model,
  cmap="viridis",
- vmin=-3,
- vmax=-1,
+ vmin=vmin,
+ vmax=vmax,
  s=3,
  rasterized=True,
  )

src/psfmachine/superstamp.py

@@ -15,7 +15,8 @@
 from astropy.wcs import WCS
 import astropy.units as u
 
-from .ffi import FFIMachine, _get_sources, _do_image_cutout
+from .ffi import FFIMachine, _get_sources
+from .utils import _do_image_cutout
 
 __all__ = ["SSMachine"]
 

src/psfmachine/tpf.py

@@ -52,9 +52,10 @@ def __init__(
  pos_corr2=None,
  focus_mask=None,
  tpf_meta=None,
- time_corrector="pos_corr",
  cartesian_knot_spacing="sqrt",
  bkg_subtracted=True,
+ sparse_dist_lim=40,
+ sources_flux_column="phot_g_mean_flux",
  ):
  super().__init__(
  time=time,
@@ -73,6 +74,8 @@ def __init__(
  time_radius=time_radius,
  rmin=rmin,
  rmax=rmax,
+ sparse_dist_lim=sparse_dist_lim,
+ sources_flux_column=sources_flux_column,
  )
  self.tpfs = tpfs
  # match cadences
@@ -93,7 +96,6 @@ def __init__(
  self.pos_corr1 = pos_corr1
  self.pos_corr2 = pos_corr2
  self.tpf_meta = tpf_meta
- self.time_corrector = time_corrector
  self.cartesian_knot_spacing = cartesian_knot_spacing
  self.bkg_subtracted = bkg_subtracted
  # limiting mag values to consider contamination for kepler & tess
@@ -153,7 +155,8 @@ def remove_background_model(
  )
  return
  # invert maks to get bkg pixels from TPFs
- self._get_source_mask()
+ if not hasattr(self, "source_mask"):
+ self._get_source_mask(correct_centroid_offset=True)
  bkg_pixel_mask = ~np.asarray(
  (self.source_mask.todense()).sum(axis=0).astype(bool)
  ).ravel()
@@ -854,7 +857,8 @@ def get_source_centroids(self, method="poscor"):
  def from_TPFs(
  tpfs,
  magnitude_limit=18,
- dr=2,
+ dr=3,
+ sources=None,
  time_mask=None,
  apply_focus_mask=True,
  renormalize_tpf_bkg=False,
@@ -882,7 +886,9 @@ def from_TPFs(
  magnitude_limit : float
  Limiting magnitude to query Gaia catalog.
  dr : int
- Gaia data release to be use, default is 2, options are DR2 and EDR3
+ Gaia data release to be use, default is 3, options are DR2 and EDR3
+ sources : pandas.DataFrame
+ Catalog with sources to be extracted by PSFMachine
  time_mask : boolean array
  Mask to be applied to discard cadences if needed.
  apply_focus_mask : boolean
@@ -1008,9 +1014,10 @@ def from_TPFs(
  saturated_mask,
  )
 
- sources = _get_coord_and_query_gaia(
- tpfs, magnitude_limit, dr=dr, ra=query_ra, dec=query_dec, rad=query_rad
- )
+ if sources is None:
+ sources = _get_coord_and_query_gaia(
+ tpfs, magnitude_limit, dr=dr, ra=query_ra, dec=query_dec, rad=query_rad
+ )
 
  def get_tpf2source():
  tpf2source = []
@@ -1052,6 +1059,20 @@ def get_tpf2source():
  sources["tpf_id"] = None
  sources.loc[idx[match], "tpf_id"] = np.asarray(tpf_meta["targetid"])[match]
 
+ # select flux column o be used as estimate.
+ if tpf_meta["mission"][0] == "TESS":
+ sources_flux_column = "phot_rp_mean_flux"
+ if np.isnan(sources[sources_flux_column].values).any():
+ log.warning(
+ f"Some sources do not have values in 'gaia.{sources_flux_column}''."
+ f" Will default to 'gaia.phot_g_mean_flux' for missing values."
+ )
+ sources[sources_flux_column].fillna(
+ sources["phot_g_mean_flux"], inplace=True
+ )
+ else:
+ sources_flux_column = "phot_g_mean_flux"
+
  # return a Machine object
  return TPFMachine(
  tpfs=tpfs,
@@ -1070,6 +1091,7 @@ def get_tpf2source():
  tpf_meta=tpf_meta,
  time_mask=time_mask,
  bkg_subtracted=not renormalize_tpf_bkg,
+ sources_flux_column=sources_flux_column,
  **kwargs,
  )
 
@@ -1121,7 +1143,7 @@ def _parse_TPFs(tpfs, renormalize_tpf_bkg=True, **kwargs):
 
  elif isinstance(tpfs[0], lk.TessTargetPixelFile):
  qual_mask = lk.utils.TessQualityFlags.create_quality_mask(
- tpfs[0].quality, lk.utils.TessQualityFlags.DEFAULT_BITMASK
+ tpfs[0].quality, 1 | 2 | 4 | 8 | 32 | 128
  )
  qual_mask &= (np.abs(tpfs[0].pos_corr1) < 5) & (np.abs(tpfs[0].pos_corr2) < 5)
  focus_mask = np.ones(len(tpfs[0].time), bool)[qual_mask]

src/psfmachine/machine.py

@@ -59,6 +59,7 @@ def __init__(
  rmax=16,
  cut_r=6,
  sparse_dist_lim=40,
+ sources_flux_column="phot_g_mean_flux",
  ):
  """
  Parameters
@@ -104,6 +105,9 @@ def __init__(
  Radial distance used to include pixels around sources when creating delta
  arrays (dra, ddec, r, and phi) as sparse matrices for efficiency.
  Default is 40" (recommended for kepler). (arcseconds)
+ sources_flux_column : str
+ Column name in `sources` table to be used as flux estimate. For Kepler data
+ gaia.phot_g_mean_flux is recommended, for TESS use gaia.phot_rp_mean_flux.
 
  Attributes
  ----------
@@ -172,6 +176,8 @@ def __init__(
  self.quiet = False
  self.contaminant_mag_limit = None
 
+ self.source_flux_estimates = np.copy(self.sources[sources_flux_column].values)
+
  if time_mask is None:
  self.time_mask = np.ones(len(time), bool)
  else:
@@ -231,7 +237,7 @@ def _create_delta_arrays(self, centroid_offset=[0, 0]):
  self.r = np.hypot(self.dra, self.ddec).to("arcsec")
  self.phi = np.arctan2(self.ddec, self.dra)
 
- def _create_delta_sparse_arrays(self, dist_lim=40, centroid_offset=[0, 0]):
+ def _create_delta_sparse_arrays(self, centroid_offset=[0, 0]):
  """
  Creates dra, ddec, r and phi arrays as sparse arrays to be used for dense data,
  e.g. Kepler FFIs or cluster fields. Assuming that there is no flux information
@@ -243,8 +249,6 @@ def _create_delta_sparse_arrays(self, dist_lim=40, centroid_offset=[0, 0]):
 
  Parameters
  ----------
- dist_lim : float
- Distance limit (in arcsecds) at which pixels are keep.
  centroid_offset : list
  Centroid offset for [ra, dec] to be included in dra and ddec computation.
  Default is [0, 0].
@@ -254,7 +258,7 @@ def _create_delta_sparse_arrays(self, dist_lim=40, centroid_offset=[0, 0]):
  if (centroid_offset[0] == centroid_offset[1] == 0) or (
  np.maximum(*np.abs(centroid_offset)) > 4 / 3600
  ):
- # iterate over sources to only keep pixels within dist_lim
+ # iterate over sources to only keep pixels within self.sparse_dist_lim
  # this is inefficient, could be done in a tiled manner? only for squared data
  dra, ddec, sparse_mask = [], [], []
  for i in tqdm(
@@ -328,17 +332,21 @@ def _get_source_mask(
  correct_centroid_offset=True,
  plot=False,
  ):
- """Find the pixel mask that identifies pixels with contributions from ANY NUMBER of Sources
+ """Find the pixel mask that identifies pixels with contributions from ANY
+ NUMBER of Sources
 
- Fits a simple polynomial model to the log of the pixel flux values, in radial dimension and source flux,
- to find the optimum circular apertures for every source.
+ Fits a simple polynomial model to the log of the pixel flux values, in radial
+ dimension and source flux, to find the optimum circular apertures for every
+ source.
 
  Parameters
  ----------
  upper_radius_limit: float
- The radius limit at which we assume there is no flux from a source of any brightness (arcsec)
+ The radius limit at which we assume there is no flux from a source of
+ any brightness (arcsec)
  lower_radius_limit: float
- The radius limit at which we assume there is flux from a source of any brightness (arcsec)
+ The radius limit at which we assume there is flux from a source of any
+ brightness (arcsec)
  upper_flux_limit: float
  The flux at which we assume as source is saturated
  lower_flux_limit: float
@@ -349,12 +357,6 @@ def _get_source_mask(
  plot: bool
  Whether to show diagnostic plot. Default is False
  """
-
- if not hasattr(self, "source_flux_estimates"):
- # gaia estimate flux values per pixel to be used as flux priors
- self.source_flux_estimates = np.copy(
- np.asarray(self.sources.phot_g_mean_flux)
- )
  # We will use the radius a lot, this is for readibility
  # don't do if sparse array
  if isinstance(self.r, u.quantity.Quantity):
@@ -365,14 +367,15 @@ def _get_source_mask(
  # The average flux, which we assume is a good estimate of the whole stack of images
  max_flux = np.nanmax(self.flux[self.time_mask], axis=0)
 
- # Mask out sources that are above the flux limit, and pixels above the radius limit
+ # Mask out sources that are above the flux limit, and pixels above the
+ # radius limit. This is a good estimation only when using phot_g_mean_flux
  source_rad = 0.5 * np.log10(self.source_flux_estimates) ** 1.5 + 3
  # temp_mask for the sparse array case should also be a sparse matrix. Then it is
  # applied to r, max_flux, and, source_flux_estimates to be used later.
  # Numpy array case:
  if not isinstance(r, sparse.csr_matrix):
  # First we make a guess that each source has exactly the gaia flux
- source_flux_estimates = np.asarray(self.sources.phot_g_mean_flux)[
+ source_flux_estimates = np.asarray(self.source_flux_estimates)[
  :, None
  ] * np.ones((self.nsources, self.npixels))
  temp_mask = (r < source_rad[:, None]) & (
@@ -528,7 +531,7 @@ def _get_source_mask(
  if plot:
  k = np.isfinite(f_temp_mask)
  # Make a nice diagnostic plot
- fig, ax = plt.subplots(1, 2, figsize=(8, 3), facecolor="white")
+ fig, ax = plt.subplots(1, 2, figsize=(11, 3), facecolor="white")
 
  ax[0].scatter(r_temp_mask[k], f_temp_mask[k], s=0.4, c="k", label="Data")
  ax[0].scatter(r_temp_mask[k], A[k].dot(w), c="r", s=0.4, label="Model")
@@ -573,8 +576,7 @@ def _get_uncontaminated_pixel_mask(self):
  # we flag sources fainter than mag_limit as non-contaminant
  if isinstance(self.contaminant_mag_limit, (float, int)):
  aux = self.source_mask.multiply(
- self.sources.phot_g_mean_mag.values[:, None]
- < self.contaminant_mag_limit
+ self.source_flux_estimates[:, None] < self.contaminant_mag_limit
  )
  aux.eliminate_zeros()
  self.uncontaminated_source_mask = aux.multiply(
@@ -957,16 +959,10 @@ def build_shape_model(
  Keyword arguments to be passed to `_get_source_mask()`
  """
 
- # gaia estimate flux values per pixel to be used as flux priors
- self.source_flux_estimates = np.copy(np.asarray(self.sources.phot_g_mean_flux))
-
  # Mask of shape nsources x number of pixels, one where flux from a
  # source exists
- if not hasattr(self, "source_mask"):
- self._get_source_mask(**kwargs)
- # Mask of shape npixels (maybe by nt) where not saturated, not faint,
- # not contaminated etc
- self._get_uncontaminated_pixel_mask()
+ # if not hasattr(self, "source_mask"):
+ self._get_source_mask(**kwargs)
 
  # for iter in range(niters):
  flux_estimates = self.source_flux_estimates[:, None]
@@ -1346,24 +1342,23 @@ def get_psf_metrics(self, npoints_per_pixel=10):
  )
  self.perturbed_std = np.nanstd(perturbed_lcs, axis=0)
 
- def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
+ def plot_shape_model(self, frame_index="mean", bin_data=False):
  """
  Diagnostic plot of shape model.
 
  Parameters
  ----------
- radius : float
- Radius (in arcseconds) limit to be shown in the figure.
  frame_index : string or int
  The frame index used to plot the shape model, if "mean" then use the
  mean value across time
+ bin_data : bool
+ Bin or not the pixel data in a 2D historgram, default is False.
 
  Returns
  -------
  fig : matplotlib.Figure
  Figure.
  """
-
  if frame_index == "mean":
  mean_f = np.log10(
  self.uncontaminated_source_mask.astype(float)
@@ -1386,6 +1381,9 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
  dx = dx.data * u.deg.to(u.arcsecond)
  dy = dy.data * u.deg.to(u.arcsecond)
 
+ radius = np.maximum(np.abs(dx).max(), np.abs(dy).max()) * 1.1
+ vmin, vmax = np.nanpercentile(mean_f, [5, 90])
+
  if bin_data:
  nbins = 30 if mean_f.shape[0] <= 5e3 else 90
  _, dx, dy, mean_f, _ = threshold_bin(
@@ -1394,7 +1392,14 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
 
  fig, ax = plt.subplots(3, 2, figsize=(9, 10.5), constrained_layout=True)
  im = ax[0, 0].scatter(
- dx, dy, c=mean_f, cmap="viridis", vmin=-3, vmax=-1, s=3, rasterized=True
+ dx,
+ dy,
+ c=mean_f,
+ cmap="viridis",
+ vmin=vmin,
+ vmax=vmax,
+ s=3,
+ rasterized=True,
  )
  ax[0, 0].set(
  ylabel=r'$\delta y$ ["]',
@@ -1418,7 +1423,7 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
 
  phi, r = np.arctan2(dy, dx), np.hypot(dx, dy)
  im = ax[0, 1].scatter(
- phi, r, c=mean_f, cmap="viridis", vmin=-3, vmax=-1, s=3, rasterized=True
+ phi, r, c=mean_f, cmap="viridis", vmin=vmin, vmax=vmax, s=3, rasterized=True
  )
  ax[0, 1].set(
  ylabel='$r$ ["]',
@@ -1448,8 +1453,8 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
  r,
  c=model,
  cmap="viridis",
- vmin=-3,
- vmax=-1,
+ vmin=vmin,
+ vmax=vmax,
  s=3,
  rasterized=True,
  )
@@ -1464,8 +1469,8 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
  dy,
  c=model,
  cmap="viridis",
- vmin=-3,
- vmax=-1,
+ vmin=vmin,
+ vmax=vmax,
  s=3,
  rasterized=True,
  )

src/psfmachine/superstamp.py

@@ -15,7 +15,8 @@
 from astropy.wcs import WCS
 import astropy.units as u
 
-from .ffi import FFIMachine, _get_sources, _do_image_cutout
+from .ffi import FFIMachine, _get_sources
+from .utils import _do_image_cutout
 
 __all__ = ["SSMachine"]
 

src/psfmachine/tpf.py

@@ -52,9 +52,10 @@ def __init__(
  pos_corr2=None,
  focus_mask=None,
  tpf_meta=None,
- time_corrector="pos_corr",
  cartesian_knot_spacing="sqrt",
  bkg_subtracted=True,
+ sparse_dist_lim=40,
+ sources_flux_column="phot_g_mean_flux",
  ):
  super().__init__(
  time=time,
@@ -73,6 +74,8 @@ def __init__(
  time_radius=time_radius,
  rmin=rmin,
  rmax=rmax,
+ sparse_dist_lim=sparse_dist_lim,
+ sources_flux_column=sources_flux_column,
  )
  self.tpfs = tpfs
  # match cadences
@@ -93,7 +96,6 @@ def __init__(
  self.pos_corr1 = pos_corr1
  self.pos_corr2 = pos_corr2
  self.tpf_meta = tpf_meta
- self.time_corrector = time_corrector
  self.cartesian_knot_spacing = cartesian_knot_spacing
  self.bkg_subtracted = bkg_subtracted
  # limiting mag values to consider contamination for kepler & tess
@@ -153,7 +155,8 @@ def remove_background_model(
  )
  return
  # invert maks to get bkg pixels from TPFs
- self._get_source_mask()
+ if not hasattr(self, "source_mask"):
+ self._get_source_mask(correct_centroid_offset=True)
  bkg_pixel_mask = ~np.asarray(
  (self.source_mask.todense()).sum(axis=0).astype(bool)
  ).ravel()
@@ -854,7 +857,8 @@ def get_source_centroids(self, method="poscor"):
  def from_TPFs(
  tpfs,
  magnitude_limit=18,
- dr=2,
+ dr=3,
+ sources=None,
  time_mask=None,
  apply_focus_mask=True,
  renormalize_tpf_bkg=False,
@@ -882,7 +886,9 @@ def from_TPFs(
  magnitude_limit : float
  Limiting magnitude to query Gaia catalog.
  dr : int
- Gaia data release to be use, default is 2, options are DR2 and EDR3
+ Gaia data release to be use, default is 3, options are DR2 and EDR3
+ sources : pandas.DataFrame
+ Catalog with sources to be extracted by PSFMachine
  time_mask : boolean array
  Mask to be applied to discard cadences if needed.
  apply_focus_mask : boolean
@@ -1008,9 +1014,10 @@ def from_TPFs(
  saturated_mask,
  )
 
- sources = _get_coord_and_query_gaia(
- tpfs, magnitude_limit, dr=dr, ra=query_ra, dec=query_dec, rad=query_rad
- )
+ if sources is None:
+ sources = _get_coord_and_query_gaia(
+ tpfs, magnitude_limit, dr=dr, ra=query_ra, dec=query_dec, rad=query_rad
+ )
 
  def get_tpf2source():
  tpf2source = []
@@ -1052,6 +1059,20 @@ def get_tpf2source():
  sources["tpf_id"] = None
  sources.loc[idx[match], "tpf_id"] = np.asarray(tpf_meta["targetid"])[match]
 
+ # select flux column o be used as estimate.
+ if tpf_meta["mission"][0] == "TESS":
+ sources_flux_column = "phot_rp_mean_flux"
+ if np.isnan(sources[sources_flux_column].values).any():
+ log.warning(
+ f"Some sources do not have values in 'gaia.{sources_flux_column}''."
+ f" Will default to 'gaia.phot_g_mean_flux' for missing values."
+ )
+ sources[sources_flux_column].fillna(
+ sources["phot_g_mean_flux"], inplace=True
+ )
+ else:
+ sources_flux_column = "phot_g_mean_flux"
+
  # return a Machine object
  return TPFMachine(
  tpfs=tpfs,
@@ -1070,6 +1091,7 @@ def get_tpf2source():
  tpf_meta=tpf_meta,
  time_mask=time_mask,
  bkg_subtracted=not renormalize_tpf_bkg,
+ sources_flux_column=sources_flux_column,
  **kwargs,
  )
 
@@ -1121,7 +1143,7 @@ def _parse_TPFs(tpfs, renormalize_tpf_bkg=True, **kwargs):
 
  elif isinstance(tpfs[0], lk.TessTargetPixelFile):
  qual_mask = lk.utils.TessQualityFlags.create_quality_mask(
- tpfs[0].quality, lk.utils.TessQualityFlags.DEFAULT_BITMASK
+ tpfs[0].quality, 1 | 2 | 4 | 8 | 32 | 128
  )
  qual_mask &= (np.abs(tpfs[0].pos_corr1) < 5) & (np.abs(tpfs[0].pos_corr2) < 5)
  focus_mask = np.ones(len(tpfs[0].time), bool)[qual_mask]

src/psfmachine/machine.py

@@ -59,6 +59,7 @@ def __init__(
  rmax=16,
  cut_r=6,
  sparse_dist_lim=40,
+ sources_flux_column="phot_g_mean_flux",
  ):
  """
  Parameters
@@ -104,6 +105,9 @@ def __init__(
  Radial distance used to include pixels around sources when creating delta
  arrays (dra, ddec, r, and phi) as sparse matrices for efficiency.
  Default is 40" (recommended for kepler). (arcseconds)
+ sources_flux_column : str
+ Column name in `sources` table to be used as flux estimate. For Kepler data
+ gaia.phot_g_mean_flux is recommended, for TESS use gaia.phot_rp_mean_flux.
 
  Attributes
  ----------
@@ -172,6 +176,8 @@ def __init__(
  self.quiet = False
  self.contaminant_mag_limit = None
 
+ self.source_flux_estimates = np.copy(self.sources[sources_flux_column].values)
+
  if time_mask is None:
  self.time_mask = np.ones(len(time), bool)
  else:
@@ -231,7 +237,7 @@ def _create_delta_arrays(self, centroid_offset=[0, 0]):
  self.r = np.hypot(self.dra, self.ddec).to("arcsec")
  self.phi = np.arctan2(self.ddec, self.dra)
 
- def _create_delta_sparse_arrays(self, dist_lim=40, centroid_offset=[0, 0]):
+ def _create_delta_sparse_arrays(self, centroid_offset=[0, 0]):
  """
  Creates dra, ddec, r and phi arrays as sparse arrays to be used for dense data,
  e.g. Kepler FFIs or cluster fields. Assuming that there is no flux information
@@ -243,8 +249,6 @@ def _create_delta_sparse_arrays(self, dist_lim=40, centroid_offset=[0, 0]):
 
  Parameters
  ----------
- dist_lim : float
- Distance limit (in arcsecds) at which pixels are keep.
  centroid_offset : list
  Centroid offset for [ra, dec] to be included in dra and ddec computation.
  Default is [0, 0].
@@ -254,7 +258,7 @@ def _create_delta_sparse_arrays(self, dist_lim=40, centroid_offset=[0, 0]):
  if (centroid_offset[0] == centroid_offset[1] == 0) or (
  np.maximum(*np.abs(centroid_offset)) > 4 / 3600
  ):
- # iterate over sources to only keep pixels within dist_lim
+ # iterate over sources to only keep pixels within self.sparse_dist_lim
  # this is inefficient, could be done in a tiled manner? only for squared data
  dra, ddec, sparse_mask = [], [], []
  for i in tqdm(
@@ -328,17 +332,21 @@ def _get_source_mask(
  correct_centroid_offset=True,
  plot=False,
  ):
- """Find the pixel mask that identifies pixels with contributions from ANY NUMBER of Sources
+ """Find the pixel mask that identifies pixels with contributions from ANY
+ NUMBER of Sources
 
- Fits a simple polynomial model to the log of the pixel flux values, in radial dimension and source flux,
- to find the optimum circular apertures for every source.
+ Fits a simple polynomial model to the log of the pixel flux values, in radial
+ dimension and source flux, to find the optimum circular apertures for every
+ source.
 
  Parameters
  ----------
  upper_radius_limit: float
- The radius limit at which we assume there is no flux from a source of any brightness (arcsec)
+ The radius limit at which we assume there is no flux from a source of
+ any brightness (arcsec)
  lower_radius_limit: float
- The radius limit at which we assume there is flux from a source of any brightness (arcsec)
+ The radius limit at which we assume there is flux from a source of any
+ brightness (arcsec)
  upper_flux_limit: float
  The flux at which we assume as source is saturated
  lower_flux_limit: float
@@ -349,12 +357,6 @@ def _get_source_mask(
  plot: bool
  Whether to show diagnostic plot. Default is False
  """
-
- if not hasattr(self, "source_flux_estimates"):
- # gaia estimate flux values per pixel to be used as flux priors
- self.source_flux_estimates = np.copy(
- np.asarray(self.sources.phot_g_mean_flux)
- )
  # We will use the radius a lot, this is for readibility
  # don't do if sparse array
  if isinstance(self.r, u.quantity.Quantity):
@@ -365,14 +367,15 @@ def _get_source_mask(
  # The average flux, which we assume is a good estimate of the whole stack of images
  max_flux = np.nanmax(self.flux[self.time_mask], axis=0)
 
- # Mask out sources that are above the flux limit, and pixels above the radius limit
+ # Mask out sources that are above the flux limit, and pixels above the
+ # radius limit. This is a good estimation only when using phot_g_mean_flux
  source_rad = 0.5 * np.log10(self.source_flux_estimates) ** 1.5 + 3
  # temp_mask for the sparse array case should also be a sparse matrix. Then it is
  # applied to r, max_flux, and, source_flux_estimates to be used later.
  # Numpy array case:
  if not isinstance(r, sparse.csr_matrix):
  # First we make a guess that each source has exactly the gaia flux
- source_flux_estimates = np.asarray(self.sources.phot_g_mean_flux)[
+ source_flux_estimates = np.asarray(self.source_flux_estimates)[
  :, None
  ] * np.ones((self.nsources, self.npixels))
  temp_mask = (r < source_rad[:, None]) & (
@@ -528,7 +531,7 @@ def _get_source_mask(
  if plot:
  k = np.isfinite(f_temp_mask)
  # Make a nice diagnostic plot
- fig, ax = plt.subplots(1, 2, figsize=(8, 3), facecolor="white")
+ fig, ax = plt.subplots(1, 2, figsize=(11, 3), facecolor="white")
 
  ax[0].scatter(r_temp_mask[k], f_temp_mask[k], s=0.4, c="k", label="Data")
  ax[0].scatter(r_temp_mask[k], A[k].dot(w), c="r", s=0.4, label="Model")
@@ -573,8 +576,7 @@ def _get_uncontaminated_pixel_mask(self):
  # we flag sources fainter than mag_limit as non-contaminant
  if isinstance(self.contaminant_mag_limit, (float, int)):
  aux = self.source_mask.multiply(
- self.sources.phot_g_mean_mag.values[:, None]
- < self.contaminant_mag_limit
+ self.source_flux_estimates[:, None] < self.contaminant_mag_limit
  )
  aux.eliminate_zeros()
  self.uncontaminated_source_mask = aux.multiply(
@@ -957,16 +959,10 @@ def build_shape_model(
  Keyword arguments to be passed to `_get_source_mask()`
  """
 
- # gaia estimate flux values per pixel to be used as flux priors
- self.source_flux_estimates = np.copy(np.asarray(self.sources.phot_g_mean_flux))
-
  # Mask of shape nsources x number of pixels, one where flux from a
  # source exists
- if not hasattr(self, "source_mask"):
- self._get_source_mask(**kwargs)
- # Mask of shape npixels (maybe by nt) where not saturated, not faint,
- # not contaminated etc
- self._get_uncontaminated_pixel_mask()
+ # if not hasattr(self, "source_mask"):
+ self._get_source_mask(**kwargs)
 
  # for iter in range(niters):
  flux_estimates = self.source_flux_estimates[:, None]
@@ -1346,24 +1342,23 @@ def get_psf_metrics(self, npoints_per_pixel=10):
  )
  self.perturbed_std = np.nanstd(perturbed_lcs, axis=0)
 
- def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
+ def plot_shape_model(self, frame_index="mean", bin_data=False):
  """
  Diagnostic plot of shape model.
 
  Parameters
  ----------
- radius : float
- Radius (in arcseconds) limit to be shown in the figure.
  frame_index : string or int
  The frame index used to plot the shape model, if "mean" then use the
  mean value across time
+ bin_data : bool
+ Bin or not the pixel data in a 2D historgram, default is False.
 
  Returns
  -------
  fig : matplotlib.Figure
  Figure.
  """
-
  if frame_index == "mean":
  mean_f = np.log10(
  self.uncontaminated_source_mask.astype(float)
@@ -1386,6 +1381,9 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
  dx = dx.data * u.deg.to(u.arcsecond)
  dy = dy.data * u.deg.to(u.arcsecond)
 
+ radius = np.maximum(np.abs(dx).max(), np.abs(dy).max()) * 1.1
+ vmin, vmax = np.nanpercentile(mean_f, [5, 90])
+
  if bin_data:
  nbins = 30 if mean_f.shape[0] <= 5e3 else 90
  _, dx, dy, mean_f, _ = threshold_bin(
@@ -1394,7 +1392,14 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
 
  fig, ax = plt.subplots(3, 2, figsize=(9, 10.5), constrained_layout=True)
  im = ax[0, 0].scatter(
- dx, dy, c=mean_f, cmap="viridis", vmin=-3, vmax=-1, s=3, rasterized=True
+ dx,
+ dy,
+ c=mean_f,
+ cmap="viridis",
+ vmin=vmin,
+ vmax=vmax,
+ s=3,
+ rasterized=True,
  )
  ax[0, 0].set(
  ylabel=r'$\delta y$ ["]',
@@ -1418,7 +1423,7 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
 
  phi, r = np.arctan2(dy, dx), np.hypot(dx, dy)
  im = ax[0, 1].scatter(
- phi, r, c=mean_f, cmap="viridis", vmin=-3, vmax=-1, s=3, rasterized=True
+ phi, r, c=mean_f, cmap="viridis", vmin=vmin, vmax=vmax, s=3, rasterized=True
  )
  ax[0, 1].set(
  ylabel='$r$ ["]',
@@ -1448,8 +1453,8 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
  r,
  c=model,
  cmap="viridis",
- vmin=-3,
- vmax=-1,
+ vmin=vmin,
+ vmax=vmax,
  s=3,
  rasterized=True,
  )
@@ -1464,8 +1469,8 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
  dy,
  c=model,
  cmap="viridis",
- vmin=-3,
- vmax=-1,
+ vmin=vmin,
+ vmax=vmax,
  s=3,
  rasterized=True,
  )

src/psfmachine/superstamp.py

@@ -15,7 +15,8 @@
 from astropy.wcs import WCS
 import astropy.units as u
 
-from .ffi import FFIMachine, _get_sources, _do_image_cutout
+from .ffi import FFIMachine, _get_sources
+from .utils import _do_image_cutout
 
 __all__ = ["SSMachine"]
 

src/psfmachine/tpf.py

@@ -52,9 +52,10 @@ def __init__(
  pos_corr2=None,
  focus_mask=None,
  tpf_meta=None,
- time_corrector="pos_corr",
  cartesian_knot_spacing="sqrt",
  bkg_subtracted=True,
+ sparse_dist_lim=40,
+ sources_flux_column="phot_g_mean_flux",
  ):
  super().__init__(
  time=time,
@@ -73,6 +74,8 @@ def __init__(
  time_radius=time_radius,
  rmin=rmin,
  rmax=rmax,
+ sparse_dist_lim=sparse_dist_lim,
+ sources_flux_column=sources_flux_column,
  )
  self.tpfs = tpfs
  # match cadences
@@ -93,7 +96,6 @@ def __init__(
  self.pos_corr1 = pos_corr1
  self.pos_corr2 = pos_corr2
  self.tpf_meta = tpf_meta
- self.time_corrector = time_corrector
  self.cartesian_knot_spacing = cartesian_knot_spacing
  self.bkg_subtracted = bkg_subtracted
  # limiting mag values to consider contamination for kepler & tess
@@ -153,7 +155,8 @@ def remove_background_model(
  )
  return
  # invert maks to get bkg pixels from TPFs
- self._get_source_mask()
+ if not hasattr(self, "source_mask"):
+ self._get_source_mask(correct_centroid_offset=True)
  bkg_pixel_mask = ~np.asarray(
  (self.source_mask.todense()).sum(axis=0).astype(bool)
  ).ravel()
@@ -854,7 +857,8 @@ def get_source_centroids(self, method="poscor"):
  def from_TPFs(
  tpfs,
  magnitude_limit=18,
- dr=2,
+ dr=3,
+ sources=None,
  time_mask=None,
  apply_focus_mask=True,
  renormalize_tpf_bkg=False,
@@ -882,7 +886,9 @@ def from_TPFs(
  magnitude_limit : float
  Limiting magnitude to query Gaia catalog.
  dr : int
- Gaia data release to be use, default is 2, options are DR2 and EDR3
+ Gaia data release to be use, default is 3, options are DR2 and EDR3
+ sources : pandas.DataFrame
+ Catalog with sources to be extracted by PSFMachine
  time_mask : boolean array
  Mask to be applied to discard cadences if needed.
  apply_focus_mask : boolean
@@ -1008,9 +1014,10 @@ def from_TPFs(
  saturated_mask,
  )
 
- sources = _get_coord_and_query_gaia(
- tpfs, magnitude_limit, dr=dr, ra=query_ra, dec=query_dec, rad=query_rad
- )
+ if sources is None:
+ sources = _get_coord_and_query_gaia(
+ tpfs, magnitude_limit, dr=dr, ra=query_ra, dec=query_dec, rad=query_rad
+ )
 
  def get_tpf2source():
  tpf2source = []
@@ -1052,6 +1059,20 @@ def get_tpf2source():
  sources["tpf_id"] = None
  sources.loc[idx[match], "tpf_id"] = np.asarray(tpf_meta["targetid"])[match]
 
+ # select flux column o be used as estimate.
+ if tpf_meta["mission"][0] == "TESS":
+ sources_flux_column = "phot_rp_mean_flux"
+ if np.isnan(sources[sources_flux_column].values).any():
+ log.warning(
+ f"Some sources do not have values in 'gaia.{sources_flux_column}''."
+ f" Will default to 'gaia.phot_g_mean_flux' for missing values."
+ )
+ sources[sources_flux_column].fillna(
+ sources["phot_g_mean_flux"], inplace=True
+ )
+ else:
+ sources_flux_column = "phot_g_mean_flux"
+
  # return a Machine object
  return TPFMachine(
  tpfs=tpfs,
@@ -1070,6 +1091,7 @@ def get_tpf2source():
  tpf_meta=tpf_meta,
  time_mask=time_mask,
  bkg_subtracted=not renormalize_tpf_bkg,
+ sources_flux_column=sources_flux_column,
  **kwargs,
  )
 
@@ -1121,7 +1143,7 @@ def _parse_TPFs(tpfs, renormalize_tpf_bkg=True, **kwargs):
 
  elif isinstance(tpfs[0], lk.TessTargetPixelFile):
  qual_mask = lk.utils.TessQualityFlags.create_quality_mask(
- tpfs[0].quality, lk.utils.TessQualityFlags.DEFAULT_BITMASK
+ tpfs[0].quality, 1 | 2 | 4 | 8 | 32 | 128
  )
  qual_mask &= (np.abs(tpfs[0].pos_corr1) < 5) & (np.abs(tpfs[0].pos_corr2) < 5)
  focus_mask = np.ones(len(tpfs[0].time), bool)[qual_mask]

src/psfmachine/machine.py

@@ -59,6 +59,7 @@ def __init__(
  rmax=16,
  cut_r=6,
  sparse_dist_lim=40,
+ sources_flux_column="phot_g_mean_flux",
  ):
  """
  Parameters
@@ -104,6 +105,9 @@ def __init__(
  Radial distance used to include pixels around sources when creating delta
  arrays (dra, ddec, r, and phi) as sparse matrices for efficiency.
  Default is 40" (recommended for kepler). (arcseconds)
+ sources_flux_column : str
+ Column name in `sources` table to be used as flux estimate. For Kepler data
+ gaia.phot_g_mean_flux is recommended, for TESS use gaia.phot_rp_mean_flux.
 
  Attributes
  ----------
@@ -172,6 +176,8 @@ def __init__(
  self.quiet = False
  self.contaminant_mag_limit = None
 
+ self.source_flux_estimates = np.copy(self.sources[sources_flux_column].values)
+
  if time_mask is None:
  self.time_mask = np.ones(len(time), bool)
  else:
@@ -231,7 +237,7 @@ def _create_delta_arrays(self, centroid_offset=[0, 0]):
  self.r = np.hypot(self.dra, self.ddec).to("arcsec")
  self.phi = np.arctan2(self.ddec, self.dra)
 
- def _create_delta_sparse_arrays(self, dist_lim=40, centroid_offset=[0, 0]):
+ def _create_delta_sparse_arrays(self, centroid_offset=[0, 0]):
  """
  Creates dra, ddec, r and phi arrays as sparse arrays to be used for dense data,
  e.g. Kepler FFIs or cluster fields. Assuming that there is no flux information
@@ -243,8 +249,6 @@ def _create_delta_sparse_arrays(self, dist_lim=40, centroid_offset=[0, 0]):
 
  Parameters
  ----------
- dist_lim : float
- Distance limit (in arcsecds) at which pixels are keep.
  centroid_offset : list
  Centroid offset for [ra, dec] to be included in dra and ddec computation.
  Default is [0, 0].
@@ -254,7 +258,7 @@ def _create_delta_sparse_arrays(self, dist_lim=40, centroid_offset=[0, 0]):
  if (centroid_offset[0] == centroid_offset[1] == 0) or (
  np.maximum(*np.abs(centroid_offset)) > 4 / 3600
  ):
- # iterate over sources to only keep pixels within dist_lim
+ # iterate over sources to only keep pixels within self.sparse_dist_lim
  # this is inefficient, could be done in a tiled manner? only for squared data
  dra, ddec, sparse_mask = [], [], []
  for i in tqdm(
@@ -328,17 +332,21 @@ def _get_source_mask(
  correct_centroid_offset=True,
  plot=False,
  ):
- """Find the pixel mask that identifies pixels with contributions from ANY NUMBER of Sources
+ """Find the pixel mask that identifies pixels with contributions from ANY
+ NUMBER of Sources
 
- Fits a simple polynomial model to the log of the pixel flux values, in radial dimension and source flux,
- to find the optimum circular apertures for every source.
+ Fits a simple polynomial model to the log of the pixel flux values, in radial
+ dimension and source flux, to find the optimum circular apertures for every
+ source.
 
  Parameters
  ----------
  upper_radius_limit: float
- The radius limit at which we assume there is no flux from a source of any brightness (arcsec)
+ The radius limit at which we assume there is no flux from a source of
+ any brightness (arcsec)
  lower_radius_limit: float
- The radius limit at which we assume there is flux from a source of any brightness (arcsec)
+ The radius limit at which we assume there is flux from a source of any
+ brightness (arcsec)
  upper_flux_limit: float
  The flux at which we assume as source is saturated
  lower_flux_limit: float
@@ -349,12 +357,6 @@ def _get_source_mask(
  plot: bool
  Whether to show diagnostic plot. Default is False
  """
-
- if not hasattr(self, "source_flux_estimates"):
- # gaia estimate flux values per pixel to be used as flux priors
- self.source_flux_estimates = np.copy(
- np.asarray(self.sources.phot_g_mean_flux)
- )
  # We will use the radius a lot, this is for readibility
  # don't do if sparse array
  if isinstance(self.r, u.quantity.Quantity):
@@ -365,14 +367,15 @@ def _get_source_mask(
  # The average flux, which we assume is a good estimate of the whole stack of images
  max_flux = np.nanmax(self.flux[self.time_mask], axis=0)
 
- # Mask out sources that are above the flux limit, and pixels above the radius limit
+ # Mask out sources that are above the flux limit, and pixels above the
+ # radius limit. This is a good estimation only when using phot_g_mean_flux
  source_rad = 0.5 * np.log10(self.source_flux_estimates) ** 1.5 + 3
  # temp_mask for the sparse array case should also be a sparse matrix. Then it is
  # applied to r, max_flux, and, source_flux_estimates to be used later.
  # Numpy array case:
  if not isinstance(r, sparse.csr_matrix):
  # First we make a guess that each source has exactly the gaia flux
- source_flux_estimates = np.asarray(self.sources.phot_g_mean_flux)[
+ source_flux_estimates = np.asarray(self.source_flux_estimates)[
  :, None
  ] * np.ones((self.nsources, self.npixels))
  temp_mask = (r < source_rad[:, None]) & (
@@ -528,7 +531,7 @@ def _get_source_mask(
  if plot:
  k = np.isfinite(f_temp_mask)
  # Make a nice diagnostic plot
- fig, ax = plt.subplots(1, 2, figsize=(8, 3), facecolor="white")
+ fig, ax = plt.subplots(1, 2, figsize=(11, 3), facecolor="white")
 
  ax[0].scatter(r_temp_mask[k], f_temp_mask[k], s=0.4, c="k", label="Data")
  ax[0].scatter(r_temp_mask[k], A[k].dot(w), c="r", s=0.4, label="Model")
@@ -573,8 +576,7 @@ def _get_uncontaminated_pixel_mask(self):
  # we flag sources fainter than mag_limit as non-contaminant
  if isinstance(self.contaminant_mag_limit, (float, int)):
  aux = self.source_mask.multiply(
- self.sources.phot_g_mean_mag.values[:, None]
- < self.contaminant_mag_limit
+ self.source_flux_estimates[:, None] < self.contaminant_mag_limit
  )
  aux.eliminate_zeros()
  self.uncontaminated_source_mask = aux.multiply(
@@ -957,16 +959,10 @@ def build_shape_model(
  Keyword arguments to be passed to `_get_source_mask()`
  """
 
- # gaia estimate flux values per pixel to be used as flux priors
- self.source_flux_estimates = np.copy(np.asarray(self.sources.phot_g_mean_flux))
-
  # Mask of shape nsources x number of pixels, one where flux from a
  # source exists
- if not hasattr(self, "source_mask"):
- self._get_source_mask(**kwargs)
- # Mask of shape npixels (maybe by nt) where not saturated, not faint,
- # not contaminated etc
- self._get_uncontaminated_pixel_mask()
+ # if not hasattr(self, "source_mask"):
+ self._get_source_mask(**kwargs)
 
  # for iter in range(niters):
  flux_estimates = self.source_flux_estimates[:, None]
@@ -1346,24 +1342,23 @@ def get_psf_metrics(self, npoints_per_pixel=10):
  )
  self.perturbed_std = np.nanstd(perturbed_lcs, axis=0)
 
- def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
+ def plot_shape_model(self, frame_index="mean", bin_data=False):
  """
  Diagnostic plot of shape model.
 
  Parameters
  ----------
- radius : float
- Radius (in arcseconds) limit to be shown in the figure.
  frame_index : string or int
  The frame index used to plot the shape model, if "mean" then use the
  mean value across time
+ bin_data : bool
+ Bin or not the pixel data in a 2D historgram, default is False.
 
  Returns
  -------
  fig : matplotlib.Figure
  Figure.
  """
-
  if frame_index == "mean":
  mean_f = np.log10(
  self.uncontaminated_source_mask.astype(float)
@@ -1386,6 +1381,9 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
  dx = dx.data * u.deg.to(u.arcsecond)
  dy = dy.data * u.deg.to(u.arcsecond)
 
+ radius = np.maximum(np.abs(dx).max(), np.abs(dy).max()) * 1.1
+ vmin, vmax = np.nanpercentile(mean_f, [5, 90])
+
  if bin_data:
  nbins = 30 if mean_f.shape[0] <= 5e3 else 90
  _, dx, dy, mean_f, _ = threshold_bin(
@@ -1394,7 +1392,14 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
 
  fig, ax = plt.subplots(3, 2, figsize=(9, 10.5), constrained_layout=True)
  im = ax[0, 0].scatter(
- dx, dy, c=mean_f, cmap="viridis", vmin=-3, vmax=-1, s=3, rasterized=True
+ dx,
+ dy,
+ c=mean_f,
+ cmap="viridis",
+ vmin=vmin,
+ vmax=vmax,
+ s=3,
+ rasterized=True,
  )
  ax[0, 0].set(
  ylabel=r'$\delta y$ ["]',
@@ -1418,7 +1423,7 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
 
  phi, r = np.arctan2(dy, dx), np.hypot(dx, dy)
  im = ax[0, 1].scatter(
- phi, r, c=mean_f, cmap="viridis", vmin=-3, vmax=-1, s=3, rasterized=True
+ phi, r, c=mean_f, cmap="viridis", vmin=vmin, vmax=vmax, s=3, rasterized=True
  )
  ax[0, 1].set(
  ylabel='$r$ ["]',
@@ -1448,8 +1453,8 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
  r,
  c=model,
  cmap="viridis",
- vmin=-3,
- vmax=-1,
+ vmin=vmin,
+ vmax=vmax,
  s=3,
  rasterized=True,
  )
@@ -1464,8 +1469,8 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
  dy,
  c=model,
  cmap="viridis",
- vmin=-3,
- vmax=-1,
+ vmin=vmin,
+ vmax=vmax,
  s=3,
  rasterized=True,
  )

src/psfmachine/superstamp.py

@@ -15,7 +15,8 @@
 from astropy.wcs import WCS
 import astropy.units as u
 
-from .ffi import FFIMachine, _get_sources, _do_image_cutout
+from .ffi import FFIMachine, _get_sources
+from .utils import _do_image_cutout
 
 __all__ = ["SSMachine"]
 

src/psfmachine/tpf.py

@@ -52,9 +52,10 @@ def __init__(
  pos_corr2=None,
  focus_mask=None,
  tpf_meta=None,
- time_corrector="pos_corr",
  cartesian_knot_spacing="sqrt",
  bkg_subtracted=True,
+ sparse_dist_lim=40,
+ sources_flux_column="phot_g_mean_flux",
  ):
  super().__init__(
  time=time,
@@ -73,6 +74,8 @@ def __init__(
  time_radius=time_radius,
  rmin=rmin,
  rmax=rmax,
+ sparse_dist_lim=sparse_dist_lim,
+ sources_flux_column=sources_flux_column,
  )
  self.tpfs = tpfs
  # match cadences
@@ -93,7 +96,6 @@ def __init__(
  self.pos_corr1 = pos_corr1
  self.pos_corr2 = pos_corr2
  self.tpf_meta = tpf_meta
- self.time_corrector = time_corrector
  self.cartesian_knot_spacing = cartesian_knot_spacing
  self.bkg_subtracted = bkg_subtracted
  # limiting mag values to consider contamination for kepler & tess
@@ -153,7 +155,8 @@ def remove_background_model(
  )
  return
  # invert maks to get bkg pixels from TPFs
- self._get_source_mask()
+ if not hasattr(self, "source_mask"):
+ self._get_source_mask(correct_centroid_offset=True)
  bkg_pixel_mask = ~np.asarray(
  (self.source_mask.todense()).sum(axis=0).astype(bool)
  ).ravel()
@@ -854,7 +857,8 @@ def get_source_centroids(self, method="poscor"):
  def from_TPFs(
  tpfs,
  magnitude_limit=18,
- dr=2,
+ dr=3,
+ sources=None,
  time_mask=None,
  apply_focus_mask=True,
  renormalize_tpf_bkg=False,
@@ -882,7 +886,9 @@ def from_TPFs(
  magnitude_limit : float
  Limiting magnitude to query Gaia catalog.
  dr : int
- Gaia data release to be use, default is 2, options are DR2 and EDR3
+ Gaia data release to be use, default is 3, options are DR2 and EDR3
+ sources : pandas.DataFrame
+ Catalog with sources to be extracted by PSFMachine
  time_mask : boolean array
  Mask to be applied to discard cadences if needed.
  apply_focus_mask : boolean
@@ -1008,9 +1014,10 @@ def from_TPFs(
  saturated_mask,
  )
 
- sources = _get_coord_and_query_gaia(
- tpfs, magnitude_limit, dr=dr, ra=query_ra, dec=query_dec, rad=query_rad
- )
+ if sources is None:
+ sources = _get_coord_and_query_gaia(
+ tpfs, magnitude_limit, dr=dr, ra=query_ra, dec=query_dec, rad=query_rad
+ )
 
  def get_tpf2source():
  tpf2source = []
@@ -1052,6 +1059,20 @@ def get_tpf2source():
  sources["tpf_id"] = None
  sources.loc[idx[match], "tpf_id"] = np.asarray(tpf_meta["targetid"])[match]
 
+ # select flux column o be used as estimate.
+ if tpf_meta["mission"][0] == "TESS":
+ sources_flux_column = "phot_rp_mean_flux"
+ if np.isnan(sources[sources_flux_column].values).any():
+ log.warning(
+ f"Some sources do not have values in 'gaia.{sources_flux_column}''."
+ f" Will default to 'gaia.phot_g_mean_flux' for missing values."
+ )
+ sources[sources_flux_column].fillna(
+ sources["phot_g_mean_flux"], inplace=True
+ )
+ else:
+ sources_flux_column = "phot_g_mean_flux"
+
  # return a Machine object
  return TPFMachine(
  tpfs=tpfs,
@@ -1070,6 +1091,7 @@ def get_tpf2source():
  tpf_meta=tpf_meta,
  time_mask=time_mask,
  bkg_subtracted=not renormalize_tpf_bkg,
+ sources_flux_column=sources_flux_column,
  **kwargs,
  )
 
@@ -1121,7 +1143,7 @@ def _parse_TPFs(tpfs, renormalize_tpf_bkg=True, **kwargs):
 
  elif isinstance(tpfs[0], lk.TessTargetPixelFile):
  qual_mask = lk.utils.TessQualityFlags.create_quality_mask(
- tpfs[0].quality, lk.utils.TessQualityFlags.DEFAULT_BITMASK
+ tpfs[0].quality, 1 | 2 | 4 | 8 | 32 | 128
  )
  qual_mask &= (np.abs(tpfs[0].pos_corr1) < 5) & (np.abs(tpfs[0].pos_corr2) < 5)
  focus_mask = np.ones(len(tpfs[0].time), bool)[qual_mask]

src/psfmachine/machine.py

@@ -59,6 +59,7 @@ def __init__(
  rmax=16,
  cut_r=6,
  sparse_dist_lim=40,
+ sources_flux_column="phot_g_mean_flux",
  ):
  """
  Parameters
@@ -104,6 +105,9 @@ def __init__(
  Radial distance used to include pixels around sources when creating delta
  arrays (dra, ddec, r, and phi) as sparse matrices for efficiency.
  Default is 40" (recommended for kepler). (arcseconds)
+ sources_flux_column : str
+ Column name in `sources` table to be used as flux estimate. For Kepler data
+ gaia.phot_g_mean_flux is recommended, for TESS use gaia.phot_rp_mean_flux.
 
  Attributes
  ----------
@@ -172,6 +176,8 @@ def __init__(
  self.quiet = False
  self.contaminant_mag_limit = None
 
+ self.source_flux_estimates = np.copy(self.sources[sources_flux_column].values)
+
  if time_mask is None:
  self.time_mask = np.ones(len(time), bool)
  else:
@@ -231,7 +237,7 @@ def _create_delta_arrays(self, centroid_offset=[0, 0]):
  self.r = np.hypot(self.dra, self.ddec).to("arcsec")
  self.phi = np.arctan2(self.ddec, self.dra)
 
- def _create_delta_sparse_arrays(self, dist_lim=40, centroid_offset=[0, 0]):
+ def _create_delta_sparse_arrays(self, centroid_offset=[0, 0]):
  """
  Creates dra, ddec, r and phi arrays as sparse arrays to be used for dense data,
  e.g. Kepler FFIs or cluster fields. Assuming that there is no flux information
@@ -243,8 +249,6 @@ def _create_delta_sparse_arrays(self, dist_lim=40, centroid_offset=[0, 0]):
 
  Parameters
  ----------
- dist_lim : float
- Distance limit (in arcsecds) at which pixels are keep.
  centroid_offset : list
  Centroid offset for [ra, dec] to be included in dra and ddec computation.
  Default is [0, 0].
@@ -254,7 +258,7 @@ def _create_delta_sparse_arrays(self, dist_lim=40, centroid_offset=[0, 0]):
  if (centroid_offset[0] == centroid_offset[1] == 0) or (
  np.maximum(*np.abs(centroid_offset)) > 4 / 3600
  ):
- # iterate over sources to only keep pixels within dist_lim
+ # iterate over sources to only keep pixels within self.sparse_dist_lim
  # this is inefficient, could be done in a tiled manner? only for squared data
  dra, ddec, sparse_mask = [], [], []
  for i in tqdm(
@@ -328,17 +332,21 @@ def _get_source_mask(
  correct_centroid_offset=True,
  plot=False,
  ):
- """Find the pixel mask that identifies pixels with contributions from ANY NUMBER of Sources
+ """Find the pixel mask that identifies pixels with contributions from ANY
+ NUMBER of Sources
 
- Fits a simple polynomial model to the log of the pixel flux values, in radial dimension and source flux,
- to find the optimum circular apertures for every source.
+ Fits a simple polynomial model to the log of the pixel flux values, in radial
+ dimension and source flux, to find the optimum circular apertures for every
+ source.
 
  Parameters
  ----------
  upper_radius_limit: float
- The radius limit at which we assume there is no flux from a source of any brightness (arcsec)
+ The radius limit at which we assume there is no flux from a source of
+ any brightness (arcsec)
  lower_radius_limit: float
- The radius limit at which we assume there is flux from a source of any brightness (arcsec)
+ The radius limit at which we assume there is flux from a source of any
+ brightness (arcsec)
  upper_flux_limit: float
  The flux at which we assume as source is saturated
  lower_flux_limit: float
@@ -349,12 +357,6 @@ def _get_source_mask(
  plot: bool
  Whether to show diagnostic plot. Default is False
  """
-
- if not hasattr(self, "source_flux_estimates"):
- # gaia estimate flux values per pixel to be used as flux priors
- self.source_flux_estimates = np.copy(
- np.asarray(self.sources.phot_g_mean_flux)
- )
  # We will use the radius a lot, this is for readibility
  # don't do if sparse array
  if isinstance(self.r, u.quantity.Quantity):
@@ -365,14 +367,15 @@ def _get_source_mask(
  # The average flux, which we assume is a good estimate of the whole stack of images
  max_flux = np.nanmax(self.flux[self.time_mask], axis=0)
 
- # Mask out sources that are above the flux limit, and pixels above the radius limit
+ # Mask out sources that are above the flux limit, and pixels above the
+ # radius limit. This is a good estimation only when using phot_g_mean_flux
  source_rad = 0.5 * np.log10(self.source_flux_estimates) ** 1.5 + 3
  # temp_mask for the sparse array case should also be a sparse matrix. Then it is
  # applied to r, max_flux, and, source_flux_estimates to be used later.
  # Numpy array case:
  if not isinstance(r, sparse.csr_matrix):
  # First we make a guess that each source has exactly the gaia flux
- source_flux_estimates = np.asarray(self.sources.phot_g_mean_flux)[
+ source_flux_estimates = np.asarray(self.source_flux_estimates)[
  :, None
  ] * np.ones((self.nsources, self.npixels))
  temp_mask = (r < source_rad[:, None]) & (
@@ -528,7 +531,7 @@ def _get_source_mask(
  if plot:
  k = np.isfinite(f_temp_mask)
  # Make a nice diagnostic plot
- fig, ax = plt.subplots(1, 2, figsize=(8, 3), facecolor="white")
+ fig, ax = plt.subplots(1, 2, figsize=(11, 3), facecolor="white")
 
  ax[0].scatter(r_temp_mask[k], f_temp_mask[k], s=0.4, c="k", label="Data")
  ax[0].scatter(r_temp_mask[k], A[k].dot(w), c="r", s=0.4, label="Model")
@@ -573,8 +576,7 @@ def _get_uncontaminated_pixel_mask(self):
  # we flag sources fainter than mag_limit as non-contaminant
  if isinstance(self.contaminant_mag_limit, (float, int)):
  aux = self.source_mask.multiply(
- self.sources.phot_g_mean_mag.values[:, None]
- < self.contaminant_mag_limit
+ self.source_flux_estimates[:, None] < self.contaminant_mag_limit
  )
  aux.eliminate_zeros()
  self.uncontaminated_source_mask = aux.multiply(
@@ -957,16 +959,10 @@ def build_shape_model(
  Keyword arguments to be passed to `_get_source_mask()`
  """
 
- # gaia estimate flux values per pixel to be used as flux priors
- self.source_flux_estimates = np.copy(np.asarray(self.sources.phot_g_mean_flux))
-
  # Mask of shape nsources x number of pixels, one where flux from a
  # source exists
- if not hasattr(self, "source_mask"):
- self._get_source_mask(**kwargs)
- # Mask of shape npixels (maybe by nt) where not saturated, not faint,
- # not contaminated etc
- self._get_uncontaminated_pixel_mask()
+ # if not hasattr(self, "source_mask"):
+ self._get_source_mask(**kwargs)
 
  # for iter in range(niters):
  flux_estimates = self.source_flux_estimates[:, None]
@@ -1346,24 +1342,23 @@ def get_psf_metrics(self, npoints_per_pixel=10):
  )
  self.perturbed_std = np.nanstd(perturbed_lcs, axis=0)
 
- def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
+ def plot_shape_model(self, frame_index="mean", bin_data=False):
  """
  Diagnostic plot of shape model.
 
  Parameters
  ----------
- radius : float
- Radius (in arcseconds) limit to be shown in the figure.
  frame_index : string or int
  The frame index used to plot the shape model, if "mean" then use the
  mean value across time
+ bin_data : bool
+ Bin or not the pixel data in a 2D historgram, default is False.
 
  Returns
  -------
  fig : matplotlib.Figure
  Figure.
  """
-
  if frame_index == "mean":
  mean_f = np.log10(
  self.uncontaminated_source_mask.astype(float)
@@ -1386,6 +1381,9 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
  dx = dx.data * u.deg.to(u.arcsecond)
  dy = dy.data * u.deg.to(u.arcsecond)
 
+ radius = np.maximum(np.abs(dx).max(), np.abs(dy).max()) * 1.1
+ vmin, vmax = np.nanpercentile(mean_f, [5, 90])
+
  if bin_data:
  nbins = 30 if mean_f.shape[0] <= 5e3 else 90
  _, dx, dy, mean_f, _ = threshold_bin(
@@ -1394,7 +1392,14 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
 
  fig, ax = plt.subplots(3, 2, figsize=(9, 10.5), constrained_layout=True)
  im = ax[0, 0].scatter(
- dx, dy, c=mean_f, cmap="viridis", vmin=-3, vmax=-1, s=3, rasterized=True
+ dx,
+ dy,
+ c=mean_f,
+ cmap="viridis",
+ vmin=vmin,
+ vmax=vmax,
+ s=3,
+ rasterized=True,
  )
  ax[0, 0].set(
  ylabel=r'$\delta y$ ["]',
@@ -1418,7 +1423,7 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
 
  phi, r = np.arctan2(dy, dx), np.hypot(dx, dy)
  im = ax[0, 1].scatter(
- phi, r, c=mean_f, cmap="viridis", vmin=-3, vmax=-1, s=3, rasterized=True
+ phi, r, c=mean_f, cmap="viridis", vmin=vmin, vmax=vmax, s=3, rasterized=True
  )
  ax[0, 1].set(
  ylabel='$r$ ["]',
@@ -1448,8 +1453,8 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
  r,
  c=model,
  cmap="viridis",
- vmin=-3,
- vmax=-1,
+ vmin=vmin,
+ vmax=vmax,
  s=3,
  rasterized=True,
  )
@@ -1464,8 +1469,8 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
  dy,
  c=model,
  cmap="viridis",
- vmin=-3,
- vmax=-1,
+ vmin=vmin,
+ vmax=vmax,
  s=3,
  rasterized=True,
  )

src/psfmachine/superstamp.py

@@ -15,7 +15,8 @@
 from astropy.wcs import WCS
 import astropy.units as u
 
-from .ffi import FFIMachine, _get_sources, _do_image_cutout
+from .ffi import FFIMachine, _get_sources
+from .utils import _do_image_cutout
 
 __all__ = ["SSMachine"]
 

src/psfmachine/tpf.py

@@ -52,9 +52,10 @@ def __init__(
  pos_corr2=None,
  focus_mask=None,
  tpf_meta=None,
- time_corrector="pos_corr",
  cartesian_knot_spacing="sqrt",
  bkg_subtracted=True,
+ sparse_dist_lim=40,
+ sources_flux_column="phot_g_mean_flux",
  ):
  super().__init__(
  time=time,
@@ -73,6 +74,8 @@ def __init__(
  time_radius=time_radius,
  rmin=rmin,
  rmax=rmax,
+ sparse_dist_lim=sparse_dist_lim,
+ sources_flux_column=sources_flux_column,
  )
  self.tpfs = tpfs
  # match cadences
@@ -93,7 +96,6 @@ def __init__(
  self.pos_corr1 = pos_corr1
  self.pos_corr2 = pos_corr2
  self.tpf_meta = tpf_meta
- self.time_corrector = time_corrector
  self.cartesian_knot_spacing = cartesian_knot_spacing
  self.bkg_subtracted = bkg_subtracted
  # limiting mag values to consider contamination for kepler & tess
@@ -153,7 +155,8 @@ def remove_background_model(
  )
  return
  # invert maks to get bkg pixels from TPFs
- self._get_source_mask()
+ if not hasattr(self, "source_mask"):
+ self._get_source_mask(correct_centroid_offset=True)
  bkg_pixel_mask = ~np.asarray(
  (self.source_mask.todense()).sum(axis=0).astype(bool)
  ).ravel()
@@ -854,7 +857,8 @@ def get_source_centroids(self, method="poscor"):
  def from_TPFs(
  tpfs,
  magnitude_limit=18,
- dr=2,
+ dr=3,
+ sources=None,
  time_mask=None,
  apply_focus_mask=True,
  renormalize_tpf_bkg=False,
@@ -882,7 +886,9 @@ def from_TPFs(
  magnitude_limit : float
  Limiting magnitude to query Gaia catalog.
  dr : int
- Gaia data release to be use, default is 2, options are DR2 and EDR3
+ Gaia data release to be use, default is 3, options are DR2 and EDR3
+ sources : pandas.DataFrame
+ Catalog with sources to be extracted by PSFMachine
  time_mask : boolean array
  Mask to be applied to discard cadences if needed.
  apply_focus_mask : boolean
@@ -1008,9 +1014,10 @@ def from_TPFs(
  saturated_mask,
  )
 
- sources = _get_coord_and_query_gaia(
- tpfs, magnitude_limit, dr=dr, ra=query_ra, dec=query_dec, rad=query_rad
- )
+ if sources is None:
+ sources = _get_coord_and_query_gaia(
+ tpfs, magnitude_limit, dr=dr, ra=query_ra, dec=query_dec, rad=query_rad
+ )
 
  def get_tpf2source():
  tpf2source = []
@@ -1052,6 +1059,20 @@ def get_tpf2source():
  sources["tpf_id"] = None
  sources.loc[idx[match], "tpf_id"] = np.asarray(tpf_meta["targetid"])[match]
 
+ # select flux column o be used as estimate.
+ if tpf_meta["mission"][0] == "TESS":
+ sources_flux_column = "phot_rp_mean_flux"
+ if np.isnan(sources[sources_flux_column].values).any():
+ log.warning(
+ f"Some sources do not have values in 'gaia.{sources_flux_column}''."
+ f" Will default to 'gaia.phot_g_mean_flux' for missing values."
+ )
+ sources[sources_flux_column].fillna(
+ sources["phot_g_mean_flux"], inplace=True
+ )
+ else:
+ sources_flux_column = "phot_g_mean_flux"
+
  # return a Machine object
  return TPFMachine(
  tpfs=tpfs,
@@ -1070,6 +1091,7 @@ def get_tpf2source():
  tpf_meta=tpf_meta,
  time_mask=time_mask,
  bkg_subtracted=not renormalize_tpf_bkg,
+ sources_flux_column=sources_flux_column,
  **kwargs,
  )
 
@@ -1121,7 +1143,7 @@ def _parse_TPFs(tpfs, renormalize_tpf_bkg=True, **kwargs):
 
  elif isinstance(tpfs[0], lk.TessTargetPixelFile):
  qual_mask = lk.utils.TessQualityFlags.create_quality_mask(
- tpfs[0].quality, lk.utils.TessQualityFlags.DEFAULT_BITMASK
+ tpfs[0].quality, 1 | 2 | 4 | 8 | 32 | 128
  )
  qual_mask &= (np.abs(tpfs[0].pos_corr1) < 5) & (np.abs(tpfs[0].pos_corr2) < 5)
  focus_mask = np.ones(len(tpfs[0].time), bool)[qual_mask]

src/psfmachine/machine.py

@@ -59,6 +59,7 @@ def __init__(
  rmax=16,
  cut_r=6,
  sparse_dist_lim=40,
+ sources_flux_column="phot_g_mean_flux",
  ):
  """
  Parameters
@@ -104,6 +105,9 @@ def __init__(
  Radial distance used to include pixels around sources when creating delta
  arrays (dra, ddec, r, and phi) as sparse matrices for efficiency.
  Default is 40" (recommended for kepler). (arcseconds)
+ sources_flux_column : str
+ Column name in `sources` table to be used as flux estimate. For Kepler data
+ gaia.phot_g_mean_flux is recommended, for TESS use gaia.phot_rp_mean_flux.
 
  Attributes
  ----------
@@ -172,6 +176,8 @@ def __init__(
  self.quiet = False
  self.contaminant_mag_limit = None
 
+ self.source_flux_estimates = np.copy(self.sources[sources_flux_column].values)
+
  if time_mask is None:
  self.time_mask = np.ones(len(time), bool)
  else:
@@ -231,7 +237,7 @@ def _create_delta_arrays(self, centroid_offset=[0, 0]):
  self.r = np.hypot(self.dra, self.ddec).to("arcsec")
  self.phi = np.arctan2(self.ddec, self.dra)
 
- def _create_delta_sparse_arrays(self, dist_lim=40, centroid_offset=[0, 0]):
+ def _create_delta_sparse_arrays(self, centroid_offset=[0, 0]):
  """
  Creates dra, ddec, r and phi arrays as sparse arrays to be used for dense data,
  e.g. Kepler FFIs or cluster fields. Assuming that there is no flux information
@@ -243,8 +249,6 @@ def _create_delta_sparse_arrays(self, dist_lim=40, centroid_offset=[0, 0]):
 
  Parameters
  ----------
- dist_lim : float
- Distance limit (in arcsecds) at which pixels are keep.
  centroid_offset : list
  Centroid offset for [ra, dec] to be included in dra and ddec computation.
  Default is [0, 0].
@@ -254,7 +258,7 @@ def _create_delta_sparse_arrays(self, dist_lim=40, centroid_offset=[0, 0]):
  if (centroid_offset[0] == centroid_offset[1] == 0) or (
  np.maximum(*np.abs(centroid_offset)) > 4 / 3600
  ):
- # iterate over sources to only keep pixels within dist_lim
+ # iterate over sources to only keep pixels within self.sparse_dist_lim
  # this is inefficient, could be done in a tiled manner? only for squared data
  dra, ddec, sparse_mask = [], [], []
  for i in tqdm(
@@ -328,17 +332,21 @@ def _get_source_mask(
  correct_centroid_offset=True,
  plot=False,
  ):
- """Find the pixel mask that identifies pixels with contributions from ANY NUMBER of Sources
+ """Find the pixel mask that identifies pixels with contributions from ANY
+ NUMBER of Sources
 
- Fits a simple polynomial model to the log of the pixel flux values, in radial dimension and source flux,
- to find the optimum circular apertures for every source.
+ Fits a simple polynomial model to the log of the pixel flux values, in radial
+ dimension and source flux, to find the optimum circular apertures for every
+ source.
 
  Parameters
  ----------
  upper_radius_limit: float
- The radius limit at which we assume there is no flux from a source of any brightness (arcsec)
+ The radius limit at which we assume there is no flux from a source of
+ any brightness (arcsec)
  lower_radius_limit: float
- The radius limit at which we assume there is flux from a source of any brightness (arcsec)
+ The radius limit at which we assume there is flux from a source of any
+ brightness (arcsec)
  upper_flux_limit: float
  The flux at which we assume as source is saturated
  lower_flux_limit: float
@@ -349,12 +357,6 @@ def _get_source_mask(
  plot: bool
  Whether to show diagnostic plot. Default is False
  """
-
- if not hasattr(self, "source_flux_estimates"):
- # gaia estimate flux values per pixel to be used as flux priors
- self.source_flux_estimates = np.copy(
- np.asarray(self.sources.phot_g_mean_flux)
- )
  # We will use the radius a lot, this is for readibility
  # don't do if sparse array
  if isinstance(self.r, u.quantity.Quantity):
@@ -365,14 +367,15 @@ def _get_source_mask(
  # The average flux, which we assume is a good estimate of the whole stack of images
  max_flux = np.nanmax(self.flux[self.time_mask], axis=0)
 
- # Mask out sources that are above the flux limit, and pixels above the radius limit
+ # Mask out sources that are above the flux limit, and pixels above the
+ # radius limit. This is a good estimation only when using phot_g_mean_flux
  source_rad = 0.5 * np.log10(self.source_flux_estimates) ** 1.5 + 3
  # temp_mask for the sparse array case should also be a sparse matrix. Then it is
  # applied to r, max_flux, and, source_flux_estimates to be used later.
  # Numpy array case:
  if not isinstance(r, sparse.csr_matrix):
  # First we make a guess that each source has exactly the gaia flux
- source_flux_estimates = np.asarray(self.sources.phot_g_mean_flux)[
+ source_flux_estimates = np.asarray(self.source_flux_estimates)[
  :, None
  ] * np.ones((self.nsources, self.npixels))
  temp_mask = (r < source_rad[:, None]) & (
@@ -528,7 +531,7 @@ def _get_source_mask(
  if plot:
  k = np.isfinite(f_temp_mask)
  # Make a nice diagnostic plot
- fig, ax = plt.subplots(1, 2, figsize=(8, 3), facecolor="white")
+ fig, ax = plt.subplots(1, 2, figsize=(11, 3), facecolor="white")
 
  ax[0].scatter(r_temp_mask[k], f_temp_mask[k], s=0.4, c="k", label="Data")
  ax[0].scatter(r_temp_mask[k], A[k].dot(w), c="r", s=0.4, label="Model")
@@ -573,8 +576,7 @@ def _get_uncontaminated_pixel_mask(self):
  # we flag sources fainter than mag_limit as non-contaminant
  if isinstance(self.contaminant_mag_limit, (float, int)):
  aux = self.source_mask.multiply(
- self.sources.phot_g_mean_mag.values[:, None]
- < self.contaminant_mag_limit
+ self.source_flux_estimates[:, None] < self.contaminant_mag_limit
  )
  aux.eliminate_zeros()
  self.uncontaminated_source_mask = aux.multiply(
@@ -957,16 +959,10 @@ def build_shape_model(
  Keyword arguments to be passed to `_get_source_mask()`
  """
 
- # gaia estimate flux values per pixel to be used as flux priors
- self.source_flux_estimates = np.copy(np.asarray(self.sources.phot_g_mean_flux))
-
  # Mask of shape nsources x number of pixels, one where flux from a
  # source exists
- if not hasattr(self, "source_mask"):
- self._get_source_mask(**kwargs)
- # Mask of shape npixels (maybe by nt) where not saturated, not faint,
- # not contaminated etc
- self._get_uncontaminated_pixel_mask()
+ # if not hasattr(self, "source_mask"):
+ self._get_source_mask(**kwargs)
 
  # for iter in range(niters):
  flux_estimates = self.source_flux_estimates[:, None]
@@ -1346,24 +1342,23 @@ def get_psf_metrics(self, npoints_per_pixel=10):
  )
  self.perturbed_std = np.nanstd(perturbed_lcs, axis=0)
 
- def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
+ def plot_shape_model(self, frame_index="mean", bin_data=False):
  """
  Diagnostic plot of shape model.
 
  Parameters
  ----------
- radius : float
- Radius (in arcseconds) limit to be shown in the figure.
  frame_index : string or int
  The frame index used to plot the shape model, if "mean" then use the
  mean value across time
+ bin_data : bool
+ Bin or not the pixel data in a 2D historgram, default is False.
 
  Returns
  -------
  fig : matplotlib.Figure
  Figure.
  """
-
  if frame_index == "mean":
  mean_f = np.log10(
  self.uncontaminated_source_mask.astype(float)
@@ -1386,6 +1381,9 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
  dx = dx.data * u.deg.to(u.arcsecond)
  dy = dy.data * u.deg.to(u.arcsecond)
 
+ radius = np.maximum(np.abs(dx).max(), np.abs(dy).max()) * 1.1
+ vmin, vmax = np.nanpercentile(mean_f, [5, 90])
+
  if bin_data:
  nbins = 30 if mean_f.shape[0] <= 5e3 else 90
  _, dx, dy, mean_f, _ = threshold_bin(
@@ -1394,7 +1392,14 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
 
  fig, ax = plt.subplots(3, 2, figsize=(9, 10.5), constrained_layout=True)
  im = ax[0, 0].scatter(
- dx, dy, c=mean_f, cmap="viridis", vmin=-3, vmax=-1, s=3, rasterized=True
+ dx,
+ dy,
+ c=mean_f,
+ cmap="viridis",
+ vmin=vmin,
+ vmax=vmax,
+ s=3,
+ rasterized=True,
  )
  ax[0, 0].set(
  ylabel=r'$\delta y$ ["]',
@@ -1418,7 +1423,7 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
 
  phi, r = np.arctan2(dy, dx), np.hypot(dx, dy)
  im = ax[0, 1].scatter(
- phi, r, c=mean_f, cmap="viridis", vmin=-3, vmax=-1, s=3, rasterized=True
+ phi, r, c=mean_f, cmap="viridis", vmin=vmin, vmax=vmax, s=3, rasterized=True
  )
  ax[0, 1].set(
  ylabel='$r$ ["]',
@@ -1448,8 +1453,8 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
  r,
  c=model,
  cmap="viridis",
- vmin=-3,
- vmax=-1,
+ vmin=vmin,
+ vmax=vmax,
  s=3,
  rasterized=True,
  )
@@ -1464,8 +1469,8 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
  dy,
  c=model,
  cmap="viridis",
- vmin=-3,
- vmax=-1,
+ vmin=vmin,
+ vmax=vmax,
  s=3,
  rasterized=True,
  )

src/psfmachine/superstamp.py

@@ -15,7 +15,8 @@
 from astropy.wcs import WCS
 import astropy.units as u
 
-from .ffi import FFIMachine, _get_sources, _do_image_cutout
+from .ffi import FFIMachine, _get_sources
+from .utils import _do_image_cutout
 
 __all__ = ["SSMachine"]
 

src/psfmachine/tpf.py

@@ -52,9 +52,10 @@ def __init__(
  pos_corr2=None,
  focus_mask=None,
  tpf_meta=None,
- time_corrector="pos_corr",
  cartesian_knot_spacing="sqrt",
  bkg_subtracted=True,
+ sparse_dist_lim=40,
+ sources_flux_column="phot_g_mean_flux",
  ):
  super().__init__(
  time=time,
@@ -73,6 +74,8 @@ def __init__(
  time_radius=time_radius,
  rmin=rmin,
  rmax=rmax,
+ sparse_dist_lim=sparse_dist_lim,
+ sources_flux_column=sources_flux_column,
  )
  self.tpfs = tpfs
  # match cadences
@@ -93,7 +96,6 @@ def __init__(
  self.pos_corr1 = pos_corr1
  self.pos_corr2 = pos_corr2
  self.tpf_meta = tpf_meta
- self.time_corrector = time_corrector
  self.cartesian_knot_spacing = cartesian_knot_spacing
  self.bkg_subtracted = bkg_subtracted
  # limiting mag values to consider contamination for kepler & tess
@@ -153,7 +155,8 @@ def remove_background_model(
  )
  return
  # invert maks to get bkg pixels from TPFs
- self._get_source_mask()
+ if not hasattr(self, "source_mask"):
+ self._get_source_mask(correct_centroid_offset=True)
  bkg_pixel_mask = ~np.asarray(
  (self.source_mask.todense()).sum(axis=0).astype(bool)
  ).ravel()
@@ -854,7 +857,8 @@ def get_source_centroids(self, method="poscor"):
  def from_TPFs(
  tpfs,
  magnitude_limit=18,
- dr=2,
+ dr=3,
+ sources=None,
  time_mask=None,
  apply_focus_mask=True,
  renormalize_tpf_bkg=False,
@@ -882,7 +886,9 @@ def from_TPFs(
  magnitude_limit : float
  Limiting magnitude to query Gaia catalog.
  dr : int
- Gaia data release to be use, default is 2, options are DR2 and EDR3
+ Gaia data release to be use, default is 3, options are DR2 and EDR3
+ sources : pandas.DataFrame
+ Catalog with sources to be extracted by PSFMachine
  time_mask : boolean array
  Mask to be applied to discard cadences if needed.
  apply_focus_mask : boolean
@@ -1008,9 +1014,10 @@ def from_TPFs(
  saturated_mask,
  )
 
- sources = _get_coord_and_query_gaia(
- tpfs, magnitude_limit, dr=dr, ra=query_ra, dec=query_dec, rad=query_rad
- )
+ if sources is None:
+ sources = _get_coord_and_query_gaia(
+ tpfs, magnitude_limit, dr=dr, ra=query_ra, dec=query_dec, rad=query_rad
+ )
 
  def get_tpf2source():
  tpf2source = []
@@ -1052,6 +1059,20 @@ def get_tpf2source():
  sources["tpf_id"] = None
  sources.loc[idx[match], "tpf_id"] = np.asarray(tpf_meta["targetid"])[match]
 
+ # select flux column o be used as estimate.
+ if tpf_meta["mission"][0] == "TESS":
+ sources_flux_column = "phot_rp_mean_flux"
+ if np.isnan(sources[sources_flux_column].values).any():
+ log.warning(
+ f"Some sources do not have values in 'gaia.{sources_flux_column}''."
+ f" Will default to 'gaia.phot_g_mean_flux' for missing values."
+ )
+ sources[sources_flux_column].fillna(
+ sources["phot_g_mean_flux"], inplace=True
+ )
+ else:
+ sources_flux_column = "phot_g_mean_flux"
+
  # return a Machine object
  return TPFMachine(
  tpfs=tpfs,
@@ -1070,6 +1091,7 @@ def get_tpf2source():
  tpf_meta=tpf_meta,
  time_mask=time_mask,
  bkg_subtracted=not renormalize_tpf_bkg,
+ sources_flux_column=sources_flux_column,
  **kwargs,
  )
 
@@ -1121,7 +1143,7 @@ def _parse_TPFs(tpfs, renormalize_tpf_bkg=True, **kwargs):
 
  elif isinstance(tpfs[0], lk.TessTargetPixelFile):
  qual_mask = lk.utils.TessQualityFlags.create_quality_mask(
- tpfs[0].quality, lk.utils.TessQualityFlags.DEFAULT_BITMASK
+ tpfs[0].quality, 1 | 2 | 4 | 8 | 32 | 128
  )
  qual_mask &= (np.abs(tpfs[0].pos_corr1) < 5) & (np.abs(tpfs[0].pos_corr2) < 5)
  focus_mask = np.ones(len(tpfs[0].time), bool)[qual_mask]

src/psfmachine/machine.py

@@ -59,6 +59,7 @@ def __init__(
  rmax=16,
  cut_r=6,
  sparse_dist_lim=40,
+ sources_flux_column="phot_g_mean_flux",
  ):
  """
  Parameters
@@ -104,6 +105,9 @@ def __init__(
  Radial distance used to include pixels around sources when creating delta
  arrays (dra, ddec, r, and phi) as sparse matrices for efficiency.
  Default is 40" (recommended for kepler). (arcseconds)
+ sources_flux_column : str
+ Column name in `sources` table to be used as flux estimate. For Kepler data
+ gaia.phot_g_mean_flux is recommended, for TESS use gaia.phot_rp_mean_flux.
 
  Attributes
  ----------
@@ -172,6 +176,8 @@ def __init__(
  self.quiet = False
  self.contaminant_mag_limit = None
 
+ self.source_flux_estimates = np.copy(self.sources[sources_flux_column].values)
+
  if time_mask is None:
  self.time_mask = np.ones(len(time), bool)
  else:
@@ -231,7 +237,7 @@ def _create_delta_arrays(self, centroid_offset=[0, 0]):
  self.r = np.hypot(self.dra, self.ddec).to("arcsec")
  self.phi = np.arctan2(self.ddec, self.dra)
 
- def _create_delta_sparse_arrays(self, dist_lim=40, centroid_offset=[0, 0]):
+ def _create_delta_sparse_arrays(self, centroid_offset=[0, 0]):
  """
  Creates dra, ddec, r and phi arrays as sparse arrays to be used for dense data,
  e.g. Kepler FFIs or cluster fields. Assuming that there is no flux information
@@ -243,8 +249,6 @@ def _create_delta_sparse_arrays(self, dist_lim=40, centroid_offset=[0, 0]):
 
  Parameters
  ----------
- dist_lim : float
- Distance limit (in arcsecds) at which pixels are keep.
  centroid_offset : list
  Centroid offset for [ra, dec] to be included in dra and ddec computation.
  Default is [0, 0].
@@ -254,7 +258,7 @@ def _create_delta_sparse_arrays(self, dist_lim=40, centroid_offset=[0, 0]):
  if (centroid_offset[0] == centroid_offset[1] == 0) or (
  np.maximum(*np.abs(centroid_offset)) > 4 / 3600
  ):
- # iterate over sources to only keep pixels within dist_lim
+ # iterate over sources to only keep pixels within self.sparse_dist_lim
  # this is inefficient, could be done in a tiled manner? only for squared data
  dra, ddec, sparse_mask = [], [], []
  for i in tqdm(
@@ -328,17 +332,21 @@ def _get_source_mask(
  correct_centroid_offset=True,
  plot=False,
  ):
- """Find the pixel mask that identifies pixels with contributions from ANY NUMBER of Sources
+ """Find the pixel mask that identifies pixels with contributions from ANY
+ NUMBER of Sources
 
- Fits a simple polynomial model to the log of the pixel flux values, in radial dimension and source flux,
- to find the optimum circular apertures for every source.
+ Fits a simple polynomial model to the log of the pixel flux values, in radial
+ dimension and source flux, to find the optimum circular apertures for every
+ source.
 
  Parameters
  ----------
  upper_radius_limit: float
- The radius limit at which we assume there is no flux from a source of any brightness (arcsec)
+ The radius limit at which we assume there is no flux from a source of
+ any brightness (arcsec)
  lower_radius_limit: float
- The radius limit at which we assume there is flux from a source of any brightness (arcsec)
+ The radius limit at which we assume there is flux from a source of any
+ brightness (arcsec)
  upper_flux_limit: float
  The flux at which we assume as source is saturated
  lower_flux_limit: float
@@ -349,12 +357,6 @@ def _get_source_mask(
  plot: bool
  Whether to show diagnostic plot. Default is False
  """
-
- if not hasattr(self, "source_flux_estimates"):
- # gaia estimate flux values per pixel to be used as flux priors
- self.source_flux_estimates = np.copy(
- np.asarray(self.sources.phot_g_mean_flux)
- )
  # We will use the radius a lot, this is for readibility
  # don't do if sparse array
  if isinstance(self.r, u.quantity.Quantity):
@@ -365,14 +367,15 @@ def _get_source_mask(
  # The average flux, which we assume is a good estimate of the whole stack of images
  max_flux = np.nanmax(self.flux[self.time_mask], axis=0)
 
- # Mask out sources that are above the flux limit, and pixels above the radius limit
+ # Mask out sources that are above the flux limit, and pixels above the
+ # radius limit. This is a good estimation only when using phot_g_mean_flux
  source_rad = 0.5 * np.log10(self.source_flux_estimates) ** 1.5 + 3
  # temp_mask for the sparse array case should also be a sparse matrix. Then it is
  # applied to r, max_flux, and, source_flux_estimates to be used later.
  # Numpy array case:
  if not isinstance(r, sparse.csr_matrix):
  # First we make a guess that each source has exactly the gaia flux
- source_flux_estimates = np.asarray(self.sources.phot_g_mean_flux)[
+ source_flux_estimates = np.asarray(self.source_flux_estimates)[
  :, None
  ] * np.ones((self.nsources, self.npixels))
  temp_mask = (r < source_rad[:, None]) & (
@@ -528,7 +531,7 @@ def _get_source_mask(
  if plot:
  k = np.isfinite(f_temp_mask)
  # Make a nice diagnostic plot
- fig, ax = plt.subplots(1, 2, figsize=(8, 3), facecolor="white")
+ fig, ax = plt.subplots(1, 2, figsize=(11, 3), facecolor="white")
 
  ax[0].scatter(r_temp_mask[k], f_temp_mask[k], s=0.4, c="k", label="Data")
  ax[0].scatter(r_temp_mask[k], A[k].dot(w), c="r", s=0.4, label="Model")
@@ -573,8 +576,7 @@ def _get_uncontaminated_pixel_mask(self):
  # we flag sources fainter than mag_limit as non-contaminant
  if isinstance(self.contaminant_mag_limit, (float, int)):
  aux = self.source_mask.multiply(
- self.sources.phot_g_mean_mag.values[:, None]
- < self.contaminant_mag_limit
+ self.source_flux_estimates[:, None] < self.contaminant_mag_limit
  )
  aux.eliminate_zeros()
  self.uncontaminated_source_mask = aux.multiply(
@@ -957,16 +959,10 @@ def build_shape_model(
  Keyword arguments to be passed to `_get_source_mask()`
  """
 
- # gaia estimate flux values per pixel to be used as flux priors
- self.source_flux_estimates = np.copy(np.asarray(self.sources.phot_g_mean_flux))
-
  # Mask of shape nsources x number of pixels, one where flux from a
  # source exists
- if not hasattr(self, "source_mask"):
- self._get_source_mask(**kwargs)
- # Mask of shape npixels (maybe by nt) where not saturated, not faint,
- # not contaminated etc
- self._get_uncontaminated_pixel_mask()
+ # if not hasattr(self, "source_mask"):
+ self._get_source_mask(**kwargs)
 
  # for iter in range(niters):
  flux_estimates = self.source_flux_estimates[:, None]
@@ -1346,24 +1342,23 @@ def get_psf_metrics(self, npoints_per_pixel=10):
  )
  self.perturbed_std = np.nanstd(perturbed_lcs, axis=0)
 
- def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
+ def plot_shape_model(self, frame_index="mean", bin_data=False):
  """
  Diagnostic plot of shape model.
 
  Parameters
  ----------
- radius : float
- Radius (in arcseconds) limit to be shown in the figure.
  frame_index : string or int
  The frame index used to plot the shape model, if "mean" then use the
  mean value across time
+ bin_data : bool
+ Bin or not the pixel data in a 2D historgram, default is False.
 
  Returns
  -------
  fig : matplotlib.Figure
  Figure.
  """
-
  if frame_index == "mean":
  mean_f = np.log10(
  self.uncontaminated_source_mask.astype(float)
@@ -1386,6 +1381,9 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
  dx = dx.data * u.deg.to(u.arcsecond)
  dy = dy.data * u.deg.to(u.arcsecond)
 
+ radius = np.maximum(np.abs(dx).max(), np.abs(dy).max()) * 1.1
+ vmin, vmax = np.nanpercentile(mean_f, [5, 90])
+
  if bin_data:
  nbins = 30 if mean_f.shape[0] <= 5e3 else 90
  _, dx, dy, mean_f, _ = threshold_bin(
@@ -1394,7 +1392,14 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
 
  fig, ax = plt.subplots(3, 2, figsize=(9, 10.5), constrained_layout=True)
  im = ax[0, 0].scatter(
- dx, dy, c=mean_f, cmap="viridis", vmin=-3, vmax=-1, s=3, rasterized=True
+ dx,
+ dy,
+ c=mean_f,
+ cmap="viridis",
+ vmin=vmin,
+ vmax=vmax,
+ s=3,
+ rasterized=True,
  )
  ax[0, 0].set(
  ylabel=r'$\delta y$ ["]',
@@ -1418,7 +1423,7 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
 
  phi, r = np.arctan2(dy, dx), np.hypot(dx, dy)
  im = ax[0, 1].scatter(
- phi, r, c=mean_f, cmap="viridis", vmin=-3, vmax=-1, s=3, rasterized=True
+ phi, r, c=mean_f, cmap="viridis", vmin=vmin, vmax=vmax, s=3, rasterized=True
  )
  ax[0, 1].set(
  ylabel='$r$ ["]',
@@ -1448,8 +1453,8 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
  r,
  c=model,
  cmap="viridis",
- vmin=-3,
- vmax=-1,
+ vmin=vmin,
+ vmax=vmax,
  s=3,
  rasterized=True,
  )
@@ -1464,8 +1469,8 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
  dy,
  c=model,
  cmap="viridis",
- vmin=-3,
- vmax=-1,
+ vmin=vmin,
+ vmax=vmax,
  s=3,
  rasterized=True,
  )

src/psfmachine/superstamp.py

@@ -15,7 +15,8 @@
 from astropy.wcs import WCS
 import astropy.units as u
 
-from .ffi import FFIMachine, _get_sources, _do_image_cutout
+from .ffi import FFIMachine, _get_sources
+from .utils import _do_image_cutout
 
 __all__ = ["SSMachine"]
 

src/psfmachine/tpf.py

@@ -52,9 +52,10 @@ def __init__(
  pos_corr2=None,
  focus_mask=None,
  tpf_meta=None,
- time_corrector="pos_corr",
  cartesian_knot_spacing="sqrt",
  bkg_subtracted=True,
+ sparse_dist_lim=40,
+ sources_flux_column="phot_g_mean_flux",
  ):
  super().__init__(
  time=time,
@@ -73,6 +74,8 @@ def __init__(
  time_radius=time_radius,
  rmin=rmin,
  rmax=rmax,
+ sparse_dist_lim=sparse_dist_lim,
+ sources_flux_column=sources_flux_column,
  )
  self.tpfs = tpfs
  # match cadences
@@ -93,7 +96,6 @@ def __init__(
  self.pos_corr1 = pos_corr1
  self.pos_corr2 = pos_corr2
  self.tpf_meta = tpf_meta
- self.time_corrector = time_corrector
  self.cartesian_knot_spacing = cartesian_knot_spacing
  self.bkg_subtracted = bkg_subtracted
  # limiting mag values to consider contamination for kepler & tess
@@ -153,7 +155,8 @@ def remove_background_model(
  )
  return
  # invert maks to get bkg pixels from TPFs
- self._get_source_mask()
+ if not hasattr(self, "source_mask"):
+ self._get_source_mask(correct_centroid_offset=True)
  bkg_pixel_mask = ~np.asarray(
  (self.source_mask.todense()).sum(axis=0).astype(bool)
  ).ravel()
@@ -854,7 +857,8 @@ def get_source_centroids(self, method="poscor"):
  def from_TPFs(
  tpfs,
  magnitude_limit=18,
- dr=2,
+ dr=3,
+ sources=None,
  time_mask=None,
  apply_focus_mask=True,
  renormalize_tpf_bkg=False,
@@ -882,7 +886,9 @@ def from_TPFs(
  magnitude_limit : float
  Limiting magnitude to query Gaia catalog.
  dr : int
- Gaia data release to be use, default is 2, options are DR2 and EDR3
+ Gaia data release to be use, default is 3, options are DR2 and EDR3
+ sources : pandas.DataFrame
+ Catalog with sources to be extracted by PSFMachine
  time_mask : boolean array
  Mask to be applied to discard cadences if needed.
  apply_focus_mask : boolean
@@ -1008,9 +1014,10 @@ def from_TPFs(
  saturated_mask,
  )
 
- sources = _get_coord_and_query_gaia(
- tpfs, magnitude_limit, dr=dr, ra=query_ra, dec=query_dec, rad=query_rad
- )
+ if sources is None:
+ sources = _get_coord_and_query_gaia(
+ tpfs, magnitude_limit, dr=dr, ra=query_ra, dec=query_dec, rad=query_rad
+ )
 
  def get_tpf2source():
  tpf2source = []
@@ -1052,6 +1059,20 @@ def get_tpf2source():
  sources["tpf_id"] = None
  sources.loc[idx[match], "tpf_id"] = np.asarray(tpf_meta["targetid"])[match]
 
+ # select flux column o be used as estimate.
+ if tpf_meta["mission"][0] == "TESS":
+ sources_flux_column = "phot_rp_mean_flux"
+ if np.isnan(sources[sources_flux_column].values).any():
+ log.warning(
+ f"Some sources do not have values in 'gaia.{sources_flux_column}''."
+ f" Will default to 'gaia.phot_g_mean_flux' for missing values."
+ )
+ sources[sources_flux_column].fillna(
+ sources["phot_g_mean_flux"], inplace=True
+ )
+ else:
+ sources_flux_column = "phot_g_mean_flux"
+
  # return a Machine object
  return TPFMachine(
  tpfs=tpfs,
@@ -1070,6 +1091,7 @@ def get_tpf2source():
  tpf_meta=tpf_meta,
  time_mask=time_mask,
  bkg_subtracted=not renormalize_tpf_bkg,
+ sources_flux_column=sources_flux_column,
  **kwargs,
  )
 
@@ -1121,7 +1143,7 @@ def _parse_TPFs(tpfs, renormalize_tpf_bkg=True, **kwargs):
 
  elif isinstance(tpfs[0], lk.TessTargetPixelFile):
  qual_mask = lk.utils.TessQualityFlags.create_quality_mask(
- tpfs[0].quality, lk.utils.TessQualityFlags.DEFAULT_BITMASK
+ tpfs[0].quality, 1 | 2 | 4 | 8 | 32 | 128
  )
  qual_mask &= (np.abs(tpfs[0].pos_corr1) < 5) & (np.abs(tpfs[0].pos_corr2) < 5)
  focus_mask = np.ones(len(tpfs[0].time), bool)[qual_mask]

src/psfmachine/machine.py

@@ -59,6 +59,7 @@ def __init__(
  rmax=16,
  cut_r=6,
  sparse_dist_lim=40,
+ sources_flux_column="phot_g_mean_flux",
  ):
  """
  Parameters
@@ -104,6 +105,9 @@ def __init__(
  Radial distance used to include pixels around sources when creating delta
  arrays (dra, ddec, r, and phi) as sparse matrices for efficiency.
  Default is 40" (recommended for kepler). (arcseconds)
+ sources_flux_column : str
+ Column name in `sources` table to be used as flux estimate. For Kepler data
+ gaia.phot_g_mean_flux is recommended, for TESS use gaia.phot_rp_mean_flux.
 
  Attributes
  ----------
@@ -172,6 +176,8 @@ def __init__(
  self.quiet = False
  self.contaminant_mag_limit = None
 
+ self.source_flux_estimates = np.copy(self.sources[sources_flux_column].values)
+
  if time_mask is None:
  self.time_mask = np.ones(len(time), bool)
  else:
@@ -231,7 +237,7 @@ def _create_delta_arrays(self, centroid_offset=[0, 0]):
  self.r = np.hypot(self.dra, self.ddec).to("arcsec")
  self.phi = np.arctan2(self.ddec, self.dra)
 
- def _create_delta_sparse_arrays(self, dist_lim=40, centroid_offset=[0, 0]):
+ def _create_delta_sparse_arrays(self, centroid_offset=[0, 0]):
  """
  Creates dra, ddec, r and phi arrays as sparse arrays to be used for dense data,
  e.g. Kepler FFIs or cluster fields. Assuming that there is no flux information
@@ -243,8 +249,6 @@ def _create_delta_sparse_arrays(self, dist_lim=40, centroid_offset=[0, 0]):
 
  Parameters
  ----------
- dist_lim : float
- Distance limit (in arcsecds) at which pixels are keep.
  centroid_offset : list
  Centroid offset for [ra, dec] to be included in dra and ddec computation.
  Default is [0, 0].
@@ -254,7 +258,7 @@ def _create_delta_sparse_arrays(self, dist_lim=40, centroid_offset=[0, 0]):
  if (centroid_offset[0] == centroid_offset[1] == 0) or (
  np.maximum(*np.abs(centroid_offset)) > 4 / 3600
  ):
- # iterate over sources to only keep pixels within dist_lim
+ # iterate over sources to only keep pixels within self.sparse_dist_lim
  # this is inefficient, could be done in a tiled manner? only for squared data
  dra, ddec, sparse_mask = [], [], []
  for i in tqdm(
@@ -328,17 +332,21 @@ def _get_source_mask(
  correct_centroid_offset=True,
  plot=False,
  ):
- """Find the pixel mask that identifies pixels with contributions from ANY NUMBER of Sources
+ """Find the pixel mask that identifies pixels with contributions from ANY
+ NUMBER of Sources
 
- Fits a simple polynomial model to the log of the pixel flux values, in radial dimension and source flux,
- to find the optimum circular apertures for every source.
+ Fits a simple polynomial model to the log of the pixel flux values, in radial
+ dimension and source flux, to find the optimum circular apertures for every
+ source.
 
  Parameters
  ----------
  upper_radius_limit: float
- The radius limit at which we assume there is no flux from a source of any brightness (arcsec)
+ The radius limit at which we assume there is no flux from a source of
+ any brightness (arcsec)
  lower_radius_limit: float
- The radius limit at which we assume there is flux from a source of any brightness (arcsec)
+ The radius limit at which we assume there is flux from a source of any
+ brightness (arcsec)
  upper_flux_limit: float
  The flux at which we assume as source is saturated
  lower_flux_limit: float
@@ -349,12 +357,6 @@ def _get_source_mask(
  plot: bool
  Whether to show diagnostic plot. Default is False
  """
-
- if not hasattr(self, "source_flux_estimates"):
- # gaia estimate flux values per pixel to be used as flux priors
- self.source_flux_estimates = np.copy(
- np.asarray(self.sources.phot_g_mean_flux)
- )
  # We will use the radius a lot, this is for readibility
  # don't do if sparse array
  if isinstance(self.r, u.quantity.Quantity):
@@ -365,14 +367,15 @@ def _get_source_mask(
  # The average flux, which we assume is a good estimate of the whole stack of images
  max_flux = np.nanmax(self.flux[self.time_mask], axis=0)
 
- # Mask out sources that are above the flux limit, and pixels above the radius limit
+ # Mask out sources that are above the flux limit, and pixels above the
+ # radius limit. This is a good estimation only when using phot_g_mean_flux
  source_rad = 0.5 * np.log10(self.source_flux_estimates) ** 1.5 + 3
  # temp_mask for the sparse array case should also be a sparse matrix. Then it is
  # applied to r, max_flux, and, source_flux_estimates to be used later.
  # Numpy array case:
  if not isinstance(r, sparse.csr_matrix):
  # First we make a guess that each source has exactly the gaia flux
- source_flux_estimates = np.asarray(self.sources.phot_g_mean_flux)[
+ source_flux_estimates = np.asarray(self.source_flux_estimates)[
  :, None
  ] * np.ones((self.nsources, self.npixels))
  temp_mask = (r < source_rad[:, None]) & (
@@ -528,7 +531,7 @@ def _get_source_mask(
  if plot:
  k = np.isfinite(f_temp_mask)
  # Make a nice diagnostic plot
- fig, ax = plt.subplots(1, 2, figsize=(8, 3), facecolor="white")
+ fig, ax = plt.subplots(1, 2, figsize=(11, 3), facecolor="white")
 
  ax[0].scatter(r_temp_mask[k], f_temp_mask[k], s=0.4, c="k", label="Data")
  ax[0].scatter(r_temp_mask[k], A[k].dot(w), c="r", s=0.4, label="Model")
@@ -573,8 +576,7 @@ def _get_uncontaminated_pixel_mask(self):
  # we flag sources fainter than mag_limit as non-contaminant
  if isinstance(self.contaminant_mag_limit, (float, int)):
  aux = self.source_mask.multiply(
- self.sources.phot_g_mean_mag.values[:, None]
- < self.contaminant_mag_limit
+ self.source_flux_estimates[:, None] < self.contaminant_mag_limit
  )
  aux.eliminate_zeros()
  self.uncontaminated_source_mask = aux.multiply(
@@ -957,16 +959,10 @@ def build_shape_model(
  Keyword arguments to be passed to `_get_source_mask()`
  """
 
- # gaia estimate flux values per pixel to be used as flux priors
- self.source_flux_estimates = np.copy(np.asarray(self.sources.phot_g_mean_flux))
-
  # Mask of shape nsources x number of pixels, one where flux from a
  # source exists
- if not hasattr(self, "source_mask"):
- self._get_source_mask(**kwargs)
- # Mask of shape npixels (maybe by nt) where not saturated, not faint,
- # not contaminated etc
- self._get_uncontaminated_pixel_mask()
+ # if not hasattr(self, "source_mask"):
+ self._get_source_mask(**kwargs)
 
  # for iter in range(niters):
  flux_estimates = self.source_flux_estimates[:, None]
@@ -1346,24 +1342,23 @@ def get_psf_metrics(self, npoints_per_pixel=10):
  )
  self.perturbed_std = np.nanstd(perturbed_lcs, axis=0)
 
- def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
+ def plot_shape_model(self, frame_index="mean", bin_data=False):
  """
  Diagnostic plot of shape model.
 
  Parameters
  ----------
- radius : float
- Radius (in arcseconds) limit to be shown in the figure.
  frame_index : string or int
  The frame index used to plot the shape model, if "mean" then use the
  mean value across time
+ bin_data : bool
+ Bin or not the pixel data in a 2D historgram, default is False.
 
  Returns
  -------
  fig : matplotlib.Figure
  Figure.
  """
-
  if frame_index == "mean":
  mean_f = np.log10(
  self.uncontaminated_source_mask.astype(float)
@@ -1386,6 +1381,9 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
  dx = dx.data * u.deg.to(u.arcsecond)
  dy = dy.data * u.deg.to(u.arcsecond)
 
+ radius = np.maximum(np.abs(dx).max(), np.abs(dy).max()) * 1.1
+ vmin, vmax = np.nanpercentile(mean_f, [5, 90])
+
  if bin_data:
  nbins = 30 if mean_f.shape[0] <= 5e3 else 90
  _, dx, dy, mean_f, _ = threshold_bin(
@@ -1394,7 +1392,14 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
 
  fig, ax = plt.subplots(3, 2, figsize=(9, 10.5), constrained_layout=True)
  im = ax[0, 0].scatter(
- dx, dy, c=mean_f, cmap="viridis", vmin=-3, vmax=-1, s=3, rasterized=True
+ dx,
+ dy,
+ c=mean_f,
+ cmap="viridis",
+ vmin=vmin,
+ vmax=vmax,
+ s=3,
+ rasterized=True,
  )
  ax[0, 0].set(
  ylabel=r'$\delta y$ ["]',
@@ -1418,7 +1423,7 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
 
  phi, r = np.arctan2(dy, dx), np.hypot(dx, dy)
  im = ax[0, 1].scatter(
- phi, r, c=mean_f, cmap="viridis", vmin=-3, vmax=-1, s=3, rasterized=True
+ phi, r, c=mean_f, cmap="viridis", vmin=vmin, vmax=vmax, s=3, rasterized=True
  )
  ax[0, 1].set(
  ylabel='$r$ ["]',
@@ -1448,8 +1453,8 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
  r,
  c=model,
  cmap="viridis",
- vmin=-3,
- vmax=-1,
+ vmin=vmin,
+ vmax=vmax,
  s=3,
  rasterized=True,
  )
@@ -1464,8 +1469,8 @@ def plot_shape_model(self, radius=20, frame_index="mean", bin_data=False):
  dy,
  c=model,
  cmap="viridis",
- vmin=-3,
- vmax=-1,
+ vmin=vmin,
+ vmax=vmax,
  s=3,
  rasterized=True,
  )

src/psfmachine/superstamp.py

@@ -15,7 +15,8 @@
 from astropy.wcs import WCS
 import astropy.units as u
 
-from .ffi import FFIMachine, _get_sources, _do_image_cutout
+from .ffi import FFIMachine, _get_sources
+from .utils import _do_image_cutout
 
 __all__ = ["SSMachine"]
 

src/psfmachine/tpf.py

@@ -52,9 +52,10 @@ def __init__(
  pos_corr2=None,
  focus_mask=None,
  tpf_meta=None,
- time_corrector="pos_corr",
  cartesian_knot_spacing="sqrt",
  bkg_subtracted=True,
+ sparse_dist_lim=40,
+ sources_flux_column="phot_g_mean_flux",
  ):
  super().__init__(
  time=time,
@@ -73,6 +74,8 @@ def __init__(
  time_radius=time_radius,
  rmin=rmin,
  rmax=rmax,
+ sparse_dist_lim=sparse_dist_lim,
+ sources_flux_column=sources_flux_column,
  )
  self.tpfs = tpfs
  # match cadences
@@ -93,7 +96,6 @@ def __init__(
  self.pos_corr1 = pos_corr1
  self.pos_corr2 = pos_corr2
  self.tpf_meta = tpf_meta
- self.time_corrector = time_corrector
  self.cartesian_knot_spacing = cartesian_knot_spacing
  self.bkg_subtracted = bkg_subtracted
  # limiting mag values to consider contamination for kepler & tess
@@ -153,7 +155,8 @@ def remove_background_model(
  )
  return
  # invert maks to get bkg pixels from TPFs
- self._get_source_mask()
+ if not hasattr(self, "source_mask"):
+ self._get_source_mask(correct_centroid_offset=True)
  bkg_pixel_mask = ~np.asarray(
  (self.source_mask.todense()).sum(axis=0).astype(bool)
  ).ravel()
@@ -854,7 +857,8 @@ def get_source_centroids(self, method="poscor"):
  def from_TPFs(
  tpfs,
  magnitude_limit=18,
- dr=2,
+ dr=3,
+ sources=None,
  time_mask=None,
  apply_focus_mask=True,
  renormalize_tpf_bkg=False,
@@ -882,7 +886,9 @@ def from_TPFs(
  magnitude_limit : float
  Limiting magnitude to query Gaia catalog.
  dr : int
- Gaia data release to be use, default is 2, options are DR2 and EDR3
+ Gaia data release to be use, default is 3, options are DR2 and EDR3
+ sources : pandas.DataFrame
+ Catalog with sources to be extracted by PSFMachine
  time_mask : boolean array
  Mask to be applied to discard cadences if needed.
  apply_focus_mask : boolean
@@ -1008,9 +1014,10 @@ def from_TPFs(
  saturated_mask,
  )
 
- sources = _get_coord_and_query_gaia(
- tpfs, magnitude_limit, dr=dr, ra=query_ra, dec=query_dec, rad=query_rad
- )
+ if sources is None:
+ sources = _get_coord_and_query_gaia(
+ tpfs, magnitude_limit, dr=dr, ra=query_ra, dec=query_dec, rad=query_rad
+ )
 
  def get_tpf2source():
  tpf2source = []
@@ -1052,6 +1059,20 @@ def get_tpf2source():
  sources["tpf_id"] = None
  sources.loc[idx[match], "tpf_id"] = np.asarray(tpf_meta["targetid"])[match]
 
+ # select flux column o be used as estimate.
+ if tpf_meta["mission"][0] == "TESS":
+ sources_flux_column = "phot_rp_mean_flux"
+ if np.isnan(sources[sources_flux_column].values).any():
+ log.warning(
+ f"Some sources do not have values in 'gaia.{sources_flux_column}''."
+ f" Will default to 'gaia.phot_g_mean_flux' for missing values."
+ )
+ sources[sources_flux_column].fillna(
+ sources["phot_g_mean_flux"], inplace=True
+ )
+ else:
+ sources_flux_column = "phot_g_mean_flux"
+
  # return a Machine object
  return TPFMachine(
  tpfs=tpfs,
@@ -1070,6 +1091,7 @@ def get_tpf2source():
  tpf_meta=tpf_meta,
  time_mask=time_mask,
  bkg_subtracted=not renormalize_tpf_bkg,
+ sources_flux_column=sources_flux_column,
  **kwargs,
  )
 
@@ -1121,7 +1143,7 @@ def _parse_TPFs(tpfs, renormalize_tpf_bkg=True, **kwargs):
 
  elif isinstance(tpfs[0], lk.TessTargetPixelFile):
  qual_mask = lk.utils.TessQualityFlags.create_quality_mask(
- tpfs[0].quality, lk.utils.TessQualityFlags.DEFAULT_BITMASK
+ tpfs[0].quality, 1 | 2 | 4 | 8 | 32 | 128
  )
  qual_mask &= (np.abs(tpfs[0].pos_corr1) < 5) & (np.abs(tpfs[0].pos_corr2) < 5)
  focus_mask = np.ones(len(tpfs[0].time), bool)[qual_mask]