Colorize atoms by velocity along path#

The example below demonstrates how to colorize atoms based on their velocity norm along a path (trajectory). After each step the viewport’s capture is saved into a file. This script uses a default color palette for colorization.

Note

To see how the default color palettes look like check User Guide: Default color palettes

Colorize atoms based on their velocity norm along a path using a color palette#

def get_range_of_velocities(path):
    """ Returns the range of velocity norm """
    minV = maxV = SBVelocity3().norm()
    if path:
        if path.hasVelocityData():
            velocityData = path.getVelocityData()
            if len(velocityData):
                if len(velocityData[0]):
                    minV = maxV = velocityData[0][0].norm()
                # go through frames
                for velocities in velocityData:
                    # go through velocities per frame
                    for velocity in velocities:
                        vn = velocity.norm()
                        if vn > maxV: maxV = vn
                        elif vn < minV: minV = vn

    return [minV, maxV]

# get path
pathIndexer = SAMSON.getNodes('node.type path')

if len(pathIndexer):
    # get the path from the user
    ret, capture_filepath = SAMSON.getPathFromUser(dialogTitle = 'Choose folder to save captures...')
else:
    ret = False

if ret and len(pathIndexer):
    # get the 1st path
    path = pathIndexer[0]
    if path.hasVelocityData():
        # get all atoms in the path
        atomIndexer = path.getAtomIndexer()

        minVelNorm, maxVelNorm = get_range_of_velocities(path)
        minVelNorm = SBVelocity3().norm() # set the min value to 0

        # choose a default color palette
        palette = SBPaletteDefaultPalette.divergingHCLBlue2Red

        # make the operation undoable
        with SAMSON.holding("Colorize atoms"):
            for step in range(0, len(path)):

                path.currentStep = step

                # colorize each atom according to its coordinate
                for atom in atomIndexer:
                    velocity = path.getVelocity(step, atom)
                    velNorm = velocity.norm()
                    # get the color 'intensity' value, which should be in the [0, 1] range
                    h = ((velNorm - minVelNorm) / (maxVelNorm - minVelNorm)).value
                    # get RGB color from the palette
                    color = palette.getColor(h)
                    # set the color of the node
                    atom.setColor(color)
                
                SAMSON.captureViewportToFile(
                    filename = f"{capture_filepath}/capture-{step}.png",
                    width = 1200, height = 800,
                    transparentBackground = False, usePathTracing = False, showProgressBar = False)
                
        

Note

You can then create a movie from the saved capture images - see the Presentation: play path example.