<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 Tutorial 2
Step1:  Parametrization (recap of Tutorial 1)
We'll rerun the parametrization performed in Tutorial 1.
Step2:  Highlighting atoms and CG beads
To link atoms to coarse-grained beads, it's useful to have a look at a few Cg_molecule attributes
Step3:  We'll want to color atoms according to beads, so let's define a color palette
Step4:  rdkit offers convenient rendering functions to highlight particular atoms and bonds. We will color heavy atoms involved in a particular CG bead. First let's focus on atoms
Step5:  Same for bonds. rdkit here comes to the rescue, because it keeps track of which heavy atoms are connected by bonds.
Step6:  We now have all the necessary ingredients to render a molecule with the appropriate atom and bond highlighting. Here we use the rdMolDraw2D from rdkit. We'll save this to a file, and render the file in the notebook
<ASSISTANT_TASK:>
Python Code:
import auto_martini as am
import numpy as np
from rdkit import Chem
from rdkit.Chem.Draw import IPythonConsole
from IPython.display import Image
import rdkit
from rdkit.Chem import Draw
from rdkit.Chem import AllChem
from rdkit.Chem import rdDepictor
from rdkit.Chem.Draw import rdMolDraw2D
print(rdkit.__version__)
Explanation: Tutorial 2: Highlighting
Following the parametrization of phenylmethanimine in Tutorial 1, let's visualize the correspondance between atoms and coarse-grained beads. 
First let's import a number of packages, especially auto_martini, various rdkit submodules, as well as IPython.display.Image to visualize an image in the notebook, and numpy.
End of explanation
smiles = "N=Cc1ccccc1"
mol = Chem.MolFromSmiles(smiles)
Chem.AddHs(mol)
AllChem.EmbedMolecule(mol)
mol
# Load the molecule in auto_martini and coarse-grain it
mol_am, _ = am.topology.gen_molecule_smi(smiles)
cg = am.solver.Cg_molecule(mol_am, "PHM")
Explanation: Parametrization (recap of Tutorial 1)
We'll rerun the parametrization performed in Tutorial 1.
End of explanation
print("list_heavyatom_names:",cg.list_heavyatom_names)
print("cg_bead_names:       ",cg.cg_bead_names)
print("atom_partitioning:   ",cg.atom_partitioning)
Explanation: Highlighting atoms and CG beads
To link atoms to coarse-grained beads, it's useful to have a look at a few Cg_molecule attributes:
- cg.list_heavyatom_names: lists the chemical elements of each heavy atom
- cg.cg_bead_names: lists the CG bead names
- cg.atom_partitioning: dictionary relating atom ID to CG bead ID.
We'll mostly make use of the last one, here
End of explanation
colors = [(0.121, 0.466, 0.705), (1.0, 0.498, 0.0549), (0.172, 0.627, 0.172)]
Explanation: We'll want to color atoms according to beads, so let's define a color palette
End of explanation
hit_ats = list(cg.atom_partitioning.keys())
atom_cols = {}
for i, at in enumerate(hit_ats):
    atom_cols[at] = colors[cg.atom_partitioning[i]%4]
Explanation: rdkit offers convenient rendering functions to highlight particular atoms and bonds. We will color heavy atoms involved in a particular CG bead. First let's focus on atoms: for each atom we want the CG bead ID it maps to. This is simply obtained from the keys() of the atom_partinioning dictionary. As for the colors, we can use the colors array defined above.
End of explanation
hit_bonds = []
bond_cols = {}
for i, at in enumerate(hit_ats):
    for j, att in enumerate(hit_ats):
        if i > j and atom_cols[i] == atom_cols[j] and mol.GetBondBetweenAtoms(i,j):
            b_idx = mol.GetBondBetweenAtoms(i,j).GetIdx()
            hit_bonds.append(b_idx)
            bond_cols[b_idx] = colors[cg.atom_partitioning[i]%4]
Explanation: Same for bonds. rdkit here comes to the rescue, because it keeps track of which heavy atoms are connected by bonds.
End of explanation
d = rdMolDraw2D.MolDraw2DCairo(400,400)
rdMolDraw2D.PrepareAndDrawMolecule(d, mol, 
                                   highlightAtoms=hit_ats,
                                  highlightAtomColors=atom_cols,
                                  highlightBonds=hit_bonds,
                                  highlightBondColors=bond_cols)
d.DrawMolecule(mol)
d.FinishDrawing()
d.WriteDrawingText('phm_highlight.png')
Image('phm_highlight.png')
Explanation: We now have all the necessary ingredients to render a molecule with the appropriate atom and bond highlighting. Here we use the rdMolDraw2D from rdkit. We'll save this to a file, and render the file in the notebook:
End of explanation
<END_TASK>