Skip to content

Commit 2974405

Browse files
eparish1eparish1
committed
Support empty linear-affine depends_on and add diffusion example
1 parent a1a6a41 commit 2974405

File tree

4 files changed

+371
-27
lines changed

4 files changed

+371
-27
lines changed

‎docs/source/examples.rst‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,3 +65,35 @@ Tips
6565
- The script uses a long training horizon. For a quick smoke test, lower
6666
``num-epochs`` and increase ``batch-size`` near the top of the script.
6767
- Output PDFs are written in the example directory.
68+
69+
Transient heat diffusion example
70+
--------------------------------
71+
72+
This example solves transient heat diffusion on the unit square with zero
73+
Dirichlet boundaries and constant forcing, builds a POD-reduced state, and
74+
trains a **linear** structure-preserving NN-OpInf model with:
75+
76+
- ``LinearAffineSpdTensorOperator(acts_on=x, depends_on=(), positive=False)``
77+
for dissipative diffusion dynamics
78+
- ``VectorOffsetOperator`` for the constant forcing term
79+
80+
Training uses ADAM with LBFGS acceleration:
81+
(``training_settings["optimizer"] = "ADAM"``,
82+
``training_settings["LBFGS-acceleration"] = True``).
83+
84+
Run it:
85+
86+
.. code-block:: bash
87+
88+
python examples/diffusion/heat_diffusion_end_to_end.py --kappa 0.75 --forcing 1.0
89+
90+
Key outputs:
91+
92+
- Plot: ``examples/diffusion/heat_diffusion_solution.pdf``
93+
- Trained models: ``examples/diffusion/ml-models/``
94+
95+
Code
96+
----
97+
98+
.. literalinclude:: ../../examples/diffusion/heat_diffusion_end_to_end.py
99+
:language: python

‎examples/diffusion/heat_diffusion_end_to_end.py‎

Lines changed: 205 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,205 @@
1+
import argparse
2+
import os
3+
4+
import numpy as np
5+
import torch
6+
from matplotlib import pyplot as plt
7+
8+
import nnopinf
9+
import nnopinf.models as models
10+
import nnopinf.training
11+
import nnopinf.training.trainers
12+
13+
14+
def laplacian_dirichlet(u, n, h):
15+
"""Five-point Laplacian on an n x n interior grid with zero Dirichlet boundaries."""
16+
u2d = u.reshape(n, n)
17+
upad = np.pad(u2d, ((1, 1), (1, 1)), mode="constant")
18+
lap = (
19+
upad[2:, 1:-1]
20+
+ upad[:-2, 1:-1]
21+
+ upad[1:-1, 2:]
22+
+ upad[1:-1, :-2]
23+
- 4.0 * upad[1:-1, 1:-1]
24+
) / (h**2)
25+
return lap.reshape(-1)
26+
27+
28+
class HeatDiffusionFOM:
29+
def __init__(self, n, forcing_value):
30+
self.n = n
31+
self.h = 1.0 / (n + 1)
32+
self.x = np.linspace(self.h, 1.0 - self.h, n)
33+
xx, yy = np.meshgrid(self.x, self.x, indexing="ij")
34+
self.u0 = (
35+
np.sin(np.pi * xx) * np.sin(np.pi * yy)
36+
+ 0.25 * np.sin(2.0 * np.pi * xx) * np.sin(np.pi * yy)
37+
).reshape(-1)
38+
self.forcing = forcing_value * np.ones(self.n * self.n)
39+
40+
def velocity(self, u, kappa):
41+
return kappa * laplacian_dirichlet(u, self.n, self.h) + self.forcing
42+
43+
def solve(self, dt, end_time, kappa):
44+
u = self.u0.copy()
45+
t = 0.0
46+
rk4const = np.array([1.0 / 4.0, 1.0 / 3.0, 1.0 / 2.0, 1.0])
47+
snapshots = []
48+
time = []
49+
while t <= end_time + 0.5 * dt:
50+
snapshots.append(u.copy())
51+
time.append(t)
52+
u0 = u.copy()
53+
for i in range(4):
54+
f = self.velocity(u, kappa)
55+
u = u0 + dt * rk4const[i] * f
56+
t += dt
57+
return np.array(snapshots).T, np.array(time)
58+
59+
60+
class NnOpInfRom:
61+
def __init__(self, model):
62+
self.model_ = model
63+
self.inputs_ = {}
64+
65+
def velocity(self, u):
66+
self.inputs_["x"] = torch.tensor(u[None], dtype=torch.float64)
67+
return self.model_.forward(self.inputs_)[0].detach().numpy()
68+
69+
def solve(self, u0, dt, end_time):
70+
u = u0.copy()
71+
t = 0.0
72+
rk4const = np.array([1.0 / 4.0, 1.0 / 3.0, 1.0 / 2.0, 1.0])
73+
snapshots = []
74+
time = []
75+
while t <= end_time + 0.5 * dt:
76+
snapshots.append(u.copy())
77+
time.append(t)
78+
u0_l = u.copy()
79+
for i in range(4):
80+
f = self.velocity(u)
81+
u = u0_l + dt * rk4const[i] * f
82+
t += dt
83+
return np.array(snapshots).T, np.array(time)
84+
85+
86+
def flatten_for_training(q):
87+
return np.reshape(q, (q.shape[0], q.shape[1] * q.shape[2])).T
88+
89+
90+
def main():
91+
parser = argparse.ArgumentParser(
92+
description="Transient heat diffusion on the unit square with a linear SPD operator."
93+
)
94+
parser.add_argument("--grid-size", type=int, default=32, help="Interior points per direction.")
95+
parser.add_argument("--dt", type=float, default=2.0e-4, help="Time step for RK4.")
96+
parser.add_argument("--end-time", type=float, default=1.0, help="Final simulation time.")
97+
parser.add_argument("--rom-dim", type=int, default=20, help="Requested POD rank.")
98+
parser.add_argument("--kappa", type=float, default=0.75, help="Diffusion coefficient.")
99+
parser.add_argument("--forcing", type=float, default=1.0, help="Constant forcing amplitude.")
100+
parser.add_argument("--num-epochs", type=int, default=500, help="Training epochs.")
101+
parser.add_argument("--tr-delta0", type=float, default=1.0, help="Initial trust-region radius.")
102+
parser.add_argument("--tr-cg-max-iters", type=int, default=200, help="Maximum CG iterations per TR step.")
103+
parser.add_argument("--tr-batch-size", type=int, default=50, help="Batch size used by TR optimizer.")
104+
args = parser.parse_args()
105+
106+
torch.manual_seed(1)
107+
np.random.seed(1)
108+
109+
output_dir = os.path.dirname(os.path.abspath(__file__))
110+
model_dir = os.path.join(output_dir, "ml-models")
111+
os.makedirs(model_dir, exist_ok=True)
112+
113+
fom = HeatDiffusionFOM(args.grid_size, args.forcing)
114+
u, _ = fom.solve(args.dt, args.end_time, args.kappa)
115+
snapshots_all = u[..., None]
116+
117+
# POD basis from all training trajectories.
118+
snapshots_matrix = np.reshape(
119+
snapshots_all, (snapshots_all.shape[0], snapshots_all.shape[1] * snapshots_all.shape[2])
120+
)
121+
phi, singular_values, _ = np.linalg.svd(snapshots_matrix, full_matrices=False)
122+
relative_energy = np.cumsum(singular_values**2) / np.sum(singular_values**2)
123+
pod_rank = np.searchsorted(relative_energy, 0.999999999) + 1
124+
rom_dim = min(args.rom_dim, int(pod_rank))
125+
phi = phi[:, :rom_dim]
126+
print(f"Using ROM dimension K={rom_dim}")
127+
128+
# Reduced snapshots and time derivatives.
129+
uhat = np.einsum("ij,ikn->jkn", phi, snapshots_all)
130+
uhat_dot = (uhat[:, 2:, :] - uhat[:, :-2, :]) / (2.0 * args.dt)
131+
uhat = uhat[:, 1:-1, :]
132+
133+
x_input = nnopinf.variables.Variable(size=rom_dim, name="x", normalization_strategy="MaxAbs")
134+
target = nnopinf.variables.Variable(size=rom_dim, name="y", normalization_strategy="MaxAbs")
135+
136+
diffusion_operator = nnopinf.operators.LinearAffineSpdTensorOperator(
137+
acts_on=x_input, depends_on=(), positive=False
138+
)
139+
forcing_operator = nnopinf.operators.VectorOffsetOperator(n_outputs=rom_dim)
140+
model = models.Model([diffusion_operator, forcing_operator])
141+
142+
training_settings = nnopinf.training.get_default_settings()
143+
training_settings["output-path"] = model_dir
144+
training_settings["optimizer"] = "ADAM"
145+
training_settings["batch-size"] = 5000
146+
training_settings["num-epochs"] = args.num_epochs
147+
training_settings["weight-decay"] = 1.0e-6
148+
training_settings["LBFGS-acceleration"] = True
149+
150+
x_input.set_data(flatten_for_training(uhat))
151+
target.set_data(flatten_for_training(uhat_dot))
152+
153+
print("Training linear SPD diffusion model")
154+
nnopinf.training.trainers.train(
155+
model, variables=[x_input], y=target, training_settings=training_settings
156+
)
157+
158+
# Evaluate for the same diffusion coefficient used to generate training data.
159+
test_kappa = args.kappa
160+
u_fom, t = fom.solve(args.dt, args.end_time, test_kappa)
161+
u0_r = phi.T @ fom.u0
162+
163+
rom = NnOpInfRom(model)
164+
u_rom_r, _ = rom.solve(u0_r, args.dt, args.end_time)
165+
u_rom = phi @ u_rom_r
166+
167+
relative_error = np.linalg.norm(u_rom - u_fom) / np.linalg.norm(u_fom)
168+
print(f"Relative trajectory error (kappa={test_kappa:.3f}): {relative_error:.4e}")
169+
170+
u_fom_final = u_fom[:, -1].reshape(args.grid_size, args.grid_size)
171+
u_rom_final = u_rom[:, -1].reshape(args.grid_size, args.grid_size)
172+
diff_final = np.abs(u_fom_final - u_rom_final)
173+
174+
fig, axes = plt.subplots(1, 3, figsize=(12, 3.8), constrained_layout=True)
175+
extent = (0.0, 1.0, 0.0, 1.0)
176+
vmin = min(np.min(u_fom_final), np.min(u_rom_final))
177+
vmax = max(np.max(u_fom_final), np.max(u_rom_final))
178+
179+
im0 = axes[0].imshow(u_fom_final, origin="lower", extent=extent, cmap="viridis", vmin=vmin, vmax=vmax)
180+
axes[0].set_title("FOM final state")
181+
axes[0].set_xlabel("x")
182+
axes[0].set_ylabel("y")
183+
184+
im1 = axes[1].imshow(u_rom_final, origin="lower", extent=extent, cmap="viridis", vmin=vmin, vmax=vmax)
185+
axes[1].set_title("NN-OpInf final state")
186+
axes[1].set_xlabel("x")
187+
axes[1].set_ylabel("y")
188+
189+
im2 = axes[2].imshow(diff_final, origin="lower", extent=extent, cmap="magma")
190+
axes[2].set_title("Absolute error")
191+
axes[2].set_xlabel("x")
192+
axes[2].set_ylabel("y")
193+
194+
fig.colorbar(im0, ax=axes[0], fraction=0.046, pad=0.04)
195+
fig.colorbar(im1, ax=axes[1], fraction=0.046, pad=0.04)
196+
fig.colorbar(im2, ax=axes[2], fraction=0.046, pad=0.04)
197+
fig.suptitle(f"Transient heat diffusion, kappa={test_kappa:.2f}, rel. error={relative_error:.2e}")
198+
199+
fig.savefig(os.path.join(output_dir, "heat_diffusion_solution.pdf"))
200+
fig.savefig(os.path.join(output_dir, "heat_diffusion_solution.png"), dpi=200)
201+
plt.close(fig)
202+
203+
204+
if __name__ == "__main__":
205+
main()

‎nnopinf/operators.py‎

Lines changed: 51 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1538,10 +1538,12 @@ def __init__(self,n_outputs,acts_on,depends_on,name="LinearAffineTensorOperator"
15381538
self.acts_on_name_ = acts_on.get_name()
15391539
self.acts_on_size_ = acts_on.get_size()
15401540
self.depends_on_names_ = []
1541-
self.n_inputs_ = 0
1541+
n_inputs_total = 0
15421542
for i in range(len(depends_on)):
15431543
self.depends_on_names_.append(depends_on[i].get_name())
1544-
self.n_inputs_ += depends_on[i].get_size()
1544+
n_inputs_total += depends_on[i].get_size()
1545+
self.has_affine_inputs_ = (n_inputs_total > 0)
1546+
self.n_inputs_ = (n_inputs_total if self.has_affine_inputs_ else 1)
15451547
# Learnable tensor -- unstructured
15461548
self.T = nn.Parameter(torch.randn([n_outputs,self.acts_on_size_,self.n_inputs_]))
15471549
self.name_ = name
@@ -1563,7 +1565,10 @@ def forward(self,inputs,return_jacobian=False):
15631565

15641566
# Separate out states and parameters
15651567
x = inputs[self.acts_on_name_]
1566-
mu = inputs_to_tensor(inputs,self.depends_on_names_)
1568+
if self.has_affine_inputs_:
1569+
mu = inputs_to_tensor(inputs,self.depends_on_names_)
1570+
else:
1571+
mu = torch.ones((x.shape[0], 1), dtype=x.dtype, device=x.device)
15671572
# Compute product (T*mu)x
15681573
result = torch.einsum('ijp,bj,bp->bi',self.T,x,mu)
15691574

@@ -1577,12 +1582,15 @@ def forward(self,inputs,return_jacobian=False):
15771582
def set_scalings(self,input_scalings_dict,output_scaling):
15781583
with torch.no_grad():
15791584
self.scalings_set_ = True
1580-
input_scalings = None
1581-
for input_arg in self.depends_on_names_:
1582-
if input_scalings is None:
1583-
input_scalings = input_scalings_dict[input_arg]
1584-
else:
1585-
input_scalings = torch.cat( (input_scalings,input_scalings_dict[input_arg]),0)
1585+
if self.has_affine_inputs_:
1586+
input_scalings = None
1587+
for input_arg in self.depends_on_names_:
1588+
if input_scalings is None:
1589+
input_scalings = input_scalings_dict[input_arg]
1590+
else:
1591+
input_scalings = torch.cat( (input_scalings,input_scalings_dict[input_arg]),0)
1592+
else:
1593+
input_scalings = torch.ones_like(self.T[0, 0, :])
15861594

15871595
#for i in range(0,self.T.shape[-1]):
15881596
self.input_scalings_ = input_scalings
@@ -1627,10 +1635,12 @@ def __init__(self,acts_on,depends_on,skew=True,name="LinearAffineSkewTensorOpera
16271635
self.n_outputs_ = self.acts_on_size_
16281636

16291637
self.depends_on_names_ = []
1630-
self.n_inputs_ = 0
1638+
n_inputs_total = 0
16311639
for i in range(len(depends_on)):
16321640
self.depends_on_names_.append(depends_on[i].get_name())
1633-
self.n_inputs_ += depends_on[i].get_size()
1641+
n_inputs_total += depends_on[i].get_size()
1642+
self.has_affine_inputs_ = (n_inputs_total > 0)
1643+
self.n_inputs_ = (n_inputs_total if self.has_affine_inputs_ else 1)
16341644

16351645
# set dims and indices based on skew or sym
16361646
if not skew:
@@ -1667,7 +1677,10 @@ def forward(self,inputs, return_jacobian=False):
16671677
assert return_jacobian == False, "Return jacobian currently not implemented for linear skew operator"
16681678
# Separate out states and parameters
16691679
x = inputs[self.acts_on_name_]
1670-
mu = inputs_to_tensor(inputs,self.depends_on_names_)
1680+
if self.has_affine_inputs_:
1681+
mu = inputs_to_tensor(inputs,self.depends_on_names_)
1682+
else:
1683+
mu = torch.ones((x.shape[0], 1), dtype=x.dtype, device=x.device)
16711684
mu = mu / self.input_scalings_
16721685
# Fill tensor with learnable parameters
16731686
S = torch.zeros(self.n_outputs_,self.n_outputs_,self.n_inputs_)
@@ -1685,12 +1698,15 @@ def forward(self,inputs, return_jacobian=False):
16851698
def set_scalings(self,input_scalings_dict,output_scaling):
16861699
with torch.no_grad():
16871700
self.scalings_set_ = True
1688-
input_scalings = None
1689-
for input_arg in self.depends_on_names_:
1690-
if input_scalings is None:
1691-
input_scalings = input_scalings_dict[input_arg]
1692-
else:
1693-
input_scalings = torch.cat( (input_scalings,input_scalings_dict[input_arg]),0)
1701+
if self.has_affine_inputs_:
1702+
input_scalings = None
1703+
for input_arg in self.depends_on_names_:
1704+
if input_scalings is None:
1705+
input_scalings = input_scalings_dict[input_arg]
1706+
else:
1707+
input_scalings = torch.cat( (input_scalings,input_scalings_dict[input_arg]),0)
1708+
else:
1709+
input_scalings = torch.ones_like(self.input_scalings_)
16941710
# Update initial layer weights
16951711
self.input_scalings_[:] = input_scalings
16961712
self.output_scalings_[:] = output_scaling
@@ -1727,10 +1743,12 @@ def __init__(self,acts_on,depends_on,positive=True,name="LinearAffineSpdTensorOp
17271743
self.n_outputs_ = self.acts_on_size_
17281744

17291745
self.depends_on_names_ = []
1730-
self.n_inputs_ = 0
1746+
n_inputs_total = 0
17311747
for i in range(len(depends_on)):
17321748
self.depends_on_names_.append(depends_on[i].get_name())
1733-
self.n_inputs_ += depends_on[i].get_size()
1749+
n_inputs_total += depends_on[i].get_size()
1750+
self.has_affine_inputs_ = (n_inputs_total > 0)
1751+
self.n_inputs_ = (n_inputs_total if self.has_affine_inputs_ else 1)
17341752

17351753
# set dims and indices
17361754
ldim = int(self.n_outputs_*(self.n_outputs_-1)/2)
@@ -1766,7 +1784,10 @@ def forward(self,inputs, return_jacobian=False):
17661784

17671785
# Separate out states and parameters
17681786
x = inputs[self.acts_on_name_]
1769-
mu = inputs_to_tensor(inputs,self.depends_on_names_)
1787+
if self.has_affine_inputs_:
1788+
mu = inputs_to_tensor(inputs,self.depends_on_names_)
1789+
else:
1790+
mu = torch.ones((x.shape[0], 1), dtype=x.dtype, device=x.device)
17701791
mu = mu / self.input_scalings_
17711792

17721793
# Fill tensor with learnable parameters
@@ -1787,12 +1808,15 @@ def forward(self,inputs, return_jacobian=False):
17871808
def set_scalings(self,input_scalings_dict,output_scaling):
17881809
with torch.no_grad():
17891810
self.scalings_set_ = True
1790-
input_scalings = None
1791-
for input_arg in self.depends_on_names_:
1792-
if input_scalings is None:
1793-
input_scalings = input_scalings_dict[input_arg]
1794-
else:
1795-
input_scalings = torch.cat( (input_scalings,input_scalings_dict[input_arg]),0)
1811+
if self.has_affine_inputs_:
1812+
input_scalings = None
1813+
for input_arg in self.depends_on_names_:
1814+
if input_scalings is None:
1815+
input_scalings = input_scalings_dict[input_arg]
1816+
else:
1817+
input_scalings = torch.cat( (input_scalings,input_scalings_dict[input_arg]),0)
1818+
else:
1819+
input_scalings = torch.ones_like(self.input_scalings_)
17961820
# Update initial layer weights
17971821
self.input_scalings_[:] = input_scalings
17981822
self.output_scalings_[:] = output_scaling

0 commit comments

Comments
 (0)