code: def KMeans_inside_dataset(
      num_ts_=1,
      num_groups=16,
      context_length=72,
      prediction_length=24,
      file_name="default",
  ):
      dataset = get_dataset("traffic")
      dataset_group = [[] for i in range(num_groups)]
      whole_data = []
      ret = dict()
      it = iter(dataset.train)
      # num_ts = int(dataset.metadata.feat_static_cat[0].cardinality)
      num_ts = num_ts_
      len_sample = context_length + prediction_length
      index = 0
      feature = torch.Tensor([])
      for i in range(num_ts):
          train_entry = next(it)
          target = train_entry["target"]
  
          for ts_sample_start in range(
              0, len(target) - len_sample, prediction_length
          ):
              ts_slice = target[
                  ts_sample_start : ts_sample_start + len_sample
              ]
              feature = torch.cat(
                  (
                      feature,
                      torch.Tensor(
                          [
                              ts_slice.mean(),
                              ts_slice.var(),
                              index % 7,
                              index // 90,
                          ]
                      ),
                  )
              )
              index += 1
      feature = feature.reshape(index, 4)
      feature = _get_pre_features(feature).contiguous()
      # print(feature)
      # import pdb;pdb.set_trace()
      cl, c = KMeans(feature, num_groups)
      it = iter(dataset.train)
      sample_id = 0
      for i in range(num_ts):
          train_entry = next(it)
          target = train_entry["target"]
          unsplit_start = train_entry["start"]
          for ts_sample_start in range(
              0, len(target) - len_sample, prediction_length
          ):
              ts_slice = target[
                  ts_sample_start : ts_sample_start + len_sample
              ]
              gid = cl[sample_id]
              dataset_group[gid].append(
                  {
                      "target": ts_slice,
                      "start": unsplit_start,
                      "feat_static_cat": train_entry["feat_static_cat"],
                  }
              )
              whole_data.append(
                  {
                      "target": ts_slice,
                      "start": unsplit_start,
                      "feat_static_cat": train_entry["feat_static_cat"],
                  }
              )
              unsplit_start += pd.Timedelta(hours=prediction_length)
              sample_id += 1
      print(len(whole_data))
      ret["group_ratio"] = [
          len(i) / len(whole_data) for i in dataset_group
      ]
      print(ret["group_ratio"])
      random.shuffle(whole_data)
      ret["whole_data"] = ListDataset(
          whole_data, freq=dataset.metadata.freq
      )
      group_data_list = []
      for group in dataset_group:
          random.shuffle(group)
          group_data_list.append(
              ListDataset(group, freq=dataset.metadata.freq)
          )
      ret["group_data"] = group_data_list
      print("write whole data")
      with open(
          "synthetic_" + file_name + "_whole_data.csv", "wb"
      ) as output:
          pickle.dump(ret["whole_data"], output)
      print("write group data")
      with open(
          "synthetic_" + file_name + "_group_data.csv", "wb"
      ) as output:
          pickle.dump(ret, output)
      return True
  
file_path: gluonts/src/gluonts/nursery/SCott/dataset_tools/algo_clustering.py
imports: from pts.dataset.repository.datasets import get_dataset, dataset_recipes
  from pts.dataset.utils import to_pandas
  import numpy
  import pandas
  from pts.dataset import ListDataset
  import torch
  from pts.model.simple_feedforward import SimpleFeedForwardEstimator
  import pickle
  import json
  import random
  import time
  from matplotlib import pyplot
  from pykeops.torch import LazyTensor
label_desc: Use image node_id 3 for calling a global function with example usage: KMeans_inside_dataset(num_ts_, num_groups, context_length, prediction_length, file_name) and returns: True