From d33e4d11202e27563cbfde156c5d0ca27de95c43 Mon Sep 17 00:00:00 2001
From: Takahiro Morishita <takahiro@ipac.caltech.edu>
Date: Fri, 26 Jan 2024 17:04:45 -0800
Subject: [PATCH] shuffle func; normal prior for AV

---
 gsf/fitting.py            | 22 ++++++++++++++------
 gsf/posterior_flexible.py | 44 +++++++++++++++++++++++++++++++++++----
 2 files changed, 56 insertions(+), 10 deletions(-)

diff --git a/gsf/fitting.py b/gsf/fitting.py
index 0781adc..25161e9 100644
--- a/gsf/fitting.py
+++ b/gsf/fitting.py
@@ -1798,8 +1798,12 @@ def prepare_class(self, add_fir=None):
 
     def get_shuffle(self, out, nshuf=3.0, amp=1e-4):
         '''
+        amp : amplitude, 0 to 1.
+        
         Shuffles initial parameters of each walker, to give it extra randomeness.
         '''
+        if amp>1:
+            amp = 1
         pos = np.zeros((self.nwalk, self.ndim), 'float')
         for ii in range(pos.shape[0]):
             aa = 0
@@ -1807,14 +1811,14 @@ def get_shuffle(self, out, nshuf=3.0, amp=1e-4):
                 if out.params[key].vary == True:
                     pos[ii,aa] += out.params[key].value
                     # This is critical to avoid parameters fall on the boundary.
-                    delpar = (out.params[key].max-out.params[key].min)/1000
+                    delpar = (out.params[key].max-out.params[key].min) * amp/2.
                     # or not,
-                    delpar = 0
+                    # delpar = 0
                     if np.random.uniform(0,1) > (1. - 1./self.ndim):        
-                        pos[ii,aa] = np.random.uniform(out.params[key].min+delpar, out.params[key].max-delpar)
+                        pos[ii,aa] = np.random.uniform(out.params[key].value-delpar, out.params[key].value+delpar)
                     else:
-                        if pos[ii,aa]<out.params[key].min+delpar or pos[ii,aa]>out.params[key].max-delpar:
-                            pos[ii,aa] = np.random.uniform(out.params[key].min+delpar, out.params[key].max-delpar)
+                        if pos[ii,aa]<out.params[key].min or pos[ii,aa]>out.params[key].max:
+                            pos[ii,aa] = np.random.uniform(out.params[key].value-delpar, out.params[key].value+delpar)
 
                     aa += 1
         return pos
@@ -1965,10 +1969,16 @@ def main(self, cornerplot:bool=True, specplot=1, sigz=1.0, ezmin=0.01, ferr=0,
                     pos = self.get_shuffle(out, amp=amp_shuffle)
                 else:
                     pos = amp_shuffle * np.random.randn(self.nwalk, self.ndim)
+                    pos += self.get_shuffle(out, amp=0)
+                    # Check boundary;
                     aa = 0
                     for aatmp,key in enumerate(out.params.valuesdict()):
                         if out.params[key].vary:
-                            pos[:,aa] += out.params[key].value
+                            con = (out.params[key].min > pos[:,aa])
+                            pos[:,aa][con] = out.params[key].min
+                            con = (out.params[key].max < pos[:,aa])
+                            pos[:,aa][con] = out.params[key].max
+                            # pos[:,aa] = out.params[key].value
                             aa += 1
 
                 if self.f_zeus:
diff --git a/gsf/posterior_flexible.py b/gsf/posterior_flexible.py
index ce9e8c5..cc577b1 100644
--- a/gsf/posterior_flexible.py
+++ b/gsf/posterior_flexible.py
@@ -6,7 +6,7 @@
 from numpy import exp as np_exp
 from numpy import log as np_log
 from scipy.special import erf
-from scipy.stats import lognorm
+from scipy.stats import lognorm,norm
 
 class Post():
     '''
@@ -290,8 +290,12 @@ def lnprob_emcee(self, pos, pars, fy:float, ey:float, wht:float, NR:float, f_fir
 
         # lognormal-prior for any params;
         for ii,key_param in enumerate(self.mb.key_params_prior):
-            sigma = self.mb.key_params_prior_sigma[ii]
-            respr += self.get_lognormal_prior(vals, key_param, sigma=sigma, mu=0)
+            if key_param[:2] == 'AV':
+                sigma = self.mb.key_params_prior_sigma[ii]
+                respr += self.get_normal_prior(vals, key_param, sigma=sigma, mu=0)
+            else:
+                sigma = self.mb.key_params_prior_sigma[ii]
+                respr += self.get_lognormal_prior(vals, key_param, sigma=sigma, mu=0)
 
         # Prior for emission line template??;
         # Still in experiment;
@@ -366,4 +370,36 @@ def get_lognormal_prior(self, vals, key_param, mu=0, sigma=100.0, check_prior=Fa
                 plt.plot(yy, np.log(self.mb.prior[key_param].pdf(yy)))
                 plt.show()
 
-        return np.log(self.mb.prior[key_param].pdf(y))
\ No newline at end of file
+        respr = np.log(self.mb.prior[key_param].pdf(y))
+
+        if not np.isfinite(respr):
+            respr = -1e10
+
+        return respr
+    
+    def get_normal_prior(self, vals, key_param, mu=0, sigma=100.0, check_prior=False):
+        '''
+        '''
+        y = vals[key_param]
+
+        if self.mb.prior == None:
+            self.mb.prior = {}
+            for key_param_tmp in self.mb.fit_params:
+                self.mb.prior[key_param_tmp] = None
+
+        if self.mb.prior[key_param] == None:
+            self.mb.logger.info('Using normal prior for %s'%key_param)
+            self.mb.prior[key_param] = norm()
+            if check_prior:
+                import matplotlib.pyplot as plt
+                plt.close()
+                yy = np.arange(-2,2,0.1)
+                plt.plot(yy, np.log(self.mb.prior[key_param].pdf((yy-mu) * np.sqrt(2) / sigma)))
+                plt.show()
+
+        respr = np.log(self.mb.prior[key_param].pdf((y-mu) * np.sqrt(2) / sigma))
+
+        if not np.isfinite(respr):
+            respr = -1e10
+
+        return respr
\ No newline at end of file