B-spline based shape function

1 year ago · 8573369193
9 changed files with 333 additions and 10 deletions
--- a/.gitignore
+++ b/.gitignore
@ -2,9 +2,9 @@
 .DS_Store
 # jupyter dev
 .ipynb_checkpoints
-*/.ipynb_checkpoints/*
+.ipynb_checkpoints/
 # vscode
-*/.vscode/
+.vscode/
 # python
 __pycache__
 # results folder as cache for data check
--- a/models/ANN_b_spline.py
+++ b/models/ANN_b_spline.py
@ -0,0 +1,38 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class ANN_b_spline_Model(nn.Module):
+    def __init__(self,in_dim=25+8,l1=36,l2=36*2,l3=36*3,l4=36*4,l5=36*5,l6=36*6,l7=36*5,l8=36*4,out_dim=36*2):
+        super().__init__()        
+        self.fc1=nn.Linear(in_dim,l1)
+        self.fc2=nn.Linear(l1,l2)
+        self.fc3=nn.Linear(l2,l3)
+        self.fc4=nn.Linear(l3,l4)
+        
+        self.fc5=nn.Linear(l4,l5)
+        self.fc6=nn.Linear(l5,l6)
+        self.fc7=nn.Linear(l6,l7)
+        self.fc8=nn.Linear(l7,l8)
+
+        
+        self.out=nn.Linear(l8,out_dim) # -> 6*6*2
+
+
+    def forward(self,x):
+        density = x[:25]
+        displace = x[25:]
+        
+        x = F.relu(self.fc1(x))
+        x = F.relu(self.fc2(x))
+        x = F.relu(self.fc3(x))
+        x = F.relu(self.fc4(x))
+
+        x = F.relu(self.fc5(x))
+        x = F.relu(self.fc6(x))
+        x = F.relu(self.fc7(x))
+        x = F.relu(self.fc8(x))
+
+        x = self.out(x)
+        
+        return x
--- a/options/base_options.py
+++ b/options/base_options.py
@ -21,14 +21,14 @@ class BaseOptions():
        # basic parameters
        parser.add_argument('--dataroot', type=str, default='./datasets', help='root path to datasets')
        parser.add_argument('--name', type=str, default='experiment_name', help='name of the experiment. It decides where to store samples and models')
-        parser.add_argument('--gpu_id', type=str, default='0', help='gpu ids: e.g. 0, 1, ... . use -1 for CPU')
-        parser.add_argument('--device', type=str, default='cuda:0', help='generate device with gpu_id and usable user device')
+        parser.add_argument('--gpu_id', type=str, default='-1', help='gpu ids: e.g. 0, 1, ... . use -1 for CPU')
+        parser.add_argument('--device', type=str, default='cuda', help='generate device with gpu_id and usable user device: e.g. cuda, mps, cpu')
        parser.add_argument('--checkpoints_dir', type=str, default='./checkpoints', help='models are saved here')
        parser.add_argument('--ms_ratio', type=int, default=5, help='multiscale ratio')
        parser.add_argument('--results_dir', type=str, default='./results', help='saves results here.')

        # model parameters
-        parser.add_argument('--model', type=str, default='ANN', help='chooses which model to use. [ANN | CNN | AutoEncoder]')
+        parser.add_argument('--model', type=str, default='ANN_b_spline', help='chooses which model to use. [ANN | CNN | AutoEncoder]')
        # dataset parameters
        # parser.add_argument('--resolution', type=str, default='180_60', help='data resolution. nelx_nely here')
        parser.add_argument('--nelx', type=int, default=180, help='num of elements on x-axis')
@ -106,8 +106,10 @@ class BaseOptions():
        self.print_options(opt)

        # set device with gpu id
-        if opt.gpu_id == -1:
+        if opt.gpu_id == '-1' or opt.device == 'cpu':
            opt.device = 'cpu'
+        elif opt.device == 'mps':
+            opt.device = f'mps:{opt.gpu_id}'
        else:
            opt.device = f'cuda:{opt.gpu_id}' if torch.cuda.is_available() else 'cpu'

--- a/train.py
+++ b/train.py
@ -10,8 +10,9 @@ import torch.nn.functional as F
 from utils.data_standardizer import standardization
 from utils.data_loader import data_loader_new
 from options.train_options import TrainOptions
-
+from utils.b_spline.surface_inter import surface_inter
 from models.ANN import ANN_Model
+from models.ANN_b_spline import ANN_b_spline_Model

 def train(X, Y, opt):
    if opt.is_standard:
@ -44,11 +45,18 @@ def train(X, Y, opt):
        for batch_idx, data_batch in enumerate(X_train):
            model.train() # 启用 batch normalization 和 dropout

-            pred_ShapeFunction = model(data_batch)
-            pred_U=pred_ShapeFunction.reshape(N,8) @ data_batch[25:]
+            # 线性插值
+            # pred_ShapeFunction = model(data_batch)
+            # pred_U=pred_ShapeFunction.reshape(N,8) @ data_batch[25:]
+            # loss=loss_function(pred_U,Y_train[batch_idx,:])
+
+            # B spline
+            control_points = model(data_batch)
+            pred_U = surface_inter(control_points)
            loss=loss_function(pred_U,Y_train[batch_idx,:])
+            
            # print(loss.item())
-            # loss.requires_grad_(True)
+            loss.requires_grad_(True) # 梯度不更新

            optimizer.zero_grad()
            loss.backward()
--- a/utils/b_spline/BaseFunction.py
+++ b/utils/b_spline/BaseFunction.py
@ -0,0 +1,28 @@
+def BaseFunction(i, k, u, knot):
+    '''
+    Calculate base function using Cox-deBoor function.
+    :param i: index of base function
+    :param k: order (degree + 1)
+    :param u: parameter
+    :param knot: knot vector
+    :return: base function
+    '''
+    Nik_u = 0
+    if k == 1:
+        if u >= knot[i] and u < knot[i + 1]:
+            Nik_u = 1.0
+        else:
+            Nik_u = 0.0
+    else:
+        length1 = knot[i + k - 1] - knot[i]
+        length2 = knot[i + k] - knot[i + 1]
+        if not length1 and not length2:
+            Nik_u = 0
+        elif not length1:
+            Nik_u = (knot[i + k] - u) / length2 * BaseFunction(i + 1, k - 1, u, knot)
+        elif not length2:
+            Nik_u = (u - knot[i]) / length1 * BaseFunction(i, k - 1, u, knot)
+        else:
+            Nik_u = (u - knot[i]) / length1 * BaseFunction(i, k - 1, u, knot) + \
+                    (knot[i + k] - u) / length2 * BaseFunction(i + 1, k - 1, u, knot)
+    return Nik_u
--- a/utils/b_spline/bspline_curve.py
+++ b/utils/b_spline/bspline_curve.py
@ -0,0 +1,101 @@
+import utils.b_spline.BaseFunction as bf
+import numpy as np
+
+
+def curve_interpolation(D, N, k, param, knot):
+    '''
+    Given a set of N data points, D0, D1, ..., Dn and a degree k,
+    find a B-spline curve of degree k defined by N control points
+    that passes all data points in the given order.
+    :param D: data points (N x 2)
+    :param N: the number of data points
+    :param k: degree
+    :param param: parameters
+    :param knot: knot vector
+    :return: control points (N x 2)
+    '''
+    Nik = np.zeros((N, N))
+
+    for i in range(N):
+        for j in range(N):
+            Nik[i][j] = bf.BaseFunction(j, k+1, param[i], knot)
+    Nik[N-1][N-1] = 1
+    print(Nik)
+    Nik_inv = np.linalg.inv(Nik)
+    print(Nik_inv)
+    P = []
+    for i in range(len(D)):
+        P.append(np.dot(Nik_inv, D[i]).tolist())
+    print(P)
+    return P
+
+
+def curve_approximation(D, N, H, k, param, knot):
+    '''
+    Given a set of N data points, D0, D1, ..., Dn, a degree k,
+    and a number H, where N > H > k >= 1, find a B-spline curve
+    of degree k defined by H control points that satisfies the
+    following conditions:
+        1. this curve contains the first and last data points;
+        2. this curve approximates the data polygon in the sense
+        of least square;
+    :param D: data points (N x 2)
+    :param H: the number of control points
+    :param k: degree
+    :param param: parameters
+    :param knot: knot vector
+    :return: control points (H x 2)
+    '''
+    P = []
+    if H >= N or H <= k:
+        print('Parameter H is out of range')
+        return P
+
+    for idx in range(len(D)):
+        P_ = np.zeros((1, H))
+        P_[0][0] = D[idx][0]
+        P_[0][H-1] = D[idx][N-1]
+        Qk = np.zeros((N - 2, 1))
+        Nik = np.zeros((N, H))
+        for i in range(N):
+            for j in range(H):
+                Nik[i][j] = bf.BaseFunction(j, k+1, param[i], knot)
+        # print(Nik)
+
+        for j in range(1, N - 1):
+            Qk[j - 1] = D[idx][j] - Nik[j][0] * P_[0][0] - Nik[j][H - 1] * P_[0][H - 1]
+
+        N_part = Nik[1: N - 1, 1: H - 1]
+        # print(N_part)
+        Q = np.dot(N_part.transpose(), Qk)
+        # print(Q)
+        M = np.dot(np.transpose(N_part), N_part)
+        P_[0][1:H - 1] = np.dot(np.linalg.inv(M), Q).transpose()
+        P.append(P_.tolist()[0])
+
+    return P
+
+
+def curve(P, N, k, param, knot):
+    '''
+    Calculate B-spline curve.
+    :param P: Control points
+    :param N: the number of control points
+    :param k: degree
+    :param param: parameters
+    :param knot: knot vector
+    :return: data point on the b-spline curve
+    '''
+    Nik = np.zeros((len(param), N))
+
+    for i in range(len(param)):
+        for j in range(N):
+            Nik[i][j] = bf.BaseFunction(j, k+1, param[i], knot)
+    Nik[len(param)-1][N - 1] = 1
+    # print(Nik)
+    # D = np.dot(Nik, P)
+    D = []
+    for i in range(len(P)):
+        D.append(np.dot(Nik, P[i]).tolist())
+    # print(D)
+    return D
--- a/utils/b_spline/bspline_surface.py
+++ b/utils/b_spline/bspline_surface.py
@ -0,0 +1,55 @@
+import numpy as np
+import utils.b_spline.parameter_selection as ps
+import utils.b_spline.bspline_curve as bc
+
+
+def surface(P, p, q, piece_uv, knot_uv):
+    '''
+    Calculate points on the surface.
+    :param P: control points
+    :param p: degree of u direction (u)
+    :param q: degree of v direction (v)
+    :param piece_uv: the number of points on u/v direction
+    :return: data points on the surface
+    '''
+    P_X = P[0]
+    P_Y = P[1]
+    P_Z = P[2]
+    M = len(P_X)
+    N = len(P_X[0])
+
+    param_u = np.linspace(0, 1, piece_uv[0])
+    param_v = np.linspace(0, 1, piece_uv[1])
+
+    Nik_u = np.zeros((piece_uv[0], M)).tolist()
+    Nik_v = np.zeros((piece_uv[1], N)).tolist()
+
+    D_tmp = [np.zeros((piece_uv[0], N)).tolist(),
+             np.zeros((piece_uv[0], N)).tolist(),
+             np.zeros((piece_uv[0], N)).tolist()]
+
+    knot_u = knot_uv[0]
+    for i in range(N):
+        P_control_X = [x[i] for x in P_X]
+        P_control_Y = [y[i] for y in P_Y]
+        P_control_Z = [z[i] for z in P_Z]
+        P_control = [P_control_X, P_control_Y, P_control_Z]
+        D_col = bc.curve(P_control, M, p, param_u, knot_u)
+        for j in range(len(D_col[0])):
+            D_tmp[0][j][i] = D_col[0][j]
+            D_tmp[1][j][i] = D_col[1][j]
+            D_tmp[2][j][i] = D_col[2][j]
+
+    D = [np.zeros((piece_uv[0], piece_uv[1])).tolist(),
+         np.zeros((piece_uv[0], piece_uv[1])).tolist(),
+         np.zeros((piece_uv[0], piece_uv[1])).tolist()]
+
+    knot_v = knot_uv[1]
+    for i in range(piece_uv[0]):
+        P_control = [D_tmp[0][i], D_tmp[1][i], D_tmp[2][i]]
+        D_ = bc.curve(P_control, N, q, param_v, knot_v)
+        D[0][i] = D_[0]
+        D[1][i] = D_[1]
+        D[2][i] = D_[2]
+
+    return D
--- a/utils/b_spline/parameter_selection.py
+++ b/utils/b_spline/parameter_selection.py
@ -0,0 +1,71 @@
+import numpy as np
+
+
+def uniform_spaced(n):
+    '''
+    Calculate parameters using the uniform spaced method.
+    :param n: the number of the data points
+    :return: parameters
+    '''
+    parameters = np.linspace(0, 1, n)
+    return parameters
+
+
+def chord_length(n, P):
+    '''
+    Calculate parameters using the chord length method.
+    :param n: the number of the data points
+    :param P: data points
+    :return: parameters
+    '''
+    parameters = np.zeros((1, n))
+    for i in range(1, n):
+        dis = 0
+        for j in range(len(P)):
+            dis = dis + (P[j][i] - P[j][i-1])**2
+        dis = np.sqrt(dis)
+        parameters[0][i] = parameters[0][i-1] + dis
+    for i in range(1, n):
+        parameters[0][i] = parameters[0][i]/parameters[0][n-1]
+    return parameters[0]
+
+
+def centripetal(n, P):
+    '''
+    Calculate parameters using the centripetal method.
+    :param n: the number of data points
+    :param P: data points
+    :return: parameters
+    '''
+    a = 0.5
+    parameters = np.zeros((1, n))
+    for i in range(1, n):
+        dis = 0
+        for j in range(len(P)):
+            dis = dis + (P[j][i]-P[j][i-1])**2
+        dis = np.sqrt(dis)
+        parameters[0][i] = parameters[0][i-1] + np.power(dis, a)
+    for i in range(1, n):
+        parameters[0][i] = parameters[0][i] / parameters[0][n-1]
+    return parameters[0]
+
+
+def knot_vector(param, k, N):
+    '''
+    Generate knot vector.
+    :param param: parameters
+    :param k: degree
+    :param N: the number of data points
+    :return: knot vector
+    '''
+    m = N + k
+    knot = np.zeros((1, m+1))
+    for i in range(k + 1):
+        knot[0][i] = 0
+    for i in range(m - k, m + 1):
+        knot[0][i] = 1
+    for i in range(k + 1, m - k):
+        for j in range(i - k, i):
+            knot[0][i] = knot[0][i] + param[j]
+        knot[0][i] = knot[0][i] / k
+    return knot[0]
--- a/utils/b_spline/surface_inter.py
+++ b/utils/b_spline/surface_inter.py
@ -0,0 +1,20 @@
+import utils.b_spline.parameter_selection as ps
+import numpy as np
+import utils.b_spline.bspline_curve as bc
+import utils.b_spline.bspline_surface as bs
+import matplotlib.pyplot as plt
+from mpl_toolkits.mplot3d import Axes3D
+import torch
+
+def surface_inter(P_control):
+    k = q = 2
+    piece_uv = [6, 6]
+    knot_uv = [[0,0,0, 0.3,0.5,0.7, 1,1,1], [0,0,0, 0.3,0.5,0.7, 1,1,1]]
+    
+    P_control = torch.cat((P_control.reshape(2,6,6), torch.zeros(1, 6, 6)), dim=0)
+    P_control = P_control.detach().numpy().tolist()
+    
+    P_piece = bs.surface(P_control, k, q, piece_uv, knot_uv)
+    P_piece = torch.Tensor(P_piece)[:2,:,:].permute(1,2,0).reshape(72)
+    
+    return P_piece