From 837560bce36a1f5f917ca7a778ac15d3908db169 Mon Sep 17 00:00:00 2001
From: Peter Fackeldey <peter.fackeldey@rwth-aachen.de>
Date: Tue, 30 Apr 2024 12:26:05 +0200
Subject: [PATCH] fix imports

---
 src/evermore/visualization.py | 51 +++++++++++++++++++++++------------
 1 file changed, 34 insertions(+), 17 deletions(-)

diff --git a/src/evermore/visualization.py b/src/evermore/visualization.py
index 82834cf..3d0a386 100644
--- a/src/evermore/visualization.py
+++ b/src/evermore/visualization.py
@@ -7,8 +7,25 @@
 import jax.tree_util as jtu
 from jaxtyping import Array, PyTree
 
-import evermore as evm
 from evermore.custom_types import ModifierLike, PDFLike
+from evermore.effect import (
+    AsymmetricExponential,
+    Effect,
+    Identity,
+    Linear,
+    VerticalTemplateMorphing,
+)
+from evermore.modifier import (
+    BooleanMask,
+    Compose,
+    Modifier,
+    Transform,
+    TransformOffset,
+    TransformScale,
+    Where,
+)
+from evermore.parameter import NormalParameter, Parameter
+from evermore.pdf import Normal, Poisson
 
 __all__ = [
     "display",
@@ -21,22 +38,22 @@ def __dir__():
 
 EVERMORE_CLASSES = set(
     {
-        evm.Parameter,
-        evm.NormalParameter,
-        evm.effect.Effect,
-        evm.effect.Identity,
-        evm.effect.Linear,
-        evm.effect.AsymmetricExponential,
-        evm.effect.VerticalTemplateMorphing,
-        evm.modifier.Modifier,
-        evm.modifier.Compose,
-        evm.modifier.Where,
-        evm.modifier.BooleanMask,
-        evm.modifier.Transform,
-        evm.modifier.TransformScale,
-        evm.modifier.TransformOffset,
-        evm.pdf.Normal,
-        evm.pdf.Poisson,
+        Parameter,
+        NormalParameter,
+        Effect,
+        Identity,
+        Linear,
+        AsymmetricExponential,
+        VerticalTemplateMorphing,
+        Modifier,
+        Compose,
+        Where,
+        BooleanMask,
+        Transform,
+        TransformScale,
+        TransformOffset,
+        Normal,
+        Poisson,
         ModifierLike,
         PDFLike,
     }