Expression graph autodiff extension

.gitignore

@@ -102,3 +102,6 @@ venv.bak/
 
 # mypy
 .mypy_cache/
+
+# Expression graph extension generated outputs
+math_expr/extension_outputs/

math_expr/EXTENSION_README.md

@@ -0,0 +1,91 @@
+# Expression Graph Extension
+
+This extension adds a symbolic expression graph layer on top of the existing
+`math_expr` graph language. It parses `.graph` programs into expression trees,
+converts those trees into differentiable PyTorch modules, and includes examples
+and experiments that demonstrate automatic differentiation through the parsed
+structure.
+
+## Architecture
+
+- `expression_graph.py` parses the existing DSL into `Node` trees and provides
+  string, edge-list, Mermaid, Graphviz, and optional NetworkX views.
+- `nn_nodes.py` contains PyTorch modules for the supported operations:
+  `sin`, `exp`, `ln`, `sinc`, `poly`, sum, and composition.
+- `module_from_tree.py` converts parsed expression trees into `torch.nn.Module`
+  instances.
+- `executor_adapter.py` bridges `.graph` files to the original executor without
+  modifying `executor.py`.
+- `examples/` contains small executable demonstrations.
+- `experiments/` contains the validation scripts used to reproduce the extension
+  experiments.
+
+Generated files are written to `math_expr/extension_outputs/`, which is ignored
+by git.
+
+## Dependencies
+
+The extension examples and tests require:
+
+```text
+numpy
+matplotlib
+pillow
+torch
+pytest
+```
+
+Optional visualization helpers:
+
+```text
+graphviz
+networkx
+```
+
+The experiments include a small self-contained Snake-style activation so they do
+not require the external `snake` package.
+
+## Run Tests
+
+From the repository root:
+
+```powershell
+python -m pytest math_expr/test_expression_graph.py -q
+```
+
+## Run Examples
+
+From the repository root:
+
+```powershell
+python math_expr/examples/exp_expression_graph_demo.py
+python math_expr/examples/exp_expression_graph_ad.py
+python math_expr/examples/exp_executor_adapter_demo.py
+python math_expr/examples/exp_expression_graph_visualize.py
+```
+
+The visualization example writes Mermaid files and, when Graphviz is available,
+PNG files under:
+
+```text
+math_expr/extension_outputs/figures/
+```
+
+## Run Experiments
+
+From the repository root:
+
+```powershell
+python math_expr/experiments/baseline_sin_plus_ln.py
+python math_expr/experiments/snake_sin_plus_ln_ln_focused.py
+python math_expr/experiments/snake_exp_plus_exp2_exp_focused.py
+```
+
+The experiments create their output directories automatically and write figures
+and metric files under:
+
+```text
+math_expr/extension_outputs/
+```
+
+These outputs are reproducible artifacts and should not be committed.

math_expr/examples/exp_executor_adapter_demo.py

@@ -0,0 +1,45 @@
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import numpy as np
+
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+try:
+    from math_expr.executor_adapter import evaluate_graph_with_tree, load_program_from_graph_file
+    from math_expr.expression_graph import tree_to_edges, tree_to_string
+except ImportError:  # pragma: no cover - supports running from this directory
+    from executor_adapter import evaluate_graph_with_tree, load_program_from_graph_file
+    from expression_graph import tree_to_edges, tree_to_string
+
+
+MATH_EXPR_DIR = Path(__file__).resolve().parents[1]
+DEMO_GRAPH = MATH_EXPR_DIR / "graphs" / "sine_in_exp.graph"
+
+
+def main() -> None:
+    program = load_program_from_graph_file(DEMO_GRAPH)
+    x = np.linspace(-1.0, 1.0, 5)
+    result, tree = evaluate_graph_with_tree(DEMO_GRAPH, x=x)
+
+    print("Input .graph program:")
+    print(program.strip())
+    print()
+    print("Pretty expression:")
+    print(tree_to_string(tree))
+    print()
+    print("Tree edges:")
+    for parent, child in tree_to_edges(tree):
+        print(f"{parent} -> {child}")
+    print()
+    print("Original executor result successfully obtained:")
+    print(result.execution_successful)
+
+
+if __name__ == "__main__":
+    main()

math_expr/examples/exp_expression_graph_ad.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+import torch
+
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+try:
+    from math_expr.expression_graph import Node, parse_program_to_tree, tree_to_string
+    from math_expr.module_from_tree import node_to_module
+except ImportError:  # pragma: no cover - supports running from this directory
+    from expression_graph import Node, parse_program_to_tree, tree_to_string
+    from module_from_tree import node_to_module
+
+
+def evaluate_tree_torch(node: Node, x: torch.Tensor) -> torch.Tensor:
+    op_type = node.op_type.upper()
+    if op_type == "X":
+        return x
+    if op_type == "SUM":
+        return sum(evaluate_tree_torch(child, x) for child in node.children)
+
+    child_value = evaluate_tree_torch(node.children[0], x)
+    if op_type == "SIN":
+        return torch.sin(child_value)
+    if op_type == "EXP":
+        return torch.exp(child_value)
+    if op_type in {"LN", "LOG"}:
+        return torch.log(child_value)
+    if op_type == "SINC":
+        return torch.sinc(child_value)
+    if op_type == "POLY":
+        result = torch.zeros_like(child_value)
+        for power, coefficient in enumerate(node.params or []):
+            result = result + float(coefficient) * child_value.pow(power)
+        return result
+
+    raise ValueError(f"Unsupported operation for torch evaluation: {node.op_type}")
+
+
+def main() -> None:
+    program = "sin(), exp()"
+    tree = parse_program_to_tree(program)
+    model = node_to_module(tree)
+
+    x = torch.linspace(-1.0, 1.0, steps=5, requires_grad=True)
+    y = model(x)
+    loss = (y**2).mean()
+    loss.backward()
+
+    print("Expression:", tree_to_string(tree))
+    print("Loss:", loss.item())
+    print("x.grad:", x.grad)
+    print("Gradients flow through the module graph that mirrors the expression tree.")
+
+
+if __name__ == "__main__":
+    main()

math_expr/examples/exp_expression_graph_demo.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+try:
+    from math_expr.expression_graph import parse_program_to_tree, tree_to_edges, tree_to_string
+except ImportError:  # pragma: no cover - supports running from this directory
+    from expression_graph import parse_program_to_tree, tree_to_edges, tree_to_string
+
+
+EXAMPLES = [
+    "sin(), exp()",
+    "sin(), poly([0,3]), exp()\nsinc()",
+    "poly( [0, 0.2] )",
+]
+
+
+def main() -> None:
+    for program in EXAMPLES:
+        tree = parse_program_to_tree(program)
+        print("Program:")
+        print(program)
+        print()
+        print("Expression:")
+        print(tree_to_string(tree))
+        print()
+        print("Edges:")
+        for parent, child in tree_to_edges(tree):
+            print(f"{parent} -> {child}")
+        print()
+
+
+if __name__ == "__main__":
+    main()

math_expr/examples/exp_expression_graph_visualize.py

@@ -0,0 +1,99 @@
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+
+REPO_ROOT = Path(__file__).resolve().parents[2]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+try:
+    from math_expr.expression_graph import (
+        export_tree_to_dot,
+        parse_program_to_tree,
+        tree_to_mermaid,
+        tree_to_string,
+    )
+except ImportError:  # pragma: no cover - supports running from this directory
+    from expression_graph import (
+        export_tree_to_dot,
+        parse_program_to_tree,
+        tree_to_mermaid,
+        tree_to_string,
+    )
+
+
+OUTPUT_DIR = Path(__file__).resolve().parents[1] / "extension_outputs" / "figures"
+
+EXAMPLES = [
+    ("sin_exp", "sin(), exp()", "sin_exp_tree.png"),
+    (
+        "sin_poly_exp_plus_sinc",
+        "sin(), poly([0,3]), exp()\nsinc()",
+        "sin_poly_exp_plus_sinc_tree.png",
+    ),
+    ("poly", "poly([0, 0.2])", "poly_tree.png"),
+]
+
+
+def ad_flow_mermaid_for_sin_exp() -> str:
+    return "\n".join(
+        [
+            "graph TD",
+            '  X["x"]',
+            '  EXP["exp"]',
+            '  SIN["sin"]',
+            '  Y["output"]',
+            '  LOSS["loss"]',
+            "  X --> EXP --> SIN --> Y --> LOSS",
+            "  LOSS -. backward / gradients .-> Y",
+            "  Y -.-> SIN",
+            "  SIN -.-> EXP",
+            "  EXP -.-> X",
+        ]
+    )
+
+
+def main() -> None:
+    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+
+    for name, program, png_name in EXAMPLES:
+        tree = parse_program_to_tree(program)
+        mermaid = tree_to_mermaid(tree)
+        mermaid_path = OUTPUT_DIR / f"{name}_tree.mmd"
+
+        print("Program:")
+        print(program)
+        print()
+        print("Expression:")
+        print(tree_to_string(tree))
+        print()
+        print("Mermaid:")
+        print(mermaid)
+        print()
+
+        mermaid_path.write_text(mermaid + "\n", encoding="utf-8")
+
+        try:
+            generated_path = export_tree_to_dot(
+                tree,
+                str(OUTPUT_DIR / png_name),
+                title=tree_to_string(tree),
+            )
+            print(f"Graphviz PNG: {generated_path}")
+        except RuntimeError as exc:
+            print(f"Warning: {exc}")
+        print()
+
+    ad_flow = ad_flow_mermaid_for_sin_exp()
+    ad_flow_path = OUTPUT_DIR / "sin_exp_ad_flow.mmd"
+    ad_flow_path.write_text(ad_flow + "\n", encoding="utf-8")
+    print("AD flow Mermaid:")
+    print(ad_flow)
+    print()
+    print(f"AD flow saved to: {ad_flow_path}")
+
+
+if __name__ == "__main__":
+    main()