Add gridded psf wrapper [WIP]

crowdsource/crowdsource_base.py

@@ -296,12 +296,13 @@ def fit_once(im, x, y, psfs, weight=None,
     repeat = 1 if not psfderiv else 3
     nskypar = nskyx * nskyy
     npixim = im.shape[0]*im.shape[1]
-    zsz = (repeat*numpy.sum(sz*sz) + nskypar*npixim).astype('i4')
-    if zsz >= 2**32:
+    zsz = (repeat*numpy.sum(sz*sz) + nskypar*npixim).astype('i8')
+    if zsz >= 2**33:
         raise ValueError(
-            'Number of pixels being fit is too large (>2**32); '
+            'Number of pixels being fit is too large (>2**33); '
             'failing early.  This usually indicates a problem with '
             'the choice of PSF size & too many sources.')
+    # DEBUG    print(f"zsz = {zsz}, sz={sz}, nskypar={nskypar}, npixim={npixim}, repeat={repeat}")
     xloc = numpy.zeros(zsz, dtype='i4')
     values = numpy.zeros(len(xloc), dtype='f4')
     colnorm = numpy.zeros(len(x)*repeat+nskypar, dtype='f4')
@@ -746,7 +747,8 @@ def fit_im(im, psf, weight=None, dq=None, psfderiv=True,
            maxstars=40000, derivcentroids=False,
            ntilex=1, ntiley=1, fewstars=100, threshold=5,
            ccd=None, plot=False, titer_thresh=2, blendthreshu=2,
-           psfvalsharpcutfac=0.7, psfsharpsat=0.7):
+           psfvalsharpcutfac=0.7, psfsharpsat=0.7,
+           add_new_peaks_every_other_iteration=False):
     if isinstance(weight, int):
         weight = numpy.ones_like(im)*weight
 
@@ -769,7 +771,7 @@ def fit_im(im, psf, weight=None, dq=None, psfderiv=True,
         titer += 1
         hsky = sky_im(im-model, weight=weight, npix=20)
         lsky = sky_im(im-model, weight=weight, npix=50*roughfwhm)
-        if titer != lastiter:
+        if titer != lastiter and not ((titer % 2 == 0) and add_new_peaks_every_other_iteration):
             # in first passes, do not split sources!
             blendthresh = blendthreshu if titer < titer_thresh else 0.2
             xn, yn = peakfind(im-model-hsky,
@@ -824,10 +826,6 @@ def fit_im(im, psf, weight=None, dq=None, psfderiv=True,
             print("Starting subregion iterations")
         for (bdxf, bdxl, bdxaf, bdxal, bdyf, bdyl, bdyaf, bdyal) in (
                 subregions(im.shape, ntilex, ntiley)):
-            if verbose:
-                print(f"Subregion iteration {subreg_iter} starting; "
-                      f"dt={time.time()-t0}", flush=True)
-                subreg_iter += 1
             mbda = in_bounds(xa, ya, [bdxaf-0.5, bdxal-0.5],
                              [bdyaf-0.5, bdyal-0.5])
             mbd = in_bounds(xa, ya, [bdxf-0.5, bdxl-0.5],
@@ -870,6 +868,11 @@ def fit_im(im, psf, weight=None, dq=None, psfderiv=True,
             import gc
             gc.collect()
 
+            if verbose:
+                print(f"Subregion iteration {subreg_iter} finished; "
+                      f"dt={time.time()-t0}", flush=True)
+                subreg_iter += 1
+
         centroids = compute_centroids(xa, ya, psfs, flux, im-(sky+msky),
                                       im-model-sky,
                                       weight, derivcentroids=derivcentroids)
@@ -908,6 +911,8 @@ def fit_im(im, psf, weight=None, dq=None, psfderiv=True,
         # i.e., 50k saturates, so we can cut there.
         brightenough = (guessflux/fluxunc > threshold*3/5.) | (guessflux > 1e5)
         isolatedenough = cull_near(xa, ya, guessflux)
+        if verbose:
+            print(f"threshold={threshold}, median fluxunc={numpy.nanmedian(fluxunc)}, median flux={numpy.nanmedian(guessflux)} time={time.time()-t0}")
 
         keep = brightenough & isolatedenough
         xa, ya = (c[keep] for c in (xa, ya))

crowdsource/psf.py

@@ -1222,3 +1222,66 @@ def wise_psf_fit(x, y, xcen, ycen, stamp, imstamp, modstamp,
         return SimplePSF(newstamp)
     else:
         return GridInterpPSF(newstamp, psfstamp[1], psfstamp[2])
+
+
+class WrappedPSFModel(SimplePSF):
+    """
+    wrapper for photutils GriddedPSFModel
+    """
+    def __init__(self, psfgridmodel, stampsz=19):
+        self.psfgridmodel = psfgridmodel
+        self.default_stampsz = stampsz
+
+    def __call__(self, col, row, stampsz=None, deriv=False):
+
+        if stampsz is None:
+            stampsz = self.default_stampsz
+
+        parshape = numpy.broadcast(col, row).shape
+        tparshape = parshape if len(parshape) > 0 else (1,)
+
+        # numpy uses row, column notation
+        rows, cols = np.indices((stampsz, stampsz)) - (np.array([stampsz, stampsz])-1)[:, None, None] / 2.
+
+        # explicitly broadcast
+        col = np.atleast_1d(col)
+        row = np.atleast_1d(row)
+        #rows = rows[:, :, None] + row[None, None, :]
+        #cols = cols[:, :, None] + col[None, None, :]
+
+        # photutils seems to use column, row notation
+        #stamps = self.psfgridmodel.evaluate(cols, rows, 1, col, row)
+        # it returns something in (nstamps, row, col) shape
+        # pretty sure that ought to be (col, row, nstamps) for crowdsource
+        stamps = []
+        for i in range(len(col)):
+            stamps.append(self.psfgridmodel.evaluate(cols+col[i], rows+row[i], 1, col[i], row[i]))
+
+        # swap (z,y,x) -> (z,x,y)
+        stamps = np.transpose(stamps, axes=(0,2,1))
+
+        if deriv:
+            dpsfdrow, dpsfdcol = np.gradient(stamps, axis=(1, 2))
+
+        ret = stamps
+        if parshape != tparshape:
+            ret = ret.reshape(stampsz, stampsz)
+            if deriv:
+                dpsfdrow = dpsfdrow.reshape(stampsz, stampsz)
+                dpsfdcol = dpsfdcol.reshape(stampsz, stampsz)
+        if deriv:
+            ret = (ret, dpsfdcol, dpsfdrow)
+
+        return ret
+
+
+    def render_model(self, col, row, stampsz=None):
+        """
+        this function likely does nothing?
+        """
+        if stampsz is not None:
+            self.stampsz = stampsz
+
+        rows, cols = np.indices(self.stampsz, dtype=float) - (np.array(self.stampsz)-1)[:, None, None] / 2.
+
+        return self.psfgridmodel.evaluate(cols, rows, 1, col, row).T.squeeze()