More speed / efficiency for new nemoModel

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

Hide file tree

Showing 2 changed files with 11 additions and 8 deletions.
diff --git a/nemo/maps.py b/nemo/maps.py
@@ -1651,7 +1651,7 @@ def makeModelImage(shape, wcs, catalog, beamFileName, obsFreqGHz = None, GNFWPar
     
     """
 
-    modelMap=np.zeros(shape, dtype = np.float32)
+    modelMap=enmap.zeros(shape, dtype = np.float32) #np.zeros(shape, dtype = np.float32)
 
     if type(catalog) == str:
         catalog=atpy.Table().read(catalog)
@@ -1751,12 +1751,13 @@ def makeModelImage(shape, wcs, catalog, beamFileName, obsFreqGHz = None, GNFWPar
                     y0ToInsert=row['y_c']*1e-4  # or fixed_y_c...
                 theta500Arcmin=signals.calcTheta500Arcmin(z, M500, cosmoModel)
                 maxSizeDeg=5*(theta500Arcmin/60)
-                signalMap=makeClusterSignalMap(z, M500, modelMap.shape, wcs, RADeg = row['RADeg'],
+                # 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)
-                modelMap=modelMap+signalMap
+                                                cosmoModel = cosmoModel, omap = modelMap)
+                # modelMap=modelMap+signalMap
             if obsFreqGHz is not None:
                 modelMap=convertToDeltaT(modelMap, obsFrequencyGHz = obsFreqGHz,
                                          TCMBAlpha = TCMBAlpha, z = z)

diff --git a/nemo/signals.py b/nemo/signals.py
@@ -619,7 +619,7 @@ def makeBeamModelSignalMap(degreesMap, wcs, beam, amplitude = None):
 
 #------------------------------------------------------------------------------------------------------------
 def _paintSignalMap(shape, wcs, tckP, beam = None, RADeg = None, decDeg = None, amplitude = None,
-                    maxSizeDeg = 10.0, convolveWithBeam = True, vmin = 1e-12):
+                    maxSizeDeg = 10.0, convolveWithBeam = True, vmin = 1e-12, omap = None):
     """Use Sigurd's fast object painter to paint given signal into map.
 
     Notes:
@@ -663,7 +663,8 @@ def _paintSignalMap(shape, wcs, tckP, beam = None, RADeg = None, decDeg = None,
         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)
+                                    rmax = np.radians(maxSizeDeg), prof_equi = False,
+                                    omap = omap)
 
     if rprof[0] < 0:
         signalMap=signalMap*-1
@@ -673,7 +674,8 @@ def _paintSignalMap(shape, wcs, tckP, beam = None, RADeg = None, decDeg = None,
 #------------------------------------------------------------------------------------------------------------
 def makeArnaudModelSignalMap(z, M500, shape, wcs, beam = None, RADeg = None, decDeg = None,\
                              GNFWParams = 'default', amplitude = None, maxSizeDeg = 15.0,\
-                             convolveWithBeam = True, cosmoModel = None, painter = 'pixell'):
+                             convolveWithBeam = True, cosmoModel = None, painter = 'pixell',
+                             omap = None):
     """Makes a 2d signal only map containing an Arnaud model cluster.
     
     Args:
@@ -735,7 +737,7 @@ def makeArnaudModelSignalMap(z, M500, shape, wcs, beam = None, RADeg = None, dec
     elif painter == 'pixell': # New method - using Sigurd's object painter
         signalMap=_paintSignalMap(shape, wcs, tckP, beam = beam, RADeg = RADeg, decDeg = decDeg,
                                   amplitude = amplitude, maxSizeDeg = maxSizeDeg,
-                                  convolveWithBeam = convolveWithBeam)
+                                  convolveWithBeam = convolveWithBeam, omap = omap)
     else:
         raise Exception("'painter' must be 'legacy' or 'pixell' (given '%s')." % (painter))