Merge pull request #425 from MilesCranmer/pre-commit

.github/ISSUE_TEMPLATE/feature_request.yml

@@ -19,4 +19,3 @@ body:
     attributes:
       value: |
         Be sure to check out the [PySR forums](https://github.com/MilesCranmer/PySR/discussions) to chat with other users about PySR use-cases!
-

.github/workflows/CI.yml

@@ -32,7 +32,7 @@ jobs:
         julia-version: ['1.9']
         python-version: ['3.10']
         os: [ubuntu-latest]
-    
+
     steps:
       - uses: actions/checkout@v3
       - name: "Set up Julia"
@@ -96,7 +96,7 @@ jobs:
       matrix:
         python-version: ['3.9']
         os: ['ubuntu-latest']
-    
+
     steps:
       - uses: actions/checkout@v3
       - name: "Cache conda"
@@ -129,7 +129,7 @@ jobs:
 
   coveralls:
     name: Indicate completion to coveralls.io
-    needs: 
+    needs:
       - test
     runs-on: ubuntu-latest
     defaults:

.github/workflows/CI_Windows.yml

@@ -32,7 +32,7 @@ jobs:
         julia-version: ['1.9']
         python-version: ['3.10']
         os: [windows-latest]
-    
+
     steps:
       - uses: actions/checkout@v3
       - name: "Set up Julia"

.github/workflows/CI_conda_forge.yml

@@ -23,7 +23,7 @@ jobs:
         python-version: ['3.8', '3.9', '3.10', '3.11']
         os: ['ubuntu-latest', 'macos-latest']
         use-mamba: [true, false]
-    
+
     steps:
       - name: "Set up Conda"
         uses: conda-incubator/setup-miniconda@v2

.github/workflows/CI_docker_large_nightly.yml

@@ -22,8 +22,8 @@ jobs:
         python-version: ['3.10']
         os: [ubuntu-latest]
         arch: ['linux/amd64', 'linux/arm64']
-      
-    
+
+
     steps:
       - uses: actions/checkout@v3
       - name: Set up QEMU

.github/workflows/CI_large_nightly.yml

@@ -26,7 +26,7 @@ jobs:
         julia-version: ['1.6', '1.8', '1.9']
         python-version: ['3.7', '3.8', '3.9', '3.10', '3.11']
         os: [ubuntu-latest, macos-latest, windows-latest]
-    
+
     steps:
       - uses: actions/checkout@v3
       - name: "Set up Julia"

.github/workflows/CI_mac.yml

@@ -32,7 +32,7 @@ jobs:
         julia-version: ['1.9']
         python-version: ['3.10']
         os: [macos-latest]
-    
+
     steps:
       - uses: actions/checkout@v3
       - name: "Set up Julia"

.github/workflows/codeql-analysis.yml

@@ -37,11 +37,11 @@ jobs:
         # If you wish to specify custom queries, you can do so here or in a config file.
         # By default, queries listed here will override any specified in a config file.
         # Prefix the list here with "+" to use these queries and those in the config file.
-        
+
         # Details on CodeQL's query packs refer to : https://docs.github.com/en/code-security/code-scanning/automatically-scanning-your-code-for-vulnerabilities-and-errors/configuring-code-scanning#using-queries-in-ql-packs
         # queries: security-extended,security-and-quality
 
-        
+
     # Autobuild attempts to build any compiled languages  (C/C++, C#, or Java).
     # If this step fails, then you should remove it and run the build manually (see below)
     - name: Autobuild
@@ -50,7 +50,7 @@ jobs:
     # ℹ️ Command-line programs to run using the OS shell.
     # 📚 See https://docs.github.com/en/actions/using-workflows/workflow-syntax-for-github-actions#jobsjob_idstepsrun
 
-    #   If the Autobuild fails above, remove it and uncomment the following three lines. 
+    #   If the Autobuild fails above, remove it and uncomment the following three lines.
     #   modify them (or add more) to build your code if your project, please refer to the EXAMPLE below for guidance.
 
     # - run: |

.github/workflows/docker_deploy.yml

@@ -9,7 +9,7 @@ on:
     tags:
       - "v*.*.*"
   workflow_dispatch:
-  
+
 
 jobs:
   docker:

.github/workflows/docs.yml

@@ -18,7 +18,7 @@ jobs:
     defaults:
       run:
         shell: bash
-    
+
     steps:
       - uses: actions/checkout@v3
       - name: "Set up Python"
@@ -33,4 +33,4 @@ jobs:
       - name: "Build API docs"
         run: cd docs && ./gen_docs.sh
       - name: "Deploy documentation"
-        run: mkdocs gh-deploy --force
+        run: mkdocs gh-deploy --force

.github/workflows/update_backend.yml

@@ -48,7 +48,7 @@ jobs:
           CURRENT_PYSR_PATCH_VERSION=$(python -c 'import pysr; print(pysr.version.__version__.split(".")[-1], end="")' 2>/dev/null)
           NEW_PYSR_PATCH_VERSION=$((CURRENT_PYSR_PATCH_VERSION + 1))
           sed -i "s/^__version__ = .*/__version__ = \"$(python -c 'import pysr; print(".".join(pysr.version.__version__.split(".")[:-1]), end="")' 2>/dev/null).${NEW_PYSR_PATCH_VERSION}\"/" pysr/version.py
-          
+
           # Set SymbolicRegression.jl version:
           sed -i "s/^__symbolic_regression_jl_version__ = .*/__symbolic_regression_jl_version__ = \"${{ steps.get-latest.outputs.version }}\"/" pysr/version.py
 

.pre-commit-config.yaml

@@ -0,0 +1,31 @@
+repos:
+  # General linting
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v3.2.0
+    hooks:
+      - id: trailing-whitespace
+      - id: end-of-file-fixer
+      - id: check-yaml
+      - id: check-added-large-files
+  # General formatting
+  - repo: https://github.com/psf/black
+    rev: 23.3.0
+    hooks:
+      - id: black
+      - id: black-jupyter
+  # Stripping notebooks
+  - repo: https://github.com/kynan/nbstripout
+    rev: 0.6.1
+    hooks:
+      - id: nbstripout
+  # Unused imports
+  - repo: https://github.com/hadialqattan/pycln
+    rev: "v2.2.2"
+    hooks:
+      - id: pycln
+  # Sorted imports
+  - repo: https://github.com/PyCQA/isort
+    rev: "5.12.0"
+    hooks:
+      - id: isort
+        additional_dependencies: [toml]

CONTRIBUTORS.md

@@ -121,4 +121,4 @@ Thanks for being part of the PySR community!
 <!-- prettier-ignore-end -->
 
 <!-- ALL-CONTRIBUTORS-LIST:END -->
-</div>
+</div>

README.md

@@ -155,7 +155,7 @@ The PySR build in conda includes all required dependencies, so you can install i
 conda install -c conda-forge pysr
 ```
 
-from within your target conda environment. 
+from within your target conda environment.
 
 However, note that the conda install does not support precompilation of Julia libraries, so the
 start time may be slightly slower as the JIT-compilation will be running.
@@ -305,7 +305,7 @@ model = PySRRegressor(
     # ^ 2 populations per core, so one is always running.
     population_size=50,
     # ^ Slightly larger populations, for greater diversity.
-    ncyclesperiteration=500, 
+    ncyclesperiteration=500,
     # ^ Generations between migrations.
     niterations=10000000,  # Run forever
     early_stop_condition=(

benchmarks/README.md

@@ -21,7 +21,7 @@ v0.3.6 | 25900
 v0.3.7 | 26600
 v0.3.8 | 7470
 v0.3.9 | 6760
-v0.3.10 | 
+v0.3.10 |
 v0.3.11 | 19500
 v0.3.12 | 19000
 v0.3.13 | 15200

benchmarks/hyperparamopt.py

@@ -1,13 +1,15 @@
 """Start a hyperoptimization from a single node"""
-import sys
-import numpy as np
 import pickle as pkl
-from pysr import PySRRegressor
+import sys
+
 import hyperopt
-from hyperopt import hp, fmin, tpe, Trials
+import numpy as np
+from hyperopt import Trials, fmin, hp, tpe
 from hyperopt.fmin import generate_trials_to_calculate
 from space import *
 
+from pysr import PySRRegressor
+
 # Change the following code to your file
 ################################################################################
 TRIALS_FOLDER = "trials2"

benchmarks/print_best_model.py

@@ -1,12 +1,11 @@
 """Print the best model parameters and loss"""
-import sys
-import numpy as np
 import pickle as pkl
-import hyperopt
-from hyperopt import hp, fmin, tpe, Trials
-from space import space
 from pprint import PrettyPrinter
 
+import hyperopt
+import numpy as np
+from hyperopt import Trials, fmin, hp, tpe
+from space import space
 
 # Change the following code to your file
 ################################################################################
@@ -51,7 +50,6 @@ import glob
 path = TRIALS_FOLDER + "/*.pkl"
 files = 0
 for fname in glob.glob(path):
-
     trials_obj = pkl.load(open(fname, "rb"))
     n_trials = trials_obj["n"]
     trials_obj = trials_obj["trials"]

benchmarks/space.py

@@ -1,5 +1,5 @@
 import numpy as np
-from hyperopt import hp, fmin, tpe, Trials
+from hyperopt import Trials, fmin, hp, tpe
 
 binary_operators = ["*", "/", "+", "-"]
 unary_operators = ["sin", "cos", "exp", "log"]

datasets/FeynmanEquations.csv

@@ -98,4 +98,4 @@ III.15.14,10,96,m,(h/(2*pi))**2/(2*E_n*d**2),3,h,1,5,E_n,1,5,d,1,5,,,,,,,,,,,,,,
 III.15.27,10,97,k,2*pi*alpha/(n*d),3,alpha,1,5,n,1,5,d,1,5,,,,,,,,,,,,,,,,,,,,,
 III.17.37,10,98,f,beta*(1+alpha*cos(theta)),3,beta,1,5,alpha,1,5,theta,1,5,,,,,,,,,,,,,,,,,,,,,
 III.19.51,10,99,E_n,-m*q**4/(2*(4*pi*epsilon)**2*(h/(2*pi))**2)*(1/n**2),5,m,1,5,q,1,5,h,1,5,n,1,5,epsilon,1,5,,,,,,,,,,,,,,,
-III.21.20,10,100,j,-rho_c_0*q*A_vec/m,4,rho_c_0,1,5,q,1,5,A_vec,1,5,m,1,5,,,,,,,,,,,,,,,,,,
+III.21.20,10,100,j,-rho_c_0*q*A_vec/m,4,rho_c_0,1,5,q,1,5,A_vec,1,5,m,1,5,,,,,,,,,,,,,,,,,,

docs/.gitignore

@@ -1,4 +1,4 @@
 build
 api.md
 index.md.bak
-papers.md
+papers.md

docs/_api.md

@@ -6,7 +6,7 @@ Let's look at them below.
 PARAMSKEY
 
 ## PySRRegressor Functions
-        
+
 ::: pysr.PySRRegressor.fit
     options:
         show_root_heading: true
@@ -60,5 +60,3 @@ PARAMSKEY
         show_root_heading: true
         heading_level: 3
         show_root_full_path: false
-
-

docs/backend.md

@@ -12,7 +12,7 @@ Generally you can do this as follows:
 git clone https://github.com/MilesCranmer/SymbolicRegression.jl
 ```
 2. Edit the source code in `src/` to your requirements:
-    -  The documentation for the backend is given [here](https://astroautomata.com/SymbolicRegression.jl/dev/). 
+    -  The documentation for the backend is given [here](https://astroautomata.com/SymbolicRegression.jl/dev/).
     - Throughout the package, you will often see template functions which typically use a symbol `T` (such as in the string `where {T<:Real}`). Here, `T` is simply the datatype of the input data and stored constants, such as `Float32` or `Float64`. Writing functions in this way lets us write functions generic to types, while still having access to the specific type specified at compilation time.
     - Expressions are stored as binary trees, using the `Node{T}` type, described [here](https://astroautomata.com/SymbolicRegression.jl/dev/types/#SymbolicRegression.CoreModule.EquationModule.Node).
     - Parts of the code which are typically edited by users include:
@@ -26,4 +26,4 @@ git clone https://github.com/MilesCranmer/SymbolicRegression.jl
 
 If you get comfortable enough with the backend, you might consider using the Julia package directly: the API is given on the [SymbolicRegression.jl documentation](https://astroautomata.com/SymbolicRegression.jl/dev/).
 
-If you make a change that you think could be useful to other users, don't hesitate to open a pull request on either the PySR or SymbolicRegression.jl repositories! Contributions are very appreciated.
+If you make a change that you think could be useful to other users, don't hesitate to open a pull request on either the PySR or SymbolicRegression.jl repositories! Contributions are very appreciated.

docs/gen_param_docs.py

@@ -1,13 +1,14 @@
 # Load YAML file param_groupings.yml:
-from pathlib import Path
-from yaml import safe_load
+import re
 import sys
 
+from docstring_parser import parse
+from yaml import safe_load
+
 sys.path.append("..")
+
+
 from pysr import PySRRegressor
-import pysr
-import re
-from docstring_parser import parse
 
 found_params = []
 

docs/generate_papers.py

@@ -1,7 +1,8 @@
 """This script generates the papers.md file from the papers.yml file."""
-import yaml
 from pathlib import Path
 
+import yaml
+
 data_file = "papers.yml"
 papers_header = Path("stylesheets") / "papers_header.txt"
 output_file = "papers.md"
@@ -49,7 +50,7 @@ with open(output_file, "w") as f:
 
 <center>
 {authors}
-    
+
 <small>{affiliations}</small>
 </center>
 

docs/operators.md

@@ -64,5 +64,3 @@ instead of `1.5e3`, if you write any constant numbers.
 Your operator should work with the entire real line (you can use
 abs(x) for operators requiring positive input - see `log_abs`); otherwise
 the search code will experience domain errors.
-
-

docs/options.md

@@ -265,7 +265,7 @@ PySRRegressor(..., loss="loss(x, y) = abs(x * y)")
 With weights:
 
 ```python
-model = PySRRegressor(..., loss="myloss(x, y, w) = w * abs(x - y)") 
+model = PySRRegressor(..., loss="myloss(x, y, w) = w * abs(x - y)")
 model.fit(..., weights=weights)
 ```
 

docs/papers.yml

@@ -151,7 +151,6 @@ papers:
     abstract: "We present an approach for using machine learning to automatically discover the governing equations and hidden properties of real physical systems from observations. We train a \"graph neural network\" to simulate the dynamics of our solar system's Sun, planets, and large moons from 30 years of trajectory data. We then use symbolic regression to discover an analytical expression for the force law implicitly learned by the neural network, which our results showed is equivalent to Newton's law of gravitation. The key assumptions that were required were translational and rotational equivariance, and Newton's second and third laws of motion. Our approach correctly discovered the form of the symbolic force law. Furthermore, our approach did not require any assumptions about the masses of planets and moons or physical constants. They, too, were accurately inferred through our methods. Though, of course, the classical law of gravitation has been known since Isaac Newton, our result serves as a validation that our method can discover unknown laws and hidden properties from observed data. More broadly this work represents a key step toward realizing the potential of machine learning for accelerating scientific discovery."
     image: rediscovering_gravity.png
     date: 2022-02-04
-    link: https://arxiv.org/abs/2202.02306
   - title: (Thesis) On Neural Differential Equations - Section 6.1
     authors:
       - Patrick Kidger (1)

docs/requirements.txt

@@ -1,4 +1,4 @@
 mkdocs-material
 mkdocs-autorefs
 mkdocstrings[python]
-docstring_parser
+docstring_parser

docs/stylesheets/extra.css

@@ -2,4 +2,4 @@
     --md-primary-fg-color:        #C13245;
     --md-primary-fg-color--light: #D35364;
     --md-primary-fg-color--dark:  #982736;
-}
+}

docs/stylesheets/papers_header.txt

@@ -6,4 +6,3 @@ These are sorted by the date of release, with most recent papers at the top.
 
 If you have used PySR in your research,
 please submit a pull request to add your paper to [this file](https://github.com/MilesCranmer/PySR/blob/master/docs/papers.yml).
-

docs/tuning.md

@@ -17,7 +17,7 @@ I run from IPython (Jupyter Notebooks don't work as well[^1]) on the head node o
 5. Set `ncyclesperiteration` to maybe `5000` or so, until the head node occupation is under `10%`.
 6. Set `constraints` and `nested_constraints` as strict as possible. These can help quite a bit with exploration. Typically, if I am using `pow`, I would set `constraints={"pow": (9, 1)}`, so that power laws can only have a variable or constant as their exponent. If I am using `sin` and `cos`, I also like to set `nested_constraints={"sin": {"sin": 0, "cos": 0}, "cos": {"sin": 0, "cos": 0}}`, so that sin and cos can't be nested, which seems to happen frequently. (Although in practice I would just use `sin`, since the search could always add a phase offset!)
 7. Set `maxsize` a bit larger than the final size you want. e.g., if you want a final equation of size `30`, you might set this to `35`, so that it has a bit of room to explore.
-8. Set `maxdepth` strictly, but leave a bit of room for exploration. e.g., if you want a final equation limited to a depth of `5`, you might set this to `6` or `7`, so that it has a bit of room to explore. 
+8. Set `maxdepth` strictly, but leave a bit of room for exploration. e.g., if you want a final equation limited to a depth of `5`, you might set this to `6` or `7`, so that it has a bit of room to explore.
 9.  Set `parsimony` equal to about the minimum loss you would expect, divided by 5-10. e.g., if you expect the final equation to have a loss of `0.001`, you might set `parsimony=0.0001`.
 10. Set `weight_optimize` to some larger value, maybe `0.001`. This is very important if `ncyclesperiteration` is large, so that optimization happens more frequently.
 11. Set `turbo` to `True`. This may or not work, if there's an error just turn it off (some operators are not SIMD-capable). If it does work, it should give you a nice 20% speedup.
@@ -31,7 +31,7 @@ Some things I try out to see if they help:
 2. Try setting `adaptive_parsimony_scaling` a bit larger, maybe up to `1000`.
 3. Sometimes I try using `warmup_maxsize_by`. This is useful if you find that the search finds a very complex equation very quickly, and then gets stuck. It basically forces it to start at the simpler equations and build up complexity slowly.
 4. Play around with different losses:
-    - I typically try `L2DistLoss()` and `L1DistLoss()`. L1 loss is more robust to outliers compared to L2 (L1 finds the median, while L2 finds the mean of a random variable), so is often a good choice for a noisy dataset. 
+    - I typically try `L2DistLoss()` and `L1DistLoss()`. L1 loss is more robust to outliers compared to L2 (L1 finds the median, while L2 finds the mean of a random variable), so is often a good choice for a noisy dataset.
     - I might also provide the `weights` parameter to `fit` if there is some reasonable choice of weighting. For example, maybe I know the signal-to-noise of a particular row of `y` - I would set that SNR equal to the weights. Or, perhaps I do some sort of importance sampling, and weight the rows by importance.
 
 Very rarely I might also try tuning the mutation weights, the crossover probability, or the optimization parameters. I never use `denoise` or `select_k_features` as I find they aren't very useful.

environment.yml

@@ -10,4 +10,4 @@ dependencies:
   - pyjulia
   - openlibm
   - openspecfun
-  - click
+  - click

mkdocs.yml

@@ -13,7 +13,7 @@ theme:
       toggle:
         icon: material/toggle-switch-off-outline
         name: Switch to light mode
-    
+
 
   features:
     - navigation.expand

pyproject.toml

@@ -0,0 +1,2 @@
+[tool.isort]
+profile = "black"

pysr/__init__.py

@@ -1,14 +1,23 @@
 from . import sklearn_monkeypatch
-from .version import __version__
-from .sr import (
-    pysr,
-    PySRRegressor,
-    best,
-    best_tex,
-    best_callable,
-    best_row,
-)
-from .julia_helpers import install
-from .feynman_problems import Problem, FeynmanProblem
 from .export_jax import sympy2jax
 from .export_torch import sympy2torch
+from .feynman_problems import FeynmanProblem, Problem
+from .julia_helpers import install
+from .sr import PySRRegressor, best, best_callable, best_row, best_tex, pysr
+from .version import __version__
+
+__all__ = [
+    "sklearn_monkeypatch",
+    "sympy2jax",
+    "sympy2torch",
+    "FeynmanProblem",
+    "Problem",
+    "install",
+    "PySRRegressor",
+    "best",
+    "best_callable",
+    "best_row",
+    "best_tex",
+    "pysr",
+    "__version__",
+]

pysr/_cli/main.py

@@ -1,4 +1,5 @@
 import click
+
 from ..julia_helpers import install
 
 

pysr/export_jax.py

@@ -1,7 +1,4 @@
-import functools as ft
 import sympy
-import string
-import random
 
 # Special since need to reduce arguments.
 MUL = 0

pysr/export_latex.py

@@ -1,9 +1,9 @@
 """Functions to help export PySR equations to LaTeX."""
+from typing import List
+
+import pandas as pd
 import sympy
 from sympy.printing.latex import LatexPrinter
-import pandas as pd
-from typing import List
-import warnings
 
 
 class PreciseLatexPrinter(LatexPrinter):

pysr/export_numpy.py

@@ -1,8 +1,9 @@
 """Code for exporting discovered expressions to numpy"""
+import warnings
+
 import numpy as np
 import pandas as pd
 from sympy import lambdify
-import warnings
 
 
 class CallableEquation:

pysr/export_torch.py

@@ -5,6 +5,7 @@
 
 import collections as co
 import functools as ft
+
 import sympy
 
 

pysr/feynman_problems.py

@@ -1,8 +1,10 @@
-import numpy as np
 import csv
-from .sr import pysr, best
-from pathlib import Path
 from functools import partial
+from pathlib import Path
+
+import numpy as np
+
+from .sr import best, pysr
 
 PKG_DIR = Path(__file__).parents[1]
 FEYNMAN_DATASET = PKG_DIR / "datasets" / "FeynmanEquations.csv"
@@ -118,6 +120,7 @@ def do_feynman_experiments_parallel(
     data_dir=FEYNMAN_DATASET,
 ):
     import multiprocessing as mp
+
     from tqdm import tqdm
 
     problems = mk_problems(first=first, gen=True, dp=dp, data_dir=data_dir)

pysr/julia_helpers.py

@@ -1,12 +1,13 @@
 """Functions for initializing the Julia environment and installing deps."""
-import sys
+import os
 import subprocess
+import sys
 import warnings
 from pathlib import Path
-import os
+
 from julia.api import JuliaError
 
-from .version import __version__, __symbolic_regression_jl_version__
+from .version import __symbolic_regression_jl_version__, __version__
 
 juliainfo = None
 julia_initialized = False

pysr/sr.py

@@ -1,40 +1,37 @@
 """Define the PySRRegressor scikit-learn interface."""
 import copy
-from io import StringIO
 import os
-import sys
-import numpy as np
-import pandas as pd
-import sympy
-from sympy import sympify
+import pickle as pkl
 import re
-import tempfile
 import shutil
-from pathlib import Path
-import pickle as pkl
-from datetime import datetime
+import sys
+import tempfile
 import warnings
+from datetime import datetime
+from io import StringIO
 from multiprocessing import cpu_count
-from sklearn.base import BaseEstimator, RegressorMixin, MultiOutputMixin
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import sympy
+from sklearn.base import BaseEstimator, MultiOutputMixin, RegressorMixin
 from sklearn.utils import check_array, check_consistent_length, check_random_state
-from sklearn.utils.validation import (
-    _check_feature_names_in,
-    check_is_fitted,
-)
+from sklearn.utils.validation import _check_feature_names_in, check_is_fitted
+from sympy import sympify
 
+from .deprecated import make_deprecated_kwargs_for_pysr_regressor
+from .export_latex import generate_multiple_tables, generate_single_table, to_latex
+from .export_numpy import CallableEquation
 from .julia_helpers import (
-    init_julia,
-    _process_julia_project,
-    is_julia_version_greater_eq,
     _escape_filename,
+    _load_backend,
     _load_cluster_manager,
+    _process_julia_project,
     _update_julia_project,
-    _load_backend,
+    init_julia,
+    is_julia_version_greater_eq,
 )
-from .export_numpy import CallableEquation
-from .export_latex import generate_single_table, generate_multiple_tables, to_latex
-from .deprecated import make_deprecated_kwargs_for_pysr_regressor
-
 
 Main = None  # TODO: Rename to more descriptive name like "julia_runtime"
 
@@ -2454,7 +2451,7 @@ def idx_model_selection(equations: pd.DataFrame, model_selection: str) -> int:
 def _denoise(X, y, Xresampled=None, random_state=None):
     """Denoise the dataset using a Gaussian process."""
     from sklearn.gaussian_process import GaussianProcessRegressor
-    from sklearn.gaussian_process.kernels import RBF, WhiteKernel, ConstantKernel
+    from sklearn.gaussian_process.kernels import RBF, ConstantKernel, WhiteKernel
 
     gp_kernel = RBF(np.ones(X.shape[1])) + WhiteKernel(1e-1) + ConstantKernel()
     gpr = GaussianProcessRegressor(

pysr/test/__init__.py

@@ -1,5 +1,7 @@
 from .test import runtests
+from .test_cli import runtests as runtests_cli
 from .test_env import runtests as runtests_env
 from .test_jax import runtests as runtests_jax
 from .test_torch import runtests as runtests_torch
-from .test_cli import runtests as runtests_cli
+
+__all__ = ["runtests", "runtests_env", "runtests_jax", "runtests_torch", "runtests_cli"]

pysr/test/test.py

@@ -1,28 +1,28 @@
+import inspect
 import os
+import pickle as pkl
+import tempfile
 import traceback
-import inspect
 import unittest
+import warnings
+from pathlib import Path
+
 import numpy as np
+import pandas as pd
+import sympy
 from sklearn import model_selection
 from sklearn.utils.estimator_checks import check_estimator
-import sympy
-import pandas as pd
-import warnings
-import pickle as pkl
-import tempfile
-from pathlib import Path
 
-from .. import julia_helpers
-from .. import PySRRegressor
+from .. import PySRRegressor, julia_helpers
+from ..export_latex import to_latex
 from ..sr import (
-    run_feature_selection,
-    _handle_feature_selection,
-    _csv_filename_to_pkl_filename,
-    idx_model_selection,
     _check_assertions,
+    _csv_filename_to_pkl_filename,
+    _handle_feature_selection,
     _process_constraints,
+    idx_model_selection,
+    run_feature_selection,
 )
-from ..export_latex import to_latex
 
 DEFAULT_PARAMS = inspect.signature(PySRRegressor.__init__).parameters
 DEFAULT_NITERATIONS = DEFAULT_PARAMS["niterations"].default

pysr/test/test_cli.py

@@ -1,5 +1,7 @@
 import unittest
+
 from click import testing as click_testing
+
 from .._cli.main import pysr
 
 

pysr/test/test_env.py

@@ -1,7 +1,7 @@
 """Contains tests for creating and initializing custom Julia projects."""
 
-import unittest
 import os
+import unittest
 from tempfile import TemporaryDirectory
 
 from .. import julia_helpers