diff --git a/src/evermore/effect.py b/src/evermore/effect.py index bee7bce..7c37597 100644 --- a/src/evermore/effect.py +++ b/src/evermore/effect.py @@ -9,7 +9,7 @@ from evermore.custom_types import OffsetAndScale from evermore.parameter import Parameter -from evermore.visualization import SupportTreescope +from evermore.visualization import SupportsTreescope __all__ = [ "Effect", @@ -24,7 +24,7 @@ def __dir__(): return __all__ -class Effect(eqx.Module, SupportTreescope): +class Effect(eqx.Module, SupportsTreescope): @abc.abstractmethod def __call__(self, parameter: PyTree[Parameter], hist: Array) -> OffsetAndScale: ... diff --git a/src/evermore/modifier.py b/src/evermore/modifier.py index 22e96c0..9da092b 100644 --- a/src/evermore/modifier.py +++ b/src/evermore/modifier.py @@ -15,7 +15,7 @@ from evermore.effect import DEFAULT_EFFECT from evermore.parameter import Parameter from evermore.util import tree_stack -from evermore.visualization import SupportTreescope +from evermore.visualization import SupportsTreescope if TYPE_CHECKING: from evermore.effect import Effect @@ -59,7 +59,7 @@ def __matmul__(self: ModifierLike, other: ModifierLike) -> Compose: return Compose(self, other) -class ModifierBase(ApplyFn, MatMulCompose, AbstractModifier, SupportTreescope): +class ModifierBase(ApplyFn, MatMulCompose, AbstractModifier, SupportsTreescope): """ This serves as a base class for all modifiers. It automatically implements the __call__ method to apply the scale factors to the hist array diff --git a/src/evermore/parameter.py b/src/evermore/parameter.py index dbb7816..79db732 100644 --- a/src/evermore/parameter.py +++ b/src/evermore/parameter.py @@ -12,7 +12,7 @@ from evermore.custom_types import PDFLike from evermore.pdf import Normal, Poisson from evermore.util import filter_tree_map -from evermore.visualization import SupportTreescope +from evermore.visualization import SupportsTreescope if TYPE_CHECKING: from evermore.modifier import Modifier @@ -34,7 +34,7 @@ def __dir__(): return __all__ -class Parameter(eqx.Module, SupportTreescope): +class Parameter(eqx.Module, SupportsTreescope): """ Implementation of a general Parameter class. The class is used to define the parameters of a statistical model. Key is the value attribute, which holds the actual value of the parameter. In additon, diff --git a/src/evermore/pdf.py b/src/evermore/pdf.py index fbbf6ee..3067b7e 100644 --- a/src/evermore/pdf.py +++ b/src/evermore/pdf.py @@ -8,7 +8,7 @@ from jax.scipy.special import gammaln, xlogy from jaxtyping import Array, PRNGKeyArray -from evermore.visualization import SupportTreescope +from evermore.visualization import SupportsTreescope __all__ = [ "PDF", @@ -21,7 +21,7 @@ def __dir__(): return __all__ -class PDF(eqx.Module, SupportTreescope): +class PDF(eqx.Module, SupportsTreescope): @abstractmethod def log_prob(self, x: Array) -> Array: ... diff --git a/src/evermore/staterror.py b/src/evermore/staterror.py index bcf1916..8c79186 100644 --- a/src/evermore/staterror.py +++ b/src/evermore/staterror.py @@ -14,7 +14,7 @@ from evermore.parameter import NormalParameter, Parameter from evermore.pdf import Poisson from evermore.util import sum_over_leaves -from evermore.visualization import SupportTreescope +from evermore.visualization import SupportsTreescope __all__ = [ "StatErrors", @@ -25,7 +25,7 @@ def __dir__(): return __all__ -class StatErrors(eqx.Module, SupportTreescope): +class StatErrors(eqx.Module, SupportsTreescope): """ Create staterror (barlow-beeston) parameters. diff --git a/src/evermore/visualization.py b/src/evermore/visualization.py index b8d9845..f878b75 100644 --- a/src/evermore/visualization.py +++ b/src/evermore/visualization.py @@ -12,7 +12,7 @@ ) -class SupportTreescope: +class SupportsTreescope: def __treescope_repr__( self, path: str, @@ -56,7 +56,6 @@ def handle_evermore_classes( # get colors for the background of the tree node def _treescope_color(node) -> str: """Returns the color of the tree node.""" - from treescope import formatting_util type_string = type(node).__module__ + "." + type(node).__qualname__ return formatting_util.color_from_string(type_string)