PySR

PySR's main interface follows the scikit-learn style. You can use PySRRegressor to perform symbolic regression by defining binary and unary operators, custom loss functions, and SymPy mappings for the discovered equations.

To use custom operators defined in Julia syntax, provide them in unary_operators and map them to Python/SymPy functions using extra_sympy_mappings.

import numpy as np
from pysr import PySRRegressor

# Generate test data
X = 2 * np.random.randn(100, 5)
y = 2.5382 * np.cos(X[:, 3]) + X[:, 0] ** 2 - 0.5

# Initialize and train the model
model = PySRRegressor(
    maxsize=20,
    niterations=40,
    binary_operators=["+", "*"],
    unary_operators=[
        "cos",
        "exp",
        "sin",
        "inv(x) = 1/x",
    ],
    extra_sympy_mappings={"inv": lambda x: 1 / x},
    elementwise_loss="loss(prediction, target) = (prediction - target)^2",
)

model.fit(X, y)

# Predict using the best equation
predictions = model.predict(X)

# Print learned equations
print(model)

PySR can run across multiple nodes in an existing Slurm allocation by setting cluster_manager="slurm".

Requirements

1. Resource Allocation: Request resources using sbatch or salloc.
2. Process Count: The procs argument in PySRRegressor must equal the total number of tasks in your Slurm allocation. This is calculated as --nodes multiplied by --ntasks-per-node (or the total --ntasks).
3. Execution: Run the Python script once inside the allocation. Do not wrap the Python command in srun.

PySR will automatically launch Julia workers across the allocation using SlurmClusterManager.jl.

### Example Slurm Job Script (`pysr_job.sh`)
```bash
#!/bin/bash
#SBATCH --job-name=pysr
#SBATCH --partition=normal
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=3
#SBATCH --time=01:00:00

set -euo pipefail
python pysr_script.py

Example Python Script (pysr_script.py)

import numpy as np
from pysr import PySRRegressor

X = np.random.RandomState(0).randn(1000, 2)
y = X[:, 0] + 2 * X[:, 1]

model = PySRRegressor(
    niterations=200,
    populations=2,
    parallelism="multiprocessing",
    cluster_manager="slurm",
    procs=6,  # Must match total tasks (2 nodes * 3 tasks/node)
)
model.fit(X, y)
print(model)

Submission

sbatch pysr_job.sh

You can pass a custom objective as a snippet of Julia code to PySRRegressor(loss_function=...). This allows for symbolic manipulations or custom functional forms that do not need to be differentiable.

Important Limitations: When using a custom objective that performs symbolic manipulations (e.g., manually splitting a tree to form a rational function), standard PySR functionalities like .sympy() and .predict() will not work because the SymPy parser cannot account for your manual manipulations. You will need to handle these manually.

To implement a custom objective, the Julia function must accept (tree, dataset::Dataset{T,L}, options) and return a scalar of type L.

objective = """
function my_custom_objective(tree, dataset::Dataset{T,L}, options) where {T,L}
    # Require root node to be binary, so we can split it,
    # otherwise return a large loss:
    tree.degree != 2 && return L(Inf)

    P = tree.l
    Q = tree.r

    # Evaluate numerator:
    P_prediction, flag = eval_tree_array(P, dataset.X, options)
    !flag && return L(Inf)

    # Evaluate denominator:
    Q_prediction, flag = eval_tree_array(Q, dataset.X, options)
    !flag && return L(Inf)

    # Impose functional form:
    prediction = P_prediction ./ Q_prediction

    diffs = prediction .- dataset.y

    return sum(diffs .^ 2) / length(diffs)
end
"""

model = PySRRegressor(
    niterations=100,
    binary_operators=["*", "+", "-"],
    loss_function=objective,
)

PySR provides several ways to persist and export your symbolic models:

Saving and Resuming

• Automatic Output: Every fit writes outputs/<run_id>/hall_of_fame.csv (continuously updated) and checkpoint.pkl.
• Reloading: Use PySRRegressor.from_file(run_directory=...) to reload a model. Note that pickles are version-locked to the PySR version used to create them.
• Warm Starting: Set warm_start=True to continue evolution from the previous call's populations within the same process. Search-space parameters (operators, maxsize, expression_spec, etc.) must remain fixed, but you can change the loss or weights between fits.

Exporting Equations

Use these methods to export the best equation (index i) to different formats:

• model.sympy(i): Export to SymPy.
• model.latex(i): Export to LaTeX.
• model.latex_table(): Export a table of equations.
• model.jax(i): Returns a dictionary {'callable', 'parameters'} for differentiable JAX use.
• model.pytorch(i): Returns a trainable PyTorch module.

Note: Custom operators require extra_sympy_mappings, extra_jax_mappings, or extra_torch_mappings to be supplied during export/reload.

PySR supports three parallelism modes. Choose the one that fits your hardware and workload:

• parallelism="multithreading" (Default): Best for laptops and single nodes. To set the thread count, you must set the environment variable PYTHON_JULIACALL_THREADS=<n> before importing pysr. Note that JULIA_NUM_THREADS is not the correct variable for use with juliacall.
• parallelism="multiprocessing": Higher startup cost but faster steady-state performance. Can span multiple nodes using cluster_manager="slurm". If using Slurm, launch the script once on one node and let it spawn workers; do not wrap the execution in srun.
• parallelism="serial": Required for full reproducibility.

Optimization Tips:

• Keep populations at approximately 2-3x the number of threads/cores to ensure workers always have work (default populations=31 is usually sufficient).
• If the coordinating thread is saturated on many-core machines, increase ncycles_per_iteration to reduce communication frequency.

For structural rules that constraints cannot express (e.g., forbidding specific variable placements or operator-specific rules), implement logic inside a custom loss_function by traversing the expression tree.

Tree Traversal API:

• node.degree: 0 (leaf), 1 (unary), 2 (binary).
• node.l, node.r: Children subtrees.
• node.constant: Boolean indicating if the leaf is a constant.
• node.val: The constant's value.
• node.feature: 1-based feature index.
• node.op: 1-based index into the operator list passed to PySRRegressor.

Recommended Pattern: Use the functional forms any(f, tree), all(f, tree), count(f, tree), etc., for efficiency. Use a 'count-then-penalize' approach: count violations and return a large finite penalty before performing the numerical evaluation.

objective = """
function constrained_loss(tree, dataset::Dataset{T,L}, options) where {T,L}
    idx_pow = 3   # position of ^ in binary_operators below (1-indexed)
    n_bad = count(tree) do node
        node.degree == 2 && node.op == idx_pow &&
            any(c -> !(c.degree == 0 && c.constant && 0 <= c.val <= 1), node.r)
    end
    n_bad > 0 && return L(10_000 * n_bad)
    prediction, valid = eval_tree_array(tree, dataset.X, options)
    !valid && return L(Inf)
    return sum(i -> abs2(prediction[i] - dataset.y[i]), eachindex(prediction)) / dataset.n
end
"""
model = PySRRegressor(binary_operators=["+", "*", "^"], loss_function=objective)