(cplx)
| 15 | |
| 16 | |
| 17 | def to_tensor(cplx): |
| 18 | hc, lc = cplx.get_heavy_chain(), cplx.get_light_chain() |
| 19 | antigen_chains = cplx.get_antigen_chains(interface_only=True, cdr=None) |
| 20 | |
| 21 | # prepare input |
| 22 | chain_lists, begins = [], [] |
| 23 | chain_lists.append([hc]) |
| 24 | begins.append(VOCAB.BOH) |
| 25 | chain_lists.append([lc]) |
| 26 | begins.append(VOCAB.BOL) |
| 27 | chain_lists.append(antigen_chains) |
| 28 | begins.append(VOCAB.BOA) |
| 29 | |
| 30 | X, S, = [], [] |
| 31 | chain_start_ends = [] # tuples of [start, end) |
| 32 | |
| 33 | # format input, box is begin of chain x |
| 34 | corrupted_idx = [] |
| 35 | for chains, box in zip(chain_lists, begins): |
| 36 | # judge if the chain has length |
| 37 | skip = True |
| 38 | for chain in chains: |
| 39 | if len(chain): |
| 40 | skip = False |
| 41 | break |
| 42 | if skip: |
| 43 | continue |
| 44 | X.append([(0, 0, 0) for _ in range(4)]) # begin symbol is global symbol, update coordination afterwards |
| 45 | S.append(VOCAB.symbol_to_idx(box)) |
| 46 | start = len(X) |
| 47 | for chain in chains: |
| 48 | for i in range(len(chain)): # some chains do not participate |
| 49 | residue = chain.get_residue(i) |
| 50 | coord = residue.get_coord_map() |
| 51 | x = [] |
| 52 | for atom in ['N', 'CA', 'C', 'O']: |
| 53 | if atom in coord: |
| 54 | x.append(coord[atom]) |
| 55 | else: |
| 56 | coord[atom] = (0, 0, 0) |
| 57 | x.append((0, 0, 0)) |
| 58 | corrupted_idx.append(len(X)) |
| 59 | # print_log(f'Missing backbone atom coordination: {atom}', level='WARN') |
| 60 | |
| 61 | X.append(np.array(x)) |
| 62 | S.append(VOCAB.symbol_to_idx(residue.get_symbol())) |
| 63 | X[start - 1] = np.mean(X[start:], axis=0) # coordinate of global node |
| 64 | chain_start_ends.append((start - 1, len(X))) |
| 65 | |
| 66 | # deal with corrupted coordinates |
| 67 | for i in corrupted_idx: |
| 68 | l, r = i - 1, i + 1 |
| 69 | if l > 0 and r < len(X): # if at start / end, then leave it be |
| 70 | X[i] = (X[l] + X[r]) / 2 |
| 71 | |
| 72 | # set CDR pos for heavy chain |
| 73 | offset = S.index(VOCAB.symbol_to_idx(VOCAB.BOH)) + 1 |
| 74 | L = ['0' for _ in X] |
no test coverage detected