Wavefield "c" Value about fbpinns HOT 1 OPEN

import jax.numpy as jnp
import numpy as np

from fbpinns.domains import RectangularDomainND
from fbpinns.problems import WaveEquationConstantVelocity3D
from fbpinns.networks import FCN
from fbpinns.constants import Constants
from fbpinns.trainers import PINNTrainer


class WaveEquationInterfaceVelocity3D(WaveEquationConstantVelocity3D):

    @staticmethod
    def init_params(c0=1, c1=2,
                    source=np.array([[-0.5, 0., 0.1, 1.]])):

        static_params = {
            "dims":(1,3),
            "c0":c0,
            "c1":c1,
            "c_fn":WaveEquationInterfaceVelocity3D.c_fn,# velocity function
            "source":jnp.array(source),# location, width and amplitude of initial gaussian sources (k, 4)
            }
        return static_params, {}
    
    @staticmethod
    def c_fn(all_params, x_batch):
        "Computes the velocity model"
        
        p = all_params["static"]["problem"]
        c0, c1 = p["c0"], p["c1"]
        
        x = x_batch[:,0:1]
        c = c0 + (c1-c0)*(1+jnp.tanh(x/0.1))/2
        
        return c

c = Constants(
    domain=RectangularDomainND,
    domain_init_kwargs=dict(
        xmin=np.array([-1,-1,0]),
        xmax=np.array([1,1,1]),
    ),
    problem=WaveEquationInterfaceVelocity3D,
    problem_init_kwargs=dict(
        c0=1, 
        c1=2,
    ),
    network=FCN,
    network_init_kwargs=dict(
        layer_sizes=[3,32,32,32,1],
    ),
    decomposition_init_kwargs=dict(
        unnorm=(0.,0.1),
    ),
    ns=((50,50,50),),
    n_test=(100,100,5),
    n_steps=15000,
    clear_output=True,
)

run = PINNTrainer(c)
_ = run.train()