Added surface plot as option for 2d models.

src/probabilistic_model/learning/jpt/jpt.py

@@ -326,9 +326,10 @@ def construct_impurity(self) -> Impurity:
                         n_num_vars_total, numeric_features, symbolic_features, symbols, max_variances,
                         dependency_indices)
 
-    def plot(self, number_of_samples: int = 1000) -> List:
+
+    def plot(self, number_of_samples: int = 1000, surface=True) -> List:
         try:
-            super().plot()
+            return super().plot(number_of_samples, surface)
         except NotImplementedError:
             return self.plot_univariate_distributions()
 

src/probabilistic_model/probabilistic_model.py

@@ -1,14 +1,18 @@
 from __future__ import annotations
 
 import abc
+import itertools
+import math
 
+import numpy as np
 from random_events.interval import closed, SimpleInterval
 from random_events.product_algebra import *
 from random_events.set import *
 from random_events.variable import *
 
 from .constants import *
 from .error import IntractableError, UndefinedOperationError
+import tqdm
 
 # Type definitions
 FullEvidenceType = np.array  # [Union[float, int, SetElement]]
@@ -294,15 +298,21 @@ def plotly_layout_3d(self) -> Dict[str, Any]:
                 "scene": {"xaxis": {"title": self.variables[0].name}, "yaxis": {"title": self.variables[1].name},
                           "zaxis": {"title": self.variables[2].name}}}
 
-    def plot(self, number_of_samples: int = 1000) -> List:
+    def plot(self, number_of_samples: int = 1000, surface=False) -> List:
         """
         Generate traces that can be plotted with plotly.
+
+        :param number_of_samples: The number of samples to draw.
+        :param surface: If True, plot the model as a surface plot.
         :return: The traces.
         """
         if len(self.variables) == 1:
             return self.plot_1d(number_of_samples)
         elif len(self.variables) == 2:
-            return self.plot_2d(number_of_samples)
+            if surface:
+                return self.plot_2d_surface(number_of_samples)
+            else:
+                return self.plot_2d(number_of_samples)
         elif len(self.variables) == 3:
             return self.plot_3d(number_of_samples)
         else:
@@ -406,7 +416,7 @@ def plot_2d(self, number_of_samples: int) -> List:
         :param number_of_samples: The number of samples to draw.
         :return: The traces.
         """
-        samples = self.sample(number_of_samples)
+        samples = self.sample(math.ceil(math.sqrt(number_of_samples)))
         likelihood = self.likelihood(samples)
         expectation = self.expectation(self.variables)
         likelihood_trace = go.Scatter(x=samples[:, 0], y=samples[:, 1], mode="markers", marker=dict(color=likelihood),
@@ -415,6 +425,120 @@ def plot_2d(self, number_of_samples: int) -> List:
                                        mode="markers", marker=dict(color=EXPECTATION_TRACE_COLOR), name=EXPECTATION_TRACE_NAME)
         return [likelihood_trace, expectation_trace] + self.multivariate_mode_traces()
 
+    def plot_2d_surface(self, number_of_samples: int) -> List:
+        """
+        Plot a two-dimensional model as a surface plot.
+
+        :param number_of_samples: The number of samples to draw.
+        :return: The traces.
+        """
+        support = self.support
+        samples = self.sample(number_of_samples)
+        likelihood = self.likelihood(samples)
+        max_likelihood = max(likelihood)
+
+
+        support_trace = self.bounding_box_trace_of_simple_event(support.bounding_box(), samples, 0.)
+        support_trace.showscale = False
+        support_trace.cmin = 0
+        support_trace.cmax = max_likelihood
+
+        expectation_trace = self.expectation_trace_2d_surface(max_likelihood * SCALING_FACTOR_FOR_EXPECTATION_IN_PLOT)
+
+        traces = [support_trace, expectation_trace]
+
+        first = True
+        for simple_set in tqdm.tqdm(support.simple_sets):
+            for i1, i2 in itertools.product(*simple_set.values()):
+                simple_event = SimpleEvent({self.variables[0]: i1, self.variables[1]: i2})
+                trace = self.plot_2d_surface_of_simple_event(simple_event, samples)
+                if not first:
+                    trace.showscale = False
+                    first = False
+                trace.cmin = 0
+                trace.cmax = max_likelihood
+                traces.append(trace)
+
+        return traces
+
+    def expectation_trace_2d_surface(self, height: float) -> go.Scatter3d:
+        expectation = self.expectation(self.variables)
+        x = expectation[self.variables[0]]
+        y = expectation[self.variables[1]]
+        return go.Scatter3d(x=[x, x], y=[y, y],
+                            z=[0, height], mode="lines+markers", name=EXPECTATION_TRACE_NAME,)
+
+
+    def bounding_box_trace_of_simple_event(self, simple_event: SimpleEvent, samples: np.array, fill_value=0.) -> go.Surface:
+        """
+        Create a bounding box trace for a simple event.
+        :param simple_event: The simple event.
+        :param samples: The samples to read from if bounds are infinite.
+        :param fill_value: The height of the box.
+
+        :return: The trace.
+        """
+        x_variable = self.variables[0]
+        y_variable = self.variables[1]
+        min_x = simple_event[x_variable].simple_sets[0].lower
+        max_x = simple_event[x_variable].simple_sets[-1].upper
+        min_y = simple_event[y_variable].simple_sets[0].lower
+        max_y = simple_event[y_variable].simple_sets[-1].upper
+
+        min_x = min_x if min_x > -np.inf else min(samples[:, 0])
+        min_x = np.nextafter(min_x, -np.inf)
+        max_x = max_x if max_x < np.inf else max(samples[:, 0])
+        max_x = np.nextafter(max_x, np.inf)
+
+        min_y = min_y if min_y > -np.inf else min(samples[:, 1])
+        min_y = np.nextafter(min_y, -np.inf)
+        max_y = max_y if max_y < np.inf else max(samples[:, 1])
+        max_y = np.nextafter(max_y, np.inf)
+
+        return go.Surface(x=[min_x, max_x], y=[min_y, max_y], z=[[fill_value, fill_value], [fill_value, fill_value]],
+                          showscale=False)
+
+
+    def plot_2d_surface_of_simple_event(self, simple_event: SimpleEvent, samples: np.array):
+        # filter samples by this event
+        samples_of_this_event = [s for s in samples if simple_event.contains(s)]
+
+        if len(samples_of_this_event) == 0:
+            return go.Surface()
+
+        samples_of_this_event = np.stack(samples_of_this_event, axis=0)
+
+        x_variable = self.variables[0]
+        y_variable = self.variables[1]
+
+        x_support: SimpleInterval = simple_event[x_variable].simple_sets[0]
+        y_support: SimpleInterval = simple_event[y_variable].simple_sets[0]
+
+        # create border points
+        min_x = x_support.lower if x_support.lower > -np.inf else min(samples_of_this_event[:, 0])
+        min_x_next_after = np.nextafter(min_x, -np.inf)
+        max_x = x_support.upper if x_support.upper < np.inf else max(samples_of_this_event[:, 0])
+        max_x_next_after = np.nextafter(max_x, np.inf)
+        min_y = y_support.lower if y_support.lower > -np.inf else min(samples_of_this_event[:, 1])
+        min_y_next_after = np.nextafter(min_y, -np.inf)
+        max_y = y_support.upper if y_support.upper < np.inf else max(samples_of_this_event[:, 1])
+        max_y_next_after = np.nextafter(max_y, np.inf)
+
+        # create x axis
+        x = samples_of_this_event[:, 0]
+        x = np.append(x, [min_x, max_x, min_x_next_after, max_x_next_after])
+        x.sort()
+
+        # create y axis
+        y = samples_of_this_event[:, 1]
+        y = np.append(y, [min_y, max_y, min_y_next_after, max_y_next_after])
+        y.sort()
+
+        meshgrid = np.array(np.meshgrid(x, y)).T.reshape(-1, 2)
+        likelihood = self.likelihood(meshgrid).reshape(len(x), len(y))
+        trace = go.Surface(z=likelihood, x=x, y=y)
+        return trace
+
     def plot_3d(self, number_of_samples: int) -> List:
         """
         Plot a three-dimensional model using samples.

test/test_jpt/test_jpt.py

@@ -13,15 +13,17 @@
 from random_events.interval import closed
 from random_events.product_algebra import SimpleEvent
 from random_events.set import SetElement
-from random_events.variable import Variable
+from random_events.variable import Variable, Continuous
 
+from probabilistic_model.distributions import GaussianDistribution
 from probabilistic_model.learning.jpt.jpt import JPT
 from probabilistic_model.learning.jpt.variables import (ScaledContinuous, infer_variables_from_dataframe, Integer,
                                                         Symbolic)
 from probabilistic_model.learning.nyga_distribution import NygaDistribution
 from probabilistic_model.probabilistic_circuit.nx.distributions import IntegerDistribution, \
-    SymbolicDistribution
-from probabilistic_model.probabilistic_circuit.nx.probabilistic_circuit import SumUnit, ProbabilisticCircuit
+    SymbolicDistribution, UnivariateContinuousLeaf
+from probabilistic_model.probabilistic_circuit.nx.probabilistic_circuit import SumUnit, ProbabilisticCircuit, \
+    ProductUnit
 
 
 class SymbolEnum(SetElement):
@@ -310,3 +312,38 @@ def test_serialization(self):
         model_ = JPT.from_json(model_)
         self.assertEqual(model, model_)
         file.close()
+
+import plotly.graph_objects as go
+
+
+class GaussianJPTTestCase(unittest.TestCase):
+
+    x: Continuous
+    y: Continuous
+
+    data: pd.DataFrame
+    multivariate_normal: ProbabilisticCircuit
+
+    @classmethod
+    def setUpClass(cls):
+        np.random.seed(69)
+
+        prod = ProductUnit()
+        prod.add_subcircuit(UnivariateContinuousLeaf(GaussianDistribution(Continuous("x"), 2, 4)))
+        prod.add_subcircuit(UnivariateContinuousLeaf(GaussianDistribution(Continuous("y"), 2, 4)))
+        cls.multivariate_normal = prod.probabilistic_circuit
+        samples = cls.multivariate_normal.sample(1000)
+        cls.data = pd.DataFrame(samples, columns=[v.name for v in cls.multivariate_normal.variables])
+
+        cls.x, cls.y = infer_variables_from_dataframe(cls.data, scale_continuous_types=False,
+                                                      min_samples_per_quantile=20)
+
+    def test_plot_2d_jpt(self):
+        model = JPT([self.x, self.y], min_samples_leaf=0.9)
+        model.fit(self.data)
+        fig = go.Figure(model.plot(500, surface=True), model.plotly_layout())
+        # fig.show()
+
+    def test_plot_2d_gaussian(self):
+        fig = go.Figure(self.multivariate_normal.plot(500, surface=True), self.multivariate_normal.plotly_layout())
+        # fig.show()

test/test_nx/test_probabilistic_circuit.py

@@ -66,12 +66,12 @@ def test_conditioning(self):
         self.assertAlmostEqual(prob, 0.375)
         conditional.plot_structure()
         # plt.show()
-    #
-    # def test_plot(self):
-    #     self.model.plot_structure()
-    #     plt.show()
-    #     fig = go.Figure(self.model.plot(600))
-    #     fig.show()
+
+    def test_plot(self):
+        self.model.plot_structure()
+        # plt.show()
+        fig = go.Figure(self.model.plot(600, surface=True))
+        # fig.show()
 
 
 if __name__ == '__main__':