Is it possible to implement numba for poisson solver operator?

[ArtyomGubin]

Hello!
I've been trying to solve this system in 2D:

On each time step firstly I need to solve biharmonic equation for $$\Psi$$, so I thought to use poisson solver twice on right hand side of the first PDE. Numpy implementation was too slow, so I've tried to impement numba:

from math import pi
from pde import PDE, CartesianGrid, MemoryStorage, ScalarField, plot_kymograph
from pde import FieldCollection, PDEBase, PlotTracker, ScalarField, UnitGrid
# HERE I IMPORT THIS METHOD
from pde.grids.operators.cartesian import make_poisson_solver
import pde
import numpy as np
import numba as nb

class myPDE(PDEBase):
    def __init__(
        self, Pe=1, Le=0.25, Cp=10, bc_phi="neumann", bc_Temp=[{"value": 2/3}, {"derivative": 0}]
        ):
        super().__init__()
        self.Pe = Pe # Pecke number
        self.Le = Le # Lewis number
        self.Cp = Cp # Heat capacity
        self.bc_phi = bc_phi # boundary condition for concentration
        self.bc_Temp = bc_Temp # boundary condition for temperature
        
    def evolution_rate(self, state, t=0):
        phi, T = state
        
        # Find stream function
        dphidx = phi.gradient(bc={'value':0}).to_scalar(0)
        dphidy = phi.gradient(bc={'value':0}).to_scalar(1)
        dfdx = ((1/T)*phi.laplace(self.bc_phi)).gradient(bc={'value':0}).to_scalar(0)
        dfdy = ((1/T)*phi.laplace(self.bc_phi)).gradient(bc={'value':0}).to_scalar(0)
        
        rhs = -2 * (
            dfdx * dphidy - dfdy*dphidx
            )
        
        psi = pde.solve_poisson_equation(rhs, {"value":0})
        psi = pde.solve_poisson_equation(psi, {"value":0})
        
        # Find velocity field
        vx = psi.gradient(bc={'value':0}).to_scalar(1)
        vy = -psi.gradient(bc={'value':0}).to_scalar(0)
        v_vector = pde.VectorField(state.grid, [vx.data,vy.data])
        
        # Mass flux
        J = (-phi.gradient(self.bc_phi) 
             + phi*(1-phi) * ( (4/T)*phi.gradient(self.bc_phi) 
                              + (2/T)*phi.laplace(self.bc_phi).gradient(self.bc_phi)))
        
        # Rate
        dT_dt = (
            - self.Pe * v_vector @ T.gradient(self.bc_Temp) 
            + self.Le*T.laplace(self.bc_Temp) 
            + (4/self.Cp) * phi.gradient(self.bc_phi) @ J
            )
        dphi_dt = (
            - self.Pe * v_vector @ phi.gradient(self.bc_phi)
            - J.divergence(self.bc_phi)
            )
        return FieldCollection([dphi_dt, dT_dt])
    
    def _make_pde_rhs_numba(self, state):
        """Nunmba-compiled implementation of the PDE."""
        Pe = self.Pe # Pecke number
        Le = self.Le # Lewis number
        Cp = self.Cp # Heat capacity
        bc_phi = self.bc_phi # boundary condition for concentration
        bc_Temp = self.bc_Temp # boundary condition for temperature
        
        laplace_phi = state.grid.make_operator("laplace", bc=self.bc_phi)
        laplace_T = state.grid.make_operator("laplace", bc=self.bc_Temp)
        gradient_phi = state.grid.make_operator("gradient", bc=self.bc_phi)
        gradient_T = state.grid.make_operator("gradient", bc=self.bc_Temp)
        d_dx = state.grid.make_operator('d_dx', bc="value")
        d_dy = state.grid.make_operator('d_dy', bc="value")
        d_dx_phi = state.grid.make_operator('d_dx', bc=self.bc_phi)
        d_dy_phi = state.grid.make_operator('d_dy', bc=self.bc_phi)
        dot = pde.VectorField(state.grid).make_dot_operator()
        div = state.grid.make_operator("divergence", bc=self.bc_phi)
        
        # AND USE IT HERE
        ddsolve = make_poisson_solver(state.grid)
        
        @pde.tools.numba.jit
        def pde_rhs(state, t):
            phi, T = state
            
            # Stream function
            dphidx = d_dx_phi(phi)
            dphidy = d_dy_phi(phi)
            dfdx = d_dx((1/T)*phi.laplace(self.bc_phi))
            dfdy = d_dy((1/T)*phi.laplace(self.bc_phi))
            
            rhs = -2 * (
                dfdx * dphidy - dfdy*dphidx
                )
            
            psi = ddsolve(rhs)
            psi = ddsolve(psi)
            
            # Velocity field
            vx = d_dy(psi)
            vy = -d_dx(psi)
            v_vector = pde.VectorField(state.grid, [vx.data,vy.data])
            
            # Mass flux
            J = (-gradient_phi(phi)
                 + phi*(1-phi) * ( (4/T) * gradient_phi(phi) 
                                  + (2/T)*gradient_phi(laplace_phi(phi)))
                 )
            
            # PDE
            dT_dt = (
                - self.Pe * dot(v_vector, gradient_T(T))
                + self.Le * laplace_T(T)
                + (4/self.Cp) * dot(gradient_phi(phi), J)
                )
            
            dphi_dt = (
                - self.Pe * dot(v_vector, gradient_phi(phi))
                - div(J)
                )
            return FieldCollection([dphi_dt, dT_dt])
        
        return pde_rhs

eq = myPDE()
grid = CartesianGrid([[dx, L], [dx,L]], [Nx])
phi = ScalarField.random_uniform(grid, 0.495, 0.505)
T = ScalarField(grid, 4/3)
state = FieldCollection([phi, T])
storage = MemoryStorage()
trackers = ["progress", 
            storage.tracker(dt_save),
            pde.PrintTracker(interrupts=pde.RealtimeInterrupts(duration=5))
            ]
sol = eq.solve(state, t_range=t_end, dt=1e-4, solver='explicit', tracker=trackers)

But when I try to start calculations I get this error:

  File ~\anaconda3\envs\py-pde_env\lib\site-packages\spyder_kernels\py3compat.py:356 in compat_exec
    exec(code, globals, locals)

  File c:\users\lenovo\onedrive\рабочий стол\диплом\источники\система pde с уравнением к-х\pypde_segal_2d.py:168
    sol = eq.solve(state, t_range=t_end, dt=1e-4, solver='explicit', tracker=trackers)

  File ~\anaconda3\envs\py-pde_env\lib\site-packages\pde\pdes\base.py:654 in solve
    final_state = controller.run(state, dt)

  File ~\anaconda3\envs\py-pde_env\lib\site-packages\pde\solvers\controller.py:409 in run
    return self._run_serial(state, dt)

  File ~\anaconda3\envs\py-pde_env\lib\site-packages\pde\solvers\controller.py:324 in _run_serial
    self._run_main_process(state, dt)

  File ~\anaconda3\envs\py-pde_env\lib\site-packages\pde\solvers\controller.py:179 in _run_main_process
    stepper = self.solver.make_stepper(state=state, dt=dt)

  File ~\anaconda3\envs\py-pde_env\lib\site-packages\pde\solvers\base.py:665 in make_stepper
    return super().make_stepper(state, dt)

  File ~\anaconda3\envs\py-pde_env\lib\site-packages\pde\solvers\base.py:389 in make_stepper
    fixed_stepper = self._make_fixed_stepper(state, dt_float)

  File ~\anaconda3\envs\py-pde_env\lib\site-packages\pde\solvers\base.py:329 in _make_fixed_stepper
    single_step = self._make_single_step_fixed_dt(state, dt)

  File ~\anaconda3\envs\py-pde_env\lib\site-packages\pde\solvers\explicit.py:151 in _make_single_step_fixed_dt
    return self._make_single_step_fixed_euler(state, dt)

  File ~\anaconda3\envs\py-pde_env\lib\site-packages\pde\solvers\explicit.py:91 in _make_single_step_fixed_euler
    rhs_pde = self._make_pde_rhs(state, backend=self.backend)

  File ~\anaconda3\envs\py-pde_env\lib\site-packages\pde\solvers\base.py:258 in _make_pde_rhs
    rhs = self.pde.make_pde_rhs(state, backend=backend)  # type: ignore

  File ~\anaconda3\envs\py-pde_env\lib\site-packages\pde\pdes\base.py:306 in make_pde_rhs
    rhs = self._make_pde_rhs_numba_cached(state, **kwargs)

  File ~\anaconda3\envs\py-pde_env\lib\site-packages\pde\pdes\base.py:275 in _make_pde_rhs_numba_cached
    rhs = self._make_pde_rhs_numba(state, **kwargs)

  File c:\users\lenovo\onedrive\рабочий стол\диплом\источники\система pde с уравнением к-х\pypde_segal_2d.py:100 in _make_pde_rhs_numba
    ddsolve = make_poisson_solver(state.grid)

  File ~\anaconda3\envs\py-pde_env\lib\site-packages\pde\grids\operators\cartesian.py:1303 in make_poisson_solver
    matrix, vector = _get_laplace_matrix(bcs)

  File ~\anaconda3\envs\py-pde_env\lib\site-packages\pde\grids\operators\cartesian.py:248 in _get_laplace_matrix
    dim = bcs.grid.dim

AttributeError: 'CartesianGrid' object has no attribute 'grid'

What am I doing wrong?

[david-zwicker]

Unfortunately, py-pde is currently not very well suited for these kind of problems. The poisson is not very fast and might also not be particularly stable. In particular, it currently does not support numba.