heatmap script

This commit is contained in:
2026-06-18 20:29:46 +02:00
parent 289d2350c6
commit d78926f440
2 changed files with 81 additions and 0 deletions
+2
View File
@@ -1 +1,3 @@
/target
/_reference
/_venv
+79
View File
@@ -0,0 +1,79 @@
#!/usr/bin/env python3
"""Render a sensor dump (one bracketed row of numbers per line) as a heatmap.
Usage:
python heatmap.py [input.txt] [-o out.png] [--cmap inferno]
Each line is expected to look like:
[55, 192, 131, ..., 138]
Brackets/commas are optional, so plain whitespace-separated rows also work.
"""
import argparse
import re
import sys
def parse_grid(text):
"""Parse lines of numbers into a 2D list of floats."""
grid = []
for line in text.splitlines():
# strip brackets and split on commas/whitespace
nums = re.findall(r"-?\d+(?:\.\d+)?", line)
if nums:
grid.append([float(n) for n in nums])
if not grid:
raise ValueError("No numeric data found in input")
width = len(grid[0])
if any(len(row) != width for row in grid):
raise ValueError("Rows have inconsistent lengths")
return grid
def main():
parser = argparse.ArgumentParser(description="Render a number grid as a heatmap.")
parser.add_argument("input", nargs="?", default="output.txt",
help="input file (default: output.txt)")
parser.add_argument("-o", "--output",
help="save figure to this path instead of showing it")
parser.add_argument("--cmap", default="inferno", help="matplotlib colormap")
parser.add_argument("--no-annot", action="store_true",
help="do not print the value inside each cell")
args = parser.parse_args()
import matplotlib.pyplot as plt
try:
with open(args.input) as f:
grid = parse_grid(f.read())
except (OSError, ValueError) as e:
sys.exit(f"error: {e}")
rows, cols = len(grid), len(grid[0])
fig, ax = plt.subplots(figsize=(max(6, cols * 0.5), max(5, rows * 0.5)))
im = ax.imshow(grid, cmap=args.cmap, interpolation="nearest")
fig.colorbar(im, ax=ax, label="value")
ax.set_title(f"{args.input} ({rows}x{cols})")
ax.set_xticks(range(cols))
ax.set_yticks(range(rows))
if not args.no_annot:
vmin, vmax = min(map(min, grid)), max(map(max, grid))
mid = (vmin + vmax) / 2
for r in range(rows):
for c in range(cols):
v = grid[r][c]
ax.text(c, r, f"{v:g}", ha="center", va="center", fontsize=6,
color="white" if v < mid else "black")
fig.tight_layout()
if args.output:
fig.savefig(args.output, dpi=150)
print(f"saved {args.output}")
else:
plt.show()
if __name__ == "__main__":
main()