forked from idrl/idrlnet
105 lines
3.4 KiB
Python
105 lines
3.4 KiB
Python
import idrlnet.shortcut as sc
|
|
import sympy as sp
|
|
import matplotlib.pyplot as plt
|
|
import matplotlib.tri as tri
|
|
import numpy as np
|
|
|
|
x, y = sp.symbols("x y")
|
|
temp = sp.Symbol("temp")
|
|
temp_range = {temp: (-0.2, 0.2)}
|
|
rec = sc.Rectangle((-1.0, -1.0), (1.0, 1.0))
|
|
|
|
|
|
@sc.datanode
|
|
class Right(sc.SampleDomain):
|
|
# Due to `name` is not specified, Right will be the name of datanode automatically
|
|
def sampling(self, *args, **kwargs):
|
|
points = rec.sample_boundary(
|
|
1000, sieve=(sp.Eq(x, 1.0)), param_ranges=temp_range
|
|
)
|
|
constraints = sc.Variables({"T": 0.0})
|
|
return points, constraints
|
|
|
|
|
|
@sc.datanode
|
|
class Left(sc.SampleDomain):
|
|
# Due to `name` is not specified, Left will be the name of datanode automatically
|
|
def sampling(self, *args, **kwargs):
|
|
points = rec.sample_boundary(
|
|
1000, sieve=(sp.Eq(x, -1.0)), param_ranges=temp_range
|
|
)
|
|
constraints = sc.Variables({"T": temp})
|
|
return points, constraints
|
|
|
|
|
|
@sc.datanode(name="up_down")
|
|
class UpDownBoundaryDomain(sc.SampleDomain):
|
|
def sampling(self, *args, **kwargs):
|
|
points = rec.sample_boundary(
|
|
1000, sieve=((x > -1.0) & (x < 1.0)), param_ranges=temp_range
|
|
)
|
|
constraints = sc.Variables({"normal_gradient_T": 0.0})
|
|
return points, constraints
|
|
|
|
|
|
@sc.datanode(name="heat_domain")
|
|
class HeatDomain(sc.SampleDomain):
|
|
def __init__(self):
|
|
self.points = 1000
|
|
|
|
def sampling(self, *args, **kwargs):
|
|
points = rec.sample_interior(self.points, param_ranges=temp_range)
|
|
constraints = sc.Variables({"diffusion_T": 1.0})
|
|
return points, constraints
|
|
|
|
|
|
net = sc.get_net_node(
|
|
inputs=("x", "y", "temp"), outputs=("T",), name="net1", arch=sc.Arch.mlp
|
|
)
|
|
pde = sc.DiffusionNode(T="T", D=1.0, Q=0.0, dim=2, time=False)
|
|
grad = sc.NormalGradient("T", dim=2, time=False)
|
|
s = sc.Solver(
|
|
sample_domains=(HeatDomain(), Left(), Right(), UpDownBoundaryDomain()),
|
|
netnodes=[net],
|
|
pdes=[pde, grad],
|
|
max_iter=3000,
|
|
)
|
|
s.solve()
|
|
|
|
|
|
def infer_temp(temp_num, file_suffix=None):
|
|
temp_range[temp] = temp_num
|
|
s.set_domain_parameter("heat_domain", {"points": 10000})
|
|
coord = s.infer_step({"heat_domain": ["x", "y", "T"]})
|
|
num_x = coord["heat_domain"]["x"].cpu().detach().numpy().ravel()
|
|
num_y = coord["heat_domain"]["y"].cpu().detach().numpy().ravel()
|
|
num_Tp = coord["heat_domain"]["T"].cpu().detach().numpy().ravel()
|
|
|
|
# Ground truth
|
|
num_T = -(num_x + 1 + temp_num) * (num_x - 1.0) / 2
|
|
|
|
fig, ax = plt.subplots(1, 3, figsize=(10, 3))
|
|
triang_total = tri.Triangulation(num_x, num_y)
|
|
ax[0].tricontourf(triang_total, num_Tp, 100, cmap="hot", vmin=-0.2, vmax=1.21 / 2)
|
|
ax[0].axis("off")
|
|
ax[0].set_title(f"prediction($T_l={temp_num:.2f}$)")
|
|
ax[1].tricontourf(triang_total, num_T, 100, cmap="hot", vmin=-0.2, vmax=1.21 / 2)
|
|
ax[1].axis("off")
|
|
ax[1].set_title(f"ground truth($T_l={temp_num:.2f}$)")
|
|
ax[2].tricontourf(
|
|
triang_total, np.abs(num_T - num_Tp), 100, cmap="hot", vmin=0, vmax=1.21 / 2
|
|
)
|
|
ax[2].axis("off")
|
|
ax[2].set_title("absolute error")
|
|
if file_suffix is None:
|
|
plt.savefig(f"poisson_{temp_num:.2f}.png", dpi=300, bbox_inches="tight")
|
|
plt.show()
|
|
else:
|
|
plt.savefig(f"poisson_{file_suffix}.png", dpi=300, bbox_inches="tight")
|
|
plt.show()
|
|
|
|
|
|
for i in range(41):
|
|
temp_num = i / 100 - 0.2
|
|
infer_temp(temp_num, file_suffix=i)
|