Long PDE input / using user_func

[gnorman7]

I'm working with a very long string input for a PDE, and I am wondering if there are any performance changes if I were to use user_func and evaluate the string myself (still using eval and numpy). Specifically, I have my own complicated PDE given by a neural network:

$$ \mathscr{N} \left ( u, \frac{du}{dx}, \frac{d^2 u}{dx^2} \right ) $$

I can convert this to an explicit string representation that looks something like this, but much longer and with more composition:
N_string = "0.5*sin(0.1*u+0.2*d_dx(u)+0.3*laplace(u)+0.4)+1.0*sin(0.6*u+0.7*d_dx(u)+0.8*laplace(u)+0.9)"

If I evaluate a similar string using numpy values for the variables with eval, it runs very fast (less than 1 second). However, as I scale up the size of my neural network that describes the PDE via a string from 600 characters to 1200 characters, the run takes 5 times as long (~30 seconds to ~150 seconds). Any larger and the run does not progress past Initializing: 0% (from the tracker=["progress",storage.tracker(1)] solver option, the progress bar shows Initializing: 0% 0/10.0 [00:00<?, ?it/s]). Trying with 1800 characters, there was no progress after 1500 seconds. I understand that a lot goes on under-the-hood, and making the PDE function more complicated may incur extra calls to the function itself, but I would not expect such a substantial change, as the 600 character version worked in less than a minute.

After implementing this string description of the PDE, I encountered the user_func option, but haven't found any examples on how to use it, and I don't know if it would help, as I would just use eval anyways.

Thus, my questions are the following:

• Would switching to user_func help fix this issue? And if so, are there any resources I've missed that can point me to how to interact with the u, d_dx(u), and laplace(u) inputs?
• From the documenation, it seems that the string is evaluated using Python's eval() with numpy variables (or maybe numba) for u, d_dx(u), and laplace(u). What are some other considerations I may be missing when I go from just doing my own test eval with random values to using the string within eq.solve? For instance, if d_dx(u) appears multiple times, will this incur extra costs? Is there a symbolic solve attempt that scales based on the string size?

I've included my code below, with a simplified N_string. The real N_string I want to use is 80000 characters. Also for determination of the grid size and time step, I'm using options that worked for Burgers' equation, and my complicated neural network should exhibit similar behavior to Burgers' equation.

from pde import PDE, CartesianGrid, MemoryStorage, ScalarField, plot_kymograph

dt = 0.01

N_string = "0.5*sin(0.1*u+0.2*d_dx(u)+0.3*laplace(u)+0.4)+1.0*sin(0.6*u+0.7*d_dx(u)+0.8*laplace(u)+0.9)"
# N_string = "-u*d_dx(u)+0.1*laplace(u)"

eq = PDE({"u": N_string})
grid = CartesianGrid([[-8.0, 8.0]], [64], periodic=True)
state = ScalarField.from_expression(grid, "exp(-(x+2.0)**2.0)")

# solve the equation and store the trajectory
storage = MemoryStorage()
eq.solve(state, t_range=10, tracker=["progress",storage.tracker(1)], dt=dt)
pde_out = storage.data

plot_kymograph(storage)

[david-zwicker]

This is an interesting use-case and I must admit that we have never tested in detail how the performance of the PDE class scales with the length of the input string. I would not be surprised if it does not perform well for your situation. Basically, the string is parsed by sympy, converted into a python function using lambdify, and compiled by numba for speed. However, every occurrence of laplace(u) and d_dx(u) will trigger a separate evaluation of the differential operators, since variable reuse has not been implemented yet (since it's a bit of a tricky thing to do right). Consequently, I suspect that using the automatic infrastructure provided by the PDE class is less than ideal.

A suitable alternative might be to do implement the class defining the PDE yourself. The documentation shows the first steps for how to do this. You could use the infrastructure of py-pde to evaluate laplace(u) and d_dx(u) at first and then call your long string separately – either using eval or also using ScalarExpression from py-pde to define a function of three variables (u, du_dx, and d2u_dx2).

[gnorman7]

For reference, this was able to totally fix the run time issue, and the 80000 character string, with 625 calls to $u$, $\frac{du}{dx}$, $\frac{d^2u}{dx^2}$ each ran in less than a minute. For any machine learning folks, I generated the string as the explicit representation of a multi-layer perceptron with 2 hidden layers of 25 units and sine activation functions. In hindsight, given this implementation, I would have been tempted to keep the neural network as a PyTorch model (rather than a string), but this would involve a difficult conversion from PyTorch Tensor datatypes to py-pde ScalarField. For anyone else interested in why I would ever use a neural network to define a PDE (a valid question), I refer to this paper:
Raissi, Maziar. "Deep hidden physics models: Deep learning of nonlinear partial differential equations." The Journal of Machine Learning Research 19.1 (2018): 932-955.

In addition to just replacing some variables in N_string, my code implementing this was the following:

class StringPDE(PDEBase):
    def __init__(self, N_string):
        self.N_string = N_string

    def evolution_rate(self, state, t=0):
        du_dx = state.gradient(bc='periodic').to_scalar()
        d2u_dx2 = state.laplace(bc='periodic')
        du_dt = eval(self.N_string)
        return du_dt

dt = 0.01

# Boundary and Initial Conditions
grid = CartesianGrid([[-8.0, 8.0]], [64], periodic=True)
# state = ScalarField.from_expression(grid, "sin(x)")
state = ScalarField.from_expression(grid, "exp(-(x+2.0)**2.0)")

# solve the equation and store the trajectory
eq = StringPDE(N_string)
storage = MemoryStorage()
eq.solve(state, t_range=10, tracker=storage.tracker(0.1), dt=dt)

plot_kymograph(storage)

[david-zwicker]

Great to hear that things worked out! If you still need to improve performance, you can look into defining a numba-compiled right hand side by implementing the _make_pde_rhs_numba method; see documentation. However, this might be more challenging and thus not worth the effort.

[gnorman7]

Thank you for the suggestion! In hindsight, it likely would have been worthwhile to implement numba-compilation, as this approach (converting a neural network to an explicit string and evaluating the string) took a few minutes, while a simpler PDE string took less than 10 seconds.

Because of changes on the neural-network side of things, I also implemented this without converting to a string, interfacing directly with my PyTorch model while it is still on my GPU. I've included the major function doing that below. The only other change is changing the eq=StringPDE(N_string) line to be eq=PyTorchPDE(model). Doing this makes the solve time with a neural network about double what a simple PDE is. I think numba-compilation could still be added on top to reduce the time, but for my application, 10 seconds is a great solve time.

class PyTorchPDE(PDEBase):
    def __init__(self, N_dnn):
        # This is a PyTorch model describing the PDE with three units: u, u_x, u_xx
        self.N = N_dnn

    def evolution_rate(self, state, t=0):
        du_dx = state.gradient(bc='periodic').to_scalar()
        d2u_dx2 = state.laplace(bc='periodic')

        # Convert to PyTorch input
        state_t = torch.tensor(state.data,dtype=torch.float32,device=device)
        du_dx_t = torch.tensor(du_dx.data,dtype=torch.float32,device=device)
        d2u_dx2_t = torch.tensor(d2u_dx2.data,dtype=torch.float32,device=device)
        inp = torch.vstack((state_t,du_dx_t,d2u_dx2_t)).T

        # Evaluate model and store into the same data structure as expected.
        du_dt = state.copy()
        with torch.no_grad():
            pytorch_output = self.N(inp).cpu().numpy().squeeze()
            du_dt.data = pytorch_output

        return du_dt