Source code for pynbody.plot.sph

"""

sph
===

routines for plotting smoothed quantities

"""

import pylab as p
import matplotlib
import numpy as np
from .. import sph, config
from .. import units as _units


[docs]def sideon_image(sim, *args, **kwargs): """ Rotate the simulation so that the disc of the passed halo is side-on, then make an SPH image by passing the parameters into the function image For a description of keyword arguments see :func:`~pynbody.plot.sph.image`. """ from ..analysis import angmom with angmom.sideon(sim): return image(sim, *args, **kwargs)
[docs]def faceon_image(sim, *args, **kwargs): """ Rotate the simulation so that the disc of the passed halo is face-on, then make an SPH image by passing the parameters into the function image For a description of keyword arguments see :func:`~pynbody.plot.sph.image`. """ from ..analysis import angmom with angmom.faceon(sim): return image(sim, *args, **kwargs)
[docs]def velocity_image(sim, width="10 kpc", vector_color='black', edgecolor='black', quiverkey_bg_color=None, vector_resolution=40, scale=None, mode='quiver', key_x=0.3, key_y=0.9, key_color='white', key_length="100 km s**-1", quiverkey=True, density=1.0, vector_qty='vel', **kwargs): """ Make an SPH image of the given simulation with velocity vectors overlaid on top. For a description of additional keyword arguments see :func:`~pynbody.plot.sph.image`, or see the `tutorial <http://pynbody.github.io/pynbody/tutorials/pictures.html#velocity-vectors>`_. **Keyword arguments:** *vector_color* (black): The color for the velocity vectors *edgecolor* (black): edge color used for the lines - using a color other than black for the *vector_color* and a black *edgecolor* can result in poor readability in pdfs *vector_resolution* (40): How many vectors in each dimension (default is 40x40) *quiverkey_bg_color* (none): The color for the legend (scale) background *scale* (None): The length of a vector that would result in a displayed length of the figure width/height. *mode* ('quiver'): make a 'quiver' or 'stream' plot *key_x* (0.3): Display x (width) position for the vector key (quiver mode only) *key_y* (0.9): Display y (height) position for the vector key (quiver mode only) *key_color* (white): Color for the vector key (quiver mode only) *key_length* (100 km/s): Velocity to use for the vector key (quiver mode only) *density* (1.0): Density of stream lines (stream mode only) *quiverkey* (True): Whether or not to inset the key *vector_qty* ('vel'): The name of the vector field to plot """ subplot = kwargs.get('subplot', False) av_z = kwargs.get('av_z',None) if subplot: p = subplot else: import matplotlib.pylab as p vx_name, vy_name, _ = sim._array_name_ND_to_1D(vector_qty) vx = image(sim, qty=vx_name, width=width, log=False, resolution=vector_resolution, noplot=True,av_z=av_z) vy = image(sim, qty=vy_name, width=width, log=False, resolution=vector_resolution, noplot=True,av_z=av_z) key_unit = _units.Unit(key_length) if isinstance(width, str) or issubclass(width.__class__, _units.UnitBase): if isinstance(width, str): width = _units.Unit(width) width = width.in_units(sim['pos'].units, **sim.conversion_context()) width = float(width) pixel_size = width / vector_resolution X, Y = np.meshgrid(np.arange(-width / 2 + pixel_size/2, width / 2 + pixel_size/2, pixel_size ), np.arange(-width / 2 + pixel_size/2, width / 2 + pixel_size/2, pixel_size)) im = image(sim, width=width, **kwargs) if mode == 'quiver': if scale is None: Q = p.quiver(X, Y, vx, vy, color=vector_color, edgecolor=edgecolor) else: Q = p.quiver(X, Y, vx, vy, scale=_units.Unit(scale).in_units( sim['vel'].units), color=vector_color, edgecolor=edgecolor) if quiverkey: qk = p.quiverkey(Q, key_x, key_y, key_unit.in_units(sim['vel'].units, **sim.conversion_context()), r"$\mathbf{" + key_unit.latex() + "}$", labelcolor=key_color, color=key_color, fontproperties={'size': 24}) if quiverkey_bg_color is not None: qk.text.set_backgroundcolor(quiverkey_bg_color) elif mode == 'stream' : Q = p.streamplot(X, Y, vx, vy, color=vector_color, density=density) # RS - if a axis object is passed in, use the right limit call if subplot: p.set_xlim(-width/2, width/2) p.set_ylim(-width/2, width/2) else: p.xlim(-width/2, width/2) p.ylim(-width/2, width/2) return im
[docs]def volume(sim, qty='rho', width=None, resolution=200, color=(1.0,1.0,1.0),vmin=None,vmax=None, dynamic_range=4.0,log=True, create_figure=True): """Create a volume rendering of the given simulation using mayavi. **Keyword arguments:** *qty* (rho): The name of the array to interpolate *width* (None): The width of the cube to generate, centered on the origin *resolution* (200): The number of elements along each side of the cube *color* (white): The color of the volume rendering. The value of each voxel is used to set the opacity. *vmin* (None): The value for zero opacity (calculated using dynamic_range if None) *vmax* (None): The value for full opacity (calculated from the maximum value in the region if None) *dynamic_range*: The dynamic range to use if vmin and vmax are not specified *log* (True): log-scale the image before passing to mayavi *create_figure* (True): create a new mayavi figure before rendering """ import mayavi from mayavi import mlab from tvtk.util.ctf import PiecewiseFunction, ColorTransferFunction if create_figure: fig = mlab.figure(size=(500,500),bgcolor=(0,0,0)) grid_data = sph.to_3d_grid(sim,qty=qty,nx=resolution, x2=None if width is None else width/2) if log: grid_data = np.log10(grid_data) if vmin is None: vmin = grid_data.max()-dynamic_range if vmax is None: vmax = grid_data.max() else: if vmin is None: vmin = np.min(grid_data) if vmax is None: vmax = np.max(grid_data) grid_data[grid_data<vmin]=vmin grid_data[grid_data>vmax]=vmax otf = PiecewiseFunction() otf.add_point(vmin,0.0) otf.add_point(vmax,1.0) sf = mayavi.tools.pipeline.scalar_field(grid_data) V = mlab.pipeline.volume(sf,color=color,vmin=vmin,vmax=vmax) V.trait_get('volume_mapper')['volume_mapper'].blend_mode = 'maximum_intensity' if color is None: ctf = ColorTransferFunction() ctf.add_rgb_point(vmin,107./255,124./255,132./255) ctf.add_rgb_point(vmin+(vmax-vmin)*0.8,200./255,178./255,164./255) ctf.add_rgb_point(vmin+(vmax-vmin)*0.9,1.0,210./255,149./255) ctf.add_rgb_point(vmax,1.0,222./255,141./255) V._volume_property.set_color(ctf) V._ctf = ctf V.update_ctf = True V._otf = otf V._volume_property.set_scalar_opacity(otf) return V
[docs]def contour(*args, **kwargs): """ Make an SPH image of the given simulation and render it as contours. nlevels and levels are passed to pyplot's contour command. Other arguments are as for *image*. """ import copy kwargs_image = copy.copy(kwargs) nlevels = kwargs_image.pop('nlevels',None) levels = kwargs_image.pop('levels',None) width = kwargs_image.get('width','10 kpc') kwargs_image['noplot']=True im = image(*args, **kwargs_image) res = im.shape units = kwargs_image.get('units',None) if isinstance(width, str) or issubclass(width.__class__, _units.UnitBase): if isinstance(width, str): width = _units.Unit(width) sim = args[0] width = width.in_units(sim['pos'].units, **sim.conversion_context()) width = float(width) x,y = np.meshgrid(np.linspace(-width/2,width/2,res[0]),np.linspace(-width/2,width/2,res[0])) p.contour(x,y,im,nlevels=nlevels,levels=levels)
def _units_imply_projection(sim, qty, units): try: sim[qty].units.ratio(units, **sim[qty].conversion_context()) # if this fails, perhaps we're requesting a projected image? return False except _units.UnitsException: # if the following fails, there's no interpretation this routine # can cope with. The error will be allowed to propagate. sim[qty].units.ratio( units / (sim['x'].units), **sim[qty].conversion_context()) return True
[docs]def image(sim, qty='rho', width="10 kpc", resolution=500, units=None, log=True, vmin=None, vmax=None, av_z=False, filename=None, z_camera=None, clear=True, cmap=None, title=None, qtytitle=None, show_cbar=True, subplot=False, noplot=False, ret_im=False, fill_nan=True, fill_val=0.0, linthresh=None, **kwargs): """ Make an SPH image of the given simulation. **Keyword arguments:** *qty* (rho): The name of the array to interpolate *width* (10 kpc): The overall width and height of the plot. If ``width`` is a float or an int, then it is assumed to be in units of ``sim['pos']``. It can also be passed in as a string indicating the units, i.e. '10 kpc', in which case it is converted to units of ``sim['pos']``. *resolution* (500): The number of pixels wide and tall *units* (None): The units of the output *av_z* (False): If True, the requested quantity is averaged down the line of sight (default False: image is generated in the thin plane z=0, unless output units imply an integral down the line of sight). If a string, the requested quantity is averaged down the line of sight weighted by the av_z array (e.g. use 'rho' for density-weighted quantity; the default results when av_z=True are volume-weighted). *z_camera* (None): If set, a perspective image is rendered. See :func:`pynbody.sph.image` for more details. *filename* (None): if set, the image will be saved in a file *clear* (True): whether to call clf() on the axes first *cmap* (None): user-supplied colormap instance *title* (None): plot title *qtytitle* (None): colorbar quantity title *show_cbar* (True): whether to plot the colorbar *subplot* (False): the user can supply a AxesSubPlot instance on which the image will be shown *noplot* (False): do not display the image, just return the image array *ret_im* (False): return the image instance returned by imshow *num_threads* (None) : if set, specify the number of threads for the multi-threaded routines; otherwise the pynbody.config default is used *fill_nan* (True): if any of the image values are NaN, replace with fill_val *fill_val* (0.0): the fill value to use when replacing NaNs *linthresh* (None): if the image has negative and positive values and a log scaling is requested, the part between `-linthresh` and `linthresh` is shown on a linear scale to avoid divergence at 0 """ if not noplot: import matplotlib.pylab as plt global config if not noplot: if subplot: p = subplot else: p = plt if qtytitle is None: qtytitle = qty if isinstance(units, str): units = _units.Unit(units) if isinstance(width, str) or issubclass(width.__class__, _units.UnitBase): if isinstance(width, str): width = _units.Unit(width) width = width.in_units(sim['pos'].units, **sim.conversion_context()) width = float(width) kernel = sph.Kernel() perspective = z_camera is not None if perspective and not av_z: kernel = sph.Kernel2D() is_projected = False if units is not None: is_projected = _units_imply_projection(sim, qty, units) if is_projected: kernel = sph.Kernel2D() if av_z: if isinstance(kernel, sph.Kernel2D): raise _units.UnitsException( "Units already imply projected image; can't also average over line-of-sight!") else: kernel = sph.Kernel2D() if units is not None: aunits = units * sim['z'].units else: aunits = None if isinstance(av_z, str): if units is not None: aunits = units * sim[av_z].units * sim['z'].units sim["__prod"] = sim[av_z] * sim[qty] qty = "__prod" else: av_z = "__one" sim["__one"] = np.ones_like(sim[qty]) sim["__one"].units = "1" im = sph.render_image(sim, qty, width / 2, resolution, out_units=aunits, kernel=kernel, z_camera=z_camera, **kwargs) im2 = sph.render_image(sim, av_z, width / 2, resolution, kernel=kernel, z_camera=z_camera, **kwargs) top = sim.ancestor try: del top["__one"] except KeyError: pass try: del top["__prod"] except KeyError: pass im = im / im2 else: im = sph.render_image(sim, qty, width / 2, resolution, out_units=units, kernel=kernel, z_camera=z_camera, **kwargs) if fill_nan: im[np.isnan(im)] = fill_val if not noplot: # set the log or linear normalizations if log: try: im[np.where(im == 0)] = abs(im[np.where(abs(im != 0))]).min() except ValueError: raise ValueError("Failed to make a sensible logarithmic image. This probably means there are no particles in the view.") # check if there are negative values -- if so, use the symmetric # log normalization if (vmin is None and (im < 0).any() ) or ((vmin is not None) and vmin<0): # need to set the linear regime around zero -- set to by # default start at 1/1000 of the log range if linthresh is None: linthresh = np.nanmax(abs(im)) / 1000. norm = matplotlib.colors.SymLogNorm( linthresh, vmin=vmin, vmax=vmax) else: norm = matplotlib.colors.LogNorm(vmin=vmin, vmax=vmax) else: norm = matplotlib.colors.Normalize(vmin=vmin, vmax=vmax) # # do the actual plotting # if clear and not subplot: p.clf() if ret_im: return p.imshow(im[::-1, :].view(np.ndarray), extent=(-width / 2, width / 2, -width / 2, width / 2), vmin=vmin, vmax=vmax, cmap=cmap, norm = norm) ims = p.imshow(im[::-1, :].view(np.ndarray), extent=(-width / 2, width / 2, -width / 2, width / 2), vmin=vmin, vmax=vmax, cmap=cmap, norm = norm) u_st = sim['pos'].units.latex() if not subplot: plt.xlabel("$x/%s$" % u_st) plt.ylabel("$y/%s$" % u_st) else: p.set_xlabel("$x/%s$" % u_st) p.set_ylabel("$y/%s$" % u_st) if units is None: units = im.units if units.latex() == "": units="" else: units = "$"+units.latex()+"$" if show_cbar: plt.colorbar(ims).set_label(qtytitle+"/"+units) # colorbar doesn't work wtih subplot: mappable is NoneType # elif show_cbar: # import matplotlib.pyplot as mpl # if qtytitle: mpl.colorbar().set_label(qtytitle) # else: mpl.colorbar().set_label(units) if title is not None: if not subplot: p.title(title) else: p.set_title(title) if filename is not None: p.savefig(filename) plt.draw() # plt.show() - removed by AP on 30/01/2013 - this should not be here as # for some systems you don't get back to the command prompt return im
def image_radial_profile(im, bins=100): xsize, ysize = np.shape(im) x = np.arange(-xsize / 2, xsize / 2) y = np.arange(-ysize / 2, ysize / 2) xs, ys = np.meshgrid(x, y) rs = np.sqrt(xs ** 2 + ys ** 2) hist, bin_edges = np.histogram(rs, bins=bins) inds = np.digitize(rs.flatten(), bin_edges) ave_vals = np.zeros(bin_edges.size) max_vals = np.zeros(bin_edges.size) min_vals = np.zeros(bin_edges.size) for i in np.arange(bin_edges.size): try: min_vals[i] = np.min(10 ** (im.flatten()[np.where(inds == i)])) except ValueError: min_vals[i] = float('nan') ave_vals[i] = np.mean(10 ** (im.flatten()[np.where(inds == i)])) try: max_vals[i] = np.max(10 ** (im.flatten()[np.where(inds == i)])) except ValueError: max_vals[i] = float('nan') return ave_vals, min_vals, max_vals, bin_edges