debug & optimize disperse

ObservationSim/MockObject/SkybackgroundMap.py

@@ -10,6 +10,8 @@ import galsim
 
 import os
 
+import time
+
 try:
     import importlib.resources as pkg_resources
 except ImportError:
@@ -95,7 +97,7 @@ def calculateSkyMap_split_g(skyMap=None, blueLimit=4200, redLimit=6500, skyfn='s
                                 tar_spec=spec,
                                 band_start=tbstart, band_end=tbend,
                                 conf=conf2,
-                                flat_cube=flat_cube, ignoreBeam=['D','E'])
+                                flat_cube=None, ignoreBeam=['D','E'])
 
                 spec_orders = sdp.compute_spec_orders()
 
@@ -147,12 +149,13 @@ def calculateSkyMap_split_g(skyMap=None, blueLimit=4200, redLimit=6500, skyfn='s
         y_len = skyMap.shape[0]
         x_len = skyMap.shape[1]
         delt_x = 500
-        delt_y = 500
+        delt_y = y_len
 
         sub_y_start_arr = np.arange(0, y_len, delt_y)
         sub_y_end_arr = sub_y_start_arr + delt_y
         sub_y_end_arr[-1] = min(sub_y_end_arr[-1], y_len)
-
+        
+        delt_x = split_pos-0
         sub_x_start_arr = np.arange(0, split_pos, delt_x)
         sub_x_end_arr = sub_x_start_arr + delt_x
         sub_x_end_arr[-1] = min(sub_x_end_arr[-1], split_pos)
@@ -166,7 +169,8 @@ def calculateSkyMap_split_g(skyMap=None, blueLimit=4200, redLimit=6500, skyfn='s
             for j,k2 in enumerate(sub_x_start_arr):
                 sub_x_s = k2
                 sub_x_e = sub_x_end_arr[j]
-
+                print(i,j,sub_y_s, sub_y_e,sub_x_s,sub_x_e)
+                T1 = time.time()
                 skyImg_sub = galsim.Image(skyImg.array[sub_y_s:sub_y_e, sub_x_s:sub_x_e])
                 origin_sub = [sub_y_s, sub_x_s]
                 sub_x_center = (sub_x_s + sub_x_e) / 2.
@@ -175,7 +179,7 @@ def calculateSkyMap_split_g(skyMap=None, blueLimit=4200, redLimit=6500, skyfn='s
                                 tar_spec=spec,
                                 band_start=tbstart, band_end=tbend,
                                 conf=conf1,
-                                flat_cube=flat_cube, ignoreBeam=['D','E'])
+                                flat_cube=None, ignoreBeam=['B','C','D','E'])
 
                 spec_orders = sdp.compute_spec_orders()
 
@@ -189,7 +193,12 @@ def calculateSkyMap_split_g(skyMap=None, blueLimit=4200, redLimit=6500, skyfn='s
                     if bounds.area() == 0:
                         continue
                     fImg[bounds] = fImg[bounds] + ssImg[bounds]
+        
+                T2 = time.time()
+
+                print('time: %s ms'% ((T2 - T1)*1000))
 
+        delt_x = x_len-split_pos
         sub_x_start_arr = np.arange(split_pos, x_len, delt_x)
         sub_x_end_arr = sub_x_start_arr + delt_x
         sub_x_end_arr[-1] = min(sub_x_end_arr[-1], x_len)
@@ -203,6 +212,9 @@ def calculateSkyMap_split_g(skyMap=None, blueLimit=4200, redLimit=6500, skyfn='s
             for j, k2 in enumerate(sub_x_start_arr):
                 sub_x_s = k2
                 sub_x_e = sub_x_end_arr[j]
+                print(i,j,sub_y_s, sub_y_e,sub_x_s,sub_x_e)
+                
+                T1 = time.time()
 
                 skyImg_sub = galsim.Image(skyImg.array[sub_y_s:sub_y_e, sub_x_s:sub_x_e])
                 origin_sub = [sub_y_s, sub_x_s]
@@ -212,7 +224,7 @@ def calculateSkyMap_split_g(skyMap=None, blueLimit=4200, redLimit=6500, skyfn='s
                                 tar_spec=spec,
                                 band_start=tbstart, band_end=tbend,
                                 conf=conf2,
-                                flat_cube=flat_cube, ignoreBeam=['D','E'])
+                                flat_cube=None, ignoreBeam=['B','C','D','E'])
 
                 spec_orders = sdp.compute_spec_orders()
 
@@ -226,6 +238,9 @@ def calculateSkyMap_split_g(skyMap=None, blueLimit=4200, redLimit=6500, skyfn='s
                     if bounds.area() == 0:
                         continue
                     fImg[bounds] = fImg[bounds] + ssImg[bounds]
+                T2 = time.time()
+
+                print('time: %s ms'% ((T2 - T1)*1000))
 
     if isAlongY == 1:
         fimg, tmx, tmy = rotate90(array_orig=fImg.array, xc=0, yc=0, isClockwise=0)

ObservationSim/MockObject/SpecDisperser/disperse_c/disperse.pyx

@@ -54,43 +54,61 @@ def disperse_grism_object(np.ndarray[FTYPE_t, ndim=2] flam,
     nk = len(idxl)
     nl = len(full)
 
-    nlamb = len(wlambda)
-    flat_eff_all = np.ones([nl, nlamb])*ysens
-
     if (flat is not None):
-        nflat = len(flat)
-        if(nflat == nl):
-            if(len(flat[0])==4):
-                lambda_co = np.zeros([4,nlamb])
-                lambda_co[0] = lambda_co[0] + 1
-                lambda_co[1] = wlambda*1
-                lambda_co[2] = wlambda*lambda_co[1]
-                lambda_co[3] = wlambda*lambda_co[2]
-                flat_eff_all = np.dot(flat, lambda_co)*ysens
-
-    for i in range(0-x0[1], x0[1]):
-        if (x0[1]+i < 0) | (x0[1]+i >= shd[1]):
-            continue
-
-        for j in range(0-x0[0], x0[0]):
-            if (x0[0]+j < 0) | (x0[0]+j >= shd[0]):
+        nlamb = len(wlambda)
+        flat_lamb_min = 2500
+        flat_lamb_max = 10000
+        flat_eff_all = np.ones([nl, nlamb])*ysens
+        lambda_co = np.zeros([len(flat[0]),nlamb])
+        lambda_co[0] = lambda_co[0] + 1
+        lambda_co[1] = (wlambda-flat_lamb_min)/(flat_lamb_max-flat_lamb_min)
+        lambda_co[2] = lambda_co[1]*lambda_co[1]
+        lambda_co[3] = lambda_co[1]*lambda_co[2]
+        flat_eff_all = np.dot(flat, lambda_co)*ysens
+
+        for i in range(0-x0[1], x0[1]):
+            if (x0[1]+i < 0) | (x0[1]+i >= shd[1]):
                 continue
 
-            fl_ij = flam[x0[0]+j, x0[1]+i] #/1.e-17
-            if (fl_ij == 0):
-                continue
+            for j in range(0-x0[0], x0[0]):
+                if (x0[0]+j < 0) | (x0[0]+j >= shd[0]):
+                    continue
+
+                fl_ij = flam[x0[0]+j, x0[1]+i] #/1.e-17
+                if (fl_ij == 0):
+                    continue
 
-            for k in range(nk):
-                k1 = idxl[k]+j*shg[1]+i
-                if (k1 >= 0) & (k1 < nl):
-                    flux_factor = flat_eff_all[k1,k]
-                    full[k1] += fl_ij*yfrac[k]*flux_factor
+                for k in range(nk):
+                    k1 = idxl[k]+j*shg[1]+i
+                    if (k1 >= 0) & (k1 < nl):
+                        full[k1] += fl_ij*yfrac[k]*flat_eff_all[k1,k]
 
-                k2 = idxl[k]+(j-1)*shg[1]+i
-                if (k2 >= 0) & (k2 < nl):
-                    flux_factor = flat_eff_all[k2,k]
-                    full[k2] += fl_ij*(1-yfrac[k])*flux_factor
+                    k2 = idxl[k]+(j-1)*shg[1]+i
+                    if (k2 >= 0) & (k2 < nl):
+                        full[k2] += fl_ij*(1-yfrac[k])*flat_eff_all[k2,k]
     
+    else:
+        for i in range(0-x0[1], x0[1]):
+            if (x0[1]+i < 0) | (x0[1]+i >= shd[1]):
+                continue
+
+            for j in range(0-x0[0], x0[0]):
+                if (x0[0]+j < 0) | (x0[0]+j >= shd[0]):
+                    continue
+
+                fl_ij = flam[x0[0]+j, x0[1]+i] #/1.e-17
+                if (fl_ij == 0):
+                    continue
+
+                for k in range(nk):
+                    k1 = idxl[k]+j*shg[1]+i
+                    if (k1 >= 0) & (k1 < nl):
+                        full[k1] += fl_ij*yfrac[k]*ysens[k]
+
+                    k2 = idxl[k]+(j-1)*shg[1]+i
+                    if (k2 >= 0) & (k2 < nl):
+                        full[k2] += fl_ij*(1-yfrac[k])*ysens[k]
+
     return True
 
 @cython.boundscheck(False)