Using new parquet in train

paper_eval.py

@@ -1,25 +1,30 @@
 # presto_pretrain_finetune, but in a notebook
 import argparse
+import gc
 import json
 import logging
+from glob import glob
 from pathlib import Path
 from typing import Optional, cast
 
 import pandas as pd
 import torch
 import xarray as xr
+from tqdm.auto import tqdm
 
 from presto.dataset import WorldCerealBase
 from presto.eval import WorldCerealEval
 from presto.presto import Presto
 from presto.utils import (
     DEFAULT_SEED,
+    NODATAVALUE,
     config_dir,
     data_dir,
     default_model_path,
     device,
     initialize_logging,
     plot_spatial,
+    process_parquet,
     seed_everything,
     timestamp_dirname,
 )
@@ -41,12 +46,20 @@
 argparser.add_argument("--num_workers", type=int, default=4)
 argparser.add_argument("--wandb", dest="wandb", action="store_true")
 argparser.add_argument("--wandb_org", type=str, default="nasa-harvest")
-argparser.add_argument("--parquet_file", type=str, default="rawts-monthly_calval.parquet")
+# argparser.add_argument("--parquet_file", type=str, default="rawts-monthly_calval.parquet")
+argparser.add_argument(
+    "--parquet_file",
+    type=str,
+    default="/vitodata/worldcereal/features/preprocessedinputs-monthly-nointerp/\
+worldcereal_training_data.parquet",
+)
 argparser.add_argument("--val_samples_file", type=str, default="cropland_test_split_samples.csv")
 argparser.add_argument("--train_only_samples_file", type=str, default="train_only_samples.csv")
 argparser.add_argument("--warm_start", dest="warm_start", action="store_true")
+argparser.add_argument("--augment", dest="augment", action="store_true")
 argparser.set_defaults(wandb=False)
 argparser.set_defaults(warm_start=True)
+argparser.set_defaults(augment=False)
 args = argparser.parse_args().__dict__
 
 model_name = args["model_name"]
@@ -79,6 +92,7 @@
 parquet_file: str = args["parquet_file"]
 val_samples_file: str = args["val_samples_file"]
 train_only_samples_file: str = args["train_only_samples_file"]
+augment: bool = args["augment"]
 
 dekadal = False
 if "10d" in parquet_file:
@@ -87,9 +101,18 @@
 path_to_config = config_dir / "default.json"
 model_kwargs = json.load(Path(path_to_config).open("r"))
 
-logger.info("Setting up dataloaders")
-
-df = pd.read_parquet(data_dir / parquet_file)
+logger.info("Reading dataset")
+files = sorted(glob(f"{parquet_file}/**/*.parquet"))[:10]
+df_list = []
+for f in tqdm(files):
+    _data = pd.read_parquet(f, engine="fastparquet")
+    _data_pivot = process_parquet(_data)
+    _data_pivot.reset_index(inplace=True)
+    df_list.append(_data_pivot)
+df = pd.concat(df_list)
+df = df.fillna(NODATAVALUE)
+del df_list
+gc.collect()
 
 logger.info("Setting up model")
 if warm_start:
@@ -104,13 +127,18 @@
     best_model_path = None
 model.to(device)
 
-model_modes = ["Random Forest", "Regression", "CatBoostClassifier"]
+# model_modes = ["Random Forest", "Regression", "CatBoostClassifier"]
+model_modes = ["CatBoostClassifier"]
 
 # 1. Using the provided split
 val_samples_df = pd.read_csv(data_dir / val_samples_file)
 train_df, test_df = WorldCerealBase.split_df(df, val_sample_ids=val_samples_df.sample_id.tolist())
 full_eval = WorldCerealEval(
-    train_df, test_df, spatial_inference_savedir=model_logging_dir, dekadal=dekadal
+    train_df,
+    test_df,
+    spatial_inference_savedir=model_logging_dir,
+    dekadal=dekadal,
+    augment=augment,
 )
 results, finetuned_model = full_eval.finetuning_results(model, sklearn_model_modes=model_modes)
 logger.info(json.dumps(results, indent=2))

presto/dataops.py

@@ -77,6 +77,8 @@
 NUM_TIMESTEPS = 12
 NUM_ORG_BANDS = len(BANDS)
 TIMESTEPS_IDX = list(range(NUM_TIMESTEPS))
+NODATAVALUE = 65535
+MIN_EDGE_BUFFER = 2  # Min amount of timesteps to include before/after the valid position
 
 NORMED_BANDS = [x for x in BANDS if x != "B9"]
 NUM_BANDS = len(NORMED_BANDS)

presto/dataset.py

@@ -16,6 +16,8 @@
 from .dataops import (
     BANDS,
     BANDS_GROUPS_IDX,
+    MIN_EDGE_BUFFER,
+    NODATAVALUE,
     NORMED_BANDS,
     S1_S2_ERA5_SRTM,
     DynamicWorld2020_2021,
@@ -32,7 +34,7 @@
 
 
 class WorldCerealBase(Dataset):
-    _NODATAVALUE = 65535
+    # _NODATAVALUE = 65535
     NUM_TIMESTEPS = 12
     BAND_MAPPING = {
         "OPTICAL-B02-ts{}-10m": "B2",
@@ -63,26 +65,97 @@ def target_crop(row_d: Dict) -> int:
         # by default, we predict crop vs non crop
         return int(row_d["LANDCOVER_LABEL"] == 11)
 
+    @classmethod
+    def get_timestep_positions(cls, row_d: Dict, augment: bool = False) -> List[int]:
+        available_timesteps = int(row_d["available_timesteps"])
+        valid_position = int(row_d["valid_position"])
+
+        if not augment:
+            #  check if the valid position is too close to the start_date and force shifting it
+            if valid_position < cls.NUM_TIMESTEPS // 2:
+                center_point = cls.NUM_TIMESTEPS // 2
+            #  or too close to the end_date
+            elif valid_position > (available_timesteps - cls.NUM_TIMESTEPS // 2):
+                center_point = available_timesteps - cls.NUM_TIMESTEPS // 2
+            else:
+                # Center the timesteps around the valid position
+                center_point = valid_position
+        else:
+            # Shift the center point but make sure the resulting range
+            # well includes the valid position
+
+            min_center_point = max(
+                cls.NUM_TIMESTEPS // 2,
+                valid_position + MIN_EDGE_BUFFER - cls.NUM_TIMESTEPS // 2,
+            )
+            max_center_point = min(
+                available_timesteps - cls.NUM_TIMESTEPS // 2,
+                valid_position - MIN_EDGE_BUFFER + cls.NUM_TIMESTEPS // 2,
+            )
+
+            center_point = np.random.randint(
+                min_center_point, max_center_point + 1
+            )  # max_center_point included
+
+        last_timestep = min(available_timesteps, center_point + cls.NUM_TIMESTEPS // 2)
+        first_timestep = max(0, last_timestep - cls.NUM_TIMESTEPS)
+        timestep_positions = list(range(first_timestep, last_timestep))
+
+        if len(timestep_positions) != cls.NUM_TIMESTEPS:
+            raise ValueError(
+                f"Acquired timestep positions do not have correct length: \
+required {cls.NUM_TIMESTEPS}, got {len(timestep_positions)}"
+            )
+        assert (
+            valid_position in timestep_positions
+        ), f"Valid position {valid_position} not in timestep positions {timestep_positions}"
+        return timestep_positions
+
     @classmethod
     def row_to_arrays(
-        cls, row: pd.Series, target_function: Callable[[Dict], int]
+        cls, row: pd.Series, target_function: Callable[[Dict], int], augment: bool = False
     ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, float, int]:
         # https://stackoverflow.com/questions/45783891/is-there-a-way-to-speed-up-the-pandas-getitem-getitem-axis-and-get-label
         # This is faster than indexing the series every time!
         row_d = pd.Series.to_dict(row)
 
         latlon = np.array([row_d["lat"], row_d["lon"]], dtype=np.float32)
-        month = datetime.strptime(row_d["start_date"], "%Y-%m-%d").month - 1
+
+        timestep_positions = cls.get_timestep_positions(row_d, augment=augment)
+
+        if cls.NUM_TIMESTEPS == 12:
+            initial_start_date_position = pd.to_datetime(row_d["start_date"]).month
+        elif cls.NUM_TIMESTEPS > 12:
+            # get the correct index of the start_date based on NUM_TIMESTEPS`
+            # e.g. if NUM_TIMESTEPS is 36 (dekadal setup), we should take the correct
+            # 10-day interval that the start_date falls into
+            # TODO: 1) this needs to go into a separate function
+            # 2) definition of valid_position and timestep_ind
+            #  should also be changed accordingly
+            year = pd.to_datetime(row_d["start_date"]).year
+            year_dates = pd.date_range(start=f"{year}-01-01", end=f"{year}-12-31")
+            bins = pd.cut(year_dates, bins=cls.NUM_TIMESTEPS, labels=False)
+            initial_start_date_position = bins[
+                np.where(year_dates == pd.to_datetime(row_d["start_date"]))[0][0]
+            ]
+        else:
+            raise ValueError(
+                f"NUM_TIMESTEPS must be at least 12. Currently it is {cls.NUM_TIMESTEPS}"
+            )
+
+        # make sure that month for encoding gets shifted according to
+        # the selected timestep positions. Also ensure circular indexing
+        month = (initial_start_date_position - 1 + timestep_positions[0]) % cls.NUM_TIMESTEPS
 
         eo_data = np.zeros((cls.NUM_TIMESTEPS, len(BANDS)))
         # an assumption we make here is that all timesteps for a token
         # have the same masking
         mask = np.zeros((cls.NUM_TIMESTEPS, len(BANDS_GROUPS_IDX)))
         for df_val, presto_val in cls.BAND_MAPPING.items():
-            values = np.array([float(row_d[df_val.format(t)]) for t in range(cls.NUM_TIMESTEPS)])
+            values = np.array([float(row_d[df_val.format(t)]) for t in timestep_positions])
             # this occurs for the DEM values in one point in Fiji
-            values = np.nan_to_num(values, nan=cls._NODATAVALUE)
-            idx_valid = values != cls._NODATAVALUE
+            values = np.nan_to_num(values, nan=NODATAVALUE)
+            idx_valid = values != NODATAVALUE
             if presto_val in ["VV", "VH"]:
                 # convert to dB
                 idx_valid = idx_valid & (values > 0)
@@ -97,8 +170,8 @@ def row_to_arrays(
             eo_data[:, BANDS.index(presto_val)] = values * idx_valid
         for df_val, presto_val in cls.STATIC_BAND_MAPPING.items():
             # this occurs for the DEM values in one point in Fiji
-            values = np.nan_to_num(row_d[df_val], nan=cls._NODATAVALUE)
-            idx_valid = values != cls._NODATAVALUE
+            values = np.nan_to_num(row_d[df_val], nan=NODATAVALUE)
+            idx_valid = values != NODATAVALUE
             eo_data[:, BANDS.index(presto_val)] = values * idx_valid
             mask[:, IDX_TO_BAND_GROUPS[presto_val]] += ~idx_valid
 
@@ -129,7 +202,7 @@ def normalize_and_mask(cls, eo: np.ndarray):
         keep_indices = [idx for idx, val in enumerate(BANDS) if val != "B9"]
         normed_eo = S1_S2_ERA5_SRTM.normalize(eo)
         # TODO: fix this. For now, we replicate the previous behaviour
-        normed_eo = np.where(eo[:, keep_indices] != cls._NODATAVALUE, normed_eo, 0)
+        normed_eo = np.where(eo[:, keep_indices] != NODATAVALUE, normed_eo, 0)
         return normed_eo
 
     @staticmethod
@@ -258,6 +331,7 @@ def __init__(
         years_to_remove: Optional[List[int]] = None,
         target_function: Optional[Callable[[Dict], int]] = None,
         balance: bool = False,
+        augment: bool = False,
         mask_ratio: float = 0.0,
     ):
         dataframe = dataframe.loc[~dataframe.LANDCOVER_LABEL.isin(self.FILTER_LABELS)]
@@ -274,6 +348,9 @@ def __init__(
             dataframe = dataframe[(~dataframe.end_date.dt.year.isin(years_to_remove))]
         self.target_function = target_function if target_function is not None else self.target_crop
         self._class_weights: Optional[np.ndarray] = None
+        self.augment = augment
+        if augment:
+            logger.info("Augmentation is enabled. The valid_date position will be shifted.")
         self.mask_ratio = mask_ratio
         self.mask_params = MaskParamsNoDw(
             (
@@ -367,7 +444,9 @@ def __getitem__(self, idx):
         # Get the sample
         df_index = self.indices[idx]
         row = self.df.iloc[df_index, :]
-        eo, mask_per_token, latlon, _, target = self.row_to_arrays(row, self.target_function)
+        eo, mask_per_token, latlon, _, target = self.row_to_arrays(
+            row, self.target_function, self.augment
+        )
         if self.mask_ratio > 0:
             mask_per_token, eo, _, _ = self.mask_params.mask_data(eo, mask_per_token)
         mask_per_variable = np.repeat(mask_per_token, BAND_EXPANSION, axis=1)
@@ -382,7 +461,24 @@ def __getitem__(self, idx):
 
 
 class WorldCerealInferenceDataset(Dataset):
-    _NODATAVALUE = 65535
+    # _NODATAVALUE = 65535
+    Y = "worldcereal_cropland"
+    BAND_MAPPING = {
+        "B02": "B2",
+        "B03": "B3",
+        "B04": "B4",
+        "B05": "B5",
+        "B06": "B6",
+        "B07": "B7",
+        "B08": "B8",
+        # B8A is missing
+        "B11": "B11",
+        "B12": "B12",
+        "VH": "VH",
+        "VV": "VV",
+        "precipitation-flux": "total_precipitation",
+        "temperature-mean": "temperature_2m",
+    }
     Y = "WORLDCEREAL_TEMPORARYCROPS_2021"
 
     def __init__(self, path_to_files: Path = data_dir / "inference_areas"):
@@ -408,7 +504,7 @@ def _extract_eo_data(cls, inarr: xr.DataArray) -> Tuple[np.ndarray, np.ndarray]:
 
         # Handle NaN values in Presto compatible way
         inarr = inarr.astype(np.float32)
-        inarr = inarr.fillna(65535)
+        inarr = inarr.fillna(NODATAVALUE)
 
         eo_data = np.zeros((num_pixels, num_timesteps, len(BANDS)))
         mask = np.zeros((num_pixels, num_timesteps, len(BANDS_GROUPS_IDX)))
@@ -420,7 +516,7 @@ def _extract_eo_data(cls, inarr: xr.DataArray) -> Tuple[np.ndarray, np.ndarray]:
                     0,
                     1,
                 )
-                idx_valid = values != cls._NODATAVALUE
+                idx_valid = values != NODATAVALUE
                 values = cls._preprocess_band_values(values, presto_band)
                 eo_data[:, :, BANDS.index(presto_band)] = values * idx_valid
                 mask[:, :, IDX_TO_BAND_GROUPS[presto_band]] += ~idx_valid
@@ -552,7 +648,7 @@ def nc_to_arrays(
         months = cls._extract_months(inarr)
 
         if cls.Y not in ds:
-            target = np.ones_like(months) * cls._NODATAVALUE
+            target = np.ones_like(months) * NODATAVALUE
         else:
             target = rearrange(inarr.sel(bands=cls.Y).values, "t x y -> (x y) t")