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]