Merge pull request #54 from mirochaj/jwst

a few additions to GalaxyEnsemble machinery

ares/analysis/GalaxyPopulation.py

@@ -41,7 +41,8 @@
 datasets_lf = ('bouwens2015', 'finkelstein2015', 'bowler2020', 'stefanon2019',
     'mclure2013', 'parsa2016', 'atek2015',  'alavi2016',
     'reddy2009', 'weisz2014', 'bouwens2017', 'oesch2018', 'oesch2013',
-    'oesch2014', 'vanderburg2010', 'morishita2018', 'rojasruiz2020')
+    'oesch2014', 'vanderburg2010', 'morishita2018', 'rojasruiz2020',
+    'harikane2022')
 datasets_smf = ('song2016', 'stefanon2017', 'duncan2014', 'tomczak2014',
 	'moustakas2013', 'mortlock2011', 'marchesini2009_10', 'perez2008')
 datasets_mzr = ('sanders2015',)
@@ -53,7 +54,7 @@
  'dropouts': ('parsa2016', 'bouwens2015',
     'finkelstein2015', 'bowler2020','stefanon2019', 'mclure2013',
     'vanderburg2010', 'reddy2009', 'oesch2018', 'oesch2013', 'oesch2014',
-    'morishita2018', 'rojasruiz2020'),
+    'morishita2018', 'rojasruiz2020', 'harikane2022'),
  'lensing': ('alavi2016', 'atek2015', 'bouwens2017'),
  'local': ('weisz2014,'),
  'all': datasets_lf,

ares/physics/HaloMassFunction.py

@@ -1421,7 +1421,7 @@ def tab_prefix_hmf(self, with_size=False):
 
     def SaveHMF(self, fn=None, clobber=False, destination=None, fmt='hdf5',
         save_MAR=True):
-        self.save(fn=fn, clobber=clobber, destination=destination, fm=fmt,
+        self.save(fn=fn, clobber=clobber, destination=destination, fmt=fmt,
             save_MAR=save_MAR)
 
     def save(self, fn=None, clobber=False, destination=None, fmt='hdf5',

ares/populations/GalaxyEnsemble.py

@@ -272,6 +272,72 @@ def _cache_halos(self):
     def _cache_halos(self, value):
         self._cache_halos_ = value
 
+    def _get_target_halos(self, raw):
+        Nz = raw['z'].size
+        tvol = self.pf['pop_target_volume']
+        tred = self.pf['pop_target_redshift']
+        tseed = self.pf['pop_target_seed']
+
+        ##
+        # Modify raw histories if volume provided.
+        assert tred is not None
+
+        # Here, each element corresponds to a distinct halo.
+        # Will add Poisson noise only.
+
+        # Loop over bins, draw halos, then add noise etc.
+        iztred = np.argmin(np.abs(tred - raw['z']))
+
+        # Mean number of halos in each mass bin at target redshift
+        Nlam = raw['nh'][:,iztred] * tvol
+
+        np.random.seed(seed=tseed)
+
+        # Actual number in target volume
+        Nact = np.random.poisson(lam=Nlam, size=Nlam.size)
+
+        Mmin = self.guide.Mmin(raw['z'])
+
+        # Loop over mass bins, create halo histories.
+        # Need to create dictionary containing 'Mh', 'MAR', 'nh', etc.,
+        # where each is 2-D array in (halo number, time)
+        ct = 0
+        for i, Nbin in enumerate(Nact):
+
+            # Skip Mh < Mmin for efficiency.
+            if np.all(raw['Mh'][i,iztred:] < Mmin[iztred:]):
+                continue
+
+            if Nbin == 0:
+                continue
+
+            _Mh = raw['Mh'][i,:]
+            _MAR = raw['MAR'][i,:]
+
+            _MhT = np.tile(_Mh, (Nbin, 1))
+            _MART = np.tile(_MAR, (Nbin, 1))
+
+            if ct == 0:
+                print("# Generating {} individual halos from {} mass bins...".format(
+                    Nact[i:].sum(), Nact.size))
+
+                Mh = _MhT
+                MAR = _MART
+            else:
+                Mh = np.vstack((Mh, _MhT))
+                MAR = np.vstack((MAR, _MART))
+
+            ct += 1
+
+        nh = np.ones(Mh.shape)
+
+        out = raw.copy()
+        out['Mh'] = Mh
+        out['nh'] = nh
+        out['MAR'] = MAR
+
+        return out
+
     def _gen_halo_histories(self):
         """
         From a set of smooth halo assembly histories, build a bigger set
@@ -281,9 +347,13 @@ def _gen_halo_histories(self):
         if hasattr(self, '_cache_halos_'):
             return self._cache_halos
 
+        thin = self.pf['pop_thin_hist']
+
         raw = self.load()
 
-        thin = self.pf['pop_thin_hist']
+        if self.pf['pop_target_volume'] is not None:
+            assert thin in [None, 0, 1, False]
+            raw = self._get_target_halos(raw)
 
         if thin > 1:
             assert self.pf['pop_histories'] is None, \
@@ -1031,10 +1101,16 @@ def _gen_prescribed_galaxy_histories(self, zstop=0):
         if 'SFR' in halos:
             SFR = halos['SFR'][:,-1::-1]
         else:
-            iz = np.argmin(np.abs(6. - z))
-            SFR = self.guide.get_fstar(z=z2d, Mh=Mh)
-            np.multiply(SFR, MAR, out=SFR)
-            SFR *= fb
+            if self.pf['pop_sfr'] is not None:
+                SFR = np.ones_like(MAR) * self.pf['pop_sfr']
+                sigma_mar = self.pf['pop_scatter_mar']
+                np.random.seed(self.pf['pop_scatter_mar_seed'])
+                noise = self.get_noise_lognormal(SFR, sigma_mar)
+                SFR += noise
+            else:
+                SFR = self.guide.get_fstar(z=z2d, Mh=Mh)
+                np.multiply(SFR, MAR, out=SFR)
+                SFR *= fb
 
         ##
         # Duty cycle effects
@@ -1105,12 +1181,21 @@ def _gen_prescribed_galaxy_histories(self, zstop=0):
             Mmin = self.halos.VirialMass(z2d, self.pf['pop_Tmin'])
             above_Mmin = Mh >= Mmin
 
+        # Cut out Mh < Mmin galaxies
+        if self.pf['pop_Mmax'] is not None:
+            above_Mmax = Mh >= self.pf['pop_Mmax']
+        else:
+            above_Mmax = np.zeros_like(Mh)
+
         # Bye bye guys
         SFR[above_Mmin==False] = 0
+        SFR[above_Mmax==True] = 0
 
         ##
         # Introduce some by-hand quenching.
-        if self.pf['pop_quench'] is not None:
+        qmethod = self.pf['pop_quench_method']
+        #print('HELLO', qmethod, self.pf['pop_quench'])
+        if (self.pf['pop_quench'] is not None) and (qmethod == 'zreion'):
             zreion = self.pf['pop_quench']
             if type(zreion) in [np.ndarray, np.ma.core.MaskedArray]:
                 assert zreion.size == Mh.size, \
@@ -1166,6 +1251,23 @@ def _gen_prescribed_galaxy_histories(self, zstop=0):
                 if sum(SFR[ok==1] == 0) < sum(is_quenched[ok==1]):
                     raise ValueError(err)
 
+        elif (self.pf['pop_quench'] is not None) and \
+           (qmethod in ['Mh', 'mass', 'sfr']):
+
+            is_quenched = np.zeros_like(Mh)
+
+            if qmethod in ['mass', 'Mh']:
+                is_quenched[Mh > self.pf['pop_quench']] = 1
+            elif qmethod == 'sfr':
+                is_quenched[SFR > self.pf['pop_quench']] = 1
+            else:
+                raise NotImplemented('help')
+
+            # Bye bye guys
+            SFR[is_quenched==True] = 0
+
+            print('quenched', is_quenched.sum() / float(is_quenched.size))
+
 
         # Stellar mass should have zeros padded at the 0th time index
         Ms = np.hstack((zeros_like_Mh,
@@ -1181,19 +1283,24 @@ def _gen_prescribed_galaxy_histories(self, zstop=0):
         ##
         if have_dust:
 
-            delay = self.pf['pop_dust_yield_delay']
+            if self.pf['pop_dust_yield_delay'] not in [0, None]:
+                have_delay = True
+            else:
+                have_delay = False
+
+            delay = self.guide.dust_yield_delay(z=z2d, Mh=Mh)
 
             if np.all(fg == 0):
-                if type(fd) in [int, float, np.float64] and delay == 0:
+                if type(fd) in [int, float, np.float64] and (not have_delay):
                     Md = fd * fZy * Ms
                 else:
 
-                    if delay > 0:
+                    if have_delay:
 
                         assert np.allclose(np.diff(dt_myr), 0.0,
                             rtol=1e-5, atol=1e-5)
 
-                        shift = int(delay // dt_myr[0])
+                        shift = np.array(delay // dt_myr[0], dtype=int)
 
                         # Need to fix so Mh-dep fd can still work.
                         assert type(fd) in [int, float, np.float64]
@@ -1361,6 +1468,11 @@ def _gen_prescribed_galaxy_histories(self, zstop=0):
             else:
                 pos = None
 
+        # Can null dust reddening by hand if we want.
+        if np.isfinite(self.pf['pop_dust_kill_redshift']):
+            Sd[np.argwhere(z2d > self.pf['pop_dust_kill_redshift'])] = 0.0
+            Md[np.argwhere(z2d > self.pf['pop_dust_kill_redshift'])] = 0.0
+
         del z2d
 
         # Pack up
@@ -1578,6 +1690,36 @@ def get_field(self, z, field):
         iz = np.argmin(np.abs(z - self.histories['z']))
         return self.histories[field][:,iz]
 
+    def get_ages(self, z, mass_frac=0.5):
+        """
+        Get age of galaxies (in Myr), assumed to be time since they assembled
+        `mass_frac` times their mass at redshift `z`.
+        """
+
+        iz = np.argmin(np.abs(z - self.histories['z']))
+
+        Ms = self.get_field(z, 'Ms')
+        ages = np.zeros_like(Ms)
+
+        tasc = self.histories['t'][-1::-1]
+        zasc = self.histories['z'][-1::-1]
+        for i, _z_ in enumerate(zasc):
+            if _z_ <= z:
+                continue
+
+            Ms_z_ = self.get_field(_z_, 'Ms')
+            is_double = Ms_z_ < mass_frac * Ms
+            is_new = ages == 0
+
+            ok = is_double * is_new
+            dt = self.histories['t'][iz] - tasc[i]
+            ages[ok==1] = dt
+
+            if np.all(ages > 0):
+                break
+
+        return ages
+
     def StellarMassFunction(self, z, bins=None, units='dex'):
         return self.get_smf(z, bins=bins, units=units)
 
@@ -2800,11 +2942,20 @@ def get_lf(self, z, bins=None, use_mags=True, wave=1600.,
             assert y[yok==1].min() < x.min()
             assert y[yok==1].max() > x.max()
 
-        hist, bin_histedges = np.histogram(y[yok==1],
-            weights=w[yok==1], bins=bin_c2e(x), density=True)
+        #if self.pf['pop_fobsc']:
+        #fobsc = (1. - self.guide.fobsc(z=z, Mh=self.halos.tab_M))
 
-        N = np.sum(w[yok==1])
-        phi = hist * N
+        if np.all(w[yok==1] == 1):
+            hist, bin_histedges = np.histogram(y[yok==1],
+                bins=bin_c2e(x), density=False)
+            dbin = bin_histedges[1] - bin_histedges[0]
+            phi = hist / self.pf['pop_target_volume'] / dbin
+        else:
+            hist, bin_histedges = np.histogram(y[yok==1],
+                weights=w[yok==1], bins=bin_c2e(x), density=True)
+
+            N = np.sum(w[yok==1])
+            phi = hist * N
 
         #self._cache_lf_[(z, wave)] = x, phi
 

ares/util/SetDefaultParameterValues.py

@@ -513,6 +513,11 @@ def PopulationParameters():
 
     "pop_type": 'galaxy',
 
+    "pop_target_volume": None,
+    "pop_target_redshift": None,
+    "pop_target_density": 0.,
+    "pop_target_seed": None,
+
     "pop_tunnel": None,
 
     "pop_sfr_model": 'fcoll', # or sfrd-func, sfrd-tab, sfe-func, sfh-tab, rates,
@@ -651,8 +656,6 @@ def PopulationParameters():
     "pop_fstar_max": 1.0,
     "pop_fstar_negligible": 1e-5, # relative to maximum
 
-    "pop_sfr": None,
-
     "pop_facc": 0.0,
     "pop_fsmooth": 1.0,
 
@@ -677,7 +680,7 @@ def PopulationParameters():
     "pop_aging": False,
     "pop_enrichment": False,
     "pop_quench": None,
-    "pop_quench_by": 'mass',
+    "pop_quench_method": 'zreion',
     "pop_flag_sSFR": None,
     "pop_flag_tau": None,
 

input/litdata/furlanetto2017.py

@@ -6,14 +6,14 @@
 
  'pop_fstar': None,
  'pop_fstar_max': 0.1,            # fstar <= this value
- 
+
  # SFE (through mass loading factor)
  'pop_sfr_model': 'mlf-func',
  'pop_mlf': 'pq[0]',
  'pq_func[0]': 'pl_evolN',
  'pq_func_var[0]': 'Mh',
  'pq_func_var2[0]': '1+z',
- 
+
  ##
  # Steve's Equation 13.
  ##
@@ -22,8 +22,8 @@
  'pq_func_par2[0]': -2./3.,
  'pq_func_par3[0]': 9.,
  'pq_func_par4[0]': -1.,
- 
- 'pop_L1600_per_sfr': 1e-28,
+
+ 'pop_lum_per_sfr': 1e-28,
 
 }
 
@@ -44,7 +44,7 @@
  'pq_func[1]': 'pl_evolN',
  'pq_func_var[1]': 'Mh',
  'pq_func_var2[1]': '1+z',
- 
+
  'pq_val_ceil[1]': 1.0,           # fshock <= 1
 
  # Steve's Equation 6 (from Faucher-Giguere+ 2011)
@@ -54,6 +54,3 @@
  'pq_func_par3[1]': 4.,
  'pq_func_par4[1]': 0.38,
 }
-
-
-