Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Make package installation automatic and deprecate manual installation #132

Merged
merged 2 commits into from
Dec 10, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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
Loading