How to create a Z-value matrix for a given meshgrid, and x, y, z columns?

Question

Let's say we would like to do a heatmap from three 1D arrays x, y, z. The answer from Plotting a heat map from three lists: X, Y, Intensity works, but the Z = np.array(z).reshape(len(y), len(x)) line highly depends from the order in which the z-values have been added to the list.

As an example, the 2 following tests give the exact same plot, whereas it should not. Indeed:

• in test1, z=2 should be for x=100, y=7.

• in test2, z=2 should be for x=102, y=5.

How should we create the Z matrix in the function heatmap, so that it's not dependent on the order in which z-values are added?

import numpy as np
import matplotlib.pyplot as plt

def heatmap(x, y, z):
    z = np.array(z)
    x = np.unique(x)
    y = np.unique(y)
    X, Y = np.meshgrid(x, y)
    Z = np.array(z).reshape(len(y), len(x))
    plt.pcolormesh(X, Y, Z)
    plt.show()

### TEST 1
x, y, z = [], [], []
k = 0
for i in range(100, 120):
    for j in range(5, 15):
        x.append(i)
        y.append(j)
        z.append(k)
        k += 1

heatmap(x, y, z)


### TEST 2
x, y, z = [], [], []
k = 0
for j in range(5, 15):
    for i in range(100, 120):
        x.append(i)
        y.append(j)
        z.append(k)
        k += 1

heatmap(x, y, z)

Edit: Example 2: Let's say

x = [0, 2, 1, 2, 0, 1]
y = [3, 4, 4, 3, 4, 3]
z = [1, 2, 3, 4, 5, 6]

There should be a non-ambiguous way to go from the 3 arrays x, y, z to a heatmap-plottable meshgrid + a z-value matrix, even if x and y are in random order.

In this example x and y are in no particular order, quite random. How to do this?

Here a reshape like Z = np.array(z).reshape(len(y), len(x)) would be in wrong order.

Answer 1

If you look into np.meshgrid, you will see that there are two indexing schemes, "xy" (the default) and "ij". When using "xy" indexing, X (the np.meshgrid result) increases along the columns and Y increases along the rows. This is the opposite of "ij" indexing where X increase along the rows and Y across the columns.

In your Test 1, if you were to reshape x using the typical c-style ordering (np.reshape(x, (20, 10))), you'd see that the resulting array increases along the rows, so it is using "ij" (and y increases along the columns). In your Test 2, the reshape (np.reshape(x, (10, 20))) would result in the reshaped x increasing along the columns, so it is using "xy" indexing.

That said, you can adjust your heatmap function call to take a parameter for which indexing to use when calling np.meshgrid and also reshape Z to be the same shape as X/Y.

def heatmap(x, y, z, indexing="xy"):
    x = np.unique(x)
    y = np.unique(y)
    X, Y = np.meshgrid(x, y, indexing=indexing)
    Z = np.asarray(z).reshape(X.shape)
    fig, ax = plt.subplots()
    p = ax.pcolormesh(X, Y, Z)
    fig.show()
    return p, fig, ax

# test 1 heatmap call
heatmap(x, y, z, "ij")

# test 2 heatmap call
heatmap(x, y, z, "xy")

Edit:

To handle the case with random-ordered data, you can deal with it by sorting along x and y in the proper order to achieve the desired indexing result.

import numpy as np
import matplotlib.pyplot as plt

def random_to_meshgrid(x, y, z, indexing="xy"):
    x = np.asarray(x)
    y = np.asarray(y)
    z = np.asarray(z)

    nx_unique = len(np.unique(x))
    ny_unique = len(np.unique(y))

    if indexing == "xy":
        shape = (ny_unique, nx_unique)

        x_sort = np.argsort(x)
        x = x[x_sort]
        y = y[x_sort]
        z = z[x_sort]

        y_sort = np.argsort(y)
        x = x[y_sort]
        y = y[y_sort]
        z = z[y_sort]

    elif indexing == "ij":
        shape = (nx_unique, ny_unique)

        y_sort = np.argsort(y)
        x = x[y_sort]
        y = y[y_sort]
        z = z[y_sort]

        x_sort = np.argsort(x)
        x = x[x_sort]
        y = y[x_sort]
        z = z[x_sort]
    else:
        raise ValueError(f"Undefined indexing '{indexing}'.")

    X = x.reshape(shape)
    Y = y.reshape(shape)
    Z = z.reshape(shape)

    return X, Y, Z

def heatmap(x, y, z, indexing="xy"):
    X, Y, Z = random_to_meshgrid(x, y, z, indexing=indexing)
    fig, ax = plt.subplots()
    p = ax.pcolormesh(X, Y, Z)
    fig.show()
    return p, fig, ax

You can test it on your sample input like so:

x = [0, 2, 1, 2, 0, 1]
y = [3, 4, 4, 3, 4, 3]
z = [1, 2, 3, 4, 5, 6]

# xy
X, Y, Z = random_to_meshgrid(x, y, z, indexing="xy")

# ij
X, Y, Z = random_to_meshgrid(x, y, z, indexing="ij")