Cbd abdot

holodeck/accretion.py

@@ -10,7 +10,7 @@
 import os
 import numpy as np
 from scipy.interpolate import RectBivariateSpline
-from holodeck import _PATH_DATA
+from holodeck import _PATH_DATA, utils
 from holodeck.constants import SPLC, EDDT
 
 
@@ -37,7 +37,7 @@ class Accretion:
     """
 
     def __init__(self, accmod='Basic', f_edd=0.01, mdot_ext=None, eccen=0.0,
-                 subpc=True, **kwargs):
+                 subpc=True, edd_lim=None, evol_mass=True, **kwargs):
         """ Initializes the Accretion class.
 
         Parameters
@@ -51,6 +51,7 @@ def __init__(self, accmod='Basic', f_edd=0.01, mdot_ext=None, eccen=0.0,
         eccen : float, optional
             Eccentricity value.
         subpc : boolean, optional
+        edd_lim : None, or numerical value. This limits the accretion to a factor edd_lim the Eddington limit.
 
         Other Parameters
         ----------------
@@ -62,6 +63,32 @@ def __init__(self, accmod='Basic', f_edd=0.01, mdot_ext=None, eccen=0.0,
         self.mdot_ext = mdot_ext
         self.eccen = eccen
         self.subpc = subpc
+        self.edd_lim = edd_lim
+        self.evol_mass = evol_mass
+
+        """ PRE-CALCULATE PREFERENTIAL ACCRETION FUNCTION """
+        if self.accmod == 'Siwek22':
+            """ interpolate to get lambda at [q,e] """
+            def lambda_qe_interp_2d(fp="data/preferential_accretion/siwek+22/", es=[0.0,0.2,0.4,0.6,0.8]):
+                all_lambdas = []
+                for e in es:
+                    fname = 'preferential_accretion/siwek+22/lambda_e=%.2f.txt' % e
+                    fname = os.path.join(_PATH_DATA, fname)
+                    lambda_e = np.loadtxt(fname)
+                    qs = lambda_e[:, 0]
+                    lambdas = lambda_e[:, 1]
+                    all_lambdas.append(lambdas)
+                # True values of q, e
+                x = qs
+                y = es
+                # Populate the true values of q, e in a meshgrid
+                X, Y = np.meshgrid(x, y)
+                Z = all_lambdas
+
+                return(RectBivariateSpline(np.array(x), np.array(y),
+                                           np.array(Z).T, kx=1, ky=1))
+
+            self.swk_acc = lambda_qe_interp_2d()
 
     def mdot_eddington(self, mass, eps=0.1):
         """Calculate the total accretion rate based on masses and a fraction of the Eddington limit.
@@ -92,34 +119,43 @@ def mdot_eddington(self, mass, eps=0.1):
         >>> print(acc.mdot_eddington(mass))
 
         """
+        from holodeck.constants import SPLC, EDDT
+
         # medd = (4.*np.pi*NWTG*MPRT)/(eps*SPLC*SIGMA_T) * self.mtot
         medd = self.f_edd * (EDDT/(eps * SPLC**2)) * mass
         return(medd)
 
-    def mdot_total(self, evol, step):
+    def mdot_total(self, evol, bin, step):
+        from holodeck.constants import PC
+
         if self.mdot_ext is not None:
             """ accretion rates have been supplied externally """
-            mdot = self.mdot_ext[:,step-1]
+            raise NotImplementedError("THIS HASNT BEEN UPDATED TO NEW TIME-BASED INTEGRATION!")
+            mdot = self.mdot_ext[:, step-1]
+
         else:
             """ Get accretion rates as a fraction (f_edd in self._acc) of the
                 Eddington limit from current BH masses """
-            total_bh_masses = np.sum(evol.mass[:, step-1, :], axis=1)
+            if bin == None: 
+                total_bh_masses = np.sum(evol.mass[:, step, :], axis=-1)
+            else:
+                total_bh_masses = np.sum(evol.mass[step, :])
             mdot = self.mdot_eddington(total_bh_masses)
 
         """ Calculate individual accretion rates """
         if self.subpc:
             """ Indices where separation is less than or equal to a parsec """
-            ind_sepa = evol.sepa[:, step] <= PC
-        else:
-            """ Indices where separation is less than or equal to 100 kilo-parsec """
-            ind_sepa = evol.sepa[:, step] <= 10**5 * PC
+            if bin == None:
+                #we are using old evol method, i.e. calculating mdot for all binaries at a given step
+                mdot[evol.sepa[:, step]>PC] = 0.0
+            else:
+                #new evol method, iterating over individual binaries, so mdot is just a float
+                if evol.sepa[step] > PC:
+                    mdot = 0.0
 
-        """ Set total accretion rates to 0 when separation is larger than 1pc or 10kpc,
-            depending on subpc switch applied to accretion instance """
-        mdot[~ind_sepa] = 0
-        return(mdot)
+        return mdot
 
-    def pref_acc(self, mdot, evol, step):
+    def pref_acc(self, mdot, evol, bin, step):
         """
         Contains a variety of accretion models to choose from to calculate primary vs secondary
         accretion rates.
@@ -156,104 +192,101 @@ def pref_acc(self, mdot, evol, step):
         accessed.
 
         """
-        m1 = evol.mass[:, step - 1, 0]
-        m2 = evol.mass[:, step - 1, 1]
+
+        if bin == None:
+            m1, m2 = utils.m1m2_ordered(evol.mass[:,step].T[0], evol.mass[:,step].T[1])
+        else:
+            m1, m2 = utils.m1m2_ordered(*evol.mass[step])
+
 
         if self.accmod == 'Siwek22':
-            # Calculate the mass ratio
-            q_b = m2 / m1
-            # secondary and primary may swap indices
-            # need to account for that and reverse the mass ratio
-            inds_rev = q_b > 1
-            q_b[inds_rev] = 1./q_b[inds_rev]
             """if evol has an eccentricity distribution,
                we use it, if not, we set each eccentricity to
                the value specified in __init__ """
-            if evol.eccen is not None:
-                e_b = evol.eccen[:, step-1]
+            if bin == None:
+                e_b = evol.eccen[:, step] if (evol.eccen is not None) else None
             else:
-                e_b = self.eccen
-            """ Now interpolate to get lambda at [q,e] """
-            def lambda_qe_interp_2d(fp="data/preferential_accretion/siwek+22/", es=[0.0,0.2,0.4,0.6,0.8]):
-                all_lambdas = []
-                for e in es:
-                    fname = 'preferential_accretion/siwek+22/lambda_e=%.2f.txt' % e
-                    fname = os.path.join(_PATH_DATA, fname)
-                    lambda_e = np.loadtxt(fname)
-                    qs = lambda_e[:, 0]
-                    lambdas = lambda_e[:, 1]
-                    all_lambdas.append(lambdas)
-                # True values of q, e
-                x = qs
-                y = es
-                # Populate the true values of q, e in a meshgrid
-                X, Y = np.meshgrid(x, y)
-                Z = all_lambdas
+                e_b = evol.eccen[step] if (evol.eccen is not None) else None
 
-                return(RectBivariateSpline(np.array(x), np.array(y),
-                                           np.array(Z).T, kx=1, ky=1))
-
-            lamb_interp = lambda_qe_interp_2d()
+            lamb_interp = self.swk_acc
             # Need to use RectBivariateSpline.ev to evaluate the interpolation
             # at points, allowing q_b and e_b to be in non-ascending order
+            q_b = m2 / m1
             lamb_qe = lamb_interp.ev(q_b, e_b)
             mdot_1 = 1./(np.array(lamb_qe) + 1.) * mdot
             mdot_2 = np.array(lamb_qe)/(np.array(lamb_qe) + 1.) * mdot
 
-            # After calculating the primary and secondary accretion rates,
-            # they need to be placed at the correct index into `mdot_arr`, to
-            # account for primary/secondary being at 0-th OR 1-st index
-            mdot_arr = np.zeros(np.shape(evol.mass[:, step-1, :]))
-            inds_m1_primary = m1 >= m2  # where first mass is actually primary
-            inds_m2_primary = m2 >= m1  # where second mass is actually primary
-            mdot_arr[:, 0][inds_m1_primary] = mdot_1[inds_m1_primary]
-            mdot_arr[:, 0][~inds_m1_primary] = mdot_2[~inds_m1_primary]
-            mdot_arr[:, 1][inds_m2_primary] = mdot_1[inds_m2_primary]
-            mdot_arr[:, 1][~inds_m2_primary] = mdot_2[~inds_m2_primary]
-            # `mdot_arr` is then passed to `_take_next_step()` in evolution.py
-            return(mdot_arr)
-
         if self.accmod == 'Basic':
             mdot_1 = mdot_2 = 0.5 * mdot
             mdot_arr = np.array([mdot_1, mdot_2]).T
-            return(mdot_arr)
 
         if self.accmod == 'Proportional':
-            """ Get primary and secondary masses """
-            m1 = self.m1
-            m2 = self.m2
             # Calculate ratio of accretion rates so that mass ratio stays
             # constant through accretion
             mdot_ratio = m2/(m1 + m2)
             mdot_2 = mdot_ratio * mdot
             mdot_1 = (1. - mdot_ratio) * mdot
-            mdot_arr = np.array([mdot_1, mdot_2]).T
-            return(mdot_arr)
 
         if self.accmod == 'Primary':
             mdot_ratio = 0.
             mdot_2 = mdot_ratio * mdot
             mdot_1 = (1. - mdot_ratio) * mdot
-            mdot_arr = np.array([mdot_1, mdot_2]).T
-            return(mdot_arr)
 
         if self.accmod == 'Secondary':
             mdot_ratio = 1.
             mdot_2 = mdot_ratio * mdot
             mdot_1 = (1. - mdot_ratio) * mdot
-            mdot_arr = np.array([mdot_1, mdot_2]).T
-            return(mdot_arr)
 
         if self.accmod == 'Duffell':
             # Taken from Paul's paper: http://arxiv.org/abs/1911.05506
             def f(q): return 1./(0.1 + 0.9 * q)
             q = m2/m1
             mdot_1 = mdot/(1. + f(q))
             mdot_2 = f(q) * mdot_1
-            mdot_arr = np.array([mdot_1, mdot_2]).T
-            return(mdot_arr)
 
-        # Catch any weirdness if no model is selected.
+        if self.edd_lim is not None:
+            #need to limit the accretion rate to edd_lim times the Eddington limit
+            medd_1 = self.edd_lim * (self.mdot_eddington(m1)/self.f_edd)
+            mdot_1 = np.minimum(np.array(medd_1), np.array(mdot_1))
+            medd_2 = self.edd_lim * (self.mdot_eddington(m2)/self.f_edd)
+            mdot_2 = np.minimum(np.array(medd_2), np.array(mdot_2))
+
+        # After calculating the primary and secondary accretion rates,
+        # they need to be placed at the correct index into `mdot_arr`, to
+        # account for primary/secondary being at 0-th OR 1-st index
+        if bin == None:
+            inds_m1_primary = evol.mass[:,step].T[0] >= evol.mass[:,step].T[1]
+            mdot_arr = np.zeros(np.shape(evol.mass[:, step-1, :]))
+            mdot_arr[:, 0][inds_m1_primary] = mdot_1[inds_m1_primary]
+            mdot_arr[:, 0][~inds_m1_primary] = mdot_2[~inds_m1_primary]
+            inds_m2_primary = evol.mass[:,step-1].T[1] >= evol.mass[:,step-1].T[0]
+            mdot_arr[:, 1][inds_m2_primary] = mdot_1[inds_m2_primary]
+            mdot_arr[:, 1][~inds_m2_primary] = mdot_2[~inds_m2_primary]
+        else:
+            if evol.mass[step, 0] >= evol.mass[step, 1]:
+                mdot_arr = [mdot_1, mdot_2]
+            else:
+                mdot_arr = [mdot_2, mdot_1]
+
+        # Catch any case where no model is selected.
         if self.accmod is None:
             raise TypeError("'None' value provided for accretion model." +
-                            "An accretion model is required.")
+
+        return np.asarray(mdot_arr)
+
+    def ebeq(self, qb):
+        import pickle as pkl
+        fp_ebeq_pkl = 'cbd_torques/siwek+23/eb_eq_arr.pkl'
+        fp_ebeq_pkl = os.path.join(_PATH_DATA, fp_ebeq_pkl)
+        fp_ebeq = open(fp_ebeq_pkl, 'rb')
+        ebeq_dict = pkl.load(fp_ebeq)
+
+        all_qb = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
+
+        all_ebeq = []
+
+        for q in all_qb: 
+            all_ebeq.append(ebeq_dict['q=%.2f_sum_grav_acc' %q])
+
+        return(np.interp(qb, all_qb, all_ebeq))
+

holodeck/cyutils.pxd

@@ -2,4 +2,7 @@
 """
 
 cdef double _interp_between_vals(double xnew, double xl, double xr, double yl, double yr)
-cdef double interp_at_index(int idx, double xnew, double[:] xold, double[:] yold)
+cdef double interp_at_index(int idx, double xnew, double[:] xold, double[:] yold)
+cdef double gw_freq_dist_func__scalar_scalar(int nn, double ee)
+
+

holodeck/discrete/evolution.py

@@ -777,7 +777,8 @@ def _init_step_zero(self):
                 setattr(self, f"_dedt_{ii}", np.zeros_like(self.dadt))
 
         # ---- Initialize hardening rate at first step
-        dadt_init, dedt_init = self._hardening_rate(step=0)
+        step0 = 0
+        dadt_init, dedt_init = self._hardening_rate(step0)
 
         self.dadt[:, 0] = dadt_init
         if (pop.eccen is not None):

holodeck/discrete/population.py

@@ -302,7 +302,7 @@ class Pop_Illustris(_Population_Discrete):
 
     """
 
-    def __init__(self, fname=None, **kwargs):
+    def __init__(self, select=None, fname=None, **kwargs):
         """Initialize a binary population using data in the given filename.
 
         Parameters
@@ -314,6 +314,8 @@ def __init__(self, fname=None, **kwargs):
             Additional keyword-arguments passed to `super().__init__`.
 
         """
+        self._select = select
+
         if fname is None:
             fname = _DEF_ILLUSTRIS_FNAME
             fname = os.path.join(_PATH_DATA, fname)
@@ -347,6 +349,15 @@ def _init(self):
         # Get the stellar mass, and take that as bulge mass
         self.mbulge = data['SubhaloMassInRadType'][:, st_idx, :]   #: Stellar mass / stellar-bulge mass [grams]
         self.vdisp = data['SubhaloVelDisp']    #: Velocity dispersion of galaxy [?cm/s?]
+
+        if self._select is not None:
+            print(f"\n=====\nWARNING SELECTING ONLY {self._select} BINARIES FROM DATA\n=====\n")
+            names = ['sepa', 'mass', 'scafa', 'mbulge', 'vdisp']
+            for nam in names:
+                vals = getattr(self, nam)
+                vals = vals[:self._select]
+                setattr(self, nam, vals)
+
         return