add KAN back

KindXiaoming · Jul 8, 2024 · 8075a4e · 8075a4e
1 parent 50529bb
commit 8075a4e
Show file tree

Hide file tree

Showing 10 changed files with 29 additions and 10 deletions.
diff --git a/kan/MultKAN.py b/kan/MultKAN.py
@@ -448,7 +448,7 @@ def get_range(self, l, i, j, verbose=True):
             print('y range: [' + '%.2f' % y_min, ',', '%.2f' % y_max, ']')
         return x_min, x_max, y_min, y_max
 
-    def plot(self, folder="./figures", beta=3, mask=False, metric='fa', scale=0.5, tick=False, sample=False, in_vars=None, out_vars=None, title=None):
+    def plot(self, folder="./figures", beta=3, mask=False, metric='fa', scale=0.5, tick=False, sample=False, in_vars=None, out_vars=None, title=None, varscale=1.0):
 
         global Symbol
 
@@ -694,7 +694,7 @@ def score2alpha(score):
             n = self.width_in[0]
             for i in range(n):
                 if isinstance(in_vars[i], sympy.Expr):
-                    plt.gcf().get_axes()[0].text(1 / (2 * (n)) + i / (n), -0.1, f'${latex(in_vars[i])}$', fontsize=40 * scale, horizontalalignment='center', verticalalignment='center')
+                    plt.gcf().get_axes()[0].text(1 / (2 * (n)) + i / (n), -0.1, f'${latex(in_vars[i])}$', fontsize=40 * scale * varscale, horizontalalignment='center', verticalalignment='center')
                 else:
                     plt.gcf().get_axes()[0].text(1 / (2 * (n)) + i / (n), -0.1, in_vars[i], fontsize=40 * scale, horizontalalignment='center', verticalalignment='center')
 
@@ -706,7 +706,7 @@ def score2alpha(score):
                 if isinstance(out_vars[i], sympy.Expr):
                     plt.gcf().get_axes()[0].text(1 / (2 * (n)) + i / (n), (y0+z0) * (len(self.width) - 1) + 0.15, f'${latex(out_vars[i])}$', fontsize=40 * scale, horizontalalignment='center', verticalalignment='center')
                 else:
-                    plt.gcf().get_axes()[0].text(1 / (2 * (n)) + i / (n), (y0+z0) * (len(self.width) - 1) + 0.15, out_vars[i], fontsize=40 * scale, horizontalalignment='center', verticalalignment='center')
+                    plt.gcf().get_axes()[0].text(1 / (2 * (n)) + i / (n), (y0+z0) * (len(self.width) - 1) + 0.15, out_vars[i], fontsize=40 * scale * varscale, horizontalalignment='center', verticalalignment='center')
 
         if title != None:
             plt.gcf().get_axes()[0].text(0.5, (y0+z0) * (len(self.width) - 1) + 0.3, title, fontsize=40 * scale, horizontalalignment='center', verticalalignment='center')

diff --git a/kan/__init__.py b/kan/__init__.py
@@ -1 +1,2 @@
 from .MultKAN import *
+from .KAN import *
diff --git a/kan/utils.py b/kan/utils.py
@@ -3,7 +3,7 @@
 from sklearn.linear_model import LinearRegression
 import sympy
 import yaml
-
+from sympy.utilities.lambdify import lambdify
 
 # sigmoid = sympy.Function('sigmoid')
 # name: (torch implementation, sympy implementation)
@@ -303,3 +303,21 @@ def ex_round(ex1, n_digit):
         if isinstance(a, sympy.Float):
             ex2 = ex2.subs(a, round(a, n_digit))
     return ex2
+
+
+def augment_input(orig_vars, aux_vars, x):
+
+    # if x is a tensor
+    if isinstance(x, torch.Tensor):
+
+        for aux_var in aux_vars:
+            func = lambdify(orig_vars, aux_var,'numpy') # returns a numpy-ready function
+            aux_value = torch.from_numpy(func(*[x[:,[i]].numpy() for i in range(len(orig_vars))]))
+            x = torch.cat([x, aux_value], dim=1)
+
+    # if x is a dataset
+    elif isinstance(x, dict):
+        x['train_input'] = augment_input(orig_vars, aux_vars, x['train_input'])
+        x['test_input'] = augment_input(orig_vars, aux_vars, x['test_input'])
+
+    return x
diff --git a/pde/Saloniki_June24.jpeg b/pde/Saloniki_June24.jpeg
diff --git a/sr/figures/sp_0_0_0.png b/sr/figures/sp_0_0_0.png
diff --git a/sr/figures/sp_0_0_1.png b/sr/figures/sp_0_0_1.png
diff --git a/sr/figures/sp_0_1_0.png b/sr/figures/sp_0_1_0.png
diff --git a/sr/figures/sp_0_1_1.png b/sr/figures/sp_0_1_1.png
diff --git a/tutorials/.ipynb_checkpoints/MultKAN_tutorial-checkpoint.ipynb b/tutorials/.ipynb_checkpoints/MultKAN_tutorial-checkpoint.ipynb
@@ -88,7 +88,7 @@
   {
    "cell_type": "code",
    "execution_count": 4,
-   "id": "5a8b091f",
+   "id": "f6900184",
    "metadata": {},
    "outputs": [
     {
@@ -106,7 +106,7 @@
   {
    "cell_type": "code",
    "execution_count": 5,
-   "id": "0ba4fcef",
+   "id": "c3a785ae",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -119,7 +119,7 @@
   {
    "cell_type": "code",
    "execution_count": 6,
-   "id": "60c879b7",
+   "id": "10d710ec",
    "metadata": {},
    "outputs": [
     {

diff --git a/tutorials/MultKAN_tutorial.ipynb b/tutorials/MultKAN_tutorial.ipynb
@@ -88,7 +88,7 @@
   {
    "cell_type": "code",
    "execution_count": 4,
-   "id": "5a8b091f",
+   "id": "f6900184",
    "metadata": {},
    "outputs": [
     {
@@ -106,7 +106,7 @@
   {
    "cell_type": "code",
    "execution_count": 5,
-   "id": "0ba4fcef",
+   "id": "c3a785ae",
    "metadata": {},
    "outputs": [],
    "source": [
@@ -119,7 +119,7 @@
   {
    "cell_type": "code",
    "execution_count": 6,
-   "id": "60c879b7",
+   "id": "10d710ec",
    "metadata": {},
    "outputs": [
     {