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