def test__process(self): df = pd.DataFrame() df["x"] = [1, 1, 2, 2] df["y"] = [1, 2, 1, 1] df["z"] = [1, 2, 3, 4] scatter = BubbleScatter(df, "x", "y", z="z", z_aggregation="sum", max_bubble_size=1) scatter._process() true_x_ticks = [0, 1] true_y_ticks = [0, 1] true_x_labels = [1, 2] true_y_labels = [1, 2] true_bubble_sizes = [0.25, 0.25, 0.5] true_bubble_colors = np.array([[1], [2], [7]]) true_x_values = np.array([[0.0], [0.0], [1.0]]) true_y_values = np.array([[0.0], [1.0], [0.0]]) self.assertListEqual(scatter.x_ticks, true_x_ticks) self.assertListEqual(scatter.y_ticks, true_y_ticks) self.assertListEqual(scatter.x_labels, true_x_labels) self.assertListEqual(scatter.y_labels, true_y_labels) self.assertListEqual(scatter.bubble_sizes, true_bubble_sizes) self.assertTrue(all(scatter.bubble_colors == true_bubble_colors)) self.assertTrue(all(scatter.x_values == true_x_values)) self.assertTrue(all(scatter.y_values == true_y_values))
def test__get_encoder_categories_with_bins_no_categories(self): df = pd.DataFrame(columns=["x", "y"]) scatter = BubbleScatter(df, "x", "y", xbins=3, ybins=3) x_categories = scatter._get_encoder_categories("x") y_categories = scatter._get_encoder_categories("y") self.assertEqual(x_categories, "auto") self.assertEqual(y_categories, "auto")
def test__calc_labels_y_with_categories(self): df = pd.DataFrame(columns=["x", "y"]) scatter = BubbleScatter(df, "x", "y", y_categories=["a", "b"]) result = scatter._calc_labels("y", ["a", "b", "c"]) true_result = ["a", "b"] self.assertListEqual(result, true_result)
def test__transform_with_int_ybins(self): df = pd.DataFrame() df["x"] = [0, 1, 2, 3, 4] df["y"] = [0, 1, 2, 3, 4] scatter = BubbleScatter(df, "x", "y", ybins=2) result = scatter._transform("y") true_result = pd.Series([1.0, 1.0, 1.0, 3.0, 3.0], name="y") self.assertTrue(result.equals(true_result))
def test__transform_no_ybins(self): df = pd.DataFrame() df["x"] = [1, 2, 3] df["y"] = [1, 2, 3] scatter = BubbleScatter(df, "x", "y") result = scatter._transform("y") true_result = df["x"] self.assertTrue(result.equals(true_result))
def test_plot_show_true(self): df = pd.DataFrame() df["x"] = [1, 1, 2, 2] df["y"] = [1, 2, 1, 1] scatter = BubbleScatter(df, "x", "y") result = scatter.plot(show=True) self.assertIsNone(result)
def test_plot_show_false(self): df = pd.DataFrame() df["x"] = [1, 1, 2, 2] df["y"] = [1, 2, 1, 1] scatter = BubbleScatter(df, "x", "y") result = scatter.plot(show=False) self.assertIsInstance(result, plt.Figure)
def test__get_encoder_categories_no_bins_with_categories(self): df = pd.DataFrame(columns=["x", "y"]) scatter = BubbleScatter(df, "x", "y", x_categories=["a", "b", "c"], y_categories=["q", "w", "e"]) x_categories = scatter._get_encoder_categories("x") y_categories = scatter._get_encoder_categories("y") self.assertListEqual(x_categories, ["a", "b", "c"]) self.assertListEqual(y_categories, ["q", "w", "e"])
def test__calc_bubble_colors(self): df = pd.DataFrame() df["x"] = [1, 1, 2, 2] df["y"] = [1, 2, 1, 1] df["z"] = [1, 2, 3, 4] scatter = BubbleScatter(df, "x", "y", z="z", z_aggregation="sum") index = pd.MultiIndex(levels=[[1, 2], [1, 2]], codes=[[1, 0, 0], [0, 1, 0]], names=["x", "y"]) true_result = np.array([[7], [2], [1]]) result = scatter._calc_bubble_colors(df, index) self.assertTrue(all(result == true_result))
def test__calc_xy_values(self): index = pd.MultiIndex(levels=[[1, 2], [1, 2]], codes=[[0, 0, 1], [0, 1, 0]], names=['x', 'y']) true_result = pd.DataFrame({"x": [1, 1, 2], "y": [1, 2, 1]}) result = BubbleScatter._calc_xy_values(index) self.assertTrue(result.equals(true_result))
def test__calc_bubble_sizes_conditional(self): """Test bubble sizes for conditional distribution (x|y).""" df = pd.DataFrame() df["x"] = [1, 1, 2, 2] df["y"] = [1, 2, 1, 1] scatter = BubbleScatter(df, "x", "y", joint_distribution=False, max_bubble_size=1) index = pd.MultiIndex(levels=[[1, 2], [1, 2]], codes=[[0, 0, 1], [0, 1, 0]], names=["x", "y"]) true_result = pd.Series([0.5, 0.5, 1.0], index=index) result = scatter._calc_bubble_sizes(df) self.assertTrue(result.equals(true_result))
def test__calc_ticks(self): result = BubbleScatter._calc_ticks(["a", "b", "c"]) true_result = [0, 1, 2] self.assertListEqual(result, true_result)