How to change atom colors from plot3d (nglview) in ball-and-stick based on species position and type

[ahmedabdelkawy]

@freyso had the question and also then the answer on how to custom color atoms in nglview while using ball and stick

• I added couple of lines below to further explain what is going on

The create_colorscheme creates a dictionary and a color scheme label for the indices of the atoms with specific colors and a label (from the getcolor function) that is passed to the nglview.color._ColorScheme directly and then the entire nglview color scheme is passed to add_ball_and_stick function

def getcolor(atom):
    """ Take a pyiron atom, return the custom color based on type and height"""
    if atom.symbol == "Si":
        if atom.position[2] > 8.5:
            return '#0000aa'
        else:
            return '#c0c0ff'
    elif atom.symbol == "B":
        return 'green'
    else:
        return None #'white'

def create_colorscheme (struct, func):
    """ Create a NGLview colorscheme from a structure and a color-defining python function"""
    colors = {}
    for i,a in enumerate(struct):
        col = func(a)
        if col is None:
            col = 'white'
        if col not in colors:
            colors[col] = []
        colors[col].append (i)
    idxlist = { col : ",".join ([str(i) for i in idx]) for col,idx in colors.items () }
    label = hashlib.md5("|".join([col + ":" + idx for col,idx in idxlist.items ()]).encode ("utf-8")).hexdigest ()
    return nv.color._ColorScheme([ [col, "@" + idx] for col,idx in idxlist.items ()], label=label)

import nglview as nv
import hashlib
from pyiron import Project

pr = Project('custom_atoms_colors')
struct = pr.create.structure.surface('Si', 'diamond100', (2,2,10),vacuum=10,a=4)
struct[-1] = 'B'
repeat=(2,2,1)
repstr = struct.repeat(repeat)
view=repstr.plot3d (view_plane=((7,1.5,-1),(0,1,0)))
view.center ()
view.clear ()
view.add_ball_and_stick(aspectRatio=3.5,color=create_colorscheme(repstr,getcolor ))
view.parameters = {"cameraType": "orthographic"}
view.control.zoom (0.2)
view