Merge pull request #33 from mirochaj/dm_extension_branch

a few additions to the LWB solver to make debugging and testing easier

ares/populations/GalaxyCohort.py

@@ -1482,7 +1482,8 @@ def _tab_Mmax(self):
                 e_limit = None
 
             if (t_limit is not None) or (m_limit is not None) or \
-               (e_limit is not None) or (T_limit is not None) or (a_limit is not None):
+               (e_limit is not None) or (T_limit is not None) or \
+               (a_limit is not None):
 
                 M0x = self.pf['pop_initial_Mh']
                 if (M0x == 0) or (M0x == 1):

ares/simulations/MetaGalacticBackground.py

@@ -15,12 +15,12 @@
 import scipy
 import numpy as np
 from ..static import Grid
-from ..util.Math import smooth
 from ..util.Pickling import write_pickle_file
 from types import FunctionType
 from ..util import ParameterFile
 from ..phenom import Madau1995
-from ..util.Math import interp1d
+from ..util.Misc import split_by_sign
+from ..util.Math import interp1d, smooth
 from ..solvers import UniformBackground
 from ..analysis.MetaGalacticBackground import MetaGalacticBackground \
     as AnalyzeMGB
@@ -895,11 +895,11 @@ def _is_Mmin_converged(self, include_pops):
 
         # Instance of a population that "feels" the feedback.
         # Need for (1) initial _Mmin_pre value, and (2) setting ceiling
-        pop_fb = self.pops[self._lwb_sources[0]]
+        pid = self.pf['feedback_LW_sfrd_popid']
+        pop_fb = self.pops[self._lwb_sources.index(pid)]
 
         # Don't re-load Mmin guesses after first iteration
         if self.pf['feedback_LW_guesses'] is not None and self.count > 1:
-            pid = self.pf['feedback_LW_sfrd_popid']
             self.pops[pid]._loaded_guesses = True
             print('turning off ModelSet load', self.count, pid, self.pops[pid]._loaded_guesses)
 
@@ -909,15 +909,18 @@ def _is_Mmin_converged(self, include_pops):
             has_guess = False
             if self.pf['feedback_LW_guesses'] is not None:
                 has_guess = True
-                pid = self.pf['feedback_LW_sfrd_popid']
+
                 #_z_guess, _Mmin_guess = guess
                 self._Mmin_pre = self.pops[pid].Mmin(zarr)
+                self._Mmax_pre = self.pops[pid]._tab_Mmax
 
             else:
                 self._Mmin_pre = np.min([self.pops[idnum].Mmin(zarr) \
                     for idnum in self._lwb_sources], axis=0)
+                self._Mmax_pre = self.pops[pid]._tab_Mmax
 
             self._Mmin_bank = [self._Mmin_pre.copy()]
+            self._Mmax_bank = [self._Mmax_pre.copy()]
             self._Jlw_bank = [Jlw]
 
             self.pf['feedback_LW_guesses'] = None
@@ -1062,7 +1065,7 @@ def _is_Mmin_converged(self, include_pops):
            (self.count - mdel) % mfreq == 0:
             _Mmin_next = np.sqrt(np.product(self._Mmin_bank[-2:], axis=0))
         elif (self.count > 1) and (self.pf['feedback_LW_softening'] is not None):
-            if self.pf['feedback_LW_softening'] == 'sqrt':
+            if self.pf['feedback_LW_softening'] in ['sqrt', 'gmean']:
                 _Mmin_next = np.sqrt(Mnext * self._Mmin_pre)
             elif self.pf['feedback_LW_softening'] == 'mean':
                 _Mmin_next = np.mean([Mnext, self._Mmin_pre], axis=0)
@@ -1074,8 +1077,61 @@ def _is_Mmin_converged(self, include_pops):
         else:
             _Mmin_next = Mnext
 
+        ##
+        # Dealing with Mmin fluctuations
+        if self.pf['feedback_LW_Mmin_monotonic'] and \
+           (self.count % self.pf['feedback_LW_Mmin_afreq'] == 0):
+
+            # Detect wiggles: expected behaviour is dMmin/dz < 0.
+            # Zero pad at the end to recover correct length.
+            dMmindz = np.concatenate((np.diff(_Mmin_next), [0]))
+
+            # Break into positive and negative chunks.
+            _x, _y = split_by_sign(zarr, dMmindz)
+
+            # Go in order of ascending time (descending z)
+            x = _x[-1::-1]
+            y = _y[-1::-1]
+
+            j = 0 # start index of a given chunk within full array of Mmin
+            zrev = zarr[-1::-1]
+            Mmin_new = _Mmin_next[-1::-1]
+            for i, _chunk_ in enumerate(y):
+                l = len(_chunk_)
+
+                if (i == 0) or (_chunk_[0] <= 0):
+                    j += l
+                    continue
+
+                # If we're here, it means Mmin is falling with redshift.
+
+                k = j + l
+
+                # Just replace with Mmin value before it last declined.
+                # Guaranteed to be OK since we iterate from high-z to low.
+                if self.pf['feedback_LW_Mmin_monotonic'] == 1:
+                    Mmin_new[j:k] = Mmin_new[j-1]
+                else:
+                    raise NotImplemented('help')
+                    # Interpolate to next point where Mmin > Mmin_problem
+                    dx = zrev[j-1] - zrev[j-2]
+                    dy = Mmin_new[j-1] - Mmin_new[j-2]
+                    m = 2 * dy / dx
+
+                    print('hey', dx, dy, m)
+                    Mmin_guess = m * (zrev[k-1] - zrev[j-1]) + Mmin_new[j-1]
+                    #while Mmin_guess > Mmin_new[k]:
+                    #    m *= 0.9
+                    #    Mmin_guess = m * (zrev[k] - zrev[j-1]) + Mmin_new[j-1]
+
+                    Mmin_new[j:k] = Mmin_guess#m * (zrev[j:k] - zrev[j-1]) + Mmin_new[j-1]
+
+                j += l
+
+            _Mmin_next = Mmin_new[-1::-1]
+
         # Detect ripples first and only do this if we see some?
-        if (self.pf['feedback_LW_Mmin_smooth'] > 0) and \
+        elif (self.pf['feedback_LW_Mmin_smooth'] > 0) and \
            (self.count % self.pf['feedback_LW_Mmin_afreq'] == 0):
 
             s = self.pf['feedback_LW_Mmin_smooth']
@@ -1089,7 +1145,7 @@ def _is_Mmin_converged(self, include_pops):
 
             _Mmin_next = 10**np.interp(zarr, ztmp, Ms)
 
-        if (self.pf['feedback_LW_Mmin_fit'] > 0) and \
+        elif (self.pf['feedback_LW_Mmin_fit'] > 0) and \
            (self.count % self.pf['feedback_LW_Mmin_afreq'] == 0):
             order = self.pf['feedback_LW_Mmin_fit']
             _Mmin_next = 10**np.polyval(np.polyfit(zarr, np.log10(_Mmin_next), order), zarr)
@@ -1111,6 +1167,9 @@ def _is_Mmin_converged(self, include_pops):
         self._Mmin_now = Mmin.copy()
         # Save for our records (read: debugging)
 
+        # Save Mmax too
+        self._Mmax_now = self.pops[pid]._tab_Mmax.copy()
+
         ##
         # Compare Mmin of last two iterations.
         ##
@@ -1191,6 +1250,7 @@ def _is_Mmin_converged(self, include_pops):
 
         if not converged:
             self._Mmin_bank.append(self._Mmin_now.copy())
+            self._Mmax_bank.append(self._Mmax_now.copy())
             self._Jlw_bank.append(Jlw)
 
         return converged

ares/util/SetDefaultParameterValues.py

@@ -259,6 +259,7 @@ def PhysicsParameters():
     'feedback_LW_softening': 'sqrt',
     'feedback_LW_tol_zrange': (0, np.inf),
 
+    'feedback_LW_Mmin_monotonic': False,
     'feedback_LW_Mmin_smooth': 0,
     'feedback_LW_Mmin_fit': 0,
     'feedback_LW_Mmin_afreq': 0,
@@ -1287,7 +1288,7 @@ def HaloMassFunctionParameters():
 
     # If a new tab_MAR should be computed when using the PCA
     "hmf_gen_MAR":False,
-    
+
     "filter_params" : None,
 
     "hmf_MAR_from_CDM": True,