Source code for scm.conformers.plams.plot

import scm.plams
from scm.conformers.plams.interface import ConformersJob
import matplotlib.pyplot as plt
import numpy as np


def get_main_results(job: ConformersJob, temperature=298, unit="kcal/mol", return_molecules: bool = True):
    molecules = None
    if return_molecules:
        molecules = job.results.get_conformers()
    energies = job.results.get_relative_energies(unit)
    populations = job.results.get_boltzmann_distribution(temperature)

    return molecules, energies, populations


def plot_conformers(job: ConformersJob, indices=None, temperature=298, unit="kcal/mol", lowest=True):
    """
    Function for plotting conformers in a Jupyter notebook

    job: ConformersJob
        Finished ConformersJob

    indices: None, int or list of int
        If None, will plot at most 3 conformers.

        If int, will plot at most the given number of conformers.

        If list of int (zero-based indices), plot those conformers.

    temperature: float
        Temperature for relative population (printed above the figure).

    unit: str
        Unit for relative energies (printed above the figure)

    lowest: bool
        Only used if ``indices`` is an integer. If True, plot the N lowest energy conformers. If False, plot conformers evenly distributed from the most stable to the least stable.

    """
    molecules, energies, populations = get_main_results(job, unit=unit)

    if isinstance(indices, int):
        N_plot = min(indices, len(energies))
        if lowest:
            indices = list(range(N_plot))
        else:
            indices = np.linspace(0, len(energies) - 1, N_plot, dtype=np.int32)
    if indices is None:
        indices = list(range(min(3, len(energies))))

    fig, axes = plt.subplots(1, len(indices), figsize=(12, 3))
    if len(indices) == 1:
        axes = [axes]

    for ax, i in zip(axes, indices):
        mol = molecules[i]
        E = energies[i]
        population = populations[i]

        scm.plams.plot_molecule(mol, ax=ax)
        ax.set_title(f"#{i+1}\nΔE = {E:.2f} {unit}\nPop.: {population:.3f} (T = {temperature} K)")

    return axes