multistage.plot_2d_solution

multistage.plot_2d_solution(net, u_true_fn=None, resolution=256, figname='2d_solution.pdf')Source

Plots solution, truth, and error on a uniform 2D grid.

Parameters:

• net (eqx.Module) – The neural network model.

• u_true_fn (callable) – Function that returns the true value of the solution.

• resolution (int) – Grid resolution for the plot.

• figname (str) – Filename to save the plot.

Returns:

fig, ax – Matplotlib fig and ax.