Skip to content

Dataset Submodule

SRToolkit.dataset

Datasets and benchmarks for symbolic regression.

A dataset (SR_dataset) represents a single equation with associated data, a symbol library, and evaluation constraints. A benchmark (SR_benchmark) is a collection of datasets.

Modules:

Name Description
sr_dataset

SR_dataset — wraps input data and evaluation settings for a single equation discovery problem.

sr_benchmark

SR_benchmark — manages a collection of datasets.

feynman

Feynman — 100-equation physics benchmark.

nguyen

Nguyen — 10-equation polynomial/trig benchmark.

Feynman

Feynman(
    dataset_directory: str = os.path.join(
        user_data_dir("SRToolkit"), "feynman"
    ),
)

Bases: SR_benchmark

The Feynman symbolic regression benchmark.

Contains 100 physics equations with up to 9 variables. Data is downloaded on first use from the SymbolicRegressionToolkit repository (10,000 samples per dataset instead of the original 1,000,000 from the paper).

For more information about the Feynman benchmark, see: https://doi.org/10.1126/sciadv.aay2631

Examples:

>>> benchmark = Feynman()
>>> len(benchmark.list_datasets(verbose=False))
100
>>> X, y = benchmark.resample('I.16.6', n=500, seed=0)
>>> X.shape
(500, 3)
>>> y.shape
(500,)

Parameters:

Name Type Description Default
dataset_directory str

Directory where dataset files are stored or will be downloaded to. Defaults to the platform-appropriate user data directory (e.g. ~/.local/share/SRToolkit/feynman on Linux).

join(user_data_dir('SRToolkit'), 'feynman')
Source code in SRToolkit/dataset/feynman.py
def __init__(self, dataset_directory: str = os.path.join(user_data_dir("SRToolkit"), "feynman")):
    super().__init__("feynman", dataset_directory)
    self._populate()

resample

resample(
    dataset_name: str, n: int, seed: Optional[int] = None
) -> Tuple[np.ndarray, np.ndarray]

Generate fresh data for a dataset by sampling new inputs and evaluating the ground truth.

Variable bounds are taken from _BOUNDS.

Examples:

>>> benchmark = Feynman()
>>> X, y = benchmark.resample('I.16.6', n=200, seed=42)
>>> X.shape
(200, 3)

Parameters:

Name Type Description Default
dataset_name str

Name of the dataset to resample.

required
n int

Number of new samples to generate.

required
seed Optional[int]

Random seed for reproducibility.

None

Returns:

Type Description
Tuple[ndarray, ndarray]

A tuple (X, y) of numpy arrays with shapes (n, n_vars) and (n,).

Raises:

Type Description
ValueError

If the dataset has no ground truth expression.

Source code in SRToolkit/dataset/feynman.py
def resample(self, dataset_name: str, n: int, seed: Optional[int] = None) -> Tuple[np.ndarray, np.ndarray]:
    """
    Generate fresh data for a dataset by sampling new inputs and evaluating the ground truth.

    Variable bounds are taken from ``_BOUNDS``.

    Examples:
        >>> benchmark = Feynman()
        >>> X, y = benchmark.resample('I.16.6', n=200, seed=42)
        >>> X.shape
        (200, 3)

    Args:
        dataset_name: Name of the dataset to resample.
        n: Number of new samples to generate.
        seed: Random seed for reproducibility.

    Returns:
        A tuple ``(X, y)`` of numpy arrays with shapes ``(n, n_vars)`` and ``(n,)``.

    Raises:
        ValueError: If the dataset has no ground truth expression.
    """
    info = self.datasets[dataset_name]
    if info.get("ground_truth") is None:
        raise ValueError(f"Dataset '{dataset_name}' has no ground truth expression — cannot compute y.")
    bounds = _BOUNDS[dataset_name]
    lb = np.array([b[0] for b in bounds], dtype=float)
    ub = np.array([b[1] for b in bounds], dtype=float)
    rng = np.random.default_rng(seed)
    X_new = rng.uniform(lb, ub, size=(n, len(bounds)))
    sl = SymbolLibrary.from_dict(info["symbol_library"])
    f = expr_to_executable_function(info["ground_truth"], sl)
    y_new = f(X_new, np.array([]))
    return X_new, y_new

Nguyen

Nguyen(
    dataset_directory: str = os.path.join(user_data_dir("SRToolkit"), "nguyen"),
)

Bases: SR_benchmark

The Nguyen symbolic regression benchmark.

Contains 10 expressions without constant parameters (first 4 are polynomials, first 8 use one variable, last 2 use two variables). The benchmark ships with pre-generated data.

For more information about the Nguyen benchmark, see: https://doi.org/10.1007/s10710-010-9121-2

Examples:

>>> benchmark = Nguyen()
>>> len(benchmark.list_datasets(verbose=False))
10

Parameters:

Name Type Description Default
dataset_directory str

Directory where dataset files are stored or will be downloaded to. Defaults to the platform-appropriate user data directory (e.g. ~/.local/share/SRToolkit/nguyen on Linux).

join(user_data_dir('SRToolkit'), 'nguyen')
Source code in SRToolkit/dataset/nguyen.py
def __init__(self, dataset_directory: str = os.path.join(user_data_dir("SRToolkit"), "nguyen")):
    super().__init__("Nguyen", dataset_directory)
    self._populate()

resample

resample(
    dataset_name: str, n: int, seed: Optional[int] = None
) -> Tuple[np.ndarray, np.ndarray]

Generate fresh data for a dataset by sampling new inputs and evaluating the ground truth.

Variable bounds are taken from _BOUNDS.

Examples:

>>> benchmark = Nguyen('data/nguyen/')
>>> X, y = benchmark.resample('NG-1', n=200, seed=42)
>>> X.shape
(200, 1)

Parameters:

Name Type Description Default
dataset_name str

Name of the dataset to resample.

required
n int

Number of new samples to generate.

required
seed Optional[int]

Random seed for reproducibility.

None

Returns:

Type Description
Tuple[ndarray, ndarray]

A tuple (X, y) of numpy arrays with shapes (n, n_vars) and (n,).

Raises:

Type Description
ValueError

If the dataset has no ground truth expression.

Source code in SRToolkit/dataset/nguyen.py
def resample(self, dataset_name: str, n: int, seed: Optional[int] = None) -> Tuple[np.ndarray, np.ndarray]:
    """
    Generate fresh data for a dataset by sampling new inputs and evaluating the ground truth.

    Variable bounds are taken from ``_BOUNDS``.

    Examples:
        >>> benchmark = Nguyen('data/nguyen/')
        >>> X, y = benchmark.resample('NG-1', n=200, seed=42)
        >>> X.shape
        (200, 1)

    Args:
        dataset_name: Name of the dataset to resample.
        n: Number of new samples to generate.
        seed: Random seed for reproducibility.

    Returns:
        A tuple ``(X, y)`` of numpy arrays with shapes ``(n, n_vars)`` and ``(n,)``.

    Raises:
        ValueError: If the dataset has no ground truth expression.
    """
    info = self.datasets[dataset_name]
    if info.get("ground_truth") is None:
        raise ValueError(f"Dataset '{dataset_name}' has no ground truth expression — cannot compute y.")
    bounds = _BOUNDS[dataset_name]
    lb = np.array([b[0] for b in bounds], dtype=float)
    ub = np.array([b[1] for b in bounds], dtype=float)
    rng = np.random.default_rng(seed)
    X_new = rng.uniform(lb, ub, size=(n, len(bounds)))
    sl = SymbolLibrary.from_dict(info["symbol_library"])
    f = expr_to_executable_function(info["ground_truth"], sl)
    y_new = f(X_new, np.array([]))
    return X_new, y_new

SR_benchmark

SR_benchmark(
    benchmark_name: str,
    base_dir: str,
    datasets: Optional[List[Union[SR_dataset, Tuple[str, SR_dataset]]]] = None,
    metadata: Optional[Dict[str, Any]] = None,
)

A named, persistent collection of symbolic regression datasets.

Examples:

>>> from SRToolkit.dataset import Feynman
>>> benchmark = Feynman() # Feynman is a specific instance of SR_benchmark with additional functionality
>>> len(benchmark.list_datasets(verbose=False))
100

Parameters:

Name Type Description Default
benchmark_name str

Name of this benchmark.

required
base_dir str

Directory where dataset files are stored or will be written.

required
datasets Optional[List[Union[SR_dataset, Tuple[str, SR_dataset]]]]

Initial datasets to add. Each element can be an SR_dataset instance (auto-named as "<benchmark_name>_<index>") or a (name, SR_dataset) tuple.

None
metadata Optional[Dict[str, Any]]

Optional dictionary of benchmark-level metadata (e.g. citation, description).

None

Raises:

Type Description
ValueError

If any element of datasets is not an SR_dataset or a valid (name, SR_dataset) tuple.

Source code in SRToolkit/dataset/sr_benchmark.py
def __init__(
    self,
    benchmark_name: str,
    base_dir: str,
    datasets: Optional[List[Union[SR_dataset, Tuple[str, SR_dataset]]]] = None,
    metadata: Optional[Dict[str, Any]] = None,
):
    """
    A named, persistent collection of symbolic regression datasets.

    Examples:
        >>> from SRToolkit.dataset import Feynman
        >>> benchmark = Feynman() # Feynman is a specific instance of SR_benchmark with additional functionality
        >>> len(benchmark.list_datasets(verbose=False))
        100

    Args:
        benchmark_name: Name of this benchmark.
        base_dir: Directory where dataset files are stored or will be written.
        datasets: Initial datasets to add. Each element can be an
            [SR_dataset][SRToolkit.dataset.sr_dataset.SR_dataset] instance (auto-named as
            ``"<benchmark_name>_<index>"``) or a ``(name, SR_dataset)`` tuple.
        metadata: Optional dictionary of benchmark-level metadata (e.g. citation, description).

    Raises:
        ValueError: If any element of ``datasets`` is not an
            [SR_dataset][SRToolkit.dataset.sr_dataset.SR_dataset] or a valid ``(name, SR_dataset)`` tuple.
    """
    self.benchmark_name = benchmark_name
    self.base_dir = base_dir
    self.datasets: Dict[str, Dict[str, Any]] = {}
    self.metadata = {} if metadata is None else metadata
    if datasets is not None:
        for i, dataset in enumerate(datasets):
            if isinstance(dataset, SR_dataset):
                self.add_dataset_instance(benchmark_name + "_" + str(i + 1), dataset)
            elif isinstance(dataset, tuple) and isinstance(dataset[0], str) and isinstance(dataset[1], SR_dataset):
                self.add_dataset_instance(dataset[0], dataset[1])
            else:
                raise ValueError(
                    "[SR_benchmark] Dataset inside the datasets argument must be either a tuple "
                    "(name, SR_dataset) or a SR_dataset instance."
                )

add_dataset_instance

add_dataset_instance(dataset_name: str, dataset: SR_dataset)

Adds an instance of the SR_dataset class to the benchmark.

Examples:

>>> from SRToolkit.dataset import Feynman
>>> benchmark = Feynman() # Feynman is a specific instance of SR_benchmark with additional functionality
>>> dataset = benchmark.create_dataset('I.16.6')
>>> isinstance(dataset, SR_dataset)
True
>>> bm = SR_benchmark("BM", "data/bm")
>>> bm.add_dataset_instance("I.16.6", dataset)

Parameters:

Name Type Description Default
dataset_name str

The name of the dataset.

required
dataset SR_dataset

An instance of the SR_dataset class.

required

Raises:

Type Description
Exception

If the dataset name already exists in the benchmark.

Source code in SRToolkit/dataset/sr_benchmark.py
def add_dataset_instance(self, dataset_name: str, dataset: SR_dataset):
    """
    Adds an instance of the SR_dataset class to the benchmark.

    Examples:
        >>> from SRToolkit.dataset import Feynman
        >>> benchmark = Feynman() # Feynman is a specific instance of SR_benchmark with additional functionality
        >>> dataset = benchmark.create_dataset('I.16.6')
        >>> isinstance(dataset, SR_dataset)
        True
        >>> bm = SR_benchmark("BM", "data/bm")
        >>> bm.add_dataset_instance("I.16.6", dataset)

    Args:
         dataset_name: The name of the dataset.
         dataset: An instance of the SR_dataset class.

    Raises:
        Exception: If the dataset name already exists in the benchmark.
    """
    if dataset_name in self.datasets:
        raise ValueError(f"Dataset {dataset_name} already exists in the benchmark.")
    else:
        self.datasets[dataset_name] = {}
    self.datasets[dataset_name]["sr_dataset"] = dataset
    self.datasets[dataset_name]["num_variables"] = dataset.X.shape[1]

add_dataset

add_dataset(
    dataset: Union[str, ndarray, Tuple[ndarray, ndarray]],
    symbol_library: SymbolLibrary,
    dataset_name: Optional[str] = None,
    ranking_function: str = "rmse",
    max_evaluations: int = -1,
    ground_truth: Optional[Union[List[str], Node, ndarray]] = None,
    original_equation: Optional[str] = None,
    success_threshold: Optional[float] = None,
    seed: Optional[int] = None,
    dataset_metadata: Optional[dict] = None,
    **kwargs: Unpack[EstimationSettings],
)

Adds a dataset to the benchmark.

Examples:

>>> from SRToolkit.dataset import Feynman
>>> fey_benchmark = Feynman() # Feynman is a specific instance of SR_benchmark with additional functionality
>>> benchmark = SR_benchmark("BM", "data/bm")
>>> benchmark.add_dataset(fey_benchmark.base_dir+"/I.14.3.npz", SymbolLibrary.default_symbols(3),
...       dataset_name="I.14.3", ranking_function="rmse", ground_truth = ["X_0", "*", "X_1", "*", "X_2"],
...       original_equation="U = m*g*z", max_evaluations=100000, max_expr_length=50,
...       success_threshold=1e-7, dataset_metadata={}, constant_bounds=(-5.0, 5.0),
...       seed = 42)
>>> len(benchmark.list_datasets(verbose=False))
1

Parameters:

Name Type Description Default
dataset Union[str, ndarray, Tuple[ndarray, ndarray]]

Data used in the dataset. Can be: - A string representing the path to a NumPy archive (.npz) containing the dataset. It should either the absolute path to the data, path relative to the base_dir 'base_dir'/'dataset', or empty, in that case the dataset will be loaded from 'base_dir'/'dataset_name'.npz. The .npz file must contain the features (saved in 'X') and if 'rmse' is used as the ranking function, the target (saved in 'y'). - A 2d numpy array containing the features (X). If 'rmse' is used as the ranking function, ground truth should also be provided to calculate the target (y). Once added, the data will be saved at 'base_dir'/'dataset_name'.npz. - A tuple containing the features (X) and the target (y). If 'bed' is used as the ranking function, the target will be ignored. Once added, the data will be saved at 'base_dir'/'dataset_name'.npz.

required
symbol_library SymbolLibrary

The symbol library to use.

required
dataset_name Optional[str]

The name of the dataset. If None, a name will be generated automatically as 'benchmark_name'_'index+1'.

None
ranking_function str

The ranking function used during evaluation. Can be: 'rmse', 'bed'.

'rmse'
max_evaluations int

The maximum number of expressions to evaluate. Less than 0 means no limit.

-1
ground_truth Optional[Union[List[str], Node, ndarray]]

Ground truth expression as a token list in infix notation, a Node tree, or a numpy behavior array. Required when ranking_function="bed".

None
original_equation Optional[str]

Human-readable string of the original equation.

None
success_threshold Optional[float]

Error threshold below which an expression is considered successful. If None, no threshold is applied.

None
seed Optional[int]

Random seed for reproducibility. None means no seed is set.

None
dataset_metadata Optional[dict]

Optional dictionary of dataset-level metadata (merged with benchmark metadata).

None
**kwargs Unpack[EstimationSettings]

Optional estimation settings passed to SR_evaluator. Supported keys: method, tol, gtol, max_iter, constant_bounds, initialization, max_constants, max_expr_length, num_points_sampled, bed_X, num_consts_sampled, domain_bounds.

{}

Raises:

Type Description
ValueError

When BED ranking function is used but ground truth is not provided. When dataset is given as a string (directory) that doesn't exist, is not a valid .npz file, or is a .npz file that doesn't contain one array for the BED ranking function (X) or two array for the RMSE ranking function (X, y). When the argument dataset is an array, ranking function RMSE and there is no ground truth or the expression given as the ground truth cannot be evaluated...

Source code in SRToolkit/dataset/sr_benchmark.py
def add_dataset(
    self,
    dataset: Union[str, np.ndarray, Tuple[np.ndarray, np.ndarray]],
    symbol_library: SymbolLibrary,
    dataset_name: Optional[str] = None,
    ranking_function: str = "rmse",
    max_evaluations: int = -1,
    ground_truth: Optional[Union[List[str], Node, np.ndarray]] = None,
    original_equation: Optional[str] = None,
    success_threshold: Optional[float] = None,
    seed: Optional[int] = None,
    dataset_metadata: Optional[dict] = None,
    **kwargs: Unpack[EstimationSettings],
):
    """
    Adds a dataset to the benchmark.

    Examples:
        >>> from SRToolkit.dataset import Feynman
        >>> fey_benchmark = Feynman() # Feynman is a specific instance of SR_benchmark with additional functionality
        >>> benchmark = SR_benchmark("BM", "data/bm")
        >>> benchmark.add_dataset(fey_benchmark.base_dir+"/I.14.3.npz", SymbolLibrary.default_symbols(3),
        ...       dataset_name="I.14.3", ranking_function="rmse", ground_truth = ["X_0", "*", "X_1", "*", "X_2"],
        ...       original_equation="U = m*g*z", max_evaluations=100000, max_expr_length=50,
        ...       success_threshold=1e-7, dataset_metadata={}, constant_bounds=(-5.0, 5.0),
        ...       seed = 42)
        >>> len(benchmark.list_datasets(verbose=False))
        1

    Args:
        dataset: Data used in the dataset. Can be:
            - A string representing the path to a NumPy archive (.npz) containing the dataset. It should either
            the absolute path to the data, path relative to the base_dir 'base_dir'/'dataset', or empty, in that
            case the dataset will be loaded from 'base_dir'/'dataset_name'.npz. The .npz file must contain the
            features (saved in 'X') and if 'rmse' is used as the ranking function, the target (saved in 'y').
            - A 2d numpy array containing the features (X). If 'rmse' is used as the ranking function, ground truth
            should also be provided to calculate the target (y). Once added, the data will be saved at
            'base_dir'/'dataset_name'.npz.
            - A tuple containing the features (X) and the target (y). If 'bed' is used as the ranking function,
            the target will be ignored. Once added, the data will be saved at 'base_dir'/'dataset_name'.npz.
        symbol_library: The symbol library to use.
        dataset_name: The name of the dataset. If None, a name will be generated automatically as
            'benchmark_name'_'index+1'.
        ranking_function: The ranking function used during evaluation. Can be: 'rmse', 'bed'.
        max_evaluations: The maximum number of expressions to evaluate. Less than 0 means no limit.
        ground_truth: Ground truth expression as a token list in infix notation, a
            [Node][SRToolkit.utils.expression_tree.Node] tree, or a numpy behavior array. Required when
            ``ranking_function="bed"``.
        original_equation: Human-readable string of the original equation.
        success_threshold: Error threshold below which an expression is considered successful. If ``None``,
            no threshold is applied.
        seed: Random seed for reproducibility. ``None`` means no seed is set.
        dataset_metadata: Optional dictionary of dataset-level metadata (merged with benchmark metadata).
        **kwargs: Optional estimation settings passed to
            [SR_evaluator][SRToolkit.evaluation.sr_evaluator.SR_evaluator].
            Supported keys: ``method``, ``tol``, ``gtol``, ``max_iter``, ``constant_bounds``,
            ``initialization``, ``max_constants``, ``max_expr_length``, ``num_points_sampled``,
            ``bed_X``, ``num_consts_sampled``, ``domain_bounds``.

    Raises:
        ValueError: When BED ranking function is used but ground truth is not provided. When dataset is given as
            a string (directory) that doesn't exist, is not a valid .npz file, or is a .npz file that doesn't
            contain one array for the BED ranking function (X) or two array for the RMSE ranking function (X, y).
            When the argument dataset is an array, ranking function RMSE and there is no ground truth or the
            expression given as the ground truth cannot be evaluated...
    """
    if dataset_name is None:
        dataset_name = f"{self.benchmark_name}_{len(self.datasets) + 1}"

    self.datasets[dataset_name] = {}
    self.datasets[dataset_name]["format_version"] = 1
    self.datasets[dataset_name]["dataset_name"] = dataset_name
    self.datasets[dataset_name]["symbol_library"] = symbol_library.to_dict()
    self.datasets[dataset_name]["ranking_function"] = ranking_function
    self.datasets[dataset_name]["max_evaluations"] = max_evaluations

    self.datasets[dataset_name]["success_threshold"] = success_threshold

    self.datasets[dataset_name]["seed"] = seed
    merged_metadata = copy.deepcopy(self.metadata)
    if dataset_metadata:
        merged_metadata.update(dataset_metadata)
    self.datasets[dataset_name]["dataset_metadata"] = merged_metadata

    if "bed_X" in kwargs and kwargs["bed_X"] is not None:
        kwargs["bed_X"] = kwargs["bed_X"].tolist()

    self.datasets[dataset_name]["kwargs"] = kwargs
    self.datasets[dataset_name]["original_equation"] = original_equation
    self.datasets[dataset_name]["ground_truth"] = ground_truth

    if ground_truth is None:
        if ranking_function == "bed":
            raise ValueError("[SR_benchmark.add_dataset] For 'bed' ranking, the ground truth must be provided. ")
        else:
            warnings.warn(
                "[SR_benchmark.add_dataset] 'ground_truth' argument not provided. We recommend providing it "
                "for more transparent evaluation."
            )
    else:
        if original_equation is None:
            if isinstance(ground_truth, str):
                self.datasets[dataset_name]["original_equation"] = "y = " + ground_truth
            elif isinstance(ground_truth, list):
                self.datasets[dataset_name]["original_equation"] = "y = " + "".join(ground_truth)

    if isinstance(dataset, str):
        dataset_path = None
        if os.path.exists(dataset):
            dataset_path = dataset
        elif dataset != "" and os.path.exists(f"{self.base_dir}/{dataset}"):
            dataset_path = f"{self.base_dir}/{dataset}"
        elif os.path.exists(f"{self.base_dir}/{dataset_name}.npz"):
            dataset_path = f"{self.base_dir}/{dataset_name}.npz"

        if dataset_path is None:
            error_msg = (
                f"[SR_benchmark.add_dataset] Could not find the dataset file. "
                f"Expected locations:\n"
                f"- Absolute path: '{dataset}'\n"
                f"- Relative to base_dir: '{self.base_dir}/{dataset}'\n"
                f"- NPZ with the name of the dataset in base_dir: '{self.base_dir}/{dataset_name}.npz'"
            )
            raise FileNotFoundError(error_msg)

        self.datasets[dataset_name]["dataset_path"] = dataset_path

        try:
            data = np.load(self.datasets[dataset_name]["dataset_path"], allow_pickle=False)
        except IOError as e:
            error_msg = (
                f"[SR_benchmark.add_dataset] Could not load dataset from path '{self.datasets[dataset_name]}' "
                f"using np.load. The file may be corrupt or not a valid NumPy archive (.npz). "
                f"Original error: {e}"
            )
            raise IOError(error_msg) from e

        if ranking_function == "rmse":
            if not (isinstance(data, np.lib.npyio.NpzFile) and "X" in data and "y" in data):
                error_msg = (
                    f"[SR_benchmark.add_dataset] For 'rmse' ranking, the dataset file "
                    f"('{self.datasets[dataset_name]['dataset_path']}') must be a .npz NumPy archive containing "
                    f"both 'X' (features) and 'y' (targets). It should be created via `np.savez(path, X=X, y=y)`."
                )
                raise ValueError(error_msg)

        elif ranking_function == "bed":
            if not (isinstance(data, np.lib.npyio.NpzFile) and "X" in data):
                error_msg = (
                    f"[SR_benchmark.add_dataset] For 'bed' ranking, the dataset file "
                    f"('{self.datasets[dataset_name]['dataset_path']}') must be a .npz NumPy archive "
                    f"containing 'X' (features). It should be created via `np.savez(path, X=X)`."
                )
                raise ValueError(error_msg)

        num_variables = data["X"].shape[1]

    elif isinstance(dataset, np.ndarray):
        if ranking_function == "rmse" and ground_truth is not None:
            if isinstance(ground_truth, np.ndarray):
                raise ValueError(
                    "[SR_benchmark.add_dataset] For 'rmse' ranking, the ground truth must be a string or a SRToolkit.utils.Node object. "
                )
            try:
                expr = expr_to_executable_function(ground_truth, symbol_library)
                y = expr(dataset, None)
            except Exception as e:
                raise Exception(
                    f"[SR_benchmark.add_dataset] Could not evaluate the ground truth. Original error: {e}"
                )
            if not os.path.isdir(self.base_dir):
                os.makedirs(self.base_dir)
            np.savez(f"{self.base_dir}/{dataset_name}.npz", X=dataset, y=y, allow_pickle=False)
        elif ranking_function == "rmse" and ground_truth is None:
            raise ValueError(
                "[SR_benchmark.add_dataset] For 'rmse' ranking, if the dataset argument is a numpy "
                "array, the ground truth must be provided in order for the target values to be "
                "calculated."
            )
        elif ranking_function == "bed":
            if not os.path.isdir(self.base_dir):
                os.makedirs(self.base_dir)
            np.savez(f"{self.base_dir}/{dataset_name}.npz", X=dataset, allow_pickle=False)

        self.datasets[dataset_name]["dataset_path"] = f"{self.base_dir}/{dataset_name}.npz"
        num_variables = dataset.shape[1]

    elif isinstance(dataset, tuple):
        if not isinstance(dataset[0], np.ndarray) or not isinstance(dataset[1], np.ndarray):
            raise ValueError(
                "[SR_benchmark.add_dataset] When dataset argument is provided as a tuple, both "
                "values must be a numpy array. The first array represents the features ('X'), "
                "the second array represents the targets ('y')."
            )
        if ranking_function == "bed":
            warnings.warn(
                "[SR_benchmark.add_dataset] 'bed' ranking only utilizes the array with features. "
                "Array with targets will be ignored."
            )
        if not os.path.isdir(self.base_dir):
            os.makedirs(self.base_dir)
        np.savez(f"{self.base_dir}/{dataset_name}.npz", X=dataset[0], y=dataset[1], allow_pickle=False)
        self.datasets[dataset_name]["dataset_path"] = f"{self.base_dir}/{dataset_name}.npz"
        num_variables = dataset[0].shape[1]

    else:
        raise ValueError(
            "[SR_benchmark.add_dataset] The dataset argument must be a string, a numpy array, "
            "or a tuple containing two numpy arrays."
        )

    self.datasets[dataset_name]["num_variables"] = num_variables

create_dataset

create_dataset(dataset_name: str) -> SR_dataset

Creates an instance of a dataset from the given dataset name.

Examples:

>>> from SRToolkit.dataset import Feynman
>>> benchmark = Feynman() # Feynman is a specific instance of SR_benchmark with additional functionality
>>> dataset = benchmark.create_dataset('I.16.6')
>>> dataset.X.shape
(10000, 3)

Parameters:

Name Type Description Default
dataset_name str

The name of the dataset to create.

required

Returns:

Type Description
SR_dataset

A SR_dataset instance containing the data, ground truth expression, and metadata for the given dataset.

Raises:

Type Description
ValueError

If the dataset name is not found in the available datasets.

Source code in SRToolkit/dataset/sr_benchmark.py
def create_dataset(self, dataset_name: str) -> SR_dataset:
    """
    Creates an instance of a dataset from the given dataset name.

    Examples:
        >>> from SRToolkit.dataset import Feynman
        >>> benchmark = Feynman() # Feynman is a specific instance of SR_benchmark with additional functionality
        >>> dataset = benchmark.create_dataset('I.16.6')
        >>> dataset.X.shape
        (10000, 3)

    Args:
        dataset_name: The name of the dataset to create.

    Returns:
        A SR_dataset instance containing the data, ground truth expression, and metadata for the given dataset.

    Raises:
        ValueError: If the dataset name is not found in the available datasets.
    """
    if dataset_name in self.datasets:
        if "sr_dataset" in self.datasets[dataset_name]:
            return self.datasets[dataset_name]["sr_dataset"]
        else:
            try:
                return SR_dataset.from_dict(self.datasets[dataset_name])
            except Exception as e:
                raise ValueError(
                    f"[SR_benchmark.create_dataset] Could not create SR_dataset from the given "
                    f"given dictionary. Original error: {e}"
                )

    else:
        raise ValueError(f"Dataset {dataset_name} not found")

list_datasets

list_datasets(verbose: bool = True, num_variables: int = -1) -> List[str]

Lists the available datasets.

Examples:

>>> from SRToolkit.dataset import Feynman
>>> benchmark = Feynman() # Feynman is a specific instance of SR_benchmark with additional functionality
>>> len(benchmark.list_datasets(num_variables=2, verbose=False))
15
>>> datasets_with_8_vars = benchmark.list_datasets(num_variables=8, verbose=False)
>>> datasets_with_8_vars[0]
'II.36.38'

Parameters:

Name Type Description Default
verbose bool

If True, also prints a description of each dataset.

True
num_variables int

If not -1, only return datasets with this many input variables.

-1

Returns:

Type Description
List[str]

A list of dataset names.

Source code in SRToolkit/dataset/sr_benchmark.py
def list_datasets(self, verbose: bool = True, num_variables: int = -1) -> List[str]:
    """
    Lists the available datasets.

    Examples:
        >>> from SRToolkit.dataset import Feynman
        >>> benchmark = Feynman() # Feynman is a specific instance of SR_benchmark with additional functionality
        >>> len(benchmark.list_datasets(num_variables=2, verbose=False))
        15
        >>> datasets_with_8_vars = benchmark.list_datasets(num_variables=8, verbose=False)
        >>> datasets_with_8_vars[0]
        'II.36.38'

    Args:
        verbose: If ``True``, also prints a description of each dataset.
        num_variables: If not ``-1``, only return datasets with this many input variables.

    Returns:
        A list of dataset names.
    """
    datasets = [
        dataset_name
        for dataset_name in self.datasets
        if num_variables < 0 or self.datasets[dataset_name]["num_variables"] == num_variables
    ]
    datasets = sorted(
        datasets,
        key=lambda dataset_name: (
            self.datasets[dataset_name]["num_variables"],
            dataset_name,
        ),
    )

    if verbose:
        part1 = []
        part2 = []
        part3 = []
        max_length_1 = 0
        max_length_2 = 0
        for d in datasets:
            if self.datasets[d]["num_variables"] == 1:
                variable_str = "1 variable"
            elif self.datasets[d]["num_variables"] < 1:
                variable_str = "Amount of variables unknown"
            else:
                variable_str = f"{self.datasets[d]['num_variables']} variables"
            part1.append(d + ":")
            part2.append(variable_str)
            part3.append(self.datasets[d]["original_equation"])
            if len(d) + 1 > max_length_1:
                max_length_1 = len(d) + 1
            if len(variable_str) > max_length_2:
                max_length_2 = len(variable_str)

        for p1, p2, p3 in zip(part1, part2, part3):
            print(f"{p1:<{max_length_1}} {p2:<{max_length_2}}, Expression: {p3}")
    return datasets

save_benchmark

save_benchmark()

Saves the benchmark to <base_dir>/dataset_info.json.

The JSON file stores dataset metadata and paths to data files; the data arrays themselves are not embedded. Use load_benchmark to restore the benchmark.

Examples:

>>> from SRToolkit.dataset import Feynman
>>> benchmark = Feynman() # Feynman is a specific instance of SR_benchmark with additional functionality
>>> benchmark.save_benchmark()
Source code in SRToolkit/dataset/sr_benchmark.py
def save_benchmark(self):
    """
    Saves the benchmark to ``<base_dir>/dataset_info.json``.

    The JSON file stores dataset metadata and paths to data files; the data arrays themselves are
    not embedded. Use [load_benchmark][SRToolkit.dataset.sr_benchmark.SR_benchmark.load_benchmark]
    to restore the benchmark.

    Examples:
        >>> from SRToolkit.dataset import Feynman
        >>> benchmark = Feynman() # Feynman is a specific instance of SR_benchmark with additional functionality
        >>> benchmark.save_benchmark()
    """
    datasets = []
    for dataset_name, dataset_info in self.datasets.items():
        if "sr_dataset" in dataset_info:
            datasets.append({"name": dataset_name, "info": dataset_info["sr_dataset"].to_dict(self.base_dir)})
        else:
            datasets.append({"name": dataset_name, "info": dataset_info})

    output = {"datasets": datasets, "metadata": self.metadata, "name": self.benchmark_name}

    with open(f"{self.base_dir}/dataset_info.json", "w") as f:
        json.dump(output, f)

load_benchmark staticmethod

load_benchmark(base_dir: str) -> SR_benchmark

Loads a benchmark stored at the base directory, returning an instance of SR_benchmark.

Examples:

>>> from SRToolkit.dataset import Feynman
>>> b1 = Feynman("data/feynman") # Feynman is a specific instance of SR_benchmark with additional functionality
>>> b1.save_benchmark()
>>> b2 = SR_benchmark.load_benchmark('data/feynman')
>>> len(b1.list_datasets(verbose=False))
100
>>> len(b2.list_datasets(verbose=False))
100
>>> dataset_name = b2.list_datasets(verbose=False)[0]
>>> dataset = b2.create_dataset(dataset_name)
>>> rmse = dataset.create_evaluator().evaluate_expr(dataset.ground_truth)
>>> bool(rmse < dataset.success_threshold)
True

Parameters:

Name Type Description Default
base_dir str

Directory containing the dataset_info.json file previously written by save_benchmark.

required

Returns:

Type Description
SR_benchmark

An SR_benchmark instance with all datasets restored from the saved JSON.

Raises:

Type Description
FileNotFoundError

If dataset_info.json does not exist in base_dir.

JSONDecodeError

If the JSON file is malformed.

Source code in SRToolkit/dataset/sr_benchmark.py
@staticmethod
def load_benchmark(base_dir: str) -> "SR_benchmark":
    """
    Loads a benchmark stored at the base directory, returning an instance of SR_benchmark.

    Examples:
        >>> from SRToolkit.dataset import Feynman
        >>> b1 = Feynman("data/feynman") # Feynman is a specific instance of SR_benchmark with additional functionality
        >>> b1.save_benchmark()
        >>> b2 = SR_benchmark.load_benchmark('data/feynman')
        >>> len(b1.list_datasets(verbose=False))
        100
        >>> len(b2.list_datasets(verbose=False))
        100
        >>> dataset_name = b2.list_datasets(verbose=False)[0]
        >>> dataset = b2.create_dataset(dataset_name)
        >>> rmse = dataset.create_evaluator().evaluate_expr(dataset.ground_truth)
        >>> bool(rmse < dataset.success_threshold)
        True

    Args:
        base_dir: Directory containing the ``dataset_info.json`` file previously written by
            [save_benchmark][SRToolkit.dataset.sr_benchmark.SR_benchmark.save_benchmark].

    Returns:
        An [SR_benchmark][SRToolkit.dataset.sr_benchmark.SR_benchmark] instance with all datasets restored from the saved JSON.

    Raises:
        FileNotFoundError: If ``dataset_info.json`` does not exist in ``base_dir``.
        json.JSONDecodeError: If the JSON file is malformed.
    """
    with open(f"{base_dir}/dataset_info.json", "r") as f:
        data = json.load(f)

    datasets = {}
    for dataset_info in data["datasets"]:
        datasets[dataset_info["name"]] = dataset_info["info"]

    benchmark = SR_benchmark(data["name"], base_dir, metadata=data["metadata"])
    benchmark.datasets = datasets
    return benchmark

SR_dataset

SR_dataset(
    X: ndarray,
    symbol_library: SymbolLibrary,
    ranking_function: str = "rmse",
    y: Optional[ndarray] = None,
    max_evaluations: int = -1,
    ground_truth: Optional[Union[List[str], Node, ndarray]] = None,
    original_equation: Optional[str] = None,
    success_threshold: Optional[float] = None,
    seed: Optional[int] = None,
    dataset_metadata: Optional[dict] = None,
    dataset_name: str = "unnamed",
    **kwargs: Unpack[EstimationSettings],
)

Wraps input data and evaluation settings for a single symbolic regression problem.

Examples:

>>> X = np.array([[1, 2], [3, 4], [5, 6]])
>>> dataset = SR_dataset(X, SymbolLibrary.default_symbols(2), ground_truth=["X_0", "+", "X_1"],
...     y=np.array([3, 7, 11]), max_evaluations=10000, original_equation="z = x + y", success_threshold=1e-6)
>>> evaluator = dataset.create_evaluator()
>>> bool(evaluator.evaluate_expr(["sin", "(", "X_0", ")"]) < dataset.success_threshold)
False
>>> bool(evaluator.evaluate_expr(["u-", "C", "*", "X_1", "+", "X_0"]) < dataset.success_threshold)
True

Parameters:

Name Type Description Default
X ndarray

Input data of shape (n_samples, n_features) used to evaluate expressions.

required
symbol_library SymbolLibrary

The symbol library defining the token vocabulary.

required
ranking_function str

Ranking function to use. "rmse" calculates the error between ground truth values and expression outputs with fitted free parameters. "bed" is a stochastic measure of behavioral distance between expressions; it is less sensitive to overfitting and focuses more on structure identification.

'rmse'
y Optional[ndarray]

Target values used for parameter estimation when ranking_function="rmse".

None
max_evaluations int

Maximum number of expressions to evaluate. Values less than 0 mean no limit.

-1
ground_truth Optional[Union[List[str], Node, ndarray]]

The ground truth expression, as a list of tokens in infix notation, a Node tree, or a numpy array of behavior vectors (see create_behavior_matrix).

None
original_equation Optional[str]

Human-readable string of the original equation (e.g. "z = x + y").

None
success_threshold Optional[float]

Error threshold below which an expression is considered successful. If None, no threshold is applied.

None
seed Optional[int]

Random seed for reproducibility. None means no seed is set.

None
dataset_metadata Optional[dict]

Optional dictionary of metadata about the dataset (e.g. citation, variable names).

None
dataset_name str

Name for this dataset. Defaults to "unnamed".

'unnamed'
**kwargs Unpack[EstimationSettings]

Optional estimation settings passed to SR_evaluator. Supported keys: method, tol, gtol, max_iter, constant_bounds, initialization, max_constants, max_expr_length, num_points_sampled, bed_X, num_consts_sampled, domain_bounds.

{}
Source code in SRToolkit/dataset/sr_dataset.py
def __init__(
    self,
    X: np.ndarray,
    symbol_library: SymbolLibrary,
    ranking_function: str = "rmse",
    y: Optional[np.ndarray] = None,
    max_evaluations: int = -1,
    ground_truth: Optional[Union[List[str], Node, np.ndarray]] = None,
    original_equation: Optional[str] = None,
    success_threshold: Optional[float] = None,
    seed: Optional[int] = None,
    dataset_metadata: Optional[dict] = None,
    dataset_name: str = "unnamed",
    **kwargs: Unpack[EstimationSettings],
) -> None:
    """
    Wraps input data and evaluation settings for a single symbolic regression problem.

    Examples:
        >>> X = np.array([[1, 2], [3, 4], [5, 6]])
        >>> dataset = SR_dataset(X, SymbolLibrary.default_symbols(2), ground_truth=["X_0", "+", "X_1"],
        ...     y=np.array([3, 7, 11]), max_evaluations=10000, original_equation="z = x + y", success_threshold=1e-6)
        >>> evaluator = dataset.create_evaluator()
        >>> bool(evaluator.evaluate_expr(["sin", "(", "X_0", ")"]) < dataset.success_threshold)
        False
        >>> bool(evaluator.evaluate_expr(["u-", "C", "*", "X_1", "+", "X_0"]) < dataset.success_threshold)
        True

    Args:
        X: Input data of shape ``(n_samples, n_features)`` used to evaluate expressions.
        symbol_library: The symbol library defining the token vocabulary.
        ranking_function: Ranking function to use. ``"rmse"`` calculates the error between ground truth
            values and expression outputs with fitted free parameters. ``"bed"`` is a stochastic measure of
            behavioral distance between expressions; it is less sensitive to overfitting and focuses more on
            structure identification.
        y: Target values used for parameter estimation when ``ranking_function="rmse"``.
        max_evaluations: Maximum number of expressions to evaluate. Values less than 0 mean no limit.
        ground_truth: The ground truth expression, as a list of tokens in infix notation, a
            [Node][SRToolkit.utils.expression_tree.Node] tree, or a numpy array of behavior vectors
            (see [create_behavior_matrix][SRToolkit.utils.measures.create_behavior_matrix]).
        original_equation: Human-readable string of the original equation (e.g. ``"z = x + y"``).
        success_threshold: Error threshold below which an expression is considered successful. If ``None``,
            no threshold is applied.
        seed: Random seed for reproducibility. ``None`` means no seed is set.
        dataset_metadata: Optional dictionary of metadata about the dataset (e.g. citation, variable names).
        dataset_name: Name for this dataset. Defaults to ``"unnamed"``.
        **kwargs: Optional estimation settings passed to
            [SR_evaluator][SRToolkit.evaluation.sr_evaluator.SR_evaluator].
            Supported keys: ``method``, ``tol``, ``gtol``, ``max_iter``, ``constant_bounds``,
            ``initialization``, ``max_constants``, ``max_expr_length``, ``num_points_sampled``,
            ``bed_X``, ``num_consts_sampled``, ``domain_bounds``.
    """
    self.X = X
    self.symbol_library = symbol_library
    self.y = y
    self.max_evaluations = max_evaluations
    self.success_threshold = success_threshold
    self.ranking_function = ranking_function
    self.ground_truth = ground_truth
    self.original_equation = original_equation
    self.kwargs = kwargs
    self.dataset_name = dataset_name

    # See if symbols contain a symbol for constants
    symbols_metadata = self.symbol_library.symbols.values()
    self.contains_constants = any([symbol["type"] == "const" for symbol in symbols_metadata])

    self.seed = seed
    self.dataset_metadata = dataset_metadata

evaluate_approach

evaluate_approach(
    sr_approach: SR_approach,
    num_experiments: int = 1,
    top_k: int = 20,
    initial_seed: Optional[int] = None,
    results: Optional[SR_results] = None,
    callbacks: Optional[
        Union[SRCallbacks, CallbackDispatcher, List[SRCallbacks]]
    ] = None,
    verbose: bool = True,
    adaptation_path: Optional[str] = None,
) -> SR_results

Evaluates an SR approach on this dataset.

Examples:

>>> X = np.array([[1, 2], [3, 4], [5, 6]])
>>> dataset = SR_dataset(X, SymbolLibrary.default_symbols(2), ground_truth=["X_0", "+", "X_1"],
...     y=np.array([3, 7, 11]), max_evaluations=10000, original_equation="z = x + y")
>>> results = dataset.evaluate_approach(my_approach, num_experiments=5)

Parameters:

Name Type Description Default
sr_approach SR_approach

The SR approach to evaluate.

required
num_experiments int

Number of independent experiments (runs) to perform.

1
top_k int

Number of top expressions to retain per experiment.

20
initial_seed Optional[int]

Seed for random number generation. If None, the dataset seed is used.

None
results Optional[SR_results]

Existing SR_results object to append results to. If None, a new one is created.

None
callbacks Optional[Union[SRCallbacks, CallbackDispatcher, List[SRCallbacks]]]

Optional list of SRCallbacks, SRCallbacks, or CallbackDispatcher for monitoring and controlling the search.

None
verbose bool

If True, prints progress for each experiment.

True
adaptation_path Optional[str]

Path to save/load the adapted state for approaches with adaptation_scope="once". If the file already exists it is loaded directly, skipping adaptation. If it does not exist, the approach is adapted and the state is saved to this path. If None, adaptation runs without saving.

None

Returns:

Type Description
SR_results

An SR_results object containing results from all experiments.

Source code in SRToolkit/dataset/sr_dataset.py
def evaluate_approach(
    self,
    sr_approach: SR_approach,
    num_experiments: int = 1,
    top_k: int = 20,
    initial_seed: Optional[int] = None,
    results: Optional[SR_results] = None,
    callbacks: Optional[Union[SRCallbacks, CallbackDispatcher, List[SRCallbacks]]] = None,
    verbose: bool = True,
    adaptation_path: Optional[str] = None,
) -> SR_results:
    """
    Evaluates an SR approach on this dataset.

    Examples:
        >>> X = np.array([[1, 2], [3, 4], [5, 6]])
        >>> dataset = SR_dataset(X, SymbolLibrary.default_symbols(2), ground_truth=["X_0", "+", "X_1"],
        ...     y=np.array([3, 7, 11]), max_evaluations=10000, original_equation="z = x + y")
        >>> results = dataset.evaluate_approach(my_approach, num_experiments=5)  # doctest: +SKIP

    Args:
        sr_approach: The SR approach to evaluate.
        num_experiments: Number of independent experiments (runs) to perform.
        top_k: Number of top expressions to retain per experiment.
        initial_seed: Seed for random number generation. If ``None``, the dataset seed is used.
        results: Existing [SR_results][SRToolkit.evaluation.sr_evaluator.SR_results] object to append
            results to. If ``None``, a new one is created.
        callbacks: Optional list of [SRCallbacks][SRToolkit.evaluation.callbacks.SRCallbacks], [SRCallbacks][SRToolkit.evaluation.callbacks.SRCallbacks], or
            [CallbackDispatcher][SRToolkit.evaluation.callbacks.CallbackDispatcher] for monitoring
            and controlling the search.
        verbose: If ``True``, prints progress for each experiment.
        adaptation_path: Path to save/load the adapted state for approaches with
            ``adaptation_scope="once"``. If the file already exists it is loaded directly,
            skipping adaptation. If it does not exist, the approach is adapted and the state
            is saved to this path. If ``None``, adaptation runs without saving.

    Returns:
        An [SR_results][SRToolkit.evaluation.sr_evaluator.SR_results] object containing results from all experiments.
    """
    if initial_seed is None:
        seed = self.seed
    else:
        seed = initial_seed

    if results is None:
        results = SR_results()

    if isinstance(callbacks, SRCallbacks):
        dispatcher = CallbackDispatcher(callbacks=[callbacks])
        callbacks = dispatcher
    elif isinstance(callbacks, list):
        if len(callbacks) == 0:
            callbacks = None
        else:
            callbacks = CallbackDispatcher(callbacks=callbacks)

    dataset_name = self.dataset_name

    if sr_approach.adaptation_scope == "once":
        if adaptation_path is not None and os.path.exists(adaptation_path):
            sr_approach.load_adapted_state(adaptation_path)
        else:
            sr_approach.adapt(self.X, self.symbol_library)
            if adaptation_path is not None:
                dir_name = os.path.dirname(adaptation_path)
                if dir_name:
                    os.makedirs(dir_name, exist_ok=True)
                sr_approach.save_adapted_state(adaptation_path)

    for experiment in range(num_experiments):
        if verbose:
            print(f"Running experiment {experiment + 1}/{num_experiments}")
        if seed is not None:
            seed += 1

        event = ExperimentEvent(
            dataset_name=dataset_name,
            approach_name=sr_approach.name,
            success_threshold=self.success_threshold,
            max_evaluations=self.max_evaluations,
            seed=seed,
        )
        if callbacks is not None:
            callbacks.on_experiment_start(event)

        sr_approach.prepare()

        if sr_approach.adaptation_scope == "experiment":
            sr_approach.adapt(self.X, self.symbol_library)

        evaluator = self.create_evaluator(seed=seed)
        evaluator._experiment_id = f"{dataset_name}_{sr_approach.name}_{seed}"
        evaluator.set_callbacks(callbacks)
        sr_approach.search(evaluator, seed)
        results += evaluator.get_results(sr_approach.name, top_k)

        if callbacks is not None:
            callbacks.on_experiment_end(event, results.results[-1])
    return results

create_evaluator

create_evaluator(
    metadata: Optional[Dict[str, Any]] = None, seed: Optional[int] = None
) -> SR_evaluator

Creates an instance of the SR_evaluator class from this dataset.

Examples:

>>> X = np.array([[1, 2], [3, 4], [5, 6]])
>>> dataset = SR_dataset(X, SymbolLibrary.default_symbols(2), ground_truth=["X_0", "+", "X_1"],
...     y=np.array([3, 7, 11]), max_evaluations=10000, original_equation="z = x + y", success_threshold=1e-6)
>>> evaluator = dataset.create_evaluator()
>>> float(evaluator.evaluate_expr(["sin", "(", "X_0", ")"]))
8.05645397...
>>> float(evaluator.evaluate_expr(["X_1", "+", "X_0"]))
0.0...

Parameters:

Name Type Description Default
metadata Optional[Dict[str, Any]]

Optional dictionary of metadata to attach to the evaluator (e.g. model name, seed). Dataset metadata is merged in automatically.

None
seed Optional[int]

Seed for the random number generator. If None, the dataset seed is used.

None

Returns:

Type Description
SR_evaluator

A configured SR_evaluator ready to evaluate expressions against this dataset.

Raises:

Type Description
Exception

If SR_evaluator cannot be instantiated with the current dataset settings.

Source code in SRToolkit/dataset/sr_dataset.py
def create_evaluator(self, metadata: Optional[Dict[str, Any]] = None, seed: Optional[int] = None) -> SR_evaluator:
    """
    Creates an instance of the SR_evaluator class from this dataset.

    Examples:
        >>> X = np.array([[1, 2], [3, 4], [5, 6]])
        >>> dataset = SR_dataset(X, SymbolLibrary.default_symbols(2), ground_truth=["X_0", "+", "X_1"],
        ...     y=np.array([3, 7, 11]), max_evaluations=10000, original_equation="z = x + y", success_threshold=1e-6)
        >>> evaluator = dataset.create_evaluator()
        >>> float(evaluator.evaluate_expr(["sin", "(", "X_0", ")"]))  # doctest: +ELLIPSIS
        8.05645397...
        >>> float(evaluator.evaluate_expr(["X_1", "+", "X_0"]))  # doctest: +ELLIPSIS
        0.0...

    Args:
        metadata: Optional dictionary of metadata to attach to the evaluator (e.g. model name, seed).
            Dataset metadata is merged in automatically.
        seed: Seed for the random number generator. If ``None``, the dataset seed is used.

    Returns:
        A configured [SR_evaluator][SRToolkit.evaluation.sr_evaluator.SR_evaluator] ready to evaluate expressions against this dataset.

    Raises:
        Exception: If [SR_evaluator][SRToolkit.evaluation.sr_evaluator.SR_evaluator] cannot be
            instantiated with the current dataset settings.
    """
    if metadata is None:
        metadata = dict()
    if self.dataset_metadata is not None:
        metadata["dataset_metadata"] = self.dataset_metadata
    metadata["dataset_name"] = self.dataset_name

    if seed is None:
        seed = self.seed

    try:
        return SR_evaluator(
            X=self.X,
            y=self.y,
            max_evaluations=self.max_evaluations,
            success_threshold=self.success_threshold,
            ranking_function=self.ranking_function,
            ground_truth=self.ground_truth,
            symbol_library=self.symbol_library,
            seed=seed,
            metadata=metadata,
            **self.kwargs,
        )
    except Exception as e:
        raise e

__str__

__str__() -> str

Returns a string describing this dataset.

The string describes the target expression, symbols that should be used, and the success threshold. It also includes any constraints that should be followed when evaluating a model on this dataset. These constraints include the maximum number of expressions to evaluate, the maximum length of the expression, and the maximum number of constants allowed in the expression. If the symbol library contains a symbol for constants, the string also includes the range of constants.

For other metadata, please refer to the attribute self.dataset_metadata.

Examples:

>>> X = np.array([[1, 2], [3, 4], [5, 6]])
>>> dataset = SR_dataset(X, SymbolLibrary.default_symbols(2), ground_truth=["X_0", "+", "X_1"],
...     y=np.array([3, 7, 11]), max_evaluations=10000, original_equation="z = x + y", success_threshold=1e-6)
>>> str(dataset)
'Dataset for target expression z = x + y. When evaluating your model on this dataset, you should limit your generative model to only produce expressions using the following symbols: +, -, *, /, ^, u-, sqrt, sin, cos, exp, tan, arcsin, arccos, arctan, sinh, cosh, tanh, floor, ceil, ln, log, ^-1, ^2, ^3, ^4, ^5, pi, e, C, X_0, X_1.\nExpressions will be ranked based on the RMSE ranking function.\nExpressions are deemed successful if the root mean squared error is less than 1e-06. However, we advise that you check the best performing expressions manually to ensure they are correct.\nDataset uses the default limitations (extra arguments) from the SR_evaluator.The expressions in the dataset can contain constants/free parameters.\nFor other metadata, please refer to the attribute self.dataset_metadata.'

Returns:

Type Description
str

A string describing this dataset.

Source code in SRToolkit/dataset/sr_dataset.py
def __str__(self) -> str:
    r"""
    Returns a string describing this dataset.

    The string describes the target expression, symbols that should be used,
    and the success threshold. It also includes any constraints that should
    be followed when evaluating a model on this dataset. These constraints include the maximum
    number of expressions to evaluate, the maximum length of the expression,
    and the maximum number of constants allowed in the expression. If the
    symbol library contains a symbol for constants, the string also includes
    the range of constants.

    For other metadata, please refer to the attribute self.dataset_metadata.

    Examples:
        >>> X = np.array([[1, 2], [3, 4], [5, 6]])
        >>> dataset = SR_dataset(X, SymbolLibrary.default_symbols(2), ground_truth=["X_0", "+", "X_1"],
        ...     y=np.array([3, 7, 11]), max_evaluations=10000, original_equation="z = x + y", success_threshold=1e-6)
        >>> str(dataset)
        'Dataset for target expression z = x + y. When evaluating your model on this dataset, you should limit your generative model to only produce expressions using the following symbols: +, -, *, /, ^, u-, sqrt, sin, cos, exp, tan, arcsin, arccos, arctan, sinh, cosh, tanh, floor, ceil, ln, log, ^-1, ^2, ^3, ^4, ^5, pi, e, C, X_0, X_1.\nExpressions will be ranked based on the RMSE ranking function.\nExpressions are deemed successful if the root mean squared error is less than 1e-06. However, we advise that you check the best performing expressions manually to ensure they are correct.\nDataset uses the default limitations (extra arguments) from the SR_evaluator.The expressions in the dataset can contain constants/free parameters.\nFor other metadata, please refer to the attribute self.dataset_metadata.'

    Returns:
        A string describing this dataset.
    """
    description = f"Dataset for target expression {self.original_equation}."
    description += (
        f" When evaluating your model on this dataset, you should limit your generative model to only "
        f"produce expressions using the following symbols: {str(self.symbol_library)}.\nExpressions will be "
        f"ranked based on the {self.ranking_function.upper()} ranking function.\n"
    )

    if self.success_threshold is not None:
        description += (
            "Expressions are deemed successful if the root mean squared error is less than "
            f"{self.success_threshold}. However, we advise that you check the best performing "
            f"expressions manually to ensure they are correct.\n"
        )

    if len(self.kwargs) == 0:
        description += "Dataset uses the default limitations (extra arguments) from the SR_evaluator."
    else:
        limitations = "Non default limitations (extra arguments) from the SR_evaluators are:"
        for key, value in self.kwargs.items():
            limitations += f" {key}={value}, "
        limitations = limitations[:-2] + ".\n"
        description += limitations

    if self.contains_constants:
        description += "The expressions in the dataset can contain constants/free parameters.\n"

    description += "For other metadata, please refer to the attribute self.dataset_metadata."

    return description

to_dict

to_dict(base_path: str) -> dict

Creates a dictionary representation of this dataset. This is mainly used for saving the dataset to disk.

Examples:

>>> import tempfile
>>> X = np.array([[1, 2], [3, 4], [5, 6]])
>>> dataset = SR_dataset(X, SymbolLibrary.default_symbols(2), ground_truth=["X_0", "+", "X_1"],
...     y=np.array([3, 7, 11]), max_evaluations=10000, original_equation="z = x + y", success_threshold=1e-6)
>>> dataset.to_dict("data/example_ds")

Parameters:

Name Type Description Default
base_path str

The path to the directory where the data in the dataset should be saved.

required

Returns:

Type Description
dict

A dictionary representation of this dataset.

Source code in SRToolkit/dataset/sr_dataset.py
def to_dict(self, base_path: str) -> dict:
    r"""
    Creates a dictionary representation of this dataset. This is mainly used for saving the dataset to disk.

    Examples:
        >>> import tempfile
        >>> X = np.array([[1, 2], [3, 4], [5, 6]])
        >>> dataset = SR_dataset(X, SymbolLibrary.default_symbols(2), ground_truth=["X_0", "+", "X_1"],
        ...     y=np.array([3, 7, 11]), max_evaluations=10000, original_equation="z = x + y", success_threshold=1e-6)
        >>> dataset.to_dict("data/example_ds")  # doctest: +SKIP

    Args:
        base_path: The path to the directory where the data in the dataset should be saved.

    Returns:
        A dictionary representation of this dataset.
    """
    output = {
        "format_version": 1,
        "symbol_library": self.symbol_library.to_dict(),
        "ranking_function": self.ranking_function,
        "max_evaluations": self.max_evaluations,
        "success_threshold": self.success_threshold,
        "original_equation": self.original_equation,
        "seed": self.seed,
        "dataset_metadata": self.dataset_metadata,
        "dataset_name": self.dataset_name,
    }

    if self.kwargs is not None and "bed_X" in self.kwargs and isinstance(self.kwargs["bed_X"], np.ndarray):
        self.kwargs["bed_X"] = self.kwargs["bed_X"].tolist()

    output["kwargs"] = self.kwargs

    if not os.path.isdir(base_path):
        os.makedirs(base_path)

    if self.ground_truth is None:
        output["ground_truth"] = None
    else:
        if isinstance(self.ground_truth, list):
            output["ground_truth"] = self.ground_truth
        elif isinstance(self.ground_truth, Node):
            output["ground_truth"] = self.ground_truth.to_list()
        elif isinstance(self.ground_truth, np.ndarray) and not os.path.exists(
            f"{base_path}/{self.dataset_name}_gt.npy"
        ):
            np.save(f"{base_path}/{self.dataset_name}_gt.npy", self.ground_truth)
            output["ground_truth"] = f"{base_path}/{self.dataset_name}_gt.npy"

    if not os.path.exists(f"{base_path}/{self.dataset_name}.npz"):
        if self.y is None:
            np.savez(f"{base_path}/{self.dataset_name}.npz", X=self.X)
        else:
            np.savez(f"{base_path}/{self.dataset_name}.npz", X=self.X, y=self.y)
    output["dataset_path"] = f"{base_path}/{self.dataset_name}.npz"

    return output

from_dict staticmethod

from_dict(d: dict) -> SR_dataset

Creates an instance of the SR_dataset class from its dictionary representation. This is mainly used for loading the dataset from disk.

Examples:

>>> import tempfile, os
>>> tmpdir = tempfile.mkdtemp()
>>> X = np.array([[1, 2], [3, 4], [5, 6]])
>>> ds = SR_dataset(X, SymbolLibrary.default_symbols(2), ground_truth=["X_0", "+", "X_1"],
...     y=np.array([3, 7, 11]), max_evaluations=10000, original_equation="z = x + y", success_threshold=1e-6)
>>> dataset_dict = ds.to_dict(tmpdir)
>>> dataset = SR_dataset.from_dict(dataset_dict)
>>> dataset.X.shape
(3, 2)

Parameters:

Name Type Description Default
d dict

Dictionary representation of the dataset, as produced by to_dict.

required

Returns:

Type Description
SR_dataset

A new SR_dataset instance.

Raises:

Type Description
ValueError

If the format_version in d is not supported.

Exception

If the dataset file or ground truth file cannot be loaded.

Source code in SRToolkit/dataset/sr_dataset.py
@staticmethod
def from_dict(d: dict) -> "SR_dataset":
    """
    Creates an instance of the SR_dataset class from its dictionary representation. This is mainly used for
    loading the dataset from disk.

    Examples:
        >>> import tempfile, os
        >>> tmpdir = tempfile.mkdtemp()
        >>> X = np.array([[1, 2], [3, 4], [5, 6]])
        >>> ds = SR_dataset(X, SymbolLibrary.default_symbols(2), ground_truth=["X_0", "+", "X_1"],
        ...     y=np.array([3, 7, 11]), max_evaluations=10000, original_equation="z = x + y", success_threshold=1e-6)
        >>> dataset_dict = ds.to_dict(tmpdir)
        >>> dataset = SR_dataset.from_dict(dataset_dict)
        >>> dataset.X.shape
        (3, 2)

    Args:
        d: Dictionary representation of the dataset, as produced by
            [to_dict][SRToolkit.dataset.sr_dataset.SR_dataset.to_dict].

    Returns:
        A new [SR_dataset][SRToolkit.dataset.sr_dataset.SR_dataset] instance.

    Raises:
        ValueError: If the ``format_version`` in ``d`` is not supported.
        Exception: If the dataset file or ground truth file cannot be loaded.
    """
    if d.get("format_version", 1) != 1:
        raise ValueError(
            f"[SR_dataset.from_dict] Unsupported format_version: {d.get('format_version')!r}. Expected 1."
        )
    try:
        data = np.load(d["dataset_path"])
        X = data["X"]
        if "y" in data:
            y = data["y"]
        else:
            y = None
    except Exception:
        raise Exception(f"[SR_dataset.from_dict] Could not load dataset from {d['dataset_path']}")

    if "ground_truth" in d and (isinstance(d["ground_truth"], list) or d["ground_truth"] is None):
        ground_truth = d["ground_truth"]
    else:
        try:
            ground_truth = np.load(d["ground_truth"])
        except Exception:
            raise Exception(f"[SR_dataset.from_dict] Could not load ground truth from {d['ground_truth']}")

    if "bed_X" in d["kwargs"] and d["kwargs"]["bed_X"] is not None:
        d["kwargs"]["bed_X"] = np.array(d["kwargs"]["bed_X"])

    try:
        return SR_dataset(
            X,
            SymbolLibrary.from_dict(d["symbol_library"]),
            ranking_function=d["ranking_function"],
            y=y,
            max_evaluations=d["max_evaluations"],
            ground_truth=ground_truth,
            original_equation=d["original_equation"],
            success_threshold=d["success_threshold"],
            seed=d["seed"],
            dataset_metadata=d["dataset_metadata"],
            dataset_name=d["dataset_name"],
            **d["kwargs"],
        )
    except Exception as e:
        raise Exception(f"[SR_dataset.from_dict] Error creating dataset: {e}")