| |
| """Neural property-prediction tool for estimating monotonicity, convexity, periodicity, and separability.""" |
| from __future__ import annotations |
|
|
| import logging |
| import os |
| import numpy as np |
| from pathlib import Path |
| from typing import Dict, Any, List, Literal |
| from ..utils import download_model, get_default, tag2ansi |
| from .base_tool import BaseTool, ToolMetadata, ToolRunAbort, is_numeric_array |
|
|
| _logger = logging.getLogger(f"sr_harness.{__name__}") |
| _PROJECT_ROOT = Path(__file__).resolve().parents[3] |
| _MODEL_CACHE_DIR = _PROJECT_ROOT / ".cache" / "nn_tools" |
| _MONO_LABELS = {0: "non-monotonic", 1: "increasing", 2: "decreasing", 3: "constant"} |
| _CONV_LABELS = {0: "non-convex/concave", 1: "convex", 2: "concave", 3: "affine"} |
| _PERIOD_LABELS = {0: "non-periodic", 1: "periodic"} |
| _SEP_LABELS = {0: "not multiplicatively separable", 1: "multiplicatively separable"} |
| _COMBO_OPS = { |
| "*": lambda a, b: a * b, |
| "+": lambda a, b: a + b, |
| "-": lambda a, b: a - b, |
| "/": lambda a, b: np.where(np.abs(b) > 1e-12, a / b, np.nan), |
| } |
| _cached_model = {} |
|
|
|
|
| def _load_model(checkpoint_path: str, device: str = "cpu"): |
| """Load model from checkpoint with caching to avoid repeated I/O.""" |
| import torch |
|
|
| cache_key = (checkpoint_path, device) |
| if cache_key in _cached_model: |
| return _cached_model[cache_key] |
|
|
| import sys |
| exp_path = str(_PROJECT_ROOT / "src" / "experimental") |
| if exp_path not in sys.path: |
| sys.path.insert(0, exp_path) |
|
|
| import argparse |
| from nn_tools.models import FloatEmbedder, DataEmbedder, PropertyPredictionModel |
|
|
| ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False) |
| saved_args = ckpt["args"] |
| if isinstance(saved_args, dict): |
| saved_args = argparse.Namespace(**saved_args) |
|
|
| float_emb = FloatEmbedder(d_model=saved_args.d_model).to(device) |
| data_emb = DataEmbedder( |
| d_model=saved_args.d_model, |
| pooling=saved_args.data_pooling, |
| float_embedder=float_emb, |
| ).to(device) |
| model = PropertyPredictionModel(saved_args).to(device) |
|
|
| model.load_state_dict(ckpt["model"]) |
| float_emb.load_state_dict(ckpt["float_embedder"]) |
| data_emb.load_state_dict(ckpt["data_embedder"]) |
|
|
| model.eval() |
| float_emb.eval() |
| data_emb.eval() |
|
|
| result = (model, float_emb, data_emb, saved_args) |
| _cached_model[cache_key] = result |
| _logger.info(f"Loaded property prediction model from {checkpoint_path}") |
| return result |
|
|
|
|
| def _predict(model, float_emb, data_emb, data_arr: np.ndarray, device: str) -> dict: |
| """Run model inference on a single (sample_num, max_var+1) data array. |
| |
| Returns raw logits dict with keys: monotonicity, convexity, periodicity, |
| multiplicative_separable. |
| """ |
| import torch |
| data_tensor = torch.from_numpy(data_arr).unsqueeze(0) |
| with torch.no_grad(): |
| B, S = data_tensor.shape[:2] |
| val_emb = float_emb(data_tensor.to(device)).flatten(0, 1) |
| d_emb = data_emb.pool(val_emb).reshape(B, S, -1) |
| out = model(d_emb) |
| return {k: v[0].cpu().numpy() for k, v in out.items()} |
|
|
|
|
| def _decode_per_var(logits: np.ndarray, label_map: dict, n_vars: int) -> list: |
| """Decode per-variable logits into list of (pred_label, confidence).""" |
| probs = _softmax(logits) |
| preds = logits.argmax(axis=-1) |
| results = [] |
| for i in range(n_vars): |
| pred = int(preds[i]) |
| conf = float(probs[i, pred]) |
| results.append((label_map[pred], round(conf, 3))) |
| return results |
|
|
|
|
| |
| |
| class PropertyPredictorTool(BaseTool): |
| """Neural-network tool for predicting per-variable monotonicity, convexity, periodicity, and multiplicative separability.""" |
| metadata = ToolMetadata(name="predict_property") |
|
|
| MODEL_TYPE: Literal[ |
| "property-scratch", |
| "property-finetune", |
| "property-scratch-v5", |
| ] = "property-scratch-v5" |
| MODEL_RELEASE_TAG: str = "sr-agent-models-v5" |
| DEVICE: str = "cpu" |
| AUTO_DOWNLOAD: bool = True |
|
|
| def execute(self) -> Dict[str, Any]: |
| """Predict heuristic mathematical properties of the target with respect to each variable and pairwise combination. |
| |
| Uses a pretrained Transformer to estimate: |
| - **Monotonicity**: increasing, decreasing, constant, or non-monotonic per variable. |
| - **Convexity**: convex, concave, affine, or non-convex/concave per variable. |
| - **Periodicity**: heuristic periodicity score per variable (note: modulated terms like ``x * sin(w * t)`` may appear non-periodic marginally). |
| - **Multiplicative Separability**: whether ``y`` factors into ``f(x1) * g(x2) * ...``. |
| Also evaluates pairwise combinations (``xi*xj``, ``xi+xj``, ``xi-xj``, ``xi/xj``) to surface hidden composite symmetries. |
| |
| Returns: |
| Dict[str, Any]: Predicted property labels and confidence scores for individual variables and pairwise combinations. |
| """ |
| data = self.context.data |
| target_name = self.context.target |
| exceptions = [] |
|
|
| |
| default_ckpt_path = _MODEL_CACHE_DIR / f"{self.MODEL_TYPE}.pth" |
| ckpt_path = Path(os.getenv("SR_AGENT_PROPERTY_MODEL_CHECKPOINT") or default_ckpt_path).expanduser() |
| repo = get_default("repo") |
| release_url = f"https://github.com/{repo}/releases/tag/{self.MODEL_RELEASE_TAG}" |
| if not ckpt_path.exists() and self.AUTO_DOWNLOAD: |
| _logger.info( |
| f"Model checkpoint {ckpt_path} not found, attempting to download from GitHub..." |
| f"You can also download it manually if it takes too long or fails. To do so, visit the GitHub Release page\n" |
| f" [blue bold]{release_url}[reset]\n" |
| f"and download the asset named\n" |
| f" [blue bold]{self.MODEL_TYPE!r}[reset]\n" |
| f"Then save it as\n" |
| f" [blue bold]{ckpt_path}[reset]\n" |
| ) |
| download_model(name=self.MODEL_TYPE, checkpoint=ckpt_path, release_tag=self.MODEL_RELEASE_TAG) |
| if not ckpt_path.exists(): |
| raise ToolRunAbort(tag2ansi( |
| f"Model checkpoint not found: {ckpt_path}.\n" |
| f"Automatic download was not available. To use predict_property, please visit the GitHub Release page\n" |
| f" [blue bold]{release_url}[reset]\n" |
| f"and download the asset named\n" |
| f" [blue bold]{self.MODEL_TYPE!r}[reset]\n" |
| f"Then save it as\n" |
| f" [blue bold]{ckpt_path}[reset]\n" |
| "Alternatively, you can ban this predict_property tool to avoid this error." |
| )) |
|
|
| model, float_emb, data_emb, saved_args = _load_model(ckpt_path, self.DEVICE) |
| max_var_num = saved_args.max_var_num |
|
|
| input_vars = [v for v in data if v != target_name and is_numeric_array(data[v])] |
| if len(input_vars) > max_var_num: |
| input_vars = input_vars[:max_var_num] |
| exceptions.append(f"Too many variables ({len(data) - 1}), only the first {max_var_num} are analyzed: {input_vars}") |
|
|
| n_vars = len(input_vars) |
| if not is_numeric_array(data[target_name]): |
| raise ValueError(f"Target variable '{target_name}' did not produce numeric values.") |
| y_arr = np.asarray(data[target_name], dtype=np.float32).flatten() |
| n_samples = len(y_arr) |
|
|
| sample_num = min(n_samples, 200) |
| rng = np.random.default_rng(42) |
| idx = rng.choice(n_samples, sample_num, replace=False) if n_samples > sample_num else np.arange(n_samples) |
|
|
| var_arrays = {} |
| for var in input_vars: |
| var_arrays[var] = np.asarray(data[var], dtype=np.float32).flatten()[idx] |
|
|
| |
| data_arr = np.zeros((sample_num, max_var_num + 1), dtype=np.float32) |
| for i, var in enumerate(input_vars): |
| data_arr[:, i] = var_arrays[var] |
| data_arr[:, -1] = y_arr[idx] |
|
|
| out = _predict(model, float_emb, data_emb, data_arr, self.DEVICE) |
|
|
| mono_decoded = _decode_per_var(out["monotonicity"], _MONO_LABELS, n_vars) |
| conv_decoded = _decode_per_var(out["convexity"], _CONV_LABELS, n_vars) |
| period_decoded = _decode_per_var(out["periodicity"], _PERIOD_LABELS, n_vars) |
|
|
| sep_logits = out["multiplicative_separable"] |
| sep_probs = _softmax(sep_logits) |
| sep_pred = int(sep_logits.argmax()) |
| sep_conf = float(sep_probs[sep_pred]) |
|
|
| monotonicity = {} |
| convexity = {} |
| periodicity = {} |
| for i, var in enumerate(input_vars): |
| monotonicity[var] = {"prediction": mono_decoded[i][0], "confidence": mono_decoded[i][1]} |
| convexity[var] = {"prediction": conv_decoded[i][0], "confidence": conv_decoded[i][1]} |
| periodicity[var] = {"prediction": period_decoded[i][0], "confidence": period_decoded[i][1]} |
|
|
| separability = { |
| "prediction": _SEP_LABELS[sep_pred], |
| "confidence": round(sep_conf, 3), |
| } |
|
|
| |
| combo_results = {} |
| if n_vars >= 2: |
| for i in range(n_vars): |
| for j in range(i + 1, n_vars): |
| vi, vj = input_vars[i], input_vars[j] |
| xi, xj = var_arrays[vi], var_arrays[vj] |
| for op_name, op_fn in _COMBO_OPS.items(): |
| combo_name = f"{vi}{op_name}{vj}" |
| combo_val = op_fn(xi, xj).astype(np.float32) |
| if not np.all(np.isfinite(combo_val)): |
| valid = np.isfinite(combo_val) |
| if valid.sum() < 20: |
| continue |
| combo_val = combo_val[valid] |
| y_combo = y_arr[idx][valid] |
| combo_s = len(combo_val) |
| else: |
| y_combo = y_arr[idx] |
| combo_s = sample_num |
|
|
| remaining_vars = [input_vars[k] for k in range(n_vars) if k != i and k != j] |
| n_combo_vars = 1 + len(remaining_vars) |
|
|
| combo_arr = np.zeros((combo_s, max_var_num + 1), dtype=np.float32) |
| combo_arr[:, 0] = combo_val[:combo_s] |
| for ri, rv in enumerate(remaining_vars): |
| if not np.all(np.isfinite(combo_val)): |
| combo_arr[:, 1 + ri] = var_arrays[rv][valid][:combo_s] |
| else: |
| combo_arr[:, 1 + ri] = var_arrays[rv][:combo_s] |
| combo_arr[:, -1] = y_combo[:combo_s] |
|
|
| combo_out = _predict(model, float_emb, data_emb, combo_arr, self.DEVICE) |
|
|
| combo_mono = _decode_per_var(combo_out["monotonicity"], _MONO_LABELS, 1) |
| combo_conv = _decode_per_var(combo_out["convexity"], _CONV_LABELS, 1) |
| combo_period = _decode_per_var(combo_out["periodicity"], _PERIOD_LABELS, 1) |
|
|
| combo_results[combo_name] = { |
| "monotonicity": {"prediction": combo_mono[0][0], "confidence": combo_mono[0][1]}, |
| "convexity": {"prediction": combo_conv[0][0], "confidence": combo_conv[0][1]}, |
| "periodicity": {"prediction": combo_period[0][0], "confidence": combo_period[0][1]}, |
| } |
|
|
| result = { |
| "target": target_name, |
| "monotonicity": monotonicity, |
| "convexity": convexity, |
| "periodicity": periodicity, |
| "multiplicative_separable": separability, |
| "n_variables_analyzed": n_vars, |
| "n_samples_used": sample_num, |
| "exceptions": exceptions, |
| } |
| if combo_results: |
| result["variable_combinations"] = combo_results |
|
|
| return result |
|
|
| @classmethod |
| def format_result_dict(cls, result: Dict[str, Any]) -> str: |
| """Format neural property predictions and search hints for the Agent. |
| |
| Args: |
| result: Result dictionary returned by :meth:`execute`. |
| |
| Returns: |
| str: Formatted property report with confidence scores and modeling hints. |
| """ |
| if "error" in result: |
| return f"Error: {result['error']}" |
|
|
| lines = [ |
| f"Heuristic neural-network property predictions for target {result['target']!r} " |
| f"(analyzed {result['n_variables_analyzed']} variables using {result['n_samples_used']} sampled rows):", |
| "Model scores below are uncalibrated confidence-like scores, not statistical proof. Treat low-score predictions as uncertain.", |
| "", |
| "=== Per-Variable Properties ===", |
| "", |
| "Monotonicity (whether y is monotonic w.r.t. each variable):", |
| ] |
| for var, info in result["monotonicity"].items(): |
| label = info['prediction'] if info['confidence'] >= 0.6 else f"uncertain; top label is {info['prediction']}" |
| lines.append(f" {var}: {label} (model score: {info['confidence']:.1%})") |
|
|
| lines.append("") |
| lines.append("Convexity (whether y is convex/concave w.r.t. each variable):") |
| for var, info in result["convexity"].items(): |
| label = info['prediction'] if info['confidence'] >= 0.6 else f"uncertain; top label is {info['prediction']}" |
| lines.append(f" {var}: {label} (model score: {info['confidence']:.1%})") |
|
|
| lines.append("") |
| lines.append("Periodicity (whether y is periodic w.r.t. each variable):") |
| for var, info in result["periodicity"].items(): |
| label = info['prediction'] if info['confidence'] >= 0.6 else f"uncertain; top label is {info['prediction']}" |
| lines.append(f" {var}: {label} (model score: {info['confidence']:.1%})") |
| lines.append( |
| "Caution: a 'non-periodic' label cannot exclude modulated oscillations such as " |
| "x * sin(omega * t); test such interactions against data when scientifically plausible." |
| ) |
|
|
| sep = result["multiplicative_separable"] |
| lines.append("") |
| sep_label = sep['prediction'] if sep['confidence'] >= 0.6 else f"uncertain; top label is {sep['prediction']}" |
| lines.append(f"Multiplicative Separability: {sep_label} (model score: {sep['confidence']:.1%})") |
|
|
| |
| if any(info["prediction"] == "periodic" and info["confidence"] >= 0.6 for info in result["periodicity"].values()): |
| lines.append("") |
| lines.append("Hint: Periodic variables suggest the formula likely contains trigonometric functions " |
| "(sin, cos) of those variables.") |
|
|
| if (result["multiplicative_separable"]["prediction"] == "multiplicatively separable" |
| and result["multiplicative_separable"]["confidence"] >= 0.6): |
| lines.append("") |
| lines.append("Hint: The formula may be expressible as a product of functions of individual variables, " |
| "e.g., y = f(x1) * g(x2). Try decomposing the problem.") |
|
|
| |
| combo = result.get("variable_combinations", {}) |
| if combo: |
| lines.append("") |
| lines.append("=== Variable Combination Properties ===") |
| lines.append("(Analyzing xi*xj, xi+xj, xi-xj, xi/xj to detect combination-level properties)") |
| lines.append("") |
|
|
| interesting = [] |
| for cname, cinfo in combo.items(): |
| flags = [] |
| if cinfo["periodicity"]["prediction"] == "periodic" and cinfo["periodicity"]["confidence"] >= 0.6: |
| flags.append(f"periodic ({cinfo['periodicity']['confidence']:.1%})") |
| if (cinfo["monotonicity"]["prediction"] in ("increasing", "decreasing") |
| and cinfo["monotonicity"]["confidence"] >= 0.6): |
| flags.append(f"{cinfo['monotonicity']['prediction']} ({cinfo['monotonicity']['confidence']:.1%})") |
| if (cinfo["convexity"]["prediction"] in ("convex", "concave") |
| and cinfo["convexity"]["confidence"] >= 0.6): |
| flags.append(f"{cinfo['convexity']['prediction']} ({cinfo['convexity']['confidence']:.1%})") |
| if flags: |
| interesting.append((cname, flags, cinfo)) |
|
|
| if interesting: |
| for cname, flags, cinfo in interesting: |
| lines.append(f" {cname}: {', '.join(flags)}") |
| lines.append("") |
| lines.append("Hint: Variable combinations with detected properties suggest the formula may contain " |
| "functions of those combinations. For example:") |
| for cname, flags, cinfo in interesting: |
| if (cinfo["periodicity"]["prediction"] == "periodic" |
| and cinfo["periodicity"]["confidence"] >= 0.6): |
| lines.append(f" - {cname} is periodic → consider sin({cname}), cos({cname})") |
| if (cinfo["monotonicity"]["prediction"] in ("increasing", "decreasing") |
| and cinfo["monotonicity"]["confidence"] >= 0.6): |
| lines.append(f" - {cname} is {cinfo['monotonicity']['prediction']} → y may depend on {cname} directly or via a monotonic function") |
| if (cinfo["convexity"]["prediction"] in ("convex", "concave") |
| and cinfo["convexity"]["confidence"] >= 0.6): |
| lines.append(f" - {cname} is {cinfo['convexity']['prediction']} → consider {cname}², sqrt({cname}), exp({cname}), etc.") |
| else: |
| lines.append(" No notable properties detected in variable combinations.") |
| lines.append("") |
| lines.append("Note: Even though no combination properties were detected, complex variable " |
| "combinations (e.g., sin(xi-xj), cos(xi*xj)) should still be explored " |
| "as the model may not capture all patterns.") |
|
|
| if result["exceptions"]: |
| lines.append("") |
| lines.append("Exceptions / Warnings:") |
| for ex in result["exceptions"]: |
| lines.append(f" - {ex}") |
|
|
| return "\n".join(lines) |
|
|
|
|
| def _softmax(logits: np.ndarray) -> np.ndarray: |
| e = np.exp(logits - logits.max(axis=-1, keepdims=True)) |
| return e / e.sum(axis=-1, keepdims=True) |
|
|
