From 3af7c4f663932235f8bf2552cbac39319b80dc0a Mon Sep 17 00:00:00 2001 From: Lilith Orion Hafner Date: Wed, 6 Dec 2023 08:25:25 -0600 Subject: [PATCH 1/2] Make package installation automatic and deprecate manual installation --- README.md | 38 +++------- diffeqpy/__init__.py | 79 ++++++++++++++------- diffeqpy/amdgpu.py | 2 +- diffeqpy/cuda.py | 2 +- diffeqpy/de.py | 2 +- diffeqpy/{ => deprecated}/install.jl | 0 diffeqpy/{ => deprecated}/install_amdgpu.jl | 0 diffeqpy/{ => deprecated}/install_cuda.jl | 0 diffeqpy/{ => deprecated}/install_metal.jl | 0 diffeqpy/{ => deprecated}/install_oneapi.jl | 0 diffeqpy/metal.py | 2 +- diffeqpy/oneapi.py | 2 +- 12 files changed, 68 insertions(+), 59 deletions(-) rename diffeqpy/{ => deprecated}/install.jl (100%) rename diffeqpy/{ => deprecated}/install_amdgpu.jl (100%) rename diffeqpy/{ => deprecated}/install_cuda.jl (100%) rename diffeqpy/{ => deprecated}/install_metal.jl (100%) rename diffeqpy/{ => deprecated}/install_oneapi.jl (100%) diff --git a/README.md b/README.md index 846ba0a..881456e 100644 --- a/README.md +++ b/README.md @@ -31,14 +31,6 @@ To install diffeqpy, use pip: pip install diffeqpy ``` -To install Julia packages required (and Julia if needed) for diffeqpy, open up Python -interpreter then run: - -```pycon ->>> import diffeqpy ->>> diffeqpy.install() -``` - and you're good! ## Collab Notebook Examples @@ -171,7 +163,7 @@ sol = de.solve(prob) #### Limitations `de.jit`, uses ModelingToolkit.jl's `modelingtoolkitize` internally and some -restrictions apply. Not all models can be jitted. See the +restrictions apply. Not all models can be jitted. See the [`modelingtoolkitize` documentation](https://docs.sciml.ai/ModelingToolkit/stable/tutorials/modelingtoolkitize/#What-is-modelingtoolkitize?) for more info. @@ -552,29 +544,19 @@ sol = de.solve(ensembleprob,de.Tsit5(),de.EnsembleSerial(),trajectories=10000,sa ``` To add GPUs to the mix, we need to bring in [DiffEqGPU](https://github.com/SciML/DiffEqGPU.jl). -The `diffeqpy.install_cuda()` will install CUDA for you and bring all of the bindings into the returned object: +The command `from diffeqpy import cuda` will install CUDA for you and bring all of the bindings into the returned object: -```py -diffeqpy.install_cuda() -``` - -then run the cuda import: - -```py -from diffeqpy import cuda -``` - -#### Note: `diffeqpy.install_cuda()` and `from diffeqpy import cuda` can take awhile to run the first time as it installs the drivers! +#### Note: `from diffeqpy import cuda` can take awhile to run the first time as it installs the drivers! Now we simply use `EnsembleGPUKernel(cuda.CUDABackend())` with a -GPU-specialized ODE solver `cuda.GPUTsit5()` to solve 10,000 ODEs on the GPU in +GPU-specialized ODE solver `cuda.GPUTsit5()` to solve 10,000 ODEs on the GPU in parallel: ```py sol = de.solve(ensembleprob,cuda.GPUTsit5(),cuda.EnsembleGPUKernel(cuda.CUDABackend()),trajectories=10000,saveat=0.01) ``` -For the full list of choices for specialized GPU solvers, see +For the full list of choices for specialized GPU solvers, see [the DiffEqGPU.jl documentation](https://docs.sciml.ai/DiffEqGPU/stable/manual/ensemblegpukernel/). Note that `EnsembleGPUArray` can be used as well, like: @@ -598,20 +580,20 @@ ensemble generation: ```py import numpy as np from scipy.integrate import odeint - + def lorenz(state, t, sigma, beta, rho): x, y, z = state - + dx = sigma * (y - x) dy = x * (rho - z) - y dz = x * y - beta * z - + return [dx, dy, dz] sigma = 10.0 beta = 8.0 / 3.0 -rho = 28.0 -p = (sigma, beta, rho) +rho = 28.0 +p = (sigma, beta, rho) y0 = [1.0, 1.0, 1.0] t = np.arange(0.0, 100.0, 0.01) diff --git a/diffeqpy/__init__.py b/diffeqpy/__init__.py index fa4efde..34dd0ee 100644 --- a/diffeqpy/__init__.py +++ b/diffeqpy/__init__.py @@ -1,17 +1,58 @@ -import os import shutil -import subprocess -import sys - from jill.install import install_julia -script_dir = os.path.dirname(os.path.realpath(__file__)) - +# juliacall must be loaded after `_ensure_julia_installed()` is run, +# so this import is in `load_julia_packages()` +# from juliacall import Main def _find_julia(): # TODO: this should probably fallback to query jill return shutil.which("julia") +def _ensure_julia_installed(): + if not _find_julia(): + print("No Julia version found. Installing Julia.") + install_julia() + if not _find_julia(): + raise RuntimeError( + "Julia installed with jill but `julia` binary cannot be found in the path" + ) + +# TODO: upstream this function or an alternative into juliacall +def load_julia_packages(*names): + """ + Load Julia packages and return references to them, automatically installing julia and + the packages as necessary. + """ + # This is terrifying to many people. However, it seems SciML takes pragmatic approach. + _ensure_julia_installed() + + script = """import Pkg + Pkg.activate(\"diffeqpy\", shared=true) + try + import {0} + catch e + e isa ArgumentError || throw(e) + Pkg.add([{1}]) + import {0} + end + {0}""".format(", ".join(names), ", ".join(f'"{name}"' for name in names)) + + # Unfortunately, `seval` doesn't support multi-line strings + # https://github.com/JuliaPy/PythonCall.jl/issues/433 + script = script.replace("\n", ";") + + # Must be loaded after `_ensure_julia_installed()` + from juliacall import Main + return Main.seval(script) + + + +# Deprecated (julia and packages now auto-install) +import os +import subprocess +import sys +script_dir = os.path.dirname(os.path.realpath(__file__)) def install(*, confirm=False): """ @@ -28,7 +69,7 @@ def install(*, confirm=False): ) env = os.environ.copy() env["PYTHON"] = sys.executable - subprocess.check_call([julia, os.path.join(script_dir, "install.jl")], env=env) + subprocess.check_call([julia, os.path.join(script_dir, "deprecated/install.jl")], env=env) def install_cuda(): julia = _find_julia() @@ -38,7 +79,7 @@ def install_cuda(): ) env = os.environ.copy() env["PYTHON"] = sys.executable - subprocess.check_call([julia, os.path.join(script_dir, "install_cuda.jl")], env=env) + subprocess.check_call([julia, os.path.join(script_dir, "deprecated/install_cuda.jl")], env=env) def install_amdgpu(): julia = _find_julia() @@ -48,17 +89,17 @@ def install_amdgpu(): ) env = os.environ.copy() env["PYTHON"] = sys.executable - subprocess.check_call([julia, os.path.join(script_dir, "install_amdgpu.jl")], env=env) + subprocess.check_call([julia, os.path.join(script_dir, "deprecated/install_amdgpu.jl")], env=env) def install_metal(): julia = _find_julia() if not julia: raise RuntimeError( - "Julia must be installed before adding Metal. Please run `diffeqpy.install()` first" + "Julia must be installed before adding Metal. Please run `deprecated/diffeqpy.install()` first" ) env = os.environ.copy() env["PYTHON"] = sys.executable - subprocess.check_call([julia, os.path.join(script_dir, "install_metal.jl")], env=env) + subprocess.check_call([julia, os.path.join(script_dir, "deprecated/install_metal.jl")], env=env) def install_oneapi(): julia = _find_julia() @@ -68,18 +109,4 @@ def install_oneapi(): ) env = os.environ.copy() env["PYTHON"] = sys.executable - subprocess.check_call([julia, os.path.join(script_dir, "install_oneapi.jl")], env=env) - -def _ensure_installed(*kwargs): - if not _find_julia(): - # TODO: this should probably ensure that packages are installed too - install(*kwargs) - -# TODO: upstream this function or an alternative into juliacall -def load_julia_packages(names): - # This is terrifying to many people. However, it seems SciML takes pragmatic approach. - _ensure_installed() - - # Must be loaded after `_ensure_installed()` - from juliacall import Main - return Main.seval(f"import Pkg; Pkg.activate(\"diffeqpy\", shared=true); import {names}; {names}") + subprocess.check_call([julia, os.path.join(script_dir, "deprecated/install_oneapi.jl")], env=env) diff --git a/diffeqpy/amdgpu.py b/diffeqpy/amdgpu.py index f6e0bb2..e4b7a7b 100644 --- a/diffeqpy/amdgpu.py +++ b/diffeqpy/amdgpu.py @@ -1,6 +1,6 @@ import sys from . import load_julia_packages -amdgpu, _ = load_julia_packages("DiffEqGPU, AMDGPU") +amdgpu, _ = load_julia_packages("DiffEqGPU", "AMDGPU") from juliacall import Main amdgpu.AMDGPUBackend = Main.seval("AMDGPU.AMDGPUBackend") # kinda hacky sys.modules[__name__] = amdgpu # mutate myself diff --git a/diffeqpy/cuda.py b/diffeqpy/cuda.py index 4cb0365..a5e55e9 100644 --- a/diffeqpy/cuda.py +++ b/diffeqpy/cuda.py @@ -1,6 +1,6 @@ import sys from . import load_julia_packages -cuda, _ = load_julia_packages("DiffEqGPU, CUDA") +cuda, _ = load_julia_packages("DiffEqGPU", "CUDA") from juliacall import Main cuda.CUDABackend = Main.seval("CUDA.CUDABackend") # kinda hacky sys.modules[__name__] = cuda # mutate myself diff --git a/diffeqpy/de.py b/diffeqpy/de.py index b049fa8..ae787c7 100644 --- a/diffeqpy/de.py +++ b/diffeqpy/de.py @@ -1,6 +1,6 @@ import sys from . import load_julia_packages -de, _ = load_julia_packages("DifferentialEquations, ModelingToolkit") +de, _ = load_julia_packages("DifferentialEquations", "ModelingToolkit") from juliacall import Main de.jit = Main.seval("jit(x) = typeof(x).name.wrapper(ModelingToolkit.modelingtoolkitize(x), x.u0, x.tspan, x.p)") # kinda hackey de.jit32 = Main.seval("jit(x) = typeof(x).name.wrapper(ModelingToolkit.modelingtoolkitize(x), Float32.(x.u0), Float32.(x.tspan), Float32.(x.p))") # kinda hackey diff --git a/diffeqpy/install.jl b/diffeqpy/deprecated/install.jl similarity index 100% rename from diffeqpy/install.jl rename to diffeqpy/deprecated/install.jl diff --git a/diffeqpy/install_amdgpu.jl b/diffeqpy/deprecated/install_amdgpu.jl similarity index 100% rename from diffeqpy/install_amdgpu.jl rename to diffeqpy/deprecated/install_amdgpu.jl diff --git a/diffeqpy/install_cuda.jl b/diffeqpy/deprecated/install_cuda.jl similarity index 100% rename from diffeqpy/install_cuda.jl rename to diffeqpy/deprecated/install_cuda.jl diff --git a/diffeqpy/install_metal.jl b/diffeqpy/deprecated/install_metal.jl similarity index 100% rename from diffeqpy/install_metal.jl rename to diffeqpy/deprecated/install_metal.jl diff --git a/diffeqpy/install_oneapi.jl b/diffeqpy/deprecated/install_oneapi.jl similarity index 100% rename from diffeqpy/install_oneapi.jl rename to diffeqpy/deprecated/install_oneapi.jl diff --git a/diffeqpy/metal.py b/diffeqpy/metal.py index 80913e6..c9c85fe 100644 --- a/diffeqpy/metal.py +++ b/diffeqpy/metal.py @@ -1,6 +1,6 @@ import sys from . import load_julia_packages -metal, _ = load_julia_packages("DiffEqGPU, Metal") +metal, _ = load_julia_packages("DiffEqGPU", "Metal") from juliacall import Main metal.MetalBackend = Main.seval("Metal.MetalBackend") # kinda hacky sys.modules[__name__] = metal # mutate myself \ No newline at end of file diff --git a/diffeqpy/oneapi.py b/diffeqpy/oneapi.py index 5033b11..abbea52 100644 --- a/diffeqpy/oneapi.py +++ b/diffeqpy/oneapi.py @@ -1,6 +1,6 @@ import sys from . import load_julia_packages -oneapi, _ = load_julia_packages("DiffEqGPU, oneAPI") +oneapi, _ = load_julia_packages("DiffEqGPU", "oneAPI") from juliacall import Main oneapi.oneAPIBackend = Main.seval("oneAPI.oneAPIBackend") # kinda hacky sys.modules[__name__] = oneapi # mutate myself From b2f6e7a0c536fe19ae498e90758f5709b34360ba Mon Sep 17 00:00:00 2001 From: Lilith Orion Hafner Date: Wed, 6 Dec 2023 09:15:38 -0600 Subject: [PATCH 2/2] Still load PythonCall --- diffeqpy/de.py | 2 +- diffeqpy/ode.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/diffeqpy/de.py b/diffeqpy/de.py index ae787c7..d201365 100644 --- a/diffeqpy/de.py +++ b/diffeqpy/de.py @@ -1,6 +1,6 @@ import sys from . import load_julia_packages -de, _ = load_julia_packages("DifferentialEquations", "ModelingToolkit") +de, _, _ = load_julia_packages("DifferentialEquations", "ModelingToolkit", "PythonCall") from juliacall import Main de.jit = Main.seval("jit(x) = typeof(x).name.wrapper(ModelingToolkit.modelingtoolkitize(x), x.u0, x.tspan, x.p)") # kinda hackey de.jit32 = Main.seval("jit(x) = typeof(x).name.wrapper(ModelingToolkit.modelingtoolkitize(x), Float32.(x.u0), Float32.(x.tspan), Float32.(x.p))") # kinda hackey diff --git a/diffeqpy/ode.py b/diffeqpy/ode.py index ce861b7..69f2301 100644 --- a/diffeqpy/ode.py +++ b/diffeqpy/ode.py @@ -1,3 +1,4 @@ import sys from . import load_julia_packages -sys.modules[__name__] = load_julia_packages("OrdinaryDiffEq") # mutate myself +ode, _ = load_julia_packages("OrdinaryDiffEq", "PythonCall") +sys.modules[__name__] = ode # mutate myself