<commit_before>import numpy as np
from scipy.ndimage.measurements import center_of_mass
from astropy.stats import sigma_clip
import matplotlib.pyplot as pl
import matplotlib.cm as cm
from matplotlib.colors import SymLogNorm
from Model_objs import *
from Data_objs import *
from uvimage import uvimageslow
from calc_likelihood import create_modelimage,fft_interpolate
from modelcal import model_cal
from GenerateLensingGrid import GenerateLensingGrid
from utils import *
import copy

def plot_images(data,mcmcresult,returnimages=False,plotcombined=False,plotall=False,
                  imsize=512,pixsize=0.2,taper=0.,**kwargs):
      """
      Create a four-panel figure from data and chains,
      showing data, best-fit model, residuals, high-res image.

      Inputs:
      data:
            The visdata object(s) to be imaged. If more than
            one is passed, they'll be imaged/differenced separately.
      chains:
            A result from running LensModelMCMC.

      returnimages: bool, optional
            If True, will also return a list of numpy arrays
            containing the imaged data, interpolated model, and full-res model.
            Default is False.
      
      plotcombined: bool, optional, default False
            If True, plots only a single row of images, after combining all datasets
            in `data' and applying any necessary rescaling/shifting contained in
            `mcmcresult'.
      
      plotall: bool, optional, default False
            If True, plots each dataset in its own row *and* the combined dataset in
            an extra row.

      Returns:
      If returnimages is False:
      f, axarr:
            A matplotlib figure and array of Axes objects.

      If returnimages is True:
      f,axarr,imagelist:
            Same as above; imagelist is a list of length (# of datasets),
            containing three arrays each, representing the data,
            interpolated model, and full-resolution model.
      """

      limits = kwargs.pop('limits',
            [-imsize*pixsize/2.,+imsize*pixsize/2.,-imsize*pixsize/2.,+imsize*pixsize/2.])
      cmap = kwargs.pop('cmap',cm.Greys)
      mapcontours = kwargs.pop('mapcontours',np.delete(np.arange(-21,22,3),7))
      rescontours = kwargs.pop('rescontours',np.array([-6,-5,-4,-3,-2,2,3,4,5,6]))
      level = kwargs.pop('level',None)
      logmodel = kwargs.pop('logmodel',False)

      datasets = list(np.array([data]).flatten())

      # shorthand for later
      c = copy.deepcopy(mcmcresult['chains'])

      # Set up to create the model image. We'll assume the best-fit values are all the medians.
      lens,source = copy.deepcopy(mcmcresult['lens_p0']), copy.deepcopy(mcmcresult['source_p0'])
      for i,ilens in enumerate(lens):
            if ilens.__class__.__name__ == 'SIELens':
                  for key in ['x','y','M','e','PA']:
                        if not vars(ilens)[key]['fixed']:
                              ilens.__dict__[key]['value'] = np.median(c[key+'L'+str(i)])
      # now do the source(s)
      for i,src in enumerate(source): # Source is a list of source objects
            if src.__class__.__name__ == 'GaussSource':
                  for key in ['xoff','yoff','flux','width']:
                        if not vars(src)[key]['fixed']:
                              src.__dict__[key]['value'] = np.median(c[key+'S'+str(i)])
            elif src.__class__.__name__ == 'SersicSource':
                  for key in ['xoff','yoff','flux','reff','index','axisratio','PA']:
                        if not vars(src)[key]['fixed']:
                              src.__dict__[key]['value'] = np.median(c[key+'S'+str(i)])
            elif src.__class__.__name__ == 'PointSource':
                  for key in ['xoff','yoff','flux']:
                        if not vars(src)[key]['fixed']:
                              src.__dict__[key]['value'] = np.median(c[key+'S'+str(i)])
      
      # now do shear, if any
      if 'shear_p0' in mcmcresult.keys():
            shear = mcmcresult['shear_p0']
            for key in ['shear','shearangle']:
                  if not vars(shear)[key]['fixed']:
                        shear.__dict__[key]['value'] = np.median(c[key])
      else: shear = None 

      # Any amplitude scaling or astrometric shifts
      scaleamp = np.ones(len(datasets))
      if 'scaleamp' in mcmcresult.keys():
            for i,t in enumerate(mcmcresult['scaleamp']): # only matters if >1 datasets
                  if i==0: pass
                  elif t: scaleamp[i] = np.median(c['ampscale_dset'+str(i)])
                  else: pass
      shiftphase = np.zeros((len(datasets),2))
      if 'shiftphase' in mcmcresult.keys():
            for i,t in enumerate(mcmcresult['shiftphase']):
                  if i==0: pass # only matters if >1 datasets
                  elif t:
                        shiftphase[i][0] = np.median(c['astromshift_x_dset'+str(i)])
                        shiftphase[i][1] = np.median(c['astromshift_y_dset'+str(i)])
                  else: pass # no shifting

      sourcedatamap = mcmcresult['sourcedatamap'] if 'sourcedatamap' in mcmcresult.keys() else None
      modelcal = mcmcresult['modelcal'] if 'modelcal' in mcmcresult.keys() else [False]*len(datasets)

      
      if plotall:
            f,axarr = pl.subplots(len(datasets)+1,4,figsize=(14,4*(len(datasets)+1)))
            axarr = np.atleast_2d(axarr)
            images = [[] for _ in range(len(datasets)+1)]
      elif plotcombined:
            f,axarr = pl.subplots(1,4,figsize=(12,3))
            axarr = np.atleast_2d(axarr)
            images = [[]]
      else:
            f,axarr = pl.subplots(len(datasets),4,figsize=(14,4*len(datasets)))
            axarr = np.atleast_2d(axarr)
            images = [[] for _ in range(len(datasets))] # effing mutable lists.
            
      plotdata,plotinterp = [],[]
      
      for i,dset in enumerate(datasets):
            # Get us some coordinates.
            xmap,ymap,xemit,yemit,ix = GenerateLensingGrid(datasets,mcmcresult['xmax'],\
                  mcmcresult['highresbox'],mcmcresult['fieldres'],mcmcresult['emitres'])
            <commit_msg>better way of retrieving best-fit params to generate model images<commit_after>import numpy as np
from scipy.ndimage.measurements import center_of_mass
from astropy.stats import sigma_clip
import matplotlib.pyplot as pl
import matplotlib.cm as cm
from matplotlib.colors import SymLogNorm
from Model_objs import *
from Data_objs import *
from uvimage import uvimageslow
from calc_likelihood import create_modelimage,fft_interpolate
from modelcal import model_cal
from GenerateLensingGrid import GenerateLensingGrid
from utils import *
import copy

def plot_images(data,mcmcresult,returnimages=False,plotcombined=False,plotall=False,
                  imsize=512,pixsize=0.2,taper=0.,**kwargs):
      """
      Create a four-panel figure from data and chains,
      showing data, best-fit model, residuals, high-res image.

      Inputs:
      data:
            The visdata object(s) to be imaged. If more than
            one is passed, they'll be imaged/differenced separately.
      chains:
            A result from running LensModelMCMC.

      returnimages: bool, optional
            If True, will also return a list of numpy arrays
            containing the imaged data, interpolated model, and full-res model.
            Default is False.
      
      plotcombined: bool, optional, default False
            If True, plots only a single row of images, after combining all datasets
            in `data' and applying any necessary rescaling/shifting contained in
            `mcmcresult'.
      
      plotall: bool, optional, default False
            If True, plots each dataset in its own row *and* the combined dataset in
            an extra row.

      Returns:
      If returnimages is False:
      f, axarr:
            A matplotlib figure and array of Axes objects.

      If returnimages is True:
      f,axarr,imagelist:
            Same as above; imagelist is a list of length (# of datasets),
            containing three arrays each, representing the data,
            interpolated model, and full-resolution model.
      """

      limits = kwargs.pop('limits',
            [-imsize*pixsize/2.,+imsize*pixsize/2.,-imsize*pixsize/2.,+imsize*pixsize/2.])
      cmap = kwargs.pop('cmap',cm.Greys)
      mapcontours = kwargs.pop('mapcontours',np.delete(np.arange(-21,22,3),7))
      rescontours = kwargs.pop('rescontours',np.array([-6,-5,-4,-3,-2,2,3,4,5,6]))
      level = kwargs.pop('level',None)
      logmodel = kwargs.pop('logmodel',False)

      datasets = list(np.array([data]).flatten())

      # shorthand for later
      c = copy.deepcopy(mcmcresult['chains'])
      """
      # Set up to create the model image. We'll assume the best-fit values are all the medians.
      lens,source = copy.deepcopy(mcmcresult['lens_p0']), copy.deepcopy(mcmcresult['source_p0'])
      for i,ilens in enumerate(lens):
            if ilens.__class__.__name__ == 'SIELens':
                  for key in ['x','y','M','e','PA']:
                        if not vars(ilens)[key]['fixed']:
                              ilens.__dict__[key]['value'] = np.median(c[key+'L'+str(i)])
      # now do the source(s)
      for i,src in enumerate(source): # Source is a list of source objects
            if src.__class__.__name__ == 'GaussSource':
                  for key in ['xoff','yoff','flux','width']:
                        if not vars(src)[key]['fixed']:
                              src.__dict__[key]['value'] = np.median(c[key+'S'+str(i)])
            elif src.__class__.__name__ == 'SersicSource':
                  for key in ['xoff','yoff','flux','reff','index','axisratio','PA']:
                        if not vars(src)[key]['fixed']:
                              src.__dict__[key]['value'] = np.median(c[key+'S'+str(i)])
            elif src.__class__.__name__ == 'PointSource':
                  for key in ['xoff','yoff','flux']:
                        if not vars(src)[key]['fixed']:
                              src.__dict__[key]['value'] = np.median(c[key+'S'+str(i)])
      
      # now do shear, if any
      if 'shear_p0' in mcmcresult.keys():
            shear = mcmcresult['shear_p0']
            for key in ['shear','shearangle']:
                  if not vars(shear)[key]['fixed']:
                        shear.__dict__[key]['value'] = np.median(c[key])
      else: shear = None 

      # Any amplitude scaling or astrometric shifts
      scaleamp = np.ones(len(datasets))
      if 'scaleamp' in mcmcresult.keys():
            for i,t in enumerate(mcmcresult['scaleamp']): # only matters if >1 datasets
                  if i==0: pass
                  elif t: scaleamp[i] = np.median(c['ampscale_dset'+str(i)])
                  else: pass
      shiftphase = np.zeros((len(datasets),2))
      if 'shiftphase' in mcmcresult.keys():
            for i,t in enumerate(mcmcresult['shiftphase']):
                  if i==0: pass # only matters if >1 datasets
                  elif t:
                        shiftphase[i][0] = np.median(c['astromshift_x_dset'+str(i)])
                        shiftphase[i][1] = np.median(c['astromshift_y_dset'+str(i)])
                  else: pass # no shifting

      sourcedatamap = mcmcresult['sourcedatamap'] if 'sourcedatamap' in mcmcresult.keys() else None
      modelcal = mcmcresult['modelcal'] if 'modelcal' in mcmcresult.keys() else [False]*len(datasets)
      """

      # Now all these things are saved separately, don't have to do the BS above
      lens = mcmcresult['best-fit']['lens']
      source=mcmcresult['best-fit']['source']
      shear =mcmcresult['best-fit']['shear'] if 'shear' in mcmcresult['best-fit'].keys() else None
      scaleamp = mcmcresult['best-fit']['scaleamp'] if 'scaleamp' in mcmcresult['best-fit'].keys() else np.ones(len(datasets))
      shiftphase=mcmcresult['best-fit']['shiftphase'] if 'shiftphase' in mcmcresult['best-fit'].keys() else np.zeros((len(datasets),2))
      sourcedatamap = mcmcresult['sourcedatamap'] if 'sourcedatamap' in mcmcresult.keys() else None
      modelcal = mcmcresult['modelcal'] if 'modelcal' in mcmcresult.keys() else [False]*len(datasets)
      
      
      if plotall:
            f,axarr = pl.subplots(len(datasets)+1,4,figsize=(14,4*(len(datasets)+1)))
            axarr = np.atleast_2d(axarr)
            images = [[] for _ in range(len(datasets)+1)]
      elif plotcombined:
            f,axarr = pl.subplots(1,4,figsize=(12,3))
            axarr = np.atleast_2d(axarr)
            images = [[]]
      else:
            f,axarr = pl.subplots(len(datasets),4,figsize=(14,4*len(datasets)))
            axarr = np.atleast_2d(axarr)
            images = [[] for _ in range(len(datasets))] # effing mutable lists.
            
      plotdata,plotinterp = [],[]
      
      for i,dset in enumerate(datasets):
            # Get us some coordinates.
            xmap,ymap,xemit,yemit,ix = GenerateLensingGrid(datasets,mcmcresult['xmax'],\
                  mcmcresult['highresbox'],mcmcresult['fieldres'],mcmcresult['emitres'])
            
{'fraction_non_alphanumeric': 0.12499048489000533, 'fraction_numerical': 0.008601659435183071, 'mean_word_length': 2.337058674117348, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2587868', 'n_tokens_mistral': 4129, 'n_tokens_neox': 3975, 'n_words': 975}