Use 4 corners for supersampled_prf (#8)

* Moved _update_coordinates

* Updated Kepler _update_coordinates

* fixed imports

* Fixed pixel offset

* modified test

* Modify tests

* Revert pixel offset

* Updated changelog

CHANGES.rst

@@ -2,6 +2,7 @@
 =====================
 
 - Added the option to use the initial TESS commissioning PRF files [#9]
+- Modified tessprf.py and keplerprf.py to use only the closest PRF measurements to create the supersampled PRF [#8]
 
 1.0.3 (2024-07-30)
 ==================

src/lkprf/keplerprf.py

@@ -7,6 +7,7 @@
 import warnings
 
 from .prfmodel import PRF
+from scipy.interpolate import RectBivariateSpline
 
 
 class KeplerPRF(PRF):
@@ -59,3 +60,38 @@ def update_coordinates(self, targets: List[Tuple], shape: Tuple):
             )
         self._update_coordinates(targets=targets, shape=shape)
         return
+
+    def _update_coordinates(self, targets: List[Tuple], shape: Tuple):
+        row, column = self._unpack_targets(targets)
+        # Set the row and column for the model
+        row, column = np.atleast_1d(row), np.atleast_1d(column)
+        row = np.min([np.max([row, row**0 * 20], axis=0), row**0 * 1044], axis=0).mean()
+        column = np.min(
+            [np.max([column, column**0 * 12], axis=0), column**0 * 1112], axis=0
+        ).mean()
+
+        # interpolate the calibrated PRF shape to the target position
+        min_prf_weight = 1e-6
+        rowdim, coldim = shape[0], shape[1]
+        ref_column = column + 0.5 * coldim
+        ref_row = row + 0.5 * rowdim
+        supersamp_prf = np.zeros(self.PRFdata.shape[1:], dtype="float32")
+
+        # Find the 3 measurements nearest the desired locations
+        # Kepler has 5 measurements, so nearest 3 triangulates the measurement
+        prf_weights = [
+            np.sqrt(
+                (ref_column - self.crval1p[i]) ** 2 + (ref_row - self.crval2p[i]) ** 2) 
+                for i in range(self.PRFdata.shape[0])
+                ]
+        idx = np.argpartition(prf_weights, 3)[:3]
+
+        for i in idx:
+            if prf_weights[i] < min_prf_weight:
+                prf_weights[i] = min_prf_weight
+            supersamp_prf += self.PRFdata[i] / prf_weights[i]
+
+
+        supersamp_prf /= np.nansum(supersamp_prf) * self.cdelt1p[0] * self.cdelt2p[0]
+        self.interpolate = RectBivariateSpline(self.PRFrow, self.PRFcol, supersamp_prf)
+        return

src/lkprf/prfmodel.py

@@ -4,7 +4,6 @@
 from typing import Tuple, List
 import numpy.typing as npt
 import numpy as np
-from scipy.interpolate import RectBivariateSpline
 
 
 class PRF(ABC):
@@ -83,8 +82,8 @@ def _evaluate(
         c1, c2 = int(np.floor(self.PRFcol[0])), int(np.ceil(self.PRFcol[-1]))
         # Position in the PRF model for each source position % 1
         delta_row, delta_col = (
-            np.arange(r1, r2)[:, None] + 0.5 - np.atleast_1d(target_row) % 1,
-            np.arange(c1, c2)[:, None] + 0.5 - np.atleast_1d(target_column) % 1,
+            np.arange(r1, r2)[:, None]  - np.atleast_1d(target_row) % 1,
+            np.arange(c1, c2)[:, None] - np.atleast_1d(target_column) % 1,
         )
 
         # prf model for each source, downsampled to pixel grid
@@ -212,15 +211,13 @@ def _prepare_prf(self):
         )
         PRFdata /= PRFdata.sum(axis=(1, 2))[:, None, None]
 
-        PRFcol = np.arange(0, np.shape(PRFdata[0])[1] + 0)
-        PRFrow = np.arange(0, np.shape(PRFdata[0])[0] + 0)
+        PRFcol = np.arange(0.5, np.shape(PRFdata[0])[1] + 0.5)
+        PRFrow = np.arange(0.5, np.shape(PRFdata[0])[0] + 0.5)
 
         # Shifts pixels so it is in pixel units centered on 0
         PRFcol = (PRFcol - np.size(PRFcol) / 2) * cdelt1p[0]
         PRFrow = (PRFrow - np.size(PRFrow) / 2) * cdelt2p[0]
 
-        PRFcol += 0.5
-        PRFrow += 0.5
 
         (
             self.PRFrow,
@@ -232,39 +229,7 @@ def _prepare_prf(self):
             self.cdelt2p,
         ) = (PRFrow, PRFcol, PRFdata, crval1p, crval2p, cdelt1p, cdelt2p)
 
-    def _update_coordinates(self, targets: List[Tuple], shape: Tuple):
-        row, column = self._unpack_targets(targets)
-        # Set the row and column for the model
-        row, column = np.atleast_1d(row), np.atleast_1d(column)
-        row = np.min([np.max([row, row**0 * 20], axis=0), row**0 * 1044], axis=0).mean()
-        column = np.min(
-            [np.max([column, column**0 * 12], axis=0), column**0 * 1112], axis=0
-        ).mean()
-
-        # interpolate the calibrated PRF shape to the target position
-        min_prf_weight = 1e-6
-        rowdim, coldim = shape[0], shape[1]
-        ref_column = column + 0.5 * coldim
-        ref_row = row + 0.5 * rowdim
-        supersamp_prf = np.zeros(self.PRFdata.shape[1:], dtype="float32")
-
-        # Find the 4 measurements nearest the desired locations
-        prf_weights = [
-            np.sqrt(
-                (ref_column - self.crval1p[i]) ** 2 + (ref_row - self.crval2p[i]) ** 2) 
-                for i in range(self.PRFdata.shape[0])
-                ]
-        idx = np.argpartition(prf_weights, 4)[:4]
-
-        for i in idx:
-            if prf_weights[i] < min_prf_weight:
-                prf_weights[i] = min_prf_weight
-            supersamp_prf += self.PRFdata[i] / prf_weights[i]
-
-
-        supersamp_prf /= np.nansum(supersamp_prf) * self.cdelt1p[0] * self.cdelt2p[0]
-        self.interpolate = RectBivariateSpline(self.PRFrow, self.PRFcol, supersamp_prf)
-        return
+
 
     @abstractmethod
     def update_coordinates(self, targets, shape):

src/lkprf/tessprf.py

@@ -7,6 +7,7 @@
 import warnings
 
 from .prfmodel import PRF
+from scipy.interpolate import RectBivariateSpline
 
 
 class TESSPRF(PRF):
@@ -36,7 +37,7 @@ def update_coordinates(self, targets: List[Tuple], shape: Tuple):
                 LKPRFWarning,
             )
 
-        if ((np.atleast_1d(column) + shape[1]) > 2093).any():
+        if ((np.atleast_1d(column) + shape[1]) > 2092).any():
             warnings.warn(
                 "`targets` contains collateral pixels: Column(s) > 2093 ",
                 LKPRFWarning,
@@ -53,3 +54,39 @@ def update_coordinates(self, targets: List[Tuple], shape: Tuple):
             )
         self._update_coordinates(targets=targets, shape=shape)
         return
+
+
+    def _update_coordinates(self, targets: List[Tuple], shape: Tuple):
+        row, column = self._unpack_targets(targets)
+        # Set the row and column for the model. 
+        # Disallows models in the collateral pixels
+        row, column = np.atleast_1d(row), np.atleast_1d(column)
+        row = np.min([np.max([row, row**0 * 0], axis=0), row**0 * 2048], axis=0).mean()
+        column = np.min(
+            [np.max([column, column**0 * 45], axis=0), column**0 * 2092], axis=0
+        ).mean()
+
+        # interpolate the calibrated PRF shape to the target position
+        min_prf_weight = 1e-6
+        rowdim, coldim = shape[0], shape[1]
+        ref_column = column + 0.5 * coldim
+        ref_row = row + 0.5 * rowdim
+        supersamp_prf = np.zeros(self.PRFdata.shape[1:], dtype="float32")
+
+        # Find the 4 measurements nearest the desired locations
+        prf_weights = [
+            np.sqrt(
+                (ref_column - self.crval1p[i]) ** 2 + (ref_row - self.crval2p[i]) ** 2) 
+                for i in range(self.PRFdata.shape[0])
+                ]
+        idx = np.argpartition(prf_weights, 4)[:4]
+
+        for i in idx:
+            if prf_weights[i] < min_prf_weight:
+                prf_weights[i] = min_prf_weight
+            supersamp_prf += self.PRFdata[i] / prf_weights[i]
+
+
+        supersamp_prf /= np.nansum(supersamp_prf) * self.cdelt1p[0] * self.cdelt2p[0]
+        self.interpolate = RectBivariateSpline(self.PRFrow, self.PRFcol, supersamp_prf)
+        return

tests/test_prf.py

@@ -18,16 +18,16 @@ def test_prfs():
         origin = (0, 0)
         shape = (21, 21)
         ar = prf.evaluate(targets=targets, origin=origin, shape=shape)
-        R, C = np.mgrid[: ar.shape[1], : ar.shape[2]] + 0.5
-        assert np.isclose(np.average(R.ravel(), weights=ar[0].ravel()), 10.5, atol=0.05)
-        assert np.isclose(np.average(C.ravel(), weights=ar[0].ravel()), 10.5, atol=0.05)
+        R, C = np.mgrid[: ar.shape[1], : ar.shape[2]] 
+        assert np.isclose(np.average(R.ravel(), weights=ar[0].ravel()), targets[0][1], atol=0.15)
+        assert np.isclose(np.average(C.ravel(), weights=ar[0].ravel()), targets[0][0], atol=0.15)
         assert np.isclose(ar[0].sum(), 1)
         assert ar.shape == (1, *shape)
         if not is_github_actions():
             fig, ax = plt.subplots(figsize=(6, 5))
             im = ax.pcolormesh(
-                np.arange(origin[1], origin[1] + shape[1]) + 0.5,
-                np.arange(origin[0], origin[0] + shape[0]) + 0.5,
+                np.arange(origin[1], origin[1] + shape[1]) ,
+                np.arange(origin[0], origin[0] + shape[0]) ,
                 ar.sum(axis=0),
                 cmap="Greys_r",
                 vmin=0,
@@ -36,8 +36,8 @@ def test_prfs():
             cbar = plt.colorbar(im, ax=ax)
             # ax.scatter(np.asarray(targets)[:, 1], np.asarray(targets)[:, 0], c="r")
             ax.set(
-                xlim=(origin[1], origin[1] + shape[1] - 1),
-                ylim=(origin[0], origin[0] + shape[0] - 1),
+                xlim=(origin[1] - 0.5, origin[1] + shape[1] - 0.5),
+                ylim=(origin[0] - 0.5, origin[0] + shape[0] - 0.5),
                 title=f"{prf.mission} PRF Example",
                 aspect="equal",
                 xlabel="Column [pixel]",
@@ -51,7 +51,7 @@ def test_prfs():
         assert not np.isclose(ar[0].sum(), 1)
         assert not np.isclose(ar[0].sum(), 0)
 
-        ar = prf.gradient(targets=(10, 10), origin=(0, 0), shape=(21, 21))
+        ar = prf.gradient(targets=targets, origin=origin, shape=shape)
         assert isinstance(ar, tuple)
         assert ar[0].shape == (1, *shape)
         for a in ar: