Custom carbon colours in py3Dmol

 Due to NGLView (the Python module) having a frozen older IPywidget version it breaks Colab and the major change for the latter library was a year ago (early 2023), so I am forced to revisit old code and switch to py3Dmol in my Colab demos. Today I figured out how to use custom carbon colours.

Given a hex code for a colour (#📕📕📗📗📘📘), in PyMOL one can set it to the carbons via color 0x📕📕📗📗📘📘, element C and 👾👾, which is a bit cryptic (0x is for hex numbers): so even with a common program that's a desktop binary and not a JS widget in Python the synthax can be tricky. In py3Dmol it is very tricky. If the color is a CSS built-in, in py3Dmol / 3Dmol.js name+'Carbon' will work. Example:

import py3Dmol
from rdkit import Chem
from rdkit.Chem import AllChem

# A fun 3D molecule (if you don't like the look of a tropane/adamantane ring please go to A&E / ER)
atropine = Chem.MolFromSmiles('CN(C)C(=O)OC1C2CC3CC1CC(C2)(C3)OC(=O)C')
AllChem.EmbedMolecule(atropine)

viewer = py3Dmol.view(width=350, height=350)
viewer.addModel(Chem.MolToMolBlock(atropine), 'sdf')
viewer.setStyle({'model':-1}, {'stick':{'colorscheme':'coralCarbon', 'opacity': 1}})
viewer.zoomTo()
viewer.show()

If a hex code is passed it will not. The 3Dmol.js documentation talks of passing a JS function to Python to define custom colour schemes here. So how does py3Dmol inject JS?

This depends on whether a viewer is displayed or not. There are three string attributes startjs, endjs, and updatejs. The method show calls _make_html, which concatenates startjs and endjs, and adds it via IPython.display.publish_display_data —which is really cool and I only learnt about it this way. The update method does the same but on the updatejs string.
So to inject JS in a yet to be shown viewer one can do:

viewer = py3Dmol.view(width=350, height=350)
viewer.addModel(Chem.MolToMolBlock(atropine), 'sdf')
viewer.startjs += '''\n
let customColorize = function(atom){
      // attribute elem is from https://3dmol.csb.pitt.edu/doc/AtomSpec.html#elem
      if (atom.elem === 'C'){
          return "#00FF00"
      }else{
          return $3Dmol.getColorFromStyle(atom, {colorscheme: "whiteCarbon"});
      }
  }    
\n'''
viewer.setStyle({'model':-1}, {'stick':{'colorfunc': 'customColorize', 'opacity': 0.7}})
# make it a function not a string "customColorize"
viewer.startjs = viewer.startjs.replace('"customColorize"', 'customColorize')
viewer.zoomTo()
viewer.show()

In the above, a trick happens: customColorize is added as a string, but the code is changed to make it a variable. Actually this is not the only case where this happens: the code will have a viewer_UNIQUEID main namespace variable going on, which gets replaced by the value of viewer.uniqueid. Main namespace pollution is frowned upon by purists, but this is very handy for debugging as one can make the dev console pop-up and type window.viewer_👾👾👾👾👾 and do whatever test!

Footnote

The Pythonic repo of py3Dmol is on GitHub (https://github.com/3dmol/3Dmol.js/tree/master/py3Dmol). But at the time, I was unable to find it, so I used this function:

def display_source(function):
    """
    Display the source code formatted
    """
    import inspect
    from IPython.display import HTML, display
    from pygments import highlight
    from pygments.lexers import PythonLexer
    from pygments.formatters import HtmlFormatter
    code:str = inspect.getsource(function)
    html:str = highlight(code, PythonLexer(), HtmlFormatter(style='colorful'))
    stylesheet:str = f""
    display(HTML(f"{stylesheet}{html}"))
    
display_source(type(viewer))