synkit 0.0.1__py3-none-any.whl

synkit/Vis/graph_visualizer.py

@@ -0,0 +1,286 @@
+"""
+This module comprises several functions adapted from the work of Klaus Weinbauer.
+The original code can be found at his GitHub repository: https://github.com/klausweinbauer/FGUtils.
+Adaptations were made to enhance functionality and integrate with other system components.
+"""
+
+import networkx as nx
+from rdkit import Chem
+from rdkit.Chem import rdDepictor
+
+import matplotlib.pyplot as plt
+from typing import Dict, Optional
+
+from synkit.IO.graph_to_mol import GraphToMol
+
+
+class GraphVisualizer:
+    def __init__(
+        self,
+        node_attributes: Dict[str, str] = {
+            "element": "element",
+            "charge": "charge",
+            "atom_map": "atom_map",
+        },
+        edge_attributes: Dict[str, str] = {"order": "order"},
+    ):
+        self.node_attributes = node_attributes
+        self.edge_attributes = edge_attributes
+
+    def _get_its_as_mol(self, its: nx.Graph) -> Optional[Chem.Mol]:
+        """
+        Convert a graph representation of an intermediate transition state into an RDKit molecule.
+
+        Parameters:
+        - its (nx.Graph): The graph to convert.
+
+        Returns:
+        - Chem.Mol or None: The RDKit molecule if conversion is successful, None otherwise.
+        """
+        _its = its.copy()
+        for n in _its.nodes():
+            _its.nodes[n]["atom_map"] = n  #
+        for u, v in _its.edges():
+            _its[u][v]["order"] = 1
+        return GraphToMol(self.node_attributes, self.edge_attributes).graph_to_mol(
+            _its, False, False
+        )
+
+    def plot_its(
+        self,
+        its: nx.Graph,
+        ax: plt.Axes,
+        use_mol_coords: bool = True,
+        title: Optional[str] = None,
+        node_color: str = "#FFFFFF",
+        node_size: int = 500,
+        edge_color: str = "#000000",
+        edge_weight: float = 2.0,
+        show_atom_map: bool = False,
+        use_edge_color: bool = False,
+        symbol_key: str = "element",
+        bond_key: str = "order",
+        aam_key: str = "atom_map",
+        standard_order_key: str = "standard_order",
+        font_size: int = 12,
+        rule: bool = False,  # New option to remove edges with specific colors
+    ):
+        """
+        Plot an intermediate transition state (ITS) graph on a given Matplotlib axes with various customizations.
+
+        Parameters:
+        - its (nx.Graph): The graph representing the intermediate transition state.
+        - ax (plt.Axes): The matplotlib axes to draw the graph on.
+        - use_mol_coords (bool): Use molecular coordinates for node positions if True, else use a spring layout.
+        - title (Optional[str]): Title for the graph. If None, no title is set.
+        - node_color (str): Color code for the graph nodes.
+        - node_size (int): Size of the graph nodes.
+        - edge_color (str): Default color code for the graph edges if not using conditional coloring.
+        - edge_weight (float): Thickness of the graph edges.
+        - show_atom_map (bool): If True, displays atom mapping numbers alongside symbols.
+        - use_edge_color (bool): If True, colors edges based on their 'standard_order' attribute.
+        - symbol_key (str): Key to access the symbol attribute in the node's data.
+        - bond_key (str): Key to access the bond type attribute in the edge's data.
+        - aam_key (str): Key to access the atom mapping number in the node's data.
+        - standard_order_key (str): Key to determine the edge color conditionally.
+        - font_size (int): Font size for labels and edge labels.
+        - rule (bool): If True, removes edges with a specific color before plotting.
+
+        Returns:
+        - None
+        """
+        bond_char = {None: "∅", 0: "∅", 1: "—", 2: "=", 3: "≡", 1.5: ":"}
+
+        positions = self._calculate_positions(its, use_mol_coords)
+
+        ax.axis("equal")
+        ax.axis("off")
+        if title:
+            ax.set_title(title)
+
+        # Conditional edge coloring based on 'standard_order'
+        if use_edge_color:
+            edge_colors = [
+                (
+                    "red"
+                    if data.get(standard_order_key, 0) > 0
+                    else "green" if data.get(standard_order_key, 0) < 0 else "black"
+                )
+                for _, _, data in its.edges(data=True)
+            ]
+        else:
+            edge_colors = edge_color
+
+        # If rule=True, remove edges with specific colors (red/green/black)
+        if rule:
+            # Get the edges that have the colors red, green, or black
+            edges_to_remove = [
+                edge
+                for edge, color in zip(its.edges(), edge_colors)
+                if color in ["red", "green", "black"]
+            ]
+            its.remove_edges_from(edges_to_remove)
+
+            # Recalculate edge_colors after removal of edges
+            if use_edge_color:
+                edge_colors = [
+                    (
+                        "red"
+                        if data.get(standard_order_key, 0) > 0
+                        else "green" if data.get(standard_order_key, 0) < 0 else "black"
+                    )
+                    for _, _, data in its.edges(data=True)
+                ]
+            else:
+                edge_colors = edge_color
+
+        # Plot the remaining graph
+        nx.draw_networkx_edges(
+            its, positions, edge_color=edge_colors, width=edge_weight, ax=ax
+        )
+        nx.draw_networkx_nodes(
+            its, positions, node_color=node_color, node_size=node_size, ax=ax
+        )
+
+        # Adjust labels to optionally show atom mapping numbers
+        labels = {
+            n: (
+                f"{d[symbol_key]} ({d.get(aam_key, '')})"
+                if show_atom_map
+                else f"{d[symbol_key]}"
+            )
+            for n, d in its.nodes(data=True)
+        }
+        edge_labels = self._determine_edge_labels(its, bond_char, bond_key)
+
+        nx.draw_networkx_labels(
+            its, positions, labels=labels, font_size=font_size, ax=ax
+        )
+        nx.draw_networkx_edge_labels(
+            its, positions, edge_labels=edge_labels, font_size=font_size, ax=ax
+        )
+
+    def _calculate_positions(self, its: nx.Graph, use_mol_coords: bool) -> dict:
+        if use_mol_coords:
+            mol = self._get_its_as_mol(its)
+            positions = {}
+            rdDepictor.Compute2DCoords(mol)
+            for i, atom in enumerate(mol.GetAtoms()):
+                aam = atom.GetAtomMapNum()
+                apos = mol.GetConformer().GetAtomPosition(i)
+                positions[aam] = [apos.x, apos.y]
+        else:
+            positions = nx.spring_layout(its)
+        return positions
+
+    def _determine_edge_labels(
+        self, its: nx.Graph, bond_char: dict, bond_key: str
+    ) -> dict:
+        edge_labels = {}
+        for u, v, data in its.edges(data=True):
+            bond_codes = data.get(bond_key, (0, 0))
+            bc1, bc2 = bond_char.get(bond_codes[0], "∅"), bond_char.get(
+                bond_codes[1], "∅"
+            )
+            if bc1 != bc2:
+                edge_labels[(u, v)] = f"({bc1},{bc2})"
+        return edge_labels
+
+    def plot_as_mol(
+        self,
+        g: nx.Graph,
+        ax: plt.Axes,
+        use_mol_coords: bool = True,
+        node_color: str = "#FFFFFF",
+        node_size: int = 500,
+        edge_color: str = "#000000",
+        edge_width: float = 2.0,
+        label_color: str = "#000000",
+        font_size: int = 12,
+        show_atom_map: bool = False,
+        bond_char: Dict[Optional[int], str] = None,
+        symbol_key: str = "element",
+        bond_key: str = "order",
+        aam_key: str = "atom_map",
+    ) -> None:
+        """
+        Plots a molecular graph on a given Matplotlib axes using either molecular coordinates
+        or a networkx layout.
+
+        Parameters:
+        - g (nx.Graph): The molecular graph to be plotted.
+        - ax (plt.Axes): Matplotlib axes where the graph will be plotted.
+        - use_mol_coords (bool, optional): Use molecular coordinates if True, else use networkx layout.
+        - node_color (str, optional): Color code for the nodes.
+        - node_size (int, optional): Size of the nodes.
+        - edge_color (str, optional): Color code for the edges.
+        - label_color (str, optional): Color for node labels.
+        - font_size (int, optional): Font size for labels.
+        - bond_char (Dict[Optional[int], str], optional): Dictionary mapping bond types to characters.
+        - symbol_key (str, optional): Node attribute key for element symbols.
+        - bond_key (str, optional): Edge attribute key for bond types.
+
+        Returns:
+        - None
+        """
+
+        # Set default bond characters if not provided
+        if bond_char is None:
+            bond_char = {None: "∅", 1: "—", 2: "=", 3: "≡", 1.5: ":"}
+
+        # Determine positions based on use_mol_coords flag
+        if use_mol_coords:
+            mol = GraphToMol(self.node_attributes, self.edge_attributes).graph_to_mol(
+                g, False
+            )  # This function needs to be defined or imported
+            positions = {}
+            rdDepictor.Compute2DCoords(mol)
+            for atom in mol.GetAtoms():
+                aidx = atom.GetIdx()
+                atom_map = atom.GetAtomMapNum()
+                apos = mol.GetConformer().GetAtomPosition(aidx)
+                positions[atom_map] = [apos.x, apos.y]
+        else:
+            positions = nx.spring_layout(g)  # Optionally provide a layout configuration
+
+        ax.axis("equal")
+        ax.axis("off")
+
+        # Drawing elements on the plot
+        nx.draw_networkx_edges(
+            g, positions, edge_color=edge_color, width=edge_width, ax=ax
+        )
+        nx.draw_networkx_nodes(
+            g, positions, node_color=node_color, node_size=node_size, ax=ax
+        )
+
+        # Preparing labels
+        labels = {}
+        for n, d in g.nodes(data=True):
+            charge = d.get("charge", 0)
+            if charge == 0:
+                charge = ""
+            elif charge > 0:
+                charge = f"{charge}+" if charge > 1 else "+"
+            else:
+                charge = f"{-charge}-" if charge < -1 else "-"
+            label = f"{d.get(symbol_key, '')}{charge}"
+            if show_atom_map:
+                label += f" ({d.get(aam_key, '')})"
+            labels[n] = label
+        edge_labels = {
+            (u, v): bond_char.get(d[bond_key], "∅") for u, v, d in g.edges(data=True)
+        }
+
+        # Drawing labels
+        nx.draw_networkx_labels(
+            g,
+            positions,
+            labels=labels,
+            font_color=label_color,
+            font_size=font_size,
+            ax=ax,
+        )
+        nx.draw_networkx_edge_labels(
+            g, positions, edge_labels=edge_labels, font_color=label_color, ax=ax
+        )

synkit/Vis/pdf_writer.py

@@ -0,0 +1,143 @@
+"""
+This module comprises several functions adapted from the work of Klaus Weinbauer.
+The original code can be found at his GitHub repository: https://github.com/klausweinbauer/FGUtils.
+Adaptations were made to enhance functionality and integrate with other system components.
+"""
+
+import matplotlib.pyplot as plt
+from matplotlib.backends.backend_pdf import PdfPages
+from typing import List, Callable, Union, Tuple, Optional
+import tqdm
+
+
+class PdfWriter:
+    """
+    A utility class to create PDF reports with plots from a list of figures or dynamically generated plots.
+
+    Parameters:
+    - file (str): The file name of the output PDF.
+    - plot_fn (Optional[Callable], optional): Function to create a plot for a single data entry or row.
+        Expected interface: `plot_fn(data_entry, axis, **kwargs)`. Default is None.
+    - plot_per_row (bool, optional): If True, calls `plot_fn` for an entire row instead of individual subplots.
+        Default is False.
+    - max_pages (int, optional): Maximum number of pages to create. Default is 999.
+    - rows (int, optional): Number of plot rows per page. Default is 7.
+    - cols (int, optional): Number of plot columns per page. Default is 2.
+    - pagesize (Tuple[float, float], optional): Size of a single page (in inches). Default is (21, 29.7).
+    - width_ratios (Optional[List[float]], optional): Column width ratios. Default is None.
+    - show_progress (bool, optional): If True, displays a progress bar using `tqdm`. Default is True.
+    """
+
+    def __init__(
+        self,
+        file: str,
+        plot_fn: Optional[Callable] = None,
+        plot_per_row: bool = False,
+        max_pages: int = 999,
+        rows: int = 7,
+        cols: int = 2,
+        pagesize: Tuple[float, float] = (21, 29.7),
+        width_ratios: Optional[List[float]] = None,
+        show_progress: bool = True,
+    ):
+        self.pdf_pages = PdfPages(file)
+        self.plot_fn = plot_fn
+        self.plot_per_row = plot_per_row
+        self.max_pages = max_pages
+        self.rows = rows
+        self.cols = cols
+        self.pagesize = pagesize
+        self.width_ratios = width_ratios
+        self.show_progress = show_progress
+
+    def plot(self, data: Union[List[plt.Figure], List], **kwargs):
+        """
+        Generate plots from data or save pre-generated figures to the PDF.
+
+        Parameters:
+        - data (Union[List[matplotlib.figure.Figure], List]): Input data or list of figures.
+          If a list of figures, they are saved directly. Otherwise, the `plot_fn` is called for each data entry.
+        - **kwargs: Additional keyword arguments passed to `plot_fn`.
+
+        Returns:
+        - None
+        """
+        # Case 1: Pre-generated figures
+        if all(isinstance(item, plt.Figure) for item in data):
+            for fig in tqdm.tqdm(
+                data, disable=not self.show_progress, desc="Saving Figures"
+            ):
+                self.save_figure(fig)
+            return
+
+        # Case 2: Generate plots dynamically using `plot_fn`
+        if self.plot_fn is None:
+            raise ValueError(
+                "plot_fn must be provided when input is not a list of figures."
+            )
+
+        if not isinstance(data, list):
+            raise ValueError(
+                "Data must be a list or a list of matplotlib.figure.Figure."
+            )
+
+        plots_per_page = self.rows if self.plot_per_row else self.rows * self.cols
+        max_plots = self.max_pages * plots_per_page
+        step = max(len(data) / max_plots, 1)
+        pages = int((len(data) / step + plots_per_page - 1) // plots_per_page)
+
+        for p in tqdm.tqdm(
+            range(pages), disable=not self.show_progress, desc="Generating Pages"
+        ):
+            fig, ax = plt.subplots(
+                self.rows,
+                self.cols,
+                figsize=self.pagesize,
+                squeeze=False,
+                gridspec_kw=(
+                    {"width_ratios": self.width_ratios} if self.width_ratios else None
+                ),
+            )
+            done = False
+            for r in range(self.rows):
+                if self.plot_per_row:
+                    _idx = int((p * self.rows + r) * step)
+                    if _idx >= len(data):
+                        done = True
+                        break
+                    self.plot_fn(data[_idx], ax[r, :], index=_idx, **kwargs)
+                else:
+                    for c in range(self.cols):
+                        _idx = int((p * plots_per_page + r * self.cols + c) * step)
+                        if _idx >= len(data):
+                            done = True
+                            break
+                        self.plot_fn(data[_idx], ax[r, c], index=_idx, **kwargs)
+            plt.tight_layout()
+            self.pdf_pages.savefig(fig, bbox_inches="tight", pad_inches=1)
+            plt.close(fig)
+            if done:
+                break
+
+    def save_figure(self, figure: plt.Figure):
+        """
+        Save a pre-generated matplotlib figure directly to the PDF.
+
+        Parameters:
+        - figure (matplotlib.figure.Figure): The figure to save.
+
+        Returns:
+        - None
+        """
+        if not isinstance(figure, plt.Figure):
+            raise ValueError("Input must be a matplotlib.figure.Figure.")
+        self.pdf_pages.savefig(figure, bbox_inches="tight", pad_inches=1)
+
+    def close(self):
+        """
+        Close the PDF file, ensuring all pages are written.
+
+        Returns:
+        - None
+        """
+        self.pdf_pages.close()

synkit/Vis/rsmi_to_fig.py

@@ -0,0 +1,169 @@
+import networkx as nx
+import matplotlib.pyplot as plt
+from typing import Union, Tuple
+
+from synkit.Vis.graph_visualizer import GraphVisualizer
+from synkit.IO.chem_converter import rsmi_to_graph
+from synkit.ITS.its_construction import ITSConstruction
+from synkit.IO.gml_to_nx import GMLToNX
+
+vis_graph = GraphVisualizer()
+
+
+def three_graph_vis(
+    input: Union[str, Tuple[nx.Graph, nx.Graph, nx.Graph]],
+    sanitize: bool = False,
+    figsize: Tuple[int, int] = (18, 5),
+    orientation: str = "horizontal",
+    show_titles: bool = True,
+    show_atom_map: bool = False,
+    titles: Tuple[str, str, str] = (
+        "Reactants",
+        "Imaginary Transition State",
+        "Products",
+    ),
+    add_gridbox: bool = False,
+    rule: bool = False,
+) -> plt.Figure:
+    """
+    Visualize three related graphs (reactants, imaginary transition state, and products)
+    side by side or vertically in a single figure.
+
+    Parameters:
+    - input (Union[str, Tuple[nx.Graph, nx.Graph, nx.Graph]]): Either
+    a reaction SMILES stringor a tuple of three NetworkX graphs
+    (reactants, products, ITS).
+    - sanitize (bool, optional): If True, sanitizes the input molecule.
+    Default is False.
+    - figsize (Tuple[int, int], optional): The size of the Matplotlib figure.
+    Default is (18, 5).
+    - orientation (str, optional): Layout of the subplots; 'horizontal' or 'vertical'.
+    Default is 'horizontal'.
+    - show_titles (bool, optional): If True, adds titles to each subplot.
+    Default is True.
+    - titles (Tuple[str, str, str], optional): Custom titles for each subplot.
+    Default is ('Reactants', 'Imaginary Transition State', 'Products').
+    - add_gridbox (bool, optional): If True, adds a gridbox cover for each subplot
+    (rectangular frame). Default is False.
+
+    Returns:
+    - plt.Figure: The Matplotlib figure containing the three subplots.
+    """
+    try:
+        # Parse input to determine graphs
+        if isinstance(input, str):
+            r, p = rsmi_to_graph(input, light_weight=True, sanitize=sanitize)
+            its = ITSConstruction().ITSGraph(r, p)
+        elif isinstance(input, tuple) and len(input) == 3:
+            r, p, its = input
+        else:
+            raise ValueError(
+                "Input must be a reaction SMILES string or a tuple of three graphs (r, p, its)."
+            )
+
+        # Set up subplots
+        if orientation == "horizontal":
+            fig, ax = plt.subplots(1, 3, figsize=figsize)
+        elif orientation == "vertical":
+            fig, ax = plt.subplots(3, 1, figsize=figsize)
+        else:
+            raise ValueError("Orientation must be 'horizontal' or 'vertical'.")
+
+        # Plot the graphs
+        vis_graph.plot_as_mol(
+            r,
+            ax[0],
+            show_atom_map=show_atom_map,
+            font_size=12,
+            node_size=800,
+            edge_width=2.0,
+        )
+        if show_titles:
+            ax[0].set_title(titles[0])
+
+        vis_graph.plot_its(
+            its, ax[1], use_edge_color=True, show_atom_map=show_atom_map, rule=rule
+        )
+        if show_titles:
+            ax[1].set_title(titles[1])
+
+        vis_graph.plot_as_mol(
+            p,
+            ax[2],
+            show_atom_map=show_atom_map,
+            font_size=12,
+            node_size=800,
+            edge_width=2.0,
+        )
+        if show_titles:
+            ax[2].set_title(titles[2])
+
+        # Add gridbox frame around each subplot if requested
+        if add_gridbox:
+            for a in ax:
+                # Make sure the grid is on top of the plot
+                a.set_axisbelow(False)
+
+                # Add a rectangular frame (gridbox) with thicker borders
+                for spine in a.spines.values():
+                    spine.set_visible(True)
+                    spine.set_linewidth(2)
+                    spine.set_color("black")
+
+                # Make gridlines lighter and under the plot elements
+                a.grid(
+                    True,
+                    which="both",
+                    axis="both",
+                    linestyle="--",
+                    color="gray",
+                    alpha=0.5,
+                )
+
+        return fig
+
+    except Exception as e:
+        raise RuntimeError(f"An error occurred during visualization: {str(e)}")
+
+
+def rule_visualize(gml, rule=True, titles=None):
+    """
+    Visualizes a reaction network from GML data with optional edge filtering (rule)
+    and custom titles.
+
+    Parameters:
+    - gml (str): GML format string representing the reaction data.
+    - rule (bool): If True, applies the rule to filter edges
+    (e.g., removing edges based on color).
+    - titles (list, optional): List of titles for the subplots.
+    Defaults to ['L', 'K', 'R'].
+
+    Returns:
+    - plt.Figure: Matplotlib figure containing the visualized reaction network.
+    """
+    try:
+        # Transform GML to NetworkX graphs
+        r, p, its = GMLToNX(gml).transform()
+
+        # If no titles are provided, default to ['L', 'K', 'R']
+        if titles is None:
+            titles = ["L", "K", "R"]
+
+        # Ensure titles match the number of graphs (3)
+        if len(titles) != 3:
+            raise ValueError(
+                "The titles list must contain exactly three titles for the three graphs."
+            )
+
+        # Call the `three_graph_vis` function with the transformed graphs and rule filtering
+        return three_graph_vis(
+            (r, p, its),
+            add_gridbox=True,  # Add the gridbox around the plot
+            titles=titles,  # Pass the titles for the subplots
+            rule=rule,  # Apply the rule filtering based on the value of `rule`
+        )
+
+    except Exception as e:
+        raise RuntimeError(
+            f"An error occurred during the visualization process: {str(e)}"
+        )