User defined function inside numba accelerated PDE
SalvadorBrandolin opened this issue · comments
Can't resolve a system of partial differential equations with numba acelleration.
Going to solve PDE system
dCa/dt = - Fv * dCa/dV - k * Ca
dCb/dt = - Fv * dCb/dV + k * Ca
where the term k * Ca comes from a pre-defined user function.
The PDE system is implemented using a custom PDE class.
The resolution of the problem using the method evolution_rate(self, state, t=0) works (Code below).
Turn reactor.resolve(Numba=False), at the end, to not use numba.
Thank you very much! :)
from pde import (PDEBase, FieldCollection, CartesianGrid, ScalarField,
ImplicitSolver, Controller
)
from pde.tools.numba import jit
from numba import njit
import matplotlib.pyplot as plt
from pde import config
config['numba.debug'] = True
# =============================================================================
# Clases
# -----------------------------------------------------------------------------
class ReactiveSystem:
def __init__(self, ks, inflows, reactions):
self.ks = ks
self.inflows = inflows
self.reactions = reactions
self.net_flow = sum(inflows)
@jit
def react(self, state):
return [
r(self.ks, state) for r in self.reactions
]
def __len__(self):
return len(self.reactions)
class ReactorSolver(PDEBase):
def __init__(
self, reactive_system, V, bcs,
n_points=10, nt=10, t_end=10, t=0, max_it=1000, tol=1e-5
):
# Number of reactives (hence, balances)
self.n_reactives = len(reactive_system)
# System dimension
grid = CartesianGrid([[0, V]], n_points)
states = [
ScalarField(grid, data=0)
for i in range(self.n_reactives)
]
self.state = FieldCollection(states)
# Border conditions
self.bcs = bcs
# Reactive system
self.reactive_system = reactive_system
# Time ranges
self.nt = nt
self.t_end = t_end
self.dt = t_end/nt
# Solver tolerance
self.max_iter = max_it
self.tol = tol
def evolution_rate(self, state, t=0):
# Agregué el índice cero en el gradiente, hay que indicar que es
# la derivada en X más allá de que sea la única variable
# Calculate reaction rates
n_reactions = len(self.reactive_system)
ks = self.reactive_system.ks
reactions = [
react(ks, state) for react in self.reactive_system.reactions
]
# Net flow (of course it shouldn't be like this)
net_flow = self.reactive_system.net_flow
# Whole balances
system = [
- net_flow * state[i].gradient(self.bcs[i]) # Flow balance
+ reactions[i] # Reactions
for i in range(self.n_reactives)
]
return FieldCollection(system)
def _make_pde_rhs_numba(self, state):
#attributes locally available
#Net flow (this is bad implemented - supposed to be vol flow)
net_flow = self.reactive_system.net_flow
#Number of reactives
n_reactives = len(self.reactive_system)
#kinetic constans
ks = self.reactive_system.ks
#Reaction function of each compound
reactions = [
react(ks, state).data for react in self.reactive_system.reactions
]
#Operands definition
gradient_c = [
state.grid.make_operator("gradient", bc=self.bcs[i]) for i
in range(n_reactives)
]
@jit
def pde_rhs(state_data, t=0):
""" compiled helper function evaluating right hand side """
state_grad = []
for i in range(n_reactives):
gradient = gradient_c[i]
state_grad.append(gradient(state_data[i]))
#Reaction function of each compound
"""system = [
- net_flow * state_grad[i] # Flow balance
+ reactions[i] # Reactions
for i in range(n_reactives)
]"""
system = []
for i in range(n_reactives):
system.append(- net_flow * state_grad[i])
return system
return FieldCollection(pde_rhs)
def resolve(self, numba = False):
if numba == False:
solver = ImplicitSolver(self, self.max_iter, self.tol,
backend='numpy')
controller = Controller(solver, t_range=self.t_end,
tracker=None)
return controller.run(self.state)
else:
solver = ImplicitSolver(self, self.max_iter, self.tol,
backend='numba')
controller = Controller(solver, t_range=self.t_end,
tracker=None)
return controller.run(self.state)
# =============================================================================
# =============================================================================
# Planteo de problema
# -----------------------------------------------------------------------------
# Variables generales
kr = 1
t0 = 0
tf = 10
V = 1
Fv = 1
Ca0 = 1
# Definicion de sistema reactivo
rs = ReactiveSystem(
ks=[kr],
inflows=[Fv, 0],
reactions=[
lambda ks, concentrations: -ks[0] * concentrations[0],
lambda ks, concentrations: ks[0] * concentrations[0],
]
)
# Border conditions
## - x0 = CaO
## - xf = dCa/dx = 0
bc_a = [
{"value": Ca0},
{"derivative": 0}
]
bc_b = [
{"value": 0},
{"derivative": 0}
]
# Reactor
reactor = ReactorSolver(
reactive_system=rs,
V=V,
bcs=[bc_a, bc_b],
n_points = 100,
t_end=10,
nt = 1000,
)
# ==============================================================================
results = reactor.resolve(numba=True)
plt.plot(results[0].data)
plt.plot(results[1].data)
plt.show()
I'm not sure what issue you're reporting. What is the expected behavior and what error are you seeing?
If you've a question about how to use py-pde, you might want to post in the discussion forum.
The error obtained from the code above is:
Traceback (most recent call last):
File "/mnt/g/Mi unidad/pde_testing/resolviendo.py", line 193, in
results = reactor.resolve(numba=True)
File "/mnt/g/Mi unidad/pde_testing/resolviendo.py", line 144, in resolve
return controller.run(self.state)
File "/home/salvadorbrandolin/.local/lib/python3.10/site-packages/pde/solvers/controller.py", line 164, in run
stepper = self.solver.make_stepper(state=state, dt=dt)
File "/home/salvadorbrandolin/.local/lib/python3.10/site-packages/pde/solvers/implicit.py", line 85, in make_stepper
rhs = self._make_pde_rhs(state, backend=self.backend)
File "/home/salvadorbrandolin/.local/lib/python3.10/site-packages/pde/solvers/base.py", line 117, in _make_pde_rhs
rhs = self.pde.make_pde_rhs(state, backend=backend)
File "/home/salvadorbrandolin/.local/lib/python3.10/site-packages/pde/pdes/base.py", line 239, in make_pde_rhs
rhs = self._make_pde_rhs_numba_cached(state)
File "/home/salvadorbrandolin/.local/lib/python3.10/site-packages/pde/pdes/base.py", line 206, in _make_pde_rhs_numba_cached
rhs = self._make_pde_rhs_numba(state)
File "/mnt/g/Mi unidad/pde_testing/resolviendo.py", line 130, in _make_pde_rhs_numba
return FieldCollection(pde_rhs)
File "/home/salvadorbrandolin/.local/lib/python3.10/site-packages/pde/fields/collection.py", line 68, in init
if len(fields) == 0:
TypeError: object of type 'CPUDispatcher' has no len()
the expected behavior is that the problem can be solved with numba accelerated PDE as explained in the documentation:
https://py-pde.readthedocs.io/en/latest/manual/advanced_usage.html#numba-accelerated-pdes
In the PDE custom class, the evolution_rate() method is defined to solve the PDE without numba, and works.
Also the method _make_pde_rhs_numba() is defined, and the numba solver doesn't work.
Thank you
This does not seem to be an issue in py-pde, but you're for some reason trying to construct a FieldCollection with a function (pde_rhs) as an argument on line 130. I don't know what you're trying to achieve and thus unfortunately cannot help you.
Ok, thank you