<jupyter_start><jupyter_text><jupyter_code>!pip install nangs
# autoreload nangs
%reload_ext autoreload
%autoreload 2
%matplotlib inline

#imports

import numpy as np 
import matplotlib.pyplot as plt 
import torch
import pandas as pd
import nangs
from nangs import *

device = "cuda" if torch.cuda.is_available() else "cpu"

nangs.__version__, torch.__version__
device

def f(x,y):
  beta=0.01
  k = 2+beta*torch.cos(10*2*3**x)*torch.cos(10*2*3**y)
  
  return k

def f1(x,y):
  beta=0.01
  k = 2+beta*np.cos(10*2*3**x)*np.cos(10*2*3**y)
  
  return k
alpha=0.0
class Difussion2d(PDE):
    def computePDELoss(self, inputs, outputs):
        p = outputs[:, 0]
        grads = self.computeGrads(p, inputs)
        dpdx, dpdy = grads[:, 0], grads[:, 1]
        x,y = inputs[:,0],inputs[:,1]
        # second order derivatives        
        dp2dx2 = self.computeGrads(f(x,y)*dpdx, inputs)[:, 0] 
        dp2dy2 = self.computeGrads(f(x,y)*dpdy, inputs)[:, 1]
        return {'loss':alpha - (dp2dx2 + dp2dy2)}
class NeumannX(Neumann):
    def computeBocoLoss(self, inputs, outputs):
        p = outputs[:, 0]
        grads = self.computeGrads(p, inputs)
        dpdx = grads[:, 0]
        return {'gradX': dpdx}
class NeumannY(Neumann):
    def computeBocoLoss(self, inputs, outputs):
        p = outputs[:, 0]
        grads = self.computeGrads(p, inputs)
        dpdy = grads[:, 1]
        return {'gradX': dpdy}
pde = Difussion2d(inputs=('x', 'y'), outputs='p')
# mesh 

x = np.linspace(0,1,250)
y = np.linspace(0,1,250)
t = np.linspace(0,1,20)

mesh = Mesh({'x': x, 'y': y}, device=device)
pde.set_mesh(mesh)
_x, _y = np.meshgrid(x, y)
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
from matplotlib import cm
fig = plt.figure()
ax = fig.gca(projection='3d')
surf = ax.plot_surface(_x, _y, f1(_x, _y).reshape((_y.shape[0],_x.shape[1])), cmap=cm.coolwarm, linewidth=0, antialiased=False)
fig.colorbar(surf, shrink=0.5, aspect=5)
plt.show()
from nangs import Dirichlet
t0 = np.array([0])
#initial_condition = Dirichlet({'x': x, 'y': y}, {'p': p0.reshape(-1)}, device=device, name="initial")
#pde.add_boco(initial_condition)
p0=np.ones(len(y))
p1=np.zeros(len(y))
#left = NeumannX({'x': np.array([0]), 'y': y, 't':t}, name='left', device=device)
left = Dirichlet({'x': x[0], 'y': y}, {'p': p0.reshape(-1)}, device=device, name="left")

pde.add_boco(left)
#right = NeumannX({'x': np.array([1]), 'y': y, 't':t}, name='right', device=device)
right = Dirichlet({'x': x[-1], 'y': y}, {'p': p1.reshape(-1)}, device=device, name="right")
pde.add_boco(right)

top = NeumannY({'x': x, 'y': y[-1]}, name='top', device=device)
pde.add_boco(top)


bottom = NeumannY({'x': x, 'y': y[-1]}, name='bottom', device=device)

pde.add_boco(bottom)
from nangs import MLP

BATCH_SIZE = 512
LR = 1e-2
EPOCHS = 11
NUM_LAYERS = 2
NUM_HIDDEN = 128

mlp = MLP(len(pde.inputs), len(pde.outputs), NUM_LAYERS, NUM_HIDDEN).to(device)
optimizer = torch.optim.Adam(mlp.parameters())
scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=LR, pct_start=0.1, total_steps=EPOCHS)
pde.inputs
pde.compile(mlp, optimizer, scheduler)
%time hist = pde.solve(EPOCHS, BATCH_SIZE)
t = 0.01
x = np.linspace(0,1,100)
y = np.linspace(0,1,100)

_x, _y = np.meshgrid(x, y)
eval_mesh = Mesh({'x': x, 'y': y}, device=device)
p = pde.eval(eval_mesh)
p = p.cpu().numpy()
p.shape
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
from matplotlib import cm
fig = plt.figure()
ax = fig.gca(projection='3d')
surf = ax.plot_surface(_x, _y, p.reshape((_y.shape[0],_x.shape[1])), cmap=cm.coolwarm, linewidth=0, antialiased=False)
fig.colorbar(surf, shrink=0.5, aspect=5)
plt.show()
p1=p.reshape((_y.shape[0],_x.shape[1]))
plt.plot(p1[0])
plt.plot(p1[1])
plt.plot(p1[2])
plt.plot(p1[50])
plt.plot(p1[-1])
p = pde.eval(eval_mesh)
grads = pde.computeGrads(p, pde.inputs)
dpdx, dpdy = grads[:, 0], grads[:, 1]
p
from google.colab import drive
drive.mount('/content/gdrive')
ls '/content/gdrive/MyDrive'<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.19025118129818452, 'fraction_numerical': 0.03481720964934096, 'mean_word_length': 4.696883852691218, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '399907', 'license': 'Apache-2.0', 'n_tokens_mistral': 1805, 'n_tokens_neox': 1727, 'n_words': 353, 'path': 'Darcy2D.ipynb', 'repo_name': 'adantra/nangs', 'size': '4021'}