plot_loss_learner_2d.py

import matplotlib.tri as mtri
from matplotlib import pyplot as plt


def plot_loss(learner, cmap="viridis", xy_size=(10, 10)):
    fig, ax = plt.subplots(figsize=xy_size)
    x_size, y_size = xy_size
    ip = learner.ip()
    zs = learner.loss_per_triangle(ip)
    tri = ip.tri
    xs, ys = tri.points.T
    x_size, y_size = xy_size
    triang = mtri.Triangulation(x_size * xs, y_size * ys, triangles=tri.vertices)
    ax.triplot(triang, c="k", lw=0.3, alpha=1, zorder=2), (ip.values, triang)
    im = ax.tripcolor(triang, zs.flatten(), zorder=0, cmap=cmap)
    plt.colorbar(im)

Example result: