code: # %% [markdown]
  # # Imports
  import json
  import os
  import warnings
  from operator import itemgetter
  from pathlib import Path
  
  import matplotlib.pyplot as plt
  import numpy as np
  import pandas as pd
  import seaborn as sns
  from joblib import Parallel, delayed
  from joblib.parallel import Parallel, delayed
  from sklearn.metrics import adjusted_rand_score
  import networkx as nx
  
  from graspy.cluster import GaussianCluster, AutoGMMCluster
  from graspy.embed import AdjacencySpectralEmbed, OmnibusEmbed
  from graspy.models import DCSBMEstimator, SBMEstimator
  from graspy.plot import heatmap, pairplot
  from graspy.utils import binarize, cartprod, get_lcc, pass_to_ranks
  from src.data import load_everything
  from src.utils import export_skeleton_json, savefig
  from src.visualization import clustergram, palplot, sankey
  from src.hierarchy import signal_flow
  
  warnings.simplefilter("ignore", category=FutureWarning)
  
  
  FNAME = os.path.basename(__file__)[:-3]
  print(FNAME)
  
  
  # %% [markdown]
  # # Parameters
  BRAIN_VERSION = "2019-12-09"
  GRAPH_TYPES = ["Gad", "Gaa", "Gdd", "Gda"]
  GRAPH_TYPE_LABELS = [r"A $\to$ D", r"A $\to$ A", r"D $\to$ D", r"D $\to$ A"]
  N_GRAPH_TYPES = len(GRAPH_TYPES)
  
  SAVEFIGS = True
  DEFAULT_FMT = "png"
  DEFUALT_DPI = 150
  
  SAVESKELS = False
  
  MIN_CLUSTERS = 8
  MAX_CLUSTERS = 8
  N_INIT = 50
  PTR = True
  ONLY_RIGHT = True
  
  embed = "LSE"
  cluster = "GMM"
  n_components = 4
  if cluster == "GMM":
      gmm_params = {"n_init": N_INIT, "covariance_type": "all"}
  elif cluster == "AutoGMM":
      gmm_params = {"max_agglom_size": None}
  
  np.random.seed(23409857)
  
  
  def stashfig(name, **kws):
      if SAVEFIGS:
          savefig(name, foldername=FNAME, fmt=DEFAULT_FMT, dpi=DEFUALT_DPI, **kws)
  
  
  def stashskel(name, ids, colors, palette=None, **kws):
      if SAVESKELS:
          return export_skeleton_json(
              name, ids, colors, palette=palette, foldername=FNAME, **kws
          )
  
  
  def ase(adj, n_components):
      if PTR:
          adj = pass_to_ranks(adj)
      ase = AdjacencySpectralEmbed(n_components=n_components)
      latent = ase.fit_transform(adj)
      latent = np.concatenate(latent, axis=-1)
      return latent
  
  
  def to_laplace(graph, form="DAD", regularizer=None):
      r"""
      A function to convert graph adjacency matrix to graph laplacian. 
      Currently supports I-DAD, DAD, and R-DAD laplacians, where D is the diagonal
      matrix of degrees of each node raised to the -1/2 power, I is the 
      identity matrix, and A is the adjacency matrix.
      
      R-DAD is regularized laplacian: where :math:`D_t = D + regularizer*I`.
      Parameters
      ----------
      graph: object
          Either array-like, (n_vertices, n_vertices) numpy array,
          or an object of type networkx.Graph.
      form: {'I-DAD' (default), 'DAD', 'R-DAD'}, string, optional
          
          - 'I-DAD'
              Computes :math:`L = I - D*A*D`
          - 'DAD'
              Computes :math:`L = D*A*D`
          - 'R-DAD'
              Computes :math:`L = D_t*A*D_t` where :math:`D_t = D + regularizer*I`
      regularizer: int, float or None, optional (default=None)
          Constant to be added to the diagonal of degree matrix. If None, average 
          node degree is added. If int or float, must be >= 0. Only used when 
          ``form`` == 'R-DAD'.
      Returns
      -------
      L: numpy.ndarray
          2D (n_vertices, n_vertices) array representing graph 
          laplacian of specified form
      References
      ----------
      .. [1] Qin, Tai, and Karl Rohe. "Regularized spectral clustering
             under the degree-corrected stochastic blockmodel." In Advances
             in Neural Information Processing Systems, pp. 3120-3128. 2013
      """
      valid_inputs = ["I-DAD", "DAD", "R-DAD"]
      if form not in valid_inputs:
          raise TypeError("Unsuported Laplacian normalization")
  
      A = graph
  
      in_degree = np.sum(A, axis=0)
      out_degree = np.sum(A, axis=1)
  
      # regularize laplacian with parameter
      # set to average degree
      if form == "R-DAD":
          if regularizer is None:
              regularizer = 1
          elif not isinstance(regularizer, (int, float)):
              raise TypeError(
                  "Regularizer must be a int or float, not {}".format(type(regularizer))
              )
          elif regularizer < 0:
              raise ValueError("Regularizer must be greater than or equal to 0")
          regularizer = regularizer * np.mean(out_degree)
  
          in_degree += regularizer
          out_degree += regularizer
  
      with np.errstate(divide="ignore"):
          in_root = 1 / np.sqrt(in_degree)  # this is 10x faster than ** -0.5
          out_root = 1 / np.sqrt(out_degree)
  
      in_root[np.isinf(in_root)] = 0
      out_root[np.isinf(out_root)] = 0
  
      in_root = np.diag(in_root)  # just change to sparse diag for sparse support
      out_root = np.diag(out_root)
  
      if form == "I-DAD":
          L = np.diag(in_degree) - A
          L = in_root @ L @ in_root
      elif form == "DAD" or form == "R-DAD":
          L = out_root @ A @ in_root
      # return symmetrize(L, method="avg")  # sometimes machine prec. makes this necessary
      return L
  
  
  def lse(adj, n_components, regularizer=None):
      if PTR:
          adj = pass_to_ranks(adj)
      lap = to_laplace(adj, form="R-DAD")
      ase = AdjacencySpectralEmbed(n_components=n_components)
      latent = ase.fit_transform(lap)
      latent = np.concatenate(latent, axis=-1)
      return latent
  
  
  def omni(adjs, n_components):
      if PTR:
          adjs = [pass_to_ranks(a) for a in adjs]
      omni = OmnibusEmbed(n_components=n_components // len(adjs))
      latent = omni.fit_transform(adjs)
      latent = np.concatenate(latent, axis=-1)  # first is for in/out
      latent = np.concatenate(latent, axis=-1)  # second is for concat. each graph
      return latent
  
  
  def ase_concatenate(adjs, n_components):
      if PTR:
          adjs = [pass_to_ranks(a) for a in adjs]
      ase = AdjacencySpectralEmbed(n_components=n_components // len(adjs))
      graph_latents = []
      for a in adjs:
          latent = ase.fit_transform(a)
          latent = np.concatenate(latent, axis=-1)
          graph_latents.append(latent)
      latent = np.concatenate(graph_latents, axis=-1)
      return latent
  
  
  def sub_ari(known_inds, true_labels, pred_labels):
      true_known_labels = true_labels[known_inds]
      pred_known_labels = pred_labels[known_inds]
      ari = adjusted_rand_score(true_known_labels, pred_known_labels)
      return ari
  
  
  # Set up plotting constants
  plt.style.use("seaborn-white")
  sns.set_palette("deep")
  sns.set_context("talk", font_scale=1)
  
  
  # %% [markdown]
  # # Load the data
  
  
  adj, class_labels, side_labels, skeleton_labels = load_everything(
      "Gad",
      version=BRAIN_VERSION,
      return_keys=["Merge Class", "Hemisphere"],
      return_ids=True,
  )
  
  
  # select the right hemisphere
  if ONLY_RIGHT:
      side = "right hemisphere"
      right_inds = np.where(side_labels == "R")[0]
      adj = adj[np.ix_(right_inds, right_inds)]
      class_labels = class_labels[right_inds]
      skeleton_labels = skeleton_labels[right_inds]
  else:
      side = "full brain"
  
  # sort by number of synapses
  degrees = adj.sum(axis=0) + adj.sum(axis=1)
  sort_inds = np.argsort(degrees)[::-1]
  adj = adj[np.ix_(sort_inds, sort_inds)]
  class_labels = class_labels[sort_inds]
  skeleton_labels = skeleton_labels[sort_inds]
  
  # remove disconnected nodes
  adj, lcc_inds = get_lcc(adj, return_inds=True)
  class_labels = class_labels[lcc_inds]
  skeleton_labels = skeleton_labels[lcc_inds]
  
  # remove pendants
  degrees = np.count_nonzero(adj, axis=0) + np.count_nonzero(adj, axis=1)
  not_pendant_mask = degrees != 1
  not_pendant_inds = np.array(range(len(degrees)))[not_pendant_mask]
  adj = adj[np.ix_(not_pendant_inds, not_pendant_inds)]
  class_labels = class_labels[not_pendant_inds]
  skeleton_labels = skeleton_labels[not_pendant_inds]
  
  # plot degree sequence
  d_sort = np.argsort(degrees)[::-1]
  degrees = degrees[d_sort]
  plt.figure(figsize=(10, 5))
  sns.scatterplot(x=range(len(degrees)), y=degrees, s=30, linewidth=0)
  
  known_inds = np.where(class_labels != "Unk")[0]
  
  
  # %% [markdown]
  # # Run clustering using LSE on the sum graph
  
  n_verts = adj.shape[0]
  
  
  latent = lse(adj, n_components, regularizer=None)
  pairplot(latent, labels=class_labels, title=embed)
  
  k_list = list(range(MIN_CLUSTERS, MAX_CLUSTERS + 1))
  n_runs = len(k_list)
  out_dicts = []
  
  bin_adj = binarize(adj)
  
  last_pred_labels = np.zeros(n_verts)
  
  if cluster == "GMM":
      ClusterModel = GaussianCluster
  elif cluster == "AutoGMM":
      ClusterModel = AutoGMMCluster
  
  for k in k_list:
      run_name = f"k = {k}, {cluster}, {embed}, {side} (A to D), PTR, raw"
      print(run_name)
      print()
  
      # Do clustering
      # TODO: make this autogmm instead
      gmm = ClusterModel(min_components=k, max_components=k, **gmm_params)
      gmm.fit(latent)
      pred_labels = gmm.predict(latent)
  
      # Score unsupervised metrics
      base_dict = {
          "K": k,
          "Cluster": cluster,
          "Embed": embed,
          "Method": f"{cluster} o {embed}",
      }
  
      # GMM likelihood
      score = gmm.model_.score(latent)
      temp_dict = base_dict.copy()
      temp_dict["Metric"] = "GMM likelihood"
      temp_dict["Score"] = score
      out_dicts.append(temp_dict)
  
      # GMM BIC
      score = gmm.model_.bic(latent)
      temp_dict = base_dict.copy()
      temp_dict["Metric"] = "GMM BIC"
      temp_dict["Score"] = score
      out_dicts.append(temp_dict)
  
      # SBM likelihood
      sbm = SBMEstimator(directed=True, loops=False)
      sbm.fit(bin_adj, y=pred_labels)
      score = sbm.score(bin_adj)
      temp_dict = base_dict.copy()
      temp_dict["Metric"] = "SBM likelihood"
      temp_dict["Score"] = score
      out_dicts.append(temp_dict)
  
      # DCSBM likelihood
      dcsbm = DCSBMEstimator(directed=True, loops=False)
      dcsbm.fit(bin_adj, y=pred_labels)
      score = dcsbm.score(bin_adj)
      temp_dict = base_dict.copy()
      temp_dict["Metric"] = "DCSBM likelihood"
      temp_dict["Score"] = score
      out_dicts.append(temp_dict)
  
      # ARI of the subset with labels
      score = sub_ari(known_inds, class_labels, pred_labels)
      temp_dict = base_dict.copy()
      temp_dict["Metric"] = "Simple ARI"
      temp_dict["Score"] = score
      out_dicts.append(temp_dict)
  
      # ARI vs K - 1
      score = adjusted_rand_score(last_pred_labels, pred_labels)
      temp_dict = base_dict.copy()
      temp_dict["Metric"] = "K-1 ARI"
      temp_dict["Score"] = score
      out_dicts.append(temp_dict)
      last_pred_labels = pred_labels
  
      save_name = f"k{k}-{cluster}-{embed}-right-ad-PTR-raw"
  
      # Plot embedding
      # pairplot(latent, labels=pred_labels, title=run_name)
      # stashfig("latent-" + save_name)
  
      # Plot everything else
      clustergram(adj, class_labels, pred_labels)
      stashfig("clustergram-" + save_name)
  
      # New plot
      # - Compute signal flow
      # - Get the centroid of each cluster and project to 1d
      # - Alternatively, just take the first dimension
      # - For each cluster plot as a node
  
      # output skeletons
      if SAVESKELS:
          _, colormap, pal = stashskel(
              save_name, skeleton_labels, pred_labels, palette="viridis", multiout=True
          )
  
          palplot(k, cmap="viridis")
          stashfig("palplot-" + save_name)
  
          # save dict colormapping
          filename = (
              Path("./maggot_models/notebooks/outs")
              / Path(FNAME)
              / str("colormap-" + save_name + ".json")
          )
          with open(filename, "w") as fout:
              json.dump(colormap, fout)
  
          stashskel(
              save_name, skeleton_labels, pred_labels, palette="viridis", multiout=False
          )
  
  # %% [markdown]
  # #  Plot results of unsupervised metrics
  
  result_df = pd.DataFrame(out_dicts)
  fg = sns.FacetGrid(result_df, col="Metric", col_wrap=3, sharey=False, height=4)
  fg.map(sns.lineplot, "K", "Score")
  stashfig(f"metrics-{cluster}-{embed}-right-ad-PTR-raw")
  
  
  # Modifications i need to make to the above
  # - Increase the height of the sankey diagram overall
  # - Look into color maps that could be better
  # - Color the cluster labels by what gets written to the JSON
  # - Plot the clusters as nodes in a small network
  
  # %% [markdown]
  # # try graph flow
  
  
  node_signal_flow = signal_flow(adj)
  mean_sf = np.zeros(k)
  for i in np.unique(pred_labels):
      inds = np.where(pred_labels == i)[0]
      mean_sf[i] = np.mean(node_signal_flow[inds])
  
  cluster_mean_latent = gmm.model_.means_[:, 0]
  block_probs = SBMEstimator().fit(bin_adj, y=pred_labels).block_p_
  block_prob_df = pd.DataFrame(data=block_probs, index=range(k), columns=range(k))
  block_g = nx.from_pandas_adjacency(block_prob_df, create_using=nx.DiGraph)
  plt.figure(figsize=(10, 10))
  # don't ever let em tell you you're too pythonic
  pos = dict(zip(range(k), zip(cluster_mean_latent, mean_sf)))
  # nx.draw_networkx_nodes(block_g, pos=pos)
  labels = nx.get_edge_attributes(block_g, "weight")
  # nx.draw_networkx_edge_labels(block_g, pos, edge_labels=labels)
  from matplotlib.cm import ScalarMappable
  import matplotlib as mpl
  
  norm = mpl.colors.LogNorm(vmin=0.01, vmax=0.1)
  
  sm = ScalarMappable(cmap="Reds", norm=norm)
  cmap = sm.to_rgba(np.array(list(labels.values())) + 0.01)
  nx.draw_networkx(
      block_g,
      pos,
      edge_cmap="Reds",
      edge_color=cmap,
      connectionstyle="arc3,rad=0.2",
      width=1.5,
  )
  
  # %% [markdown]
  # # signal flow marginals
  
  signal_flow_marginal(adj, pred_labels)
  
  # %% [markdown]
  # #
  
  
  def signal_flow_marginal(adj, labels, col_wrap=5, palette="tab20"):
      sf = signal_flow(adj)
      uni_labels = np.unique(labels)
      medians = []
      for i in uni_labels:
          inds = np.where(labels == i)[0]
          medians.append(np.median(sf[inds]))
      sort_inds = np.argsort(medians)[::-1]
      col_order = uni_labels[sort_inds]
      plot_df = pd.DataFrame()
      plot_df["Signal flow"] = sf
      plot_df["Class"] = labels
      fg = sns.FacetGrid(
          plot_df,
          col="Class",
          aspect=1.5,
          palette=palette,
          col_order=col_order,
          sharey=False,
          col_wrap=col_wrap,
          xlim=(-3, 3),
      )
      fg = fg.map(sns.distplot, "Signal flow")  # bins=np.linspace(-2.2, 2.2))
      fg.set(yticks=[], yticklabels=[])
      plt.tight_layout()
      return fg
  
  
  signal_flow_marginal(adj, class_labels)
  stashfig("known-class-sf-marginal")
  
  # tomorrow
  # DEFINITELY
  # run with unsupervised metrics from k=2-50
  
  # IF TIME
  # run hgmm
   # %% [markdown]
  # #
  import itertools
  import os
  import time
  from itertools import chain
  
  import colorcet as cc
  import matplotlib as mpl
  import matplotlib.pyplot as plt
  import networkx as nx
  import numpy as np
  import pandas as pd
  import seaborn as sns
  from anytree import LevelOrderGroupIter, Node, RenderTree
  from joblib import Parallel, delayed
  from mpl_toolkits.axes_grid1 import make_axes_locatable
  from sklearn.decomposition import PCA
  
  from graspy.plot import heatmap, pairplot
  from src.data import load_metagraph
  from src.graph import MetaGraph, preprocess
  from src.io import savecsv, savefig, saveskels
  from src.traverse import (
      cascades_from_node,
      generate_cascade_tree,
      generate_random_walks,
      path_to_visits,
      to_markov_matrix,
      to_path_graph,
  )
  from src.visualization import (
      CLASS_COLOR_DICT,
      barplot_text,
      draw_networkx_nice,
      draw_separators,
      matrixplot,
      remove_shared_ax,
      remove_spines,
      screeplot,
      sort_meta,
      stacked_barplot,
  )
  
  FNAME = os.path.basename(__file__)[:-3]
  print(FNAME)
  
  
  def stashfig(name, **kws):
      savefig(name, foldername=FNAME, save_on=True, **kws)
  
  
  def stashcsv(df, name, **kws):
      savecsv(df, name, foldername=FNAME, save_on=True, **kws)
  
  
  #%% Load and preprocess the data
  
  VERSION = "2020-03-09"
  print(f"Using version {VERSION}")
  
  plot_examples = False
  plot_embed = False
  plot_full_mat = False
  graph_type = "Gad"
  threshold = 0
  weight = "weight"
  mg = load_metagraph(graph_type, VERSION)
  mg = preprocess(
      mg,
      threshold=threshold,
      sym_threshold=False,
      remove_pdiff=True,
      binarize=False,
      weight=weight,
  )
  print(f"Preprocessed graph {graph_type} with threshold={threshold}, weight={weight}")
  
  # TODO update this with the mixed groups
  # TODO make these functional for selecting proper paths
  
  out_classes = [
      "O_dSEZ",
      "O_dSEZ;CN",
      "O_dSEZ;LHN",
      "O_dVNC",
      "O_dVNC;O_RG",
      "O_dVNC;CN",
      "O_RG",
      "O_dUnk",
      "O_RG-IPC",
      "O_RG-ITP",
      "O_RG-CA-LP",
  ]
  from_groups = [
      ("sens-ORN",),
      ("sens-photoRh5", "sens-photoRh6"),
      ("sens-MN",),
      ("sens-thermo",),
      ("sens-vtd",),
      ("sens-AN",),
  ]
  from_group_names = ["Odor", "Photo", "MN", "Temp", "VTD", "AN"]
  
  out_groups = [
      ("motor-mAN", "motormVAN", "motor-mPaN"),
      ("O_dSEZ", "O_dVNC;O_dSEZ", "O_dSEZ;CN", "LHN;O_dSEZ"),
      ("O_dVNC", "O_dVNC;CN", "O_RG;O_dVNC", "O_dVNC;O_dSEZ"),
      ("O_RG", "O_RG-IPC", "O_RG-ITP", "O_RG-CA-LP", "O_RG;O_dVNC"),
      ("O_dUnk",),
  ]
  out_group_names = ["Motor", "SEZ", "VNC", "RG", "dUnk"]
  
  
  from_classes = list(chain.from_iterable(from_groups))  # make this a flat list
  out_classes = list(chain.from_iterable(out_groups))
  
  class_key = "Merge Class"
  
  adj = nx.to_numpy_array(mg.g, weight=weight, nodelist=mg.meta.index.values)
  n_verts = len(adj)
  meta = mg.meta.copy()
  g = mg.g.copy()
  meta["idx"] = range(len(meta))
  
  from_inds = meta[meta[class_key].isin(from_classes)]["idx"].values
  out_inds = meta[meta[class_key].isin(out_classes)]["idx"].values
  ind_map = dict(zip(meta.index, meta["idx"]))
  g = nx.relabel_nodes(g, ind_map, copy=True)
  out_ind_map = dict(zip(out_inds, range(len(out_inds))))
  
  # %% [markdown]
  # # Use a method to generate visits
  
  path_type = "cascade"
  if path_type == "cascade":
      p = 0.01
      not_probs = (
          1 - p
      ) ** adj  # probability of none of the synapses causing postsynaptic
      probs = 1 - not_probs  # probability of ANY of the synapses firing onto next
  elif path_type == "fancy-cascade":
      alpha = 0.5
      flat = np.full(adj.shape, alpha)
      deg = meta["dendrite_input"].values
      deg[deg == 0] = 1
      flat = flat / deg[None, :]
      not_probs = np.power((1 - flat), adj)
      probs = 1 - not_probs
  
  #%%
  seed = 8888
  max_depth = 10
  n_bins = 10
  n_sims = 100
  method = "tree"
  normalize_n_source = False
  
  
  basename = f"-{graph_type}-t{threshold}-pt{path_type}-b{n_bins}-n{n_sims}-m{method}"
  basename += f"-norm{normalize_n_source}"
  basename += f"-plus-inverted"
  
  
  np.random.seed(seed)
  if method == "tree":
      seeds = np.random.choice(int(1e8), size=len(from_inds), replace=False)
      outs = Parallel(n_jobs=1, verbose=10)(
          delayed(cascades_from_node)(
              fi, probs, out_inds, max_depth, n_sims, seed, n_bins, method
          )
          for fi, seed in zip(from_inds, seeds)
      )
  elif method == "path":
      outs = []
      for start_ind in from_inds:
          temp_hist = cascades_from_node(
              start_ind, probs, out_inds, max_depth, n_sims, seed, n_bins, method
          )
          outs.append(temp_hist)
  from_hist_mat = np.concatenate(outs, axis=-1)
  
  ###
  # invert
  if method == "tree":
      seeds = np.random.choice(int(1e8), size=len(out_inds), replace=False)
      outs = Parallel(n_jobs=1, verbose=10)(
          delayed(cascades_from_node)(
              fi, probs.T, from_inds, max_depth, n_sims, seed, n_bins, method
          )
          for fi, seed in zip(out_inds, seeds)
      )
  elif method == "path":
      outs = []
      for start_ind in from_inds:
          temp_hist = cascades_from_node(
              start_ind, probs.T, out_inds, max_depth, n_sims, seed, n_bins, method
          )
          outs.append(temp_hist)
  out_hist_mat = np.concatenate(outs, axis=-1)
  
  
  # generate_cascade_paths(start_ind, probs, 1, stop_inds=out_inds, max_depth=10)
  # %% [markdown]
  # # Sort metadata
  full_hist_mat = np.concatenate((from_hist_mat, out_hist_mat), axis=1)
  hist_mat = full_hist_mat
  # row metadata
  ids = pd.Series(index=meta["idx"], data=meta.index, name="id")
  to_class = ids.map(meta["Merge Class"])
  to_class.name = "to_class"
  row_df = pd.concat([ids, to_class], axis=1)
  
  # col metadata
  orders = pd.Series(data=len(from_inds) * list(range(n_bins)), name="order")
  from_idx = pd.Series(data=np.repeat(from_inds, n_bins), name="idx")
  from_ids = from_idx.map(ids)
  from_ids.name = "id"
  from_class = from_ids.map(meta["Merge Class"])
  from_class.name = "class"
  from_col_df = pd.concat([orders, from_idx, from_ids, from_class], axis=1)
  
  orders = pd.Series(data=len(out_inds) * list(range(n_bins)), name="order")
  out_idx = pd.Series(data=np.repeat(out_inds, n_bins), name="idx")
  out_ids = out_idx.map(ids)
  out_ids.name = "id"
  out_class = out_ids.map(meta["Merge Class"])
  out_class.name = "class"
  out_col_df = pd.concat([orders, out_idx, out_ids, out_class], axis=1)
  col_df = pd.concat([from_col_df, out_col_df], axis=0, ignore_index=True)
  # %% [markdown]
  # #
  log_mat = np.log10(hist_mat + 1)
  if plot_full_mat:
      shape = log_mat.shape
      figsize = (10, 20)
      fig, ax = plt.subplots(1, 1, figsize=figsize)
      matrixplot(
          log_mat,
          ax=ax,
          col_meta=col_df,
          col_sort_class=["from_class"],
          row_meta=row_df,
          row_sort_class=["to_class"],
          plot_type="scattermap",
          sizes=(0.5, 0.5),
          tick_rot=45,
      )
      stashfig("log-full-scatter" + basename)
  
      fig, ax = plt.subplots(1, 1, figsize=figsize)
      matrixplot(
          log_mat,
          ax=ax,
          col_meta=col_df,
          col_sort_class=["from_class"],
          row_colors=CLASS_COLOR_DICT,
          row_meta=row_df,
          row_sort_class=["to_class"],
          plot_type="heatmap",
          sizes=(0.5, 0.5),
          tick_rot=45,
      )
      stashfig("log-full-heat" + basename)
  
  # %% [markdown]
  # # Screeplots
  
  if plot_embed:
      screeplot(hist_mat.astype(float), title="Raw hist mat (full)")
      stashfig("scree-raw-mat" + basename)
      screeplot(log_mat, title="Log hist mat (full)")
      stashfig("scree-log-mat" + basename)
  
  # %% [markdown]
  # # Pairplots
  if plot_embed:
      pca = PCA(n_components=6)
      embed = pca.fit_transform(log_mat)
      loadings = pca.components_.T
      pg = pairplot(
          embed,
          labels=to_class.values,
          palette=CLASS_COLOR_DICT,
          height=5,
          title="Node response embedding (log)",
      )
      pg._legend.remove()
      stashfig("node-pca-log" + basename)
      pg = pairplot(
          loadings,
          labels=from_class.values,
          height=5,
          title="Source class embedding (log)",
      )
      stashfig("source-pca-log" + basename)
  
      pca = PCA(n_components=6)
      embed = pca.fit_transform(hist_mat.astype(float))
      loadings = pca.components_.T
      pg = pairplot(
          embed,
          labels=to_class.values,
          palette=CLASS_COLOR_DICT,
          height=5,
          title="Node response embedding (raw)",
      )
      pg._legend.remove()
      stashfig("node-pca-log" + basename)
      pg = pairplot(
          loadings,
          labels=from_class.values,
          height=5,
          title="Source class embedding (raw)",
      )
      stashfig("source-pca-log" + basename)
  
  # %% [markdown]
  # # Collapse that matrix
  hist_mat = full_hist_mat
  collapsed_hist = []
  collapsed_col_df = []
  groups = from_groups + out_groups
  names = from_group_names + out_group_names
  for fg, fg_name in zip(groups, names):
      from_df = col_df[col_df["class"].isin(fg)]
      n_in_group = len(from_df)
      for order in from_df["order"].unique():
          inds = from_df[from_df["order"] == order].index
          col = hist_mat[:, inds].sum(axis=1)
          if normalize_n_source:
              col = col.astype(float)
              col /= n_in_group
          collapsed_hist.append(col)
          row = {"order": order, "class": fg_name}
          collapsed_col_df.append(row)
  
  
  collapsed_col_df = pd.DataFrame(collapsed_col_df)
  collapsed_hist = np.array(collapsed_hist).T
  log_collapsed_hist = np.log10(collapsed_hist + 1)
  
  # %% [markdown]
  # #
  if plot_embed:
      pca = PCA(n_components=6)
      embed = pca.fit_transform(log_collapsed_hist)
      loadings = pca.components_.T
      pg = pairplot(
          embed,
          labels=to_class.values,
          palette=CLASS_COLOR_DICT,
          height=5,
          title="Collapsed node response embedding (log)",
      )
      pg._legend.remove()
      stashfig("coll-node-pca-log" + basename)
      pg = pairplot(
          loadings,
          labels=collapsed_col_df["from_class"].values,
          height=5,
          title="Collapsed source class embedding (log)",
      )
      stashfig("coll-source-pca-log" + basename)
  
      pca = PCA(n_components=6)
      embed = pca.fit_transform(collapsed_hist.astype(float))
      loadings = pca.components_.T
      pg = pairplot(
          embed,
          labels=to_class.values,
          palette=CLASS_COLOR_DICT,
          height=5,
          title="Collapsed node response embedding (raw)",
      )
      pg._legend.remove()
      stashfig("coll-node-pca-log" + basename)
      pg = pairplot(
          loadings,
          labels=collapsed_col_df["from_class"].values,
          height=5,
          title="Collapsed source class embedding (raw)",
      )
      stashfig("coll-source-pca-log" + basename)
  
  # %% [markdown]
  # # Compute mean visit over all sources, for plotting
  def mean_visit(row):
      n_groups = len(row) // n_bins
      s = 0
      for i in range(n_groups):
          group = row[i * n_bins : (i + 1) * n_bins]
          for j, val in enumerate(group):
              s += j * val
      s /= row.sum()
      return s
  
  
  visits = []
  for r in collapsed_hist:
      mv = mean_visit(r)
      visits.append(mv)
  visits = np.array(visits)
  visits[np.isnan(visits)] = n_bins + 1
  row_df["visit_order"] = visits
  mean_visit_order = row_df.groupby(["to_class"])["visit_order"].mean()
  row_df["group_visit_order"] = row_df["to_class"].map(mean_visit_order)
  row_df["n_visit"] = collapsed_hist.sum(axis=1)
  # %% [markdown]
  # #
  fig, ax = plt.subplots(1, 1, figsize=(15, 15))
  sns.set_context("talk", font_scale=0.8)
  gridline_kws = dict(color="grey", linestyle="--", alpha=0.7, linewidth=0.3)
  matrixplot(
      log_collapsed_hist,
      ax=ax,
      col_meta=collapsed_col_df,
      col_sort_class=["class"],
      row_meta=row_df,
      row_sort_class=["to_class"],
      row_colors=CLASS_COLOR_DICT,
      row_class_order="group_visit_order",
      row_item_order=["visit_order"],
      plot_type="heatmap",
      tick_rot=0,
      row_ticks=False,
      gridline_kws=gridline_kws,
  )
  stashfig("collapsed-log-heat" + basename)
  
  # %% [markdown]
  # #
  sns.set_context("talk", font_scale=1)
  gridline_kws = dict(color="grey", linestyle="--", alpha=0.7, linewidth=0.3)
  
  fig, ax = plt.subplots(1, 1, figsize=(25, 15))
  ax, divider, top_cax, left_cax = matrixplot(
      log_collapsed_hist.T,
      ax=ax,
      row_meta=collapsed_col_df,
      row_sort_class=["class"],
      col_meta=row_df,
      col_sort_class=["to_class"],
      col_colors=CLASS_COLOR_DICT,
      col_class_order="group_visit_order",
      col_item_order=["visit_order"],
      plot_type="heatmap",
      tick_rot=45,
      col_ticks=False,
      gridline_kws=gridline_kws,
  )
  cax = divider.append_axes("right", size="1%", pad=0.02, sharey=ax)
  remove_shared_ax(cax)
  sns.heatmap(
      collapsed_col_df["order"][:, None], ax=cax, cbar=False, cmap="RdBu", center=0
  )
  cax.set_xticks([])
  cax.set_yticks([])
  cax.set_ylabel(r"Hops $\to$", rotation=-90, ha="center", va="center", labelpad=20)
  cax.yaxis.set_label_position("right")
  top_cax.set_yticks([0.5])
  top_cax.set_yticklabels(["Class"], va="center")
  ax.set_xlabel("Neuron")
  ax.set_ylabel("Source class")
  stashfig("collapsed-log-heat-transpose" + basename, dpi=200)
  
  fig, ax = plt.subplots(1, 1, figsize=(25, 15))
  ax, divider, top_cax, left_cax = matrixplot(
      log_collapsed_hist.T,
      ax=ax,
      row_meta=collapsed_col_df,
      row_sort_class=["class"],
      col_meta=row_df,
      col_sort_class=["to_class"],
      col_colors=CLASS_COLOR_DICT,
      col_class_order="group_visit_order",
      col_item_order=["visit_order"],
      plot_type="heatmap",
      tick_rot=45,
      col_ticks=True,
      gridline_kws=gridline_kws,
  )
  cax = divider.append_axes("right", size="1%", pad=0.02, sharey=ax)
  remove_shared_ax(cax)
  sns.heatmap(
      collapsed_col_df["order"][:, None], ax=cax, cbar=False, cmap="RdBu", center=0
  )
  cax.set_xticks([])
  cax.set_yticks([])
  cax.set_ylabel(r"Hops $\to$", rotation=-90, ha="center", va="center", labelpad=20)
  cax.yaxis.set_label_position("right")
  top_cax.set_yticks([0.5])
  top_cax.set_yticklabels(["Class"], va="center")
  ax.set_xlabel("Neuron")
  ax.set_ylabel("Source class")
  stashfig("collapsed-log-heat-transpose-labeled" + basename, dpi=200)
  
  # %% [markdown]
  # # clustermap the matrix
  
  
  sns.set_context("talk", font_scale=1)
  linkage = "average"
  metric = "euclidean"
  colors = np.vectorize(CLASS_COLOR_DICT.get)(row_df["to_class"])
  
  perm_inds, sort_collapsed_col_df = sort_meta(
      collapsed_col_df, sort_class=["from_class"]
  )
  sort_log_collapsed_hist = log_collapsed_hist[:, perm_inds]
  
  
  cg = sns.clustermap(
      data=sort_log_collapsed_hist.T,
      col_cluster=True,
      row_cluster=False,
      col_colors=colors,
      cmap="RdBu_r",
      center=0,
      cbar_pos=None,
      method=linkage,
      metric=metric,
  )
  ax = cg.ax_heatmap
  draw_separators(
      ax,
      ax_type="y",
      sort_meta=sort_collapsed_col_df,
      sort_class=["from_class"],
      tick_rot=0,
  )
  ax.xaxis.set_ticks([])
  # ax.set_ylabel(r"Visits over time $\to$")
  ax.set_xlabel("Neuron")
  ax.yaxis.tick_left()
  # ax.set_yticklabels(ax.get_yticklabels(), ha="left")
  stashfig("collapsed-log-clustermap" + basename)
  # stashfig("collapsed-log-clustermap" + basename, fmt="pdf")
  
  
  # %% [markdown]
  # # Do some plotting for illustration only
  
  
  if plot_examples:
      sns.set_context("talk")
      sns.set_palette("Set1")
      examples = [742, 605, 743, 2282, 596, 2367, 1690, 2313]
      for target_ind in examples:
          row = collapsed_hist[target_ind, :]
          perm_inds, sort_col_df = sort_meta(collapsed_col_df, sort_class=["from_class"])
          sort_row = row[perm_inds]
  
          fig, ax = plt.subplots(1, 1)
          xs = np.arange(len(sort_row)) + 0.5
          divider = make_axes_locatable(ax)
          bot_cax = divider.append_axes("bottom", size="3%", pad=0.02, sharex=ax)
          remove_shared_ax(bot_cax)
  
          ax.bar(x=xs, height=sort_row, width=0.8)
          draw_separators(
              ax, sort_meta=sort_col_df, sort_class=["from_class"], tick_rot=0
          )
          ax.set_xlim(0, len(xs))
          ax.set_ylabel("# hits @ time")
  
          sns.heatmap(
              collapsed_col_df["order"][None, :],
              ax=bot_cax,
              cbar=False,
              cmap="RdBu",
              center=0,
          )
          bot_cax.set_xticks([])
          bot_cax.set_yticks([])
          bot_cax.set_xlabel(r"Hops $\to$", x=0.1, ha="left", labelpad=-22)
          bot_cax.set_xticks([20.5, 24.5, 28.5])
          bot_cax.set_xticklabels([1, 5, 9], rotation=0)
  
          ax.spines["right"].set_visible(False)
          ax.spines["top"].set_visible(False)
          target_skid = meta.iloc[target_ind, :].name
          ax.set_title(
              f"Response for cell {target_skid} ({meta[meta['idx'] == target_ind]['Merge Class'].values[0]})"
          )
  
          stashfig(f"{target_skid}-response-hist" + basename)
  
   #%%
  from src.data import load_metagraph
  import seaborn as sns
  import numpy as np
  import pandas as pd
  import matplotlib.pyplot as plt
  
  
  mg = load_metagraph("G", "2020-01-21")
  
  is_pdiff = np.where(mg["is_pdiff"])[0]
  mg = mg.reindex(is_pdiff)
  degree_df = mg.calculate_degrees()
  plt.figure()
  melt_degree = pd.melt(
      degree_df.reset_index(),
      id_vars=["ID"],
      value_vars=["In degree", "Out degree", "Total degree"],
      value_name="Degree",
  )
  sns.stripplot(y="Degree", data=melt_degree, x="variable", jitter=0.45)
  
  plt.figure()
  melt_syns = pd.melt(
      degree_df.reset_index(),
      id_vars=["ID"],
      value_vars=["In edgesum", "Out edgesum", "Total edgesum"],
      value_name="Synapses",
  )
  sns.stripplot(y="Synapses", data=melt_syns, x="variable", jitter=0.45)
  
apis: numpy.sum numpy.diag matplotlib.pyplot.tight_layout numpy.random.seed numpy.argsort matplotlib.cm.ScalarMappable matplotlib.pyplot.style.use matplotlib.pyplot.figure matplotlib.colors.LogNorm numpy.where numpy.unique numpy.mean numpy.zeros numpy.median numpy.count_nonzero numpy.ix_ sklearn.metrics.adjusted_rand_score numpy.isinf pandas.DataFrame numpy.errstate numpy.sqrt numpy.concatenate
apis: pandas.Series numpy.vectorize pandas.DataFrame numpy.random.seed numpy.repeat matplotlib.pyplot.subplots numpy.power numpy.log10 pandas.concat numpy.isnan numpy.array numpy.concatenate numpy.full sklearn.decomposition.PCA
apis: matplotlib.pyplot.figure numpy.where