yms-kan 0.0.1__py3-none-any.whl

yms_kan/spline.py

@@ -0,0 +1,144 @@
+import torch
+
+
+def B_batch(x, grid, k=0, extend=True, device='cpu'):
+    '''
+    evaludate x on B-spline bases
+    
+    Args:
+    -----
+        x : 2D torch.tensor
+            inputs, shape (number of splines, number of samples)
+        grid : 2D torch.tensor
+            grids, shape (number of splines, number of grid points)
+        k : int
+            the piecewise polynomial order of splines.
+        extend : bool
+            If True, k points are extended on both ends. If False, no extension (zero boundary condition). Default: True
+        device : str
+            devicde
+    
+    Returns:
+    --------
+        spline values : 3D torch.tensor
+            shape (batch, in_dim, G+k). G: the number of grid intervals, k: spline order.
+      
+    Example
+    -------
+    >>> from yms_kan.spline import B_batch
+    >>> x = torch.rand(100,2)
+    >>> grid = torch.linspace(-1,1,steps=11)[None, :].expand(2, 11)
+    >>> B_batch(x, grid, k=3).shape
+    '''
+    
+    x = x.unsqueeze(dim=2)
+    grid = grid.unsqueeze(dim=0)
+    
+    if k == 0:
+        value = (x >= grid[:, :, :-1]) * (x < grid[:, :, 1:])
+    else:
+        B_km1 = B_batch(x[:,:,0], grid=grid[0], k=k - 1)
+        
+        value = (x - grid[:, :, :-(k + 1)]) / (grid[:, :, k:-1] - grid[:, :, :-(k + 1)]) * B_km1[:, :, :-1] + (
+                    grid[:, :, k + 1:] - x) / (grid[:, :, k + 1:] - grid[:, :, 1:(-k)]) * B_km1[:, :, 1:]
+    
+    # in case grid is degenerate
+    value = torch.nan_to_num(value)
+    return value
+
+
+
+def coef2curve(x_eval, grid, coef, k, device="cpu"):
+    '''
+    converting B-spline coefficients to B-spline curves. Evaluate x on B-spline curves (summing up B_batch results over B-spline basis).
+    
+    Args:
+    -----
+        x_eval : 2D torch.tensor
+            shape (batch, in_dim)
+        grid : 2D torch.tensor
+            shape (in_dim, G+2k). G: the number of grid intervals; k: spline order.
+        coef : 3D torch.tensor
+            shape (in_dim, out_dim, G+k)
+        k : int
+            the piecewise polynomial order of splines.
+        device : str
+            devicde
+        
+    Returns:
+    --------
+        y_eval : 3D torch.tensor
+            shape (batch, in_dim, out_dim)
+        
+    '''
+    
+    b_splines = B_batch(x_eval, grid, k=k)
+    y_eval = torch.einsum('ijk,jlk->ijl', b_splines, coef.to(b_splines.device))
+    
+    return y_eval
+
+
+def curve2coef(x_eval, y_eval, grid, k):
+    '''
+    converting B-spline curves to B-spline coefficients using least squares.
+    
+    Args:
+    -----
+        x_eval : 2D torch.tensor
+            shape (batch, in_dim)
+        y_eval : 3D torch.tensor
+            shape (batch, in_dim, out_dim)
+        grid : 2D torch.tensor
+            shape (in_dim, grid+2*k)
+        k : int
+            spline order
+        lamb : float
+            regularized least square lambda
+            
+    Returns:
+    --------
+        coef : 3D torch.tensor
+            shape (in_dim, out_dim, G+k)
+    '''
+    #print('haha', x_eval.shape, y_eval.shape, grid.shape)
+    batch = x_eval.shape[0]
+    in_dim = x_eval.shape[1]
+    out_dim = y_eval.shape[2]
+    n_coef = grid.shape[1] - k - 1
+    mat = B_batch(x_eval, grid, k)
+    mat = mat.permute(1,0,2)[:,None,:,:].expand(in_dim, out_dim, batch, n_coef)
+    #print('mat', mat.shape)
+    y_eval = y_eval.permute(1,2,0).unsqueeze(dim=3)
+    #print('y_eval', y_eval.shape)
+    device = mat.device
+    
+    #coef = torch.linalg.lstsq(mat, y_eval, driver='gelsy' if device == 'cpu' else 'gels').solution[:,:,:,0]
+    try:
+        coef = torch.linalg.lstsq(mat, y_eval).solution[:,:,:,0]
+    except:
+        print('lstsq failed')
+    
+    # manual psuedo-inverse
+    '''lamb=1e-8
+    XtX = torch.einsum('ijmn,ijnp->ijmp', mat.permute(0,1,3,2), mat)
+    Xty = torch.einsum('ijmn,ijnp->ijmp', mat.permute(0,1,3,2), y_eval)
+    n1, n2, n = XtX.shape[0], XtX.shape[1], XtX.shape[2]
+    identity = torch.eye(n,n)[None, None, :, :].expand(n1, n2, n, n).to(device)
+    A = XtX + lamb * identity
+    B = Xty
+    coef = (A.pinverse() @ B)[:,:,:,0]'''
+    
+    return coef
+
+
+def extend_grid(grid, k_extend=0):
+    '''
+    extend grid
+    '''
+    h = (grid[:, [-1]] - grid[:, [0]]) / (grid.shape[1] - 1)
+
+    for i in range(k_extend):
+        grid = torch.cat([grid[:, [0]] - h, grid], dim=1)
+        grid = torch.cat([grid, grid[:, [-1]] + h], dim=1)
+
+    return grid

yms_kan/tool.py

@@ -0,0 +1,304 @@
+import os
+
+import numpy as np
+import torch
+import wandb
+
+
+def initialize_results_file(results_file, result_info):
+    """
+    初始化结果文件，确保文件存在且第一行包含指定的内容。
+
+    参数:
+        results_file (str): 结果文件的路径。
+        result_info (list): 需要写入的第一行内容列表。
+    """
+    # 处理 result_info，在每个单词后添加两个空格
+    result_info_str = "  ".join(result_info) + '\n'
+    # 检查文件是否存在
+    if os.path.exists(results_file):
+        # 如果文件存在，读取第一行
+        with open(results_file, "r") as f:
+            first_line = f.readline().strip()
+        # 检查第一行是否与 result_info 一致
+        if first_line == result_info_str.strip():
+            print(f"文件 {results_file} 已存在且第一行已包含 result_info，不进行写入。")
+        else:
+            # 如果不一致，写入 result_info
+            with open(results_file, "w") as f:
+                f.write(result_info_str)
+            print(f"文件 {results_file} 已被重新初始化。")
+    else:
+        # 如果文件不存在，创建并写入 result_info
+        with open(results_file, "w") as f:
+            f.write(result_info_str)
+        print(f"文件 {results_file} 已创建并写入 result_info。")
+
+
+def write_results_file(file_path: str,
+                       data_dict: dict,
+                       column_order: list,
+                       float_precision: int = 5) -> None:
+    """
+    通用格式化文本行写入函数（支持列表形式数据）
+
+    参数：
+    file_path: 目标文件路径
+    data_dict: 包含数据的字典，键为列名，值为列表
+    column_order: 列顺序列表，元素为字典键
+    float_precision: 浮点数精度位数 (默认5位)
+    """
+    # 验证数据格式
+    rows = None
+    for key in data_dict:
+        if not isinstance(data_dict[key], list):
+            raise ValueError(f"Value for key '{key}' is not a list")
+        if rows is None:
+            rows = len(data_dict[key])
+        else:
+            if len(data_dict[key]) != rows:
+                raise ValueError("All lists in data_dict must have the same length")
+
+    # 辅助函数：格式化单个值
+    def format_value(value, column_name):
+        if isinstance(value, (int, np.integer)):
+            return f"{value:d}"
+        elif isinstance(value, (float, np.floating)):
+            if column_name in ['train_losses', 'val_losses']:
+                return f"{value:.{float_precision + 1}f}"
+            elif column_name == 'lrs':
+                return f"{value:.8f}"
+            else:
+                return f"{value:.{float_precision}f}"
+        elif isinstance(value, str):
+            return value
+        else:
+            return str(value)
+
+    # 计算列宽
+    column_widths = []
+    for col in column_order:
+        dict_key = 'val_accuracies' if col == 'accuracies' else col
+        if dict_key not in data_dict:
+            raise ValueError(f"Missing required column: {dict_key}")
+        values = data_dict[dict_key]
+
+        max_width = len(col)
+        for val in values:
+            fmt_val = format_value(val, col)
+            max_width = max(max_width, len(fmt_val))
+        column_widths.append(max_width)
+
+    # 生成格式化行
+    lines = []
+    for i in range(rows):
+        row = []
+        for j, col in enumerate(column_order):
+            dict_key = 'val_accuracies' if col == 'accuracies' else col
+            val = data_dict[dict_key][i]
+            fmt_val = format_value(val, col)
+
+            # 对齐处理
+            if j == len(column_order) - 1:
+                fmt_val = fmt_val.ljust(column_widths[j])
+            else:
+                fmt_val = fmt_val.rjust(column_widths[j])
+            row.append(fmt_val)
+        lines.append("  ".join(row) + '\n')
+
+    # 写入文件
+    with open(file_path, 'a', encoding='utf-8') as f:
+        f.writelines(lines)
+
+
+def append_to_results_file(file_path: str,
+                           data_dict: dict,
+                           column_order: list,
+                           float_precision: int = 5) -> None:
+    """
+    通用格式化文本行写入函数
+
+    参数：
+    file_path: 目标文件路径
+    data_dict: 包含数据的字典，键为列名
+    column_order: 列顺序列表，元素为字典键
+    float_precision: 浮点数精度位数 (默认5位)
+    """
+    # 检查 data_dict 中的值是否为列表
+    all_values_are_lists = all(isinstance(value, list) for value in data_dict.values())
+    if all_values_are_lists:
+        num_rows = len(next(iter(data_dict.values())))
+        # 逐行处理
+        for row_index in range(num_rows):
+            formatted_data = []
+            column_widths = []
+            for col in column_order:
+                # 处理字典键的别名
+                dict_key = 'val_accuracies' if col == 'accuracies' else col
+                # 如果键不存在，跳过该列
+                if dict_key not in data_dict:
+                    continue
+                value_list = data_dict[dict_key]
+                if row_index >= len(value_list):
+                    continue
+                value = value_list[row_index]
+
+                # 根据数据类型进行格式化
+                if isinstance(value, (int, np.integer)):
+                    fmt_value = f"{value:d}"
+                elif isinstance(value, (float, np.floating)):
+                    if col in ['train_losses', 'val_losses']:  # 如果列名是'train_losses'或'val_losses'，保留浮点数精度位数+1位
+                        fmt_value = f"{value:.{float_precision + 1}f}"
+                    elif col == 'lrs':  # 如果列名是'lrs'，保留8位小数
+                        fmt_value = f"{value:.8f}"
+                    else:
+                        fmt_value = f"{value:.{float_precision}f}"
+                elif isinstance(value, str):
+                    fmt_value = value
+                else:  # 处理其他类型转换为字符串
+                    fmt_value = str(value)
+
+                # 取列名长度和数值长度的最大值作为列宽
+                column_width = max(len(col), len(fmt_value))
+                column_widths.append(column_width)
+
+                # 应用列宽对齐
+                if col == column_order[-1]:  # 最后一列左边对齐
+                    fmt_value = fmt_value.ljust(column_width)
+                else:
+                    fmt_value = fmt_value.rjust(column_width)
+
+                formatted_data.append(fmt_value)
+
+            # 构建文本行并写入，列之间用两个空格分隔
+            if formatted_data:
+                line = "  ".join(formatted_data) + '\n'
+                with open(file_path, 'a', encoding='utf-8') as f:
+                    f.write(line)
+    else:
+        # 非列表情况，原逻辑处理
+        # 计算每列的最大宽度
+        column_widths = []
+        formatted_data = []
+        for col in column_order:
+            # 处理字典键的别名
+            dict_key = 'val_accuracies' if col == 'accuracies' else col
+            # 如果键不存在，跳过该列
+            if dict_key not in data_dict:
+                continue
+
+            value = data_dict[dict_key]
+
+            # 根据数据类型进行格式化
+            if isinstance(value, (int, np.integer)):
+                fmt_value = f"{value:d}"
+            elif isinstance(value, (float, np.floating)):
+                if col in ['train_losses', 'val_losses']:  # 如果列名是'train_losses'或'val_losses'，保留浮点数精度位数+1位
+                    fmt_value = f"{value:.{float_precision + 1}f}"
+                elif col == 'lrs':  # 如果列名是'lrs'，保留8位小数
+                    fmt_value = f"{value:.8f}"
+                else:
+                    fmt_value = f"{value:.{float_precision}f}"
+            elif isinstance(value, str):
+                fmt_value = value
+            else:  # 处理其他类型转换为字符串
+                fmt_value = str(value)
+
+            # 取列名长度和数值长度的最大值作为列宽
+            column_width = max(len(col), len(fmt_value))
+            column_widths.append(column_width)
+
+            # 应用列宽对齐
+            if col == column_order[-1]:  # 最后一列左边对齐
+                fmt_value = fmt_value.ljust(column_width)
+            else:
+                fmt_value = fmt_value.rjust(column_width)
+
+            formatted_data.append(fmt_value)
+
+        # 构建文本行并写入，列之间用两个空格分隔
+        if formatted_data:
+            line = "  ".join(formatted_data) + '\n'
+            with open(file_path, 'a', encoding='utf-8') as f:
+                f.write(line)
+
+
+# def append_to_results_file(file_path: str,
+#                            data_dict: dict,
+#                            column_order: list,
+#                            column_widths: list = None,
+#                            float_precision: int = 5) -> None:
+#     """
+#     通用格式化文本行写入函数
+#
+#     参数：
+#     file_path: 目标文件路径
+#     data_dict: 包含数据的字典，键为列名
+#     column_order: 列顺序列表，元素为字典键
+#     column_widths: 每列字符宽度列表 (可选)
+#     float_precision: 浮点数精度位数 (默认4位)
+#     """
+#     formatted_data = []
+#
+#     # 遍历指定列顺序处理数据
+#     for i, col in enumerate(column_order):
+#         # 处理字典键的别名
+#         if col == 'accuracies':
+#             dict_key = 'val_accuracies'
+#         else:
+#             dict_key = col
+#
+#         if dict_key not in data_dict:
+#             raise ValueError(f"Missing required column: {dict_key}")
+#
+#         value = data_dict[dict_key]
+#
+#         # 根据数据类型进行格式化
+#         if isinstance(value, (int, np.integer)):
+#             fmt_value = f"{value:d}"
+#         elif isinstance(value, (float, np.floating)):
+#             if col in ['train_losses', 'val_losses']:  # 如果列名是'train_losses'或'val_losses'，保留浮点数精度位数+1位
+#                 fmt_value = f"{value:.{float_precision + 1}f}"
+#             elif col == 'lr':  # 如果列名是'lr'，保留8位小数
+#                 fmt_value = f"{value:.8f}"
+#             else:
+#                 fmt_value = f"{value:.{float_precision}f}"
+#         elif isinstance(value, str):
+#             fmt_value = value
+#         else:  # 处理其他类型转换为字符串
+#             fmt_value = str(value)
+#
+#         # 应用列宽对齐
+#         if column_widths and i < len(column_widths):
+#             try:
+#                 if i == len(column_order) - 1:  # 最后一列左边对齐
+#                     fmt_value = fmt_value.ljust(column_widths[i])
+#                 else:
+#                     fmt_value = fmt_value.rjust(column_widths[i])
+#             except TypeError:  # 处理非字符串类型
+#                 if i == len(column_order) - 1:  # 最后一列左边对齐
+#                     fmt_value = str(fmt_value).ljust(column_widths[i])
+#                 else:
+#                     fmt_value = str(fmt_value).rjust(column_widths[i])
+#
+#         formatted_data.append(fmt_value)
+#
+#     # 构建文本行并写入
+#     line = '\t'.join(formatted_data) + '\n'
+#     with open(file_path, 'a', encoding='utf-8') as f:
+#         f.write(line)
+
+
+def get_wandb_key(key_path='tools/wandb_key.txt'):
+    with open(key_path, 'r', encoding='utf-8') as f:
+        key = f.read()
+    return key
+
+
+def wandb_use(project=None, name=None, key_path='tools/wandb_key.txt'):
+    run = None
+    if project is not None:
+        wandb_key = get_wandb_key(key_path)
+        wandb.login(key=wandb_key)
+        run = wandb.init(project=project, name=name)
+    return run

yms_kan/train_eval_utils.py

@@ -0,0 +1,175 @@
+import math
+import os
+import sys
+from enum import Enum, auto
+
+import numpy as np
+import torch
+from matplotlib import pyplot as plt
+from torch.optim.lr_scheduler import ReduceLROnPlateau
+from tqdm import tqdm
+
+from yms_kan import LBFGS
+
+
+class TaskType(Enum):
+    classification = auto()
+    zlm = auto()
+
+
+def train_val(model, dataset: dict, batch_size, batch_size_test, opt="LBFGS", epochs=100, lamb=0.,
+              lamb_l1=1., label=None, lamb_entropy=2., lamb_coef=0.,
+              lamb_coefdiff=0., update_grid=True, grid_update_num=10, loss_fn=None, lr=1., start_grid_update_step=-1,
+              stop_grid_update_step=100,
+              save_fig=False, in_vars=None, out_vars=None, beta=3, save_fig_freq=1, img_folder='./video',
+              singularity_avoiding=False, y_th=1000., reg_metric='edge_forward_spline_n'):
+    # result_info = ['epoch','train_losses', 'val_losses', 'regularize', 'accuracies',
+    #                'precisions', 'recalls', 'f1-scores']
+    # initialize_results_file(results_file, result_info)
+    all_predictions = []
+    all_labels = []
+    if lamb > 0. and not model.save_act:
+        print('setting lamb=0. If you want to set lamb > 0, set model.save_act=True')
+
+    old_save_act, old_symbolic_enabled = model.disable_symbolic_in_fit(lamb)
+    if label is not None:
+        label = label.to(model.device)
+
+    if loss_fn is None:
+        loss_fn = lambda x, y: torch.mean((x - y) ** 2)
+    else:
+        loss_fn = loss_fn
+
+    grid_update_freq = int(stop_grid_update_step / grid_update_num)
+
+    if opt == "Adam":
+        optimizer = torch.optim.Adam(model.get_params(), lr=lr)
+    elif opt == "LBFGS":
+        optimizer = LBFGS(model.get_params(), lr=lr, history_size=10, line_search_fn="strong_wolfe",
+                          tolerance_grad=1e-32, tolerance_change=1e-32, tolerance_ys=1e-32)
+    else:
+        optimizer = torch.optim.SGD(model.get_params(), lr=lr)
+
+    lr_scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.1, patience=5, min_lr=1e-9)
+
+    results = {'train_losses': [], 'val_losses': [], 'regularize': [], 'accuracies': [],
+               'precisions': [], 'recalls': [], 'f1-scores': []}
+
+    steps = math.ceil(dataset['train_input'].shape[0] / batch_size)
+
+    train_loss = torch.zeros(1).to(model.device)
+    reg_ = torch.zeros(1).to(model.device)
+
+    def closure():
+        nonlocal train_loss, reg_
+        optimizer.zero_grad()
+        pred = model.forward(batch_train_input, singularity_avoiding=singularity_avoiding, y_th=y_th)
+        loss = loss_fn(pred, batch_train_label)
+        if model.save_act:
+            if reg_metric == 'edge_backward':
+                model.attribute()
+            if reg_metric == 'node_backward':
+                model.node_attribute()
+            reg_ = model.get_reg(reg_metric, lamb_l1, lamb_entropy, lamb_coef, lamb_coefdiff)
+        else:
+            reg_ = torch.tensor(0.)
+        objective = loss + lamb * reg_
+        train_loss = (train_loss * batch_num + objective.detach()) / (batch_num + 1)
+        objective.backward()
+        return objective
+
+    if save_fig:
+        if not os.path.exists(img_folder):
+            os.makedirs(img_folder)
+
+    for epoch in range(epochs):
+
+        if epoch == epochs - 1 and old_save_act:
+            model.save_act = True
+
+        if save_fig and epoch % save_fig_freq == 0:
+            save_act = model.save_act
+            model.save_act = True
+
+        train_indices = np.arange(dataset['train_input'].shape[0])
+        np.random.shuffle(train_indices)
+        train_pbar = tqdm(range(steps), desc=f'Epoch {epoch + 1}/{epochs} Training', file=sys.stdout)
+        for batch_num in train_pbar:
+            step = epoch * steps + batch_num + 1
+            i = batch_num * batch_size
+            batch_train_id = train_indices[i:i + batch_size]
+            batch_train_input = dataset['train_input'][batch_train_id].to(model.device)
+            batch_train_label = dataset['train_label'][batch_train_id].to(model.device)
+
+            if step % grid_update_freq == 0 and step < stop_grid_update_step and update_grid and step >= start_grid_update_step:
+                model.update_grid(batch_train_input)
+
+            if opt == "LBFGS":
+                optimizer.step(closure)
+
+            else:
+                optimizer.zero_grad()
+                pred = model.forward(batch_train_input, singularity_avoiding=singularity_avoiding,
+                                     y_th=y_th)
+                loss = loss_fn(pred, batch_train_label)
+                if model.save_act:
+                    if reg_metric == 'edge_backward':
+                        model.attribute()
+                    if reg_metric == 'node_backward':
+                        model.node_attribute()
+                    reg_ = model.get_reg(reg_metric, lamb_l1, lamb_entropy, lamb_coef, lamb_coefdiff)
+                else:
+                    reg_ = torch.tensor(0.)
+                loss = loss + lamb * reg_
+                train_loss = (train_loss * batch_num + loss.detach()) / (batch_num + 1)
+                loss.backward()
+                optimizer.step()
+            train_pbar.set_postfix(loss=train_loss.item())
+
+        val_loss = torch.zeros(1).to(model.device)
+        with torch.no_grad():
+            test_indices = np.arange(dataset['test_input'].shape[0])
+            np.random.shuffle(test_indices)
+            test_steps = math.ceil(dataset['test_input'].shape[0] / batch_size_test)
+            test_pbar = tqdm(range(test_steps), desc=f'Epoch {epoch + 1}/{epochs} Validation', file=sys.stdout)
+            for batch_num in test_pbar:
+                i = batch_num * batch_size_test
+                batch_test_id = test_indices[i:i + batch_size_test]
+                batch_test_input = dataset['test_input'][batch_test_id].to(model.device)
+                batch_test_label = dataset['test_label'][batch_test_id].to(model.device)
+
+                outputs = model.forward(batch_test_input, singularity_avoiding=singularity_avoiding,
+                                        y_th=y_th)
+
+                loss = loss_fn(outputs, batch_test_label)
+
+                val_loss = (val_loss * batch_num + loss.detach()) / (batch_num + 1)
+                test_pbar.set_postfix(loss=loss.item(), val_loss=val_loss.item())
+                if label is not None:
+                    diffs = torch.abs(outputs - label)
+                    closest_indices = torch.argmin(diffs, dim=1)
+                    closest_values = label[closest_indices]
+                    all_predictions.extend(closest_values.detach().cpu().numpy())
+                    all_labels.extend(batch_test_label.detach().cpu().numpy())
+
+            lr_scheduler.step(val_loss)
+
+        results['train_losses'].append(train_loss.cpu().item())
+        results['val_losses'].append(val_loss.cpu().item())
+        results['regularize'].append(reg_.cpu().item())
+
+        if save_fig and epoch % save_fig_freq == 0:
+            model.plot(folder=img_folder, in_vars=in_vars, out_vars=out_vars, title="Step {}".format(epoch),
+                       beta=beta)
+            plt.savefig(img_folder + '/' + str(epoch) + '.jpg', bbox_inches='tight', dpi=100)
+            plt.close()
+            model.save_act = save_act
+
+    # append_to_results_file(results_file, results, result_info)
+    model.log_history('fit')
+    model.symbolic_enabled = old_symbolic_enabled
+    return results
+
+
+if __name__ == '__main__':
+    print(TaskType.zlm.value)