def test_gridpoint2index(self):
"""
Should convert between 1D and 3D grid indices.
"""
vm = VM()
# should return an integer index
i0 = (10, 0, 3)
idx = vm.gridpoint2index(*i0)
self.assertEqual(type(idx), int)
# should return original 3D indices
i1 = vm.index2gridpoint(idx)
self.assertEqual(type(i1), tuple)
for j in range(3):
self.assertEqual(i0[j], i1[j])
def test_insert_interface(self):
"""
Should insert an interface and handle setting flags correctly
"""
# Initialize a new model
vm = VM(r1=(0, 0, 0), r2=(50, 0, 30), dx=0.5, dy=0.5, dz=0.5)
# Should add a new, flat interface at 10 km
z0 = 10.
vm.insert_interface(z0 * np.ones((vm.nx, vm.ny)))
self.assertEqual(vm.nr, 1)
self.assertEqual(vm.rf[0].min(), z0)
self.assertEqual(vm.rf[0].max(), z0)
# New interfaces should have jp=0
self.assertEqual(vm.jp[0].min(), 0)
self.assertEqual(vm.jp[0].max(), 0)
# New layers should have ir and ij = index of new interface
self.assertEqual(vm.ir[0].min(), 0)
self.assertEqual(vm.ir[0].max(), 0)
self.assertEqual(vm.ij[0].min(), 0)
self.assertEqual(vm.ij[0].max(), 0)
# Adding a new interface should increase ir and ij of deeper layers
for z0 in [5., 15., 1., 20.]:
vm.insert_interface(z0 * np.ones((vm.nx, vm.ny)))
for iref in range(0, vm.nr):
self.assertEqual(vm.ir[iref].min(), iref)
self.assertEqual(vm.ir[iref].max(), iref)
self.assertEqual(vm.ij[iref].min(), iref)
self.assertEqual(vm.ij[iref].max(), iref)
# should take a scalar value for a constant depth interface
vm = VM(r1=(0, 0, 0), r2=(50, 0, 30), dx=0.5, dy=0.5, dz=0.5)
z0 = 10.
vm.insert_interface(z0)
self.assertEqual(vm.nr, 1)
self.assertEqual(vm.rf[0].min(), z0)
self.assertEqual(vm.rf[0].max(), z0)
def test_x2i(self):
"""
Should convert between x indices and coordinates and back
"""
r1 = (-10, -20, -30)
r2 = (10, 20, 30)
vm = VM(r1=r1, r2=r2, dx=1, dy=1, dz=1)
ix = vm.x2i([vm.r1[0]])[0]
x = vm.i2x([ix])[0]
self.assertEqual(ix, 0)
self.assertEqual(x, vm.r1[0])
ix = vm.x2i([vm.r2[0]])[0]
x = vm.i2x([ix])[0]
self.assertEqual(ix, vm.nx - 1)
self.assertEqual(x, vm.r2[0])
def test_z2i(self):
"""
Should convert between z indices and coordinates and back
"""
r1 = (-10, -20, -30)
r2 = (10, 20, 30)
vm = VM(r1=r1, r2=r2, dx=1, dy=1, dz=1)
iz = vm.z2i([vm.r1[2]])[0]
z = vm.i2z([iz])[0]
self.assertEqual(iz, 0)
self.assertEqual(z, vm.r1[2])
iz = vm.z2i([vm.r2[2]])[0]
z = vm.i2z([iz])[0]
self.assertEqual(iz, vm.nz - 1)
self.assertEqual(z, vm.r2[2])
def test_y2i(self):
"""
Should convert between y indices and coordinates and back
"""
r1 = (-10, -20, -30)
r2 = (10, 20, 30)
vm = VM(r1=r1, r2=r2, dx=1, dy=1, dz=1)
iy = vm.y2i([vm.r1[1]])[0]
y = vm.i2y([iy])[0]
self.assertEqual(iy, 0)
self.assertEqual(y, vm.r1[1])
iy = vm.y2i([vm.r2[1]])[0]
y = vm.i2y([iy])[0]
self.assertEqual(iy, vm.ny - 1)
self.assertEqual(y, vm.r2[1])
def test_get_layer_bounds(self):
"""
Should return arrays with layer top and bottom bounding surfaces
"""
# Create a simple 3D model
r1 = (0, 0, 0)
r2 = (50, 50, 20)
vm = VM(r1=r1, r2=r2, dx=2, dy=2, dz=0.2)
# Should have no layers
self.assertEqual(vm.nr, 0)
# Should return model top and bottom
z0, z1 = vm.get_layer_bounds(0)
self.assertEqual(z0.shape, (vm.nx, vm.ny))
self.assertEqual(z1.shape, (vm.nx, vm.ny))
self.assertEqual(z0.min(), r1[2])
self.assertEqual(z0.max(), r1[2])
self.assertEqual(z1.min(), r2[2])
self.assertEqual(z1.max(), r2[2])
# Insert an interfce and get bounds again
_z = 5.
vm.insert_interface(_z * np.ones((vm.nx, vm.ny)))
# Top layer should be between r1[2] and _z
z0, z1 = vm.get_layer_bounds(0)
self.assertEqual(z0.min(), r1[2])
self.assertEqual(z0.max(), r1[2])
self.assertEqual(z1.min(), _z)
self.assertEqual(z1.max(), _z)
# Bottom layer should be between _z and r2[2]
z0, z1 = vm.get_layer_bounds(1)
self.assertEqual(z0.min(), _z)
self.assertEqual(z0.max(), _z)
self.assertEqual(z1.min(), r2[2])
self.assertEqual(z1.max(), r2[2])
# Should handle pinchouts
_z1 = 2 * _z
iref = vm.insert_interface(_z1 * np.ones((vm.nx, vm.ny)))
ix = vm.xrange2i(0, 25)
vm.rf[iref, ix, :] = vm.rf[iref - 1, ix, :]
z0, z1 = vm.get_layer_bounds(1)
self.assertEqual(z0.min(), _z)
self.assertEqual(z0.max(), _z)
self.assertEqual(z1.min(), _z)
self.assertEqual(z1.max(), _z1)
def test_remove_interface(self):
"""
Should remove an interface and handle setting flags correctly
"""
# Initialize a new model
vm = VM(r1=(0, 0, 0), r2=(50, 0, 30), dx=0.5, dy=0.5, dz=0.5)
# Insert a set of interaces
nr = 0
for z0 in [1., 5., 10., 15., 20.]:
nr += 1
vm.insert_interface(z0 * np.ones((vm.nx, vm.ny)))
# Should remove interface and decrease ir and ij of deeper layers
for iref in range(0, nr):
vm.remove_interface(0)
nr -= 1
self.assertEqual(vm.nr, nr)
for _iref in range(0, vm.nr):
self.assertEqual(vm.ir[_iref].min(), _iref)
self.assertEqual(vm.ir[_iref].max(), _iref)
self.assertEqual(vm.ij[_iref].min(), _iref)
self.assertEqual(vm.ij[_iref].max(), _iref)
def setUp(self):
"""
Build test models
"""
# Create a 3D model with three layers (2 boundaries)
r1 = (0, 0, 0)
r2 = (20, 20, 30)
self.vm = VM(r1=r1, r2=r2)
self.vm.insert_interface(5)
self.vm.insert_interface(10)
# Define a path for the test case
self.pi = [0, 1, 1, 2, 2, 2, 1, 1, 0]
self.px = [1, 1, 1, 1, 1, 1, 1, 1, 1]
self.py = [1, 2, 3, 4, 5, 6, 7, 8, 9]
self.pz = [1, 6, 8, 11, 12, 11, 8, 6, 2]
self.path = zip(self.px, self.py, self.pz, self.pi)
def setUp(self):
"""
Build a test VM model
"""
self.vm = VM(r1=(0, 0, 0), r2=(100, 100, 50), dx=5, dy=5, dz=5)
self.vm.define_constant_layer_velocity(0, 1.5)
self.vm.insert_interface(3)
self.vm.insert_interface(10)
# dipping layer
dzdx = 0.1
dzdy = 0.1
rf = 25 * np.ones((self.vm.nx, self.vm.ny))
for ix in range(self.vm.nx):
for iy in range(self.vm.ny):
rf[ix, iy] += dzdx * (ix * self.vm.dx)
rf[ix, iy] += dzdy * (iy * self.vm.dy)
# XXX self.vm.insert_interface(rf)
self.vm.insert_interface(30)
ray0 = [ # x, y, z, ilyr, direction
[2, 2, 1, 0, 1],
[22, 22, 25, 2, 1],
[23, 23, 28, 2, 1],
[40, 40, 44, 3, 1],
[41, 41, 46, 3, 0],
[62, 62, 19, 2, -1],
[80, 80, 5, 1, -1],
]
ray1 = [[5, 5, 2, 0, 1], [60, 60, 45, 3, 1]] # x, y, z, ilyr, direction
ray2 = [[5, 5, 45, 3, -1], [60, 60, 2, 0, -1]] # x, y, z, ilyr, direction
# rays
self.raypaths = np.asarray([ray0, ray1, ray2])
self.vm.write("test.vm")
def dev_project_layer_velocities(self):
"""
Should map 2D model velocities to a 3D model.
"""
# Create a simple 2D model with sl set to layer IDs
vm2d = VM(r1=(-2, 0, -2), r2=(50, 0, 30), dx=1, dy=1, dz=1)
z = 3 + vm2d.x * 0.1
vm2d.insert_interface([[_z] for _z in z])
vm2d.insert_interface(10)
z = 22 + vm2d.x * -0.1
vm2d.insert_interface([[_z] for _z in z])
vm2d.sl = vm2d.layers
# Project into a 3D model
phi = 30
vm3d = VM(r1=(10, 1, -2), r2=(20, 20, 30), dx=1, dy=1, dz=1)
vm3d.insert_interface(5)
vm3d.insert_interface(10)
vm3d.insert_interface(20)
for ilyr in range(vm2d.nr + 1):
project_layer_velocities(vm2d, vm3d, phi, ilyr)
# Layers should only include values from the same 2d layers
for ilyr in range(vm2d.nr + 1):
idx = np.nonzero(vm3d.layers == ilyr)
self.assertEqual(vm3d.sl[idx].min(), ilyr)
class raygeometryTestCase(unittest.TestCase):
"""
Test cases for the raygeometry module.
"""
def setUp(self):
"""
Build test models
"""
# Create a 3D model with three layers (2 boundaries)
r1 = (0, 0, 0)
r2 = (20, 20, 30)
self.vm = VM(r1=r1, r2=r2)
self.vm.insert_interface(5)
self.vm.insert_interface(10)
# Define a path for the test case
self.pi = [0, 1, 1, 2, 2, 2, 1, 1, 0]
self.px = [1, 1, 1, 1, 1, 1, 1, 1, 1]
self.py = [1, 2, 3, 4, 5, 6, 7, 8, 9]
self.pz = [1, 6, 8, 11, 12, 11, 8, 6, 2]
self.path = zip(self.px, self.py, self.pz, self.pi)
def test_distance(self):
"""
Should calculate cartesian distance
"""
# should work in 2D
p0 = (0, 0)
p1 = (0, 1.)
self.assertEqual(distance(p0, p1), 1.)
# should work in 3D
p0 = (0, 0, 0)
p1 = (0, 0, 1)
self.assertEqual(distance(p0, p1), 1.)
def test_assign_points_to_layers(self):
"""
Should assign a layer index to each point.
"""
# Should label each point with its layer index
pi = assign_points_to_layers(self.vm, self.px, self.py, self.pz)
for _pi0, _pi in zip(self.pi, pi):
self.assertEqual(_pi0, _pi)
def test_get_indices_near_piercing(self):
"""
Should find indices for points near a piercing points.
"""
# should get indices if path crosses layer
iref = 0
ip = get_indices_near_piercing(self.pi, iref)
self.assertEqual(len(ip), 2)
self.assertEqual(ip[0][0], 0)
self.assertEqual(ip[0][1], 1)
self.assertEqual(ip[1][0], 8)
self.assertEqual(ip[1][1], 7)
# should return empty array if path does not cross interface
pi = 999 * np.ones(len(self.pi))
ip = get_indices_near_piercing(pi, iref)
self.assertEqual(len(ip), 0)
# should just return points on the downward leg
ip = get_indices_near_piercing(self.pi, iref, downward=True,
upward=False)
self.assertEqual(len(ip), 1)
self.assertEqual(ip[0][0], 0)
self.assertEqual(ip[0][1], 1)
# should just return points on the upward leg
ip = get_indices_near_piercing(self.pi, iref, downward=False,
upward=True)
self.assertEqual(len(ip), 1)
self.assertEqual(ip[0][0], 8)
self.assertEqual(ip[0][1], 7)
def test_get_piercing_points(self):
"""
Should find coordinates of piercing points.
"""
iref = 0
# indices of points near the piercing points
ip = get_indices_near_piercing(self.pi, iref)
# should find piercing points if path crosses interface
pp = get_piercing_points(self.vm, iref, self.px, self.py, self.pz,
ip)
self.assertEqual(pp[0, 0], 1.)
self.assertEqual(pp[0, 1], 1.8)
self.assertEqual(pp[0, 2], 5.)
self.assertEqual(pp[1, 0], 1.)
self.assertEqual(pp[1, 1], 8.75)
self.assertEqual(pp[1, 2], 5.)
# should return empty array if path does not cross interface
px = np.ones(len(self.pi))
py = np.linspace(1, 9, len(self.pi))
pz = 12 * np.ones(len(self.pi))
pi = assign_points_to_layers(self.vm, px, py, pz)
ip = get_indices_near_piercing(pi, iref)
pp = get_piercing_points(self.vm, 0, px, py, pz, ip)
self.assertEqual(len(pp), 0)
# should just return points on the downward leg
ip = get_indices_near_piercing(self.pi, iref, downward=True,
upward=False)
pp = get_piercing_points(self.vm, 0, self.px, self.py, self.pz, ip)
self.assertEqual(pp[0, 0], 1.)
self.assertEqual(pp[0, 1], 1.8)
self.assertEqual(pp[0, 2], 5.)
# should just return points on the upward leg
ip = get_indices_near_piercing(self.pi, iref, downward=False,
upward=True)
pp = get_piercing_points(self.vm, 0, self.px, self.py, self.pz, ip)
self.assertEqual(pp[0, 0], 1.)
self.assertEqual(pp[0, 1], 8.75)
self.assertEqual(pp[0, 2], 5.)
def test_path_in_layer(self):
"""
Should get coordinates for portion of path in layer.
"""
down_npts = [2, 4, 5]
up_npts = [2, 4, 0]
for ilyr in range(3):
d, u = get_path_in_layer(self.vm, ilyr, self.px, self.py, self.pz)
self.assertEqual(len(d), down_npts[ilyr])
def two2three(vm, sol, eol, dx=None, dy=None, phi=90.0, head_only=False):
"""
Convert a 2D model to a 3D model by giving the coordinates of the line
endpoints.
Creates a cube with edges that are oriented parallel to the coordinate
system.
:See also: :meth:`rockfish.tomography.model.VM.extrude`
Parameters
----------
vm : :class:`rockfish.tomography.model.VM`
2D model to convert to a 3D model. Must have ``vm.ny==1``.
sol, eol : (float, float)
Coordinates for the start and end of the line, respectively.
dx : float, optional
Grid spacing in the x direction. Default is to use the same x spacing
in the 2D input mdel.
dy : float, optional
Grid spacing in the y direction. Default is to set ``dy=dx``.
phi : float, optional
Angle of the extrusion direction, measured in degrees from the line
azimuth. Default is ``90.0`` (i.e., in the line-perpendicular
direction).
head_only : bool, optional
If ``True``, only calculates the size values for the 3D model and does
not perform the actual extrusion.
"""
# check that current model is 2D
assert vm.ny == 1, 'vm must be 2D with vm.ny == 1'
# set defaults
if dx is None:
dx = vm.dx
if dy is None:
dy = vm.dx
# calculate line angle
theta = np.arctan2(eol[1] - sol[1], eol[0] - sol[0]) \
+ np.deg2rad(phi - 90.)
# initialize a new model
vm1 = VM(r1=(sol[0], sol[1], vm.r1[2]),
r2=(eol[0], eol[1], vm.r2[2]),
dx=dx, dy=dy, nr=vm.nr)
if head_only:
return vm1
# build surface interpolators
x2rf = []
x2jp = []
x2ir = []
x2ij = []
xp = vm.x - vm.r1[0] # x in 2D model relative to 2D origin
for iref in range(0, vm.nr):
x2rf.append(interp1d(xp, [v[0] for v in vm.rf[iref]],
kind='linear'))
x2jp.append(interp1d(xp, [v[0] for v in vm.jp[iref]],
kind='linear'))
x2ir.append(interp1d(xp, [v[0] for v in vm.ir[iref]],
kind='nearest'))
x2ij.append(interp1d(xp, [v[0] for v in vm.ij[iref]],
kind='nearest'))
# map 2D model to the 3D model
x = vm1.x - vm1.r1[0] # x in 3D model relative to origin
for _iy in vm1.yrange2i():
_y = vm1.y[_iy] - vm1.r1[1] # y in 3D model relative to origin
a, _, _ = project_point(x, _y, theta)
a = a.clip(vm.r1[0], vm.r2[0])
nclip = len(np.nonzero(a == vm.r1[0]))\
+ len(np.nonzero(a == vm.r2[0]))
if nclip > 2:
warnings.warn('{:} points off the 2D line at y={:}'\
.format(nclip - 2, _y))
# interpolate interface values
for iref in range(0, vm.nr):
try:
vm1.rf[iref, :, _iy] = x2rf[iref](a)
vm1.jp[iref, :, _iy] = x2jp[iref](a)
vm1.ir[iref, :, _iy] = x2ir[iref](a)
vm1.ij[iref, :, _iy] = x2ij[iref](a)
except ValueError as e:
msg = e.message
msg += ' (min(a),max(a)={:},{:}; vm.r1[0]={:}; vm.r2[0]={:})'\
.format(min(a), max(a), vm.r1[0], vm.r2[0])
raise ValueError(msg)
# assign velocities
for _ix, _a in enumerate(a):
_ixp = vm.x2i([_a])[0]
vm1.sl[_ix, _iy, :] = vm.sl[_ixp, 0, :]
return vm1
class geometryTestCase(unittest.TestCase):
"""
Test cases for the geometry module
"""
def setUp(self):
"""
Build a test VM model
"""
self.vm = VM(r1=(0, 0, 0), r2=(100, 100, 50), dx=5, dy=5, dz=5)
self.vm.define_constant_layer_velocity(0, 1.5)
self.vm.insert_interface(3)
self.vm.insert_interface(10)
# dipping layer
dzdx = 0.1
dzdy = 0.1
rf = 25 * np.ones((self.vm.nx, self.vm.ny))
for ix in range(self.vm.nx):
for iy in range(self.vm.ny):
rf[ix, iy] += dzdx * (ix * self.vm.dx)
rf[ix, iy] += dzdy * (iy * self.vm.dy)
# XXX self.vm.insert_interface(rf)
self.vm.insert_interface(30)
ray0 = [ # x, y, z, ilyr, direction
[2, 2, 1, 0, 1],
[22, 22, 25, 2, 1],
[23, 23, 28, 2, 1],
[40, 40, 44, 3, 1],
[41, 41, 46, 3, 0],
[62, 62, 19, 2, -1],
[80, 80, 5, 1, -1],
]
ray1 = [[5, 5, 2, 0, 1], [60, 60, 45, 3, 1]] # x, y, z, ilyr, direction
ray2 = [[5, 5, 45, 3, -1], [60, 60, 2, 0, -1]] # x, y, z, ilyr, direction
# rays
self.raypaths = np.asarray([ray0, ray1, ray2])
self.vm.write("test.vm")
def plot_init(self):
"""
setup should make a test model
"""
if not PLOT:
return
fig = plt.figure()
ax = fig.add_subplot(111)
y0 = 50
vm = self.vm.slice_along_xy_line(self.vm.x, y0 * np.ones(self.vm.nx))
vm.plot_layers(ax=ax)
for path in self.raypaths:
_path = np.asarray(path)
ax.plot(_path[:, 0], _path[:, 2], ".-c")
plt.title("test_init: Test model at y = {:}".format(y0))
plt.show()
def test_assign_points_to_layers(self):
"""
Should determine what layers a series of points are in
"""
pi = []
for path in self.raypaths:
_path = np.asarray(path)
_pi = assign_points_to_layers(self.vm, _path[:, 0], _path[:, 1], _path[:, 2])
pi.append(_pi)
for i0, i1 in zip(_path[:, 3], _pi):
self.assertEqual(i0, i1)
if PLOT:
fig = plt.figure()
ax = fig.add_subplot(111)
y0 = 50
vm = self.vm.slice_along_xy_line(self.vm.x, y0 * np.ones(self.vm.nx))
vm.plot(ax=ax, show_grid=True)
for i in range(len(pi)):
path = np.asarray(self.raypaths[i])
ax.plot(path[:, 0], path[:, 2], "-c")
ax.scatter(path[:, 0], path[:, 2], c=pi[i], s=50)
plt.title("test_assign_points_to_layers: nodes colored by layer")
plt.show()
def test_get_points_in_layer(self):
"""
Should return points within a layer
"""
for path in self.raypaths:
for ilyr in range(self.vm.nr):
_path = np.asarray(path)
px, py, pz = _path[:, 0], _path[:, 1], _path[:, 2]
d, u = split_downup(px, py, pz)
for p in [d, u]:
px, py, pz = p[:, 0], p[:, 1], p[:, 2]
pi = assign_points_to_layers(self.vm, px, py, pz)
pi0 = np.nonzero(pi == ilyr)[0]
# should only return points that fall within a layer
px1, py1, pz1, pi1 = get_points_in_layer(self.vm, ilyr, px, py, pz, overlap=False)
self.assertEqual(len(pi0), len(px1))
self.assertEqual(len(pi0), len(py1))
self.assertEqual(len(pi0), len(pz1))
self.assertEqual(len(pi0), len(pi1))
# should include neighboring points
px2, py2, pz2, pi2 = get_points_in_layer(self.vm, ilyr, px, py, pz, overlap=True)
self.assertGreaterEqual(len(px2), len(px1))
self.assertGreaterEqual(len(py2), len(py1))
self.assertGreaterEqual(len(pz2), len(pz1))
self.assertGreaterEqual(len(pi2), len(pi1))
# should be empty or have more than 1 point
self.assertNotEqual(len(px2), 1)
self.assertNotEqual(len(py2), 1)
self.assertNotEqual(len(pz2), 1)
self.assertNotEqual(len(pi2), 1)
if PLOT:
fig = plt.figure()
ax = fig.add_subplot(111)
self.vm.plot(ax=ax)
# ax.scatter(px, pz, c=pi, marker='.', s=30, zorder=0)
ax.plot(px, pz, ".-k", lw=5, zorder=1)
ax.plot(px1, pz1, ".-m", lw=4, zorder=2)
ax.plot(px2, pz2, ".-c", lw=1, zorder=3)
plt.title("test_get_points_in_layer: ilyr={:}".format(ilyr))
plt.show()
def test_split_downup(self):
"""
Should split path into down-going and up-coming segments
"""
down = []
up = []
for path in self.raypaths:
_path = np.asarray(path)
px, py, pz = _path[:, 0], _path[:, 1], _path[:, 2]
d, u = split_downup(px, py, pz)
down.append(d)
up.append(u)
ibot = np.nonzero(_path[:, 4] == 0)
idown = np.append(np.nonzero(_path[:, 4] == 1), ibot)
iup = np.append(ibot, np.nonzero(_path[:, 4] == -1))
self.assertEqual(len(d), len(idown))
self.assertEqual(len(u), len(iup))
if PLOT:
fig = plt.figure()
ax = fig.add_subplot(111)
self.vm.plot(ax=ax, ir=False, ij=False, rf=False)
for i in range(len(self.raypaths)):
path = np.asarray(self.raypaths[i])
ax.plot(path[:, 0], path[:, 2], "-k", lw=5)
ax.plot(down[i][:, 0], down[i][:, 2], "-r")
ax.plot(up[i][:, 0], up[i][:, 2], "-g")
plt.title("test_split_downup: up and down paths")
plt.show()
def test_find_line_intersection(self):
"""
Should find where a point crosses an interface
"""
p0 = [1, 2, 1]
p1 = [82, 40, 40]
for iref in [0, 1, 2]:
x, y, z = find_line_intersection(self.vm, iref, p0, p1)
ix = self.vm.x2i(x)
iy = self.vm.y2i(y)
z0 = self.vm.rf[iref][ix, iy]
self.assertAlmostEqual(z, z0, 1)
if PLOT:
fig = plt.figure()
ax = fig.add_subplot(111)
self.vm.plot(ax=ax, show_grid=True)
ax.plot([p0[0], p1[0]], [p0[2], p1[2]], ".-c")
ax.plot(x, z, ".w")
plt.title("test_find_line_intersection: iref={:}".format(iref))
plt.show()
def test__trim_path_to_layer(self):
"""
Should trim a pre-processed path to single layer
"""
# path that crosses layer 2
ray0 = [[1, 1, 1], [40, 1, 20], [80, 1, 40]]
# path that terminates in layer 2
ray1 = [[1, 1, 1], [40, 1, 20], [80, 1, 22]]
# path that starts in layer 2
ray2 = [[1, 1, 19], [40, 1, 20], [80, 1, 40]]
# path that starts and ends in layer 2
ray3 = [[1, 1, 19], [40, 1, 20], [80, 1, 22]]
for ray in [ray0, ray1, ray2, ray3]:
px = [r[0] for r in ray]
py = [r[1] for r in ray]
pz = [r[2] for r in ray]
if PLOT:
fig = plt.figure()
ax = fig.add_subplot(111)
self.vm.plot(ax=ax, show_grid=True)
ax.plot(px, pz, ".-g")
_trim_path_ends_to_layer(self.vm, px, py, pz)
pi = assign_points_to_layers(self.vm, px, py, pz)
for i in pi:
self.assertEqual(i, 2)
if PLOT:
ax.plot(px, pz, ".-m")
plt.show()
def test_insert_intersections(self):
"""
Should add layer intersections
"""
for path in self.raypaths[0:1]:
_path = np.asarray(path)
px, py, pz = _path[:, 0], _path[:, 1], _path[:, 2]
for duplicate in [True, False]:
if PLOT:
fig = plt.figure()
ax = fig.add_subplot(111)
plt.title("duplicate = {}".format(duplicate))
self.vm.plot(ax=ax, show_grid=True)
ax.plot(px, pz, ".-g", markersize=10, lw=3)
px, py, pz, pi = insert_intersections(self.vm, px, py, pz, duplicate=duplicate)
if PLOT:
ax.plot(px, pz, ".-m")
plt.show()
self.assertEqual(len(px), len(pi))
self.assertEqual(len(py), len(pi))
self.assertEqual(len(pz), len(pi))
def test_split_path_at_interface_intersections(self):
"""
Should split a path into layer segments
"""
sym = [".-r", ".-c", ".-m", ".-w"]
for path in self.raypaths:
_path = np.asarray(path)
px, py, pz = _path[:, 0], _path[:, 1], _path[:, 2]
_px, _, _pz, _ = insert_intersections(self.vm, px, py, pz)
if PLOT:
fig = plt.figure()
ax = fig.add_subplot(111)
self.vm.plot(ax=ax, show_grid=True)
ax.plot(px, pz, ".-g", markersize=20, lw=5)
ax.plot(_px, _pz, ".-k", markersize=10, lw=3)
segments = split_path_at_interface_intersections(self.vm, px, py, pz)
if PLOT:
for i, segments in enumerate(segments):
for p in segments:
if len(p) == 0:
continue
ax.plot(p[:, 0], p[:, 2], sym[i])
plt.show()
self.assertEqual(len(segments), self.vm.nr + 1)
def test_resample_path(self):
"""
Should interpolate path
"""
if PLOT:
fig = plt.figure()
ax = fig.add_subplot(111)
for path in self.raypaths:
_path = np.asarray(path)
px, py, pz = _path[:, 0], _path[:, 1], _path[:, 2]
pxi, pyi, pzi = resample_path(px, py, pz, dx=1)
if PLOT:
self.vm.plot(ax=ax, show_grid=True)
ax.plot(px, pz, ".-g", markersize=20, lw=5)
ax.plot(pxi, pzi, ".-k", markersize=10, lw=3)
if PLOT:
plt.show()