Skip to content

Commit

Permalink
KAN implementation in the TEDEouS (#41)
Browse files Browse the repository at this point in the history
* Update models.py

* Update model.py

* Update models.py

KAN implementation was added to the models.py file. The implementation consists of two classes:
1. class KANLinear - realization of a linear layer using Kolmogorov-Arnold splines. It allows to model nonlinear dependencies between input and output data using splines.
2. class KAN - realization of multilayer network using Kolmogorov-Arnold splines.

* Create example_wave_physics_kan.py

Adding an example of solving the wave equation using KAN.

* Create example_burgers_inverse_kan.py

Adding an example of solving the Bürgers equation using KAN.

* Update model.py

Update model.py

* typo fix

* Update requirements.txt

Adding git modules: kan, efficient_kan, fast_kan

* Update example_wave_physics_kan.py

Update kan example

* Create example_wave_physics_efficient_kan.py

Create efficient kan example

* Update example_burgers_inverse_kan.py

Update burgers inverse kan example

* Create example_burgers_inverse_efficient_kan.py

Create burgers inverse example

* Revert "Create example_burgers_inverse_efficient_kan.py"

This reverts commit 8d7a28f.

* Create example_burgers_inverse_efficient_kan.py

Create burgers inverse efficient kan example

* Update requirements.txt

Update requirements

* Update requirements.txt

Update requirements

* Update utils.py

Adding condition for fast kan

* Update plot.py

Adding some attributes for different models

* Update plot.py

Adding some attributes for different models

* Update example_wave_physics_efficient_kan.py

Update efficient kan example

* Update example_wave_physics_kan.py

Update kan parameters in wave physics kan example

* Create example_wave_physics_fast_kan.py

Create wave physics fast kan example

* More robust cache

* Separate file for optional requirements

* Version incrementation

* Update models.py

Delete custom KAN

* Update example_wave_physics_kan.py

Update parameters

* Update example_wave_physics_efficient_kan.py

Update parameters

* Update example_wave_physics_fast_kan.py

Update parameters

* Update cache.py

---------

Co-authored-by: SuperSashka <heretik.unlimited@gmail.com>
  • Loading branch information
dimerf99 and SuperSashka authored Aug 20, 2024
1 parent 8512383 commit d574422
Show file tree
Hide file tree
Showing 12 changed files with 922 additions and 43 deletions.
109 changes: 109 additions & 0 deletions examples/example_burgers_inverse_efficient_kan.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,109 @@
import torch
import numpy as np
import scipy
import os
import sys

os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
sys.path.append(os.path.abspath(os.path.join(os.path.dirname( __file__ ), '..')))

from tedeous.data import Domain, Conditions, Equation
from tedeous.model import Model
from tedeous.callbacks import early_stopping, plot, inverse_task
from tedeous.optimizers.optimizer import Optimizer
from tedeous.device import solver_device
from tedeous.models import parameter_registr

import efficient_kan

solver_device('cuda')

domain = Domain()

domain.variable('x', [-1, 1], 60, dtype='float32')
domain.variable('t', [0, 1], 60, dtype='float32')

boundaries = Conditions()

data = scipy.io.loadmat(os.path.abspath(os.path.join(os.path.dirname( __file__ ), 'wolfram_sln/Burgers.mat')))

x = torch.tensor(data['x']).reshape(-1)
t = torch.tensor(data['t']).reshape(-1)

usol = data['usol']

bnd1 = torch.cartesian_prod(x, t).float()
bndval1 = torch.tensor(usol).reshape(-1, 1)

id_f = np.random.choice(len(bnd1), 2000, replace=False)

bnd1 = bnd1[id_f]
bndval1 = bndval1[id_f]

boundaries.data(bnd=bnd1, operator=None, value=bndval1)

net = efficient_kan.KAN(
[2, 100, 1],
grid_size=5,
spline_order=3,
scale_noise=0.1,
scale_base=1.0,
scale_spline=1.0,
base_activation=torch.nn.Tanh,
grid_eps=0.02,
grid_range=[-1, 1]
)

parameters = {'lam1': 2., 'lam2': 0.2} # true parameters: lam1 = 1, lam2 = -0.01*pi

parameter_registr(net, parameters)

equation = Equation()

burgers_eq = {
'du/dt**1':
{
'coeff': 1.,
'du/dt': [1],
'pow': 1,
'var': 0
},
'+u*du/dx':
{
'coeff': net.lam1,
'u*du/dx': [[None], [0]],
'pow': [1, 1],
'var': [0, 0]
},
'-mu*d2u/dx2':
{
'coeff': net.lam2,
'd2u/dx2': [0, 0],
'pow': 1,
'var': 0
}
}

equation.add(burgers_eq)

model = Model(net, domain, equation, boundaries)

model.compile('autograd', lambda_operator=1, lambda_bound=100)

img_dir = os.path.join(os.path.dirname( __file__ ), 'burgers_eq_img_efficient_kan')

cb_es = early_stopping.EarlyStopping(eps=1e-7,
loss_window=100,
no_improvement_patience=1000,
patience=3,
abs_loss=1e-5,
randomize_parameter=1e-5,
info_string_every=10)

cb_plots = plot.Plots(save_every=500, print_every=500, img_dir=img_dir)

cb_params = inverse_task.InverseTask(parameters=parameters, info_string_every=10)

optimizer = Optimizer('Adam', {'lr': 1e-4})

model.train(optimizer, 25e3, save_model=False, callbacks=[cb_es, cb_plots, cb_params])
102 changes: 102 additions & 0 deletions examples/example_burgers_inverse_kan.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,102 @@
import torch
import numpy as np
import scipy
import os
import sys

os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
sys.path.append(os.path.abspath(os.path.join(os.path.dirname( __file__ ), '..')))

from tedeous.data import Domain, Conditions, Equation
from tedeous.model import Model
from tedeous.callbacks import early_stopping, plot, inverse_task
from tedeous.optimizers.optimizer import Optimizer
from tedeous.device import solver_device
from tedeous.models import parameter_registr

import kan

solver_device('cuda')

domain = Domain()

domain.variable('x', [-1, 1], 60, dtype='float32')
domain.variable('t', [0, 1], 60, dtype='float32')

boundaries = Conditions()

data = scipy.io.loadmat(os.path.abspath(os.path.join(os.path.dirname( __file__ ), 'wolfram_sln/Burgers.mat')))

x = torch.tensor(data['x']).reshape(-1)
t = torch.tensor(data['t']).reshape(-1)

usol = data['usol']

bnd1 = torch.cartesian_prod(x, t).float()
bndval1 = torch.tensor(usol).reshape(-1, 1)

id_f = np.random.choice(len(bnd1), 2000, replace=False)

bnd1 = bnd1[id_f]
bndval1 = bndval1[id_f]

boundaries.data(bnd=bnd1, operator=None, value=bndval1)

net = kan.KAN(
width=[2, 100, 1],
base_fun='silu'
)

parameters = {'lam1': 2., 'lam2': 0.2} # true parameters: lam1 = 1, lam2 = -0.01*pi

parameter_registr(net, parameters)

equation = Equation()

burgers_eq = {
'du/dt**1':
{
'coeff': 1.,
'du/dt': [1],
'pow': 1,
'var': 0
},
'+u*du/dx':
{
'coeff': net.lam1,
'u*du/dx': [[None], [0]],
'pow': [1, 1],
'var': [0, 0]
},
'-mu*d2u/dx2':
{
'coeff': net.lam2,
'd2u/dx2': [0, 0],
'pow': 1,
'var': 0
}
}

equation.add(burgers_eq)

model = Model(net, domain, equation, boundaries)

model.compile('autograd', lambda_operator=1, lambda_bound=100)

img_dir = os.path.join(os.path.dirname( __file__ ), 'burgers_eq_img_kan')

cb_es = early_stopping.EarlyStopping(eps=1e-7,
loss_window=100,
no_improvement_patience=1000,
patience=3,
abs_loss=1e-5,
randomize_parameter=1e-5,
info_string_every=10)

cb_plots = plot.Plots(save_every=500, print_every=500, img_dir=img_dir)

cb_params = inverse_task.InverseTask(parameters=parameters, info_string_every=10)

optimizer = Optimizer('Adam', {'lr': 1e-4})

model.train(optimizer, 25e3, save_model=False, callbacks=[cb_es, cb_plots, cb_params])
Loading

0 comments on commit d574422

Please sign in to comment.