more intel vs generic mp_fourier work; factor out lanczos_shift_image…

…() function
dmuniesa · Nov 7, 2017 · 75ac41f · 75ac41f
1 parent 0ae6fe1
commit 75ac41f
Show file tree

Hide file tree

Showing 4 changed files with 55 additions and 67 deletions.
diff --git a/tractor/Makefile b/tractor/Makefile
@@ -26,7 +26,7 @@ _mp_fourier$(PYTHON_SO_EXT): mp_fourier.i setup-mpf.py
 	$(PYTHON) setup-mpf.py build_ext --inplace
 
 _intel_mp_fourier$(PYTHON_SO_EXT): mp_fourier.i setup-mpf.py
-	cp mp_fourier.i intel_mp_fourier.i
+	cat mp_fourier.i | sed s/mp_fourier/intel_mp_fourier/g > intel_mp_fourier.i
 	CC=icc $(PYTHON) setup-mpf.py build_ext --inplace
 
 mix.py _mix$(PYTHON_SO_EXT): mix.i approx3.c gauss_masked.c setup-mix.py

diff --git a/tractor/galaxy.py b/tractor/galaxy.py
@@ -16,22 +16,6 @@
 from __future__ import division
 
 import numpy as np
-# import numpy.fft_intel.libifft as m
-# def irfftn_numpy(x, s=None, axes=None):
-#     a = np.asarray(x)
-#     no_trim = (s is None) and (axes is None)
-#     s, axes = m._cook_nd_args(a, s, axes, invreal=True)
-#     la = axes[-1]
-#     ovr_x = False
-#     if len(s) > 1:
-#         if not no_trim:
-#             a = m._fix_dimensions(a, s, axes)
-#         for ii in range(len(axes)-1):
-#             a = m.ifft(a, s[ii], axes[ii], overwrite_x=ovr_x)
-#             ovr_x = True
-#     a = m.irfft_numpy(a, n = s[-1], axis=la)
-#     return a
-# m.irfftn_numpy  = irfftn_numpy
 
 from astrometry.util.miscutils import get_overlapping_region
 
@@ -40,8 +24,6 @@
 from tractor.patch import Patch, add_patches, ModelMask
 from tractor.basics import SingleProfileSource, BasicSource
 
-#from .cache import Cache
-
 debug_ps = None
 
 def get_galaxy_cache():
@@ -478,29 +460,32 @@ def run_mog(amix=None, mm=None):
             # after cutting G down to nearly its final size... tricky
             # tho
 
-            # Lanczos-3 interpolation in ~ the same way we do for
-            # pixelized PSFs.
-            from astrometry.util.miscutils import lanczos_filter
-            #from scipy.ndimage.filters import correlate1d
-            #L = 3
-            L = fft_lanczos_order
-            Lx = lanczos_filter(L, np.arange(-L, L+1) + mux)
-            Ly = lanczos_filter(L, np.arange(-L, L+1) + muy)
-            # Normalize the Lanczos interpolants (preserve flux)
-            Lx /= Lx.sum()
-            Ly /= Ly.sum()
-            #print('Lx centroid', np.sum(Lx * (np.arange(-L,L+1))))
-            #print('Ly centroid', np.sum(Ly * (np.arange(-L,L+1))))
-
-            assert(len(Lx) == 7)
-            assert(len(Ly) == 7)
-            #cx = correlate1d(G,  Lx, axis=1, mode='constant')
-            #G  = correlate1d(cx, Ly, axis=0, mode='constant')
-            #del cx
-
-            from tractor.c_mp_fourier import correlate7
-            G = np.require(G, requirements=['A'])
-            correlate7(G, Lx, Ly, work_corr7)
+            from tractor.psf import lanczos_shift_image
+            G = lanczos_shift_image(G, mux, muy)
+
+            # # Lanczos-3 interpolation in ~ the same way we do for
+            # # pixelized PSFs.
+            # from astrometry.util.miscutils import lanczos_filter
+            # #from scipy.ndimage.filters import correlate1d
+            # #L = 3
+            # L = fft_lanczos_order
+            # Lx = lanczos_filter(L, np.arange(-L, L+1) + mux)
+            # Ly = lanczos_filter(L, np.arange(-L, L+1) + muy)
+            # # Normalize the Lanczos interpolants (preserve flux)
+            # Lx /= Lx.sum()
+            # Ly /= Ly.sum()
+            # #print('Lx centroid', np.sum(Lx * (np.arange(-L,L+1))))
+            # #print('Ly centroid', np.sum(Ly * (np.arange(-L,L+1))))
+            # 
+            # assert(len(Lx) == 7)
+            # assert(len(Ly) == 7)
+            # #cx = correlate1d(G,  Lx, axis=1, mode='constant')
+            # #G  = correlate1d(cx, Ly, axis=0, mode='constant')
+            # #del cx
+            # 
+            # from tractor.c_mp_fourier import correlate7
+            # G = np.require(G, requirements=['A'])
+            # correlate7(G, Lx, Ly, work_corr7)
 
         else:
             G = np.zeros((pH,pW), np.float32)

diff --git a/tractor/psf.py b/tractor/psf.py
@@ -14,35 +14,39 @@
 from tractor import mixture_profiles as mp
 from tractor import ducks
 
+
+try:
+    from tractor import intel_mp_fourier as mp_fourier
+except:
+    try:
+        from tractor import mp_fourier as mp_fourier
+    except:
+        mp_fourier = None
+# from tractor.c_mp_fourier import correlate7, correlate7f
+
 def lanczos_shift_image(img, dx, dy):
+    from scipy.ndimage import correlate1d
     L = 3
     Lx = lanczos_filter(L, np.arange(-L, L+1) + dx)
     Ly = lanczos_filter(L, np.arange(-L, L+1) + dy)
     # Normalize the Lanczos interpolants (preserve flux)
     Lx /= Lx.sum()
     Ly /= Ly.sum()
 
-    sx      = correlate1d(img, Lx, axis=1, mode='constant')
-    shifted = correlate1d(sx,  Ly, axis=0, mode='constant')
-    assert(np.all(np.isfinite(shifted)))
-    return shifted
-
-#     assert(len(Lx) == 7)
-#     assert(len(Ly) == 7)
-#     from tractor.c_mp_fourier import correlate7, correlate7f
-#     outimg = np.empty(img.shape, np.float32)
-#     outimg[:,:] = img
-#     correlate7f(outimg, Lx, Ly, work_corr7f)
-# 
-# 
-# try:
-#     from tractor import intel_mp_fourier as mp_fourier
-# except:
-#     try:
-#         from tractor import mp_fourier
-#     except:
-#         mp_fourier = None
-# from tractor.c_mp_fourier import correlate7, correlate7f
+    #print('mp_fourier:', mp_fourier)
+    if mp_fourier is None:
+        sx     = correlate1d(img, Lx, axis=1, mode='constant')
+        outimg = correlate1d(sx,  Ly, axis=0, mode='constant')
+    else:
+        assert(len(Lx) == 7)
+        assert(len(Ly) == 7)
+        outimg = np.empty(img.shape, np.float32)
+        outimg[:,:] = img
+        mp_fourier.correlate7f(outimg, Lx, Ly, work_corr7f)
+
+    assert(np.all(np.isfinite(outimg)))
+    return outimg
+
 
 
 class HybridPSF(object):
@@ -145,19 +149,18 @@ def getPointSourcePatch(self, px, py, minval=0., modelMask=None, **kwargs):
             return Patch(x0, y0, outimg)
 
         #
+        L = 3
         padding = L
         # Create a modelMask + padding sized stamp and insert PSF image into it
         #mm = np.zeros((mh+2*padding, mw+2*padding), img.dtype)
         mm = np.zeros((mh+2*padding, mw+2*padding), np.float32)
-
         yi,yo = get_overlapping_region(my0-y0-padding, my0-y0+mh-1+padding, 0, H-1)
         xi,xo = get_overlapping_region(mx0-x0-padding, mx0-x0+mw-1+padding, 0, W-1)
         mm[yo,xo] = img[yi,xi]
-
         mm = lanczos_shift_image(mm, dx, dy)
-
         mm = mm[padding:-padding, padding:-padding]
         assert(np.all(np.isfinite(mm)))
+
         return Patch(mx0, my0, mm)
 
     def getFourierTransformSize(self, radius):

diff --git a/tractor/setup-mpf.py b/tractor/setup-mpf.py
@@ -6,7 +6,7 @@
 
 if os.environ.get('CC') == 'icc':
     mpf_module = Extension('_intel_mp_fourier',
-                           sources=['mp_fourier.i'],
+                           sources=['intel_mp_fourier.i'],
                            include_dirs=numpy_inc,
                            extra_compile_args=['-g', '-xhost', '-qopt-report=5', '-axMIC-AVX512'],
                            extra_link_args=['-g', '-lsvml']