add two functions with different shape

[ian204]

Issue solved

I solved my problem in a simple way - may be not the best but it works. See below

Hello everybody,
I am a very beginner user of Jax and I am trying to use it for solving PDEs. I have the following problem for which I have made some search and I have not found an answer. The problem is as follows.

I have a function p(x, y), the pressure, and divu(params_u, x, y), the divergence of velocity, where params_u are the weights of function u (the velocity), and I need to update p(x, y):

p(x, y) <-- p(x, y) - r*divu(params_u, x, y) (with r parameter).

I looks like a recursive function, but it is not efficient to follow this approach because otherwise I have to save all the weights params_u.

Rather I would like to save
p_old(x, y) = p(x, y) - r*divu(params_u, x, y)
and use it next step time to update p(x, y).

Is there any way to do this?

Many thanks.

[jakevdp]

Hi - thanks for the question. I'm having trouble understanding what your question is – do you think you could express it in terms of some example code, and show more explicitly what operation you want to do on which quantities?

[ian204]

Here is a piece of code that explains what I want to do.
Main function is u(params, inputs) - this function is computed somehow. The other function p(params, x) = p_prev(x) + dudx(params, x) is updated once u is computed.
Here, p_prev(x) is supposed to store the function p(params, x).

`import jax.numpy as jnp
import numpy as np
from jax import grad
from jax import random
from jax import vmap
from jax import jit
from functools import partial

def sigmoid(x):
return 0.5 * (jnp.tanh(x/2.) + 1.)

def nn(params, x): # neural network
w0 = params[:ne]
b0 = params[ne:2ne]
w1 = params[2ne:3ne]
b1 = params[3ne]
x = sigmoid(xw0 + b0)
x = sigmoid(jnp.sum(xw1) + b1)
return x

@jit
def u(params, x): # define the function
return nn(params, x)

dudx = grad(u, 1)
u_vect = vmap(u, (None, 0))
dudx_vect = vmap(dudx, (None, 0))

initialise inputs

num = 3
inputs = np.linspace(0., 1., num)

initialise params

ne = 1
key = random.PRNGKey(0)
params_u = random.normal(key, shape=(3*ne+1,))

initialise auxiliary variable pp, pp_prev

pp_prev = dudx_vect(params_u, inputs)
pp = dudx_vect(params_u, inputs)

@jit
def p(inputs, pp): # function wanted: p = p_prev + dudx
return pp

#p_vect = vmap(p, (0, None))

check; ok

print('dudx: ', dudx_vect(params_u, inputs))
print('pp = dudx: ', pp)
print('pp_prev = p: ', pp_prev)

update

pp = pp + dudx_vect(params_u, inputs)
pp_prev = pp
print('pp = pp_prev + dudx: ', pp)
print('pp_prev = pp: ', pp_prev)
print('p = p + dudx: ', p(inputs, pp))

params_u = 2.*params_u # just for updating dudx
print('new dudx: ', dudx_vect(params_u, inputs))
#pp = dudx_vect(params_u, inputs)
pp = pp_prev + dudx_vect(params_u, inputs)
print('p = p + dudx: ', p(inputs, pp))
print('pp: ', pp)
pp_prev = pp
print('pp_prev: ', pp_prev)`

[jakevdp]

It's tough to guess what the fix is without more information. Can you put together a minimal reproducible example of your problem? I'd suggest editing the original question rather than putting it in this discussion thread.