12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368 | def poisson_surface_reconstruction(P,N,gs=None,h=None,corner=None,stochastic=False,sigma_n=0.0,sigma=0.05,solve_subspace_dim=0,verbose=False,prior_fun=None):
"""
Runs Poisson Surface Reconstruction from a set of points and normals to output a scalar field that takes negative values inside the surface and positive values outside the surface.
Parameters
----------
P : (n,dim) numpy array
Coordinates of points in R^dim
N : (n,dim) numpy array
(Unit) normals at each point
gs : (dim,) numpy array
Number of grid points in each dimension
h : (dim,) numpy array
Grid spacing in each dimension
corner : (dim,) numpy array
Coordinates of the lower left corner of the grid
stochastic : bool, optional (default False)
Whether to use "Stochastic Poisson Surface Reconstruction" to output a mean and variance scalar field instead of just one scalar field
sigma_n : float, optional (default 0.0)
Noise level in the normals
sigma : float, optional (default 0.05)
Scalar global variance parameter
solve_subspace_dim : int, optional (default 0)
If > 0, use a subspace solver to solve the linear system. This is useful for large problems and essential in 3D.
verbose : bool, optional (default True)
Whether to print progress
prior_fun : function, optional (default None)
Function that takes a (m,dim) numpy array and returns a (m,dim) numpy array. This is used to specify a prior on the gradient of the scalar field (see Sec. 5 in "Stochastic Poisson Surface Reconstruction").
Returns
-------
scalar_mean : (gs[0],gs[1],...,gs[dim-1]) numpy array
Mean of the reconstructed scalar field
scalar_variance : (gs[0],gs[1],...,gs[dim-1]) numpy array
Variance of the reconstructed scalar field. This will only be part of the output if stochastic=True.
grid_vertices : list of (gs[0],gs[1],...,gs[dim-1],dim) numpy arrays
Grid vertices (each element in the list is one dimension), as in the output of np.meshgrid
Notes
-----
See [this jupyter notebook](https://colab.research.google.com/drive/1DOXbDmqzIygxoQ6LeX0Ewq7HP4201mtr?usp=sharing) for a full tutorial on how to use this function.
See also
--------
fd_interpolate, fd_grad, matrix_from_function, compactly_supported_normal, grid_neighbors
Examples
--------
```python
from scipy.stats import norm
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
from gpytoolbox.poisson_surface_reconstruction import poisson_surface_reconstruction, random_points_on_polyline, png2poly
# Generate random points on a polyline
poly = gpytoolbox.png2poly("test/unit_tests_data/illustrator.png")[0]
poly = poly - np.min(poly)
poly = poly/np.max(poly)
poly = 0.5*poly + 0.25
poly = 3*poly - 1.5
num_samples = 40
np.random.seed(2)
EC = edge_indices(poly.shape[0],closed=False)
P,I,_ = random_points_on_mesh(poly, EC, num_samples, return_indices=True)
vecs = poly[EC[:,0],:] - poly[EC[:,1],:]
vecs /= np.linalg.norm(vecs, axis=1)[:,None]
J = np.array([[0., -1.], [1., 0.]])
N = vecs @ J.T
N = N[I,:]
# Problem parameters
gs = np.array([50,50])
# Call to PSR
scalar_mean, scalar_var, grid_vertices = gpytoolbox.poisson_surface_reconstruction(P,N,gs=gs,solve_subspace_dim=0,verbose=True)
# The probability of each grid vertex being inside the shape
prob_out = 1 - norm.cdf(scalar_mean,0,np.sqrt(scalar_var))
gx = grid_vertices[0]
gy = grid_vertices[1]
# Plot mean and variance side by side with colormap
fig, ax = plt.subplots(1,3)
m0 = ax[0].pcolormesh(gx,gy,np.reshape(scalar_mean,gx.shape), cmap='RdBu',shading='gouraud', vmin=-np.max(np.abs(scalar_mean)), vmax=np.max(np.abs(scalar_mean)))
ax[0].scatter(P[:,0],P[:,1],30 + 0*P[:,0])
q0 = ax[0].quiver(P[:,0],P[:,1],N[:,0],N[:,1])
ax[0].set_title('Mean')
divider = make_axes_locatable(ax[0])
cax = divider.append_axes('right', size='5%', pad=0.05)
fig.colorbar(m0, cax=cax, orientation='vertical')
m1 = ax[1].pcolormesh(gx,gy,np.reshape(np.sqrt(scalar_var),gx.shape), cmap='plasma',shading='gouraud')
ax[1].scatter(P[:,0],P[:,1],30 + 0*P[:,0])
q1 = ax[1].quiver(P[:,0],P[:,1],N[:,0],N[:,1])
ax[1].set_title('Variance')
divider = make_axes_locatable(ax[1])
cax = divider.append_axes('right', size='5%', pad=0.05)
fig.colorbar(m1, cax=cax, orientation='vertical')
m2 = ax[2].pcolormesh(gx,gy,np.reshape(prob_out,gx.shape), cmap='viridis',shading='gouraud')
ax[2].scatter(P[:,0],P[:,1],30 + 0*P[:,0])
q2 = ax[2].quiver(P[:,0],P[:,1],N[:,0],N[:,1])
ax[2].set_title('Probability of being inside')
divider = make_axes_locatable(ax[2])
cax = divider.append_axes('right', size='5%', pad=0.05)
fig.colorbar(m2, cax=cax, orientation='vertical')
plt.show()
```
"""
# Set problem parameters to values that make sense
dim = P.shape[1]
assert(dim == N.shape[1])
if ((gs is None) and (h is None) and (corner is None)):
# Default to a grid that is 100x100x...x100
gs = 100*np.ones(dim,dtype=int)
envelope_mult = 1.5 # how tightly we want to envelope the data
if (gs is None):
assert(h is not None)
assert(corner is not None)
gs = np.floor((np.max(envelope_mult*P,axis=0) - corner)/h).astype(int)
# print(gs)
# print(gs)
elif ((h is None) or (corner is None)):
h = (np.max(envelope_mult*P,axis=0) - np.min(envelope_mult*P,axis=0))/gs
corner = np.min(envelope_mult*P,axis=0)
assert(gs.shape[0] == dim)
grid_length = gs*h
# This is grid we will obtain the final values on
grid_vertices = np.meshgrid(*[np.linspace(corner[dd], corner[dd] + (gs[dd]-1)*h[dd], gs[dd]) for dd in range(dim)])
# Kernel function for the Gaussian process
def kernel_fun(X,Y):
return compactly_supported_normal(X-Y,n=3,sigma=1.5*np.min(h))
# np.meshgrid(*[np.linspace(corner_dd[dd], corner_dd[dd] + (gs_dd[dd]-1)*h[dd], gs_dd[dd]) for dd in range(dim)])
eps = 0.000001 # very tiny value to regularize matrix rank
# Estimate sampling density at each point in P
W_weights = fd_interpolate(P,gs=(gs+1),h=h,corner=(corner-0.5*h))
image = (W_weights.T @ np.ones((N.shape[0],1)))/(np.prod(h))
image_blurred = gaussian_filter(np.reshape(image,gs+1,order='F'),sigma=3)
image_blurred_vectorized = np.reshape(image_blurred,(np.prod(gs+1),1),order='F')
sampling_density = W_weights @ image_blurred_vectorized
# Step 1: Gaussian process interpolation from N to a regular grid
if verbose:
print("Step 1: Gaussian process interpolation from N to a regular grid")
# Log time to compute the kernel matrix
import time
t0 = time.time()
means = []
covs = []
for dd in range(dim):
# Build a staggered grid in the dd-th dimension
corner_dd = corner.copy()
corner_dd[dd] += 0.5*h[dd]
gs_dd = gs.copy()
gs_dd[dd] -= 1
# generate grid vertices of dimension dim
if dim==2:
staggered_grid_vertices = np.meshgrid(*[np.linspace(corner_dd[dd], corner_dd[dd] + (gs_dd[dd]-1)*h[dd], gs_dd[dd]) for dd in range(dim)])
staggered_grid_vertices = np.array(staggered_grid_vertices).reshape(dim, -1).T
elif dim==3:
staggered_grid_vertices = np.meshgrid(*[np.linspace(corner_dd[dd], corner_dd[dd] + (gs_dd[dd]-1)*h[dd], gs_dd[dd]) for dd in range(dim)],indexing='ij')
staggered_grid_vertices = np.array(staggered_grid_vertices).reshape(dim, -1,order='F').T
if verbose:
t00 = time.time()
### Step 1.1.: Compute the matrix k1, which has size prod(gs_dd) x prod(gs_dd) and contains the kernel evaluations between all pairs of points in the staggered grid.
# We could compute this matrix easily, by running
# k1_slow = matrix_from_function(kernel_fun, staggered_grid_vertices, staggered_grid_vertices)
# However, this would be extremely slow, O(prod(gs_dd)^2). Instead, we use the fact that the kernel is compactly supported, and only compute the kernel evaluations between pairs of points that are within a second-order neighborhood of each other.
# Find the neighbors of each point in the staggered grid
neighbor_rows = grid_neighbors(gs_dd,include_diagonals=True,include_self=True, order=2)
# Find one cell with no out of bounds neighbors
min_ind = np.min(neighbor_rows,axis=0)
valid_ind = np.argwhere(min_ind>=0)[0]
# What is the center of said cell
center_sample_cell = staggered_grid_vertices[valid_ind,:]
# And the coordinates of its second-order neighbors
neighbors_sample_cell = staggered_grid_vertices[np.squeeze(neighbor_rows[:,valid_ind]),:]
# Evaluate the kernel function just for this cell
center_sample_cell_tiled = np.tile(center_sample_cell,(neighbors_sample_cell.shape[0],1))
values_sample_cell = kernel_fun(center_sample_cell_tiled,neighbors_sample_cell)
# Thanks to the grid structure, the values for every cell will be the same, so we can just tile the values_sample_cell vector
V = np.tile(values_sample_cell,(np.prod(gs_dd),1)).T
I = np.tile(np.arange(np.prod(gs_dd)), (neighbor_rows.shape[0],1))
J = neighbor_rows
# Remove out-of-bounds indices
V[J<0] = 0
J[J<0] = 0
# And build k1:
k1_fast = csr_matrix((V.ravel(), (I.ravel(), J.ravel())), shape=(np.prod(gs_dd),np.prod(gs_dd)))
k1 = k1_fast
if verbose:
print("Time to compute k1: ", time.time() - t00)
t00 = time.time()
### Step 1.2: Building k2, the matrix of kernel evaluations between the points in the staggered grid and the points in P. Once again, we could compute this easily with
# k2_slow = matrix_from_function(kernel_fun, P, staggered_grid_vertices)
# But this would be very slow, of order O(N*prod(gs_dd)). Instead, we use the grid structure to compute the cell that each point in P falls into and then only evaluate the kernel on said cell and its third-order neighborhood.
# Find the cell that P falls into
P_cells = np.floor((P - np.tile(corner_dd,(P.shape[0],1))) / np.tile(h,(P.shape[0],1))).astype(int)
if dim==2:
P_cells = np.ravel_multi_index((P_cells[:,0], P_cells[:,1]), dims=gs_dd,order='F')
else:
P_cells = np.ravel_multi_index((P_cells[:,0], P_cells[:,1], P_cells[:,2]), dims=gs_dd,order='F')
# Find the neighbors of each P-associated cell
order_neighbors = 3
neighbors = np.arange(-order_neighbors,order_neighbors+1,dtype=int)
for dd2 in range(dim-1):
neighbors_along_dim =np.arange(-order_neighbors,order_neighbors+1,dtype=int)*np.prod(gs_dd[:dd2+1])
previous_neighbors = np.kron(neighbors,np.ones(neighbors_along_dim.shape[0],dtype=int))
neighbors_along_dim_repeated = np.kron(np.ones(neighbors.shape[0],dtype=int),neighbors_along_dim)
neighbors = previous_neighbors + neighbors_along_dim_repeated
# The neighbors give us the sparsity pattern of k2
I = np.tile(np.arange(P.shape[0]),(neighbors.shape[0],1)).T
J = np.tile(P_cells,(neighbors.shape[0],1)).T + np.tile(neighbors,(P_cells.shape[0],1))
valid_indices = (J>=0)*(J<np.prod(gs_dd))
J = J[valid_indices]
I = I[valid_indices]
# We evaluate the kernel function to get k2, but we do it with a known sparsity pattern, so this is O(P.shape[0]) instead of O(P.shape[0]*prod(gs_dd)).
k2_fast = matrix_from_function(kernel_fun, P, staggered_grid_vertices,sparsity_pattern=[I,J])
k2 = k2_fast
if verbose:
print("Time to compute k2: ", time.time() - t00)
### Step 1.3: Build K3, the matrix of kernel evaluations between sample points. If we did a true GP, we would compute this with
# K3 = matrix_from_function(kernel_fun, P, P)
# In reality, we approximate K3 with a lumped covariance matrix:
# K3 = diags(np.squeeze(sampling_density)) + sigma_n*sigma_n*eye(P.shape[0])
# But we really only need its inverse:
K3_inv = diags(1/(np.squeeze(sampling_density) + sigma_n*sigma_n))
### Step 1.4: Solve for the mean and covariance of the vector field:
if (prior_fun is None):
means.append(k2.T @ K3_inv @ N[:,dd][:,None])
else:
# print(prior_fun(staggered_grid_vertices)[:,dd][:,None])
# print(prior_fun(staggered_grid_vertices)[:,dd][:,None].shape)
# print(N[:,dd][:,None].shape)
means.append(prior_fun(staggered_grid_vertices)[:,dd][:,None] + k2.T @ K3_inv @ (N[:,dd][:,None] - prior_fun(P)[:,dd][:,None]))
if stochastic:
covs.append(k1 - k2.T @ K3_inv @ k2)
# Concatenate the mean and covariance matrices
vector_mean = np.concatenate(means)
if stochastic:
vector_cov = sigma*sigma*block_diag(covs)
if verbose:
print("Total step 1 time: ", time.time() - t0)
# Step 2: Solve Poisson equation on the regular grid
if verbose:
print("Step 2: Solve Poisson equation on the regular grid")
t1 = time.time()
if verbose:
t10 = time.time()
# Get the gradient so that its transpose is the divergence
G = fd_grad(gs=gs,h=h)
# Compute the divergence
mean_divergence = G.T @ vector_mean
# Build Laplacian
L = G.T @ G
# Solve for the mean scalar field
mean_scalar, info = bicg(L + eps*eye(L.shape[0]),mean_divergence,atol=1e-10)
assert(info==0)
# Shift values of mean
W = fd_interpolate(P,gs=gs,h=h,corner=corner)
shift = np.sum((W @ mean_scalar) ) / P.shape[0]
mean_scalar = mean_scalar - shift
# ...and we're done, mean_scalar is computed
if verbose:
print("Time to compute mean PDE solution: ", time.time() - t10)
t10 = time.time()
if stochastic:
# In this case, we want to compute the variances of the scalar field
cov_divergence = G.T @ vector_cov @ G
if (solve_subspace_dim>0):
# Project the covariance matrix onto a subspace (fast)
if verbose:
print("Solving for covariance on grid using subspace method")
t20 = time.time()
# _, vecs = eigenfunctions_laplacian(solve_subspace_dim,gs,grid_length)
vecs = grid_laplacian_eigenfunctions(solve_subspace_dim,gs,grid_length)
if verbose:
print("Time to compute eigenfunctions: ", time.time() - t20)
t20 = time.time()
vecs = np.real(vecs)
#vals = vecs.transpose() @ (L+0.0001*eye(L.shape[0])) @ vecs
vals = np.sum(np.multiply(vecs.transpose()@(L+eps*eye(L.shape[0])),vecs.transpose()),axis=1)
vals = csr_matrix(diags(vals))
if verbose:
print("Time to compute eigenvalues: ", time.time() - t20)
t20 = time.time()
B = (vecs.transpose()@cov_divergence.astype(np.float32))@vecs
if verbose:
print("Time to project problem onto subspace: ", time.time() - t20)
t20 = time.time()
D = spsolve(vals,B)
#D = np.linalg.solve(vals,B)
# cov_small = np.linalg.solve(vals,D.transpose())
cov_small = spsolve(vals,D.transpose()).astype(np.float32)
if verbose:
print("Time to solve in subspace: ", time.time() - t20)
t20 = time.time()
var_scalar = np.sum(np.multiply(vecs@cov_small,vecs),axis=1)
if verbose:
print("Time to reproject to full space", time.time() - t20)
else:
# Solve directly for the covariance on the grid (slow)
if verbose:
print("Solving for covariance directly")
lu = splu(L+eps*eye(L.shape[0]))
D = lu.solve(cov_divergence.toarray())
cov_scalar = lu.solve(D.transpose())
var_scalar = np.diag(cov_scalar)
# Shift values of covariance (Q: is a constant shift enough?)
var_scalar = var_scalar - np.min(var_scalar) + sigma_n*sigma_n + eps
if verbose:
print("Time to compute covariance PDE solution: ", time.time() - t10)
if verbose:
print("Total step 2 time: ", time.time() - t1)
print("Total time: ", time.time() - t0)
return mean_scalar, var_scalar, grid_vertices
else:
if verbose:
print("Total step 2 time: ", time.time() - t1)
print("Total time: ", time.time() - t0)
return mean_scalar, grid_vertices
|