ENH: Warn when physics-informed motion training collapses to the trivial solution

src/monai_physio/train_physicsnemo_physics_informed_motion.py

@@ -522,6 +522,9 @@ def __init__(self, log_level: int | str = logging.INFO) -> None:
         self._epoch_data_loss: Optional[torch.Tensor] = None
         self._epoch_physics_loss: Optional[torch.Tensor] = None
         self._epoch_batches = 0
+        # Consecutive epochs with physics_mean ~ 0 while lambda_physics > 0;
+        # see the collapse warning in _log_epoch.
+        self._zero_physics_streak = 0
 
     def set_mechanics(
         self,
@@ -833,6 +836,21 @@ def _log_epoch(self, context: DistributedContext, epoch: int, epochs: int) -> No
 
         divisor = max(batches, 1)
         physics_mean = physics_sum / divisor
+
+        if self.lambda_physics > 0.0 and physics_mean < 1e-9:
+            self._zero_physics_streak += 1
+        else:
+            self._zero_physics_streak = 0
+        if self._zero_physics_streak == 5 and context.is_main:
+            self.log_warning(
+                "physics loss has been ~0 for 5 consecutive epochs while "
+                "lambda_physics=%.4g; the network likely collapsed to the "
+                "trivial uniform-displacement solution described in "
+                "set_lambda_physics_warmup()'s docstring. Consider a longer "
+                "warmup or a lower lambda_physics target.",
+                self.lambda_physics,
+            )
+
         self._log_main(
             context,
             # physics/weighted in scientific notation: %f rounds anything

tests/test_physics_informed_motion.py

@@ -502,6 +502,56 @@ def test_the_epoch_log_separates_the_two_loss_terms() -> None:
     )
 
 
+def test_zero_physics_streak_warns_once_after_five_epochs() -> None:
+    """A collapsed physics term should be flagged, not silently trained through.
+
+    ``set_lambda_physics_warmup``'s docstring describes a trivial-solution
+    collapse where the physics term reads exactly zero for the rest of the
+    run. Five consecutive zero-physics epochs should raise exactly one
+    warning -- not one per epoch, and not at all before the streak or once
+    physics recovers.
+    """
+    import torch
+
+    from monai_physio.process_physicsnemo import DistributedContext
+    from monai_physio.train_physicsnemo_physics_informed_motion import (
+        TrainPhysicsNeMoPhysicsInformedMotion,
+    )
+
+    method = TrainPhysicsNeMoPhysicsInformedMotion()
+    method.lambda_physics = 0.1
+
+    warnings: list[str] = []
+    method.log_warning = lambda *args: warnings.append(str(args[0]) % args[1:])  # type: ignore[method-assign]
+    method.log_info = lambda *args: None  # type: ignore[method-assign]
+
+    context = DistributedContext(
+        device=torch.device("cpu"), rank=0, local_rank=0, world_size=1
+    )
+
+    def log_epoch_with(physics: float, epoch: int) -> None:
+        method._epoch_data_loss = torch.tensor(1.0)
+        method._epoch_physics_loss = torch.tensor(physics)
+        method._epoch_batches = 1
+        method._log_epoch(context, epoch=epoch, epochs=10)
+
+    for epoch in range(4):
+        log_epoch_with(0.0, epoch)
+    assert not warnings, "Should not warn before five consecutive zero epochs"
+
+    log_epoch_with(0.0, 4)
+    assert len(warnings) == 1, "Should warn exactly once at the fifth zero epoch"
+    assert "collapsed" in warnings[0]
+
+    log_epoch_with(0.0, 5)
+    assert len(warnings) == 1, "Should not warn again every epoch after"
+
+    log_epoch_with(1.0, 6)
+    for epoch in range(4):
+        log_epoch_with(0.0, 7 + epoch)
+    assert len(warnings) == 1, "A streak broken by recovery should not re-warn early"
+
+
 def test_bind_reference_meshes_repairs_against_template_elements(tmp_path: Any) -> None:
     """Repair must use ``self._tets``, not whatever cells the file stores.