Merge pull request #65 from simonsobs/dev

mattyowl · web-flow · commit 2db3c0076959 · 2023-09-01T01:24:56.000+02:00
Dev
diff --git a/INSTALL.rst b/INSTALL.rst
@@ -10,7 +10,7 @@ brackets, later versions also probably work):
 * Pillow (5.1.0)
 * astropy (4.0)
 * PyYAML (3.12)
-* `CCL <https://github.com/LSSTDESC/CCL>`_ (2.1 or later)
+* `CCL <https://github.com/LSSTDESC/CCL>`_ (3.0.0+)
 * mpi4py (3.0.0)
 * colorcet (1.0.0; https://github.com/bokeh/colorcet/releases)
 
diff --git a/bin/nemoMass b/bin/nemoMass
@@ -105,11 +105,7 @@ def calcMass(tab, massOptions, QFit, fRelWeightsDict, mockSurvey, otherMassEstim
     
     """
 
-    try:
-        refMassDef=ccl.halos.MassDef(massOptions['delta'], massOptions['rhoType'],
-                                     c_m_relation = massOptions['concMassRelation'])
-    except:
-        refMassDef=ccl.halos.MassDef(massOptions['delta'], massOptions['rhoType'])
+    refMassDef=ccl.halos.MassDef(massOptions['delta'], massOptions['rhoType'])
 
     if 'relativisticCorrection' not in massOptions.keys():
         massOptions['relativisticCorrection']=True
@@ -209,11 +205,9 @@ def calcMass(tab, massOptions, QFit, fRelWeightsDict, mockSurvey, otherMassEstim
                     if 'concMassRelation' not in massDefDict.keys():
                         massDefDict['concMassRelation']=None
                     thisLabel='M%d%s' % (massDefDict['delta'], massDefDict['rhoType'][0])
-                    try:
-                        thisMassDef=ccl.halos.MassDef(massDefDict['delta'], massDefDict['rhoType'], c_m_relation = massDefDict['concMassRelation'])
-                    except:
-                        thisMassDef=ccl.halos.MassDef(massDefDict['delta'], massDefDict['rhoType'])
-                    thisMass=signals.MDef1ToMDef2(resultDict[label]*1e14, row['redshift'], refMassDef, thisMassDef, mockSurvey.cosmoModel)/1e14
+                    thisMassDef=ccl.halos.MassDef(massDefDict['delta'], massDefDict['rhoType'])
+                    thisMass=signals.MDef1ToMDef2(resultDict[label]*1e14, row['redshift'], refMassDef, thisMassDef, mockSurvey.cosmoModel,
+                                                  c_m_relation = massDefDict['concMassRelation'])/1e14
                     row[thisLabel+suffix]=thisMass
                     row[thisLabel+suffix+'_errPlus']=(row[label+suffix+'_errPlus']/row[label+suffix])*row[thisLabel+suffix]
                     row[thisLabel+suffix+'_errMinus']=(row[label+suffix+'_errMinus']/row[label+suffix])*row[thisLabel+suffix]
@@ -362,8 +356,6 @@ if __name__ == '__main__':
     mockSurvey=MockSurvey.MockSurvey(minMass, areaDeg2, zMin, zMax, H0, Om0, Ob0, sigma8, ns,
                                      rhoType = massOptions['rhoType'], delta = massOptions['delta'])
         
-    # Seems like multiprocessing doesn't work under slurm...   
-    # So use MPI instead: just divide up catalog among processes
     tab.add_column(atpy.Column(np.arange(len(tab)), "sortIndex"))
     if config.MPIEnabled == True:
         numRowsPerProcess=int(np.ceil(len(tab)/config.size))
@@ -372,7 +364,7 @@ if __name__ == '__main__':
         if config.rank == config.size-1:
             endIndex=len(tab)
         tab=tab[startIndex:endIndex]
-        
+
     tab=calcMass(tab, massOptions, Q, fRelWeightsDict, mockSurvey, otherMassEstimates = otherMassEstimates)
     
     if config.MPIEnabled == True:
diff --git a/bin/nemoMock b/bin/nemoMock
@@ -54,7 +54,7 @@ def makeParser():
                         fixed_SNR > this value.", default = 4.0, type = float)
     parser.add_argument("-Q", "--Q-source", dest="QSource", help = "Source of the Q function data - either\
                         'fit' (the 'classic' method) or 'injection' (for Q based on source injection test\
-                        results.", default = 'fit')
+                        results).", default = 'fit')
     #parser.add_argument("-M", "--mpi", dest="MPIEnabled", action="store_true", help="""Enable MPI. If you 
                         #want to use this, run with e.g., mpiexec -np 4 nemoMock selFnDir mocksDir -M""", 
                         #default = False)
diff --git a/nemo/MockSurvey.py b/nemo/MockSurvey.py
@@ -63,7 +63,8 @@ class MockSurvey(object):
     """
     def __init__(self, minMass, areaDeg2, zMin, zMax, H0, Om0, Ob0, sigma8, ns, zStep = 0.01, 
                  enableDrawSample = False, delta = 500, rhoType = 'critical', 
-                 transferFunction = 'boltzmann_camb', massFunction = 'Tinker08'):
+                 transferFunction = 'boltzmann_camb', massFunction = 'Tinker08',
+                 c_m_relation = 'Bhattacharya13'):
         """Create a MockSurvey object, for performing calculations of cluster counts or generating mock
         catalogs. The Tinker et al. (2008) halo mass function is used (hardcoded at present, but in 
         principle this can easily be swapped for any halo mass function supported by CCL).
@@ -89,7 +90,9 @@ def __init__(self, minMass, areaDeg2, zMin, zMax, H0, Om0, Ob0, sigma8, ns, zSte
                 'boltzmann_camb').
             massFunction (:obj:`str`): Name of the mass function to use, currently either 'Tinker08' or
                 'Tinker10'. Mass function calculations are done by CCL.
-                
+            c_m_relation ('obj':`str`): Name of the concentration -- mass relation to assume, as understood by
+                CCL (this may be used internally for conversion between mass definitions, as needed).
+
         """
         
         if areaDeg2 == 0:
@@ -104,11 +107,8 @@ def __init__(self, minMass, areaDeg2, zMin, zMax, H0, Om0, Ob0, sigma8, ns, zSte
         
         self.delta=delta
         self.rhoType=rhoType
-        if self.delta == 200:
-            c_m_relation='Bhattacharya13'
-        else:
-            c_m_relation=None
-        self.mdef=ccl.halos.MassDef(self.delta, self.rhoType, c_m_relation = c_m_relation)
+        self.c_m_relation=c_m_relation
+        self.mdef=ccl.halos.MassDef(self.delta, self.rhoType)
         self.transferFunction=transferFunction
         self.massFuncName=massFunction
         
@@ -129,7 +129,10 @@ def __init__(self, minMass, areaDeg2, zMin, zMax, H0, Om0, Ob0, sigma8, ns, zSte
                                         self.log10M.max()+(self.log10M[1]-self.log10M[0])/2, len(self.log10M)+1)  
 
         # Below is needed for Q calc when not using M500c definition (for now at least)
-        self._M500cDef=ccl.halos.MassDef(500, "critical")
+        if self.delta != 500 and self.rhoType != 'critical':
+            self._M500cDef=ccl.halos.MassDef(500, "critical")
+            self._transToM500c=ccl.halos.mass_translator(mass_in = self.mdef, mass_out = self._M500cDef,
+                                                         concentration = self.c_m_relation)
 
         self.enableDrawSample=enableDrawSample
         self.update(H0, Om0, Ob0, sigma8, ns)
@@ -161,18 +164,16 @@ def _get_new_cosmo(self, H0, Om0, Ob0, sigma8, ns):
             self.Ob0=Ob0
             self.sigma8=sigma8
             self.ns=ns
-            self.cosmoModel = ccl.Cosmology(Omega_c=Om0-Ob0,
-                                            Omega_b=Ob0,
-                                            h=0.01*H0,
-                                            sigma8=sigma8,
-                                            n_s=ns,
-                                            transfer_function=self.transferFunction)
+            self.cosmoModel=ccl.Cosmology(Omega_c=Om0-Ob0,
+                                          Omega_b=Ob0,
+                                          h=0.01*H0,
+                                          sigma8=sigma8,
+                                          n_s=ns,
+                                          transfer_function=self.transferFunction)
             if self.massFuncName == 'Tinker10':
-                self.mfunc=ccl.halos.MassFuncTinker10(self.cosmoModel,
-                                                      self.mdef)
+                self.mfunc=ccl.halos.MassFuncTinker10(mass_def = self.mdef)
             elif self.massFuncName == 'Tinker08':
-                self.mfunc=ccl.halos.MassFuncTinker08(self.cosmoModel,
-                                                      self.mdef)
+                self.mfunc=ccl.halos.MassFuncTinker08(mass_def = self.mdef)
 
             
     def update(self, H0, Om0, Ob0, sigma8, ns):
@@ -205,11 +206,8 @@ def update(self, H0, Om0, Ob0, sigma8, ns):
                 interpLim_minLog10M500c=self.log10M.min()
                 interpLim_maxLog10M500c=self.log10M.max()
             else:
-                interpLim_minLog10M500c=np.log10(self.mdef.translate_mass(self.cosmoModel, self.M.min(), 
-                                                                          self.a[k], self._M500cDef))
-                interpLim_maxLog10M500c=np.log10(self.mdef.translate_mass(self.cosmoModel, self.M.max(), 
-                                                                          self.a[k], self._M500cDef))
-
+                interpLim_minLog10M500c=np.log10(self._transToM500c(self.cosmoModel, self.M.min(), self.a[k]))
+                interpLim_maxLog10M500c=np.log10(self._transToM500c(self.cosmoModel, self.M.max(), self.a[k]))
             zk=self.z[k]
             interpPoints=100
             fitM500s=np.power(10, np.linspace(interpLim_minLog10M500c, interpLim_maxLog10M500c, interpPoints))
@@ -251,8 +249,7 @@ def _cumulativeNumberDensity(self, z):
         """
 
         h=self.cosmoModel['h']
-        dndlnM=self.mfunc.get_mass_function(self.cosmoModel,
-                                            self.M, 1/(1+z)) / np.log(10) #/ h**3
+        dndlnM=self.mfunc(self.cosmoModel, self.M, 1/(1+z)) / np.log(10)
         dndM=dndlnM/self.M
         ngtm=integrate.cumtrapz(dndlnM[::-1], np.log(self.M), initial = 0)[::-1]
         
@@ -293,8 +290,7 @@ def _doClusterCount(self):
             zShellMin=zRange[i]
             zShellMax=zRange[i+1]
             zShellMid=(zShellMax+zShellMin)/2.
-            dndlnM=self.mfunc.get_mass_function(self.cosmoModel, self.M,
-                                                1./(1+zShellMid)) * norm_mfunc
+            dndlnM=self.mfunc(self.cosmoModel, self.M, 1./(1+zShellMid)) * norm_mfunc
             dndM = dndlnM / self.M
             n=dndM * np.gradient(self.M)
             numberDensity.append(n)
@@ -362,7 +358,7 @@ def drawSample(self, y0Noise, scalingRelationDict, QFit = None, wcs = None, phot
                    tileName = None, SNRLimit = None, makeNames = False, z = None, numDraws = None,\
                    areaDeg2 = None, applySNRCut = False, applyPoissonScatter = True,\
                    applyIntrinsicScatter = True, applyNoiseScatter = True,\
-                   applyRelativisticCorrection = True, verbose = False):
+                   applyRelativisticCorrection = True, verbose = False, biasModel = None):
         """Draw a cluster sample from the mass function, generating mock y0~ values (called `fixed_y_c` in
         Nemo catalogs) by applying the given scaling relation parameters, and then (optionally) applying
         a survey selection function.
@@ -555,9 +551,7 @@ def drawSample(self, y0Noise, scalingRelationDict, QFit = None, wcs = None, phot
             if self.delta == 500 and self.rhoType == "critical":
                 log10M500cs[mask]=log10Ms[mask]
             else:
-                log10M500cs[mask]=np.log10(self.mdef.translate_mass(self.cosmoModel, np.power(10, log10Ms[mask]),
-                                                                1/(1+zk), self._M500cDef))
-
+                log10M500cs[mask]=np.log10(self._transToM500c(self.cosmoModel, np.power(10, log10Ms[mask]), 1/(1+zk)))
             theta500s=interpolate.splev(log10M500cs[mask], self.theta500Splines[k], ext = 3)
             if QFit is not None:
                 Qs[mask]=QFit.getQ(theta500s, z = zk, tileName = tileName)
@@ -609,8 +603,15 @@ def drawSample(self, y0Noise, scalingRelationDict, QFit = None, wcs = None, phot
             tab.add_column(atpy.Column(true_y0s/1e-4, 'true_fixed_y_c'))
             tab.add_column(atpy.Column(measured_y0s/1e-4, 'fixed_y_c'))
             tab.add_column(atpy.Column(y0Noise/1e-4, 'fixed_err_y_c'))
+            tab['true_fixed_SNR']=tab['true_fixed_y_c']/tab['fixed_err_y_c']  # True truth, but pre-intrinsic and measurement scatter
+            # tab['true_fixed_SNR']=(scattered_y0s/1e-4)/tab['fixed_err_y_c']     # With intrinsic scatter, no measurement scatter
+            # tab['true_fixed_SNR']=tab['fixed_y_c']/tab['fixed_err_y_c']         # Like forced photometry case on a real map at true location
+            # Apply optimization bias first, then it'll feed through to SNR automatically
+            if biasModel is not None:
+                corrFactors=biasModel['func'](tab['true_fixed_SNR'], biasModel['params'][0], biasModel['params'][1], biasModel['params'][2])
+                tab['fixed_y_c']=tab['fixed_y_c']*corrFactors
+            tab['fixed_SNR']=tab['fixed_y_c']/tab['fixed_err_y_c']
 
-        tab.add_column(atpy.Column(measured_y0s/y0Noise, 'fixed_SNR'))
         tab.add_column(atpy.Column(zs, 'redshift'))
         tab.add_column(atpy.Column(zErrs, 'redshiftErr'))
         if photFilterLabel is not None and tileName is not None:
@@ -619,7 +620,7 @@ def drawSample(self, y0Noise, scalingRelationDict, QFit = None, wcs = None, phot
                 
         # Apply selection?
         if applySNRCut == True:
-            selMask=measured_y0s > y0Noise*SNRLimit
+            selMask=tab['fixed_SNR'] > tab['fixed_err_y_c']*SNRLimit
             tab=tab[selMask]
         t1=time.time()
 
diff --git a/nemo/completeness.py b/nemo/completeness.py
@@ -135,9 +135,10 @@ def __init__(self, selFnDir, SNRCut, configFileName = None, footprint = None, zS
                  downsampleRMS = True, applyMFDebiasCorrection = True, applyRelativisticCorrection = True,
                  setUpAreaMask = False, enableCompletenessCalc = True, delta = 500, rhoType = 'critical',
                  massFunction = 'Tinker08', maxTheta500Arcmin = None, method = 'fast',
-                 QSource = 'fit', noiseCut = None):
+                 QSource = 'fit', noiseCut = None, biasModel = None):
         
         self.SNRCut=SNRCut
+        self.biasModel=biasModel
         if footprint == 'full':
             footprint=None
         self.footprint=footprint
@@ -436,9 +437,14 @@ def update(self, H0, Om0, Ob0, sigma8, ns, scalingRelationDict = None):
                 log_y0=np.log(y0Grid)
                 compMz=np.zeros(log_y0.shape)
                 for i in range(len(RMSTab)):
+                    if self.biasModel is not None:
+                        trueSNR=y0Grid/RMSTab['y0RMS'][i]
+                        corrFactors=self.biasModel['func'](trueSNR, self.biasModel['params'][0], self.biasModel['params'][1], self.biasModel['params'][2])
+                    else:
+                        corrFactors=np.ones(y0Grid.shape)
                     # With intrinsic scatter [may not be quite right, but a good approximation]
                     totalLogErr=np.sqrt((RMSTab['y0RMS'][i]/y0Grid)**2 + self.scalingRelationDict['sigma_int']**2)
-                    sfi=stats.norm.sf(y0Lim[i], loc = y0Grid, scale = totalLogErr*y0Grid)
+                    sfi=stats.norm.sf(y0Lim[i], loc = y0Grid*corrFactors, scale = totalLogErr*(y0Grid*corrFactors))
                     compMz=compMz+sfi*areaWeights[i]
                     # No intrinsic scatter, Gaussian noise
                     #sfi=stats.norm.sf(y0Lim[i], loc = y0Grid, scale = RMSTab['y0RMS'][i])
@@ -450,6 +456,7 @@ def update(self, H0, Om0, Ob0, sigma8, ns, scalingRelationDict = None):
             compMzCube=np.array(compMzCube)
             y0GridCube=np.array(y0GridCube)
             self.compMz=np.average(compMzCube, axis = 0, weights = self.fracArea)
+            # self.compMz=np.einsum('ijk,i->jk', np.nan_to_num(compMzCube), self.fracArea) # Equivalent, for noting
             self.y0TildeGrid=np.average(y0GridCube, axis = 0, weights = self.fracArea)
 
 
@@ -466,10 +473,9 @@ def _makeSignalGrids(self, applyQ = True, tileName = None):
             zk=zRange[i]
             k=np.argmin(abs(self.mockSurvey.z-zk))
             if self.mockSurvey.delta != 500 or self.mockSurvey.rhoType != "critical":
-                log10M500s=np.log10(self.mockSurvey.mdef.translate_mass(self.mockSurvey.cosmoModel,
-                                                                        self.mockSurvey.M,
-                                                                        self.mockSurvey.a[k],
-                                                                        self.mockSurvey._M500cDef))
+                log10M500s=np.log10(self.mockSurvey._transToM500c(self.mockSurvey.cosmoModel,
+                                                                  self.mockSurvey.M,
+                                                                  self.mockSurvey.a[k]))
             else:
                 log10M500s=self.mockSurvey.log10M
             theta500s_zk=interpolate.splev(log10M500s, self.mockSurvey.theta500Splines[k])
@@ -613,7 +619,8 @@ def generateMockSample(self, mockOversampleFactor = None, applyPoissonScatter =
                                                applyPoissonScatter = applyPoissonScatter,
                                                applyIntrinsicScatter = True,
                                                applyNoiseScatter = True,
-                                               applyRelativisticCorrection = self.applyRelativisticCorrection)
+                                               applyRelativisticCorrection = self.applyRelativisticCorrection,
+                                               biasModel = self.biasModel)
             if mockTab is not None and len(mockTab) > 0:
                 mockTabsList.append(mockTab)
         tab=atpy.vstack(mockTabsList)
@@ -1349,8 +1356,9 @@ def calcCompleteness(RMSTab, SNRCut, tileName, mockSurvey, massOptions, QFit, pl
             if massOptions['delta'] == 500 and massOptions['rhoType'] == "critical":
                 log10M500cs=mockSurvey.log10M
             else:
-                log10M500cs=np.log10(mockSurvey.mdef.translate_mass(mockSurvey.cosmoModel, np.power(10, mockSurvey.log10M),
-                                                                    1/(1+zk), mockSurvey._M500cDef))
+                log10M500cs=np.log10(mockSurvey._transToM500c(mockSurvey.cosmoModel,
+                                                              np.power(10, mockSurvey.log10M),
+                                                              1/(1+zk)))
 
             theta500s_zk=interpolate.splev(log10M500cs, mockSurvey.theta500Splines[k])
             Qs_zk=QFit.getQ(theta500s_zk, z = zk, tileName = tileName)
@@ -1362,7 +1370,7 @@ def calcCompleteness(RMSTab, SNRCut, tileName, mockSurvey, massOptions, QFit, pl
             
         # For some cosmological parameters, we can still get the odd -ve y0
         y0Grid[y0Grid <= 0] = 1e-9
-        
+
         # Calculate completeness using area-weighted average
         # NOTE: RMSTab that is fed in here can be downsampled in noise resolution for speed
         areaWeights=RMSTab['areaDeg2']/RMSTab['areaDeg2'].sum()
diff --git a/nemo/pipelines.py b/nemo/pipelines.py
@@ -60,7 +60,7 @@ def filterMapsAndMakeCatalogs(config, rootOutDir = None, useCachedFilters = Fals
     """
     
     # Multi-pass pipeline, enabled with use of filterSets parameter in config file
-    if config.filterSets != [] and useCachedFilters == False:
+    if config.filterSets != [] and useCachedFilters == False and useCachedFilteredMaps == False:
         # If we wanted to save results from each step, could set-up filterSet specific diagnostics dir here
         if rootOutDir is None:
             rootOutDir=config.rootOutDir
diff --git a/nemo/signals.py b/nemo/signals.py
diff --git a/setup.cfg b/setup.cfg
