Wavefield "c" Value

[ShaikhaTheGreen]

Dear Ben,

Thank you for this creative work! I saw in the code that you can use two types of c values; one is obtained by _constant_c() and the other by _gaussian_c(). I understand the conceptual difference between them, but I wonder if I wanted to test a two wavefeild with two values of c (e.g.: c1 and c2) how can I write a function to do that?

My initial attempt was to write:

def _variable_c(x):
        "Defines a variable velocity model"
        return ((np.tanh(80*(x-2.5))+5)/4) * torch.ones((x.shape[0],1), dtype=x.dtype, device= x.device) 

I then use this assignment:

 c = _variable_c 
 c = c(torch.from_numpy(np.stack([xx,yy],-1).reshape((NX*NY,2)))).numpy().reshape(NX,NY)

The idea that I'm trying to achieve is to assign the value of c=1 to x<2.5 and c=1.5 to x>2.5 by using the "tanh" function for a smooth transition of c value. I'm getting this error:

~\AppData\Local\Temp/ipykernel_16716/2797370787.py in exact_solution(x, batch_size)
     51 
     52         c = _variable_c #_constant_c
---> 53         c = c(torch.from_numpy(np.stack([xx,yy],-1).reshape((NX*NY,2)))).numpy().reshape(NX,NY)
     54         #c(c0,torch.from_numpy(np.stack([xx,yy],-1).reshape((NX*NY,2)))).numpy().reshape(NX,NY)
     55 

ValueError: cannot reshape array of size 2000000 into shape (1000,1000)

Any ideas on how to achieve this and get around the error?

[benmoseley]

Hi @engsbk, please check out the latest FBPINN release - it is a major update and this should be easily implementable with the new fbpinns.problems.Problem class. Something like this:

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()

[WeCcRy]

can this problem use subdomain scheduling and how? I refered to the code presented in the example 3 but failed. scheduler = LineSchedulerRectangularND, scheduler_kwargs = dict( point=[0.], iaxis=0, ),

[benmoseley]

Yes, but this is a 2+1D equation so you probably want to use a schedulers.PlaneSchedulerRectangularND instead of a LineSchedulerRectangularND