Configurable script to plot compositional phase diagrams

phase_diagram.py

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.patches import Polygon, FancyArrowPatch
from typing import List, Tuple, Optional

class Structure:
    def __init__(self, name: str, composition: List[float], color: str = 'black', 
                 show_label: bool = True, label_font_size: Optional[int] = None):
        self.name = name
        self.composition = composition
        self.color = color
        self.show_label = show_label
        self.label_font_size = label_font_size

    @property
    def is_corner(self) -> bool:
        return any(x == 1 for x in self.composition)

class CompositionTriangle:
    def __init__(self, structures: List[Structure], figsize: Tuple[int, int] = (10, 8), 
                 font_size: int = 12, label_distance: float = 1.15):
        self.fig, self.ax = plt.subplots(figsize=figsize)
        self.fig.patch.set_facecolor('white')
        self.default_font_size = font_size
        plt.rcParams.update({'font.size': font_size})
        self.corners = np.array([[0, 0], [1, 0], [0.5, np.sqrt(3)/2]])
        self.structures = {}
        self.show_labels = True
        self.label_distance = label_distance
        self.show_grid = False
        self.grid_density = 10
        
        self._setup_triangle()
        for structure in structures:
            self.add_structure(structure)

    def _setup_triangle(self):
        self.ax.plot([0, 1, 0.5, 0], [0, 0, np.sqrt(3)/2, 0], 'k-', linewidth=2)
        extended_ylim = self.label_distance * np.sqrt(3)/2
        self.ax.set_xlim(-0.2, 1.2)
        self.ax.set_ylim(-0.2, extended_ylim + 0.1)
        self.ax.set_aspect('equal', adjustable='box')
        self.ax.axis('off')

    def add_structure(self, structure: Structure, marker: str = 'o', size: Optional[int] = None):
        x, y = self._barycentric_to_cartesian(structure.composition)
        self.structures[structure] = (x, y)
        
        if size is None:
            size = self._calculate_dynamic_size(x, y)
        
        self.ax.scatter(x, y, marker=marker, c=structure.color, s=size, zorder=5)
        
        if self.show_labels and structure.show_label:
            self._add_label(structure, x, y)

    def _add_label(self, structure: Structure, x: float, y: float):
        label_pos, ha, va = self._calculate_label_position(x, y, structure.composition)
        font_size = structure.label_font_size or self.default_font_size
        text = self.ax.annotate(structure.name, label_pos, ha=ha, va=va, zorder=6, fontsize=font_size)
        return text

    def _calculate_label_position(self, x: float, y: float, composition: List[float]) -> Tuple[Tuple[float, float], str, str]:
        center = np.array([0.5, np.sqrt(3)/6])
        point = np.array([x, y])
        vec = point - center
        
        edge_proximity = min(composition)
        adjusted_extension = self.label_distance * (1 + edge_proximity * 0.1)
        
        extended_point = center + adjusted_extension * vec
        
        if composition[2] > 0.8:
            return (extended_point[0], extended_point[1] + 0.02), 'center', 'bottom'
        elif y < 0.1:
            return (extended_point[0], extended_point[1] - 0.02), 'center', 'top'
        elif composition[0] > composition[1]:
            return (extended_point[0] - 0.02, extended_point[1]), 'right', 'center'
        else:
            return (extended_point[0] + 0.02, extended_point[1]), 'left', 'center'

    def _barycentric_to_cartesian(self, coords: List[float]) -> Tuple[float, float]:
        a, b, c = coords
        x = b + 0.5 * c
        y = np.sqrt(3)/2 * c
        return x, y

    def _calculate_dynamic_size(self, x: float, y: float) -> int:
        base_size = 50
        densities = [len([s for s in self.structures.values() if abs(sx-x) < 0.1 and abs(sy-y) < 0.1]) 
                     for sx, sy in self.structures.values()]
        local_density = max(densities) if densities else 1
        return max(20, base_size // local_density)

    def add_composition_point(self, composition: List[float], name: str, color: str = 'black', 
                              marker: str = 'o', size: Optional[int] = None, label_font_size: Optional[int] = None):
        structure = Structure(name, composition, color, label_font_size=label_font_size)
        self.add_structure(structure, marker, size)
        return structure  # Return the created structure for potential further use

    def add_line(self, start_structure: Structure, end_structure: Structure, style: str = '-', color: str = 'black', zorder: int = 3):
        start_x, start_y = self.structures[start_structure]
        end_x, end_y = self.structures[end_structure]
        self.ax.plot([start_x, end_x], [start_y, end_y], linestyle=style, color=color, zorder=zorder)

    def add_arrow(self, side: str, label: str, color: str = 'blue', offset: float = 0.1, 
                  length: float = 0.2, reverse: bool = False, label_font_size: Optional[int] = None, linewidth: float = 1):
        if side not in ['left', 'right', 'bottom']:
            raise ValueError("Side must be 'left', 'right', or 'bottom'")

        if side == 'left':
            mid = np.array([0.25, np.sqrt(3)/4])
            direction = np.array([0.5, np.sqrt(3)/2]) - np.array([0, 0])
            label_offset = np.array([-0.15, 0])
        elif side == 'right':
            mid = np.array([0.75, np.sqrt(3)/4])
            direction = np.array([0.5, np.sqrt(3)/2]) - np.array([1, 0])
            label_offset = np.array([0.15, 0])
        else:  # bottom
            mid = np.array([0.5, 0])
            direction = np.array([1, 0]) - np.array([0, 0])
            label_offset = np.array([0, -0.1])

        direction = direction / np.linalg.norm(direction)
        perp = np.array([-direction[1], direction[0]])
        mid = mid + perp * offset
        start = mid - direction * length / 2
        end = mid + direction * length / 2

        if reverse:
            start, end = end, start

        arrow = FancyArrowPatch(start, end, color=color, arrowstyle='->', mutation_scale=20, linewidth=linewidth, zorder=7)
        self.ax.add_patch(arrow)

        if self.show_labels:
            label_pos = mid + label_offset
            font_size = label_font_size or self.default_font_size
            self.ax.text(label_pos[0], label_pos[1], label, ha='center', va='center', color=color, zorder=8, fontsize=font_size)

    def highlight_region(self, structures: List[Structure], color: str = 'yellow', alpha: float = 0.2, zorder: int = 2):
        coords = [self.structures[s] for s in structures]
        poly = Polygon(coords, closed=True, facecolor=color, alpha=alpha, edgecolor=None, zorder=zorder)
        self.ax.add_patch(poly)

    def set_label_visibility(self, visible: bool):
        self.show_labels = visible
        self.ax.texts.clear()
        if self.show_labels:
            for structure, (x, y) in self.structures.items():
                if structure.show_label:
                    self._add_label(structure, x, y)

    def set_label_distance(self, distance: float):
        """Set the distance of labels from the triangle."""
        self.label_distance = distance
        self.set_label_visibility(self.show_labels)  # Refresh labels

    def toggle_composition_grid(self, show: bool = True, density: int = 10):
        self.show_grid = show
        self.grid_density = density
        if self.show_grid:
            self._draw_composition_grid()
        else:
            self._clear_composition_grid()

    def _draw_composition_grid(self):
        for i in range(1, self.grid_density):
            alpha = i / self.grid_density
            # Draw lines parallel to the bottom edge
            start = self._barycentric_to_cartesian([1-alpha, 0, alpha])
            end = self._barycentric_to_cartesian([0, 1-alpha, alpha])
            self.ax.plot([start[0], end[0]], [start[1], end[1]], 'k-', alpha=0.2, linewidth=0.5)
            
            # Draw lines parallel to the left edge
            start = self._barycentric_to_cartesian([1-alpha, alpha, 0])
            end = self._barycentric_to_cartesian([0, alpha, 1-alpha])
            self.ax.plot([start[0], end[0]], [start[1], end[1]], 'k-', alpha=0.2, linewidth=0.5)
            
            # Draw lines parallel to the right edge
            start = self._barycentric_to_cartesian([alpha, 1-alpha, 0])
            end = self._barycentric_to_cartesian([alpha, 0, 1-alpha])
            self.ax.plot([start[0], end[0]], [start[1], end[1]], 'k-', alpha=0.2, linewidth=0.5)

    def _clear_composition_grid(self):
        for line in self.ax.lines[:]:
            if line.get_alpha() == 0.2 and line.get_linewidth() == 0.5:
                line.remove()

    def add_title(self, title: str, font_size: Optional[int] = None):
        font_size = font_size or self.default_font_size
        self.ax.set_title(title, fontsize=font_size)

    def show(self):
        plt.tight_layout()
        plt.show()

    def save(self, filename: str, dpi: int = 300):
        plt.tight_layout()
        plt.savefig(filename, dpi=dpi, bbox_inches='tight')

# Example usage
if __name__ == "__main__":
    # Define structures
    Li7PS6 = Structure('Li₇PS₆', [0, 0, 1], color='red', label_font_size=14)
    Li5PS4Cl2 = Structure('Li₅PS₄Cl₂', [1, 0, 0], color='green', label_font_size=14)
    Li5PS4Br2 = Structure('Li₅PS₄Br₂', [0, 1, 0], color='blue', label_font_size=14)
    Li6PS5Cl = Structure('Li₆PS₅Cl', [0.5, 0, 0.5], color='purple', label_font_size=14)
    Li6PS5Br = Structure('Li₆PS₅Br', [0, 0.5, 0.5], color='orange', label_font_size=14)
    Li5PS4BrCl = Structure('Li₅PS₄BrCl', [0.5, 0.5, 0], color='brown', label_font_size=14)

    structures = [Li7PS6, Li5PS4Cl2, Li5PS4Br2, Li6PS5Cl, Li6PS5Br, Li5PS4BrCl]

    # Create the triangle with all structures and closer labels
    triangle = CompositionTriangle(structures, label_distance=1.05)

    # Add a new point inside the triangle
    new_point = triangle.add_composition_point([0.3, 0.3, 0.4], "New Point", color='black', size=60, label_font_size=10)

    # Highlight the entire upper triangle
    triangle.highlight_region([Li7PS6, Li5PS4Cl2, Li5PS4Br2], color='yellow', alpha=0.1)

    # Highlight a region including the new point
    triangle.highlight_region([Li7PS6, new_point, Li5PS4Cl2], color='cyan', alpha=0.1)

    # Add lines with custom colors and styles
    triangle.add_line(Li7PS6, Li6PS5Cl)
    triangle.add_line(Li7PS6, Li6PS5Br)
    triangle.add_line(Li5PS4Cl2, Li5PS4Br2)
    triangle.add_line(Li6PS5Cl, Li5PS4BrCl, style='--')
    triangle.add_line(Li6PS5Br, Li5PS4BrCl, style='--')
    triangle.add_line(Li6PS5Cl, Li6PS5Br, style='--')

    # Add arrows for substitutions
    triangle.add_arrow('left', 'Cl Substitution', color='blue', offset=0.22, length=0.4, reverse=True, label_font_size=12)
    triangle.add_arrow('right', 'Br Substitution', color='blue', offset=-0.22, length=0.4, reverse=True, label_font_size=12)

    # Add an arrow to the new point
    triangle.add_arrow('bottom', 'New Point', color='red', offset=-0.1, length=0.2)

    # Toggle composition grid
    triangle.toggle_composition_grid(True, density=5)

    # Add a title
    triangle.add_title('Li-PS-Cl-Br Composition Triangle', font_size=16)

    # Show the diagram
    triangle.show()

    # Optionally, save the diagram
    # triangle.save('composition_triangle.png')