''' 
In this example we generate synthetic scattered data with added noise
and then fit it with a smoothed RBF interpolant. The interpolant in
this example is equivalent to a thin plate spline.
'''
import numpy as np
from rbf.interpolate import RBFInterpolant
import rbf.basis
import matplotlib.pyplot as plt
np.random.seed(1)

basis = rbf.basis.phs2
order = 1  

x_obs = np.random.random((100,2)) # observation points
u_obs = np.sin(2*np.pi*x_obs[:,0])*np.cos(2*np.pi*x_obs[:,1]) # signal
u_obs += np.random.normal(0.0,0.2,100) # add noise to signal
I = RBFInterpolant(x_obs,u_obs,sigma=0.1,basis=basis,order=order)
vals = np.linspace(0,1,200)
x_itp = np.reshape(np.meshgrid(vals,vals),(2,200*200)).T # interp points
u_itp = I(x_itp) # evaluate the interpolant
# plot the results
plt.tripcolor(x_itp[:,0],x_itp[:,1],u_itp,vmin=-1.1,vmax=1.1,cmap='viridis')
plt.scatter(x_obs[:,0],x_obs[:,1],s=100,c=u_obs,vmin=-1.1,vmax=1.1,
            cmap='viridis',edgecolor='k')
plt.xlim((0.05,0.95))
plt.ylim((0.05,0.95))
plt.colorbar()
plt.tight_layout()
plt.savefig('../figures/interpolate.a.png')
plt.show()

''' 
In this example we solve the Poisson equation over an L-shaped domain 
with fixed boundary conditions. We use the multiquadratic RBF (`mq`) 
'''
import numpy as np
from rbf.basis import mq
from rbf.geometry import contains
from rbf.nodes import min_energy_nodes
import matplotlib.pyplot as plt

# Define the problem domain with line segments.
vert = np.array([[0.0,0.0],[2.0,0.0],[2.0,1.0],
                 [1.0,1.0],[1.0,2.0],[0.0,2.0]])
smp = np.array([[0,1],[1,2],[2,3],[3,4],[4,5],[5,0]])
N = 500 # total number of nodes
nodes,idx,_ = min_energy_nodes(N,vert,smp) # generate nodes
eps = 5.0  # shape parameter
# create "left hand side" matrix
A = np.empty((N,N))
A[idx['interior']]  = mq(nodes[idx['interior']],nodes,eps=eps,diff=[2,0])
A[idx['interior']] += mq(nodes[idx['interior']],nodes,eps=eps,diff=[0,2])
A[idx['boundary:all']]  = mq(nodes[idx['boundary:all']],nodes,eps=eps)
# create "right hand side" vector
d = np.empty(N)
d[idx['interior']] = -1.0 # forcing term
d[idx['boundary:all']] = 0.0 # boundary condition
# Solve for the RBF coefficients
coeff = np.linalg.solve(A,d)
# interpolate the solution on a grid
xg,yg = np.meshgrid(np.linspace(-0.05,2.05,400),np.linspace(-0.05,2.05,400))
points = np.array([xg.flatten(),yg.flatten()]).T
u = mq(points,nodes,eps=eps).dot(coeff) # evaluate at the interp points
u[~contains(points,vert,smp)] = np.nan # mask outside points
ug = u.reshape((400,400)) # fold back into a grid
# make a contour plot of the solution
fig,ax = plt.subplots()
p = ax.contourf(xg,yg,ug,np.linspace(0.0,0.16,9),cmap='viridis')
ax.plot(nodes[:,0],nodes[:,1],'ko',markersize=4)
for s in smp:
  ax.plot(vert[s,0],vert[s,1],'k-',lw=2)

ax.set_aspect('equal')
fig.colorbar(p,ax=ax)
fig.tight_layout()
plt.show()

''' 
In this example we solve the Poisson equation over an L-shaped domain
with fixed boundary conditions. We use the RBF-FD method. The RBF-FD
method is preferable over the spectral RBF method because it is
scalable and does not require the user to specify a shape parameter
(assuming that we use odd order polyharmonic splines to generate the
weights).
'''
import numpy as np
from rbf.fd import weight_matrix, add_rows
from rbf.basis import phs3
from rbf.geometry import contains
from rbf.nodes import min_energy_nodes
import matplotlib.pyplot as plt
from scipy.sparse import csc_matrix
from scipy.sparse.linalg import spsolve
from scipy.interpolate import LinearNDInterpolator

# Define the problem domain with line segments.
vert = np.array([[0.0, 0.0], [2.0, 0.0], [2.0, 1.0],
                 [1.0, 1.0], [1.0, 2.0], [0.0, 2.0]])
smp = np.array([[0, 1], [1, 2], [2, 3], [3, 4], [4, 5], [5, 0]])

N = 500 # total number of nodes.

n = 20 # stencil size. Increase this will generally improve accuracy

basis = phs3 # radial basis function used to compute the weights. Odd
             # order polyharmonic splines (e.g., phs3) have always
             # performed well for me and they do not require the user
             # to tune a shape parameter. Use higher order
             # polyharmonic splines for higher order PDEs.

order = 2 # Order of the added polynomials. This should be at least as
          # large as the order of the PDE being solved (2 in this
          # case). Larger values may improve accuracy

# generate nodes
nodes, idx, _ = min_energy_nodes(N, vert, smp) 

# create the "left hand side" matrix. 
# create the component which evaluates the PDE
A_interior = weight_matrix(nodes[idx['interior']], nodes,
                           diffs=[[2, 0], [0, 2]], n=n, 
                         basis=basis, order=order)
# create the component for the fixed boundary conditions
A_boundary = weight_matrix(nodes[idx['boundary:all']], nodes, 
                           diffs=[0, 0]) 
# Add the components to the corresponding rows of `A`
A = csc_matrix((N, N))
A = add_rows(A,A_interior,idx['interior'])
A = add_rows(A,A_boundary,idx['boundary:all'])

# create "right hand side" vector
d = np.zeros((N,))
d[idx['interior']] = -1.0
d[idx['boundary:all']] = 0.0

# find the solution at the nodes
u_soln = spsolve(A, d) 

# interpolate the solution on a grid
xg, yg = np.meshgrid(np.linspace(-0.05, 2.05, 400), 
                     np.linspace(-0.05, 2.05, 400))
points = np.array([xg.flatten(), yg.flatten()]).T                    
u_itp = LinearNDInterpolator(nodes, u_soln)(points)
# mask points outside of the domain
u_itp[~contains(points, vert, smp)] = np.nan 
ug = u_itp.reshape((400, 400)) # fold back into a grid
# make a contour plot of the solution
fig, ax = plt.subplots()
p = ax.contourf(xg, yg, ug, np.linspace(-1e-6, 0.16, 9), cmap='viridis')
ax.plot(nodes[:, 0], nodes[:, 1], 'ko', markersize=4)
for s in smp:
ax.plot(vert[s, 0], vert[s, 1], 'k-', lw=2)

ax.set_aspect('equal')
fig.colorbar(p, ax=ax)
fig.tight_layout()
plt.show()