wmaee.scopes.md

 1"""Analysis and plotting of (AI)MD runs"""
 2
 3import numpy as np
 4import pandas as pd
 5import matplotlib.pyplot as plt
 6from typing import Optional, Tuple, List
 7
 8def plot_MD(df: pd.DataFrame, 
 9            time: Optional[str] = None,
10            timestep: int = 1,
11            grid: Tuple[int, int] = (1, 1),
12            props: List[str] = [],
13            labels: Optional[List[str]] = None,
14            show: bool = True,
15            **kwargs) -> Tuple[plt.Figure, np.ndarray]:
16    """
17    Plots multiple properties from a DataFrame in a grid layout.
18    
19    Parameters
20    ----------
21    df : pd.DataFrame
22        The DataFrame containing the data to plot.
23    time : Optional[str], optional
24        The column name to use as the time axis. If None, a default range is used, by default None.
25    timestep : int, optional
26        The timestep interval for the default time axis, by default 1.
27    grid : Tuple[int, int], optional
28        The grid size for the subplots (rows, columns), by default (1, 1).
29    props : List[str], optional
30        The list of DataFrame columns to plot, by default [].
31    labels : Optional[List[str]], optional
32        The labels for the y-axis of each subplot, by default None.
33    show : bool, optional
34        Whether to display the plot, by default True.
35    **kwargs
36        Additional keyword arguments passed to `plt.subplots`.
37
38    Returns
39    -------
40    Tuple[plt.Figure, np.ndarray]
41        The figure and axes array of the created subplots.
42    """
43    
44    # Determine the time axis values
45    if time is None:
46        time = np.arange(len(df)) * timestep
47    else:
48        time = df[time]    
49    
50    # Create subplots with the specified grid size
51    fig, axs = plt.subplots(*grid, squeeze=False, sharex=True, **kwargs)
52            
53    # Iterate through the grid and plot the data
54    for i in range(grid[0]):
55        for j in range(grid[1]):
56            n = i * grid[1] + j
57            if n >= len(props):
58                break
59            # Plot the data for the nth property
60            axs[i, j].plot(time, df[props[n]])
61            # Set the y-axis label if provided
62            if labels is not None:
63                axs[i, j].set_ylabel(labels[n])
64            # Set the x-axis label for the last row
65            axs[-1, j].set_xlabel('step')
66        else:
67            continue
68        break
69
70    # Adjust subplot spacing
71    fig.subplots_adjust(hspace=0)
72    # Show the plot if requested
73    if show:
74        plt.show()
75    
76    return fig, axs