Added tests for reading configuration file

pyproject.toml

@@ -36,10 +36,13 @@ PyYAML = "^6.0"
 scikit-learn = "^1.0.2"
 xgboost = "^1.6.1"
 pandas = "^1.4.3"
+pydantic = "^2.6.4"
+pydantic-yaml = "^1.2.1"
 
 [tool.poetry.dev-dependencies]
 pytest = "7.1.2"
 pytest-cov = "^4.1.0"
+pytest-subtests = "^0.12.1"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]

src/pd_dwi/config/config.py

@@ -0,0 +1,99 @@
+from enum import Enum
+from typing import List, Dict, Any, Optional, Set
+
+from pydantic import BaseModel, Field, PositiveInt, NonNegativeInt, root_validator, model_validator
+
+
+class Labels(BaseModel):
+    negative: str = Field(min_length=1)
+    positive: str = Field(min_length=1)
+
+
+class TimePoint(Enum):
+    T0 = 'T0'
+    T1 = 'T1'
+    T2 = 'T2'
+
+
+class Modality(Enum):
+    ADC0100 = 'ADC 0100'
+    ADC0100600800 = 'ADC 0100600800'
+    ADC100600800 = 'ADC 100600800'
+    F = 'F'
+
+
+class Mask(Enum):
+    DWI = 'DWI MASK'
+
+
+class Dataset(BaseModel):
+    labels: Labels
+    time_points: Set[TimePoint] = Field(min_items=1)
+    modalities: Set[Modality] = Field(min_items=1)
+    masks: Set[Mask]
+
+
+class RadiomicsFeaturesEncoder(BaseModel):
+    image: Modality
+    mask: Mask
+    time_points: Set[TimePoint] = Field(min_items=1)
+
+
+class RadiomicsFeaturesTransformer(BaseModel):
+    encoders: List[RadiomicsFeaturesEncoder]
+    engine: Dict[str, Any]
+
+
+class FeatureTransformer(BaseModel):
+    radiomics: RadiomicsFeaturesTransformer
+
+
+class FeatureSelection(BaseModel):
+    k: int
+
+
+class Classifier(BaseModel):
+    module: str
+    parameters: Dict[str, Any]
+
+
+class Pipeline(BaseModel):
+    features_transformer: FeatureTransformer
+    feature_selection: FeatureSelection
+    classifier: Classifier
+
+
+class GridSearchParamGrid(BaseModel):
+    classifier: Dict[str, Any] = None
+    feature_selection: Dict[str, Any] = None
+
+
+class GridSearch(BaseModel):
+    verbose: NonNegativeInt
+    scoring: str = Field(default='roc_auc')
+    cv: PositiveInt = Field(default=5)
+    param_grid: GridSearchParamGrid
+
+
+class ModelConfig(BaseModel):
+    dataset: Dataset
+    pipeline: Pipeline
+    grid_search_cv: GridSearch = None
+
+    class Config:
+        frozen = True
+        extra = 'forbid'
+
+    @model_validator(mode='after')
+    def validate_encoders_dataset(self):
+        modalities = {e.image for e in self.pipeline.features_transformer.radiomics.encoders}
+
+        if not modalities.issubset(self.dataset.modalities):
+            raise ValueError("Encoders contain modalities that are not available in dataset")
+
+        masks = {e.mask for e in self.pipeline.features_transformer.radiomics.encoders}
+
+        if not masks.issubset(self.dataset.masks):
+            raise ValueError("Encoders contain masks that are not available in dataset")
+
+        return self

src/pd_dwi/configurations/xgboost/F.yaml → src/pd_dwi/config/samples/F.yaml

@@ -16,6 +16,7 @@ dataset:
 
 pipeline:
   features_transformer:
+    radiomics:
       encoders: # encoder names format is {time point}_{image}
         - image: F
           mask: DWI MASK

src/pd_dwi/configurations/xgboost/adc100600800.yaml → src/pd_dwi/config/samples/adc100600800.yaml

@@ -16,6 +16,7 @@ dataset:
 
 pipeline:
   features_transformer:
+    radiomics:
       encoders: # encoder names format is {time point}_{image}
         - image: ADC 100600800
           mask: DWI MASK

src/pd_dwi/config/utils.py

@@ -0,0 +1,15 @@
+from yaml import load, FullLoader
+
+from pd_dwi.config.config import ModelConfig
+
+
+def read_config(config):
+    if hasattr(config, 'read'):
+        config = load(config, Loader=FullLoader)
+    elif config.endswith('.yaml') or config.endswith('.yml'):
+        config = load(open(config), Loader=FullLoader)
+    else:
+        raise NotImplementedError()
+
+    ModelConfig.model_validate(config)
+    return config

src/pd_dwi/model.py

@@ -1,11 +1,15 @@
+from pathlib import Path
 from pickle import dump, load as pkl_load, HIGHEST_PROTOCOL
 
 import pandas as pd
 from jsonschema.validators import validate
+from pydantic_yaml import parse_yaml_raw_as
 from sklearn.metrics import roc_auc_score
 from sklearn.model_selection import GridSearchCV
 from yaml import FullLoader, load
 
+from pd_dwi.config.config import ModelConfig
+from pd_dwi.config.utils import read_config
 from pd_dwi.dataset import create_dataset, validate_dataset
 from pd_dwi.training_utils import create_model_from_config
 
@@ -17,14 +21,7 @@ def __init__(self, config=None, model_obj=None):
 
     @classmethod
     def from_config(cls, config):
-        if hasattr(config, 'read'):
-            config = load(config, Loader=FullLoader)
-        elif not isinstance(config, dict):
-            raise NotImplementedError()
-
-        validate(instance=config, schema=load(open('./configurations/schema.yaml'), Loader=FullLoader))
-
-        return cls(config=config)
+        return cls(config=read_config(config))
 
     def save(self, path):
         assert self.model is not None