added code for the first time

data_formatter/base.py

@@ -0,0 +1,213 @@
+'''Defines a generic data formatter for CGM data sets.'''
+import sys
+import warnings
+import numpy as np
+import pandas as pd
+import sklearn.preprocessing
+import data_formatter.types as types
+import data_formatter.utils as utils
+
+DataTypes = types.DataTypes
+InputTypes = types.InputTypes
+
+dict_data_type = {'categorical': DataTypes.CATEGORICAL,
+                  'real_valued': DataTypes.REAL_VALUED,
+                  'date': DataTypes.DATE}
+dict_input_type = {'target': InputTypes.TARGET,
+                   'observed_input': InputTypes.OBSERVED_INPUT,      
+                   'known_input': InputTypes.KNOWN_INPUT,
+                   'static_input': InputTypes.STATIC_INPUT,
+                   'id': InputTypes.ID,
+                   'time': InputTypes.TIME}
+
+
+class DataFormatter:
+  # Defines and formats data.
+
+  def __init__(self, cnf):
+    """Initialises formatter."""
+    # load parameters from the config file
+    self.params = cnf
+    # write progress to file if specified
+
+    # load column definition
+    print('-'*32)
+    print('Loading column definition...')
+    self.__process_column_definition()
+
+    # check that column definition is valid
+    print('Checking column definition...')
+    self.__check_column_definition()
+
+    # load data
+    # check if data table has index col: -1 if not, index >= 0 if yes
+    print('Loading data...')
+    self.params['index_col'] = False if self.params['index_col'] == -1 else self.params['index_col']
+    # read data table
+    self.data = pd.read_csv(self.params['data_csv_path'], index_col=self.params['index_col'])
+
+    # drop columns / rows
+    print('Dropping columns / rows...')
+    self.__drop()
+
+    # check NA values
+    print('Checking for NA values...')
+    self.__check_nan()
+
+    # set data types in DataFrame to match column definition
+    print('Setting data types...')
+    self.__set_data_types()
+
+    # drop columns / rows
+    print('Dropping columns / rows...')
+    self.__drop()
+
+    # encode
+    print('Encoding data...')
+    self._encoding_params = self.params['encoding_params']
+    self.__encode()
+
+    # interpolate
+    print('Interpolating data...')
+    self._interpolation_params = self.params['interpolation_params']
+    self._interpolation_params['interval_length'] = self.params['observation_interval']
+    self.__interpolate()
+    
+    # split data
+    print('Splitting data...')
+    self._split_params = self.params['split_params']
+    self._split_params['max_length_input'] = self.params['max_length_input']
+    self.__split_data()
+
+    # scale
+    print('Scaling data...')
+    self._scaling_params = self.params['scaling_params']
+    self.__scale()
+
+    print('Data formatting complete.')
+    print('-'*32)
+
+
+  def __process_column_definition(self):
+    self._column_definition = []
+    for col in self.params['column_definition']:
+      self._column_definition.append((col['name'], 
+                                      dict_data_type[col['data_type']], 
+                                      dict_input_type[col['input_type']]))
+
+  def __check_column_definition(self):
+    # check that there is unique ID column
+    assert len([col for col in self._column_definition if col[2] == InputTypes.ID]) == 1, 'There must be exactly one ID column.'
+    # check that there is unique time column
+    assert len([col for col in self._column_definition if col[2] == InputTypes.TIME]) == 1, 'There must be exactly one time column.'
+    # check that there is at least one target column
+    assert len([col for col in self._column_definition if col[2] == InputTypes.TARGET]) >= 1, 'There must be at least one target column.'
+  
+  def __set_data_types(self):
+    # set time column as datetime format in pandas
+    for col in self._column_definition:
+      if col[1] == DataTypes.DATE:
+        self.data[col[0]] = pd.to_datetime(self.data[col[0]])
+      if col[1] == DataTypes.CATEGORICAL:
+        self.data[col[0]] = self.data[col[0]].astype('category')
+      if col[1] == DataTypes.REAL_VALUED:
+        self.data[col[0]] = self.data[col[0]].astype(np.float32)
+
+  def __check_nan(self):
+    # delete rows where target, time, or id are na
+    self.data = self.data.dropna(subset=[col[0] 
+                                  for col in self._column_definition 
+                                  if col[2] in [InputTypes.TARGET, InputTypes.TIME, InputTypes.ID]])
+    # assert that there are no na values in the data
+    assert self.data.isna().sum().sum() == 0, 'There are NA values in the data even after dropping with missing time, glucose, or id.'
+
+  def __drop(self):
+    # drop columns that are not in the column definition
+    self.data = self.data[[col[0] for col in self._column_definition]]
+    # drop rows based on conditions set in the formatter
+    if self.params['drop'] is not None:
+      if self.params['drop']['rows'] is not None:
+        # drop row at indices in the list self.params['drop']['rows']
+        self.data = self.data.drop(self.params['drop']['rows'])
+        self.data = self.data.reset_index(drop=True)
+      if self.params['drop']['columns'] is not None:
+        for col in self.params['drop']['columns'].keys():
+          # drop rows where specified columns have values in the list self.params['drop']['columns'][col]
+          self.data = self.data.loc[~self.data[col].isin(self.params['drop']['columns'][col])].copy()
+  
+  def __interpolate(self):
+    self.data, self._column_definition = utils.interpolate(self.data, 
+                                                           self._column_definition, 
+                                                           **self._interpolation_params)
+
+  def __split_data(self):
+    if self.params['split_params']['test_percent_subjects'] == 0 or \
+        self.params['split_params']['length_segment'] == 0:
+      print('\tNo splitting performed since test_percent_subjects or length_segment is 0.')
+      self.train_idx, self.val_idx, self.test_idx, self.test_idx_ood = None, None, None, None
+      self.train_data, self.val_data, self.test_data = self.data, None, None
+    else:
+      assert self.params['split_params']['length_segment'] > self.params['length_pred'], \
+        'length_segment for test / val must be greater than length_pred.'
+      self.train_idx, self.val_idx, self.test_idx, self.test_idx_ood = utils.split(self.data, 
+                                                                                  self._column_definition, 
+                                                                                  **self._split_params)
+      self.train_data, self.val_data, self.test_data = self.data.iloc[self.train_idx], \
+                                                        self.data.iloc[self.val_idx], \
+                                                          self.data.iloc[self.test_idx + self.test_idx_ood]
+
+  def __encode(self):
+    self.data, self._column_definition, self.encoders = utils.encode(self.data, 
+                                                                     self._column_definition,
+                                                                     **self._encoding_params)
+  
+  def __scale(self):
+    self.train_data, self.val_data, self.test_data, self.scalers = utils.scale(self.train_data, 
+                                                                               self.val_data, 
+                                                                               self.test_data, 
+                                                                               self._column_definition, 
+                                                                               **self.params['scaling_params'])
+
+  def reshuffle(self, seed):
+    stdout = sys.stdout
+    f = open(self.study_file, 'a')
+    sys.stdout = f
+    self.params['split_params']['random_state'] = seed
+    # split data
+    self.train_idx, self.val_idx, self.test_idx, self.test_idx_ood = utils.split(self.data, 
+                                                                                  self._column_definition, 
+                                                                                  **self._split_params)
+    self.train_data, self.val_data, self.test_data = self.data.iloc[self.train_idx], \
+                                                      self.data.iloc[self.val_idx], \
+                                                        self.data.iloc[self.test_idx+self.test_idx_ood]
+    # re-scale data
+    self.train_data, self.val_data, self.test_data, self.scalers = utils.scale(self.train_data, 
+                                                                               self.val_data, 
+                                                                               self.test_data, 
+                                                                               self._column_definition, 
+                                                                               **self.params['scaling_params'])
+    sys.stdout = stdout
+    f.close()
+    
+  def get_column(self, column_name):
+    # write cases for time, id, target, future, static, dynamic covariates
+    if column_name == 'time':
+      return [col[0] for col in self._column_definition if col[2] == InputTypes.TIME][0]
+    elif column_name == 'id':
+      return [col[0] for col in self._column_definition if col[2] == InputTypes.ID][0]
+    elif column_name == 'sid':
+      return [col[0] for col in self._column_definition if col[2] == InputTypes.SID][0]
+    elif column_name == 'target':
+      return [col[0] for col in self._column_definition if col[2] == InputTypes.TARGET]
+    elif column_name == 'future_covs':
+      future_covs = [col[0] for col in self._column_definition if col[2] == InputTypes.KNOWN_INPUT] 
+      return future_covs if len(future_covs) > 0 else None
+    elif column_name == 'static_covs':
+      static_covs = [col[0] for col in self._column_definition if col[2] == InputTypes.STATIC_INPUT]
+      return static_covs if len(static_covs) > 0 else None
+    elif column_name == 'dynamic_covs':
+      dynamic_covs = [col[0] for col in self._column_definition if col[2] == InputTypes.OBSERVED_INPUT]
+      return dynamic_covs if len(dynamic_covs) > 0 else None
+    else:
+      raise ValueError('Column {} not found.'.format(column_name))
+  

data_formatter/types.py

@@ -0,0 +1,19 @@
+'''Defines data and input types of each column in the dataset.'''
+
+import enum
+
+class DataTypes(enum.IntEnum):
+  """Defines numerical types of each column."""
+  REAL_VALUED = 0
+  CATEGORICAL = 1
+  DATE = 2
+
+class InputTypes(enum.IntEnum):
+  """Defines input types of each column."""
+  TARGET = 0
+  OBSERVED_INPUT = 1
+  KNOWN_INPUT = 2
+  STATIC_INPUT = 3
+  ID = 4  # Single column used as an entity identifier
+  SID = 5  # Single column used as a segment identifier
+  TIME = 6  # Single column exclusively used as a time index

data_formatter/utils.py

@@ -0,0 +1,323 @@
+# coding=utf-8
+# Copyright 2019 The Google Research Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Lint as: python3
+"""Generic helper functions used across codebase."""
+import warnings
+from collections import namedtuple
+from datetime import datetime
+import os
+import math
+import pathlib
+import torch
+import numpy as np
+import pandas as pd
+pd.options.mode.chained_assignment = None
+from typing import List, Tuple
+from sklearn import preprocessing
+
+import data_formatter
+from data_formatter import types
+
+DataTypes = types.DataTypes
+InputTypes = types.InputTypes
+MINUTE = 60
+
+# OS related functions.
+def create_folder_if_not_exist(directory):
+  """Creates folder if it doesn't exist.
+
+  Args:
+    directory: Folder path to create.
+  """
+  # Also creates directories recursively
+  pathlib.Path(directory).mkdir(parents=True, exist_ok=True)
+
+
+def csv_path_to_folder(path: str):
+    return "/".join(path.split('/')[:-1]) + "/"
+
+
+def interpolate(data: pd.DataFrame, 
+                column_definition: List[Tuple[str, DataTypes, InputTypes]],
+                gap_threshold: int = 0, 
+                min_drop_length: int = 0, 
+                interval_length: int = 0):
+  """Interpolates missing values in data.
+
+  Args:
+    df: Dataframe to interpolate on. Sorted by id and then time (a DateTime object).
+    column_definition: List of tuples describing columns (column_name, data_type, input_type).
+    gap_threshold: Number in minutes, maximum allowed gap for interpolation.
+    min_drop_length: Number of points, minimum number within an interval to interpolate.
+    interval_length: Number in minutes, length of interpolation.
+
+  Returns:
+    data: DataFrame with missing values interpolated and 
+          additional column ('segment') indicating continuous segments.
+    column_definition: Updataed list of tuples (column_name, data_type, input_type).
+  """
+  # select all real-valued columns that are not id, time, or static
+  interpolation_columns = [column_name for column_name, data_type, input_type in column_definition if 
+    data_type == DataTypes.REAL_VALUED and 
+    input_type not in set([InputTypes.ID, InputTypes.TIME, InputTypes.STATIC_INPUT])]
+  # select all other columns except time
+  constant_columns = [column_name for column_name, data_type, input_type in column_definition if
+    input_type not in set([InputTypes.TIME])]
+  constant_columns += ['id_segment']
+
+  # get id and time columns
+  id_col = [column_name for column_name, data_type, input_type in column_definition if input_type == InputTypes.ID][0]
+  time_col = [column_name for column_name, data_type, input_type in column_definition if input_type == InputTypes.TIME][0]
+  
+  # round to minute
+  data[time_col] = data[time_col].dt.round('1min')
+  # count dropped segments
+  dropped_segments = 0
+  # count number of values that are interpolated
+  interpolation_count = 0
+  # store final output
+  output = []
+  for id, id_data in data.groupby(id_col):
+    # sort values 
+    id_data.sort_values(time_col, inplace=True)
+    # get time difference between consecutive rows
+    lag = (id_data[time_col].diff().dt.total_seconds().fillna(0) / 60.0).astype(int)
+    # if lag > gap_threshold
+    id_segment = (lag > gap_threshold).cumsum()
+    id_data['id_segment'] = id_segment
+    for segment, segment_data in id_data.groupby('id_segment'):
+      # if segment is too short, then we don't interpolate
+      if len(segment_data) < min_drop_length:
+        dropped_segments += 1
+        continue
+      
+      # find and print duplicated times
+      duplicates = segment_data.duplicated(time_col, keep=False)
+      if duplicates.any():
+        print(segment_data[duplicates])
+        raise ValueError('Duplicate times in segment {} of id {}'.format(segment, id))
+
+      # reindex at interval_length minute intervals
+      segment_data = segment_data.set_index(time_col)
+      index_new = pd.date_range(start = segment_data.index[0], 
+                            end = segment_data.index[-1], 
+                            freq = interval_length)
+      index_union = index_new.union(segment_data.index)
+      segment_data = segment_data.reindex(index_union)
+      # count nan values in interpolation columns
+      interpolation_count += segment_data[interpolation_columns[0]].isna().sum()
+      # interpolate
+      segment_data[interpolation_columns] = segment_data[interpolation_columns].interpolate(method='index')
+      # fill constant columns with last value
+      segment_data[constant_columns] = segment_data[constant_columns].ffill()
+      # delete rows not conforming to frequency
+      segment_data = segment_data.reindex(index_new)
+      # reset index, make the time a column with name time_col
+      segment_data = segment_data.reset_index().rename(columns={'index': time_col})
+      # set the id_segment to position in output
+      segment_data['id_segment'] = len(output)
+      # add to output
+      output.append(segment_data)
+  # print number of dropped segments and number of segments
+  print('\tDropped segments: {}'.format(dropped_segments))
+  print('\tExtracted segments: {}'.format(len(output)))
+  # concat all segments and reset index
+  output = pd.concat(output)
+  output.reset_index(drop=True, inplace=True)
+  # count number of interpolated values
+  print('\tInterpolated values: {}'.format(interpolation_count))
+  print('\tPercent of values interpolated: {:.2f}%'.format(interpolation_count / len(output) * 100))
+  # add id_segment column to column_definition as ID
+  column_definition += [('id_segment', DataTypes.CATEGORICAL, InputTypes.SID)]
+
+  return output, column_definition
+
+def create_index(time_col: pd.Series, interval_length: int):
+  """Creates a new index at interval_length minute intervals.
+
+  Args:
+    time_col: Series of times.
+    interval_length: Number in minutes, length of interpolation.
+
+  Returns:
+    index: New index.
+  """
+  # margin of error
+  eps = pd.Timedelta('1min')
+  new_time_col = [time_col.iloc[0]]
+  for time in time_col.iloc[1:]:
+    if time - new_time_col[-1] <= pd.Timedelta(interval_length) + eps:
+      new_time_col.append(time)
+    else:
+      filler = new_time_col[-1] + pd.Timedelta(interval_length)
+      while filler < time:
+        new_time_col.append(filler)
+        filler += pd.Timedelta(interval_length)
+      new_time_col.append(time)
+  return pd.to_datetime(new_time_col)
+
+def split(df: pd.DataFrame, 
+          column_definition: List[Tuple[str, DataTypes, InputTypes]],
+          test_percent_subjects: float, 
+          length_segment: int,
+          max_length_input: int,
+          random_state: int = 42):
+  """Splits data into train, validation and test sets.
+
+  Args: 
+    df: Dataframe to split.
+    column_definition: List of tuples describing columns (column_name, data_type, input_type).
+    test_percent_subjects: Float number from [0, 1], percentage of subjects to use for test set.
+    length_segment: Number of points, length of segments saved for validation / test sets.
+    max_length_input: Number of points, maximum length of input sequences for models.
+    random_state: Number, Random state for reproducibility.
+
+  Returns:
+    train_idx: Training set indices.
+    val_idx: Validation set indices.
+    test_idx: Test set indices.
+  """
+  # set random state
+  np.random.seed(random_state)
+  # get id and id_segment columns
+  id_col = [column_name for column_name, data_type, input_type in column_definition if input_type == InputTypes.ID][0]
+  id_segment_col = [column_name for column_name, data_type, input_type in column_definition if input_type == InputTypes.SID][0]
+  # get unique ids
+  ids = df[id_col].unique()
+
+  # select some subjects for test data set
+  test_ids = np.random.choice(ids, math.ceil(len(ids) * test_percent_subjects), replace=False)
+  test_idx_ood = list(df[df[id_col].isin(test_ids)].index)
+  # get the remaning data for training and validation
+  df = df[~df[id_col].isin(test_ids)]
+
+  # iterate through subjects and split into train, val and test
+  train_idx = []; val_idx = []; test_idx = []
+  for id, id_data in df.groupby(id_col):
+    segment_ids = id_data[id_segment_col].unique()
+    if len(segment_ids) >= 2:
+      train_idx += list(id_data.loc[id_data[id_segment_col].isin(segment_ids[:-2])].index)
+      penultimate_segment = id_data[id_data[id_segment_col] == segment_ids[-2]]
+      last_segment = id_data[id_data[id_segment_col] == segment_ids[-1]]
+      if len(last_segment) >= max_length_input + 3 * length_segment:
+        train_idx += list(penultimate_segment.index)
+        train_idx += list(last_segment.iloc[:-2*length_segment].index)
+        val_idx += list(last_segment.iloc[-2*length_segment-max_length_input:-length_segment].index)
+        test_idx += list(last_segment.iloc[-length_segment-max_length_input:].index)
+      elif len(last_segment) >= max_length_input + 2 * length_segment:
+        train_idx += list(penultimate_segment.index)
+        val_idx += list(last_segment.iloc[:-length_segment].index)
+        test_idx += list(last_segment.iloc[-length_segment-max_length_input:].index)
+      else:
+        test_idx += list(last_segment.index)
+        if len(penultimate_segment) >= max_length_input + 2 * length_segment:
+          val_idx += list(penultimate_segment.iloc[-length_segment-max_length_input:].index)
+          train_idx += list(penultimate_segment.iloc[:-length_segment].index)
+        else:
+          train_idx += list(penultimate_segment.index)
+    else:
+      if len(id_data) >= max_length_input + 3 * length_segment:
+        train_idx += list(id_data.iloc[:-2*length_segment].index)
+        val_idx += list(id_data.iloc[-2*length_segment-max_length_input:-length_segment].index)
+        test_idx += list(id_data.iloc[-length_segment-max_length_input:].index)
+      elif len(id_data) >= max_length_input + 2 * length_segment:
+        train_idx += list(id_data.iloc[:-length_segment].index)
+        test_idx += list(id_data.iloc[-length_segment-max_length_input:].index)
+      else:
+        train_idx += list(id_data.index)
+  total_len = len(train_idx) + len(val_idx) + len(test_idx) + len(test_idx_ood)
+  print('\tTrain: {} ({:.2f}%)'.format(len(train_idx), len(train_idx) / total_len * 100))
+  print('\tVal: {} ({:.2f}%)'.format(len(val_idx), len(val_idx) / total_len * 100))
+  print('\tTest: {} ({:.2f}%)'.format(len(test_idx), len(test_idx) / total_len * 100))
+  print('\tTest OOD: {} ({:.2f}%)'.format(len(test_idx_ood), len(test_idx_ood) / total_len * 100))
+  return train_idx, val_idx, test_idx, test_idx_ood
+
+def encode(df: pd.DataFrame, 
+          column_definition: List[Tuple[str, DataTypes, InputTypes]],
+          date: List,):
+  """Encodes categorical columns.
+
+  Args: 
+    df: Dataframe to split.
+    column_definition: List of tuples describing columns (column_name, data_type, input_type).
+    date: List of str, list containing date info to extract.
+
+  Returns:
+    df: Dataframe with encoded columns.
+    column_definition: Updated list of tuples containing column name and types.
+    encoders: dictionary containing encoders.
+  """ 
+  encoders = {}
+  new_columns = []
+  for i in range(len(column_definition)):
+    column, column_type, input_type = column_definition[i]
+    if column_type == DataTypes.DATE:
+      for extract_col in date:
+        df[column + '_' + extract_col] = getattr(df[column].dt, extract_col)
+        df[column + '_' + extract_col] = df[column + '_' + extract_col].astype(np.float32)
+        new_columns.append((column + '_' + extract_col, DataTypes.REAL_VALUED, InputTypes.KNOWN_INPUT))
+    elif column_type == DataTypes.CATEGORICAL:
+      encoders[column] = preprocessing.LabelEncoder()
+      df[column] = encoders[column].fit_transform(df[column]).astype(np.float32)
+      column_definition[i] = (column, DataTypes.REAL_VALUED, input_type)
+    else:
+      continue
+  column_definition += new_columns
+  # print updated column definition
+  print('\tUpdated column definition:')
+  for column, column_type, input_type in column_definition:
+    print('\t\t{}: {} ({})'.format(column, 
+                                   DataTypes(column_type).name, 
+                                   InputTypes(input_type).name))
+  return df, column_definition, encoders
+  
+def scale(train_data: pd.DataFrame,
+          val_data: pd.DataFrame,
+          test_data: pd.DataFrame,
+          column_definition: List[Tuple[str, DataTypes, InputTypes]],
+          scaler: str):
+  """Scales numerical data.
+
+  Args:
+    train_data: pd.Dataframe, DataFrame of training data.
+    val_data: pd.Dataframe, DataFrame of validation data.
+    test_data: pd.Dataframe, DataFrame of testing data.
+    column_definition: List of tuples describing columns (column_name, data_type, input_type).
+    scaler: String, scaler to use.
+  
+  Returns:
+    train_data: pd.Dataframe, DataFrame of scaled training data.
+    val_data: pd.Dataframe, DataFrame of scaled validation data.
+    test_data: pd.Dataframe, DataFrame of scaled testing data.
+    scalers: dictionary index by column names containing scalers.
+  """
+  # select all real-valued columns
+  columns_to_scale = [column for column, data_type, input_type in column_definition if data_type == DataTypes.REAL_VALUED]
+  # handle no scaling case
+  if scaler == 'None':
+    print('\tNo scaling applied')
+    return train_data, val_data, test_data, None
+  scalers = {}
+  for column in columns_to_scale:
+    scaler_column = getattr(preprocessing, scaler)()
+    train_data[column] = scaler_column.fit_transform(train_data[column].values.reshape(-1, 1))
+    # handle empty validation and test sets
+    val_data[column] = scaler_column.transform(val_data[column].values.reshape(-1, 1)) if val_data.shape[0] > 0 else val_data[column]
+    test_data[column] = scaler_column.transform(test_data[column].values.reshape(-1, 1)) if test_data.shape[0] > 0 else test_data[column]
+    scalers[column] = scaler_column
+  # print columns that were scaled
+  print('\tScaled columns: {}'.format(columns_to_scale))
+  return train_data, val_data, test_data, scalers

environment.yaml

@@ -0,0 +1,28 @@
+name: glucose_genie
+channels:
+  - conda-forge
+  - defaults
+dependencies:
+  - python=3.11
+  - gradio
+  - seaborn
+  - pytorch
+  - optuna
+  - numpy<2.0.0
+  - tensorboard
+  - pip:
+    - fastapi
+    - uvicorn
+    - thefuzz
+    - pycomfort>=0.0.15
+    - polars>=1.3.0
+    - hybrid_search>=0.0.15
+    - psutil #compartibility
+    - httpx
+    - just-agents>=0.1.0
+    - FlagEmbedding
+    - typer
+    - darts==0.29.0
+    - pmdarima==2.0.4
+    - numpy==1.26.4
+    - peft

files/config.yaml

@@ -0,0 +1,81 @@
+data_csv_path: ./raw_data/anton.csv
+drop: null
+ds_name: livia_mini
+index_col: -1
+observation_interval: 5min
+
+column_definition:
+  - data_type: categorical
+    input_type: id
+    name: id
+  - data_type: date
+    input_type: time
+    name: time
+  - data_type: real_valued
+    input_type: target
+    name: gl
+
+encoding_params:
+  date:
+    - day
+    - month
+    - year
+    - hour
+    - minute
+    - second
+
+# NA values abbreviation
+nan_vals: null
+
+# Interpolation parameters
+interpolation_params:
+  gap_threshold: 45 # in minutes
+  min_drop_length: 240 # in number of points (20 hrs)
+
+scaling_params:
+  scaler: None
+
+split_params:
+  length_segment: 13
+  random_state: 0
+  test_percent_subjects: 0.1
+
+
+# Splitting parameters
+#split_params:
+#  test_percent_subjects: .1
+#  length_segment: 240
+#  random_state: 0
+
+# Model params
+max_length_input: 192
+length_pred: 12
+
+transformer:
+  batch_size: 32
+  d_model: 96
+  dim_feedforward: 448
+  dropout: 0.10161152207464333
+  in_len: 96
+  lr: 0.000840888489686657
+  lr_epochs: 16
+  max_grad_norm: 0.6740479322943925
+  max_samples_per_ts: 50
+  n_heads: 4
+  num_decoder_layers: 1
+  num_encoder_layers: 4
+
+transformer_covariates:
+  batch_size: 32
+  d_model: 128
+  dim_feedforward: 160
+  dropout: 0.044926981080245884
+  in_len: 108
+  lr: 0.00029632347559614453
+  lr_epochs: 20
+  max_grad_norm: 0.8890169619043728
+  max_samples_per_ts: 50
+  n_heads: 2
+  num_decoder_layers: 2
+  num_encoder_layers: 2
+

format_dexcom.py

@@ -0,0 +1,152 @@
+import pandas as pd
+from pathlib import Path
+import typer
+
+
+def process_csv(
+        input_dir: Path = typer.Argument( help="Directory containing the input CSV files."),
+        output_file: Path = typer.Argument( help="Path to save the processed CSV file."),
+        event_type_filter: str = typer.Option('egv', help="Event type to filter by."),
+        drop_duplicates: bool = typer.Option(True, help="Whether to drop duplicate timestamps."),
+        time_diff_minutes: int = typer.Option(1, help="Minimum time difference in minutes to keep a row."),
+        chunk_size: int = typer.Option(1000, help="Chunk size for the 'id' column increment. Set to 0 or None for a single id."),
+) -> pd.DataFrame:
+
+    # Read CSV file into a DataFrame
+    filename=input_dir
+    df = pd.read_csv(filename, low_memory=False)
+
+
+    # Filter by Event Type and Event Subtype
+    df = df[df['Event Type'].str.lower() == event_type_filter]
+    df = df[df['Event Subtype'].isna()]
+
+    # List of columns to keep
+    columns_to_keep = [
+        'Index',
+        'Timestamp (YYYY-MM-DDThh:mm:ss)',
+        'Glucose Value (mg/dL)',
+    ]
+
+    # Keep only the specified columns
+    df = df[columns_to_keep]
+
+    # Rename columns
+    column_rename = {
+        'Index': 'id',
+        'Timestamp (YYYY-MM-DDThh:mm:ss)': 'time',
+        'Glucose Value (mg/dL)': 'gl'
+    }
+    df = df.rename(columns=column_rename)
+
+    
+    # Handle id assignment based on chunk_size
+    if chunk_size is None or chunk_size == 0:
+        df['id'] = 1  # Assign the same id to all rows
+    else:
+        df['id'] = ((df.index // chunk_size) % (df.index.max() // chunk_size + 1)).astype(int)
+
+    # Convert timestamp to datetime
+    df['time'] = pd.to_datetime(df['time'])
+
+    # Calculate time difference and keep rows with at least the specified time difference
+    df['time_diff'] = df['time'].diff()
+    df = df[df['time_diff'].isna() | (df['time_diff'] >= pd.Timedelta(minutes=time_diff_minutes))]
+
+    # Drop the temporary time_diff column
+    df = df.drop(columns=['time_diff'])
+
+    # Ensure glucose values are in float64
+    df['gl'] = df['gl'].astype('float64')
+
+    # Optionally drop duplicate rows based on time
+    if drop_duplicates:
+        df = df.drop_duplicates(subset=['time'], keep='first')
+
+    # Write the modified dataframe to a new CSV file
+    df.to_csv(output_file, index=False)
+
+    typer.echo("CSV files have been successfully merged, modified, and saved.")
+
+    return df
+
+
+
+
+def process_multiple_csv(
+        input_dir: Path = typer.Argument('./raw_data/livia_unmerged', help="Directory containing the input CSV files."),
+        output_file: Path = typer.Argument('./raw_data/livia_unmerged/livia_mini.csv', help="Path to save the processed CSV file."),
+        event_type_filter: str = typer.Option('egv', help="Event type to filter by."),
+        drop_duplicates: bool = typer.Option(True, help="Whether to drop duplicate timestamps."),
+        time_diff_minutes: int = typer.Option(1, help="Minimum time difference in minutes to keep a row."),
+        chunk_size: int = typer.Option(1000, help="Chunk size for the 'id' column increment. Set to 0 or None for a single id."),
+):
+    # Get all the CSV files in the specified directory
+    all_files = list(input_dir.glob("*.csv"))
+
+    # List to store the DataFrames
+    df_list = []
+
+    # Read each CSV file into a DataFrame and append to the list
+    for filename in all_files:
+        df = pd.read_csv(filename, low_memory=False)
+        df_list.append(df)
+
+    # Concatenate all DataFrames in the list
+    combined_df = pd.concat(df_list, ignore_index=True)
+
+    # Filter by Event Type and Event Subtype
+    combined_df = combined_df[combined_df['Event Type'].str.lower() == event_type_filter]
+    combined_df = combined_df[combined_df['Event Subtype'].isna()]
+
+    # List of columns to keep
+    columns_to_keep = [
+        'Index',
+        'Timestamp (YYYY-MM-DDThh:mm:ss)',
+        'Glucose Value (mg/dL)',
+    ]
+
+    # Keep only the specified columns
+    combined_df = combined_df[columns_to_keep]
+
+    # Rename columns
+    column_rename = {
+        'Index': 'id',
+        'Timestamp (YYYY-MM-DDThh:mm:ss)': 'time',
+        'Glucose Value (mg/dL)': 'gl'
+    }
+    combined_df = combined_df.rename(columns=column_rename)
+
+    # Sort the combined DataFrame by timestamp
+    combined_df = combined_df.sort_values('time')
+
+    # Handle id assignment based on chunk_size
+    if chunk_size is None or chunk_size == 0:
+        combined_df['id'] = 1  # Assign the same id to all rows
+    else:
+        combined_df['id'] = ((combined_df.index // chunk_size) % (combined_df.index.max() // chunk_size + 1)).astype(int)
+
+    # Convert timestamp to datetime
+    combined_df['time'] = pd.to_datetime(combined_df['time'])
+
+    # Calculate time difference and keep rows with at least the specified time difference
+    combined_df['time_diff'] = combined_df['time'].diff()
+    combined_df = combined_df[combined_df['time_diff'].isna() | (combined_df['time_diff'] >= pd.Timedelta(minutes=time_diff_minutes))]
+
+    # Drop the temporary time_diff column
+    combined_df = combined_df.drop(columns=['time_diff'])
+
+    # Ensure glucose values are in float64
+    combined_df['gl'] = combined_df['gl'].astype('float64')
+
+    # Optionally drop duplicate rows based on time
+    if drop_duplicates:
+        combined_df = combined_df.drop_duplicates(subset=['time'], keep='first')
+
+    # Write the modified dataframe to a new CSV file
+    combined_df.to_csv(output_file, index=False)
+
+    typer.echo("CSV files have been successfully merged, modified, and saved.")
+
+if __name__ == "__main__":
+    typer.run(process_csv)

gluformer/attention.py

@@ -0,0 +1,70 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from math import sqrt
+
+class CausalConv1d(torch.nn.Conv1d):
+  def __init__(self,
+                in_channels,
+                out_channels,
+                kernel_size,
+                stride=1,
+                dilation=1,
+                groups=1,
+                bias=True):
+    self.__padding = (kernel_size - 1) * dilation
+
+    super(CausalConv1d, self).__init__(
+        in_channels,
+        out_channels,
+        kernel_size=kernel_size,
+        stride=stride,
+        padding=self.__padding,
+        dilation=dilation,
+        groups=groups,
+        bias=bias)
+
+  def forward(self, input):
+    result = super(CausalConv1d, self).forward(input)
+    if self.__padding != 0:
+        return result[:, :, :-self.__padding]
+    return result
+
+class TriangularCausalMask():
+    def __init__(self, b, n, device="cpu"):
+        mask_shape = [b, 1, n, n]
+        with torch.no_grad():
+            self._mask = torch.triu(torch.ones(mask_shape, dtype=torch.bool), diagonal=1).to(device)
+
+    @property
+    def mask(self):
+        return self._mask
+
+class MultiheadAttention(nn.Module):
+  def __init__(self, d_model, n_heads, d_keys, mask_flag, r_att_drop=0.1):
+    super(MultiheadAttention, self).__init__()
+    self.h, self.d, self.mask_flag= n_heads, d_keys, mask_flag
+    self.proj_q = nn.Linear(d_model, self.h * self.d)
+    self.proj_k = nn.Linear(d_model, self.h * self.d)
+    self.proj_v = nn.Linear(d_model, self.h * self.d)
+    self.proj_out = nn.Linear(self.h * self.d, d_model)
+    self.dropout = nn.Dropout(r_att_drop) 
+
+  def forward(self, q, k, v):
+    b, n_q, n_k, h, d = q.size(0), q.size(1), k.size(1), self.h, self.d
+
+    q, k, v = self.proj_q(q), self.proj_k(k), self.proj_v(v)        # b, n_*, h*d
+    q, k, v = map(lambda x: x.reshape(b, -1, h, d), [q, k, v])      # b, n_*, h, d
+    scores = torch.einsum('bnhd,bmhd->bhnm', (q,k))                 # b, h, n_q, n_k
+    
+    if self.mask_flag:
+      att_mask = TriangularCausalMask(b, n_q, device=q.device)
+      scores.masked_fill_(att_mask.mask, -np.inf)
+
+    att = F.softmax(scores / (self.d ** .5), dim=-1)                # b, h, n_q, n_k
+    att = self.dropout(att)
+    att_out = torch.einsum('bhnm,bmhd->bnhd', (att,v))              # b, n_q, h, d
+    att_out = att_out.reshape(b, -1, h*d)                           # b, n_q, h*d
+    out = self.proj_out(att_out)                                    # b, n_q, d_model
+    return out

gluformer/decoder.py

@@ -0,0 +1,50 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .attention import *
+
+class DecoderLayer(nn.Module):
+  def __init__(self, self_att, cross_att, d_model, d_fcn,
+                r_drop, activ="relu"):
+    super(DecoderLayer, self).__init__()
+
+    self.self_att = self_att
+    self.cross_att = cross_att
+
+    self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_fcn, kernel_size=1)
+    self.conv2 = nn.Conv1d(in_channels=d_fcn, out_channels=d_model, kernel_size=1)
+
+    self.norm1 = nn.LayerNorm(d_model)
+    self.norm2 = nn.LayerNorm(d_model)
+    self.norm3 = nn.LayerNorm(d_model)
+
+    self.dropout = nn.Dropout(r_drop)
+    self.activ = F.relu if activ == "relu" else F.gelu
+
+  def forward(self, x_dec, x_enc):
+    x_dec = x_dec + self.self_att(x_dec, x_dec, x_dec)
+    x_dec = self.norm1(x_dec)
+
+    x_dec = x_dec + self.cross_att(x_dec, x_enc, x_enc)
+    res = x_dec = self.norm2(x_dec)
+
+    res = self.dropout(self.activ(self.conv1(res.transpose(-1,1))))
+    res = self.dropout(self.conv2(res).transpose(-1,1))
+    
+    return self.norm3(x_dec+res)
+
+class Decoder(nn.Module):
+  def __init__(self, layers, norm_layer=None):
+    super(Decoder, self).__init__()
+    self.layers = nn.ModuleList(layers)
+    self.norm = norm_layer
+
+  def forward(self, x_dec, x_enc):
+    for layer in self.layers:
+      x_dec = layer(x_dec, x_enc)
+
+    if self.norm is not None:
+      x_dec = self.norm(x_dec)
+
+    return x_dec

gluformer/embed.py

@@ -0,0 +1,69 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+
+class PositionalEmbedding(nn.Module):
+  def __init__(self, d_model, max_len=5000):
+    super(PositionalEmbedding, self).__init__()
+    # Compute the positional encodings once in log space.
+    pos_emb = torch.zeros(max_len, d_model)
+    pos_emb.require_grad = False
+
+    position = torch.arange(0, max_len).unsqueeze(1)
+    div_term = (torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model)).exp()
+
+    pos_emb[:, 0::2] = torch.sin(position * div_term)
+    pos_emb[:, 1::2] = torch.cos(position * div_term)
+
+    pos_emb = pos_emb.unsqueeze(0)
+    self.register_buffer('pos_emb', pos_emb)
+
+  def forward(self, x):
+    return self.pos_emb[:, :x.size(1)]
+
+class TokenEmbedding(nn.Module):
+  def __init__(self, d_model):
+    super(TokenEmbedding, self).__init__()
+    D_INP = 1 # one sequence
+    self.conv = nn.Conv1d(in_channels=D_INP, out_channels=d_model, 
+                          kernel_size=3, padding=1, padding_mode='circular')
+    # nn.init.kaiming_normal_(self.conv.weight, mode='fan_in', nonlinearity='leaky_relu')
+
+  def forward(self, x):
+    x = self.conv(x.transpose(-1, 1)).transpose(-1, 1)
+    return x
+
+class TemporalEmbedding(nn.Module):
+  def __init__(self, d_model, num_features):
+    super(TemporalEmbedding, self).__init__()
+    self.embed = nn.Linear(num_features, d_model)
+  
+  def forward(self, x):
+    return self.embed(x)
+
+class SubjectEmbedding(nn.Module):
+  def __init__(self, d_model, num_features):
+    super(SubjectEmbedding, self).__init__()
+    self.id_embedding = nn.Linear(num_features, d_model)
+
+  def forward(self, x):
+    embed_x = self.id_embedding(x)
+
+    return embed_x
+
+class DataEmbedding(nn.Module):
+  def __init__(self, d_model, r_drop, num_dynamic_features, num_static_features):
+    super(DataEmbedding, self).__init__()
+    # note: d_model // 2 == 0
+    self.value_embedding = TokenEmbedding(d_model)
+    self.time_embedding = TemporalEmbedding(d_model, num_dynamic_features) # alternative: TimeFeatureEmbedding
+    self.positional_embedding = PositionalEmbedding(d_model)
+    self.subject_embedding = SubjectEmbedding(d_model, num_static_features)
+    self.dropout = nn.Dropout(r_drop)
+
+  def forward(self, x_id, x, x_mark):
+    x =  self.value_embedding(x) + self.positional_embedding(x) + self.time_embedding(x_mark)
+    x_id = self.subject_embedding(x_id)
+    x = torch.cat((x_id.unsqueeze(1), x), dim = 1)
+    return self.dropout(x)

gluformer/encoder.py

@@ -0,0 +1,67 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .attention import *
+
+class ConvLayer(nn.Module):
+  def __init__(self, d_model):
+    super(ConvLayer, self).__init__()
+    self.downConv = nn.Conv1d(in_channels=d_model, out_channels=d_model,
+                              kernel_size=3, padding=1, padding_mode='circular')
+    self.norm = nn.BatchNorm1d(d_model)
+    self.activ = nn.ELU()
+    self.maxPool = nn.MaxPool1d(kernel_size=3, stride=2, padding=1)
+
+  def forward(self, x):
+    x = self.downConv(x.transpose(-1, 1))
+    x = self.norm(x)
+    x = self.activ(x)
+    x = self.maxPool(x)
+    x = x.transpose(-1,1)
+    return x
+
+class EncoderLayer(nn.Module):
+  def __init__(self, att, d_model, d_fcn, r_drop, activ="relu"):
+    super(EncoderLayer, self).__init__()
+    
+    self.att = att
+    self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_fcn, kernel_size=1)
+    self.conv2 = nn.Conv1d(in_channels=d_fcn, out_channels=d_model, kernel_size=1)
+    self.norm1 = nn.LayerNorm(d_model)
+    self.norm2 = nn.LayerNorm(d_model)
+    self.dropout = nn.Dropout(r_drop)
+    self.activ = F.relu if activ == "relu" else F.gelu
+
+  def forward(self, x):
+    new_x = self.att(x, x, x)
+    x = x + self.dropout(new_x)
+
+    res = x = self.norm1(x)
+    res = self.dropout(self.activ(self.conv1(res.transpose(-1,1))))
+    res = self.dropout(self.conv2(res).transpose(-1,1))
+
+    return self.norm2(x+res)
+
+class Encoder(nn.Module):
+  def __init__(self, enc_layers, conv_layers=None, norm_layer=None):
+    super(Encoder, self).__init__()
+    self.enc_layers = nn.ModuleList(enc_layers)
+    self.conv_layers = nn.ModuleList(conv_layers) if conv_layers is not None else None
+    self.norm = norm_layer
+
+  def forward(self, x):
+    # x [B, L, D]
+    if self.conv_layers is not None:
+        for enc_layer, conv_layer in zip(self.enc_layers, self.conv_layers):
+            x = enc_layer(x)
+            x = conv_layer(x)
+        x = self.enc_layers[-1](x)
+    else:
+        for enc_layer in self.enc_layers:
+            x = enc_layer(x)
+
+    if self.norm is not None:
+        x = self.norm(x)
+
+    return x

gluformer/model.py

@@ -0,0 +1,334 @@
+import os
+import sys
+from tqdm import tqdm
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .embed import *
+from .attention import *
+from .encoder import *
+from .decoder import *
+from .variance import *
+
+############################################
+# Added for GluNet package
+############################################
+import optuna
+import darts
+from torch.utils.tensorboard import SummaryWriter
+sys.path.append(os.path.join(os.path.dirname(__file__), '../..'))
+from glucose.gluformer.utils.training import ExpLikeliLoss, \
+                                         EarlyStop, \
+                                         modify_collate, \
+                                         adjust_learning_rate
+from glucose.utils.darts_dataset import SamplingDatasetDual
+############################################
+
+class Gluformer(nn.Module):
+  def __init__(self, d_model, n_heads, d_fcn, r_drop, 
+                activ, num_enc_layers, num_dec_layers, 
+                distil, len_seq, len_pred, num_dynamic_features,
+                num_static_features, label_len):
+    super(Gluformer, self).__init__()
+    # Set prediction length
+    self.len_pred = len_pred
+    self.label_len = label_len
+    # Embedding
+    # note: d_model // 2 == 0
+    self.enc_embedding = DataEmbedding(d_model, r_drop, num_dynamic_features, num_static_features)
+    self.dec_embedding = DataEmbedding(d_model, r_drop, num_dynamic_features, num_static_features)
+    # Encoding
+    self.encoder = Encoder(
+      [
+        EncoderLayer(
+          att=MultiheadAttention(d_model=d_model, n_heads=n_heads, 
+                                  d_keys=d_model//n_heads, mask_flag=False, 
+                                  r_att_drop=r_drop),
+          d_model=d_model,
+          d_fcn=d_fcn,
+          r_drop=r_drop,
+          activ=activ) for l in range(num_enc_layers)
+      ],
+      [
+        ConvLayer(
+          d_model) for l in range(num_enc_layers-1)
+      ] if distil else None, 
+      norm_layer=torch.nn.LayerNorm(d_model)
+    )
+
+    # Decoding
+    self.decoder = Decoder(
+      [
+        DecoderLayer(
+          self_att=MultiheadAttention(d_model=d_model, n_heads=n_heads, 
+                                  d_keys=d_model//n_heads, mask_flag=True, 
+                                  r_att_drop=r_drop),
+          cross_att=MultiheadAttention(d_model=d_model, n_heads=n_heads, 
+                                  d_keys=d_model//n_heads, mask_flag=False, 
+                                  r_att_drop=r_drop),
+          d_model=d_model,
+          d_fcn=d_fcn,
+          r_drop=r_drop,
+          activ=activ) for l in range(num_dec_layers)
+      ], 
+      norm_layer=torch.nn.LayerNorm(d_model)
+    )
+    
+    # Output
+    D_OUT = 1
+    self.projection = nn.Linear(d_model, D_OUT, bias=True)
+
+    # Train variance
+    self.var = Variance(d_model, r_drop, len_seq)
+
+  def forward(self, x_id, x_enc, x_mark_enc, x_dec, x_mark_dec):
+    enc_out = self.enc_embedding(x_id, x_enc, x_mark_enc)
+    var_out = self.var(enc_out)
+    enc_out = self.encoder(enc_out)
+
+    dec_out = self.dec_embedding(x_id, x_dec, x_mark_dec)
+    dec_out = self.decoder(dec_out, enc_out)
+    dec_out = self.projection(dec_out)
+    
+    return dec_out[:, -self.len_pred:, :], var_out # [B, L, D], log variance
+  
+  ############################################
+  # Added for GluNet package
+  ############################################
+  def fit(self, 
+          train_dataset: SamplingDatasetDual,
+          val_dataset: SamplingDatasetDual,
+          learning_rate: float = 1e-3,
+          batch_size: int = 32,
+          epochs: int = 100,
+          num_samples: int = 100,
+          device: str = 'cuda',
+          model_path: str = None,
+          trial: optuna.trial.Trial = None,
+          logger: SummaryWriter = None,):
+    """
+    Fit the model to the data, using Optuna for hyperparameter tuning.
+    
+    Parameters
+    ----------
+    train_dataset: SamplingDatasetPast
+      Training dataset.
+    val_dataset: SamplingDatasetPast
+      Validation dataset.
+    learning_rate: float
+      Learning rate for Adam.
+    batch_size: int
+      Batch size.
+    epochs: int
+      Number of epochs.
+    num_samples: int
+      Number of samples for infinite mixture
+    device: str
+      Device to use.
+    model_path: str
+      Path to save the model.
+    trial: optuna.trial.Trial
+      Trial for hyperparameter tuning.
+    logger: SummaryWriter
+      Tensorboard logger for logging.
+    """
+    # create data loaders, optimizer, loss, and early stopping
+    collate_fn_custom = modify_collate(num_samples)
+    train_loader = torch.utils.data.DataLoader(train_dataset, 
+                                               batch_size=batch_size,
+                                               shuffle=True,
+                                               drop_last=True,
+                                               collate_fn=collate_fn_custom)
+    val_loader = torch.utils.data.DataLoader(val_dataset,
+                                             batch_size=batch_size,
+                                             shuffle=True,
+                                             drop_last=True,
+                                             collate_fn=collate_fn_custom)
+    criterion = ExpLikeliLoss(num_samples)
+    optimizer = torch.optim.Adam(self.parameters(), lr=learning_rate, betas=(0.1, 0.9))
+    scaler = torch.cuda.amp.GradScaler()
+    early_stop = EarlyStop(patience=10, delta=0.001)
+    self.to(device)
+    # train and evaluate the model
+    for epoch in range(epochs):
+      train_loss = []
+      for i, (past_target_series, 
+              past_covariates, 
+              future_covariates, 
+              static_covariates, 
+              future_target_series) in enumerate(train_loader):
+        # zero out gradient
+        optimizer.zero_grad()
+        # reshape static covariates to be [batch_size, num_static_covariates]
+        static_covariates = static_covariates.reshape(-1, static_covariates.shape[-1])
+        # create decoder input: pad with zeros the prediction sequence
+        dec_inp = torch.cat([past_target_series[:, -self.label_len:, :], 
+                             torch.zeros([
+                                          past_target_series.shape[0], 
+                                          self.len_pred, 
+                                          past_target_series.shape[-1]
+                                          ])], 
+                             dim=1)
+        future_covariates = torch.cat([past_covariates[:, -self.label_len:, :],
+                                        future_covariates], dim=1)
+        # move to device
+        dec_inp = dec_inp.to(device)
+        past_target_series = past_target_series.to(device)
+        past_covariates = past_covariates.to(device)
+        future_covariates = future_covariates.to(device)
+        static_covariates = static_covariates.to(device)
+        future_target_series = future_target_series.to(device)
+        # forward pass with autograd
+        with torch.cuda.amp.autocast():
+          pred, logvar = self(static_covariates,
+                              past_target_series,
+                              past_covariates,
+                              dec_inp,
+                              future_covariates)
+          loss = criterion(pred, future_target_series, logvar)
+        # backward pass
+        scaler.scale(loss).backward()
+        scaler.step(optimizer)
+        scaler.update()
+        # log loss
+        if logger is not None:
+          logger.add_scalar('train_loss', loss.item(), epoch * len(train_loader) + i)
+        train_loss.append(loss.item())
+      # log loss
+      if logger is not None:
+        logger.add_scalar('train_loss_epoch', np.mean(train_loss), epoch)
+      # evaluate the model
+      val_loss = []
+      with torch.no_grad():
+        for i, (past_target_series, 
+                past_covariates, 
+                future_covariates, 
+                static_covariates, 
+                future_target_series) in enumerate(val_loader):
+          # reshape static covariates to be [batch_size, num_static_covariates]
+          static_covariates = static_covariates.reshape(-1, static_covariates.shape[-1])
+          # create decoder input
+          dec_inp = torch.cat([past_target_series[:, -self.label_len:, :], 
+                              torch.zeros([
+                                            past_target_series.shape[0], 
+                                            self.len_pred, 
+                                            past_target_series.shape[-1]
+                                            ])], 
+                              dim=1)
+          future_covariates = torch.cat([past_covariates[:, -self.label_len:, :],
+                                          future_covariates], dim=1)
+          # move to device
+          dec_inp = dec_inp.to(device)
+          past_target_series = past_target_series.to(device)
+          past_covariates = past_covariates.to(device)
+          future_covariates = future_covariates.to(device)
+          static_covariates = static_covariates.to(device)
+          future_target_series = future_target_series.to(device)
+          # forward pass
+          pred, logvar = self(static_covariates,
+                              past_target_series,
+                              past_covariates,
+                              dec_inp,
+                              future_covariates)
+          loss = criterion(pred, future_target_series, logvar)
+          val_loss.append(loss.item())
+          # log loss
+          if logger is not None:
+            logger.add_scalar('val_loss', loss.item(), epoch * len(val_loader) + i)
+      # log loss
+      logger.add_scalar('val_loss_epoch', np.mean(val_loss), epoch)
+      # check early stopping
+      early_stop(np.mean(val_loss), self, model_path)
+      if early_stop.stop:
+        break
+      # check pruning 
+      if trial is not None:
+        trial.report(np.mean(val_loss), epoch)
+        if trial.should_prune():
+          raise optuna.exceptions.TrialPruned()
+    # load best model
+    if model_path is not None:  
+      self.load_state_dict(torch.load(model_path))
+
+  def predict(self, test_dataset: SamplingDatasetDual, 
+              batch_size: int = 32,
+              num_samples: int = 100,
+              device: str = 'cuda'):
+    """
+    Predict the future target series given the supplied samples from the dataset.
+
+    Parameters
+    ----------
+    test_dataset : SamplingDatasetInferenceDual
+        The dataset to use for inference.
+    batch_size : int, optional
+        The batch size to use for inference, by default 32
+    num_samples : int, optional
+        The number of samples to use for inference, by default 100
+    
+    Returns
+    -------
+    Predictions
+        The predicted future target series in shape n x len_pred x num_samples, where
+        n is total number of predictions.
+    Logvar
+        The logvariance of the predicted future target series in shape n x len_pred.
+    """
+    # define data loader
+    collate_fn_custom = modify_collate(num_samples)
+    test_loader = torch.utils.data.DataLoader(test_dataset,
+                                              batch_size=batch_size,
+                                              shuffle=False,
+                                              drop_last=False,
+                                              collate_fn=collate_fn_custom)
+    # predict
+    self.train()
+    # move to device
+    self.to(device)
+    predictions = []; logvars = []
+    for i, (past_target_series,
+            historic_future_covariates,
+            future_covariates,
+            static_covariates) in enumerate(test_loader):
+      # reshape static covariates to be [batch_size, num_static_covariates]
+      static_covariates = static_covariates.reshape(-1, static_covariates.shape[-1])
+      # create decoder input
+      dec_inp = torch.cat([past_target_series[:, -self.label_len:, :], 
+                          torch.zeros([
+                                        past_target_series.shape[0], 
+                                        self.len_pred, 
+                                        past_target_series.shape[-1]
+                                        ])], 
+                          dim=1)
+      future_covariates = torch.cat([historic_future_covariates[:, -self.label_len:, :],
+                                      future_covariates], dim=1)
+      # move to device
+      dec_inp = dec_inp.to(device)
+      past_target_series = past_target_series.to(device)
+      historic_future_covariates = historic_future_covariates.to(device)
+      future_covariates = future_covariates.to(device)
+      static_covariates = static_covariates.to(device)
+      # forward pass
+      pred, logvar = self(static_covariates,
+                          past_target_series,
+                          historic_future_covariates,
+                          dec_inp,
+                          future_covariates)
+      # transfer in numpy and arrange sample along last axis
+      pred = pred.cpu().detach().numpy()
+      logvar = logvar.cpu().detach().numpy()
+      pred = pred.transpose((1, 0, 2)).reshape((pred.shape[1], -1, num_samples)).transpose((1, 0, 2))
+      logvar = logvar.transpose((1, 0, 2)).reshape((logvar.shape[1], -1, num_samples)).transpose((1, 0, 2))
+      predictions.append(pred)
+      logvars.append(logvar)
+    predictions = np.concatenate(predictions, axis=0)
+    logvars = np.concatenate(logvars, axis=0)
+    return predictions, logvars
+    
+  ############################################
+
+
+
+

gluformer/utils/collate.py

@@ -0,0 +1,84 @@
+r""""Contains definitions of the methods used by the _BaseDataLoaderIter workers to
+collate samples fetched from dataset into Tensor(s).
+These **needs** to be in global scope since Py2 doesn't support serializing
+static methods.
+"""
+
+import torch
+import re
+import collections
+
+np_str_obj_array_pattern = re.compile(r'[SaUO]')
+
+
+def default_convert(data):
+    r"""Converts each NumPy array data field into a tensor"""
+    elem_type = type(data)
+    if isinstance(data, torch.Tensor):
+        return data
+    elif elem_type.__module__ == 'numpy' and elem_type.__name__ != 'str_' \
+            and elem_type.__name__ != 'string_':
+        # array of string classes and object
+        if elem_type.__name__ == 'ndarray' \
+                and np_str_obj_array_pattern.search(data.dtype.str) is not None:
+            return data
+        return torch.as_tensor(data)
+    elif isinstance(data, collections.abc.Mapping):
+        return {key: default_convert(data[key]) for key in data}
+    elif isinstance(data, tuple) and hasattr(data, '_fields'):  # namedtuple
+        return elem_type(*(default_convert(d) for d in data))
+    elif isinstance(data, collections.abc.Sequence) and not isinstance(data, str):
+        return [default_convert(d) for d in data]
+    else:
+        return data
+
+
+default_collate_err_msg_format = (
+    "default_collate: batch must contain tensors, numpy arrays, numbers, "
+    "dicts or lists; found {}")
+
+
+def default_collate(batch):
+    r"""Puts each data field into a tensor with outer dimension batch size"""
+
+    elem = batch[0]
+    elem_type = type(elem)
+    if isinstance(elem, torch.Tensor):
+        out = None
+        if torch.utils.data.get_worker_info() is not None:
+            # If we're in a background process, concatenate directly into a
+            # shared memory tensor to avoid an extra copy
+            numel = sum(x.numel() for x in batch)
+            storage = elem.storage()._new_shared(numel)
+            out = elem.new(storage)
+        return torch.stack(batch, 0, out=out)
+    elif elem_type.__module__ == 'numpy' and elem_type.__name__ != 'str_' \
+            and elem_type.__name__ != 'string_':
+        if elem_type.__name__ == 'ndarray' or elem_type.__name__ == 'memmap':
+            # array of string classes and object
+            if np_str_obj_array_pattern.search(elem.dtype.str) is not None:
+                raise TypeError(default_collate_err_msg_format.format(elem.dtype))
+
+            return default_collate([torch.as_tensor(b) for b in batch])
+        elif elem.shape == ():  # scalars
+            return torch.as_tensor(batch)
+    elif isinstance(elem, float):
+        return torch.tensor(batch, dtype=torch.float64)
+    elif isinstance(elem, int):
+        return torch.tensor(batch)
+    elif isinstance(elem, str):
+        return batch
+    elif isinstance(elem, collections.abc.Mapping):
+        return {key: default_collate([d[key] for d in batch]) for key in elem}
+    elif isinstance(elem, tuple) and hasattr(elem, '_fields'):  # namedtuple
+        return elem_type(*(default_collate(samples) for samples in zip(*batch)))
+    elif isinstance(elem, collections.abc.Sequence):
+        # check to make sure that the elements in batch have consistent size
+        it = iter(batch)
+        elem_size = len(next(it))
+        if not all(len(elem) == elem_size for elem in it):
+            raise RuntimeError('each element in list of batch should be of equal size')
+        transposed = zip(*batch)
+        return [default_collate(samples) for samples in transposed]
+
+    raise TypeError(default_collate_err_msg_format.format(elem_type))

gluformer/utils/evaluation.py

@@ -0,0 +1,81 @@
+import sys
+import os
+import yaml
+import random
+from typing import Any, \
+                   BinaryIO, \
+                   Callable, \
+                   Dict, \
+                   List, \
+                   Optional, \
+                   Sequence, \
+                   Tuple, \
+                   Union
+
+import numpy as np 
+import scipy as sp
+import pandas as pd
+import torch
+
+# import data formatter
+sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
+
+def test(series: np.ndarray, 
+         forecasts: np.ndarray, 
+         var: np.ndarray,
+         cal_thresholds: Optional[np.ndarray] = np.linspace(0, 1, 11),
+    ):
+    """
+    Test the (rescaled to original scale) forecasts on the series.
+
+    Parameters
+    ----------
+    series
+        The target time series of shape (n, t), 
+        where t is length of prediction.
+    forecasts
+        The forecasted means of mixture components of shape (n, t, k),
+        where k is the number of mixture components.
+    var
+        The forecasted variances of mixture components of shape (n, 1, k),
+        where k is the number of mixture components.
+    metric
+        The metric or metrics to use for backtesting.
+    cal_thresholds
+        The thresholds to use for computing the calibration error.
+
+    Returns
+    -------
+    np.ndarray
+        Error array. Array of shape (n, p) 
+        where n = series.shape[0] = forecasts.shape[0] and p = len(metric). 
+    float
+        The estimated log-likelihood of the model on the data.
+    np.ndarray
+        The ECE for each time point in the forecast.
+    """
+    # compute errors: 1) get samples 2) compute errors using median
+    samples = np.random.normal(loc=forecasts[..., None],
+                               scale=np.sqrt(var)[..., None],
+                               size=(forecasts.shape[0], 
+                                     forecasts.shape[1], 
+                                     forecasts.shape[2],
+                                     30))
+    samples = samples.reshape(samples.shape[0], samples.shape[1], -1)
+    mse = np.mean((series.squeeze() - forecasts.mean(axis=-1))**2, axis=-1)
+    mae = np.mean(np.abs(series.squeeze() - forecasts.mean(axis=-1)), axis=-1)
+    errors = np.stack([mse, mae], axis=-1)
+     
+    # compute likelihood
+    log_likelihood = sp.special.logsumexp((forecasts - series)**2 / (2 * var) - 
+                                           0.5 * np.log(2 * np.pi * var), axis=-1)
+    log_likelihood = np.mean(log_likelihood)
+    
+    # compute calibration error: 
+    cal_error = np.zeros(forecasts.shape[1])
+    for p in cal_thresholds:
+        q = np.quantile(samples, p, axis=-1)
+        est_p = np.mean(series.squeeze() <= q, axis=0)
+        cal_error += (est_p - p) ** 2
+        
+    return errors, log_likelihood, cal_error

gluformer/utils/training.py

@@ -0,0 +1,80 @@
+from typing import List, Callable, Any, Tuple
+
+import numpy as np
+import torch
+from sympy import pprint
+from torch import nn, Tensor
+
+from .collate import default_collate
+
+class EarlyStop:
+  def __init__(self, patience: int, delta: float):
+    self.patience: int = patience
+    self.delta: float = delta
+    self.counter: int = 0
+    self.best_loss: float = np.Inf
+    self.stop: bool = False
+
+  def __call__(self, loss: float, model: nn.Module, path: str) -> None:
+    if loss < self.best_loss:
+      self.best_loss = loss
+      self.counter = 0
+      torch.save(model.state_dict(), path)
+    elif loss > self.best_loss + self.delta:
+      self.counter = self.counter + 1
+      if self.counter >= self.patience:
+        self.stop = True
+
+class ExpLikeliLoss(nn.Module):
+  def __init__(self, num_samples: int = 100):
+    super(ExpLikeliLoss, self).__init__()
+    self.num_samples: int = num_samples
+
+  def forward(self, pred: Tensor, true: Tensor, logvar: Tensor) -> Tensor:
+    b, l, d = pred.size(0), pred.size(1), pred.size(2)
+    true = true.transpose(0,1).reshape(l, -1, self.num_samples).transpose(0, 1)
+    pred = pred.transpose(0,1).reshape(l, -1, self.num_samples).transpose(0, 1)
+    logvar = logvar.reshape(-1, self.num_samples)
+
+    loss = torch.mean((-1) * torch.logsumexp((-l / 2) * logvar + (-1 / (2 * torch.exp(logvar))) * torch.sum((true - pred) ** 2, dim=1), dim=1))
+    return loss
+
+def modify_collate(num_samples: int) -> Callable[[List[Any]], Any]:
+  def wrapper(batch: List[Any]) -> Any:
+    batch_rep = [sample for sample in batch for _ in range(num_samples)]
+    result = default_collate(batch_rep)
+    return result
+  return wrapper
+
+def adjust_learning_rate(model_optim: torch.optim.Optimizer, epoch: int, lr: float) -> None:
+  lr = lr * (0.5 ** epoch)
+  print("Learning rate halving...")
+  print(f"New lr: {lr:.7f}")
+  for param_group in model_optim.param_groups:
+    param_group['lr'] = lr
+
+def process_batch(
+        subj_id: Tensor,
+        batch_x: Tensor,
+        batch_y: Tensor,
+        batch_x_mark: Tensor,
+        batch_y_mark: Tensor,
+        len_pred: int,
+        len_label: int,
+        model: nn.Module,
+        device: torch.device
+) -> Tuple[Tensor, Tensor, Tensor]:
+  subj_id = subj_id.long().to(device)
+  batch_x = batch_x.float().to(device)
+  batch_y = batch_y.float()
+  batch_x_mark = batch_x_mark.float().to(device)
+  batch_y_mark = batch_y_mark.float().to(device)
+
+  true = batch_y[:, -len_pred:, :].to(device)
+
+  dec_inp = torch.zeros([batch_y.shape[0], len_pred, batch_y.shape[-1]], dtype=torch.float, device=device)
+  dec_inp = torch.cat([batch_y[:, :len_label, :].to(device), dec_inp], dim=1)
+
+  pred, logvar = model(subj_id, batch_x, batch_x_mark, dec_inp, batch_y_mark)
+
+  return pred, true, logvar

gluformer/variance.py

@@ -0,0 +1,24 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class Variance(nn.Module):
+  def __init__(self, d_model, r_drop, len_seq):
+    super(Variance, self).__init__()
+
+    self.proj1 = nn.Linear(d_model, 1)
+    self.dropout = nn.Dropout(r_drop) 
+    self.activ1 = nn.ReLU()
+    # + 1 (for seq) for embedded person token
+    self.proj2 = nn.Linear(len_seq+1, 1)
+    self.activ2 = nn.Tanh()
+
+  def forward(self, x):
+    x = self.proj1(x)
+    x = self.activ1(x)
+    x = self.dropout(x)
+    x = x.transpose(-1, 1)
+    x = self.proj2(x)
+    # scale to [-10, 10] range
+    x = 10 * self.activ2(x)
+    return x

main.py

@@ -0,0 +1,8 @@
+import gradio as gr
+from tools import *
+
+
+def gradio_output():
+    return (predict_glucose_tool())
+
+gr.Interface(fn=gradio_output).launch()

tools.py

@@ -0,0 +1,198 @@
+import sys
+import os
+import pickle
+import gzip
+from pathlib import Path
+
+import seaborn as sns
+import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib.colors as mcolors
+from matplotlib.figure import Figure
+import torch
+from scipy import stats
+
+from gluformer.model import Gluformer
+from utils.darts_processing import *
+from utils.darts_dataset import *
+
+
+import hashlib
+from urllib.parse import urlparse
+
+import numpy as np
+import typer
+
+
+glucose = Path(os.path.abspath(__file__)).parent.resolve()
+file_directory = glucose / "files"
+
+
+def plot_forecast(forecasts: np.ndarray, scalers: Any, dataset_test_glufo: Any, filename: str):
+    filename=filename
+    forecasts = (forecasts - scalers['target'].min_) / scalers['target'].scale_
+
+    trues = [dataset_test_glufo.evalsample(i) for i in range(len(dataset_test_glufo))]
+    trues = scalers['target'].inverse_transform(trues)
+
+    trues = [ts.values() for ts in trues]  # Convert TimeSeries to numpy arrays
+    trues = np.array(trues)
+
+    inputs = [dataset_test_glufo[i][0] for i in range(len(dataset_test_glufo))]
+    inputs = (np.array(inputs) - scalers['target'].min_) / scalers['target'].scale_
+
+    # Plot settings
+    colors = ['#00264c', '#0a2c62', '#14437f', '#1f5a9d', '#2973bb', '#358ad9', '#4d9af4', '#7bb7ff', '#add5ff', '#e6f3ff']
+    cmap = mcolors.LinearSegmentedColormap.from_list('my_colormap', colors)
+    sns.set_theme(style="whitegrid")
+
+    # Generate the plot
+    fig, ax = plt.subplots(figsize=(10, 6))
+
+
+    # Select a specific sample to plot
+    ind = 30  # Example index
+
+    samples = np.random.normal(
+        loc=forecasts[ind, :],  # Mean (center) of the distribution
+        scale=0.1,  # Standard deviation (spread) of the distribution
+        size=(forecasts.shape[1], forecasts.shape[2])
+    )
+    #samples = samples.reshape(samples.shape[0], samples.shape[1], -1)
+    #print ("samples",samples.shape)
+
+    # Plot predictive distribution
+    for point in range(samples.shape[0]):
+        kde = stats.gaussian_kde(samples[point,:])
+        maxi, mini = 1.2 * np.max(samples[point, :]), 0.8 * np.min(samples[point, :])
+        y_grid = np.linspace(mini, maxi, 200)
+        x = kde(y_grid)
+        ax.fill_betweenx(y_grid, x1=point, x2=point - x * 15,
+                         alpha=0.7,
+                         edgecolor='black',
+                         color=cmap(point / samples.shape[0]))
+
+    # Plot median
+    forecast = samples[:, :]
+    median = np.quantile(forecast, 0.5, axis=-1)
+    ax.plot(np.arange(12), median, color='red', marker='o')
+
+    # Plot true values
+    ax.plot(np.arange(-12, 12), np.concatenate([inputs[ind, -12:], trues[ind, :]]), color='blue')
+
+    # Add labels and title
+    ax.set_xlabel('Time (in 5 minute intervals)')
+    ax.set_ylabel('Glucose (mg/dL)')
+    ax.set_title(f'Gluformer Prediction with Gradient for dateset')
+
+    # Adjust font sizes
+    ax.xaxis.label.set_fontsize(16)
+    ax.yaxis.label.set_fontsize(16)
+    ax.title.set_fontsize(18)
+    for item in ax.get_xticklabels() + ax.get_yticklabels():
+        item.set_fontsize(14)
+
+    # Save figure
+    plt.tight_layout()
+    where = file_directory /filename
+    plt.savefig(str(where), dpi=300, bbox_inches='tight')
+
+    return where,ax
+
+
+
+def generate_filename_from_url(url: str, extension: str = "png") -> str:
+    """
+    :param url:
+    :param extension:
+    :return:
+    """
+    # Extract the last segment of the URL
+    last_segment = urlparse(url).path.split('/')[-1]
+
+    # Compute the hash of the URL
+    url_hash = hashlib.md5(url.encode('utf-8')).hexdigest()
+
+    # Create the filename
+    filename = f"{last_segment.replace('.','_')}_{url_hash}.{extension}"
+
+    return filename
+
+
+
+def predict_glucose_tool(url: str= 'https://huggingface.co/datasets/Livia-Zaharia/glucose_processed/blob/main/livia_mini.csv',
+                        model: str = 'https://huggingface.co/Livia-Zaharia/gluformer_models/blob/main/gluformer_1samples_10000epochs_10heads_32batch_geluactivation_livia_mini_weights.pth'
+                    ) -> Figure:
+    """
+    Function to predict future glucose of user. It receives URL with users csv. It will run an ML and will return URL with predictions that user can open on her own..
+    :param url: of the csv file with glucose values
+    :param model: model that is used to predict the glucose
+    :param explain if it should give both url and explanation
+    :param if the person is diabetic when doing prediction and explanation
+    :return:
+    """
+
+    formatter, series, scalers = load_data(url=str(url), config_path=file_directory / "config.yaml", use_covs=True,
+                                           cov_type='dual',
+                                           use_static_covs=True)
+
+    filename = generate_filename_from_url(url)
+
+    formatter.params['gluformer'] = {
+        'in_len': 96,  # example input length, adjust as necessary
+        'd_model': 512,  # model dimension
+        'n_heads': 10,  # number of attention heads##############################################################################
+        'd_fcn': 1024,  # fully connected layer dimension
+        'num_enc_layers': 2,  # number of encoder layers
+        'num_dec_layers': 2,  # number of decoder layers
+        'length_pred': 12  # prediction length, adjust as necessary
+    }
+
+    num_dynamic_features = series['train']['future'][-1].n_components
+    num_static_features = series['train']['static'][-1].n_components
+
+    glufo = Gluformer(
+        d_model=formatter.params['gluformer']['d_model'],
+        n_heads=formatter.params['gluformer']['n_heads'],
+        d_fcn=formatter.params['gluformer']['d_fcn'],
+        r_drop=0.2,
+        activ='gelu',
+        num_enc_layers=formatter.params['gluformer']['num_enc_layers'],
+        num_dec_layers=formatter.params['gluformer']['num_dec_layers'],
+        distil=True,
+        len_seq=formatter.params['gluformer']['in_len'],
+        label_len=formatter.params['gluformer']['in_len'] // 3,
+        len_pred=formatter.params['length_pred'],
+        num_dynamic_features=num_dynamic_features,
+        num_static_features=num_static_features
+    )
+    weights = gr.Interface.load(model)
+    assert f"weights for {model} should exist", weights.exists()
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    glufo.load_state_dict(torch.load(str(weights), map_location=torch.device(device), weights_only=False))
+
+    # Define dataset for inference
+    dataset_test_glufo = SamplingDatasetInferenceDual(
+        target_series=series['test']['target'],
+        covariates=series['test']['future'],
+        input_chunk_length=formatter.params['gluformer']['in_len'],
+        output_chunk_length=formatter.params['length_pred'],
+        use_static_covariates=True,
+        array_output_only=True
+    )
+
+    forecasts, _ = glufo.predict(
+        dataset_test_glufo,
+        batch_size=16,####################################################
+        num_samples=10,
+        device='cpu'
+    )
+    figure_path, result = plot_forecast(forecasts, scalers, dataset_test_glufo,filename)
+    
+    return result
+
+
+
+if __name__ == "__main__":
+    predict_glucose_tool()

utils/darts_dataset.py

@@ -0,0 +1,881 @@
+import sys
+import os
+import yaml
+import random
+from typing import Any, BinaryIO, Callable, Dict, List, Optional, Sequence, Tuple, Union
+
+import numpy as np 
+from scipy import stats
+import pandas as pd
+import darts
+
+from darts import models
+from darts import metrics
+from darts import TimeSeries
+from darts.dataprocessing.transformers import Scaler
+from pytorch_lightning.callbacks import Callback
+
+# for darts dataset
+from darts.logging import get_logger, raise_if_not
+
+from darts.utils.data.training_dataset import PastCovariatesTrainingDataset, \
+                                              DualCovariatesTrainingDataset, \
+                                              MixedCovariatesTrainingDataset
+from darts.utils.data.inference_dataset import PastCovariatesInferenceDataset, \
+                                                DualCovariatesInferenceDataset, \
+                                                MixedCovariatesInferenceDataset
+from darts.utils.data.utils import CovariateType
+
+# import data formatter
+sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
+from data_formatter.base import *
+
+def get_valid_sampling_locations(target_series: Union[TimeSeries, Sequence[TimeSeries]],
+                                 output_chunk_length: int = 12,
+                                 input_chunk_length: int = 12,
+                                 random_state: Optional[int] = 0,
+                                 max_samples_per_ts: Optional[int] = None):
+    """
+    Get valid sampling indices data for the model.
+
+    Parameters
+    ----------
+    target_series
+        The target time series.
+    output_chunk_length
+        The length of the output chunk.
+    input_chunk_length
+        The length of the input chunk.
+    use_static_covariates
+        Whether to use static covariates.
+    max_samples_per_ts
+        The maximum number of samples per time series.
+    """
+    random.seed(random_state)
+    valid_sampling_locations = {}
+    total_length = input_chunk_length + output_chunk_length
+    for id, series in enumerate(target_series):
+        num_entries = len(series)
+        if num_entries >= total_length:
+            valid_sampling_locations[id] = [i for i in range(num_entries - total_length + 1)]
+    if max_samples_per_ts is not None:
+        updated_sampling_locations = {}
+        for id, locations in valid_sampling_locations.items():
+            if len(locations) > max_samples_per_ts:
+                updated_sampling_locations[id] = random.sample(locations, max_samples_per_ts)
+            else:
+                updated_sampling_locations[id] = locations
+        valid_sampling_locations = updated_sampling_locations
+            
+    return valid_sampling_locations
+
+class SamplingDatasetPast(PastCovariatesTrainingDataset):
+    def __init__(
+        self,
+        target_series: Union[TimeSeries, Sequence[TimeSeries]],
+        covariates: Optional[Union[TimeSeries, Sequence[TimeSeries]]] = None,
+        output_chunk_length: int = 12,
+        input_chunk_length: int = 12,
+        use_static_covariates: bool = True,
+        random_state: Optional[int] = 0,
+        max_samples_per_ts: Optional[int] = None,
+        remove_nan: bool = False,
+    ) -> None:
+        """
+        Parameters
+        ----------
+        target_series
+            One or a sequence of target `TimeSeries`.
+        covariates:
+            Optionally, one or a sequence of `TimeSeries` containing past-observed covariates. If this parameter is set,
+            the provided sequence must have the same length as that of `target_series`. Moreover, all
+            covariates in the sequence must have a time span large enough to contain all the required slices.
+            The joint slicing of the target and covariates is relying on the time axes of both series.
+        output_chunk_length
+            The length of the "output" series emitted by the model
+        input_chunk_length
+            The length of the "input" series fed to the model
+        use_static_covariates
+            Whether to use/include static covariate data from input series.
+        random_state
+            The random state to use for sampling.
+        max_samples_per_ts
+            The maximum number of samples to be drawn from each time series. If None, all samples will be drawn.
+        remove_nan
+            Whether to remove None from the output. E.g. if no covariates are provided, the covariates output will be None 
+            or (optionally) removed from the __getitem__ output.
+        """
+        super().__init__()
+        self.remove_nan = remove_nan
+
+        self.target_series = (
+            [target_series] if isinstance(target_series, TimeSeries) else target_series
+        )
+        self.covariates = (
+            [covariates] if isinstance(covariates, TimeSeries) else covariates
+        )
+
+        # checks
+        raise_if_not(
+            covariates is None or len(self.target_series) == len(self.covariates),
+            "The provided sequence of target series must have the same length as "
+            "the provided sequence of covariate series.",
+        )
+
+        # get valid sampling locations
+        self.valid_sampling_locations = get_valid_sampling_locations(target_series,
+                                                                     output_chunk_length,
+                                                                     input_chunk_length,
+                                                                     random_state,
+                                                                     max_samples_per_ts)
+        
+        # set parameters
+        self.output_chunk_length = output_chunk_length
+        self.input_chunk_length = input_chunk_length
+        self.total_length = input_chunk_length + output_chunk_length
+        self.total_number_samples = sum([len(v) for v in self.valid_sampling_locations.values()])
+        self.use_static_covariates = use_static_covariates
+
+    def __len__(self):
+        """
+        Returns the total number of possible (input, target) splits.
+        """
+        return self.total_number_samples
+
+    def __getitem__(self, idx: int):
+        # get idx of target series
+        target_idx = 0
+        while idx >= len(self.valid_sampling_locations[target_idx]):
+            idx -= len(self.valid_sampling_locations[target_idx])
+            target_idx += 1
+        # get sampling location within the target series
+        sampling_location = self.valid_sampling_locations[target_idx][idx]
+        # get target series
+        target_series = self.target_series[target_idx].values()
+        past_target_series = target_series[sampling_location : sampling_location + self.input_chunk_length]
+        future_target_series = target_series[sampling_location + self.input_chunk_length : sampling_location + self.total_length]
+        # get covariates
+        if self.covariates is not None:
+            covariates = self.covariates[target_idx].values()
+            covariates = covariates[sampling_location : sampling_location + self.input_chunk_length]
+        else:
+            covariates = None
+        # get static covariates
+        if self.use_static_covariates:
+            static_covariates = self.target_series[target_idx].static_covariates_values(copy=True)
+        else:
+            static_covariates = None
+        
+        # return elements that are not None
+        if self.remove_nan:
+            out = []
+            out += [past_target_series] if past_target_series is not None else []
+            out += [covariates] if covariates is not None else []
+            out += [static_covariates] if static_covariates is not None else []
+            out += [future_target_series] if future_target_series is not None else []
+            return tuple(out)
+        else:
+            return tuple([past_target_series, 
+                          covariates, 
+                          static_covariates, 
+                          future_target_series])
+    
+class SamplingDatasetDual(DualCovariatesTrainingDataset):
+    def __init__(
+        self,
+        target_series: Union[TimeSeries, Sequence[TimeSeries]],
+        covariates: Optional[Union[TimeSeries, Sequence[TimeSeries]]] = None,
+        output_chunk_length: int = 12,
+        input_chunk_length: int = 12,
+        use_static_covariates: bool = True,
+        random_state: Optional[int] = 0,
+        max_samples_per_ts: Optional[int] = None,
+        remove_nan: bool = False,
+    ) -> None:
+        """
+        Parameters
+        ----------
+        target_series
+            One or a sequence of target `TimeSeries`.
+        covariates:
+            Optionally, one or a sequence of `TimeSeries` containing future-known covariates. If this parameter is set,
+            the provided sequence must have the same length as that of `target_series`. Moreover, all
+            covariates in the sequence must have a time span large enough to contain all the required slices.
+            The joint slicing of the target and covariates is relying on the time axes of both series.
+        output_chunk_length
+            The length of the "output" series emitted by the model
+        input_chunk_length
+            The length of the "input" series fed to the model
+        use_static_covariates
+            Whether to use/include static covariate data from input series.
+        random_state
+            The random state to use for sampling.
+        max_samples_per_ts
+            The maximum number of samples to be drawn from each time series. If None, all samples will be drawn.
+        remove_nan
+            Whether to remove None from the output. E.g. if no covariates are provided, the covariates output will be None 
+            or (optionally) removed from the __getitem__ output.
+        """
+        super().__init__()
+        self.remove_nan = remove_nan
+
+        self.target_series = (
+            [target_series] if isinstance(target_series, TimeSeries) else target_series
+        )
+        self.covariates = (
+            [covariates] if isinstance(covariates, TimeSeries) else covariates
+        )
+
+        # checks
+        raise_if_not(
+            covariates is None or len(self.target_series) == len(self.covariates),
+            "The provided sequence of target series must have the same length as "
+            "the provided sequence of covariate series.",
+        )
+
+        # get valid sampling locations
+        self.valid_sampling_locations = get_valid_sampling_locations(target_series,
+                                                                     output_chunk_length,
+                                                                     input_chunk_length,
+                                                                     random_state,
+                                                                     max_samples_per_ts,)
+        
+        # set parameters
+        self.output_chunk_length = output_chunk_length
+        self.input_chunk_length = input_chunk_length
+        self.total_length = input_chunk_length + output_chunk_length
+        self.total_number_samples = sum([len(v) for v in self.valid_sampling_locations.values()])
+        self.use_static_covariates = use_static_covariates
+
+    def __len__(self):
+        """
+        Returns the total number of possible (input, target) splits.
+        """
+        return self.total_number_samples
+
+    def __getitem__(self, idx: int):
+        # get idx of target series
+        target_idx = 0
+        while idx >= len(self.valid_sampling_locations[target_idx]):
+            idx -= len(self.valid_sampling_locations[target_idx])
+            target_idx += 1
+        # get sampling location within the target series
+        sampling_location = self.valid_sampling_locations[target_idx][idx]
+        # get target series
+        target_series = self.target_series[target_idx].values()
+        past_target_series = target_series[sampling_location : sampling_location + self.input_chunk_length]
+        future_target_series = target_series[sampling_location + self.input_chunk_length : sampling_location + self.total_length]
+        # get covariates
+        if self.covariates is not None:
+            covariates = self.covariates[target_idx].values()
+            past_covariates = covariates[sampling_location : sampling_location + self.input_chunk_length]
+            future_covariates = covariates[sampling_location + self.input_chunk_length : sampling_location + self.total_length]
+        else:
+            past_covariates = None
+            future_covariates = None
+        # get static covariates
+        if self.use_static_covariates:
+            static_covariates = self.target_series[target_idx].static_covariates_values(copy=True)
+        else:
+            static_covariates = None
+        
+        # return elements that are not None
+        if self.remove_nan:
+            out = []
+            out += [past_target_series] if past_target_series is not None else []
+            out += [past_covariates] if past_covariates is not None else []
+            out += [future_covariates] if future_covariates is not None else []
+            out += [static_covariates] if static_covariates is not None else []
+            out += [future_target_series] if future_target_series is not None else []
+            return tuple(out)
+        else:
+            return tuple([past_target_series, 
+                          past_covariates, 
+                          future_covariates, 
+                          static_covariates, 
+                          future_target_series])
+    
+class SamplingDatasetMixed(MixedCovariatesTrainingDataset):
+    def __init__(
+        self,
+        target_series: Union[TimeSeries, Sequence[TimeSeries]],
+        past_covariates: Optional[Union[TimeSeries, Sequence[TimeSeries]]] = None,
+        future_covariates: Optional[Union[TimeSeries, Sequence[TimeSeries]]] = None,
+        output_chunk_length: int = 12,
+        input_chunk_length: int = 12,
+        use_static_covariates: bool = True,
+        random_state: Optional[int] = 0,
+        max_samples_per_ts: Optional[int] = None,
+        remove_nan: bool = False,
+    ) -> None:
+        """
+        Parameters
+        ----------
+        target_series
+            One or a sequence of target `TimeSeries`.
+        past_covariates
+            Optionally, one or a sequence of `TimeSeries` containing past-observed covariates. If this parameter is set,
+            the provided sequence must have the same length as that of `target_series`. Moreover, all
+            covariates in the sequence must have a time span large enough to contain all the required slices.
+            The joint slicing of the target and covariates is relying on the time axes of both series.
+        future_covariates
+            Optionally, one or a sequence of `TimeSeries` containing future-known covariates. This has to follow
+            the same constraints as `past_covariates`.
+        output_chunk_length
+            The length of the "output" series emitted by the model
+        input_chunk_length
+            The length of the "input" series fed to the model
+        use_static_covariates
+            Whether to use/include static covariate data from input series.
+        random_state
+            The random state to use for sampling.
+        max_samples_per_ts
+            The maximum number of samples to be drawn from each time series. If None, all samples will be drawn.
+        remove_nan
+            Whether to remove None from the output. E.g. if no covariates are provided, the covariates output will be None 
+            or (optionally) removed from the __getitem__ output.
+        """
+        super().__init__()
+        self.remove_nan = remove_nan
+
+        self.target_series = (
+            [target_series] if isinstance(target_series, TimeSeries) else target_series
+        )
+        self.past_covariates = (
+            [past_covariates] if isinstance(past_covariates, TimeSeries) else past_covariates
+        )
+        self.future_covariates = (
+            [future_covariates] if isinstance(future_covariates, TimeSeries) else future_covariates
+        )
+
+        # checks
+        raise_if_not(
+            future_covariates is None or len(self.target_series) == len(self.future_covariates),
+            "The provided sequence of target series must have the same length as "
+            "the provided sequence of covariate series.",
+        )
+        raise_if_not(
+            past_covariates is None or len(self.target_series) == len(self.past_covariates),
+            "The provided sequence of target series must have the same length as "
+            "the provided sequence of covariate series.",
+        )
+
+        # get valid sampling locations
+        self.valid_sampling_locations = get_valid_sampling_locations(target_series,
+                                                                     output_chunk_length,
+                                                                     input_chunk_length,
+                                                                     random_state,
+                                                                     max_samples_per_ts,)
+        
+        # set parameters
+        self.output_chunk_length = output_chunk_length
+        self.input_chunk_length = input_chunk_length
+        self.total_length = input_chunk_length + output_chunk_length
+        self.total_number_samples = sum([len(v) for v in self.valid_sampling_locations.values()])
+        self.use_static_covariates = use_static_covariates
+
+    def __len__(self):
+        """
+        Returns the total number of possible (input, target) splits.
+        """
+        return self.total_number_samples
+
+    def __getitem__(self, idx: int):
+        # get idx of target series
+        target_idx = 0
+        while idx >= len(self.valid_sampling_locations[target_idx]):
+            idx -= len(self.valid_sampling_locations[target_idx])
+            target_idx += 1
+        # get sampling location within the target series
+        sampling_location = self.valid_sampling_locations[target_idx][idx]
+        # get target series
+        target_series = self.target_series[target_idx].values()
+        past_target_series = target_series[sampling_location : sampling_location + self.input_chunk_length]
+        future_target_series = target_series[sampling_location + self.input_chunk_length : sampling_location + self.total_length]
+        # get past covariates
+        if self.past_covariates is not None:
+            past_covariates = self.past_covariates[target_idx].values()
+            past_covariates = past_covariates[sampling_location : sampling_location + self.input_chunk_length]
+        else:
+            past_covariates = None
+        # get future covariates
+        if self.future_covariates is not None:
+            future_covariates = self.future_covariates[target_idx].values()
+            historic_future_covariates = future_covariates[sampling_location : sampling_location + self.input_chunk_length]
+            future_covariates = future_covariates[sampling_location + self.input_chunk_length : sampling_location + self.total_length]
+        else:
+            future_covariates = None
+            historic_future_covariates = None
+        # get static covariates
+        if self.use_static_covariates:
+            static_covariates = self.target_series[target_idx].static_covariates_values(copy=True)
+        else:
+            static_covariates = None
+
+        # return elements that are not None
+        if self.remove_nan:
+            out = []
+            out += [past_target_series] if past_target_series is not None else []
+            out += [past_covariates] if past_covariates is not None else []
+            out += [historic_future_covariates] if historic_future_covariates is not None else []
+            out += [future_covariates] if future_covariates is not None else []
+            out += [static_covariates] if static_covariates is not None else []
+            out += [future_target_series] if future_target_series is not None else []
+            return tuple(out)
+        else:
+            return tuple([past_target_series, 
+                          past_covariates, 
+                          historic_future_covariates, 
+                          future_covariates, 
+                          static_covariates, 
+                          future_target_series])
+
+class SamplingDatasetInferenceMixed(MixedCovariatesInferenceDataset):
+    def __init__(
+        self,
+        target_series: Union[TimeSeries, Sequence[TimeSeries]],
+        past_covariates: Optional[Union[TimeSeries, Sequence[TimeSeries]]] = None,
+        future_covariates: Optional[Union[TimeSeries, Sequence[TimeSeries]]] = None,
+        n: int = 1,
+        input_chunk_length: int = 12,
+        output_chunk_length: int = 1,
+        use_static_covariates: bool = True,
+        random_state: Optional[int] = 0,
+        max_samples_per_ts: Optional[int] = None,
+        array_output_only: bool = False,
+    ):
+        """
+        Parameters
+        ----------
+        target_series
+            One or a sequence of target `TimeSeries`.
+        past_covariates
+            Optionally, one or a sequence of `TimeSeries` containing past-observed covariates. If this parameter is set,
+            the provided sequence must have the same length as that of `target_series`. Moreover, all
+            covariates in the sequence must have a time span large enough to contain all the required slices.
+            The joint slicing of the target and covariates is relying on the time axes of both series.
+        future_covariates
+            Optionally, one or a sequence of `TimeSeries` containing future-known covariates. This has to follow
+            the same constraints as `past_covariates`.
+        n
+            Number of predictions into the future, could be greater than the output chunk length, in which case, the model
+            will be called autorregressively.
+        output_chunk_length
+            The length of the "output" series emitted by the model
+        input_chunk_length
+            The length of the "input" series fed to the model
+        use_static_covariates
+            Whether to use/include static covariate data from input series.
+        random_state
+            The random state to use for sampling.
+        max_samples_per_ts
+            The maximum number of samples to be drawn from each time series. If None, all samples will be drawn.
+        array_output_only
+            Whether __getitem__ returns only the arrays or adds the full `TimeSeries` object to the output tuple
+            This may cause problems with the torch collate and loader functions but works for Darts.
+        """
+        super().__init__(target_series = target_series,
+                         past_covariates = past_covariates,
+                         future_covariates = future_covariates,
+                         n = n,
+                         input_chunk_length = input_chunk_length,
+                         output_chunk_length = output_chunk_length,)
+
+        self.target_series = (
+            [target_series] if isinstance(target_series, TimeSeries) else target_series
+        )
+        self.past_covariates = (
+            [past_covariates] if isinstance(past_covariates, TimeSeries) else past_covariates
+        )
+        self.future_covariates = (
+            [future_covariates] if isinstance(future_covariates, TimeSeries) else future_covariates
+        )
+
+        # checks
+        raise_if_not(
+            future_covariates is None or len(self.target_series) == len(self.future_covariates),
+            "The provided sequence of target series must have the same length as "
+            "the provided sequence of covariate series.",
+        )
+        raise_if_not(
+            past_covariates is None or len(self.target_series) == len(self.past_covariates),
+            "The provided sequence of target series must have the same length as "
+            "the provided sequence of covariate series.",
+        )
+
+        # get valid sampling locations
+        self.valid_sampling_locations = get_valid_sampling_locations(target_series,
+                                                                     output_chunk_length,
+                                                                     input_chunk_length,
+                                                                     random_state,
+                                                                     max_samples_per_ts,)
+        
+        # set parameters
+        self.output_chunk_length = output_chunk_length
+        self.input_chunk_length = input_chunk_length
+        self.total_length = input_chunk_length + output_chunk_length
+        self.total_number_samples = sum([len(v) for v in self.valid_sampling_locations.values()])
+        self.use_static_covariates = use_static_covariates
+        self.array_output_only = array_output_only
+
+    def __len__(self):
+        """
+        Returns the total number of possible (input, target) splits.
+        """
+        return self.total_number_samples
+
+    def __getitem__(self, idx: int):
+        # get idx of target series
+        target_idx = 0
+        while idx >= len(self.valid_sampling_locations[target_idx]):
+            idx -= len(self.valid_sampling_locations[target_idx])
+            target_idx += 1
+        # get sampling location within the target series
+        sampling_location = self.valid_sampling_locations[target_idx][idx]
+        # get target series
+        target_series = self.target_series[target_idx]
+        past_target_series_with_time = target_series[sampling_location : sampling_location + self.input_chunk_length]
+        past_end = past_target_series_with_time.time_index[-1]
+        target_series = self.target_series[target_idx].values()
+        past_target_series = target_series[sampling_location : sampling_location + self.input_chunk_length]
+        # get past covariates
+        if self.past_covariates is not None:
+            past_covariates = self.past_covariates[target_idx].values()
+            past_covariates = past_covariates[sampling_location : sampling_location + self.input_chunk_length]
+            future_past_covariates = past_covariates[sampling_location + self.input_chunk_length : sampling_location + self.total_length]
+        else:
+            past_covariates = None
+            future_past_covariates = None
+        # get future covariates
+        if self.future_covariates is not None:
+            future_covariates = self.future_covariates[target_idx].values()
+            historic_future_covariates = future_covariates[sampling_location : sampling_location + self.input_chunk_length]
+            future_covariates = future_covariates[sampling_location + self.input_chunk_length : sampling_location + self.total_length]
+        else:
+            future_covariates = None
+            historic_future_covariates = None
+        # get static covariates
+        if self.use_static_covariates:
+            static_covariates = self.target_series[target_idx].static_covariates_values(copy=True)
+        else:
+            static_covariates = None
+        # whether to remove Timeseries and None and return only arrays   
+        
+        if self.array_output_only:
+            out = []
+            out += [past_target_series] if past_target_series is not None else []
+            out += [past_covariates] if past_covariates is not None else []
+            out += [historic_future_covariates] if historic_future_covariates is not None else []
+            out += [future_covariates] if future_covariates is not None else []
+            out += [future_past_covariates] if future_past_covariates is not None else []
+            out += [static_covariates] if static_covariates is not None else []
+            return tuple(out)
+        else:
+            return tuple([past_target_series,
+                          past_covariates,
+                          historic_future_covariates,
+                          future_covariates,
+                          future_past_covariates,
+                          static_covariates,
+                          past_target_series_with_time,
+                          past_end + past_target_series_with_time.freq
+                          ])
+
+    def evalsample(
+            self, idx: int
+        ) -> TimeSeries:
+        """
+        Returns the future target series at the given index.
+        """
+        # get idx of target series
+        target_idx = 0
+        while idx >= len(self.valid_sampling_locations[target_idx]):
+            idx -= len(self.valid_sampling_locations[target_idx])
+            target_idx += 1
+        # get sampling location within the target series
+        sampling_location = self.valid_sampling_locations[target_idx][idx]
+        # get target series
+        target_series = self.target_series[target_idx][sampling_location + self.input_chunk_length : sampling_location + self.total_length]
+
+        return target_series
+
+class SamplingDatasetInferencePast(PastCovariatesInferenceDataset):
+    def __init__(
+        self,
+        target_series: Union[TimeSeries, Sequence[TimeSeries]],
+        covariates: Optional[Union[TimeSeries, Sequence[TimeSeries]]] = None,
+        n: int = 1,
+        input_chunk_length: int = 12,
+        output_chunk_length: int = 1,
+        use_static_covariates: bool = True,
+        random_state: Optional[int] = 0,
+        max_samples_per_ts: Optional[int] = None,
+        array_output_only: bool = False,
+    ):
+        """
+        Parameters
+        ----------
+        target_series
+            One or a sequence of target `TimeSeries`.
+        past_covariates
+            Optionally, one or a sequence of `TimeSeries` containing past-observed covariates. If this parameter is set,
+            the provided sequence must have the same length as that of `target_series`. Moreover, all
+            covariates in the sequence must have a time span large enough to contain all the required slices.
+            The joint slicing of the target and covariates is relying on the time axes of both series.
+        n
+            Number of predictions into the future, could be greater than the output chunk length, in which case, the model
+            will be called autorregressively.
+        output_chunk_length
+            The length of the "output" series emitted by the model
+        input_chunk_length
+            The length of the "input" series fed to the model
+        use_static_covariates
+            Whether to use/include static covariate data from input series.
+        random_state
+            The random state to use for sampling.
+        max_samples_per_ts
+            The maximum number of samples to be drawn from each time series. If None, all samples will be drawn.
+        array_output_only
+            Whether __getitem__ returns only the arrays or adds the full `TimeSeries` object to the output tuple
+            This may cause problems with the torch collate and loader functions but works for Darts.
+        """
+        super().__init__(target_series = target_series,
+                         covariates = covariates,
+                         n = n,
+                         input_chunk_length = input_chunk_length,
+                         output_chunk_length = output_chunk_length,)
+
+        self.target_series = (
+            [target_series] if isinstance(target_series, TimeSeries) else target_series
+        )
+        self.covariates = (
+            [covariates] if isinstance(covariates, TimeSeries) else covariates
+        )
+
+        raise_if_not(
+            covariates is None or len(self.target_series) == len(self.covariates),
+            "The provided sequence of target series must have the same length as "
+            "the provided sequence of covariate series.",
+        )
+
+        # get valid sampling locations
+        self.valid_sampling_locations = get_valid_sampling_locations(target_series,
+                                                                     output_chunk_length,
+                                                                     input_chunk_length,
+                                                                     random_state,
+                                                                     max_samples_per_ts,)
+        
+        # set parameters
+        self.output_chunk_length = output_chunk_length
+        self.input_chunk_length = input_chunk_length
+        self.total_length = input_chunk_length + output_chunk_length
+        self.total_number_samples = sum([len(v) for v in self.valid_sampling_locations.values()])
+        self.use_static_covariates = use_static_covariates
+        self.array_output_only = array_output_only
+
+    def __len__(self):
+        """
+        Returns the total number of possible (input, target) splits.
+        """
+        return self.total_number_samples
+
+    def __getitem__(self, idx: int):
+        # get idx of target series
+        target_idx = 0
+        while idx >= len(self.valid_sampling_locations[target_idx]):
+            idx -= len(self.valid_sampling_locations[target_idx])
+            target_idx += 1
+        # get sampling location within the target series
+        sampling_location = self.valid_sampling_locations[target_idx][idx]
+        # get target series
+        target_series = self.target_series[target_idx]
+        past_target_series_with_time = target_series[sampling_location : sampling_location + self.input_chunk_length]
+        past_end = past_target_series_with_time.time_index[-1]
+        target_series = self.target_series[target_idx].values()
+        past_target_series = target_series[sampling_location : sampling_location + self.input_chunk_length]
+        # get past covariates
+        if self.covariates is not None:
+            past_covariates = self.covariates[target_idx].values()
+            past_covariates = past_covariates[sampling_location : sampling_location + self.input_chunk_length]
+            future_past_covariates = past_covariates[sampling_location + self.input_chunk_length : sampling_location + self.total_length]
+        else:
+            past_covariates = None
+            future_past_covariates = None
+        # get static covariates
+        if self.use_static_covariates:
+            static_covariates = self.target_series[target_idx].static_covariates_values(copy=True)
+        else:
+            static_covariates = None
+        # return arrays or arrays with TimeSeries
+        if self.array_output_only:
+            out = []
+            out += [past_target_series] if past_target_series is not None else []
+            out += [past_covariates] if past_covariates is not None else []
+            out += [future_past_covariates] if future_past_covariates is not None else []
+            out += [static_covariates] if static_covariates is not None else []
+            return tuple(out)
+        else:
+            return tuple([past_target_series,
+                          past_covariates,
+                          future_past_covariates,
+                          static_covariates,
+                          past_target_series_with_time,
+                          past_end + past_target_series_with_time.freq])
+
+    def evalsample(
+            self, idx: int
+        ) -> TimeSeries:
+        """
+        Returns the future target series at the given index.
+        """
+        # get idx of target series
+        target_idx = 0
+        while idx >= len(self.valid_sampling_locations[target_idx]):
+            idx -= len(self.valid_sampling_locations[target_idx])
+            target_idx += 1
+        # get sampling location within the target series
+        sampling_location = self.valid_sampling_locations[target_idx][idx]
+        # get target series
+        target_series = self.target_series[target_idx][sampling_location + self.input_chunk_length : sampling_location + self.total_length]
+
+        return target_series
+
+class SamplingDatasetInferenceDual(DualCovariatesInferenceDataset):
+    def __init__(
+        self,
+        target_series: Union[TimeSeries, Sequence[TimeSeries]],
+        covariates: Optional[Union[TimeSeries, Sequence[TimeSeries]]] = None,
+        n: int = 12,
+        input_chunk_length: int = 12,
+        output_chunk_length: int = 1,
+        use_static_covariates: bool = True,
+        random_state: Optional[int] = 0,
+        max_samples_per_ts: Optional[int] = None,
+        array_output_only: bool = False,
+    ):
+        """
+        Parameters
+        ----------
+        target_series
+            One or a sequence of target `TimeSeries`.
+        covariates
+            Optionally, some future-known covariates that are used for predictions. This argument is required
+            if the model was trained with future-known covariates.
+        n
+            Number of predictions into the future, could be greater than the output chunk length, in which case, the model
+            will be called autorregressively.
+        output_chunk_length
+            The length of the "output" series emitted by the model
+        input_chunk_length
+            The length of the "input" series fed to the model
+        use_static_covariates
+            Whether to use/include static covariate data from input series.
+        random_state
+            The random state to use for sampling.
+        max_samples_per_ts
+            The maximum number of samples to be drawn from each time series. If None, all samples will be drawn.
+        array_output_only
+            Whether __getitem__ returns only the arrays or adds the full `TimeSeries` object to the output tuple
+            This may cause problems with the torch collate and loader functions but works for Darts.
+        """
+        super().__init__(target_series = target_series,
+                         covariates = covariates,
+                         n = n,
+                         input_chunk_length = input_chunk_length,
+                         output_chunk_length = output_chunk_length,)
+
+        self.target_series = (
+            [target_series] if isinstance(target_series, TimeSeries) else target_series
+        )
+        self.covariates = (
+            [covariates] if isinstance(covariates, TimeSeries) else covariates
+        )
+
+        raise_if_not(
+            covariates is None or len(self.target_series) == len(self.covariates),
+            "The provided sequence of target series must have the same length as "
+            "the provided sequence of covariate series.",
+        )
+
+        # get valid sampling locations
+        self.valid_sampling_locations = get_valid_sampling_locations(target_series,
+                                                                     output_chunk_length,
+                                                                     input_chunk_length,
+                                                                     random_state,
+                                                                     max_samples_per_ts,)
+        
+        # set parameters
+        self.output_chunk_length = output_chunk_length
+        self.input_chunk_length = input_chunk_length
+        self.total_length = input_chunk_length + output_chunk_length
+        self.total_number_samples = sum([len(v) for v in self.valid_sampling_locations.values()])
+        self.use_static_covariates = use_static_covariates
+        self.array_output_only = array_output_only
+
+    def __len__(self):
+        """
+        Returns the total number of possible (input, target) splits.
+        """
+        return self.total_number_samples
+
+    def __getitem__(self, idx: int):
+        # get idx of target series
+        target_idx = 0
+        while idx >= len(self.valid_sampling_locations[target_idx]):
+            idx -= len(self.valid_sampling_locations[target_idx])
+            target_idx += 1
+        # get sampling location within the target series
+        sampling_location = self.valid_sampling_locations[target_idx][idx]
+        # get target series
+        target_series = self.target_series[target_idx]
+        past_target_series_with_time = target_series[sampling_location : sampling_location + self.input_chunk_length]
+        past_end = past_target_series_with_time.time_index[-1]
+        target_series = self.target_series[target_idx].values()
+        past_target_series = target_series[sampling_location : sampling_location + self.input_chunk_length]
+        # get past covariates
+        if self.covariates is not None:
+            future_covariates = self.covariates[target_idx].values()
+            historic_future_covariates = future_covariates[sampling_location : sampling_location + self.input_chunk_length]
+            future_covariates = future_covariates[sampling_location + self.input_chunk_length : sampling_location + self.total_length]
+        else:
+            historic_future_covariates = None
+            future_covariates = None
+        # get static covariates
+        if self.use_static_covariates:
+            static_covariates = self.target_series[target_idx].static_covariates_values(copy=True)
+        else:
+            static_covariates = None
+        # return arrays or arrays with TimeSeries
+        if self.array_output_only:
+            out = []
+            out += [past_target_series] if past_target_series is not None else []
+            out += [historic_future_covariates] if historic_future_covariates is not None else []
+            out += [future_covariates] if future_covariates is not None else []
+            out += [static_covariates] if static_covariates is not None else []
+            return tuple(out)
+        else:
+            return tuple([past_target_series,
+                            historic_future_covariates,
+                            future_covariates,
+                            static_covariates,
+                            past_target_series_with_time,
+                            past_end + past_target_series_with_time.freq,])
+        
+    def evalsample(
+            self, idx: int
+        ) -> TimeSeries:
+        """
+        Returns the future target series at the given index.
+        """
+        # get idx of target series
+        target_idx = 0
+        while idx >= len(self.valid_sampling_locations[target_idx]):
+            idx -= len(self.valid_sampling_locations[target_idx])
+            target_idx += 1
+        # get sampling location within the target series
+        sampling_location = self.valid_sampling_locations[target_idx][idx]
+        # get target series
+        target_series = self.target_series[target_idx][sampling_location + self.input_chunk_length : sampling_location + self.total_length]
+
+        return target_series
+        

utils/darts_evaluation.py

@@ -0,0 +1,280 @@
+import sys
+import os
+import yaml
+import random
+from typing import Any, BinaryIO, Callable, Dict, List, Optional, Sequence, Tuple, Union
+
+import numpy as np 
+from scipy import stats
+import pandas as pd
+import darts
+
+from darts import models
+from darts import metrics
+from darts import TimeSeries
+
+# import data formatter
+sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
+from data_formatter.base import *
+from utils.darts_processing import *
+
+def _get_values(
+    series: TimeSeries, stochastic_quantile: Optional[float] = 0.5
+) -> np.ndarray:
+    """
+    Returns the numpy values of a time series.
+    For stochastic series, return either all sample values with (stochastic_quantile=None) or the quantile sample value
+    with (stochastic_quantile {>=0,<=1})
+    """
+    if series.is_deterministic:
+        series_values = series.univariate_values()
+    else:  # stochastic
+        if stochastic_quantile is None:
+            series_values = series.all_values(copy=False)
+        else:
+            series_values = series.quantile_timeseries(
+                quantile=stochastic_quantile
+            ).univariate_values()
+    return series_values
+
+def _get_values_or_raise(
+    series_a: TimeSeries,
+    series_b: TimeSeries,
+    intersect: bool,
+    stochastic_quantile: Optional[float] = 0.5,
+    remove_nan_union: bool = False,
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Returns the processed numpy values of two time series. Processing can be customized with arguments
+    `intersect, stochastic_quantile, remove_nan_union`.
+
+    Raises a ValueError if the two time series (or their intersection) do not have the same time index.
+
+    Parameters
+    ----------
+    series_a
+        A univariate deterministic ``TimeSeries`` instance (the actual series).
+    series_b
+        A univariate (deterministic or stochastic) ``TimeSeries`` instance (the predicted series).
+    intersect
+        A boolean for whether or not to only consider the time intersection between `series_a` and `series_b`
+    stochastic_quantile
+        Optionally, for stochastic predicted series, return either all sample values with (`stochastic_quantile=None`)
+        or any deterministic quantile sample values by setting `stochastic_quantile=quantile` {>=0,<=1}.
+    remove_nan_union
+        By setting `remove_non_union` to True, remove all indices from `series_a` and `series_b` which have a NaN value
+        in either of the two input series.
+    """
+    series_a_common = series_a.slice_intersect(series_b) if intersect else series_a
+    series_b_common = series_b.slice_intersect(series_a) if intersect else series_b
+
+    series_a_det = _get_values(series_a_common, stochastic_quantile=stochastic_quantile)
+    series_b_det = _get_values(series_b_common, stochastic_quantile=stochastic_quantile)
+
+    if not remove_nan_union:
+        return series_a_det, series_b_det
+
+    b_is_deterministic = bool(len(series_b_det.shape) == 1)
+    if b_is_deterministic:
+        isnan_mask = np.logical_or(np.isnan(series_a_det), np.isnan(series_b_det))
+    else:
+        isnan_mask = np.logical_or(
+            np.isnan(series_a_det), np.isnan(series_b_det).any(axis=2).flatten()
+        )
+    return np.delete(series_a_det, isnan_mask), np.delete(
+        series_b_det, isnan_mask, axis=0
+    )
+
+def rescale_and_backtest(series: Union[TimeSeries, 
+                                    Sequence[TimeSeries]],
+                         forecasts: Union[TimeSeries, 
+                                            Sequence[TimeSeries],
+                                            Sequence[Sequence[TimeSeries]]], 
+                         metric: Union[
+                                    Callable[[TimeSeries, TimeSeries], float],
+                                    List[Callable[[TimeSeries, TimeSeries], float]],
+                                ], 
+                         scaler: Callable[[TimeSeries], TimeSeries] = None,
+                         reduction: Union[Callable[[np.ndarray], float], None] = np.mean,
+                         likelihood: str = "GaussianMean",
+                         cal_thresholds: Optional[np.ndarray] = np.linspace(0, 1, 11),
+                        ):
+    """
+    Backtest the historical forecasts (as provided by Darts) on the series.
+
+    Parameters
+    ----------
+    series
+        The target time series.
+    forecasts
+        The forecasts.
+    scaler
+        The scaler used to scale the series.
+    metric
+        The metric or metrics to use for backtesting.
+    reduction
+        The reduction to apply to the metric.
+    likelihood
+        The likelihood to use for evaluating the model.
+    cal_thresholds
+        The thresholds to use for computing the calibration error.
+
+    Returns
+    -------
+    np.ndarray
+        Error array. If the reduction is none, array is of shape (n, p)
+        where n is the total number of samples (forecasts) and p is the number of metrics.
+        If the reduction is not none, array is of shape (k, p), where k is the number of series. 
+    float
+        The estimated log-likelihood of the model on the data.
+    np.ndarray
+        The ECE for each time point in the forecast.
+    """
+    series = [series] if isinstance(series, TimeSeries) else series
+    forecasts = [forecasts] if isinstance(forecasts, TimeSeries) else forecasts
+    metric = [metric] if not isinstance(metric, list) else metric
+
+    # compute errors: 1) reverse scaling forecasts and true values, 2)compute errors
+    backtest_list = []
+    for idx in range(len(series)):
+        if scaler is not None:
+            series[idx] = scaler.inverse_transform(series[idx])
+            forecasts[idx] = [scaler.inverse_transform(f) for f in forecasts[idx]]
+        errors = [
+            [metric_f(series[idx], f) for metric_f in metric]
+            if len(metric) > 1
+            else metric[0](series[idx], f)
+            for f in forecasts[idx]
+        ]
+        if reduction is None:
+            backtest_list.append(np.array(errors))
+        else:
+            backtest_list.append(reduction(np.array(errors), axis=0))
+    backtest_list = np.vstack(backtest_list)
+    
+    if likelihood == "GaussianMean":
+        # compute likelihood
+        est_var = []
+        for idx, target_ts in enumerate(series):
+            est_var += [metrics.mse(target_ts, f) for f in forecasts[idx]]
+        est_var = np.mean(est_var)
+        forecast_len = forecasts[0][0].n_timesteps
+        log_likelihood = -0.5*forecast_len - 0.5*np.log(2*np.pi*est_var)
+
+        # compute calibration error: 1) cdf values 2) compute calibration error
+        # compute the cdf values
+        cdf_vals = []
+        for idx in range(len(series)):
+            for forecast in forecasts[idx]:
+                y_true, y_pred = _get_values_or_raise(series[idx], 
+                                                      forecast, 
+                                                      intersect=True, 
+                                                      remove_nan_union=True)
+                y_true, y_pred = y_true.flatten(), y_pred.flatten()
+                cdf_vals.append(stats.norm.cdf(y_true, loc=y_pred, scale=np.sqrt(est_var)))
+        cdf_vals = np.vstack(cdf_vals)
+        # compute the prediction calibration
+        cal_error = np.zeros(forecasts[0][0].n_timesteps)
+        for p in cal_thresholds:
+            est_p = (cdf_vals <= p).astype(float)
+            est_p = np.mean(est_p, axis=0)
+            cal_error += (est_p - p) ** 2
+
+    return backtest_list, log_likelihood, cal_error
+
+def rescale_and_test(series: Union[TimeSeries, 
+                                   Sequence[TimeSeries]],
+                    forecasts: Union[TimeSeries, 
+                                    Sequence[TimeSeries]], 
+                    metric: Union[
+                            Callable[[TimeSeries, TimeSeries], float],
+                            List[Callable[[TimeSeries, TimeSeries], float]],
+                        ], 
+                    scaler: Callable[[TimeSeries], TimeSeries] = None,
+                    likelihood: str = "GaussianMean",
+                    cal_thresholds: Optional[np.ndarray] = np.linspace(0, 1, 11),
+                ):
+    """
+    Test the forecasts on the series.
+
+    Parameters
+    ----------
+    series
+        The target time series.
+    forecasts
+        The forecasts.
+    scaler
+        The scaler used to scale the series.
+    metric
+        The metric or metrics to use for backtesting.
+    reduction
+        The reduction to apply to the metric.
+    likelihood
+        The likelihood to use for evaluating the likelihood and calibration of model.
+    cal_thresholds
+        The thresholds to use for computing the calibration error.
+
+    Returns
+    -------
+    np.ndarray
+        Error array. If the reduction is none, array is of shape (n, p)
+        where n is the total number of samples (forecasts) and p is the number of metrics.
+        If the reduction is not none, array is of shape (k, p), where k is the number of series. 
+    float
+        The estimated log-likelihood of the model on the data.
+    np.ndarray
+        The ECE for each time point in the forecast.
+    """
+    series = [series] if isinstance(series, TimeSeries) else series
+    forecasts = [forecasts] if isinstance(forecasts, TimeSeries) else forecasts
+    metric = [metric] if not isinstance(metric, list) else metric
+
+    # compute errors: 1) reverse scaling forecasts and true values, 2)compute errors
+    series = scaler.inverse_transform(series)
+    forecasts = scaler.inverse_transform(forecasts)
+    errors = [
+        [metric_f(t, f) for metric_f in metric]
+        if len(metric) > 1
+        else metric[0](t, f)
+        for (t, f) in zip(series, forecasts)
+        ]
+    errors = np.array(errors)
+
+    if likelihood == "GaussianMean":        
+        # compute likelihood
+        est_var = [metrics.mse(t, f) for (t, f) in zip(series, forecasts)]
+        est_var = np.mean(est_var)
+        forecast_len = forecasts[0].n_timesteps
+        log_likelihood = -0.5*forecast_len - 0.5*np.log(2*np.pi*est_var)
+
+        # compute calibration error: 1) cdf values 2) compute calibration error
+        # compute the cdf values
+        cdf_vals = []
+        for t, f in zip(series, forecasts):
+            t, f = _get_values_or_raise(t, f, intersect=True, remove_nan_union=True)
+            t, f = t.flatten(), f.flatten()
+            cdf_vals.append(stats.norm.cdf(t, loc=f, scale=np.sqrt(est_var)))
+        cdf_vals = np.vstack(cdf_vals)
+        # compute the prediction calibration
+        cal_error = np.zeros(forecasts[0].n_timesteps)
+        for p in cal_thresholds:
+            est_p = (cdf_vals <= p).astype(float)
+            est_p = np.mean(est_p, axis=0)
+            cal_error += (est_p - p) ** 2
+
+    if likelihood == "Quantile":
+        # no likelihood since we don't have a parametric model
+        log_likelihood = 0
+
+        # compute calibration error: 1) get quantiles 2) compute calibration error
+        cal_error = np.zeros(forecasts[0].n_timesteps)
+        for p in cal_thresholds:
+            est_p = 0
+            for t, f in zip(series, forecasts):
+                q = f.quantile(p)
+                t, q = _get_values_or_raise(t, q, intersect=True, remove_nan_union=True)
+                t, q = t.flatten(), q.flatten()
+                est_p += (t <= q).astype(float)
+            est_p = (est_p / len(series)).flatten()
+            cal_error += (est_p - p) ** 2
+        
+    return errors, log_likelihood, cal_error

utils/darts_processing.py

@@ -0,0 +1,367 @@
+import sys
+import os
+import yaml
+import random
+from typing import Any, BinaryIO, Callable, Dict, List, Optional, Sequence, Tuple, Union
+from pathlib import Path
+import numpy as np 
+from scipy import stats
+import pandas as pd
+import darts
+
+from darts import models
+from darts import metrics
+from darts import TimeSeries
+from darts.dataprocessing.transformers import Scaler
+from pytorch_lightning.callbacks import Callback
+from sympy import pprint
+
+# import data formatter
+sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
+from data_formatter.base import *
+
+pd.set_option('display.width', None)   # Set display width to None to avoid truncation
+pd.set_option('display.max_columns', None)  # Display all columns
+
+def make_series(data: Dict[str, pd.DataFrame],
+                time_col: str,
+                group_col: str,
+                value_cols: Dict[str, List[str]],
+                include_sid: bool = False,
+                verbose: bool = False
+                ) -> Dict[str, darts.TimeSeries]:
+    """
+    Makes the TimeSeries from the data.
+
+    Parameters
+    ----------
+    data
+        dict of train, val, test dataframes
+    time_col
+        name of time column
+    group_col
+        name of group column
+    value_cols
+        dict with key specifying the type of covariate and value specifying the list of columns.
+    include_sid
+        whether to include segment id as static covariate
+    
+    Returns
+    -------
+    series: Dict[str, Dict[str, darts.TimeSeries]]
+        dict of train, val, test splits of target and covariates TimeSeries objects
+    scalers: Dict[str, darts.preprocessing.Scaler]
+        dict of scalers for target and covariates
+    """
+    series = {i: {j: None for j in value_cols} for i in data.keys()}
+    scalers = {}
+    for key, df in data.items():
+        
+        for name, cols in value_cols.items():
+            # Adjust display settings
+            if verbose:
+                print(f"DATAFRAME for key {key} in NAME {name} and COLS {cols} and GROUP_COL {group_col}")
+                pprint(df.head(1))
+            series[key][name] = TimeSeries.from_group_dataframe(df = df,
+                                                                group_cols = group_col,
+                                                                time_col = time_col,
+                                                                value_cols = cols) if cols is not None else None
+            if series[key][name] is not None and include_sid is False:
+                for i in range(len(series[key][name])):
+                    series[key][name][i] = series[key][name][i].with_static_covariates(None)
+            if cols is not None: 
+                if key == 'train':
+                    scalers[name] = ScalerCustom()
+                    series[key][name] = scalers[name].fit_transform(series[key][name])
+                else:
+                    series[key][name] = scalers[name].transform(series[key][name])
+            else:
+                scalers[name] = None
+    return series, scalers
+
+def load_data(url: str,
+              config_path: Path,
+              use_covs: bool = False,
+              cov_type: str = 'past',
+              use_static_covs: bool = False, seed = 0):
+    """
+    Load data according to the specified config file and covert to Darts TimeSeries objects.
+
+    Parameters
+    ----------
+    seed: int
+        Random seed for data splitting.
+    study_file: str
+        Path to the study file.
+    dataset: str
+        Name of the dataset.
+    use_covs: bool
+        Whether to use covariates.
+    cov_type: str
+        Type of covariates to use. Can be 'past' or 'mixed' or 'dual'.
+    use_static_covs: bool
+        Whether to use static covariates.
+    
+    Returns
+    -------
+    formatter: DataFormatter
+        Data formatter object.
+    series: Dict[str, Dict[str, TimeSeries]]
+       First dictionary specified the split, second dictionary specifies the type of series (target or covariate).
+    scalers: Dict[str, Scaler]
+       Dictionary of scalers with key indicating the type of series (target or covariate).
+    """
+
+
+    """
+    config={
+        'data_csv_path':f'{url}',
+        'drop': None,
+        'ds_name': 'livia_mini',
+        'index_col': -1,
+        'observation_interval': '5min',
+        'column_definition': {
+            {'data_type': 'categorical',
+             'input_type':'id',
+             'name':'id'
+             },
+            {'date_type':'date',
+             'input_type':'time',
+             'name':'time'
+             },
+            {'date_type':'real_valued',
+             'input_type':'target',
+             'name':'gl'
+             }
+        },
+        'encoding_params':{'date':['day','month','year','hour','minute','second']
+                           },
+        'nan_vals':None,
+        'interpolation_params':{'gap_threshold': 45,
+                                'min_drop_length': 240
+                                },
+        'scaling_params':{'scaler':None
+                          },
+        'split_params':{'length_segment': 13,
+                        'random_state':seed,
+                        'test_percent_subjects': 0.1
+                        },
+        'max_length_input': 192,
+        'length_pred': 12,
+        'params':{
+            'gluformer':{'in_len': 96,
+                         'd_model': 512,
+                         'n_heads': 10,
+                         'd_fcn': 1024,
+                         'num_enc_layers': 2,
+                         'num_dec_layers': 2,
+                         'length_pred': 12
+                         }
+        }
+    }
+    """
+    with config_path.open("r") as f:
+        config = yaml.safe_load(f)
+    config["data_csv_path"] = url
+
+    formatter = DataFormatter(config)
+    #assert dataset is not None, 'dataset must be specified in the load_data call'
+    assert use_covs is not None, 'use_covs must be specified in the load_data call'
+
+    # convert to series
+    time_col = formatter.get_column('time')
+    group_col = formatter.get_column('sid')
+    target_col = formatter.get_column('target')
+    static_cols = formatter.get_column('static_covs')
+    static_cols = static_cols + [formatter.get_column('id')] if static_cols is not None else [formatter.get_column('id')]
+    dynamic_cols = formatter.get_column('dynamic_covs')
+    future_cols = formatter.get_column('future_covs')
+
+    data = {'train': formatter.train_data,
+            'val': formatter.val_data,
+            'test': formatter.test_data.loc[~formatter.test_data.index.isin(formatter.test_idx_ood)],
+            'test_ood': formatter.test_data.loc[formatter.test_data.index.isin(formatter.test_idx_ood)]}
+    value_cols = {'target': target_col,
+                  'static': static_cols,
+                  'dynamic': dynamic_cols,
+                  'future': future_cols}
+    # build series
+    series, scalers = make_series(data,
+                                    time_col,
+                                    group_col,
+                                    value_cols)
+    if not use_covs:
+        # set dynamic and future covariates to None
+        for split in ['train', 'val', 'test', 'test_ood']:
+            for cov in ['dynamic', 'future']:
+                series[split][cov] = None
+    elif use_covs and cov_type == 'mixed':
+        pass # this is the default for make_series()
+    elif use_covs and cov_type == 'past':
+        # use future covariates as dynamic (past) covariates
+        if series['train']['dynamic'] is None:
+            for split in ['train', 'val', 'test', 'test_ood']:
+                series[split]['dynamic'] = series[split]['future']
+        else:
+            for split in ['train', 'val', 'test', 'test_ood']:
+                for i in range(len(series[split]['future'])):
+                    series[split]['dynamic'][i] = series[split]['dynamic'][i].concatenate(series[split]['future'][i], axis=1)
+        # erase future covariates
+        for split in ['train', 'val', 'test', 'test_ood']:
+            series[split]['future'] = None
+    elif use_covs and cov_type == 'dual':
+        # erase dynamic (past) covariates
+        for split in ['train', 'val', 'test', 'test_ood']:
+            series[split]['dynamic'] = None
+    
+    if use_static_covs:
+        # attach static covariates to series
+        for split in ['train', 'val', 'test', 'test_ood']:
+            for i in range(len(series[split]['target'])):
+                static_covs = series[split]['static'][i][0].pd_dataframe()
+                series[split]['target'][i] = series[split]['target'][i].with_static_covariates(static_covs)
+
+    return formatter, series, scalers
+
+def reshuffle_data(formatter: DataFormatter, 
+                   seed: int = 0, 
+                   use_covs: bool = None,
+                   cov_type: str = 'past',
+                   use_static_covs: bool = False,):
+    """
+    Reshuffle data according to the seed and covert to Darts TimeSeries objects.
+    
+    Parameters
+    ----------
+    formatter: DataFormatter
+        Data formatter object containing the data
+    seed: int
+        Random seed for data splitting.
+    use_covs: bool
+        Whether to use covariates.
+    cov_type: str
+        Type of covariates to use. Can be 'past' or 'mixed' or 'dual'.
+    use_static_covs: bool
+        Whether to use static covariates.
+    
+    Returns
+    -------
+    formatter: DataFormatter
+        Reshuffled data formatter object.
+    series: Dict[str, Dict[str, TimeSeries]]
+         First dictionary specified the split, second dictionary specifies the type of series (target or covariate).
+    scalers: Dict[str, Scaler]
+        Dictionary of scalers with key indicating the type of series (target or covariate).
+    """
+    # reshuffle
+    formatter.reshuffle(seed)
+    assert use_covs is not None, 'use_covs must be specified in the reshuffle_data call'
+
+    # convert to series
+    time_col = formatter.get_column('time')
+    group_col = formatter.get_column('sid')
+    target_col = formatter.get_column('target')
+    static_cols = formatter.get_column('static_covs')
+    static_cols = static_cols + [formatter.get_column('id')] if static_cols is not None else [formatter.get_column('id')]
+    dynamic_cols = formatter.get_column('dynamic_covs')
+    future_cols = formatter.get_column('future_covs')
+
+    # build series
+    series, scalers = make_series({'train': formatter.train_data,
+                                    'val': formatter.val_data,
+                                    'test': formatter.test_data.loc[~formatter.test_data.index.isin(formatter.test_idx_ood)],
+                                    'test_ood': formatter.test_data.loc[formatter.test_data.index.isin(formatter.test_idx_ood)]},
+                                    time_col,
+                                    group_col,
+                                    {'target': target_col,
+                                    'static': static_cols,
+                                    'dynamic': dynamic_cols,
+                                    'future': future_cols})
+    
+    if not use_covs:
+        # set dynamic and future covariates to None
+        for split in ['train', 'val', 'test', 'test_ood']:
+            for cov in ['dynamic', 'future']:
+                series[split][cov] = None
+    elif use_covs and cov_type == 'past':
+        # use future covariates as dynamic covariates
+        if series['train']['dynamic'] is None:
+            for split in ['train', 'val', 'test', 'test_ood']:
+                series[split]['dynamic'] = series[split]['future']
+        # or attach them to dynamic covariates
+        else:
+            for split in ['train', 'val', 'test', 'test_ood']:
+                for i in range(len(series[split]['future'])):
+                    series[split]['dynamic'][i] = series[split]['dynamic'][i].concatenate(series[split]['future'][i], axis=1)
+    elif use_covs and cov_type == 'dual':
+        # set dynamic covariates to None, because they are not supported
+        for split in ['train', 'val', 'test', 'test_ood']:
+            series[split]['dynamic'] = None
+    
+    if use_static_covs:
+        # attach static covariates to series
+        for split in ['train', 'val', 'test', 'test_ood']:
+            for i in range(len(series[split]['target'])):
+                static_covs = series[split]['static'][i][0].pd_dataframe()
+                series[split]['target'][i] = series[split]['target'][i].with_static_covariates(static_covs)
+    
+    return formatter, series, scalers
+
+class ScalerCustom:
+    '''
+    Min-max scaler for TimeSeries that fits on all sequences simultaenously.
+    Default Darts scaler fits one scaler per sequence in the list.
+
+    Attributes
+    ----------
+    scaler: Scaler
+        Darts scaler object.
+    min_: np.ndarray
+        Per feature adjustment for minimum (see Scikit-learn).
+    scale_: np.ndarray
+        Per feature relative scaling of the data (see Scikit-learn).
+    '''
+    def __init__(self):
+        self.scaler = Scaler()
+        self.min_ = None
+        self.scale_ = None 
+
+    def fit(self, time_series: Union[List[TimeSeries], TimeSeries]) -> None:
+        
+        if isinstance(time_series, list):
+
+            # extract series as Pandas dataframe
+            df = pd.concat([ts.pd_dataframe() for ts in time_series])
+            value_cols = df.columns
+            df.reset_index(inplace=True)
+            # create new equally spaced time grid
+            df['new_time'] = pd.date_range(start=df['time'].min(), periods=len(df), freq='1h')
+            # fit scaler
+            series = TimeSeries.from_dataframe(df, time_col='new_time', value_cols=value_cols)
+            series = self.scaler.fit(series)
+        else:
+            series = self.scaler.fit(time_series)
+        # extract min and scale
+        self.min_ = self.scaler._fitted_params[0].min_
+        self.scale_ = self.scaler._fitted_params[0].scale_
+
+    def transform(self, time_series: Union[List[TimeSeries], TimeSeries]) -> Union[List[TimeSeries], TimeSeries]:
+        if isinstance(time_series, list):
+            # transform one by one
+            series = [self.scaler.transform(ts) for ts in time_series]
+        else:
+            series = self.scaler.transform(time_series)
+        return series
+    
+    def inverse_transform(self, time_series: Union[List[TimeSeries], TimeSeries]) -> Union[List[TimeSeries], TimeSeries]:
+        if isinstance(time_series, list):
+            # transform one by one
+            series = [self.scaler.inverse_transform(ts) for ts in time_series]
+        else:
+            series = self.scaler.inverse_transform(time_series)
+        return series
+    
+    def fit_transform(self, time_series: Union[List[TimeSeries], TimeSeries]) -> Union[List[TimeSeries], TimeSeries]:
+        self.fit(time_series)
+        series = self.transform(time_series)
+        return series

utils/darts_training.py

@@ -0,0 +1,114 @@
+import sys
+import os
+import yaml
+import random
+from typing import Any, BinaryIO, Callable, Dict, List, Optional, Sequence, Tuple, Union
+
+import numpy as np 
+from scipy import stats
+import pandas as pd
+import darts
+
+from darts import models
+from darts import metrics
+from darts import TimeSeries
+from pytorch_lightning.callbacks import Callback
+from darts.logging import get_logger, raise_if_not
+
+# for optuna callback
+import warnings
+import optuna
+from optuna.storages._cached_storage import _CachedStorage
+from optuna.storages._rdb.storage import RDBStorage
+# Define key names of `Trial.system_attrs`.
+_PRUNED_KEY = "ddp_pl:pruned"
+_EPOCH_KEY = "ddp_pl:epoch"
+with optuna._imports.try_import() as _imports:
+    import pytorch_lightning as pl
+    from pytorch_lightning import LightningModule
+    from pytorch_lightning import Trainer
+    from pytorch_lightning.callbacks import Callback
+if not _imports.is_successful():
+    Callback = object  # type: ignore  # NOQA
+    LightningModule = object  # type: ignore  # NOQA
+    Trainer = object  # type: ignore  # NOQA
+
+def print_callback(study, trial, study_file=None):
+    # write output to a file
+    with open(study_file, "a") as f:
+        f.write(f"Current value: {trial.value}, Current params: {trial.params}\n")
+        f.write(f"Best value: {study.best_value}, Best params: {study.best_trial.params}\n")
+
+def early_stopping_check(study, 
+                         trial, 
+                         study_file, 
+                         early_stopping_rounds=10):
+      """
+      Early stopping callback for Optuna.
+      This function checks the current trial number and the best trial number.
+      """
+      current_trial_number = trial.number
+      best_trial_number = study.best_trial.number
+      should_stop = (current_trial_number - best_trial_number) >= early_stopping_rounds
+      if should_stop:
+          with open(study_file, 'a') as f:
+              f.write('\nEarly stopping at trial {} (best trial: {})'.format(current_trial_number, best_trial_number))
+          study.stop()
+
+class LossLogger(Callback):
+    def __init__(self):
+        self.train_loss = []
+        self.val_loss = []
+
+    def on_train_epoch_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
+        self.train_loss.append(float(trainer.callback_metrics["train_loss"]))
+
+    def on_validation_epoch_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
+        self.val_loss.append(float(trainer.callback_metrics["val_loss"]))
+
+class PyTorchLightningPruningCallback(Callback):
+    """PyTorch Lightning callback to prune unpromising trials.
+    See `the example <https://github.com/optuna/optuna-examples/blob/
+    main/pytorch/pytorch_lightning_simple.py>`__
+    if you want to add a pruning callback which observes accuracy.
+    Args:
+        trial:
+            A :class:`~optuna.trial.Trial` corresponding to the current evaluation of the
+            objective function.
+        monitor:
+            An evaluation metric for pruning, e.g., ``val_loss`` or
+            ``val_acc``. The metrics are obtained from the returned dictionaries from e.g.
+            ``pytorch_lightning.LightningModule.training_step`` or
+            ``pytorch_lightning.LightningModule.validation_epoch_end`` and the names thus depend on
+            how this dictionary is formatted.
+    """
+
+    def __init__(self, trial: optuna.trial.Trial, monitor: str) -> None:
+        super().__init__()
+
+        self._trial = trial
+        self.monitor = monitor
+
+    def on_validation_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
+        # When the trainer calls `on_validation_end` for sanity check,
+        # do not call `trial.report` to avoid calling `trial.report` multiple times
+        # at epoch 0. The related page is
+        # https://github.com/PyTorchLightning/pytorch-lightning/issues/1391.
+        if trainer.sanity_checking:
+            return
+
+        epoch = pl_module.current_epoch
+
+        current_score = trainer.callback_metrics.get(self.monitor)
+        if current_score is None:
+            message = (
+                "The metric '{}' is not in the evaluation logs for pruning. "
+                "Please make sure you set the correct metric name.".format(self.monitor)
+            )
+            warnings.warn(message)
+            return
+
+        self._trial.report(current_score, step=epoch)
+        if self._trial.should_prune():
+            message = "Trial was pruned at epoch {}.".format(epoch)
+            raise optuna.TrialPruned(message)