"""MeshGraphNet inference method for PhysicsNeMo mesh-stage models."""
from __future__ import annotations
import logging
from pathlib import Path
from typing import TYPE_CHECKING, cast
import numpy as np
from .infer_physicsnemo_base import InferPhysicsNeMoBase
from .process_physicsnemo import ProcessPhysicsNemo
if TYPE_CHECKING: # typed for mypy; imported lazily at runtime
import torch
[docs]
class InferPhysicsNeMoMGN(InferPhysicsNeMoBase):
"""Predict mesh stages with a trained PhysicsNeMo MeshGraphNet.
The shared graph topology and edge features are loaded from the tensors the
training workflow saved next to the checkpoint.
"""
[docs]
def __init__(
self, model_tag: str = "mgn", log_level: int | str = logging.INFO
) -> None:
"""Initialize the MeshGraphNet inference method.
Args:
model_tag: Prefix of the checkpoint files to load
(``<model_tag>_stage_model.pt``), e.g.
``"physics_informed_motion"`` for weights written by
:class:`monai_physio.TrainPhysicsNeMoPhysicsInformedMotion`.
Default: ``"mgn"``.
log_level: Logging level. Default: ``logging.INFO``.
"""
super().__init__(model_tag=model_tag, log_level=log_level)
[docs]
def build_model(self, meta: dict) -> torch.nn.Module:
MeshGraphNet = ProcessPhysicsNemo.import_meshgraphnet()
num_layers = int(meta.get("num_layers", 2))
hidden_dim = int(meta["hidden_dim"])
model = MeshGraphNet(
input_dim_nodes=int(meta["in_features"]),
input_dim_edges=int(meta.get("input_dim_edges", 4)),
output_dim=int(meta["n_target"]),
processor_size=int(meta["processor_size"]),
hidden_dim_processor=hidden_dim,
hidden_dim_node_encoder=hidden_dim,
num_layers_node_encoder=num_layers,
hidden_dim_node_decoder=hidden_dim,
num_layers_node_decoder=num_layers,
hidden_dim_edge_encoder=hidden_dim,
num_layers_edge_encoder=num_layers,
num_layers_edge_processor=num_layers,
num_layers_node_processor=num_layers,
aggregation="mean",
num_processor_checkpoint_segments=int(
meta.get("num_processor_checkpoint_segments", 0)
),
)
return cast("torch.nn.Module", model)
[docs]
def load_artifacts(
self, model_directory: Path, n_points: int, device: torch.device
) -> None:
import torch
from torch_geometric.data import Data
edge_index = torch.load(
str(model_directory / "shared_edge_index.pt"),
map_location="cpu",
weights_only=True,
)
edge_feats = torch.load(
str(model_directory / "shared_edge_features.pt"),
map_location="cpu",
weights_only=True,
)
self._shared_edge_index = edge_index
self._shared_edge_feats = edge_feats.to(device)
self._shared_graph = Data(edge_index=edge_index, num_nodes=n_points).to(device)
[docs]
def predict(self, node_feats: np.ndarray) -> np.ndarray:
import torch
nf = torch.from_numpy(node_feats.astype(np.float32)).to(self._device)
with torch.no_grad():
pred = self._model(nf, self._shared_edge_feats, self._shared_graph)
return np.asarray(pred.cpu().numpy(), dtype=np.float32)