Dataset step size

.gitignore

@@ -5,3 +5,10 @@ docs/_static/
 docs/_templates
 folder
 test
+tutorials/figures/
+tutorials/model_ckpt/model_oct0.ckpt
+tutorials/continual_learning_structural.py
+tutorials/model_ckpt/model_oct1.ckpt
+tutorials/model_ckpt/model_oct2.ckpt
+tutorials/continual_learning_step_by_stop.ipynb
+tutorials/continual_learning_structural.ipynb

kan/KAN.py

@@ -760,7 +760,7 @@ def score2alpha(score):
             plt.gcf().get_axes()[0].text(0.5, y0 * (len(self.width) - 1) + 0.2, title, fontsize=40 * scale, horizontalalignment='center', verticalalignment='center')
 
     def train(self, dataset, opt="LBFGS", steps=100, log=1, lamb=0., lamb_l1=1., lamb_entropy=2., lamb_coef=0., lamb_coefdiff=0., update_grid=True, grid_update_num=10, loss_fn=None, lr=1., stop_grid_update_step=50, batch=-1,
-              small_mag_threshold=1e-16, small_reg_factor=1., metrics=None, sglr_avoid=False, save_fig=False, in_vars=None, out_vars=None, beta=3, save_fig_freq=1, img_folder='./video', device='cpu'):
+              small_mag_threshold=1e-16, small_reg_factor=1., metrics=None, sglr_avoid=False, save_fig=False, in_vars=None, out_vars=None, beta=3, save_fig_freq=1, img_folder='./video', device='cpu', custom_lr_param = []):
         '''
         training
 
@@ -800,6 +800,8 @@ def train(self, dataset, opt="LBFGS", steps=100, log=1, lamb=0., lamb_l1=1., lam
                 device   
             save_fig_freq : int
                 save figure every (save_fig_freq) step
+            custom_lr_param: list
+                custom learning rate parameters for Adam optimizer
 
         Returns:
         --------
@@ -851,7 +853,10 @@ def nonlinear(x, th=small_mag_threshold, factor=small_reg_factor):
         grid_update_freq = int(stop_grid_update_step / grid_update_num)
 
         if opt == "Adam":
-            optimizer = torch.optim.Adam(self.parameters(), lr=lr)
+            if custom_lr_param == []:
+                optimizer = torch.optim.Adam(self.parameters(), lr=lr)
+            else:
+                optimizer = torch.optim.Adam(custom_lr_param, lr=lr)
         elif opt == "LBFGS":
             optimizer = LBFGS(self.parameters(), lr=lr, history_size=10, line_search_fn="strong_wolfe", tolerance_grad=1e-32, tolerance_change=1e-32, tolerance_ys=1e-32)
 

kan/utils.py

@@ -2,6 +2,7 @@
 import torch
 from sklearn.linear_model import LinearRegression
 import sympy
+import pdb
 
 # sigmoid = sympy.Function('sigmoid')
 # name: (torch implementation, sympy implementation)
@@ -39,6 +40,7 @@ def create_dataset(f,
                    ranges = [-1,1],
                    train_num=1000, 
                    test_num=1000,
+                   step_size=None,
                    normalize_input=False,
                    normalize_label=False,
                    device='cpu',
@@ -56,6 +58,8 @@ def create_dataset(f,
             the number of training samples. Default: 1000.
         test_num : int
             the number of test samples. Default: 1000.
+        step_size : float
+            step_size within the specified range. Default: None.
         normalize_input : bool
             If True, apply normalization to inputs. Default: False.
         normalize_label : bool
@@ -64,6 +68,7 @@ def create_dataset(f,
             device. Default: 'cpu'.
         seed : int
             random seed. Default: 0.
+
 
     Returns:
     --------
@@ -87,13 +92,36 @@ def create_dataset(f,
     else:
         ranges = np.array(ranges)
 
-    train_input = torch.zeros(train_num, n_var)
-    test_input = torch.zeros(test_num, n_var)
-    for i in range(n_var):
-        train_input[:,i] = torch.rand(train_num,)*(ranges[i,1]-ranges[i,0])+ranges[i,0]
-        test_input[:,i] = torch.rand(test_num,)*(ranges[i,1]-ranges[i,0])+ranges[i,0]
-
-
+
+    if step_size is None:
+        train_input = torch.zeros(train_num, n_var)
+        test_input = torch.zeros(test_num, n_var)
+        for i in range(n_var):
+            train_input[:,i] = torch.rand(train_num,)*(ranges[i,1]-ranges[i,0])+ranges[i,0]
+            test_input[:,i] = torch.rand(test_num,)*(ranges[i,1]-ranges[i,0])+ranges[i,0]
+    else:
+        # generate warning that if step_size is provided, test_num is only the required argument
+        if test_num is not None:
+            print('Warning: if step_size is provided, test_num is only the required argument. train_num equals to all possible combinations minus test_num.')
+        def generate_grid(dimensions, ranges, step_size):
+            # Generate an array of values for each dimension
+            axisvalues = []
+            for _ in range(dimensions):
+                axisvalues.append(np.arange(ranges[_][0], ranges[_][1], step_size))
+
+            # Create the grid using numpy's meshgrid function
+            grid = np.meshgrid(*axisvalues, indexing='ij')
+
+            # Reshape the grid to list all possible combinations of coordinates
+            grid = np.stack(grid, axis=-1).reshape(-1, dimensions)
+
+            return grid
+
+        all_data = torch.from_numpy(generate_grid(n_var, ranges, step_size)).float()
+        train_input = all_data[np.random.choice(all_data.shape[0], all_data.shape[0]-test_num, replace=False)]
+        test_input = all_data[np.random.choice(all_data.shape[0], test_num, replace=False)]
+
+
     train_label = f(train_input)
     test_label = f(test_input)