From 3e3405ff997d3a2abf82dbdda0e43acf26e5308e Mon Sep 17 00:00:00 2001 From: jrzkaminski Date: Fri, 5 Jul 2024 18:39:53 +0300 Subject: [PATCH] Some refactoring --- bamt/core/nodes/root_nodes/continuous_node.py | 10 -------- bamt/core/nodes/root_nodes/discrete_node.py | 23 ++++++++++++++++++- 2 files changed, 22 insertions(+), 11 deletions(-) diff --git a/bamt/core/nodes/root_nodes/continuous_node.py b/bamt/core/nodes/root_nodes/continuous_node.py index 82edcf5..9e07fd9 100644 --- a/bamt/core/nodes/root_nodes/continuous_node.py +++ b/bamt/core/nodes/root_nodes/continuous_node.py @@ -42,13 +42,3 @@ def __str__(self) -> str: str: The string representation of the node. """ return "Continuous Node with " + str(self._distribution) - - @property - def distribution(self) -> ContinuousDistribution: - """ - Get the continuous distribution of this node. - - Returns: - ContinuousDistribution: The continuous distribution. - """ - return self._distribution diff --git a/bamt/core/nodes/root_nodes/discrete_node.py b/bamt/core/nodes/root_nodes/discrete_node.py index d9b04c8..fca7ccf 100644 --- a/bamt/core/nodes/root_nodes/discrete_node.py +++ b/bamt/core/nodes/root_nodes/discrete_node.py @@ -1,6 +1,27 @@ +from typing import Optional + from .root_node import RootNode +from bamt.core.node_models import EmpiricalDistribution class DiscreteNode(RootNode): - def __init__(self): + def __init__(self, distribution: Optional[EmpiricalDistribution] = None): + """ + Initialize the DisscreteNode with an optional EmpiricalDistribution. + + Args: + distribution (Optional[ContinuousDistribution]): A ContinuousDistribution object. + """ super().__init__() + self._distribution = ( + distribution if distribution is not None else EmpiricalDistribution + ) + + def __str__(self) -> str: + """ + Return the string representation of the Discrete node. + + Returns: + str: The string representation of the node. + """ + return "Discrete Node with " + str(self._distribution)