postprocessor.py

import os
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
 
class PostProcessor:
    """
    Loads multiple simulation CSV outputs and provides visualization tools.
 
    CSV dictionary example:
        {"phi": "phi.csv", "efield": "efield.csv", "x": "x.csv", "time": "time.csv", "dx": "dx.csv"}
 
    Node-based fields are plotted at nodes (x.csv).  
    Cell-based fields are plotted at cell centers (dx.csv used to compute positions).
    """
    def __init__(self, folder="./output", csv_dict=None, cell_fields=None):
        """
        Parameters
        ----------
        folder : str
            Relative path to folder containing the CSV files.
        csv_dict : dict
            Dictionary mapping attribute names -> CSV filenames.
            Special keys:
                - "x": spatial node positions (single column)
                - "time": time array (single column)
                - "dx": cell sizes (single column)
        cell_fields : list
            List of names that are defined on cells (not nodes).
        """
        self.colormap = plt.cm.viridis
        plt.rcParams.update({'font.size':15,'lines.linewidth':2.5,'figure.autolayout':True})
 
        self.folder = folder
        if csv_dict is None:
            raise ValueError("You must provide a csv_dict mapping names to CSV files.")
        if cell_fields is None:
            cell_fields = []
 
        self._names = list(csv_dict.keys())
        self._cell_fields = cell_fields
 
        # Initialize default x, time, dx arrays
        self.x = None
        self.time = None
        self.dx = None
        self.x_cell = None
 
        # Load CSV files
        for name, fname in csv_dict.items():
            path = os.path.join(folder, fname)
            if not os.path.exists(path):
                raise FileNotFoundError(f"CSV file not found: {path}")
            print(f"[PostProcessor] Loading {name} from {path}")
            data = np.loadtxt(path, skiprows=1) # first row includes metadata comment
 
            # Flatten for 1D arrays
            if name in ("x", "time", "dx"):
                data = data.flatten()
                setattr(self, name, data)
            else:
                setattr(self, name, data)
 
        # Set default x and time if not provided
        first_field = next((n for n in self._names if n not in ("x", "time", "dx")), None)
        if first_field is None:
            raise ValueError("No field CSVs found (excluding 'x', 'time', 'dx').")
 
        data_field = getattr(self, first_field)
 
        if self.x is None:
            self.x = np.arange(data_field.shape[1])
        if self.time is None:
            self.time = np.arange(data_field.shape[0])
 
        # Compute cell-centered positions if dx is provided
        if hasattr(self, "dx") and self.dx is not None:
            if len(self.dx) != len(self.x)-1:
                raise ValueError("dx length must be number of nodes - 1")
            self.x_cell = self.x[:-1] + 0.5*self.dx
 
    # ------------------------------------------------------------
    # Determine plotting x-axis based on node vs cell field
    # ------------------------------------------------------------
    def _get_x_for_field(self, name):
        if name in self._cell_fields:
            if self.x_cell is None:
                raise ValueError(f"Field '{name}' is cell-based but dx/x not defined")
            return self.x_cell
        else:
            return self.x
 
    # ------------------------------------------------------------
    # Plot a single timestep (step for cell fields)
    # ------------------------------------------------------------
    def plot_timestep(self, name, t, ax=None):
        if not hasattr(self, name):
            raise AttributeError(f"No field named '{name}' loaded.")
        data = getattr(self, name)
 
        if t < 0 or t >= data.shape[0]:
            raise IndexError(f"Timestep {t} is out of range")
 
        x_axis = self._get_x_for_field(name)
 
        if ax is None:
            fig, ax = plt.subplots()
 
        if name in self._cell_fields:
            # Plot as constant per cell using step function
            x_step = np.zeros(len(x_axis)+1)
            x_step[:-1] = x_axis
            # extend last cell to keep value constant
            last_dx = self.dx[-1] if self.dx is not None else x_axis[-1]-x_axis[-2]
            x_step[-1] = x_axis[-1] + last_dx
 
            y_step = np.zeros(len(x_axis)+1)
            y_step[:-1] = data[t, :]
            y_step[-1] = data[t, -1]
 
            ax.step(x_step, y_step, where='post', lw=2)
        else:
            ax.plot(x_axis, data[t, :], lw=2)
 
        ax.set_title(f"{name} at t = {self.time[t]:.3e}")
        ax.set_xlabel("x-position")
        ax.set_ylabel(name)
        ax.grid(True)
 
        if ax is None:
            plt.show()
 
    # ------------------------------------------------------------
    # Animate a field (step for cell fields)
    # ------------------------------------------------------------
    def animate(self, name, step=1, interval=80):
        if not hasattr(self, name):
            raise AttributeError(f"No field named '{name}' loaded.")
        data = getattr(self, name)
 
        x_axis = self._get_x_for_field(name)
        num_nodes = len(x_axis)
 
        fig, ax = plt.subplots()
        line, = ax.plot([], [], lw=2)
 
        # Extend last cell for cell fields
        last_dx = self.dx[-1] if (name in self._cell_fields and self.dx is not None) else (x_axis[-1]-x_axis[-2])
        ax.set_xlim(np.min(x_axis), np.max(x_axis) + last_dx)
        ymin, ymax = np.min(data), np.max(data)
        ax.set_ylim(ymin, ymax)
        ax.set_xlabel("x-position")
        ax.set_ylabel(name)
        ax.grid(True)
 
        frames = range(0, data.shape[0], step)
 
        def init():
            line.set_data([], [])
            return line,
 
        def update(frame):
            if name in self._cell_fields:
                x_step = np.zeros(len(x_axis)+1)
                x_step[:-1] = x_axis
                x_step[-1] = x_axis[-1] + last_dx
                y_step = np.zeros(len(x_axis)+1)
                y_step[:-1] = data[frame, :]
                y_step[-1] = data[frame, -1]
                line.set_data(x_step, y_step)
            else:
                line.set_data(x_axis, data[frame])
            ax.set_title(f"{name} at t = {self.time[frame]:.3e}")
            return line,
 
        anim = FuncAnimation(
            fig, update, frames=frames,
            init_func=init, blit=True, interval=interval
        )
 
        plt.show()
        return anim
 
    def plot_columns_over_time(self, name: str, columns=None):
        """
        Plot selected columns of a variable over time.
 
        Parameters
        ----------
        name : str
            Name of the variable loaded from CSV.
        columns : list or array-like, optional
            Indices of columns to plot. If None, all columns are plotted.
        """
        if not hasattr(self, name):
            raise AttributeError(f"No field named '{name}' loaded.")
 
        data = getattr(self, name)
 
        if self.time is None:
            raise ValueError("Time vector not loaded. Provide a 'time' CSV in initialization.")
 
        n_cols = data.shape[1]
 
        # If no specific columns given → use all
        if columns is None:
            columns = list(range(n_cols))
        else:
            # Convert to list and validate
            columns = list(columns)
            for c in columns:
                if c < 0 or c >= n_cols:
                    raise IndexError(f"Column index {c} out of range (0 to {n_cols-1}).")
 
        # Colormap for selected columns only
        colors = self.colormap(np.linspace(0, 1, len(columns)))
 
        plt.figure(figsize=(8, 5))
 
        for idx, col in enumerate(columns):
            plt.plot(self.time, data[:, col], label=f"Col {col}", color=colors[idx])
 
        plt.xlabel("Time")
        plt.ylabel(name)
        plt.title(f"{name}: selected columns vs. time")
        plt.legend()
        plt.grid(True)
        plt.tight_layout()
        plt.show()
 
    def animate_histogram(self, name: str, step=1, interval=80):
        """
        Create an animated histogram where each row of `data` is one timestep.
 
        Parameters
        ----------
        data : np.ndarray
            2D array of shape (timesteps, bins).
        bin_edges : np.ndarray, optional
            1D array specifying the histogram bin edges. If None, bins = [0, 1, ..., N].
        interval : int
            Delay between frames in milliseconds.
        repeat : bool
            Whether the animation should loop.
        """
        if not hasattr(self, name):
            raise AttributeError(f"No field named '{name}' loaded.")
        data = getattr(self, name)
        if self.time is None:
            raise ValueError("Time vector not loaded. Provide a 'time' CSV in initialization.")
 
        timesteps, nbins = data.shape
        bin_edges = np.arange(nbins + 1)
 
        fig, ax = plt.subplots()
        bars = ax.bar(bin_edges[:-1], data[0], width=np.diff(bin_edges), align="edge")
 
        ax.set_xlim(bin_edges[0], bin_edges[-1])
        ax.set_ylim(0, 1.1 * data.max())
        ax.set_xlabel("Bin")
        ax.set_ylabel("Counts")
        ax.set_title("Animated Histogram")
 
        def update(frame):
            for b, h in zip(bars, data[frame]):
                b.set_height(h)
            ax.set_title(f"Timestep {frame}")
            return bars
 
        ani = FuncAnimation(
            fig, update, frames=timesteps, interval=interval, blit=False
        )
 
        plt.show()
        return ani