Estimating Parameter Depending only on Space in the Wave Equation

[msomeya1]

Dear Professor Lu Lu and DeepXDE Community,

First, I would like to express my gratitude for developing this great library. It has been very helpful in conducting my Master's thesis research!

I am currently working on the inverse problem of the 2D wave equation φ_tt = α^2 * (φ_xx + φ_zz), where φ = φ(x, z, t) is the potential and α = α(x, z) is the speed of sound. My goal is to estimate the spatially dependent parameter α based on observed data for the potential φ.

There are two variables to be estimated here, φ and α.
The point is that φ depends on (x,z,t), but α depends only on (x,z) and not on t.
(In other words, PDE with heterogeneous domain.)
I found a similar issue #178, and realized that I can use apply_output_transform to modify the structure of a Neural Network.

The original answer in #178 suggested using tf.layers.dense, but I encountered an error AttributeError: module 'tensorflow' has no attribute 'layers'. As an alternative, I used tf.keras.layers.Dense (I am using TensorFlow as the backend);

def pde(x, y):
    alpha = y[:, 1:2]
    phi_xx = dde.grad.hessian(y, x, component=0, i=0, j=0)
    phi_zz = dde.grad.hessian(y, x, component=0, i=1, j=1)
    phi_tt = dde.grad.hessian(y, x, component=0, i=2, j=2)
    return alpha**2 * (phi_xx + phi_zz) - phi_tt

def output_transform(inputs, outputs):
    phi = outputs[:, 0:1]

    x = inputs[:, 0:1]
    z = inputs[:, 1:2]
    xz = tf.concat([x, z], axis=1)

    l1 = tf.keras.layers.Dense(50, activation="tanh")(xz)
    l2 = tf.keras.layers.Dense(50, activation="tanh")(l1)
    l3 = tf.keras.layers.Dense(50, activation="tanh")(l2)
    l4 = tf.keras.layers.Dense(50, activation="tanh")(l3)
    alpha = tf.keras.layers.Dense(1)(l4)

    return tf.concat([phi, alpha], axis=1)

NN = dde.nn.FNN(
  [3] + [50]*4 + [2], "tanh", "Glorot uniform"
)
NN.apply_output_transform(output_transform)

However, this resulted in an error:

ValueError: tf.function-decorated function tried to create variables on non-first call.

What can I do to resolve this problem?
I have only recently started studying machine learning, and I don't know much about TensorFlow.
Your insights would be greatly appreciated. Thank you in advance.

[lululxvi]

I believe the answers in #178 are based on TensorFlow v1. Your code looks OK, and I am not sure yet what the problem is.