Factor of 100 speed gain in new nemoModel

simonsobs · Dec 12, 2023 · ba00792 · ba00792
1 parent 788300f
commit ba00792
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 8 deletions.
diff --git a/nemo/maps.py b/nemo/maps.py
@@ -1752,11 +1752,11 @@ def makeModelImage(shape, wcs, catalog, beamFileName, obsFreqGHz = None, GNFWPar
                 theta500Arcmin=signals.calcTheta500Arcmin(z, M500, cosmoModel)
                 maxSizeDeg=5*(theta500Arcmin/60)
                 # Updated in place
-                modelMap=makeClusterSignalMap(z, M500, modelMap.shape, wcs, RADeg = row['RADeg'],
-                                                decDeg = row['decDeg'], beam = beam,
-                                                GNFWParams = GNFWParams, amplitude = y0ToInsert,
-                                                maxSizeDeg = maxSizeDeg, convolveWithBeam = True,
-                                                cosmoModel = cosmoModel, omap = modelMap)
+                makeClusterSignalMap(z, M500, modelMap.shape, wcs, RADeg = row['RADeg'],
+                                     decDeg = row['decDeg'], beam = beam,
+                                     GNFWParams = GNFWParams, amplitude = y0ToInsert,
+                                     maxSizeDeg = maxSizeDeg, convolveWithBeam = True,
+                                     cosmoModel = cosmoModel, omap = modelMap)
                 # modelMap=modelMap+signalMap
             if obsFreqGHz is not None:
                 modelMap=convertToDeltaT(modelMap, obsFrequencyGHz = obsFreqGHz,

diff --git a/nemo/signals.py b/nemo/signals.py
@@ -662,9 +662,22 @@ def _paintSignalMap(shape, wcs, tckP, beam = None, RADeg = None, decDeg = None,
     else:
         amps=np.array([amp], dtype = dtype)
 
-    signalMap=pointsrcs.sim_objects(shape, wcs.AWCS, poss, amps, (r, abs(rprof)), vmin = vmin,
-                                    rmax = np.radians(maxSizeDeg), prof_equi = False,
-                                    omap = omap)
+    if omap is not None:
+        ra1=RADeg-(maxSizeDeg/2)/np.cos(np.radians(decDeg))
+        ra0=RADeg+(maxSizeDeg/2)/np.cos(np.radians(decDeg))
+        dec1=decDeg-maxSizeDeg/2
+        dec0=decDeg+maxSizeDeg/2
+        clip=astImages.clipUsingRADecCoords(omap, wcs, ra1, ra0, dec0, dec1)
+        modelClip=pointsrcs.sim_objects(clip['data'].shape, clip['wcs'].AWCS, poss, amps, (r, abs(rprof)), vmin = vmin,
+                                        rmax = np.radians(maxSizeDeg), #prof_equi = False,
+                                        pixwin = False)
+        xMin, xMax, yMin, yMax=clip['clippedSection']
+        omap[yMin:yMax, xMin:xMax]=omap[yMin:yMax, xMin:xMax]+modelClip
+        signalMap=omap
+    else:
+        signalMap=pointsrcs.sim_objects(shape, wcs.AWCS, poss, amps, (r, abs(rprof)), vmin = vmin,
+                                        rmax = np.radians(maxSizeDeg), #prof_equi = False,
+                                        pixwin = False)
 
     if rprof[0] < 0:
         signalMap=signalMap*-1