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/README.md b/README.md index 62f575626..aad6ae367 100644 --- a/README.md +++ b/README.md @@ -73,7 +73,27 @@ 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. + +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.", ), ]