PDB ファイルを読み込みコンタクトマップをミリ秒で計算する

タンパク質の構造解析をする際、各アミノ酸残基が近接しているかどうかを記述するコンタクトマップが役立つ場合があります。通例、コンタクトマップは各アミノ酸の Cα の三次元座標を用いて、互いの距離が一定以下（8Åとすることが多い）であれば近接していると判断し作成されます。

基本的には全残基の組み合わせ（N^2 オーダー）だけ距離計算をする必要があり、下手に実装すると相応の時間がかかってしまいます。今回は、このコンタクトマップ計算を numpy ベースで高速に実行するコードをメモ程度に残します。

(2020.02.18追記)
その後 scipy.spatial.distance.cdist の存在を知り、圧倒的に手軽かつ多機能で高速だったので、「scipy.spatial.distance.cdist を使うと」欄に内容を追記しておきました。完全に無知だった。。。。。本記事はいちおう残しておきます。

使用するデータ

適当な PDB ファイルを用意しておきます。

wget http://www.ebi.ac.uk/pdbe/entry-files/download/pdb1nd4.ent

このタンパク質は配列長 264 ですが、結晶構造からは始めの 9 つのアミノ酸が欠損しています。つまり、255 残基の 3 次元座標が保存されています。

コンタクトマップ計算

PDB ファイル読み込みのために biopython を使います。このライブラリには Bio.PDB.NeighborSearch という近傍探索コードが含まれているのですが、今回はこれを使わずスクラッチ実装します。

import Bio.PDB

structure = Bio.PDB.PDBParser().get_structure("1nd4", "./pdb1nd4.ent")
chain = next(structure.get_chains())

この PDB ファイルには2つの chain が含まれており、コンタクトマップ計算の事例として1つ目のものを取得しています。

コンタクトマップ計算は、次のように実装しました。

from typing import Tuple
import numpy as np

def compute_distance_and_contact_map(
    chain: Bio.PDB.Chain.Chain,
    thr: float
) -> Tuple[np.ndarray]:
    coords = np.array(
        [r["CA"].coord for r
         in Bio.PDB.Selection.unfold_entities(chain, "R")
         if r.full_id[-1][0] == " "])
    distance_map = np.linalg.norm(coords[:, None, :] - coords, axis=-1)
    contact_map = (distance_map <= thr).astype(int)
    
    return distance_map, contact_map

実質的には 1 行で計算が実行されています。

1-4 行目は、chain から水分子等の不要な情報を除外し、各アミノ酸残基の Cα の座標を抽出しています。5 行目が distance map 計算、6 行目が contact map 計算の実体で、後者は前者に対して適当な閾値を設けることで得られます。

もともと coords は (n_residues, 3) という shape であり、distance_map 計算時に (n_residues, 1, 3) と余分な次元を追加することで全 Cα の組み合わせを作ることができます（numpy のブロードキャスト）。距離は np.linalg.norm に計算させることで高速化が見込めます。

実際、ipython 上で実行時間を計測すると

%timeit compute_distance_and_contact_map(chain, 8)
# 2.77 ms ± 26.4 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

となり、約 2.8ms で計算が完了しています。なお、先ほどのように None を入れずに各残基について明示的にループを回すと、約10倍ほど実行時間が延びることも確認しています。

scipy.spatial.distance.cdist を使うと

ここまでは numpy を使った自作関数について書いてきましたが、実は scipy に実装されている距離計算関数を使うこともできます。特に scipy.spatial.distance.cdist は、ユークリッド距離以外にも複数種の距離を指定して計算できる便利な関数で、実行も速く良いです。

今回の例では、

%timeit cdist(coords, coords, metric="euclidean")
# 392 µs ± 64.2 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

となり、容易にマイクロ秒が達成されています。こちらを使うのが良さそうです。

なお scipy.distance.distance_matrix という関数もありますが、こちらは低速でした。

%timeit distance_matrix(coords, coords)
# 3.64 ms ± 436 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

なぜ cdist が速いか？ということですが、おそらくC言語で実装されているからだと思います。
https://github.com/scipy/scipy/blob/v1.4.1/scipy/spatial/src/distance_wrap.c#L256
が実体で、scipy ライブラリのビルド時に setup.py 経由でコンパイルされ、
https://github.com/scipy/scipy/blob/v1.4.1/scipy/spatial/distance.py#L2778
で python 側から呼ばれていますね。

検証環境

• Ubuntu 18.04
• Python 3.7.6

下記ライブラリは conda 経由でインストールした。

• biopython 1.74 py37h7b6447c_0
• ipython 7.11.1 py37h39e3cac_0
• numpy 1.17.4 py37hc1035e2_0
• numpy-base 1.17.4 py37hde5b4d6_0
• blas 1.0 mkl
• mkl 2019.4 243
• scipy 1.4.1