Merge remote-tracking branch 'refs/remotes/origin/feature_grit_prNov25' into feature_grit_prNov25

Author: tolhq

examples/config/GRIT_PF_datakit_case14.yaml

@@ -5,7 +5,8 @@ task:
   task_name: PowerFlow
 data:
   baseMVA: 100
-  mask_value: 0.0
+  mask_type: rnd        # or determinstic
+  mask_ratio: 0.5       # for random masking only
   normalization: HeteroDataMVANormalizer
   networks:
   - case14_ieee

gridfm_graphkit/datasets/masking.py

@@ -33,6 +33,60 @@
 from torch_geometric.nn import MessagePassing
 
 
+class AddRandomHeteroMask(BaseTransform):
+    """Creates random masks for self-supervised pretraining on heterogeneous power grid graphs.
+
+    Each selected feature dimension is independently masked per node/edge with
+    probability ``mask_ratio``.  Masked bus features: VM, VA, QG.  Masked gen
+    features: PG.  Masked branch features: P_E, Q_E.
+
+    The output ``data.mask_dict`` has the same structure as the deterministic
+    PF / OPF masks so that downstream losses (``MaskedBusMSE``, ``MaskedGenMSE``,
+    ``PBELoss``, etc.) work without modification.
+    """
+
+    def __init__(self, mask_ratio=0.5):
+        super().__init__()
+        self.mask_ratio = mask_ratio
+
+    def forward(self, data):
+        bus_x = data.x_dict["bus"]
+        gen_x = data.x_dict["gen"]
+
+        # Bus type indicators (needed by losses and test metrics)
+        mask_PQ = bus_x[:, PQ_H] == 1
+        mask_PV = bus_x[:, PV_H] == 1
+        mask_REF = bus_x[:, REF_H] == 1
+
+        # Random bus mask on variable features the model reconstructs
+        mask_bus = torch.zeros_like(bus_x, dtype=torch.bool)
+        n_bus = bus_x.size(0)
+        for feat_idx in (VM_H, VA_H, QG_H):
+            mask_bus[:, feat_idx] = torch.rand(n_bus) < self.mask_ratio
+
+        # Random gen mask on PG
+        mask_gen = torch.zeros_like(gen_x, dtype=torch.bool)
+        mask_gen[:, PG_H] = torch.rand(gen_x.size(0)) < self.mask_ratio
+
+        # Random branch mask on flow features
+        branch_attr = data.edge_attr_dict[("bus", "connects", "bus")]
+        mask_branch = torch.zeros_like(branch_attr, dtype=torch.bool)
+        n_edge = branch_attr.size(0)
+        for feat_idx in (P_E, Q_E):
+            mask_branch[:, feat_idx] = torch.rand(n_edge) < self.mask_ratio
+
+        data.mask_dict = {
+            "bus": mask_bus,
+            "gen": mask_gen,
+            "branch": mask_branch,
+            "PQ": mask_PQ,
+            "PV": mask_PV,
+            "REF": mask_REF,
+        }
+
+        return data
+
+
 class AddPFHeteroMask(BaseTransform):
     """Creates masks for a heterogeneous power flow graph."""
 

gridfm_graphkit/datasets/task_transforms.py

@@ -8,6 +8,7 @@
 from gridfm_graphkit.datasets.masking import (
     AddOPFHeteroMask,
     AddPFHeteroMask,
+    AddRandomHeteroMask,
     SimulateMeasurements,
 )
 from gridfm_graphkit.io.registries import TRANSFORM_REGISTRY
@@ -20,7 +21,13 @@ def __init__(self, args):
 
         transforms.append(RemoveInactiveBranches())
         transforms.append(RemoveInactiveGenerators())
-        transforms.append(AddPFHeteroMask())
+
+        mask_type = getattr(args.data, "mask_type", None)
+        if mask_type == "rnd":
+            transforms.append(AddRandomHeteroMask(mask_ratio=args.data.mask_ratio))
+        else:
+            transforms.append(AddPFHeteroMask())
+
         transforms.append(ApplyMasking(args=args))
 
         # Pass the list of transforms to Compose
@@ -34,7 +41,13 @@ def __init__(self, args):
 
         transforms.append(RemoveInactiveBranches())
         transforms.append(RemoveInactiveGenerators())
-        transforms.append(AddOPFHeteroMask())
+
+        mask_type = getattr(args.data, "mask_type", None)
+        if mask_type == "rnd":
+            transforms.append(AddRandomHeteroMask(mask_ratio=args.data.mask_ratio))
+        else:
+            transforms.append(AddOPFHeteroMask())
+
         transforms.append(ApplyMasking(args=args))
 
         # Pass the list of transforms to Compose

scripts/benchmark_model_inference.py

@@ -16,8 +16,8 @@
 OUT_DIR=../scripts
 mkdir $OUT_DIR
 
-python benchmark_model_inference.py --model hetero --config $CONF_PATH/case30_ieee_base.yaml --num_nodes 30 --num_edges 82 --num_gens 6 --iterations 20 --output_csv $OUT_DIR/case30.csv || true
-python benchmark_model_inference.py --model hetero --config $CONF_PATH/case118_ieee_base.yaml --num_nodes 118 --num_edges 372 --num_gens 54 --iterations 20 --output_csv $OUT_DIR/case118.csv || true
+python benchmark_model_inference.py --model hetero --config $CONF_PATH/HGNS_PF_datakit_case30.yaml --num_nodes 30 --num_edges 82 --num_gens 6 --iterations 20 --output_csv $OUT_DIR/case30.csv || true
+python benchmark_model_inference.py --model hetero --config $CONF_PATH/HGNS_PF_datakit_case118.yaml --num_nodes 118 --num_edges 372 --num_gens 54 --iterations 20 --output_csv $OUT_DIR/case118.csv || true
 
 ######################################
 
@@ -117,11 +117,11 @@
 
 if MODEL_TYPE == "grit":
     # Positional encoding config (only GRIT uses these)
+    # Read enablement and dimensions from data config (canonical source).
     RRWP_ENABLED = getattr(config_args.data.posenc_RRWP, "enable", False) if hasattr(config_args.data, "posenc_RRWP") else False
     RRWP_KSTEPS = getattr(config_args.data.posenc_RRWP, "ksteps", 21) if RRWP_ENABLED else 0
-    RWSE_ENABLED = hasattr(config_args.model, "encoder") and getattr(config_args.model.encoder, "node_encoder", False) \
-                   and "RWSE" in getattr(config_args.model.encoder, "node_encoder_name", "")
-    RWSE_TIMES = getattr(config_args.model.encoder.posenc_RWSE.kernel, "times", 21) if RWSE_ENABLED else 0
+    RWSE_ENABLED = hasattr(config_args.data, "posenc_RWSE") and getattr(config_args.data.posenc_RWSE, "enable", False)
+    RWSE_TIMES = getattr(config_args.data.posenc_RWSE.kernel, "times", 21) if RWSE_ENABLED else 0
 else:
     RRWP_ENABLED = False
     RRWP_KSTEPS = 0