Show neighbours in nglview

Nglview is a really nice Python library which encodes a widget to show a NGL viewport, a JS 3D protein viewer used until recently by the PDB. One annoying feature is that one cannot select neighbours as easily as say PyMOL's "select byres HEM around 3".  But it is possible and here is how.

There are two ways.
The first is to generate a selection string of the neighbours* with, say, PyMOL or PyRosetta, and use that. The other is to use NGL's inbuild functionality even if it is not exposed to Python.
* Neighbours or "neighbors" if you either are American or have muddled your brian by coding too much like me...

The latter consists in:

# load a structure:
import nglview as nv
view:nv.NGLWidget = nv.show_pdbid('1MBN')

# tamper with the JS
view = nv.show_pdbid('1MBN')
def add_neighbors(view, selection:str, comp_id:int=0, radius:float=5, style:str = 'hyperball', color:str='gainsboro'):
    view._js(f"""const comp = this.stage.compList[{comp_id}]
                 const target_sele = new NGL.Selection('{selection}');
                 const radius = {radius};
                 const neigh_atoms = comp.structure.getAtomSetWithinSelection( target_sele, radius );
                 const resi_atoms = comp.structure.getAtomSetWithinGroup( neigh_atoms );
                 comp.addRepresentation( "{style}", {{sele: resi_atoms.toSeleString(),
                 									  colorValue: "{color}",
                                                      multipleBond: true
                                                                   }});
                 comp.addRepresentation( "contact", {{sele: resi_atoms.toSeleString()}});
            """)
# show the viewer (via the IPython.display.display than wraps cell runs)
view

This snippet is nothing more than the example snippet in the manual of ngl.js, but with JavaScript's infamous this. Compare:

stage.loadFile( "rcsb://3pqr" ).then( function( o ){
    // get all atoms within 5 Angstrom of retinal
    var selection = new NGL.Selection( "RET" );
    var radius = 5;
    var atomSet = o.structure.getAtomSetWithinSelection( selection, radius );
    // expand selection to complete groups
    var atomSet2 = o.structure.getAtomSetWithinGroup( atomSet );
    o.addRepresentation( "licorice", { sele: atomSet2.toSeleString() } );
    o.addRepresentation( "cartoon" );
    o.autoView();
} );

Looking back at the PyMOL select byres HEM expand 3, there are two parts: HEM expand 3, which will select the neighbouring atoms (or "neighboring atoms" if you will), and byres sele which expands the selection of atoms to complete residues. This is the same thing that happens in the snippet above. PyRosetta does it differently as it is residue-centric as opposed to atom-centric because residues are defined by residue types and it is not possible to have a random extra atom here and there.

The major difference between the snippet from the JavaScript manual and the one for NGLView is that in the former an object in the local namespace is called stage and the function loadFile returns a JavaScript Promise (because IO operations are classically asynchronous in JavaScript) whose result is a component, whereas in the Pythonically excuted JS code the object this refers to a stage object. The object this is similar to self in Python, but a lot weirder, for example it can get unbound and end up being the global namespace window or a function has its own this (as a class is a fancy function after all in JavaScript) unless an arrow function is used or the method bind is called on the offending function.

Excuting JavaScript

In a previous blog post I discuss JS in colab and Jupyter, but here we are dealing with a well made widget, so things are nice and tidy!

Firstly, adding a IPython.display.display(IPython.display.Javascript(js_code) to the Python code has the problem that the code will be run in the global namespace and the stage object will not be available or would NGL for that matter, hence why the widget needs to be used.

nv.NGLWidget is a subclass of ipywidgets.DOMWidget and you may notice that an instance will have an attribute comm and a method send, which is how the JS–Python dialogue happens. However, it is extremely tedious and weird and what can be done depends on how the frontend listens to changes, generally via a call to the method listenTo in initialize or render —see documentation for more.
Luckly this is abstracted out and a bit of reverse engineering can reveal how to use it to one's advantage as seen in my snippet with the call to the private method _js

GitHub searches

There are two ways to figure out how to use a piece of code. Starting from the bedrock and going down into the specifics or starting from a specific method and seeing how it works.

Top down

The first place to check for a given functionality in any codebase on GitHub is to search in Issues. The next is to search in the repository for pieces of code you'd expect, like self.send: in nglview's GitHub this leads us to base.py as one would expect were one finds out that the send and its tratlet buddies are neatly wrapped up in the method _call, which is next to the method _js, which I used in my snippet. Now, there may be a better way to do this so searching the _call would be the next step. This is where the top-down and bottom-up approach converge.

Bottom up

Given a specific method which is believed to do the functionality one wants to replicate one can see its code block in a Jupyter notebook with the handy module inspect. For example:

import inspect
print_code = lambda fun: print(inspect.getsource(fun))
print_code(view.component_0.add_representation)

Does a call to self._view.add_representation, checking view.add_representation, one finds:

self._remote_call('addRepresentation',
                          target='compList',
                          args=[
                              d['type'],
                          ],
                          kwargs=params)

In JS addRepresentation is a function of the component, thus revealing how the Python to JavaScript dialogue. The target='compList', is promising. However, one key part of the JS neighbourhood selection snippet is instantiating a Selection class, so a dictionary mocking it may work, but would be overly complicated. Searching the GitHub code I could not find a call to NGL.Selection in the widget typescript file or in the python module, so appears to be a dead-end. Similarly, this is pretty common outcome with Python modules that wrap C++ code with Boost or pybind11: one often finds cases where in C++ there's a variable that does not make it to the Python side and the best bet it try a different route. Luckily we have in our case the _js which suits us perfectly!