plot.py

aimsprop.plot

Functions

plot_population(filename: str, bundle: Bundle, bundles: list, time_units: str = 'au', legend_loc = 'right', state_colors: list = None, plot_total: bool = True, clf: bool = True, tmax: float = None) -> Figure

Plot the AIMS state populations.

Parameters:

Name	Type	Description	Default
filename	str	the output PDF file	required
bundle	Bundle	the Bundle to plot populations for	required
bundles	list	a list of Bundle objects, one for each IC. This is used for the "spaghetti" plots.	required
time_units	str	"au" or "fs"	'au'
legend_loc	legend location	location indicator for legend	'right'
state_colors	list	list of colors to use for state plotting. None defaults to interpolation on jet colormap. For two or three states, many people prefer ['r', 'b', 'g'] or similar.	None
plot_total	bool	Plot total population?	True
clf	bool	clear plot or not?	True

Returns:

Type	Description
Figure	plt handle for further modification saves figure to filename

Source code in aimsprop/plot.py

def plot_population(
    filename: str,
    bundle: Bundle,
    bundles: list,
    time_units: str = "au",
    legend_loc="right",
    state_colors: list = None,
    plot_total: bool = True,
    clf: bool = True,
    tmax: float = None,
) -> Figure:
    """Plot the AIMS state populations.

    Params:
        filename: the output PDF file
        bundle: the Bundle to plot populations for
        bundles: a list of Bundle objects, one for each
            IC. This is used for the "spaghetti" plots.
        time_units: "au" or "fs"
        legend_loc (legend location): location indicator for legend
        state_colors (list of colors): list of colors to use for state
            plotting. None defaults to interpolation on jet colormap. For two
            or three states, many people prefer ['r', 'b', 'g'] or similar.
        plot_total: Plot total population?
        clf: clear plot or not?
    Returns:
         plt handle for further modification
         saves figure to filename
    """

    if time_units == "au":
        time_scale = 1.0
    elif time_units == "fs":
        time_scale = 1.0 / 41.3413745758  # TODO: standardize this
    else:
        raise ValueError("Unknown time units: %s" % time_units)

    if state_colors:
        colors = state_colors
    else:
        cmap = matplotlib.cm.get_cmap("jet")
        colors = [
            cmap(float(x) / (len(bundle.Is) - 1))
            for x in reversed(list(range(len(bundle.Is))))
        ]

    if clf:
        plt.clf()

    # Spaghetti plots from bundles
    Imap = {I: Iind for Iind, I in enumerate(bundle.Is)}
    for bundle2 in bundles:
        ts = np.array(bundle2.ts)
        for I in bundle2.Is:
            bundle3 = bundle2.subset_by_I(I)
            t3s = bundle3.ts
            w3s = [sum(frame.w for frame in bundle3.subset_by_t(t).frames) for t in t3s]
            plt.plot(time_scale * t3s, w3s, "-", color=colors[Imap[I]], linewidth=1.0)

    # State populations
    ts = np.array(bundle.ts)
    populations = pop.compute_population(bundle)
    for Iind, I in enumerate(sorted(populations.keys())):
        plt.plot(
            time_scale * ts,
            populations[I],
            "-",
            color=colors[Iind],
            linewidth=2.0,
            label="State=%d" % I,
        )

    # Total population (to check norm, state cutoffs, etc)
    if plot_total:
        total = np.zeros_like(ts)
        for population in list(populations.values()):
            total += population
        plt.plot(time_scale * ts, total, "-", color="k", linewidth=2.0, label="Total")

    if tmax is None:
        tmax = max(ts) * time_scale

    plt.xlabel("t [%s]" % time_units)
    plt.ylabel("Population [-]")
    plt.axis([time_scale * min(ts), tmax, -0.1, 1.1])
    plt.legend(loc=legend_loc)
    plt.tight_layout()
    plt.savefig(filename)

    return plt

plot_scalar(filename: str, bundle: Bundle, key: str, ylabel: str = None, time_units: str = 'au', legend_loc = 1, state_colors: list = None, plot_average: bool = True, clf: bool = True) -> Figure

Plot an AIMS scalar property (e.g., a "spaghetti plot" for a bond distance).

Parameters:

Name	Type	Description	Default
filename	str	the output PDF file	required
bundle	Bundle	the Bundle to plot properties for	required
key	str	the key of the property	required
ylabel	str	the label of the y-axis (defaults to key)	None
time_units	str	"au" or "fs"	'au'
legend_loc	legend location	location indicator for legend	1
state_colors	list	list of colors to use for state plotting. None defaults to interpolation on jet colormap. For two or three states, many people prefer ['r', 'b', 'g'] or similar.	None
plot_average	bool	Plot averages for each state and total?	True
clf	bool	clear plot or not?	True

Returns:

Type	Description
Figure	plt handle for further modification saves figure to filename

Source code in aimsprop/plot.py

def plot_scalar(
    filename: str,
    bundle: Bundle,
    key: str,
    ylabel: str = None,
    time_units: str = "au",
    legend_loc=1,
    state_colors: list = None,
    plot_average: bool = True,
    clf: bool = True,
) -> Figure:
    """Plot an AIMS scalar property (e.g., a "spaghetti plot" for a bond distance).

    Params:
        filename: the output PDF file
        bundle: the Bundle to plot properties for
        key: the key of the property
        ylabel: the label of the y-axis (defaults to key)
        time_units: "au" or "fs"
        legend_loc (legend location): location indicator for legend
        state_colors: list of colors to use for state
            plotting. None defaults to interpolation on jet colormap. For two
            or three states, many people prefer ['r', 'b', 'g'] or similar.
        plot_average: Plot averages for each state and total?
        clf : clear plot or not?
    Returns:
        plt handle for further modification
        saves figure to filename
    """

    if time_units == "au":
        time_scale = 1.0
    elif time_units == "fs":
        time_scale = 1.0 / 41.3413745758  # TODO: standardize this
    else:
        raise ValueError("Unknown time units: %s" % time_units)

    if state_colors:
        colors = state_colors
    else:
        cmap = matplotlib.cm.get_cmap("jet")
        colors = [
            cmap(float(x) / (len(bundle.Is) - 1))
            for x in reversed(list(range(len(bundle.Is))))
        ]

    if clf:
        plt.clf()
    if plot_average:
        # Plot average
        plt.plot(
            time_scale * np.array(bundle.ts),
            bundle.extract_property(key),
            "-k",
            linewidth=3.0,
        )
        # Plot state averages
        for Iind, I in enumerate(bundle.Is):
            bundle2 = bundle.subset_by_I(I)
            color = colors[Iind]
            plt.plot(
                time_scale * np.array(bundle2.ts),
                bundle2.extract_property(key),
                "-",
                color=color,
                linewidth=2.0,
            )
    # Plot individual frames
    for Iind, I in enumerate(bundle.Is):
        bundle2 = bundle.subset_by_I(I)
        color = colors[Iind]
        for lind, label in enumerate(bundle2.labels):
            bundle3 = bundle2.subset_by_label(label)
            plt.plot(
                time_scale * np.array(bundle3.ts),
                bundle3.extract_property(key),
                "-",
                color=color,
                linewidth=1.0,
                label=("State=%d" % I if lind == 0 else None),
            )

    plt.xlabel("t [%s]" % time_units)
    plt.ylabel(ylabel if ylabel else key)
    plt.axis("tight")  # TODO: Does not seem to respect this
    plt.legend(loc=legend_loc)
    plt.savefig(filename)
    return plt

plot_vector(filename: str, bundle: Bundle, key: str, y: ndarray, ylabel: str = None, time_units: str = 'au', diff: bool = False, cmap = <matplotlib.colors.LinearSegmentedColormap object at 0x116765610>, levels: ndarray = None, nlevel: int = 65, twosided: bool = True, clf: bool = True) -> Figure

Plot an AIMS vector property (e.g., a heatmap of a UED or PES signal).

Parameters:

Name	Type	Description	Default
filename	str	the output PDF file	required
bundle	Bundle	the Bundle to plot properties for	required
key	str	the key of the property	required
y	ndarray	the indices of the y axis (e.g., R or Q or something like it)	required
ylabel	str	the label of the y-axis (defaults to key)	None
time_units	str	"au" or "fs"	'au'
diff	bool	is this a difference property from t=0?	False
cmap		the desired colormap	<matplotlib.colors.LinearSegmentedColormap object at 0x116765610>
levels	ndarray	the explicitly desired contour levels (1st priority).	None
nlevel	int	number of evenly spaced contour levels to saturate data (2nd priority).	65
twosided	bool	is the colormap two-sided (used only with nlevel)	True
clf	bool	clear plot or not?	True

Returns:

Type	Description
Figure	plt handle for further modification saves figure to filename

Source code in aimsprop/plot.py

def plot_vector(
    filename: str,
    bundle: Bundle,
    key: str,
    y: np.ndarray,
    ylabel: str = None,
    time_units: str = "au",
    diff: bool = False,
    cmap=plt.cm.bwr,
    levels: np.ndarray = None,
    nlevel: int = 65,
    twosided: bool = True,
    clf: bool = True,
) -> Figure:
    """Plot an AIMS vector property (e.g., a heatmap of a UED or PES signal).

    Params:
        filename: the output PDF file
        bundle: the Bundle to plot properties for
        key: the key of the property
        y: the indices of the y axis (e.g., R or Q or something
            like it)
        ylabel: the label of the y-axis (defaults to key)
        time_units: "au" or "fs"
        diff: is this a difference property from t=0?
        cmap: the desired colormap
        levels: the explicitly desired contour levels (1st
            priority).
        nlevel: number of evenly spaced contour levels to
            saturate data (2nd priority).
        twosided: is the colormap two-sided (used only with nlevel)
        clf: clear plot or not?
    Returns:
        plt handle for further modification
        saves figure to filename
    """

    if time_units == "au":
        time_scale = 1.0
    elif time_units == "fs":
        time_scale = 1.0 / 41.3413745758  # TODO: standardize this
    else:
        raise ValueError("Unknown time units: %s" % time_units)

    # Extract data to plot
    ts = np.array(bundle.ts)
    ys = y
    Vs = bundle.extract_property(key)

    # Difference properties
    if diff:
        Vs -= np.outer(np.ones_like(ts), Vs[0, :])

    # Levels and color ticks
    if levels is None:
        vmax = np.max(np.abs(Vs))
        if twosided:
            levels = np.linspace(-vmax, +vmax, nlevel)
            cticks = [-int(vmax), 0, +int(vmax)]
        else:
            levels = np.linspace(0, +vmax, nlevel)
            cticks = [0, +int(vmax)]
    else:
        vmax = np.max(levels)
        if twosided:
            cticks = [-int(vmax), 0, +int(vmax)]
        else:
            cticks = [0, +int(vmax)]

    if clf:
        plt.clf()
    hs = plt.contourf(
        time_scale * ts, ys, Vs.T, levels=levels, cmap=cmap, extend="both"
    )
    for h in hs.collections:
        h.set_edgecolor("face")
    plt.colorbar(ticks=cticks)
    plt.xlabel("t [%s]" % time_units)
    plt.ylabel(ylabel if ylabel else key)
    plt.axis("tight")  # TODO: Does not seem to respect this
    plt.savefig(filename)
    return plt