text: #!/usr/bin/env python
  """
  Tree Tracing
  
  Here we trace a tree generated by Bayesian teaching
  of a matrix of size (10*10) or (20*20)
  """
  
  import os
  
  import pickle
  
  from datetime import datetime as dt
  import numpy as np
  
  import torch
  # import torch.multiprocessing as mp
  import multiprocessing as mp
  
  CURDIR = os.path.abspath(os.curdir)
  # os.chdir("../Stability/")
  import sinkhorn_torch as sk
  # os.chdir(CURDIR)
  
  FL = torch.float64
  LOG_FILES = {}
  
  class TreeTracing:
      """
      Tracing rooted tree branched by choosing different data to teach.
      The tree is a rooted d+1 valence tree. Level can be infinite, but
      we fix a max depth
  
      Strategy is to do a 2-layer calculation, use multiprocessing in the second
  
      Theoretically, all the matrices and priors are instances
      of torch.Tensor
      """
  
      @staticmethod
      def gen_perturb_default(mat_teach, prior_teach, *args):
          """
          Default Perturbations
          """
          n_row, n_col = mat_teach.shape
          perturb = args[0]
          mat_learn = mat_teach.clone()
          prior_learn = prior_teach.clone()
          index_row, index_col = np.random.choice(n_row), np.random.choice(n_col)
          mat_learn[index_row, index_col] += perturb
          prior_learn[np.random.choice(n_col)] += perturb
          # We need to perturb it in some way
  
          mat_learn = sk.col_normalize(mat_learn, torch.ones(n_col, dtype=FL))
          prior_learn /= torch.sum(prior_learn)
          return mat_learn, prior_learn
  
  
      def __init__(self, init_depth=2,
                   max_depth=5, name="test",
                   prefix=CURDIR, **args):
          """
          Initialization
          """
          self.correct_hypo = None
          arg_processor = {"device": self.set_device,
                           "mat_teach": self.set_mat_teach,
                           "mat_learn": self.set_mat_learn,
                           "prior_teach": self.set_prior_teach,
                           "prior_learn": self.set_prior_learn,
                           "correct_hypo": self.set_correct_hypo,}
          self.set_depth(init_depth, max_depth)
          self.set_prefix(prefix+"/"+name)
          for key, value in args.items():
              arg_processor[key](value)
  
          # self.log_files = LOG_FILES
          # self.log_files is a collection of file handles, structure:
          # log_files = {branch_id, [list of length=self.max_depth+1, all file handles ]}
  
      def set_depth(self, init_depth, max_depth):
          """
          Set init_depth and max_depth
          """
          self.init_depth, self.max_depth = init_depth, max_depth
  
  
      def set_prefix(self, prefix):
          """
          set prefix and timestamp
          """
          self.prefix = prefix
          self.timestamp = dt.today().strftime("%Y-%m-%d_%H:%M:%S.%f")
  
      def set_mat_teach(self, mat_teach):
          """
          Set teaching matrix (accurate one)
          """
          self.mat_teach = mat_teach
          self.n_row, self.n_col = mat_teach.shape
  
  
      def set_mat_learn(self, mat_learn):
          """
          Set learning matrix (accurate one)
          """
          self.mat_learn = mat_learn
  
      def set_prior_learn(self, prior_learn):
          """
          Set learning prior (accurate one)
          """
          self.prior_learn = prior_learn
  
      def set_prior_teach(self, prior_teach):
          """
          Set teaching prior (accurate one)
          """
          self.prior_teach = prior_teach
  
      def set_correct_hypo(self, correct_hypo):
          """
          Set the correct hypothesis
          """
          self.correct_hypo = correct_hypo
  
      def set_device(self, device):
          """
          Set Device
          """
          self.device = device
  
      def init_default(self, n_row=10, n_col=10, perturb=0.1,
                       generate_perturbation=gen_perturb_default.__func__):
          """
          Initialization
          """
          self.n_row, self.n_col = n_row, n_col
          mat_teach = torch.distributions.dirichlet.Dirichlet(torch.ones([n_row, n_col],
                                                                         dtype=FL)).sample().T
          prior_teach = torch.distributions.dirichlet.Dirichlet(torch.ones(n_row, dtype=FL)).sample()
  
          mat_learn, prior_learn = generate_perturbation(mat_teach, prior_teach, perturb)
  
          with open(self.prefix+"_setup.log", "wb") as file_ptr:
              pickle.dump({"mat_teach": mat_teach,
                           "mat_learn": mat_learn,
                           "prior_teach": prior_teach,
                           "prior_learn": prior_learn},
                          file_ptr)
          self.set_correct_hypo(0)
          self.set_mat_teach(mat_teach)
          self.set_mat_learn(mat_learn)
          self.set_prior_teach(prior_teach)
          self.set_prior_learn(prior_learn)
  
      def init_log(self, branch_id, start_depth, target_depth, write_head=False):
          """
          In the first round : write_head=True
          others : write_head=False
  
          In fact, each process may have different log_files, so just work with their own, no harm
          """
          LOG_FILES[branch_id] = [None, ] * (self.max_depth + 1)
          for layer in range(start_depth+(0 if write_head else 1), target_depth+1):
              try:
                  LOG_FILES[branch_id][layer] = \
                      open(self.prefix+"_branch_"+str(branch_id)+"_layer_"+str(layer)+".log", "wb")
              except IOError:
                  print("Error in open file at branch", branch_id, "layer", layer)
  
      def finish_log(self, branch_id):
          """
          Close all files in a branch, when finish a whole branch
          """
          for handle in LOG_FILES[branch_id]:
              if handle is not None:
                  handle.close()
          # del LOG_FILES[branch_id]
          print(list(map(lambda f: (f.closed if f is not None else None), LOG_FILES[branch_id])))
          print("Branch", branch_id, "finished.")
  
      def write_log(self, branch_id, layer, data):
          """
          Safely write data to correct position
          """
          # assert branch_id in self.log_files.keys()
          # if we do things correct, no need to check this
  
          # pickle.dump(data, self.log_files[branch_id][layer])
          pickle.dump(data, LOG_FILES[branch_id][layer])
          # print("Data written to branch:", branch_id, "\t layer:", layer)
          # this is also for logs... we may use os.async to guarantee written.
  
      def single_round(self, root_data):
          """
          Single round entry in multiprocessing
          may do some root node job before calling self.dfs()
          root_data = {"branch_id": id of current branch, 0 is with grand root (layer one),
                       "start_depth": root = 0,
                       "target_depth": init_depth or max_depth,
                       "data": {"prob_bi": ,
                                "prob_scbi": ,
                                "prior_teach": current priors,
                                "prior_learn":   ,
                                "prior_bayes":   ,},
                       "write_head": whether branch root is written
          }
          """
          self.init_log(root_data["branch_id"],
                        root_data["start_depth"],
                        root_data["target_depth"],
                        root_data["write_head"])
  
          self.dfs(root_data["start_depth"],
                   root_data["target_depth"],
                   root_data["data"],
                   root_data["branch_id"],
                   root_data["write_head"])
  
          self.finish_log(root_data["branch_id"])
  
          return root_data["branch_id"]
  
      def dfs(self, current_depth, target_depth, data, branch_id, write=True):
          """
          Depth-First-Search
          save data at first when reaching this node
          data = {"prob_scbi": probability of getting here through scbi,
                  "prob_bi": probability of bi
                  "teach": prior_teach
                  "learn":
                  "bayes":}
          """
          if write:
              self.write_log(branch_id, current_depth, data)
  
          if current_depth == target_depth:
              return
  
          mat_teach = sk.sinkhorn_torch(self.mat_teach, col_sum=data["teach"]*self.n_row)
          mat_learn = sk.sinkhorn_torch(self.mat_learn, col_sum=data["learn"]*self.n_row)
          mat_bayes = sk.row_normalize(data["bayes"]*self.mat_teach.clone(),
                                       torch.ones(self.n_row, dtype=FL))
          for teach_d in range(self.n_row):
              prob_bi = data["prob_bi"] * self.mat_teach[teach_d, self.correct_hypo]
              prob_scbi = data["prob_scbi"] * mat_teach[teach_d, self.correct_hypo]
  
              self.dfs(current_depth+1,
                       target_depth,
                       {"prob_scbi": prob_scbi.item(),
                        "prob_bi": prob_bi.item(),
                        "teach": mat_teach[teach_d],
                        "learn": mat_learn[teach_d],
                        "bayes": mat_bayes[teach_d]},
                       branch_id)
  
      def main(self, num_branch=40):
          """
          Main entry
          """
          if self.correct_hypo is None:
              self.init_default()
  
          # >>> Layer 1 DFS
          root_datum = {"branch_id"   : 0,
                        "start_depth" : 0,
                        "target_depth": self.init_depth,
                        "data"        : {"prob_scbi": 1.,
                                         "prob_bi"  : 1.,
                                         "teach"    : self.prior_teach,
                                         "learn"    : self.prior_learn,
                                         "bayes"    : self.prior_teach,},
                        "write_head"  : True,}
          self.single_round(root_datum)
          # <<< Layer 1 DFS
  
          # >>> Layer 2 DFS
          root_data = []
          with open(self.prefix+"_branch_0_layer_"+str(self.init_depth)+".log",
                    "rb") as file_ptr:
              for branch_id in range(1, int(self.n_row ** self.init_depth) + 1):
                  root_data += [{"branch_id"   : branch_id,
                                 "start_depth" : self.init_depth,
                                 "target_depth": self.max_depth,
                                 "data"        : pickle.load(file_ptr),
                                 "write_head"  : False}, ]
          # possible error: we trust that file_ptr has just exact data we need!
  
          # print(root_data)
  
          # pool.map(self.test, range(10))
          for i in range(int(len(root_data)/num_branch)+1):
              pool = mp.Pool(num_branch)
              pool.map(self.single_round, root_data[i * num_branch : min((i+1) * num_branch, len(root_data))])
              pool.close()
          # <<< Layer 2 DFS
  
  
  if __name__ == '__main__':
  
      MODEL = TreeTracing(2, 5, name="20by20")
      MODEL.init_default(20, 20, 0.05)
      MODEL.main(10)
  