### Install jaxKAN Tutorial Dependencies via PyPI (Bash)

Source: https://github.com/srigas/jaxkan/blob/main/docs/installation.rst

This command installs jaxKAN along with all additional dependencies required to run the provided tutorials. It uses the `[doc]` option with pip to ensure a complete environment for learning and development.

```bash
pip install jaxkan[doc]
```

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

### Install jaxKAN from Source for Development (Bash)

Source: https://github.com/srigas/jaxkan/blob/main/docs/installation.rst

These commands clone the jaxKAN repository from GitHub and install it in editable mode from the local source directory using pip. This method is primarily used for installing pre-release versions or for contributing to the project's development.

```bash
git clone https://github.com/srigas/jaxkan.git
cd jaxkan
pip install .
```

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

### Install jaxKAN CPU Version via PyPI (Bash)

Source: https://github.com/srigas/jaxkan/blob/main/docs/installation.rst

This command installs the default CPU-only version of jaxKAN and its JAX dependency from PyPI using pip. It is suitable for all platforms and does not require specific hardware drivers.

```bash
pip install jaxkan
```

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

### Defining Optax Adam Optimizer for KAN - Python

Source: https://github.com/srigas/jaxkan/blob/main/docs/tutorials/Tutorial 1 - Getting Started.ipynb

This snippet defines an Adam optimizer using `optax` with a specified `learning_rate`. It then wraps the KAN model and the `optax` optimizer within a `flax.nnx.Optimizer` object, preparing the model for the training process by linking it with the chosen optimization algorithm.

```python
opt_type = optax.adam(learning_rate=0.001)

optimizer = nnx.Optimizer(model, opt_type)
```

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

### Importing Libraries for KAN Model Development - Python

Source: https://github.com/srigas/jaxkan/blob/main/docs/tutorials/Tutorial 1 - Getting Started.ipynb

This snippet imports all necessary libraries for building and training a KAN model. It includes `jaxkan` for the KAN implementation, `jax` and `jax.numpy` for numerical operations, `sklearn` for data splitting and metrics, `flax.nnx` for model and optimizer management, `optax` for optimization, `matplotlib.pyplot` for plotting, and `numpy` for general array operations.

```python
from jaxkan.KAN import KAN

import jax
import jax.numpy as jnp

from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error

from flax import nnx
import optax

import matplotlib.pyplot as plt
import numpy as np
```

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

### Install jaxKAN GPU Version via PyPI (Bash)

Source: https://github.com/srigas/jaxkan/blob/main/docs/installation.rst

This command installs the GPU-enabled version of jaxKAN and its JAX dependency from PyPI using pip. It is highly recommended for performance but requires NVIDIA drivers and is primarily supported on Linux or Windows WSL2.

```bash
pip install jaxkan[gpu]
```

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

### Generating Synthetic Dataset - Python

Source: https://github.com/srigas/jaxkan/blob/main/docs/tutorials/Tutorial 1 - Getting Started.ipynb

This snippet executes the `generate_data` function to create the synthetic dataset `X` (input features) and `y` (output labels). It uses a fixed `seed` for reproducibility, `minval` and `maxval` to define the range of input values, and `num_samples` to specify the dataset size.

```python
seed = 42

X, y = generate_data(minval=-1, maxval=1, num_samples=1000, seed=seed)
```

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

### Splitting Data into Training and Test Sets - Python

Source: https://github.com/srigas/jaxkan/blob/main/docs/tutorials/Tutorial 1 - Getting Started.ipynb

This snippet uses `sklearn.model_selection.train_test_split` to divide the generated dataset `X` and `y` into training and testing subsets. A `test_size` of 0.2 means 20% of the data is allocated for testing, and `random_state` ensures reproducibility of the split. It then prints the shapes of the resulting training and test sets.

```python
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=seed)

print("Training set size:", X_train.shape)
print("Test set size:", X_test.shape)
```

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

### Defining JIT-Compiled KAN Training Step - Python

Source: https://github.com/srigas/jaxkan/blob/main/docs/tutorials/Tutorial 1 - Getting Started.ipynb

This snippet defines a single training step function, `train_step`, decorated with `@nnx.jit` for JIT compilation to improve performance. Inside, it defines a `loss_fn` that calculates the Mean Squared Error (MSE) between the model's predictions and the true labels. It then computes the loss and gradients using `nnx.value_and_grad` and updates the model's parameters via the optimizer.

```python
# Define train loop
@nnx.jit
def train_step(model, optimizer, X_train, y_train):

    def loss_fn(model):
        residual = model(X_train) - y_train
        loss = jnp.mean((residual)**2)

        return loss
    
    loss, grads = nnx.value_and_grad(loss_fn)(model)
    optimizer.update(grads)
    
    return loss
```

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

### Initializing Chebyshev KAN Model - Python

Source: https://github.com/srigas/jaxkan/blob/main/docs/tutorials/Tutorial 1 - Getting Started.ipynb

This snippet initializes a KAN model using the `jaxkan.KAN` class. It defines the `layer_dims` based on input, hidden, and output dimensions, specifies `chebyshev` as the `layer_type`, and sets `required_parameters` for the Chebyshev layer, including `D` (degree of Chebyshev polynomials) and `flavor`. A `seed` is provided for reproducibility.

```python
# Initialize a KAN model
n_in = X_train.shape[1]
n_out = y_train.shape[1]
n_hidden = 6

layer_dims = [n_in, n_hidden, n_hidden, n_out]
req_params = {'D': 5, 'flavor': 'exact'}

model = KAN(layer_dims = layer_dims,
            layer_type = 'chebyshev',
            required_parameters = req_params,
            seed = 42
           )

print(model)
```

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

### Re-initializing KAN Model and Optimizer (Python)

Source: https://github.com/srigas/jaxkan/blob/main/docs/tutorials/Tutorial 4 - Adaptive State Transition.ipynb

This code re-initializes a KAN model with specified dimensions, type, parameters, and a random seed. It then re-initializes the nnx.Optimizer for this new model, preparing for a new training run, potentially with different settings or a fresh start.

```Python
# Re-initialize the model
model = KAN(layer_dims = layer_dims,
            layer_type = 'base',
            required_parameters = req_params,
            seed = seed
           )

# Re-initialize the optimizer
optimizer = nnx.Optimizer(model, opt_type)
```

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

### Defining Target Function and Data Generation - Python

Source: https://github.com/srigas/jaxkan/blob/main/docs/tutorials/Tutorial 1 - Getting Started.ipynb

This snippet defines the target function `f(x,y) = x^2 + 2*exp(y)` that the KAN model will attempt to fit. It also includes `generate_data` function which creates synthetic input features (x1, x2) and corresponding output labels (y) based on the defined function, using JAX's random number generation for reproducibility and uniform distribution.

```python
def f(x,y):
    return x**2 + 2*jnp.exp(y)

def generate_data(minval=-1, maxval=1, num_samples=1000, seed=42):
    key = jax.random.PRNGKey(seed)
    x_key, y_key = jax.random.split(key)

    x1 = jax.random.uniform(x_key, shape=(num_samples,), minval=minval, maxval=maxval)
    x2 = jax.random.uniform(y_key, shape=(num_samples,), minval=minval, maxval=maxval)

    y = f(x1, x2).reshape(-1, 1)
    X = jnp.stack([x1, x2], axis=1)
    
    return X, y
```

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

### Executing KAN Model Training Loop - Python

Source: https://github.com/srigas/jaxkan/blob/main/docs/tutorials/Tutorial 1 - Getting Started.ipynb

This snippet executes the training loop for the KAN model over a specified number of epochs (`num_epochs`). In each epoch, it calls the `train_step` function to perform a forward pass, calculate loss, compute gradients, and update model parameters. The training loss for each epoch is recorded in the `train_losses` array.

```python
# Initialize train_losses
num_epochs = 2000
train_losses = jnp.zeros((num_epochs,))

for epoch in range(num_epochs):
    # Calculate the loss
    loss = train_step(model, optimizer, X_train, y_train)
    
    # Append the loss
    train_losses = train_losses.at[epoch].set(loss)
```

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

### Plotting KAN Training Loss Over Epochs - Python

Source: https://github.com/srigas/jaxkan/blob/main/docs/tutorials/Tutorial 1 - Getting Started.ipynb

This snippet generates a plot of the training loss over epochs. It uses `matplotlib.pyplot` to visualize the `train_losses` array, applying a logarithmic scale to the y-axis to better observe convergence. The plot helps confirm that the model is learning and the backpropagation process is working correctly.

```python
plt.figure(figsize=(7, 4))

plt.plot(np.array(train_losses), label='Train Loss', marker='o', color='#25599c', markersize=1)

plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.title('Training Loss Over Epochs')
plt.yscale('log')

plt.legend()
plt.grid(True, which='both', linestyle='--', linewidth=0.5) 

plt.show()
```

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

### Configuring Optimizer and Grid Updates for KAN Training in Python

Source: https://github.com/srigas/jaxkan/blob/main/docs/tutorials/Tutorial 4 - Adaptive State Transition.ipynb

This snippet sets up the optimization process for the KAN model. It initializes an `optax.adam` optimizer with a fixed `learning_rate` of 0.06. An `nnx.Optimizer` is then created, linking the model to the chosen optimization algorithm. The `grid_upds` dictionary defines the schedule for adaptive grid extensions, specifying the epoch at which the grid size `G` should transition to a new value.

```python
opt_type = optax.adam(learning_rate=0.06)

optimizer = nnx.Optimizer(model, opt_type)

grid_upds = {0 : 3, 200 : 6, 400 : 10, 600 : 24}
```

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

### Importing Libraries for KAN Model Training in Python

Source: https://github.com/srigas/jaxkan/blob/main/docs/tutorials/Tutorial 4 - Adaptive State Transition.ipynb

This snippet imports all necessary libraries for building, training, and evaluating a KAN model. It includes `jaxkan` for KAN-specific functionalities, `jax` and `jax.numpy` for numerical operations, `sklearn` for data splitting and metrics, `flax` and `optax` for model definition and optimization, and `matplotlib` and `numpy` for plotting and general array operations.

```python
from jaxkan.KAN import KAN
from jaxkan.utils.general import adam_transition

import jax
import jax.numpy as jnp

from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error

from flax import nnx
import optax

import matplotlib.pyplot as plt
import numpy as np
```

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

### Initializing Optax Adam Optimizer

Source: https://github.com/srigas/jaxkan/blob/main/docs/tutorials/Tutorial 6 - DIY KANs.ipynb

This snippet sets up the optimizer for training the `CustomKAN` model. It uses `optax.adam` with a specified learning rate of 0.001, and then wraps the model and optimizer type within `nnx.Optimizer`. This prepares the model for gradient-based updates during the training loop.

```python
opt_type = optax.adam(learning_rate=0.001)

optimizer = nnx.Optimizer(model, opt_type)
```

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

### Importing Libraries and Setting up Source Path in Python

Source: https://github.com/srigas/jaxkan/blob/main/misc/IEEE paper results/Section III/State Transition.ipynb

This snippet imports necessary libraries for JAX, Flax, Optax, and custom KAN/PIKAN modules. It also dynamically adds the '/src' directory to the system path, ensuring that custom modules can be imported correctly.

```python
import jax
import jax.numpy as jnp

import optax
from flax import linen as nn

import sys
import os

import time

# Add /src to path
path_to_src = os.path.abspath(os.path.join(os.getcwd(), '../../../src'))
if path_to_src not in sys.path:
    sys.path.append(path_to_src)

from KAN import KAN
from PIKAN import *

import numpy as np

import matplotlib.pyplot as plt
```

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

### Initializing Optax Adam Optimizer for KAN Model in Python

Source: https://github.com/srigas/jaxkan/blob/main/docs/tutorials/Tutorial 2 - Classification.ipynb

This code initializes an Adam optimizer from the Optax library with a specified learning rate of 0.001. It then wraps the KAN model and the optimizer type into a `flax.nnx.Optimizer` instance, preparing it for the training process.

```python
opt_type = optax.adam(learning_rate=0.001)

optimizer = nnx.Optimizer(model, opt_type)
```

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

### Evaluating KAN Model Performance with MSE - Python

Source: https://github.com/srigas/jaxkan/blob/main/docs/tutorials/Tutorial 1 - Getting Started.ipynb

This snippet evaluates the trained KAN model's performance on the test set. It generates predictions `y_pred` using the `model` on `X_test` and then calculates the Mean Squared Error (MSE) between these predictions and the true test labels `y_test` using `sklearn.metrics.mean_squared_error`. The calculated MSE is then printed to assess the model's fitting accuracy.

```python
y_pred = model(X_test)
mse = mean_squared_error(y_test, y_pred)

print(f"The MSE of the fit is {mse:.5f}")
```

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

### Generating Dataset for KAN Training in Python

Source: https://github.com/srigas/jaxkan/blob/main/docs/tutorials/Tutorial 4 - Adaptive State Transition.ipynb

This snippet executes the `generate_data` function to create the dataset for training and testing the KAN model. It sets a random `seed` for reproducibility and generates 4000 samples of `X` (input features) and `y` (target values) within the `[-1, 1]^4` domain, which will be used for curve fitting.

```python
seed = 42

X, y = generate_data(minval=-1, maxval=1, num_samples=4000, seed=seed)
```

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

### Importing Libraries for Adaptive PIKAN Training - Python

Source: https://github.com/srigas/jaxkan/blob/main/docs/tutorials/Tutorial 5 - Adaptive PIKANs.ipynb

This snippet imports all necessary libraries and modules for implementing adaptive PIKANs. It includes `jaxkan` for KAN models and utilities, `jax` and `jax.numpy` for numerical operations, `flax.nnx` and `optax` for neural network components and optimizers, and `matplotlib.pyplot` and `numpy` for data handling and plotting.

```python
from jaxkan.KAN import KAN

import jax
import jax.numpy as jnp

from jaxkan.utils.PIKAN import sobol_sample, gradf, get_adaptive_loss
from jaxkan.utils.general import adam_transition

from flax import nnx
import optax

import matplotlib.pyplot as plt
import numpy as np
```

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

### Visualizing KAN Model Predictions vs. Ground Truth - Python

Source: https://github.com/srigas/jaxkan/blob/main/docs/tutorials/Tutorial 1 - Getting Started.ipynb

This snippet creates a scatter plot comparing the model's predicted values (`y_pred`) against the actual ground truth values (`y_test`) from the test set. It also includes a diagonal line representing a 'perfect fit' for visual comparison. This plot helps to visually assess the model's accuracy and how well its predictions align with the true values.

```python
plt.figure(figsize=(7, 4))
plt.scatter(y_test, y_pred, alpha=0.7, color='#a3630f', marker='x', label='Predicted vs Actual')
plt.plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], color='#25599c', linestyle='--', label='Perfect Fit')
plt.xlabel('Actual Values')
plt.ylabel('Predicted Values')
plt.title('Model Predictions vs Ground Truth')
plt.legend()
plt.grid(alpha=0.3)
plt.show()
```

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

### Importing Libraries for Custom KAN Development

Source: https://github.com/srigas/jaxkan/blob/main/docs/tutorials/Tutorial 6 - DIY KANs.ipynb

This snippet imports necessary libraries for defining a custom KAN, performing numerical operations, data splitting, model optimization, and plotting. Key imports include `RBFLayer` and `SineLayer` from `jaxkan.layers`, `jax` and `jax.numpy` for array manipulation, `sklearn` for data utilities, `flax.nnx` for neural network modules, `optax` for optimizers, and `matplotlib` for visualization.

```python
from jaxkan.layers.RBF import RBFLayer
from jaxkan.layers.Sine import SineLayer

import jax
import jax.numpy as jnp

from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error

from flax import nnx
import optax

from typing import List

import matplotlib.pyplot as plt
import numpy as np
```

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

### Plotting Training Loss (High Resolution Configuration)

Source: https://github.com/srigas/jaxkan/blob/main/misc/IEEE paper results/Section III/State Transition.ipynb

Visualizes the training loss over epochs for the high-resolution training run. Similar to the base configuration plot, this helps in analyzing the model's convergence, but specifically for the more advanced training setup with adaptation and Nesterov acceleration.

```python
plt.figure(figsize=(10, 6))

plt.plot(np.array(losses_full_res), label='Train Loss', marker='o', color='blue', markersize=1)

plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.title('Training Loss Over Epochs')
plt.yscale('log')  # Set y-axis to logarithmic scale

plt.legend()
plt.grid(True, which='both', linestyle='--', linewidth=0.5) 

plt.show()
```

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

### Training PyKAN Model for Diffusion Equation

Source: https://github.com/srigas/jaxkan/blob/main/misc/IEEE paper results/Section II/pykan results.ipynb

This snippet defines the training loop for a PyKAN model to solve a diffusion equation. It includes the setup of the optimizer, definition of PDE and boundary condition losses, and the backpropagation step. It calculates derivatives using `batch_jacobian` for the PDE loss.

```python
grid_vals = [3] # Grid sizes to use
init_lr = 0.001
num_epochs = 50000

t_losses = []

tstart_time = time.time()
for grid in grid_vals:
    if grid == grid_vals[0]:
        model = pykan(width=[2,6,6,1], grid=3, k=3, grid_eps=0.05, symbolic_enabled=False, device=device)
        model.update_grid_from_samples(tcollocs)
        model = model.to(device)
    else:
        model = pykan(width=[2,6,6,1], grid=grid, k=3, grid_eps=0.05, symbolic_enabled=False, device=device).initialize_from_another_model(model, collocs)
        model = model.to(device)

    # Define the training loop
    def train():
    
        optimizer = torch.optim.Adam(model.parameters(), lr=init_lr)

        # Parameters
        D = torch.tensor(0.15, device=device, requires_grad=False)
        w_pde = torch.tensor(1.0, device=device, requires_grad=False)
        w_bcs = [torch.tensor(1.0, device=device, requires_grad=False),
                 torch.tensor(1.0, device=device, requires_grad=False),
                 torch.tensor(1.0, device=device, requires_grad=False)]

        def closure():
            optimizer.zero_grad()
            
            # The kan output has shape (batch, kan_output)
            u = model(tcollocs)

            # Define u_t and u_x functions
            u_t_fun = lambda x: batch_jacobian(model, x, create_graph=True)[:,0,0].unsqueeze(1)
            u_x_fun = lambda x: batch_jacobian(model, x, create_graph=True)[:,0,1].unsqueeze(1)

            # Define u_xx function
            u_xx_fun = lambda x: batch_jacobian(u_x_fun, x, create_graph=True)[:,0,1].unsqueeze(1)

            # Get u_t and u_x with shapes (batch,1)
            u_t = u_t_fun(tcollocs)
            #u_x = u_x_fun(tcollocs)
            # Get u_xx with shape (batch,1)
            u_xx = u_xx_fun(tcollocs)

            # PDE Loss
            source=torch.exp(-tcollocs[:,[0]])*(-torch.sin(torch.pi*tcollocs[:,[1]])+torch.pi**2*torch.sin(torch.pi*tcollocs[:,[1]]))
            pde_loss = w_pde*torch.mean((u_t - u_xx-source)**2)
            
            # Boundary Conditions Loss
            bc_loss = 0.0
            for idx, _ in enumerate(tbc_collocs):
                # Get model's prediction
                pred = model(tbc_collocs[idx])
                # Retrieve true values from boundary conditions
                true = tbc_data[idx]
                # MSE Loss
                bc_loss += w_bcs[idx]*torch.mean((pred-true)**2)

            total_loss = pde_loss + bc_loss
            total_loss.backward()
            return total_loss
            
        for step in range(num_epochs):
    
            optimizer.step(closure)

            # Get Loss
            current_loss = closure().item()
            
t_losses.append(current_loss)
    
train()

tend_time = time.time()
telapsed = tend_time - tstart_time
print(f"Total Time: {telapsed} s")
print(f"Average time per iteration: {telapsed/num_epochs:.2f} s")
```

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

### Defining a Custom KAN Class with Mixed Layers

Source: https://github.com/srigas/jaxkan/blob/main/docs/tutorials/Tutorial 6 - DIY KANs.ipynb

This class, `CustomKAN`, inherits from `nnx.Module` and demonstrates how to construct a KAN using a sequence of `RBFLayer` and `SineLayer` instances. The `__init__` method initializes these layers based on provided input and output dimensions, while the `__call__` method defines the forward pass, including an intermediate GELU activation. This allows for fine-grained control over the network architecture.

```python
class CustomKAN(nnx.Module):
    def __init__(self, rbf_layers, sine_layers, add_bias = True, seed = 42):

        self.RLayers = [
                RBFLayer(
                    n_in = rbf_layers[i],
                    n_out = rbf_layers[i + 1],
                    D = 5,
                    kernel = {"type":"gaussian"},
                    add_bias = add_bias,
                    seed = seed
                )
                for i in range(len(rbf_layers)-1)
            ]

        self.SLayers = [
                SineLayer(
                    n_in = sine_layers[i],
                    n_out = sine_layers[i + 1],
                    D = 8,
                    add_bias = add_bias,
                    seed = seed
                )
                for i in range(len(sine_layers)-1)
            ]

    def __call__(self, x):
        for layer in self.RLayers:
            x = layer(x)

        # Apply a GELU transformation
        x = nnx.gelu(x, approximate=True)

        for layer in self.SLayers:
            x = layer(x)

        return x
```

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

### Evaluating KAN Model Performance with F1-Score in Python

Source: https://github.com/srigas/jaxkan/blob/main/docs/tutorials/Tutorial 2 - Classification.ipynb

This code evaluates the trained KAN model's performance on the test set. It obtains logits from the model, applies a sigmoid activation to get probabilities, converts probabilities to binary predictions, and then calculates the F1-Score against the true labels to quantify classification accuracy.

```python
logits = model(X_test)
y_pred = np.array((nnx.sigmoid(logits) > 0.5).astype(int))
score = f1_score(y_pred, y_test)

print(f"The F1-Score of the fit is {100*score:.3f}%")
```

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

### Evaluating Model Performance with MSE

Source: https://github.com/srigas/jaxkan/blob/main/docs/tutorials/Tutorial 6 - DIY KANs.ipynb

This snippet evaluates the trained `CustomKAN` model on the test dataset. It performs a forward pass to get predictions (`y_pred`) for `X_test` and then calculates the Mean Squared Error (MSE) between the predicted values and the actual test labels (`y_test`) using `sklearn.metrics.mean_squared_error`. The calculated MSE is then printed to the console.

```python
y_pred = model(X_test)
mse = mean_squared_error(y_test, y_pred)

print(f"The MSE of the fit is {mse:.5f}")
```

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

### Initializing JAX-KAN Model and Training Parameters (Python)

Source: https://github.com/srigas/jaxkan/blob/main/misc/IEEE paper results/Section II/jaxKAN results.ipynb

This code initializes the Kernel Adaptive Network (KAN) model with specified layer dimensions and hyperparameters. It also sets up learning rate schedules, grid adaptation and extension parameters for the KAN, and defines global and local loss weights for the training process.

```python
# Initialize model
key = jax.random.PRNGKey(0)

layer_dims = [2, 6, 6, 1]
model = KAN(layer_dims=layer_dims, k=3, const_spl=False, const_res=False, add_bias=True, grid_e=0.05)
variables = model.init(key, jnp.ones([1, 2]))

# Define learning rates for scheduler
lr_vals = dict()
lr_vals['init_lr'] = 0.001
lr_vals['scales'] = {0 : 1.0}

# Define epochs for grid adaptation
adapt_every = 150
adapt_stop = 5000
grid_adapt = [i * adapt_every for i in range(1, (adapt_stop // adapt_every) + 1)]
grid_adapt = []

# Define epochs for grid extension, along with grid sizes
grid_extend = {0 : 3}

# Define global loss weights
glob_w = [jnp.array(1.0, dtype=float), jnp.array(1.0, dtype=float), jnp.array(1.0, dtype=float),
          jnp.array(1.0, dtype=float), jnp.array(1.0, dtype=float)]

# Initialize RBA weights
loc_w = [jnp.ones((collocs.shape[0],1)), jnp.ones((BC1_colloc.shape[0],1)),
         jnp.ones((BC2_colloc.shape[0],1)), jnp.ones((BC3_colloc.shape[0],1)), jnp.ones((BC4_colloc.shape[0],1))]
loc_w = None
```

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

### Defining Global Loss Weights for PIKAN Training - Python

Source: https://github.com/srigas/jaxkan/blob/main/docs/tutorials/Tutorial 5 - Adaptive PIKANs.ipynb

This snippet defines global weights for the total loss function, allowing for prioritization of different loss terms (e.g., PDE loss vs. boundary condition losses). In this example, all weights are set to 1.0, indicating equal importance for each component of the combined loss.

```python
# Define global loss weights
glob_w = [jnp.array(1.0, dtype=float), jnp.array(1.0, dtype=float), jnp.array(1.0, dtype=float), jnp.array(1.0, dtype=float)]
```

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

### Initializing Custom KAN Model Instance

Source: https://github.com/srigas/jaxkan/blob/main/docs/tutorials/Tutorial 6 - DIY KANs.ipynb

This code initializes an instance of the `CustomKAN` class. It defines the architecture by specifying the number of hidden units (`n_hidden`) and the input/output dimensions for the RBF and Sine layers, ensuring the dimensions match for sequential processing. An assertion is included to validate the layer connectivity.

```python
n_hidden = 5
rbf_layers = [X_train.shape[1], n_hidden, n_hidden]
sine_layers = [n_hidden, n_hidden, y_train.shape[1]]

# Sanity check for dimensions matching
assert rbf_layers[-1] == sine_layers[0]

model = CustomKAN(rbf_layers, sine_layers, True, 42)
```

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

### Visualizing Allen-Cahn Equation Solution with Matplotlib

Source: https://github.com/srigas/jaxkan/blob/main/misc/IEEE paper results/Section II/jaxKAN results.ipynb

This code generates a 2D plot of the Allen-Cahn equation solution obtained from the trained JAX-KAN model. It defines a grid for time and space, applies the model to these coordinates to get the output, reshapes the results, and then uses 'pcolormesh' to visualize the solution, including labels and a color bar.

```Python
N_t, N_x = 101, 201

t = np.linspace(0.0, 1.0, N_t)
x = np.linspace(-1.0, 1.0, N_x)
T, X = np.meshgrid(t, x, indexing='ij')
coords = np.stack([T.flatten(), X.flatten()], axis=1)

output, _ = model.apply(variables, jnp.array(coords))
resplot = np.array(output).reshape(N_t, N_x)

plt.figure(figsize=(10, 5))
plt.pcolormesh(T, X, resplot, shading='auto', cmap='Spectral_r')
plt.colorbar()

plt.title('Solution of Allen-Cahn Equation')
plt.xlabel('t')

plt.ylabel('x')

plt.tight_layout()
plt.show()
```

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

### Configuring and Calling the `train_PIKAN` Utility Function

Source: https://github.com/srigas/jaxkan/blob/main/docs/tutorials/Tutorial 5 - Adaptive PIKANs.ipynb

This snippet demonstrates how to configure parameters and call the high-level `train_PIKAN` utility function from `jaxkan.utils.PIKAN` for adaptive PIKAN training. It sets the total number of epochs, defines learning rate schedules (`lr_vals`) with initial learning rate and scaling factors at specific epochs, and specifies grid extension epochs along with their corresponding grid sizes (`grid_extend`). This encapsulates the complex training logic into a single function call.

```python
from jaxkan.utils.PIKAN import train_PIKAN

num_epochs = 50000

lr_vals = dict()
lr_vals['init_lr'] = 0.001
lr_vals['scales'] = {0 : 1.0, 15_000 : 0.6, 25_000 : 0.8}

# Define epochs for grid extension, along with grid sizes
grid_extend = {0 : 3, 5000 : 8, 20_000 : 12}
```

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

### Visualizing Allen-Cahn Equation Solution with Matplotlib and KAN

Source: https://github.com/srigas/jaxkan/blob/main/misc/IEEE paper results/Section II/pykan results.ipynb

This snippet visualizes the learned solution of the Allen-Cahn equation by the trained KAN model. It creates a 2D grid of time (t) and space (x) coordinates, feeds them to the `model` to get predictions, and then uses `pcolormesh` from Matplotlib to create a heatmap of the solution, showing its evolution over time and space.

```python
N_t, N_x = 100, 256

tt = torch.linspace(0, 1.0, N_t, device='cpu', requires_grad=False)
tx = torch.linspace(-1.0, 1.0, N_x, device='cpu', requires_grad=False)
tT, tX = torch.meshgrid(tt, tx, indexing='ij')
tcoords = torch.stack([tT.flatten(), tX.flatten()], dim=1).double().to(device)

tresplot = model(tcoords)[:,0].cpu().detach().reshape(N_t, N_x)
plt.figure(figsize=(10, 5))
plt.pcolormesh(tT, tX, tresplot, shading='auto', cmap='Spectral_r')
plt.colorbar()

plt.title('Solution of Allen Cahn Equation')
plt.xlabel('t')

plt.ylabel('x')

plt.tight_layout()
plt.show()
```

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

### Initializing Optax Adam Optimizer for PINN Training in Python

Source: https://github.com/srigas/jaxkan/blob/main/docs/tutorials/Tutorial 7 - Inverse PINN Problem.ipynb

This code initializes an Adam optimizer from the Optax library with a specified learning rate of 0.001. It then wraps the `MyKAN` model with this optimizer using `nnx.Optimizer`, preparing the model's parameters for gradient-based updates during the training process.

```python
opt_type = optax.adam(learning_rate=0.001)

optimizer = nnx.Optimizer(model, opt_type)
```

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

### Defining and Training KAN Model for PDE with Custom Loss in PyTorch

Source: https://github.com/srigas/jaxkan/blob/main/misc/IEEE paper results/Section II/pykan results.ipynb

This snippet defines the training loop for a `pykan` model, including model initialization, optimizer setup, and a custom loss function. The loss function combines a PDE loss (Allen-Cahn equation) and boundary condition losses, calculated using `batch_jacobian` for automatic differentiation of derivatives. The model is trained for a specified number of epochs, and the total loss is tracked.

```python
grid_vals = [3] # Grid sizes to use
init_lr = 0.001
num_epochs = 50000

t_losses = []

tstart_time = time.time()
for grid in grid_vals:
    if grid == grid_vals[0]:
        model = pykan(width=[2,6,6,1], grid=3, k=3, grid_eps=0.05, symbolic_enabled=False, device=device)
        model.update_grid_from_samples(tcollocs)
        model = model.to(device)
    else:
        model = pykan(width=[2,6,6,1], grid=grid, k=3, grid_eps=0.05, symbolic_enabled=False, device=device).initialize_from_another_model(model, collocs)
        model = model.to(device)

    # Define the training loop
    def train():
    
        optimizer = torch.optim.Adam(model.parameters(), lr=init_lr)

        # Parameters
        D = torch.tensor(0.001, device=device, requires_grad=False)
        c = torch.tensor(5.0, device=device, requires_grad=False)
        w_pde = torch.tensor(1.0, device=device, requires_grad=False)
        w_bcs = [torch.tensor(1.0, device=device, requires_grad=False),
                 torch.tensor(1.0, device=device, requires_grad=False),
                 torch.tensor(1.0, device=device, requires_grad=False)]

        def closure():
            optimizer.zero_grad()
            
            # The kan output has shape (batch, kan_output)
            u = model(tcollocs)

            # Define u_t and u_x functions
            u_t_fun = lambda x: batch_jacobian(model, x, create_graph=True)[:,0,0].unsqueeze(1)
            u_x_fun = lambda x: batch_jacobian(model, x, create_graph=True)[:,0,1].unsqueeze(1)

            # Define u_xx function
            u_xx_fun = lambda x: batch_jacobian(u_x_fun, x, create_graph=True)[:,0,1].unsqueeze(1)

            # Get u_t and u_x with shapes (batch,1)
            u_t = u_t_fun(tcollocs)
            #u_x = u_x_fun(tcollocs)
            # Get u_xx with shape (batch,1)
            u_xx = u_xx_fun(tcollocs)

            # PDE Loss
            pde_loss = w_pde*torch.mean((u_t - D*u_xx-c*(u-u**3))**2)
            
            # Boundary Conditions Loss
            bc_loss = 0.0
            for idx, _ in enumerate(tbc_collocs):
                # Get model's prediction
                pred = model(tbc_collocs[idx])
                # Retrieve true values from boundary conditions
                true = tbc_data[idx]
                # MSE Loss
                bc_loss += w_bcs[idx]*torch.mean((pred-true)**2)

            total_loss = pde_loss + bc_loss
            total_loss.backward()
            return total_loss
            
        for step in range(num_epochs):
    
            optimizer.step(closure)

            # Get Loss
            current_loss = closure().item()
            
            t_losses.append(current_loss)
    
    train()

tend_time = time.time()
print(f"Total Time: {telapsed} s")
print(f"Average time per iteration: {telapsed/num_epochs:.2f} s")
```

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

### Executing Baseline PIKAN Model Training - Python

Source: https://github.com/srigas/jaxkan/blob/main/misc/IEEE paper results/Section III/Adaptive Training Results/4. Allen-Cahn Equation.ipynb

This code block initiates the training process for the PIKAN model using the `train_PIKAN` function. It specifies the total number of training epochs and passes all necessary components: the initialized model, its variables, the PDE loss function, collocation points, boundary conditions and data, global loss weights, and learning rate values. This setup executes a standard, non-adaptive training run.

```Python
num_epochs = 100000

model, variables, train_losses = train_PIKAN(model, variables, pde_loss, collocs, bc_collocs, bc_data, glob_w=glob_w, 
                                             lr_vals=lr_vals, adapt_state=False, loc_w=None, nesterov=False, 
                                             num_epochs=num_epochs, grid_extend=grid_extend, grid_adapt=[], 
                                             colloc_adapt={'epochs' : []})
```

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

### Initializing MyKAN Model Instance for PINN in Python

Source: https://github.com/srigas/jaxkan/blob/main/docs/tutorials/Tutorial 7 - Inverse PINN Problem.ipynb

This snippet initializes an instance of the `MyKAN` model. It sets up the network architecture with input dimensions derived from the collocation points, a single output, and two hidden layers, preparing the KAN for training within the PINN framework.

```python
# Initialize a MyKAN model instance
n_in = collocs.shape[1]
n_out = 1
n_hidden = 6

layer_dims = [n_in, n_hidden, n_hidden, n_out]

model = MyKAN(layer_dims = layer_dims, k = 3, G = 5, add_bias = True, seed = 42)
```

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

### Visualizing KAN Training Loss Over Epochs (Python)

Source: https://github.com/srigas/jaxkan/blob/main/docs/tutorials/Tutorial 4 - Adaptive State Transition.ipynb

This snippet generates a plot to visualize the training loss over epochs. It uses matplotlib.pyplot to create a line plot of train_losses on a logarithmic scale, adding labels, a title, a legend, and a grid for clarity.

```Python
plt.figure(figsize=(7, 4))

plt.plot(np.array(train_losses), label='Train Loss', marker='o', color='#25599c', markersize=1)

plt.xlabel('Epochs')
plt.ylabel('Loss')
plt.title('Training Loss Over Epochs')
plt.yscale('log')

plt.legend()
plt.grid(True, which='both', linestyle='--', linewidth=0.5) 

plt.show()
```

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

### Defining and Executing KAN Model Training Loop with PDE and BC Losses in PyTorch

Source: https://github.com/srigas/jaxkan/blob/main/misc/IEEE paper results/Section II/pykan results.ipynb

This comprehensive snippet defines the training process for a `pykan` model, including model initialization, optimizer setup, and a custom loss function. The `closure` function calculates PDE and boundary condition losses by computing derivatives using `batch_jacobian`, combining them into a total loss. The optimizer then performs steps based on this loss over a specified number of epochs, with training time tracked.

```Python
grid_vals = [3] # Grid sizes to use
init_lr = 0.001
num_epochs = 50000

t_losses = []

tstart_time = time.time()
for grid in grid_vals:
    if grid == grid_vals[0]:
        model = pykan(width=[2,6,6,1], grid=3, k=3, grid_eps=0.05, symbolic_enabled=False, device=device)
        model.update_grid_from_samples(tcollocs)
        model = model.to(device)
    else:
        model = pykan(width=[2,6,6,1], grid=grid, k=3, grid_eps=0.05, symbolic_enabled=False, device=device).initialize_from_another_model(model, collocs)
        model = model.to(device)

    # Define the training loop
    def train():
    
        optimizer = torch.optim.Adam(model.parameters(), lr=init_lr)

        # Parameters
        a1 = torch.tensor(1.0, device=device, requires_grad=False)
        a2 = torch.tensor(4.0, device=device, requires_grad=False)
        k = torch.tensor(1.0, device=device, requires_grad=False)
        w_pde = torch.tensor(1.0, device=device, requires_grad=False)
        w_bcs = [torch.tensor(1.0, device=device, requires_grad=False),
                 torch.tensor(1.0, device=device, requires_grad=False),
                 torch.tensor(1.0, device=device, requires_grad=False),
                 torch.tensor(1.0, device=device, requires_grad=False)]

        def closure():
            optimizer.zero_grad()
            
            # The kan output has shape (batch, kan_output)
            u = model(tcollocs)

            # Define u_t and u_x functions
            u_x_fun = lambda x: batch_jacobian(model, x, create_graph=True)[:,0,0].unsqueeze(1)
            u_y_fun = lambda x: batch_jacobian(model, x, create_graph=True)[:,0,1].unsqueeze(1)

            # Define u_xx function
            u_xx_fun = lambda x: batch_jacobian(u_x_fun, x, create_graph=True)[:,0,0].unsqueeze(1)
            u_yy_fun = lambda x: batch_jacobian(u_y_fun, x, create_graph=True)[:,0,1].unsqueeze(1)

            # Get u_t and u_x with shapes (batch,1)
            u_xx = u_xx_fun(tcollocs)
            #u_x = u_x_fun(tcollocs)
            # Get u_xx with shape (batch,1)
            u_yy = u_yy_fun(tcollocs)

            # PDE Loss
            sina1sina2=torch.sin(a1*torch.pi*tcollocs[:,[0]])*torch.sin(a2*torch.pi*tcollocs[:,[1]])
            q= -(a1*torch.pi)**2*sina1sina2 -(a2*torch.pi)**2*sina1sina2 +k*sina1sina2
            
            pde_loss = w_pde*torch.mean((u_xx + u_yy +k**2 *u - q)**2)
            
            # Boundary Conditions Loss
            bc_loss = 0.0
            for idx, _ in enumerate(tbc_collocs):
                # Get model's prediction
                pred = model(tbc_collocs[idx])
                # Retrieve true values from boundary conditions
                true = tbc_data[idx]
                # MSE Loss
                bc_loss += w_bcs[idx]*torch.mean((pred-true)**2)

            total_loss = pde_loss + bc_loss
            total_loss.backward()
            return total_loss
            
        for step in range(num_epochs):
    
            optimizer.step(closure)

            # Get Loss
            current_loss = closure().item()
            
            t_losses.append(current_loss)
    
    train()

tend_time = time.time()
telapsed = tend_time - tstart_time
print(f"Total Time: {telapsed} s")
print(f"Average time per iteration: {telapsed/num_epochs:.2f} s")
```

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

### Importing Libraries and Setting up Source Path in Python

Source: https://github.com/srigas/jaxkan/blob/main/misc/IEEE paper results/Section III/Adaptive Loss.ipynb

This snippet imports necessary libraries for JAX, Optax, Flax, system operations, time, and scientific computing. It also dynamically adds the project's 'src' directory to the Python system path, enabling the import of custom KAN and PIKAN modules required for the neural network implementation.

```python
import jax
import jax.numpy as jnp

import optax
from flax import linen as nn

import sys
import os

import time
import scipy

# Add /src to path
path_to_src = os.path.abspath(os.path.join(os.getcwd(), '../../../src'))
if path_to_src not in sys.path:
    sys.path.append(path_to_src)

from KAN import KAN
from PIKAN import *

import numpy as np

import matplotlib.pyplot as plt
```

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

### Initializing KAN Model and Training Parameters for Adaptation Only in Python

Source: https://github.com/srigas/jaxkan/blob/main/misc/IEEE paper results/Section III/State Transition.ipynb

This snippet initializes a KAN model and sets up training parameters specifically for an 'adaptation only' strategy. It defines learning rates, a single grid extension epoch, and a specific epoch for grid adaptation, along with global loss weights.

```python
# Initialize model
layer_dims = [2, 8, 8, 1]
model = KAN(layer_dims=layer_dims, k=3, const_spl=False, const_res=False, add_bias=True, grid_e=0.05)
variables = model.init(jax.random.PRNGKey(0), jnp.ones([1, 2]))

# Define learning rates for scheduler
lr_vals = dict()
lr_vals['init_lr'] = 0.003
lr_vals['scales'] = {0 : 1.0}

# Define epochs for grid extension, along with grid sizes
grid_extend = {0 : 4}

grid_adapt = [750]

# Define global loss weights
glob_w = [jnp.array(1.0, dtype=float), jnp.array(1.0, dtype=float), jnp.array(1.0, dtype=float), jnp.array(1.0, dtype=float)]
```

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

### Defining Grid Extension and Adaptation Epochs for KAN Training - Python

Source: https://github.com/srigas/jaxkan/blob/main/docs/tutorials/Tutorial 5 - Adaptive PIKANs.ipynb

This snippet configures the epochs at which the KAN's grid will be extended and adapted during training. `grid_extend` specifies the initial grid size and subsequent extensions at particular epoch milestones, while `grid_adapt` is left empty, indicating no standalone grid adaptation steps are planned for this specific training run.

```python
# Define epochs for grid extension, along with grid sizes
grid_extend = {0 : 3, 5000 : 8, 20_000 : 12}

# Define epochs for grid adaptation - in this case we opt not to
grid_adapt = []
```

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

### Configuring Adam Optimizer with Piecewise Constant Learning Rate Schedule - Python

Source: https://github.com/srigas/jaxkan/blob/main/docs/tutorials/Tutorial 5 - Adaptive PIKANs.ipynb

This snippet sets up the training optimizer and its learning rate schedule. It defines the total number of training epochs, specifies initial learning rate and scaling factors at different epoch boundaries, and then configures an `optax.adam` optimizer with Nesterov momentum, using a `piecewise_constant_schedule` for adaptive learning rate control.

```python
# Training epochs
num_epochs = 50000

# Define scheduler learning rates
lr_vals = dict()
lr_vals['init_lr'] = 0.001
lr_vals['scales'] = {0 : 1.0, 15_000 : 0.6, 25_000 : 0.8}

# Setup scheduler for optimizer
scheduler = optax.piecewise_constant_schedule(
    init_value=lr_vals['init_lr'],
    boundaries_and_scales=lr_vals['scales']
)

# We will also use the Nesterov version of Adam
opt_type = optax.adam(learning_rate=scheduler, nesterov=True)

# Define the optimizer
optimizer = nnx.Optimizer(model, opt_type)
```

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

### Initializing KAN Model and Training Parameters (Base)

Source: https://github.com/srigas/jaxkan/blob/main/misc/IEEE paper results/Section III/State Transition.ipynb

Initializes a KAN model with specified layer dimensions and hyperparameters. It also sets up learning rate schedules, grid extension epochs, and global loss weights for the training process. This configuration serves as a baseline for subsequent training.

```python
# Initialize model
layer_dims = [2, 8, 8, 1]
model = KAN(layer_dims=layer_dims, k=3, const_spl=False, const_res=False, add_bias=True, grid_e=0.05)
variables = model.init(jax.random.PRNGKey(0), jnp.ones([1, 2]))

# Define learning rates for scheduler
lr_vals = dict()
lr_vals['init_lr'] = 0.003
lr_vals['scales'] = {0 : 1.0}

# Define epochs for grid extension, along with grid sizes
grid_extend = {0 : 4}

# Define global loss weights
glob_w = [jnp.array(1.0, dtype=float), jnp.array(1.0, dtype=float), jnp.array(1.0, dtype=float), jnp.array(1.0, dtype=float)]
```