Solving coupled PDE with numba acceleration

[ScottDemarest]

I am attempting to solve a coupled system of two PDEs and I was able to do so by following the example provided in this repo. I then tried applying numba acceleration like in the docs by returning a tuple in the pde_rhs() function. But if I return a tuple with two entries, I get the following error:

TypeError: can't multiply sequence by non-int of type 'float'

In the case of a coupled system, what is pde_rhs() supposed to return?

I can provide sample code if necessary.

[david-zwicker]

It sounds as if you want to solve a coupled PDE, e.g., describing the dynamics of two fields, simultaneously. In the pure python version, you would use a FieldCollection (instead of two separate ScalarFields) to describe the state. The FieldCollection has a single numpy array that contains the field values of all fields. It is this numpy array that needs to be returned by your pde_rhs() function (and not a tuple). If you post concrete code, it might be easier to see where the problem might lie.

[ScottDemarest]

You are right that I want the dynamics of two scalar fields simultaneously. Below is my code. It works fine if I comment out the _make_pde_rhs_numba which is the case in the example below. But when I try using that method, that's where I get issues. I was able to get numba acceleration for simpler examples, but now that I am coupling two fields, that's where I am having trouble.

Thank you so much for your time by the way. I am just starting out with simulation and this has been a very good gateway into it.

from math import pi
from pde.tools.numba import jit
from pde import CartesianGrid, MemoryStorage, PDEBase, ScalarField, plot_kymograph, movie_scalar, PlotTracker, VectorField, FieldCollection
from pde import config
config["numba.multithreading_threshold"] = 64

name = "rmhd_pulses_long3"
x_lims = [0, 10]
num_points = [512]
initial_state_p = ".5*sin(2*pi*x/2)*exp(-((x-6)/3)**2)"
initial_state_m = ".5*cos(2*pi*x/5)*exp(-((x-4)/1)**2)"
t_range = 200
dt = .001
mu = 0.002
Va = 0#.5

class RMHD_1D(PDEBase):

    def evolution_rate(self, state, t=0):
        """implement the python version of the evolution equation"""
        zp, zm = state
        grad_x_p = zp.gradient("auto_periodic_neumann")[0]
        grad_x_m = zm.gradient("auto_periodic_neumann")[0]

        zpt = (Va - zm)*grad_x_p + mu*zp.laplace("auto_periodic_neumann")
        zmt = (-Va - zp)*grad_x_m + mu*zm.laplace("auto_periodic_neumann")

        return FieldCollection([zpt, zmt])
'''
    def _make_pde_rhs_numba(self, state):
        """ the numba-accelerated evolution equation """
        zp, zm = state
        # create operators
        laplace_u_p = zp.grid.make_operator("laplace", bc="auto_periodic_neumann")
        gradient_u_p = zp.grid.make_operator("gradient", bc="auto_periodic_neumann")
        laplace_u_m = zm.grid.make_operator("laplace", bc="auto_periodic_neumann")
        gradient_u_m = zm.grid.make_operator("gradient", bc="auto_periodic_neumann")

        @jit
        def pde_rhs(state_data, t=0):
            """ compiled helper function evaluating right hand side """
            zpd, zmd = state_data
            zpt = (Va-zmd)*gradient_u_p(zpd)[0] + mu*laplace_u_p(zpd)
            zmt = (-Va-zpd)*gradient_u_m(zmd)[0] + mu*laplace_u_m(zmd)
            
            return (zpt, zmt)

        return pde_rhs
'''

grid = CartesianGrid([x_lims], num_points, periodic=True)
storage = MemoryStorage()

zpi = ScalarField.from_expression(grid, initial_state_p)
zmi = ScalarField.from_expression(grid, initial_state_m)
state = FieldCollection([zpi, zmi])
eq = RMHD_1D()
#method = 'scipy'
eq.solve(state, t_range=t_range, dt = dt, tracker=["progress", storage.tracker(1.0)])

[david-zwicker]

Thanks for the code – here is an untested suggestion for the method using numba:

def _make_pde_rhs_numba(self, state):
        """ the numba-accelerated evolution equation """
        zp, zm = state
        # create operators
        laplace_u_p = zp.grid.make_operator("laplace", bc="auto_periodic_neumann")
        gradient_u_p = zp.grid.make_operator("gradient", bc="auto_periodic_neumann")
        laplace_u_m = zm.grid.make_operator("laplace", bc="auto_periodic_neumann")
        gradient_u_m = zm.grid.make_operator("gradient", bc="auto_periodic_neumann")

        @jit
        def pde_rhs(state_data, t=0):
            """ compiled helper function evaluating right hand side """
            zpd, zmd = state_data
            res = np.empty_like(state_data)
            res[0] = (Va-zmd)*gradient_u_p(zpd)[0] + mu*laplace_u_p(zpd)
            res[1] = (-Va-zpd)*gradient_u_m(zmd)[0] + mu*laplace_u_m(zmd)
            
            return res

        return pde_rhs

Alternatively, you could use the PDE class, e.g., like in this example.

[ScottDemarest]

It works! Thank you very much. For large grids and time spans, it dramatically speeds up my sims.

On a related note, if I wanted to add an explicitly time dependent term to my PDE, can I do that by simply adding a numpy array to the res[0] and res[1] terms?

[david-zwicker]

You can manipulate res like any other numpy array using all the functions that numba can compile. In particular, these expressions can involve the time t, which is available as a local variable.

[ScottDemarest]

Ok thank you. I will try it out later. I believe this discussion is closed. Thanks again