Overlay a number of colormaps with matplotlib - python

I have a total of 16 color maps which look as follows:
Is there a way to overlay all the maps while retaining their color? That means, I want to obtain a final image which consists of 16 distributions with different colors. I've been searching a lot, but unfortunately didn't find anything good yet. Thank's a lot!
For reproduction, the code looks as follows:
import torch
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import colors
def softmax(logit_map):
bn, kn, h, w = logit_map.shape
map_norm = F.softmax(logit_map.reshape(bn, kn, -1), dim=2).reshape(bn, kn, h, w)
return map_norm
def get_mu_and_prec(part_maps, device, scal):
"""
Calculate mean for each channel of part_maps
:param part_maps: tensor of part map activations [bn, n_part, h, w]
:return: mean calculated on a grid of scale [-1, 1]
"""
bn, nk, h, w = part_maps.shape
y_t = torch.linspace(-1., 1., h).reshape(h, 1).repeat(1, w).unsqueeze(-1)
x_t = torch.linspace(-1., 1., w).reshape(1, w).repeat(h, 1).unsqueeze(-1)
meshgrid = torch.cat((y_t, x_t), dim=-1).to(device) # 64 x 64 x 2
mu = torch.einsum('ijl, akij -> akl', meshgrid, part_maps) # bn x nk x 2
mu_out_prod = torch.einsum('akm,akn->akmn', mu, mu)
mesh_out_prod = torch.einsum('ijm,ijn->ijmn', meshgrid, meshgrid)
stddev = torch.einsum('ijmn,akij->akmn', mesh_out_prod, part_maps) - mu_out_prod
a_sq = stddev[:, :, 0, 0]
a_b = stddev[:, :, 0, 1]
b_sq_add_c_sq = stddev[:, :, 1, 1]
eps = 1e-12
a = torch.sqrt(a_sq + eps) # Σ = L L^T Prec = Σ^-1 = L^T^-1 * L^-1 ->looking for L^-1 but first L = [[a, 0], [b, c]
b = a_b / (a + eps)
c = torch.sqrt(b_sq_add_c_sq - b ** 2 + eps)
z = torch.zeros_like(a)
det = (a * c).unsqueeze(-1).unsqueeze(-1)
row_1 = torch.cat((c.unsqueeze(-1), z.unsqueeze(-1)), dim=-1).unsqueeze(-2)
row_2 = torch.cat((-b.unsqueeze(-1), a.unsqueeze(-1)), dim=-1).unsqueeze(-2)
L_inv = scal / (det + eps) * torch.cat((row_1, row_2), dim=-2) # L^⁻1 = 1/(ac)* [[c, 0], [-b, a]
return mu, L_inv
def get_heat_map(mu, L_inv, device):
h, w, nk = 64, 64, L_inv.shape[1]
y_t = torch.linspace(-1., 1., h).reshape(h, 1).repeat(1, w).unsqueeze(-1)
x_t = torch.linspace(-1., 1., w).reshape(1, w).repeat(h, 1).unsqueeze(-1)
y_t_flat = y_t.reshape(1, 1, 1, -1)
x_t_flat = x_t.reshape(1, 1, 1, -1)
mesh = torch.cat((y_t_flat, x_t_flat), dim=-2).to(device)
dist = mesh - mu.unsqueeze(-1)
proj_precision = torch.einsum('bnik, bnkf -> bnif', L_inv, dist) ** 2 # tf.matmul(precision, dist)**2
proj_precision = torch.sum(proj_precision, -2) # sum x and y axis
heat = 1 / (1 + proj_precision)
heat = heat.reshape(-1, nk, h, w) # bn number parts width height
return heat
color_list = ['black', 'gray', 'brown', 'chocolate', 'orange', 'gold', 'olive', 'lawngreen', 'aquamarine',
'dodgerblue', 'midnightblue', 'mediumpurple', 'indigo', 'magenta', 'pink', 'springgreen']
fmap = torch.randn(1, 16, 64, 64)
fmap_norm = softmax(fmap)
mu, L_inv = get_mu_and_prec(fmap_norm, 'cpu', scal=5.)
heat_map = get_heat_map(mu, L_inv, "cpu")
for i in range(16):
cmap = colors.LinearSegmentedColormap.from_list('my_colormap',
['white', color_list[i]],
256)
plt.imshow(heat_map[0][i].numpy(), cmap=cmap)
plt.show()

Related

Is this code for Gaussian Mixture Model correct?

I am learning GMM to do color segmentation. I found a good resource online with the following GMM code:
import matplotlib.pyplot as plt
from matplotlib import style
style.use('fivethirtyeight')
import numpy as np
from scipy.stats import norm
np.random.seed(0)
X = np.linspace(-5, 5, num=20)
X0 = X * np.random.rand(len(X)) + 15 # Create data cluster 1
X1 = X * np.random.rand(len(X)) - 15 # Create data cluster 2
X2 = X * np.random.rand(len(X)) # Create data cluster 3
X_tot = np.stack((X0, X1, X2)).flatten() # Combine the clusters to get the random datapoints from above
class GM1D:
def __init__(self, X, iterations):
self.iterations = iterations
self.X = X
self.mu = None
self.pi = None
self.var = None
def run(self):
self.mu = [-8, 8, 5]
self.pi = [1 / 3, 1 / 3, 1 / 3]
self.var = [5, 3, 1]
for iter in range(self.iterations):
r = np.zeros((len(X_tot), 3))
for c, g, p in zip(range(3), [norm(loc=self.mu[0], scale=self.var[0]),
norm(loc=self.mu[1], scale=self.var[1]),
norm(loc=self.mu[2], scale=self.var[2])], self.pi):
r[:, c] = p * g.pdf(X_tot) # Write the probability that x belongs to gaussian c in column c.
for i in range(len(r)):
r[i] = r[i] / (np.sum(self.pi) * np.sum(r, axis=1)[i])
fig = plt.figure(figsize=(10, 10))
ax0 = fig.add_subplot(111)
for i in range(len(r)):
ax0.scatter(self.X[i], 0, c=np.array([r[i][0], r[i][1], r[i][2]]), s=100)
for g, c in zip([norm(loc=self.mu[0], scale=self.var[0]).pdf(np.linspace(-20, 20, num=60)),
norm(loc=self.mu[1], scale=self.var[1]).pdf(np.linspace(-20, 20, num=60)),
norm(loc=self.mu[2], scale=self.var[2]).pdf(np.linspace(-20, 20, num=60))], ['r', 'g', 'b']):
ax0.plot(np.linspace(-20, 20, num=60), g, c=c)
m_c = []
for c in range(len(r[0])):
m = np.sum(r[:, c])
m_c.append(m) # For each cluster c, calculate the m_c and add it to the list m_c
for k in range(len(m_c)):
self.pi[k] = (m_c[k] / np.sum(m_c)) # For each cluster c, calculate the fraction of points pi_c which belongs to cluster c
self.mu = np.sum(self.X.reshape(len(self.X), 1) * r, axis=0) / m_c
var_c = []
for c in range(len(r[0])):
var_c.append((1 / m_c[c]) * np.dot(((np.array(r[:, c]).reshape(60, 1)) * (self.X.reshape(len(self.X), 1) - self.mu[c])).T, (self.X.reshape(len(self.X), 1) - self.mu[c])))
plt.show()
GM1D = GM1D(X_tot, 10)
GM1D.run()
Now my understanding is that in the Maximation step of EM, we must update the gaussian parameters (covariance matrix, mean and size of the gaussian(pi_c) ) In the code given above, I can see the values of pi_c and mean(mu) being updated but I don't think the value of the covariance matrix is updated. However, when I run the code it seems to be working (?). Can someone please help me determine if the code is correct or not. The code is from the following resource

Rotating a crystal in NumPy

I want to rotate a 5-atom crystal defined by X,Y,Z coordinates of atoms by a random angle. My initial idea was to use an external package to generate a rotation matrix (https://github.com/qobilidop/randrot) and then multiplying this matrix by a vector, which defines the coordinates of a single atom. However, that did not work at all and all the atoms got dispersed. Here's a function I wrote for that purpose:
def rotation():
crystal = []
rotmat = np.asarray(randrot.generate(3)) #generates 3x3 rotation matrix
for x,y,z in zip(new_x, new_y, new_z):
vec = np.array([x,y,z])
rot = vec.dot(rotmat)
for elem in rot:
crystal.append(elem)
return np.array(crystal).reshape([5,3])
rotated = rotation()
ax.scatter(rotated[0], rotated[1], rotated[2], marker='.', s=100, color='green')
Here's how it looks (red is the initial placement, green is after rotation):
pyplot
Here is an example code that rotates given 3d points about a randomly generated rotation matrix, rotation matrix creation is taken from another answer.
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import math
# taken from https://stackoverflow.com/questions/6802577/rotation-of-3d-vector
def rotation_matrix(axis, theta):
"""
Return the rotation matrix associated with counterclockwise rotation about
the given axis by theta radians.
"""
axis = np.asarray(axis)
axis = axis / math.sqrt(np.dot(axis, axis))
a = math.cos(theta / 2.0)
b, c, d = -axis * math.sin(theta / 2.0)
aa, bb, cc, dd = a * a, b * b, c * c, d * d
bc, ad, ac, ab, bd, cd = b * c, a * d, a * c, a * b, b * d, c * d
return np.array([[aa + bb - cc - dd, 2 * (bc + ad), 2 * (bd - ac)],
[2 * (bc - ad), aa + cc - bb - dd, 2 * (cd + ab)],
[2 * (bd + ac), 2 * (cd - ab), aa + dd - bb - cc]])
# initial xyz coordinates
xs = [0, 1, 1, 1, 1, -1, -1, -1, -1]
ys = [0, 1, 1, -1, -1, 1, 1, -1, -1]
zs = [0, 1, -1, 1, -1, 1, -1, 1, -1]
atoms_initial = np.array([xs, ys, zs]).T
# specify rotation matrix parameters
# let us generate a random axis and angle for rotation
rotation_axis = np.random.uniform(low=0, high=1, size=3) # three numbers between 0 and 1
rotation_angle = np.random.uniform(low=0, high=2*np.pi, size=1) # random number between 0 and 2pi
print("Rotation axis:{}, rotation angle:{} radians".format(rotation_axis, rotation_angle))
# create our rotation matrix
rotmat = rotation_matrix(rotation_axis, rotation_angle)
# apply rotation matrix to our points
atoms_rotated = np.dot(atoms_initial, rotmat)
# draw
fig = plt.figure()
ax = Axes3D(fig)
ax.scatter(atoms_initial[:,0], atoms_initial[:,1], atoms_initial[:,2], marker='.', s=100, color='red')
ax.scatter(atoms_rotated[:,0], atoms_rotated[:,1], atoms_rotated[:,2], marker='.', s=100, color="green")
plt.show()

Using matplotlib to generate 3D surfaces from 2D matrices

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)

Adding an offset to one dimension of a tensor, the offset being the entry of another tensor

I am looking for an efficient way to code the sample below with TensorFlow. Below, it is coded in the most silly way with numpy.
The idea is to offset a tensor of size (batch_size, height, width, channels) on the height dimension using another tensor value. In other terms:
tensor_2[i, j, k, l] = tensor_1[i, j + tensor_offset[i, j, k, l], k, l]
Here's the code I'm using:
import numpy as np
import time
begin = time.time()
b, h, w ,c = 5, 256, 512, 20
offset = np.random.rand(b, h, w , c).astype(int)
image = np.ones((b, h, w ,c))
label = np.ones((b, h, w ,c))
label_offset = np.zeros ((b, h, w,c ))
loss = 0
count = 0
for i in range(b):
for j in range (h):
for k in range (w):
for l in range (c):
offset_ = j + np.int(offset [i,j,k,l])
if offset_ > 255:
pass
else:
label_offset[i,j,k,l] = label [i,offset_,k,l]
loss =+ label_offset[i,j,k,l]*np.log(image [i,j,k,l])
count=+1
loss = loss/count
end = time.time()
print ('duree:', end - begin)
You can do it using gather_nd, with indices[i, j, k, l]= [i, j+ tensor_offset [i,j,k,l], k, l] (so indices would be of dimension 5).
You can build it like that:
import numpy as np
import tensorflow as tf
b, h, w ,c = 2, 11, 13, 7
final_shape = [b, h, w, c]
offset = np.random.randint(low=0, high=h, size=final_shape)
image = np.random.randint(low=0, high=1000, size=final_shape)
input_tensor = tf.constant(image)
m1 = tf.transpose(tf.reshape(tf.tile(tf.range(b), multiples=[h * w * c]), [h, w, c, b]), perm=[3, 0, 1, 2])
m2 = tf.transpose(tf.reshape(tf.tile(tf.range(h), multiples=[b * w * c]), [b, w, c, h]), perm=[0, 3, 1, 2]) + offset
not_too_big = tf.less(m2, h)
m2_safe = tf.mod(m2, h) # Makes sure we don't go too far in the original array
m3 = tf.transpose(tf.reshape(tf.tile(tf.range(w), multiples=[b * h * c]), [b, h, c, w]), perm=[0, 1, 3, 2])
m4 = tf.reshape(tf.tile(tf.range(c), multiples=[b * h * w]), [b, h, w, c]) # No transposition needed here
indices = tf.stack([m1, m2_safe, m3, m4], axis=-1)
tmp = tf.gather_nd(input_tensor, indices)
output = tf.multiply(tmp, tf.cast(not_too_big, tmp.dtype)) # Sets all the values corresponding to j+offset>h to 0
EDIT: this works for me with the transposition.
#gdelab
I have made the following improvement to your code:
def tensor_offset (input_tensor, offset_tensor, batch, nbcl):
b, h, w ,c = batch, 256,512,nbcl
m = tf.reshape(tf.tile(tf.range(b), multiples=[w*h*c]), [h,w,c, b])
m1 = tf.reshape(tf.tile(tf.range(h), multiples=[w*b*c]), [b,w,c,h])
m2 = tf.reshape(tf.tile(tf.range(w), multiples=[h*b*c]), [b,h,c,w])
m2 = m2 +tf.transpose(tf.cast(offset_tensor,tf.int32),perm=[0, 1, 3, 2])
m3 = tf.reshape(tf.tile(tf.range(c), multiples=[h*b*w]), [b,h,w,c])
indices = tf.stack([tf.transpose(m,perm=[3,0,1,2]), tf.transpose(m1,perm=[0, 3, 1, 2]), tf.transpose(m2,perm=[0, 1, 3, 2]),m3], axis=-1)
paddings = tf.constant([[0, 0], [0, 0], [0,100], [0,0]])
output = tf.gather_nd(tf.pad(input_tensor, paddings), indices)
return output

Python 3D polynomial surface fit, order dependent

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

Categories