feat(l2gprediction): add model as attribute

src/gentropy/dataset/l2g_prediction.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import TYPE_CHECKING
 
 import pyspark.sql.functions as f
@@ -29,6 +29,8 @@ class L2GPrediction(Dataset):
     confidence of the prediction that a gene is causal to an association.
     """
 
+    model: LocusToGeneModel | None = field(default=None, repr=False)
+
     @classmethod
     def get_schema(cls: type[L2GPrediction]) -> StructType:
         """Provides the schema for the L2GPrediction dataset.
@@ -82,7 +84,6 @@ def from_credible_set(
             .fill_na()
             .select_features(l2g_model.features_list)
         )
-
         return l2g_model.predict(fm, session)
 
     def to_disease_target_evidence(
@@ -172,4 +173,5 @@ def add_locus_to_gene_features(
                 aggregated_features, on=["studyLocusId", "geneId"], how="left"
             ),
             _schema=self.get_schema(),
+            model=self.model,
         )

src/gentropy/method/l2g/model.py

@@ -175,6 +175,7 @@ def predict(
         return L2GPrediction(
             _df=session.spark.createDataFrame(feature_matrix_pdf.filter(output_cols)),
             _schema=L2GPrediction.get_schema(),
+            model=self,
         )
 
     def save(self: LocusToGeneModel, path: str) -> None: