add docstring

Signed-off-by: ytl0623 <[email protected]>

Author: ytl0623

monai/losses/unified_focal_loss.py

@@ -73,7 +73,7 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
 
         if not self.include_background:
             if n_pred_ch == 1:
-                warnings.warn("single channel prediction, `include_background=False` ignored.")
+                warnings.warn("single channel prediction, `include_background=False` ignored.", stacklevel=2)
             else:
                 # if skipping background, removing first channel
                 y_true = y_true[:, 1:]
@@ -110,7 +110,6 @@ def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
                 loss_list.append(1 - dice_class[:, i])
             else:
                 # Foreground classes: apply focal modulation
-                # Original logic: (1 - dice) * (1 - dice)^(-gamma) -> (1-dice)^(1-gamma)
                 back_dice = torch.clamp(1 - dice_class[:, i], min=self.epsilon)
                 loss_list.append(back_dice * torch.pow(back_dice, -self.gamma))
 
@@ -176,8 +175,11 @@ def __init__(
     def forward(self, y_pred: torch.Tensor, y_true: torch.Tensor) -> torch.Tensor:
         """
         Args:
-            y_pred: (BNH[WD]) Logits (raw scores).
+            y_pred: (BNH[WD]) Logits (raw scores, not probabilities).
+                Do not pass pre-activated inputs; activation is applied internally.
             y_true: (BNH[WD]) Ground truth labels.
+        Returns:
+            torch.Tensor: Weighted combination of focal loss and asymmetric focal Tversky loss.
         """
         focal_loss = self.focal_loss(y_pred, y_true)
 

tests/losses/test_unified_focal_loss.py

@@ -72,6 +72,14 @@ class TestAsymmetricUnifiedFocalLoss(unittest.TestCase):
 
     @parameterized.expand(TEST_CASES)
     def test_result(self, input_param, input_data, expected_val):
+        """
+        Test AsymmetricUnifiedFocalLoss with various configurations.
+        
+        Args:
+            input_param: Dict of loss constructor parameters (use_softmax, include_background, etc.).
+            input_data: Dict containing y_pred (logits) and y_true (ground truth) tensors.
+            expected_val: Expected loss value.
+        """
         loss = AsymmetricUnifiedFocalLoss(**input_param)
         result = loss(**input_data)
         np.testing.assert_allclose(result.detach().cpu().numpy(), expected_val, atol=1e-4, rtol=1e-4)