MoleculeNet Classification Data : Modularization, dynamic splits, Clean Code

.vscode/settings.json

@@ -2,7 +2,7 @@
     "python.testing.unittestArgs": [
         "-v",
         "-s",
-        "./tests",
+        "./tests/unit",
         "-p",
         "test*.py"
     ],

chebai/callbacks/epoch_metrics.py

@@ -1,5 +1,8 @@
+import warnings
+
 import torch
 import torchmetrics
+from sklearn.metrics import roc_auc_score
 
 
 def custom_reduce_fx(input: torch.Tensor) -> torch.Tensor:
@@ -179,3 +182,83 @@ def compute(self) -> torch.Tensor:
 
         balanced_acc = (tpr + tnr) / 2
         return torch.mean(balanced_acc)
+
+
+class HiMolMacroAUROC(torchmetrics.Metric):
+    """
+    Macro-averaged multilabel AUROC that exactly replicates HiMol's eval():
+      - missing labels (== ignore_index) are excluded per-task before scoring
+      - any task that, after exclusion, has only one class present is
+        dropped from BOTH the sum and the divisor of the macro average
+        (instead of being scored as 0.0 and diluting the mean, which is
+        torchmetrics' default MultilabelAUROC behavior)
+
+    Assumes preds/target are shape (N, num_labels), target values in {0, 1},
+    with `ignore_index` marking missing/unlabeled entries.
+
+    References:
+        https://github.com/ZangXuan/HiMol/blob/ffdcb247b361a1f85ddb741862cff25e4a3b3341/finetune/optimization.py#L95-L104
+    """
+
+    full_state_update = False
+    is_differentiable = False
+    higher_is_better = True
+
+    def __init__(self, num_labels: int, ignore_index: int = 0, **kwargs):
+        super().__init__(**kwargs)
+        self.num_labels = num_labels
+        self.ignore_index = ignore_index
+
+        self.add_state("preds", default=[], dist_reduce_fx="cat")
+        self.add_state("target", default=[], dist_reduce_fx="cat")
+
+    def update(self, preds: torch.Tensor, target: torch.Tensor) -> None:
+        if preds.shape != target.shape:
+            raise ValueError(
+                f"preds/target shape mismatch: {preds.shape} vs {target.shape}"
+            )
+        if preds.ndim != 2 or preds.shape[1] != self.num_labels:
+            raise ValueError(
+                f"expected shape (N, {self.num_labels}), got {tuple(preds.shape)}"
+            )
+
+        self.preds.append(preds.detach().cpu())
+        self.target.append(target.detach().cpu())
+
+    def compute(self) -> torch.Tensor:
+        preds = torch.cat(self.preds, dim=0).cpu().numpy()
+        target = torch.cat(self.target, dim=0).cpu().numpy()
+
+        roc_list = []
+        n_skipped = 0
+
+        for i in range(self.num_labels):
+            col_target = target[:, i]
+            col_preds = preds[:, i]
+
+            valid = col_target != self.ignore_index
+            col_target = col_target[valid]
+            col_preds = col_preds[valid]
+
+            # need at least one of each class to define AUC
+            if (
+                len(col_target) == 0
+                or (col_target == 0).sum() == 0
+                or (col_target == 1).sum() == 0
+            ):
+                n_skipped += 1
+                continue
+
+            roc_list.append(roc_auc_score(col_target, col_preds))
+
+        if n_skipped > 0:
+            warnings.warn(
+                f"{n_skipped}/{self.num_labels} labels skipped (missing or single-class "
+                f"after masking). Macro AUROC computed over {len(roc_list)} labels.",
+                stacklevel=2,
+            )
+
+        if len(roc_list) == 0:
+            return torch.tensor(float("nan"))
+
+        return torch.tensor(sum(roc_list) / len(roc_list))

chebai/cli.py

@@ -75,6 +75,8 @@ def call_data_methods(data: Type[XYBaseDataModule]):
                 "mse",
                 "rmse",
                 "r2",
+                "pr-auc",
+                "himol-marco-roc-auc",
             ):
                 # When using lightning > 2.5.1 then need to uncomment all metrics that are not used
                 # for average in ("mse", "rmse","r2"): # for regression

chebai/loss/asymmetric_loss.py

@@ -9,10 +9,10 @@ class AsymmetricLoss(nn.Module):
     Asymmetric Loss for multi-label and single-label classification tasks.
 
     Implementation adapted from: https://github.com/Alibaba-MIIL/ASL
-    
+
     Asymmetric Loss from: "Asymmetric Loss For Multi-Label Classification"
     https://openaccess.thecvf.com/content/ICCV2021/papers/Ben-Baruch_Asymmetric_Loss_For_Multi-Label_Classification_ICCV_2021_paper.pdf
-    
+
     Args:
         gamma_neg (float): Negative focusing parameter. Default is 1.
         gamma_pos (float): Positive focusing parameter. Default is 1.
@@ -47,12 +47,12 @@ def __init__(
     def forward(self, inputs, targets, **kwargs):
         """
         Forward pass to compute the Asymmetric Loss.
-        
+
         Args:
             inputs: Predictions (logits) from the model.
             targets: Ground truth labels.
             **kwargs: Additional keyword arguments (for compatibility with training framework).
-        
+
         Returns:
             Loss tensor with the reduction option applied.
         """
@@ -68,7 +68,7 @@ def forward(self, inputs, targets, **kwargs):
     def _multi_label_asymmetric_loss(self, x, y):
         """
         Standard asymmetric loss for multi-label classification.
-        
+
         Parameters
         ----------
         x: input logits
@@ -112,26 +112,30 @@ def _multi_label_asymmetric_loss(self, x, y):
     def _single_label_asymmetric_loss(self, inputs, target):
         """
         Asymmetric loss for single-label classification problems.
-        
+
         "input" dimensions: - (batch_size, number_classes)
         "target" dimensions: - (batch_size)
         """
         log_preds = self.logsoftmax(inputs)
-        self.targets_classes = torch.zeros_like(inputs).scatter_(1, target.long().unsqueeze(1), 1)
+        self.targets_classes = torch.zeros_like(inputs).scatter_(
+            1, target.long().unsqueeze(1), 1
+        )
 
         # ASL weights
         targets = self.targets_classes
         anti_targets = 1 - targets
         xs_pos = torch.exp(log_preds)
         xs_neg = 1 - xs_pos
         grad_ctx = (
-                torch.no_grad() if self.disable_torch_grad_focal_loss else nullcontext()
-            )
+            torch.no_grad() if self.disable_torch_grad_focal_loss else nullcontext()
+        )
         with grad_ctx:
             xs_pos = xs_pos * targets
             xs_neg = xs_neg * anti_targets
-            asymmetric_w = torch.pow(1 - xs_pos - xs_neg,
-                                 self.gamma_pos * targets + self.gamma_neg * anti_targets)
+            asymmetric_w = torch.pow(
+                1 - xs_pos - xs_neg,
+                self.gamma_pos * targets + self.gamma_neg * anti_targets,
+            )
             log_preds = log_preds * asymmetric_w
 
             # loss calculation

chebai/loss/base.py

@@ -0,0 +1,29 @@
+from typing import Optional
+
+import torch
+
+
+class BCELogitLossWithValidLabels(torch.nn.BCEWithLogitsLoss):
+    def __init__(self, **kwargs):
+        kwargs["reduction"] = "none"
+        super().__init__(**kwargs)
+
+    def forward(
+        self,
+        input: torch.Tensor,
+        target: torch.Tensor,
+        valid_label_mask: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        loss_mat = super().forward(input, target)
+
+        if valid_label_mask is None:
+            return loss_mat.mean()
+
+        loss_mat = torch.where(
+            valid_label_mask,
+            loss_mat,
+            torch.zeros_like(loss_mat),
+        )
+
+        return loss_mat.sum() / valid_label_mask.sum().clamp_min(1)

chebai/loss/bce_weighted.py

@@ -3,11 +3,12 @@
 
 import torch
 
+from chebai.loss.base import BCELogitLossWithValidLabels
 from chebai.preprocessing.datasets.base import XYBaseDataModule
 from chebai.preprocessing.datasets.chebi import _ChEBIDataExtractor
 
 
-class BCEWeighted(torch.nn.BCEWithLogitsLoss):
+class BCEWeighted(BCELogitLossWithValidLabels):
     """
     BCEWithLogitsLoss with weights automatically computed according to the beta parameter.
 
@@ -104,10 +105,9 @@ def forward(
             torch.Tensor: The computed loss.
         """
         self.set_pos_weight(input)
-        return super().forward(input, target)
+        return super().forward(input, target, **kwargs)
 
 
-class UnWeightedBCEWithLogitsLoss(torch.nn.BCEWithLogitsLoss):
+class UnWeightedBCEWithLogitsLoss(BCELogitLossWithValidLabels):
     def forward(self, input, target, **kwargs):
-        # As the custom passed kwargs are not used in BCEWithLogitsLoss, we can ignore them
-        return super().forward(input, target)
+        return super().forward(input, target, **kwargs)

chebai/models/base.py

@@ -295,6 +295,9 @@ def _execute(
                     model_output, labels, data.get("loss_kwargs", dict())
                 )
                 loss_kwargs = dict()
+                loss_kwargs["valid_label_mask"] = loss_kwargs_candidates[
+                    "valid_label_mask"
+                ]
                 if self.pass_loss_kwargs:
                     loss_kwargs = loss_kwargs_candidates
                     loss_kwargs["current_epoch"] = self.trainer.current_epoch

chebai/preprocessing/collate.py

@@ -61,10 +61,13 @@ def __call__(self, data: List[Union[Dict, Tuple]]) -> XYData:
         """
         Collate ragged data samples (i.e., samples of unequal size, such as molecular sequences) into a batch.
 
-        Handles both fully and partially labeled data, where some samples may have `None` as their label. The indices
-        of non-null labels are stored in the `non_null_labels` field, which is used to filter out predictions for
-        unlabeled data during evaluation (e.g., F1, MSE). For models supporting partially labeled data, this method
-        ensures alignment between features and labels. Missing labels are passed as a loss keyword.
+        Handles both fully and partially labeled data by use of the following fields in the returned XYData:
+
+        `non_null_labels`: Stores batch row indices of samples where the whole `labels` field is not None, like [0, 2].
+            - Example: [[True, False], None, [False, None]] would result in `non_null_labels` = [0, 2].
+            - This is used to filter out predictions for unlabeled samples during evaluation.
+
+        `valid_label_mask`: Stores a per-sample, per-label-position boolean mask for valid entries inside a label row, like the True in [1, None, 0].
 
         Args:
             data (List[Union[Dict, Tuple]]): List of ragged data samples. Each sample can be a dictionary or tuple
@@ -81,38 +84,33 @@ def __call__(self, data: List[Union[Dict, Tuple]]) -> XYData:
         if isinstance(data[0], tuple):
             # For legacy data
             x, y, idents = zip(*data)
-            missing_labels = None
         else:
             x, y, idents = zip(
                 *((d["features"], d["labels"], d.get("ident")) for d in data)
             )
-            missing_labels = [
-                d.get(
-                    "missing_labels",
-                    [False for _ in y[0]] if y[0] is not None else [False],
-                )
-                for d in data
-            ]
 
+        valid_label_mask = self._get_valid_label_mask(y)
+
+        # Typical y: ([True, False], None, [True, None], [True])
         if any(x is not None for x in y):
-            # If any label is not None: (None, None, `1`, None)
+            # If any label is not None: (None, None, `[True, None]`, None)
             if any(x is None for x in y):
-                # If any label is None: (`None`, `None`, 1, `None`)
+                # If any label is None: (`None`, [True, False], [True], [False])
                 non_null_labels = [i for i, r in enumerate(y) if r is not None]
                 y = self.process_label_rows(
                     tuple(ye for i, ye in enumerate(y) if i in non_null_labels)
                 )
                 loss_kwargs["non_null_labels"] = non_null_labels
             else:
-                # If all labels are not None: (`0`, `2`, `1`, `3`)
+                # If all labels are not None: (`[True, False]`, `[False, True, True]`, `[False]`, `[True]`)
                 y = self.process_label_rows(y)
 
         else:
             # If all labels are None : (`None`, `None`, `None`, `None`)
             y = None
             loss_kwargs["non_null_labels"] = []
 
-        loss_kwargs["missing_labels"] = torch.tensor(missing_labels)
+        loss_kwargs["valid_label_mask"] = valid_label_mask
         # Calculate the lengths of each sequence, create a binary mask for valid (non-padded) positions
         lens = torch.tensor(list(map(len, x)))
         model_kwargs["mask"] = torch.arange(max(lens))[None, :] < lens[:, None]
@@ -146,3 +144,26 @@ def process_label_rows(self, labels: Tuple) -> torch.Tensor:
             ],
             batch_first=True,
         )
+
+    def _get_valid_label_mask(self, y: Tuple) -> torch.Tensor | None:
+        # Compute the per-sample, per-label-position boolean mask for unknown entries
+        # (e.g., the None in [1, None, 0]) on the *original* labels, before any
+        # filtering/padding is applied to `y`. Rows whose entire label is None are
+        # represented as all-False rows of the maximum label length.
+        if any(labels is not None for labels in y):
+            max_label_len = max(len(labels) for labels in y if labels is not None)
+            valid_label_mask = pad_sequence(
+                [
+                    torch.tensor([label is not None for label in labels])
+                    if labels is not None
+                    else torch.zeros(max_label_len, dtype=torch.bool)
+                    for labels in y
+                ],
+                batch_first=True,
+            )
+            if (~valid_label_mask).sum() != 0:
+                # If there are any invalid labels, return the valid_label_mask
+                # Else, return None to indicate that all labels are valid (no None entries).
+                return valid_label_mask
+
+        return None