PDEs in Spherical Coordinates

[Rohan131200]

Hi! At the outset, thanks for the neurodiffeq libraries! They're really useful and interesting! Now, I am trying to solve the Spherical Laplace Equation with boundary conditions: u(r=0)=u(r=1)=0 for all theta and phi. Following is the code for the same

import numpy as np
import matplotlib.pyplot as plt
import torch
from neurodiffeq import diff 
from neurodiffeq.networks import FCNN 
from neurodiffeq.conditions import DirichletBVPSpherical
from neurodiffeq.solvers import SolverSpherical
from neurodiffeq.monitors import MonitorSpherical
from neurodiffeq.generators import Generator3D
%matplotlib notebook

laplace = lambda u, r, theta, phi: [
diff(((r**2)*diff(u,r,order=1)), r, order=1)/r**2 + 
diff((np.sin(theta))*diff(u,theta,order=1), theta, order=1)/((r**2)*(np.sin(theta))) +
diff(u,phi,order=2)/(r*np.sin(theta))**2
]

conditions = [
    DirichletBVPSpherical(r_0=0.0,f=0.0,r_1=1.0,g=0.0)
]

nets = [
FCNN(n_input_units=3, n_output_units=1, hidden_units=[512]),
]

monitor=MonitorSpherical(r_min=0.0,r_max=1.0,check_every=10,shape=(10,10,10),r_scale='linear',theta_min=0,theta_max=np.pi,phi_min=0,phi_max=2*np.pi)
monitor_callback = monitor.to_callback()

solver = SolverSpherical(
    pde_system=laplace,
    conditions=conditions,
    r_min=0.0,
    r_max=1.0,
    nets=nets,
    train_generator=Generator3D(grid=(10, 10, 10), xyz_min=(0.0, 0.0, 0.0), xyz_max=(1.0, 1.0, 1.0), method='equally-spaced'),
    valid_generator=Generator3D(grid=(10, 10, 10), xyz_min=(0.0, 0.0, 0.0), xyz_max=(1.0, 1.0, 1.0), method='equally-spaced-noisy'),
)

solver.fit(max_epochs=200, callbacks=[monitor_callback])

solution_neural_net_laplace = solver.get_solution()

I'm getting the following error

mat1 and mat2 shapes cannot be multiplied (1000x1 and 3x512)

A similar error arises when I try to get the Schwarzschild solution using Einstein's equations. I tried to find the source of this error by reading the documentation but couldn't find it. I'd appreciate any help in resolving this error. Thanks in advance!

[shuheng-liu]

HI @Rohan131200 , there are a couple of things to watch out for

1. When defining the PDE with the lambda functions, each input is a torch.Tenosr. So you need to change all occurence np.sin in your laplace to torch.sin.
2. When defining your boundary conditions, f and g are both functions of theta and phi so you need to use

   conditions = [DirichletBVPSpherical(r_0=0.0, f=lambda theta, phi: 0.0, r_1=1.0, g=lambda theta, phi: 0.0)]

   instead of

   conditions = [DirichletBVPSpherical(r_0=0.0, f=0.0, r_1=1.0, g=0.0)]

3. Note that theta=0 implies torch.sin(theta) = 0. So, if you include the points on the z-axis where (theta=0) in your domain, you will have NaN in the Laplacian operator. Subsequently, the loss will be NaN, which, after the first optimization step, makes all the weights in the network become NaN.
4. Thanks to your report, I just found a bug in the spherical solver. I'm going to fix it soon. Before that, you could to pass in enforcer=lambda net, cond, coords: cond.enforce(net, *coords) when instantiating the solver to make everything run smoothly.

[shuheng-liu]

This is the modified version, note that I used GeneratorSpherical instead of Generator3D to make sure no points on the z-axis (theta=0) are ever sampled.

import numpy as np
import matplotlib.pyplot as plt
import torch
from neurodiffeq import diff
from neurodiffeq.networks import FCNN
from neurodiffeq.conditions import DirichletBVPSpherical
from neurodiffeq.solvers import SolverSpherical
from neurodiffeq.monitors import MonitorSpherical
from neurodiffeq.generators import Generator3D, GeneratorSpherical

laplace = lambda u, r, theta, phi: [
    diff(((r ** 2) * diff(u, r, order=1)), r, order=1) / r ** 2 +
    diff((torch.sin(theta)) * diff(u, theta, order=1), theta, order=1) / ((r ** 2) * (torch.sin(theta))) +
    diff(u, phi, order=2) / (r * torch.sin(theta)) ** 2
]

conditions = [
    DirichletBVPSpherical(r_0=0.0, f=lambda theta, phi: 0.0, r_1=1.0, g=lambda theta, phi: 0.0)
]

nets = [
    FCNN(n_input_units=3, n_output_units=1, hidden_units=[512]),
]

monitor = MonitorSpherical(r_min=0.1, r_max=1.0, check_every=10, shape=(10, 10, 10), r_scale='linear', theta_min=0,
                           theta_max=np.pi, phi_min=0, phi_max=2 * np.pi)
monitor_callback = monitor.to_callback()

solver = SolverSpherical(
    pde_system=laplace,
    conditions=conditions,
    r_min=0.1,
    r_max=1.0,
    nets=nets,
    train_generator=GeneratorSpherical(size=1024, r_min=0.1, r_max=1.0, method='equally-spaced-noisy'),
    valid_generator=GeneratorSpherical(size=1024, r_min=0.1, r_max=1.0, method='equally-spaced-noisy'),
    enforcer=lambda net, cond, coords: cond.enforce(net, *coords),
)

solver.fit(max_epochs=2, callbacks=[monitor_callback])
solution_neural_net_laplace = solver.get_solution()

[shuheng-liu]

The enforcer issue should have just been fixed by 143df46, which will be released in v0.4.1

Before we release v0.4.1 on PyPI, you can install the latest dev branch using

pip install -U git+https://github.com/NeuroDiffGym/neurodiffeq.git@v0.4.1

After we release v0.4.1, you can simply run

pip install -U neurodiffeq

[Rohan131200]

@shuheng-liu Thanks for the detailed explanation! I'm quite new to neural networks and neurodiffeq which is why I was facing these issues. I'll be sure to keep these points in mind. Thanks!