multistage.plot_2d_residual

multistage.plot_2d_residual(net, residual_fun, resolution=256, figname='2d_residual.pdf', cmap='seismic')Source

Plots the signed PDE residual on a uniform 2D grid with a linear scale.

Parameters:

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

• residual_fun (callable) – Function that returns (u, f), where f is the residual.

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

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

• cmap (str) – Matplotlib colormap.

Returns:

fig, ax – Matplotlib fig, ax.