18846156192

Committed 27 Oct 2025 03:18PM UTC coverage: 45.557% (-0.2%) from 45.806%

Build # 18846156192

Build Type

push

github

Committed by

web-flow

Commit Message

Merge pull request #71 from aangress/normalization-feature

Option to scale images to same range across multiple models

Coverage Stats

0 of 17 new or added lines in 1 file covered. (0.0%)

1420 of 3117 relevant lines covered (45.56%)

0.46 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

0.0

/coolest/api/plot_util.py

__author__ = 'aymgal', 'Giorgos Vernardos'


import numpy as np
import warnings
import matplotlib
import matplotlib.pyplot as plt
from matplotlib import ticker
from mpl_toolkits.axes_grid1 import make_axes_locatable
from scipy.spatial import Voronoi, voronoi_plot_2d
from matplotlib.colors import Normalize, LogNorm, TwoSlopeNorm
from matplotlib.cm import ScalarMappable
from mpl_toolkits.axes_grid1.anchored_artists import AnchoredSizeBar
from coolest.api.composable_models import ComposableLensModel
import os
import tarfile
import tempfile
import io
from coolest.api import util
from coolest.api.analysis import Analysis
from coolest.api.plotting import ModelPlotter, ParametersPlotter


def plot_voronoi(ax, x, y, z, neg_values_as_bad=False, 
                 norm=None, cmap=None, zmin=None, zmax=None, 
                 edgecolor=None, zorder=1):

    if cmap is None:
        cmap = plt.get_cmap('inferno')
    if norm is None:
        if zmin is None:
            zmin = np.min(z)
        if zmax is None:
            zmax = np.max(z)
        norm = Normalize(zmin, zmax)

    # get voronoi regions
    voronoi_points = np.column_stack((x,y))
    vor = Voronoi(voronoi_points)
    new_regions, vertices = voronoi_finite_polygons_2d(vor)
    
    # get cell colors
    m = matplotlib.cm.ScalarMappable(norm=norm, cmap=cmap)

    # plot voronoi points
    #point_colors = [ m.to_rgba(v) for v in z ]
    #ax.scatter(voronoi_points[:,0],voronoi_points[:,1],c=point_colors)

    # plot voronoi cells
    for i, region in enumerate(new_regions):
        polygon = vertices[region]
        z_i = z[i]
        if neg_values_as_bad is True and z_i < 0.:
            cell_color = m.to_rgba(np.nan)
        else:
            cell_color = m.to_rgba(z_i)
        ax.fill(*zip(*polygon), facecolor=cell_color, edgecolor=edgecolor, zorder=zorder)
    return m


def voronoi_finite_polygons_2d(vor,radius=None):
    """
    Reconstruct infinite voronoi regions in a 2D diagram to finite
    regions.
    This function is taken from: https://gist.github.com/pv/8036995

    Parameters
    ----------
    vor : Voronoi
        Input diagram
    radius : float, optional
        Distance to 'points at infinity'.

    Returns
    -------
    regions : list of tuples
        Indices of vertices in each revised Voronoi regions.
    vertices : list of tuples
        Coordinates for revised Voronoi vertices. Same as coordinates
        of input vertices, with 'points at infinity' appended to the
        end.
    """

    if vor.points.shape[1] != 2:
        raise ValueError("Requires 2D input")

    new_regions = []
    new_vertices = vor.vertices.tolist()

    center = vor.points.mean(axis=0)
    if radius is None:
        radius = vor.points.ptp().max()

    # Construct a map containing all ridges for a given point
    all_ridges = {}
    for (p1, p2), (v1, v2) in zip(vor.ridge_points, vor.ridge_vertices):
        all_ridges.setdefault(p1, []).append((p2, v1, v2))
        all_ridges.setdefault(p2, []).append((p1, v1, v2))

    # Reconstruct infinite regions
    for p1, region in enumerate(vor.point_region):
        vertices = vor.regions[region]

        if all(v >= 0 for v in vertices):
            # finite region
            new_regions.append(vertices)
            continue

        # reconstruct a non-finite region
        ridges = all_ridges[p1]
        new_region = [v for v in vertices if v >= 0]

        for p2, v1, v2 in ridges:
            if v2 < 0:
                v1, v2 = v2, v1
            if v1 >= 0:
                # finite ridge: already in the region
                continue

            # Compute the missing endpoint of an infinite ridge

            t = vor.points[p2] - vor.points[p1] # tangent
            t /= np.linalg.norm(t)
            n = np.array([-t[1], t[0]])  # normal

            midpoint = vor.points[[p1, p2]].mean(axis=0)
            direction = np.sign(np.dot(midpoint - center, n)) * n
            far_point = vor.vertices[v2] + direction * radius

            new_region.append(len(new_vertices))
            new_vertices.append(far_point.tolist())

        # sort region counterclockwise
        vs = np.asarray([new_vertices[v] for v in new_region])
        c = vs.mean(axis=0)
        angles = np.arctan2(vs[:,1] - c[1], vs[:,0] - c[0])
        new_region = np.array(new_region)[np.argsort(angles)]

        # finish
        new_regions.append(new_region.tolist())
        
    return new_regions, np.asarray(new_vertices)


def std_colorbar(mappable, label=None, fontsize=12, label_kwargs={}, **colorbar_kwargs):
    cb = plt.colorbar(mappable, **colorbar_kwargs)
    if label is not None:
        colorbar_kwargs.pop('label', None)
        cb.set_label(label, fontsize=fontsize, **label_kwargs)
    return cb

def std_colorbar_residuals(mappable, res_map, vmin, vmax, label=None, fontsize=12, 
                           label_kwargs={}, **colorbar_kwargs):
    if res_map.min() < vmin and res_map.max() > vmax:
        cb_extend = 'both'
    elif res_map.min() < vmin:
        cb_extend = 'min'
    elif res_map.max() > vmax:
        cb_extend = 'max'
    else:
        cb_extend = 'neither'
    colorbar_kwargs.update({'extend': cb_extend})
    return std_colorbar(mappable, label=label, fontsize=fontsize, 
                        label_kwargs=label_kwargs, **colorbar_kwargs)

def nice_colorbar(mappable, ax=None, position='right', pad=0.1, size='5%', label=None, fontsize=12, 
                  invisible=False, 
                  #max_nbins=None,
                  divider_kwargs={}, colorbar_kwargs={}, label_kwargs={}):
    divider_kwargs.update({'position': position, 'pad': pad, 'size': size})
    if ax is None:
        ax = mappable.axes
    divider = make_axes_locatable(ax)
    cax = divider.append_axes(**divider_kwargs)
    if invisible:
        cax.axis('off')
        return None
    cb = plt.colorbar(
        mappable, cax=cax, 
        **colorbar_kwargs
    )
    if label is not None:
        colorbar_kwargs.pop('label', None)
        cb.set_label(label, fontsize=fontsize, **label_kwargs)
    if position == 'top':
        cax.xaxis.set_ticks_position('top')
    # if max_nbins is not None: # TODO: this leads to strange results
    #     # cb.locator = ticker.LogLocator(subs=range(10))
    #     # cb.update_ticks()
    return cb

def nice_colorbar_residuals(mappable, res_map, vmin, vmax, ax=None, position='right', pad=0.1, size='5%', 
                            invisible=False, label=None, fontsize=12,
                            divider_kwargs={}, colorbar_kwargs={}, label_kwargs={}):
    if res_map.min() < vmin and res_map.max() > vmax:
        cb_extend = 'both'
    elif res_map.min() < vmin:
        cb_extend = 'min'
    elif res_map.max() > vmax:
        cb_extend = 'max'
    else:
        cb_extend = 'neither'
    colorbar_kwargs.update({'extend': cb_extend})
    return nice_colorbar(mappable, ax=ax, position=position, pad=pad, size=size, label=label, fontsize=fontsize,
                  invisible=invisible, colorbar_kwargs=colorbar_kwargs, label_kwargs=label_kwargs,
                  divider_kwargs=divider_kwargs)

def cax_colorbar(fig, cax, norm=None, cmap=None, mappable=None, label=None, fontsize=12, orientation='horizontal', label_kwargs={}):
    if mappable is None:
        mappable = ScalarMappable(norm=norm, cmap=cmap)
    cb = fig.colorbar(mappable=mappable, orientation=orientation, cax=cax)
    if label is not None:
        cb.set_label(label, fontsize=fontsize, **label_kwargs)

def scale_bar(ax, size, unit_suffix='"', loc='lower left', color='#FFFFFFBB', fontsize=12):
    if size == int(size):
        label = f"{int(size)}"
    else:
        label = f"{size:.1f}"
    label += unit_suffix
    artist = AnchoredSizeBar(
        ax.transData,
        size, label,
        loc=loc, label_top=True,
        pad=0.8, sep=5, 
        color=color, fontproperties=dict(size=fontsize),
        frameon=False, size_vertical=0,
    )
    ax.add_artist(artist)
    return artist

def plot_regular_grid(ax, title, image_, neg_values_as_bad=False, xylim=None, **imshow_kwargs):
    if neg_values_as_bad:
        image = np.copy(image_)
        image[image < 0] = np.nan
    else:
        image = image_
    im = ax.imshow(image, **imshow_kwargs)
    im.set_rasterized(True)
    set_xy_limits(ax, xylim)
    ax.xaxis.set_major_locator(plt.MaxNLocator(3))
    ax.yaxis.set_major_locator(plt.MaxNLocator(3))
    ax.set_title(title)
    return ax, im

def plot_irregular_grid(ax, title, points, xylim, neg_values_as_bad=False,
                            norm=None, cmap=None, plot_points=False):
    x, y, z = points
    im = plot_voronoi(ax, x, y, z, neg_values_as_bad=neg_values_as_bad, 
                      norm=norm, cmap=cmap, zorder=1)
    ax.set_aspect('equal', 'box')
    set_xy_limits(ax, xylim)
    ax.xaxis.set_major_locator(plt.MaxNLocator(3))
    ax.yaxis.set_major_locator(plt.MaxNLocator(3))
    if plot_points:
        ax.scatter(x, y, s=5, c='white', marker='.', alpha=0.4, zorder=2)
    ax.set_title(title)
    return ax, im

def set_xy_limits(ax, xylim):
    if xylim is None: return  # do nothing
    if isinstance(xylim, (int, float)):
        xylim_ = [-xylim, xylim, -xylim, xylim]
    elif isinstance(xylim, (list, tuple)) and len(xylim) == 2:
        xylim_ = [xylim[0], xylim[1], xylim[0], xylim[1]]
    elif not isinstance(xylim, (list, tuple)) and len(xylim) != 4:
        raise ValueError("`xylim` argument should be a single number, a 2-tuple or a 4-tuple.")
    else:
        xylim_ = xylim
    ax.set_xlim(xylim_[0], xylim_[1])
    ax.set_ylim(xylim_[2], xylim_[3])

def panel_label(ax, text, color, fontsize, alpha=0.8, loc='upper left'):
    if loc == 'upper left':
        x, y, ha, va = 0.03, 0.97, 'left', 'top'
    elif loc == 'lower left':
        x, y, ha, va = 0.03, 0.03, 'left', 'bottom'
    elif loc == 'upper right':
        x, y, ha, va = 0.97, 0.97, 'right', 'top'
    elif loc == 'lower right':
        x, y, ha, va = 0.97, 0.03, 'right', 'bottom'
    ax.text(x, y, text, color=color, fontsize=fontsize, alpha=alpha, 
            ha=ha, va=va, transform=ax.transAxes)

def normalize_across_images(plotter_list, data_model_specifier, kwargs_source = None, kwargs_lens_mass = None,
                            supersampling=5, convolved=True, super_convolution=True):
    """Calculate the vmin and vmax to normalize the colormap across multiple coolest objects

    Parameters
    ----------
    plotter_list: list
        List of ModelPlotter objects. May be acquired from MultiModelPlotter object
    data_model_specifier: list
        List of 0s and 1s; 0 = data, 1 = model. Specifies which set of pixel values should be used
        when finding global minima and maxima -- data or model
    kwargs_source: dict
        Dictionary with "entity_selection" key, same as used in MultiModelPlotters.
        "Entity_selection" contains list of lists. Selects source entities.
        Insert dummy None values into dictionary for data.
    kwargs_lens_mass: dict
        Dictionary with "entity_selection" key, same as used in MultiModelPlotters.
        "Entity_selection" contains list of lists. Selects lens mass entities.
        Insert dummy None values into dictionary for data.
    supersampling: int
        Model image generation param
    convolved: bool
        Model image generation param
    super_convolution: bool
        Model image generation param
    
    
    Returns
    -------
    vmin: float
        global min value across all coolest objects in plotter_list for the specified data/models
    vmax: float
        global max value across all coolest objects in plotter_list for the specified data/models    
    """
    
    mins = []
    maxes = []
    ks_arr = kwargs_source['entity_selection']
    km_arr = kwargs_lens_mass['entity_selection']
    for plotter, d_or_f, ks, km in zip(plotter_list, data_model_specifier, ks_arr, km_arr):
        # Check if we are finding extrema for data or model
        if d_or_f == 0:
            image = plotter.coolest.observation.pixels.get_pixels(directory=plotter._directory)
        elif d_or_f == 1:
            lens_model = ComposableLensModel(plotter.coolest, plotter._directory,
                                         kwargs_selection_source=dict(entity_selection=ks),
                                         kwargs_selection_lens_mass=dict(entity_selection=km))
            image, _ = lens_model.model_image(supersampling, convolved, super_convolution)
        # Find min and max and append
        mins.append(np.min(image))
        maxes.append(np.max(image))
    vmin = min(mins)
    vmax = max(maxes)
    return vmin, vmax

    
def dmr_corner(tar_path, output_dir = None):
    """Given .tar.gz COOLEST file, plots and optionally saves DMR and corner plots for COOLEST file. Returns dictionary of important extracted information.

    Parameters
    ----------
    tar_path : string
        Path to .tar.gz COOLEST file
    output_dir : string, optional
        Path to automatically save DMR and corner plot to if specified, by default None
    
    Returns
    -------
    results: dictionary
        Contains useful information about the COOLEST objects
    """
    
    results = {}
    with tempfile.TemporaryDirectory() as tmpdir:
        # Extract tar.gz archive
        with tarfile.open(tar_path, "r:gz") as tar:
            tar.extractall(path=tmpdir)

        
        # Get path to the extracted JSON file
        extracted_items = os.listdir(tmpdir)
        extracted_path = os.path.join(tmpdir, extracted_items[0])
        if os.path.isdir(extracted_path):
            extracted_files = os.listdir(extracted_path)
        else:
            extracted_files = extracted_items
            extracted_path = tmpdir  # fallback
        
        json_files = [f for f in extracted_files if f.endswith(".json")]
        if not json_files:
            raise ValueError("No .json file found in archive.")
        
        json_path = os.path.join(extracted_path, json_files[0])
        target_path = os.path.splitext(json_path)[0]

        # Load COOLEST object
        coolest_1 = util.get_coolest_object(target_path, verbose=False)

        # Run analysis
        analysis = Analysis(coolest_1, target_path, supersampling=5)

        coord_orig = util.get_coordinates(coolest_1)
        coord_src = coord_orig.create_new_coordinates(pixel_scale_factor=0.1, grid_shape=(1.42, 1.42))

        # Extract values
        r_eff_source = analysis.effective_radius_light(center=(0,0), coordinates=coord_src, outer_radius=1., entity_selection=[2])
        einstein_radius = analysis.effective_einstein_radius(entity_selection=[0,1])

        results['r_eff_source'] = r_eff_source
        results['einstein_radius'] = einstein_radius
        results['lensing_entities'] = [type(le).__name__ for le in coolest_1.lensing_entities]
        results['source_light_model'] = [type(m).__name__ for m in coolest_1.lensing_entities[2].light_model]

        ### DMR Plot
        norm = Normalize(-0.005, 0.05)
        fig, axes = plt.subplots(2, 2, constrained_layout=True)
        splotter = ModelPlotter(coolest_1, coolest_directory=os.path.dirname(target_path))

        splotter.plot_data_image(axes[0, 0], norm=norm)
        axes[0, 0].set_title("Observed Data")

        splotter.plot_model_image(
            axes[0, 1],
            supersampling=5, convolved=True,
            kwargs_source=dict(entity_selection=[2]),
            kwargs_lens_mass=dict(entity_selection=[0, 1]),
            norm=norm
        )
        axes[0, 1].text(0.05, 0.05, f"$\\theta_{{\\rm E}}$ = {einstein_radius:.2f}\"", color='white', fontsize=12,
                        transform=axes[0, 1].transAxes)
        axes[0, 1].set_title("Image Model")

        splotter.plot_model_residuals(axes[1, 0], supersampling=5, add_chi2_label=True, chi2_fontsize=12,
                                      kwargs_source=dict(entity_selection=[2]),
                                      kwargs_lens_mass=dict(entity_selection=[0, 1]))
        axes[1, 0].set_title("Normalized Residuals")

        splotter.plot_surface_brightness(axes[1, 1], kwargs_light=dict(entity_selection=[2]),
                                         norm=norm, coordinates=coord_src)
        axes[1, 1].text(0.05, 0.05, f"$\\theta_{{\\rm eff}}$ = {r_eff_source:.2f}\"", color='white', fontsize=12,
                        transform=axes[1, 1].transAxes)
        axes[1, 1].set_title("Surface Brightness")

        for ax in axes[1]:
            ax.set_xlabel(r"$x$ (arcsec)")
            ax.set_ylabel(r"$y$ (arcsec)")
            
        if output_dir is not None:
            dmr_plot_path = os.path.join(output_dir, "dmr_plot.png")
            plt.savefig(dmr_plot_path, format='png', bbox_inches='tight')
            results['dmr_plot'] = dmr_plot_path
        plt.show()
        plt.close()

        

        
        ### Corner Plot
        truth = coolest_1
        # Only creates corner plot if sampling method was used to create lens model
        # Otherwise, no chains available for corner plot!
        if 'chain_file_name' in truth.meta.keys():
            free_pars = truth.lensing_entities.get_parameter_ids()[:-2]
            reorder = [2, 3, 4, 5, 6, 0, 1]
            pars = [free_pars[i] for i in reorder]
            results['free_parameters'] = pars
    
            param_plotter = ParametersPlotter(
                pars, [truth],
                coolest_directories=[os.path.dirname(target_path)],
                coolest_names=["Smooth source"],
                ref_coolest_objects=[truth],
                colors=['#7FB6F5', '#E03424'],
            )
    
            settings = {
                "ignore_rows": 0.0,
                "fine_bins_2D": 800,
                "smooth_scale_2D": 0.5,
                "mult_bias_correction_order": 5
            }
            param_plotter.init_getdist(settings_mcsamples=settings)
            param_plotter.plot_triangle_getdist(filled_contours=True, subplot_size=3)
            if output_dir is not None:
                corner_plot_path = os.path.join(output_dir, "corner_plot.png")
                plt.savefig(corner_plot_path, format='png', bbox_inches='tight')
                results['corner_plot'] = corner_plot_path
            plt.close()
    
            
        
    return results

1	__author__ = 'aymgal', 'Giorgos Vernardos'	×
2
3
4	import numpy as np	×
5	import warnings	×
6	import matplotlib	×
7	import matplotlib.pyplot as plt	×
8	from matplotlib import ticker	×
9	from mpl_toolkits.axes_grid1 import make_axes_locatable	×
10	from scipy.spatial import Voronoi, voronoi_plot_2d	×
11	from matplotlib.colors import Normalize, LogNorm, TwoSlopeNorm	×
12	from matplotlib.cm import ScalarMappable	×
13	from mpl_toolkits.axes_grid1.anchored_artists import AnchoredSizeBar	×
NEW 14	from coolest.api.composable_models import ComposableLensModel	×
15	import os	×
16	import tarfile	×
17	import tempfile	×
18	import io	×
19	from coolest.api import util	×
20	from coolest.api.analysis import Analysis	×
21	from coolest.api.plotting import ModelPlotter, ParametersPlotter	×
22
23
24	def plot_voronoi(ax, x, y, z, neg_values_as_bad=False,	×
25	norm=None, cmap=None, zmin=None, zmax=None,
26	edgecolor=None, zorder=1):
27
28	if cmap is None:	×
29	cmap = plt.get_cmap('inferno')	×
30	if norm is None:	×
31	if zmin is None:	×
32	zmin = np.min(z)	×
33	if zmax is None:	×
34	zmax = np.max(z)	×
35	norm = Normalize(zmin, zmax)	×
36
37	# get voronoi regions
38	voronoi_points = np.column_stack((x,y))	×
39	vor = Voronoi(voronoi_points)	×
40	new_regions, vertices = voronoi_finite_polygons_2d(vor)	×
41
42	# get cell colors
43	m = matplotlib.cm.ScalarMappable(norm=norm, cmap=cmap)	×
44
45	# plot voronoi points
46	#point_colors = [ m.to_rgba(v) for v in z ]
47	#ax.scatter(voronoi_points[:,0],voronoi_points[:,1],c=point_colors)
48
49	# plot voronoi cells
50	for i, region in enumerate(new_regions):	×
51	polygon = vertices[region]	×
52	z_i = z[i]	×
53	if neg_values_as_bad is True and z_i < 0.:	×
54	cell_color = m.to_rgba(np.nan)	×
55	else:
56	cell_color = m.to_rgba(z_i)	×
57	ax.fill(zip(polygon), facecolor=cell_color, edgecolor=edgecolor, zorder=zorder)	×
58	return m	×
59
60
61	def voronoi_finite_polygons_2d(vor,radius=None):	×
62	"""
63	Reconstruct infinite voronoi regions in a 2D diagram to finite
64	regions.
65	This function is taken from: https://gist.github.com/pv/8036995
66
67	Parameters
68	----------
69	vor : Voronoi
70	Input diagram
71	radius : float, optional
72	Distance to 'points at infinity'.
73
74	Returns
75	-------
76	regions : list of tuples
77	Indices of vertices in each revised Voronoi regions.
78	vertices : list of tuples
79	Coordinates for revised Voronoi vertices. Same as coordinates
80	of input vertices, with 'points at infinity' appended to the
81	end.
82	"""
83
84	if vor.points.shape[1] != 2:	×
85	raise ValueError("Requires 2D input")	×
86
87	new_regions = []	×
88	new_vertices = vor.vertices.tolist()	×
89
90	center = vor.points.mean(axis=0)	×
91	if radius is None:	×
92	radius = vor.points.ptp().max()	×
93
94	# Construct a map containing all ridges for a given point
95	all_ridges = {}	×
96	for (p1, p2), (v1, v2) in zip(vor.ridge_points, vor.ridge_vertices):	×
97	all_ridges.setdefault(p1, []).append((p2, v1, v2))	×
98	all_ridges.setdefault(p2, []).append((p1, v1, v2))	×
99
100	# Reconstruct infinite regions
101	for p1, region in enumerate(vor.point_region):	×
102	vertices = vor.regions[region]	×
103
104	if all(v >= 0 for v in vertices):	×
105	# finite region
106	new_regions.append(vertices)	×
107	continue	×
108
109	# reconstruct a non-finite region
110	ridges = all_ridges[p1]	×
111	new_region = [v for v in vertices if v >= 0]	×
112
113	for p2, v1, v2 in ridges:	×
114	if v2 < 0:	×
115	v1, v2 = v2, v1	×
116	if v1 >= 0:	×
117	# finite ridge: already in the region
118	continue	×
119
120	# Compute the missing endpoint of an infinite ridge
121
122	t = vor.points[p2] - vor.points[p1] # tangent	×
123	t /= np.linalg.norm(t)	×
124	n = np.array([-t[1], t[0]]) # normal	×
125
126	midpoint = vor.points[[p1, p2]].mean(axis=0)	×
127	direction = np.sign(np.dot(midpoint - center, n)) * n	×
128	far_point = vor.vertices[v2] + direction * radius	×
129
130	new_region.append(len(new_vertices))	×
131	new_vertices.append(far_point.tolist())	×
132
133	# sort region counterclockwise
134	vs = np.asarray([new_vertices[v] for v in new_region])	×
135	c = vs.mean(axis=0)	×
136	angles = np.arctan2(vs[:,1] - c[1], vs[:,0] - c[0])	×
137	new_region = np.array(new_region)[np.argsort(angles)]	×
138
139	# finish
140	new_regions.append(new_region.tolist())	×
141
142	return new_regions, np.asarray(new_vertices)	×
143
144
145	def std_colorbar(mappable, label=None, fontsize=12, label_kwargs={}, **colorbar_kwargs):	×
146	cb = plt.colorbar(mappable, **colorbar_kwargs)	×
147	if label is not None:	×
148	colorbar_kwargs.pop('label', None)	×
149	cb.set_label(label, fontsize=fontsize, **label_kwargs)	×
150	return cb	×
151
152	def std_colorbar_residuals(mappable, res_map, vmin, vmax, label=None, fontsize=12,	×
153	label_kwargs={}, **colorbar_kwargs):
154	if res_map.min() < vmin and res_map.max() > vmax:	×
155	cb_extend = 'both'	×
156	elif res_map.min() < vmin:	×
157	cb_extend = 'min'	×
158	elif res_map.max() > vmax:	×
159	cb_extend = 'max'	×
160	else:
161	cb_extend = 'neither'	×
162	colorbar_kwargs.update({'extend': cb_extend})	×
163	return std_colorbar(mappable, label=label, fontsize=fontsize,	×
164	label_kwargs=label_kwargs, **colorbar_kwargs)
165
166	def nice_colorbar(mappable, ax=None, position='right', pad=0.1, size='5%', label=None, fontsize=12,	×
167	invisible=False,
168	#max_nbins=None,
169	divider_kwargs={}, colorbar_kwargs={}, label_kwargs={}):
170	divider_kwargs.update({'position': position, 'pad': pad, 'size': size})	×
171	if ax is None:	×
172	ax = mappable.axes	×
173	divider = make_axes_locatable(ax)	×
174	cax = divider.append_axes(**divider_kwargs)	×
175	if invisible:	×
176	cax.axis('off')	×
177	return None	×
178	cb = plt.colorbar(	×
179	mappable, cax=cax,
180	**colorbar_kwargs
181	)
182	if label is not None:	×
183	colorbar_kwargs.pop('label', None)	×
184	cb.set_label(label, fontsize=fontsize, **label_kwargs)	×
185	if position == 'top':	×
186	cax.xaxis.set_ticks_position('top')	×
187	# if max_nbins is not None: # TODO: this leads to strange results
188	# # cb.locator = ticker.LogLocator(subs=range(10))
189	# # cb.update_ticks()
190	return cb	×
191
192	def nice_colorbar_residuals(mappable, res_map, vmin, vmax, ax=None, position='right', pad=0.1, size='5%',	×
193	invisible=False, label=None, fontsize=12,
194	divider_kwargs={}, colorbar_kwargs={}, label_kwargs={}):
195	if res_map.min() < vmin and res_map.max() > vmax:	×
196	cb_extend = 'both'	×
197	elif res_map.min() < vmin:	×
198	cb_extend = 'min'	×
199	elif res_map.max() > vmax:	×
200	cb_extend = 'max'	×
201	else:
202	cb_extend = 'neither'	×
203	colorbar_kwargs.update({'extend': cb_extend})	×
204	return nice_colorbar(mappable, ax=ax, position=position, pad=pad, size=size, label=label, fontsize=fontsize,	×
205	invisible=invisible, colorbar_kwargs=colorbar_kwargs, label_kwargs=label_kwargs,
206	divider_kwargs=divider_kwargs)
207
208	def cax_colorbar(fig, cax, norm=None, cmap=None, mappable=None, label=None, fontsize=12, orientation='horizontal', label_kwargs={}):	×
209	if mappable is None:	×
210	mappable = ScalarMappable(norm=norm, cmap=cmap)	×
211	cb = fig.colorbar(mappable=mappable, orientation=orientation, cax=cax)	×
212	if label is not None:	×
213	cb.set_label(label, fontsize=fontsize, **label_kwargs)	×
214
215	def scale_bar(ax, size, unit_suffix='"', loc='lower left', color='#FFFFFFBB', fontsize=12):	×
216	if size == int(size):	×
217	label = f"{int(size)}"	×
218	else:
219	label = f"{size:.1f}"	×
220	label += unit_suffix	×
221	artist = AnchoredSizeBar(	×
222	ax.transData,
223	size, label,
224	loc=loc, label_top=True,
225	pad=0.8, sep=5,
226	color=color, fontproperties=dict(size=fontsize),
227	frameon=False, size_vertical=0,
228	)
229	ax.add_artist(artist)	×
230	return artist	×
231
232	def plot_regular_grid(ax, title, image_, neg_values_as_bad=False, xylim=None, **imshow_kwargs):	×
233	if neg_values_as_bad:	×
234	image = np.copy(image_)	×
235	image[image < 0] = np.nan	×
236	else:
237	image = image_	×
238	im = ax.imshow(image, **imshow_kwargs)	×
239	im.set_rasterized(True)	×
240	set_xy_limits(ax, xylim)	×
241	ax.xaxis.set_major_locator(plt.MaxNLocator(3))	×
242	ax.yaxis.set_major_locator(plt.MaxNLocator(3))	×
243	ax.set_title(title)	×
244	return ax, im	×
245
246	def plot_irregular_grid(ax, title, points, xylim, neg_values_as_bad=False,	×
247	norm=None, cmap=None, plot_points=False):
248	x, y, z = points	×
249	im = plot_voronoi(ax, x, y, z, neg_values_as_bad=neg_values_as_bad,	×
250	norm=norm, cmap=cmap, zorder=1)
251	ax.set_aspect('equal', 'box')	×
252	set_xy_limits(ax, xylim)	×
253	ax.xaxis.set_major_locator(plt.MaxNLocator(3))	×
254	ax.yaxis.set_major_locator(plt.MaxNLocator(3))	×
255	if plot_points:	×
256	ax.scatter(x, y, s=5, c='white', marker='.', alpha=0.4, zorder=2)	×
257	ax.set_title(title)	×
258	return ax, im	×
259
260	def set_xy_limits(ax, xylim):	×
261	if xylim is None: return # do nothing	×
262	if isinstance(xylim, (int, float)):	×
263	xylim_ = [-xylim, xylim, -xylim, xylim]	×
264	elif isinstance(xylim, (list, tuple)) and len(xylim) == 2:	×
265	xylim_ = [xylim[0], xylim[1], xylim[0], xylim[1]]	×
266	elif not isinstance(xylim, (list, tuple)) and len(xylim) != 4:	×
267	raise ValueError("`xylim` argument should be a single number, a 2-tuple or a 4-tuple.")	×
268	else:
269	xylim_ = xylim	×
270	ax.set_xlim(xylim_[0], xylim_[1])	×
271	ax.set_ylim(xylim_[2], xylim_[3])	×
272
273	def panel_label(ax, text, color, fontsize, alpha=0.8, loc='upper left'):	×
274	if loc == 'upper left':	×
275	x, y, ha, va = 0.03, 0.97, 'left', 'top'	×
276	elif loc == 'lower left':	×
277	x, y, ha, va = 0.03, 0.03, 'left', 'bottom'	×
278	elif loc == 'upper right':	×
279	x, y, ha, va = 0.97, 0.97, 'right', 'top'	×
280	elif loc == 'lower right':	×
281	x, y, ha, va = 0.97, 0.03, 'right', 'bottom'	×
282	ax.text(x, y, text, color=color, fontsize=fontsize, alpha=alpha,	×
283	ha=ha, va=va, transform=ax.transAxes)
284
NEW 285	def normalize_across_images(plotter_list, data_model_specifier, kwargs_source = None, kwargs_lens_mass = None,	×
286	supersampling=5, convolved=True, super_convolution=True):
287	"""Calculate the vmin and vmax to normalize the colormap across multiple coolest objects
288
289	Parameters
290	----------
291	plotter_list: list
292	List of ModelPlotter objects. May be acquired from MultiModelPlotter object
293	data_model_specifier: list
294	List of 0s and 1s; 0 = data, 1 = model. Specifies which set of pixel values should be used
295	when finding global minima and maxima -- data or model
296	kwargs_source: dict
297	Dictionary with "entity_selection" key, same as used in MultiModelPlotters.
298	"Entity_selection" contains list of lists. Selects source entities.
299	Insert dummy None values into dictionary for data.
300	kwargs_lens_mass: dict
301	Dictionary with "entity_selection" key, same as used in MultiModelPlotters.
302	"Entity_selection" contains list of lists. Selects lens mass entities.
303	Insert dummy None values into dictionary for data.
304	supersampling: int
305	Model image generation param
306	convolved: bool
307	Model image generation param
308	super_convolution: bool
309	Model image generation param
310
311
312	Returns
313	-------
314	vmin: float
315	global min value across all coolest objects in plotter_list for the specified data/models
316	vmax: float
317	global max value across all coolest objects in plotter_list for the specified data/models
318	"""
319
NEW 320	mins = []	×
NEW 321	maxes = []	×
NEW 322	ks_arr = kwargs_source['entity_selection']	×
NEW 323	km_arr = kwargs_lens_mass['entity_selection']	×
NEW 324	for plotter, d_or_f, ks, km in zip(plotter_list, data_model_specifier, ks_arr, km_arr):	×
325	# Check if we are finding extrema for data or model
NEW 326	if d_or_f == 0:	×
NEW 327	image = plotter.coolest.observation.pixels.get_pixels(directory=plotter._directory)	×
NEW 328	elif d_or_f == 1:	×
NEW 329	lens_model = ComposableLensModel(plotter.coolest, plotter._directory,	×
330	kwargs_selection_source=dict(entity_selection=ks),
331	kwargs_selection_lens_mass=dict(entity_selection=km))
NEW 332	image, _ = lens_model.model_image(supersampling, convolved, super_convolution)	×
333	# Find min and max and append
NEW 334	mins.append(np.min(image))	×
NEW 335	maxes.append(np.max(image))	×
NEW 336	vmin = min(mins)	×
NEW 337	vmax = max(maxes)	×
NEW 338	return vmin, vmax	×
339
340
341	def dmr_corner(tar_path, output_dir = None):	×
342	"""Given .tar.gz COOLEST file, plots and optionally saves DMR and corner plots for COOLEST file. Returns dictionary of important extracted information.
343
344	Parameters
345	----------
346	tar_path : string
347	Path to .tar.gz COOLEST file
348	output_dir : string, optional
349	Path to automatically save DMR and corner plot to if specified, by default None
350
351	Returns
352	-------
353	results: dictionary
354	Contains useful information about the COOLEST objects
355	"""
356
357	results = {}	×
358	with tempfile.TemporaryDirectory() as tmpdir:	×
359	# Extract tar.gz archive
360	with tarfile.open(tar_path, "r:gz") as tar:	×
361	tar.extractall(path=tmpdir)	×
362
363
364	# Get path to the extracted JSON file
365	extracted_items = os.listdir(tmpdir)	×
366	extracted_path = os.path.join(tmpdir, extracted_items[0])	×
367	if os.path.isdir(extracted_path):	×
368	extracted_files = os.listdir(extracted_path)	×
369	else:
370	extracted_files = extracted_items	×
371	extracted_path = tmpdir # fallback	×
372
373	json_files = [f for f in extracted_files if f.endswith(".json")]	×
374	if not json_files:	×
375	raise ValueError("No .json file found in archive.")	×
376
377	json_path = os.path.join(extracted_path, json_files[0])	×
378	target_path = os.path.splitext(json_path)[0]	×
379
380	# Load COOLEST object
381	coolest_1 = util.get_coolest_object(target_path, verbose=False)	×
382
383	# Run analysis
384	analysis = Analysis(coolest_1, target_path, supersampling=5)	×
385
386	coord_orig = util.get_coordinates(coolest_1)	×
387	coord_src = coord_orig.create_new_coordinates(pixel_scale_factor=0.1, grid_shape=(1.42, 1.42))	×
388
389	# Extract values
390	r_eff_source = analysis.effective_radius_light(center=(0,0), coordinates=coord_src, outer_radius=1., entity_selection=[2])	×
391	einstein_radius = analysis.effective_einstein_radius(entity_selection=[0,1])	×
392
393	results['r_eff_source'] = r_eff_source	×
394	results['einstein_radius'] = einstein_radius	×
395	results['lensing_entities'] = [type(le).__name__ for le in coolest_1.lensing_entities]	×
396	results['source_light_model'] = [type(m).__name__ for m in coolest_1.lensing_entities[2].light_model]	×
397
398	### DMR Plot
399	norm = Normalize(-0.005, 0.05)	×
400	fig, axes = plt.subplots(2, 2, constrained_layout=True)	×
401	splotter = ModelPlotter(coolest_1, coolest_directory=os.path.dirname(target_path))	×
402
403	splotter.plot_data_image(axes[0, 0], norm=norm)	×
404	axes[0, 0].set_title("Observed Data")	×
405
406	splotter.plot_model_image(	×
407	axes[0, 1],
408	supersampling=5, convolved=True,
409	kwargs_source=dict(entity_selection=[2]),
410	kwargs_lens_mass=dict(entity_selection=[0, 1]),
411	norm=norm
412	)
413	axes[0, 1].text(0.05, 0.05, f"$\\theta_{{\\rm E}}$ = {einstein_radius:.2f}\"", color='white', fontsize=12,	×
414	transform=axes[0, 1].transAxes)
415	axes[0, 1].set_title("Image Model")	×
416
417	splotter.plot_model_residuals(axes[1, 0], supersampling=5, add_chi2_label=True, chi2_fontsize=12,	×
418	kwargs_source=dict(entity_selection=[2]),
419	kwargs_lens_mass=dict(entity_selection=[0, 1]))
420	axes[1, 0].set_title("Normalized Residuals")	×
421
422	splotter.plot_surface_brightness(axes[1, 1], kwargs_light=dict(entity_selection=[2]),	×
423	norm=norm, coordinates=coord_src)
424	axes[1, 1].text(0.05, 0.05, f"$\\theta_{{\\rm eff}}$ = {r_eff_source:.2f}\"", color='white', fontsize=12,	×
425	transform=axes[1, 1].transAxes)
426	axes[1, 1].set_title("Surface Brightness")	×
427
428	for ax in axes[1]:	×
429	ax.set_xlabel(r"$x$ (arcsec)")	×
430	ax.set_ylabel(r"$y$ (arcsec)")	×
431
432	if output_dir is not None:	×
433	dmr_plot_path = os.path.join(output_dir, "dmr_plot.png")	×
434	plt.savefig(dmr_plot_path, format='png', bbox_inches='tight')	×
435	results['dmr_plot'] = dmr_plot_path	×
436	plt.show()	×
437	plt.close()	×
438
439
440
441
442	### Corner Plot
443	truth = coolest_1	×
444	# Only creates corner plot if sampling method was used to create lens model
445	# Otherwise, no chains available for corner plot!
446	if 'chain_file_name' in truth.meta.keys():	×
447	free_pars = truth.lensing_entities.get_parameter_ids()[:-2]	×
448	reorder = [2, 3, 4, 5, 6, 0, 1]	×
449	pars = [free_pars[i] for i in reorder]	×
450	results['free_parameters'] = pars	×
451
452	param_plotter = ParametersPlotter(	×
453	pars, [truth],
454	coolest_directories=[os.path.dirname(target_path)],
455	coolest_names=["Smooth source"],
456	ref_coolest_objects=[truth],
457	colors=['#7FB6F5', '#E03424'],
458	)
459
460	settings = {	×
461	"ignore_rows": 0.0,
462	"fine_bins_2D": 800,
463	"smooth_scale_2D": 0.5,
464	"mult_bias_correction_order": 5
465	}
466	param_plotter.init_getdist(settings_mcsamples=settings)	×
467	param_plotter.plot_triangle_getdist(filled_contours=True, subplot_size=3)	×
468	if output_dir is not None:	×
469	corner_plot_path = os.path.join(output_dir, "corner_plot.png")	×
470	plt.savefig(corner_plot_path, format='png', bbox_inches='tight')	×
471	results['corner_plot'] = corner_plot_path	×
472	plt.close()	×
473
474
475
476	return results	×

aymgal / COOLEST / 18846156192

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous