Skip to content

Commit

Permalink
Merge pull request #132 from LilithHafner/lh/install
Browse files Browse the repository at this point in the history
Make package installation automatic and deprecate manual installation
  • Loading branch information
ChrisRackauckas authored Dec 10, 2023
2 parents e280355 + b2f6e7a commit b448769
Show file tree
Hide file tree
Showing 13 changed files with 70 additions and 60 deletions.
38 changes: 10 additions & 28 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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:
Expand All @@ -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)
Expand Down
79 changes: 53 additions & 26 deletions diffeqpy/__init__.py
Original file line number Diff line number Diff line change
@@ -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):
"""
Expand All @@ -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()
Expand All @@ -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()
Expand All @@ -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()
Expand All @@ -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)
2 changes: 1 addition & 1 deletion diffeqpy/amdgpu.py
Original file line number Diff line number Diff line change
@@ -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
2 changes: 1 addition & 1 deletion diffeqpy/cuda.py
Original file line number Diff line number Diff line change
@@ -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
2 changes: 1 addition & 1 deletion diffeqpy/de.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down
File renamed without changes.
File renamed without changes.
File renamed without changes.
File renamed without changes.
File renamed without changes.
2 changes: 1 addition & 1 deletion diffeqpy/metal.py
Original file line number Diff line number Diff line change
@@ -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
3 changes: 2 additions & 1 deletion diffeqpy/ode.py
Original file line number Diff line number Diff line change
@@ -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
2 changes: 1 addition & 1 deletion diffeqpy/oneapi.py
Original file line number Diff line number Diff line change
@@ -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

0 comments on commit b448769

Please sign in to comment.