How to plot a contour plot (density) of a data file with 3 columns (x, y, density) with the script automatically picking the data array dimension?

I have a gnuplot script I use for plotting data that is as follows: (simplified)

0 0 1.9
0 1 1.92
0 2 1.93
0 3 1.98
0 4 1.89

1 0 1.8
1 1 1.83
1 2 1.79
1 3 1.8
1 4 1.86

2 0 1.5
2 1 1.6
2 2 1.4
2 3 1.55
2 4 1.62

3 0 1.4
3 1 1.35
3 2 1.36
3 3 1.34
3 4 1.3

4 0 1.1
4 1 1.2
4 2 1.15
4 3 1.05
4 4 1.06

The first column is the x-coordinate value, the second one is the y-coordinate and the third column is the value of the density.
So the plot for this data needs to be a contour plot of 5x5..

Right now I use the following gnuplot script that was given to me to plot this:

reset

set term png size 2048, 2048
set output "n2d.png"
filename = "n2d.dat"   ## electron density ne (x,y)
set cbrange [-1 : 2] 
set size ratio 1
set palette defined (-1 "black", -0.5 "blue", 0 "white", 1 "red", 1.25 "black")
set pm3d map
splot filename notitle
set ticslevel 0

It picks up the data array size automatically without having to specify, which makes it general for whether the data is 5x5 or 128x128 or something else.

I tried to produce a matplotlib script that does exactly this but it is not general, I am having to specify the size of the array. The script is:

import matplotlib
import matplotlib.pyplot as plt
import numpy as np


data = np.loadtxt("n2d.dat", usecols=range(0,3), dtype=np.float32)
datamod = np.transpose( data[:,2].reshape(5,5) ) 
#reshapes the data into a 2D array of 5x5 with all the density values
density = plt.imshow(datamod, cmap='viridis', origin='lower', interpolation='none')
plt.colorbar(density)
plt.savefig("density.png", dpi=300)

The python script does not have the same generality as the gnuplot script, i am having to specify the array size to be 5x5 with "reshape(5,5)".

How do I achieve the same generality as the gnuplot script with matplotlib?

I think this can do what you want:

import matplotlib
import matplotlib.pyplot as plt
import numpy as np


data = np.loadtxt("n2d.dat", usecols=range(0,3), dtype=np.float32)

nx = len(np.unique(data[:,0]))
ny = len(np.unique(data[:,1]))
assert nx*ny == len(data)

datamod = np.transpose( data[:,2].reshape(nx, ny) ) 

#reshapes the data into a 2D array of 5x5 with all the density values
density = plt.imshow(datamod, cmap='viridis', origin='lower', interpolation='none')
plt.colorbar(density)
plt.savefig("density.png", dpi=300)

Note the assert statement: it basically verifies that every (x, y) pair is unique, that is, there are no duplicate coordinates (possibly with conflicting values) in the input data file. Arguably, you could replace this with an exception:

if nx * ny != len(data):
    raise ValueError("non-unique coordinates in the input data")

since an assert statement is generally seen as a way to verify that programmatically, something really shouldn't occur. Here, it's used for verification of the input data, which is outside the of the programmers control.

There is also the finicky issue of using unique with floating point numbers. As long as the data is just read in, equal coordinates should be tested equal, since they'll hold the same representation bitwise for their values, and thus unique values will indeed be uniquely counted. If that is not the case, the assert statement is likely to fail.

NB: this case also handles non-square n by m data. If you know the data is n by n, you can make things simpler, because then you know that n = round(np.sqrt(len(data))). (The rounding may hide non-square lengths, so you'd still need to verify that n * n == len(data).)