Add context to plot_corner to allow the user to customise the plot

pyfstat.py

@@ -643,8 +643,10 @@ class MCMCSearch(BaseSearchClass):
         self.save_data(sampler, samples, lnprobs, lnlikes)
 
     def plot_corner(self, figsize=(7, 7),  tglitch_ratio=False,
-                    add_prior=False, nstds=None, label_offset=0.4, **kwargs):
+                    add_prior=False, nstds=None, label_offset=0.4,
+                    dpi=300, rc_context={}, **kwargs):
 
+        with plt.rc_context(rc_context):
             fig, axes = plt.subplots(self.ndim, self.ndim,
                                      figsize=figsize)
 
@@ -700,7 +702,7 @@ class MCMCSearch(BaseSearchClass):
                 self.add_prior_to_corner(axes, samples_plt)
 
             fig_triangle.savefig('{}/{}_corner.png'.format(
-            self.outdir, self.label))
+                self.outdir, self.label), dpi=dpi)
 
     def add_prior_to_corner(self, axes, samples):
         for i, key in enumerate(self.theta_keys):
@@ -757,6 +759,8 @@ class MCMCSearch(BaseSearchClass):
             return lambda x: logunif(x, kwargs['lower'], kwargs['upper'])
         elif kwargs['type'] == 'halfnorm':
             return lambda x: halfnorm(x, kwargs['loc'], kwargs['scale'])
+        elif kwargs['type'] == 'neghalfnorm':
+            return lambda x: halfnorm(-x, kwargs['loc'], kwargs['scale'])
         elif kwargs['type'] == 'norm':
             return lambda x: -0.5*((x - kwargs['loc'])**2/kwargs['scale']**2
                                    + np.log(2*np.pi*kwargs['scale']**2))
@@ -773,6 +777,9 @@ class MCMCSearch(BaseSearchClass):
         if dist_type == "halfnorm":
             return np.abs(np.random.normal(loc=kwargs['loc'],
                                            scale=kwargs['scale']))
+        if dist_type == "neghalfnorm":
+            return -1 * np.abs(np.random.normal(loc=kwargs['loc'],
+                                                scale=kwargs['scale']))
         if dist_type == "lognorm":
             return np.random.lognormal(
                 mean=kwargs['loc'], sigma=kwargs['scale'])