Sharing a rewrite of deepxde based on jax ecosystem: pinnx

[chaoming0625]

Hello everyone, I am pleased to introduce our recent rewrite of the DeepXDE framework: the PINNx framework. We have just successfully tested all the examples. In these examples, problems can be defined as variables, and derivatives are more intuitive to obtain. Below is a simple example:

def pde(x, y):
    jacobian = net.jacobian(x)
    hessian = net.hessian(x)

    u_vel, v_vel, p = y['u_vel'], y['v_vel'], y['p']
    u_vel_x = jacobian['u_vel']['x']
    u_vel_y = jacobian['u_vel']['y']
    u_vel_xx = hessian['u_vel']['x']['x']
    u_vel_yy = hessian['u_vel']['y']['y']

    v_vel_x = jacobian['v_vel']['x']
    v_vel_y = jacobian['v_vel']['y']
    v_vel_xx = hessian['v_vel']['x']['x']
    v_vel_yy = hessian['v_vel']['y']['y']

    p_x = jacobian['p']['x']
    p_y = jacobian['p']['y']

    momentum_x = u_vel * u_vel_x + v_vel * u_vel_y + p_x - 1 / Re * (u_vel_xx + u_vel_yy)
    momentum_y = u_vel * v_vel_x + v_vel * v_vel_y + p_y - 1 / Re * (v_vel_xx + v_vel_yy)
    continuity = u_vel_x + v_vel_y

    return momentum_x, momentum_y, continuity

For more examples, please refer to https://github.com/chaobrain/pinnx/tree/main/examples.

We have just completed testing all the examples in PINNx, and most examples in DeepXDE are now supported. However, since the framework was developed in a relatively short period, there may still be some issues. We welcome feedback from the DeepXDE community.

Best

[lululxvi]

It would be great if you would be interested in contributing JAX code to DeepXDE.

[chaoming0625]

I am really interested in this. Howover, it seems that pinnx and DeepXDE show many architecture differences due to the autograd of jax. I am wondering how can I merge the latest pinnx into DeepXDE.

[lululxvi]

The good solution might be developing a unified way for DeepXDE and pinnx. Or if necessary, we can modify the current DeepXDE implementation, or add another autograd implementation.