94 lines
3.5 KiB
Python
94 lines
3.5 KiB
Python
""" Module for plotting. """
|
|
import matplotlib
|
|
matplotlib.use('Qt5Agg')
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
import torch
|
|
|
|
from pina import LabelTensor
|
|
from pina import PINN
|
|
from .problem import Problem2D, Problem1D, TimeDependentProblem
|
|
#from pina.tdproblem1d import TimeDepProblem1D
|
|
|
|
|
|
class Plotter:
|
|
|
|
def _plot_2D(self, obj, method='contourf'):
|
|
"""
|
|
"""
|
|
if not isinstance(obj, PINN):
|
|
raise RuntimeError
|
|
|
|
res = 256
|
|
pts = obj.problem.spatial_domain.discretize(res, 'grid')
|
|
grids_container = [
|
|
pts[:, 0].reshape(res, res),
|
|
pts[:, 1].reshape(res, res),
|
|
]
|
|
pts = LabelTensor(torch.tensor(pts), obj.problem.input_variables)
|
|
predicted_output = obj.model(pts.tensor)
|
|
|
|
if hasattr(obj.problem, 'truth_solution'):
|
|
truth_output = obj.problem.truth_solution(*pts.tensor.T).float()
|
|
fig, axes = plt.subplots(nrows=1, ncols=3, figsize=(16, 6))
|
|
|
|
cb = getattr(axes[0], method)(*grids_container, predicted_output.tensor.reshape(res, res).detach())
|
|
fig.colorbar(cb, ax=axes[0])
|
|
cb = getattr(axes[1], method)(*grids_container, truth_output.reshape(res, res).detach())
|
|
fig.colorbar(cb, ax=axes[1])
|
|
cb = getattr(axes[2], method)(*grids_container, (truth_output-predicted_output.tensor.float().flatten()).detach().reshape(res, res))
|
|
fig.colorbar(cb, ax=axes[2])
|
|
else:
|
|
fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(8, 6))
|
|
cb = getattr(axes, method)(*grids_container, predicted_output.tensor.reshape(res, res).detach())
|
|
fig.colorbar(cb, ax=axes)
|
|
|
|
|
|
def _plot_1D_TimeDep(self, obj, method='contourf'):
|
|
"""
|
|
"""
|
|
if not isinstance(obj, PINN):
|
|
raise RuntimeError
|
|
|
|
res = 256
|
|
grids_container = np.meshgrid(
|
|
obj.problem.spatial_domain.discretize(res, 'grid'),
|
|
obj.problem.temporal_domain.discretize(res, 'grid'),
|
|
)
|
|
pts = np.hstack([
|
|
grids_container[0].reshape(-1, 1),
|
|
grids_container[1].reshape(-1, 1),
|
|
])
|
|
pts = LabelTensor(torch.tensor(pts), obj.problem.input_variables)
|
|
predicted_output = obj.model(pts.tensor)
|
|
|
|
if hasattr(obj.problem, 'truth_solution'):
|
|
truth_output = obj.problem.truth_solution(*pts.tensor.T).float()
|
|
fig, axes = plt.subplots(nrows=1, ncols=3, figsize=(16, 6))
|
|
|
|
cb = getattr(axes[0], method)(*grids_container, predicted_output.tensor.reshape(res, res).detach())
|
|
fig.colorbar(cb, ax=axes[0])
|
|
cb = getattr(axes[1], method)(*grids_container, truth_output.reshape(res, res).detach())
|
|
fig.colorbar(cb, ax=axes[1])
|
|
cb = getattr(axes[2], method)(*grids_container, (truth_output-predicted_output.tensor.float().flatten()).detach().reshape(res, res))
|
|
fig.colorbar(cb, ax=axes[2])
|
|
else:
|
|
fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(8, 6))
|
|
cb = getattr(axes, method)(*grids_container, predicted_output.tensor.reshape(res, res).detach())
|
|
fig.colorbar(cb, ax=axes)
|
|
|
|
|
|
|
|
def plot(self, obj, filename=None):
|
|
"""
|
|
"""
|
|
if isinstance(obj.problem, (TimeDependentProblem, Problem1D)): # time-dep 1D
|
|
self._plot_1D_TimeDep(obj, method='pcolor')
|
|
elif isinstance(obj.problem, Problem2D): # 2D
|
|
self._plot_2D(obj, method='pcolor')
|
|
|
|
if filename:
|
|
plt.savefig(filename)
|
|
else:
|
|
plt.show()
|