update evermore visualization with penzai

docs/building_blocks.md

@@ -125,7 +125,7 @@ Correlate a Parameter
 :::{admonition} Inspect `evm.Parameters` with `penzai`
 :class: tip dropdown
 
-Inspect a (PyTree of) `evm.Parameters` with [`penzai`'s treescope](https://penzai.readthedocs.io/en/stable/notebooks/treescope_prettyprinting.html) visualization in IPython or Colab notebooks.
+Inspect a (PyTree of) `evm.Parameters` with [`penzai`'s treescope](https://penzai.readthedocs.io/en/stable/notebooks/treescope_prettyprinting.html) visualization in IPython or Colab notebooks (see <project:#penzai-visualization> for more information).
 You can even add custom visualizers, such as:
 
 ```{code-block} python
@@ -134,7 +134,9 @@ import evermore as evm
 
 tree = {"a": evm.NormalParameter(), "b": evm.NormalParameter()}
 
-evm.visualization.display(tree)
+with pz.ts.active_autovisualizer.set_scoped(pz.ts.ArrayAutovisualizer()):
+    pz_tree = evm.visualization.convert_tree_to_penzai(tree)
+    pz.ts.display(pz_tree)
 ```
 :::
 
@@ -273,6 +275,7 @@ Multiple modifiers should be combined using `evm.modifier.Compose` or the `@` op
 import jax
 import jax.numpy as jnp
 import evermore as evm
+from penzai import pz
 
 
 jax.config.update("jax_enable_x64", True)
@@ -290,5 +293,7 @@ modifier2 = param.scale_log(up=1.1, down=0.9)
 (modifier1 @ modifier2)(jnp.array([10, 20, 30]))
 # -> Array([10.259877, 20.500944, 30.760822], dtype=float32)
 
-evm.visualization.display(modifier1 @ modifier2)
+with pz.ts.active_autovisualizer.set_scoped(pz.ts.ArrayAutovisualizer()):
+    pz_tree = evm.visualization.convert_tree_to_penzai(modifier1 @ modifier2)
+    pz.ts.display(pz_tree)
 ```

docs/tips_and_tricks.md

@@ -15,16 +15,17 @@ kernelspec:
 
 Here are some advanced tips and tricks.
 
-
+(penzai-visualization)=
 ## penzai Visualization
 
-Use `penzai` to visualize evermore components!
+evermore components can be visualized with `penzai`. Convert the corresponding PyTree using `evermore.visualization.convert_tree_to_penzai` and use it with `penzai`. In IPython notebooks you can display the tree using `penzai.ts.display`.
 
 ```{code-cell} ipython3
 import jax
 import jax.numpy as jnp
 import evermore as evm
 import equinox as eqx
+from penzai import pz
 
 jax.config.update("jax_enable_x64", True)
 
@@ -49,10 +50,20 @@ composition = evm.modifier.Compose(
     evm.Modifier(parameter=sigma1, effect=evm.effect.AsymmetricExponential(up=1.2, down=0.8)),
 )
 
-evm.visualization.display(composition)
+with pz.ts.active_autovisualizer.set_scoped(pz.ts.ArrayAutovisualizer()):
+    pz_tree = evm.visualization.convert_tree_to_penzai(composition)
+    pz.ts.display(pz_tree)
 ```
 
+You can also save the tree to an HTML file.
+```{code-cell} python
+with pz.ts.active_autovisualizer.set_scoped(pz.ts.ArrayAutovisualizer()):
+    pz_tree = evm.visualization.convert_tree_to_penzai(composition)
+    contents = pz.ts.render_to_html(pz_tree, roundtrip_mode=True)
 
+with open("composition.html", "w") as f:
+    f.write(contents)
+```
 
 ## Parameter Partitioning
 
@@ -110,7 +121,9 @@ rng_keys = jax.random.split(rng_key, 100)
 
 vec_sample = jax.vmap(evm.parameter.sample, in_axes=(None, 0))
 
-evm.visualization.display(vec_sample(params, rng_keys))
+with pz.ts.active_autovisualizer.set_scoped(pz.ts.ArrayAutovisualizer()):
+    pz_tree = evm.visualization.convert_tree_to_penzai(vec_sample(params, rng_keys))
+    pz.ts.display(pz_tree)
 ```
 
 Many minimizers from the JAX ecosystem are e.g. batchable (`optax`, `optimistix`), which allows you vectorize _full fits_, e.g., for embarrassingly parallel likleihood profiles.

src/evermore/visualization.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import dataclasses
-import importlib.util
 import threading
 from typing import Any
 
@@ -29,7 +28,7 @@
 from evermore.pdf import Normal, Poisson
 
 __all__ = [
-    "display",
+    "convert_tree_to_penzai",
 ]
 
 
@@ -69,47 +68,7 @@ class EvermoreClassesContext(threading.local):
 )
 
 
-def display(tree: PyTree) -> None:
-    """
-    Visualize PyTrees of evermore components with penzai in a notebook.
-
-    Usage:
-
-        .. code-block:: python
-
-            import evermore as evm
-
-
-            tree = ...
-            evm.visualization.display(tree)
-    """
-    penzai_installed = importlib.util.find_spec("penzai") is not None
-
-    if not penzai_installed:
-        msg = "install 'penzai' with:\n\n"
-        msg += "\tpython -m pip install penzai[notebook]"
-        raise ModuleNotFoundError(msg)
-
-    try:
-        from IPython import get_ipython
-
-        in_ipython = get_ipython() is not None
-    except ImportError:
-        in_ipython = False
-
-    if not in_ipython:
-        print(tree)
-        return
-
-    # now we can pretty-print
-    from penzai import pz
-
-    with pz.ts.active_autovisualizer.set_scoped(pz.ts.ArrayAutovisualizer()):
-        pz_tree = convert_tree(tree)
-        pz.ts.display(pz_tree)
-
-
-def convert_tree(tree: PyTree) -> PyTree:
+def convert_tree_to_penzai(tree: PyTree) -> PyTree:
     from functools import partial
 
     for cls in Context.cls_types: