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__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()