text: #!/usr/bin/env python
  # -*- coding: utf-8 -*-
  # Python version: 3.6
  
  
  import numpy as np
  from torchvision import datasets, transforms
  
  
  def mnist_iid(dataset, num_users):
      """
      Sample I.I.D. client data from MNIST dataset
      :param dataset:
      :param num_users:
      :return: dict of image index
      """
      num_items = int(len(dataset)/num_users)
      dict_users, all_idxs = {}, [i for i in range(len(dataset))]
      for i in range(num_users):
          dict_users[i] = set(np.random.choice(all_idxs, num_items,
                                               replace=False))
          all_idxs = list(set(all_idxs) - dict_users[i])
      return dict_users
  
  
  def mnist_noniid(dataset, num_users):
      """
      Sample non-I.I.D client data from MNIST dataset
      :param dataset:
      :param num_users:
      :return:
      """
      # 60,000 training imgs -->  200 imgs/shard X 300 shards
      num_shards, num_imgs = 200, 300
      idx_shard = [i for i in range(num_shards)]
      dict_users = {i: np.array([]) for i in range(num_users)}
      idxs = np.arange(num_shards*num_imgs)
      labels = dataset.train_labels.numpy()
  
      # sort labels
      idxs_labels = np.vstack((idxs, labels))
      idxs_labels = idxs_labels[:, idxs_labels[1, :].argsort()]
      idxs = idxs_labels[0, :]
  
      # divide and assign 2 shards/client
      for i in range(num_users):
          rand_set = set(np.random.choice(idx_shard, 2, replace=False))
          idx_shard = list(set(idx_shard) - rand_set)
          for rand in rand_set:
              dict_users[i] = np.concatenate(
                  (dict_users[i], idxs[rand*num_imgs:(rand+1)*num_imgs]), axis=0)
      return dict_users
  
  
  def mnist_noniid_unequal(dataset, num_users):
      """
      Sample non-I.I.D client data from MNIST dataset s.t clients
      have unequal amount of data
      :param dataset:
      :param num_users:
      :returns a dict of clients with each clients assigned certain
      number of training imgs
      """
      # 60,000 training imgs --> 50 imgs/shard X 1200 shards
      num_shards, num_imgs = 1200, 50
      idx_shard = [i for i in range(num_shards)]
      dict_users = {i: np.array([]) for i in range(num_users)}
      idxs = np.arange(num_shards*num_imgs)
      labels = dataset.train_labels.numpy()
  
      # sort labels
      idxs_labels = np.vstack((idxs, labels))
      idxs_labels = idxs_labels[:, idxs_labels[1, :].argsort()]
      idxs = idxs_labels[0, :]
  
      # Minimum and maximum shards assigned per client:
      min_shard = 1
      max_shard = 30
  
      # Divide the shards into random chunks for every client
      # s.t the sum of these chunks = num_shards
      random_shard_size = np.random.randint(min_shard, max_shard+1,
                                            size=num_users)
      random_shard_size = np.around(random_shard_size /
                                    sum(random_shard_size) * num_shards)
      random_shard_size = random_shard_size.astype(int)
  
      # Assign the shards randomly to each client
      if sum(random_shard_size) > num_shards:
  
          for i in range(num_users):
              # First assign each client 1 shard to ensure every client has
              # atleast one shard of data
              rand_set = set(np.random.choice(idx_shard, 1, replace=False))
              idx_shard = list(set(idx_shard) - rand_set)
              for rand in rand_set:
                  dict_users[i] = np.concatenate(
                      (dict_users[i], idxs[rand*num_imgs:(rand+1)*num_imgs]),
                      axis=0)
  
          random_shard_size = random_shard_size-1
  
          # Next, randomly assign the remaining shards
          for i in range(num_users):
              if len(idx_shard) == 0:
                  continue
              shard_size = random_shard_size[i]
              if shard_size > len(idx_shard):
                  shard_size = len(idx_shard)
              rand_set = set(np.random.choice(idx_shard, shard_size,
                                              replace=False))
              idx_shard = list(set(idx_shard) - rand_set)
              for rand in rand_set:
                  dict_users[i] = np.concatenate(
                      (dict_users[i], idxs[rand*num_imgs:(rand+1)*num_imgs]),
                      axis=0)
      else:
  
          for i in range(num_users):
              shard_size = random_shard_size[i]
              rand_set = set(np.random.choice(idx_shard, shard_size,
                                              replace=False))
              idx_shard = list(set(idx_shard) - rand_set)
              for rand in rand_set:
                  dict_users[i] = np.concatenate(
                      (dict_users[i], idxs[rand*num_imgs:(rand+1)*num_imgs]),
                      axis=0)
  
          if len(idx_shard) > 0:
              # Add the leftover shards to the client with minimum images:
              shard_size = len(idx_shard)
              # Add the remaining shard to the client with lowest data
              k = min(dict_users, key=lambda x: len(dict_users.get(x)))
              rand_set = set(np.random.choice(idx_shard, shard_size,
                                              replace=False))
              idx_shard = list(set(idx_shard) - rand_set)
              for rand in rand_set:
                  dict_users[k] = np.concatenate(
                      (dict_users[k], idxs[rand*num_imgs:(rand+1)*num_imgs]),
                      axis=0)
  
      return dict_users
  
  
  def cifar_iid(dataset, num_users):
      """
      用户抽取训练数据 CIFAR10 I.I.D.
      Sample I.I.D. client data from CIFAR10 dataset
      :param dataset: 训练集
      :param num_users: 用户数量
      :return: dict of image index  字典，key是用户编号比如1到10，value是一个set，是数据集中的数据的编号，如 3，14，17...
      """
      num_items = int(len(dataset)/num_users)  # 用户平均获取的数据集大小
      dict_users, all_idxs = {}, [i for i in range(len(dataset))]  # all_idxs是数据集每个数据的位置
      for i in range(num_users):  # 对每个用户
          dict_users[i] = set(np.random.choice(all_idxs, num_items,
                                               replace=False))  # 从all_idxs（所有数据的位置）选择num_items（用户平均数量）个，replace=False不取重复值
          all_idxs = list(set(all_idxs) - dict_users[i])  # 总集合-选完的集合，避免重复
      return dict_users
  
  
  def cifar_noniid(dataset, num_users):
      """
      Sample non-I.I.D client data from CIFAR10 dataset
      用户抽取训练数据 CIFAR10 non-I.I.D
      :param dataset: 训练集
      :param num_users: 用户数量
      :return:
      """
      # 50,000 training imgs -->  200 imgs/shard X 250 shards   五万张图片分为200个分片，每个分片有250张图片
      num_shards, num_imgs = 200, 250
      idx_shard = [i for i in range(num_shards)]  # idx_shard [0,1,2,3...,199]
      dict_users = {i: np.array([]) for i in range(num_users)}  # {0:[],1:[],2:[],...,9:[]}
      idxs = np.arange(num_shards*num_imgs)  # [0,1,2,...,200*250]
      # labels = dataset.train_labels.numpy()
      labels = np.array(dataset.targets)  # []标签
  
      # sort labels
      idxs_labels = np.vstack((idxs, labels))  # [[位置],[标签]]，如[[0,1,2,3],["d","c","a","b"]]
      idxs_labels = idxs_labels[:, idxs_labels[1, :].argsort()]  # argsort对二维数组按照labels进行了排序，idxs也随着标签移动了[[2,3,1,0],["a","b","c","d"]]，把相同标签的聚在一起，分片时保证了非独立同分布
      idxs = idxs_labels[0, :]  # 位置取出来，即[2,3,1,0]
  
      # divide and assign
      for i in range(num_users):  # 对每个用户：
          rand_set = set(np.random.choice(idx_shard, 2, replace=False))  # idx_shard 200里面选择2个
          idx_shard = list(set(idx_shard) - rand_set)  # idx_shard 减去选择过的值
          for rand in rand_set:
              dict_users[i] = np.concatenate(  # concatenate合并
                  (dict_users[i], idxs[rand*num_imgs:(rand+1)*num_imgs]), axis=0)  # 取两个分片，分片内几乎都是同一类的数据
      return dict_users  # 字典，key是用户编号，value是数据在数据集的位置，这是非独立同分布
  
  
  if __name__ == '__main__':
      dataset_train = datasets.MNIST('./data/mnist/', train=True, download=True,
                                     transform=transforms.Compose([
                                         transforms.ToTensor(),
                                         transforms.Normalize((0.1307,),
                                                              (0.3081,))
                                     ]))
      num = 100
      d = mnist_noniid(dataset_train, num)
  