182 lines
6.5 KiB
Python
182 lines
6.5 KiB
Python
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 (number of splines, number of B-spline bases (coeffcients), number of samples). The numbef of B-spline bases = number of grid points + k - 1.
|
|
|
|
Example
|
|
-------
|
|
>>> num_spline = 5
|
|
>>> num_sample = 100
|
|
>>> num_grid_interval = 10
|
|
>>> k = 3
|
|
>>> x = torch.normal(0,1,size=(num_spline, num_sample))
|
|
>>> grids = torch.einsum('i,j->ij', torch.ones(num_spline,), torch.linspace(-1,1,steps=num_grid_interval+1))
|
|
>>> B_batch(x, grids, k=k).shape
|
|
torch.Size([5, 13, 100])
|
|
'''
|
|
|
|
'''# x shape: (size, x); grid shape: (size, grid)
|
|
def extend_grid(grid, k_extend=0):
|
|
# pad k to left and right
|
|
# grid shape: (batch, 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)
|
|
grid = grid.to(device)
|
|
return grid
|
|
|
|
if extend == True:
|
|
grid = extend_grid(grid, k_extend=k)
|
|
|
|
grid = grid.unsqueeze(dim=2).to(device)
|
|
x = x.unsqueeze(dim=1).to(device)
|
|
|
|
if k == 0:
|
|
value = (x >= grid[:, :-1]) * (x < grid[:, 1:])
|
|
else:
|
|
B_km1 = B_batch(x[:, 0], grid=grid[:, :, 0], k=k - 1, extend=False, device=device)
|
|
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:]'''
|
|
|
|
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 (number of splines, number of samples)
|
|
grid : 2D torch.tensor)
|
|
shape (number of splines, number of grid points)
|
|
coef : 2D torch.tensor)
|
|
shape (number of splines, number of coef params). number of coef params = number of grid intervals + k
|
|
k : int
|
|
the piecewise polynomial order of splines.
|
|
device : str
|
|
devicde
|
|
|
|
Returns:
|
|
--------
|
|
y_eval : 2D torch.tensor
|
|
shape (number of splines, number of samples)
|
|
|
|
Example
|
|
-------
|
|
>>> num_spline = 5
|
|
>>> num_sample = 100
|
|
>>> num_grid_interval = 10
|
|
>>> k = 3
|
|
>>> x_eval = torch.normal(0,1,size=(num_spline, num_sample))
|
|
>>> grids = torch.einsum('i,j->ij', torch.ones(num_spline,), torch.linspace(-1,1,steps=num_grid_interval+1))
|
|
>>> coef = torch.normal(0,1,size=(num_spline, num_grid_interval+k))
|
|
>>> coef2curve(x_eval, grids, coef, k=k).shape
|
|
torch.Size([5, 100])
|
|
'''
|
|
# x_eval: (size, batch), grid: (size, grid), coef: (size, coef)
|
|
# coef: (size, coef), B_batch: (size, coef, batch), summer over coef
|
|
|
|
b_splines = B_batch(x_eval, grid, k=k) # (batch, in_dim, n_coef)
|
|
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 (number of splines, number of samples)
|
|
y_eval : 2D torch.tensor
|
|
shape (number of splines, number of samples)
|
|
grid : 2D torch.tensor
|
|
shape (number of splines, number of grid points)
|
|
k : int
|
|
the piecewise polynomial order of splines.
|
|
device : str
|
|
devicde
|
|
|
|
Example
|
|
-------
|
|
>>> num_spline = 5
|
|
>>> num_sample = 100
|
|
>>> num_grid_interval = 10
|
|
>>> k = 3
|
|
>>> x_eval = torch.normal(0,1,size=(num_spline, num_sample))
|
|
>>> y_eval = torch.normal(0,1,size=(num_spline, num_sample))
|
|
>>> grids = torch.einsum('i,j->ij', torch.ones(num_spline,), torch.linspace(-1,1,steps=num_grid_interval+1))
|
|
torch.Size([5, 13])
|
|
'''
|
|
'''
|
|
# x_eval: (size, batch); y_eval: (size, batch); grid: (size, grid); k: scalar
|
|
mat = B_batch(x_eval, grid, k, device=device).permute(0, 2, 1)
|
|
# coef = torch.linalg.lstsq(mat, y_eval.unsqueeze(dim=2)).solution[:, :, 0]
|
|
coef = torch.linalg.lstsq(mat.to(device), y_eval.unsqueeze(dim=2).to(device),
|
|
driver='gelsy' if device == 'cpu' else 'gels').solution[:, :, 0]'''
|
|
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, device=device).permute(0, 2, 1)
|
|
mat = B_batch(x_eval, grid, k) # (batch, in_dim, G+k)
|
|
mat = mat.permute(1,0,2)[:,None,:,:].expand(in_dim, out_dim, batch, n_coef) # (in_dim, out_dim, batch, n_coef)
|
|
# coef shape: (in_dim, outdim, G+k)
|
|
y_eval = y_eval.permute(1,2,0).unsqueeze(dim=3) # y_eval: (in_dim, out_dim, batch, 1)
|
|
#print(mat)
|
|
device = mat.device
|
|
coef = torch.linalg.lstsq(mat, y_eval,
|
|
driver='gelsy' if device == 'cpu' else 'gels').solution[:,:,:,0]
|
|
return coef
|
|
|
|
|
|
def extend_grid(grid, k_extend=0):
|
|
# pad k to left and right
|
|
# grid shape: (batch, 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 |