diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 908b73ee2..57c89a459 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -110,6 +110,33 @@ jobs: - name: "Run development test" run: PYSR_TEST_JULIA_VERSION=${{ matrix.julia-version }} PYSR_TEST_PYTHON_VERSION=${{ matrix.python-version }} python -m pysr test dev + rust_backend_adapter: + name: Test Rust backend adapter + runs-on: ubuntu-latest + defaults: + run: + shell: bash + strategy: + matrix: + python-version: ['3.13'] + + steps: + - uses: actions/checkout@v6 + - name: "Set up Python" + uses: actions/setup-python@v6 + with: + python-version: ${{ matrix.python-version }} + cache: pip + - name: "Install PySR test dependencies" + run: | + python -m pip install --upgrade pip + pip install -e . + python -m pip install pytest + - name: "Assert default import is Julia-free" + run: python -c "import sys; import pysr; assert 'juliacall' not in sys.modules; from pysr import PySRRegressor; assert PySRRegressor().backend == 'auto'" + - name: "Run Rust backend adapter tests" + run: python -m pysr test rust + conda_test: runs-on: ${{ matrix.os }} defaults: diff --git a/.github/workflows/CI_rust_backend_connector.yml b/.github/workflows/CI_rust_backend_connector.yml new file mode 100644 index 000000000..9b2002bdd --- /dev/null +++ b/.github/workflows/CI_rust_backend_connector.yml @@ -0,0 +1,80 @@ +name: Rust backend connector + +on: + pull_request: + branches: + - master + paths: + - pysr_rust_backend/** + - .github/workflows/CI_rust_backend_connector.yml + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +env: + CARGO_TERM_COLOR: always + +jobs: + rust-backend-connector: + name: Rust backend connector + runs-on: ubuntu-latest + strategy: + matrix: + python-version: ["3.9", "3.12", "3.13"] + steps: + - name: Checkout + uses: actions/checkout@v7 + + - name: Install Rust + uses: actions-rust-lang/setup-rust-toolchain@v1 + + - name: Rust cache + uses: Swatinem/rust-cache@v2 + + - name: Set up Python + uses: actions/setup-python@v6 + with: + python-version: ${{ matrix.python-version }} + + - name: Install uv + uses: astral-sh/setup-uv@v7 + + - name: Install connector + run: uv pip install --system -e pysr_rust_backend + + - name: Run connector tests + run: python -m unittest discover pysr_rust_backend/tests + + - name: Build connector wheel + working-directory: pysr_rust_backend + run: uvx maturin build --out ../dist + + - name: Test built wheel + working-directory: pysr_rust_backend + run: | + python -m venv "$RUNNER_TEMP/pysr-rust-backend-wheel" + . "$RUNNER_TEMP/pysr-rust-backend-wheel/bin/activate" + python -m pip install --upgrade pip + python -m pip install ../dist/*.whl + python -c "import numpy as np, symbolic_regression_rs; X=np.linspace(-1,1,24,dtype=np.float32).reshape(-1,1); y=X[:,0]; out=symbolic_regression_rs.search(X,y,options={'seed':0,'niterations':1,'populations':1,'population_size':16,'ncycles_per_iteration':20,'maxsize':10,'maxdepth':10,'deterministic':True,'progress':False},operators={2:['+','sub','*','/']},variable_names=['x0']); assert out['hall_of_fame']" + + - name: Build source distribution + working-directory: pysr_rust_backend + run: uvx maturin sdist --out ../dist + + - name: Test source distribution + working-directory: pysr_rust_backend + run: | + python -m venv "$RUNNER_TEMP/pysr-rust-backend-sdist" + . "$RUNNER_TEMP/pysr-rust-backend-sdist/bin/activate" + python -m pip install --upgrade pip + python -m pip install ../dist/*.tar.gz + python -c "import numpy as np, symbolic_regression_rs; X=np.linspace(-1,1,24,dtype=np.float32).reshape(-1,1); y=X[:,0]; out=symbolic_regression_rs.search(X,y,options={'seed':0,'niterations':1,'populations':1,'population_size':16,'ncycles_per_iteration':20,'maxsize':10,'maxdepth':10,'deterministic':True,'progress':False},operators={2:['+','sub','*','/']},variable_names=['x0']); assert out['hall_of_fame']" + + - name: cargo check + run: cargo check --manifest-path pysr_rust_backend/Cargo.toml + + - name: cargo fmt --check + run: cargo fmt --check --manifest-path pysr_rust_backend/Cargo.toml diff --git a/README.md b/README.md index 62f575626..d83ecefe4 100644 --- a/README.md +++ b/README.md @@ -73,7 +73,28 @@ You can install PySR with pip: pip install pysr ``` -Julia dependencies will be installed at first import. +If you plan to use the experimental Rust adapter, select the Rust dependency at install time: + +```bash +pip install "pysr[rust]" +``` + +Julia dependencies will be installed when the Julia backend is used. + +PySR also includes an experimental runtime switch for the Rust adapter: + +```python +PySRRegressor(backend="rust") +``` + +The Rust backend adapter expects the optional `symbolic_regression_rs` Python +module from the `pysr-rust-backend` package. It is installed automatically by +the `rust` extra. In this draft branch, the connector package source lives in +[`pysr_rust_backend`](./pysr_rust_backend). + +The default `backend="auto"` uses the Rust backend when `pysr-rust-backend` is +installed and falls back to Julia otherwise. Set `backend="julia"` to force the +full-featured and supported Julia backend. ### Conda diff --git a/benchmarks/README.md b/benchmarks/README.md index b8a5221c0..9a8186680 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -1,5 +1,24 @@ # Benchmark 1 +## Backend comparison + +To compare the experimental Rust backend against the default Julia backend on +the same generated dataset, run: + +```bash +python benchmarks/compare_backends.py --niterations 3 --repeats 3 +``` + +Use `--mode subprocess` to run each trial in a fresh Python process, which is +useful when measuring startup-heavy workflows: + +```bash +python benchmarks/compare_backends.py --mode subprocess --niterations 1 --repeats 3 +``` + +The script reports wall-clock timing plus the best loss, prediction MSE, and +selected equation for each backend. + The following benchmarks were ran with this command on a node on CCA's BNL cluster (40-cores). At no time was the node fully busy. The tags were put into the file `tags.txt`, and the `benchmark.sh` was copied to the root folder. This is the command used: ```bash diff --git a/benchmarks/compare_backends.py b/benchmarks/compare_backends.py new file mode 100644 index 000000000..497dc059d --- /dev/null +++ b/benchmarks/compare_backends.py @@ -0,0 +1,407 @@ +"""Benchmark PySR Julia and Rust backends on the same synthetic dataset. + +Examples +-------- +Quick smoke benchmark: + + python benchmarks/compare_backends.py --niterations 1 --repeats 1 + +Measure each backend in a fresh Python process: + + python benchmarks/compare_backends.py --mode subprocess --repeats 3 + +Longer run with JSON output: + + python benchmarks/compare_backends.py --niterations 10 --repeats 5 --json-output results.json +""" + +from __future__ import annotations + +import argparse +import csv +import json +import statistics +import subprocess +import sys +import tempfile +import time +from collections.abc import Iterable +from pathlib import Path +from typing import Any + +import numpy as np + +BACKENDS = ("julia", "rust") + + +def make_dataset( + *, + samples: int, + features: int, + seed: int, + noise: float, +) -> tuple[np.ndarray, np.ndarray]: + if features < 2: + raise ValueError("features must be at least 2 for the default benchmark target") + + rng = np.random.default_rng(seed) + X = rng.uniform(-2.0, 2.0, size=(samples, features)).astype(np.float32) + y = ( + 1.5 * np.sin(X[:, 0]) + + 0.75 * X[:, 1] * X[:, 1] + - 0.25 * X[:, 0] * X[:, 1] + + 0.5 + ).astype(np.float32) + if noise > 0: + y = y + rng.normal(0.0, noise, size=samples).astype(np.float32) + return X, y + + +def build_model(backend: str, args: argparse.Namespace, *, seed: int): + from pysr import PySRRegressor + + return PySRRegressor( + backend=backend, + binary_operators=["+", "-", "*", "/"], + unary_operators=["sin", "cos", "exp"], + niterations=args.niterations, + populations=args.populations, + population_size=args.population_size, + ncycles_per_iteration=args.ncycles_per_iteration, + maxsize=args.maxsize, + maxdepth=args.maxdepth, + deterministic=not args.non_deterministic, + random_state=seed, + parallelism=args.parallelism, + progress=False, + verbosity=0, + temp_equation_file=True, + model_selection="accuracy", + ) + + +def run_fit_once( + backend: str, + args: argparse.Namespace, + *, + repeat_index: int, + include_import: bool = False, +) -> dict[str, Any]: + X, y = make_dataset( + samples=args.samples, + features=args.features, + seed=args.seed, + noise=args.noise, + ) + + if include_import: + start = time.perf_counter() + else: + from pysr import PySRRegressor # noqa: F401 + + start = time.perf_counter() + + model = build_model(backend, args, seed=args.seed + repeat_index) + model.fit(X, y) + seconds = time.perf_counter() - start + + prediction = model.predict(X) + mse = float(np.mean((prediction - y) ** 2)) + best = model.get_best() + + return { + "backend": backend, + "mode": "fit", + "repeat": repeat_index, + "seconds": seconds, + "best_loss": float(best["loss"]), + "mse": mse, + "equation": str(best["equation"]), + "n_equations": int(len(model.equations_)), + "backend_version": ( + getattr(model, "rust_backend_version_", None) if backend == "rust" else None + ), + } + + +def worker_command( + args: argparse.Namespace, backend: str, output_path: Path, repeat: int +) -> list[str]: + script = Path(__file__).resolve() + command = [ + sys.executable, + str(script), + "--_worker-backend", + backend, + "--_worker-output", + str(output_path), + "--_worker-repeat", + str(repeat), + "--samples", + str(args.samples), + "--features", + str(args.features), + "--seed", + str(args.seed), + "--noise", + str(args.noise), + "--niterations", + str(args.niterations), + "--populations", + str(args.populations), + "--population-size", + str(args.population_size), + "--ncycles-per-iteration", + str(args.ncycles_per_iteration), + "--maxsize", + str(args.maxsize), + "--maxdepth", + str(args.maxdepth), + "--parallelism", + args.parallelism, + ] + if args.non_deterministic: + command.append("--non-deterministic") + return command + + +def run_subprocess_once( + backend: str, + args: argparse.Namespace, + *, + repeat_index: int, +) -> dict[str, Any]: + with tempfile.TemporaryDirectory() as tmpdir: + output_path = Path(tmpdir) / "result.json" + command = worker_command(args, backend, output_path, repeat_index) + start = time.perf_counter() + completed = subprocess.run( + command, + cwd=Path(__file__).resolve().parents[1], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + seconds = time.perf_counter() - start + + if completed.returncode != 0: + raise RuntimeError( + f"{backend!r} subprocess benchmark failed with exit code " + f"{completed.returncode}\nSTDOUT:\n{completed.stdout}\nSTDERR:\n{completed.stderr}" + ) + result = json.loads(output_path.read_text(encoding="utf-8")) + result["mode"] = "subprocess" + result["worker_seconds"] = result["seconds"] + result["seconds"] = seconds + return result + + +def summarize(results: list[dict[str, Any]]) -> list[dict[str, Any]]: + rows = [] + keys = sorted({(row["backend"], row["mode"]) for row in results}) + for backend, mode in keys: + group = [ + row for row in results if row["backend"] == backend and row["mode"] == mode + ] + seconds = [row["seconds"] for row in group] + losses = [row["best_loss"] for row in group] + mses = [row["mse"] for row in group] + rows.append( + { + "backend": backend, + "mode": mode, + "repeats": len(group), + "mean_seconds": statistics.fmean(seconds), + "median_seconds": statistics.median(seconds), + "min_seconds": min(seconds), + "mean_best_loss": statistics.fmean(losses), + "mean_mse": statistics.fmean(mses), + "last_equation": group[-1]["equation"], + } + ) + return rows + + +def format_seconds(value: float) -> str: + return f"{value:9.3f}" + + +def print_summary(results: list[dict[str, Any]]) -> None: + rows = summarize(results) + print() + print("Backend timing summary") + print( + "backend mode repeats mean_s median_s min_s " + "mean_loss mean_mse last_equation" + ) + print("-" * 108) + for row in rows: + print( + f"{row['backend']:<11} {row['mode']:<13} {row['repeats']:>7} " + f"{format_seconds(row['mean_seconds'])} " + f"{format_seconds(row['median_seconds'])} " + f"{format_seconds(row['min_seconds'])} " + f"{row['mean_best_loss']:11.4g} " + f"{row['mean_mse']:11.4g} " + f"{row['last_equation']}" + ) + + +def write_json(path: Path, results: list[dict[str, Any]]) -> None: + payload = { + "results": results, + "summary": summarize(results), + } + path.write_text(json.dumps(payload, indent=2), encoding="utf-8") + + +def write_csv(path: Path, results: list[dict[str, Any]]) -> None: + fieldnames = [ + "backend", + "mode", + "repeat", + "seconds", + "worker_seconds", + "best_loss", + "mse", + "equation", + "n_equations", + "backend_version", + ] + with path.open("w", newline="", encoding="utf-8") as f: + writer = csv.DictWriter(f, fieldnames=fieldnames, extrasaction="ignore") + writer.writeheader() + writer.writerows(results) + + +def run_requested_benchmarks(args: argparse.Namespace) -> list[dict[str, Any]]: + modes = ["fit", "subprocess"] if args.mode == "both" else [args.mode] + results = [] + + for mode in modes: + for backend in args.backends: + for warmup_index in range(args.warmups): + if mode == "fit": + run_fit_once( + backend, + args, + repeat_index=-(warmup_index + 1), + include_import=False, + ) + else: + run_subprocess_once( + backend, + args, + repeat_index=-(warmup_index + 1), + ) + + for repeat_index in range(args.repeats): + if mode == "fit": + result = run_fit_once( + backend, + args, + repeat_index=repeat_index, + include_import=False, + ) + else: + result = run_subprocess_once( + backend, + args, + repeat_index=repeat_index, + ) + result["mode"] = mode + results.append(result) + print( + f"{backend:<5} {mode:<10} repeat={repeat_index} " + f"seconds={result['seconds']:.3f} " + f"loss={result['best_loss']:.4g} " + f"equation={result['equation']}" + ) + return results + + +def parse_args(argv: Iterable[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--backends", nargs="+", choices=BACKENDS, default=list(BACKENDS) + ) + parser.add_argument("--mode", choices=["fit", "subprocess", "both"], default="fit") + parser.add_argument("--repeats", type=int, default=1) + parser.add_argument("--warmups", type=int, default=0) + parser.add_argument("--samples", type=int, default=256) + parser.add_argument("--features", type=int, default=4) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--noise", type=float, default=0.0) + parser.add_argument("--niterations", type=int, default=3) + parser.add_argument("--populations", type=int, default=4) + parser.add_argument("--population-size", type=int, default=64) + parser.add_argument("--ncycles-per-iteration", type=int, default=100) + parser.add_argument("--maxsize", type=int, default=20) + parser.add_argument("--maxdepth", type=int, default=10) + parser.add_argument( + "--parallelism", + choices=["serial", "multithreading"], + default="serial", + help="Use serial by default so deterministic=True is accepted by both backends.", + ) + parser.add_argument( + "--non-deterministic", + action="store_true", + help="Disable deterministic=True. Useful when benchmarking multithreading.", + ) + parser.add_argument("--json-output", type=Path) + parser.add_argument("--csv-output", type=Path) + + parser.add_argument("--_worker-backend", choices=BACKENDS, help=argparse.SUPPRESS) + parser.add_argument("--_worker-output", type=Path, help=argparse.SUPPRESS) + parser.add_argument("--_worker-repeat", type=int, default=0, help=argparse.SUPPRESS) + args = parser.parse_args(argv) + + if args.repeats < 1: + parser.error("--repeats must be at least 1") + if args.warmups < 0: + parser.error("--warmups must be non-negative") + if args.samples < 2: + parser.error("--samples must be at least 2") + if args.features < 2: + parser.error("--features must be at least 2") + if args._worker_backend and args._worker_output is None: + parser.error("--_worker-output is required with --_worker-backend") + return args + + +def main(argv: Iterable[str] | None = None) -> int: + args = parse_args(argv) + + if args._worker_backend: + result = run_fit_once( + args._worker_backend, + args, + repeat_index=args._worker_repeat, + include_import=True, + ) + args._worker_output.write_text(json.dumps(result), encoding="utf-8") + return 0 + + print("Dataset target: y = 1.5*sin(x0) + 0.75*x1^2 - 0.25*x0*x1 + 0.5") + print( + f"samples={args.samples} features={args.features} " + f"niterations={args.niterations} repeats={args.repeats} " + f"mode={args.mode}" + ) + + results = run_requested_benchmarks(args) + print_summary(results) + + if args.json_output is not None: + write_json(args.json_output, results) + print(f"\nWrote JSON results to {args.json_output}") + if args.csv_output is not None: + write_csv(args.csv_output, results) + print(f"Wrote CSV results to {args.csv_output}") + + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/docs/src/backend.md b/docs/src/backend.md index 801f353cb..e782d2640 100644 --- a/docs/src/backend.md +++ b/docs/src/backend.md @@ -1,5 +1,41 @@ # Customization +## Backend selection + +PySR uses the Julia backend by default: + +```python +PySRRegressor(backend="julia") +``` + +An experimental Rust backend adapter is also available: + +```python +PySRRegressor(backend="rust") +``` + +The Rust backend is intended for vanilla symbolic regression with builtin +operators and faster Julia-free startup. It requires an optional Python package +exposing `symbolic_regression_rs.search`. +Install it with: + +```bash +pip install "pysr[rust]" +``` + +| Feature | `backend="julia"` | `backend="rust"` | +| --- | --- | --- | +| Builtin operators | Yes | Yes | +| Custom Julia operators/losses | Yes | No | +| Template and parametric expression specs | Yes | No | +| Operator and nested constraints | Yes | No | +| Units and dimensional constraints | Yes | No | +| Cluster managers and Julia extensions | Yes | No | + +Use `backend="julia"` when you need the full PySR feature set. + +## Julia backend customization + If you have explored the [options](options.md) and [PySRRegressor reference](api.md), and still haven't figured out how to specify a constraint or objective required for your problem, you might consider editing the backend. The backend of PySR is written as a pure Julia package under the name [SymbolicRegression.jl](https://github.com/astroautomata/SymbolicRegression.jl). This package is accessed with [`juliacall`](https://github.com/JuliaPy/PythonCall.jl), which allows us to transfer objects back and forth between the Python and Julia runtimes. diff --git a/docs/src/options.md b/docs/src/options.md index 7a8f5c11e..17f53c645 100644 --- a/docs/src/options.md +++ b/docs/src/options.md @@ -19,6 +19,7 @@ may find useful include: - [Exporting to numpy, pytorch, and jax](#exporting-to-numpy-pytorch-and-jax) - [Loss functions](#loss) - [Model loading](#model-loading) +- [Backend selection](#backend-selection) These are described below. Also check out the [tuning page](tuning.md) for workflow tips. @@ -32,6 +33,25 @@ at the end of every iteration, which is `.hall_of_fame_{date_time}.csv` by default. It also prints the equations to stdout. +## Backend selection + +`PySRRegressor` uses `backend="julia"` by default, which provides the full +SymbolicRegression.jl feature set. An experimental `backend="rust"` option is +available for vanilla symbolic regression through the optional +`symbolic_regression_rs` Python package. + +Install the Rust backend dependency via: + +```bash +pip install "pysr[rust]" +``` + +`backend="rust"` currently supports builtin operators and a reduced set of +search options. +It rejects Julia-specific features such as custom Julia operators, custom Julia +losses, templates, operator constraints, units, cluster managers, and Julia +extensions. + ## Model selection By default, `PySRRegressor` uses `model_selection='best'` diff --git a/pyproject.toml b/pyproject.toml index fccd37f7a..ed2ced6bf 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,6 +28,7 @@ dependencies = [ ] [project.optional-dependencies] +rust = ["pysr-rust-backend>=0.1.0,<0.2.0"] docs = [ "docstring-parser>=0.17.0", "pyyaml>=6.0.0", diff --git a/pysr/__init__.py b/pysr/__init__.py index 3e3f7c8be..9074a6512 100644 --- a/pysr/__init__.py +++ b/pysr/__init__.py @@ -14,11 +14,6 @@ beartype_this_package() -# This must be imported as early as possible to prevent -# library linking issues caused by numpy/pytorch/etc. importing -# old libraries: -from .julia_import import jl, SymbolicRegression # isort:skip - # Get the version using importlib.metadata (Python >= 3.8 is required): from importlib.metadata import PackageNotFoundError, version @@ -32,7 +27,6 @@ ParametricExpressionSpec, TemplateExpressionSpec, ) -from .julia_extensions import load_all_packages from .logger_specs import AbstractLoggerSpec, TensorBoardLoggerSpec from .sr import PySRRegressor @@ -42,6 +36,23 @@ # package is not installed __version__ = "unknown" + +def __getattr__(name: str): + if name in {"jl", "SymbolicRegression"}: + # Kept lazy so importing PySRRegressor can remain Julia-free. + from .julia_import import SymbolicRegression, jl + + value = {"jl": jl, "SymbolicRegression": SymbolicRegression}[name] + globals()[name] = value + return value + if name == "load_all_packages": + from .julia_extensions import load_all_packages + + globals()[name] = load_all_packages + return load_all_packages + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + __all__ = [ "jl", "SymbolicRegression", diff --git a/pysr/_cli/main.py b/pysr/_cli/main.py index cd81710ff..1e834ff3e 100644 --- a/pysr/_cli/main.py +++ b/pysr/_cli/main.py @@ -11,6 +11,7 @@ runtests_autodiff, runtests_dev, runtests_jax, + runtests_rust, runtests_slurm, runtests_startup, runtests_torch, @@ -46,11 +47,22 @@ def pysr(context): ) def _install(julia_project, quiet, precompile): warnings.warn( - "This command is deprecated. Julia dependencies are now installed at first import." + "This command is deprecated. Julia dependencies are now installed when " + "the Julia backend is used." ) -TEST_OPTIONS = {"main", "jax", "torch", "autodiff", "cli", "dev", "startup", "slurm"} +TEST_OPTIONS = { + "main", + "jax", + "torch", + "autodiff", + "cli", + "dev", + "startup", + "slurm", + "rust", +} @pysr.command("test") @@ -65,7 +77,7 @@ def _install(julia_project, quiet, precompile): def _tests(tests, expressions): """Run parts of the PySR test suite. - Choose from main, jax, torch, autodiff, cli, dev, startup, and slurm. + Choose from main, jax, torch, autodiff, cli, dev, startup, slurm, and rust. You can give multiple tests, separated by commas. """ test_cases = [] @@ -87,6 +99,8 @@ def _tests(tests, expressions): test_cases.extend(runtests_startup(just_tests=True)) elif test == "slurm": test_cases.extend(runtests_slurm(just_tests=True)) + elif test == "rust": + test_cases.extend(runtests_rust(just_tests=True)) else: warnings.warn(f"Invalid test {test}. Skipping.") diff --git a/pysr/backends/__init__.py b/pysr/backends/__init__.py new file mode 100644 index 000000000..c166d2b49 --- /dev/null +++ b/pysr/backends/__init__.py @@ -0,0 +1,5 @@ +"""Backend adapters for PySR search implementations.""" + +from __future__ import annotations + +__all__ = ["rust"] diff --git a/pysr/backends/base.py b/pysr/backends/base.py new file mode 100644 index 000000000..f045225d0 --- /dev/null +++ b/pysr/backends/base.py @@ -0,0 +1,15 @@ +"""Shared backend adapter helpers.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import pandas as pd + + +@dataclass(frozen=True) +class BackendSearchResult: + """Normalized search result returned by non-Julia backends.""" + + hall_of_fame: pd.DataFrame + backend_version: str | None = None diff --git a/pysr/backends/rust.py b/pysr/backends/rust.py new file mode 100644 index 000000000..d33802771 --- /dev/null +++ b/pysr/backends/rust.py @@ -0,0 +1,316 @@ +"""Adapter for the optional Rust symbolic regression backend.""" + +from __future__ import annotations + +from importlib import import_module +from numbers import Real +from pathlib import Path +from typing import TYPE_CHECKING, Any + +import numpy as np +import pandas as pd + +from .base import BackendSearchResult + +if TYPE_CHECKING: + from numpy import ndarray + + from pysr.sr import _DynamicallySetParams + + +_RUST_OPERATOR_ALIASES = { + "-": "sub", + "square": "abs2", +} + +_RUST_BUILTIN_OPERATORS = { + "+", + "*", + "/", + "sub", + "sin", + "cos", + "tan", + "exp", + "log", + "sqrt", + "abs", + "abs2", +} + + +def _rust_only_error(param_name: str) -> NotImplementedError: + return NotImplementedError( + f"`{param_name}` is not supported by `backend='rust'` yet. " + "Use `backend='julia'` for the full PySR feature set." + ) + + +def _normalize_operator_name(operator: str) -> str: + return _RUST_OPERATOR_ALIASES.get(operator, operator) + + +def _is_nonnegative_integer(value: Any) -> bool: + return ( + isinstance(value, Real) + and not isinstance(value, bool) + and float(value).is_integer() + and float(value) >= 0 + ) + + +def _validate_rust_backend_request(model: Any, weights, category) -> None: + from pysr.expression_specs import ExpressionSpec + + if model.nout_ != 1: + raise _rust_only_error("multi-output regression") + if weights is not None: + raise _rust_only_error("weights") + if not isinstance(model.expression_spec_, ExpressionSpec): + raise _rust_only_error("expression_spec") + if category is not None: + raise _rust_only_error("category") + if model.loss_function is not None: + raise _rust_only_error("loss_function") + if model.loss_function_expression is not None: + raise _rust_only_error("loss_function_expression") + if model.elementwise_loss not in (None, "L2DistLoss()"): + raise _rust_only_error("elementwise_loss") + if model.constraints is not None: + raise _rust_only_error("constraints") + if model.nested_constraints is not None: + raise _rust_only_error("nested_constraints") + if model.X_units_ is not None or model.y_units_ is not None: + raise _rust_only_error("units") + if model.fast_cycle: + raise _rust_only_error("fast_cycle") + if model.turbo: + raise _rust_only_error("turbo") + if model.bumper: + raise _rust_only_error("bumper") + if model.autodiff_backend is not None: + raise _rust_only_error("autodiff_backend") + if model.cluster_manager is not None: + raise _rust_only_error("cluster_manager") + if model.worker_imports is not None: + raise _rust_only_error("worker_imports") + if model.logger_spec is not None: + raise _rust_only_error("logger_spec") + if model.warm_start: + raise _rust_only_error("warm_start") + if model.guesses is not None: + raise _rust_only_error("guesses") + if model.optimizer_algorithm != "BFGS": + raise _rust_only_error("optimizer_algorithm") + if model.precision not in (32, 64): + raise _rust_only_error("precision=16") + if model.complexity_mapping is not None: + raise _rust_only_error("complexity_mapping") + if model.complexity_of_operators is not None: + raise _rust_only_error("complexity_of_operators") + if model.complexity_of_constants is not None and not _is_nonnegative_integer( + model.complexity_of_constants + ): + raise _rust_only_error("complexity_of_constants") + if model.complexity_of_variables_ is not None and not _is_nonnegative_integer( + model.complexity_of_variables_ + ): + raise _rust_only_error("complexity_of_variables") + if model.dimensional_constraint_penalty is not None: + raise _rust_only_error("dimensional_constraint_penalty") + if model.dimensionless_constants_only: + raise _rust_only_error("dimensionless_constants_only") + if isinstance(model.early_stop_condition, str): + raise _rust_only_error("string early_stop_condition") + if model.parallelism not in (None, "serial", "multithreading"): + raise _rust_only_error("parallelism") + if model.procs is not None: + raise _rust_only_error("procs") + if model.heap_size_hint_in_bytes is not None: + raise _rust_only_error("heap_size_hint_in_bytes") + if model.worker_timeout is not None: + raise _rust_only_error("worker_timeout") + + +def _load_rust_module(): + try: + return import_module("symbolic_regression_rs") + except ImportError as exc: + raise ImportError( + "The Rust backend requires the optional `symbolic_regression_rs` " + "module from the `pysr-rust-backend` package. Install PySR with " + "`pip install 'pysr[rust]'`, or use `backend='julia'`." + ) from exc + + +def _build_rust_operators(operators: dict[int, list[str]]) -> dict[int, list[str]]: + rust_operators: dict[int, list[str]] = {} + for arity, op_list in operators.items(): + if arity not in (1, 2): + raise _rust_only_error(f"operators with arity {arity}") + rust_op_list = [] + for op in op_list: + if "(" in op: + raise _rust_only_error("inline custom operators") + rust_op = _normalize_operator_name(op) + if rust_op not in _RUST_BUILTIN_OPERATORS: + raise ValueError( + f"`backend='rust'` does not recognize operator {op!r}. " + "Use a Rust builtin operator or `backend='julia'` for " + "custom operators." + ) + rust_op_list.append(rust_op) + rust_operators[arity] = rust_op_list + return rust_operators + + +def _build_rust_options(model: Any, runtime_params: Any, seed: int, X: np.ndarray): + from pysr.sr import _get_batch_size + + batching = model.batching is True or (model.batching == "auto" and len(X) > 1000) + batch_size = _get_batch_size(len(X), runtime_params.batch_size) + + options: dict[str, Any] = { + "seed": int(seed), + "niterations": int(model.niterations), + "populations": int(model.populations), + "population_size": int(model.population_size), + "ncycles_per_iteration": int(model.ncycles_per_iteration), + "batch_size": int(batch_size), + "maxsize": int(model.maxsize), + "maxdepth": int(runtime_params.maxdepth), + "warmup_maxsize_by": float(runtime_params.warmup_maxsize_by), + "parsimony": float(model.parsimony), + "adaptive_parsimony_scaling": float(model.adaptive_parsimony_scaling), + "crossover_probability": float(model.crossover_probability), + "perturbation_factor": float(model.perturbation_factor), + "probability_negate_constant": float(model.probability_negate_constant), + "tournament_selection_n": int(model.tournament_selection_n), + "tournament_selection_p": float(model.tournament_selection_p), + "alpha": float(model.alpha), + "optimizer_nrestarts": int(model.optimizer_nrestarts), + "optimizer_probability": float(model.optimize_probability), + "optimizer_iterations": int(model.optimizer_iterations), + "optimizer_f_calls_limit": int(model.optimizer_f_calls_limit or 10_000), + "fraction_replaced": float(model.fraction_replaced), + "fraction_replaced_hof": float(model.fraction_replaced_hof), + "topn": int(model.topn), + "print_precision": int(model.print_precision), + "max_evals": int(model.max_evals or 0), + "timeout_in_seconds": float(model.timeout_in_seconds or 0.0), + "use_frequency": bool(model.use_frequency), + "use_frequency_in_tournament": bool(model.use_frequency_in_tournament), + "skip_mutation_failures": bool(model.skip_mutation_failures), + "annealing": bool(model.annealing), + "should_optimize_constants": bool(model.should_optimize_constants), + "migration": bool(model.migration), + "hof_migration": bool(model.hof_migration), + "should_simplify": bool(model.should_simplify), + "batching": bool(batching), + "deterministic": bool(model.deterministic), + "parallelism": model.parallelism or "multithreading", + "progress": bool(runtime_params.progress and model.verbosity > 0), + "mutation_weights": { + "mutate_constant": float(model.weight_mutate_constant), + "mutate_operator": float(model.weight_mutate_operator), + "mutate_feature": float(model.weight_mutate_feature), + "swap_operands": float(model.weight_swap_operands), + "rotate_tree": float(model.weight_rotate_tree), + "add_node": float(model.weight_add_node), + "insert_node": float(model.weight_insert_node), + "delete_node": float(model.weight_delete_node), + "simplify": float(model.weight_simplify), + "randomize": float(model.weight_randomize), + "do_nothing": float(model.weight_do_nothing), + "optimize": float(model.weight_optimize), + }, + } + if model.complexity_of_constants is not None: + options["complexity_of_constants"] = int(model.complexity_of_constants) + if model.complexity_of_variables_ is not None: + options["complexity_of_variables"] = int(model.complexity_of_variables_) + if model.early_stop_condition is not None: + options["early_stop_condition"] = float(model.early_stop_condition) + return options + + +def _normalize_hall_of_fame(raw_result: dict[str, Any]) -> BackendSearchResult: + rows = raw_result.get("hall_of_fame") + if not rows: + raise RuntimeError("Rust backend did not return any hall-of-fame equations.") + + df = pd.DataFrame(rows) + df = df.rename( + columns={ + "Complexity": "complexity", + "Loss": "loss", + "Equation": "equation", + } + ) + required = {"complexity", "loss", "equation"} + missing = required.difference(df.columns) + if missing: + raise RuntimeError( + "Rust backend hall-of-fame output is missing required columns: " + + ", ".join(sorted(missing)) + ) + + df = df.loc[:, ["complexity", "loss", "equation"]].copy() + df["complexity"] = df["complexity"].astype(int) + df["loss"] = df["loss"].astype(float) + df["equation"] = df["equation"].astype(str) + df = df.sort_values(["complexity", "loss"], kind="mergesort") + df = df.drop_duplicates(subset=["complexity"], keep="first") + df = df.reset_index(drop=True) + return BackendSearchResult( + hall_of_fame=df, + backend_version=raw_result.get("backend_version"), + ) + + +def _write_hall_of_fame(model: Any, hall_of_fame: pd.DataFrame) -> None: + equation_file = model.get_equation_file() + Path(equation_file).parent.mkdir(parents=True, exist_ok=True) + hall_of_fame.to_csv(equation_file, index=False) + + +def run_rust_backend( + model: Any, + X: "ndarray", + y: "ndarray", + runtime_params: "_DynamicallySetParams", + *, + weights, + category, + seed: int, +): + _validate_rust_backend_request(model, weights, category) + rust_operators = _build_rust_operators(runtime_params.operators) + rust = _load_rust_module() + + if np.issubdtype(np.asarray(X).dtype, np.complexfloating): + raise _rust_only_error("complex input data") + + np_dtype = model._get_precision_mapped_dtype(np.asarray(X)) + if np_dtype not in (np.float32, np.float64): + raise _rust_only_error("non-real precision") + + X_rust = np.ascontiguousarray(np.asarray(X, dtype=np_dtype)) + y_rust = np.ascontiguousarray(np.asarray(y, dtype=np_dtype)) + options = _build_rust_options(model, runtime_params, seed, X_rust) + + raw_result = rust.search( + X_rust, + y_rust, + options=options, + operators=rust_operators, + variable_names=[str(v) for v in model.feature_names_in_], + ) + result = _normalize_hall_of_fame(raw_result) + + model.rust_state_ = raw_result + model.rust_backend_version_ = result.backend_version + model.equation_file_contents_ = [result.hall_of_fame] + _write_hall_of_fame(model, result.hall_of_fame) + model.equations_ = model.get_hof() + return model diff --git a/pysr/deprecated.py b/pysr/deprecated.py index 8905f2823..3b7fb09ce 100644 --- a/pysr/deprecated.py +++ b/pysr/deprecated.py @@ -2,23 +2,24 @@ import warnings -from .julia_import import jl - def install(*args, **kwargs): del args, kwargs warnings.warn( "The `install` function has been removed. " - "PySR now uses the `juliacall` package to install its dependencies automatically at import time. ", + "PySR now uses the `juliacall` package to install its dependencies " + "automatically when the Julia backend is used.", FutureWarning, ) def init_julia(*args, **kwargs): del args, kwargs + from .julia_import import jl + warnings.warn( "The `init_julia` function has been removed. " - "Julia is now initialized automatically at import time.", + "Julia is now initialized automatically when the Julia backend is used.", FutureWarning, ) return jl diff --git a/pysr/expression_specs.py b/pysr/expression_specs.py index bd8d5ec71..1fceae36e 100644 --- a/pysr/expression_specs.py +++ b/pysr/expression_specs.py @@ -11,10 +11,10 @@ import pandas as pd from .export import add_export_formats -from .julia_helpers import jl_array -from .julia_import import AnyValue, SymbolicRegression, jl from .utils import ArrayLike +AnyValue = Any + try: from typing import TypeAlias except ImportError: @@ -88,6 +88,8 @@ class ExpressionSpec(AbstractExpressionSpec): """The default expression specification, with no special behavior.""" def julia_expression_spec(self): + from .julia_import import SymbolicRegression + return SymbolicRegression.ExpressionSpec() def create_exports( @@ -252,6 +254,8 @@ def _get_cache_key(self): ) def julia_expression_spec(self): + from .julia_import import SymbolicRegression + key = self._get_cache_key() if key in self._spec_cache: return self._spec_cache[key] @@ -267,6 +271,8 @@ def julia_expression_spec(self): return result def _call_template_macro(self): + from .julia_import import jl + return jl.seval(self._template_macro_str()) def _template_macro_str(self): @@ -282,6 +288,8 @@ def _template_macro_str(self): """) def julia_expression_options(self): + from .julia_import import jl + f_combine = jl.seval(self.combine) creator = jl.seval(""" function _pysr_create_template_structure( @@ -378,6 +386,8 @@ def __init__(self, max_parameters: int): self.max_parameters = max_parameters def julia_expression_spec(self): + from .julia_import import SymbolicRegression + return SymbolicRegression.ParametricExpressionSpec( max_parameters=self.max_parameters, warn=False ) @@ -402,6 +412,8 @@ def __init__(self, expression): self.expression = expression def __call__(self, X: np.ndarray, *args): + from .julia_helpers import jl_array + raw_output = self.expression(jl_array(X.T), *args) return np.array(raw_output).T diff --git a/pysr/logger_specs.py b/pysr/logger_specs.py index af8a68550..a4693c5ed 100644 --- a/pysr/logger_specs.py +++ b/pysr/logger_specs.py @@ -4,8 +4,7 @@ from dataclasses import dataclass from typing import Any -from .julia_helpers import jl_array, jl_dict -from .julia_import import AnyValue, jl +AnyValue = Any class AbstractLoggerSpec(ABC): @@ -47,6 +46,8 @@ class TensorBoardLoggerSpec(AbstractLoggerSpec): overwrite: bool = False def create_logger(self) -> AnyValue: + from .julia_import import jl + # We assume that TensorBoardLogger is already imported via `julia_extensions.py` make_logger = jl.seval(""" function make_logger(log_dir::AbstractString, overwrite::Bool, log_interval::Int) @@ -61,6 +62,9 @@ def create_logger(self) -> AnyValue: return make_logger(log_dir, self.overwrite, self.log_interval) def write_hparams(self, logger: AnyValue, hparams: dict[str, Any]) -> None: + from .julia_helpers import jl_array, jl_dict + from .julia_import import jl + base_logger = jl.SymbolicRegression.get_logger(logger) writer = jl.seval("TensorBoardLogger.write_hparams!") jl_clean_hparams = jl_dict( @@ -81,5 +85,7 @@ def write_hparams(self, logger: AnyValue, hparams: dict[str, Any]) -> None: ) def close(self, logger: AnyValue) -> None: + from .julia_import import jl + base_logger = jl.SymbolicRegression.get_logger(logger) jl.close(base_logger) diff --git a/pysr/param_groupings.yml b/pysr/param_groupings.yml index f4b769c27..3f865fa64 100644 --- a/pysr/param_groupings.yml +++ b/pysr/param_groupings.yml @@ -1,4 +1,6 @@ - The Algorithm: + - Backend: + - backend - Creating the Search Space: - binary_operators - unary_operators diff --git a/pysr/sr.py b/pysr/sr.py index f09b8bc50..69b700692 100644 --- a/pysr/sr.py +++ b/pysr/sr.py @@ -3,6 +3,7 @@ from __future__ import annotations import copy +import importlib.util import logging import os import pickle as pkl @@ -42,17 +43,6 @@ parametric_expression_deprecation_warning, ) from .feature_selection import run_feature_selection -from .julia_extensions import load_required_packages -from .julia_helpers import ( - _escape_filename, - _load_cluster_manager, - jl_array, - jl_deserialize, - jl_is_function, - jl_named_tuple, - jl_serialize, -) -from .julia_import import AnyValue, SymbolicRegression, VectorValue, jl from .logger_specs import AbstractLoggerSpec from .utils import ( ArrayLike, @@ -76,6 +66,9 @@ from typing_extensions import List +AnyValue = Any +VectorValue = Any + ALREADY_RAN = False pysr_logger = logging.getLogger(__name__) @@ -150,6 +143,8 @@ def _maybe_create_inline_operators( is_user_defined_operator = "(" in op if is_user_defined_operator: + from .julia_import import jl + jl.seval(op) # Cut off from the first non-alphanumeric char: first_non_char = [j for j, char in enumerate(op) if char == "("][0] @@ -245,6 +240,8 @@ def _validate_elementwise_loss( If the probe fails, it raises a `ValueError` describing the expected signature. """ + from .julia_helpers import jl_is_function + from .julia_import import jl # This can be either a LossFunctions.jl object (e.g. `L2DistLoss()`) or a Julia function. # Only validate arity when the evaluated object is actually a function. @@ -277,6 +274,9 @@ def _validate_custom_objective( signature, other_alternative=None, ) -> None: + from .julia_helpers import jl_is_function + from .julia_import import jl + if not jl_is_function(custom_objective): raise ValueError(f"`{knob}` must evaluate to a callable Julia function.") @@ -387,6 +387,13 @@ class PySRRegressor(MultiOutputMixin, RegressorMixin, BaseEstimator): `'best'` selects the candidate model with the highest score among expressions with a loss better than at least 1.5x the most accurate model. + backend : {"auto", "julia", "rust"} + Search backend to use. The default `"auto"` backend uses the optional + Rust backend if its `symbolic_regression_rs` module is installed, + otherwise it falls back to the full-featured Julia backend. The + optional `"rust"` backend targets vanilla symbolic regression through + the external `symbolic_regression_rs` package and supports a smaller + feature set. binary_operators : list[str] List of strings for binary operators used in the search. See the [operators page](https://ai.damtp.cam.ac.uk/pysr/operators/) @@ -949,6 +956,7 @@ def __init__( self, model_selection: Literal["best", "accuracy", "score"] = "best", *, + backend: Literal["auto", "julia", "rust"] = "auto", binary_operators: list[str] | None = None, unary_operators: list[str] | None = None, operators: dict[int, list[str]] | None = None, @@ -1061,6 +1069,7 @@ def __init__( **kwargs, ): # Hyperparameters + self.backend = backend # - Model search parameters self.model_selection = model_selection self.binary_operators = binary_operators @@ -1295,6 +1304,9 @@ def from_file( with open(pkl_filename, "rb") as f: model = cast("PySRRegressor", pkl.load(f)) + if not hasattr(model, "backend"): + model.backend = "julia" + # Update any parameters if necessary, such as # extra_sympy_mappings: model.set_params(**pysr_kwargs) @@ -1498,11 +1510,15 @@ def equations(self): # pragma: no cover @property def julia_options_(self): """The deserialized julia options.""" + from .julia_helpers import jl_deserialize + return jl_deserialize(self.julia_options_stream_) @property def julia_state_(self): """The deserialized state.""" + from .julia_helpers import jl_deserialize + return cast( Union[Tuple[VectorValue, AnyValue], None], jl_deserialize(self.julia_state_stream_), @@ -1578,6 +1594,25 @@ def equation_file_(self): "instead. For loading, you should pass `run_directory`." ) + def _resolve_backend(self) -> Literal["julia", "rust"]: + """Resolve the runtime backend without importing Julia.""" + if self.backend == "rust": + return "rust" + if self.backend == "julia": + return "julia" + if self.backend != "auto": + raise ValueError("`backend` must be one of 'auto', 'julia', or 'rust'.") + + if "symbolic_regression_rs" in sys.modules: + return "rust" + try: + rust_spec = importlib.util.find_spec("symbolic_regression_rs") + except (ImportError, ValueError): + rust_spec = None + if rust_spec is not None: + return "rust" + return "julia" + def _setup_equation_file(self): """Set the pathname of the output directory.""" if self.warm_start and ( @@ -1600,11 +1635,18 @@ def _setup_equation_file(self): if self.output_directory is None else self.output_directory ) - self.run_id_ = ( - cast(str, SymbolicRegression.SearchUtilsModule.generate_run_id()) - if self.run_id is None - else self.run_id - ) + if self.run_id is None: + backend = getattr(self, "backend_", self.backend) + if backend == "rust": + self.run_id_ = next(tempfile._get_candidate_names()) + else: + from .julia_import import SymbolicRegression + + self.run_id_ = cast( + str, SymbolicRegression.SearchUtilsModule.generate_run_id() + ) + else: + self.run_id_ = self.run_id if self.temp_equation_file: assert self.output_directory is None @@ -1627,6 +1669,8 @@ def _validate_and_modify_params(self) -> _DynamicallySetParams: """ # Immutable parameter validation # Ensure instance parameters are allowable values: + if self.backend not in ("auto", "julia", "rust"): + raise ValueError("`backend` must be one of 'auto', 'julia', or 'rust'.") # Validate operators vs binary_operators/unary_operators mutual exclusion if self.operators is not None: @@ -1997,7 +2041,7 @@ def _run( seed: int, ): """ - Run the symbolic regression fitting process on the julia backend. + Run the symbolic regression fitting process on the configured backend. Parameters ---------- @@ -2017,7 +2061,7 @@ def _run( argument should be a list of integers representing the category of each sample in `X`. seed : int - Random seed for julia backend process. + Random seed passed to the backend runtime. Returns ------- @@ -2027,8 +2071,32 @@ def _run( Raises ------ ImportError - Raised when the julia backend fails to import a package. + Raised when the selected backend fails to import a required package. """ + backend = getattr(self, "backend_", self.backend) + if backend == "rust": + from .backends.rust import run_rust_backend + + return run_rust_backend( + self, + X, + y, + runtime_params, + weights=weights, + category=category, + seed=seed, + ) + + from .julia_extensions import load_required_packages + from .julia_helpers import ( + _escape_filename, + _load_cluster_manager, + jl_array, + jl_is_function, + jl_serialize, + ) + from .julia_import import SymbolicRegression, jl + # Need to be global as we don't want to recreate/reinstate julia for # every new instance of PySRRegressor global ALREADY_RAN @@ -2487,6 +2555,7 @@ def fit( self.X_units_ = None self.y_units_ = None + self.backend_ = self._resolve_backend() self._setup_equation_file() self._clear_equation_file_contents() @@ -2684,6 +2753,8 @@ def predict( X = X.astype(self._get_precision_mapped_dtype(X)) if category is not None: + from .julia_helpers import jl_array + offset_for_julia_indexing = 1 args: tuple = ( jl_array((category + offset_for_julia_indexing).astype(np.int64)), @@ -3085,6 +3156,8 @@ def _prepare_guesses_for_julia(guesses, nout) -> VectorValue | None: jl_guesses: VectorValue | None Julia-compatible guesses array or None if no guesses provided """ + from .julia_helpers import jl_array, jl_named_tuple + if guesses is None: return None diff --git a/pysr/test/__init__.py b/pysr/test/__init__.py index 0997991fd..8f65f1fe8 100644 --- a/pysr/test/__init__.py +++ b/pysr/test/__init__.py @@ -1,11 +1,56 @@ -from .test_autodiff import runtests as runtests_autodiff -from .test_cli import get_runtests as get_runtests_cli -from .test_dev import runtests as runtests_dev -from .test_jax import runtests as runtests_jax -from .test_main import runtests -from .test_slurm import runtests as runtests_slurm -from .test_startup import runtests as runtests_startup -from .test_torch import runtests as runtests_torch +def runtests(*args, **kwargs): + from .test_main import runtests as _runtests + + return _runtests(*args, **kwargs) + + +def runtests_jax(*args, **kwargs): + from .test_jax import runtests as _runtests + + return _runtests(*args, **kwargs) + + +def runtests_torch(*args, **kwargs): + from .test_torch import runtests as _runtests + + return _runtests(*args, **kwargs) + + +def runtests_autodiff(*args, **kwargs): + from .test_autodiff import runtests as _runtests + + return _runtests(*args, **kwargs) + + +def get_runtests_cli(*args, **kwargs): + from .test_cli import get_runtests as _get_runtests + + return _get_runtests(*args, **kwargs) + + +def runtests_startup(*args, **kwargs): + from .test_startup import runtests as _runtests + + return _runtests(*args, **kwargs) + + +def runtests_dev(*args, **kwargs): + from .test_dev import runtests as _runtests + + return _runtests(*args, **kwargs) + + +def runtests_slurm(*args, **kwargs): + from .test_slurm import runtests as _runtests + + return _runtests(*args, **kwargs) + + +def runtests_rust(*args, **kwargs): + from .test_rust_backend import runtests as _runtests + + return _runtests(*args, **kwargs) + __all__ = [ "runtests", @@ -16,4 +61,5 @@ "runtests_startup", "runtests_dev", "runtests_slurm", + "runtests_rust", ] diff --git a/pysr/test/test_cli.py b/pysr/test/test_cli.py index f374ddbda..e50552939 100644 --- a/pysr/test/test_cli.py +++ b/pysr/test/test_cli.py @@ -3,76 +3,74 @@ from click import testing as click_testing +from .._cli.main import pysr -def get_runtests(): - # Lazy load to avoid circular imports. - - from .._cli.main import pysr - - class TestCli(unittest.TestCase): - # TODO: Include test for custom project here. - def setUp(self): - self.cli_runner = click_testing.CliRunner() - - def test_help_on_all_commands(self): - expected = dedent(""" - Usage: pysr [OPTIONS] COMMAND [ARGS]... - - Options: - --help Show this message and exit. - - Commands: - install DEPRECATED (dependencies are now installed at import). - test Run parts of the PySR test suite. - """) - result = self.cli_runner.invoke(pysr, ["--help"]) - self.assertEqual(result.output.strip(), expected.strip()) - self.assertEqual(result.exit_code, 0) - def test_help_on_install(self): - expected = dedent(""" - Usage: pysr install [OPTIONS] +class TestCli(unittest.TestCase): + # TODO: Include test for custom project here. + def setUp(self): + self.cli_runner = click_testing.CliRunner() - DEPRECATED (dependencies are now installed at import). + def test_help_on_all_commands(self): + expected = dedent(""" + Usage: pysr [OPTIONS] COMMAND [ARGS]... Options: - -p, --project TEXT - -q, --quiet Disable logging. - --precompile - --no-precompile - --help Show this message and exit. - """) - result = self.cli_runner.invoke(pysr, ["install", "--help"]) - self.assertEqual(result.output.strip(), expected.strip()) - self.assertEqual(result.exit_code, 0) + --help Show this message and exit. + + Commands: + install DEPRECATED (dependencies are now installed at import). + test Run parts of the PySR test suite. + """) + result = self.cli_runner.invoke(pysr, ["--help"]) + self.assertEqual(result.output.strip(), expected.strip()) + self.assertEqual(result.exit_code, 0) + + def test_help_on_install(self): + expected = dedent(""" + Usage: pysr install [OPTIONS] + + DEPRECATED (dependencies are now installed at import). + + Options: + -p, --project TEXT + -q, --quiet Disable logging. + --precompile + --no-precompile + --help Show this message and exit. + """) + result = self.cli_runner.invoke(pysr, ["install", "--help"]) + self.assertEqual(result.output.strip(), expected.strip()) + self.assertEqual(result.exit_code, 0) + + def test_help_on_test(self): + expected = dedent(""" + Usage: pysr test [OPTIONS] TESTS + + Run parts of the PySR test suite. + + Choose from main, jax, torch, autodiff, cli, dev, startup, slurm, and rust. + You can give multiple tests, separated by commas. + + Options: + -k TEXT Filter expressions to select specific tests. + --help Show this message and exit. + """) + result = self.cli_runner.invoke(pysr, ["test", "--help"]) + self.assertEqual(result.output.strip(), expected.strip()) + self.assertEqual(result.exit_code, 0) + + +def runtests(just_tests=False): + """Run all tests in cliTest.py.""" + if just_tests: + return [TestCli] + loader = unittest.TestLoader() + suite = unittest.TestSuite() + suite.addTests(loader.loadTestsFromTestCase(TestCli)) + runner = unittest.TextTestRunner() + return runner.run(suite) - def test_help_on_test(self): - expected = dedent(""" - Usage: pysr test [OPTIONS] TESTS - - Run parts of the PySR test suite. - - Choose from main, jax, torch, autodiff, cli, dev, startup, and slurm. You can - give multiple tests, separated by commas. - - Options: - -k TEXT Filter expressions to select specific tests. - --help Show this message and exit. - """) - result = self.cli_runner.invoke(pysr, ["test", "--help"]) - self.assertEqual(result.output.strip(), expected.strip()) - self.assertEqual(result.exit_code, 0) - - def runtests(just_tests=False): - """Run all tests in cliTest.py.""" - tests = [TestCli] - if just_tests: - return tests - loader = unittest.TestLoader() - suite = unittest.TestSuite() - for test in tests: - suite.addTests(loader.loadTestsFromTestCase(test)) - runner = unittest.TextTestRunner() - return runner.run(suite) +def get_runtests(): return runtests diff --git a/pysr/test/test_dev.py b/pysr/test/test_dev.py index ce9675900..01e61bf5f 100644 --- a/pysr/test/test_dev.py +++ b/pysr/test/test_dev.py @@ -45,7 +45,12 @@ def test_simple_change_to_backend(self): cwd=repo_root, ) self.assertEqual(test_result.returncode, 0) - self.assertEqual(test_result.stdout.decode("utf-8").strip(), "2.3") + stdout_lines = [ + line.strip() + for line in test_result.stdout.decode("utf-8").splitlines() + if line.strip() + ] + self.assertEqual(stdout_lines[-1], "2.3") def runtests(just_tests=False): diff --git a/pysr/test/test_rust_backend.py b/pysr/test/test_rust_backend.py new file mode 100644 index 000000000..60948c3e6 --- /dev/null +++ b/pysr/test/test_rust_backend.py @@ -0,0 +1,458 @@ +import importlib.util +import pickle +import subprocess +import sys +import tempfile +import types +import unittest +import warnings +from pathlib import Path +from unittest.mock import patch + +import numpy as np + +from pysr import PySRRegressor +from pysr.expression_specs import ParametricExpressionSpec, TemplateExpressionSpec +from pysr.logger_specs import TensorBoardLoggerSpec + + +class TestRustBackend(unittest.TestCase): + def tearDown(self): + sys.modules.pop("symbolic_regression_rs", None) + + def test_importing_regressor_does_not_import_juliacall(self): + code = ( + "import sys; " + "from pysr import PySRRegressor; " + "assert 'juliacall' not in sys.modules; " + "assert PySRRegressor().backend == 'auto'" + ) + result = subprocess.run( + [sys.executable, "-c", code], + cwd=".", + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + self.assertEqual(result.returncode, 0, result.stderr) + + def test_auto_backend_resolves_to_rust_when_module_is_available(self): + self._install_fake_rust_module() + + X = np.linspace(-1, 1, 16, dtype=np.float32).reshape(-1, 1) + y = X[:, 0] + model = PySRRegressor( + binary_operators=["+"], + niterations=1, + populations=1, + population_size=16, + deterministic=True, + random_state=0, + progress=False, + temp_equation_file=True, + model_selection="accuracy", + ) + + model.fit(X, y) + + self.assertEqual(model.backend, "auto") + self.assertEqual(model.backend_, "rust") + calls = sys.modules["symbolic_regression_rs"].calls + self.assertEqual(len(calls), 1) + + def test_auto_backend_resolves_to_julia_when_module_is_unavailable(self): + sys.modules.pop("symbolic_regression_rs", None) + model = PySRRegressor() + + with patch("pysr.sr.importlib.util.find_spec", return_value=None): + self.assertEqual(model._resolve_backend(), "julia") + + def test_rust_backend_uses_optional_search_module(self): + self._install_fake_rust_module() + + X = np.linspace(-1, 1, 16, dtype=np.float32).reshape(-1, 1) + y = X[:, 0] + model = PySRRegressor( + backend="rust", + binary_operators=["+", "-"], + niterations=1, + populations=1, + population_size=16, + deterministic=True, + random_state=0, + parallelism="serial", + progress=False, + temp_equation_file=True, + model_selection="accuracy", + complexity_of_constants=2, + ) + + model.fit(X, y, complexity_of_variables=3) + + calls = sys.modules["symbolic_regression_rs"].calls + self.assertEqual(len(calls), 1) + self.assertEqual(calls[0]["operators"], {2: ["+", "sub"]}) + self.assertEqual(calls[0]["variable_names"], ["x0"]) + self.assertEqual(calls[0]["options"]["niterations"], 1) + self.assertEqual(calls[0]["options"]["batch_size"], 16) + self.assertEqual(calls[0]["options"]["parallelism"], "serial") + self.assertEqual(calls[0]["options"]["complexity_of_constants"], 2) + self.assertEqual(calls[0]["options"]["complexity_of_variables"], 3) + self.assertEqual(model.rust_backend_version_, "test") + self.assertLess(model.get_best()["loss"], 1e-8) + self.assertEqual(str(model.sympy()), "x0") + self.assertIsInstance(model.latex(), str) + np.testing.assert_allclose(model.predict(X), y) + self.assertIn("lambda_format", model.equations_.columns) + + def test_rust_backend_checkpoint_and_csv_loading(self): + self._install_fake_rust_module() + + X = np.linspace(-1, 1, 16, dtype=np.float32).reshape(-1, 1) + y = X[:, 0] + with tempfile.TemporaryDirectory() as tmpdir: + model = PySRRegressor( + backend="rust", + binary_operators=["+", "-"], + niterations=1, + populations=1, + population_size=16, + deterministic=True, + random_state=0, + progress=False, + output_directory=tmpdir, + run_id="rust-run", + model_selection="accuracy", + ) + model.fit(X, y) + + run_directory = Path(tmpdir) / "rust-run" + equation_file = run_directory / "hall_of_fame.csv" + checkpoint_file = run_directory / "checkpoint.pkl" + self.assertTrue(equation_file.exists()) + self.assertTrue(checkpoint_file.exists()) + + loaded = PySRRegressor.from_file(run_directory=run_directory) + np.testing.assert_allclose(loaded.predict(X), y) + self.assertEqual(loaded.backend, "rust") + + checkpoint_file.unlink() + loaded_from_csv = PySRRegressor.from_file( + run_directory=run_directory, + backend="rust", + binary_operators=["+", "-"], + n_features_in=1, + ) + np.testing.assert_allclose(loaded_from_csv.predict(X), y) + + def test_rust_backend_pickle_roundtrip(self): + self._install_fake_rust_module() + + X = np.linspace(-1, 1, 16, dtype=np.float32).reshape(-1, 1) + y = X[:, 0] + model = PySRRegressor( + backend="rust", + binary_operators=["+", "-"], + niterations=1, + populations=1, + population_size=16, + deterministic=True, + random_state=0, + progress=False, + temp_equation_file=True, + model_selection="accuracy", + ) + model.fit(X, y) + + loaded = pickle.loads(pickle.dumps(model)) + + self.assertEqual(loaded.backend, "rust") + np.testing.assert_allclose(loaded.predict(X), y) + + def _install_fake_rust_module(self): + calls = [] + + def search(X, y, *, options, operators, variable_names): + calls.append( + { + "X": X, + "y": y, + "options": options, + "operators": operators, + "variable_names": variable_names, + } + ) + return { + "backend_version": "test", + "hall_of_fame": [ + {"complexity": 1, "loss": 0.0, "equation": variable_names[0]}, + ], + } + + fake_module = types.SimpleNamespace(search=search, __version__="test") + fake_module.calls = calls + sys.modules["symbolic_regression_rs"] = fake_module + + def _assert_unsupported_before_import( + self, + unsupported_name, + *, + model_kwargs=None, + fit_kwargs=None, + X=None, + y=None, + ): + model_kwargs = model_kwargs or {} + fit_kwargs = fit_kwargs or {} + if X is None: + X = np.ones((4, 1), dtype=np.float32) + if y is None: + y = X[:, 0] + model = PySRRegressor( + backend="rust", + niterations=1, + populations=1, + progress=False, + temp_equation_file=True, + **model_kwargs, + ) + + sys.modules.pop("symbolic_regression_rs", None) + with patch( + "pysr.backends.rust.import_module", + side_effect=AssertionError("Rust module should not be imported"), + ): + with self.assertRaisesRegex(NotImplementedError, unsupported_name): + with warnings.catch_warnings(): + warnings.simplefilter("ignore", FutureWarning) + model.fit(X, y, **fit_kwargs) + + def test_rust_backend_missing_optional_dependency_errors(self): + X = np.ones((4, 1), dtype=np.float32) + y = X[:, 0] + model = PySRRegressor( + backend="rust", + niterations=1, + populations=1, + progress=False, + temp_equation_file=True, + ) + + with patch("pysr.backends.rust.import_module", side_effect=ImportError): + with self.assertRaisesRegex(ImportError, "pysr\\[rust\\]"): + model.fit(X, y) + + def test_rust_backend_maps_square_to_abs2(self): + self._install_fake_rust_module() + + X = np.ones((4, 1), dtype=np.float32) + y = X[:, 0] + model = PySRRegressor( + backend="rust", + unary_operators=["square"], + niterations=1, + populations=1, + progress=False, + temp_equation_file=True, + ) + + model.fit(X, y) + + calls = sys.modules["symbolic_regression_rs"].calls + self.assertEqual(calls[0]["operators"][1], ["abs2"]) + + def test_rust_backend_uses_pysr_batch_size_defaults(self): + self._install_fake_rust_module() + + X = np.linspace(-1, 1, 1200, dtype=np.float32).reshape(-1, 1) + y = X[:, 0] + model = PySRRegressor( + backend="rust", + batching=True, + binary_operators=["+"], + niterations=1, + populations=1, + progress=False, + temp_equation_file=True, + ) + + model.fit(X, y) + + calls = sys.modules["symbolic_regression_rs"].calls + self.assertTrue(calls[0]["options"]["batching"]) + self.assertEqual(calls[0]["options"]["batch_size"], 128) + + def test_rust_backend_rejects_unsupported_builtin_operator(self): + X = np.ones((4, 1), dtype=np.float32) + y = X[:, 0] + model = PySRRegressor( + backend="rust", + unary_operators=["cube"], + niterations=1, + populations=1, + progress=False, + temp_equation_file=True, + ) + + with self.assertRaisesRegex(ValueError, "cube"): + model.fit(X, y) + + def test_rust_backend_rejects_inline_custom_operator_before_import(self): + X = np.ones((4, 1), dtype=np.float32) + y = X[:, 0] + model = PySRRegressor( + backend="rust", + unary_operators=["inv(x) = 1 / x"], + extra_sympy_mappings={"inv": lambda x: 1 / x}, + niterations=1, + populations=1, + progress=False, + temp_equation_file=True, + ) + + with self.assertRaisesRegex(NotImplementedError, "inline custom operators"): + model.fit(X, y) + + def test_rust_backend_rejects_custom_loss_before_import(self): + unsupported_cases = [ + ( + "loss_function", + {"loss_function": "loss(tree, dataset, options) = 0.0"}, + {}, + ), + ( + "loss_function_expression", + {"loss_function_expression": "loss(expr, dataset, options) = 0.0"}, + {}, + ), + ("elementwise_loss", {"elementwise_loss": "L1DistLoss()"}, {}), + ("constraints", {"constraints": {"*": (2, 2)}}, {}), + ("complexity_mapping", {"complexity_mapping": "x -> 1"}, {}), + ("complexity_of_operators", {"complexity_of_operators": {"*": 2}}, {}), + ("complexity_of_constants", {"complexity_of_constants": 1.5}, {}), + ( + "complexity_of_variables", + {}, + {"complexity_of_variables": [1]}, + ), + ("nested_constraints", {"nested_constraints": {"sin": {"sin": 0}}}, {}), + ( + "dimensional_constraint_penalty", + {"dimensional_constraint_penalty": 1000.0}, + {}, + ), + ( + "dimensionless_constants_only", + {"dimensionless_constants_only": True}, + {}, + ), + ( + "expression_spec", + { + "expression_spec": TemplateExpressionSpec( + "f(x0)", expressions=["f"], variable_names=["x0"] + ) + }, + {}, + ), + ( + "expression_spec", + {"expression_spec": ParametricExpressionSpec(max_parameters=1)}, + {"category": np.zeros(4, dtype=np.int64)}, + ), + ("fast_cycle", {"fast_cycle": True}, {}), + ("turbo", {"turbo": True}, {}), + ("bumper", {"bumper": True}, {}), + ("autodiff_backend", {"autodiff_backend": "Zygote"}, {}), + ("cluster_manager", {"cluster_manager": "multiprocessing"}, {}), + ("worker_imports", {"worker_imports": ["SomePackage"]}, {}), + ("logger_spec", {"logger_spec": TensorBoardLoggerSpec()}, {}), + ("warm_start", {"warm_start": True}, {}), + ("guesses", {"guesses": ["x0"]}, {}), + ("optimizer_algorithm", {"optimizer_algorithm": "NelderMead"}, {}), + ("precision=16", {"precision": 16}, {}), + ( + "string early_stop_condition", + {"early_stop_condition": "stop_if(loss, complexity) = loss < 1e-6"}, + {}, + ), + ("parallelism", {"parallelism": "multiprocessing"}, {}), + ("procs", {"procs": 2}, {}), + ("heap_size_hint_in_bytes", {"heap_size_hint_in_bytes": 1_000_000}, {}), + ("worker_timeout", {"worker_timeout": 10.0}, {}), + ( + "units", + {}, + {"X_units": ["m"], "y_units": "m"}, + ), + ( + "category", + {}, + {"category": np.zeros(4, dtype=np.int64)}, + ), + ( + "weights", + {}, + {"weights": np.ones(4, dtype=np.float32)}, + ), + ] + + for unsupported_name, model_kwargs, fit_kwargs in unsupported_cases: + with self.subTest(unsupported_name=unsupported_name): + self._assert_unsupported_before_import( + unsupported_name, + model_kwargs=model_kwargs, + fit_kwargs=fit_kwargs, + ) + + def test_rust_backend_rejects_multi_output_before_import(self): + X = np.ones((4, 1), dtype=np.float32) + y = np.ones((4, 2), dtype=np.float32) + + self._assert_unsupported_before_import( + "multi-output regression", + X=X, + y=y, + ) + + def test_rust_backend_real_wrapper_smoke_when_installed(self): + if importlib.util.find_spec("symbolic_regression_rs") is None: + self.skipTest("symbolic_regression_rs is not installed") + + pre_modules = set(sys.modules) + X = np.linspace(-1, 1, 32, dtype=np.float32).reshape(-1, 1) + y = X[:, 0] + model = PySRRegressor( + backend="rust", + binary_operators=["+", "-", "*"], + niterations=1, + populations=1, + population_size=16, + ncycles_per_iteration=20, + maxsize=10, + maxdepth=10, + deterministic=True, + random_state=0, + progress=False, + temp_equation_file=True, + model_selection="accuracy", + ) + + model.fit(X, y) + + self.assertLess(model.get_best()["loss"], 1e-3) + np.testing.assert_allclose(model.predict(X), y, atol=0.05) + self.assertIsInstance(model.latex(), str) + self.assertNotIn("juliacall", set(sys.modules) - pre_modules) + + +def runtests(just_tests=False): + tests = [TestRustBackend] + if just_tests: + return tests + suite = unittest.TestSuite() + loader = unittest.TestLoader() + for test in tests: + suite.addTests(loader.loadTestsFromTestCase(test)) + runner = unittest.TextTestRunner() + return runner.run(suite) diff --git a/pysr/test/test_startup.py b/pysr/test/test_startup.py index 8bd81cd93..1c8af1915 100644 --- a/pysr/test/test_startup.py +++ b/pysr/test/test_startup.py @@ -110,15 +110,15 @@ def test_warm_start_from_file(self): def test_bad_startup_options(self): warning_tests = [ dict( - code='import os; os.environ["PYTHON_JULIACALL_HANDLE_SIGNALS"] = "no"; import pysr', + code='import os; os.environ["PYTHON_JULIACALL_HANDLE_SIGNALS"] = "no"; from pysr import jl', msg="PYTHON_JULIACALL_HANDLE_SIGNALS environment variable is set", ), dict( - code='import os; os.environ["PYTHON_JULIACALL_THREADS"] = "1"; import pysr', + code='import os; os.environ["PYTHON_JULIACALL_THREADS"] = "1"; from pysr import jl', msg="PYTHON_JULIACALL_THREADS environment variable is set", ), dict( - code="import juliacall; import pysr", + code="import juliacall; from pysr import jl", msg="juliacall module already imported.", ), ] diff --git a/pysr_rust_backend/Cargo.lock b/pysr_rust_backend/Cargo.lock new file mode 100644 index 000000000..5b781cdb8 --- /dev/null +++ b/pysr_rust_backend/Cargo.lock @@ -0,0 +1,440 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + +[[package]] +name = "bumpalo" +version = "3.20.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "crossbeam-deque" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9dd111b7b7f7d55b72c0a6ae361660ee5853c9af73f70c3c2ef6858b950e2e51" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" + +[[package]] +name = "dynamic_expressions" +version = "0.10.1" +source = "git+https://github.com/astroautomata/symbolic_regression.rs#60a4be904265ecd945be6a02e9f0a4b689801eb3" +dependencies = [ + "ndarray", + "num-traits", + "paste", + "rustc-hash", +] + +[[package]] +name = "either" +version = "1.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" + +[[package]] +name = "fastrand" +version = "2.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6" + +[[package]] +name = "futures-core" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" + +[[package]] +name = "futures-task" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "037711b3d59c33004d3856fbdc83b99d4ff37a24768fa1be9ce3538a1cde4393" + +[[package]] +name = "futures-util" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" +dependencies = [ + "futures-core", + "futures-task", + "pin-project-lite", + "slab", +] + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "js-sys" +version = "0.3.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53b44bfcdb3f8d5837a46dae1ca9660a837176eee74a28b229bc626816589102" +dependencies = [ + "cfg-if", + "futures-util", + "wasm-bindgen", +] + +[[package]] +name = "libc" +version = "0.2.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[package]] +name = "matrixmultiply" +version = "0.3.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a06de3016e9fae57a36fd14dba131fccf49f74b40b7fbdb472f96e361ec71a08" +dependencies = [ + "autocfg", + "rawpointer", +] + +[[package]] +name = "ndarray" +version = "0.17.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "520080814a7a6b4a6e9070823bb24b4531daac8c4627e08ba5de8c5ef2f2752d" +dependencies = [ + "matrixmultiply", + "num-complex", + "num-integer", + "num-traits", + "portable-atomic", + "portable-atomic-util", + "rawpointer", +] + +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-integer" +version = "0.1.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + +[[package]] +name = "numpy" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a5b15d63a5ff39e378daed0e1340d3a5964703ea9712eb09a0dc66fade996f4" +dependencies = [ + "libc", + "ndarray", + "num-complex", + "num-integer", + "num-traits", + "pyo3", + "pyo3-build-config", + "rustc-hash", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "portable-atomic" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" + +[[package]] +name = "portable-atomic-util" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a106d1259c23fac8e543272398ae0e3c0b8d33c88ed73d0cc71b0f1d902618" +dependencies = [ + "portable-atomic", +] + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "pyo3" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd274650b21d4bfc26a0a47587962c1edb425f69287324355cd040c3ea66071c" +dependencies = [ + "libc", + "once_cell", + "portable-atomic", + "pyo3-build-config", + "pyo3-ffi", + "pyo3-macros", +] + +[[package]] +name = "pyo3-build-config" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c5e2a7d2f0d013342f295c048ad19237add5154a55b1c5a254c0ec93d4109078" +dependencies = [ + "target-lexicon", +] + +[[package]] +name = "pyo3-ffi" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca85c467da1bbc8d866eea5deff9cf29ea5f7785054a17da36e65bda9c05845b" +dependencies = [ + "libc", + "pyo3-build-config", +] + +[[package]] +name = "pyo3-macros" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ac53762fd065daa3194dd09337a38bd793a188100fd1a9304c4ab312d901771" +dependencies = [ + "proc-macro2", + "pyo3-macros-backend", + "quote", + "syn", +] + +[[package]] +name = "pyo3-macros-backend" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ca3a1557399783172dc5bf39cfca835157732532cba56b71d2292161e53b362" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "quote" +version = "1.0.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfbc457d0c7a0759a614551b11a6409e5951f6c7537be1f1b7682b9ae9230368" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "rawpointer" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" + +[[package]] +name = "rayon" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + +[[package]] +name = "rustc-hash" +version = "2.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94300abf3f1ae2e2b8ffb7b58043de3d399c73fa6f4b73826402a5c457614dbe" + +[[package]] +name = "rustversion" +version = "1.0.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "symbolic_regression" +version = "0.12.2" +source = "git+https://github.com/astroautomata/symbolic_regression.rs#60a4be904265ecd945be6a02e9f0a4b689801eb3" +dependencies = [ + "dynamic_expressions", + "fastrand", + "ndarray", + "num-traits", + "rayon", + "web-time", +] + +[[package]] +name = "symbolic_regression_rs_py" +version = "0.1.0" +dependencies = [ + "dynamic_expressions", + "ndarray", + "num-traits", + "numpy", + "pyo3", + "rayon", + "symbolic_regression", +] + +[[package]] +name = "syn" +version = "2.0.118" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b9ae57f904213ebb649ce6895b8a66c66f0203b9319718f69a5612a065b1422" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "target-lexicon" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "wasm-bindgen" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b067c0c11094aef6b7a801c1e34a26affafdf3d051dba08456b868789aaf9a4" +dependencies = [ + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "167ce5e579f6bcf889c4f7175a8a5a585de84e8ff93976ce393efa5f2837aab1" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3997c7839262f4ef12cf90b818d6340c18e80f263f1a94bf157d0ec4420380e" +dependencies = [ + "bumpalo", + "proc-macro2", + "quote", + "syn", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.126" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc1b4cb0cc549fcf58d7dfc081778139b3d283a081644e833e84682ad71cea24" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "web-time" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a6580f308b1fad9207618087a65c04e7a10bc77e02c8e84e9b00dd4b12fa0bb" +dependencies = [ + "js-sys", + "wasm-bindgen", +] diff --git a/pysr_rust_backend/Cargo.toml b/pysr_rust_backend/Cargo.toml new file mode 100644 index 000000000..87a262de0 --- /dev/null +++ b/pysr_rust_backend/Cargo.toml @@ -0,0 +1,22 @@ +[package] +name = "symbolic_regression_rs_py" +version = "0.1.0" +edition = "2021" +license = "Apache-2.0" +description = "Experimental Rust backend package for PySR." +homepage = "https://github.com/astroautomata/PySR" +repository = "https://github.com/astroautomata/PySR" +publish = false + +[lib] +name = "symbolic_regression_rs" +crate-type = ["cdylib"] + +[dependencies] +dynamic_expressions = { version = "0.10.1", git = "https://github.com/astroautomata/symbolic_regression.rs", package = "dynamic_expressions", default-features = false } +ndarray = "0.17" +num-traits = "0.2" +numpy = "0.29" +pyo3 = { version = "0.29", features = ["extension-module"] } +rayon = "1" +symbolic_regression = { version = "0.12.2", git = "https://github.com/astroautomata/symbolic_regression.rs", package = "symbolic_regression", default-features = false } diff --git a/pysr_rust_backend/README.md b/pysr_rust_backend/README.md new file mode 100644 index 000000000..c20853822 --- /dev/null +++ b/pysr_rust_backend/README.md @@ -0,0 +1,98 @@ +# pysr-rust-backend + +Experimental Rust backend package for PySR. + +The distribution name is `pysr-rust-backend`; it installs the +`symbolic_regression_rs` import module consumed by PySR's optional Rust backend. +It deliberately exposes a small wire contract first: Python arrays and plain +dictionaries in, plain dictionaries out. + +> This package is experimental and should be treated as a PySR backend boundary, +> not a stable standalone Python API, while the Rust engine API is still +> stabilizing. + +This draft places the Python/Rust connector package inside the PySR repository +for review. The Rust search engine crates are still consumed from +`astroautomata/symbolic_regression.rs`. + +## Install For Development + +From the workspace root: + +```bash +uv pip install --python /path/to/python -e pysr_rust_backend +``` + +Or, from this directory: + +```bash +python -m pip install -e . +``` + +## Python API + +```python +import numpy as np +import symbolic_regression_rs + +X = np.linspace(-1, 1, 64, dtype=np.float32).reshape(-1, 1) +y = X[:, 0] + +out = symbolic_regression_rs.search( + X, + y, + options={ + "seed": 0, + "niterations": 1, + "populations": 1, + "population_size": 16, + "ncycles_per_iteration": 20, + "deterministic": True, + "progress": False, + }, + operators={2: ["+", "sub", "*"]}, + variable_names=["x0"], +) + +print(out["hall_of_fame"]) +``` + +The return value is a dictionary: + +```python +{ + "backend_version": "0.1.0", + "hall_of_fame": [ + {"complexity": 1, "loss": 0.0, "equation": "x0"}, + ], +} +``` + +## Current Limits + +- Dense NumPy arrays only. +- `float32` and `float64` only. +- Single-output regression only. +- Builtin Rust operators only, with operator arities validated from the + `operators={arity: [names]}` mapping. +- Unknown or unsupported option keys raise `ValueError`. +- No sample weights, non-MSE losses, per-variable/per-operator complexities, + constraints, nested constraints, Python callbacks, custom losses, custom + operators, GPU, warm start, or richer PySR output yet. + +## Test + +Install the package first, then run: + +```bash +python -m unittest discover pysr_rust_backend/tests +cargo check --manifest-path pysr_rust_backend/Cargo.toml +cargo test --manifest-path pysr_rust_backend/Cargo.toml +``` + +## Publish + +The repository includes a manual `Publish Python wrapper` GitHub Actions +workflow. It builds Linux, Windows, and macOS wheels for Python 3.9 through +3.13 plus a source distribution from a selected ref, uploads the artifacts, and +publishes to TestPyPI or PyPI through trusted publishing. diff --git a/pysr_rust_backend/pyproject.toml b/pysr_rust_backend/pyproject.toml new file mode 100644 index 000000000..1f71fb5b2 --- /dev/null +++ b/pysr_rust_backend/pyproject.toml @@ -0,0 +1,35 @@ +[build-system] +requires = ["maturin>=1.8,<2"] +build-backend = "maturin" + +[project] +name = "pysr-rust-backend" +version = "0.1.0" +description = "Experimental Rust backend package for PySR." +readme = "README.md" +requires-python = ">=3.9" +dependencies = ["numpy>=1.23"] +license = { text = "Apache-2.0" } +keywords = ["symbolic-regression", "genetic-programming", "pysr", "rust-backend"] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: Apache Software License", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", + "Programming Language :: Rust", + "Topic :: Scientific/Engineering :: Artificial Intelligence", +] + +[project.urls] +Homepage = "https://github.com/astroautomata/PySR" +Repository = "https://github.com/astroautomata/PySR" +Issues = "https://github.com/astroautomata/PySR/issues" + +[tool.maturin] +module-name = "symbolic_regression_rs" +features = ["pyo3/extension-module"] diff --git a/pysr_rust_backend/src/lib.rs b/pysr_rust_backend/src/lib.rs new file mode 100644 index 000000000..6375e45cf --- /dev/null +++ b/pysr_rust_backend/src/lib.rs @@ -0,0 +1,538 @@ +use std::ops::AddAssign; + +use dynamic_expressions::operator_enum::presets::{BuiltinOpsF32, BuiltinOpsF64}; +use dynamic_expressions::strings::{string_tree, StringTreeOptions}; +use dynamic_expressions::Operators; +use ndarray::{Array1, Array2}; +use num_traits::{Float, FromPrimitive, ToPrimitive}; +use numpy::{Element, PyReadonlyArray1, PyReadonlyArray2}; +use pyo3::exceptions::{PyRuntimeError, PyTypeError, PyValueError}; +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyList, PyModule}; +use symbolic_regression::{equation_search, Dataset, EarlyStop, Options, OutputStyle}; + +const MAX_ARITY: usize = 3; +const SUPPORTED_OPTIONS: &[&str] = &[ + "seed", + "niterations", + "populations", + "population_size", + "ncycles_per_iteration", + "batch_size", + "complexity_of_constants", + "complexity_of_variables", + "maxsize", + "maxdepth", + "warmup_maxsize_by", + "parsimony", + "adaptive_parsimony_scaling", + "crossover_probability", + "perturbation_factor", + "probability_negate_constant", + "tournament_selection_n", + "tournament_selection_p", + "alpha", + "optimizer_nrestarts", + "optimizer_probability", + "optimizer_iterations", + "optimizer_f_calls_limit", + "fraction_replaced", + "fraction_replaced_hof", + "topn", + "print_precision", + "max_evals", + "timeout_in_seconds", + "use_frequency", + "use_frequency_in_tournament", + "skip_mutation_failures", + "annealing", + "should_optimize_constants", + "migration", + "hof_migration", + "progress", + "should_simplify", + "batching", + "deterministic", + "parallelism", + "early_stop_condition", + "mutation_weights", +]; +const SUPPORTED_MUTATION_WEIGHTS: &[&str] = &[ + "mutate_constant", + "mutate_operator", + "mutate_feature", + "swap_operands", + "rotate_tree", + "add_node", + "insert_node", + "delete_node", + "simplify", + "randomize", + "do_nothing", + "optimize", + "form_connection", + "break_connection", +]; + +#[pyfunction] +#[pyo3(signature = (x, y, *, options, operators, variable_names))] +fn search( + py: Python<'_>, + x: &Bound<'_, PyAny>, + y: &Bound<'_, PyAny>, + options: &Bound<'_, PyDict>, + operators: &Bound<'_, PyDict>, + variable_names: Vec, +) -> PyResult> { + if let (Ok(x), Ok(y)) = ( + x.extract::>(), + y.extract::>(), + ) { + return run_search::(py, x, y, options, operators, variable_names); + } + + if let (Ok(x), Ok(y)) = ( + x.extract::>(), + y.extract::>(), + ) { + return run_search::(py, x, y, options, operators, variable_names); + } + + Err(PyTypeError::new_err( + "X and y must be NumPy arrays with matching float32 or float64 dtypes", + )) +} + +fn run_search( + py: Python<'_>, + x: PyReadonlyArray2<'_, T>, + y: PyReadonlyArray1<'_, T>, + options_dict: &Bound<'_, PyDict>, + operators_dict: &Bound<'_, PyDict>, + variable_names: Vec, +) -> PyResult> +where + T: Float + + AddAssign + + FromPrimitive + + ToPrimitive + + std::fmt::Display + + Send + + Sync + + Element + + 'static, + Ops: dynamic_expressions::OperatorSet + Send + Sync, +{ + let dataset = build_dataset(x, y, variable_names)?; + let operator_names = extract_operator_names(operators_dict)?; + let operators = + Operators::::from_names::(operator_names.iter().map(String::as_str)) + .map_err(|err| PyValueError::new_err(format!("invalid operator selection: {err}")))?; + + let mut options = Options:: { + operators, + ..Default::default() + }; + apply_options(&mut options, options_dict)?; + let parallelism = extract_parallelism(options_dict)?; + + let result = if parallelism.as_deref() == Some("serial") { + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(1) + .build() + .map_err(|err| { + PyRuntimeError::new_err(format!("failed to create serial thread pool: {err}")) + })?; + py.detach(|| pool.install(|| equation_search::(&dataset, &options))) + } else { + py.detach(|| equation_search::(&dataset, &options)) + }; + + let rows = PyList::empty(py); + for member in result.hall_of_fame.pareto_front() { + let row = PyDict::new(py); + row.set_item("complexity", member.complexity)?; + row.set_item( + "loss", + member + .loss + .to_f64() + .ok_or_else(|| PyRuntimeError::new_err("failed to convert loss to float"))?, + )?; + row.set_item( + "equation", + string_tree( + &member.expr, + StringTreeOptions { + variable_names: Some(dataset.variable_names.as_slice()), + pretty: false, + print_precision: Some(options.print_precision), + }, + ), + )?; + rows.append(row)?; + } + + let out = PyDict::new(py); + out.set_item("backend_version", env!("CARGO_PKG_VERSION"))?; + out.set_item("hall_of_fame", rows)?; + Ok(out.into_any().unbind()) +} + +fn build_dataset( + x: PyReadonlyArray2<'_, T>, + y: PyReadonlyArray1<'_, T>, + variable_names: Vec, +) -> PyResult> +where + T: Float + FromPrimitive + Element, +{ + let x_view = x.as_array(); + let y_view = y.as_array(); + let shape = x_view.shape(); + let n_rows = shape[0]; + let n_features = shape[1]; + + if y_view.len() != n_rows { + return Err(PyValueError::new_err(format!( + "X has {n_rows} rows but y has {} elements", + y_view.len() + ))); + } + if !variable_names.is_empty() && variable_names.len() != n_features { + return Err(PyValueError::new_err(format!( + "expected {n_features} variable names, got {}", + variable_names.len() + ))); + } + + let mut rust_x = Array2::::zeros((n_features, n_rows)); + for row in 0..n_rows { + for feature in 0..n_features { + rust_x[(feature, row)] = x_view[(row, feature)]; + } + } + let rust_y = Array1::from_iter(y_view.iter().copied()); + + Ok(Dataset::with_weights_and_names( + rust_x, + rust_y, + None, + variable_names, + )) +} + +fn extract_operator_names(operators: &Bound<'_, PyDict>) -> PyResult> { + let mut out = Vec::new(); + for (arity_obj, item) in operators.iter() { + let arity: usize = arity_obj.extract()?; + if arity > MAX_ARITY { + return Err(PyValueError::new_err(format!( + "operators only support arities up to {MAX_ARITY}, got {arity}" + ))); + } + let names: Vec = item.extract()?; + for name in &names { + let expected_arity = operator_arity(name).ok_or_else(|| { + PyValueError::new_err(format!("unsupported builtin operator {name:?}")) + })?; + if arity != expected_arity { + return Err(PyValueError::new_err(format!( + "operator {name:?} has arity {expected_arity} but was provided under arity {arity}" + ))); + } + } + out.extend(names); + } + if out.is_empty() { + return Err(PyValueError::new_err("at least one operator is required")); + } + Ok(out) +} + +fn operator_arity(name: &str) -> Option { + match name { + "sin" | "cos" | "tan" | "exp" | "log" | "sqrt" | "abs" | "abs2" => Some(1), + "+" | "*" | "/" | "sub" => Some(2), + _ => None, + } +} + +fn extract_parallelism(options: &Bound<'_, PyDict>) -> PyResult> { + let Some(value) = options.get_item("parallelism")? else { + return Ok(None); + }; + let parallelism: String = value.extract()?; + match parallelism.as_str() { + "serial" | "multithreading" => Ok(Some(parallelism)), + other => Err(PyValueError::new_err(format!( + "parallelism must be 'serial' or 'multithreading', got {other:?}" + ))), + } +} + +fn apply_options( + options: &mut Options, + source: &Bound<'_, PyDict>, +) -> PyResult<()> { + options.output_style = OutputStyle::Plain; + validate_known_keys(source, SUPPORTED_OPTIONS, "options")?; + + assign_u64(source, "seed", &mut options.seed)?; + assign_usize(source, "niterations", &mut options.niterations)?; + assign_usize(source, "populations", &mut options.populations)?; + assign_usize(source, "population_size", &mut options.population_size)?; + assign_usize( + source, + "ncycles_per_iteration", + &mut options.ncycles_per_iteration, + )?; + assign_usize(source, "batch_size", &mut options.batch_size)?; + assign_u16( + source, + "complexity_of_constants", + &mut options.complexity_of_constants, + )?; + assign_u16( + source, + "complexity_of_variables", + &mut options.complexity_of_variables, + )?; + assign_usize(source, "maxsize", &mut options.maxsize)?; + assign_usize(source, "maxdepth", &mut options.maxdepth)?; + assign_f32(source, "warmup_maxsize_by", &mut options.warmup_maxsize_by)?; + assign_f64(source, "parsimony", &mut options.parsimony)?; + assign_f64( + source, + "adaptive_parsimony_scaling", + &mut options.adaptive_parsimony_scaling, + )?; + assign_f64( + source, + "crossover_probability", + &mut options.crossover_probability, + )?; + assign_f64( + source, + "perturbation_factor", + &mut options.perturbation_factor, + )?; + assign_f64( + source, + "probability_negate_constant", + &mut options.probability_negate_constant, + )?; + assign_usize( + source, + "tournament_selection_n", + &mut options.tournament_selection_n, + )?; + assign_f32( + source, + "tournament_selection_p", + &mut options.tournament_selection_p, + )?; + assign_f64(source, "alpha", &mut options.alpha)?; + assign_usize( + source, + "optimizer_nrestarts", + &mut options.optimizer_nrestarts, + )?; + assign_f64( + source, + "optimizer_probability", + &mut options.optimizer_probability, + )?; + assign_usize( + source, + "optimizer_iterations", + &mut options.optimizer_iterations, + )?; + assign_usize( + source, + "optimizer_f_calls_limit", + &mut options.optimizer_f_calls_limit, + )?; + assign_f64(source, "fraction_replaced", &mut options.fraction_replaced)?; + assign_f64( + source, + "fraction_replaced_hof", + &mut options.fraction_replaced_hof, + )?; + assign_usize(source, "topn", &mut options.topn)?; + assign_usize(source, "print_precision", &mut options.print_precision)?; + assign_u64(source, "max_evals", &mut options.max_evals)?; + assign_f64( + source, + "timeout_in_seconds", + &mut options.timeout_in_seconds, + )?; + + assign_bool(source, "use_frequency", &mut options.use_frequency)?; + assign_bool( + source, + "use_frequency_in_tournament", + &mut options.use_frequency_in_tournament, + )?; + assign_bool( + source, + "skip_mutation_failures", + &mut options.skip_mutation_failures, + )?; + assign_bool(source, "annealing", &mut options.annealing)?; + assign_bool( + source, + "should_optimize_constants", + &mut options.should_optimize_constants, + )?; + assign_bool(source, "migration", &mut options.migration)?; + assign_bool(source, "hof_migration", &mut options.hof_migration)?; + assign_bool(source, "progress", &mut options.progress)?; + assign_bool(source, "should_simplify", &mut options.should_simplify)?; + assign_bool(source, "batching", &mut options.batching)?; + assign_bool(source, "deterministic", &mut options.deterministic)?; + + if let Some(value) = source.get_item("early_stop_condition")? { + let threshold_f64: f64 = value.extract()?; + let threshold = T::from_f64(threshold_f64).ok_or_else(|| { + PyValueError::new_err("early_stop_condition must be representable as backend float") + })?; + options.early_stop_condition = Some(EarlyStop::below(threshold)); + } + + if let Some(value) = source.get_item("mutation_weights")? { + let weights = value.cast::()?; + validate_known_keys(weights, SUPPORTED_MUTATION_WEIGHTS, "mutation_weights")?; + assign_f64( + weights, + "mutate_constant", + &mut options.mutation_weights.mutate_constant, + )?; + assign_f64( + weights, + "mutate_operator", + &mut options.mutation_weights.mutate_operator, + )?; + assign_f64( + weights, + "mutate_feature", + &mut options.mutation_weights.mutate_feature, + )?; + assign_f64( + weights, + "swap_operands", + &mut options.mutation_weights.swap_operands, + )?; + assign_f64( + weights, + "rotate_tree", + &mut options.mutation_weights.rotate_tree, + )?; + assign_f64(weights, "add_node", &mut options.mutation_weights.add_node)?; + assign_f64( + weights, + "insert_node", + &mut options.mutation_weights.insert_node, + )?; + assign_f64( + weights, + "delete_node", + &mut options.mutation_weights.delete_node, + )?; + assign_f64(weights, "simplify", &mut options.mutation_weights.simplify)?; + assign_f64( + weights, + "randomize", + &mut options.mutation_weights.randomize, + )?; + assign_f64( + weights, + "do_nothing", + &mut options.mutation_weights.do_nothing, + )?; + assign_f64(weights, "optimize", &mut options.mutation_weights.optimize)?; + assign_f64( + weights, + "form_connection", + &mut options.mutation_weights.form_connection, + )?; + assign_f64( + weights, + "break_connection", + &mut options.mutation_weights.break_connection, + )?; + } + + Ok(()) +} + +fn validate_known_keys( + source: &Bound<'_, PyDict>, + allowed: &[&str], + context: &str, +) -> PyResult<()> { + let mut unknown = Vec::new(); + for key in source.keys() { + let key: String = key.extract()?; + if !allowed.contains(&key.as_str()) { + unknown.push(key); + } + } + if unknown.is_empty() { + return Ok(()); + } + unknown.sort(); + Err(PyValueError::new_err(format!( + "unsupported or unknown {context}: {}", + unknown.join(", ") + ))) +} + +fn assign_usize(source: &Bound<'_, PyDict>, name: &str, target: &mut usize) -> PyResult<()> { + if let Some(value) = source.get_item(name)? { + *target = value.extract()?; + } + Ok(()) +} + +fn assign_u16(source: &Bound<'_, PyDict>, name: &str, target: &mut u16) -> PyResult<()> { + if let Some(value) = source.get_item(name)? { + *target = value.extract()?; + } + Ok(()) +} + +fn assign_u64(source: &Bound<'_, PyDict>, name: &str, target: &mut u64) -> PyResult<()> { + if let Some(value) = source.get_item(name)? { + *target = value.extract()?; + } + Ok(()) +} + +fn assign_f32(source: &Bound<'_, PyDict>, name: &str, target: &mut f32) -> PyResult<()> { + if let Some(value) = source.get_item(name)? { + *target = value.extract()?; + } + Ok(()) +} + +fn assign_f64(source: &Bound<'_, PyDict>, name: &str, target: &mut f64) -> PyResult<()> { + if let Some(value) = source.get_item(name)? { + *target = value.extract()?; + } + Ok(()) +} + +fn assign_bool(source: &Bound<'_, PyDict>, name: &str, target: &mut bool) -> PyResult<()> { + if let Some(value) = source.get_item(name)? { + *target = value.extract()?; + } + Ok(()) +} + +#[pymodule] +fn symbolic_regression_rs(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add("__version__", env!("CARGO_PKG_VERSION"))?; + m.add_function(wrap_pyfunction!(search, m)?)?; + Ok(()) +} diff --git a/pysr_rust_backend/tests/test_search.py b/pysr_rust_backend/tests/test_search.py new file mode 100644 index 000000000..a0ee62cb4 --- /dev/null +++ b/pysr_rust_backend/tests/test_search.py @@ -0,0 +1,130 @@ +import unittest + +import numpy as np +import symbolic_regression_rs + + +class TestSearch(unittest.TestCase): + def _search(self, X, y, *, options=None, operators=None, variable_names=None): + search_options = { + "seed": 0, + "niterations": 1, + "populations": 1, + "population_size": 16, + "ncycles_per_iteration": 20, + "maxsize": 10, + "maxdepth": 10, + "deterministic": True, + "progress": False, + } + if options is not None: + search_options.update(options) + return symbolic_regression_rs.search( + X, + y, + options=search_options, + operators=operators or {2: ["+", "sub", "*"]}, + variable_names=variable_names or ["x0"], + ) + + def test_search_float32_returns_hall_of_fame(self): + X = np.linspace(-1, 1, 32, dtype=np.float32).reshape(-1, 1) + y = X[:, 0] + + out = self._search(X, y) + + self.assertEqual(out["backend_version"], symbolic_regression_rs.__version__) + self.assertTrue(out["hall_of_fame"]) + best_loss = min(row["loss"] for row in out["hall_of_fame"]) + self.assertLessEqual(best_loss, 1e-6) + for row in out["hall_of_fame"]: + self.assertIsInstance(row["complexity"], int) + self.assertIsInstance(row["loss"], float) + self.assertIsInstance(row["equation"], str) + + def test_search_float64_returns_hall_of_fame(self): + X = np.linspace(-1, 1, 24, dtype=np.float64).reshape(-1, 1) + y = X[:, 0] + + out = self._search(X, y) + + self.assertTrue(out["hall_of_fame"]) + + def test_search_rejects_mismatched_rows(self): + X = np.ones((4, 1), dtype=np.float32) + y = np.ones(3, dtype=np.float32) + + with self.assertRaisesRegex(ValueError, "X has 4 rows"): + self._search(X, y) + + def test_search_rejects_bad_variable_names_length(self): + X = np.ones((4, 2), dtype=np.float32) + y = np.ones(4, dtype=np.float32) + + with self.assertRaisesRegex(ValueError, "expected 2 variable names"): + self._search(X, y, variable_names=["x0"]) + + def test_search_rejects_unknown_operator(self): + X = np.ones((4, 1), dtype=np.float32) + y = np.ones(4, dtype=np.float32) + + with self.assertRaisesRegex(ValueError, "unsupported builtin operator"): + self._search(X, y, operators={1: ["not_an_operator"]}) + + def test_search_rejects_operator_under_wrong_arity(self): + X = np.ones((4, 1), dtype=np.float32) + y = np.ones(4, dtype=np.float32) + + with self.assertRaisesRegex(ValueError, "has arity 2"): + self._search(X, y, operators={1: ["+"]}) + + def test_search_accepts_pysr_division_token(self): + X = np.linspace(1, 2, 16, dtype=np.float32).reshape(-1, 1) + y = X[:, 0] + + out = self._search(X, y, operators={2: ["/"]}) + + self.assertTrue(out["hall_of_fame"]) + + def test_search_accepts_serial_parallelism(self): + X = np.linspace(-1, 1, 24, dtype=np.float32).reshape(-1, 1) + y = X[:, 0] + + out = self._search(X, y, options={"parallelism": "serial"}) + + self.assertTrue(out["hall_of_fame"]) + + def test_search_rejects_unknown_parallelism(self): + X = np.ones((4, 1), dtype=np.float32) + y = np.ones(4, dtype=np.float32) + + with self.assertRaisesRegex(ValueError, "parallelism"): + self._search(X, y, options={"parallelism": "multiprocessing"}) + + def test_search_rejects_unknown_option(self): + X = np.ones((4, 1), dtype=np.float32) + y = np.ones(4, dtype=np.float32) + + with self.assertRaisesRegex(ValueError, "unsupported or unknown options"): + self._search(X, y, options={"not_a_real_option": 1}) + + def test_search_rejects_unknown_mutation_weight(self): + X = np.ones((4, 1), dtype=np.float32) + y = np.ones(4, dtype=np.float32) + + with self.assertRaisesRegex( + ValueError, "unsupported or unknown mutation_weights" + ): + self._search(X, y, options={"mutation_weights": {"not_a_real_weight": 1.0}}) + + def test_search_accepts_nondefault_seed(self): + X = np.linspace(-1, 1, 24, dtype=np.float32).reshape(-1, 1) + y = X[:, 0] + + out = self._search(X, y, options={"seed": 123}) + + self.assertTrue(out["hall_of_fame"]) + + +if __name__ == "__main__": + unittest.main()