Related
I'm interested in plotting a real-valued function f(x,y,z)=a, where (x,y,z) is a 3D point on the sphere and a is a real number. I calculate the Cartesian coordinates of the points of the sphere as follows, but I have no clue on how to visualize the value of f on each of those points.
import plotly.graph_objects as go
import numpy as np
fig = go.Figure(layout=go.Layout(title=go.layout.Title(text=title), hovermode=False))
# Create mesh grid for spherical coordinates
phi, theta = np.mgrid[0.0:np.pi:100j, 0.0:2.0 * np.pi:100j]
# Get Cartesian mesh grid
x = np.sin(phi) * np.cos(theta)
y = np.sin(phi) * np.sin(theta)
z = np.cos(phi)
# Plot sphere surface
self.fig.add_surface(x=x, y=y, z=z, opacity=0.35)
fig.show()
I would imagine/expect/like a visualization like this
Additionally, I also have the gradient of f calculated in closed-form (i.e., for each (x,y,z) I calculate the 3D-dimensional gradient of f). Is there a way of plotting this vector field, similarly to what is shown in the figure above?
Here's an answer that's far from perfect, but hopefully that's enough for you to build on.
For the sphere itself, I don't know of any "shortcut" to do something like that in plotly, so my approach is simply to manually create a sphere mesh. Generating the vertices is simple, for example like you did - the slightly more tricky part is figuring out the vertex indices for the triangles (which depends on the vertex generation scheme). There are various algorithms to do that smoothly (i.e. generating a sphere with no "tip"), I hacked something crude just for the demonstration. Then we can use the Mesh3d object to display the sphere along with the intensities and your choice of colormap:
N = 100 # Sphere resolution (both rings and segments, can be separated to different constants)
theta, z = np.meshgrid(np.linspace(-np.pi, np.pi, N), np.linspace(-1, 1, N))
r = np.sqrt(1 - z ** 2)
x = r * np.cos(theta)
y = r * np.sin(theta)
x = x.ravel()
y = y.ravel()
z = z.ravel()
# Triangle indices
indices = np.arange(N * (N - 1) - 1)
i1 = np.concatenate([indices, (indices // N + 1) * N + (indices + 1) % N])
i2 = np.concatenate([indices + N, indices // N * N + (indices + 1) % N])
i3 = np.concatenate([(indices // N + 1) * N + (indices + 1) % N, indices])
# Point intensity function
def f(x, y, z):
return (np.cos(x * 2) + np.sin(y ** 2) + np.sin(z) + 3) / 6
fig = go.Figure(data=[
go.Mesh3d(
x=x,
y=y,
z=z,
colorbar_title='f(x, y, z)',
colorscale=[[0, 'gold'],
[0.5, 'mediumturquoise'],
[1, 'magenta']],
intensity = f(x, y, z),
i = i1,
j = i2,
k = i3,
name='y',
showscale=True
)
])
fig.show()
This yields the following interactive plot:
To add the vector field you can use the Cone plot; this requires some tinkering because when I simply draw the cones at the same x, y, z position as the sphere, some of the cones are partially or fully occluded by the sphere. So I generate another sphere, with a slightly larger radius, and place the cones there. I also played with some lighting parameters to make it black like in your example. The full code looks like this:
N = 100 # Sphere resolution (both rings and segments, can be separated to different constants)
theta, z = np.meshgrid(np.linspace(-np.pi, np.pi, N), np.linspace(-1, 1, N))
r = np.sqrt(1 - z ** 2)
x = r * np.cos(theta)
y = r * np.sin(theta)
x = x.ravel()
y = y.ravel()
z = z.ravel()
# Triangle indices
indices = np.arange(N * (N - 1) - 1)
i1 = np.concatenate([indices, (indices // N + 1) * N + (indices + 1) % N])
i2 = np.concatenate([indices + N, indices // N * N + (indices + 1) % N])
i3 = np.concatenate([(indices // N + 1) * N + (indices + 1) % N, indices])
# Point intensity function
def f(x, y, z):
return (np.cos(x * 2) + np.sin(y ** 2) + np.sin(z) + 3) / 6
# Vector field function
def grad_f(x, y, z):
return np.stack([np.cos(3 * y + 5 * x),
np.sin(z * y),
np.cos(4 * x - 3 * y + z * 7)], axis=1)
# Second sphere for placing cones
N2 = 50 # Smaller resolution (again rings and segments combined)
R2 = 1.05 # Slightly larger radius
theta2, z2 = np.meshgrid(np.linspace(-np.pi, np.pi, N2), np.linspace(-R2, R2, N2))
r2 = np.sqrt(R2 ** 2 - z2 ** 2)
x2 = r2 * np.cos(theta2)
y2 = r2 * np.sin(theta2)
x2 = x2.ravel()
y2 = y2.ravel()
z2 = z2.ravel()
uvw = grad_f(x2, y2, z2)
fig = go.Figure(data=[
go.Mesh3d(
x=x,
y=y,
z=z,
colorbar_title='f(x, y, z)',
colorscale=[[0, 'gold'],
[0.5, 'mediumturquoise'],
[1, 'magenta']],
intensity = f(x, y, z),
i = i1,
j = i2,
k = i3,
name='y',
showscale=True
),
go.Cone(
x=x2, y=y2, z=z2, u=uvw[:, 0], v=uvw[:, 1], w=uvw[:, 2], sizemode='absolute', sizeref=2, anchor='tail',
lighting_ambient=0, lighting_diffuse=0, opacity=.2
)
])
fig.show()
And yields this plot:
Hope this helps. There are a lot of tweaks to the display, and certainly better ways to construct a sphere mesh (e.g. see this article), so there should be a lot of freedom there (albeit at the cost of some work).
Good luck!
I have a radially symmetric function evaluated on a 3D Cartesian grid. How can I numerically calculate the radial derivative of the function?
For a simple example (spherical Gaussian), calculate derivatives df/dx, df/dy and df/dz:
# Parameters
start = 0
end = 5
n = 20
# Variables
x = np.linspace(start, end, num=n)
y = np.linspace(start, end, num=n)
z = np.linspace(start, end, num=n)
dx = (end - start) / n
dy = (end - start) / n
dz = (end - start) / n
x_grid, y_grid, z_grid = np.meshgrid(x, y, z)
eval_xyz = np.exp(-(x_grid ** 2 + y_grid ** 2 + z_grid ** 2))
# Allocate
df_dx = np.zeros((n, n, n))
df_dy = np.zeros((n, n, n))
df_dz = np.zeros((n, n, n))
# Calculate Cartesian gradient numerically
for x in range(eval_xyz.shape[0] - 1):
for y in range(eval_xyz.shape[1] - 1):
for z in range(eval_xyz.shape[2] - 1):
df_dx[x, y, z] = (eval_xyz[x + 1, y, z] - eval_xyz[x, y, z]) / dx
df_dy[x, y, z] = (eval_xyz[x, y + 1, z] - eval_xyz[x, y, z]) / dy
df_dz[x, y, z] = (eval_xyz[x, y, z + 1] - eval_xyz[x, y, z]) / dz
Is it then possible to easily calculate the radial derivative df/dr from the Cartesian derivatives?
The trick is to express the radial derivatives as sum of Cartesian derivatives, taking into account theta and phi at each point which can be expressed in Cartesian coordiantes as:
The code therefore becomes:
theta_val = theta(i * dx, j * dy, k * dz)
phi_val = phi(i * dx, j * dy)
df_dr[i, j, k] = df_dx[i, j, k] * np.sin(theta_val) * np.cos(phi_val) \
+ df_dy[i, j, k] * np.sin(theta_val) * np.sin(phi_val) \
+ df_dz[i, j, k] * np.cos(theta_val)
Where theta and phi are calculated carefully to deal with divide by zero
def theta(x, y, z):
if x == 0 and y == 0 and z == 0:
return 0
elif z == 0:
return np.pi / 2
elif x == 0 and y == 0:
return 0
else:
return np.arctan(np.sqrt(x ** 2 + y ** 2) / z)
def phi(x, y):
if x == 0 and y == 0:
return 0
elif x == 0:
return np.pi / 2
elif y == 0:
return 0
else:
return math.atan2(y, x)
Your own answer is a step in the right direction, but there are some issues both in the answer and in the code generating the Cartesian derivatives.
These lines have a problem:
x = np.linspace(start, end, num=n)
dx = (end - start) / n
The step size is actually (end-start)/(n-1).
Here:
x_grid, y_grid, z_grid = np.meshgrid(x, y, z)
df_dx[x, y, z] = (eval_xyz[x + 1, y, z] - eval_xyz[x, y, z]) / dx
you fell in the trap of meshgrid's default setting: meshgrid(np.arange(n1), np.arange(n2)) will return arrays in the shape (n2, n1) unless you add the parameter indexing='ij'. Because you have size n in all dimensions, you will not get indexing errors to alert you, but you might be spending a lot of time trying to debug why the numbers make no sense.
When you manipulate multidimensional arrays, it's a good idea to set the sizes in different directions to slightly different values, so that you can easily check that the array shapes are what you want them to be.
Also, you should generally evaluate the derivative as (f[i+1]-f[i-1])/(2*dx), which is correct up to the second order in x.
for x in range(eval_xyz.shape[0] - 1):
for y in range(eval_xyz.shape[1] - 1):
for z in range(eval_xyz.shape[2] - 1):
When working with numpy, you should always try to vectorize operations rather than writing out for loops that potentially need to iterate over thousands of elements.
Here is code that calculates the Cartesian derivative and then the radial derivative.
import numpy as np
def get_cartesian_gradient(f, xyzsteps):
"""For f shape (nx, ny, nz), return gradient as (3, nx, ny, nz) shape.
xyzsteps is a (3,) array.
Note: edge points of the gradient array are set to NaN.
(Exercise for the reader to implement those).
"""
fshape = f.shape
grad = np.full((3,) + fshape, np.nan, dtype=np.float64)
sl, sm, sr = slice(0, -2), slice(1, -1), slice(2, None)
# Note: multiplying is faster than dividing.
grad[0, sm, sm, sm] = (f[sr, sm, sm] - f[sl, sm, sm]) * (0.5/xyzsteps[0])
grad[1, sm, sm, sm] = (f[sm, sr, sm] - f[sm, sl, sm]) * (0.5/xyzsteps[1])
grad[2, sm, sm, sm] = (f[sm, sm, sr] - f[sm, sm, sl]) * (0.5/xyzsteps[2])
return grad
def get_dfdr_from_cartesian(grad, x1s, y1s, z1s):
"""Return df/dr array from gradient(f).
grad.shape must be (3, nx, ny, nz)
return shape (nx, ny, nz).
"""
_, nx, ny, nz = grad.shape
# we need sin(theta), cos(theta), sin(phi), and cos(phi)
# rxy: shape (nx, ny, 1)
rxy = np.sqrt(x1s.reshape(-1, 1, 1)**2 + y1s.reshape(1, -1, 1)**2)
# r: shape (nx, ny, nz)
r = np.sqrt(rxy**2 + z1s.reshape(1, 1, -1)**2)
# change zeros to NaN
r = np.where(r==0, np.nan, r)
rxy = np.where(rxy==0, np.nan, rxy)
cos_theta = z1s.reshape(1, 1, -1) / r
sin_theta = rxy / r
cos_phi = x1s.reshape(-1, 1, 1) / rxy
sin_phi = y1s.reshape(1, -1, 1) / rxy
# and the derivative
dfdr = (grad[0]*cos_phi + grad[1]*sin_phi)*sin_theta + grad[2]*cos_theta
return dfdr
x1s = np.linspace(-1, 1, 19)
y1s = np.linspace(-1, 1, 21)
z1s = np.linspace(-1, 1, 23)
xs, ys, zs = np.meshgrid(x1s, y1s, z1s, indexing='ij')
xyzsteps = [x1s[1]-x1s[0], y1s[1]-y1s[0], z1s[1]-z1s[0]]
def func(x, y, z):
return x**2 + y**2 + z**2
def dfdr_analytical(x, y, z):
r = np.sqrt(x**2 + y**2 + z**2)
return 2*r
# grad has shape (3, nx, ny, nz)
grad = get_cartesian_gradient(func(xs, ys, zs), xyzsteps)
dfdr = get_dfdr_from_cartesian(grad, x1s, y1s, z1s)
# test
diff = dfdr - dfdr_analytical(xs, ys, zs)
assert np.nanmax(np.abs(diff)) < 1e-14
Note that I've chosen to return NaN values for points on the z-axis, because df/dr is not defined there unless f(x,y,z) is rotationally symmetric around the z-axis and has df/dr=0 in all directions. This is something that is not guaranteed for an arbitrary dataset.
The reason for replacing zeros in the denominators by np.nan using np.where is because dividing by zero will give warning messages, whereas dividing by nan won't.
Using a 2d matrix in python, how can I create a 3d surface plot, where columns=x, rows=y and the values are the heights in z?
I can't understand how to creat 3D surface plot using matplotlib.
Maybe it's different from MatLab.
example:
from pylab import *
from mpl_toolkits.mplot3d import Axes3D
def p(eps=0.9, lmd=1, err=10e-3, m=60, n=40):
delta_phi = 2 * np.pi / m
delta_lmd = 2 / n
k = 1
P0 = np.zeros([m + 1, n + 1])
P = np.zeros([m + 1, n + 1])
GAP = 1
while GAP >= err:
k = k + 1
for i in range(0, m):
for j in range(0, n):
if (i == 1) or (j == 1) or (i == m + 1) or (i == n + 1):
P[i,j] = 0
else:
A = (1+eps*np.cos((i+1/2)*delta_phi))**3
B = (1+eps*np.cos((i-1/2)*delta_phi))**3
C = (lmd*delta_phi/delta_lmd)**2 * (1+eps*np.cos((i)*delta_phi))**3
D = C
E = A + B + C + D
F = 3*delta_phi*((1+eps*np.cos((i+1/2)*delta_phi))-(1+eps*np.cos((i-1/2)*delta_phi)))
P[i,j] = (A*P[i+1,j] + B*P[i-1,j] + C*P[i,j+1] + D*P[i,j-1] - F)/E
if P[i,j] < 0:
P[i,j] = 0
S = P.sum() - P0.sum()
T = P.sum()
GAP = S / T
P0 = P.copy()
return P, k
def main():
start = time.time()
eps = 0.9
lmd = 1
err = 10e-8
m = 60
n = 40
P, k = p()
fig = figure()
ax = Axes3D(fig)
X = np.linspace(0, 2*np.pi, m+1)
Y = np.linspace(-1, 1, n+1)
X, Y = np.meshgrid(X, Y)
#Z = P[0:m, 0:n]
#Z = Z.reshape(X.shape)
ax.set_xticks([0, np.pi/2, np.pi, np.pi*1.5, 2*np.pi])
ax.set_yticks([-1, -0.5, 0, 0.5, 1])
ax.plot_surface(X, Y, P)
show()
if __name__ == '__main__':
main()
ValueError: shape mismatch: objects cannot be broadcast to a single
shape
And the pic
pic by matplotlic
And I also use MatLab to generate,the pic:
pic by MatLab
I should think this is a problem of getting the notaton straight. A m*n matrix is a matrix with m rows and n columns. Hence Y should be of length m and X of length n, such that after meshgridding X,Y and P all have shape (m,n).
At this point there would be no need to reshape of reindex and just plotting
ax.plot_surface(X, Y, P)
would give your the desired result.
Let's assume if you have a matrix mat.
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits import mplot3d
h, w = mat.shape
plt.figure(figsize=(16, 8))
ax = plt.axes(projection='3d')
X, Y = np.meshgrid(np.arange(w), np.arange(h))
ax.plot_surface(X, Y, mat, rstride=1, cstride=1, cmap='viridis', edgecolor='none', antialiased=False)
I've trying to simulate a 2D Sérsic profile and then testing an extraction routine on it. However, when I do a test by extracting all the points lying along an ellipse supposedly aligned with an image, I get a periodic function. It is meant to be a straight line since all points along the ellipse should have equal intensity, although there will be a small amount of deviation due to rounding errors in the rough coordinate estimation (get_I()).
from __future__ import division
import numpy as np
import matplotlib.pyplot as plt
from scipy.interpolate import NearestNDInterpolator
def rotate(x, y, angle):
x1 = x*np.cos(angle) + y*np.sin(angle)
y1 = y*np.cos(angle) - x*np.sin(angle)
return x1, y1
def sersic_1d(R, mu0, h, n, zp=0):
exponent = (R / h) ** (1 / n)
I0 = np.exp((zp - mu0) / 2.5)
return I0 * np.exp(-1.* exponent)
def sersic_2d(x, y, e, i, mu0, h, n, zp=0):
xp, yp = rotate(x, y, i)
alpha = np.arctan2(yp, xp * (1-e))
a = xp / np.cos(alpha)
b = a * (1 - e)
# R2 = (a*a) + ((1 - (e*e)) * yp*yp)
return sersic_1d(a, mu0, h, n, zp)
def ellipse(x0, y0, a, e, i, theta):
b = a * (1 - e)
x = a * np.cos(theta)
y = b * np.sin(theta)
x, y = rotate(x, y, i)
return x + x0, y + y0
def get_I(x, y, Z):
return Z[np.round(x).astype(int), np.round(y).astype(int)]
if __name__ == '__main__':
n = np.linspace(-100,100,1000)
nx, ny = np.meshgrid(n, n)
Z = sersic_2d(nx, ny, 0.5, 0., 0, 50, 1, 25)
theta = np.linspace(0, 2*np.pi, 1000.)
a = 100.
e = 0.5
i = np.pi / 4.
x, y = ellipse(0, 0, a, e, i, theta)
I = get_I(x, y, Z)
plt.plot(I)
# plt.imshow(Z)
plt.show()
However, What I actually get is a massive periodic function. I've checked the alignment and it's correct and the float-> int rounding errors can't account for this kind of shift?
Any ideas?
There are two things that strike me as odd, one of which for sure is not what you wanted, the other I'm not sure about because astronomy is not my field of expertise.
The first is in your function get_I:
def get_I(x, y, Z):
return Z[np.round(x).astype(int), np.round(y).astype(int)]
When you call that function, x an y outline an ellipse, with its center at the origin (0,0). That means x and y both become negative at some point. The indexing you perfom in that function will then take values from the array's last elements, because Z[0,0] is in fact the top left corner of the image (which you plotted, but commented), while Z[-1, -1] is the bottom right corner. What you want is to take the values of Z that are on the ellipse contour, but both have to have the same center. To do that, you would first make sure you use an uneven amount of samples for n (which ultimately defines the shape of Z) and second, you would add an indexing offset:
def get_I(x, y, Z):
offset = Z.shape[0]//2
return Z[np.round(y).astype(int) + offset, np.round(x).astype(int) + offset]
...
n = np.linspace(-100,100,1001) # changed from 1000 to 1001 to ensure a point of origin is present and that the image exhibits point symmetry
Also notice that I changed the order of y and x in get_I: that's because you first index along the rows (for which we usually take the y-coordinate) and only then along the columns (which map to the x-coordinate in most conventions).
The second item that struck me as unusual is that your ellipse has its axes at an angle of pi/4 with respect to the horizontal axis, whereas your sersic (which maps to the 2D array of Z) does not have a tilt at all.
Changing all that, I end up with this code:
from __future__ import division
import numpy as np
import matplotlib.pyplot as plt
def rotate(x, y, angle):
x1 = x*np.cos(angle) + y*np.sin(angle)
y1 = y*np.cos(angle) - x*np.sin(angle)
return x1, y1
def sersic_1d(R, mu0, h, n, zp=0):
exponent = (R / h) ** (1 / n)
I0 = np.exp((zp - mu0) / 2.5)
return I0 * np.exp(-1.* exponent)
def sersic_2d(x, y, e, ang, mu0, h, n, zp=0):
xp, yp = rotate(x, y, ang)
alpha = np.arctan2(yp, xp * (1-e))
a = xp / np.cos(alpha)
b = a * (1 - e)
return sersic_1d(a, mu0, h, n, zp)
def ellipse(x0, y0, a, e, i, theta):
b = a * (1 - e) # half of a
x = a * np.cos(theta)
y = b * np.sin(theta)
x, y = rotate(x, y, i) # rotated by 45deg
return x + x0, y + y0
def get_I(x, y, Z):
offset = Z.shape[0]//2
return Z[np.round(y).astype(int) + offset, np.round(x).astype(int) + offset]
#return Z[np.round(y).astype(int), np.round(x).astype(int)]
if __name__ == '__main__':
n = np.linspace(-100,100,1001) # changed
nx, ny = np.meshgrid(n, n)
ang = 0;#np.pi / 4.
Z = sersic_2d(nx, ny, 0.5, ang=0, mu0=0, h=50, n=1, zp=25)
f, ax = plt.subplots(1,2)
dn = n[1]-n[0]
ax[0].imshow(Z, cmap='gray', aspect='equal', extent=[-100-dn/2, 100+dn/2, -100-dn/2, 100+dn/2])
theta = np.linspace(0, 2*np.pi, 1000.)
a = 20. # decreased long axis of ellipse to see the intensity-map closer to the "center of the galaxy"
e = 0.5
x, y = ellipse(0,0, a, e, ang, theta)
I = get_I(x, y, Z)
ax[0].plot(x,y) # easier to see where you want the intensities
ax[1].plot(I)
plt.show()
and this image:
The intensity variations look like quantisation noise to me, with the exception of the peaks, which are due to the asymptote in sersic_1d.
I am currently working with astronomical data among which I have comet images. I would like to remove the background sky gradient in these images due to the time of capture (twilight). The first program I developed to do so took user selected points from Matplotlib's "ginput" (x,y) pulled the data for each coordinate (z) and then gridded the data in a new array with SciPy's "griddata."
Since the background is assumed to vary only slightly, I would like to fit a 3d low order polynomial to this set of (x,y,z) points. However, the "griddata" does not allow for an input order:
griddata(points,values, (dimension_x,dimension_y), method='nearest/linear/cubic')
Any ideas on another function that may be used or a method for developing a leas-squares fit that will allow me to control the order?
Griddata uses a spline fitting. A 3rd order spline is not the same thing as a 3rd order polynomial (instead, it's a different 3rd order polynomial at every point).
If you just want to fit a 2D, 3rd order polynomial to your data, then do something like the following to estimate the 16 coefficients using all of your data points.
import itertools
import numpy as np
import matplotlib.pyplot as plt
def main():
# Generate Data...
numdata = 100
x = np.random.random(numdata)
y = np.random.random(numdata)
z = x**2 + y**2 + 3*x**3 + y + np.random.random(numdata)
# Fit a 3rd order, 2d polynomial
m = polyfit2d(x,y,z)
# Evaluate it on a grid...
nx, ny = 20, 20
xx, yy = np.meshgrid(np.linspace(x.min(), x.max(), nx),
np.linspace(y.min(), y.max(), ny))
zz = polyval2d(xx, yy, m)
# Plot
plt.imshow(zz, extent=(x.min(), y.max(), x.max(), y.min()))
plt.scatter(x, y, c=z)
plt.show()
def polyfit2d(x, y, z, order=3):
ncols = (order + 1)**2
G = np.zeros((x.size, ncols))
ij = itertools.product(range(order+1), range(order+1))
for k, (i,j) in enumerate(ij):
G[:,k] = x**i * y**j
m, _, _, _ = np.linalg.lstsq(G, z)
return m
def polyval2d(x, y, m):
order = int(np.sqrt(len(m))) - 1
ij = itertools.product(range(order+1), range(order+1))
z = np.zeros_like(x)
for a, (i,j) in zip(m, ij):
z += a * x**i * y**j
return z
main()
The following implementation of polyfit2d uses the available numpy methods numpy.polynomial.polynomial.polyvander2d and numpy.polynomial.polynomial.polyval2d
#!/usr/bin/env python3
import unittest
def polyfit2d(x, y, f, deg):
from numpy.polynomial import polynomial
import numpy as np
x = np.asarray(x)
y = np.asarray(y)
f = np.asarray(f)
deg = np.asarray(deg)
vander = polynomial.polyvander2d(x, y, deg)
vander = vander.reshape((-1,vander.shape[-1]))
f = f.reshape((vander.shape[0],))
c = np.linalg.lstsq(vander, f)[0]
return c.reshape(deg+1)
class MyTest(unittest.TestCase):
def setUp(self):
return self
def test_1(self):
self._test_fit(
[-1,2,3],
[ 4,5,6],
[[1,2,3],[4,5,6],[7,8,9]],
[2,2])
def test_2(self):
self._test_fit(
[-1,2],
[ 4,5],
[[1,2],[4,5]],
[1,1])
def test_3(self):
self._test_fit(
[-1,2,3],
[ 4,5],
[[1,2],[4,5],[7,8]],
[2,1])
def test_4(self):
self._test_fit(
[-1,2,3],
[ 4,5],
[[1,2],[4,5],[0,0]],
[2,1])
def test_5(self):
self._test_fit(
[-1,2,3],
[ 4,5],
[[1,2],[4,5],[0,0]],
[1,1])
def _test_fit(self, x, y, c, deg):
from numpy.polynomial import polynomial
import numpy as np
X = np.array(np.meshgrid(x,y))
f = polynomial.polyval2d(X[0], X[1], c)
c1 = polyfit2d(X[0], X[1], f, deg)
np.testing.assert_allclose(c1,
np.asarray(c)[:deg[0]+1,:deg[1]+1],
atol=1e-12)
unittest.main()
According to the principle of Least squares, and imitate Kington's style,
while move argument m to argument m_1 and argument m_2.
import numpy as np
import matplotlib.pyplot as plt
import itertools
# w = (Phi^T Phi)^{-1} Phi^T t
# where Phi_{k, j + i (m_2 + 1)} = x_k^i y_k^j,
# t_k = z_k,
# i = 0, 1, ..., m_1,
# j = 0, 1, ..., m_2,
# k = 0, 1, ..., n - 1
def polyfit2d(x, y, z, m_1, m_2):
# Generate Phi by setting Phi as x^i y^j
nrows = x.size
ncols = (m_1 + 1) * (m_2 + 1)
Phi = np.zeros((nrows, ncols))
ij = itertools.product(range(m_1 + 1), range(m_2 + 1))
for h, (i, j) in enumerate(ij):
Phi[:, h] = x ** i * y ** j
# Generate t by setting t as Z
t = z
# Generate w by solving (Phi^T Phi) w = Phi^T t
w = np.linalg.solve(Phi.T.dot(Phi), (Phi.T.dot(t)))
return w
# t' = Phi' w
# where Phi'_{k, j + i (m_2 + 1)} = x'_k^i y'_k^j
# t'_k = z'_k,
# i = 0, 1, ..., m_1,
# j = 0, 1, ..., m_2,
# k = 0, 1, ..., n' - 1
def polyval2d(x_, y_, w, m_1, m_2):
# Generate Phi' by setting Phi' as x'^i y'^j
nrows = x_.size
ncols = (m_1 + 1) * (m_2 + 1)
Phi_ = np.zeros((nrows, ncols))
ij = itertools.product(range(m_1 + 1), range(m_2 + 1))
for h, (i, j) in enumerate(ij):
Phi_[:, h] = x_ ** i * y_ ** j
# Generate t' by setting t' as Phi' w
t_ = Phi_.dot(w)
# Generate z_ by setting z_ as t_
z_ = t_
return z_
if __name__ == "__main__":
# Generate x, y, z
n = 100
x = np.random.random(n)
y = np.random.random(n)
z = x ** 2 + y ** 2 + 3 * x ** 3 + y + np.random.random(n)
# Generate w
w = polyfit2d(x, y, z, m_1=3, m_2=2)
# Generate x', y', z'
n_ = 1000
x_, y_ = np.meshgrid(np.linspace(x.min(), x.max(), n_),
np.linspace(y.min(), y.max(), n_))
z_ = np.zeros((n_, n_))
for i in range(n_):
z_[i, :] = polyval2d(x_[i, :], y_[i, :], w, m_1=3, m_2=2)
# Plot
plt.imshow(z_, extent=(x_.min(), y_.max(), x_.max(), y_.min()))
plt.scatter(x, y, c=z)
plt.show()
If anyone is looking for fitting a polynomial of a specific order (rather than polynomials where the highest power is equal to order, you can make this adjustment to the accepted answer's polyfit and polyval:
instead of:
ij = itertools.product(range(order+1), range(order+1))
which, for order=2 gives [(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (1, 2), (2, 0), (2, 1), (2, 2)] (aka up to a 4th degree polynomial), you can use
def xy_powers(order):
powers = itertools.product(range(order + 1), range(order + 1))
return [tup for tup in powers if sum(tup) <= order]
This returns [(0, 0), (0, 1), (0, 2), (1, 0), (1, 1), (2, 0)] for order=2