def __init__(self, shape=(600, 600)):
self.format = None
self.shape = shape
self.cache1 = {}
self.cache2 = {}
self.F = {}
self.cut = Cut()
self.periodic = False
def zoom(self, center=None, size=None):
if center is None:
center = self.cut.center
if size is None:
size = self.cut.size
self.cut.take(Cut(center=center, size=size))
self.invalidate()
def unfold(self, M):
self.gas.unfold(M, self.boxsize)
self.star.unfold(M, self.boxsize)
self.boxsize = self.bh.unfold(M, self.boxsize)
self.cut.take(
Cut(xcut=[0, self.boxsize[0]],
ycut=[0, self.boxsize[1]],
zcut=[0, self.boxsize[2]]))
def __init__(self, shape = (600,600)):
self.format = None
self.shape = shape
self.cache1 = {}
self.cache2 = {}
self.F = {}
self.cut = Cut()
self.periodic = False
def crop_snapshot(center, size, map, format, snapname, output):
cut = Cut(center=center, size=size)
mesh = Meshmap(map)
gas = Field(
components={
'mass': 'f4',
'rho': 'f4',
'ie': 'f4',
'xHI': 'f4',
'ye': 'f4',
'id': 'u8',
'sfr': 'f4',
'met': 'f4',
'sml': 'f4'
})
bh = Field(components={
'mass': 'f4',
'bhmdot': 'f4',
'bhmass': 'f4',
'id': 'u8'
})
star = Field(components={
'mass': 'f4',
'sft': 'f4',
'met': 'f4',
'id': 'u8'
})
fids = mesh.cut2fid(cut)
print fids
snapshots = [Snapshot(snapname % fid, format) for fid in fids]
gas.take_snapshots(snapshots, 0, cut=cut)
star.take_snapshots(snapshots, 4, cut=cut)
bh.take_snapshots(snapshots, 5, cut=cut)
left = center - size * 0.5
gas['locations'][:, :] -= left[newaxis, :]
star['locations'][:, :] -= left[newaxis, :]
bh['locations'][:, :] -= left[newaxis, :]
print bh['id']
out = Snapshot(output, format, create=True)
for f in snapshots[0].header.dtype.names:
out.header[f] = snapshots[0].header[f]
out.header['unused'][0:3] = center
out.header['boxsize'] = size[0]
gas.dump_snapshots([out], 0)
star.dump_snapshots([out], 4)
bh.dump_snapshots([out], 5)
def use(self,
snapname,
format,
components={},
bhcomponents={
'bhmass': 'f4',
'bhmdot': 'f4',
'id': 'u8'
},
starcomponents={
'sft': 'f4',
'mass': 'f4'
},
gas=0,
halo=1,
disk=2,
bulge=3,
star=4,
bh=5,
cut=None,
gascomponents=None,
periodic=True):
if gascomponents is not None:
self.components = gascomponents
else:
self.components = components
self.components['mass'] = 'f4'
self.components['sml'] = 'f4'
self.snapname = snapname
self.format = format
self.F['gas'] = Field(components=self.components)
self.F['bh'] = Field(components=bhcomponents)
self.F['star'] = Field(components=starcomponents)
self.F['halo'] = Field(components=self.components)
self.F['disk'] = Field(components=self.components)
self.F['bulge'] = Field(components=self.components)
self.ptype = {
"gas": gas,
"halo": halo,
"disk": disk,
"bulge": bulge,
"star": star,
"bh": bh,
}
try:
snapname = self.snapname % 0
except TypeError:
snapname = self.snapname
snap = Snapshot(snapname, self.format)
self.gas.init_from_snapshot(snap)
self.bh.init_from_snapshot(snap)
self.star.init_from_snapshot(snap)
self.halo.init_from_snapshot(snap)
self.disk.init_from_snapshot(snap)
self.bulge.init_from_snapshot(snap)
self.C = snap.C
if cut is not None:
self.cut.take(cut)
else:
try:
boxsize = snap.C['boxsize']
self.cut.take(
Cut(xcut=[0, boxsize],
ycut=[0, boxsize],
zcut=[0, boxsize]))
except:
pass
self.boxsize = ones(3) * snap.C['boxsize']
self.redshift = snap.C['redshift']
self.invalidate()
self.periodic = periodic
def slice(self, z, thickness):
cut = Cut(center=self.cut.center, size=self.cut.size)
cut['z'] = [z - thickness / 2.0, z + thickness / 2.0]
self.cut.take(cut)
self.invalidate()
class GaplotContext(object):
def __init__(self, shape=(600, 600)):
self.format = None
self.shape = shape
self.cache1 = {}
self.cache2 = {}
self.F = {}
self.cut = Cut()
self.periodic = False
def invalidate(self):
self.cache1.clear()
self.cache2.clear()
def zoom(self, center=None, size=None):
if center is None:
center = self.cut.center
if size is None:
size = self.cut.size
self.cut.take(Cut(center=center, size=size))
self.invalidate()
def autozoom(self, ftype):
if self.F[ftype].numpoints > 0:
min = self.F[ftype]['locations'].min(axis=0)
max = self.F[ftype]['locations'].max(axis=0)
self.zoom(center=(min + max) * 0.5, size=(max - min))
def slice(self, z, thickness):
cut = Cut(center=self.cut.center, size=self.cut.size)
cut['z'] = [z - thickness / 2.0, z + thickness / 2.0]
self.cut.take(cut)
self.invalidate()
def unfold(self, M):
self.gas.unfold(M, self.boxsize)
self.star.unfold(M, self.boxsize)
self.boxsize = self.bh.unfold(M, self.boxsize)
self.cut.take(
Cut(xcut=[0, self.boxsize[0]],
ycut=[0, self.boxsize[1]],
zcut=[0, self.boxsize[2]]))
@property
def pixel_area(self):
return (self.cut.size[0] * (self.cut.size[1] * 1.0) /
(self.shape[0] * self.shape[1]))
def use(self,
snapname,
format,
components={},
bhcomponents={
'bhmass': 'f4',
'bhmdot': 'f4',
'id': 'u8'
},
starcomponents={
'sft': 'f4',
'mass': 'f4'
},
gas=0,
halo=1,
disk=2,
bulge=3,
star=4,
bh=5,
cut=None,
gascomponents=None,
periodic=True):
if gascomponents is not None:
self.components = gascomponents
else:
self.components = components
self.components['mass'] = 'f4'
self.components['sml'] = 'f4'
self.snapname = snapname
self.format = format
self.F['gas'] = Field(components=self.components)
self.F['bh'] = Field(components=bhcomponents)
self.F['star'] = Field(components=starcomponents)
self.F['halo'] = Field(components=self.components)
self.F['disk'] = Field(components=self.components)
self.F['bulge'] = Field(components=self.components)
self.ptype = {
"gas": gas,
"halo": halo,
"disk": disk,
"bulge": bulge,
"star": star,
"bh": bh,
}
try:
snapname = self.snapname % 0
except TypeError:
snapname = self.snapname
snap = Snapshot(snapname, self.format)
self.gas.init_from_snapshot(snap)
self.bh.init_from_snapshot(snap)
self.star.init_from_snapshot(snap)
self.halo.init_from_snapshot(snap)
self.disk.init_from_snapshot(snap)
self.bulge.init_from_snapshot(snap)
self.C = snap.C
if cut is not None:
self.cut.take(cut)
else:
try:
boxsize = snap.C['boxsize']
self.cut.take(
Cut(xcut=[0, boxsize],
ycut=[0, boxsize],
zcut=[0, boxsize]))
except:
pass
self.boxsize = ones(3) * snap.C['boxsize']
self.redshift = snap.C['redshift']
self.invalidate()
self.periodic = periodic
@property
def extent(self):
return (self.cut['x'][0], self.cut['x'][1], self.cut['y'][0],
self.cut['y'][1])
@property
def halo(self):
return self.F['halo']
@halo.setter
def halo(self, value):
self.F['halo'] = value
print 'halo set'
self.invalidate()
@property
def bulge(self):
return self.F['bulge']
@bulge.setter
def bulge(self, value):
self.F['bulge'] = value
print 'bulge set'
self.invalidate()
@property
def disk(self):
return self.F['disk']
@disk.setter
def disk(self, value):
self.F['disk'] = value
print 'disk set'
self.invalidate()
@property
def gas(self):
return self.F['gas']
@gas.setter
def gas(self, value):
self.F['gas'] = value
print 'gas set'
self.invalidate()
@property
def bh(self):
return self.F['bh']
@bh.setter
def bh(self, value):
self.F['bh'] = value
print 'bh set'
self.invalidate()
@property
def star(self):
return self.F['star']
@star.setter
def star(self, value):
self.F['star'] = value
print 'star set'
self.invalidate()
def read(self,
fids=None,
use_gas=True,
use_bh=True,
use_star=True,
use_halo=False,
use_disk=False,
use_bulge=False,
numthreads=None):
if fids is not None:
snapnames = [self.snapname % i for i in fids]
elif '%d' in self.snapname:
snapnames = [self.snapname % i for i in range(self.C['Nfiles'])]
else:
snapnames = [self.snapname]
snapshots = [Snapshot(snapname, self.format) for snapname in snapnames]
if use_gas:
self.gas.take_snapshots(snapshots,
ptype=self.ptype['gas'],
nthreads=numthreads,
cut=self.cut)
if use_halo:
self.halo.take_snapshots(snapshots,
ptype=self.ptype['halo'],
nthreads=numthreads,
cut=self.cut)
if use_disk:
self.disk.take_snapshots(snapshots,
ptype=self.ptype['disk'],
nthreads=numthreads,
cut=self.cut)
if use_bulge:
self.bulge.take_snapshots(snapshots,
ptype=self.ptype['bulge'],
nthreads=numthreads,
cut=self.cut)
if use_bh:
self.bh.take_snapshots(snapshots,
ptype=self.ptype['bh'],
nthreads=numthreads,
cut=self.cut)
if use_star:
self.star.take_snapshots(snapshots,
ptype=self.ptype['star'],
nthreads=numthreads,
cut=self.cut)
self.invalidate()
def radial_mean(self,
component,
weightcomponent='mass',
bins=100,
min=None,
max=None,
std=False,
origin=None):
from numpy import histogram
if origin is None: origin = self.cut.center
d = self.gas.dist(origin=origin)
print 'radial_mean', origin
if min is not None and max is not None: range = (min, max)
else: range = None
if weightcomponent is not None:
m = self.gas[weightcomponent]
mx = m * self.gas[component]
mass, bins = histogram(d, range=range, bins=bins, weights=m)
value, bins = histogram(d, range=range, bins=bins, weights=mx)
value /= mass
if std:
mxx = mx * self.gas[component]
std, bins = histogram(d, range=range, bins=bins, weights=mxx)
std = sqrt(abs(std / mass - value**2))
return bins[:-1], value, std
return bins[:-1], value
else:
mx = self.gas[component]
value, bins = histogram(d, range=range, bins=bins, weights=mx)
value = value.cumsum()
r2, r1 = bins[1:], bins[:-1]
r = .5 * (r1 + r2)
V = 4 / 3.0 * 3.14 * (r**3)
value /= V
if std:
# possion error
std, bins = histogram(d, range=range, bins=bins)
return bins[:-1], value, value * sqrt(std) / std
return bins[:-1], value
def rotate(self, *args, **kwargs):
kwargs['origin'] = self.cut.center
self.gas.rotate(*args, **kwargs)
self.bh.rotate(*args, **kwargs)
self.star.rotate(*args, **kwargs)
self.invalidate()
def vector(self, ftype, component, grids=(20, 20), quick=True):
xs, xstep = linspace(self.cut['x'][0],
self.cut['x'][1],
grids[0],
endpoint=False,
retstep=True)
ys, ystep = linspace(self.cut['y'][0],
self.cut['y'][1],
grids[1],
endpoint=False,
retstep=True)
X, Y = meshgrid(xs + xstep / 2.0, ys + ystep / 2.0)
q = zeros(shape=(grids[0], grids[1], 3), dtype='f4')
mass = zeros(shape=(q.shape[0], q.shape[1]), dtype='f4')
raise NotImplemented('fix me')
print 'num particles rastered', self.mraster(q, ftype, mass, component,
quick)
q[:, :, 0] /= mass[:, :]
q[:, :, 1] /= mass[:, :]
q[:, :, 2] /= mass[:, :]
return X, Y, q.transpose((1, 0, 2))
def raster(self, ftype, component, quick=True):
""" ftype is amongst 'gas', 'star', those with has an sml """
field = self.F[ftype]
result = zeros(dtype='f4', shape=(self.shape[0], self.shape[1]))
ccode.rasterize(field,
result,
component,
xrange=self.cut[0],
yrange=self.cut[1],
zrange=self.cut[2],
quick=quick)
return result
def wraster(self, ftype, component, weightcomponent, quick=True):
field = self.F[ftype]
old = field[component].copy()
try:
if len(field[component].shape) == 1:
field[component][:] *= field[weightcomponent][:]
else:
field[component][:, 0] *= field[weightcomponent][:]
field[component][:, 1] *= field[weightcomponent][:]
field[component][:, 2] *= field[weightcomponent][:]
result = zeros(dtype='f4', shape=(self.shape[0], self.shape[1]))
mass = zeros(dtype='f4', shape=(self.shape[0], self.shape[1]))
ccode.rasterize(field, [mass, result],
[weightcomponent, component],
xrange=self.cut[0],
yrange=self.cut[1],
zrange=self.cut[2],
quick=quick)
return result, mass
except Exception as e:
raise e
finally:
field[component] = old
def linefield(self, ftype, component, src, dir, length=None):
""" return a line of sight of a field, pos, value """
f = self.F[ftype]
tree = f.tree
if length is None: length = self.boxsize[0]
pars = tree.trace(src, dir, length)
pos = f['locations'][pars, :]
sml = f['sml'][pars]
mass = f['mass'][pars]
comp = f[component][pars] * mass
if component == 'mass':
Lmass = zeros(shape=1024, dtype='f8')
ccode.scanline(locations=pos,
sml=sml,
targets=[Lmass],
values=[mass],
src=asarray(src),
dir=asarray(dir),
L=length)
return linspace(0, length, 1024), Lmass
else:
Larray = zeros(shape=1024, dtype='f8')
Lmass = zeros(shape=1024, dtype='f8')
ccode.scanline(locations=pos,
sml=sml,
targets=[Lmass, Larray],
values=[mass, comp],
src=asarray(src),
dir=asarray(dir),
L=length)
Larray /= Lmass
return linspace(0, length, 1024), Larray
def camera(self, ftype, component, camera, weightcomponent=None):
# from gaepsi.tools.simplegl import GLTrans
# t = GLTrans()
# t.lookat(target=target, pos=pos, up=up)
# t.perspective(near=near, far=far, up=up, pos=pos, fov=camera.fov / 360. * 3.1415, aspect=1.0 * raster[0].shape[0] / raster[0].shape[1])
f = self.F[ftype]
if weightcomponent is None:
sph = [float32(f[component])]
raster = [zeros(dtype='f8', shape=(self.shape[0], self.shape[1]))]
else:
m = float32(f[weightcomponent])
t = float32(f[component] * m)
sph = [m, t]
raster = [
zeros(dtype='f8', shape=(self.shape[0], self.shape[1])),
zeros(dtype='f8', shape=(self.shape[0], self.shape[1])),
]
if self.periodic: boxsize = asarray(self.boxsize)
else: boxsize = None
ccode.camera(raster=raster,
sph=sph,
locations=f['locations'],
sml=f['sml'],
near=camera.near,
far=camera.far,
Fov=camera.fov / 360. * 3.1415,
dim=asarray(raster[0].shape),
target=asarray(camera.target),
up=asarray(camera.up),
pos=asarray(camera.pos),
mask=None,
boxsize=boxsize)
if weightcomponent is None:
return raster[0], None
return raster[1], raster[0]
def projfield(self,
ftype,
component,
weightcomponent=None,
use_cache=True):
"""raster a field. ftype can be gas or star.
there are two types of components, depending if
weightcommponent is given.
If not given, assume component is 'intensive',
If given assume component is 'extensive'.
Extensive component:
After dividing a particle in two, an extensive quantity is divided by two.
Examples: mass, sfr.
Interpolation formula (of component A)
rho_A(x) = sum_i A_i W_i(x, x_i, h_i)
raster returns the integral of the density rho_A with in a pixel.
divided by the area of the pixel it becomes the column density of A.
the pixel integral of rho_A is saved in the cache, if cache enabled.
returns the pixel mean column density.
Intensive component:
After dividing a particle in two, an intensive quantity doesn't change.
Examples: T, xHI.
An intensive component has to be weighted by an extensive component.
For example, weighting the temperature T by total mass or xHI * mass.
Interpolation formula (of component A, weighted by M)
A(x) = sum_i A_i M_i / rho_M(x_i) W_i(x, x_i, h_i)
raster is invoked so that it returns an approximation of the
pixel integral of the M weighted A, as well as the pixel
integral of M. pixel mean is then obtained by dividing the two.
returns the pixel weighted integral per pixel, and the integrated weight per pixel.
"""
if use_cache:
if ftype not in self.cache1:
self.cache1[ftype] = {}
if ftype not in self.cache2:
self.cache2[ftype] = {}
if weightcomponent is None:
if use_cache and component in self.cache1[ftype]:
return self.cache1[ftype][component] / self.pixel_area, None
integral = self.raster(ftype, component, quick=False)
if use_cache:
self.cache1[ftype][component] = integral.copy()
integral /= self.pixel_area
columnden = integral
return columnden, None
else:
cache2name = component + weightcomponent
if use_cache and cache2name in self.cache2[
ftype] and weightcomponent in self.cache1[ftype]:
return self.cache2[ftype][cache2name].copy(
) / self.pixel_area, self.cache1[ftype][weightcomponent].copy(
) / self.pixel_area
integral, weight_int = self.wraster(ftype,
component,
weightcomponent,
quick=False)
if use_cache:
self.cache2[ftype][cache2name] = integral.copy()
self.cache1[ftype][weightcomponent] = weight_int.copy()
integral /= self.pixel_area
weight_int /= self.pixel_area
return integral, weight_int
def pointshow(self,
ax,
ftype,
component='bhmass',
radius=(4, 1),
logscale=True,
labelfmt=None,
labelcolor='white',
vmin=None,
vmax=None,
count=-1,
*args,
**kwargs):
if 'camera' in kwargs:
return self.bhshow_camera(ax=ax,
component=component,
radius=radius,
logscale=logscale,
vmin=vmin,
vmax=vmax,
count=count,
*args,
**kwargs)
from matplotlib.collections import CircleCollection
mask = self.cut.select(self.F[ftype]['locations'])
X = self.F[ftype]['locations'][mask, 0]
Y = self.F[ftype]['locations'][mask, 1]
bhmass = self.F[ftype][component][mask]
if bhmass.size == 0: return
R = bhmass
if count > 0:
ind = (-R).argsort()
X = X[ind[0:count]]
Y = Y[ind[0:count]]
R = R[ind[0:count]]
if vmax is None:
vmax = R.max()
if vmin is None:
vmin = ccode.pmin.reduce(R)
print 'bhshow, vmax, vmin =', vmax, vmin
if logscale:
R = log10(R)
Nm = Normalize(vmax=log10(vmax), vmin=log10(vmin), clip=True)
else:
Nm = Normalize(vmax=vmax, vmin=vmin, clip=True)
if bhmass.size > 1:
R = Nm(R)
if R.min() == R.max():
R = ones(R.shape)
if isscalar(radius):
R *= radius
else:
rmin = min(radius)
rmax = max(radius)
R *= (rmax - rmin)
R += rmin
print 'R max, R min', R.min(), R.max()
print R
# R.clip(min=4, max=radius**2)
if not 'marker' in kwargs:
use_spike = True
else:
use_spike = False
if not 'edgecolor' in kwargs:
kwargs['edgecolor'] = 'green'
if not 'facecolor' in kwargs:
kwargs['facecolor'] = (0, 1, 0, 0.0)
if use_spike:
c = SpikeCollection(X, Y, R, **kwargs)
ax.add_collection(c)
else:
ax.scatter(X, Y, s=R**2, **kwargs)
if labelfmt:
if count > 0:
ID = ID[ind[0:count]]
ID = self.F[type]['id'][mask]
for x, y, id in zip(X, Y, ID):
rat = random.random() * 360
if rat > 90 and rat <= 180:
trat = rat + 180
elif rat > 180 and rat <= 270:
trat = rat - 180
else:
trat = rat
if rat < 315 and rat > 135:
dir = 1
rat -= 180
else:
dir = 0
ax.text(x,
y,
labelfmt % id,
withdash=True,
dashlength=radius * 10,
dashpush=radius,
dashdirection=dir,
rotation=trat,
dashrotation=rat,
color=labelcolor)
def bhshow(self,
ax,
component='bhmass',
radius=(4, 1),
logscale=True,
labelfmt=None,
labelcolor='white',
vmin=None,
vmax=None,
count=-1,
*args,
**kwargs):
return self.pointshow(ax, 'bh', component, radius, logscale, labelfmt,
labelcolor, vmin, vmax, count, *args, **kwargs)
# col = CircleCollection(offsets=zip(X.flat,Y.flat), sizes=(R * radius)**2, edgecolor='green', facecolor='none', transOffset=gca().transData)
# ax.add_collection(col)
def bhshow_camera(self,
ax,
component='bhmdot',
radius=4,
logscale=True,
vmin=None,
vmax=None,
count=-1,
camera=None,
*args,
**kwargs):
bhlum = self.bh.cosmology.QSObol(self.bh[component], 'blue')
good = bhlum > 0
if len(bhlum) > 0 and good.sum() > 0:
bhlum0 = bhlum[good]
bhpos0 = self.bh['locations'][good]
if self.periodic:
bhpos = zeros((27, bhpos0.shape[0], 3))
bhlum = zeros((27, bhlum0.shape[0]))
for im in range(27):
i, j, k = im / 9 - 1, (im % 9) / 3 - 1, im % 3 - 1
bhpos[
im, :, :] = bhpos0[:, :] + array([i, j, k]) * self.boxsize
bhlum[im, :] = bhlum0[:]
bhpos.shape = -1, 3
bhlum.shape = -1
else:
bhpos = bhpos0
bhlum = bhlum0
bhpos = camera.map(bhpos, self.shape)
bhdist = bhpos[:, 2]
bhlum /= 4 * 3.1416 * bhdist**2
bhlum /= (self.bh.cosmology.units.W / self.bh.cosmology.units.METER**2)
print 'distance of bhs', bhdist.min(), bhdist.max()
print 'lum of bhs', bhlum.min(), bhlum.max()
# apparent maginitute
mag = ((-26.74 + 2.5 * log10(1400 / bhlum)))
ind = mag.argsort()
if count > -1 and count < len(mag):
magcut = mag[ind[count]]
else:
magcut = 9999
# 8 is the naked eye limit
selected = (bhpos[:, 0] > 0) & (bhpos[:, 0] < self.shape[0])
selected &= (bhpos[:, 1] > 0) & (bhpos[:, 1] < self.shape[1])
selected &= (bhpos[:, 2] > 0.0) & (mag < magcut)
print 'mag', mag.min(), mag.max()
bhpos = bhpos[selected]
bhlum = bhlum[selected]
print 'bhs shown', selected.sum()
mag = -mag[selected]
mag -= mag.min()
diff = mag.max()
if diff > 0:
mag /= diff
else:
mag[:] = 1
bhx = bhpos[:, 0]
bhy = bhpos[:, 1]
print mag
t = 0.05
marker = ((-1, 0), (-t, t), (0, 1), (t, t), (1, 0), (t, -t), (0, -1),
(-t, -t)), 0
color = (0.8, 1, 1, 0.7)
ecolor = (1., 1, 0.8, 0.2)
if isscalar(radius):
mag *= radius
else:
rmin = min(radius)
rmax = max(radius)
mag *= (rmax - rmin)
mag += rmin
c = SpikeCollection(bhx, bhy, radius=mag, color=color)
ax.add_collection(c)
# ax.scatter(x=bhx, y=bhy, s=mag * 10**2, marker=marker, edgecolor=ecolor, color=color)
def velshow(self, ax, ftype, relative=False, color='cyan', alpha=0.8):
X, Y, vel = self.vector(ftype, 'vel', grids=(20, 20), quick=False)
field = self.F[ftype]
mean = None
if not type(relative) == bool or relative == True:
if type(relative) == bool:
mean = average(field['vel'], weights=field['mass'],
axis=0)[0:2]
else:
mean = asarray(relative)
vel[:, :, 0] -= mean[0]
vel[:, :, 1] -= mean[1]
print 'mean = ', mean
print 'max component of velocity', abs(vel).max()
sorted = absolute(vel.ravel())
sorted.sort()
scale = sorted[sorted.size * 9 / 10] * 20
print X.shape, vel[:, :, 0].shape
ax.quiver(X,
Y,
vel[:, :, 0],
vel[:, :, 1],
width=0.003,
scale=scale,
scale_units='width',
angles='xy',
color=color,
alpha=alpha)
if mean != None:
ax.quiver(self.cut.center[0],
self.cut.center[1],
mean[0],
mean[1],
scale=scale,
scale_units='width',
width=0.01,
angles='xy',
color=(0, 0, 0, 0.5))
def starshow_poor(self, ax, *args, **kwargs):
star = self.star
mask = self.cut.select(star['locations'])
X = star['locations'][mask, 0]
Y = star['locations'][mask, 1]
if not 'color' in kwargs:
kwargs['color'] = 'white'
if not 'alpha' in kwargs:
kwargs['alpha'] = 0.5
ax.plot(X, Y, ', ', *args, **kwargs)
def reset_view(self, ax, camera=False):
if camera:
left, right = 0, self.shape[0]
bottom, top = 0, self.shape[1]
else:
left, right = self.cut['x']
bottom, top = self.cut['y']
print left, right, bottom, top
ax.set_xlim(left, right)
ax.set_ylim(bottom, top)
def makeT(self, Xh=0.76):
"""T will be in Kelvin"""
gas = self.gas
C = gas.cosmology
gas['T'] = zeros(dtype='f4', shape=gas.numpoints)
C.ie2T(ie=gas['ie'], ye=gas['ye'], Xh=Xh, out=gas['T'])
gas['T'] *= C.units.TEMPERATURE
def mergeBHs(self, threshold=1.0):
bh = self.bh
if bh.numpoints == 0: return
pos = bh['locations']
posI = tile(pos, pos.shape[0]).reshape(pos.shape[0], pos.shape[0], 3)
posJ = posI.transpose((1, 0, 2))
d = posI - posJ
d = sqrt((d**2).sum(axis=2))
mask = d < threshold
mergeinto = zeros(bh.numpoints, dtype='i4')
for i in range(bh.numpoints):
mergeinto[i] = nonzero(mask[i])[0][0]
g, ind = unique(mergeinto, return_index=True)
comp = 'bhmass'
bak = bh[comp].copy()
for i in range(bh.numpoints):
bh[comp][i] = bak[mask[i]].sum()
try:
comp = 'bhmdot'
bak = bh[comp].copy()
for i in range(bh.numpoints):
bh[comp][i] = bak[mask[i]].sum()
except KeyError:
pass
comp = 'locations'
#def sel(comp):
bak = bh[comp].copy()
for i in range(bh.numpoints):
bh[comp][i] = bak[mask[i]][0]
bh.numpoints = len(ind)
for comp in bh.names:
bh[comp] = bh[comp][ind]
def fieldshow(self,
ax,
ftype,
component,
mode=None,
camera=None,
vmin=None,
vmax=None,
logscale=True,
cmap=pygascmap,
gamma=1.0,
mmin=None,
mmax=None,
over=None,
return_raster=False,
logweight=None,
weightcomponent='mass',
use_cache=True):
if camera is None:
if mode is None: # extensive, mass->density, sfr->sf density,
todraw, mass = self.projfield(ftype=ftype,
component=component,
use_cache=use_cache)
else: # intensive, decorate the result aftwards
todraw, mass = self.projfield(ftype=ftype,
component=component,
weightcomponent=weightcomponent,
use_cache=use_cache)
else:
if mode is None: # extensive, mass->density, sfr->sf density,
todraw, mass = self.camera(ftype=ftype,
component=component,
camera=camera)
else: # intensive, decorate the result aftwards
todraw, mass = self.camera(ftype=ftype,
component=component,
weightcomponent=weightcomponent,
camera=camera)
if camera is None:
if logweight is None: logweight = True
image = self.render(todraw=todraw,
mass=mass,
mode=mode,
vmin=vmin,
vmax=vmax,
logscale=logscale,
logweight=logweight,
cmap=cmap,
gamma=gamma,
mmin=mmin,
mmax=mmax,
over=over)
ret = ax.imshow(image.transpose((1, 0, 2)),
origin='lower',
extent=self.extent,
vmin=vmin,
vmax=vmax,
cmap=cmap)
else:
if logweight is None: logweight = False
image = self.render(todraw=todraw,
mass=mass,
mode=mode,
vmin=vmin,
vmax=vmax,
logscale=logscale,
logweight=logweight,
cmap=cmap,
gamma=gamma,
mmin=mmin,
mmax=mmax,
over=over)
ret = ax.imshow(image.transpose((1, 0, 2)),
origin='lower',
extent=(0, self.shape[0], 0, self.shape[1]),
vmin=vmin,
vmax=vmax,
cmap=cmap)
return ret
def fieldcontour(self,
ax,
ftype,
component,
camera=None,
levels=None,
logweight=None,
weightcomponent='mass',
use_cache=True,
colors='k',
linewidth=2,
linestyle=['solid']):
if camera is None:
todraw, mass = self.projfield(ftype=ftype,
component=component,
weightcomponent=weightcomponent,
use_cache=use_cache)
else:
todraw, mass = self.camera(ftype=ftype,
component=component,
weightcomponent=weightcomponent,
camera=camera)
if camera is None:
if logweight is None: logweight = True
ret = ax.contourf((todraw / mass).T,
extent=self.extent,
colors=colors,
linewidth=linewidth,
levels=levels,
linestyle=linestyle)
else:
if logweight is None: logweight = False
ret = ax.contourf((todraw / mass).T,
extent=(0, self.shape[0], 0, self.shape[1]),
colors=colors,
linewidth=linewidth,
levels=levels,
linestyle=linestyle)
return ret
def render(self,
todraw,
mass,
mode=None,
vmin=None,
vmax=None,
logscale=True,
logweight=True,
cmap=pygascmap,
gamma=1.0,
mmin=None,
mmax=None,
over=None,
composite=None):
""" vmin can be a string, eg '40 dB', so the vmin=vmax/1e4 """
if mode == 'intensity' or mode == 'mean':
todraw /= mass
if logscale:
if vmin is not None and not isinstance(vmin, basestring):
vmin = 10**vmin
if vmax is not None:
vmax = 10**vmax
if vmax is None:
if mode == 'intensity' or mode == 'mean' or over is None:
vmax = fmax.reduce(todraw.flat)
else:
sort = todraw.ravel().argsort()
ind = sort[(len(sort) - 1) * (1.0 - over)]
vmax = todraw.ravel()[ind]
del sort
if vmin is None:
if logscale:
vmin = ccode.pmin.reduce(todraw.flat)
else:
vmin = fmin.reduce(todraw.flat)
elif isinstance(vmin, basestring):
db = int(vmin[:-3])
vmin = vmax / 10**(db / 10.)
if logscale:
log10(todraw, todraw)
vmin = log10(vmin)
vmax = log10(vmax)
print 'vmax, vmin =', vmax, vmin
image = zeros(dtype=('u1', 4), shape=todraw.shape)
# gacolor is much faster then matplotlib's normalize and uses less memory(4times fewer).
if not hasattr(cmap, 'table'):
cmap = gacmap(cmap)
ccode.color(image,
todraw,
max=vmax,
min=vmin,
logscale=False,
colormap=cmap)
del todraw
if mode == 'intensity':
weight = mass
if logweight:
# in camera mode do not use logscale for mass.
weight.clip(ccode.pmin.reduce(weight.ravel()), inf, weight)
log10(weight, weight)
print 'weight', weight.mean(), weight.max(), weight.min()
if mmin is None:
mmin = weight.min()
if mmax is None:
if over is None:
mmax = weight.max()
else:
sort = weight.ravel().argsort()
ind = sort[(len(sort) - 1) * (1.0 - over)]
mmax = weight.ravel()[ind]
print 'mmax, mmin', mmax, mmin
weight -= mmin
weight /= (mmax - mmin)
weight.clip(0, 1, weight)
else:
weight = image[:, :, 3] / 255.0
weight **= gamma
if composite is not None:
alpha = composite[:, :, 3] / 255.0
multiply(1.0 - weight[:, :], alpha[:, :], alpha[:, :])
multiply(image[:, :, 0:3], weight[:, :, newaxis], image[:, :, 0:3])
multiply(composite[:, :, 0:3], alpha[:, :, newaxis],
composite[:, :, 0:3])
add(image[:, :, 0:3], composite[:, :, 0:3], image[:, :, 0:3])
image[:, :, 3] = (alpha[:, :] + weight[:, :]) * 255
else:
#multiply(image[:, :, 0:3], weight[:, :, newaxis], image[:, :, 0:3])
#image[:, :, 3] = 255
#weight **= 0.33333333333
multiply(255.9999,
weight[:, :],
out=image[:, :, 3],
casting='unsafe')
# print 'alpha', image[:, :, 3].ravel().min(), image[:,:,3].ravel().max()
return image
def decorate(self,
ax,
frameon=True,
title=None,
bgcolor='k',
color='w',
fontsize='small'):
cut = self.cut
ax.set_axis_bgcolor(bgcolor)
if frameon:
ax.ticklabel_format(axis='x', useOffset=cut.center[0])
ax.ticklabel_format(axis='y', useOffset=cut.center[1])
ax.set_xticks(linspace(cut['x'][0], cut['x'][1], 5))
ax.set_yticks(linspace(cut['y'][0], cut['y'][1], 5))
#ax.set_title(title)
if (title is not None):
ax.set_title(title,
position=(0.1, 0.9),
fontsize=fontsize,
color=color,
transform=ax.transAxes)
else:
ax.axison = False
if (title is not None):
ax.set_title(title,
position=(0.1, 0.9),
fontsize=fontsize,
color=color,
transform=ax.transAxes)
ax.figure.set_facecolor(bgcolor)
def drawscale(self, ax, color='white', fontsize=None):
from mpl_toolkits.axes_grid.anchored_artists import AnchoredSizeBar
l = (self.cut.size[0]) * 0.2
l = l // 10**int(log10(l)) * 10**int(log10(l))
if l > 500:
l /= 1000.0
l = int(l + 0.5)
text = r"%g Mpc/h" % l
l *= 1000.0
else:
text = r"%g Kpc/h" % l
print text
b = AnchoredSizeBar(ax.transData,
l,
text,
loc=8,
pad=0.1,
borderpad=0.5,
sep=5,
frameon=False)
for r in b.size_bar.findobj(Rectangle):
r.set_edgecolor(color)
for t in b.txt_label.findobj(Text):
t.set_color(color)
if fontsize is not None:
t.set_fontsize(fontsize)
ax.add_artist(b)
def circle(self, ax, *args, **kwargs):
from matplotlib.patches import Circle
c = Circle(*args, **kwargs)
ax.add_patch(c)
class GaplotContext(object):
def __init__(self, shape = (600,600)):
self.format = None
self.shape = shape
self.cache1 = {}
self.cache2 = {}
self.F = {}
self.cut = Cut()
self.periodic = False
def invalidate(self):
self.cache1.clear()
self.cache2.clear()
def zoom(self, center=None, size=None):
if center is None:
center = self.cut.center
if size is None:
size = self.cut.size
self.cut.take(Cut(center=center, size=size))
self.invalidate()
def autozoom(self, ftype):
if self.F[ftype].numpoints > 0:
min=self.F[ftype]['locations'].min(axis=0)
max=self.F[ftype]['locations'].max(axis=0)
self.zoom(center=(min+max) * 0.5, size=(max - min))
def slice(self, z, thickness):
cut = Cut(center=self.cut.center, size=self.cut.size)
cut['z'] = [z - thickness / 2.0, z + thickness / 2.0]
self.cut.take(cut)
self.invalidate()
def unfold(self, M):
self.gas.unfold(M, self.boxsize)
self.star.unfold(M, self.boxsize)
self.boxsize = self.bh.unfold(M, self.boxsize)
self.cut.take(Cut(xcut=[0, self.boxsize[0]],
ycut=[0, self.boxsize[1]],
zcut=[0, self.boxsize[2]]))
@property
def pixel_area(self):
return (self.cut.size[0] * (self.cut.size[1] * 1.0)/(self.shape[0] * self.shape[1]))
def use(self, snapname, format, components={},
bhcomponents={'bhmass':'f4', 'bhmdot':'f4', 'id':'u8'},
starcomponents={'sft':'f4', 'mass':'f4'}, gas=0, halo=1, disk=2, bulge=3, star=4, bh=5, cut=None, gascomponents=None, periodic=True):
if gascomponents is not None:
self.components = gascomponents
else:
self.components = components
self.components['mass'] = 'f4'
self.components['sml'] = 'f4'
self.snapname = snapname
self.format = format
self.F['gas'] = Field(components=self.components)
self.F['bh'] = Field(components=bhcomponents)
self.F['star'] = Field(components=starcomponents)
self.F['halo'] = Field(components=self.components)
self.F['disk'] = Field(components=self.components)
self.F['bulge'] = Field(components=self.components)
self.ptype = {
"gas": gas,
"halo": halo,
"disk": disk,
"bulge": bulge,
"star": star,
"bh": bh,
}
try:
snapname = self.snapname % 0
except TypeError:
snapname = self.snapname
snap = Snapshot(snapname, self.format)
self.gas.init_from_snapshot(snap)
self.bh.init_from_snapshot(snap)
self.star.init_from_snapshot(snap)
self.halo.init_from_snapshot(snap)
self.disk.init_from_snapshot(snap)
self.bulge.init_from_snapshot(snap)
self.C = snap.C
if cut is not None:
self.cut.take(cut)
else:
try:
boxsize = snap.C['boxsize']
self.cut.take(Cut(xcut=[0, boxsize], ycut=[0, boxsize], zcut=[0, boxsize]))
except:
pass
self.boxsize = ones(3) * snap.C['boxsize']
self.redshift = snap.C['redshift']
self.invalidate()
self.periodic = periodic
@property
def extent(self):
return (self.cut['x'][0], self.cut['x'][1], self.cut['y'][0], self.cut['y'][1])
@property
def halo(self):
return self.F['halo']
@halo.setter
def halo(self, value):
self.F['halo'] = value
print 'halo set'
self.invalidate()
@property
def bulge(self):
return self.F['bulge']
@bulge.setter
def bulge(self, value):
self.F['bulge'] = value
print 'bulge set'
self.invalidate()
@property
def disk(self):
return self.F['disk']
@disk.setter
def disk(self, value):
self.F['disk'] = value
print 'disk set'
self.invalidate()
@property
def gas(self):
return self.F['gas']
@gas.setter
def gas(self, value):
self.F['gas'] = value
print 'gas set'
self.invalidate()
@property
def bh(self):
return self.F['bh']
@bh.setter
def bh(self, value):
self.F['bh'] = value
print 'bh set'
self.invalidate()
@property
def star(self):
return self.F['star']
@star.setter
def star(self, value):
self.F['star'] = value
print 'star set'
self.invalidate()
def read(self, fids=None, use_gas=True, use_bh=True, use_star=True, use_halo=False, use_disk=False, use_bulge=False, numthreads=None):
if fids is not None:
snapnames = [self.snapname % i for i in fids]
elif '%d' in self.snapname:
snapnames = [self.snapname % i for i in range(self.C['Nfiles'])]
else:
snapnames = [self.snapname]
snapshots = [Snapshot(snapname, self.format) for snapname in snapnames]
if use_gas:
self.gas.take_snapshots(snapshots, ptype = self.ptype['gas'], nthreads=numthreads, cut=self.cut)
if use_halo:
self.halo.take_snapshots(snapshots, ptype = self.ptype['halo'], nthreads=numthreads, cut=self.cut)
if use_disk:
self.disk.take_snapshots(snapshots, ptype = self.ptype['disk'], nthreads=numthreads, cut=self.cut)
if use_bulge:
self.bulge.take_snapshots(snapshots, ptype = self.ptype['bulge'], nthreads=numthreads, cut=self.cut)
if use_bh:
self.bh.take_snapshots(snapshots, ptype = self.ptype['bh'], nthreads=numthreads, cut=self.cut)
if use_star:
self.star.take_snapshots(snapshots, ptype = self.ptype['star'], nthreads=numthreads, cut=self.cut)
self.invalidate()
def radial_mean(self, component, weightcomponent='mass', bins=100, min=None, max=None, std=False, origin=None):
from numpy import histogram
if origin is None: origin = self.cut.center
d = self.gas.dist(origin=origin)
print 'radial_mean', origin
if min is not None and max is not None: range=(min, max)
else: range= None
if weightcomponent is not None:
m = self.gas[weightcomponent]
mx = m * self.gas[component]
mass, bins = histogram(d, range=range, bins=bins, weights=m)
value, bins = histogram(d, range=range, bins=bins, weights=mx)
value /= mass
if std:
mxx = mx * self.gas[component]
std, bins = histogram(d, range=range, bins=bins, weights=mxx)
std = sqrt(abs(std / mass - value ** 2))
return bins[:-1], value, std
return bins[:-1], value
else:
mx = self.gas[component]
value, bins = histogram(d, range=range, bins=bins, weights=mx)
value = value.cumsum()
r2,r1 = bins[1:], bins[:-1]
r = .5 * (r1 + r2)
V = 4 / 3.0 * 3.14 * (r ** 3)
value /= V
if std:
# possion error
std, bins = histogram(d, range=range, bins=bins)
return bins[:-1], value, value * sqrt(std) / std
return bins[:-1], value
def rotate(self, *args, **kwargs):
kwargs['origin'] = self.cut.center
self.gas.rotate(*args, **kwargs)
self.bh.rotate(*args, **kwargs)
self.star.rotate(*args, **kwargs)
self.invalidate()
def vector(self, ftype, component, grids=(20,20), quick=True):
xs,xstep = linspace(self.cut['x'][0], self.cut['x'][1], grids[0], endpoint=False,retstep=True)
ys,ystep = linspace(self.cut['y'][0], self.cut['y'][1], grids[1], endpoint=False,retstep=True)
X,Y=meshgrid(xs+ xstep/2.0,ys+ystep/2.0)
q = zeros(shape=(grids[0],grids[1],3), dtype='f4')
mass = zeros(shape = (q.shape[0], q.shape[1]), dtype='f4')
raise NotImplemented('fix me')
print 'num particles rastered', self.mraster(q, ftype, mass, component, quick)
q[:,:,0]/=mass[:,:]
q[:,:,1]/=mass[:,:]
q[:,:,2]/=mass[:,:]
return X,Y,q.transpose((1,0,2))
def raster(self, ftype, component, quick=True):
""" ftype is amongst 'gas', 'star', those with has an sml """
field = self.F[ftype]
result = zeros(dtype='f4', shape = (self.shape[0], self.shape[1]))
ccode.rasterize(field, result, component, xrange=self.cut[0], yrange=self.cut[1], zrange=self.cut[2], quick=quick)
return result
def wraster(self, ftype, component, weightcomponent, quick=True):
field = self.F[ftype]
old = field[component].copy()
try:
if len(field[component].shape) == 1:
field[component][:] *= field[weightcomponent][:]
else:
field[component][:,0] *= field[weightcomponent][:]
field[component][:,1] *= field[weightcomponent][:]
field[component][:,2] *= field[weightcomponent][:]
result = zeros(dtype='f4', shape = (self.shape[0], self.shape[1]))
mass = zeros(dtype='f4', shape = (self.shape[0], self.shape[1]))
ccode.rasterize(field, [mass, result], [weightcomponent, component], xrange=self.cut[0], yrange=self.cut[1], zrange=self.cut[2], quick=quick)
return result, mass
except Exception as e:
raise e
finally:
field[component] = old
def linefield(self, ftype, component, src, dir, length=None):
""" return a line of sight of a field, pos, value """
f = self.F[ftype]
tree = f.tree
if length is None: length = self.boxsize[0]
pars = tree.trace(src, dir, length)
pos = f['locations'][pars, :]
sml = f['sml'][pars]
mass = f['mass'][pars]
comp = f[component][pars] * mass
if component == 'mass':
Lmass = zeros(shape=1024, dtype='f8')
ccode.scanline(locations = pos, sml = sml, targets = [Lmass], values = [mass], src=asarray(src), dir=asarray(dir), L=length)
return linspace(0, length, 1024), Lmass
else:
Larray = zeros(shape=1024, dtype='f8')
Lmass = zeros(shape=1024, dtype='f8')
ccode.scanline(locations = pos, sml = sml, targets = [Lmass, Larray], values = [mass, comp], src=asarray(src), dir=asarray(dir), L=length)
Larray /= Lmass
return linspace(0, length, 1024), Larray
def camera(self, ftype, component, camera, weightcomponent=None):
# from gaepsi.tools.simplegl import GLTrans
# t = GLTrans()
# t.lookat(target=target, pos=pos, up=up)
# t.perspective(near=near, far=far, up=up, pos=pos, fov=camera.fov / 360. * 3.1415, aspect=1.0 * raster[0].shape[0] / raster[0].shape[1])
f = self.F[ftype]
if weightcomponent is None:
sph = [float32(f[component])]
raster = [zeros(dtype='f8', shape = (self.shape[0], self.shape[1]))]
else:
m = float32(f[weightcomponent])
t = float32(f[component] * m)
sph = [m, t]
raster = [ zeros(dtype='f8', shape = (self.shape[0], self.shape[1])),
zeros(dtype='f8', shape = (self.shape[0], self.shape[1])), ]
if self.periodic: boxsize = asarray(self.boxsize)
else: boxsize = None
ccode.camera(raster=raster, sph=sph, locations=f['locations'],
sml=f['sml'], near=camera.near, far=camera.far, Fov=camera.fov / 360. * 3.1415,
dim=asarray(raster[0].shape), target=asarray(camera.target), up = asarray(camera.up),
pos=asarray(camera.pos), mask=None, boxsize=boxsize)
if weightcomponent is None:
return raster[0], None
return raster[1], raster[0]
def projfield(self, ftype, component, weightcomponent=None, use_cache=True):
"""raster a field. ftype can be gas or star.
there are two types of components, depending if
weightcommponent is given.
If not given, assume component is 'intensive',
If given assume component is 'extensive'.
Extensive component:
After dividing a particle in two, an extensive quantity is divided by two.
Examples: mass, sfr.
Interpolation formula (of component A)
rho_A(x) = sum_i A_i W_i(x, x_i, h_i)
raster returns the integral of the density rho_A with in a pixel.
divided by the area of the pixel it becomes the column density of A.
the pixel integral of rho_A is saved in the cache, if cache enabled.
returns the pixel mean column density.
Intensive component:
After dividing a particle in two, an intensive quantity doesn't change.
Examples: T, xHI.
An intensive component has to be weighted by an extensive component.
For example, weighting the temperature T by total mass or xHI * mass.
Interpolation formula (of component A, weighted by M)
A(x) = sum_i A_i M_i / rho_M(x_i) W_i(x, x_i, h_i)
raster is invoked so that it returns an approximation of the
pixel integral of the M weighted A, as well as the pixel
integral of M. pixel mean is then obtained by dividing the two.
returns the pixel weighted integral per pixel, and the integrated weight per pixel.
"""
if use_cache :
if ftype not in self.cache1:
self.cache1[ftype] = {}
if ftype not in self.cache2:
self.cache2[ftype] = {}
if weightcomponent is None:
if use_cache and component in self.cache1[ftype]:
return self.cache1[ftype][component] / self.pixel_area, None
integral = self.raster(ftype, component, quick=False)
if use_cache:
self.cache1[ftype][component] = integral.copy()
integral /= self.pixel_area
columnden = integral
return columnden, None
else:
cache2name = component + weightcomponent
if use_cache and cache2name in self.cache2[ftype] and weightcomponent in self.cache1[ftype]:
return self.cache2[ftype][cache2name].copy() / self.pixel_area, self.cache1[ftype][weightcomponent].copy() / self.pixel_area
integral, weight_int = self.wraster(ftype, component, weightcomponent, quick=False)
if use_cache :
self.cache2[ftype][cache2name] = integral.copy()
self.cache1[ftype][weightcomponent] = weight_int.copy()
integral /= self.pixel_area
weight_int /= self.pixel_area
return integral, weight_int
def pointshow(self, ax, ftype, component='bhmass', radius=(4, 1), logscale=True, labelfmt=None, labelcolor='white', vmin=None, vmax=None, count=-1, *args, **kwargs):
if 'camera' in kwargs:
return self.bhshow_camera(ax=ax, component=component, radius=radius, logscale=logscale, vmin=vmin, vmax=vmax, count=count, *args, **kwargs)
from matplotlib.collections import CircleCollection
mask = self.cut.select(self.F[ftype]['locations'])
X = self.F[ftype]['locations'][mask,0]
Y = self.F[ftype]['locations'][mask,1]
bhmass = self.F[ftype][component][mask]
if bhmass.size == 0: return
R = bhmass
if count > 0:
ind = (-R).argsort()
X = X[ind[0:count]]
Y = Y[ind[0:count]]
R = R[ind[0:count]]
if vmax is None:
vmax = R.max()
if vmin is None:
vmin = ccode.pmin.reduce(R)
print 'bhshow, vmax, vmin =', vmax, vmin
if logscale:
R = log10(R)
Nm = Normalize(vmax=log10(vmax), vmin=log10(vmin), clip=True)
else:
Nm = Normalize(vmax=vmax, vmin=vmin, clip=True)
if bhmass.size > 1:
R = Nm(R)
if R.min() == R.max():
R = ones(R.shape)
if isscalar(radius):
R*=radius
else:
rmin = min(radius)
rmax = max(radius)
R *= (rmax - rmin)
R += rmin
print 'R max, R min', R.min(), R.max()
print R
# R.clip(min=4, max=radius**2)
if not 'marker' in kwargs:
use_spike = True
else:
use_spike = False
if not 'edgecolor' in kwargs:
kwargs['edgecolor'] = 'green'
if not 'facecolor' in kwargs:
kwargs['facecolor'] = (0, 1, 0, 0.0)
if use_spike:
c = SpikeCollection(X, Y, R, **kwargs)
ax.add_collection(c)
else:
ax.scatter(X,Y, s=R**2, **kwargs)
if labelfmt:
if count > 0:
ID = ID[ind[0:count]]
ID = self.F[type]['id'][mask]
for x,y,id in zip(X,Y,ID):
rat = random.random() * 360
if rat > 90 and rat <= 180:
trat = rat + 180
elif rat > 180 and rat <=270:
trat = rat - 180
else:
trat = rat
if rat < 315 and rat > 135:
dir = 1
rat -= 180
else:
dir = 0
ax.text(x,y, labelfmt % id, withdash=True,
dashlength=radius * 10,
dashpush=radius,
dashdirection=dir,
rotation=trat,
dashrotation=rat,
color=labelcolor
)
def bhshow(self, ax, component='bhmass', radius=(4, 1), logscale=True, labelfmt=None, labelcolor='white', vmin=None, vmax=None, count=-1, *args, **kwargs):
return self.pointshow(ax, 'bh', component, radius, logscale, labelfmt, labelcolor, vmin, vmax, count, *args, **kwargs)
# col = CircleCollection(offsets=zip(X.flat,Y.flat), sizes=(R * radius)**2, edgecolor='green', facecolor='none', transOffset=gca().transData)
# ax.add_collection(col)
def bhshow_camera(self, ax, component='bhmdot', radius=4, logscale=True, vmin=None, vmax=None, count=-1, camera=None, *args, **kwargs):
bhlum = self.bh.cosmology.QSObol(self.bh[component], 'blue')
good = bhlum > 0
if len(bhlum) > 0 and good.sum() > 0:
bhlum0 = bhlum[good]
bhpos0 = self.bh['locations'][good]
if self.periodic:
bhpos = zeros((27, bhpos0.shape[0], 3))
bhlum = zeros((27, bhlum0.shape[0]))
for im in range(27):
i, j, k = im / 9 - 1, (im % 9) / 3 -1, im % 3 -1
bhpos[im, :, :] = bhpos0[:, :] + array([i, j, k]) * self.boxsize
bhlum[im, :] = bhlum0[:]
bhpos.shape = -1, 3
bhlum.shape = -1
else:
bhpos = bhpos0
bhlum = bhlum0
bhpos = camera.map(bhpos, self.shape)
bhdist = bhpos[:, 2]
bhlum /= 4 * 3.1416 * bhdist ** 2
bhlum /= (self.bh.cosmology.units.W / self.bh.cosmology.units.METER ** 2 )
print 'distance of bhs', bhdist.min(), bhdist.max()
print 'lum of bhs', bhlum.min(), bhlum.max()
# apparent maginitute
mag = ((-26.74 + 2.5 * log10(1400 / bhlum)))
ind = mag.argsort()
if count > -1 and count < len(mag):
magcut = mag[ind[count]]
else:
magcut = 9999
# 8 is the naked eye limit
selected = (bhpos[:, 0] > 0) & (bhpos[:, 0] < self.shape[0])
selected &= (bhpos[:, 1] > 0) & (bhpos[:, 1] < self.shape[1])
selected &= (bhpos[:, 2] > 0.0)& (mag < magcut)
print 'mag', mag.min(), mag.max()
bhpos = bhpos[selected]
bhlum = bhlum[selected]
print 'bhs shown', selected.sum()
mag = - mag[selected]
mag -= mag.min()
diff = mag.max()
if diff > 0:
mag /= diff
else: mag[:] = 1
bhx = bhpos[:, 0]
bhy = bhpos[:, 1]
print mag
t = 0.05
marker = ((-1, 0), (-t, t), (0, 1), (t, t), (1, 0), (t, -t), (0, -1), (-t, -t)), 0
color=(0.8, 1, 1, 0.7)
ecolor=(1., 1, 0.8, 0.2)
if isscalar(radius):
mag*=radius
else:
rmin = min(radius)
rmax = max(radius)
mag *= (rmax - rmin)
mag += rmin
c = SpikeCollection(bhx, bhy, radius=mag, color=color)
ax.add_collection(c)
# ax.scatter(x=bhx, y=bhy, s=mag * 10**2, marker=marker, edgecolor=ecolor, color=color)
def velshow(self, ax, ftype, relative=False, color='cyan', alpha=0.8):
X,Y,vel = self.vector(ftype, 'vel', grids=(20,20), quick=False)
field = self.F[ftype]
mean = None
if not type(relative) == bool or relative == True:
if type(relative) == bool:
mean = average(field['vel'], weights=field['mass'], axis=0)[0:2]
else:
mean = asarray(relative)
vel[:,:,0]-=mean[0]
vel[:,:,1]-=mean[1]
print 'mean = ', mean
print 'max component of velocity', abs(vel).max()
sorted = absolute(vel.ravel())
sorted.sort()
scale = sorted[sorted.size * 9 / 10] * 20
print X.shape, vel[:,:,0].shape
ax.quiver(X,Y, vel[:,:,0], vel[:,:,1], width=0.003, scale=scale, scale_units='width', angles='xy', color=color,alpha=alpha)
if mean != None:
ax.quiver(self.cut.center[0], self.cut.center[1], mean[0], mean[1], scale=scale, scale_units='width', width=0.01, angles='xy', color=(0,0,0,0.5))
def starshow_poor(self, ax, *args, **kwargs):
star = self.star
mask = self.cut.select(star['locations'])
X = star['locations'][mask,0]
Y = star['locations'][mask,1]
if not 'color' in kwargs:
kwargs['color'] = 'white'
if not 'alpha' in kwargs:
kwargs['alpha'] = 0.5
ax.plot(X, Y, ', ', *args, **kwargs)
def reset_view(self, ax, camera=False):
if camera:
left, right = 0, self.shape[0]
bottom, top = 0, self.shape[1]
else:
left,right =self.cut['x']
bottom, top = self.cut['y']
print left, right, bottom, top
ax.set_xlim(left, right)
ax.set_ylim(bottom, top)
def makeT(self, Xh = 0.76):
"""T will be in Kelvin"""
gas =self.gas
C = gas.cosmology
gas['T'] = zeros(dtype='f4', shape=gas.numpoints)
C.ie2T(ie = gas['ie'], ye = gas['ye'], Xh = Xh, out = gas['T'])
gas['T'] *= C.units.TEMPERATURE
def mergeBHs(self, threshold=1.0):
bh = self.bh
if bh.numpoints == 0: return
pos = bh['locations']
posI = tile(pos, pos.shape[0]).reshape(pos.shape[0], pos.shape[0], 3)
posJ = posI.transpose((1,0,2))
d = posI - posJ
d = sqrt((d ** 2).sum(axis=2))
mask = d < threshold
mergeinto = zeros(bh.numpoints, dtype='i4')
for i in range(bh.numpoints): mergeinto[i] = nonzero(mask[i])[0][0]
g, ind = unique(mergeinto, return_index=True)
comp='bhmass'
bak = bh[comp].copy()
for i in range(bh.numpoints): bh[comp][i] = bak[mask[i]].sum()
try :
comp='bhmdot'
bak = bh[comp].copy()
for i in range(bh.numpoints): bh[comp][i] = bak[mask[i]].sum()
except KeyError: pass
comp='locations'
#def sel(comp):
bak = bh[comp].copy()
for i in range(bh.numpoints): bh[comp][i] = bak[mask[i]][0]
bh.numpoints = len(ind)
for comp in bh.names: bh[comp] = bh[comp][ind]
def fieldshow(self, ax, ftype, component, mode=None, camera=None,
vmin=None, vmax=None,
logscale=True, cmap=pygascmap,
gamma=1.0, mmin=None, mmax=None, over=None,
return_raster=False, logweight=None, weightcomponent='mass', use_cache=True):
if camera is None:
if mode is None: # extensive, mass->density, sfr->sf density,
todraw, mass = self.projfield(ftype=ftype, component=component, use_cache=use_cache)
else : # intensive, decorate the result aftwards
todraw, mass = self.projfield(ftype=ftype, component=component, weightcomponent=weightcomponent, use_cache=use_cache)
else:
if mode is None: # extensive, mass->density, sfr->sf density,
todraw, mass = self.camera(ftype=ftype, component=component, camera=camera)
else : # intensive, decorate the result aftwards
todraw, mass = self.camera(ftype=ftype, component=component, weightcomponent=weightcomponent, camera=camera)
if camera is None:
if logweight is None: logweight=True
image = self.render(todraw=todraw, mass=mass, mode=mode, vmin=vmin, vmax=vmax, logscale=logscale, logweight=logweight, cmap=cmap, gamma=gamma, mmin=mmin, mmax=mmax, over=over)
ret = ax.imshow(image.transpose((1,0,2)), origin='lower',
extent=self.extent, vmin=vmin, vmax=vmax, cmap=cmap)
else:
if logweight is None: logweight=False
image = self.render(todraw=todraw, mass=mass, mode=mode, vmin=vmin, vmax=vmax, logscale=logscale, logweight=logweight, cmap=cmap, gamma=gamma, mmin=mmin, mmax=mmax, over=over)
ret = ax.imshow(image.transpose((1,0,2)), origin='lower',
extent=(0, self.shape[0], 0, self.shape[1]), vmin=vmin, vmax=vmax, cmap=cmap)
return ret
def fieldcontour(self, ax, ftype, component, camera=None,
levels=None, logweight=None, weightcomponent='mass', use_cache=True, colors='k', linewidth=2, linestyle=['solid']):
if camera is None:
todraw, mass = self.projfield(ftype=ftype, component=component, weightcomponent=weightcomponent, use_cache=use_cache)
else:
todraw, mass = self.camera(ftype=ftype, component=component, weightcomponent=weightcomponent, camera=camera)
if camera is None:
if logweight is None: logweight=True
ret = ax.contourf((todraw/mass).T, extent=self.extent, colors=colors, linewidth=linewidth, levels=levels, linestyle=linestyle)
else:
if logweight is None: logweight=False
ret = ax.contourf((todraw/mass).T, extent=(0, self.shape[0], 0, self.shape[1]), colors=colors, linewidth=linewidth, levels=levels, linestyle=linestyle)
return ret
def render(self, todraw, mass, mode=None, vmin=None, vmax=None, logscale=True, logweight=True, cmap=pygascmap, gamma=1.0, mmin=None, mmax=None, over=None, composite=None):
""" vmin can be a string, eg '40 dB', so the vmin=vmax/1e4 """
if mode == 'intensity' or mode == 'mean':
todraw /= mass
if logscale:
if vmin is not None and not isinstance(vmin, basestring):
vmin = 10 ** vmin
if vmax is not None:
vmax = 10 ** vmax
if vmax is None:
if mode == 'intensity' or mode == 'mean' or over is None:
vmax = fmax.reduce(todraw.flat)
else:
sort = todraw.ravel().argsort()
ind = sort[(len(sort) - 1) * (1.0 - over)]
vmax = todraw.ravel()[ind]
del sort
if vmin is None:
if logscale:
vmin = ccode.pmin.reduce(todraw.flat)
else:
vmin = fmin.reduce(todraw.flat)
elif isinstance(vmin, basestring):
db = int(vmin[:-3])
vmin = vmax / 10 ** (db/10.)
if logscale:
log10(todraw, todraw)
vmin = log10(vmin)
vmax = log10(vmax)
print 'vmax, vmin =' , vmax, vmin
image = zeros(dtype = ('u1', 4), shape = todraw.shape)
# gacolor is much faster then matplotlib's normalize and uses less memory(4times fewer).
if not hasattr(cmap, 'table'):
cmap = gacmap(cmap)
ccode.color(image, todraw, max = vmax, min = vmin, logscale=False, colormap = cmap)
del todraw
if mode == 'intensity':
weight = mass
if logweight:
# in camera mode do not use logscale for mass.
weight.clip(ccode.pmin.reduce(weight.ravel()), inf, weight)
log10(weight, weight)
print 'weight', weight.mean(), weight.max(), weight.min()
if mmin is None:
mmin = weight.min()
if mmax is None:
if over is None:
mmax = weight.max()
else:
sort = weight.ravel().argsort()
ind = sort[(len(sort) - 1) * (1.0 - over)]
mmax = weight.ravel()[ind]
print 'mmax, mmin', mmax, mmin
weight -= mmin
weight /= (mmax - mmin)
weight.clip(0, 1, weight)
else:
weight = image[:, :, 3] / 255.0
weight **= gamma
if composite is not None:
alpha = composite[:, :, 3] / 255.0
multiply(1.0 - weight[:, :], alpha[:, :], alpha[:, :])
multiply(image[:, :, 0:3], weight[:, :, newaxis], image[:, :, 0:3])
multiply(composite[:, :, 0:3], alpha[:, :, newaxis], composite[:, :, 0:3])
add(image[:, :, 0:3], composite[:, :, 0:3], image[:, :, 0:3])
image[:, :, 3] = (alpha[:, :] + weight[:, :]) * 255
else:
#multiply(image[:, :, 0:3], weight[:, :, newaxis], image[:, :, 0:3])
#image[:, :, 3] = 255
#weight **= 0.33333333333
multiply(255.9999, weight[:, :], out=image[:,:,3], casting='unsafe')
# print 'alpha', image[:, :, 3].ravel().min(), image[:,:,3].ravel().max()
return image
def decorate(self, ax, frameon=True, title=None, bgcolor='k', color='w', fontsize='small'):
cut = self.cut
ax.set_axis_bgcolor(bgcolor)
if frameon :
ax.ticklabel_format(axis='x', useOffset=cut.center[0])
ax.ticklabel_format(axis='y', useOffset=cut.center[1])
ax.set_xticks(linspace(cut['x'][0], cut['x'][1], 5))
ax.set_yticks(linspace(cut['y'][0], cut['y'][1], 5))
#ax.set_title(title)
if(title is not None):
ax.set_title(title, position=(0.1, 0.9), fontsize=fontsize, color=color, transform=ax.transAxes)
else :
ax.axison = False
if(title is not None):
ax.set_title(title, position=(0.1, 0.9), fontsize=fontsize, color=color, transform=ax.transAxes)
ax.figure.set_facecolor(bgcolor)
def drawscale(self, ax, color='white', fontsize=None):
from mpl_toolkits.axes_grid.anchored_artists import AnchoredSizeBar
l = (self.cut.size[0]) * 0.2
l = l // 10 ** int(log10(l)) * 10 ** int(log10(l))
if l > 500 :
l/=1000.0
l = int(l+0.5)
text = r"%g Mpc/h" % l
l *= 1000.0
else:
text = r"%g Kpc/h" %l
print text
b = AnchoredSizeBar(ax.transData, l, text, loc = 8,
pad=0.1, borderpad=0.5, sep=5, frameon=False)
for r in b.size_bar.findobj(Rectangle):
r.set_edgecolor(color)
for t in b.txt_label.findobj(Text):
t.set_color(color)
if fontsize is not None:
t.set_fontsize(fontsize)
ax.add_artist(b)
def circle(self, ax, *args, **kwargs):
from matplotlib.patches import Circle
c = Circle(*args, **kwargs)
ax.add_patch(c)
