Merge pull request #28 from simonsobs/v0.3.0-beta-fixes

V0.3.0 beta fixes

bin/nemoCosmo

@@ -6,7 +6,7 @@ Example of using nemo selection function stuff for cosmology.
 
 NOTE: relies on global variables:
 - selFn
-- tab
+- tab (and related stuff on a grid)
 
 Using Cobaya as the sampler.
 
@@ -36,9 +36,36 @@ import getdist.plots as gdplt
 # Extreme debugging (better for multiprocessing stuff to crash where we can see it)
 import warnings
 #warnings.filterwarnings("error")
+
+#-------------------------------------------------------------------------------------------------------------
+def lnprob_yc(H0, Om0, sigma8, Ob0, ns, tenToA0, B0, Mpivot, sigma_int):
+    """Log likelihood function for use with Cobaya.
+        
+    """
+
+    # This can fail if wander outside parameter space where mass function is defined
+    try:
+        selFn.update(H0, Om0, Ob0, sigma8, ns, scalingRelationDict = {'tenToA0': tenToA0, 'B0': B0, 
+                                                                      'Mpivot': Mpivot, 'sigma_int': sigma_int})
+    except:
+        return -np.inf
 
+    mockTab=selFn.generateMockSample()
+    mockGrid, mock_log10ycBinEdges, mock_zBinEdges=np.histogram2d(np.log10(mockTab['fixed_y_c']), mockTab['redshift'], 
+                                                                  bins = [log10ycBinEdges, selFn.mockSurvey.zBinEdges])
+    mockGrid=mockGrid/selFn.mockOversampleFactor
+
+    # Poisson probability in (log10(y0~), z) grid
+    mask=np.logical_and(np.greater(obsGrid, 0), np.greater(mockGrid, 0))
+    if mask.sum() > 0:
+        lnlike=np.sum(obsGrid[mask]*np.log(mockGrid[mask])-mockGrid[mask]-np.log(factorial(obsGrid[mask])))
+    else:
+        lnlike=-np.inf
+
+    return lnlike
+
 #-------------------------------------------------------------------------------------------------------------
-def lnprob(H0, Om0, sigma8, Ob0, ns, tenToA0, B0, Mpivot, sigma_int):
+def lnprob_M500(H0, Om0, sigma8, Ob0, ns, tenToA0, B0, Mpivot, sigma_int):
     """Log likelihood function for use with Cobaya.
         
     """
@@ -87,7 +114,7 @@ def makeGetDistPlot(cosmoOutDir):
 
 #------------------------------------------------------------------------------------------------------------
 if __name__ == '__main__':
-
+    
     defaultMaxSamples=3000
     defaultNumToBurn=50
     parser=argparse.ArgumentParser("nemoCosmo")
@@ -114,6 +141,15 @@ if __name__ == '__main__':
     parser.add_argument("-b", "--burn", dest = "numToBurn", help = """Number of samples to burn. If given,
                         overrides the value given in the Cobaya configuration file.""", 
                         default = defaultNumToBurn, type = int)
+    parser.add_argument("-t", "--test-likelihood", dest = "testLikelihood", help = """Run a test of the
+                        likelihood, varying one parameter at a time, instead of running Cobaya,
+                        producing diagnostic plots called 'testLikelihood_paramName.png' in the current
+                        working directory. This test assumes you are running on a mock catalog with 
+                        cosmology and scaling relation parameters that match the WebSky simulations.""",
+                        action = "store_true", default = False)
+    parser.add_argument("-L", "--likelihood", dest = "likelihoodType", help = """The likelihood function
+                        to use. Options are 'mass', or 'yc' (latter is even more experimental).""", 
+                        default = "mass")
     args = parser.parse_args()
 
     tabFileName=args.catalogFileName
@@ -124,6 +160,13 @@ if __name__ == '__main__':
     cosmoOutDir=args.cosmoOutDir
     numToBurn=args.numToBurn
 
+    if args.likelihoodType == "mass":
+        lnprob=lnprob_M500
+    elif args.likelihoodType == "yc":
+        lnprob=lnprob_yc
+    else:
+        raise Exception("Didn't understand likelihoodType - use either 'mass' or 'yc'")
+
     tab=atpy.Table().read(tabFileName)
     if 'redshift' not in tab.keys():
         raise Exception("no 'redshift' column in catalog")
@@ -173,13 +216,96 @@ if __name__ == '__main__':
         info['sampler']['mcmc']['max_samples']=maxSamples
 
     print("Setting up SNR > %.2f selection function (footprint: %s)" % (SNRCut, footprintLabel))
-    selFn=completeness.SelFn(selFnDir, SNRCut, footprintLabel = footprintLabel, zStep = 0.1)
+    selFn=completeness.SelFn(selFnDir, SNRCut, footprintLabel = footprintLabel, zStep = 0.1,
+                             enableDrawSample = True, mockOversampleFactor = 10)
     tab=tab[np.where(tab['fixed_SNR'] > SNRCut)]
 
     # NOTE: Since we use such coarse z binning, we completely ignore photo-z errors here
     # Otherwise, signals.calcPM500 throws exception when z grid is not fine enough
     tab['redshiftErr'][:]=0.0
-
+
+    # Define binning on (log10(fixed_y_c), redshift) grid
+    log10ycBinEdges=np.linspace(np.log10(tab['fixed_y_c'].min()*0.8), np.log10(tab['fixed_y_c'].max()*1.1), 30)
+    obsGrid, obs_log10ycBinEdges, obs_zBinEdges=np.histogram2d(np.log10(tab['fixed_y_c']), tab['redshift'], 
+                                                               bins = [log10ycBinEdges, selFn.mockSurvey.zBinEdges])
+
+    # Testing
+    if args.testLikelihood == True:
+
+        def makeTestPlot(pRange, probs, var, label):
+            plt.plot(pRange, probs)
+            minP=min(probs)*1.1
+            maxP=max(probs)*0.9
+            plt.plot([var]*10, np.linspace(minP, maxP, 10), 'k--')
+            plt.xlabel(label)
+            plt.ylabel("lnprob")
+            plt.ylim(minP, maxP)
+            plt.savefig("testLikelihood_%s_%s.png" % (label, args.likelihoodType))
+            plt.close()
+
+        H0, Om0, Ob0, sigma8, ns = 68.0, 0.31, 0.049, 0.81, 0.965       # WebSky cosmology
+        tenToA0, B0, Mpivot, sigma_int = 2.65e-05, 0.0, 3.0e+14, 0.2    # WebSky scaling relation
+        #lnprob(H0, Om0, sigma8, Ob0, ns, p, B0, Mpivot, sigma_int)
+        testsToRun=['H0', 'Om0', 'sigma8', 'tenToA0']
+        print("Running likelihood tests")
+        if 'H0' in testsToRun:
+            label="H0"; var=H0
+            pRange=np.linspace(60, 80, 21)
+            probs=[]
+            count=0
+            t0=time.time()
+            print("%s:" % (label))
+            for p in pRange:
+                count=count+1
+                print("    %d/%d" % (count, len(pRange)))
+                probs.append(lnprob(p, Om0, sigma8, Ob0, ns, tenToA0, B0, Mpivot, sigma_int))
+            t1=time.time()
+            print("time per step = %.3f" % ((t1-t0)/len(probs)))
+            makeTestPlot(pRange, probs, var, label)
+        if 'Om0' in testsToRun:
+            label="Om0"; var=Om0
+            pRange=np.linspace(0.2, 0.4, 21)
+            probs=[]
+            count=0
+            t0=time.time()
+            print("%s:" % (label))
+            for p in pRange:
+                count=count+1
+                print(count, len(pRange))
+                probs.append(lnprob(H0, p, sigma8, Ob0, ns, tenToA0, B0, Mpivot, sigma_int))
+            t1=time.time()
+            print("time per step = %.3f" % ((t1-t0)/len(probs)))
+            makeTestPlot(pRange, probs, var, label)
+        if 'sigma8' in testsToRun:
+            label="sigma8"; var=sigma8
+            pRange=np.linspace(0.7, 0.9, 21)
+            probs=[]
+            count=0
+            t0=time.time()
+            print("%s:" % (label))
+            for p in pRange:
+                count=count+1
+                print(count, len(pRange))
+                probs.append(lnprob(H0, Om0, p, Ob0, ns, tenToA0, B0, Mpivot, sigma_int))
+            t1=time.time()
+            print("time per step = %.3f" % ((t1-t0)/len(probs)))
+            makeTestPlot(pRange, probs, var, label)
+        if 'tenToA0' in testsToRun:
+            label="tenToA0"; var=tenToA0
+            pRange=np.linspace(2.0e-5, 6.0e-05, 21)
+            probs=[]
+            count=0
+            t0=time.time()
+            print("%s:" % (label))
+            for p in pRange:
+                count=count+1
+                print(count, len(pRange))
+                probs.append(lnprob(H0, Om0, sigma8, Ob0, ns, p, B0, Mpivot, sigma_int))
+            t1=time.time()
+            print("time per step = %.3f" % ((t1-t0)/len(probs)))
+            makeTestPlot(pRange, probs, var, label)
+        sys.exit()
+
     # Run Cobaya...
     updated_info, products=run(info)
 

bin/nemoModel

@@ -21,8 +21,10 @@ if __name__ == '__main__':
 
     parser=argparse.ArgumentParser("nemoModel")
     parser.add_argument("catalogFileName", help = """A catalog (FITS table format) as produced by nemo.""")
-    parser.add_argument("templateFileName", help = """A FITS image file. The output sky model image will
-                        have the same pixelization.""")
+    parser.add_argument("maskFileName", help = """A FITS image file, containing a mask of the desired sky
+                        area. Non-zero values in the mask are used to define tiles (typically 10 x 5 deg),
+                        which are processed in parallel if MPI is enabled (see -M switch). The output sky
+                        model image will have the same pixelization as the mask image.""")
     parser.add_argument("beamFileName", help = """A file containing the beam profile, in the standard format
                         used by ACT.""")
     parser.add_argument("outputFileName", help = """The name of the output file that will contain the sky
@@ -49,18 +51,20 @@ if __name__ == '__main__':
         strictMPIExceptions=True
 
     # Create a stub config (then we can use existing start-up stuff to dish out the tiles)
+    print(">>> Setting up ...")
     parDict={}
     mapDict={}
-    mapDict['mapFileName']=args.templateFileName
+    mapDict['mapFileName']=args.maskFileName
     mapDict['obsFreqGHz']=args.obsFreqGHz
     mapDict['beamFileName']=args.beamFileName
     parDict['unfilteredMaps']=[mapDict]
     parDict['mapFilters']=[]
-    if args.MPIEnabled == True:
-        parDict['makeTileDir']=True
-        parDict['tileOverlapDeg']=1.0
-        parDict['tileDefLabel']='auto'
-        parDict['tileDefinitions']={'targetTileWidthDeg': 10.0, 'targetTileHeightDeg': 5.0}
+    #if args.MPIEnabled == True:
+    parDict['makeTileDir']=True
+    parDict['tileOverlapDeg']=1.0
+    parDict['tileDefLabel']='auto'
+    parDict['tileDefinitions']={'mask': args.maskFileName,
+                                'targetTileWidthDeg': 10.0, 'targetTileHeightDeg': 5.0}
 
     config=startUp.NemoConfig(parDict, MPIEnabled = args.MPIEnabled, divideTilesByProcesses = True,
                               makeOutputDirs = False, setUpMaps = True, writeTileDir = False, 
@@ -69,32 +73,55 @@ if __name__ == '__main__':
     tab=atpy.Table().read(args.catalogFileName)
 
     # Build a dictionary containing the model images which we'll stich together at the end
+    print(">>> Building models in tiles ...")
     modelsDict={}
     for tileName in config.tileNames:
-        shape=(config.tileCoordsDict[tileName]['clippedSection'][3], 
-               config.tileCoordsDict[tileName]['clippedSection'][1])
+        print("... %s ..." % (tileName))
+        shape=(config.tileCoordsDict[tileName]['clippedSection'][3]-config.tileCoordsDict[tileName]['clippedSection'][2], 
+               config.tileCoordsDict[tileName]['clippedSection'][1]-config.tileCoordsDict[tileName]['clippedSection'][0])
         wcs=astWCS.WCS(config.tileCoordsDict[tileName]['header'], mode = 'pyfits')
         try:
             modelsDict[tileName]=maps.makeModelImage(shape, wcs, tab, args.beamFileName, 
-                                                     obsFreqGHz = args.obsFreqGHz)
+                                                     obsFreqGHz = args.obsFreqGHz,
+                                                     validAreaSection = config.tileCoordsDict[tileName]['areaMaskInClipSection'])
         except:
             raise Exception("makeModelImage failed on tile '%s'" % (tileName))
 
     # Gathering
+    #if config.MPIEnabled == True:
+        #config.comm.barrier()
+        #gathered_modelsDicts=config.comm.gather(modelsDict, root = 0)
+        #if config.rank == 0:
+            #print("... gathered sky model tiles ...")
+            #for tileModelDict in gathered_modelsDicts:
+                #for tileName in tileModelDict.keys():
+                    #modelsDict[tileName]=tileModelDict[tileName]
+
+    # We can't just gather as that triggers a 32-bit overflow (?)
+    # So, sending one object at a time
     if config.MPIEnabled == True:
         config.comm.barrier()
-        gathered_modelsDicts=config.comm.gather(modelsDict, root = 0)
-        if config.rank == 0:
-            print("... gathered sky model tiles ...")
+        if config.rank > 0:
+            print("... rank = %d sending sky model tiles ..." % (config.rank))
+            config.comm.send(modelsDict, dest = 0)
+        elif config.rank == 0:
+            print("... gathering sky model tiles ...")
+            gathered_modelsDicts=[]
+            gathered_modelsDicts.append(modelsDict)
+            for source in range(1, config.size):
+                gathered_modelsDicts.append(config.comm.recv(source = source))
             for tileModelDict in gathered_modelsDicts:
                 for tileName in tileModelDict.keys():
                     modelsDict[tileName]=tileModelDict[tileName]
 
     # Stitching
+    print(">>> Stitching tiles ...")
     if config.rank == 0:
         d=np.zeros([config.origWCS.header['NAXIS2'], config.origWCS.header['NAXIS1']])
         wcs=config.origWCS
         for tileName in modelsDict.keys():
+            print("... %s ..." % (tileName))
             minX, maxX, minY, maxY=config.tileCoordsDict[tileName]['clippedSection']
-        d[minY:maxY, minX:maxX]=modelsDict[tileName]
+            if modelsDict[tileName] is not None:
+                d[minY:maxY, minX:maxX]=d[minY:maxY, minX:maxX]+modelsDict[tileName]
         astImages.saveFITS(args.outputFileName, d, wcs)

docs/scripts.rst

@@ -44,6 +44,17 @@ in the ``nemo`` configuration file. Forced photometry on sources of any kind can
 method.
 
 
+Making model images
+-------------------
+
+The nemoModel script can make model images from a Nemo catalog. If run using MPI, it will break the map
+into tiles to speed up processing, based on the given mask file. For example,::
+
+    mpiexec nemoModel S18d_202003_test/S18d_202003_test_optimalCatalog.fits maps/Mar2020/AdvACTSurveyMask_v7_S18.fits Beams/190809/b20190809_s16_pa2_f150_nohwp_night_beam_profile_jitter_cmb.txt testModel_f150.fits -f 149.6 -M
+
+will make a model image over the whole AdvACT footprint.
+
+
 Using the selection function
 ----------------------------
 

examples/AdvACT/S18d_202006.yml

@@ -51,7 +51,7 @@ massOptions: {tenToA0: 4.95e-5,
               Mpivot: 3.0e+14, 
               sigma_int: 0.2,
               relativisticCorrection: True,
-              rescaleFactor: 0.69, 
+              rescaleFactor: 0.70,
               rescaleFactorErr: 0.07,
               redshiftCatalog: "S18Clusters/AdvACT_confirmed.fits"}
 

examples/SOSims/cobayaConf_SOSims_fixedScalingRelation.yml

@@ -0,0 +1,39 @@
+params:
+  ns: 0.965
+  Ob0: 0.049
+  tenToA0: 2.65e-05
+  B0: 0.0
+  Mpivot: 300000000000000.0
+  sigma_int: 0.2
+  H0:
+    latex: H_0
+    prior:
+      dist: norm
+      loc: 70.0
+      scale: 4.0
+    proposal: 5.0
+  Om0:
+    latex: \Omega_{\rm m0}
+    prior:
+      max: 0.5
+      min: 0.1
+    proposal: 0.05
+    ref:
+      dist: norm
+      loc: 0.3
+      scale: 0.1
+  sigma8:
+    latex: \sigma_8
+    prior:
+      max: 0.9
+      min: 0.6
+    proposal: 0.02
+    ref:
+      dist: norm
+      loc: 0.8
+      scale: 0.1
+sampler:
+  mcmc:
+    burn_in: 50
+    max_samples: 3000
+    max_tries: .inf

examples/SOSims/slurm_small_cosmo.sh

@@ -0,0 +1,9 @@
+#!/bin/sh
+#SBATCH --nodes=4
+#SBATCH --ntasks-per-node=20
+#SBATCH --mem=64000
+#SBATCH --time=23:59:00
+
+source /home/mjh/SETUP_CONDA.sh
+time mpiexec nemoCosmo mocks_small/mockCatalog_1.fits MFMF_SOSim_3freq_small/selFn -c cobayaConf_SOSims_fixedScalingRelation.yml -S 5 -o small_cosmoResults1
+