Merge branch 'main' of github.com:aymgal/LensModelAPI

coolest/api/plotting.py

@@ -1,7 +1,9 @@
 __author__ = 'XXX'
 
 
-from coolest.io.analysis import Analysis
+from coolest.api.analysis import Analysis
+from coolest.api.util import read_json_param
+import matplotlib.pyplot as plt
 
 
 class Plotting(object):
@@ -15,3 +17,92 @@ def __init__(self, coolest_file_path):
 
     def plot_summary(self):
         pass
+
+class Comparison_analytical(object):
+    """
+    Handles plot of analytical models in a comparative way
+    """
+    def __init__(self,coolest_file_list, nickname_file_list, posterior_bool_list):
+        self.file_names = nickname_file_list
+        self.posterior_bool_list = posterior_bool_list
+        self.param_lens, self.param_source = read_json_param(coolest_file_list,self.file_names, lens_light=False)
+
+    def plotting_routine(self,param_dict,idx_file=0):
+        """
+        plot the parameters
+
+        INPUT
+        -----
+        param_dict: dict, organized dictonnary with all parameters results of the different files
+        idx_file: int, chooses the file on which the choice of plotted parameters will be made
+        (not very clear: basically in file 0 you may have a sersic fit and in file 1 sersic+shapelets. If you choose
+         idx_file=0, you will plot the sersic results of both file. If you choose idx_file=1, you will plot all the
+         sersic and shapelets parameters when available)
+        """
+
+        #find the numer of parameters to plot and define a nice looking figure
+        number_param = len(param_dict[self.file_names[idx_file]])
+        unused_figs = []
+        if number_param <= 4:
+            print('so few parameters not implemented yet')
+        else:
+            if number_param % 4 == 0:
+                num_lines = int(number_param / 4.)
+            else:
+                num_lines = int(number_param / 4.) + 1
+
+                for idx in range(3):
+                    if (number_param + idx) % 4 != 0:
+                        unused_figs.append(-idx - 1)
+                    else:
+                        break
+
+        f, ax = plt.subplots(num_lines, 4, figsize=(4 * 3.5, 2.5 * num_lines))
+        markers = ['*', '.', 's', '^','<','>','v','p','P','X','D','1','2','3','4','+']
+        #may find a better way to define markers but right now, it is sufficient
+
+        for j, file_name in enumerate(self.file_names):
+            i = 0
+            result = param_dict[file_name]
+            for key in result.keys():
+                idx_line = int(i / 4.)
+                idx_col = i % 4
+                p = result[key]
+                m = markers[j]
+                if self.posterior_bool_list[j]:
+                    # UNCOMMENT IF NO ERROR BARS AVAILABLE ON SHEAR
+                    #             if (j== 1) and (key=='SHEAR_0_gamma_ext' or key == 'SHEAR_0_phi_ext'):
+                    #                 ax[idx_line,idx_col].plot(j,p['point_estimate'],marker=m,ls='',label=file_name)
+                    #                 i+=1
+                    #                 continue
+
+                    #trick to plot correct error bars if close to the +180/-180 edge
+                    if (key == 'SHEAR_0_phi_ext' or key == 'PEMD_0_phi'):
+                        if p['percentile_16th'] > p['median']:
+                            p['percentile_16th'] -= 180.
+                        if p['percentile_84th'] < p['median']:
+                            p['percentile_84th'] += 180.
+                    ax[idx_line, idx_col].errorbar(j, p['median'], [[p['median'] - p['percentile_16th']],
+                                                                    [p['percentile_84th'] - p['median']]],
+                                                   marker=m, ls='', label=file_name)
+                else:
+                    ax[idx_line, idx_col].plot(j, p['point_estimate'], marker=m, ls='', label=file_name)
+
+                if j == 0:
+                    ax[idx_line, idx_col].get_xaxis().set_visible(False)
+                    ax[idx_line, idx_col].set_ylabel(p['latex_str'], fontsize=12)
+                    ax[idx_line, idx_col].tick_params(axis='y', labelsize=12)
+                i += 1
+
+        ax[0, 0].legend()
+        for idx in unused_figs:
+            ax[-1, idx].axis('off')
+        plt.tight_layout()
+        plt.show()
+        return f,ax
+    def plot_source(self,idx_file=0):
+        f,ax = self.plotting_routine(self.param_source,idx_file)
+        return f,ax
+    def plot_lens(self,idx_file=0):
+        f,ax = self.plotting_routine(self.param_lens,idx_file)
+        return f,ax

coolest/api/profiles/light.py

@@ -0,0 +1,16 @@
+__author__ = 'lynevdv'
+
+
+import numpy as np
+from coolest.api.profiles import util
+
+
+class Sersic(object):
+
+    """Elliptical Sersic"""
+
+    def surface_brightness(self, x, y, I_eff=1., theta_eff=2., n=4., phi=0., q=1., center_x=0., center_y=0.):
+        """Returns the surface brightness at the given position (x, y)"""
+        x_t, y_t = util.shift_rotate_elliptical(x, y, phi, q, center_x, center_y)
+        bn = 1.9992*n - 0.3271
+        return I_eff * np.exp( - bn * ( (np.sqrt(x_t**2+y_t**2) / theta_eff )**(1./n) -1. ) )