Commit 0538b4fd authored by Marius Kriegerowski's avatar Marius Kriegerowski
Browse files

wip: 3d plotting tests

parent df292896
import math
import numpy as num
import logging
from pyrocko.guts import Object, Float
from pyrocko import orthodrome as od
from lassie import receiver
logger = logging.getLogger('lassie.grid')
guts_prefix = 'lassie'
......@@ -108,6 +111,7 @@ class Carthesian3DGrid(Grid):
amax=None,
z_slice=None,
cmap=None,
projection='2d',
system='latlon'):
if system == 'latlon':
......@@ -121,14 +125,48 @@ class Carthesian3DGrid(Grid):
a3d = a.reshape((nz, ny, nx))
if z_slice is not None:
iz = num.argmin(num.abs(z-z_slice))
a2d = a3d[iz, :, :]
else:
a2d = num.max(a3d, axis=0)
try:
import vtk
use_vtk = True
except ImportError as e:
logger.warn(e)
use_vtk = False
if projection == '2d' or use_vtk:
if z_slice is not None:
iz = num.argmin(num.abs(z-z_slice))
a2d = a3d[iz, :, :]
else:
a2d = num.max(a3d, axis=0)
axes.pcolormesh(y, x, a2d.T, vmin=amin, vmax=amax, cmap=cmap)
axes.pcolormesh(y, x, a2d.T, vmin=amin, vmax=amax, cmap=cmap)
if projection == '3d':
#if use_vtk:
if True:
from lassie import vtk_graph
actors = vtk_graph.grid_actors(a3d.T, x, y, z, normalize=True)
vtk_graph.render_actors(actors)
else:
iy = len(y)
ix = len(x)
iz = len(z)
x = num.repeat(x, iy*iz)
y = num.repeat(num.tile(y, ix), iz)
z = num.tile(z, ix*iy)
a3d /= num.max(a3d)
#a3d = num.exp(a3d)
iplot = num.where(a3d.T!=num.nan)#>num.median(a3d))
iflat = num.ravel(iplot)
axes.scatter(x[iflat], y[iflat], z[iflat], s=a3d.T[iplot]*20.)
#, c=a3d[iflat],
#vmin=num.min(a3d),
#vmax=num.max(a3d),
#cmap=cmap)
axes.invert_zaxis()
axes.set_xlabel('E [m]')
axes.set_ylabel('N [m]')
axes.set_ylabel('D [m]')
def geometrical_normalization(grid, receivers):
distances = grid.distances(receivers)
......
......@@ -93,11 +93,14 @@ def plot_detection(
grid, receivers, frames, tmin_frames, deltat_cf, imax, iframe,
fsmooth_min, xpeak, ypeak, zpeak, tr_stackmax, tpeaks, apeaks,
detector_threshold, pdata, trs_raw, fmin, fmax, idetection,
grid_station_shift_max,
grid_station_shift_max, projection='3d',
movie=False):
import matplotlib
matplotlib.use('TkAgg')
from matplotlib import pyplot as plt
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D
from pyrocko.cake_plot import mpl_init, labelspace, colors, \
str_to_mpl_color as scolor
......@@ -107,7 +110,7 @@ def plot_detection(
fig = plt.figure(figsize=(16, 9))
axes = plt.subplot2grid((2, 3), (0, 2), aspect=1.0)
axes = plt.subplot2grid((2, 3), (0, 2), aspect=1.0, projection=projection)
labelspace(axes)
axes2 = plt.subplot2grid((2, 3), (1, 2))
......@@ -285,13 +288,16 @@ def plot_detection(
else:
frame = num.max(frames[:, iframe_min:iframe_max+1], axis=1)
grid.plot(axes, frame, amin=0.0, amax=amax, cmap=cmap, system=system)
plot_receivers(axes, receivers, system=system)
axes.plot(
ypeak, xpeak, '*',
ms=20.,
mec='black',
mfc='white')
grid.plot(axes, frame, amin=0.0, amax=amax, cmap=cmap, system=system,
projection=projection)
#plot_receivers(axes, receivers, system=system)
#axes.plot(
# ypeak, xpeak, zpeak, '*',
# ms=20.,
# mec='black',
# mfc='white')
#axes.scatter(
# ypeak, xpeak, zpeak, '*')
fig.tight_layout()
......
import vtk
import numpy as num
from vtk.util import numpy_support
def numpy_to_vtk(a, nc=3):
flattened = a.flatten(order='F')
data = numpy_support.numpy_to_vtk(flattened, deep=True)
data.SetNumberOfComponents(nc)
return data
def grid_actors_points(pnts, xg, yg, zg, normalize=True):
'''
:param pnts: matrix of size (len(xg), len(yg), len(zg)) with scales
:param xg, yg, zg: coordinates vectors.'''
actors = []
if normalize:
pntsmax = num.max(pnts)
else:
pntsmax = 1.
for ix, xi in enumerate(xg):
for iy, yi in enumerate(yg):
for iz, zi in enumerate(zg):
scale = 0.2 + 4.*(pnts[ix, iy, iz]/pntsmax)
actors.append(
point_actor(num.array((xi, yi, zi)),
scale=scale))
return actors
def grid_actors(pnts, xg, yg, zg, normalize=True):
'''
:param pnts: matrix of size (len(xg), len(yg), len(zg)) with scales
:param xg, yg, zg: coordinates vectors.
grid needs to be regularly spaced!'''
actors = []
if normalize:
pntsmax = num.max(pnts)
else:
pntsmax = 1.
pnts /= pntsmax
opacities = pnts**10.
dx = xg[1] - xg[0]
dy = yg[1] - yg[0]
dz = zg[1] - zg[0]
for ix, xi in enumerate(xg):
for iy, yi in enumerate(yg):
for iz, zi in enumerate(zg):
scale = 0.2 + 4.*(pnts[ix, iy, iz]/pntsmax)
s = vtk.vtkCubeSource()
s.SetCenter(xi, yi, zi)
s.SetXLength(dx)
s.SetYLength(dy)
s.SetZLength(dz)
mapper = vtk.vtkPolyDataMapper()
mapper.SetInputConnection(s.GetOutputPort())
actor = vtk.vtkActor()
actor.GetProperty().SetOpacity(opacities[ix, iy, iz])
actor.SetMapper(mapper)
actors.append(actor)
return actors
def point_actor(pnts, scale=1):
'''pnts: numpy array of length 3
actually, it's a sphere not a point'''
data = numpy_to_vtk(pnts)
s = vtk.vtkSphereSource()
s.SetCenter(*pnts)
s.SetRadius(10.*scale)
mapper = vtk.vtkPolyDataMapper()
mapper.SetInputConnection(s.GetOutputPort())
actor = vtk.vtkActor()
actor.GetProperty().SetOpacity(scale/10.)
actor.SetMapper(mapper)
return actor
def ray_actor(pnts, opacity=1.):
'''pnts: numpy array of shape (3, X) for X is number of points.
So, it's x,y,z columns'''
data = numpy_to_vtk(pnts)
points = vtk.vtkPoints()
points.SetData(data)
lines = vtk.vtkCellArray()
npoints = pnts.shape[1]
lines.InsertNextCell(npoints)
for i in range(npoints):
lines.InsertCellPoint(i)
polygon = vtk.vtkPolyData()
polygon.SetPoints(points)
polygon.SetLines(lines)
# standard stuff:
polygonMapper = vtk.vtkPolyDataMapper()
if vtk.VTK_MAJOR_VERSION <= 5:
polygonMapper.SetInputConnection(polygon.GetProducerPort())
else:
polygonMapper.SetInputData(polygon)
polygonMapper.Update()
polygonActor = vtk.vtkActor()
polygonActor.GetProperty().SetPointSize(20)
p = polygonActor.GetProperty()
p.SetPointSize(20)
p.SetOpacity(opacity)
polygonActor.GetProperty().SetPointSize(20)
polygonActor.SetMapper(polygonMapper)
return polygonActor
def render_actors(actors):
''' Take a list of actors and render them.'''
ren1 = vtk.vtkRenderer()
for a in actors:
ren1.AddActor(a)
ren1.SetBackground(0.1, 0.2, 0.4)
ren1.ResetCamera()
renWin = vtk.vtkRenderWindow()
renWin.AddRenderer(ren1)
renWin.SetSize(300, 300)
iren = vtk.vtkRenderWindowInteractor()
iren.SetRenderWindow(renWin)
iren.Initialize()
iren.Start()
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment