fix(l2g): calculate_feature_missingness_rate counts features annotated with 0 as incomplete

src/otg/dataset/l2g_feature_matrix.py

@@ -109,7 +109,12 @@ def calculate_feature_missingness_rate(
             raise ValueError("No features found")
 
         return {
-            feature: (self._df.filter(self._df[feature].isNull()).count() / total_count)
+            feature: (
+                self._df.filter(
+                    (self._df[feature].isNull()) | (self._df[feature] == 0)
+                ).count()
+                / total_count
+            )
             for feature in self.features_list
         }
 

src/otg/method/l2g/evaluator.py

@@ -4,18 +4,17 @@
 import itertools
 from typing import TYPE_CHECKING, Any, Dict
 
-import wandb
 from pyspark import keyword_only
 from pyspark.ml.evaluation import (
     BinaryClassificationEvaluator,
     Evaluator,
     MulticlassClassificationEvaluator,
 )
 from pyspark.ml.param import Param, Params, TypeConverters
+from wandb.sdk.wandb_run import Run
 
 if TYPE_CHECKING:
     from pyspark.sql import DataFrame
-    from wandb.wandb_run import Run
 
 
 class WandbEvaluator(Evaluator):
@@ -124,11 +123,11 @@ def getspark_ml_evaluator(self: WandbEvaluator) -> Evaluator:
         """
         return self.getOrDefault(self.spark_ml_evaluator)
 
-    def getwandb_run(self: WandbEvaluator) -> wandb.sdk.wandb_run.Run:
+    def getwandb_run(self: WandbEvaluator) -> Run:
         """Get the wandb_run parameter.
 
         Returns:
-            wandb.sdk.wandb_run.Run: Wandb run object.
+            Run: Wandb run object.
         """
         return self.getOrDefault(self.wandb_run)
 

src/otg/method/l2g/model.py

@@ -5,14 +5,15 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Type
 
-import wandb
 from pyspark.ml import Pipeline, PipelineModel
 from pyspark.ml.evaluation import (
     BinaryClassificationEvaluator,
     MulticlassClassificationEvaluator,
 )
 from pyspark.ml.feature import StringIndexer, VectorAssembler
 from pyspark.ml.tuning import ParamGridBuilder
+from wandb.data_types import Table
+from wandb.sdk import init as wandb_init
 from wandb.wandb_run import Run
 from xgboost.spark.core import SparkXGBClassifierModel
 
@@ -126,7 +127,7 @@ def log_to_wandb(
         ## Track feature importance
         wandb_run.log({"importances": self.get_feature_importance()})
         ## Track training set
-        training_table = wandb.Table(dataframe=training_data.df.toPandas())
+        training_table = Table(dataframe=training_data.df.toPandas())
         wandb_run.log({"trainingSet": training_table})
         # Count number of positive and negative labels
         gs_counts_dict = {
@@ -224,7 +225,7 @@ def evaluate(
         )
 
         if wandb_run_name and training_data:
-            run = wandb.init(
+            run = wandb_init(
                 project=self.wandb_l2g_project_name,
                 config=hyperparameters,
                 name=wandb_run_name,

tests/dataset/test_l2g.py

@@ -149,3 +149,27 @@ def test_remove_false_negatives(spark: SparkSession) -> None:
     )
 
     assert observed_df.collect() == expected_df.collect()
+
+
+def test_calculate_feature_missingness_rate(spark: SparkSession) -> None:
+    """Test L2GFeatureMatrix.calculate_feature_missingness_rate."""
+    fm = L2GFeatureMatrix(
+        _df=spark.createDataFrame(
+            [
+                (1, "gene1", 100.0, None),
+                (2, "gene2", 1000.0, 0.0),
+            ],
+            "studyLocusId LONG, geneId STRING, distanceTssMean DOUBLE, distanceTssMinimum DOUBLE",
+        ),
+        _schema=L2GFeatureMatrix.get_schema(),
+    )
+
+    expected_missingness = {"distanceTssMean": 0.0, "distanceTssMinimum": 1.0}
+    observed_missingness = fm.calculate_feature_missingness_rate()
+    assert isinstance(observed_missingness, dict)
+    assert len(observed_missingness) == len(
+        fm.features_list  # type: ignore
+    ), "Missing features in the missingness rate dictionary."
+    assert (
+        observed_missingness == expected_missingness
+    ), "Missingness rate is incorrect."