diff --git a/prolintpy/vis/show_contact_projection.py b/prolintpy/vis/show_contact_projection.py index ad8a2ef..f6f0dc6 100644 --- a/prolintpy/vis/show_contact_projection.py +++ b/prolintpy/vis/show_contact_projection.py @@ -7,7 +7,7 @@ from prolintpy.utils.shift_range import shift_range -def show_contact_projection(t, bf, protein=None, ngl_repr='surface', cmap='Reds'): +def show_contact_projection(t, bf, protein=None, residue_list=None, ngl_repr='surface', cmap='Reds'): """Visualize lipid-protein contacts by mapping them onto the structure of the protein. Parameters @@ -19,6 +19,9 @@ def show_contact_projection(t, bf, protein=None, ngl_repr='surface', cmap='Reds' List of contacts. protein : ProLint.Protein + + residue_list: list + List of residues. Required when working with a subset of residues. ngl_repr: str One representation that will be used by nglview to display the protein. The following are supported @@ -45,11 +48,33 @@ def show_contact_projection(t, bf, protein=None, ngl_repr='surface', cmap='Reds' t_slice = t[0].atom_slice(indices) resseq = df.resSeq.to_list() + + + atomic_bfactors = [] + if residue_list is not None: + init, bfactors = 1, [] + for idx, residue in enumerate(residue_list): + while init < residue: + bfactors.append(0) + init += 1 + bfactors.append(bf[idx]) + init += 1 + dif = indices.size - len(bfactors) + for remaining_residues in range(dif): + bfactors.append(0) + + atomic_bfactors = [] + for atom in resseq: + atomic_bfactors.append(bfactors[atom-1]) + else: + if len(df.resSeq.unique()) != len(bf): + raise TypeError ('When projecting only a subset of residues provide a list of tuples: [(residue_id, value), ...]') + for atom in resseq: + atomic_bfactors.append(bf[atom-1]) bf_cmap = cm.get_cmap(cmap) - colors = [mpl.colors.to_hex(x) for x in bf_cmap(shift_range(bf))] - # cs = [[y, str(x+offset)] for x, y in enumerate(colors)] + colors = [mpl.colors.to_hex(x) for x in bf_cmap(shift_range(atomic_bfactors))] cs = [[y, str(resseq[x])] for x, y in enumerate(colors)] scheme = nv.color._ColorScheme(cs, 'bf') @@ -77,4 +102,3 @@ def show_contact_projection(t, bf, protein=None, ngl_repr='surface', cmap='Reds' return view -