plot.py 5.64 KB
 Martin Reinecke committed Sep 08, 2017 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 ``````from __future__ import division import numpy as np from ..import Field, RGSpace, HPSpace, GLSpace, PowerSpace import os # relevant properties: # - x/y size # - x/y/z log # - x/y/z min/max # - colorbar/colormap # - axis on/off # - title # - axis labels def _mollweide_helper(xsize): xsize = int(xsize) ysize = xsize//2 res = np.full(shape=(ysize, xsize), fill_value=np.nan, dtype=np.float64) xc = (xsize-1)*0.5 yc = (ysize-1)*0.5 u, v = np.meshgrid(np.arange(xsize), np.arange(ysize)) u = 2*(u-xc)/(xc/1.02) v = (v-yc)/(yc/1.02) mask = np.where((u*u*0.25 + v*v) <= 1.) t1 = v[mask] theta = 0.5*np.pi-( np.arcsin(2/np.pi*(np.arcsin(t1) + t1*np.sqrt((1.-t1)*(1+t1))))) phi = -0.5*np.pi*u[mask]/np.maximum(np.sqrt((1-t1)*(1+t1)), 1e-6) phi = np.where(phi < 0, phi+2*np.pi, phi) return res, mask, theta, phi def _find_closest(A, target): # A must be sorted idx = A.searchsorted(target) idx = np.clip(idx, 1, len(A)-1) left = A[idx-1] right = A[idx] idx -= target - left < right - target return idx def _makeplot(name): `````` Martin Reinecke committed Sep 08, 2017 44 `````` import matplotlib.pyplot as plt `````` Martin Reinecke committed Sep 08, 2017 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 `````` if name is None: plt.show() return extension = os.path.splitext(name)[1] if extension==".pdf": plt.savefig(name) plt.close() elif extension==".png": plt.savefig(name) plt.close() #elif extension==".html": #import mpld3 #mpld3.save_html(plt.gcf(),fileobj=name,no_extras=True) #import plotly.offline as py #import plotly.tools as tls #plotly_fig = tls.mpl_to_plotly(plt.gcf()) #py.plot(plotly_fig,filename=name) #py.plot_mpl(plt.gcf(),filename=name) #import bokeh #bokeh.mpl.to_bokeh(plt.gcf()) else: raise ValueError("file format not understood") `````` Martin Reinecke committed Sep 08, 2017 68 69 70 71 72 73 74 75 76 77 78 79 80 81 ``````def _limit_xy(**kwargs): import matplotlib.pyplot as plt x1,x2,y1,y2 = plt.axis() if (kwargs.get("xmin")) is not None: x1 = kwargs.get("xmin") if (kwargs.get("xmax")) is not None: x2 = kwargs.get("xmax") if (kwargs.get("ymin")) is not None: y1 = kwargs.get("ymin") if (kwargs.get("ymax")) is not None: y2 = kwargs.get("ymax") plt.axis((x1,x2,y1,y2)) def plot(f,**kwargs): `````` Martin Reinecke committed Sep 08, 2017 82 `````` import matplotlib.pyplot as plt `````` Martin Reinecke committed Sep 08, 2017 83 84 85 86 87 88 `````` if not isinstance(f,Field): raise TypeError("incorrect data type") if len(f.domain)!=1: raise ValueError("input field must have exactly one domain") dom = f.domain[0] `````` Martin Reinecke committed Sep 08, 2017 89 90 91 92 93 94 95 96 97 `````` fig = plt.figure() ax = fig.add_subplot(1,1,1) xsize,ysize = 6,6 if kwargs.get("xsize") is not None: xsize = kwargs.get("xsize") if kwargs.get("ysize") is not None: ysize = kwargs.get("ysize") fig.set_size_inches(xsize,ysize) `````` Martin Reinecke committed Sep 08, 2017 98 `````` `````` Martin Reinecke committed Sep 08, 2017 99 100 101 102 103 104 105 106 107 `````` if kwargs.get("title") is not None: ax.set_title(kwargs.get("title")) if kwargs.get("xlabel") is not None: ax.set_xlabel(kwargs.get("xlabel")) if kwargs.get("ylabel") is not None: ax.set_ylabel(kwargs.get("ylabel")) cmap=plt.rcParams['image.cmap'] if kwargs.get("colormap") is not None: cmap = kwargs.get("colormap") `````` Martin Reinecke committed Sep 08, 2017 108 109 110 111 112 113 114 `````` if isinstance(dom, RGSpace): if len(dom.shape)==1: npoints = dom.shape[0] dist = dom.distances[0] xcoord = np.arange(npoints,dtype=np.float64)*dist ycoord = f.val plt.plot(xcoord, ycoord) `````` Martin Reinecke committed Sep 08, 2017 115 116 `````` _limit_xy(**kwargs) _makeplot(kwargs.get("name")) `````` Martin Reinecke committed Sep 08, 2017 117 118 119 120 121 122 123 124 `````` return elif len(dom.shape)==2: nx = dom.shape[0] ny = dom.shape[1] dx = dom.distances[0] dy = dom.distances[1] xc = np.arange(nx,dtype=np.float64)*dx yc = np.arange(ny,dtype=np.float64)*dy `````` Martin Reinecke committed Sep 08, 2017 125 126 127 128 129 130 131 132 `````` im=ax.imshow(f.val,extent=[xc[0],xc[-1],yc[0],yc[-1]],vmin=kwargs.get("zmin"),vmax=kwargs.get("zmax"),cmap=cmap) #from mpl_toolkits.axes_grid1 import make_axes_locatable #divider = make_axes_locatable(ax) #cax = divider.append_axes("right", size="5%", pad=0.05) #plt.colorbar(im,cax=cax) plt.colorbar(im) _limit_xy(**kwargs) _makeplot(kwargs.get("name")) `````` Martin Reinecke committed Sep 08, 2017 133 134 135 136 137 138 139 140 `````` return elif isinstance(dom, PowerSpace): xcoord = dom.kindex ycoord = f.val plt.xscale('log') plt.yscale('log') plt.title('power') plt.plot(xcoord, ycoord) `````` Martin Reinecke committed Sep 08, 2017 141 142 `````` _limit_xy(**kwargs) _makeplot(kwargs.get("name")) `````` Martin Reinecke committed Sep 08, 2017 143 144 145 146 147 148 149 150 151 152 153 154 `````` return elif isinstance(dom, HPSpace): import pyHealpix xsize = 800 res, mask, theta, phi = _mollweide_helper(xsize) ptg = np.empty((phi.size, 2), dtype=np.float64) ptg[:, 0] = theta ptg[:, 1] = phi base = pyHealpix.Healpix_Base(int(np.sqrt(f.val.size//12)), "RING") res[mask] = f.val[base.ang2pix(ptg)] plt.axis('off') `````` Martin Reinecke committed Sep 08, 2017 155 156 157 `````` plt.imshow(res,vmin=kwargs.get("zmin"),vmax=kwargs.get("zmax"),cmap=cmap) plt.colorbar(orientation="horizontal") _makeplot(kwargs.get("name")) `````` Martin Reinecke committed Sep 08, 2017 158 159 160 161 162 163 164 165 166 167 168 169 170 `````` return elif isinstance(dom, GLSpace): import pyHealpix xsize = 800 res, mask, theta, phi = _mollweide_helper(xsize) ra = np.linspace(0, 2*np.pi, dom.nlon+1) dec = pyHealpix.GL_thetas(dom.nlat) ilat = _find_closest(dec, theta) ilon = _find_closest(ra, phi) ilon = np.where(ilon == dom.nlon, 0, ilon) res[mask] = f.val[ilat*dom.nlon + ilon] plt.axis('off') `````` Martin Reinecke committed Sep 08, 2017 171 172 173 `````` plt.imshow(res,vmin=kwargs.get("zmin"),vmax=kwargs.get("zmax"),cmap=cmap) plt.colorbar(orientation="horizontal") _makeplot(kwargs.get("name")) `````` Martin Reinecke committed Sep 08, 2017 174 175 176 `````` return raise ValueError("Field type not(yet) supported")``````