ternary_plot.py

from typing import List

import matplotlib.pyplot as plt
import mpltern
import pandas as pd
from spt.plot_params import set_roman_plot_params


def plot_ternary(
    element_list: List[str],
    max_natom: int = 8,
):
    """单个子三元系绘制"""
    # TO DO: 添加并设置 colorbar

    system_name = "".join(element_list)
    csv_fn = f"{system_name}_1_{max_natom}.csv"
    df = pd.read_csv(csv_fn)
    columns = df.columns

    mixing_energy_pa_df = df["mixing_energy_pa"]
    t0 = df.iloc[:, -3]
    l0 = df.iloc[:, -2]
    r0 = df.iloc[:, -1]

    set_roman_plot_params()
    ax = plt.subplot(projection="ternary")

    ax.scatter(t0, l0, r0, c=mixing_energy_pa_df)

    ax.set_tlabel(f"{columns[-3]}")
    ax.set_llabel(f"{columns[-2]}")
    ax.set_rlabel(f"{columns[-1]}")

    plt.savefig(f"{system_name}_ternary.png")

    print(f"\n{system_name} figure is generated!")


if __name__ == "__main__":

    element_list = ["Al", "Mo", "Zr"]

    plot_ternary(element_list=element_list)