def test_watershed_pts_3d(self):
x = np.arange(10, dtype=float)
# generate 3D grid
grid = UniformGrid3D(x, x, x)
def gauss(coors_xyz, center=[0, 0, 0]):
center = np.array(center)
return np.exp(-(np.sum((coors_xyz - center)**2 / 30, axis=-1)))
value_array = gauss(grid.points, [0, 0, 0]) + gauss(
grid.points, [7, 8, 9])
wt_obj = Watershed(grid)
wt_obj.search_watershed_pts(value_array)
assert wt_obj.n_basins == 2
def test_basin_wts2(self):
x = np.linspace(0., 2., 11)
n_x = len(x)
# generate 3D grid
grid = UniformGrid3D(x, x, x)
def gauss(coors_xyz, center=[0, 0, 0]):
center = np.array(center)
return np.exp(-(np.sum((coors_xyz - center)**2, axis=-1)))
for _ in range(5):
center = np.random.rand(3) * 0.5
value_array = gauss(grid.points, center) + gauss(
grid.points, [2., 2., 2.] - center)
wt_obj = Watershed(grid)
wt_obj.search_watershed_pts(value_array)
assert wt_obj.n_basins == 2
basin_wt1 = wt_obj.compute_basin_wts(0)
basin_wt2 = wt_obj.compute_basin_wts(1)
assert_allclose(basin_wt1 + basin_wt2, np.ones(grid.size))
sum1 = np.sum(basin_wt1 * value_array)
sum2 = np.sum(basin_wt2 * value_array)
assert_almost_equal(sum1, sum2)
def test_compute_weights_for_watershed_pts(self):
x = np.arange(10, dtype=float)
# generate 3D grid
grid = UniformGrid3D(x, x, x)
def gauss(coors_xyz, center=[0, 0, 0]):
center = np.array(center)
return np.exp(-(np.sum((coors_xyz - center)**2 / 25, axis=-1)))
value_array = gauss(grid.points, [0, 0, 0]) + gauss(
grid.points, [7, 8, 9])
wt_obj = Watershed(grid)
wt_obj.search_watershed_pts(value_array)
assert wt_obj.n_basins == 2
# wt_obj.compute_weights_for_all_watershed_pts(value_array)
# print(wt_obj.water_basin_wts.toarray())
for i in range(grid.size):
basin = wt_obj.basins[i]
assert basin != -2
if basin == -1:
bs_vals = wt_obj.water_pt_basins[i]
wts = wt_obj.water_pt_wts[i]
assert len(np.asarray(bs_vals)) == len(np.asarray(wts))
assert_almost_equal(np.sum(np.asarray(bs_vals)), 1)