Creating an affinity-matrix between protein and RNA sequences

Question

I wrote an algorithm that analyzes protein-RNA interactions and I found that the following function is the bottleneck that causes performance issues:

import numpy as np
#len(protein_sequence)~500, len(rna_sequence)~1500

def affinity_matrix(protein_sequence, rna_sequence): 
    python_matrix = [[] for _ in range(len(protein_sequence))]

    for i, AA in enumerate(protein_sequence):
        for base in rna_sequence:
            python_matrix[i].append(scales[base][AA])
            #(Where "scales" is a small dict() with the structure: 
            #scales[base][AA] = float(), with 4 bases and 20 AA, so 80 values total.)
    return(np.array(python_matrix))

I suspect 2 problems in this code, but I don't know how to solve them:

1. I am retrieving values from this scales dict millions of times, and it's said calling values from "non-static" data structures like dict() is slow. So how can I make this small dictionary "static" instead? (This dictionary is only created once, while the rna_sequence and protein_sequence will be different each time I call the function.)
2. I am building the matrix at first with pythons tools, and later I convert it into the (faster) NumPy array. Possibly it is faster to directly create it with NumPy, but I am not sure if this is possible.

I am thankful for any tips how to improve this code or references to guides that would help in this case.

Here is example data, in case you want to try out the algorithm:

rna_sequence = 'ACAGGAGGAGCCGCUCGCUGGCGGCUGAUCCAGCGUCUCCGUGACAGGCACCCUGCUCCGCCGCCACCGCCACCGCCACCGCCACCGUCGCCUUUUCUUCUUCGUCCCGGGCGGUGCGUUCCACUGCUCUGGGGCCGGCGCCGCGCCCAGUCCCGCUUCGGGCCGCAAGCCCCACCGCUCCCCUCCCCGGGCAGGGGCGCCGCGCAGCCCGCUCCCGCCGCCACCUCCUCCCCUGCCGCCCUCCUAGCCGGCAGGAAUUGCGCGACCACAGCGCCGCUCGCGUCGCCCGCAUCAGCUCAGCCCGCUGCCGCUCGGCCCUCGGCACCGCUCCGGGUCCGGCCGCCGCGCGGCCAGGGCUCCCCCUGCCCAGCGCUCCCAGGCCCCGCCACGCGUCGCCGCGCCCAGCUCCAGUCUCCCCUCCCCGGGGUCUCGCCAGCCCCUUCCUGCAGCCGCCGCCUCCGAAGGAGCGGGUCCGCCGCGGGUAACCAUGCCUAGCAAAACCAAGUACAACCUUGUGGACGAUGGGCACGACCUGCGGAUCCCCUUGCACAACGAGGACGCCUUCCAGCACGGCAUCUGCUUUGAGGCCAAGUACGUAGGAAGCCUGGACGUGCCAAGGCCCAACAGCAGGGUGGAGAUCGUGGCUGCCAUGCGCCGGAUACGGUAUGAGUUUAAAGCCAAGAACAUCAAGAAGAAGAAAGUGAGCAUUAUGGUUUCAGUGGAUGGAGUGAAAGUGAUUCUGAAGAAGAAGAAAAAGCUUCUUUUAUUGCAGAAAAAGGAAUGGACGUGGGAUGAGAGCAAGAUGCUGGUGAUGCAGGACCCCAUCUACAGGAUCUUCUAUGUCUCUCAUGAUUCCCAAGACUUGAAGAUCUUCAGCUAUAUCGCUCGAGAUGGUGCCAGCAAUAUCUUCAGGUGUAACGUCUUUAAAUCCAAGAAGAAGAGCCAAGCUAUGAGAAUCGUUCGGACGGUGGGGCAGGCCUUUGAGGUCUGCCACAAGCUGAGCCUGCAGCACACGCAGCAGAAUGCAGAUGGCCAGGAAGAUGGAGAGAGCGAGAGGAACAGCAACAGCUCAGGAGACCCAGGCCGCCAGCUCACUGGAGCCGAGAGGGCCUCCACGGCCACUGCAGAGGAGACUGACAUCGAUGCGGUGGAGGUCCCACUUCCAGGGAAUGAUGUCCUGGAAUUCAGCCGAGGUGUGACUGAUCUAGAUGCUGUAGGGAAGGAAGGAGGCUCUCACACAGGCUCCAAGGUUUCGCACCCCCAGGAGCCCAUGCUGACAGCCUCACCCAGGAUGCUGCUCCCUUCUUCUUCCUCGAAGCCUCCAGGCCUGGGCACAGAGACACCGCUGUCCACUCACCACCAGAUGCAGCUCCUCCAGCAGCUCCUCCAGCAGCAGCAGCAGCAGACACAAGUGGCUGUGGCCCAGGUACACUUGCUGAAGGACCAGUUGGCUGCUGAGGCUGCGGCGCGGCUGGAGGCCCAGGCUCGCGUGCAUCAGCUUUUGCUGCAGAACAAGGACAUGCUCCAGCACAUCUCCCUGCUGGUCAAGCAGGUGCAAGAGCUGGAACUGAAGCUGUCAGGACAGAACGCCAUGGGCUCCCAGGACAGCUUGCUGGAGAUCACCUUCCGCUCCGGAGCCCUGCCCGUGCUCUGUGACCCCACGACCCCUAAGCCAGAGGACCUGCAUUCGCCGCCGCUGGGCGCGGGCUUGGCUGACUUUGCCCACCCUGCGGGCAGCCCCUUAGGUAGGCGCGACUGCUUGGUGAAGCUGGAGUGCUUUCGCUUUCUUCCGCCCGAGGACACCCCGCCCCCAGCGCAGGGCGAGGCGCUCCUGGGCGGUCUGGAGCUCAUCAAGUUCCGAGAGUCAGGCAUCGCCUCGGAGUACGAGUCCAACACGGACGAGAGCGAGGAGCGCGACUCGUGGUCCCAGGAGGAGCUGCCGCGCCUGCUGAAUGUCCUGCAGAGGCAGGAACUGGGCGACGGCCUGGAUGAUGAGAUCGCCGUGUAGGUGCCGAGGGCGAGGAGAUGGAGGCGGCGGCGUGGCUGGAGGGGCCGUGUCUGGCUGCUGCCCGGGUAGGGGAUGCCCAGUGAAUGUGCACUGCCGAGGAGAAUGCCAGCCAGGGCCCGGGAGAGUGUGAGGUUUCAGGAAAGUAUUGAGAUUCUGCUUUGGAGGGUAAAGUGGGGAAGAAAUCGGAUUCCCAGAGGUGAAUCAGCUCCUCUCCUACUUGUGACUAGAGGGUGGUGGAGGUAAGGCCUUCCAGAGCCCAUGGCUUCAGGAGAGGGUCUCUCUCCAGGACUGCCAGGCUGCUGGAGGACCUGCCCCUACCUGCUGCAUCGUCAGGCUCCCACGCUUUGUCCGUGAUGCCCCCCUACCCCCUCACUCUCCCCGUCUCCAUGGUCCCGACCAGGAAGGGAAGCCAUCGGUACCUUCUCAGGUACUUUGUUUCUGGAUAUCACGAUGCUGCGAGUUGCCUAACCCUCCCCCUACCUUUAUGAGAGGAAUUCCUUCUCCAGGCCCUUGCUGAGAUUGUAGAGAUUGAGUGCUCUGGACCGCAAAAGCCAGGCUAGUCCUUGUAGGGUGAGCAUGGAAUUGGAAUGUGUCACAGUGGAUAAGCUUUUAGAGGAACUGAAUCCAAACAUUUUCUCCAGCCGGACAUUGAAUGUUGCUACAAAGGGAGCCUUGAAGCUUUAACAUGGUUCAGGCCCUUGGUGUGAGAGCCCAGGGGGAGGACAGCUUGUCUGCUGCUCCAAAUCACUUAGAUCUGAUUCCUGUUUUGAAAGUCCUGCCCUGCCUUCCUCCUGCCUGUAGCCCAGCCCAUCUAAAUGGAAGCUGGGAAUUGCCCCUCACCUCCCCUGUGUCCUGUCCAGCUGAAGCUUUUGCAGCACUUUACCUCUCUGAAAGCCCCAGAGGACCAGAGCCCCCAGCCUUACCUCUCAACCUGUCCCCUCCACUGGGCAGUGGUGGUCAGUUUUUACUGCAAAAAAAAAAAAAGAAAAAAGAGAAAGAAAAAAAAGAAUGAAUGCAAGCUGAUAGCUGAGACUGUGAGACUGUUUUUGUCCACUCUUCUGAAUCACUGCCACUUGGGUCAGGGACCACAGCCAUUGCCACCCUUGGCCCAUCUCUCUGCGUGCGUGCCUUGAGCACACAUAUAAAAAGUGCCAUGUGCAAUUGUCUUAUCUUUUAUGAUCUAGGCUUUGCCUAGGGAUCACUACUCCUUAACGGGCUGGCUGGGGCAAUGAGGAAAAGCUCCUUUGCUCCUGUAAGGCCAUAAGUGGCUGUUAACAGAUUUUCAAAUGCCUGAAGAGAUUGCUGAGACCUGCUAGAGUCAUAUGUUCGGGGAAUUAAGUCUUUAUCCUAGACAACAAGGUACAGAUGCAAACUGCAGUGUUAUUGGAGGGUCAAUCGGCAAGGAUAUGAUUAUCCCAAAAUGGAGUUCAUCGACCCUAGCUUUCCUUUAGAUUAUAUAUAAAUAAAAGUGCAGUCCUCUUCUAAUGGCCACAGUUGGUUUUCUUGUAGCCCAGAAAGUCCAAAUUAAAGGAAAUAAAUUCAGUUUUAUGUUAGCCUUCCUUGGUGCAUCAGGGUGUCAGUGGAAAUAGGAUCAGGUGGUGUGUGUGUGUGUGUUUUGUGUGUGUGUGUACACAUGUGUUUAUAUAUACAUGUGUGAGGGAAAGUGUGUACAUAUAUGUAGGAUUGUAACCAGACGGAAAAGAACGAGGAUCUCCAGGGUGUUUGAAUCAGCAACAGAUUUGUGUUUUCUAACAUGCAUUUAGUUGGAGAGGCAUGGUUCUGUUUGUUUUGUUUUGAUCUAAUUUGCCAUUGGAAAUAGGUACAGUUACACAGAGAAGGAAGAACCAGGAAAGUGAGAUCCAUGAAACUAAAUGAGCAGCUGUCAGAAUCCAGUGUGGCUGAGCCUACCUAGCUUAUGAAAUCUAACCCAGGGUUCCCUGAGUCCAAGACCACUUAGAUUAUUAAGAUUUUGAACGUCCAGAGGAGUGAAAAGUCUGUUUUCUGACGUAAGCCGGAGCUGAGGAUAAAGCCAGAGGCCAGUGGAUUAGGUGUAUGGAAUGUGGAUGGAGAGGGCUUGUGUGGGAUGUGGCCAGGGAGUGGGUGAGGAAGGCCGCUUCUAAAUGGCCUGUAAAAACUUGAGAUUGGAUAGACGAAAGGAAAUGGAGAAAUUAAAGAAUUGGAGAAACUAGUUAUCUGUGUUGCUGACUUUGGGACCCAUCCAAGACUCCUGCCCUUGGGGUGUUCCAUGGUGGUUUCUUCCUGCCUGGGCGCCACCCUUUCCCCAGUUCAGGCCCUCCCUGGAGGACUAGUUUGUGUAUUGGUAUCCUCCCCAGUGGACCCAAACCAGCGCAUACUUGGUGUGUGGAGAUGGGAGACAAAGGACAGAUCUAGGAGCCUUGAAGGAUCACCAGCCACCGACCCUCCAUCAGGGCCAACUGGGCAGGAAAGGGAACAUUGCAGACCUGAUUUCCCGACGAUGUCACCCUGUCCUCCCUCCUUGCUUCUUGCUCUGCUAACUCAACUCUGCCUUCCUCUUUUUCAUUCUUCUACUCUGCCCUAUAUGGAGGACAAAUGGACACCAGGGGUGCUAACCUUAUUGGUGCCUGCCCCAGCCUACCCCAGGUGCCAGCAGACUCUCGUGCACAGGAGGCUCCCACAGUUAUGGAGCCAGGAAAGAAUUUCUCUGCACUGGAUGGACUGUAUAUUGAGAUUAAAAAUUAUAUUCCUUAUAUUCCUGCUUAUAUCAAUGCUCUCUCUGUAAAACCUCUUCCUAGCCUCAUUUCUCUCAACUGAUCUUGUUUAGGCGUUGUAUUCCUUUUAUUUACUCUUUGCUUGACUGCUUCCUCCUAACCCUCUACCCACUAGCACUCUACUUCCUAAAGCUGUUGUGUCAUUAACUCUGUUGGAUCAACUCUCUGGGAAAAGAUUCUGUUAAUGUAAGUGCACUUACUCCCUGGAUGUUGUCACUAGUCUAGUGGCUUUUGCUAAAUAAACCUUUCUUAUUUCUA'
protein_sequence = 'MPSKTKYNLVDDGHDLRIPLHNEDAFQHGICFEAKYVGSLDVPRPNSRVEIVAAMRRIRYEFKAKNIKKKKVSIMVSVDGVKVILKKKKKLLLLQKKEWTWDESKMLVMQDPIYRIFYVSHDSQDLKIFSYIARDGASNIFRCNVFKSKKKSQAMRIVRTVGQAFEVCHKLSLQHTQQNADGQEDGESERNSNSSGDPGRQLTGAERASTATAEETDIDAVEVPLPGNDVLEFSRGVTDLDAVGKEGGSHTGSKVSHPQEPMLTASPRMLLPSSSSKPPGLGTETPLSTHHQMQLLQQLLQQQQQQTQVAVAQVHLLKDQLAAEAAARLEAQARVHQLLLQNKDMLQHISLLVKQVQELELKLSGQNAMGSQDSLLEITFRSGALPVLCDPTTPKPEDLHSPPLGAGLADFAHPAGSPLGRRDCLVKLECFRFLPPEDTPPPAQGEALLGGLELIKFRESGIASEYESNTDESEERDSWSQEELPRLLNVLQRQELGDGLDDEIAV'
scales = {'A': {'A': -0.103534,
  'C': -0.141027,
  'D': -0.057364,
  'E': 0.025673,
  'F': -0.160025,
  'G': 0.155926,
  'H': 0.114486,
  'I': -0.064125,
  'K': 0.058532,
  'L': -0.093059,
  'M': -0.093239,
  'N': 0.211717,
  'P': -0.187549,
  'Q': 0.286417,
  'R': 0.045458,
  'S': -0.068299,
  'T': -0.256675,
  'V': -0.126037,
  'W': -0.352338,
  'Y': -0.059715},
 'C': {'A': -0.170978,
  'C': 0.460587,
  'D': 0.011352,
  'E': 0.456847,
  'F': -0.157944,
  'G': 0.209027,
  'H': -0.166851,
  'I': -0.016438,
  'K': 0.093534,
  'L': 0.112175,
  'M': -0.08656,
  'N': 0.439522,
  'P': -0.189412,
  'Q': -0.031332,
  'R': 0.009869,
  'S': -0.357084,
  'T': -0.360561,
  'V': 0.248514,
  'W': 0.417484,
  'Y': 0.059507},
 'G': {'A': 0.157098,
  'C': 0.006956,
  'D': 0.047448,
  'E': -0.193184,
  'F': 0.480436,
  'G': -0.279106,
  'H': 0.144117,
  'I': 0.170305,
  'K': -0.186721,
  'L': 0.177462,
  'M': -0.139056,
  'N': -0.225754,
  'P': 0.097558,
  'Q': -0.085275,
  'R': -0.055897,
  'S': 0.273463,
  'T': 0.166763,
  'V': 0.108108,
  'W': 0.093968,
  'Y': 0.263202},
 'U': {'A': 0.130008,
  'C': -0.012462,
  'D': 0.005557,
  'E': 0.043517,
  'F': -0.269202,
  'G': -0.031819,
  'H': -0.198936,
  'I': -0.084679,
  'K': 0.096389,
  'L': -0.153234,
  'M': 0.351179,
  'N': -0.215165,
  'P': 0.240871,
  'Q': -0.065384,
  'R': 0.005856,
  'S': 0.31803,
  'T': 0.264754,
  'V': -0.078486,
  'W': 0.094917,
  'Y': -0.118544}}

Answer 1

It is likely not the dictionary look-up that kills your performance here, and more those nested for loops. Python is known for its notoriously slow (for) loops.

Since my numpy fu is a little bit weak lately, short of writing the code in C(++), I could only come up with an improvement using a nested list comprehension:

def affinity_matrix_lc(protein_sequence, rna_sequence):
    python_matrix = [[scales[base][item] for base in rna_sequence]
                     for item in protein_sequence]

    return np.array(python_matrix)

Conceptually it is very similar to RomanPerekhrest's version, but slightly easier to see what's going on I would say.

The list comprehension reduces the runtime of your example by about 30% here on my machine, with on-par performance to the one from Roman.

%timeit affinity_matrix_op(protein_sequence, rna_sequence)   # your original code
409 ms ± 3.79 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit affinity_matrix_rp(protein_sequence, rna_sequence)   # RomanPerekhrest's version
277 ms ± 5.63 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit affinity_matrix_lc(protein_sequence, rna_sequence)              
275 ms ± 2.06 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Edit - int version at OP's request:

I cobbled up a version of your code that first performs a translation to a vector of ints for both inputs. The look-up then uses numpy's indexing and broadcasting features to get the affinity matrix:

RNA_MAPPING = {'U': 0, 'G': 1, 'C': 2, 'A': 3}
PROTEIN_MAPPING = {
    'R': 0,'A': 1,'S': 2,'D': 3,'Q': 4,'N': 5,'W': 6,'V': 7,'L': 8,'K': 9,
    'H': 10,'E': 11,'C': 12,'G': 13,'Y': 14,'P': 15,'M': 16,'I': 17,'T': 18,
    'F': 19
}
LUT = np.array([
    [ 0.005856,  0.130008,  0.31803 ,  0.005557, -0.065384, -0.215165,
      0.094917, -0.078486, -0.153234,  0.096389, -0.198936,  0.043517,
     -0.012462, -0.031819, -0.118544,  0.240871,  0.351179, -0.084679,
      0.264754, -0.269202],
    [-0.055897,  0.157098,  0.273463,  0.047448, -0.085275, -0.225754,
      0.093968,  0.108108,  0.177462, -0.186721,  0.144117, -0.193184,
      0.006956, -0.279106,  0.263202,  0.097558, -0.139056,  0.170305,
      0.166763,  0.480436],
    [ 0.009869, -0.170978, -0.357084,  0.011352, -0.031332,  0.439522,
      0.417484,  0.248514,  0.112175,  0.093534, -0.166851,  0.456847,
      0.460587,  0.209027,  0.059507, -0.189412, -0.08656 , -0.016438,
     -0.360561, -0.157944],
    [ 0.045458, -0.103534, -0.068299, -0.057364,  0.286417,  0.211717,
     -0.352338, -0.126037, -0.093059,  0.058532,  0.114486,  0.025673,
     -0.141027,  0.155926, -0.059715, -0.187549, -0.093239, -0.064125,
     -0.256675, -0.160025]
])

def affinity_matrix_int(protein_sequence, rna_sequence):
    protein_sequence_int = np.array(
        [PROTEIN_MAPPING[i] for i in protein_sequence], dtype=int
    ).reshape(-1, 1)
    rna_sequence_int = np.array(
        [RNA_MAPPING[i] for i in rna_sequence], dtype=int
    ).reshape(1, -1)
    return LUT[rna_sequence_int, protein_sequence_int]

I then checked the correctness of the implementation using

>>> np.allclose(
        affinity_matrix_int(protein_sequence, rna_sequence),
        affinity_matrix_op(protein_sequence, rna_sequence)
    )
True

And did another round of timing:

%timeit affinity_matrix_int(protein_sequence, rna_sequence)             
10.9 ms ± 36.4 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

To answer your question in the comments: It turns out, this is quite a bit faster ;-)

Answer 2

Boosting performance

The initial approach creates a 2-dimensional list of empty lists beforehand and performs 2 538 096 expensive list.append operations.

To make it go faster and more efficiently we'll go through 2 steps:

• schedule a generator/iterator yielding a sequence of all the needed values
• generate numpy array from the scheduled generaor with convenient numpy.fromiter routine and immediately giving a new shape to the array with numpy.reshape routine (assuming that len(protein_sequence) points to number of rows and len(rna_sequence) is a number of columns)

---

The final optimized affinity_matrix function:

def affinity_matrix(protein_sequence, rna_sequence):
    it = (scales[base][prot] for prot in protein_sequence for base in rna_sequence)
    return np.fromiter(it, dtype=float).reshape(len(protein_sequence), len(rna_sequence))

---

Let's move to tests. I've renamed the old function to affinity_matrix_old for comparison.

>>> from timeit import timeit
>>> timeit('affinity_matrix_old(protein_sequence, rna_sequence)', setup='from __main__ import affinity_matrix_old, protein_sequence, rna_sequence', number=10)
3.1148553189996164
>>> timeit('affinity_matrix(protein_sequence, rna_sequence)', setup='from __main__ import affinity_matrix, protein_sequence, rna_sequence', number=10)
1.9052914250059985

---

The resulting array looks as:

>>> affinity_matrix(protein_sequence, rna_sequence)
array([[-0.093239, -0.08656 , -0.093239, ..., -0.08656 ,  0.351179,
        -0.093239],
       [-0.187549, -0.189412, -0.187549, ..., -0.189412,  0.240871,
        -0.187549],
       [-0.068299, -0.357084, -0.068299, ..., -0.357084,  0.31803 ,
        -0.068299],
       ...,
       [-0.064125, -0.016438, -0.064125, ..., -0.016438, -0.084679,
        -0.064125],
       [-0.103534, -0.170978, -0.103534, ..., -0.170978,  0.130008,
        -0.103534],
       [-0.126037,  0.248514, -0.126037, ...,  0.248514, -0.078486,
        -0.126037]])

Answer 3

You can speed this further in 2 ways.

local variables

a local variable lookup is faster than a global. Since you lookup scales an awful lot, this can improve performance

np.fromiter

where you can even specify the initial length to improve performance further

def affinity_matrix_generator(protein_sequence, rna_sequence, scales=scales):
    for base in rna_sequence:
        base_scale = scales[base]  # saving another dict lookup
        for protein in protein_sequence:
            yield base_scale[protein]

def affinity_matrix(protein_sequence, rna_sequence, scales=scales):
    iterator = affinity_matrix_generator(protein_sequence, rna_sequence, scales=scales)
    size = len(protein_sequence) * len(rna_sequence)
    return np.fromiter(iterator, dtype=float, count=size).reshape(len(protein_sequence), len(rna_sequence))