Feynman diagram transformer demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8161e120",
   "metadata": {},
   "source": [
    "# SYMBA Model for QED Feynamn diagram:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "dc756052",
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import layers\n",
    "from tensorflow.keras.layers import TextVectorization\n",
    "import re\n",
    "import time"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3f4f107",
   "metadata": {},
   "source": [
    "You can find the QED Feynman dataset here:\n",
    "https://drive.google.com/file/d/1dHeZcWwPqew5EM3pBr1rp9mbtZd6KMJ9/view?usp=share_link"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "6ff324d5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "data size:  257991\n"
     ]
    }
   ],
   "source": [
    "with open('fey_qed_order.txt', 'r', encoding='utf-8') as f:\n",
    "    lines = f.read().split('\\n')\n",
    "    \n",
    "\n",
    "\n",
    "text_pairs =[]\n",
    "for line in lines[: min(len(lines), len(lines)-1)]:\n",
    "    intr, amp, sqamp, t  = line.split('>')\n",
    "    sqamp = \"[start] \" + sqamp + \" [end]\"\n",
    "    text_pairs.append((intr, amp,sqamp, float(t) ))\n",
    "    \n",
    "text_pairs = list(set(text_pairs))\n",
    "random.seed(3333)\n",
    "random.shuffle(text_pairs)\n",
    "\n",
    "\n",
    "print('data size: ', len(text_pairs))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad1114d0",
   "metadata": {},
   "source": [
    "Example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "8b138d9f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fynman diagram:   \n",
      "  Connection on vertex V_0 :   Off e(X_3)   Off A(X_5) Anti Off e(V_0)  , Connection on vertex V_2 :    u(X_1)   Off u(X_4)   Off A(V_2)  , Connection on vertex V_1 :   Off e(X_2)   Off A(V_1)   Off e(V_1)  ,  \n",
      "\n",
      "Squared Amplitude:   \n",
      " [start] ([m_e^4, m_e^2, 1], [-128*(2*m_u^2 - s_14), 64*(m_u^2*(3*s_23 + 4*s_25 - 4*s_35 - 4*s_55) - 3*s_12*s_34 - 4*s_12*s_45 - 3*s_13*s_24 + 2*s_14*s_35 + 2*s_14*s_55 - 4*s_15*s_24), -64*(m_u^2*(s_23*s_33 - s_23*s_55 + 2*s_25*s_33 + 2*s_25*s_35) - s_12*s_33*s_34 - 2*s_12*s_33*s_45 + s_12*s_34*s_55 - 2*s_12*s_35*s_45 - s_13*s_24*s_33 + s_13*s_24*s_55 - 2*s_15*s_24*s_33 - 2*s_15*s_24*s_35)], [9*(m_u^2 - 2*s_14 + s_44)^2*(-m_e^2 + s_33 + 2*s_35 + s_55)^2]) [end]\n"
     ]
    }
   ],
   "source": [
    "print('Fynman diagram:  ' , '\\n', text_pairs[1][1], '\\n' )\n",
    "print('Squared Amplitude:  ', '\\n' ,text_pairs[1][2] )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "ef65484b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "data size:  257982\n"
     ]
    }
   ],
   "source": [
    "text_pairs1 = []\n",
    "\n",
    "for i in range(len(text_pairs)):\n",
    "    if  len(text_pairs[i][2]) < 1800:\n",
    "        text_pairs1.append(text_pairs[i])\n",
    "\n",
    "text_pairs = text_pairs1\n",
    "print('data size: ', len(text_pairs))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dddc3676",
   "metadata": {},
   "source": [
    "### Data preprocessing:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "6e8aaf8c",
   "metadata": {},
   "outputs": [],
   "source": [
    "#preprocessing for the amplitudes:\n",
    "\n",
    "def prepro_feyn(data):\n",
    "\n",
    "    for r in (('(', '('), (')', ')'), ('  ', ' '), (' e(', 'e(m_e,-1,' ),(' mu(', 'mu(m_mu,-1,') ,(' u(', ' u(m_u,2/3,'), (' d(', 'd(m_d,-1/3,'), (' t(', ' t(m_t,-1,') ,(' s(', 's(m_s,-1/3,'), (' tt(', ' tt(m_tt,-1,'), (' c(', 'c(m_c,2/3,'),(' b(', 'b(m_b,-1/3,'), ('Anti ', 'Anti,'), ('Off ', 'Off,'), ('  ', ' ') ): \n",
    "        data = data.replace(*r) \n",
    "        \n",
    "    return data\n",
    "\n",
    "#preprocessing for the squared amplitudes:\n",
    "\n",
    "def prepro_squared_ampl(data):\n",
    "\n",
    "    for r in (('*', '*'), (',', ' , '), ('*(', ' *( ') , ('([', '[ '), ('])', ' ]'), ('[', '[ '), (']', ' ]'), ('[ start ]', '[start]'), ('[ end ]', '[end]'), (' - ', ' -'), (' + ',' +' ) ,('/', ' / ') ,('  ', ' ')) :\n",
    "        data = data.replace(*r) \n",
    "    data = re.sub(r\"\\*(s_\\d+\\*s_\\d+)\", r\"* \\1\", data)\n",
    "    data = re.sub(r\"\\*(s_\\d+\\^\\d+\\*s_\\d+)\", r\"* \\1\", data)\n",
    "    data = re.sub(r\"\\*(m_\\w+\\^\\d+\\*s_\\d+)\", r\"* \\1\", data)\n",
    "    data = re.sub(r\"(m_\\w+\\^\\d+)\", r\" \\1 \", data)\n",
    "    data = data.replace('  ', ' ')\n",
    "    \n",
    "    \n",
    "    return data\n",
    "\n",
    "\n",
    "\n",
    "text_pairs_prep = []\n",
    "for i in range(len(text_pairs)):\n",
    "    text_pairs_prep.append((text_pairs[i][0], prepro_feyn(text_pairs[i][1]), prepro_squared_ampl(text_pairs[i][2]), text_pairs[i][3]))\n",
    "\n",
    "text_pairs = text_pairs_prep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "be922576",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fynman diagram:   \n",
      "  Connection on vertex V_0 : Offe(m_e,-1,X_3) Off,A(X_5) Anti,Offe(m_e,-1,V_0) , Connection on vertex V_2 : u(m_u,2/3,X_1) Off,u(m_u,2/3,X_4) Off,A(V_2) , Connection on vertex V_1 : Offe(m_e,-1,X_2) Off,A(V_1) Offe(m_e,-1,V_1) ,  \n",
      "\n",
      "Squared Amplitude:   \n",
      " [start] [ m_e^4 , m_e^2 , 1 ] , [ -128 *( 2* m_u^2 -s_14) , 64 *( m_u^2 *( 3*s_23 +4*s_25 -4*s_35 -4*s_55) -3* s_12*s_34 -4* s_12*s_45 -3* s_13*s_24 +2* s_14*s_35 +2* s_14*s_55 -4* s_15*s_24) , -64 *( m_u^2 *( s_23*s_33 -s_23*s_55 +2* s_25*s_33 +2* s_25*s_35) -s_12* s_33*s_34 -2* s_12*s_33*s_45 +s_12* s_34*s_55 -2* s_12*s_35*s_45 -s_13* s_24*s_33 +s_13* s_24*s_55 -2* s_15*s_24*s_33 -2* s_15*s_24*s_35) ] , [ 9 *( m_u^2 -2*s_14 +s_44)^2 *( - m_e^2 +s_33 +2*s_35 +s_55)^2 ] [end]\n"
     ]
    }
   ],
   "source": [
    "print('Fynman diagram:  ' , '\\n', text_pairs[1][1], '\\n' )\n",
    "print('Squared Amplitude:  ', '\\n' ,text_pairs[1][2] )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eed0789c",
   "metadata": {},
   "source": [
    "###  Maximum sequence length:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "01d52235",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Maximum sequence length of Feynman diagram        : 27\n",
      "Maximum sequence length of squared amplitudes: 194\n"
     ]
    }
   ],
   "source": [
    "def max_len(sq_data):\n",
    "    l = len(sq_data[sq_data.index(max(sq_data, key=len))].split())\n",
    "    return l\n",
    "\n",
    "\n",
    "feyn = [pair[1] for pair in text_pairs]\n",
    "sq_ampl= [pair[2] for pair in text_pairs]\n",
    "\n",
    "print( 'Maximum sequence length of Feynman diagram        :' ,max_len(feyn))\n",
    "print( 'Maximum sequence length of squared amplitudes:' ,max_len(sq_ampl))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0ae2b34",
   "metadata": {},
   "source": [
    "# Tokenization:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9179c6cf",
   "metadata": {},
   "source": [
    "### Split the data:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "d1c783e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_val_samples = int(0.15 * len(text_pairs))\n",
    "num_train_samples = len(text_pairs) - 2 * num_val_samples\n",
    "train_pairs = text_pairs[:num_train_samples]\n",
    "val_pairs  = text_pairs[num_train_samples : num_train_samples + num_val_samples]\n",
    "test_pairs = text_pairs[num_train_samples + num_val_samples :]\n",
    "\n",
    "train_input_texts = [pair[1] for pair in train_pairs]\n",
    "train_output_texts = [pair[2] for pair in train_pairs]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "f7641601",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-12-07 11:44:17.765714: I tensorflow/core/platform/cpu_feature_guard.cc:151] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "number of input (Feynman diagram) tokens:   257\n",
      "number of target (squared amplitude) tokens:  2230\n"
     ]
    }
   ],
   "source": [
    "vocab_size = 2230\n",
    "sequence_length = 194\n",
    "batch_size = 64\n",
    "\n",
    "input_vectorization = TextVectorization(\n",
    "    max_tokens=vocab_size, output_mode=\"int\", output_sequence_length=sequence_length, standardize=None, )\n",
    "\n",
    "output_vectorization = TextVectorization(\n",
    "    max_tokens=vocab_size,\n",
    "    output_mode=\"int\",\n",
    "    output_sequence_length=sequence_length + 1, standardize=None)\n",
    "\n",
    "input_vectorization.adapt(train_input_texts)\n",
    "output_vectorization.adapt(train_output_texts)\n",
    "\n",
    "target_tokens = output_vectorization.get_vocabulary()\n",
    "input_tokens = input_vectorization.get_vocabulary()\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "print('number of input (Feynman diagram) tokens:  ', len(input_tokens))\n",
    "print('number of target (squared amplitude) tokens: ', len(target_tokens))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "ef0b279c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def format_dataset(input_exp, target_exp):\n",
    "    input_exp = input_vectorization(input_exp)\n",
    "    target_exp = output_vectorization(target_exp)\n",
    "    return ({\"encoder_inputs\": input_exp, \"decoder_inputs\": target_exp[:, :-1],}, target_exp[:, 1:])\n",
    "\n",
    "\n",
    "def make_dataset(pairs):\n",
    "    intr, fey_texts, sqampl_texts, t = zip(*pairs)\n",
    "    fey_texts = list(fey_texts)\n",
    "    sqampl_texts = list(sqampl_texts)\n",
    "    dataset = tf.data.Dataset.from_tensor_slices((fey_texts, sqampl_texts))\n",
    "    dataset = dataset.batch(batch_size)\n",
    "    dataset = dataset.map(format_dataset)\n",
    "    return dataset.shuffle(2048).prefetch(16).cache()\n",
    "\n",
    "\n",
    "train_ds = make_dataset(train_pairs)\n",
    "val_ds = make_dataset(val_pairs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a1f7453",
   "metadata": {},
   "source": [
    "# Transformer Model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "77bc56b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TransformerEncoder(layers.Layer):\n",
    "    def __init__(self, embed_dim, dense_dim, num_heads, **kwargs):\n",
    "        super(TransformerEncoder, self).__init__(**kwargs)\n",
    "        self.embed_dim = embed_dim\n",
    "        self.dense_dim = dense_dim\n",
    "        self.num_heads = num_heads\n",
    "        self.attention = layers.MultiHeadAttention(\n",
    "            num_heads=num_heads, key_dim=embed_dim\n",
    "        )\n",
    "        self.dense_proj = keras.Sequential(\n",
    "            [layers.Dense(dense_dim, activation=\"relu\"), layers.Dense(embed_dim),]\n",
    "        )\n",
    "        self.layernorm_1 = layers.LayerNormalization()\n",
    "        self.layernorm_2 = layers.LayerNormalization()\n",
    "        self.supports_masking = True\n",
    "\n",
    "    def call(self, inputs, mask=None):\n",
    "        if mask is not None:\n",
    "            padding_mask = tf.cast(mask[:, tf.newaxis, tf.newaxis, :], dtype=\"int64\")\n",
    "        attention_output = self.attention(\n",
    "            query=inputs, value=inputs, key=inputs, attention_mask=padding_mask\n",
    "        )\n",
    "        proj_input = self.layernorm_1(inputs + attention_output)\n",
    "        proj_output = self.dense_proj(proj_input)\n",
    "        return self.layernorm_2(proj_input + proj_output)\n",
    "\n",
    "\n",
    "class PositionalEmbedding(layers.Layer):\n",
    "    def __init__(self, sequence_length, vocab_size, embed_dim, **kwargs):\n",
    "        super(PositionalEmbedding, self).__init__(**kwargs)\n",
    "        self.token_embeddings = layers.Embedding(\n",
    "            input_dim=vocab_size, output_dim=embed_dim\n",
    "        )\n",
    "        self.position_embeddings = layers.Embedding(\n",
    "            input_dim=sequence_length, output_dim=embed_dim\n",
    "        )\n",
    "        self.sequence_length = sequence_length\n",
    "        self.vocab_size = vocab_size\n",
    "        self.embed_dim = embed_dim\n",
    "\n",
    "    def call(self, inputs):\n",
    "        length = tf.shape(inputs)[-1]\n",
    "        positions = tf.range(start=0, limit=length, delta=1)\n",
    "        embedded_tokens = self.token_embeddings(inputs)\n",
    "        embedded_positions = self.position_embeddings(positions)\n",
    "        return embedded_tokens + embedded_positions\n",
    "\n",
    "    def compute_mask(self, inputs, mask=None):\n",
    "        return tf.math.not_equal(inputs, 0)\n",
    "\n",
    "\n",
    "class TransformerDecoder(layers.Layer):\n",
    "    def __init__(self, embed_dim, latent_dim, num_heads, **kwargs):\n",
    "        super(TransformerDecoder, self).__init__(**kwargs)\n",
    "        self.embed_dim = embed_dim\n",
    "        self.latent_dim = latent_dim\n",
    "        self.num_heads = num_heads\n",
    "        self.attention_1 = layers.MultiHeadAttention(\n",
    "            num_heads=num_heads, key_dim=embed_dim\n",
    "        )\n",
    "        self.attention_2 = layers.MultiHeadAttention(\n",
    "            num_heads=num_heads, key_dim=embed_dim\n",
    "        )\n",
    "        self.dense_proj = keras.Sequential(\n",
    "            [layers.Dense(latent_dim, activation=\"relu\"), layers.Dense(embed_dim),]\n",
    "        )\n",
    "        self.layernorm_1 = layers.LayerNormalization()\n",
    "        self.layernorm_2 = layers.LayerNormalization()\n",
    "        self.layernorm_3 = layers.LayerNormalization()\n",
    "        self.supports_masking = True\n",
    "\n",
    "    def call(self, inputs, encoder_outputs, mask=None):\n",
    "        causal_mask = self.get_causal_attention_mask(inputs)\n",
    "        if mask is not None:\n",
    "            padding_mask = tf.cast(mask[:, tf.newaxis, :], dtype=\"int64\")\n",
    "            padding_mask = tf.minimum(padding_mask, causal_mask)\n",
    "\n",
    "        attention_output_1 = self.attention_1(\n",
    "            query=inputs, value=inputs, key=inputs, attention_mask=causal_mask\n",
    "        )\n",
    "        out_1 = self.layernorm_1(inputs + attention_output_1)\n",
    "\n",
    "        attention_output_2 = self.attention_2(\n",
    "            query=out_1,\n",
    "            value=encoder_outputs,\n",
    "            key=encoder_outputs,\n",
    "            attention_mask=padding_mask,\n",
    "        )\n",
    "        out_2 = self.layernorm_2(out_1 + attention_output_2)\n",
    "\n",
    "        proj_output = self.dense_proj(out_2)\n",
    "        return self.layernorm_3(out_2 + proj_output)\n",
    "\n",
    "    def get_causal_attention_mask(self, inputs):\n",
    "        input_shape = tf.shape(inputs)\n",
    "        batch_size, sequence_length = input_shape[0], input_shape[1]\n",
    "        i = tf.range(sequence_length)[:, tf.newaxis]\n",
    "        j = tf.range(sequence_length)\n",
    "        mask = tf.cast(i >= j, dtype=\"int64\")\n",
    "        mask = tf.reshape(mask, (1, input_shape[1], input_shape[1]))\n",
    "        mult = tf.concat(\n",
    "            [tf.expand_dims(batch_size, -1), tf.constant([1, 1], dtype=tf.int32)],\n",
    "            axis=0,\n",
    "        )\n",
    "        return tf.tile(mask, mult)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "f4da45b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "embed_dim = 512\n",
    "latent_dim = 4096\n",
    "num_heads = 8\n",
    "\n",
    "encoder_inputs = keras.Input(shape=(None,), dtype=\"int64\", name=\"encoder_inputs\")\n",
    "x = PositionalEmbedding(sequence_length, vocab_size, embed_dim)(encoder_inputs)\n",
    "encoder_outputs = TransformerEncoder(embed_dim, latent_dim, num_heads)(x)\n",
    "encoder = keras.Model(encoder_inputs, encoder_outputs)\n",
    "\n",
    "decoder_inputs = keras.Input(shape=(None,), dtype=\"int64\", name=\"decoder_inputs\")\n",
    "encoded_seq_inputs = keras.Input(shape=(None, embed_dim), name=\"decoder_state_inputs\")\n",
    "x = PositionalEmbedding(sequence_length, vocab_size, embed_dim)(decoder_inputs)\n",
    "x = TransformerDecoder(embed_dim, latent_dim, num_heads)(x, encoded_seq_inputs)\n",
    "x = layers.Dropout(0.5)(x)\n",
    "decoder_outputs = layers.Dense(vocab_size, activation=\"softmax\")(x)\n",
    "decoder = keras.Model([decoder_inputs, encoded_seq_inputs], decoder_outputs)\n",
    "\n",
    "decoder_outputs = decoder([decoder_inputs, encoder_outputs])\n",
    "transformer = keras.Model(\n",
    "    [encoder_inputs, decoder_inputs], decoder_outputs, name=\"transformer\")\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a8eea97",
   "metadata": {},
   "source": [
    "# Training:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70d75acc",
   "metadata": {},
   "outputs": [],
   "source": [
    "epochs = 5   #at least 30 epochs\n",
    "learning_rate=0.0001\n",
    "\n",
    "opt = keras.optimizers.Adam(learning_rate=learning_rate)\n",
    "\n",
    "transformer.summary()\n",
    "transformer.compile(loss=\"sparse_categorical_crossentropy\", metrics=[\"accuracy\"], optimizer=opt)\n",
    "transformer.fit(train_ds, epochs=epochs, validation_data=val_ds)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a5025ca",
   "metadata": {},
   "source": [
    "# Inference:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "a9436c57",
   "metadata": {},
   "outputs": [],
   "source": [
    "target_tokens = output_vectorization.get_vocabulary()\n",
    "target_index_lookup = dict(zip(range(len(target_tokens)), target_tokens))\n",
    "max_decoded_sentence_length = sequence_length\n",
    "\n",
    "\n",
    "def decode_sequence(input_sentence):\n",
    "    tokenized_input_sentence = input_vectorization([input_sentence])\n",
    "    decoded_sentence = \"[start]\"\n",
    "    for i in range(max_decoded_sentence_length):\n",
    "        tokenized_target_sentence = output_vectorization([decoded_sentence])[:, :-1]\n",
    "        predictions = transformer([tokenized_input_sentence, tokenized_target_sentence])\n",
    "\n",
    "        sampled_token_index = np.argmax(predictions[0, i, :])\n",
    "        sampled_token = target_index_lookup[sampled_token_index]\n",
    "        decoded_sentence += \" \" + sampled_token\n",
    "\n",
    "        if sampled_token == \"[end]\":\n",
    "            break\n",
    "    return decoded_sentence"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3bc4d582",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_input_texts = [pair[1] for pair in test_pairs]\n",
    "test_output_texts = [pair[2] for pair in test_pairs]\n",
    "marty_time = [pair[3] for pair in test_pairs]\n",
    "\n",
    "\n",
    "\n",
    "for i in random.sample(range(0,len(test_input_texts)), 5):\n",
    "    input_sentence = test_input_texts[i]\n",
    "    start = time.process_time()\n",
    "    translated = decode_sequence(input_sentence)\n",
    "    elapsed = (time.process_time() - start)\n",
    "    print('Actual:    ', test_output_texts[i], '\\n')\n",
    "    print('Predicted: ', translated, '\\n')\n",
    "    print(test_output_texts[i]==translated, '\\n')\n",
    "    print('symba time: ', elapsed, '\\n')\n",
    "    print('Marty time: ', str(marty_time[i]), '\\n')\n",
    "    "
   ]
  }
 ],
 "metadata": {
  "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}