[FEA] Physics informed examples using the PhysicsInformer utility (#664)

* update to use newer apis, and some misc cleanup

* update the stokes example

* add pinns example

* add pi informed gnn example using least squares

---------

Author: ktangsali

examples/cfd/darcy_physics_informed/README.md

@@ -50,13 +50,15 @@ the loss function and the use of one over the other can change from case-to-case
 With this example, we intend to demonstrate both such cases so that the users can compare
 and contrast the two approaches.
 
-In this example we will also use the `PDE` class from Modulus-Sym to symbolically define
-the PDEs. This is very convinient and most natural way to define these PDEs and allows
-us to print the equations to check for correctness. This also abstracts out the
+In this example we will use the `PDE` class from Modulus-Sym to symbolically define
+the PDEs and use the `PhysicsInformer` utility to introduce the PDE
+constraints. Defining the PDEs sympolically is very convinient and most natural way to
+define these PDEs and allows us to print the equations to check for correctness.
+This also abstracts out the
 complexity of converting the equation into a pytorch representation. Modulus Sym also
 provides several complex, well tested PDEs like 3D Navier-Stokes, Linear elasticity,
 Electromagnetics, etc. pre-defined which can be used directly in physics-informing
-applications.  
+applications.
 
 ## Getting Started
 

examples/cfd/darcy_physics_informed/conf/config_deeponet.yaml

@@ -24,7 +24,7 @@ start_lr: 0.001
 gamma: 0.99948708
 max_epochs: 50
 
-phy_wt: 0.1
+physics_weight: 0.1
 
 model:
   fno:

examples/cfd/darcy_physics_informed/conf/config_pino.yaml

@@ -24,7 +24,7 @@ start_lr: 0.001
 gamma: 0.99948708
 max_epochs: 50
 
-phy_wt: 0.1
+physics_weight: 0.1
 
 model:
   fno:

examples/cfd/darcy_physics_informed/darcy_physics_informed_deeponet.py

@@ -14,29 +14,27 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from itertools import chain
+from typing import Dict
+
 import hydra
-from omegaconf import DictConfig
-import torch
-import numpy as np
 import matplotlib.pyplot as plt
-from hydra.utils import to_absolute_path
+import numpy as np
+import torch
 import torch.nn.functional as F
-from torch.utils.data import DataLoader
-from itertools import chain
-
-from modulus.models.mlp import FullyConnected
-from modulus.models.fno import FNO
+from hydra.utils import to_absolute_path
 from modulus.launch.logging import LaunchLogger
 from modulus.launch.utils.checkpoint import save_checkpoint
-
-from utils import HDF5MapStyleDataset
+from modulus.models.fno import FNO
+from modulus.models.mlp import FullyConnected
 from modulus.sym.eq.pdes.diffusion import Diffusion
-
-from modulus.sym.graph import Graph
+from modulus.sym.eq.phy_informer import PhysicsInformer
 from modulus.sym.key import Key
-from modulus.sym.node import Node
-from typing import Optional, Dict
 from modulus.sym.models.arch import Arch
+from omegaconf import DictConfig
+from torch.utils.data import DataLoader
+
+from utils import HDF5MapStyleDataset
 
 
 def validation_step(graph, dataloader, epoch):
@@ -61,14 +59,14 @@ def validation_step(graph, dataloader, epoch):
         # plotting
         fig, ax = plt.subplots(1, 3, figsize=(25, 5))
 
-        d_min = np.min(outvar[0, 0, ...])
-        d_max = np.max(outvar[0, 0, ...])
+        d_min = np.min(outvar[0, 0])
+        d_max = np.max(outvar[0, 0])
 
-        im = ax[0].imshow(outvar[0, 0, ...], vmin=d_min, vmax=d_max)
+        im = ax[0].imshow(outvar[0, 0], vmin=d_min, vmax=d_max)
         plt.colorbar(im, ax=ax[0])
-        im = ax[1].imshow(predvar[0, 0, ...], vmin=d_min, vmax=d_max)
+        im = ax[1].imshow(predvar[0, 0], vmin=d_min, vmax=d_max)
         plt.colorbar(im, ax=ax[1])
-        im = ax[2].imshow(np.abs(predvar[0, 0, ...] - outvar[0, 0, ...]))
+        im = ax[2].imshow(np.abs(predvar[0, 0] - outvar[0, 0]))
         plt.colorbar(im, ax=ax[2])
 
         ax[0].set_title("True")
@@ -113,7 +111,7 @@ def __init__(
         trunk_net=None,
         branch_net=None,
     ):
-        super(MdlsSymWrapper, self).__init__(
+        super().__init__(
             input_keys=input_keys,
             output_keys=output_keys,
         )
@@ -162,9 +160,8 @@ def main(cfg: DictConfig):
     LaunchLogger.initialize()
 
     # Use Diffusion equation for the Darcy PDE
-    darcy = Diffusion(T="u", time=False, dim=2, D="k", Q=1.0 * 4.49996e00 * 3.88433e-03)
-
-    darcy_node = darcy.make_nodes()
+    forcing_fn = 1.0 * 4.49996e00 * 3.88433e-03  # after scaling
+    darcy = Diffusion(T="u", time=False, dim=2, D="k", Q=forcing_fn)
 
     dataset = HDF5MapStyleDataset(
         to_absolute_path("./datasets/Darcy_241/train.hdf5"), device=device
@@ -204,30 +201,14 @@ def main(cfg: DictConfig):
         output_keys=[Key("k"), Key("u")],
         trunk_net=model_trunk,
         branch_net=model_branch,
-    )
-
-    nodes = darcy_node + [model.make_node(name="network", jit=False)]
-
-    # note: this example uses the Graph class from Modulus Sym to construct the
-    # computational graph. This allows you to leverage Modulus Sym's optimized
-    # derivative backend to compute the derivatives, along with other benefits like
-    # symbolic definition of PDEs and leveraging the PDEs from Modulus Sym's PDE
-    # module.
-    # For more details, refer: https://docs.nvidia.com/deeplearning/modulus/modulus-sym/api/modulus.sym.html#module-modulus.sym.graph
-    graph = Graph(
-        nodes,
-        [Key("k_prime"), Key("x"), Key("y")],
-        [Key("k"), Key("u"), Key("diffusion_u")],
-        func_arch=False,
     ).to(device)
 
-    # For pure inference (no gradients)
-    graph_infer = Graph(
-        [model.make_node(name="network", jit=False)],
-        [Key("k_prime"), Key("x"), Key("y")],
-        [Key("k"), Key("u")],  # No PDE Key
-        func_arch=False,
-    ).to(device)
+    phy_informer = PhysicsInformer(
+        required_outputs=["diffusion_u"],
+        equations=darcy,
+        grad_method="autodiff",
+        device=device,
+    )
 
     optimizer = torch.optim.Adam(
         chain(model_branch.parameters(), model_trunk.parameters()),
@@ -251,16 +232,24 @@ def main(cfg: DictConfig):
                 optimizer.zero_grad()
                 outvar = data[1]
 
+                coords = torch.stack([data[2], data[3]], dim=1).requires_grad_(True)
                 # compute forward pass
-                out = graph.forward(
+                out = model.forward(
                     {
                         "k_prime": data[0][:, 0].unsqueeze(dim=1),
-                        "x": data[2].requires_grad_(True),
-                        "y": data[3].requires_grad_(True),
+                        "x": coords[:, 0:1],
+                        "y": coords[:, 1:2],
                     }
                 )
 
-                pde_out_arr = out["diffusion_u"]
+                residuals = phy_informer.forward(
+                    {
+                        "coordinates": coords,
+                        "u": out["u"],
+                        "k": out["k"],
+                    }
+                )
+                pde_out_arr = residuals["diffusion_u"]
 
                 # Boundary condition
                 pde_out_arr = F.pad(
@@ -276,7 +265,7 @@ def main(cfg: DictConfig):
                 )
 
                 # Compute total loss
-                loss = loss_data + cfg.phy_wt * loss_pde
+                loss = loss_data + cfg.physics_weight * loss_pde
 
                 # Backward pass and optimizer and learning rate update
                 loss.backward()
@@ -289,7 +278,7 @@ def main(cfg: DictConfig):
             log.log_epoch({"Learning Rate": optimizer.param_groups[0]["lr"]})
 
         with LaunchLogger("valid", epoch=epoch) as log:
-            error = validation_step(graph_infer, validation_dataloader, epoch)
+            error = validation_step(model, validation_dataloader, epoch)
             log.log_epoch({"Validation error": error})
 
         save_checkpoint(

examples/cfd/darcy_physics_informed/darcy_physics_informed_fno.py

@@ -15,23 +15,20 @@
 # limitations under the License.
 
 import hydra
-from omegaconf import DictConfig
-import torch
-import numpy as np
-
 import matplotlib.pyplot as plt
-from hydra.utils import to_absolute_path
-import torch.nn.functional as F
-from torch.utils.data import DataLoader
+import numpy as np
+import torch
 import torch.nn.functional as F
-
-from modulus.models.fno import FNO
+from hydra.utils import to_absolute_path
 from modulus.launch.logging import LaunchLogger
 from modulus.launch.utils.checkpoint import save_checkpoint
+from modulus.models.fno import FNO
 from modulus.sym.eq.pdes.diffusion import Diffusion
+from modulus.sym.eq.phy_informer import PhysicsInformer
+from omegaconf import DictConfig
+from torch.utils.data import DataLoader
 
 from utils import HDF5MapStyleDataset
-from ops import dx, ddx
 
 
 def validation_step(model, dataloader, epoch):
@@ -53,14 +50,14 @@ def validation_step(model, dataloader, epoch):
         # plotting
         fig, ax = plt.subplots(1, 3, figsize=(25, 5))
 
-        d_min = np.min(outvar[0, 0, ...])
-        d_max = np.max(outvar[0, 0, ...])
+        d_min = np.min(outvar[0, 0])
+        d_max = np.max(outvar[0, 0])
 
-        im = ax[0].imshow(outvar[0, 0, ...], vmin=d_min, vmax=d_max)
+        im = ax[0].imshow(outvar[0, 0], vmin=d_min, vmax=d_max)
         plt.colorbar(im, ax=ax[0])
-        im = ax[1].imshow(predvar[0, 0, ...], vmin=d_min, vmax=d_max)
+        im = ax[1].imshow(predvar[0, 0], vmin=d_min, vmax=d_max)
         plt.colorbar(im, ax=ax[1])
-        im = ax[2].imshow(np.abs(predvar[0, 0, ...] - outvar[0, 0, ...]))
+        im = ax[2].imshow(np.abs(predvar[0, 0] - outvar[0, 0]))
         plt.colorbar(im, ax=ax[2])
 
         ax[0].set_title("True")
@@ -84,8 +81,8 @@ def main(cfg: DictConfig):
     LaunchLogger.initialize()
 
     # Use Diffusion equation for the Darcy PDE
-    darcy = Diffusion(T="u", time=False, dim=2, D="k", Q=1.0 * 4.49996e00 * 3.88433e-03)
-    darcy_node = darcy.make_nodes()
+    forcing_fn = 1.0 * 4.49996e00 * 3.88433e-03  # after scaling
+    darcy = Diffusion(T="u", time=False, dim=2, D="k", Q=forcing_fn)
 
     dataset = HDF5MapStyleDataset(
         to_absolute_path("./datasets/Darcy_241/train.hdf5"), device=device
@@ -110,6 +107,14 @@ def main(cfg: DictConfig):
         padding=cfg.model.fno.padding,
     ).to(device)
 
+    phy_informer = PhysicsInformer(
+        required_outputs=["diffusion_u"],
+        equations=darcy,
+        grad_method="finite_difference",
+        device=device,
+        fd_dx=1 / 240,  # Unit square with resoultion as 240
+    )
+
     optimizer = torch.optim.Adam(
         model.parameters(),
         betas=(0.9, 0.999),
@@ -135,37 +140,15 @@ def main(cfg: DictConfig):
                 # Compute forward pass
                 out = model(invar[:, 0].unsqueeze(dim=1))
 
-                dxf = 1.0 / out.shape[-2]
-                dyf = 1.0 / out.shape[-1]
-
-                # Compute gradients using finite difference
-                sol_x = dx(out, dx=dxf, channel=0, dim=1, order=1, padding="zeros")
-                sol_y = dx(out, dx=dyf, channel=0, dim=0, order=1, padding="zeros")
-                sol_x_x = ddx(out, dx=dxf, channel=0, dim=1, order=1, padding="zeros")
-                sol_y_y = ddx(out, dx=dyf, channel=0, dim=0, order=1, padding="zeros")
-
-                k_x = dx(invar, dx=dxf, channel=0, dim=1, order=1, padding="zeros")
-                k_y = dx(invar, dx=dxf, channel=0, dim=0, order=1, padding="zeros")
-
-                k, _, _ = (
-                    invar[:, 0],
-                    invar[:, 1],
-                    invar[:, 2],
-                )
-
-                pde_out = darcy_node[0].evaluate(
+                # print(out.shape, invar[:,0:1].shape)
+                residuals = phy_informer.forward(
                     {
-                        "u__x": sol_x,
-                        "u__y": sol_y,
-                        "u__x__x": sol_x_x,
-                        "u__y__y": sol_y_y,
-                        "k": k,
-                        "k__x": k_x,
-                        "k__y": k_y,
+                        "u": out,
+                        "k": invar[:, 0:1],
                     }
                 )
+                pde_out_arr = residuals["diffusion_u"]
 
-                pde_out_arr = pde_out["diffusion_u"]
                 pde_out_arr = F.pad(
                     pde_out_arr[:, :, 2:-2, 2:-2], [2, 2, 2, 2], "constant", 0
                 )
@@ -175,7 +158,7 @@ def main(cfg: DictConfig):
                 loss_data = F.mse_loss(outvar, out)
 
                 # Compute total loss
-                loss = loss_data + 1 / 240 * cfg.phy_wt * loss_pde
+                loss = loss_data + 1 / 240 * cfg.physics_weight * loss_pde
 
                 # Backward pass and optimizer and learning rate update
                 loss.backward()

examples/cfd/darcy_physics_informed/download_data.py

@@ -16,6 +16,7 @@
 
 import h5py
 import numpy as np
+
 from utils import download_FNO_dataset
 
 download_FNO_dataset("Darcy_241", outdir="datasets/")