def draw_orbitals(traj, align=True, **kwargs):
    """ Visualize trajectory with molecular orbitals

    Args:
        align (bool): Align orbital phases (i.e., multiplying by -1 as needed) to prevent sign
           flips between frames

    Returns:
        TrajectoryOrbViewer: create a trajectory visualization
    """
    from ..viewers.trajectory_viewer import TrajectoryOrbViewer
    for frame in traj:
        if 'wfn' not in frame:
            raise ValueError(
                "Can't draw orbitals - orbital information missing in at least one frame. It must be calculated with a QM method."
                )
    if align:
        traj.align_orbital_phases()
    traj._viz = TrajectoryOrbViewer(traj, **kwargs)
    return traj._viz
