Skip to content

Commit

Permalink
Some refactoring
Browse files Browse the repository at this point in the history
  • Loading branch information
jrzkaminski committed Jul 5, 2024
1 parent 9795fef commit 3e3405f
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 11 deletions.
10 changes: 0 additions & 10 deletions bamt/core/nodes/root_nodes/continuous_node.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
23 changes: 22 additions & 1 deletion bamt/core/nodes/root_nodes/discrete_node.py
Original file line number Diff line number Diff line change
@@ -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)

0 comments on commit 3e3405f

Please sign in to comment.