API Reference

PySRRegressor has many options for controlling a symbolic regression search. Let's look at them below.

PySRRegressor Parameters

The Algorithm

Creating the Search Space

• binary_operators

  List of strings for binary operators used in the search. See the operators page for more details.

  Default: ["+", "-", "*", "/"]

• unary_operators

  Operators which only take a single scalar as input. For example, "cos" or "exp".

  Default: None

• expression_spec

  The type of expression to search for. By default, this is just ExpressionSpec(). You can also use TemplateExpressionSpec(...) which allows you to specify a custom template for the expressions.

  Default: ExpressionSpec()

• maxsize

  Max complexity of an equation.

  Default: 30

• maxdepth

  Max depth of an equation. You can use both maxsize and maxdepth. maxdepth is by default not used.

  Default: None

Setting the Search Size

• niterations

  Number of iterations of the algorithm to run. The best equations are printed and migrate between populations at the end of each iteration.

  Default: 100

• populations

  Number of populations running.

  Default: 31

• population_size

  Number of individuals in each population.

  Default: 27

• ncycles_per_iteration

  Number of total mutations to run, per 10 samples of the population, per iteration.

  Default: 380

The Objective

• elementwise_loss

  String of Julia code specifying an elementwise loss function. Can either be a loss from LossFunctions.jl, or your own loss written as a function. Examples of custom written losses include: myloss(x, y) = abs(x-y) for non-weighted, or myloss(x, y, w) = w*abs(x-y) for weighted. The included losses include: Regression: LPDistLoss{P}(), L1DistLoss(), L2DistLoss() (mean square), LogitDistLoss(), HuberLoss(d), L1EpsilonInsLoss(ϵ), L2EpsilonInsLoss(ϵ), PeriodicLoss(c), QuantileLoss(τ). Classification: ZeroOneLoss(), PerceptronLoss(), L1HingeLoss(), SmoothedL1HingeLoss(γ), ModifiedHuberLoss(), L2MarginLoss(), ExpLoss(), SigmoidLoss(), DWDMarginLoss(q).

  Default: "L2DistLoss()"

• loss_function

  Alternatively, you can specify the full objective function as a snippet of Julia code, including any sort of custom evaluation (including symbolic manipulations beforehand), and any sort of loss function or regularizations. The default loss_function used in SymbolicRegression.jl is roughly equal to:

  function eval_loss(tree, dataset::Dataset{T,L}, options)::L where {T,L}
      prediction, flag = eval_tree_array(tree, dataset.X, options)
      if !flag
          return L(Inf)
      end
      return sum((prediction .- dataset.y) .^ 2) / dataset.n
  end
  

  where the example elementwise loss is mean-squared error. You may pass a function with the same arguments as this (note that the name of the function doesn't matter). Here, both prediction and dataset.y are 1D arrays of length dataset.n. If using batching, then you should add an idx argument to the function, which is nothing for non-batched, and a 1D array of indices for batched.

  Default: None

• model_selection

  Model selection criterion when selecting a final expression from the list of best expression at each complexity. Can be 'accuracy', 'best', or 'score'. 'accuracy' selects the candidate model with the lowest loss (highest accuracy). 'score' selects the candidate model with the highest score. Score is defined as the negated derivative of the log-loss with respect to complexity - if an expression has a much better loss at a slightly higher complexity, it is preferred. 'best' selects the candidate model with the highest score among expressions with a loss better than at least 1.5x the most accurate model.

  Default: 'best'

• dimensional_constraint_penalty

  Additive penalty for if dimensional analysis of an expression fails. By default, this is 1000.0.

• dimensionless_constants_only

  Whether to only search for dimensionless constants, if using units.

  Default: False

Working with Complexities

• parsimony

  Multiplicative factor for how much to punish complexity.

  Default: 0.0

• constraints

  Dictionary of int (unary) or 2-tuples (binary), this enforces maxsize constraints on the individual arguments of operators. E.g., 'pow': (-1, 1) says that power laws can have any complexity left argument, but only 1 complexity in the right argument. Use this to force more interpretable solutions.

  Default: None

• nested_constraints

  Specifies how many times a combination of operators can be nested. For example, {"sin": {"cos": 0}}, "cos": {"cos": 2}} specifies that cos may never appear within a sin, but sin can be nested with itself an unlimited number of times. The second term specifies that cos can be nested up to 2 times within a cos, so that cos(cos(cos(x))) is allowed (as well as any combination of + or - within it), but cos(cos(cos(cos(x)))) is not allowed. When an operator is not specified, it is assumed that it can be nested an unlimited number of times. This requires that there is no operator which is used both in the unary operators and the binary operators (e.g., - could be both subtract, and negation). For binary operators, you only need to provide a single number: both arguments are treated the same way, and the max of each argument is constrained.

  Default: None

• complexity_of_operators

  If you would like to use a complexity other than 1 for an operator, specify the complexity here. For example, {"sin": 2, "+": 1} would give a complexity of 2 for each use of the sin operator, and a complexity of 1 for each use of the + operator (which is the default). You may specify real numbers for a complexity, and the total complexity of a tree will be rounded to the nearest integer after computing.

  Default: None

• complexity_of_constants

  Complexity of constants.

  Default: 1

• complexity_of_variables

  Global complexity of variables. To set different complexities for different variables, pass a list of complexities to the fit method with keyword complexity_of_variables. You cannot use both.

  Default: 1

• complexity_mapping

  Alternatively, you can pass a function (a string of Julia code) that takes the expression as input and returns the complexity. Make sure that this operates on AbstractExpression (and unpacks to AbstractExpressionNode), and returns an integer.

  Default: None

• warmup_maxsize_by

  Whether to slowly increase max size from a small number up to the maxsize (if greater than 0). If greater than 0, says the fraction of training time at which the current maxsize will reach the user-passed maxsize.

  Default: 0.0

• use_frequency

  Whether to measure the frequency of complexities, and use that instead of parsimony to explore equation space. Will naturally find equations of all complexities.

  Default: True

• use_frequency_in_tournament

  Whether to use the frequency mentioned above in the tournament, rather than just the simulated annealing.

  Default: True

• adaptive_parsimony_scaling

  If the adaptive parsimony strategy (use_frequency and use_frequency_in_tournament), this is how much to (exponentially) weight the contribution. If you find that the search is only optimizing the most complex expressions while the simpler expressions remain stagnant, you should increase this value.

  Default: 1040.0

• should_simplify

  Whether to use algebraic simplification in the search. Note that only a few simple rules are implemented.

  Default: True

Mutations

• weight_add_node

  Relative likelihood for mutation to add a node.

  Default: 2.47

• weight_insert_node

  Relative likelihood for mutation to insert a node.

  Default: 0.0112

• weight_delete_node

  Relative likelihood for mutation to delete a node.

  Default: 0.870

• weight_do_nothing

  Relative likelihood for mutation to leave the individual.

  Default: 0.273

• weight_mutate_constant

  Relative likelihood for mutation to change the constant slightly in a random direction.

  Default: 0.0346

• weight_mutate_operator

  Relative likelihood for mutation to swap an operator.

  Default: 0.293

• weight_swap_operands

  Relative likehood for swapping operands in binary operators.

  Default: 0.198

• weight_rotate_tree

  How often to perform a tree rotation at a random node.

  Default: 4.26

• weight_randomize

  Relative likelihood for mutation to completely delete and then randomly generate the equation

  Default: 0.000502

• weight_simplify

  Relative likelihood for mutation to simplify constant parts by evaluation

  Default: 0.00209

• weight_optimize

  Constant optimization can also be performed as a mutation, in addition to the normal strategy controlled by optimize_probability which happens every iteration. Using it as a mutation is useful if you want to use a large ncycles_periteration, and may not optimize very often.

  Default: 0.0

• crossover_probability

  Absolute probability of crossover-type genetic operation, instead of a mutation.

  Default: 0.0259

• annealing

  Whether to use annealing.

  Default: False

• alpha

  Initial temperature for simulated annealing (requires annealing to be True).

  Default: 3.17

• perturbation_factor

  Constants are perturbed by a max factor of (perturbation_factor*T + 1). Either multiplied by this or divided by this.

  Default: 0.129

• probability_negate_constant

  Probability of negating a constant in the equation when mutating it.

  Default: 0.00743

• skip_mutation_failures

  Whether to skip mutation and crossover failures, rather than simply re-sampling the current member.

  Default: True

Tournament Selection

• tournament_selection_n

  Number of expressions to consider in each tournament.

  Default: 15

• tournament_selection_p

  Probability of selecting the best expression in each tournament. The probability will decay as p*(1-p)^n for other expressions, sorted by loss.

  Default: 0.982

Constant Optimization

• optimizer_algorithm

  Optimization scheme to use for optimizing constants. Can currently be NelderMead or BFGS.

  Default: "BFGS"

• optimizer_nrestarts

  Number of time to restart the constants optimization process with different initial conditions.

  Default: 2

• optimizer_f_calls_limit

  How many function calls to allow during optimization.

  Default: 10_000

• optimize_probability

  Probability of optimizing the constants during a single iteration of the evolutionary algorithm.

  Default: 0.14

• optimizer_iterations

  Number of iterations that the constants optimizer can take.

  Default: 8

• should_optimize_constants

  Whether to numerically optimize constants (Nelder-Mead/Newton) at the end of each iteration.

  Default: True

Migration between Populations

• fraction_replaced

  How much of population to replace with migrating equations from other populations.

  Default: 0.00036

• fraction_replaced_hof

  How much of population to replace with migrating equations from hall of fame.

  Default: 0.0614

• migration

  Whether to migrate.

  Default: True

• hof_migration

  Whether to have the hall of fame migrate.

  Default: True

• topn

  How many top individuals migrate from each population.

  Default: 12

Data Preprocessing

• denoise

  Whether to use a Gaussian Process to denoise the data before inputting to PySR. Can help PySR fit noisy data.

  Default: False

• select_k_features

  Whether to run feature selection in Python using random forests, before passing to the symbolic regression code. None means no feature selection; an int means select that many features.

  Default: None

Stopping Criteria

• max_evals

  Limits the total number of evaluations of expressions to this number.

  Default: None

• timeout_in_seconds

  Make the search return early once this many seconds have passed.

  Default: None

• early_stop_condition

  Stop the search early if this loss is reached. You may also pass a string containing a Julia function which takes a loss and complexity as input, for example: "f(loss, complexity) = (loss < 0.1) && (complexity < 10)".

  Default: None

Performance and Parallelization

• parallelism

  Parallelism to use for the search. Can be "serial", "multithreading", or "multiprocessing".

  Default: "multithreading"

• procs

  Number of processes to use for parallelism. If None, defaults to cpu_count().

  Default: None

• cluster_manager

  For distributed computing, this sets the job queue system. Set to one of "slurm", "pbs", "lsf", "sge", "qrsh", "scyld", or "htc". If set to one of these, PySR will run in distributed mode, and use procs to figure out how many processes to launch.

  Default: None

• heap_size_hint_in_bytes

  For multiprocessing, this sets the --heap-size-hint parameter for new Julia processes. This can be configured when using multi-node distributed compute, to give a hint to each process about how much memory they can use before aggressive garbage collection.

• batching

  Whether to compare population members on small batches during evolution. Still uses full dataset for comparing against hall of fame.

  Default: False

• batch_size

  The amount of data to use if doing batching.

  Default: 50

• precision

  What precision to use for the data. By default this is 32 (float32), but you can select 64 or 16 as well, giving you 64 or 16 bits of floating point precision, respectively. If you pass complex data, the corresponding complex precision will be used (i.e., 64 for complex128, 32 for complex64).

  Default: 32

• fast_cycle

  Batch over population subsamples. This is a slightly different algorithm than regularized evolution, but does cycles 15% faster. May be algorithmically less efficient.

  Default: False

• turbo

  (Experimental) Whether to use LoopVectorization.jl to speed up the search evaluation. Certain operators may not be supported. Does not support 16-bit precision floats.

  Default: False

• bumper

  (Experimental) Whether to use Bumper.jl to speed up the search evaluation. Does not support 16-bit precision floats.

  Default: False

• autodiff_backend

  Which backend to use for automatic differentiation during constant optimization. Currently only "Zygote" is supported. The default, None, uses forward-mode or finite difference.

  Default: None

Determinism

• random_state

  Pass an int for reproducible results across multiple function calls. See :term:Glossary <random_state>.

  Default: None

• deterministic

  Make a PySR search give the same result every run. To use this, you must turn off parallelism (with parallelism="serial"), and set random_state to a fixed seed.

  Default: False

• warm_start

  Tells fit to continue from where the last call to fit finished. If false, each call to fit will be fresh, overwriting previous results.

  Default: False

Monitoring

• verbosity

  What verbosity level to use. 0 means minimal print statements.

  Default: 1

• update_verbosity

  What verbosity level to use for package updates. Will take value of verbosity if not given.

  Default: None

• print_precision

  How many significant digits to print for floats.

  Default: 5

• progress

  Whether to use a progress bar instead of printing to stdout.

  Default: True

• logger_spec

  Logger specification for the Julia backend. See, for example, TensorBoardLoggerSpec.

  Default: None

• input_stream

  The stream to read user input from. By default, this is "stdin". If you encounter issues with reading from stdin, like a hang, you can simply pass "devnull" to this argument. You can also reference an arbitrary Julia object in the Main namespace.

  Default: "stdin"

Environment

• temp_equation_file

  Whether to put the hall of fame file in the temp directory. Deletion is then controlled with the delete_tempfiles parameter.

  Default: False

• tempdir

  directory for the temporary files.

  Default: None

• delete_tempfiles

  Whether to delete the temporary files after finishing.

  Default: True

• update

  Whether to automatically update Julia packages when fit is called. You should make sure that PySR is up-to-date itself first, as the packaged Julia packages may not necessarily include all updated dependencies.

  Default: False

Exporting the Results

• output_directory

  The base directory to save output files to. Files will be saved in a subdirectory according to the run ID. Will be set to outputs/ if not provided.

  Default: None

• run_id

  A unique identifier for the run. Will be generated using the current date and time if not provided.

  Default: None

• output_jax_format

  Whether to create a 'jax_format' column in the output, containing jax-callable functions and the default parameters in a jax array.

  Default: False

• output_torch_format

  Whether to create a 'torch_format' column in the output, containing a torch module with trainable parameters.

  Default: False

• extra_sympy_mappings

  Provides mappings between custom binary_operators or unary_operators defined in julia strings, to those same operators defined in sympy. E.G if unary_operators=["inv(x)=1/x"], then for the fitted model to be export to sympy, extra_sympy_mappings would be {"inv": lambda x: 1/x}.

  Default: None

• extra_torch_mappings

  The same as extra_jax_mappings but for model export to pytorch. Note that the dictionary keys should be callable pytorch expressions. For example: extra_torch_mappings={sympy.sin: torch.sin}.

  Default: None

• extra_jax_mappings

  Similar to extra_sympy_mappings but for model export to jax. The dictionary maps sympy functions to jax functions. For example: extra_jax_mappings={sympy.sin: "jnp.sin"} maps the sympy.sin function to the equivalent jax expression jnp.sin.

  Default: None

PySRRegressor Functions

fit(X, y, *, Xresampled=None, weights=None, variable_names=None, complexity_of_variables=None, X_units=None, y_units=None, category=None)

Search for equations to fit the dataset and store them in self.equations_.

Parameters:

Name	Type	Description	Default
X	ndarray | DataFrame	Training data of shape (n_samples, n_features).	required
y	ndarray | DataFrame	Target values of shape (n_samples,) or (n_samples, n_targets). Will be cast to X's dtype if necessary.	required
Xresampled	ndarray | DataFrame	Resampled training data, of shape (n_resampled, n_features), to generate a denoised data on. This will be used as the training data, rather than X.	None
weights	ndarray | DataFrame	Weight array of the same shape as y. Each element is how to weight the mean-square-error loss for that particular element of y. Alternatively, if a custom loss was set, it will can be used in arbitrary ways.	None
variable_names	list[str]	A list of names for the variables, rather than "x0", "x1", etc. If X is a pandas dataframe, the column names will be used instead of variable_names. Cannot contain spaces or special characters. Avoid variable names which are also function names in sympy, such as "N".	None
X_units	list[str]	A list of units for each variable in X. Each unit should be a string representing a Julia expression. See DynamicQuantities.jl https://symbolicml.org/DynamicQuantities.jl/dev/units/ for more information.	None
y_units	str | list[str]	Similar to X_units, but as a unit for the target variable, y. If y is a matrix, a list of units should be passed. If X_units is given but y_units is not, then y_units will be arbitrary.	None
category	list[int]	If expression_spec is a ParametricExpressionSpec, then this argument should be a list of integers representing the category of each sample.	None

Returns:

Name	Type	Description
self	object	Fitted estimator.

Source code in pysr/sr.py

def fit(
    self,
    X,
    y,
    *,
    Xresampled=None,
    weights=None,
    variable_names: ArrayLike[str] | None = None,
    complexity_of_variables: int | float | list[int | float] | None = None,
    X_units: ArrayLike[str] | None = None,
    y_units: str | ArrayLike[str] | None = None,
    category: ndarray | None = None,
) -> "PySRRegressor":
    """
    Search for equations to fit the dataset and store them in `self.equations_`.

    Parameters
    ----------
    X : ndarray | pandas.DataFrame
        Training data of shape (n_samples, n_features).
    y : ndarray | pandas.DataFrame
        Target values of shape (n_samples,) or (n_samples, n_targets).
        Will be cast to X's dtype if necessary.
    Xresampled : ndarray | pandas.DataFrame
        Resampled training data, of shape (n_resampled, n_features),
        to generate a denoised data on. This
        will be used as the training data, rather than `X`.
    weights : ndarray | pandas.DataFrame
        Weight array of the same shape as `y`.
        Each element is how to weight the mean-square-error loss
        for that particular element of `y`. Alternatively,
        if a custom `loss` was set, it will can be used
        in arbitrary ways.
    variable_names : list[str]
        A list of names for the variables, rather than "x0", "x1", etc.
        If `X` is a pandas dataframe, the column names will be used
        instead of `variable_names`. Cannot contain spaces or special
        characters. Avoid variable names which are also
        function names in `sympy`, such as "N".
    X_units : list[str]
        A list of units for each variable in `X`. Each unit should be
        a string representing a Julia expression. See DynamicQuantities.jl
        https://symbolicml.org/DynamicQuantities.jl/dev/units/ for more
        information.
    y_units : str | list[str]
        Similar to `X_units`, but as a unit for the target variable, `y`.
        If `y` is a matrix, a list of units should be passed. If `X_units`
        is given but `y_units` is not, then `y_units` will be arbitrary.
    category : list[int]
        If `expression_spec` is a `ParametricExpressionSpec`, then this
        argument should be a list of integers representing the category
        of each sample.

    Returns
    -------
    self : object
        Fitted estimator.
    """
    # Init attributes that are not specified in BaseEstimator
    if self.warm_start and hasattr(self, "julia_state_stream_"):
        pass
    else:
        if hasattr(self, "julia_state_stream_"):
            warnings.warn(
                "The discovered expressions are being reset. "
                "Please set `warm_start=True` if you wish to continue "
                "to start a search where you left off.",
            )

        self.equations_ = None
        self.nout_ = 1
        self.selection_mask_ = None
        self.julia_state_stream_ = None
        self.julia_options_stream_ = None
        self.complexity_of_variables_ = None
        self.X_units_ = None
        self.y_units_ = None

    self._setup_equation_file()
    self._clear_equation_file_contents()

    runtime_params = self._validate_and_modify_params()

    if category is not None:
        assert Xresampled is None

    if isinstance(self.expression_spec, ParametricExpressionSpec):
        assert category is not None

    # TODO: Put `category` here
    (
        X,
        y,
        Xresampled,
        weights,
        variable_names,
        complexity_of_variables,
        X_units,
        y_units,
    ) = self._validate_and_set_fit_params(
        X,
        y,
        Xresampled,
        weights,
        variable_names,
        complexity_of_variables,
        X_units,
        y_units,
    )

    if X.shape[0] > 10000 and not self.batching:
        warnings.warn(
            "Note: you are running with more than 10,000 datapoints. "
            "You should consider turning on batching (https://ai.damtp.cam.ac.uk/pysr/options/#batching). "
            "You should also reconsider if you need that many datapoints. "
            "Unless you have a large amount of noise (in which case you "
            "should smooth your dataset first), generally < 10,000 datapoints "
            "is enough to find a functional form with symbolic regression. "
            "More datapoints will lower the search speed."
        )

    random_state = check_random_state(self.random_state)  # For np random
    seed = cast(int, random_state.randint(0, 2**31 - 1))  # For julia random

    # Pre transformations (feature selection and denoising)
    X, y, variable_names, complexity_of_variables, X_units, y_units = (
        self._pre_transform_training_data(
            X,
            y,
            Xresampled,
            variable_names,
            complexity_of_variables,
            X_units,
            y_units,
            random_state,
        )
    )

    # Warn about large feature counts (still warn if feature count is large
    # after running feature selection)
    if self.n_features_in_ >= 10:
        warnings.warn(
            "Note: you are running with 10 features or more. "
            "Genetic algorithms like used in PySR scale poorly with large numbers of features. "
            "You should run PySR for more `niterations` to ensure it can find "
            "the correct variables, and consider using a larger `maxsize`."
        )

    # Assertion checks
    use_custom_variable_names = variable_names is not None
    # TODO: this is always true.

    _check_assertions(
        X,
        use_custom_variable_names,
        variable_names,
        complexity_of_variables,
        weights,
        y,
        X_units,
        y_units,
    )

    # Initially, just save model parameters, so that
    # it can be loaded from an early exit:
    if not self.temp_equation_file:
        self._checkpoint()

    # Perform the search:
    self._run(X, y, runtime_params, weights=weights, seed=seed, category=category)

    # Then, after fit, we save again, so the pickle file contains
    # the equations:
    if not self.temp_equation_file:
        self._checkpoint()

    return self

predict(X, index=None, *, category=None)

Predict y from input X using the equation chosen by model_selection.

You may see what equation is used by printing this object. X should have the same columns as the training data.

Parameters:

Name	Type	Description	Default
X	ndarray | DataFrame	Training data of shape (n_samples, n_features).	required
index	int | list[int]	If you want to compute the output of an expression using a particular row of self.equations_, you may specify the index here. For multiple output equations, you must pass a list of indices in the same order.	None
category	ndarray | None	If expression_spec is a ParametricExpressionSpec, then this argument should be a list of integers representing the category of each sample in X.	None

Returns:

Name	Type	Description
y_predicted	ndarray of shape (n_samples, nout_)	Values predicted by substituting X into the fitted symbolic regression model.

Raises:

Type	Description
ValueError	Raises if the best_equation cannot be evaluated.

Source code in pysr/sr.py

def predict(
    self,
    X,
    index: int | list[int] | None = None,
    *,
    category: ndarray | None = None,
) -> ndarray:
    """
    Predict y from input X using the equation chosen by `model_selection`.

    You may see what equation is used by printing this object. X should
    have the same columns as the training data.

    Parameters
    ----------
    X : ndarray | pandas.DataFrame
        Training data of shape `(n_samples, n_features)`.
    index : int | list[int]
        If you want to compute the output of an expression using a
        particular row of `self.equations_`, you may specify the index here.
        For multiple output equations, you must pass a list of indices
        in the same order.
    category : ndarray | None
        If `expression_spec` is a `ParametricExpressionSpec`, then this
        argument should be a list of integers representing the category
        of each sample in `X`.

    Returns
    -------
    y_predicted : ndarray of shape (n_samples, nout_)
        Values predicted by substituting `X` into the fitted symbolic
        regression model.

    Raises
    ------
    ValueError
        Raises if the `best_equation` cannot be evaluated.
    """
    check_is_fitted(
        self, attributes=["selection_mask_", "feature_names_in_", "nout_"]
    )
    best_equation = self.get_best(index=index)

    # When X is an numpy array or a pandas dataframe with a RangeIndex,
    # the self.feature_names_in_ generated during fit, for the same X,
    # will cause a warning to be thrown during _validate_data.
    # To avoid this, convert X to a dataframe, apply the selection mask,
    # and then set the column/feature_names of X to be equal to those
    # generated during fit.
    if not isinstance(X, pd.DataFrame):
        X = check_array(X)
        X = pd.DataFrame(X)
    if isinstance(X.columns, pd.RangeIndex):
        if self.selection_mask_ is not None:
            # RangeIndex enforces column order allowing columns to
            # be correctly filtered with self.selection_mask_
            X = X[X.columns[self.selection_mask_]]
        X.columns = self.feature_names_in_
    # Without feature information, CallableEquation/lambda_format equations
    # require that the column order of X matches that of the X used during
    # the fitting process. _validate_data removes this feature information
    # when it converts the dataframe to an np array. Thus, to ensure feature
    # order is preserved after conversion, the dataframe columns must be
    # reordered/reindexed to match those of the transformed (denoised and
    # feature selected) X in fit.
    X = X.reindex(columns=self.feature_names_in_)
    X = self._validate_data_X(X)
    if self.expression_spec_.evaluates_in_julia:
        # Julia wants the right dtype
        X = X.astype(self._get_precision_mapped_dtype(X))

    if category is not None:
        offset_for_julia_indexing = 1
        args: tuple = (
            jl_array((category + offset_for_julia_indexing).astype(np.int64)),
        )
    else:
        args = ()

    try:
        if isinstance(best_equation, list):
            assert self.nout_ > 1
            return np.stack(
                [
                    cast(ndarray, eq["lambda_format"](X, *args))
                    for eq in best_equation
                ],
                axis=1,
            )
        else:
            return cast(ndarray, best_equation["lambda_format"](X, *args))
    except Exception as error:
        raise ValueError(
            "Failed to evaluate the expression. "
            "If you are using a custom operator, make sure to define it in `extra_sympy_mappings`, "
            "e.g., `model.set_params(extra_sympy_mappings={'inv': lambda x: 1/x})`, where "
            "`lambda x: 1/x` is a valid SymPy function defining the operator. "
            "You can then run `model.refresh()` to re-load the expressions."
        ) from error

from_file(equation_file=None, *, run_directory, binary_operators=None, unary_operators=None, n_features_in=None, feature_names_in=None, selection_mask=None, nout=1, **pysr_kwargs) classmethod

Create a model from a saved model checkpoint or equation file.

Parameters:

Name	Type	Description	Default
run_directory	str	The directory containing outputs from a previous run. This is of the form [output_directory]/[run_id]. Default is None.	required
binary_operators	list[str]	The same binary operators used when creating the model. Not needed if loading from a pickle file.	None
unary_operators	list[str]	The same unary operators used when creating the model. Not needed if loading from a pickle file.	None
n_features_in	int	Number of features passed to the model. Not needed if loading from a pickle file.	None
feature_names_in	list[str]	Names of the features passed to the model. Not needed if loading from a pickle file.	None
selection_mask	NDArray[bool_]	If using select_k_features, you must pass model.selection_mask_ here. Not needed if loading from a pickle file.	None
nout	int	Number of outputs of the model. Not needed if loading from a pickle file. Default is 1.	1
**pysr_kwargs	dict	Any other keyword arguments to initialize the PySRRegressor object. These will overwrite those stored in the pickle file. Not needed if loading from a pickle file.	{}

Returns:

Name	Type	Description
model	PySRRegressor	The model with fitted equations.

Source code in pysr/sr.py

@classmethod
def from_file(
    cls,
    equation_file: None = None,  # Deprecated
    *,
    run_directory: PathLike,
    binary_operators: list[str] | None = None,
    unary_operators: list[str] | None = None,
    n_features_in: int | None = None,
    feature_names_in: ArrayLike[str] | None = None,
    selection_mask: NDArray[np.bool_] | None = None,
    nout: int = 1,
    **pysr_kwargs,
) -> "PySRRegressor":
    """
    Create a model from a saved model checkpoint or equation file.

    Parameters
    ----------
    run_directory : str
        The directory containing outputs from a previous run.
        This is of the form `[output_directory]/[run_id]`.
        Default is `None`.
    binary_operators : list[str]
        The same binary operators used when creating the model.
        Not needed if loading from a pickle file.
    unary_operators : list[str]
        The same unary operators used when creating the model.
        Not needed if loading from a pickle file.
    n_features_in : int
        Number of features passed to the model.
        Not needed if loading from a pickle file.
    feature_names_in : list[str]
        Names of the features passed to the model.
        Not needed if loading from a pickle file.
    selection_mask : NDArray[np.bool_]
        If using `select_k_features`, you must pass `model.selection_mask_` here.
        Not needed if loading from a pickle file.
    nout : int
        Number of outputs of the model.
        Not needed if loading from a pickle file.
        Default is `1`.
    **pysr_kwargs : dict
        Any other keyword arguments to initialize the PySRRegressor object.
        These will overwrite those stored in the pickle file.
        Not needed if loading from a pickle file.

    Returns
    -------
    model : PySRRegressor
        The model with fitted equations.
    """
    if equation_file is not None:
        raise ValueError(
            "Passing `equation_file` is deprecated and no longer compatible with "
            "the most recent versions of PySR's backend. Please pass `run_directory` "
            "instead, which contains all checkpoint files."
        )

    pkl_filename = Path(run_directory) / "checkpoint.pkl"
    if pkl_filename.exists():
        print(f"Attempting to load model from {pkl_filename}...")
        assert binary_operators is None
        assert unary_operators is None
        assert n_features_in is None
        with open(pkl_filename, "rb") as f:
            model = cast("PySRRegressor", pkl.load(f))

        # Update any parameters if necessary, such as
        # extra_sympy_mappings:
        model.set_params(**pysr_kwargs)

        if "equations_" not in model.__dict__ or model.equations_ is None:
            model.refresh()

        return model
    else:
        print(
            f"Checkpoint file {pkl_filename} does not exist. "
            "Attempting to recreate model from scratch..."
        )
        csv_filename = Path(run_directory) / "hall_of_fame.csv"
        csv_filename_bak = Path(run_directory) / "hall_of_fame.csv.bak"
        if not csv_filename.exists() and not csv_filename_bak.exists():
            raise FileNotFoundError(
                f"Hall of fame file `{csv_filename}` or `{csv_filename_bak}` does not exist. "
                "Please pass a `run_directory` containing a valid checkpoint file."
            )
        assert binary_operators is not None or unary_operators is not None
        assert n_features_in is not None
        model = cls(
            binary_operators=binary_operators,
            unary_operators=unary_operators,
            **pysr_kwargs,
        )
        model.nout_ = nout
        model.n_features_in_ = n_features_in

        if feature_names_in is None:
            model.feature_names_in_ = np.array(
                [f"x{i}" for i in range(n_features_in)]
            )
            model.display_feature_names_in_ = np.array(
                [f"x{_subscriptify(i)}" for i in range(n_features_in)]
            )
        else:
            assert len(feature_names_in) == n_features_in
            model.feature_names_in_ = feature_names_in
            model.display_feature_names_in_ = feature_names_in

        if selection_mask is None:
            model.selection_mask_ = np.ones(n_features_in, dtype=np.bool_)
        else:
            model.selection_mask_ = selection_mask

        model.refresh(run_directory=run_directory)

        return model

sympy(index=None)

Return sympy representation of the equation(s) chosen by model_selection.

Parameters:

Name	Type	Description	Default
index	int | list[int]	If you wish to select a particular equation from self.equations_, give the index number here. This overrides the model_selection parameter. If there are multiple output features, then pass a list of indices with the order the same as the output feature.	None

Returns:

Name	Type	Description
best_equation	str, list[str] of length nout_	SymPy representation of the best equation.

Source code in pysr/sr.py

def sympy(self, index: int | list[int] | None = None):
    """
    Return sympy representation of the equation(s) chosen by `model_selection`.

    Parameters
    ----------
    index : int | list[int]
        If you wish to select a particular equation from
        `self.equations_`, give the index number here. This overrides
        the `model_selection` parameter. If there are multiple output
        features, then pass a list of indices with the order the same
        as the output feature.

    Returns
    -------
    best_equation : str, list[str] of length nout_
        SymPy representation of the best equation.
    """
    if not self.expression_spec_.supports_sympy:
        raise ValueError(
            f"`expression_spec={self.expression_spec_}` does not support sympy export."
        )
    self.refresh()
    best_equation = self.get_best(index=index)
    if isinstance(best_equation, list):
        assert self.nout_ > 1
        return [eq["sympy_format"] for eq in best_equation]
    else:
        return best_equation["sympy_format"]

latex(index=None, precision=3)

Return latex representation of the equation(s) chosen by model_selection.

Parameters:

Name	Type	Description	Default
index	int | list[int]	If you wish to select a particular equation from self.equations_, give the index number here. This overrides the model_selection parameter. If there are multiple output features, then pass a list of indices with the order the same as the output feature.	None
precision	int	The number of significant figures shown in the LaTeX representation. Default is 3.	3

Returns:

Name	Type	Description
best_equation	str or list[str] of length nout_	LaTeX expression of the best equation.

Source code in pysr/sr.py

def latex(
    self, index: int | list[int] | None = None, precision: int = 3
) -> str | list[str]:
    """
    Return latex representation of the equation(s) chosen by `model_selection`.

    Parameters
    ----------
    index : int | list[int]
        If you wish to select a particular equation from
        `self.equations_`, give the index number here. This overrides
        the `model_selection` parameter. If there are multiple output
        features, then pass a list of indices with the order the same
        as the output feature.
    precision : int
        The number of significant figures shown in the LaTeX
        representation.
        Default is `3`.

    Returns
    -------
    best_equation : str or list[str] of length nout_
        LaTeX expression of the best equation.
    """
    if not self.expression_spec_.supports_latex:
        raise ValueError(
            f"`expression_spec={self.expression_spec_}` does not support latex export."
        )
    self.refresh()
    sympy_representation = self.sympy(index=index)
    if self.nout_ > 1:
        output = []
        for s in sympy_representation:
            latex = sympy2latex(s, prec=precision)
            output.append(latex)
        return output
    return sympy2latex(sympy_representation, prec=precision)

pytorch(index=None)

Return pytorch representation of the equation(s) chosen by model_selection.

Each equation (multiple given if there are multiple outputs) is a PyTorch module containing the parameters as trainable attributes. You can use the module like any other PyTorch module: module(X), where X is a tensor with the same column ordering as trained with.

Parameters:

Name	Type	Description	Default
index	int | list[int]	If you wish to select a particular equation from self.equations_, give the index number here. This overrides the model_selection parameter. If there are multiple output features, then pass a list of indices with the order the same as the output feature.	None

Returns:

Name	Type	Description
best_equation	Module	PyTorch module representing the expression.

Source code in pysr/sr.py

def pytorch(self, index=None):
    """
    Return pytorch representation of the equation(s) chosen by `model_selection`.

    Each equation (multiple given if there are multiple outputs) is a PyTorch module
    containing the parameters as trainable attributes. You can use the module like
    any other PyTorch module: `module(X)`, where `X` is a tensor with the same
    column ordering as trained with.

    Parameters
    ----------
    index : int | list[int]
        If you wish to select a particular equation from
        `self.equations_`, give the index number here. This overrides
        the `model_selection` parameter. If there are multiple output
        features, then pass a list of indices with the order the same
        as the output feature.

    Returns
    -------
    best_equation : torch.nn.Module
        PyTorch module representing the expression.
    """
    if not self.expression_spec_.supports_torch:
        raise ValueError(
            f"`expression_spec={self.expression_spec_}` does not support torch export."
        )
    self.set_params(output_torch_format=True)
    self.refresh()
    best_equation = self.get_best(index=index)
    if isinstance(best_equation, list):
        return [eq["torch_format"] for eq in best_equation]
    else:
        return best_equation["torch_format"]

jax(index=None)

Return jax representation of the equation(s) chosen by model_selection.

Each equation (multiple given if there are multiple outputs) is a dictionary containing {"callable": func, "parameters": params}. To call func, pass func(X, params). This function is differentiable using jax.grad.

Parameters:

Name	Type	Description	Default
index	int | list[int]	If you wish to select a particular equation from self.equations_, give the index number here. This overrides the model_selection parameter. If there are multiple output features, then pass a list of indices with the order the same as the output feature.	None

Returns:

Name	Type	Description
best_equation	dict[str, Any]	Dictionary of callable jax function in "callable" key, and jax array of parameters as "parameters" key.

Source code in pysr/sr.py

def jax(self, index=None):
    """
    Return jax representation of the equation(s) chosen by `model_selection`.

    Each equation (multiple given if there are multiple outputs) is a dictionary
    containing {"callable": func, "parameters": params}. To call `func`, pass
    func(X, params). This function is differentiable using `jax.grad`.

    Parameters
    ----------
    index : int | list[int]
        If you wish to select a particular equation from
        `self.equations_`, give the index number here. This overrides
        the `model_selection` parameter. If there are multiple output
        features, then pass a list of indices with the order the same
        as the output feature.

    Returns
    -------
    best_equation : dict[str, Any]
        Dictionary of callable jax function in "callable" key,
        and jax array of parameters as "parameters" key.
    """
    if not self.expression_spec_.supports_jax:
        raise ValueError(
            f"`expression_spec={self.expression_spec_}` does not support jax export."
        )
    self.set_params(output_jax_format=True)
    self.refresh()
    best_equation = self.get_best(index=index)
    if isinstance(best_equation, list):
        assert self.nout_ > 1
        return [eq["jax_format"] for eq in best_equation]
    else:
        return best_equation["jax_format"]

latex_table(indices=None, precision=3, columns=['equation', 'complexity', 'loss', 'score'])

Create a LaTeX/booktabs table for all, or some, of the equations.

Parameters:

Name	Type	Description	Default
indices	list[int] | list[list[int]]	If you wish to select a particular subset of equations from self.equations_, give the row numbers here. By default, all equations will be used. If there are multiple output features, then pass a list of lists.	None
precision	int	The number of significant figures shown in the LaTeX representations. Default is 3.	3
columns	list[str]	Which columns to include in the table. Default is ["equation", "complexity", "loss", "score"].	['equation', 'complexity', 'loss', 'score']

Returns:

Name	Type	Description
latex_table_str	str	A string that will render a table in LaTeX of the equations.

Source code in pysr/sr.py

def latex_table(
    self,
    indices: list[int] | None = None,
    precision: int = 3,
    columns: list[str] = ["equation", "complexity", "loss", "score"],
) -> str:
    """Create a LaTeX/booktabs table for all, or some, of the equations.

    Parameters
    ----------
    indices : list[int] | list[list[int]]
        If you wish to select a particular subset of equations from
        `self.equations_`, give the row numbers here. By default,
        all equations will be used. If there are multiple output
        features, then pass a list of lists.
    precision : int
        The number of significant figures shown in the LaTeX
        representations.
        Default is `3`.
    columns : list[str]
        Which columns to include in the table.
        Default is `["equation", "complexity", "loss", "score"]`.

    Returns
    -------
    latex_table_str : str
        A string that will render a table in LaTeX of the equations.
    """
    if not self.expression_spec_.supports_latex:
        raise ValueError(
            f"`expression_spec={self.expression_spec_}` does not support latex export."
        )
    self.refresh()

    if isinstance(self.equations_, list):
        if indices is not None:
            assert isinstance(indices, list)
            assert isinstance(indices[0], list)
            assert len(indices) == self.nout_

        table_string = sympy2multilatextable(
            self.equations_, indices=indices, precision=precision, columns=columns
        )
    elif isinstance(self.equations_, pd.DataFrame):
        if indices is not None:
            assert isinstance(indices, list)
            assert isinstance(indices[0], int)

        table_string = sympy2latextable(
            self.equations_, indices=indices, precision=precision, columns=columns
        )
    else:
        raise ValueError(
            "Invalid type for equations_ to pass to `latex_table`. "
            "Expected a DataFrame or a list of DataFrames."
        )

    return with_preamble(table_string)

refresh(run_directory=None)

Update self.equations_ with any new options passed.

For example, updating extra_sympy_mappings will require a .refresh() to update the equations.

Parameters:

Name	Type	Description	Default
checkpoint_file	str or Path	Path to checkpoint hall of fame file to be loaded. The default will use the set equation_file_.	required

Source code in pysr/sr.py

def refresh(self, run_directory: PathLike | None = None) -> None:
    """
    Update self.equations_ with any new options passed.

    For example, updating `extra_sympy_mappings`
    will require a `.refresh()` to update the equations.

    Parameters
    ----------
    checkpoint_file : str or Path
        Path to checkpoint hall of fame file to be loaded.
        The default will use the set `equation_file_`.
    """
    if run_directory is not None:
        self.output_directory_ = str(Path(run_directory).parent)
        self.run_id_ = Path(run_directory).name
        self._clear_equation_file_contents()
    check_is_fitted(self, attributes=["run_id_", "output_directory_"])
    self.equations_ = self.get_hof()

Expression Specifications

ExpressionSpec

The default expression specification, with no special behavior.

TemplateExpressionSpec(function_symbols, combine, num_features=None)

Spec for templated expressions.

This class allows you to specify how multiple sub-expressions should be combined in a structured way, with constraints on which variables each sub-expression can use. Pass this to PySRRegressor with the expression_spec argument when you are using the TemplateExpression expression type.

Parameters:

Name	Type	Description	Default
function_symbols	list[str]	List of symbols representing the inner expressions (e.g., ["f", "g"]). These will be used as keys in the template structure.	required
combine	str	Julia function string that defines how the sub-expressions are combined. Takes a NamedTuple of expressions and a tuple of data vectors. For example: "((; f, g), (x1, x2, x3)) -> f(x1, x2) + g(x3)^2" would constrain f to use x1,x2 and g to use x3.	required
num_features	dict[str, int]	Dictionary mapping function symbols to the number of features each can use. For example: {"f": 2, "g": 1} means f takes 2 inputs and g takes 1. If not provided, will be inferred from the combine function.	None

Examples:

# Create template that combines f(x1, x2) and g(x3):
expression_spec = TemplateExpressionSpec(
    function_symbols=["f", "g"],
    combine="((; f, g), (x1, x2, x3)) -> sin(f(x1, x2)) + g(x3)^2",
)

# Use in PySRRegressor:
model = PySRRegressor(
    expression_spec=expression_spec
)

Source code in pysr/expression_specs.py

def __init__(
    self,
    function_symbols: list[str],
    combine: str,
    num_features: dict[str, int] | None = None,
):
    self.function_symbols = function_symbols
    self.combine = combine
    self.num_features = num_features

ParametricExpressionSpec(max_parameters)

Spec for parametric expressions that vary by category.

This class allows you to specify expressions with parameters that vary across different categories in your dataset. The expression structure remains the same, but parameters are optimized separately for each category.

Parameters:

Name	Type	Description	Default
max_parameters	int	Maximum number of parameters that can appear in the expression. Each parameter will take on different values for each category in the data.	required

Examples:

For example, if we want to allow for a model with up to 2 parameters (each category can have a different value for these parameters), we can use:

model = PySRRegressor(
    expression_spec=ParametricExpressionSpec(max_parameters=2),
    binary_operators=["+", "*"],
    unary_operators=["sin"]
)
model.fit(X, y, category=category)

Source code in pysr/expression_specs.py

def __init__(self, max_parameters: int):
    self.max_parameters = max_parameters

AbstractExpressionSpec

Abstract base class describing expression types.

This basically just holds the options for the expression type, as well as explains how to parse and evaluate them.

All expression types must implement:

1. julia_expression_type(): The actual expression type, returned as a Julia object. This will get stored as expression_type in SymbolicRegression.Options.
2. julia_expression_options(): Method to create the expression options, returned as a Julia object. These will get stored as expression_options in SymbolicRegression.Options.
3. create_exports(), which will be used to create the exports of the equations, such as the executable format, the SymPy format, etc.

It may also optionally implement:

• supports_sympy, supports_torch, supports_jax, supports_latex: Whether this expression type supports the corresponding export format.

create_exports(model, equations, search_output) abstractmethod

Create additional columns in the equations dataframe.

Source code in pysr/expression_specs.py

@abstractmethod
def create_exports(
    self,
    model: PySRRegressor,
    equations: pd.DataFrame,
    search_output,
) -> pd.DataFrame:
    """Create additional columns in the equations dataframe."""
    pass  # pragma: no cover

julia_expression_options() abstractmethod

The expression options

Source code in pysr/expression_specs.py

@abstractmethod
def julia_expression_options(self) -> AnyValue:
    """The expression options"""
    pass  # pragma: no cover

julia_expression_type() abstractmethod

The expression type

Source code in pysr/expression_specs.py

@abstractmethod
def julia_expression_type(self) -> AnyValue:
    """The expression type"""
    pass  # pragma: no cover

Logger Specifications

TensorBoardLoggerSpec(log_dir='logs/run', log_interval=1, overwrite=False) dataclass

Specification for TensorBoard logger.

Parameters:

Name	Type	Description	Default
log_dir	str	Directory where TensorBoard logs will be saved. If overwrite is False, new logs will be saved to {log_dir}_1, and so on. Default is "logs/run".	'logs/run'
log_interval	int	Interval (in steps) at which logs are written. Default is 10.	1
overwrite	bool	Whether to overwrite existing logs in the directory. Default is False.	False

AbstractLoggerSpec

Abstract base class for logger specifications.

close(logger) abstractmethod

Close the logger instance.

Source code in pysr/logger_specs.py

@abstractmethod
def close(self, logger: AnyValue) -> None:
    """Close the logger instance."""
    pass  # pragma: no cover

create_logger() abstractmethod

Create a logger instance.

Source code in pysr/logger_specs.py

@abstractmethod
def create_logger(self) -> AnyValue:
    """Create a logger instance."""
    pass  # pragma: no cover

write_hparams(logger, hparams) abstractmethod

Write hyperparameters to the logger.

Source code in pysr/logger_specs.py

@abstractmethod
def write_hparams(self, logger: AnyValue, hparams: dict[str, Any]) -> None:
    """Write hyperparameters to the logger."""
    pass  # pragma: no cover