utils.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import time \n",
    "\n",
    "from sklearn.model_selection import train_test_split\n",
    "from itertools import zip_longest\n",
    "from Bio.Data import CodonTable\n",
    "from Bio.Seq import translate, IUPAC"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tested Functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def normalize(X): \n",
    "    \"\"\" normalize an array \"\"\"\n",
    "    return (X - X.mean()) / X.std()\n",
    "\n",
    "def get_wild_type_dna_sequence(): \n",
    "    return pd.read_csv(\"./data/gfp_data.csv\")[\"nucSequence\"].values[0]\n",
    "\n",
    "def dna_to_amino_acid(dna_seq):\n",
    "    return translate(dna_seq)\n",
    "\n",
    "def get_all_amino_acids(): \n",
    "    return \"*\" + IUPAC.protein.letters\n",
    "    \n",
    "def get_wild_type_amino_acid_sequence(): \n",
    "    return dna_to_amino_acid(get_wild_type_dna_sequence())\n",
    "\n",
    "def count_substring_mismatch(s1, s2): \n",
    "    \"\"\" returns the number of misaligned pairs within strings s1 and s2\"\"\"\n",
    "    return sum([i != j for i, j in zip_longest(s1, s2)])\n",
    "\n",
    "def generate_random_gfp_data_mutations(num_of_mutations_lst, num_per_mutation_count = 1000): \n",
    "    \"\"\"\n",
    "    Input: num_of_mutations_lst is a list defining the count of mutations from the base sequence we want, \n",
    "           num_per_mutation_count is an int defining how many of each mutation count do we want. Often set to the test size\n",
    "    Output: a pandas dataframe comprised of two columns: the number of mutations and the mutation dna sequence\n",
    "    \"\"\"\n",
    "    start_time = time.time()\n",
    "    total_data_points = len(num_of_mutations_lst) * num_per_mutation_count\n",
    "    assert(total_data_points < 200000)\n",
    "    wild_type_sequence = get_wild_type_dna_sequence()\n",
    "    wild_type_lst = list(wild_type_sequence)\n",
    "    mutation_lst = np.vstack([wild_type_lst] * total_data_points)\n",
    "    mutation_count_lst = np.array([[mutation_count for _ in range(num_per_mutation_count)] for mutation_count in num_of_mutations_lst]).flatten()\n",
    "    \n",
    "    np.testing.assert_array_equal(mutation_lst[0], mutation_lst[2])\n",
    "    np.testing.assert_array_equal(mutation_lst[0], wild_type_lst)\n",
    "    assert(len(mutation_lst) == total_data_points and len(mutation_count_lst) == len(mutation_lst))\n",
    "    assert(mutation_count_lst[0] == 1 and mutation_count_lst[num_per_mutation_count] == 2 and mutation_count_lst[num_per_mutation_count - 1] == 1)\n",
    "    \n",
    "    bases = \"ACTG\"\n",
    "    index = list(range(0, 4))\n",
    "    base_index_map = dict(zip(bases, index))\n",
    "    index_base_map = dict(zip(index, bases))\n",
    "\n",
    "    for i, mutation_count in enumerate(mutation_count_lst):   \n",
    "        mutation_index = np.random.choice(len(wild_type_sequence), mutation_count, replace = False).tolist()\n",
    "        for j in mutation_index: \n",
    "            k = base_index_map[mutation_lst[i, j]]\n",
    "            new_index = (k + np.random.randint(1, 4)) % 4\n",
    "            mutation_lst[i, j] = index_base_map[new_index] \n",
    "    mutation_sequence_lst = np.array([\"\".join(lst) for lst in mutation_lst])\n",
    "    mutated_df = pd.DataFrame.from_dict({'mutation_count' : mutation_count_lst, 'mutated_dna_sequence' : mutation_sequence_lst})\n",
    "    mutated_df[\"mutated_amino_acid_sequence\"] = mutated_df[\"mutated_dna_sequence\"].apply(lambda x : dna_to_amino_acid(x))\n",
    "    print(time.time() - start_time, \"seconds to generate the mutated df\")\n",
    "    return mutated_df\n",
    "\n",
    "def get_mutation(string, num_mutations, alphabet): \n",
    "    mutation = list(string)\n",
    "    num_characters = len(alphabet)\n",
    "    characters_to_index = dict(zip(alphabet, range(num_characters)))\n",
    "    index_to_characters = dict(zip(range(num_characters), alphabet))\n",
    "    indexes = np.random.choice(range(len(string)), num_mutations, replace=False)\n",
    "    for i in indexes:\n",
    "        original_c = string[i]\n",
    "        original_index = characters_to_index[original_c]\n",
    "        new_index = np.random.randint(0, num_characters)\n",
    "        while new_index == original_index: \n",
    "            new_index = np.random.randint(0, num_characters)\n",
    "        mutation[i] = index_to_characters[new_index] \n",
    "    return \"\".join(mutation)\n",
    "\n",
    "def save_mutated_gfp_data(mutated_df, path = \"./data/mutated_df.csv\"):\n",
    "    mutated_df.to_csv(path, index = None)\n",
    "    \n",
    "def load_saved_mutated_gfp_data(path = \"./data/mutated_df.csv\"): \n",
    "    return pd.read_csv(path)\n",
    "\n",
    "def get_gfp_data(amino_acid = False, gfp_data_path = \"./data/gfp_data.csv\", x_feature = \"nucSequence\", y_feature = \"medianBrightness\", normalize_y = True, test_size = 0.2, shuffle = False):\n",
    "    # returns gfp data in dna or amino acid form which is then split into train and test set\n",
    "    df = pd.read_csv(gfp_data_path, index_col = 0)\n",
    "    if amino_acid: \n",
    "        X = df[x_feature].apply(lambda x : dna_to_amino_acid(x)).values\n",
    "    else: \n",
    "        X = df[x_feature].values\n",
    "    y = df[y_feature].values\n",
    "    if normalize_y: \n",
    "        y = normalize(y)\n",
    "    return train_test_split(X, y, test_size = test_size, shuffle = shuffle)\n",
    "\n",
    "def save_gfp_data(X_train, X_test, y_train, y_test, gfp_data_path): \n",
    "    np.save(gfp_data_path + \"X_train.npy\", X_train)\n",
    "    np.save(gfp_data_path + \"X_test.npy\", X_test)\n",
    "    np.save(gfp_data_path + \"y_train.npy\", y_train)\n",
    "    np.save(gfp_data_path + \"y_test.npy\", y_test)\n",
    "\n",
    "def load_gfp_data(gfp_data_path):\n",
    "    X_train = np.load(gfp_data_path + \"X_train.npy\")\n",
    "    X_test = np.load(gfp_data_path + \"X_test.npy\")\n",
    "    y_train = np.load(gfp_data_path + \"y_train.npy\")\n",
    "    y_test = np.load(gfp_data_path + \"y_test.npy\")\n",
    "    return X_train, X_test, y_train, y_test\n",
    "\n",
    "  \n",
    "def one_hot_encode_dna_sequence(X):\n",
    "    \"\"\"\n",
    "    Input: X is a list of DNA Sequences represented by the base pairs ACTG. \n",
    "        All DNA Sequences must be the same length\n",
    "    Output: one hot encoded list of dna sequences\n",
    "    Example: one_hot_encode([\"ACT\", \"ACG\"]) = [[1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0], \n",
    "                                              [1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1]]\n",
    "    \"\"\"\n",
    "    assert(len(X) > 0)\n",
    "    assert(all([len(X[0]) == len(X[i]) for i in range(len(X))]))\n",
    "    alphabet = [\"A\", \"C\", \"T\", \"G\"]\n",
    "    alphabet_size = len(alphabet)\n",
    "    alphabet_dict = dict(zip(alphabet, range(alphabet_size)))\n",
    "    one_hot_matrix = np.zeros((len(X), alphabet_size * len(X[0]))) \n",
    "    for i, dna_sequence in enumerate(X):\n",
    "        for j, base_pair in enumerate(dna_sequence):\n",
    "            index = alphabet_dict[base_pair]\n",
    "            one_hot_matrix[i, alphabet_size * j + index] = 1.0\n",
    "    return one_hot_matrix\n",
    "\n",
    "\n",
    "def one_hot_decode_dna_sequence(X): \n",
    "    \"\"\"\n",
    "    Input: X is a one hot encoded list of DNA Sequences represented by the base pairs ACTG. \n",
    "        All DNA Sequences must be the same length\n",
    "    Output: list of dna sequences\n",
    "    Example: one_hot_decode([[1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0], \n",
    "                            [1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1]]) = [\"ACT\", \"ACG\"]\n",
    "    \"\"\"\n",
    "    assert(len(X) > 0)\n",
    "    assert(all([len(X[0]) == len(X[i]) for i in range(len(X))]))\n",
    "    alphabet = [\"A\", \"C\", \"T\", \"G\"]\n",
    "    alphabet_size = len(alphabet)\n",
    "    dna_sequences = []\n",
    "    for i, one_hot_sequence in enumerate(X): \n",
    "        dna_sequence = []\n",
    "        for j in range(0, len(one_hot_sequence), 4): \n",
    "            if one_hot_sequence[j]:\n",
    "                dna_sequence.append(\"A\")\n",
    "            elif one_hot_sequence[j + 1]: \n",
    "                dna_sequence.append(\"C\")\n",
    "            elif one_hot_sequence[j + 2]: \n",
    "                dna_sequence.append(\"T\")\n",
    "            elif one_hot_sequence[j + 3]: \n",
    "                dna_sequence.append(\"G\")\n",
    "        dna_sequences.append(\"\".join(dna_sequence))\n",
    "    return np.array(dna_sequences)\n",
    "\n",
    "def one_hot_encode(X, alphabet):\n",
    "    \"\"\"\n",
    "    Input: X is a list of sequences represented by the set of letters in alphabet \n",
    "        All sequences must be the same length\n",
    "    Output: one hot encoded list of X sequences\n",
    "    Example: one_hot_encode([\"ACT\", \"ACG\"], \"ACTG\") = [[1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0], \n",
    "                                              [1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1]]\n",
    "    \"\"\"\n",
    "    assert(len(X) > 0)\n",
    "    assert(all([len(X[0]) == len(X[i]) for i in range(len(X))]))\n",
    "    alphabet_size = len(alphabet)\n",
    "    alphabet_dict = dict(zip(alphabet, range(alphabet_size)))\n",
    "    one_hot_matrix = np.zeros((len(X), alphabet_size * len(X[0])))\n",
    "    for i, sequence in enumerate(X):\n",
    "        for j, letter in enumerate(sequence):\n",
    "            if letter not in alphabet:\n",
    "                raise KeyError(\"letter not in alphabet\")\n",
    "            index = alphabet_dict[letter]    \n",
    "            one_hot_matrix[i, alphabet_size * j + index] = 1.0\n",
    "    return one_hot_matrix\n",
    "\n",
    "def one_hot_decode(X, alphabet):\n",
    "    \"\"\"\n",
    "    Input: X is a one hot encoded list of DNA Sequences represented by the base pairs ACTG. \n",
    "        All DNA Sequences must be the same length\n",
    "    Output: list of dna sequences\n",
    "    Example: one_hot_decode([[1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0], \n",
    "                            [1, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1]], \"ACTG\") = [\"ACT\", \"ACG\"]\n",
    "    \"\"\"\n",
    "    assert(len(X) > 0)\n",
    "    assert(all([len(X[0]) == len(X[i]) for i in range(len(X))]))\n",
    "    alphabet_size = len(alphabet)\n",
    "    sequences_lst = []\n",
    "    for i, one_hot_sequence in enumerate(X): \n",
    "        sequence, sequence_len = [], len(one_hot_sequence)\n",
    "        for j in range(0, sequence_len, alphabet_size):\n",
    "            index = np.argmax(one_hot_sequence[j:j+alphabet_size])\n",
    "            sequence.append(alphabet[index])        \n",
    "        sequences_lst.append(\"\".join(sequence))\n",
    "    return sequences_lst\n",
    "\n",
    "def to_tensor(x, device = torch.device(\"cpu\")):\n",
    "    \"\"\"\n",
    "    Converts numpy array to tensor on the specific torch device\n",
    "    >>> to_tensor(np.ones((3, 3)), torch.device(\"cpu\"))\n",
    "    tensor([[1., 1., 1.],\n",
    "        [1., 1., 1.],\n",
    "        [1., 1., 1.]])\n",
    "    \"\"\"\n",
    "    assert(type(x) == np.ndarray)\n",
    "    return torch.from_numpy(x).float().to(device)\n",
    "    \n",
    "def sample_tensor_to_string(x, int_to_character, softmax=False):\n",
    "    \"\"\"\n",
    "    Samples a tensor from the probability distribution of x. \n",
    "    Input: a 2d or 1d vector of shape (characters x vocabulary_size)\n",
    "    Output: a string of length characters long \n",
    "    Note that the results are sampled each time and will be different\n",
    "    \n",
    "    >>> int_to_character = dict(zip(range(21), get_all_amino_acids()))\n",
    "    >>> sample_tensor_to_string(torch.randn(18, 21), int_to_character, softmax=True) \n",
    "    IIIMMREWQMPFQVS*TT \n",
    "    \"\"\"\n",
    "    num_characters = len(int_to_character)\n",
    "    assert(type(x) == torch.Tensor)\n",
    "    assert(x.shape[0] % num_characters == 0 or x.shape[1] % num_characters == 0)\n",
    "    x = x.reshape(-1, num_characters)\n",
    "    if softmax:\n",
    "        x = F.softmax(x, dim = -1)\n",
    "    string = []\n",
    "    for dist in x: \n",
    "        index = torch.multinomial(dist, 1).item()\n",
    "        string.append(int_to_character[index])\n",
    "    return \"\".join(string)\n",
    "\n",
    "\n",
    "def tensor_to_string(x, int_to_character):\n",
    "    \"\"\"\n",
    "    Converts tensor to string\n",
    "    Input: A sequence in tensor format\n",
    "    Output: A sequence in string format\n",
    "\n",
    "    >>> int_to_character = dict(zip(range(21), get_all_amino_acids()))\n",
    "    >>> tensor_to_string(torch.tensor([0, 0, 1, 0, 0, 0, 1, 0]))\n",
    "    \"TT\"\n",
    "    >>> tensor_to_string(torch.tensor([0.8, 0.15, 0.05, 0, 0, 0.9, 0.1, 0]))\n",
    "    \"AC\"\n",
    "    \"\"\"\n",
    "    num_characters = len(int_to_character)\n",
    "    assert(type(x) == torch.Tensor)\n",
    "    assert(len(x) % num_characters == 0)\n",
    "    x = x.reshape(-1, num_characters)\n",
    "    _, index = x.max(dim = 1)\n",
    "    return \"\".join([int_to_character[i] for i in index.numpy()])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = pd.read_csv(\"./data/gfp_data.csv\")[\"nucSequence\"].values\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_mutated_df(save_mutated_df=False): \n",
    "    num_of_mutations_lst = list(range(1, 10)) + [i * 10 for i in range(1, 6)]\n",
    "    num_per_mutation_count = 11684 #test size = len(X_test)\n",
    "    mutated_df = generate_random_gfp_data_mutations(num_of_mutations_lst=num_of_mutations_lst, num_per_mutation_count=num_per_mutation_count)\n",
    "    wild_type_dna_seq, wild_type_amino_acid_seq = get_wild_type_dna_sequence(), get_wild_type_amino_acid_sequence()\n",
    "    for mutation_count, dna_seq, amino_acid_seq in zip(mutated_df[\"mutation_count\"].values, \n",
    "                                                       mutated_df[\"mutated_dna_sequence\"].values, \n",
    "                                                       mutated_df[\"mutated_amino_acid_sequence\"].values): \n",
    "        assert(mutation_count == count_substring_mismatch(wild_type_dna_seq, dna_seq))\n",
    "        assert(mutation_count >= count_substring_mismatch(wild_type_amino_acid_seq, amino_acid_seq))\n",
    "        assert(dna_to_amino_acid(dna_seq) == amino_acid_seq)    \n",
    "    if save_mutated_df: \n",
    "        save_mutated_gfp_data(mutated_df=mutated_df)\n",
    "        load_mutated_df = load_saved_mutated_gfp_data()\n",
    "        assert(load_mutated_df.equals(mutated_df))\n",
    "\n",
    "def test_get_save_load_gfp_dna_seq():\n",
    "    X_train, X_test, y_train, y_test = get_gfp_data(amino_acid = False)\n",
    "    save_gfp_data(X_train, X_test, y_train, y_test, \"./data/gfp_dna_\")\n",
    "    train_X, test_X, train_y, test_y = load_gfp_data(\"./data/gfp_dna_\")\n",
    "    np.testing.assert_array_equal(X_train, train_X)\n",
    "    np.testing.assert_array_equal(X_test, test_X)\n",
    "    np.testing.assert_array_equal(y_train, train_y)\n",
    "    np.testing.assert_array_equal(y_test, test_y)\n",
    "    assert(X_train[0] == get_wild_type_dna_sequence())\n",
    "    assert(count_substring_mismatch(X_test[0], get_wild_type_dna_sequence()) == 2)\n",
    "    assert(X_train.shape == (46733,))\n",
    "    assert(X_test.shape == (11684,))\n",
    "    assert(y_train.shape == (46733,))\n",
    "    assert(y_test.shape == (11684,))\n",
    "    assert(isinstance(X_train[-1], str))\n",
    "    assert(isinstance(X_test[-1], str))\n",
    "    assert(isinstance(y_train[-1], float))\n",
    "    assert(isinstance(y_test[-1], float))\n",
    "    \n",
    "def test_get_save_load_gfp_amino_acid_seq():\n",
    "    X_train, X_test, y_train, y_test = get_gfp_data(amino_acid = True)\n",
    "    save_gfp_data(X_train, X_test, y_train, y_test, \"./data/gfp_amino_acid_\")\n",
    "    train_X, test_X, train_y, test_y = load_gfp_data(\"./data/gfp_amino_acid_\")\n",
    "    np.testing.assert_array_equal(X_train, train_X)\n",
    "    np.testing.assert_array_equal(X_test, test_X)\n",
    "    np.testing.assert_array_equal(y_train, train_y)\n",
    "    np.testing.assert_array_equal(y_test, test_y)\n",
    "    assert(X_train[0] == get_wild_type_amino_acid_sequence())\n",
    "    assert(count_substring_mismatch(X_test[0], get_wild_type_amino_acid_sequence()) == 2)\n",
    "    assert(X_train.shape == (46733,))\n",
    "    assert(X_test.shape == (11684,))\n",
    "    assert(y_train.shape == (46733,))\n",
    "    assert(y_test.shape == (11684,))\n",
    "    assert(isinstance(X_train[-1], str))\n",
    "    assert(isinstance(X_test[-1], str))\n",
    "    assert(isinstance(y_train[-1], float))\n",
    "    assert(isinstance(y_test[-1], float))\n",
    "\n",
    "def test_encode_decode():\n",
    "    for data_dir in [\"./data/gfp_amino_acid_\", \"./data/gfp_dna_\"]:\n",
    "        train_X, test_X, train_y, test_y = load_gfp_data(data_dir)\n",
    "        if \"dna\" in data_dir:\n",
    "            alphabet = \"ACTG\"\n",
    "        else: \n",
    "            alphabet = get_all_amino_acids()\n",
    "        one_hot_X_train, one_hot_X_test = one_hot_encode(train_X, alphabet=alphabet), one_hot_encode(test_X, alphabet=alphabet)\n",
    "        X_train, X_test = one_hot_decode(one_hot_X_train, alphabet=alphabet), one_hot_decode(one_hot_X_test, alphabet=alphabet)\n",
    "        np.testing.assert_array_equal(train_X, X_train)\n",
    "        np.testing.assert_array_equal(test_X, X_test)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_mutated_df(True)\n",
    "test_get_save_load_gfp_dna_seq()\n",
    "test_get_save_load_gfp_amino_acid_seq()\n",
    "test_encode_decode()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finished encoding data in 8.30 seconds\n"
     ]
    }
   ],
   "source": [
    "start_time = time.time()\n",
    "X_train, X_test, y_train, y_test = load_gfp_data(\"./data/gfp_amino_acid_\")\n",
    "data = np.array([x[:] for x in X_train[0:]])\n",
    "one_hot_X_train = one_hot_encode(data, get_all_amino_acids())\n",
    "z = one_hot_decode([one_hot_X_train[0]], get_all_amino_acids())[0]\n",
    "assert(z == get_wild_type_amino_acid_sequence())\n",
    "print(\"Finished encoding data in {:.2f} seconds\".format(time.time() - start_time))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(46733,) ['SKGEELFTGVVPILVELDGDVNGHEFSVSGEGEGDATYGKLTLKFICTTGKLPVPWPTLVTTLSYGVQCFSRYPDHMKQHDFFKSAMPEGYVQERTIFFRDDGNYKTRAEVKFEGDTLVNRIELKGIDFKEDGNILGHKLEYNYNSHNVYIMADKQKNGSKVNFKIRHNIEDGSAQLADHYQQNTPIGDGPVLLPDNHYLSTQSALSKDPNEKRDHMVLLEFVTAAGITHGMDELYK*'\n",
      " 'SKGGELFTGVVPILAELDGDVNGHKFCVSGEGEGDATYGKLTLKFICTTGKLPVPWPTLVTTLSYGAQCFSRCPDHMKQHDFFKSAMPEGYVRERTVFFKDDGNYKTRAEVKFEGDTLVNRIELKGIDFKEDGNILGHKLEYNYNSHNVYIMADKQKNGIKVNFKIRHNIEDGSVQLADHYQQNTPIGDGPVLLPDNHYLSTQSALSKDPNEKRDHMVLLEFVTAAGITHGMDELYK*'\n",
      " 'SKGEELFTGVVPILVELDGDVNGHKFSVSGEGEGDATYGKLTLKFICTTGKLPVPWPTLVTTLSYGVQCFSRYPDHMKQHDFFKSAMPEGYVQERTIFFKDDGNYKTRAEVKFEGDTLVNRIELKGIDFKEDGNILGHRLEYNYNSHNVYIMADKQKNGIKVNFKIRHNIEDGSVQLADHYQQSTPIGDGPVLLPDNHYLSTQSALSKDPNEKRDHMVLLEFVAAAGITHGMDELYK*'\n",
      " ...\n",
      " 'SKGEELFTGVVPILVELDGDVNGHKFSVSGEGEGDATYGKLTLKFICTTGKLPVPWPTLVTTLSYGVQCLSRYPDHMKQHDFFKSAMPEGYVQERTIFFKDDGNYKARAEVKFEGDTLVNRIELKGIDFKEDGNILGHKLEYNYNSHNVYIMADKQKNGIKVNFKIRHNIVDGSVQFADHYQQNTPIGDGPVLLPDNHYLSTQSALSKDPNEKRDHMVLLELVTAAGITHGMDELYK*'\n",
      " 'SKGEELFTGVVPILVELDGDVNGHKFSVSGEGEGDATYGKLTLKFICTTGKLPVPWPTLVTTLSYGVQCFSRYPDHMKQHDFFKAAMPEGYVQERTIFFKDDGNYKTRAEVKFEGDTLVNRIELKGIDFKEDGNILGHKLEYNYNSHNVYIMADKQKNGIKVNFKIRHNIEDGSVQLADHYQQNTPIGDGPVLLPDNHYLSTQSALSKDPNEKRDHMVLLEFVTAAGITHGMDELYK*'\n",
      " 'SKGEELFTGVVPILVELDGDVNGHKFSVSGEGEGDATYGKLTLKFICATGKLPVPWPTLVTTLSYGVQCFSRYPDHMKQHDFFKSAMPEGYVQERTIFFKDDGNYKTRAEVKFEGDALVNRIELKGIDFKEDGNILGHKLEYNYNSHNVYIMADKQKNGIKVNFKIRHNIEDGSVLLADHYQQNTPIGDGPVLLPDNHYLSTQSALSKDPNEKRDHMVLLEFVTAAGITHGMDELYK*'] (46733,) (11684,)\n"
     ]
    }
   ],
   "source": [
    "X_train, X_test, y_train, y_test = get_gfp_data(amino_acid=True, shuffle=True)\n",
    "print(X_train.shape, X_test.shape, y_train.shape, y_test.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(46733,) (11684,) (46733,) (11684,)\n"
     ]
    }
   ],
   "source": [
    "wild_type = get_wild_type_amino_acid_sequence()\n",
    "print(X_train.shape, X_test.shape, y_train.shape, y_test.shape)\n",
    "assert(wild_type in X_train)\n",
    "np.save(\"./data/gfp_amino_acid_shuffle_X_train.npy\", X_train)\n",
    "np.save(\"./data/gfp_amino_acid_shuffle_X_test.npy\", X_test)\n",
    "np.save(\"./data/gfp_amino_acid_shuffle_y_train.npy\", y_train)\n",
    "np.save(\"./data/gfp_amino_acid_shuffle_y_test.npy\", y_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['SKGEELFTGVVPILVELDGDVNGHKFSVSGEGEGDATYGELTLKFICTTGKLPVPWPTLVTTLSYGVQCFSCYPDHMKQHDFFKSAMPEGYVQERAIFFKDDGNYKTRAEVEFEGDTLVNRIELKGIDFKEDGDILGRKLEYNYNSHNVYIMADKQKNGIKVNFKIRHNIEDGSVQLADHYQQNTPIGDGPVLLPDNHYLSTQSALSIDPSEKRDHMVLLEFVTAAGITHGMDELYK*'\n",
      " 'SKGEELFTGVVPILVELDGDVNGHKFSVSGEGEGDATYGKLTLKFICTTGKLPVPWPTLVTTLSYGVQCFSRYPDHMKQHDFFKSAMPEGYVQERTTFFKDDGNYKTRAEVKFEGDTLVNRIELKGIDFKEDGNILGHKLEYNYNSHNVYIMADKQKNGIKVNFKIRHNIEDGSVQLADHYQQNTPIGDGPVLLPDNHYLSTQSALSKDPNEKRDHMVLLEFVTAAGITHGMDELYK*'\n",
      " 'SKGEELFTGVVPILVELDGDVNGHKFSVSGEGEGDATYGKLTLKFICTTGKLPVPWPTLVSTLSYGVQCFSRYPDHMKQHDFFKSAMPEGYVQERTIFFKDDGDCKTRAEV*FEGDTLVNCIELKGIDFKEDGNILGHKLEYNYNSHNVYIMADKQKNGIKVNFKIRHNIEDGSVRLADHYQQNTPIGDGPVLLTDNHYLSTQSALSKDPNEKRDHMVLLEFVTAAGITHGMDELYK*'\n",
      " ...\n",
      " 'SKGEELFTGVVPILVELDGDVNGHKFSVSGEGEGDTTYGKLTLKFICTTGKLPVPWPTLVTTLSYGVQCFSRYPDHMKQPDFFKSAMPEGYVQERTIFFKDDGNYKTRAEVKFEGDTLVNRIELKGIDFKEDGNILGHKLEYNYNSHNVYIMADKQKNGIKVNFKIRHNIEDGSVQLADHYQQNTPIGDGPVPLPDNHYLSTQSALSEDPNERRDHMALLEFVTAAGITHGMDELYK*'\n",
      " 'SKGEELFTGVVPILVELDGDVNGHKFSVSGEGEGDATYGKLTLKFICTTGKPPVPWPTLVTTLSYGVQCFSRYPDHMKQHDFFKTAMPEGYVQERTIFFKDDGNYKTRAEVKFEGDTLVNRIELKGIDLKEDGNILGHKLEYNYNSHNVYIMADKQKNGIKVNFKIRHNIEDGSVQLADHYQQNTPIGDGPVLLPDNHYLSTQSALSKDPNEKRDHMALLEFVTAAGITHGMDELYK*'\n",
      " 'SKGEELFTGVVPILVELDGDVNGHKFSVSGEGEGDATYGKLALRFICTTGKLPVPWPTLVTTLSYGVQCFSRYPDHMEQHDFFKSAMPESYVQERTIFFKDDGNYKTRAEVKFEGDTLVNRIGLKGIDFKEDGNILGHKLEYNNNSHNVYIMADKQKNGIKVNFKIRHNIEDGSVQLADHYQQNTPIGDGPVLLPDNHYLSTQSALSKDPNEKRDHMVLLEFVTAAGITHGMDELYK*']\n",
      "SKGEELFTGVVPILVELDGDVNGHKFSVSGEGEGDATYGELTLKFICTTGKLPVPWPTLVTTLSYGVQCFSCYPDHMKQHDFFKSAMPEGYVQERAIFFKDDGNYKTRAEVEFEGDTLVNRIELKGIDFKEDGDILGRKLEYNYNSHNVYIMADKQKNGIKVNFKIRHNIEDGSVQLADHYQQNTPIGDGPVLLPDNHYLSTQSALSIDPSEKRDHMVLLEFVTAAGITHGMDELYK*\n",
      "-1.3248364688852856\n",
      "(46733,) (11684,) (46733,) (11684,)\n"
     ]
    }
   ],
   "source": [
    "wild_type = get_wild_type_amino_acid_sequence()\n",
    "X_train, X_test, y_train, y_test = load_gfp_data(\"./data/gfp_amino_acid_shuffle_\")\n",
    "print(X_train)\n",
    "assert(wild_type in X_train)\n",
    "data = np.array([x[:] for x in X_train[0:]])\n",
    "one_hot_X_train = one_hot_encode(data, get_all_amino_acids())\n",
    "X_prime = one_hot_decode(one_hot_X_train, get_all_amino_acids())\n",
    "print(X_prime[0])\n",
    "print(y_train[0])\n",
    "print(X_prime.shape, X_test.shape, y_train.shape, y_test.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Average mismatch in test set: 3.72\n",
      "Average mismatch in train set: 3.71\n"
     ]
    }
   ],
   "source": [
    "for i, j in zip(X_train, X_prime):\n",
    "    assert(i == j)\n",
    "print(\"Average mismatch in test set: {0:.2f}\".format(np.mean([count_substring_mismatch(string, wild_type) for string in X_test])))\n",
    "print(\"Average mismatch in train set: {0:.2f}\".format(np.mean([count_substring_mismatch(string, wild_type) for string in X_train])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "int_to_character = dict(zip(range(21), get_all_amino_acids()))\n",
    "print(to_tensor(np.ones((3, 3)), torch.device(\"cpu\")))\n",
    "print(sample_tensor_to_string(torch.randn(18, 21), int_to_character, softmax=True))\n",
    "print(tensor_to_string(torch.randn(18 * 21), int_to_character))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}