<jupyter_start><jupyter_text>Visualizing mbuild CompoundsIn this example, we will demonstrate the `visualize` function which uses [fresnel](https://fresnel.readthedocs.io) to visualize an [mbuild Compound](https://mosdef.org/mbuild/data_structures.htmlcompound). We will show how to display a static view, interact with the visualization, and render it as a reproducible, publication quality image. We will highlight scope of the `visualize` function including the capability to color particles by name using various color schemes and visualizing the simulation box.<jupyter_code># QtCore, needed for interactive visualization, will open a separate python process
from PySide2 import QtCore

%gui qt
import random
import warnings

import fresnel
import fresnel.interact
import freud
import matplotlib.cm
import mbuild as mb
import numpy as np
import PIL
## Start building the fresnel scene
device = fresnel.Device()
scene = fresnel.Scene(device=device)

# Spheres for every particle in the system
geometry = fresnel.geometry.Sphere(scene, N=2)
geometry.position[:] = np.array([[0,0,0],[1,0,0]])
geometry.material = fresnel.material.Material(roughness=0.2)
geometry.outline_width = 0.01 * 1.0

# use color instead of material.color
geometry.material.primitive_color_mix = 1.0
geometry.color[:] = fresnel.color.linear([0.30, 0.30, 0.30])

# resize radii
geometry.radius[:] = 1.0

#geometry.material.spec_trans = 1


tracer = fresnel.tracer.Path(device=device, w=500, h=500)
scene.lights = fresnel.light.lightbox()
tracer.sample(scene, samples=64, light_samples=40)

#view = fresnel.interact.SceneView(scene)
#view<jupyter_output><empty_output><jupyter_text>The following cell contains the functions and variables we will use for visualization.<jupyter_code>cpk_colors = {
    "H": fresnel.color.linear([1.00, 1.00, 1.00]),  # white
    "C": fresnel.color.linear([0.30, 0.30, 0.30]),  # grey
    "N": fresnel.color.linear([0.13, 0.20, 1.00]),  # dark blue
    "O": fresnel.color.linear([1.00, 0.13, 0.00]),  # red
    "F": fresnel.color.linear([0.12, 0.94, 0.12]),  # green
    "Cl": fresnel.color.linear([0.12, 0.94, 0.12]),  # green
    "Br": fresnel.color.linear([0.60, 0.13, 0.00]),  # dark red
    "I": fresnel.color.linear([0.40, 0.00, 0.73]),  # dark violet
    "He": fresnel.color.linear([0.00, 1.00, 1.00]),  # cyan
    "Ne": fresnel.color.linear([0.00, 1.00, 1.00]),  # cyan
    "Ar": fresnel.color.linear([0.00, 1.00, 1.00]),  # cyan
    "Xe": fresnel.color.linear([0.00, 1.00, 1.00]),  # cyan
    "Kr": fresnel.color.linear([0.00, 1.00, 1.00]),  # cyan
    "P": fresnel.color.linear([1.00, 0.60, 0.00]),  # orange
    "S": fresnel.color.linear([1.00, 0.90, 0.13]),  # yellow
    "B": fresnel.color.linear([1.00, 0.67, 0.47]),  # peach
    "Li": fresnel.color.linear([0.47, 0.00, 1.00]),  # violet
    "Na": fresnel.color.linear([0.47, 0.00, 1.00]),  # violet
    "K": fresnel.color.linear([0.47, 0.00, 1.00]),  # violet
    "Rb": fresnel.color.linear([0.47, 0.00, 1.00]),  # violet
    "Cs": fresnel.color.linear([0.47, 0.00, 1.00]),  # violet
    "Fr": fresnel.color.linear([0.47, 0.00, 1.00]),  # violet
    "Be": fresnel.color.linear([0.00, 0.47, 0.00]),  # dark green
    "Mg": fresnel.color.linear([0.00, 0.47, 0.00]),  # dark green
    "Ca": fresnel.color.linear([0.00, 0.47, 0.00]),  # dark green
    "Sr": fresnel.color.linear([0.00, 0.47, 0.00]),  # dark green
    "Ba": fresnel.color.linear([0.00, 0.47, 0.00]),  # dark green
    "Ra": fresnel.color.linear([0.00, 0.47, 0.00]),  # dark green
    "Ti": fresnel.color.linear([0.60, 0.60, 0.60]),  # grey
    "Fe": fresnel.color.linear([0.87, 0.47, 0.00]),  # dark orange
    "default": fresnel.color.linear([0.87, 0.47, 1.00]),  # pink
}

# Made space to add more later
radii_dict = {"H": 0.05, "default": 0.06}


def distance(pos1, pos2):
    """
    Calculates euclidean distance between two points.

    Parameters
    ----------
    pos1, pos2 : ((3,) numpy.ndarray), xyz coordinates
        (2D also works)

    Returns
    -------
    float distance
    """
    return np.linalg.norm(pos1 - pos2)


def mb_to_freud_box(box):
    """  
    Convert an mbuild box object to a freud box object
    These sites are helpful as reference: 
    http://gisaxs.com/index.php/Unit_cell
    https://hoomd-blue.readthedocs.io/en/stable/box.html

    Parameters
    ----------
    box : mbuild.box.Box()

    Returns
    -------
    freud.box.Box()
    """
    Lx = box.lengths[0]
    Ly = box.lengths[1]
    Lz = box.lengths[2]
    alpha = box.angles[0]
    beta = box.angles[1]
    gamma = box.angles[2]

    frac = (
        np.cos(np.radians(alpha)) - np.cos(np.radians(beta)) * np.cos(np.radians(gamma))
    ) / np.sin(np.radians(gamma))
    c = np.sqrt(1 - np.cos(np.radians(beta)) ** 2 - frac ** 2)

    xy = np.cos(np.radians(gamma)) / np.sin(np.radians(gamma))
    xz = frac / c
    yz = np.cos(np.radians(beta)) / c

    box_list = list(box.maxs) + [xy, yz, xz]
    return freud.box.Box(*box_list)


def visualize(comp, color="cpk", scale=1.0, show_box=False):
    """
    Visualize an mbuild Compound using fresnel.
    
    Parameters
    ----------
    comp : (mbuild.Compound), compound to visualize
    color : ("cpk" or the name of a matplotlib colormap), color scheme to use (default "cpk")
    scale : (float), scaling factor for the particle, bond, and box radii (default 1.0)
    show_box : (bool), whether or not to visualize the box. The function will first check if the 
        user has assigned a box to comp.box, if not, then it will use comp.boundingbox
    
    Returns
    -------
    fresnel.Scene
    """

    ## Extract some info about the compound
    N = comp.n_particles
    particle_names = [p.name for p in comp.particles()]

    # all_bonds.shape is (nbond, 2 ends, xyz)
    all_bonds = np.stack([np.stack((i[0].pos, i[1].pos)) for i in comp.bonds()])

    N_bonds = all_bonds.shape[0]

    color_array = np.empty((N, 3), dtype="float64")
    if color == "cpk":
        # Populate the color_array with colors based on particle name
        # -- if name is not defined in the dictionary, use pink (the default)
        for i, n in enumerate(particle_names):
            try:
                color_array[i, :] = cpk_colors[n]
            except KeyError:
                color_array[i, :] = cpk_colors["default"]
    else:
        # Populate the color_array with colors based on particle name
        # choose colors evenly distributed through a matplotlib colormap
        try:
            cmap = matplotlib.cm.get_cmap(name=color)
        except ValueError:
            print(
                "The 'color' argument takes either 'cpk' or the name of a matplotlib colormap."
            )
            raise
        mapper = matplotlib.cm.ScalarMappable(
            norm=matplotlib.colors.Normalize(vmin=0, vmax=1, clip=True), cmap=cmap
        )
        particle_types = list(set(particle_names))
        N_types = len(particle_types)
        v = np.linspace(0, 1, N_types)
        # Color by typeid
        type_ids = np.array([particle_types.index(i) for i in particle_names])
        for i in range(N_types):
            color_array[type_ids == i] = fresnel.color.linear(mapper.to_rgba(v)[i])

    # Make an array of the radii based on particle name
    # -- if name is not defined in the dictionary, use default
    rad_array = np.empty((N), dtype="float64")
    for i, n in enumerate(particle_names):
        try:
            rad_array[i] = radii_dict[n] * scale
        except KeyError:
            rad_array[i] = radii_dict["default"] * scale

    ## Start building the fresnel scene
    scene = fresnel.Scene()

    # Spheres for every particle in the system
    geometry = fresnel.geometry.Sphere(scene, N=N)
    geometry.position[:] = comp.xyz
    geometry.material = fresnel.material.Material(roughness=1.0)
    geometry.outline_width = 0.01 * scale

    # use color instead of material.color
    geometry.material.primitive_color_mix = 1.0
    geometry.color[:] = color_array

    # resize radii
    geometry.radius[:] = rad_array

    # bonds
    bonds = fresnel.geometry.Cylinder(scene, N=N_bonds)
    bonds.material = fresnel.material.Material(roughness=0.5)
    bonds.outline_width = 0.01 * scale

    # bonds are white
    bond_colors = np.ones((N_bonds, 3), dtype="float64")

    bonds.material.primitive_color_mix = 1.0
    bonds.points[:] = all_bonds

    bonds.color[:] = np.stack(
        [fresnel.color.linear(bond_colors), fresnel.color.linear(bond_colors)], axis=1
    )
    bonds.radius[:] = [0.03 * scale] * N_bonds

    if show_box:
        # Use comp.box, unless it does not exist, then use comp.boundingbox
        try:
            freud_box = mb_to_freud_box(comp.box)
        except AttributeError:
            freud_box = mb_to_freud_box(comp.boundingbox)
        # Create box in fresnel
        fresnel.geometry.Box(scene, freud_box, box_radius=0.008 * scale)
    return scene<jupyter_output><empty_output><jupyter_text>The cell below creates two sample compound classes: atomistic methane and a random coarse-grain compound.<jupyter_code>class Methane(mb.Compound):
    def __init__(self):
        super(Methane, self).__init__()
        carbon = mb.Particle(name="C")
        self.add(carbon, label="C[$]")

        hydrogen = mb.Particle(name="H", pos=[0.1, 0, -0.07])
        self.add(hydrogen, label="HC[$]")

        self.add_bond((self[0], self["HC"][0]))

        self.add(mb.Particle(name="H", pos=[-0.1, 0, -0.07]), label="HC[$]")
        self.add(mb.Particle(name="H", pos=[0, 0.1, 0.07]), label="HC[$]")
        self.add(mb.Particle(name="H", pos=[0, -0.1, 0.07]), label="HC[$]")

        self.add_bond((self[0], self["HC"][1]))
        self.add_bond((self[0], self["HC"][2]))
        self.add_bond((self[0], self["HC"][3]))


class CG(mb.Compound):
    def __init__(self):
        super(CG, self).__init__()
        np.random.seed(42)
        for i in range(ord("A"), ord("Z") + 1):
            random.seed(i)
            N = random.randint(1, 4)
            for n in range(N):
                self.add(
                    mb.Particle(name=f"_{chr(i)}", pos=np.random.random(3) * 2 - 1)
                )

        for i in range(self.n_particles - 1):
            for j in range(i + 1, self.n_particles):
                dist = distance(self.xyz[i], self.xyz[j])
                if dist < 0.3:
                    self.add_bond((self[i], self[j]))<jupyter_output><empty_output><jupyter_text>Create an instance of a methane `mbuild.Compound` visualize it in a `fresnel.Scene`. By default `visualize` uses the [CPK color scheme](https://en.wikipedia.org/wiki/CPK_coloringTypical_assignments) and does not render the box. We can then use [interactive scene view](https://fresnel.readthedocs.io/en/stable/examples/02-Advanced-topics/03-Interactive-scene-view.html) to view the scene in a pop-up window, using the mouse to rotate and zoom.<jupyter_code>methane = Methane()

#methane.visualize().show();

methane.box = mb.Box(lengths=[0.5,0.5,0.5])

scene = visualize(methane, show_box=True)

view = fresnel.interact.SceneView(scene)
view<jupyter_output><empty_output><jupyter_text>After rotating the scene in the view window, the `scene.camera` will retain its position. We can then use [Python Image Library (PIL)](https://pillow.readthedocs.io/en/stable/) to save the rendering of the scene to an image file.<jupyter_code>output = fresnel.pathtrace(scene, light_samples=40, w=600, h=600)
filename = "methane.png"

image = PIL.Image.fromarray(output[:], mode='RGBA')
image.save(filename, dpi=(300, 300))

old_camera = scene.camera
print(f"Rendered {filename} with {scene.camera}")<jupyter_output><empty_output><jupyter_text>By printing out or saving the camera position, we can easily recreate this image.<jupyter_code>scene = visualize(methane, show_box=True)

print(f"Camera on initialization: {scene.camera}")
scene.camera = old_camera
print(f"Camera after reverting to previous position: {scene.camera}")
fresnel.pathtrace(scene, light_samples=40, w=600, h=600)<jupyter_output><empty_output><jupyter_text>If your system is not atomistic, the default cpk color scheme will not work well. Even though `cg` contains 26 different particle types, they are all rendered as the default pink color:<jupyter_code>cg = CG()

scene = visualize(cg, scale=2)
fresnel.pathtrace(scene, w=300, h=300, light_samples=10)

particle_types = list(set([i.name for i in cg.particles()]))
print(particle_types)<jupyter_output><empty_output><jupyter_text>The visualize function can also take in the name of any [matplotlib colormap](https://matplotlib.org/3.1.0/tutorials/colors/colormaps.html). The function will evenly distribute the color mapping across the different particle types, so that each type will have a unique color (as long as the colormap does not repeat).<jupyter_code>cg = CG()

#with warnings.catch_warnings():
#    warnings.simplefilter("ignore")
#    cg.visualize().show();

scene = visualize(cg, color="gist_rainbow", scale=2)
fresnel.pathtrace(scene, w=300, h=300, light_samples=10)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14639191052550943, 'fraction_numerical': 0.03692354833767428, 'mean_word_length': 3.208575112830432, 'pattern_counts': {'":': 34, '<': 34, '<?xml version=': 0, '>': 33, 'https://': 7, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '8', 'dir': 'jupyter-structured-clean-dedup', 'id': '57633', 'license': 'BSD-3-Clause', 'n_tokens_mistral': 4888, 'n_tokens_neox': 4494, 'n_words': 1121, 'path': '02-Advanced-topics/06-Compound-visualization.ipynb', 'repo_name': 'jennyfothergill/fresnel-examples', 'size': '13054'}