KAN with non-vectorized spline

Author: ndem0

pina/_src/model/spline.py

@@ -117,7 +117,7 @@ def __init__(self, order=4, knots=None, control_points=None):
             raise ValueError("knots must be one-dimensional.")
 
         # Check dimensionality of control points
-        if self.control_points.ndim > 1:
+        if self.control_points.ndim > 2:
             raise ValueError("control_points must be one-dimensional.")
 
         # Raise error if #knots != order + #control_points
@@ -277,9 +277,10 @@ def forward(self, x):
         :return: The output tensor.
         :rtype: torch.Tensor
         """
+        basis = self.basis(x.as_subclass(torch.Tensor))
         return torch.einsum(
-            "...bi, i -> ...b",
-            self.basis(x.as_subclass(torch.Tensor)).squeeze(-1),
+            "...bi, ...i -> ...b",
+            basis,
             self.control_points,
         )
 

pina/condition/tensor_condition.py

@@ -0,0 +1,84 @@
+"""Module for the DataCondition class."""
+
+import torch
+from torch_geometric.data import Data
+from .condition_interface import ConditionInterface
+from ..label_tensor import LabelTensor
+from ..graph import Graph
+
+
+class _TensorCondition(ConditionInterface):
+
+    __slots__ = ["input", "conditional_variables"]
+    _avail_input_cls = (torch.Tensor, LabelTensor, Data, Graph, list, tuple)
+    _avail_conditional_variables_cls = (torch.Tensor, LabelTensor)
+
+    def __new__(cls, input, conditional_variables=None):
+        """
+        Instantiate the appropriate subclass of :class:`DataCondition` based on
+        the type of ``input``.
+
+        :param input: Input data for the condition.
+        :type input: torch.Tensor | LabelTensor | Graph |
+            Data | list[Graph] | list[Data] | tuple[Graph] | tuple[Data]
+        :param conditional_variables: Conditional variables for the condition.
+        :type conditional_variables: torch.Tensor | LabelTensor, optional
+        :return: Subclass of DataCondition.
+        :rtype: pina.condition.data_condition.TensorDataCondition |
+            pina.condition.data_condition.GraphDataCondition
+
+        :raises ValueError: If input is not of type :class:`torch.Tensor`,
+            :class:`~pina.label_tensor.LabelTensor`, :class:`~pina.graph.Graph`,
+            or :class:`~torch_geometric.data.Data`.
+        """
+
+        if cls != DataCondition:
+            return super().__new__(cls)
+        if isinstance(input, (torch.Tensor, LabelTensor)):
+            subclass = TensorDataCondition
+            return subclass.__new__(subclass, input, conditional_variables)
+
+        if isinstance(input, (Graph, Data, list, tuple)):
+            cls._check_graph_list_consistency(input)
+            subclass = GraphDataCondition
+            return subclass.__new__(subclass, input, conditional_variables)
+
+        raise ValueError(
+            "Invalid input types. "
+            "Please provide either torch_geometric.data.Data or Graph objects."
+        )
+
+    def __init__(self, input, conditional_variables=None):
+        """
+        Initialize the object by storing the input and conditional
+        variables (if any).
+
+        :param input: Input data for the condition.
+        :type input: torch.Tensor | LabelTensor | Graph | Data | list[Graph] |
+            list[Data] | tuple[Graph] | tuple[Data]
+        :param conditional_variables: Conditional variables for the condition.
+        :type conditional_variables: torch.Tensor | LabelTensor
+
+        .. note::
+            If ``input`` consists of a list of :class:`~pina.graph.Graph` or
+            :class:`~torch_geometric.data.Data`, all elements must have the same
+            structure (keys and data types)
+        """
+
+        super().__init__()
+        self.input = input
+        self.conditional_variables = conditional_variables
+
+
+class TensorDataCondition(DataCondition):
+    """
+    DataCondition for :class:`torch.Tensor` or
+    :class:`~pina.label_tensor.LabelTensor` input data
+    """
+
+
+class GraphDataCondition(DataCondition):
+    """
+    DataCondition for :class:`~pina.graph.Graph` or
+    :class:`~torch_geometric.data.Data` input data
+    """

pina/model/block/__init__.py

@@ -25,6 +25,7 @@
     "RBFBlock",
     "GNOBlock",
     "PirateNetBlock",
+    "KANBlock",
 ]
 
 from pina._src.model.block.convolution_2d import ContinuousConvBlock

…l/kolmogorov_arnold_network/kan_layer.py pina/model/block/kan_block.pypina/model/kolmogorov_arnold_network/kan_layer.py renamed to pina/model/block/kan_block.py pina/model/kolmogorov_arnold_network/kan_layer.py renamed to pina/model/block/kan_block.py

@@ -5,11 +5,13 @@
 from pina.model.spline import Spline
 
 
-class KAN_layer(torch.nn.Module):
+class KANBlock(torch.nn.Module):
     """define a KAN layer using splines"""
     def __init__(self, k: int, input_dimensions: int, output_dimensions: int, inner_nodes: int, num=3, grid_eps=0.1, grid_range=[-1, 1], grid_extension=True, noise_scale=0.1, base_function=torch.nn.SiLU(), scale_base_mu=0.0, scale_base_sigma=1.0, scale_sp=1.0, sparse_init=True, sp_trainable=True, sb_trainable=True) -> None:
         """
         Initialize the KAN layer.
+
+        num è il numero di intervalli nella griglia iniziale (esclusi gli eventuali nodi di estensione)
         """
         super().__init__()
         self.k = k
@@ -27,24 +29,46 @@ def __init__(self, k: int, input_dimensions: int, output_dimensions: int, inner_
             self.mask = torch.nn.Parameter(torch.ones(input_dimensions, output_dimensions)).requires_grad_(False)        
 
         grid = torch.linspace(grid_range[0], grid_range[1], steps=self.num + 1)[None,:].expand(self.input_dimensions, self.num+1)
+        knots = torch.linspace(grid_range[0], grid_range[1], steps=self.num + 1)
 
         if grid_extension:
             h = (grid[:, [-1]] - grid[:, [0]]) / (grid.shape[1] - 1)
             for i in range(self.k):
                 grid = torch.cat([grid[:, [0]] - h, grid], dim=1)
                 grid = torch.cat([grid, grid[:, [-1]] + h], dim=1)
 
-        n_coef = grid.shape[1] - (self.k + 1)
+        n_control_points = len(knots) - (self.k )
 
-        control_points = torch.nn.Parameter(
-            torch.randn(self.input_dimensions, self.output_dimensions, n_coef) * noise_scale
-        )
+        # control_points = torch.nn.Parameter(
+        #     torch.randn(self.input_dimensions, self.output_dimensions, n_control_points) * noise_scale
+        # )
+        # print(control_points.shape)
+        spline_q = []
+        for q in range(self.output_dimensions):
+            spline_p = []
+            for p in range(self.input_dimensions):
+                spline_ = Spline(
+                    order=self.k,
+                    knots=knots,
+                    control_points=torch.randn(n_control_points)
+                )
+                spline_p.append(spline_)
+            spline_p = torch.nn.ModuleList(spline_p)
+            spline_q.append(spline_p)
+        self.spline_q = torch.nn.ModuleList(spline_q)
+
+                 
+        # control_points = torch.nn.Parameter(
+        #     torch.randn(n_control_points, self.output_dimensions) * noise_scale)
+        # print(control_points)
+        # print('uuu')
 
-        self.spline = Spline(order=self.k+1, knots=grid, control_points=control_points, grid_extension=grid_extension)
+        # self.spline = Spline(
+        #     order=self.k, knots=knots, control_points=control_points)
 
-        self.scale_base = torch.nn.Parameter(scale_base_mu * 1 / np.sqrt(input_dimensions) + \
-                         scale_base_sigma * (torch.rand(input_dimensions, output_dimensions)*2-1) * 1/np.sqrt(input_dimensions), requires_grad=sb_trainable)
-        self.scale_spline = torch.nn.Parameter(torch.ones(input_dimensions, output_dimensions) * scale_sp * 1 / np.sqrt(input_dimensions) * self.mask, requires_grad=sp_trainable)
+        # self.scale_base = torch.nn.Parameter(scale_base_mu * 1 / np.sqrt(input_dimensions) + \
+        #                  scale_base_sigma * (torch.rand(input_dimensions, output_dimensions)*2-1) * 1/np.sqrt(input_dimensions), requires_grad=sb_trainable)
+        # self.scale_spline = torch.nn.Parameter(torch.ones(input_dimensions, output_dimensions) * scale_sp * 1 / np.sqrt(input_dimensions) * self.mask, requires_grad=sp_trainable)
         self.base_function = base_function
 
     @staticmethod
@@ -76,19 +100,17 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         else:
             x_tensor = x
 
-        base = self.base_function(x_tensor)  # (batch, input_dimensions)
-        
-        basis = self.spline.basis(x_tensor, self.spline.k, self.spline.knots)
-        spline_out_per_input = torch.einsum("bil,iol->bio", basis, self.spline.control_points)
-
-        base_term = self.scale_base[None, :, :] * base[:, :, None]
-        spline_term = self.scale_spline[None, :, :] * spline_out_per_input
-        combined = base_term + spline_term
-        combined = self.mask[None,:,:] * combined
-        
-        output = torch.sum(combined, dim=1)  # (batch, output_dimensions)
-        
-        return output
+        y = []
+        for q in range(self.output_dimensions):
+            y_q = []
+            for p in range(self.input_dimensions):
+                spline_out = self.spline_q[q][p].forward(x_tensor[:, p])  # (batch, input_dimensions, output_dimensions)
+                base_out = self.base_function(x_tensor[:, p])  # (batch, input_dimensions)
+                y_q.append(spline_out + base_out)
+            y.append(torch.stack(y_q, dim=1).sum(dim=1))
+        y = torch.stack(y, dim=1)
+            
+        return y
 
     def update_grid_from_samples(self, x: torch.Tensor, mode: str = 'sample'):
         """

…kolmogorov_arnold_network/kan_network.py pina/model/kolmogorov_arnold_network.pypina/model/kolmogorov_arnold_network/kan_network.py renamed to pina/model/kolmogorov_arnold_network.py pina/model/kolmogorov_arnold_network/kan_network.py renamed to pina/model/kolmogorov_arnold_network.py

@@ -3,15 +3,20 @@
 import torch.nn as nn
 from typing import List
 
-try:
-    from .kan_layer import KAN_layer
-except ImportError:
-    from kan_layer import KAN_layer
+from pina.model.block import KANBlock
 
-class KAN_Network(torch.nn.Module):
+class KolmogorovArnoldNetwork(torch.nn.Module):
     """
-    Kolmogorov Arnold Network - A neural network using KAN layers instead of traditional MLP layers.
-    Each layer uses learnable univariate functions (B-splines + base functions) on edges.
+    Kolmogorov Arnold Network, a neural network using KAN layers instead of
+    traditional MLP layers. Each layer uses learnable univariate functions
+    (B-splines + base functions) on edges.
+
+    .. references::
+
+        Liu, Z., Wang, Y., Vaidya, S., Ruehle, F., Halverson, J., Soljačić, M.,
+        ... & Tegmark, M. (2024). Kan: Kolmogorov-arnold networks. arXiv
+        preprint arXiv:2404.19756.
+
     """
 
     def __init__(
@@ -35,19 +40,25 @@ def __init__(
     ):
         """
         Initialize the KAN network.
-        
-        Args:
-            layer_sizes: List of integers defining the size of each layer [input_dim, hidden1, hidden2, ..., output_dim]
-            k: Order of the B-spline
-            num: Number of grid points for B-splines
-            grid_eps: Epsilon for grid spacing
-            grid_range: Range for the grid [min, max]
-            grid_extension: Whether to extend the grid
-            noise_scale: Scale for initialization noise
-            base_function: Base activation function (e.g., SiLU)
-            scale_base_mu: Mean for base function scaling
-            scale_base_sigma: Std for base function scaling
-            scale_sp: Scale for spline functions
+
+        :param iterable layer_sizes: List of layer sizes including input and
+            output dimensions.
+        :param int k: Order of the B-spline.
+        :param int num: Number of grid points for B-splines.
+        :param float grid_eps: Epsilon for grid spacing.
+        :param list grid_range: Range for the grid [min, max].
+        :param bool grid_extension: Whether to extend the grid.
+        :param float noise_scale: Scale for initialization noise.
+        :param base_function: Base activation function (e.g., SiLU).
+        :param float scale_base_mu: Mean for base function scaling.
+        :param float scale_base_sigma: Std for base function scaling.
+        :param float scale_sp: Scale for spline functions.
+        :param int inner_nodes: Number of inner nodes for KAN layers.
+        :param bool sparse_init: Whether to use sparse initialization.
+        :param bool sp_trainable: Whether spline parameters are trainable.
+        :param bool sb_trainable: Whether base function parameters are
+            trainable.
+        :param bool save_act: Whether to save activations after each layer.
         """
         super().__init__()
 
@@ -62,7 +73,7 @@ def __init__(
         self.kan_layers = nn.ModuleList()
 
         for i in range(self.num_layers):
-            layer = KAN_layer(
+            layer = KANBlock(
                 k=k,
                 input_dimensions=layer_sizes[i],
                 output_dimensions=layer_sizes[i+1],
@@ -97,6 +108,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
         for i, layer in enumerate(self.kan_layers):
             current = layer(current)
+            current = torch.nn.functional.sigmoid(current)
 
             if self.save_act:
                 self.acts.append(current.detach())