diff --git a/juliapkg.json b/juliapkg.json new file mode 100644 index 000000000..15686a6a8 --- /dev/null +++ b/juliapkg.json @@ -0,0 +1,13 @@ +{ + "julia": "1.6", + "packages": { + "SymbolicRegression": { + "uuid": "8254be44-1295-4e6a-a16d-46603ac705cb", + "version": "0.23.1" + }, + "ClusterManagers": { + "uuid": "34f1f09b-3a8b-5176-ab39-66d58a4d544e", + "version": "0.4" + } + } +} diff --git a/pysr/__init__.py b/pysr/__init__.py index 1f2e9775d..3662c90e2 100644 --- a/pysr/__init__.py +++ b/pysr/__init__.py @@ -12,8 +12,7 @@ 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 +from .sr import PySRRegressor, jl from .version import __version__ __all__ = [ diff --git a/pysr/julia_helpers.py b/pysr/julia_helpers.py index e2f76090c..2ccf4a777 100644 --- a/pysr/julia_helpers.py +++ b/pysr/julia_helpers.py @@ -1,13 +1,10 @@ """Functions for initializing the Julia environment and installing deps.""" -import os -import subprocess -import sys import warnings -from pathlib import Path -from julia.api import JuliaError +import juliacall +import juliapkg -from .version import __symbolic_regression_jl_version__, __version__ +jl = juliacall.newmodule("PySR") juliainfo = None julia_initialized = False @@ -15,270 +12,9 @@ julia_activated_env = None -def _load_juliainfo(): - """Execute julia.core.JuliaInfo.load(), and store as juliainfo.""" - global juliainfo - - if juliainfo is None: - from julia.core import JuliaInfo - - try: - juliainfo = JuliaInfo.load(julia="julia") - except FileNotFoundError: - env_path = os.environ["PATH"] - raise FileNotFoundError( - f"Julia is not installed in your PATH. Please install Julia and add it to your PATH.\n\nCurrent PATH: {env_path}", - ) - - return juliainfo - - -def _get_julia_env_dir(): - # Have to manually get env dir: - try: - julia_env_dir_str = subprocess.run( - ["julia", "-e using Pkg; print(Pkg.envdir())"], - capture_output=True, - env=os.environ, - ).stdout.decode() - except FileNotFoundError: - env_path = os.environ["PATH"] - raise FileNotFoundError( - f"Julia is not installed in your PATH. Please install Julia and add it to your PATH.\n\nCurrent PATH: {env_path}", - ) - return Path(julia_env_dir_str) - - -def _set_julia_project_env(julia_project, is_shared): - if is_shared: - if is_julia_version_greater_eq(version=(1, 7, 0)): - os.environ["JULIA_PROJECT"] = "@" + str(julia_project) - else: - julia_env_dir = _get_julia_env_dir() - os.environ["JULIA_PROJECT"] = str(julia_env_dir / julia_project) - else: - os.environ["JULIA_PROJECT"] = str(julia_project) - - def _get_io_arg(quiet): io = "devnull" if quiet else "stderr" - io_arg = f"io={io}" if is_julia_version_greater_eq(version=(1, 6, 0)) else "" - return io_arg - - -def install(julia_project=None, quiet=False, precompile=None): # pragma: no cover - """ - Install PyCall.jl and all required dependencies for SymbolicRegression.jl. - - Also updates the local Julia registry. - """ - import julia - - _julia_version_assertion() - # Set JULIA_PROJECT so that we install in the pysr environment - processed_julia_project, is_shared = _process_julia_project(julia_project) - _set_julia_project_env(processed_julia_project, is_shared) - - if precompile == False: - os.environ["JULIA_PKG_PRECOMPILE_AUTO"] = "0" - - try: - julia.install(quiet=quiet) - except julia.tools.PyCallInstallError: - # Attempt to reset PyCall.jl's build: - subprocess.run( - [ - "julia", - "-e", - f'ENV["PYTHON"] = "{sys.executable}"; import Pkg; Pkg.build("PyCall")', - ], - ) - # Try installing again: - try: - julia.install(quiet=quiet) - except julia.tools.PyCallInstallError: - warnings.warn( - "PyCall.jl failed to install on second attempt. " - + "Please consult the GitHub issue " - + "https://github.com/MilesCranmer/PySR/issues/257 " - + "for advice on fixing this." - ) - - Main, init_log = init_julia(julia_project, quiet=quiet, return_aux=True) - io_arg = _get_io_arg(quiet) - - if precompile is None: - precompile = init_log["compiled_modules"] - - if not precompile: - Main.eval('ENV["JULIA_PKG_PRECOMPILE_AUTO"] = 0') - - if is_shared: - # Install SymbolicRegression.jl: - _add_sr_to_julia_project(Main, io_arg) - - Main.eval("using Pkg") - Main.eval(f"Pkg.instantiate({io_arg})") - - if precompile: - Main.eval(f"Pkg.precompile({io_arg})") - - if not quiet: - warnings.warn( - "It is recommended to restart Python after installing PySR's dependencies," - " so that the Julia environment is properly initialized." - ) - - -def _import_error(): - return """ - Required dependencies are not installed or built. Run the following command in your terminal: - python3 -m pysr install - """ - - -def _process_julia_project(julia_project): - if julia_project is None: - is_shared = True - processed_julia_project = f"pysr-{__version__}" - elif julia_project[0] == "@": - is_shared = True - processed_julia_project = julia_project[1:] - else: - is_shared = False - processed_julia_project = Path(julia_project) - return processed_julia_project, is_shared - - -def is_julia_version_greater_eq(juliainfo=None, version=(1, 6, 0)): - """Check if Julia version is greater than specified version.""" - if juliainfo is None: - juliainfo = _load_juliainfo() - current_version = ( - juliainfo.version_major, - juliainfo.version_minor, - juliainfo.version_patch, - ) - return current_version >= version - - -def _check_for_conflicting_libraries(): # pragma: no cover - """Check whether there are conflicting modules, and display warnings.""" - # See https://github.com/pytorch/pytorch/issues/78829: importing - # pytorch before running `pysr.fit` causes a segfault. - torch_is_loaded = "torch" in sys.modules - if torch_is_loaded: - warnings.warn( - "`torch` was loaded before the Julia instance started. " - "This may cause a segfault when running `PySRRegressor.fit`. " - "To avoid this, please run `pysr.julia_helpers.init_julia()` *before* " - "importing `torch`. " - "For updates, see https://github.com/pytorch/pytorch/issues/78829" - ) - - -def init_julia(julia_project=None, quiet=False, julia_kwargs=None, return_aux=False): - """Initialize julia binary, turning off compiled modules if needed.""" - global julia_initialized - global julia_kwargs_at_initialization - global julia_activated_env - - if not julia_initialized: - _check_for_conflicting_libraries() - - if julia_kwargs is None: - julia_kwargs = {"optimize": 3} - - from julia.core import JuliaInfo, UnsupportedPythonError - - _julia_version_assertion() - processed_julia_project, is_shared = _process_julia_project(julia_project) - _set_julia_project_env(processed_julia_project, is_shared) - - try: - info = JuliaInfo.load(julia="julia") - except FileNotFoundError: - env_path = os.environ["PATH"] - raise FileNotFoundError( - f"Julia is not installed in your PATH. Please install Julia and add it to your PATH.\n\nCurrent PATH: {env_path}", - ) - - if not info.is_pycall_built(): - raise ImportError(_import_error()) - - from julia.core import Julia - - try: - Julia(**julia_kwargs) - except UnsupportedPythonError: - # Static python binary, so we turn off pre-compiled modules. - julia_kwargs = {**julia_kwargs, "compiled_modules": False} - Julia(**julia_kwargs) - warnings.warn( - "Your system's Python library is static (e.g., conda), so precompilation will be turned off. For a dynamic library, try using `pyenv` and installing with `--enable-shared`: https://github.com/pyenv/pyenv/blob/master/plugins/python-build/README.md#building-with---enable-shared." - ) - - using_compiled_modules = (not "compiled_modules" in julia_kwargs) or julia_kwargs[ - "compiled_modules" - ] - - from julia import Main as _Main - - Main = _Main - - if julia_activated_env is None: - julia_activated_env = processed_julia_project - - if julia_initialized and julia_kwargs_at_initialization is not None: - # Check if the kwargs are the same as the previous initialization - init_set = set(julia_kwargs_at_initialization.items()) - new_set = set(julia_kwargs.items()) - set_diff = new_set - init_set - # Remove the `compiled_modules` key, since it is not a user-specified kwarg: - set_diff = {k: v for k, v in set_diff if k != "compiled_modules"} - if len(set_diff) > 0: - warnings.warn( - "Julia has already started. The new Julia options " - + str(set_diff) - + " will be ignored." - ) - - if julia_initialized and julia_activated_env != processed_julia_project: - Main.eval("using Pkg") - - io_arg = _get_io_arg(quiet) - # Can't pass IO to Julia call as it evaluates to PyObject, so just directly - # use Main.eval: - Main.eval( - f'Pkg.activate("{_escape_filename(processed_julia_project)}",' - f"shared = Bool({int(is_shared)}), " - f"{io_arg})" - ) - - julia_activated_env = processed_julia_project - - if not julia_initialized: - julia_kwargs_at_initialization = julia_kwargs - - julia_initialized = True - if return_aux: - return Main, {"compiled_modules": using_compiled_modules} - return Main - - -def _add_sr_to_julia_project(Main, io_arg): - Main.eval("using Pkg") - Main.eval("Pkg.Registry.update()") - Main.sr_spec = Main.PackageSpec( - name="SymbolicRegression", - url="https://github.com/MilesCranmer/SymbolicRegression.jl", - rev="v" + __symbolic_regression_jl_version__, - ) - Main.clustermanagers_spec = Main.PackageSpec( - name="ClusterManagers", - version="0.4", - ) - Main.eval(f"Pkg.add([sr_spec, clustermanagers_spec], {io_arg})") + return f"io={io}" def _escape_filename(filename): @@ -288,60 +24,17 @@ def _escape_filename(filename): return str_repr -def _julia_version_assertion(): - if not is_julia_version_greater_eq(version=(1, 6, 0)): - raise NotImplementedError( - "PySR requires Julia 1.6.0 or greater. " - "Please update your Julia installation." - ) - - -def _backend_version_assertion(Main): - try: - backend_version = Main.eval("string(SymbolicRegression.PACKAGE_VERSION)") - expected_backend_version = __symbolic_regression_jl_version__ - if backend_version != expected_backend_version: # pragma: no cover - warnings.warn( - f"PySR backend (SymbolicRegression.jl) version {backend_version} " - f"does not match expected version {expected_backend_version}. " - "Things may break. " - "Please update your PySR installation with " - "`python3 -m pysr install`." - ) - except JuliaError: # pragma: no cover +def _backend_version_assertion(): + backend_version = jl.seval("string(SymbolicRegression.PACKAGE_VERSION)") + expected_backend_version = juliapkg.status(target="SymbolicRegression").version + if backend_version != expected_backend_version: # pragma: no cover warnings.warn( - "You seem to have an outdated version of SymbolicRegression.jl. " + f"PySR backend (SymbolicRegression.jl) version {backend_version} " + f"does not match expected version {expected_backend_version}. " "Things may break. " - "Please update your PySR installation with " - "`python3 -m pysr install`." ) -def _load_cluster_manager(Main, cluster_manager): - Main.eval(f"import ClusterManagers: addprocs_{cluster_manager}") - return Main.eval(f"addprocs_{cluster_manager}") - - -def _update_julia_project(Main, is_shared, io_arg): - try: - if is_shared: - _add_sr_to_julia_project(Main, io_arg) - Main.eval("using Pkg") - Main.eval(f"Pkg.resolve({io_arg})") - except (JuliaError, RuntimeError) as e: - raise ImportError(_import_error()) from e - - -def _load_backend(Main): - try: - # Load namespace, so that various internal operators work: - Main.eval("using SymbolicRegression") - except (JuliaError, RuntimeError) as e: - raise ImportError(_import_error()) from e - - _backend_version_assertion(Main) - - # Load Julia package SymbolicRegression.jl - from julia import SymbolicRegression - - return SymbolicRegression +def _load_cluster_manager(cluster_manager): + jl.seval(f"using ClusterManagers: addprocs_{cluster_manager}") + return jl.seval(f"addprocs_{cluster_manager}") diff --git a/pysr/sr.py b/pysr/sr.py index 8bddece19..3e23a95e4 100644 --- a/pysr/sr.py +++ b/pysr/sr.py @@ -32,15 +32,7 @@ from .export_sympy import assert_valid_sympy_symbol, create_sympy_symbols, pysr2sympy from .export_torch import sympy2torch from .feature_selection import run_feature_selection -from .julia_helpers import ( - _escape_filename, - _load_backend, - _load_cluster_manager, - _process_julia_project, - _update_julia_project, - init_julia, - is_julia_version_greater_eq, -) +from .julia_helpers import _escape_filename, _load_cluster_manager, jl from .utils import ( _csv_filename_to_pkl_filename, _preprocess_julia_floats, @@ -48,8 +40,6 @@ _subscriptify, ) -Main = None # TODO: Rename to more descriptive name like "julia_runtime" - already_ran = False @@ -92,7 +82,6 @@ def _process_constraints(binary_operators, unary_operators, constraints): def _maybe_create_inline_operators( binary_operators, unary_operators, extra_sympy_mappings ): - global Main binary_operators = binary_operators.copy() unary_operators = unary_operators.copy() for op_list in [binary_operators, unary_operators]: @@ -100,7 +89,7 @@ def _maybe_create_inline_operators( is_user_defined_operator = "(" in op if is_user_defined_operator: - Main.eval(op) + jl.seval(op) # Cut off from the first non-alphanumeric char: first_non_char = [j for j, char in enumerate(op) if char == "("][0] function_name = op[:first_non_char] @@ -1528,7 +1517,6 @@ def _run(self, X, y, mutated_params, weights, seed): # Need to be global as we don't want to recreate/reinstate julia for # every new instance of PySRRegressor global already_ran - global Main # These are the parameters which may be modified from the ones # specified in init, so we define them here locally: @@ -1549,26 +1537,17 @@ def _run(self, X, y, mutated_params, weights, seed): if not already_ran and update_verbosity != 0: print("Compiling Julia backend...") - Main = init_julia(self.julia_project, julia_kwargs=julia_kwargs) - if cluster_manager is not None: - cluster_manager = _load_cluster_manager(Main, cluster_manager) + cluster_manager = _load_cluster_manager(cluster_manager) - if self.update: - _, is_shared = _process_julia_project(self.julia_project) - io = "devnull" if update_verbosity == 0 else "stderr" - io_arg = ( - f"io={io}" if is_julia_version_greater_eq(version=(1, 6, 0)) else "" - ) - _update_julia_project(Main, is_shared, io_arg) - - SymbolicRegression = _load_backend(Main) + jl.seval("using SymbolicRegression") + SymbolicRegression = jl.SymbolicRegression - Main.plus = Main.eval("(+)") - Main.sub = Main.eval("(-)") - Main.mult = Main.eval("(*)") - Main.pow = Main.eval("(^)") - Main.div = Main.eval("(/)") + jl.plus = jl.seval("(+)") + jl.sub = jl.seval("(-)") + jl.mult = jl.seval("(*)") + jl.pow = jl.seval("(^)") + jl.div = jl.seval("(/)") # TODO(mcranmer): These functions should be part of this class. binary_operators, unary_operators = _maybe_create_inline_operators( @@ -1594,7 +1573,7 @@ def _run(self, X, y, mutated_params, weights, seed): nested_constraints_str += f"({inner_k}) => {inner_v}, " nested_constraints_str += "), " nested_constraints_str += ")" - nested_constraints = Main.eval(nested_constraints_str) + nested_constraints = jl.seval(nested_constraints_str) # Parse dict into Julia Dict for complexities: if complexity_of_operators is not None: @@ -1602,13 +1581,17 @@ def _run(self, X, y, mutated_params, weights, seed): for k, v in complexity_of_operators.items(): complexity_of_operators_str += f"({k}) => {v}, " complexity_of_operators_str += ")" - complexity_of_operators = Main.eval(complexity_of_operators_str) + complexity_of_operators = jl.seval(complexity_of_operators_str) - custom_loss = Main.eval(self.loss) - custom_full_objective = Main.eval(self.full_objective) + custom_loss = jl.seval(str(self.loss) if self.loss is not None else "nothing") + custom_full_objective = jl.seval( + str(self.full_objective) if self.full_objective is not None else "nothing" + ) - early_stop_condition = Main.eval( - str(self.early_stop_condition) if self.early_stop_condition else None + early_stop_condition = jl.seval( + str(self.early_stop_condition) + if self.early_stop_condition is not None + else "nothing" ) mutation_weights = SymbolicRegression.MutationWeights( @@ -1626,9 +1609,10 @@ def _run(self, X, y, mutated_params, weights, seed): # Call to Julia backend. # See https://github.com/MilesCranmer/SymbolicRegression.jl/blob/master/src/OptionsStruct.jl + print(bin_constraints) options = SymbolicRegression.Options( - binary_operators=Main.eval(str(binary_operators).replace("'", "")), - unary_operators=Main.eval(str(unary_operators).replace("'", "")), + binary_operators=jl.seval(str(binary_operators).replace("'", "")), + unary_operators=jl.seval(str(unary_operators).replace("'", "")), bin_constraints=bin_constraints, una_constraints=una_constraints, complexity_of_operators=complexity_of_operators, diff --git a/pysr/version.py b/pysr/version.py index d88a664ec..fd86b3ee9 100644 --- a/pysr/version.py +++ b/pysr/version.py @@ -1,2 +1 @@ -__version__ = "0.16.9" -__symbolic_regression_jl_version__ = "0.23.1" +__version__ = "0.17.0" diff --git a/requirements.txt b/requirements.txt index ffb3e0d63..435e481ad 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,7 +2,8 @@ sympy>=1.0.0,<2.0.0 pandas>=0.21.0,<3.0.0 numpy>=1.13.0,<2.0.0 scikit_learn>=1.0.0,<2.0.0 -julia>=0.6.0,<0.7.0 +juliacall>=0.9.15,<0.10.0 +juliapkg>=0.1.10,<0.2.0 click>=7.0.0,<9.0.0 setuptools>=50.0.0 typing_extensions>=4.0.0,<5.0.0; python_version < "3.8"