diff --git a/.gitignore b/.gitignore index 894a44c..443b523 100644 --- a/.gitignore +++ b/.gitignore @@ -102,3 +102,6 @@ venv.bak/ # mypy .mypy_cache/ + +# Expression graph extension generated outputs +math_expr/extension_outputs/ diff --git a/math_expr/EXTENSION_README.md b/math_expr/EXTENSION_README.md new file mode 100644 index 0000000..7bf46ab --- /dev/null +++ b/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. diff --git a/math_expr/examples/exp_executor_adapter_demo.py b/math_expr/examples/exp_executor_adapter_demo.py new file mode 100644 index 0000000..2aa67cd --- /dev/null +++ b/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() diff --git a/math_expr/examples/exp_expression_graph_ad.py b/math_expr/examples/exp_expression_graph_ad.py new file mode 100644 index 0000000..38f5bf0 --- /dev/null +++ b/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() diff --git a/math_expr/examples/exp_expression_graph_demo.py b/math_expr/examples/exp_expression_graph_demo.py new file mode 100644 index 0000000..a315036 --- /dev/null +++ b/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() diff --git a/math_expr/examples/exp_expression_graph_visualize.py b/math_expr/examples/exp_expression_graph_visualize.py new file mode 100644 index 0000000..06547c3 --- /dev/null +++ b/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() diff --git a/math_expr/executor_adapter.py b/math_expr/executor_adapter.py new file mode 100644 index 0000000..dff563c --- /dev/null +++ b/math_expr/executor_adapter.py @@ -0,0 +1,114 @@ +from __future__ import annotations + +import os +import shutil +import tempfile +from contextlib import contextmanager +from dataclasses import dataclass +from os import PathLike +from pathlib import Path +from typing import Any, Iterator + + +DEFAULT_MPL_CONFIG_DIR = Path(tempfile.gettempdir()) / "math_expr_matplotlib" +DEFAULT_MPL_CONFIG_DIR.mkdir(parents=True, exist_ok=True) +os.environ.setdefault("MPLCONFIGDIR", str(DEFAULT_MPL_CONFIG_DIR)) + +try: + from math_expr.executor import execute as original_execute +except ImportError: # pragma: no cover - supports running from inside math_expr/ + from executor import execute as original_execute + +try: + from math_expr.expression_graph import Node, parse_program_to_tree +except ImportError: # pragma: no cover - supports running from inside math_expr/ + from expression_graph import Node, parse_program_to_tree + + +@dataclass(frozen=True) +class ExecutorAdapterResult: + """Execution metadata returned by the adapter wrapper.""" + + graph_path: Path + program_name: str + original_return: Any + execution_successful: bool + + +def load_program_from_graph_file(path: str | PathLike[str]) -> str: + return Path(path).read_text(encoding="utf-8") + + +def evaluate_graph_with_tree( + graph_path: str | PathLike[str], + x=None, +) -> tuple[ExecutorAdapterResult, Node]: + """ + Parse a .graph file to an expression tree and evaluate it with the original + executor without modifying the existing executor implementation. + + The original executor expects to open graphs/.graph relative to the + process cwd and writes plot files as a side effect. This wrapper creates a + temporary compatible workspace, calls the unchanged executor there, and then + returns execution metadata together with the parsed tree. + """ + graph_file = Path(graph_path).resolve() + program = load_program_from_graph_file(graph_file) + tree = parse_program_to_tree(program) + + with _original_executor_workspace(graph_file) as program_name: + original_return = original_execute(program_name, x) + + result = ExecutorAdapterResult( + graph_path=graph_file, + program_name=program_name, + original_return=original_return, + execution_successful=True, + ) + return result, tree + + +def execute_with_tree( + graph_path: str | PathLike[str], + x=None, + *, + return_tree: bool = True, +): + """ + Read a .graph file, parse it to an expression tree, then call the original + executor without changing its implementation. + + By default this returns (executor_result, tree). Pass return_tree=False for + compatibility with executor-style callers that only expect the executor + result. + """ + result, tree = evaluate_graph_with_tree(graph_path, x) + + if return_tree: + return result.original_return, tree + return result.original_return + + +@contextmanager +def _original_executor_workspace(graph_file: Path) -> Iterator[str]: + """Create the graphs/.graph layout expected by original_execute.""" + with tempfile.TemporaryDirectory(prefix="executor_adapter_") as workspace: + workspace_path = Path(workspace) + graphs_dir = workspace_path / "graphs" + mpl_config_dir = workspace_path / "matplotlib" + graphs_dir.mkdir(parents=True) + mpl_config_dir.mkdir(parents=True) + shutil.copy2(graph_file, graphs_dir / graph_file.name) + + previous_cwd = Path.cwd() + previous_mpl_config_dir = os.environ.get("MPLCONFIGDIR") + try: + os.chdir(workspace_path) + os.environ["MPLCONFIGDIR"] = str(mpl_config_dir) + yield graph_file.stem + finally: + os.chdir(previous_cwd) + if previous_mpl_config_dir is None: + os.environ.pop("MPLCONFIGDIR", None) + else: + os.environ["MPLCONFIGDIR"] = previous_mpl_config_dir diff --git a/math_expr/experiments/_common.py b/math_expr/experiments/_common.py new file mode 100644 index 0000000..acfe656 --- /dev/null +++ b/math_expr/experiments/_common.py @@ -0,0 +1,81 @@ +from __future__ import annotations + +import os +import sys +from pathlib import Path + +import torch +from torch import nn + + +MATH_EXPR_DIR = Path(__file__).resolve().parents[1] +REPO_ROOT = MATH_EXPR_DIR.parent +OUTPUT_ROOT = MATH_EXPR_DIR / "extension_outputs" +FIGURES_DIR = OUTPUT_ROOT / "figures" +RESULTS_DIR = OUTPUT_ROOT / "results" +MPL_CONFIG_DIR = OUTPUT_ROOT / "matplotlib" + + +def configure_output_environment() -> None: + if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + MPL_CONFIG_DIR.mkdir(parents=True, exist_ok=True) + os.environ.setdefault("MPLCONFIGDIR", str(MPL_CONFIG_DIR)) + + +class SnakeActivation(nn.Module): + """Small self-contained Snake-style activation used by the experiments.""" + + def __init__(self, features: int): + super().__init__() + self.alpha = nn.Parameter(torch.ones(features)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + alpha = self.alpha.clamp_min(1e-6) + return x + torch.sin(alpha * x).pow(2) / alpha + + +class FunctionApproximator(nn.Module): + def __init__(self) -> None: + super().__init__() + self.model = nn.Sequential( + nn.Linear(1, 64), + SnakeActivation(64), + nn.Linear(64, 64), + SnakeActivation(64), + nn.Linear(64, 64), + SnakeActivation(64), + nn.Linear(64, 1), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.model(x) + + +class SmallFunctionApproximator(nn.Module): + def __init__(self) -> None: + super().__init__() + self.model = nn.Sequential( + nn.Linear(1, 32), + SnakeActivation(32), + nn.Linear(32, 32), + SnakeActivation(32), + nn.Linear(32, 1), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.model(x) + + +class DoubleSin(nn.Module): + def __init__(self, small: bool = True) -> None: + super().__init__() + model_class = SmallFunctionApproximator if small else FunctionApproximator + self.model1 = model_class() + self.model2 = model_class() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.model1(x) + self.model2(2 * x) + + def get_internal(self) -> list[nn.Module]: + return [self.model1, self.model2] diff --git a/math_expr/experiments/baseline_sin_plus_ln.py b/math_expr/experiments/baseline_sin_plus_ln.py new file mode 100644 index 0000000..a97f8e3 --- /dev/null +++ b/math_expr/experiments/baseline_sin_plus_ln.py @@ -0,0 +1,198 @@ +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path + +from _common import ( + FIGURES_DIR, + RESULTS_DIR, + SmallFunctionApproximator, + configure_output_environment, +) + +configure_output_environment() + +import matplotlib.pyplot as plt +import torch +from torch import nn +from torch.utils.data import DataLoader, TensorDataset + + +FIGURE_PATH = FIGURES_DIR / "baseline_sin_plus_ln.png" +METRICS_PATH = RESULTS_DIR / "baseline_metrics.json" +COMPARISON_PATH = RESULTS_DIR / "baseline_comparison.md" +GRAPH_GUIDED_MSE_F = 0.00035406 + + +@dataclass(frozen=True) +class TrainingConfig: + x_min: float = 0.1 + x_max: float = 10.0 + train_points: int = 8192 + eval_points: int = 1000 + batch_size: int = 64 + epochs: int = 50 + lr: float = 1e-3 + seed: int = 42 + + +def build_dataset_sin_plus_ln(config: TrainingConfig) -> TensorDataset: + x = torch.linspace(config.x_min, config.x_max, config.train_points).unsqueeze(1) + y_f = torch.sin(x) + torch.log(x) + return TensorDataset(x, y_f) + + +def train_baseline( + model: SmallFunctionApproximator, + train_loader: DataLoader, + config: TrainingConfig, +) -> list[float]: + loss_fn = nn.MSELoss() + optimizer = torch.optim.Adam(model.parameters(), lr=config.lr) + history: list[float] = [] + + for epoch in range(1, config.epochs + 1): + total_loss = 0.0 + batches = 0 + + for x, y_f in train_loader: + optimizer.zero_grad() + f_hat = model(x) + loss = loss_fn(f_hat, y_f) + loss.backward() + optimizer.step() + + total_loss += loss.item() + batches += 1 + + epoch_loss = total_loss / batches + history.append(epoch_loss) + print(f"Epoch {epoch:03d}/{config.epochs} loss={epoch_loss:.6f}") + + return history + + +def evaluate_baseline( + model: SmallFunctionApproximator, + config: TrainingConfig, +) -> tuple[dict[str, float], torch.Tensor, torch.Tensor, torch.Tensor]: + model.eval() + x_eval = torch.linspace(config.x_min, config.x_max, config.eval_points).unsqueeze(1) + + with torch.no_grad(): + f_true = torch.sin(x_eval) + torch.log(x_eval) + f_hat = model(x_eval) + metrics = {"mse_f": nn.functional.mse_loss(f_hat, f_true).item()} + + return metrics, x_eval, f_true, f_hat + + +def plot_baseline_results( + x_eval: torch.Tensor, + f_true: torch.Tensor, + f_hat: torch.Tensor, + output_path: Path = FIGURE_PATH, +) -> None: + output_path.parent.mkdir(parents=True, exist_ok=True) + + x_np = x_eval.squeeze(1).cpu().numpy() + f_true_np = f_true.squeeze(1).cpu().numpy() + f_hat_np = f_hat.squeeze(1).cpu().numpy() + + fig, ax = plt.subplots(1, 1, figsize=(7.2, 4.6), dpi=160) + ax.plot(x_np, f_true_np, linewidth=2.0, label="true f(x) = sin(x) + ln(x)") + ax.plot(x_np, f_hat_np, "--", linewidth=2.0, label="predicted f_hat(x)") + ax.set_title("Baseline Neural Function Approximation") + ax.set_xlabel("x") + ax.set_ylabel("f(x)") + ax.grid(True, alpha=0.25) + ax.legend() + + fig.tight_layout() + fig.savefig(output_path, bbox_inches="tight") + plt.close(fig) + + +def save_metrics(metrics: dict[str, float], output_path: Path = METRICS_PATH) -> None: + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text(json.dumps(metrics, indent=2) + "\n", encoding="utf-8") + + +def build_comparison(metrics: dict[str, float]) -> tuple[str, str]: + baseline_mse = metrics["mse_f"] + table = "\n".join( + [ + "| Method | MSE |", + "| ------------------------------------- | ---------- |", + f"| Baseline Neural Approximation | {baseline_mse:.8f} |", + f"| Expression Graph Guided Approximation | {GRAPH_GUIDED_MSE_F:.8f} |", + ] + ) + + if baseline_mse > GRAPH_GUIDED_MSE_F: + absolute_gain = baseline_mse - GRAPH_GUIDED_MSE_F + relative_gain = absolute_gain / baseline_mse * 100.0 + interpretation = ( + "The graph-guided method improved performance by " + f"{absolute_gain:.8f} MSE, a {relative_gain:.2f}% reduction versus the baseline." + ) + elif baseline_mse < GRAPH_GUIDED_MSE_F: + absolute_delta = GRAPH_GUIDED_MSE_F - baseline_mse + relative_delta = absolute_delta / GRAPH_GUIDED_MSE_F * 100.0 + interpretation = ( + "The graph-guided method did not improve performance; the baseline was better by " + f"{absolute_delta:.8f} MSE, a {relative_delta:.2f}% reduction " + "versus the graph-guided result." + ) + else: + interpretation = "Both methods achieve similar accuracy." + + return table, interpretation + + +def save_comparison( + table: str, + interpretation: str, + output_path: Path = COMPARISON_PATH, +) -> None: + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text(f"{table}\n\n{interpretation}\n", encoding="utf-8") + + +def main() -> None: + config = TrainingConfig() + torch.manual_seed(config.seed) + + dataset = build_dataset_sin_plus_ln(config) + train_loader = DataLoader(dataset, batch_size=config.batch_size, shuffle=True) + model = SmallFunctionApproximator() + + print("Model:") + print(model) + print() + print("Training baseline with loss = MSE(f_hat, sin(x) + ln(x))") + + train_baseline(model, train_loader, config) + metrics, x_eval, f_true, f_hat = evaluate_baseline(model, config) + plot_baseline_results(x_eval, f_true, f_hat) + save_metrics(metrics) + + table, interpretation = build_comparison(metrics) + save_comparison(table, interpretation) + + print() + print("Evaluation metrics:") + print(f"mse_f={metrics['mse_f']:.8f}") + print(f"Saved figure: {FIGURE_PATH}") + print(f"Saved metrics: {METRICS_PATH}") + print(f"Saved comparison: {COMPARISON_PATH}") + print() + print(table) + print() + print("Interpretation:") + print(interpretation) + + +if __name__ == "__main__": + main() diff --git a/math_expr/experiments/snake_exp_plus_exp2_exp_focused.py b/math_expr/experiments/snake_exp_plus_exp2_exp_focused.py new file mode 100644 index 0000000..543290c --- /dev/null +++ b/math_expr/experiments/snake_exp_plus_exp2_exp_focused.py @@ -0,0 +1,235 @@ +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path + +from _common import DoubleSin, FIGURES_DIR, RESULTS_DIR, configure_output_environment + +configure_output_environment() + +import matplotlib.pyplot as plt +import torch +from torch import nn +from torch.utils.data import DataLoader, TensorDataset + +from math_expr.expression_graph import Node, parse_program_to_tree, tree_to_string + + +FIGURE_PATH = FIGURES_DIR / "snake_exp_plus_exp2_improved.png" +METRICS_PATH = RESULTS_DIR / "snake_exp_plus_exp2_metrics.json" + + +@dataclass(frozen=True) +class TrainingConfig: + x_min: float = -2.0 + x_max: float = 2.0 + train_points: int = 8192 + eval_points: int = 1000 + batch_size: int = 64 + epochs: int = 50 + lr: float = 1e-3 + lambda_exp: float = 2.0 + lambda_scaled: float = 0.5 + seed: int = 42 + + +def build_dataset_exp_plus_exp2(config: TrainingConfig) -> TensorDataset: + x = torch.linspace(config.x_min, config.x_max, config.train_points).unsqueeze(1) + y_g = torch.exp(x) + y_scaled = torch.exp(2.0 * x) + y_f = y_g + y_scaled + return TensorDataset(x, y_f, y_g, y_scaled) + + +def build_models() -> DoubleSin: + return DoubleSin(small=True) + + +def parse_and_validate_expression_graph(program: str) -> Node: + tree = parse_program_to_tree(program) + assert tree.op_type.upper() == "SUM", f"Expected SUM root, found {tree.op_type}" + assert len(tree.children) == 2, f"Expected two SUM branches, found {len(tree.children)}" + + exp_branch, scaled_branch = tree.children + assert exp_branch.op_type.upper() == "EXP", ( + f"Expected EXP first branch, found {exp_branch.op_type}" + ) + assert exp_branch.children[0].op_type.upper() == "X", ( + "Expected first branch to represent exp(x)" + ) + + assert scaled_branch.op_type.upper() == "EXP", ( + f"Expected EXP second branch, found {scaled_branch.op_type}" + ) + scaled_inner = scaled_branch.children[0] + assert scaled_inner.op_type.upper() == "POLY", ( + f"Expected POLY inside scaled branch, found {scaled_inner.op_type}" + ) + assert scaled_inner.params == [0, 2], ( + f"Expected POLY [0, 2] for 2x scaling, found {scaled_inner.params}" + ) + + print("Expression graph program:") + print(program) + print("Parsed expression tree:", tree_to_string(tree)) + print("Top-level node:", tree.op_type.upper()) + print("Branches: exp(x), exp(poly([0, 2], x))") + print() + return tree + + +def forward_components( + model: DoubleSin, + x: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + branch_exp, branch_scaled = model.get_internal() + g_hat = branch_exp(x) + g_scaled_hat = branch_scaled(2.0 * x) + f_hat = g_hat + g_scaled_hat + return f_hat, g_hat, g_scaled_hat + + +def train_exp_plus_exp2_with_exp_aux_loss( + model: DoubleSin, + train_loader: DataLoader, + config: TrainingConfig, +) -> list[dict[str, float]]: + loss_fn = nn.MSELoss() + optimizer = torch.optim.Adam(model.parameters(), lr=config.lr) + history: list[dict[str, float]] = [] + + for epoch in range(1, config.epochs + 1): + totals = { + "loss_f": 0.0, + "loss_exp": 0.0, + "loss_scaled": 0.0, + "total_loss": 0.0, + } + batches = 0 + + for x, y_f, y_g, y_scaled in train_loader: + optimizer.zero_grad() + f_hat, g_hat, g_scaled_hat = forward_components(model, x) + + loss_f = loss_fn(f_hat, y_f) + loss_exp = loss_fn(g_hat, y_g) + loss_scaled = loss_fn(g_scaled_hat, y_scaled) + total_loss = ( + loss_f + + config.lambda_exp * loss_exp + + config.lambda_scaled * loss_scaled + ) + + total_loss.backward() + optimizer.step() + + totals["loss_f"] += loss_f.item() + totals["loss_exp"] += loss_exp.item() + totals["loss_scaled"] += loss_scaled.item() + totals["total_loss"] += total_loss.item() + batches += 1 + + epoch_metrics = {name: value / batches for name, value in totals.items()} + history.append(epoch_metrics) + print( + f"Epoch {epoch:03d}/{config.epochs} " + f"loss_f={epoch_metrics['loss_f']:.6f} " + f"loss_exp={epoch_metrics['loss_exp']:.6f} " + f"loss_scaled={epoch_metrics['loss_scaled']:.6f} " + f"total={epoch_metrics['total_loss']:.6f}" + ) + + return history + + +def plot_exp_plus_exp2_results( + model: DoubleSin, + config: TrainingConfig, + output_path: Path = FIGURE_PATH, +) -> dict[str, float]: + output_path.parent.mkdir(parents=True, exist_ok=True) + model.eval() + + x_eval = torch.linspace(config.x_min, config.x_max, config.eval_points).unsqueeze(1) + with torch.no_grad(): + f_hat, g_hat, g_scaled_hat = forward_components(model, x_eval) + exp_true = torch.exp(x_eval) + exp2_true = torch.exp(2.0 * x_eval) + f_true = exp_true + exp2_true + + metrics = { + "mse_f": nn.functional.mse_loss(f_hat, f_true).item(), + "mse_exp": nn.functional.mse_loss(g_hat, exp_true).item(), + "mse_scaled": nn.functional.mse_loss(g_scaled_hat, exp2_true).item(), + } + + x_np = x_eval.squeeze(1).cpu().numpy() + f_true_np = f_true.squeeze(1).cpu().numpy() + f_hat_np = f_hat.squeeze(1).cpu().numpy() + exp_true_np = exp_true.squeeze(1).cpu().numpy() + exp_hat_np = g_hat.squeeze(1).cpu().numpy() + + fig, axes = plt.subplots(1, 2, figsize=(12, 4.6), dpi=160) + + axes[0].plot(x_np, f_true_np, linewidth=2.0, label="true f(x) = exp(x) + exp(2x)") + axes[0].plot(x_np, f_hat_np, "--", linewidth=2.0, label="predicted f_hat(x)") + axes[0].set_title("Composite Function Approximation") + axes[0].set_xlabel("x") + axes[0].set_ylabel("f(x)") + axes[0].grid(True, alpha=0.25) + axes[0].legend() + + axes[1].plot(x_np, exp_true_np, linewidth=2.0, label="true g(x) = exp(x)") + axes[1].plot(x_np, exp_hat_np, "--", linewidth=2.0, label="predicted g_hat(x)") + axes[1].set_title("Focused Secondary Function Approximation") + axes[1].set_xlabel("x") + axes[1].set_ylabel("g(x)") + axes[1].grid(True, alpha=0.25) + axes[1].legend() + + fig.tight_layout() + fig.savefig(output_path, bbox_inches="tight") + plt.close(fig) + return metrics + + +def save_metrics(metrics: dict[str, float], output_path: Path = METRICS_PATH) -> None: + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text(json.dumps(metrics, indent=2) + "\n", encoding="utf-8") + + +def main() -> None: + config = TrainingConfig() + torch.manual_seed(config.seed) + + parse_and_validate_expression_graph("exp()\nexp(), poly([0, 2])") + + dataset = build_dataset_exp_plus_exp2(config) + train_loader = DataLoader(dataset, batch_size=config.batch_size, shuffle=True) + model = build_models() + + print("Model:") + print(model) + print() + print( + "Training with total_loss = MSE(f_hat, f) " + f"+ {config.lambda_exp} * MSE(g_hat, exp) " + f"+ {config.lambda_scaled} * MSE(g_scaled_hat, exp(2x))" + ) + + train_exp_plus_exp2_with_exp_aux_loss(model, train_loader, config) + metrics = plot_exp_plus_exp2_results(model, config) + save_metrics(metrics) + + print() + print("Evaluation metrics:") + print(f"mse_f={metrics['mse_f']:.8f}") + print(f"mse_exp={metrics['mse_exp']:.8f}") + print(f"mse_scaled={metrics['mse_scaled']:.8f}") + print(f"Saved figure: {FIGURE_PATH}") + print(f"Saved metrics: {METRICS_PATH}") + + +if __name__ == "__main__": + main() diff --git a/math_expr/experiments/snake_sin_plus_ln_ln_focused.py b/math_expr/experiments/snake_sin_plus_ln_ln_focused.py new file mode 100644 index 0000000..3601c93 --- /dev/null +++ b/math_expr/experiments/snake_sin_plus_ln_ln_focused.py @@ -0,0 +1,209 @@ +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +from _common import FIGURES_DIR, SmallFunctionApproximator, configure_output_environment + +configure_output_environment() + +import matplotlib.pyplot as plt +import torch +from torch import nn +from torch.utils.data import DataLoader, TensorDataset + +from math_expr.expression_graph import Node, parse_program_to_tree, tree_to_string + + +FIGURE_PATH = FIGURES_DIR / "snake_sin_plus_ln_improved.png" + + +@dataclass(frozen=True) +class TrainingConfig: + x_min: float = 0.1 + x_max: float = 10.0 + train_points: int = 8192 + eval_points: int = 1000 + batch_size: int = 64 + epochs: int = 50 + lr: float = 1e-3 + lambda_ln: float = 2.0 + seed: int = 42 + + +class SinPlusLnModel(nn.Module): + """Two small approximators summed as f_hat = g1_hat + g2_hat.""" + + def __init__(self) -> None: + super().__init__() + self.g1 = SmallFunctionApproximator() + self.g2 = SmallFunctionApproximator() + + def forward_components( + self, + x: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + g1_hat = self.g1(x) + g2_hat = self.g2(x) + f_hat = g1_hat + g2_hat + return f_hat, g1_hat, g2_hat + + def forward(self, x: torch.Tensor) -> torch.Tensor: + f_hat, _, _ = self.forward_components(x) + return f_hat + + +def build_dataset_sin_plus_ln(config: TrainingConfig) -> TensorDataset: + x = torch.linspace(config.x_min, config.x_max, config.train_points).unsqueeze(1) + y_sin = torch.sin(x) + y_ln = torch.log(x) + y_f = y_sin + y_ln + return TensorDataset(x, y_f, y_sin, y_ln) + + +def build_models() -> SinPlusLnModel: + return SinPlusLnModel() + + +def parse_and_validate_expression_graph(program: str) -> Node: + tree = parse_program_to_tree(program) + child_ops = [child.op_type.upper() for child in tree.children] + + assert tree.op_type.upper() == "SUM", f"Expected SUM root, found {tree.op_type}" + assert child_ops == ["SIN", "LN"], f"Expected SIN and LN children, found {child_ops}" + + print("Expression graph program:") + print(program) + print("Parsed expression tree:", tree_to_string(tree)) + print("Top-level node:", tree.op_type.upper()) + print("Children:", ", ".join(child_ops)) + print() + return tree + + +def train_sin_plus_ln_with_ln_aux_loss( + model: SinPlusLnModel, + train_loader: DataLoader, + config: TrainingConfig, +) -> list[dict[str, float]]: + loss_fn = nn.MSELoss() + optimizer = torch.optim.Adam(model.parameters(), lr=config.lr) + history: list[dict[str, float]] = [] + + for epoch in range(1, config.epochs + 1): + totals = {"loss_f": 0.0, "loss_ln": 0.0, "loss_g1": 0.0, "total_loss": 0.0} + batches = 0 + + for x, y_f, y_sin, y_ln in train_loader: + optimizer.zero_grad() + f_hat, g1_hat, g2_hat = model.forward_components(x) + + loss_f = loss_fn(f_hat, y_f) + loss_ln = loss_fn(g2_hat, y_ln) + loss_g1 = loss_fn(g1_hat, y_sin) + total_loss = loss_f + config.lambda_ln * loss_ln + + total_loss.backward() + optimizer.step() + + totals["loss_f"] += loss_f.item() + totals["loss_ln"] += loss_ln.item() + totals["loss_g1"] += loss_g1.item() + totals["total_loss"] += total_loss.item() + batches += 1 + + epoch_metrics = {name: value / batches for name, value in totals.items()} + history.append(epoch_metrics) + print( + f"Epoch {epoch:03d}/{config.epochs} " + f"loss_f={epoch_metrics['loss_f']:.6f} " + f"loss_ln={epoch_metrics['loss_ln']:.6f} " + f"loss_g1={epoch_metrics['loss_g1']:.6f} " + f"total={epoch_metrics['total_loss']:.6f}" + ) + + return history + + +def plot_sin_plus_ln_results( + model: SinPlusLnModel, + config: TrainingConfig, + output_path: Path = FIGURE_PATH, +) -> dict[str, float]: + output_path.parent.mkdir(parents=True, exist_ok=True) + model.eval() + + x_eval = torch.linspace(config.x_min, config.x_max, config.eval_points).unsqueeze(1) + with torch.no_grad(): + f_hat, g1_hat, g2_hat = model.forward_components(x_eval) + f_true = torch.sin(x_eval) + torch.log(x_eval) + ln_true = torch.log(x_eval) + sin_true = torch.sin(x_eval) + + metrics = { + "mse_f": nn.functional.mse_loss(f_hat, f_true).item(), + "mse_ln": nn.functional.mse_loss(g2_hat, ln_true).item(), + "mse_g1": nn.functional.mse_loss(g1_hat, sin_true).item(), + } + + x_np = x_eval.squeeze(1).cpu().numpy() + f_true_np = f_true.squeeze(1).cpu().numpy() + f_hat_np = f_hat.squeeze(1).cpu().numpy() + ln_true_np = ln_true.squeeze(1).cpu().numpy() + ln_hat_np = g2_hat.squeeze(1).cpu().numpy() + + fig, axes = plt.subplots(1, 2, figsize=(12, 4.6), dpi=160) + + axes[0].plot(x_np, f_true_np, linewidth=2.0, label="true f(x) = sin(x) + ln(x)") + axes[0].plot(x_np, f_hat_np, "--", linewidth=2.0, label="predicted f_hat(x)") + axes[0].set_title("Composite Function Approximation") + axes[0].set_xlabel("x") + axes[0].set_ylabel("f(x)") + axes[0].grid(True, alpha=0.25) + axes[0].legend() + + axes[1].plot(x_np, ln_true_np, linewidth=2.0, label="true g2(x) = ln(x)") + axes[1].plot(x_np, ln_hat_np, "--", linewidth=2.0, label="predicted g2_hat(x)") + axes[1].set_title("Focused Secondary Function Approximation") + axes[1].set_xlabel("x") + axes[1].set_ylabel("g2(x)") + axes[1].grid(True, alpha=0.25) + axes[1].legend() + + fig.tight_layout() + fig.savefig(output_path, bbox_inches="tight") + plt.close(fig) + + print() + print("Evaluation metrics:") + print(f"mse_f={metrics['mse_f']:.8f}") + print(f"mse_ln={metrics['mse_ln']:.8f}") + print(f"mse_g1={metrics['mse_g1']:.8f}") + print(f"Saved figure: {output_path}") + return metrics + + +def main() -> None: + config = TrainingConfig() + torch.manual_seed(config.seed) + + parse_and_validate_expression_graph("sin()\nln()") + + dataset = build_dataset_sin_plus_ln(config) + train_loader = DataLoader(dataset, batch_size=config.batch_size, shuffle=True) + model = build_models() + + print("Model:") + print(model) + print() + print( + "Training with total_loss = MSE(f_hat, f) " + f"+ {config.lambda_ln} * MSE(g2_hat, ln)" + ) + + train_sin_plus_ln_with_ln_aux_loss(model, train_loader, config) + plot_sin_plus_ln_results(model, config) + + +if __name__ == "__main__": + main() diff --git a/math_expr/expression_graph.py b/math_expr/expression_graph.py new file mode 100644 index 0000000..d035445 --- /dev/null +++ b/math_expr/expression_graph.py @@ -0,0 +1,249 @@ +from __future__ import annotations + +import ast +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + + +@dataclass +class Node: + """A node in a parsed mathematical expression tree.""" + + op_type: str + params: Any = None + children: list["Node"] = field(default_factory=list) + + +def parse_program_to_tree(program: str) -> Node: + """Parse a graph DSL program into an expression tree.""" + lines = [line.strip() for line in program.splitlines() if line.strip()] + if not lines: + raise ValueError("Program is empty") + + line_nodes = [_parse_program_line(line) for line in lines] + if len(line_nodes) == 1: + return line_nodes[0] + return Node("SUM", children=line_nodes) + + +def tree_to_string(node: Node) -> str: + """Render an expression tree as a readable mathematical expression.""" + op_type = node.op_type.upper() + if op_type == "X": + return "x" + if op_type == "SUM": + return " + ".join(tree_to_string(child) for child in node.children) + if op_type == "POLY": + child = _single_child(node) + return f"poly({_format_params(node.params)}, {tree_to_string(child)})" + + child = _single_child(node) + return f"{op_type.lower()}({tree_to_string(child)})" + + +def tree_to_edges(node: Node) -> list[tuple[str, str]]: + """Return expression tree edges as parent and child labels.""" + edges = [] + for child in node.children: + edges.append((_node_label(node), _node_label(child))) + edges.extend(tree_to_edges(child)) + return edges + + +def tree_to_mermaid(node: Node) -> str: + """Render an expression tree as a Mermaid graph definition.""" + lines = ["graph TD"] + nodes: list[tuple[str, Node]] = [] + edges: list[tuple[str, str]] = [] + + def visit(current: Node, parent_id: str | None = None) -> None: + node_id = f"N{len(nodes)}" + nodes.append((node_id, current)) + if parent_id is not None: + edges.append((parent_id, node_id)) + for child in current.children: + visit(child, node_id) + + visit(node) + + for node_id, current in nodes: + lines.append(f' {node_id}["{_escape_mermaid_label(_node_label(current))}"]') + for parent_id, child_id in edges: + lines.append(f" {parent_id} --> {child_id}") + return "\n".join(lines) + + +def export_tree_to_dot( + node: Node, + output_path: str, + *, + title: str | None = None, +) -> str: + """Export an expression tree to a Graphviz-rendered PNG file.""" + try: + from graphviz import Digraph + from graphviz.backend.execute import ExecutableNotFound + except ImportError as exc: + raise RuntimeError( + "Graphviz export requires the optional Python package. Install it with " + "`pip install graphviz`. The Graphviz system binary may also be needed." + ) from exc + + output = Path(output_path) + output.parent.mkdir(parents=True, exist_ok=True) + filename = ( + str(output.with_suffix("")) + if output.suffix.lower() == ".png" + else str(output) + ) + + graph = Digraph("expression_tree", format="png") + graph.attr( + rankdir="TB", + bgcolor="white", + pad="0.25", + nodesep="0.45", + ranksep="0.65", + ) + if title: + graph.attr(label=title, labelloc="t", fontsize="18", fontname="Helvetica") + graph.attr( + "node", + shape="box", + style="rounded,filled", + fillcolor="#F7FAFC", + color="#334155", + fontname="Helvetica", + fontsize="12", + margin="0.12,0.08", + ) + graph.attr("edge", color="#64748B", arrowsize="0.75") + + counter = 0 + + def add_node(current: Node, parent_id: str | None = None) -> None: + nonlocal counter + node_id = f"N{counter}" + counter += 1 + graph.node(node_id, _node_label(current)) + if parent_id is not None: + graph.edge(parent_id, node_id) + for child in current.children: + add_node(child, node_id) + + add_node(node) + try: + return graph.render(filename=filename, cleanup=True) + except ExecutableNotFound as exc: + raise RuntimeError( + "Graphviz export requires the Graphviz system binary in addition to the " + "Python package. Install the Python package with `pip install graphviz` " + "and install Graphviz for your operating system." + ) from exc + + +def tree_to_networkx(node: Node): + """Return a NetworkX directed graph for an expression tree, if available.""" + try: + import networkx as nx + except ImportError: + return None + + graph = nx.DiGraph() + + counter = 0 + + def add_node(current: Node, parent_id: str | None = None): + nonlocal counter + node_id = f"n{counter}" + counter += 1 + graph.add_node( + node_id, + label=_node_label(current), + op_type=current.op_type, + params=current.params, + ) + if parent_id is not None: + graph.add_edge(parent_id, node_id) + for child in current.children: + add_node(child, node_id) + + add_node(node) + return graph + + +def _parse_program_line(line: str) -> Node: + calls = [_parse_call(part) for part in _split_top_level_commas(line)] + if not calls: + raise ValueError(f"No function calls found in line: {line!r}") + + current = Node("X") + for op_type, params in reversed(calls): + current = Node(op_type, params=params, children=[current]) + return current + + +def _split_top_level_commas(text: str) -> list[str]: + parts = [] + start = 0 + depth = 0 + for index, char in enumerate(text): + if char in "([": + depth += 1 + elif char in ")]": + depth -= 1 + elif char == "," and depth == 0: + part = text[start:index].strip() + if part: + parts.append(part) + start = index + 1 + + part = text[start:].strip() + if part: + parts.append(part) + return parts + + +def _parse_call(text: str) -> tuple[str, Any]: + open_index = text.find("(") + close_index = text.rfind(")") + if open_index == -1 or close_index == -1 or close_index < open_index: + raise ValueError(f"Invalid function call: {text!r}") + + fn_name = text[:open_index].strip().lower() + raw_params = text[open_index + 1:close_index].strip() + if fn_name == "log": + fn_name = "ln" + if fn_name not in {"exp", "ln", "sin", "sinc", "poly"}: + raise ValueError(f"Unsupported function: {fn_name!r}") + + params = None + if raw_params: + try: + params = ast.literal_eval(raw_params) + except (SyntaxError, ValueError) as exc: + raise ValueError(f"Invalid parameters for {fn_name}: {raw_params!r}") from exc + return fn_name.upper(), params + + +def _single_child(node: Node) -> Node: + if len(node.children) != 1: + raise ValueError(f"{node.op_type} expects exactly one child") + return node.children[0] + + +def _node_label(node: Node) -> str: + if node.op_type.upper() == "POLY" and node.params is not None: + return f"POLY {_format_params(node.params)}" + return node.op_type.upper() + + +def _format_params(params: Any) -> str: + if isinstance(params, list): + return "[" + ", ".join(str(value) for value in params) + "]" + return str(params) + + +def _escape_mermaid_label(label: str) -> str: + return label.replace("\\", "\\\\").replace('"', '\\"') diff --git a/math_expr/module_from_tree.py b/math_expr/module_from_tree.py new file mode 100644 index 0000000..2c6771f --- /dev/null +++ b/math_expr/module_from_tree.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +import torch +from torch import nn + +try: + from math_expr.expression_graph import Node + from math_expr.nn_nodes import ( + ComposeNode, + ExpNode, + LnNode, + PolyNode, + SincNode, + SinNode, + SumNode, + ) +except ImportError: # pragma: no cover - supports running examples from inside math_expr/ + from expression_graph import Node + from nn_nodes import ( + ComposeNode, + ExpNode, + LnNode, + PolyNode, + SincNode, + SinNode, + SumNode, + ) + + +class IdentityNode(nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + + +def node_to_module(node: Node) -> nn.Module: + """Convert an expression tree node into an equivalent PyTorch module.""" + op_type = node.op_type.upper() + + if op_type == "X": + return IdentityNode() + if op_type == "SUM": + return SumNode([node_to_module(child) for child in node.children]) + if op_type == "SIN": + return _wrap_unary(SinNode(), node) + if op_type == "EXP": + return _wrap_unary(ExpNode(), node) + if op_type in {"LN", "LOG"}: + return _wrap_unary(LnNode(), node) + if op_type == "SINC": + return _wrap_unary(SincNode(), node) + if op_type == "POLY": + return _wrap_unary(PolyNode(node.params or []), node) + + raise ValueError(f"Unsupported op_type for module conversion: {node.op_type}") + + +def _wrap_unary(outer: nn.Module, node: Node) -> nn.Module: + if not node.children: + return ComposeNode(outer, IdentityNode()) + if len(node.children) != 1: + raise ValueError(f"{node.op_type} expects exactly one child") + return ComposeNode(outer, node_to_module(node.children[0])) diff --git a/math_expr/nn_nodes.py b/math_expr/nn_nodes.py new file mode 100644 index 0000000..94eb8c9 --- /dev/null +++ b/math_expr/nn_nodes.py @@ -0,0 +1,57 @@ +from __future__ import annotations + +from collections.abc import Sequence + +import torch +from torch import nn + + +class SinNode(nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.sin(x) + + +class ExpNode(nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.exp(x) + + +class LnNode(nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.log(x) + + +class SincNode(nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.sinc(x) + + +class PolyNode(nn.Module): + def __init__(self, coeffs: Sequence[float] | None): + super().__init__() + self.coeffs = [float(coefficient) for coefficient in (coeffs or [])] + + def forward(self, x: torch.Tensor) -> torch.Tensor: + result = torch.zeros_like(x) + for power, coefficient in enumerate(self.coeffs): + result = result + coefficient * x.pow(power) + return result + + +class SumNode(nn.Module): + def __init__(self, children: list[nn.Module]): + super().__init__() + self.terms = nn.ModuleList(children) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return sum(child(x) for child in self.terms) + + +class ComposeNode(nn.Module): + def __init__(self, outer: nn.Module, inner: nn.Module): + super().__init__() + self.outer = outer + self.inner = inner + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.outer(self.inner(x)) diff --git a/math_expr/test_expression_graph.py b/math_expr/test_expression_graph.py new file mode 100644 index 0000000..024a648 --- /dev/null +++ b/math_expr/test_expression_graph.py @@ -0,0 +1,128 @@ +from pathlib import Path + +import torch + +import executor_adapter +from expression_graph import parse_program_to_tree, tree_to_mermaid, tree_to_string +from module_from_tree import node_to_module + + +def evaluate_tree_torch(node, x): + 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 test_parse_composition_to_string(): + tree = parse_program_to_tree("sin(), exp()") + + assert tree.op_type == "SIN" + assert tree_to_string(tree) == "sin(exp(x))" + + +def test_multiline_expression_parses_as_sum(): + tree = parse_program_to_tree("sin(), poly([0,3]), exp()\nsinc()") + + assert tree.op_type == "SUM" + assert len(tree.children) == 2 + assert "+" in tree_to_string(tree) + + +def test_poly_parameter_parsing_preserves_coefficients(): + tree = parse_program_to_tree("poly([0, 0.2, 3])") + + assert tree.op_type == "POLY" + assert tree.params == [0, 0.2, 3] + assert tree_to_string(tree) == "poly([0, 0.2, 3], x)" + + +def test_tree_to_mermaid_exports_basic_composition(): + tree = parse_program_to_tree("sin(), exp()") + mermaid = tree_to_mermaid(tree) + + assert "graph TD" in mermaid + assert '["SIN"]' in mermaid + assert '["EXP"]' in mermaid + assert '["X"]' in mermaid + + +def test_node_to_module_matches_recursive_torch_evaluator(): + tree = parse_program_to_tree("sin(), poly([1, 2]), exp()") + model = node_to_module(tree) + x = torch.linspace(-1.0, 1.0, steps=5) + + assert torch.allclose(model(x), evaluate_tree_torch(tree, x)) + + +def test_node_to_module_supports_backward_for_sum(): + tree = parse_program_to_tree("sin(), exp()\nsinc()") + model = node_to_module(tree) + x = torch.linspace(-1.0, 1.0, steps=5, requires_grad=True) + loss = model(x).mean() + + loss.backward() + + assert x.grad is not None + + +def test_evaluate_graph_with_tree_calls_original_executor_in_compatible_workspace( + tmp_path, + monkeypatch, +): + graph_file = tmp_path / "demo.graph" + graph_file.write_text("sin(), exp()", encoding="utf-8") + observed = {} + + def fake_original_execute(program_name, x=None): + observed["program_name"] = program_name + observed["graph_exists"] = Path("graphs/demo.graph").exists() + observed["x"] = x + return "original-result" + + monkeypatch.setattr(executor_adapter, "original_execute", fake_original_execute) + + result, tree = executor_adapter.evaluate_graph_with_tree(graph_file, x="sample-x") + + assert observed == { + "program_name": "demo", + "graph_exists": True, + "x": "sample-x", + } + assert result.original_return == "original-result" + assert result.execution_successful is True + assert tree_to_string(tree) == "sin(exp(x))" + + +def test_execute_with_tree_returns_executor_result_and_tree(tmp_path, monkeypatch): + graph_file = tmp_path / "demo.graph" + graph_file.write_text("sinc()", encoding="utf-8") + + monkeypatch.setattr( + executor_adapter, + "original_execute", + lambda program_name, x=None: "executor-result", + ) + + executor_result, tree = executor_adapter.execute_with_tree(graph_file) + + assert executor_result == "executor-result" + assert tree_to_string(tree) == "sinc(x)"