Merge pull request #43 from metadsl/next-sklearn

Cleanup sklearn notebook

docs/changelog.md

@@ -29,6 +29,9 @@ _This project uses semantic versioning. Before 1.0.0, this means that every brea
   - Removes custom fork of egglog now that visualizations are in core
   - Adds int and float to string functions
   - Switches `define` to `let`
+- Tidy up notebook appearence [#43](https://github.com/metadsl/egglog-python/pull/43)
+  - Display expressions as code in Jupyter notebook
+  - Display all expressions when graphing
 
 ## 0.5.1 (2023-07-18)
 

python/egglog/bindings.pyi

@@ -13,7 +13,12 @@ class EGraph:
     def run_program(self, *commands: _Command) -> list[str]: ...
     def extract_report(self) -> Optional[_ExtractReport]: ...
     def run_report(self) -> Optional[RunReport]: ...
-    def to_graphviz_string(self) -> str: ...
+    def to_graphviz_string(
+        self,
+        *,
+        max_functions: Optional[int] = None,
+        max_calls_per_function: Optional[int] = None,
+    ) -> str: ...
     def save_object(self, __o: object, /) -> _Expr: ...
     def load_object(self, __e: _Expr, /) -> object: ...
 

python/egglog/exp/array_api.py

@@ -28,11 +28,11 @@
 
 
 def extract_py(e: Expr) -> Any:
-    print(e)
+    # print(e)
     egraph.register(e)
     egraph.run((run() * 10).saturate())
     final_object = egraph.extract(e)
-    print(f"  -> {final_object}")
+    # print(f"  -> {final_object}")
     # with egraph:
     egraph.run((run(runtime_ruleset) * 10 + run() * 10).saturate())
     # Run saturation again b/c sometimes it doesn't work the first time.
@@ -55,7 +55,7 @@ def extract_py(e: Expr) -> Any:
     # final_object = egraph.extract(egraph.extract(final_object))
     # egraph.run(run(limit=10).saturate())
 
-    print(f"     -> {egraph.extract(final_object)}\n")
+    # print(f"     -> {egraph.extract(final_object)}\n")
     res = egraph.load_object(egraph.extract(final_object.to_py()))  # type: ignore[attr-defined]
     return res
 
@@ -172,7 +172,7 @@ def _isdtype(d: DType, k1: IsDtypeKind, k2: IsDtypeKind):
     ]
 
 
-assert not bool(isdtype(DType.float32, IsDtypeKind.string("integral")))
+# assert not bool(isdtype(DType.float32, IsDtypeKind.string("integral")))
 
 
 @egraph.class_

python/egglog/runtime.py

@@ -414,6 +414,11 @@ def __str__(self) -> str:
         except black.parsing.InvalidInput:
             return pretty_expr
 
+    def _ipython_display_(self) -> None:
+        from IPython.display import Code, display
+
+        display(Code(str(self), language="python"))
+
     def __dir__(self) -> Iterable[str]:
         return list(self.__egg_decls__.get_class_decl(self.__egg_typed_expr__.tp.name).methods)
 

src/egraph.rs

@@ -5,6 +5,7 @@ use crate::error::EggResult;
 use crate::py_object_sort::PyObjectSort;
 
 use egglog::sort::Sort;
+use egglog::SerializeConfig;
 use log::info;
 use pyo3::{prelude::*, PyTraverseError, PyVisit};
 use std::path::PathBuf;
@@ -81,11 +82,24 @@ impl EGraph {
     }
 
     /// Returns the EGraph as graphviz string.
-    #[pyo3(signature = ())]
-    fn to_graphviz_string(&self) -> String {
+    #[pyo3(
+        signature = (*, max_functions=None, max_calls_per_function=None),
+        text_signature = "(self, *, max_functions=None, max_calls_per_function=None)"
+    )]
+    fn to_graphviz_string(
+        &self,
+        max_functions: Option<usize>,
+        max_calls_per_function: Option<usize>,
+    ) -> String {
         info!("Getting graphviz");
         // TODO: Expose full serialized e-graph in the future
-        self.egraph.serialize_for_graphviz().to_dot()
+        let mut serialized = self.egraph.serialize(SerializeConfig {
+            max_functions,
+            max_calls_per_function,
+            include_temporary_functions: false,
+        });
+        serialized.inline_leaves();
+        serialized.to_dot()
     }
 
     /// Register a Python object with the EGraph and return the Expr which represents it.