Covariance Kernels

dist_cov.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
#  dist_cov.py
#  
#  Copyright 2015 Gregory M. Green
#  
#  This program is free software; you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation; either version 2 of the License, or
#  (at your option) any later version.
#  
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#  
#  You should have received a copy of the GNU General Public License
#  along with this program; if not, write to the Free Software
#  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
#  MA 02110-1301, USA.
#  
#  

import numpy as np
import matplotlib
matplotlib.use('Agg')
matplotlib.rc('text', usetex=True)
import matplotlib.pyplot as plt


def calc_cov(n_pix, kernel):
    x = np.arange(n_pix)
    xx,yy = np.meshgrid(x, x)
    xx = xx.flatten()
    yy = yy.flatten()
    
    d2 = (xx[:,None]-xx[None,:])**2. + (yy[:,None]-yy[None,:])**2.
    
    cov = kernel(d2)
    inv_cov = np.linalg.inv(cov)
    
    return cov, inv_cov


def make_kernel_plot(n_pix, kernel, title, fname, dpi=300):
    cov, inv_cov = calc_cov(n_pix, kernel)
    
    fig = plt.figure(figsize=(10,10), dpi=dpi)
    
    i0 = n_pix**2/2-250
    i1 = n_pix**2/2+250
    
    # Covariance matrix
    ax = fig.add_subplot(2,2,1)
    ax.imshow(cov.T[i0:i1,i0:i1], origin='lower', aspect='equal', cmap='binary',
              vmin=0., vmax=1., interpolation='nearest')
    ax.set_title(r'$\Sigma \ \left( \mathrm{zoomed \ in} \right)$', fontsize=16)
    ax.set_xlabel(r'$\mathrm{pixel \ index}$', fontsize=14)
    ax.set_ylabel(r'$\mathrm{pixel \ index}$', fontsize=14)
    
    # Inverse covariance matrix
    vmax = np.max(np.abs(inv_cov))
    ax = fig.add_subplot(2,2,2)
    ax.imshow(inv_cov.T[i0:i1,i0:i1], origin='lower', aspect='equal', cmap='bwr_r',
              vmin=-vmax, vmax=vmax, interpolation='nearest')
    ax.set_title(r'$\Sigma^{-1} \ \left( \mathrm{zoomed \ in} \right)$', fontsize=16)
    ax.set_xlabel(r'$\mathrm{pixel \ index}$', fontsize=14)
    ax.set_ylabel(r'$\mathrm{pixel \ index}$', fontsize=14)
    
    # Covariance map
    cov_map = np.reshape(cov[cov.shape[0]/2,:], (n_pix,n_pix))
    ax = fig.add_subplot(2,2,3)
    ax.imshow(cov_map.T, origin='lower', aspect='equal', cmap='binary',
              vmin=0., vmax=1., interpolation='nearest')
    ax.set_title(r'$\mathrm{Covariance \ with \ central \ pixel}$', fontsize=16)
    ax.set_xlabel(r'$x$', fontsize=14)
    ax.set_ylabel(r'$y$', fontsize=14)
    
    # Inverse covariance map
    inv_cov_map = np.reshape(inv_cov[inv_cov.shape[0]/2,:], (n_pix,n_pix))
    ax = fig.add_subplot(2,2,4)
    ax.imshow(inv_cov_map.T, origin='lower', aspect='equal', cmap='bwr_r',
              vmin=-vmax, vmax=vmax, interpolation='nearest')
    ax.set_title(r'$\mathrm{Inv. \ cov. \ with \ central \ pixel}$', fontsize=16)
    ax.set_xlabel(r'$x$', fontsize=14)
    ax.set_ylabel(r'$y$', fontsize=14)
    
    fig.subplots_adjust(hspace=0.28)
    
    fig.suptitle(title, fontsize=20)
    
    fig.savefig(fname, dpi=dpi, bbox_inches='tight')
    plt.close(fig)


def main():
    n_pix = 61
    sigma = 3
    scale = 10
    
    kernel_gauss = lambda d2: np.exp(-0.5 * d2 / sigma**2.)
    kernel_exp = lambda d2: np.exp(-np.sqrt(d2) / scale)
    kernel_sum = lambda d2: 0.5 * (kernel_gauss(d2) + kernel_exp(d2))
    
    kernels = [kernel_exp, kernel_gauss, kernel_sum]
    
    titles = (
        r'$\Sigma_{{ij}} = \exp \left( -d_{{ij}} / \ell \right), \ \ \ell = {:d}$'.format(scale),
        r'$\Sigma_{{ij}} = \exp \left( -d_{{ij}}^{{2}} / \sigma^{{2}} \right), \ \ \sigma = {:d}$'.format(sigma),
        r'$\Sigma_{{ij}} = \frac{{1}}{{2}} \left[ \exp \left( -d_{{ij}} / \ell \right) + \exp \left( -d_{{ij}}^{{2}} / \sigma^{{2}} \right) \right], \ \ \ell = {:d}, \ \sigma = {:d}$'.format(scale, sigma),
    )
    
    fnames = ['kernel_{}.png'.format(s) for s in (
        'exp_l{:d}'.format(scale), 'gauss_s{:d}'.format(sigma), 'sum_l{:d}_s{:d}'.format(scale, sigma)
    )]
    
    for t,k,fn in zip(titles, kernels, fnames):
        print 'Generating {} ...'.format(fn)
        make_kernel_plot(n_pix, k, t, fn, dpi=500)
    
    return 0

if __name__ == '__main__':
    main()