Merge pull request #25 from mhpi/multimodel_dev

Data Loader Update

Author: leoglonz

.gitignore

@@ -43,3 +43,4 @@ deltaModel/conf/observations/*
 archive/
 runs/
 results/
+validation/

deltaModel/__main__.py

@@ -4,16 +4,16 @@
 
 import hydra
 import torch
-from core.data.dataset_loading import get_dataset_dict
+from omegaconf import DictConfig
+
 from core.utils import initialize_config, print_config, set_randomseed
+from core.utils.module_loaders import load_data_loader, load_trainer
 from models.model_handler import ModelHandler as dModel
-from omegaconf import DictConfig
 from trainers.trainer import Trainer
 
 log = logging.getLogger(__name__)
 
 
-
 @hydra.main(
     version_base='1.3',
     config_path='conf/',
@@ -31,15 +31,22 @@ def main(config: DictConfig) -> None:
         print_config(config)
 
         ### Create/Load differentiable model ###
-        model = dModel(config, verbose=True) #.to(config['device'])
+        model = dModel(config, verbose=True)
 
         ### Process datasets ###
-        log.info("Processing datasets...")
-        train_dataset = get_dataset_dict(config, train=True)
-        eval_dataset = get_dataset_dict(config, train=False)
+        log.info("Loading dataset...")
+        data_loader = load_data_loader(config['data_loader'])
+        data_loader = data_loader(config, test_split=True, overwrite=False)
 
         ### Create Trainer object ###
-        trainer = Trainer(config, model, train_dataset, eval_dataset, verbose=True)
+        trainer = load_trainer(config['trainer'])
+        trainer = trainer(
+            config,
+            model,
+            data_loader.train_dataset,
+            data_loader.eval_dataset,
+            verbose=True
+        )
 
         mode = config['mode']
         if mode == 'train':

deltaModel/conf/config.py

@@ -105,7 +105,7 @@ class ObservationConfig(BaseModel):
     name: str = 'not_defined'
     train_path: str = 'not_defined'
     test_path: str = 'not_defined'
-    start_time_all: str = 'not_defined'
+    start_time: str = 'not_defined'
     end_time_all: str = 'not_defined'
     forcings_all: List[str] = Field(default_factory=list, description="List of dynamic input variables.")
     attributes_all: List[str] = Field(default_factory=list, description="List of static input variables.")
@@ -137,6 +137,9 @@ class Config(BaseModel):
     random_seed: int = 0
     device: str = 'cpu'
     gpu_id: int = 0
+    data_loader: str = 'none'
+    data_sampler: str = 'none'
+    trainer: str = 'none'
     save_path: str
     train: TrainingConfig
     test: TestingConfig
@@ -182,6 +185,9 @@ def check_device(cls, values):
             random_seed=42,
             device='cuda',
             gpu_id=0,
+            data_loader='base_loader',
+            data_sampler='base_sampler',
+            trainer='trainer',
             save_path='../results',
             train={
                 'start_time': '2000/01/01',

deltaModel/conf/config.yaml

@@ -11,6 +11,10 @@ random_seed: 111111
 device: cuda
 gpu_id: 0
 
+data_loader: base_loader
+data_sampler: base_sampler
+trainer: trainer
+
 save_path: ../results
 
 

deltaModel/conf/observations/none.yaml

@@ -2,8 +2,8 @@ name: None
 train_path: /path/to/training/data
 test_path: /path/to/validation/data
 
-start_time_all: 2000/01/01
-end_time_all: 2024/12/31
+start_time: 2000/01/01
+end_time: 2024/12/31
 
 forcings_all: [
     x1_var,

deltaModel/core/utils/Dropout.py deltaModel/core/calc/dropout.pydeltaModel/core/utils/Dropout.py renamed to deltaModel/core/calc/dropout.py

@@ -15,12 +15,10 @@
 class Metrics(BaseModel):
     """Metrics for model evaluation.
 
-    Using pydantic BaseModel for validation.
+    Using Pydantic BaseModel for validation.
     Metrics are calculated at each grid point and are listed below.
     
     Adapted from Tadd Bindas, Yalan Song, Farshid Rahmani.
-
-    Note: Considering conversion to torch.nn.module for codebase consistency.
     """
     model_config = ConfigDict(arbitrary_types_allowed=True)
     pred: npt.NDArray[np.float32]
@@ -70,9 +68,11 @@ def __init__(
 
         super().__init__(pred=pred, target=target)
 
-    def model_post_init(self, __context: Any):
+    def model_post_init(self, __context: Any) -> Any:
         """Calculate metrics.
-        
+
+        This method is called after the model is initialized.
+
         Parameters
         ----------
         __context : Any
@@ -208,7 +208,7 @@ def validate_pred(cls, metrics: Any) -> Any:
             raise ValueError(msg)
         return metrics
 
-    def calc_statistics(self, *args, **kwargs) -> Dict[str, Dict[str, float]]:
+    def calc_stats(self, *args, **kwargs) -> Dict[str, Dict[str, float]]:
         """Calculate aggregate statistics of metrics."""
         stats = {}
         model_dict = self.model_dump()
@@ -225,15 +225,15 @@ def calc_statistics(self, *args, **kwargs) -> Dict[str, Dict[str, float]]:
                 }
         return stats
 
-    def dump_agg_statistics(self, path: str) -> None:
+    def model_dump_agg_stats(self, path: str) -> None:
         """Dump aggregate statistics (median, mean, std) to json or csv.
         
         Parameters
         ----------
         path : str
             Path to save file.
         """
-        stats = self.calc_statistics()
+        stats = self.calc_stats()
 
         if path.endswith('.json'):
             with open(path, 'w') as f:
@@ -271,7 +271,7 @@ def dump_metrics(self, path: str) -> None:
         """
         # Save aggregate statistics
         save_path = os.path.join(path, 'metrics_agg.json')
-        self.dump_agg_statistics(save_path)
+        self.model_dump_agg_stats(save_path)
 
         # Save raw metrics
         save_path = os.path.join(path, f'metrics.json')
@@ -280,16 +280,6 @@ def dump_metrics(self, path: str) -> None:
         with open(save_path, "w") as f:
             json.dump(json_dat, f)        
 
-    @property
-    def ngrid(self) -> int:
-        """Calculate number of items in grid."""
-        return self.pred.shape[0]
-    
-    @property
-    def nt(self) -> int:
-        """Calculate number of time steps."""
-        return self.pred.shape[1]
-
     def tile_mean(self, data: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
         """Calculate mean of target.
 
@@ -305,6 +295,16 @@ def tile_mean(self, data: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
         """
         return np.tile(np.nanmean(data, axis=1), (self.nt, 1)).transpose()
 
+    @property
+    def ngrid(self) -> int:
+        """Calculate number of items in grid."""
+        return self.pred.shape[0]
+    
+    @property
+    def nt(self) -> int:
+        """Calculate number of time steps."""
+        return self.pred.shape[1]
+    
     @staticmethod
     def _bias(
         pred: npt.NDArray[np.float32],