class_docstr: Memory Augmented Time Series Dataset
  
  Args:
      handler (DataHandler): data handler
      segments (dict): data split segments
      seq_len (int): time series sequence length
      horizon (int): label horizon
      num_states (int): how many memory states to be added
      memory_mode (str): memory mode (daily or sample)
      batch_size (int): batch size (<0 will use daily sampling)
      n_samples (int): number of samples in the same day
      shuffle (bool): whether shuffle data
      drop_last (bool): whether drop last batch < batch_size
      input_size (int): reshape flatten rows as this input_size (backward compatibility)
class_docstr_tok: ["Memory","Augmented","Time","Series","Dataset","Args",":","handler","(","DataHandler",")",":","data","handler","segments","(","dict",")",":","data","split","segments","seq_len","(","int",")",":","time","series","sequence","length","horizon","(","int",")",":","label","horizon","num_states","(","int",")",":","how","many","memory","states","to","be","added","memory_mode","(","str",")",":","memory","mode","(","daily","or","sample",")","batch_size","(","int",")",":","batch","size","(","<","0","will","use","daily","sampling",")","n_samples","(","int",")",":","number","of","samples","in","the","same","day","shuffle","(","bool",")",":","whether","shuffle","data","drop_last","(","bool",")",":","whether","drop","last","batch","<","batch_size","input_size","(","int",")",":","reshape","flatten","rows","as","this","input_size","(","backward","compatibility",")"]
code: def setup_data(self, handler_kwargs: dict = None, **kwargs):
      super().setup_data(**kwargs)
  
      if handler_kwargs is not None:
          self.handler.setup_data(**handler_kwargs)
  
      # pre-fetch data and change index to <code, date>
      # NOTE: we will use inplace sort to reduce memory use
      try:
          df = (
              self.handler._learn.copy()
          )  # use copy otherwise recorder will fail
          # FIXME: currently we cannot support switching from `_learn` to `_infer` for inference
      except Exception:
          warnings.warn("cannot access `_learn`, will load raw data")
          df = self.handler._data.copy()
      df.index = df.index.swaplevel()
      df.sort_index(inplace=True)
  
      # convert to numpy
      self._data = df["feature"].values.astype("float32")
      np.nan_to_num(
          self._data, copy=False
      )  # NOTE: fillna in case users forget using the fillna processor
      self._label = df["label"].squeeze().values.astype("float32")
      self._index = df.index
  
      if (
          self.input_size is not None
          and self.input_size != self._data.shape[1]
      ):
          warnings.warn(
              "the data has different shape from input_size and the data will be reshaped"
          )
          assert (
              self._data.shape[1] % self.input_size == 0
          ), "data mismatch, please check `input_size`"
  
      # create batch slices
      self._batch_slices = _create_ts_slices(self._index, self.seq_len)
  
      # create daily slices
      daily_slices = {
          date: [] for date in sorted(self._index.unique(level=1))
      }  # sorted by date
      for i, (code, date) in enumerate(self._index):
          daily_slices[date].append(self._batch_slices[i])
      self._daily_slices = np.array(
          list(daily_slices.values()), dtype="object"
      )
      self._daily_index = pd.Series(
          list(daily_slices.keys())
      )  # index is the original date index
  
      # add memory (sample wise and daily)
      if self.memory_mode == "sample":
          self._memory = np.zeros(
              (len(self._data), self.num_states), dtype=np.float32
          )
      elif self.memory_mode == "daily":
          self._memory = np.zeros(
              (len(self._daily_index), self.num_states),
              dtype=np.float32,
          )
      else:
          raise ValueError(f"invalid memory_mode `{self.memory_mode}`")
  
      # padding tensor
      self._zeros = np.zeros(
          (self.seq_len, max(self.num_states, self._data.shape[1])),
          dtype=np.float32,
      )
  
imports: import copy
  import torch
  import warnings
  import numpy
  import pandas
  from qlib.data.dataset import DatasetH
label_desc: Use image node_id 2 for calling the MTSDatasetH obj's underlying member method code with example usage: obj.setup_data(handler_kwargs) without return types