representation_theory.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JqikLr4e9afF"
      },
      "source": [
        "Copyright 2021 DeepMind Technologies Limited\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\"); you may not use this file except in compliance with the License. You may obtain a copy of the License at\n",
        "\n",
        "https://www.apache.org/licenses/LICENSE-2.0\n",
        "Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MFCFPZ-izslo"
      },
      "source": [
        "For training it is strongly encouraged that a GPU is used (eg a local kernel or colab pro). Training on the free colab instance can be done, but the risk of the inactivity timeout or preemption make it less reliable since training for 20 epochs takes around an hour. Ideally 50 epochs would be used to complete training.\n",
        "\n",
        "Pretrained weights are provided so by default the colab will use these and run from start to finish in a reasonable time and replicate the results from the paper from these saved weights."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ufdlZejGDGOJ"
      },
      "outputs": [],
      "source": [
        "#@title Install modules\n",
        "from IPython.display import clear_output\n",
        "\n",
        "!pip install dm-haiku\n",
        "!pip install jax\n",
        "!pip install optax\n",
        "clear_output()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "efHiN6JwDRti"
      },
      "outputs": [],
      "source": [
        "#@title Imports\n",
        "\n",
        "import collections\n",
        "import dataclasses\n",
        "import datetime\n",
        "import enum\n",
        "import functools\n",
        "import itertools\n",
        "import json\n",
        "import os\n",
        "import pickle\n",
        "import random\n",
        "import tempfile\n",
        "from typing import Sequence\n",
        "\n",
        "import haiku as hk\n",
        "import jax\n",
        "import jax.numpy as jnp\n",
        "import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "import networkx as nx\n",
        "import numpy as np\n",
        "import optax\n",
        "import psutil\n",
        "import scipy.sparse as sp\n",
        "import seaborn as sns"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "PeoIjYst1csH"
      },
      "outputs": [],
      "source": [
        "#@title Download data\n",
        "\n",
        "DATA_DIR = tempfile.mkdtemp()\n",
        "!mkdir -p {DATA_DIR}\n",
        "\n",
        "print(f\"Copying data to {DATA_DIR} - NB this requires ~1.5G of space.\")\n",
        "!gsutil -m cp \"gs://maths_conjectures/representation_theory/*\" \"{DATA_DIR}/\"\n",
        "\n",
        "# Extract the graph data.\n",
        "GRAPH_DIR = os.path.join(DATA_DIR, \"graph_data\")\n",
        "!mkdir -p {GRAPH_DIR}\n",
        "!tar -xzf {DATA_DIR}/graph_data.tar.gz -C {GRAPH_DIR}\n",
        "\n",
        "!echo \"Files present:\"\n",
        "!ls -lh {DATA_DIR}\n",
        "!du -hs {DATA_DIR}\n",
        "\n",
        "with open(os.path.join(DATA_DIR, \"graph_index_to_node_index_to_permutation.json\"), \"rt\") as f:\n",
        "    graph_index_to_node_index_to_permutation = json.load(f)\n",
        "NUM_GRAPHS = len([f for f in os.listdir(GRAPH_DIR) if f.startswith(\"graph_\")])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "32yHJE0aErKK"
      },
      "outputs": [],
      "source": [
        "#@title Data loading functions\n",
        "\n",
        "#@markdown These functions load in previously generated Bruhat intervals (NetworkX graphs) and\n",
        "#@markdown their associated KL polynomials (a list of integer coefficients), converting them\n",
        "#@markdown into an adjacency matrix format and single label that is appropriate for our JAX model.\n",
        "#@markdown The label is the degree-label'th coeefficient from the KL polynomial. See\n",
        "#@markdown generate_graph_data() for details.\n",
        "\n",
        "train_fraction = .8\n",
        "\n",
        "\n",
        "def pad(iterable, size, padding=None):\n",
        "  return itertools.islice(pad_infinite(iterable, padding), size)\n",
        "\n",
        "\n",
        "def pad_infinite(iterable, padding=None):\n",
        "  return itertools.chain(iterable, itertools.repeat(padding))\n",
        "\n",
        "\n",
        "def convert_networkx_to_adjacency_input(graph):\n",
        "  adjacency_matrix = nx.to_scipy_sparse_matrix(graph, format='coo')\n",
        "  adjacency_matrix += sp.eye(adjacency_matrix.shape[0])\n",
        "  return adjacency_matrix\n",
        "\n",
        "\n",
        "@dataclasses.dataclass(frozen=True)\n",
        "class GraphData:\n",
        "  features: Sequence[np.ndarray]\n",
        "  labels: Sequence[np.ndarray]\n",
        "  adjacencies: Sequence[sp.csr_matrix]\n",
        "\n",
        "\n",
        "def generate_graph_data(degree_label):\n",
        "  \"\"\"Generate dataset for training GraphNet model on KL data.\n",
        "\n",
        "  This generates a dataset for training a GraphNet model.\n",
        "\n",
        "  Args:\n",
        "    degree_label: The polynomial coefficient to use as the label.\n",
        "\n",
        "  Returns:\n",
        "    An GraphData instance with features, adjacencies and labels.\n",
        "  \"\"\"\n",
        "  kls = read_kl_coefficients()\n",
        "  max_degree = max(len(kl) for kl in kls)\n",
        "\n",
        "  ys = np.array([list(pad(kl, max_degree, 0)) for kl in kls])\n",
        "  ys = ys[:, degree_label:degree_label+1]\n",
        "\n",
        "  features = []\n",
        "\n",
        "  for graph in iter_graph():\n",
        "    feat_dict = {\n",
        "        'in_centrality': nx.in_degree_centrality(graph),\n",
        "        'out_centrality': nx.out_degree_centrality(graph),\n",
        "    }\n",
        "\n",
        "    curr_feature = np.zeros((len(graph), len(feat_dict)))\n",
        "\n",
        "    for n, perm in enumerate(graph.nodes):\n",
        "      for i, (name, value) in enumerate(feat_dict.items()):\n",
        "        curr_feature[n,i] = value[perm]\n",
        "\n",
        "    features.append(curr_feature)\n",
        "  adjacencies = [convert_networkx_to_adjacency_input(g) for g in graphs]\n",
        "\n",
        "  return GraphData(features=features, labels=ys, adjacencies=adjacencies)\n",
        "\n",
        "\n",
        "@functools.lru_cache()\n",
        "def load_graphs_from_pickle():\n",
        "  assert hold_graphs_in_memory, \"Should only load data from the pickle if 'hold_graphs_in_memory' is True\"\n",
        "  with open(os.path.join(DATA_DIR, 'bruhat_data_S9.pickle'), 'rb') as ifile:\n",
        "    unused_interval_spec, unused_interval_lengths, graphs, unused_kls = pickle.load(ifile)\n",
        "  return graphs\n",
        "\n",
        "\n",
        "def iter_graph():\n",
        "  if hold_graphs_in_memory:\n",
        "    yield from load_graphs_from_pickle()\n",
        "  else:\n",
        "    for i in range(NUM_GRAPHS):\n",
        "      filename = os.path.join(GRAPH_DIR, f\"graph_{i:04d}.npz\")\n",
        "      yield nx.from_scipy_sparse_matrix(\n",
        "          sp.load_npz(filename), create_using=nx.DiGraph)\n",
        "\n",
        "\n",
        "@functools.lru_cache()\n",
        "def read_kl_coefficients():\n",
        "  with open(os.path.join(GRAPH_DIR, \"kl_coefficients.json\")) as f:\n",
        "    return json.load(f)\n",
        "\n",
        "\n",
        "def get_root_node(col):\n",
        "  return np.bincount(col).argmin()\n",
        "\n",
        "\n",
        "@dataclasses.dataclass(frozen=True)\n",
        "class InputData:\n",
        "  features: Sequence[np.ndarray]\n",
        "  labels: Sequence[np.ndarray]\n",
        "  rows: Sequence[sp.csr_matrix]\n",
        "  columns: Sequence[sp.csr_matrix]\n",
        "  root_nodes: Sequence[int]\n",
        "\n",
        "\n",
        "def load_input_data(degree_label=1):\n",
        "  \"\"\"Loads input data for the specified prediction problem.\n",
        "\n",
        "  This loads a dataset that can be used with a GraphNet model. The Bruhat\n",
        "  intervals are taken from the dataset of intervals in S9 and the label\n",
        "  is the coefficient of specified degree.\n",
        "\n",
        "  The datasets are cached, and only regenerated when not found on disk.\n",
        "\n",
        "  Args:\n",
        "    degree_label: the polynomial coefficient to use as the label.\n",
        "  Returns:\n",
        "    Three InputData instances with features, rows, cols and labels. They are the\n",
        "    full/train/test set respectively.\n",
        "  \"\"\"\n",
        "  input_data_cache_dir = os.path.join(DATA_DIR, f\"input_data_{degree_label}\")\n",
        "\n",
        "  # Extract from .tar if not already done.\n",
        "  tar_path = f\"{DATA_DIR}/S9_{degree_label}.tar.gz\"\n",
        "  cache_dir = os.path.join(DATA_DIR, f\"input_data_{degree_label}\")\n",
        "  if os.path.exists(tar_path) and not os.path.exists(cache_dir):\n",
        "    print(f\"Extracting input files from {tar_path}\")\n",
        "    !mkdir {cache_dir}\n",
        "    !tar -xzf {tar_path} -C {cache_dir}\n",
        "\n",
        "  # Load from cache for either extracted-tar or a previously computed run.\n",
        "  if os.path.exists(input_data_cache_dir):\n",
        "    print(f\"Loading np arrays from directory: '{input_data_cache_dir}'\", flush=True)\n",
        "    # Load adj\n",
        "    adjacencies = [sp.load_npz(os.path.join(input_data_cache_dir, filename))\n",
        "                   for filename in sorted(os.listdir(input_data_cache_dir))\n",
        "                   if not filename.endswith(\"arrays.npz\")]\n",
        "    # Load np arrays\n",
        "    with np.load(os.path.join(input_data_cache_dir, \"arrays.npz\")) as data:\n",
        "      ys = data[\"labels\"]\n",
        "      features = [data[f\"feature_{i:04d}\"] for i in range(len(adjacencies))]\n",
        "    print(\"Data loaded from cache.\", flush=True)\n",
        "  else:\n",
        "    print(f\"Generating data for degree_label {degree_label} and caching (~1m to generate)\", flush=True)\n",
        "    graph_data = generate_graph_data(degree_label)\n",
        "    features = graph_data.features\n",
        "    adjacencies = graph_data.adjacencies\n",
        "    ys = graph_data.labels\n",
        "\n",
        "    # Save to disk to save time in future:\n",
        "    !mkdir {input_data_cache_dir}\n",
        "    np.savez(os.path.join(input_data_cache_dir, \"arrays.npz\"),\n",
        "             **{f\"feature_{i:04d}\": f for i, f in enumerate(features)}, labels=ys)\n",
        "    for i, adj in enumerate(adjacencies):\n",
        "      sp.save_npz(os.path.join(input_data_cache_dir, f\"adj_{i:04d}.npz\"), adj)\n",
        "    print(f\"Data cached to directory {input_data_cache_dir}; future runs should be much faster!\")\n",
        "\n",
        "  rows = [sp.coo_matrix(a).row for a in adjacencies]\n",
        "  cols = [sp.coo_matrix(a).col for a in adjacencies]\n",
        "  root_nodes = [get_root_node(col) for col in cols]\n",
        "\n",
        "  num_training = int(len(ys) * train_fraction)\n",
        "\n",
        "  features_train = features[:num_training]\n",
        "  rows_train = [sp.coo_matrix(a).row for a in adjacencies[:num_training]]\n",
        "  cols_train = [sp.coo_matrix(a).col for a in adjacencies[:num_training]]\n",
        "  ys_train = ys[:num_training]\n",
        "  root_nodes_train = root_nodes[:num_training]\n",
        "\n",
        "  features_test = features[num_training:]\n",
        "  rows_test = [sp.coo_matrix(a).row for a in adjacencies[num_training:]]\n",
        "  cols_test = [sp.coo_matrix(a).col for a in adjacencies[num_training:]]\n",
        "  ys_test = ys[num_training:]\n",
        "  root_nodes_test = root_nodes[num_training:]\n",
        "  return (\n",
        "      InputData(features=features, rows=rows, columns=cols, labels=ys, root_nodes=root_nodes),\n",
        "      InputData(features=features_train, rows=rows_train, columns=cols_train, labels=ys_train, root_nodes=root_nodes_train),\n",
        "      InputData(features=features_test, rows=rows_test, columns=cols_test, labels=ys_test, root_nodes=root_nodes_test))\n",
        "\n",
        "\n",
        "#@markdown As the graphs generally do not have the same number of nodes, and because\n",
        "#@markdown JAX relies on data shapes being fixed and known upfront, we batch\n",
        "#@markdown together a set of graphs into a large batch graph that contains each\n",
        "#@markdown graph as a disconnected component.\n",
        "def batch(features, rows, cols, ys, root_nodes):\n",
        "  \"\"\"Converts a list of training examples into a batched single graph.\"\"\"\n",
        "  batch_size = len(features)\n",
        "  max_features = max(f.shape[0] for f in features)\n",
        "  b_features = np.zeros((batch_size, max_features, features[0].shape[1]))\n",
        "  b_rows = []\n",
        "  b_cols = []\n",
        "  b_ys = np.zeros((batch_size, 1))\n",
        "  b_masks = np.zeros((batch_size, max_features, 1))\n",
        "  for i in range(batch_size):\n",
        "    b_features[i, :features[i].shape[0], :] = features[i]\n",
        "    b_rows.append(rows[i] + i * max_features)\n",
        "    b_cols.append(cols[i] + i * max_features)\n",
        "    b_ys[i, 0] = ys[i, 0]\n",
        "    root_node = root_nodes[i]\n",
        "    b_masks[i, root_node, 0] = 1.0\n",
        "\n",
        "  b_features = b_features.reshape((-1, b_features.shape[-1]))\n",
        "  b_rows = np.concatenate(b_rows)\n",
        "  b_cols = np.concatenate(b_cols)\n",
        "\n",
        "  return b_features, b_rows, b_cols, b_ys, b_masks"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "_hARJYqID1-N"
      },
      "outputs": [],
      "source": [
        "#@title MPNN model definition code and helper functions\n",
        "\n",
        "\n",
        "class Direction(enum.Enum):\n",
        "  FORWARD = enum.auto()\n",
        "  BACKWARD = enum.auto()\n",
        "  BOTH = enum.auto()\n",
        "\n",
        "\n",
        "class Reduction(enum.Enum):\n",
        "  SUM = enum.auto()\n",
        "  MAX = enum.auto()\n",
        "\n",
        "\n",
        "class MPNN(hk.Module):\n",
        "  \"\"\"Sparse Message-Passing Neural Network (Gilmer et al., ICML 2017).\"\"\"\n",
        "\n",
        "  def __init__(\n",
        "      self,\n",
        "      *,\n",
        "      out_size: int,\n",
        "      mid_size: int,\n",
        "      activation,\n",
        "      direction: Direction,\n",
        "      residual: bool,\n",
        "      reduction: Reduction,\n",
        "      message_relu: bool,\n",
        "      with_bias: bool,\n",
        "  ):\n",
        "    \"\"\"Build MPNN layer.\n",
        "\n",
        "    Args:\n",
        "      out_size: Output width of the network.\n",
        "      mid_size: Width of the hidden layer.\n",
        "      activation: Activation function to use before the layer output.\n",
        "      direction: Direction of message passing. See Direction Enum.\n",
        "      residual: Whether to use resiudal connections.\n",
        "      reduction: Reduction function to aggregate messages at nodes. See\n",
        "        Reduction enum.\n",
        "      message_relu: Whether to apply a relu on each message.\n",
        "      with_bias: Whether to add biases in linear layers.\n",
        "\n",
        "    Returns:\n",
        "      The output of the MPNN layer.\n",
        "    \"\"\"\n",
        "    super().__init__(name='mpnn_aggr')\n",
        "    self.mid_size = out_size if mid_size is None else mid_size\n",
        "    self.out_size = out_size\n",
        "    self.activation = activation\n",
        "    self.direction = direction\n",
        "    self.reduction = reduction\n",
        "    self.residual = residual\n",
        "    self.message_relu = message_relu\n",
        "    self.with_bias = with_bias\n",
        "\n",
        "    @jax.jit\n",
        "    def jax_coo_sum(rows, cols, msg_in, msg_out):\n",
        "      msg_vect = msg_in[rows] + msg_out[cols]\n",
        "      if message_relu:\n",
        "        msg_vect = jax.nn.relu(msg_vect)\n",
        "      return jnp.zeros_like(msg_out).at[rows].add(msg_vect)\n",
        "\n",
        "    @jax.jit\n",
        "    def jax_coo_max(rows, cols, msg_in, msg_out):\n",
        "      msg_vect = msg_in[rows] + msg_out[cols]\n",
        "      if message_relu:\n",
        "        msg_vect = jax.nn.relu(msg_vect)\n",
        "      return jnp.zeros_like(msg_in).at[rows].max(msg_vect)\n",
        "\n",
        "    self.jax_coo_sum = jax_coo_sum\n",
        "    self.jax_coo_max = jax_coo_max\n",
        "\n",
        "  def __call__(self, features, rows, cols):\n",
        "    if self.direction == Direction.FORWARD or self.direction == Direction.BOTH:\n",
        "      m1_1 = hk.Linear(self.mid_size, with_bias=self.with_bias)\n",
        "      m2_1 = hk.Linear(self.mid_size, with_bias=self.with_bias)\n",
        "      msg_1_1 = m1_1(features)\n",
        "      msg_2_1 = m2_1(features)\n",
        "    if self.direction == Direction.BACKWARD or self.direction == Direction.BOTH:\n",
        "      m1_2 = hk.Linear(self.mid_size, with_bias=self.with_bias)\n",
        "      m2_2 = hk.Linear(self.mid_size, with_bias=self.with_bias)\n",
        "      msg_1_2 = m1_2(features)\n",
        "      msg_2_2 = m2_2(features)\n",
        "\n",
        "    o2 = hk.Linear(self.out_size, with_bias=self.with_bias)\n",
        "\n",
        "    if self.reduction == Reduction.MAX:\n",
        "      reduction = self.jax_coo_max\n",
        "    elif self.reduction == Reduction.SUM:\n",
        "      reduction = self.jax_coo_sum\n",
        "    else:\n",
        "      raise ValueError('Unknown reduction %s' % self.reduction)\n",
        "\n",
        "    if self.direction == Direction.FORWARD:\n",
        "      msgs = reduction(rows, cols, msg_1_1, msg_2_1)\n",
        "    elif self.direction == Direction.BACKWARD:\n",
        "      msgs = reduction(cols, rows, msg_1_2, msg_2_2)\n",
        "    elif self.direction == Direction.BOTH:\n",
        "      msgs_1 = reduction(rows, cols, msg_1_1, msg_2_1)\n",
        "      msgs_2 = reduction(cols, rows, msg_1_2, msg_2_2)\n",
        "      msgs = jnp.concatenate([msgs_1, msgs_2], axis=-1)\n",
        "      pass\n",
        "    else:\n",
        "      raise ValueError('Unknown direction %s' % self.direction)\n",
        "\n",
        "    h_2 = o2(msgs)\n",
        "    if self.residual:\n",
        "      o1 = hk.Linear(self.out_size, with_bias=self.with_bias)\n",
        "      h_1 = o1(features)\n",
        "      network_output = h_1 + h_2\n",
        "    else:\n",
        "      network_output = h_2\n",
        "\n",
        "    if self.activation is not None:\n",
        "      network_output = self.activation(network_output)\n",
        "\n",
        "    return network_output\n",
        "\n",
        "\n",
        "class Model:\n",
        "\n",
        "  def __init__(\n",
        "      self,\n",
        "      *,\n",
        "      num_layers: int,\n",
        "      num_features: int,\n",
        "      num_classes: int,\n",
        "      direction: Direction,\n",
        "      reduction: Reduction,\n",
        "      apply_relu_activation: bool,\n",
        "      use_mask: bool,\n",
        "      share: bool,\n",
        "      message_relu: bool,\n",
        "      with_bias: bool,\n",
        "  ):\n",
        "    \"\"\"Get the jax model function and associated functions.\n",
        "\n",
        "    Args:\n",
        "      num_layers: The number of layers in the GraphNet - equivalently the number\n",
        "        of propagation steps.\n",
        "      num_features: The dimension of the hidden layers / messages.\n",
        "      num_classes: The number of target classes.\n",
        "      direction: Edges to pass messages along, see Direction enum.\n",
        "      reduction: The reduction operation to be used to aggregate messages at\n",
        "        each node at each step. See Reduction enum.\n",
        "      apply_relu_activation: Whether to apply a relu at the end of each\n",
        "        propogration step.\n",
        "      use_mask: Boolean; should a masked prediction in central node be\n",
        "        performed?\n",
        "      share: Boolean; should the GNN layers be shared?\n",
        "      message_relu: Boolean; should a ReLU be used in the message function?\n",
        "      with_bias: Boolean; should the linear layers have bias?\n",
        "    \"\"\"\n",
        "    self._num_layers = num_layers\n",
        "    self._num_features = num_features\n",
        "    self._num_classes = num_classes\n",
        "    self._direction = direction\n",
        "    self._reduction = reduction\n",
        "    self._apply_relu_activation = apply_relu_activation\n",
        "    self._use_mask = use_mask\n",
        "    self._share = share\n",
        "    self._message_relu = message_relu\n",
        "    self._with_bias = with_bias\n",
        "\n",
        "  def _kl_net(self, features, rows, cols, batch_size, masks):\n",
        "    in_enc = hk.Linear(self._num_features)\n",
        "\n",
        "    if self._apply_relu_activation:\n",
        "      activation_fn = jax.nn.relu\n",
        "    else:\n",
        "      activation_fn = lambda net: net\n",
        "\n",
        "    gnns = []\n",
        "    for i in range(self._num_layers):\n",
        "      if i == 0 or not self._share:\n",
        "        gnns.append(\n",
        "            MPNN(\n",
        "                out_size=self._num_features,\n",
        "                mid_size=None,\n",
        "                direction=self._direction,\n",
        "                reduction=self._reduction,\n",
        "                activation=activation_fn,\n",
        "                message_relu=self._message_relu,\n",
        "                with_bias=self._with_bias,\n",
        "                residual=True))\n",
        "      else:\n",
        "        gnns.append(gnns[-1])\n",
        "\n",
        "    out_enc = hk.Linear(self._num_classes, with_bias=self._with_bias)\n",
        "\n",
        "    hiddens = []\n",
        "    hidden = in_enc(features)\n",
        "    hiddens.append(jnp.reshape(hidden, (batch_size, -1, self._num_features)))\n",
        "    for gnn in gnns:\n",
        "      hidden = gnn(hidden, rows, cols)\n",
        "      hiddens.append(jnp.reshape(hidden, (batch_size, -1, self._num_features)))\n",
        "\n",
        "    hidden = jnp.reshape(hidden, (batch_size, -1, self._num_features))\n",
        "\n",
        "    if self._use_mask:\n",
        "      h_bar = jnp.sum(hidden * masks, axis=1)\n",
        "    else:\n",
        "      h_bar = jnp.max(hidden, axis=1)\n",
        "\n",
        "    lgts = out_enc(h_bar)\n",
        "\n",
        "    return hiddens, lgts\n",
        "\n",
        "  @property\n",
        "  def net(self):\n",
        "    return hk.transform(self._kl_net)\n",
        "\n",
        "  @functools.partial(jax.jit, static_argnums=(0,))\n",
        "  def loss(self, params, features, rows, cols, ys, masks):\n",
        "    _, lgts = self.net.apply(params, None, features, rows, cols, ys.shape[0],\n",
        "                             masks)\n",
        "    return -jnp.mean(\n",
        "        jax.nn.log_softmax(lgts) *\n",
        "        jnp.squeeze(jax.nn.one_hot(ys, self._num_classes), 1))\n",
        "\n",
        "  @functools.partial(jax.jit, static_argnums=(0,))\n",
        "  def accuracy(self, params, features, rows, cols, ys, masks):\n",
        "    _, lgts = self.net.apply(params, None, features, rows, cols, ys.shape[0],\n",
        "                             masks)\n",
        "    pred = jnp.argmax(lgts, axis=-1)\n",
        "    true_vals = jnp.squeeze(ys, axis=1)\n",
        "    acc = jnp.mean(pred == true_vals)\n",
        "    return acc"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c3DB916c37E34ba5"
      },
      "source": [
        ""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "rsL8rxtaWUeD"
      },
      "outputs": [],
      "source": [
        "#@title Load data\n",
        "\n",
        "#@markdown Training this model is pretty slow - an hour or so on the free tier colab, but subject to inactivity timeouts and pre-emptions.\n",
        "\n",
        "#@markdown In order to make it possible to recreate the results from the paper reliably and quickly, we provide several helpers to either speed things up, or reduce the memory footprint:\n",
        "#@markdown * Pretrained weights - greatly speeds things up by loading the trained model parameters rather than learning from the data\n",
        "#@markdown * If you are running on a high memory machine (ie *not* on the free colab instance!) the input graph data can be loaded from a pickle (which is faster to load) and kept in memory (faster to re-use, but uses ~12Gb of memory). This makes no difference to training speed (it's only relevant for `generate_graph_data()` and `get_saliency_vectors()`).\n",
        "\n",
        "use_pretrained_weights = True  #@param{type:\"boolean\"}\n",
        "hold_graphs_in_memory = False  #@param{type:\"boolean\"}\n",
        "\n",
        "gb = 1024**3\n",
        "total_memory = psutil.virtual_memory().total / gb\n",
        "# Less than 20Gb of RAM means we need to do some things slower, but with lower memory impact - in\n",
        "# particular, we want to allow things to run on the free colab tier.\n",
        "if total_memory \u003c 20 and hold_graphs_in_memory:\n",
        "    raise RuntimeError(f\"It is unlikely your machine (with {total_memory}Gb) will have enough memory to complete the colab's execution!\")\n",
        "\n",
        "print(\"Loading input data...\")\n",
        "full_dataset, train_dataset, test_dataset = load_input_data(degree_label=4)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2dNdizZMAytc"
      },
      "source": [
        "The below section defines the model used for predicting a given KL coefficient from an adjacency representation of the Bruhat interval. The model is a version of the Message-Passing Neural Network of Gilmer et al. While there may be other models that can also effectively model this problem, this was chosen in part due to knowledge of the structure of the KL polynomials. We treat the problem of predicting a coefficient as a classification problem, with the number of classes as the largest coefficient observed in the dataset. While this ignores ordering information in the label, we are still able to achieve high accuracies and derive insights from the network."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "PxsZ50p_FU_0"
      },
      "outputs": [],
      "source": [
        "#@title Network Setup\n",
        "\n",
        "step_size = 0.001\n",
        "batch_size = 100\n",
        "\n",
        "num_classes = np.max(train_dataset.labels) + 1\n",
        "model = Model(\n",
        "    num_layers=3,\n",
        "    num_features=64,\n",
        "    num_classes=num_classes,\n",
        "    direction=Direction.BOTH,\n",
        "    reduction=Reduction.SUM,\n",
        "    apply_relu_activation=True,\n",
        "    use_mask=False,\n",
        "    share=False,\n",
        "    message_relu=True,\n",
        "    with_bias=True)\n",
        "\n",
        "loss_val_gr = jax.value_and_grad(model.loss)\n",
        "opt_init, opt_update = optax.adam(step_size)\n",
        "\n",
        "\n",
        "def train(params, opt_state, features, rows, cols, ys, masks):\n",
        "  curr_loss, gradient = loss_val_gr(params, features, rows, cols, ys, masks)\n",
        "  updates, opt_state = opt_update(gradient, opt_state)\n",
        "  new_params = optax.apply_updates(params, updates)\n",
        "  return new_params, opt_state, curr_loss\n",
        "\n",
        "\n",
        "def compute_accuracies(params_to_evaluate, dataset, batch_size=100):\n",
        "  total_correct = 0.0\n",
        "  for i in range(0, len(dataset.features), batch_size):\n",
        "    b_features, b_rows, b_cols, b_ys, b_masks = batch(\n",
        "        dataset.features[i:i + batch_size], dataset.rows[i:i + batch_size],\n",
        "        dataset.columns[i:i + batch_size], dataset.labels[i:i + batch_size],\n",
        "        dataset.root_nodes[i:i + batch_size])\n",
        "\n",
        "    accs = model.accuracy(params_to_evaluate, b_features, b_rows, b_cols, b_ys,\n",
        "                          b_masks)\n",
        "    total_correct += accs * len(dataset.features[i:i + batch_size])\n",
        "  return total_correct / len(dataset.features)\n",
        "\n",
        "\n",
        "def print_accuracies(params_to_evaluate,\n",
        "                     dataset_test,\n",
        "                     dataset_train,\n",
        "                     batch_size=100):\n",
        "  train_accuracy = compute_accuracies(\n",
        "      params_to_evaluate, dataset=train_dataset, batch_size=batch_size)\n",
        "  test_accuracy = compute_accuracies(\n",
        "      params_to_evaluate, dataset=test_dataset, batch_size=batch_size)\n",
        "\n",
        "  combined_accuracy = np.average(\n",
        "      [train_accuracy, test_accuracy],\n",
        "      weights=[len(dataset_train.features),\n",
        "               len(dataset_test.features)])\n",
        "  print(f'Train accuracy: {train_accuracy:.3f} | '\n",
        "        f'Test accuracy: {test_accuracy:.3f} | '\n",
        "        f'Combined accuracy: {combined_accuracy:.3f}')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "26DGTXc1Ayth"
      },
      "source": [
        "To replicate the Figure 3a from the paper, it is sufficient to use pre-trained set of parameters which were trained for 100 epochs on the 4th degree coefficient of S9. To do so, leave the above `use_pretrained_weights` set to `True` and the (**much** slower) training loop can be skipped.\n",
        "\n",
        "To replicate the results from scratch, set `use_pretrained_weights=False` and perform training from a fresh set of parameters. The final results should be visible after a large number of epochs, and although full convergence is usually achieved before 100 epochs it is still expected to take an hour on GPU and even longer on CPU."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "9k2gMxcPHFeX"
      },
      "outputs": [],
      "source": [
        "#@title Perform training / Load pretrained weights\n",
        "\n",
        "if use_pretrained_weights:\n",
        "  print(\"Loading pre-trained weights\")\n",
        "  flat_trained_params = jax.numpy.load(\n",
        "      os.path.join(DATA_DIR, \"trained_params.npz\"))\n",
        "  trained_params = collections.defaultdict(dict)\n",
        "  for key, array in flat_trained_params.items():\n",
        "    layer, weight_or_bias = key.split()\n",
        "    assert weight_or_bias in (\"w\", \"b\"), weight_or_bias\n",
        "    assert \"linear\" in layer, layer\n",
        "    trained_params[layer][weight_or_bias] = array\n",
        "  trained_params = dict(trained_params)\n",
        "else:\n",
        "  num_epochs = 20\n",
        "  trained_params = model.net.init(\n",
        "      jax.random.PRNGKey(42),\n",
        "      features=train_dataset.features[0],\n",
        "      rows=train_dataset.rows[0],\n",
        "      cols=train_dataset.columns[0],\n",
        "      batch_size=1,\n",
        "      masks=train_dataset.features[0][np.newaxis, :, :])\n",
        "  trained_opt_state = opt_init(trained_params)\n",
        "\n",
        "  for ep in range(1, num_epochs + 1):\n",
        "    tr_data = list(\n",
        "        zip(\n",
        "            train_dataset.features,\n",
        "            train_dataset.rows,\n",
        "            train_dataset.columns,\n",
        "            train_dataset.labels,\n",
        "            train_dataset.root_nodes,\n",
        "        ))\n",
        "    random.shuffle(tr_data)\n",
        "    features_train, rows_train, cols_train, ys_train, root_nodes_train = zip(\n",
        "        *tr_data)\n",
        "\n",
        "    features_train = list(features_train)\n",
        "    rows_train = list(rows_train)\n",
        "    cols_train = list(cols_train)\n",
        "    ys_train = np.array(ys_train)\n",
        "    root_nodes_train = list(root_nodes_train)\n",
        "\n",
        "    for i in range(0, len(features_train), batch_size):\n",
        "      b_features, b_rows, b_cols, b_ys, b_masks = batch(\n",
        "          features_train[i:i + batch_size],\n",
        "          rows_train[i:i + batch_size],\n",
        "          cols_train[i:i + batch_size],\n",
        "          ys_train[i:i + batch_size],\n",
        "          root_nodes_train[i:i + batch_size],\n",
        "      )\n",
        "\n",
        "      trained_params, trained_opt_state, curr_loss = train(\n",
        "          trained_params,\n",
        "          trained_opt_state,\n",
        "          b_features,\n",
        "          b_rows,\n",
        "          b_cols,\n",
        "          b_ys,\n",
        "          b_masks,\n",
        "      )\n",
        "\n",
        "      accs = model.accuracy(\n",
        "          trained_params,\n",
        "          b_features,\n",
        "          b_rows,\n",
        "          b_cols,\n",
        "          b_ys,\n",
        "          b_masks,\n",
        "      )\n",
        "      print(datetime.datetime.now(),\n",
        "            f\"Iteration {i:4d} | Batch loss {curr_loss:.6f}\",\n",
        "            f\"Batch accuracy {accs:.2f}\")\n",
        "\n",
        "    print(datetime.datetime.now(), f\"Epoch {ep:2d} completed!\")\n",
        "\n",
        "    # Calculate accuracy across full dataset once per epoch\n",
        "    print(datetime.datetime.now(), f\"Epoch {ep:2d}       | \", end=\"\")\n",
        "    print_accuracies(trained_params, test_dataset, train_dataset)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "_Q75g0GgCimX",
        "scrolled": false
      },
      "outputs": [],
      "source": [
        "#@title Print model accuracies\n",
        "#@markdown Baseline accuracy should be ~88%; trained accuracy should be ~98%.\n",
        "\n",
        "#@markdown If only 20 epochs are trained for (as is the default setting above\n",
        "#@markdown for training from scratch), the overall accuracy will be between\n",
        "#@markdown the two, near 95%.\n",
        "print('Baseline accuracy', 1 - np.mean(train_dataset.labels))\n",
        "print_accuracies(trained_params, test_dataset, train_dataset)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "G3B3a8KgX8o4"
      },
      "outputs": [],
      "source": [
        "#@title Calculate salience and aggregate by edge labels\n",
        "def get_salience_vectors(salience_fn, params, full_dataset):\n",
        "  salient_features_arr = []\n",
        "  for i in range(0, len(full_dataset.features), batch_size):\n",
        "    b_features, b_rows, b_cols, b_ys, b_masks = batch(\n",
        "        full_dataset.features[i:i + batch_size],\n",
        "        full_dataset.rows[i:i + batch_size],\n",
        "        full_dataset.columns[i:i + batch_size],\n",
        "        full_dataset.labels[i:i + batch_size],\n",
        "        full_dataset.root_nodes[i:i + batch_size],\n",
        "    )\n",
        "    salient_features = salience_fn(params, b_features, b_rows, b_cols, b_ys,\n",
        "                                   b_masks)\n",
        "    effective_batch_size = len(full_dataset.features[i:i + batch_size])\n",
        "    salient_features_arr.extend(\n",
        "        np.reshape(salient_features, [effective_batch_size, -1, 2]))\n",
        "  return salient_features_arr\n",
        "\n",
        "\n",
        "def aggregate_by_edges(salient_features_arr, cutoff, ys):\n",
        "  refl_count = {\n",
        "      'salient_all': collections.defaultdict(int),\n",
        "      'all': collections.defaultdict(int)\n",
        "  }\n",
        "  for graph_index, (graph, saliency, label) in enumerate(\n",
        "      zip(iter_graph(), salient_features_arr, ys)):\n",
        "    [salient_nodes] = np.where(np.linalg.norm(saliency, axis=1) \u003e cutoff)\n",
        "    subgraph = graph.subgraph(salient_nodes)\n",
        "    for reflection in get_reflections(graph_index, graph):\n",
        "      refl_count['all'][reflection] += 1\n",
        "    for reflection in get_reflections(graph_index, subgraph):\n",
        "      refl_count['salient_all'][reflection] += 1\n",
        "\n",
        "  norm_refl_mat = {}\n",
        "  for title, counts in refl_count.items():\n",
        "    reflection_mat = np.zeros((9, 9))\n",
        "    # Loop over the upper triangle.\n",
        "    for i in range(9):\n",
        "      for j in range(i + 1, 9):\n",
        "        count = counts[(i, j)] + counts[(j, i)]\n",
        "        reflection_mat[i, j] = count\n",
        "        reflection_mat[j, i] = count\n",
        "    norm_refl_mat[title] = reflection_mat / reflection_mat.sum()\n",
        "\n",
        "  return refl_count, norm_refl_mat\n",
        "\n",
        "\n",
        "def get_reflections(graph_index, graph):\n",
        "  node_index_to_permutation = graph_index_to_node_index_to_permutation[str(\n",
        "      graph_index)]\n",
        "  for permutation_x, permutation_y in graph.edges():\n",
        "    if np.isscalar(permutation_x):\n",
        "      # If the data was loaded as compressed sci-py arrays, the permutations\n",
        "      # need to be looked up by index in the data loaded separate from JSON.\n",
        "      permutation_x = node_index_to_permutation[str(permutation_x)]\n",
        "      permutation_y = node_index_to_permutation[str(permutation_y)]\n",
        "    yield tuple(i for i, (x, y) in enumerate(zip(permutation_x, permutation_y))\n",
        "                if x != y)\n",
        "\n",
        "\n",
        "print('Computing saliences...')\n",
        "salience_fn = jax.jit(jax.grad(lambda *args: jnp.sum(model.loss(*args)), 1))\n",
        "salient_features_arr = get_salience_vectors(salience_fn, trained_params,\n",
        "                                            full_dataset)\n",
        "saliencies = np.linalg.norm(\n",
        "    np.concatenate(salient_features_arr, axis=0), axis=1)\n",
        "\n",
        "print('Aggregating by edges...')\n",
        "cutoff = np.percentile(saliencies, 99)\n",
        "refl_count, norm_refl_mat = aggregate_by_edges(salient_features_arr, cutoff,\n",
        "                                               full_dataset.labels)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PyYUnkVfAytj"
      },
      "source": [
        "The final cell replicates Figure 3a from the paper - it shows the relative frequency of different edge types in salient subgraphs compared with the frequency across the full dataset."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "TOWHb6CiZsc2"
      },
      "outputs": [],
      "source": [
        "#@title Plot edge attribution\n",
        "\n",
        "font = {'family': 'normal', 'weight': 'bold', 'size': 18}\n",
        "\n",
        "matplotlib.rc('font', **font)\n",
        "sns.set_style('ticks')\n",
        "\n",
        "np.fill_diagonal(norm_refl_mat['all'], 1)  # Avoid 0/0 warning.\n",
        "change_grid = ((norm_refl_mat['salient_all'] - norm_refl_mat['all']) /\n",
        "               norm_refl_mat['all'] * 100)\n",
        "\n",
        "f, ax = plt.subplots(figsize=(10, 10))\n",
        "ax = sns.heatmap(\n",
        "    change_grid,\n",
        "    mask=np.triu(np.ones_like(change_grid)),\n",
        "    center=0,\n",
        "    square=True,\n",
        "    cmap='RdBu',\n",
        "    cbar_kws={'shrink': .82},\n",
        "    ax=ax,\n",
        "    vmin=-50,\n",
        "    vmax=50)\n",
        "\n",
        "ax.set_ylabel('1st reflection index')\n",
        "ax.set_xlabel('2nd reflection index')\n",
        "sns.despine()\n",
        "\n",
        "plt.show()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "Representation_theory",
      "private_outputs": true,
      "provenance": [
        {
          "file_id": "1UzqhYNN_d8HM0gHXtbl88yHoiTETGqQU",
          "timestamp": 1632429941242
        }
      ]
    },
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.9.7"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}