Colorize connected components by molecular weight

This example demonstrates how to colorize connected components based on their molecular weight (MW).

Example: Colorize connected components by molecular weight

"""
This script colorizes molecules based on connected components and their molecular weight (MW)
"""

import math

def get_molecular_weight(atom_indexer):
    '''
    Returns the molecular weight of atoms in the atom_indexer
    '''
    mass = SBQuantity.mass(0.0)
    for atom in atom_indexer:
        mass += atom.molecularWeight
    return mass.Da

def get_connected_components():
    '''
    Returns a list of indexers with connected components from the active document
    '''
    connected_component_list = []
    # get all atoms in the active document
    atom_indexer = SAMSON.getNodes('node.type atom')
    # go through atoms to get their connected components
    while len(atom_indexer):
        atom = atom_indexer[0]
        # get connected component for this atom
        connected_component = atom.getConnectedComponent()
        connected_component_list.append(connected_component)
        # remove atoms of this connected component from the atom indexer
        for a in connected_component:
            atom_indexer.removeNode(a)
    return connected_component_list

connected_component_list = get_connected_components()
masses = []
n_connected_components_per_MW = dict()
for connected_component in connected_component_list:
    n_atoms = len(connected_component)
    # reduce significant digits in mass
    mw = get_molecular_weight(connected_component)
    mass_floor = math.floor(mw.value * 1e2) * 1e-2
    mass_ceil = math.ceil(mw.value * 1e2) * 1e-2
    mass = 0.5 * (mass_floor + mass_ceil)
    masses.append(mass)
    n_connected_components_per_MW[mass] = n_connected_components_per_MW.get(mass, 0) + 1

# sort in ascending order
n_connected_components_per_MW = dict(sorted(n_connected_components_per_MW.items()))
# print the number of molecules per number of atoms
for key, value in n_connected_components_per_MW.items():
    print(f"{value} molecules with MW {key} Da")

# get min and max masses with only 2 significant digits
min_mass = math.floor(min(masses) * 1e2) * 1e-2
max_mass = math.ceil(max(masses) * 1e2) * 1e-2

# make the operation undoable
SAMSON.beginHolding("Colorize molecules by mass")

# colorize based on the number of atoms in a connected component
for i, connected_component in enumerate(connected_component_list):
    mass = masses[i]
    # get the color 'intensity' value, which should be in the [0, 1] range
    h = (mass - min_mass) / (max_mass - min_mass)
    # get RGB color from the palette
    color = SBPaletteDefaultPalette.flexibleDivergingHCLSpectral.getColor(h)
    for atom in connected_component:
        atom.setColor(color)

# stop holding the undoable operation
SAMSON.endHolding()