knot_theory.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "n-vSyzB-qmir"
      },
      "source": [
        "\n",
        "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": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "xQnp1V1cvdxy"
      },
      "outputs": [],
      "source": [
        "# @title Install required modules\n",
        "from IPython.display import clear_output\n",
        "\n",
        "!pip install dm-haiku\n",
        "!pip install optax\n",
        "clear_output()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "3NFhPsG4u1L1"
      },
      "outputs": [],
      "source": [
        "# @title Imports\n",
        "import tempfile\n",
        "\n",
        "import haiku as hk\n",
        "import jax\n",
        "import jax.numpy as jnp\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import optax\n",
        "import pandas as pd\n",
        "import seaborn as sns\n",
        "from sklearn.model_selection import train_test_split"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "l10N2ZbHu6Ob"
      },
      "outputs": [],
      "source": [
        "# @title Download data\n",
        "\n",
        "_, input_filename = tempfile.mkstemp()\n",
        "!gsutil cp \"gs://maths_conjectures/knot_theory/knot_theory_invariants.csv\" {input_filename}\n",
        "\n",
        "full_df = pd.read_csv(input_filename)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "C-kZ_rW-h8Nk"
      },
      "outputs": [],
      "source": [
        "# @title Load and preprocess data\n",
        "\n",
        "#@markdown The columns of the dataset which will make up the inputs to the network.\n",
        "#@markdown In other words, for a knot k, X(k) will be a vector consisting of these quantities. In this case, these are the geometric invariants of each knot.\n",
        "#@markdown For descriptions of these invariants see https://knotinfo.math.indiana.edu/\n",
        "display_name_from_short_name = {\n",
        "    'chern_simons': 'Chern-Simons',\n",
        "    'cusp_volume': 'Cusp volume',\n",
        "    'hyperbolic_adjoint_torsion_degree': 'Adjoint Torsion Degree',\n",
        "    'hyperbolic_torsion_degree': 'Torsion Degree',\n",
        "    'injectivity_radius': 'Injectivity radius',\n",
        "    'longitudinal_translation': 'Longitudinal translation',\n",
        "    'meridinal_translation_imag': 'Re(Meridional translation)',\n",
        "    'meridinal_translation_real': 'Im(Meridional translation)',\n",
        "    'short_geodesic_imag_part': 'Im(Short geodesic)',\n",
        "    'short_geodesic_real_part': 'Re(Short geodesic)',\n",
        "    'Symmetry_0': 'Symmetry: $0$',\n",
        "    'Symmetry_D3': 'Symmetry: $D_3$',\n",
        "    'Symmetry_D4': 'Symmetry: $D_4$',\n",
        "    'Symmetry_D6': 'Symmetry: $D_6$',\n",
        "    'Symmetry_D8': 'Symmetry: $D_8$',\n",
        "    'Symmetry_Z/2 + Z/2': 'Symmetry: $\\\\frac{Z}{2} + \\\\frac{Z}{2}$',\n",
        "    'volume': 'Volume',\n",
        "}\n",
        "column_names = list(display_name_from_short_name)\n",
        "target = 'signature'\n",
        "\n",
        "#@markdown Split the data into a training, a validation and a holdout test set. To check\n",
        "#@markdown the robustness of the model and any proposed relationship, the training\n",
        "#@markdown process can be repeated with multiple different train/validation/test splits.\n",
        "\n",
        "random_seed = 2 # @param {type: \"integer\"}\n",
        "random_state = np.random.RandomState(random_seed)\n",
        "train_df, validation_and_test_df = train_test_split(\n",
        "    full_df, random_state=random_state)\n",
        "validation_df, test_df = train_test_split(\n",
        "    validation_and_test_df, test_size=.5, random_state=random_state)\n",
        "\n",
        "# Find bounds for the signature in the training dataset.\n",
        "max_signature = train_df[target].max()\n",
        "min_signature = train_df[target].min()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "jmBkFuYGu_j5"
      },
      "outputs": [],
      "source": [
        "# @title Network Definition\n",
        "\n",
        "#@markdown Create a simple feedforward network, using the DM Haiku library\n",
        "#@markdown (https://github.com/deepmind/dm-haiku).\n",
        "\n",
        "#@markdown The output of the network is a predicted categorical distribution, represented\n",
        "#@markdown by a vector q, where softmax(q)[i] is the predicted probability that the\n",
        "#@markdown signature of the knot is equal to 2*i + min_signature. Note that signature is\n",
        "#@markdown always an even integer.\n",
        "\n",
        "#@markdown We take the cross entropy between this distribution and the true distribution\n",
        "#@markdown (i.e. 1 at the true value of the signature, 0 everywhere else) as the loss\n",
        "#@markdown function.\n",
        "\n",
        "\n",
        "def net_forward(inp):\n",
        "  return hk.Sequential([\n",
        "      hk.Linear(300),\n",
        "      jax.nn.sigmoid,\n",
        "      hk.Linear(300),\n",
        "      jax.nn.sigmoid,\n",
        "      hk.Linear(300),\n",
        "      jax.nn.sigmoid,\n",
        "      hk.Linear(int((max_signature - min_signature) / 2)),\n",
        "  ])(\n",
        "      inp)\n",
        "\n",
        "\n",
        "def softmax_cross_entropy(logits, labels):\n",
        "  # Labels are the true values of the signature\n",
        "  one_hot = jax.nn.one_hot((labels - min_signature) / 2, logits.shape[-1])\n",
        "  return -jnp.sum(jax.nn.log_softmax(logits) * one_hot, axis=-1)\n",
        "\n",
        "\n",
        "# The cross-entropy loss is composed with the network predictions, to define\n",
        "# `loss_fn` as a function on X and y.\n",
        "def loss_fn(inputs, labels):\n",
        "  return jnp.mean(softmax_cross_entropy(net_forward(inputs), labels))\n",
        "\n",
        "\n",
        "# Haiku network transformation steps.\n",
        "loss_fn_t = hk.without_apply_rng(hk.transform(loss_fn))\n",
        "net_forward_t = hk.without_apply_rng(hk.transform(net_forward))\n",
        "\n",
        "\n",
        "@jax.jit\n",
        "def predict(params, data_X):\n",
        "  return (np.argmax(net_forward_t.apply(params, data_X), axis=1) * 2 +\n",
        "          min_signature)\n",
        "\n",
        "\n",
        "#@markdown Calculate the mean and standard deviation over each column in the training\n",
        "#@markdown dataset. We use this to normalize each feature, this is best practice for\n",
        "#@markdown inputting features into a network, but is also very important in this case\n",
        "#@markdown to ensure the gradients used for saliency are meaningfully comparable.\n",
        "def normalize_features(df, cols, add_target=True):\n",
        "  features = df[cols]\n",
        "  sigma = features.std()\n",
        "  if any(sigma == 0):\n",
        "    print(sigma)\n",
        "    raise RuntimeError(\n",
        "        \"A poor data stratification has led to no variation in one of the data \"\n",
        "        \"splits for at least one feature (ie std=0). Restratify and try again.\")\n",
        "  mu = features.mean()\n",
        "  normed_df = (features - mu) / sigma\n",
        "  if add_target:\n",
        "    normed_df[target] = df[target]\n",
        "  return normed_df\n",
        "\n",
        "\n",
        "def get_batch(df, cols, size=None):\n",
        "  batch_df = df if size is None else df.sample(size)\n",
        "  X = batch_df[cols].to_numpy()\n",
        "  y = batch_df[target].to_numpy()\n",
        "  return X, y\n",
        "\n",
        "\n",
        "normed_train_df = normalize_features(train_df, column_names)\n",
        "normed_validation_df = normalize_features(validation_df, column_names)\n",
        "normed_test_df = normalize_features(test_df, column_names)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "_ZyvOA1g_08D"
      },
      "outputs": [],
      "source": [
        "# @title Network Setup (re-run before re-training)\n",
        "\n",
        "batch_size = 64\n",
        "learning_rate = 0.001\n",
        "num_training_steps = 10_000\n",
        "validation_interval = 100\n",
        "\n",
        "train_X, train_y = get_batch(normed_train_df, column_names, batch_size)\n",
        "\n",
        "# Pick a random seed for the network weights. To check the robustness of the\n",
        "# model, the training process can be repeated with multiple different random\n",
        "# seeds.\n",
        "rng = jax.random.PRNGKey(1)\n",
        "init_params = loss_fn_t.init(rng, train_X, train_y)\n",
        "\n",
        "opt_init, opt_update = optax.adam(learning_rate)\n",
        "opt_state = opt_init(init_params)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "sofhcPK4vZWl"
      },
      "outputs": [],
      "source": [
        "# @title Network Training\n",
        "\n",
        "# We train until the validation loss stops decreasing, checking every \u003cvalidation_interval\u003e steps,\n",
        "# up to a maximum of 10k steps.\n",
        "\n",
        "\n",
        "@jax.jit\n",
        "def update(params, opt_state, batch_X, batch_y):\n",
        "  grads = jax.grad(loss_fn_t.apply)(params, batch_X, batch_y)\n",
        "  upds, new_opt_state = opt_update(grads, opt_state)\n",
        "  new_params = optax.apply_updates(params, upds)\n",
        "  return new_params, new_opt_state\n",
        "\n",
        "\n",
        "def train(columns_to_train_on, params, opt_state, update_fn):\n",
        "  best_validation_loss = np.inf\n",
        "  for i in range(num_training_steps):\n",
        "    train_X, train_y = get_batch(normed_train_df, columns_to_train_on,\n",
        "                                 batch_size)\n",
        "    params, opt_state = update_fn(params, opt_state, train_X, train_y)\n",
        "\n",
        "    if i % validation_interval == 0:\n",
        "      # Run validation on the full validation dataset.\n",
        "      validation_X, validation_y = get_batch(normed_validation_df,\n",
        "                                             columns_to_train_on)\n",
        "      train_loss = loss_fn_t.apply(params, train_X, train_y)\n",
        "      validation_loss = loss_fn_t.apply(params, validation_X, validation_y)\n",
        "      print(f\"Step count: {i}\")\n",
        "      print(f\"Train loss: {train_loss}\")\n",
        "      print(f\"Validation loss: {validation_loss}\")\n",
        "\n",
        "      if validation_loss \u003e best_validation_loss:\n",
        "        print(\"Validation loss increased. Stopping!\")\n",
        "        return params\n",
        "      else:\n",
        "        best_validation_loss = validation_loss\n",
        "  return params\n",
        "\n",
        "\n",
        "trained_params = train(column_names, init_params, opt_state, update)\n",
        "# Print the test accuracy, i.e. the proportion of the knots for which the\n",
        "# network predicts the correct signature.\n",
        "test_X, test_y = get_batch(normed_test_df, column_names)\n",
        "\n",
        "# The final below accuracy should be in the low 80%s.\n",
        "print(\"Test Accuracy: \",\n",
        "      np.mean((predict(trained_params, test_X) - test_y) == 0))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qj0TjRTnAmGW"
      },
      "source": [
        "The below cell replicates Figure 2a from the paper, though for simplicity omitting the error bars.\n",
        "\n",
        "To compute the saliency, we take the gradient of the loss function, with respect to each of the the components of the X input to the network (i.e. the geometric invariants of each knot), and average over the dataset.\n",
        "\n",
        "We plot the feature saliencies in decreasing order. The plot (should!) show that the overall loss is influenced far more by the longitudinal translation and the real and imaginary parts of the meridional translation than any of the other invariants."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "j1CEbkUgxtGZ",
        "scrolled": true
      },
      "outputs": [],
      "source": [
        "# @title Saliency Analysis\n",
        "train_X = normalize_features(train_df, column_names, add_target=False).to_numpy()\n",
        "train_y = train_df[target].to_numpy()\n",
        "\n",
        "\n",
        "saliencies = np.mean(\n",
        "    np.abs(jax.grad(loss_fn_t.apply, 1)(trained_params, train_X, train_y)), axis=0)\n",
        "\n",
        "\n",
        "decreasing_saliency = reversed(sorted(zip(saliencies, display_name_from_short_name.values())))\n",
        "sorted_saliencies, sorted_columns = zip(*decreasing_saliency)\n",
        "\n",
        "fig, ax = plt.subplots(figsize=(6,8))\n",
        "sns.barplot(y=np.array(sorted_columns),\n",
        "            x=np.array(sorted_saliencies) / max(sorted_saliencies),\n",
        "            color=\"#0077c6\");\n",
        "\n",
        "ax.tick_params(labelsize=15);\n",
        "ax.set_ylabel('Geometric invariants X(z)', fontsize=20);\n",
        "plt.xlabel('Normalized attribution score', fontsize=20);"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "EsDbwItnlNUY"
      },
      "outputs": [],
      "source": [
        "# @title Confirming the Feature Saliency\n",
        "\n",
        "#@markdown To confirm the results of the saliency analysis, we re-train the network with\n",
        "#@markdown only these three features as input to the network.\n",
        "salient_column_names = ['longitudinal_translation',\n",
        "                        'meridinal_translation_imag',\n",
        "                        'meridinal_translation_real']\n",
        "target = 'signature'"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "SKY75X_sEsy2"
      },
      "outputs": [],
      "source": [
        "# @title Confirming the Feature Saliency: Network Setup (re-run before re-training)\n",
        "\n",
        "train_X, train_y = get_batch(normed_train_df, salient_column_names, batch_size)\n",
        "\n",
        "init_params_salient = loss_fn_t.init(rng, train_X, train_y)\n",
        "\n",
        "opt_init_salient, opt_update_salient = optax.adam(learning_rate)\n",
        "opt_state_salient = opt_init(init_params_salient)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "lGOfm424lXtR"
      },
      "outputs": [],
      "source": [
        "# @title Confirming the Feature Saliency: Network Training\n",
        "\n",
        "\n",
        "#@markdown Re-train the network using only the most salient features.\n",
        "@jax.jit\n",
        "def update_salient(params, opt_state, batch_X, batch_y):\n",
        "  grads = jax.grad(loss_fn_t.apply)(params, batch_X, batch_y)\n",
        "  upds, new_opt_state = opt_update_salient(grads, opt_state)\n",
        "  new_params = optax.apply_updates(params, upds)\n",
        "  return new_params, new_opt_state\n",
        "\n",
        "\n",
        "trained_params_salient = train(salient_column_names, init_params_salient,\n",
        "                               opt_state_salient, update_salient)\n",
        "\n",
        "#@markdown Print the test accuracy. This should be very similar to the test accuracy in\n",
        "#@markdown the case that all columns / invariants are included, demonstrating that most\n",
        "#@markdown of the predictve information about the signature is contained in the three\n",
        "#@markdown selected invariants.\n",
        "test_X, test_y = get_batch(normed_test_df, salient_column_names)\n",
        "\n",
        "#@markdown The final below accuracy should be in the low 80%s, probably 0.8 -\u003e 0.85\n",
        "print(\"Test Accuracy: \",\n",
        "      np.mean((predict(trained_params_salient, test_X) - test_y) == 0))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "ZIqvjO_O8Zti"
      },
      "outputs": [],
      "source": [
        "# @title Slope vs. Signature: Proposed Linear Relationship\n",
        "\n",
        "#@markdown The quantity we proposed was the \"natural slope\", given by\n",
        "#@markdown real(longitudinal_translation / meridinal_translation). We show that\n",
        "#@markdown this is approximately twice the signature (up to a correction term based on\n",
        "#@markdown other hyperbolic invariants) which we can check by comparing the predictions\n",
        "#@markdown made by this rule to those made by the previously trained models.\n",
        "\n",
        "\n",
        "def predict_signature_from_slope(data_X, min_signature, max_signature):\n",
        "  meridinal_translation = (\n",
        "      data_X['meridinal_translation_real'] +\n",
        "      1j * data_X['meridinal_translation_imag'])\n",
        "  slope = data_X['longitudinal_translation'] / meridinal_translation\n",
        "  return slope.real / 2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "KAwLtNFxTwCc"
      },
      "outputs": [],
      "source": [
        "# @title Proposed Linear Relationship: Scatter plot\n",
        "\n",
        "#@markdown Scatter plot of the slope against predicted signature.\n",
        "predictions = [\n",
        "    predict_signature_from_slope(x, min_signature, max_signature)\n",
        "    for _, x in test_df.iterrows()\n",
        "]\n",
        "\n",
        "fig, ax = plt.subplots(figsize=(8, 8))\n",
        "sns.scatterplot(\n",
        "    x=predictions, y=test_df[target], alpha=0.2)\n",
        "ax.set_xlabel('Predicted Signature')\n",
        "ax.set_ylabel('Signature')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "HkoLDdYV_Idv"
      },
      "outputs": [],
      "source": [
        "# @title Proposed Linear Relationship: Test Accuracy\n",
        "\n",
        "\n",
        "#@markdown In order to compute the \"test accuracy\" in the same way as before, we quantize\n",
        "#@markdown the predicted signature values to even integers, between min_signature and\n",
        "#@markdown max_signature.\n",
        "def quantize(x, min_signature, max_signature):\n",
        "  return min(max(2 * round(x / 2), min_signature), max_signature)\n",
        "\n",
        "\n",
        "quantized_predictions = [\n",
        "    quantize(x, min_signature, max_signature) for x in predictions\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZhJAExt6739Y"
      },
      "source": [
        "Now we can compute the \"test accuracy\" of this prediction.\n",
        "\n",
        "The value is slightly lower than for the trained network predictors (although still far higher than chance), but this is not unexpected.\n",
        "\n",
        "Indeed, the proposed rule gives a provable bound on the signature over all knots, instead of maximizing prediction performance over a given finite dataset, as the networks are doing. Although we do use separate training, validation and test datasets for the networks, these are all drawn from approximately the same distribution, whereas the proposed rule could be considered (in some imprecise sense) to have been \"trained\" on the set of all knots, a very different distribution.\n",
        "\n",
        "The correction term may also have some bias, for example tending to be positive more often than it is negative, information which the network predictors would be able to use to increase their prediction performance relative to that of the proposed rule."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2m-RE-PQ74GC"
      },
      "outputs": [],
      "source": [
        "#@markdown The below accuracy will probably be lower than the previous ~80%, but not by much, likely still \u003e70%\n",
        "print(\"Test Accuracy: \", np.mean(test_df[target] - quantized_predictions == 0))"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "Knot_theory",
      "private_outputs": true,
      "provenance": [
        {
          "file_id": "1tKuhjAHpuTUr5fs683eWFNPAuhmLt0uP",
          "timestamp": 1632418370828
        }
      ]
    },
    "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
}