### Initialize a Simple KAN Model

Source: https://github.com/kindxiaoming/pykan/blob/master/README.md

Start with a basic KAN configuration for initial testing. This setup uses a small width and grid size, with no regularization.

```python
KAN(width=[5,1,1], grid=3, k=3)
```

--------------------------------

### Clone and Install KAN from GitHub

Source: https://github.com/kindxiaoming/pykan/blob/master/docs/index.rst

Clone the KAN repository from GitHub and install it in editable mode. Optionally, install requirements.

```bash
git clone https://github.com/KindXiaoming/pykan.git
cd pykan
pip install -e .
# pip install -r requirements.txt # install requirements
```

--------------------------------

### Initialize KAN Model and Train

Source: https://github.com/kindxiaoming/pykan/blob/master/docs/Interp/Interp_4_feature_attribution.ipynb

Sets up the KAN model, creates a dataset, and trains the model. Ensure PyTorch and SymPy are installed.

```python
from kan import *
from sympy import *

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)

# let's construct a dataset
f = lambda x: x[:,0]**2 + 0.3*x[:,1] + 0.1*x[:,2]**3 + 0.0*x[:,3]
dataset = create_dataset(f, n_var=4, device=device)

input_vars = [r'$x_'+str(i)+'$' for i in range(4)]

model = KAN(width=[4,5,1], device=device)
model.fit(dataset, steps=40, lamb=0.001);
```

--------------------------------

### Conda Environment Setup and PyKan Installation

Source: https://github.com/kindxiaoming/pykan/blob/master/README.md

Create and activate a Conda environment for PyKan, then install the library using pip from GitHub or PyPI.

```bash
conda create --name pykan-env python=3.9.7
conda activate pykan-env
pip install git+https://github.com/KindXiaoming/pykan.git  # For GitHub installation
# or
pip install pykan  # For PyPI installation
```

--------------------------------

### Install PyKan from GitHub

Source: https://github.com/kindxiaoming/pykan/blob/master/README.md

Use this command to install the development version of PyKan directly from its GitHub repository.

```bash
git clone https://github.com/KindXiaoming/pykan.git
cd pykan
pip install -e .
```

--------------------------------

### Initialize KAN and Create Dataset

Source: https://github.com/kindxiaoming/pykan/blob/master/tutorials/API_demo/API_9_video.ipynb

Sets up the device, creates a KAN model with specified architecture and parameters, defines a function, and generates a dataset for training. Ensure `torch` is installed.

```python
from kan import *
import torch

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)

# create a KAN: 2D inputs, 1D output, and 5 hidden neurons. cubic spline (k=3), 5 grid intervals (grid=5).
model = KAN(width=[4,2,1,1], grid=3, k=3, seed=1, device=device)
f = lambda x: torch.exp((torch.sin(torch.pi*(x[:,[0]]**2+x[:,[1]]**2))+torch.sin(torch.pi*(x[:,[2]]**2+x[:,[3]]**2)))/2)
dataset = create_dataset(f, n_var=4, train_num=3000, device=device)
```

--------------------------------

### Install Python Package Requirements

Source: https://github.com/kindxiaoming/pykan/blob/master/README.md

After setting up the environment, use this command to install all necessary Python packages listed in requirements.txt.

```bash
pip install -r requirements.txt
```

--------------------------------

### Install PyKan via GitHub URL

Source: https://github.com/kindxiaoming/pykan/blob/master/README.md

Install PyKan using pip by specifying the GitHub repository URL. This is useful for getting the latest updates.

```bash
pip install git+https://github.com/KindXiaoming/pykan.git
```

--------------------------------

### Create Dataset and Initialize KAN Model

Source: https://github.com/kindxiaoming/pykan/blob/master/tutorials/API_demo/API_6_training_hyperparameter.ipynb

Initializes a KAN model and creates a dataset for training. Ensure PyTorch is installed and CUDA is available if possible.

```python
from kan import *
import torch

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)

f = lambda x: torch.exp(torch.sin(torch.pi*x[:,[0]]) + x[:,[1]]**2)
dataset = create_dataset(f, n_var=2, device=device)
dataset['train_input'].shape, dataset['train_label'].shape
```

```python
# train the model
model = KAN(width=[2,5,1], grid=5, k=3, seed=1, device=device)
model.fit(dataset, opt="LBFGS", steps=20, lamb=0.01);
model.plot()
```

--------------------------------

### Setup and Data Generation for Continual Learning

Source: https://github.com/kindxiaoming/pykan/blob/master/docs/Example/Example_8_continual_learning.ipynb

Initializes necessary libraries and generates data for a 1D function composed of five Gaussian peaks. This data will be used to train the KAN in sequential phases.

```python
from kan import *
import numpy as np
import torch
import matplotlib.pyplot as plt


datasets = []

n_peak = 5
n_num_per_peak = 100
n_sample = n_peak * n_num_per_peak

x_grid = torch.linspace(-1,1,steps=n_sample)

x_centers = 2/n_peak * (np.arange(n_peak) - n_peak/2+0.5)

x_sample = torch.stack([torch.linspace(-1/n_peak,1/n_peak,steps=n_num_per_peak)+center for center in x_centers]).reshape(-1,)

y = 0.
for center in x_centers:
    y += torch.exp(-(x_grid-center)**2*300)
    
y_sample = 0.
for center in x_centers:
    y_sample += torch.exp(-(x_sample-center)**2*300)
    

plt.plot(x_grid.detach().numpy(), y.detach().numpy())
plt.scatter(x_sample.detach().numpy(), y_sample.detach().numpy())
```

--------------------------------

### Initialize KAN Model and Plot

Source: https://github.com/kindxiaoming/pykan/blob/master/docs/API_demo/API_1_indexing.rst

Initializes a KAN model with specified width and noise scale, then plots it. Ensure PyTorch and KAN are installed.

```ipython3
from kan import *

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)

model = KAN(width=[2,3,2,1], noise_scale=0.3, device=device)
x = torch.normal(0,1,size=(100,2)).to(device)
model(x);
beta = 100
model.plot(beta=beta)
```

--------------------------------

### PDE Solver Setup and KAN Initialization

Source: https://github.com/kindxiaoming/pykan/blob/master/docs/Example/Example_7_PDE_accuracy.ipynb

Sets up the environment, defines the PDE problem, and initializes the KAN model. This includes defining the solution and source functions, sampling points, and configuring the KAN for different grid sizes.

```python
from kan import KAN, LBFGS
import torch
import matplotlib.pyplot as plt
from torch import autograd
from tqdm import tqdm

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)


dim = 2
np_i = 51 # number of interior points (along each dimension)
np_b = 51 # number of boundary points (along each dimension)
ranges = [-1, 1]


def batch_jacobian(func, x, create_graph=False):
    # x in shape (Batch, Length)
    def _func_sum(x):
        return func(x).sum(dim=0)
    return autograd.functional.jacobian(_func_sum, x, create_graph=create_graph).permute(1,0,2)

# define solution
sol_fun = lambda x: torch.sin(torch.pi*x[:,[0]])*torch.sin(torch.pi*x[:,[1]])
source_fun = lambda x: -2*torch.pi**2 * torch.sin(torch.pi*x[:,[0]])*torch.sin(torch.pi*x[:,[1]])

# interior
sampling_mode = 'mesh' # 'radnom' or 'mesh'

x_mesh = torch.linspace(ranges[0],ranges[1],steps=np_i)
y_mesh = torch.linspace(ranges[0],ranges[1],steps=np_i)
X, Y = torch.meshgrid(x_mesh, y_mesh, indexing="ij")
if sampling_mode == 'mesh':
    #mesh
    x_i = torch.stack([X.reshape(-1,), Y.reshape(-1,)]).permute(1,0)
else:
    #random
    x_i = torch.rand((np_i**2,2))*2-1
    
x_i = x_i.to(device)

# boundary, 4 sides
helper = lambda X, Y: torch.stack([X.reshape(-1,), Y.reshape(-1,)]).permute(1,0)
xb1 = helper(X[0], Y[0])
xb2 = helper(X[-1], Y[0])
xb3 = helper(X[:,0], Y[:,0])
xb4 = helper(X[:,0], Y[:,-1])
x_b = torch.cat([xb1, xb2, xb3, xb4], dim=0)

x_b = x_b.to(device)

alpha = 0.01
log = 1


grids = [5,10,20]
steps = 50

pde_losses = []
bc_losses = []
l2_losses = []

for grid in grids:
    if grid == grids[0]:
        model = KAN(width=[2,2,1], grid=grid, k=3, seed=1, device=device)
        model = model.speed()
    else:
        model.save_act = True
        model.get_act(x_i)
        model = model.refine(grid)
        model = model.speed()

```

--------------------------------

### Setup and Initialization for Physics-Informed KAN

Source: https://github.com/kindxiaoming/pykan/blob/master/docs/Community/Community_1_physics_informed_kan.ipynb

Initializes PyTorch, TensorBoard, Matplotlib, and the KAN library. Sets up the device, physical constants (rho, nu, eps), and the computational domain dimensions. Defines the grid for the simulation.

```python
import torch
from torch import autograd
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm
import matplotlib.pyplot as plt
from kan import KAN, LBFGS

device = torch.device("cpu")
print("Using device:", device)

rho = torch.tensor(1.0, device=device, requires_grad=False)
nu = torch.tensor(0.01, device=device, requires_grad=False)
olds = torch.tensor(1e-8, device=device, requires_grad=False)

width, height = 10.0, 2.0
num_points_x, num_points_y = 100, 20

x = torch.linspace(0, width, num_points_x, device=device, requires_grad=False)
y = torch.linspace(0, height, num_points_y, device=device, requires_grad=False)
X, Y = torch.meshgrid(x, y, indexing='ij')
coordinates = torch.stack([X.flatten(), Y.flatten()], dim=1).to(device)
coordinates.requires_grad = True  # Ensure coordinates require grad

model = KAN(width=[2,3,3, 3], grid=5, k=10, grid_eps=1.0,
            noise_scale_base=0.25).to(device)
```

--------------------------------

### Create dataset and train KAN model

Source: https://github.com/kindxiaoming/pykan/blob/master/docs/Interp/Interp_6_test_symmetry_NN.ipynb

This snippet sets up the device, creates a dataset for a given function, and trains a KAN model. Ensure PyTorch is installed and CUDA is available if using GPU.

```python
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)

dataset = create_dataset(f, n_var=4, device=device)
model = KAN(width=[4,5,5,1], seed=0, device=device)
model.fit(dataset, steps=100);
```

--------------------------------

### Install PyKan via PyPI

Source: https://github.com/kindxiaoming/pykan/blob/master/README.md

Install the latest stable version of PyKan from the Python Package Index (PyPI) using pip.

```bash
pip install pykan
```

--------------------------------

### Initialize KAN Model and Setup for PDE

Source: https://github.com/kindxiaoming/pykan/blob/master/docs/Example/Example_6_PDE_interpretation.ipynb

Initializes the KAN model, defines helper functions for Jacobian calculation, and sets up the true solution and source functions for the PDE. It also prepares interior and boundary sampling points.

```python
from kan import *
import matplotlib.pyplot as plt
from torch import autograd
from tqdm import tqdm

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)

dim = 2
np_i = 21 # number of interior points (along each dimension)
np_b = 21 # number of boundary points (along each dimension)
ranges = [-1, 1]

model = KAN(width=[2,2,1], grid=5, k=3, seed=1, device=device)

def batch_jacobian(func, x, create_graph=False):
    # x in shape (Batch, Length)
    def _func_sum(x):
        return func(x).sum(dim=0)
    return autograd.functional.jacobian(_func_sum, x, create_graph=create_graph).permute(1,0,2)

# define solution
sol_fun = lambda x: torch.sin(torch.pi*x[:,[0]])*torch.sin(torch.pi*x[:,[1]])
source_fun = lambda x: -2*torch.pi**2 * torch.sin(torch.pi*x[:,[0]])*torch.sin(torch.pi*x[:,[1]])

# interior
sampling_mode = 'random' # 'radnom' or 'mesh'

x_mesh = torch.linspace(ranges[0],ranges[1],steps=np_i)
y_mesh = torch.linspace(ranges[0],ranges[1],steps=np_i)
X, Y = torch.meshgrid(x_mesh, y_mesh, indexing="ij")
if sampling_mode == 'mesh':
    #mesh
    x_i = torch.stack([X.reshape(-1,), Y.reshape(-1,)]).permute(1,0)
else:
    #random
    x_i = torch.rand((np_i**2,2))*2-1
    
x_i = x_i.to(device)

# boundary, 4 sides
helper = lambda X, Y: torch.stack([X.reshape(-1,), Y.reshape(-1,)]).permute(1,0)
xb1 = helper(X[0], Y[0])
xb2 = helper(X[-1], Y[0])
xb3 = helper(X[:,0], Y[:,0])
xb4 = helper(X[:,0], Y[:,-1])
x_b = torch.cat([xb1, xb2, xb3, xb4], dim=0)

x_b = x_b.to(device)

steps = 20
alpha = 0.01
log = 1

def train():
    optimizer = LBFGS(model.parameters(), lr=1, history_size=10, line_search_fn="strong_wolfe", tolerance_grad=1e-32, tolerance_change=1e-32, tolerance_ys=1e-32)

    pbar = tqdm(range(steps), desc='description', ncols=100)

    for _ in pbar:
        def closure():
            global pde_loss, bc_loss
            optimizer.zero_grad()
            # interior loss
            sol = sol_fun(x_i)
            sol_D1_fun = lambda x: batch_jacobian(model, x, create_graph=True)[:,0,:]
            sol_D1 = sol_D1_fun(x_i)
            sol_D2 = batch_jacobian(sol_D1_fun, x_i, create_graph=True)[:,:,:]
            lap = torch.sum(torch.diagonal(sol_D2, dim1=1, dim2=2), dim=1, keepdim=True)
            source = source_fun(x_i)
            pde_loss = torch.mean((lap - source)**2)

            # boundary loss
            bc_true = sol_fun(x_b)
            bc_pred = model(x_b)
            bc_loss = torch.mean((bc_pred-bc_true)**2)

            loss = alpha * pde_loss + bc_loss
            loss.backward()
            return loss

        if _ % 5 == 0 and _ < 50:
            model.update_grid_from_samples(x_i)

        optimizer.step(closure)
        sol = sol_fun(x_i)
        loss = alpha * pde_loss + bc_loss
        l2 = torch.mean((model(x_i) - sol)**2)

        if _ % log == 0:
            pbar.set_description("pde loss: %.2e | bc loss: %.2e | l2: %.2e " % (pde_loss.cpu().detach().numpy(), bc_loss.cpu().detach().numpy(), l2.cpu().detach().numpy()))

train()

```

--------------------------------

### Initialize and Train MLP Model

Source: https://github.com/kindxiaoming/pykan/blob/master/docs/Interp/Interp_10B_swap.ipynb

Initializes and trains an MLP model using the KAN library. This example demonstrates creating a dataset from specific input-output pairs and training an MLP.

```python
# MLP
from kan import *
from kan.MLP import MLP

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)

inputs = []
for i in range(2**10):
    string = "{0:b}".format(i)
    sample = [int(string[i]) for i in range(len(string))]
    sample = (10 - len(sample)) * [0] + sample
    inputs.append(sample)
   
inputs = np.array(inputs).astype(np.float32)
labels = np.sum(inputs.reshape(2**10,5,2), axis=2) % 2
inputs = torch.tensor(inputs)
labels = torch.tensor(labels)

dataset = create_dataset_from_data(inputs, labels, device=device)

model = MLP(width=[10,20,5], seed=5, device=device)
model.fit(dataset, steps=100, lamb=2e-4, reg_metric='w');
```

--------------------------------

### Blackhole Metric Transformation Setup

Source: https://github.com/kindxiaoming/pykan/blob/master/docs/Physics/Physics_3_blackhole.rst

Sets up KAN/MLP models and defines functions for metric transformation and Schwarzschild metric calculation. This snippet is used for initializing the physics simulation environment.

```python
import torch
from kan import *
from kan.utils import batch_jacobian
from kan.MLP import MLP

torch.use_deterministic_algorithms(True)
torch.set_default_dtype(torch.float64)

#model = KAN(width=[1,5,5,1], grid=5, grid_range=[1.2,3], grid_eps=0.)
#model = KAN(width=[1,1], grid=1, grid_range=[1.2,3], grid_eps=0., base_fun='zero')
#model.speed()
model = MLP(width=[1,20,20,1])

#model = lambda x: 2 * torch.sqrt(x) + torch.log((torch.sqrt(x)-1)/(torch.sqrt(x)+1))

# Schwarzschild 2M = 1
def g(x_):
    bs = x_.shape[0]
    t = x_[:,0]
    a = torch.ones_like(t)
    x = x_[:,1]
    y = x_[:,2]
    z = x_[:,3]
    r = torch.sqrt(x**2+y**2+z**2)
    stack1 = torch.stack([torch.ones(bs,)-1/r, torch.zeros(bs,), torch.zeros(bs,), torch.zeros(bs,)])
    stack2 = torch.stack([torch.zeros(bs,), -(1+x**2/((r-1)*r**2)), -x*y/((r-1)*r**2), -x*z/((r-1)*r**2)])
    stack3 = torch.stack([torch.zeros(bs,), -x*y/((r-1)*r**2), -(1+y**2/((r-1)*r**2)), -y*z/((r-1)*r**2)])
    stack4 = torch.stack([torch.zeros(bs,), -x*z/((r-1)*r**2), -y*z/((r-1)*r**2), -(1+z**2/((r-1)*r**2))])
    gs = torch.stack([stack1, stack2, stack3, stack4]).permute(2,0,1)
    return gs

def transform_g(transform, g, x):
    jac = batch_jacobian(transform, x, create_graph=True, mode='vector')
    jac_inv = torch.inverse(jac)
    return torch.matmul(torch.matmul(jac_inv.permute(0,2,1), g(x)),jac_inv)

def transform(x):
    t = x[:,[0]]
    r = torch.linalg.norm(x[:,1:], dim=1, keepdim=True)
    tp = t + model(r)
    #u = torch.sqrt(r)
    #tp = t - (2*u+torch.log((u-1)/(u+1)))
    return torch.cat([tp, x[:,1:]], dim=1)
```

--------------------------------

### KAN PDE Solver Setup and Training

Source: https://github.com/kindxiaoming/pykan/blob/master/docs/Example/Example_7_PDE_accuracy.rst

Sets up and trains a KAN model to solve a 2D Poisson equation. It defines the solution and source functions, samples interior and boundary points, and iteratively refines the KAN grid during training with LBFGS optimization. Requires PyTorch and KAN libraries.

```ipython3
from kan import KAN, LBFGS
import torch
import matplotlib.pyplot as plt
from torch import autograd
from tqdm import tqdm

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)


dim = 2
np_i = 51 # number of interior points (along each dimension)
np_b = 51 # number of boundary points (along each dimension)
ranges = [-1, 1]


def batch_jacobian(func, x, create_graph=False):
    # x in shape (Batch, Length)
    def _func_sum(x):
        return func(x).sum(dim=0)
    return autograd.functional.jacobian(_func_sum, x, create_graph=create_graph).permute(1,0,2)

# define solution
sol_fun = lambda x: torch.sin(torch.pi*x[:,[0]])*torch.sin(torch.pi*x[:,[1]])
source_fun = lambda x: -2*torch.pi**2 * torch.sin(torch.pi*x[:,[0]])*torch.sin(torch.pi*x[:,[1]])

# interior
sampling_mode = 'mesh' # 'radnom' or 'mesh'

x_mesh = torch.linspace(ranges[0],ranges[1],steps=np_i)
y_mesh = torch.linspace(ranges[0],ranges[1],steps=np_i)
X, Y = torch.meshgrid(x_mesh, y_mesh, indexing="ij")
if sampling_mode == 'mesh':
    #mesh
    x_i = torch.stack([X.reshape(-1,), Y.reshape(-1,)]).permute(1,0)
else:
    #random
    x_i = torch.rand((np_i**2,2))*2-1
    
x_i = x_i.to(device)

# boundary, 4 sides
helper = lambda X, Y: torch.stack([X.reshape(-1,), Y.reshape(-1,)]).permute(1,0)
xb1 = helper(X[0], Y[0])
xb2 = helper(X[-1], Y[0])
xb3 = helper(X[:,0], Y[:,0])
xb4 = helper(X[:,0], Y[:,-1])
x_b = torch.cat([xb1, xb2, xb3, xb4], dim=0)

x_b = x_b.to(device)

alpha = 0.01
log = 1


grids = [5,10,20]
steps = 50

pde_losses = []
bc_losses = []
l2_losses = []

for grid in grids:
    if grid == grids[0]:
        model = KAN(width=[2,2,1], grid=grid, k=3, seed=1, device=device)
        model = model.speed()
    else:
        model.save_act = True
        model.get_act(x_i)
        model = model.refine(grid)
        model = model.speed()

    def train():
        optimizer = LBFGS(model.parameters(), lr=1, history_size=10, line_search_fn="strong_wolfe", tolerance_grad=1e-32, tolerance_change=1e-32, tolerance_ys=1e-32)

        pbar = tqdm(range(steps), desc='description', ncols=100)

        for _ in pbar:
            def closure():
                global pde_loss, bc_loss
                optimizer.zero_grad()
                # interior loss
                sol = sol_fun(x_i)
                sol_D1_fun = lambda x: batch_jacobian(model, x, create_graph=True)[:,0,:]
                sol_D1 = sol_D1_fun(x_i)
                sol_D2 = batch_jacobian(sol_D1_fun, x_i, create_graph=True)[:,:,:]
                lap = torch.sum(torch.diagonal(sol_D2, dim1=1, dim2=2), dim=1, keepdim=True)
                source = source_fun(x_i)
                pde_loss = torch.mean((lap - source)**2)

                # boundary loss
                bc_true = sol_fun(x_b)
                bc_pred = model(x_b)
                bc_loss = torch.mean((bc_pred-bc_true)**2)

                loss = alpha * pde_loss + bc_loss
                loss.backward()
                return loss

            if _ % 5 == 0 and _ < 20:
                model.update_grid_from_samples(x_i)

            optimizer.step(closure)
            sol = sol_fun(x_i)
            loss = alpha * pde_loss + bc_loss
            l2 = torch.mean((model(x_i) - sol)**2)

            if _ % log == 0:
                pbar.set_description("pde loss: %.2e | bc loss: %.2e | l2: %.2e " % (pde_loss.cpu().detach().numpy(), bc_loss.cpu().detach().numpy(), l2.cpu().detach().numpy()))

            pde_losses.append(pde_loss.cpu().detach().numpy())
            bc_losses.append(bc_loss.cpu().detach().numpy())
            l2_losses.append(l2.cpu().detach().numpy())
            
        
train()

```

--------------------------------

### Initialize KAN Model and Create Dataset

Source: https://github.com/kindxiaoming/pykan/blob/master/docs/Example/Example_1_function_fitting.ipynb

Initializes a KAN model with specified width, grid size, and k-value, and creates a dataset from a given function. Ensure PyTorch and KAN are installed and the device is set correctly.

```python
from kan import *


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)

# initialize KAN with G=3
model = KAN(width=[2,1,1], grid=3, k=3, seed=1, device=device)

# create dataset
f = lambda x: torch.exp(torch.sin(torch.pi*x[:,[0]]) + x[:,[1]]**2)
dataset = create_dataset(f, n_var=2, device=device)
```

--------------------------------

### Create and Plot a KAN Model

Source: https://github.com/kindxiaoming/pykan/blob/master/docs/API_demo/API_3_extract_activations.ipynb

Initializes a KAN model with specified dimensions and activation function properties, then plots the model. Ensure PyTorch and matplotlib are installed.

```python
from kan import *
import matplotlib.pyplot as plt

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)

# create a KAN: 2D inputs, 1D output, and 5 hidden neurons. cubic spline (k=3), 5 grid intervals (grid=5).
model = KAN(width=[2,5,1], grid=3, k=3, seed=1, device=device)
x = torch.normal(0,1,size=(100,2)).to(device)
model(x)
model.plot(beta=100)
```

--------------------------------

### Standard Way to Create Dataset

Source: https://github.com/kindxiaoming/pykan/blob/master/docs/API_demo/API_11_create_dataset.rst

Use `create_dataset` with a function that expects batched inputs and explicitly selects columns using slicing like `x[:,[0]]`. Ensure necessary imports and device setup.

```ipython3
from kan.utils import create_dataset
import torch

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)

f = lambda x: x[:,[0]] * x[:,[1]]
dataset = create_dataset(f, n_var=2, device=device)
dataset['train_label'].shape
```

--------------------------------

### Create and Train a KAN Model for Edge Pruning

Source: https://github.com/kindxiaoming/pykan/blob/master/docs/API_demo/API_7_pruning.rst

This snippet initializes a KAN model, creates a dataset, and trains the model for a few steps. It's a setup for demonstrating edge pruning.

```ipython3
from kan import *
# create a KAN: 2D inputs, 1D output, and 5 hidden neurons. cubic spline (k=3), 5 grid intervals (grid=5).
model = KAN(width=[2,5,1], grid=5, k=3, seed=1, device=device)

# create dataset f(x,y) = exp(sin(pi*x)+y^2)
f = lambda x: torch.exp(torch.sin(torch.pi*x[:,[0]]) + x[:,[1]]**2)
dataset = create_dataset(f, n_var=2, device=device)
dataset['train_input'].shape, dataset['train_label'].shape

# train the model
model.fit(dataset, opt="LBFGS", steps=6, lamb=0.01);
model(dataset['train_input'])
model.plot()
```

--------------------------------

### Create and Train a KAN Model for Combined Pruning

Source: https://github.com/kindxiaoming/pykan/blob/master/docs/API_demo/API_7_pruning.rst

Initializes and trains a KAN model, preparing it for combined node and edge pruning. This setup is similar to the initial training steps.

```ipython3
from kan import *
# create a KAN: 2D inputs, 1D output, and 5 hidden neurons. cubic spline (k=3), 5 grid intervals (grid=5).
model = KAN(width=[2,5,1], grid=5, k=3, seed=1, device=device)

# create dataset f(x,y) = exp(sin(pi*x)+y^2)
f = lambda x: torch.exp(torch.sin(torch.pi*x[:,[0]]) + x[:,[1]]**2)
dataset = create_dataset(f, n_var=2, device=device)
dataset['train_input'].shape, dataset['train_label'].shape

# train the model
model.fit(dataset, opt="LBFGS", steps=20, lamb=0.01);
model(dataset['train_input'])
model.plot()
```

--------------------------------

### Setup and Train Model with LBFGS Optimizer

Source: https://github.com/kindxiaoming/pykan/blob/master/docs/Physics/Physics_3_blackhole.ipynb

Sets up the training process using the LBFGS optimizer and a custom closure function. The closure calculates the loss based on the transformed metric tensor and performs backpropagation.

```python
from kan import *
from kan.utils import batch_jacobian, create_dataset_from_data
import numpy as np


steps = 5
log = 1
optimizer = LBFGS(model.parameters(), lr=1, history_size=10, line_search_fn="strong_wolfe", tolerance_grad=1e-32, tolerance_change=1e-32, tolerance_ys=1e-32)
#optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
#optimizer = torch.optim.Adam([], lr=1e-3)
pbar = tqdm(range(steps), desc='description', ncols=100)

n_train = 10000
W = torch.normal(0,1,size=(n_train,4))
input_ = torch.empty(n_train,4, requires_grad=False)
input_[:,0] = torch.rand(n_train, requires_grad=True)
#rs = 1.2 + 1.8 * torch.linspace(0,1,n_train)
rs = 1.2 + 1.8 * torch.linspace(0,1,n_train)
input_[:,1:] = W[:,1:]/torch.norm(W[:,1:], dim=1, keepdim=True)*torch.unsqueeze(rs, dim=1)
x = input_.detach().requires_grad_(True)

def closure():
    
    global loss
    global x
    optimizer.zero_grad()
    
    g_GP = transform_g(transform, g, x)
    num = x.shape[0]
    loss = torch.mean((g_GP[:,1:,1:] + torch.eye(3,3)[None,:,:].expand(num,3,3))**2)

    loss.backward()
    return loss

for _ in pbar:
        
    '''if _ < 50 and _ % 5 == 0:
        model.update_grid(x)'''
        
    optimizer.step(closure)
    
    if _ % log == 0:
        pbar.set_description("| loss: %.2e |" % loss.cpu().detach().numpy())
        
```

--------------------------------

### Setup and Navier-Stokes Residuals with KAN

Source: https://github.com/kindxiaoming/pykan/blob/master/docs/Community/Community_1_physics_informed_kan.rst

Initializes the KAN model, defines helper functions for computing Jacobians and Hessians, and sets up the residual calculation for the Navier-Stokes equations. Requires PyTorch and KAN libraries. Ensure coordinates require gradients for backpropagation.

```python
import torch
from torch import autograd
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm
import matplotlib.pyplot as plt
from kan import KAN, LBFGS

device = torch.device("cpu")
print("Using device:", device)

rho = torch.tensor(1.0, device=device, requires_grad=False)
nu = torch.tensor(0.01, device=device, requires_grad=False)
olds = torch.tensor(1e-8, device=device, requires_grad=False)

width, height = 10.0, 2.0
num_points_x, num_points_y = 100, 20

x = torch.linspace(0, width, num_points_x, device=device, requires_grad=False)
y = torch.linspace(0, height, num_points_y, device=device, requires_grad=False)
X, Y = torch.meshgrid(x, y, indexing='ij')
coordinates = torch.stack([X.flatten(), Y.flatten()], dim=1).to(device)
coordinates.requires_grad = True  # Ensure coordinates require grad

model = KAN(width=[2,3,3, 3], grid=5, k=10, grid_eps=1.0,
            noise_scale_base=0.25).to(device)

def batch_jacobian(func, x, create_graph=False):
    def _func_sum(x):
        return func(x).sum(dim=0)
    return autograd.functional.jacobian(_func_sum, x, create_graph=create_graph).permute(1, 0, 2)

def batch_hessian(func, x):
    jacobian = batch_jacobian(func, x, create_graph=True)
    hessians = []
    for i in range(jacobian.size(1)):
        grad = autograd.grad(jacobian[:, i].sum(), x, create_graph=True, retain_graph=True)[0]
        hessians.append(grad.unsqueeze(1))
    return torch.cat(hessians, dim=1)

def navier_stokes_residuals(coords):
    coords = coords.clone().detach().requires_grad_(True)  # Ensure coords require grad
    y_pred = model(coords)
    grads = batch_jacobian(model, coords, create_graph=True)
    hessians = batch_hessian(model, coords)

    u, v, p = y_pred[:, 0], y_pred[:, 1], y_pred[:, 2]
    u_x, u_y = grads[:, 0, 0], grads[:, 0, 1]
    v_x, v_y = grads[:, 1, 0], grads[:, 1, 1]
    p_x, p_y = grads[:, 2, 0], grads[:, 2, 1]

    u_xx, u_yy = hessians[:, 0, 0], hessians[:, 0, 1]
    v_xx, v_yy = hessians[:, 1, 0], hessians[:, 1, 1]

    continuity = u_x + v_y + eps * p
    x_momentum = u * u_x + v * u_y + (1 / rho) * p_x - nu * (u_xx + u_yy)
    y_momentum = u * v_x + v * v_y + (1 / rho) * p_y - nu * (v_xx + v_yy)

    no_slip_mask = (coords[:, 1] == 0) | (coords[:, 1] == height)
    inlet_mask = (coords[:, 0] == 0)
    outlet_mask = (coords[:, 0] == width)

    no_slip_loss = torch.mean(u[no_slip_mask] ** 2 + v[no_slip_mask] ** 2)
    inlet_loss = torch.mean((u[inlet_mask] - 1) ** 2)
    outlet_pressure_loss = torch.mean(p[outlet_mask] ** 2)

    bc_loss = no_slip_loss + inlet_loss + outlet_pressure_loss
    total_loss = torch.mean(continuity ** 2 + x_momentum ** 2 + y_momentum ** 2) + bc_loss
    return total_loss

writer = SummaryWriter()

def train():
    optimizer = LBFGS(model.parameters(), lr=1,
                      history_size=10, line_search_fn="strong_wolfe", tolerance_grad=1e-32, tolerance_change=1e-32, tolerance_ys=1e-32)
    
    steps = 200 # 20 steps are enough
    pbar = tqdm(range(steps), desc='Training Progress')

    for step in pbar:
        def closure():
            optimizer.zero_grad()
            loss = navier_stokes_residuals(coordinates)
            loss.backward()
            return loss

        optimizer.step(closure)
        if step % 5 == 0:
            current_loss = closure().item()
            pbar.set_description("Step: %d | Loss: %.3f" %
                                 (step, current_loss))
            writer.add_scalar('Loss/train', current_loss, step)

train()

writer.close()
```

--------------------------------

### Creating Video from Saved Training Frames

Source: https://github.com/kindxiaoming/pykan/blob/master/docs/API_demo/API_9_video.ipynb

This snippet uses `moviepy` to compile a sequence of JPG images saved during KAN training into a video file. Ensure `moviepy` is installed and the `image_folder` path is correct.

```python
import os
import numpy as np
import moviepy.video.io.ImageSequenceClip # moviepy == 1.0.3

video_name='video'
fps=5

fps = fps
files = os.listdir(image_folder)
train_index = []
for file in files:
    if file[0].isdigit() and file.endswith('.jpg'):
        train_index.append(int(file[:-4]))

train_index = np.sort(train_index)

image_files = [image_folder+'/'+str(train_index[index])+'.jpg' for index in train_index]

clip = moviepy.video.io.ImageSequenceClip.ImageSequenceClip(image_files, fps=fps)
clip.write_videofile(video_name+'.mp4')
```

--------------------------------

### Define Physical System Examples

Source: https://github.com/kindxiaoming/pykan/blob/master/docs/Physics/Physics_1_Lagrangian.rst

Sets up data for different physical systems (harmonic oscillator, single pendulum, relativistic mass) by defining their generalized coordinates, velocities, and accelerations. This code is used to generate training data for the KAN model.

```ipython3
from kan import *
from kan.utils import batch_jacobian, batch_hessian

torch.use_deterministic_algorithms(True)
torch.set_default_dtype(torch.float64)

seed = 0
torch.manual_seed(seed)

#example = 'harmonic_oscillator'
#example = 'single_pendulum'
example = 'relativistic_mass'

# three examples: harmonic oscillator, single pendulum, double pendulum

# dimension of q
# Lagrangian: (q, qd) -> (qdd)

if example == 'harmonic_oscillator':
    n_sample = 1000
    # harmonic oscillator
    d = 1
    q = torch.rand(size=(n_sample,1)) * 4 - 2
    qd = torch.rand(size=(n_sample,1)) * 4 - 2
    qdd = - q
    x = torch.cat([q, qd], dim=1)
    
if example == 'single_pendulum':
    n_sample = 1000
    # harmonic oscillator
    d = 1
    q = torch.rand(size=(n_sample,1)) * 4 - 2
    qd = torch.rand(size=(n_sample,1)) * 4 - 2
    qdd = - torch.sin(q)
    x = torch.cat([q, qd], dim=1)
    
if example == 'relativistic_mass':
    n_sample = 10000
    # harmonic oscillator
    d = 1
    q = torch.rand(size=(n_sample,1)) * 4 - 2
    #qd = torch.rand(size=(n_sample,1)) * 1.998 - 0.999
    #qd = 0.95 + torch.rand(size=(n_sample,1)) * 0.05
    qd = torch.rand(size=(n_sample,1)) * 2 - 1
    qdd = (1 - qd**2)**(3/2)
    x = torch.cat([q, qd], dim=1)
    
```

--------------------------------

### Initialize KAN Model and Create Dataset

Source: https://github.com/kindxiaoming/pykan/blob/master/docs/Example/Example_10_relativity-addition.ipynb

Initializes a KAN model with specified width, grid, and k values, and creates a dataset for the relativistic velocity addition function. Ensure PyTorch is installed and a CUDA-enabled GPU is available for faster computation.

```python
from kan import *

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)

# initialize KAN with G=3
model = KAN(width=[2,1,1], grid=10, k=3, device=device)

# create dataset
f = lambda x: (x[:,[0]]+x[:,[1]])/(1+x[:,[0]]*x[:,[1]])
dataset = create_dataset(f, n_var=2, ranges=[-0.9,0.9], device=device)
```

--------------------------------

### KAN Uniform Grid Update from Samples

Source: https://github.com/kindxiaoming/pykan/blob/master/docs/API_demo/API_5_grid.ipynb

Illustrates updating a KAN model's grid using samples from a normal distribution. This example uses the default `grid_eps=1`, resulting in a uniform grid spanning the range of the samples.

```python
# uniform grid
model = KAN(width=[1,1], grid=G, k=k)
print(model.act_fun[0].grid) # by default, the grid is in [-1,1]
x = torch.normal(0,1,size=(1000,1))
model.update_grid_from_samples(x)
print(model.act_fun[0].grid) # now the grid becomes in [-10,10]. We add a 0.01 margin in case x have zero variance
```

--------------------------------

### KAN Training Closure and Optimizer Setup

Source: https://github.com/kindxiaoming/pykan/blob/master/docs/Physics/Physics_1_Lagrangian.rst

Defines the training closure function that calculates the loss based on the predicted and actual accelerations derived from the KAN model's Jacobian and Hessian. It also sets up the LBFGS optimizer for training.

```ipython3
from kan import *
from kan.utils import batch_jacobian, create_dataset_from_data
import numpy as np

torch.use_deterministic_algorithms(True)

def closure():
    
    global loss
    optimizer.zero_grad()
    
    jacobian = batch_jacobian(model, x, create_graph=True)
    hessian = batch_hessian(model, x, create_graph=True)
    Lqdqd = hessian[:,d:,d:]
    Lq = jacobian[:,:d]
    Lqqd = hessian[:,d:,:d]

    Lqqd_qd_prod = torch.einsum('ijk,ik->ij', Lqqd, qd)

    qdd_pred = torch.einsum('ijk,ik->ij', torch.linalg.inv(Lqdqd), Lq - Lqqd_qd_prod)
    loss = torch.mean((qdd - qdd_pred)**2)

    loss.backward()
    return loss

steps = 20
log = 1
optimizer = LBFGS(model.parameters(), lr=1, history_size=10, line_search_fn="strong_wolfe", tolerance_grad=1e-32, tolerance_change=1e-32, tolerance_ys=1e-32)
#optimizer = torch.optim.Adam(params, lr=1e-2)
pbar = tqdm(range(steps), desc='description', ncols=100)
```

--------------------------------

### Setting up KAN and Training with Video Saving

Source: https://github.com/kindxiaoming/pykan/blob/master/docs/API_demo/API_9_video.ipynb

This snippet initializes a KAN model, creates a dataset, and then trains the model while saving training figures for video generation. Ensure the `img_folder` is correctly specified.

```python
from kan import *
import torch

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)

# create a KAN: 2D inputs, 1D output, and 5 hidden neurons. cubic spline (k=3), 5 grid intervals (grid=5).
model = KAN(width=[4,2,1,1], grid=3, k=3, seed=1, device=device)
f = lambda x: torch.exp((torch.sin(torch.pi*(x[:,[0]]**2+x[:,[1]]**2))+torch.sin(torch.pi*(x[:,[2]]**2+x[:,[3]]**2)))/2)
dataset = create_dataset(f, n_var=4, train_num=3000, device=device)

image_folder = 'video_img'

# train the model
#model.train(dataset, opt="LBFGS", steps=20, lamb=1e-3, lamb_entropy=2.);
model.fit(dataset, opt="LBFGS", steps=5, lamb=0.001, lamb_entropy=2., save_fig=True, beta=10, 
            in_vars=[r'$x_1$', r'$x_2$', r'$x_3$', r'$x_4$'],
            out_vars=[r'${\rm exp}({\rm sin}(x_1^2+x_2^2)+{\rm sin}(x_3^2+x_4^2))$'],
            img_folder=image_folder);
```