#!/usr/bin/env python3 """ Usage: python analyse_zoom_tasks.py --files file1 file2 --labels label1 label2 Running this script will produce a series of plots that compare different thread_info-step*.dat files produced by SWIFT when configured with --enable-task-debugging. This will produce a series of diagnostic plots that show the task counts, runtime and distance metrics for different types of tasks in the simulation. This script is specifically designed to work with zoom simulations, where the tasking is complicated by having several different types of cell and complex geometries. This file is part of SWIFT. Copyright (C) 2024 Will Roper (w.roper@sussex.ac.uk) All Rights Reserved. This program is free software: you can redistribute it and/or modify it under the terms of the GNU Lesser General Public License as published by the Free Software Foundation, either version 3 of the License, or (at your option) any later version. This program is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for more details. You should have received a copy of the GNU Lesser General Public License along with this program. If not, see . """ import argparse import matplotlib.gridspec as gridspec import matplotlib.patches as mpatches import matplotlib.pyplot as plt import numpy as np from task_parser import TaskParser def make_task_hist_time_split(runs, order_by_count=True, output=""): """ Plot the runtime of each task type including cell and depth information. Args: runs: Dictionary of TaskParser objects. order_by_count: If True, order the tasks by count, otherwise order by time. output: Output filepath. """ fig = plt.figure(figsize=(12, 16)) ax = fig.add_subplot(111) ax.set_xscale("log") ax.grid(True) # Combine all information into the labels labels_dict = { name: np.zeros(run.ntasks, dtype=object) for name, run in runs.items() } time_dict = {} for name, run in runs.items(): for i in range(run.ntasks): task = run.task_labels[i] labels_dict[name][i] = f"{task}:{run.tasks[i].ci_type}" if run.tasks[i].ci_subtype != "Regular": labels_dict[name][i] += f"({run.tasks[i].ci_subtype})" if "pair" in task: labels_dict[name][i] += f"->{run.tasks[i].cj_type}" if run.tasks[i].cj_subtype != "Regular": labels_dict[name][i] += f"({run.tasks[i].cj_subtype})" labels_dict[name][i] += f"@{run.tasks[i].ci_depth}" time_dict[labels_dict[name][i]] = ( time_dict.get(labels_dict[name][i], 0) + run.tasks[i].dt ) for i, (name, run) in enumerate(runs.items()): labels, counts = np.unique(labels_dict[name], return_counts=True) if order_by_count: # Sort the labels if i == 0: sinds = np.argsort(-counts) labels = labels[sinds] counts = counts[sinds] # Unpack the times times = np.array([time_dict[lab] for lab in labels]) else: # Unpack the times times = np.array([time_dict[lab] for lab in labels]) # Sort the labels by time if i == 0: sinds = np.argsort(-times) labels = labels[sinds] times = times[sinds] # Calculate positions for horizontal bars positions = np.arange(len(labels)) # Compute the width between labels width = 0.8 / (len(runs) + 1) # Create horizontal bar plot ax.barh( positions + (i * width), times, height=0.75 / len(runs), label=name, alpha=0.7, ) ax.set_yticks(np.arange(len(labels)) + 0.2) ax.set_yticklabels(labels) ax.invert_yaxis() ax.set_xlabel("Time (ms)") # Place the legend at the bottom of the plot ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.05), ncol=3) # Define the filename filename = "task_time_comp_split" if order_by_count: filename += "_count_ordered" filename += ".png" fig.tight_layout() fig.savefig(f"{output}{filename}", bbox_inches="tight") plt.close(fig) def make_task_hist_split(runs, output=""): """ Plot the count of each task type including cell and depth information. Args: runs: Dictionary of TaskParser objects. output: Output filepath. """ fig = plt.figure(figsize=(12, 16)) ax = fig.add_subplot(111) ax.set_xscale("log") ax.grid(True) # Combine all information into the labels labels_dict = { name: np.zeros(run.ntasks, dtype=object) for name, run in runs.items() } for name, run in runs.items(): for i in range(run.ntasks): task = run.task_labels[i] label = f"{task}:{run.tasks[i].ci_type}" if run.tasks[i].ci_subtype != "Regular": label += f"({run.tasks[i].ci_subtype})" if "pair" in task or "mm" in task: label += f"->{run.tasks[i].cj_type}" if run.tasks[i].cj_subtype != "Regular": label += f"({run.tasks[i].cj_subtype})" label += f"@{run.tasks[i].ci_depth}" labels_dict[name][i] = label for i, (name, run) in enumerate(runs.items()): labels, counts = np.unique(labels_dict[name], return_counts=True) # Sort the labels if i == 0: sinds = np.argsort(-counts) labels = labels[sinds] counts = counts[sinds] # Calculate positions for horizontal bars positions = np.arange(len(labels)) # Compute the width between labels width = 0.8 / (len(runs) + 1) # Create horizontal bar plot ax.barh( positions + (i * width), counts, height=0.75 / len(runs), label=name, alpha=0.7, ) ax.set_yticks(np.arange(len(labels)) + 0.2) ax.set_yticklabels(labels) ax.invert_yaxis() ax.set_xlabel("Count") # Place the legend at the bottom of the plot ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.05), ncol=3) # Define the filename filename = "task_count_comp_split.png" fig.tight_layout() fig.savefig(f"{output}{filename}", bbox_inches="tight") plt.close(fig) def make_task_hist( runs, ci_type=None, cj_type=None, ci_subtype=None, cj_subtype=None, depth=None, output="", ): """ Plot the count of each task type. If any of the filter arguments are specifie then a subset of tasks will be plotted. If all are None then all tasks will be plotted. Note that ci and ci are symmetric, i.e. ci_type=1, cj_type=2 is the same as ci_type=2, cj_type=1. Args: runs: Dictionary of TaskParser objects. ci_type: Cell type of the first cell to filter on. (Optional) cj_type: Cell type of the second cell to filter on. (Optional) ci_subtype: Cell subtype of the first cell to filter on. (Optional) cj_subtype: Cell subtype of the second cell to filter on. (Optional) depth: Depth of the tasks to filter on. (Optional) output: Output filepath. """ fig = plt.figure(figsize=(12, 6)) ax = fig.add_subplot(111) ax.set_xscale("log") ax.grid(True) # Combine all information into the labels nempty_runs = 0 need_sort = True for i, (name, run) in enumerate(runs.items()): mask = run.get_mask(ci_type, cj_type, ci_subtype, cj_subtype, depth) # Check we have something to plot if mask.sum() == 0: nempty_runs += 1 continue labels, counts = np.unique(run.task_labels[mask], return_counts=True) # Sort the labels and counts by counts in descending order if need_sort: sorted_indices = np.argsort(-counts) need_sort = False labels = labels[sorted_indices] counts = counts[sorted_indices] # Calculate positions for horizontal bars positions = np.arange(len(labels)) # Compute the width between labels width = 0.8 / (len(runs) + 1) # Create horizontal bar plot ax.barh( positions + (i * width), counts, height=0.75 / len(runs), label=name, alpha=0.7, ) # Exit if there's nothing to plot if nempty_runs == len(runs): print( f"Nothing to plot for ci_type={ci_type}, cj_type={cj_type} " f"ci_subtype={ci_subtype}, cj_subtype={cj_subtype}, depth={depth}" ) return ax.set_yticks(np.arange(len(labels)) + 0.2) ax.set_yticklabels(labels) ax.invert_yaxis() ax.set_xlabel("Count") # Place the legend at the bottom of the plot ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.15), ncol=3) # Define the filename filename = "task_count_comp" if ci_type is not None and cj_type is not None: filename += f"_types{ci_type}-{cj_type}" if ci_subtype is not None and cj_subtype is not None: filename += f"_subtypes{ci_subtype}-{cj_subtype}" if ci_type is not None and cj_type is None: filename += f"_type{ci_type}" if ci_subtype is not None and cj_subtype is None: filename += f"_subtype{ci_subtype}" if ci_type is None and cj_type is not None: filename += f"_type{cj_type}" if ci_subtype is None and cj_subtype is not None: filename += f"_subtype{cj_subtype}" if depth is not None: filename += f"_depth{depth}" fig.tight_layout() fig.savefig(f"{output}{filename}", bbox_inches="tight") plt.close(fig) def make_task_hist_time_weighted( runs, ci_type=None, cj_type=None, ci_subtype=None, cj_subtype=None, depth=None, output="", ): """ Plot the runtime of each task type. If any of the filter arguments are specifie then a subset of tasks will be plotted. If all are None then all tasks will be plotted. Note that ci and ci are symmetric, i.e. ci_type=1, cj_type=2 is the same as ci_type=2, cj_type=1. Args: runs: Dictionary of TaskParser objects. ci_type: Cell type of the first cell to filter on. (Optional) cj_type: Cell type of the second cell to filter on. (Optional) ci_subtype: Cell subtype of the first cell to filter on. (Optional) cj_subtype: Cell subtype of the second cell to filter on. (Optional) depth: Depth of the tasks to filter on. (Optional) output: Output filepath. """ fig = plt.figure(figsize=(12, 6)) ax = fig.add_subplot(111) ax.set_xscale("log") ax.grid(True) # Combine all information into the labels nempty_runs = 0 need_sort = True for i, (name, run) in enumerate(runs.items()): mask = run.get_mask(ci_type, cj_type, ci_subtype, cj_subtype, depth) # Check we have something to plot if mask.sum() == 0: nempty_runs += 1 continue # Loop over tasks collecting their runtime labels = np.unique(run.task_labels[mask]) counts = np.array( [np.sum(run.dt[mask][run.task_labels[mask] == k]) for k in labels] ) # Sort the labels and counts by counts in descending order if need_sort: sorted_indices = np.argsort(-counts) need_sort = False labels = labels[sorted_indices] counts = counts[sorted_indices] # Calculate positions for horizontal bars positions = np.arange(len(labels)) # Compute the width between labels width = 0.8 / (len(runs) + 1) # Create horizontal bar plot ax.barh( positions + (i * width), counts, height=0.75 / len(runs), label=name, alpha=0.7, ) # Exit if there's nothing to plot if nempty_runs == len(runs): print( f"Nothing to plot for ci_type={ci_type}, cj_type={cj_type} " f"ci_subtype={ci_subtype}, cj_subtype={cj_subtype}, depth={depth}" ) return ax.set_yticks(np.arange(len(labels)) + 0.2) ax.set_yticklabels(labels) ax.invert_yaxis() ax.set_xlabel("Time (ms)") # Place the legend at the bottom of the plot ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.15), ncol=3) # Define the filename filename = "task_time_comp" if ci_type is not None and cj_type is not None: filename += f"_types{ci_type}-{cj_type}" if ci_subtype is not None and cj_subtype is not None: filename += f"_subtypes{ci_subtype}-{cj_subtype}" if ci_type is not None and cj_type is None: filename += f"_type{ci_type}" if ci_subtype is not None and cj_subtype is None: filename += f"_subtype{ci_subtype}" if ci_type is None and cj_type is not None: filename += f"_type{cj_type}" if ci_subtype is None and cj_subtype is not None: filename += f"_subtype{cj_subtype}" if depth is not None: filename += f"_depth{depth}" fig.tight_layout() fig.savefig(f"{output}{filename}", bbox_inches="tight") plt.close(fig) def make_pair_mindist_plot( runs, ci_type=None, cj_type=None, ci_subtype=None, cj_subtype=None, depth=None, nbins=30, output="", ): """ Histogram of the minimum distances between cells. This will histogram the output of sqrt(cell_min_dist2). If any of the filter arguments are specifie then a subset of tasks will be plotted. If all are None then all tasks will be plotted. Note that ci and ci are symmetric, i.e. ci_type=1, cj_type=2 is the same as ci_type=2, cj_type=1. Args: runs: Dictionary of TaskParser objects. ci_type: Cell type of the first cell to filter on. (Optional) cj_type: Cell type of the second cell to filter on. (Optional) ci_subtype: Cell subtype of the first cell to filter on. (Optional) cj_subtype: Cell subtype of the second cell to filter on. (Optional) depth: Depth of the tasks to filter on. (Optional) nbins: Number of bins in the histogram. output: Output filepath. """ # Make the figure fig = plt.figure(figsize=(12, 6)) ax = fig.add_subplot(111) ax.set_yscale("log") ax.grid(True) # Collect the distances dists = {} for i, (name, run) in enumerate(runs.items()): mask = run.get_mask(ci_type, cj_type, ci_subtype, cj_subtype, depth) # Ensure we only have pair tasks (i.e. the string "pair" is in the # task label) mask = np.logical_and(mask, np.array(["pair" in t for t in run.task_labels])) # Get the distances dists[name] = run.min_dists[mask] # Collect together all the distances all_dists = np.concatenate(list(dists.values())) # Exit if there are no distances if all_dists.size == 0: print( f"No distances to plot for ci_type={ci_type}, cj_type={cj_type} " f"ci_subtype={ci_subtype}, cj_subtype={cj_subtype}, depth={depth}" ) return # Construct the bins bins = np.linspace(all_dists.min(), all_dists.max(), nbins + 1) bin_cents = (bins[:-1] + bins[1:]) / 2 # Compute histogram and plot for name in dists.keys(): linestyle = "--" if "long_range" in name else "-" H, _ = np.histogram(dists[name], bins=bins) ax.plot(bin_cents, H, label=name, linestyle=linestyle) ax.set_xlabel("sqrt(cell_min_dist2) (U_L)") ax.set_ylabel("Count") # Place the legend at the bottom of the plot ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.15), ncol=3) # Define the filename filename = "pair_min_dist_comp" if ci_type is not None and cj_type is not None: filename += f"_types{ci_type}-{cj_type}" if ci_subtype is not None and cj_subtype is not None: filename += f"_subtypes{ci_subtype}-{cj_subtype}" if ci_type is not None and cj_type is None: filename += f"_type{ci_type}" if ci_subtype is not None and cj_subtype is None: filename += f"_subtype{ci_subtype}" if ci_type is None and cj_type is not None: filename += f"_type{cj_type}" if ci_subtype is None and cj_subtype is not None: filename += f"_subtype{cj_subtype}" if depth is not None: filename += f"_depth{depth}" fig.tight_layout() fig.savefig(f"{output}{filename}", bbox_inches="tight") plt.close(fig) def make_pair_mpoledist_plot( runs, ci_type=None, cj_type=None, ci_subtype=None, cj_subtype=None, depth=None, nbins=30, output="", ): """ Histogram of the minimum distances between multipoles. This will histogram the distances between the centres of mass of the multipoles which have pair tasks. If any of the filter arguments are specifie then a subset of tasks will be plotted. If all are None then all tasks will be plotted. Note that ci and ci are symmetric, i.e. ci_type=1, cj_type=2 is the same as ci_type=2, cj_type=1. Args: runs: Dictionary of TaskParser objects. ci_type: Cell type of the first cell to filter on. (Optional) cj_type: Cell type of the second cell to filter on. (Optional) ci_subtype: Cell subtype of the first cell to filter on. (Optional) cj_subtype: Cell subtype of the second cell to filter on. (Optional) depth: Depth of the tasks to filter on. (Optional) nbins: Number of bins in the histogram. output: Output filepath. """ # Make the figure fig = plt.figure(figsize=(12, 6)) ax = fig.add_subplot(111) ax.set_yscale("log") ax.grid(True) # Collect the distances dists = {} for i, (name, run) in enumerate(runs.items()): mask = run.get_mask(ci_type, cj_type, ci_subtype, cj_subtype, depth) # Ensure we only have pair tasks (i.e. the string "pair" is in the # task label) mask = np.logical_and(mask, np.array(["pair" in t for t in run.task_labels])) # Get the distances dists[name] = run.mpole_dists[mask] # Collect together all the distances all_dists = np.concatenate(list(dists.values())) # Exit if there are no distances if all_dists.size == 0: print( f"No distances to plot for ci_type={ci_type}, cj_type={cj_type} " f"ci_subtype={ci_subtype}, cj_subtype={cj_subtype}, depth={depth}" ) return # Construct the bins bins = np.linspace(all_dists.min(), all_dists.max(), nbins + 1) bin_cents = (bins[:-1] + bins[1:]) / 2 # Compute histogram and plot for name in dists.keys(): linestyle = "--" if "long_range" in name else "-" H, _ = np.histogram(dists[name], bins=bins) ax.plot(bin_cents, H, label=name, linestyle=linestyle) ax.set_xlabel("Multipole CoM distance (U_L)") ax.set_ylabel("Count") # Place the legend at the bottom of the plot ax.legend(loc="upper center", bbox_to_anchor=(0.5, -0.15), ncol=3) # Define the filename filename = "pair_mpole_dist_comp" if ci_type is not None and cj_type is not None: filename += f"_types{ci_type}-{cj_type}" if ci_subtype is not None and cj_subtype is not None: filename += f"_subtypes{ci_subtype}-{cj_subtype}" if ci_type is not None and cj_type is None: filename += f"_type{ci_type}" if ci_subtype is not None and cj_subtype is None: filename += f"_subtype{ci_subtype}" if ci_type is None and cj_type is not None: filename += f"_type{cj_type}" if ci_subtype is None and cj_subtype is not None: filename += f"_subtype{cj_subtype}" if depth is not None: filename += f"_depth{depth}" fig.tight_layout() fig.savefig(f"{output}{filename}", bbox_inches="tight") plt.close(fig) def make_mindist_mpoledist_comp( run, ci_type=None, cj_type=None, ci_subtype=None, cj_subtype=None, depth=None, output="", ): """ Plot a scatter comparing the different distances. This function will compare the minimum distance between cells and the distance between the multipoles. Unlike other functions this function acts on a single run at a time. Args: run: A TaskParser object. ci_type: Cell type of the first cell to filter on. (Optional) cj_type: Cell type of the second cell to filter on. (Optional) ci_subtype: Cell subtype of the first cell to filter on. (Optional) cj_subtype: Cell subtype of the second cell to filter on. (Optional) depth: Depth of the tasks to filter on. (Optional) output: Output filepath. """ # Make the mask mask = run.get_mask(ci_type, cj_type, ci_subtype, cj_subtype, depth) # Ensure we only have pair tasks (i.e. the string "pair" is in the # task label) mask = np.logical_and(mask, np.array(["pair" in t for t in run.task_labels])) # Exit if there are no distances if np.sum(mask) == 0: print( f"No distances to plot for ci_type={ci_type}, cj_type={cj_type} " f"ci_subtype={ci_subtype}, cj_subtype={cj_subtype}, depth={depth}" ) return # Get the distances min_dists = run.min_dists[mask] mpole_dists = run.mpole_dists[mask] # Set up the figure fig = plt.figure(figsize=(6, 6)) ax = fig.add_subplot(111) ax.loglog() ax.grid(True) ax.scatter(min_dists, mpole_dists, marker=".", color="grey", alpha=0.7) ax.set_xlabel("Minimum distance between cells (U_L)") ax.set_ylabel("Distance between multipoles (U_L)") # Define the filename filename = "min_dist_mpole_dist_comp" if ci_type is not None and cj_type is not None: filename += f"_types{ci_type}-{cj_type}" if ci_subtype is not None and cj_subtype is not None: filename += f"_subtypes{ci_subtype}-{cj_subtype}" if ci_type is not None and cj_type is None: filename += f"_type{ci_type}" if ci_subtype is not None and cj_subtype is None: filename += f"_subtype{ci_subtype}" if ci_type is None and cj_type is not None: filename += f"_type{cj_type}" if ci_subtype is None and cj_subtype is not None: filename += f"_subtype{cj_subtype}" if depth is not None: filename += f"_depth{depth}" fig.tight_layout() # Save the figure fig.savefig(f"{output}{run.name}_{filename}.png", bbox_inches="tight") plt.close(fig) def make_task_plot( run, task_type=None, ci_type=None, cj_type=None, ci_subtype=None, cj_subtype=None, depth=None, output="", ): """ Plot tasks as a function of time per thread. This will give the same output as plot_tasks.py but offers more flexibility in the filters that can be applied. Args: run: A TaskParser object. task_type: The type of task to plot to filter on. (Optional) ci_type: Cell type of the first cell to filter on. (Optional) cj_type: Cell type of the second cell to filter on. (Optional) ci_subtype: Cell subtype of the first cell to filter on. (Optional) cj_subtype: Cell subtype of the second cell to filter on. (Optional) depth: Depth of the tasks to filter on. (Optional) """ # If we have nothing then exit and move on mask = run.get_mask(ci_type, cj_type, ci_subtype, cj_subtype, depth) ntasks_tot = mask.sum() if ntasks_tot == 0: print( f"No tasks to plot for task_type={task_type} ci_type={ci_type} " f"cj_type={cj_type} ci_subtype={ci_subtype} " f"cj_subtype={cj_subtype} depth={depth}" ) return # Get the dictionaries containing the labels, tics, tocs and colors labels, tics, tocs, colors = run.get_tasks_tictoc_by_thread( task_type, ci_type, cj_type, ci_subtype, cj_subtype, depth ) # Get the unique tasks unique_tasks, task_counts = np.unique(run.task_labels[mask], return_counts=True) ntasks = len(unique_tasks) # Set up the figure fig = plt.figure(figsize=(16, 0.15 * run.nthread)) ax = fig.add_subplot(111) ax.grid(True, zorder=0) ax.set_xlim(-run.delta_t * 0.01, run.delta_t * 1.01) ax.set_ylim(0.5, run.nthread + 1.0) # Loop over threads plotting the tasks for i in labels.keys(): # Collect the ranges and colors into lists _tictocs = [] _colors = [] # Loop over the tasks for j in range(len(labels[i])): _tictocs.append((tics[i][j], tocs[i][j] - tics[i][j])) _colors.append(colors[labels[i][j]]) ax.broken_barh( _tictocs, [i + 0.55, 0.9], facecolors=_colors, linewidth=0, zorder=1 ) # Create legend handles from sorted labels and colors sinds = np.argsort(-task_counts) sorted_labels = unique_tasks[sinds] sorted_colors = [colors[lab] for lab in sorted_labels] handles = [ mpatches.Patch(color=color, label=label) for label, color in zip(sorted_labels, sorted_colors) ] # Legend and room for it. nrow = ntasks / 8 ax.fill_between([0, 0], run.nthread, run.nthread + nrow, facecolor="white") ax.legend( handles=handles, loc="lower left", shadow=True, bbox_to_anchor=(0.0, 1.0, 1.0, 0.2), mode="expand", ncol=8, ) # Start and end of time-step ax.plot([0, 0], [0, run.nthread + nrow + 1], "k--", linewidth=1) ax.plot( [run.delta_t, run.delta_t], [0, run.nthread + nrow + 1], "k--", linewidth=1, zorder=1, ) ax.set_xlabel("Wall clock time [ms]") ax.set_ylabel("Thread ID") # Define the filename filename = "tasks" if ci_type is not None and cj_type is not None: filename += f"_types{ci_type}-{cj_type}" if ci_subtype is not None and cj_subtype is not None: filename += f"_subtypes{ci_subtype}-{cj_subtype}" if ci_type is not None and cj_type is None: filename += f"_type{ci_type}" if ci_subtype is not None and cj_subtype is None: filename += f"_subtype{ci_subtype}" if ci_type is None and cj_type is not None: filename += f"_type{cj_type}" if ci_subtype is None and cj_subtype is not None: filename += f"_subtype{cj_subtype}" if depth is not None: filename += f"_depth{depth}" if run.mpimode: filename += f"_rank{run.rank}" fig.tight_layout() # Save the figure fig.savefig(f"{output}{run.name}_{filename}.png", bbox_inches="tight") plt.close(fig) def make_task_activity_plot( run, task_type=None, ci_type=None, cj_type=None, ci_subtype=None, cj_subtype=None, depth=None, xbins=1000, sort_threads=True, output="", ): """ Histogram tasks as a function of time per thread. This will bin tasks into a grid, counting how many tasks ran in a particular bin. This is particularly useful for visualising runs containing numerous short lived tasks which will not necessarily appear in a normal task plot. Args: run: A TaskParser object. task_type: The type of task to plot to filter on. (Optional) ci_type: Cell type of the first cell to filter on. (Optional) cj_type: Cell type of the second cell to filter on. (Optional) ci_subtype: Cell subtype of the first cell to filter on. (Optional) cj_subtype: Cell subtype of the second cell to filter on. (Optional) depth: Depth of the tasks to filter on. (Optional) xbins: Number of bins in the x-axis. This defines the minimum time period that can be resolved. sort_threads: Sort the threads by the end of their tasks. (Optional) """ # If we have nothing then exit and move on ntasks_tot = run.get_mask(ci_type, cj_type, ci_subtype, cj_subtype, depth).sum() if ntasks_tot == 0: print( f"No tasks to plot for task_type={task_type} ci_type={ci_type} " f"cj_type={cj_type} ci_subtype={ci_subtype} " f"cj_subtype={cj_subtype} depth={depth}" ) return # Get the dictionaries containing the labels, tics, tocs and colors labels, tics, tocs, colors = run.get_tasks_tictoc_by_thread( task_type, ci_type, cj_type, ci_subtype, cj_subtype, depth ) # Exit if there are no tasks to plot if len(labels) == 0: print( f"No tasks to plot for task_type={task_type} ci_type={ci_type} " f"cj_type={cj_type} ci_subtype={ci_subtype} " f"cj_subtype={cj_subtype} depth={depth}" ) return # Define the grid of bins. This has shape (run.nthread, xbins) grid = np.zeros((run.nthread, xbins)) # Populate the grid thread by thread and task by task for i in labels.keys(): for j in range(len(labels[i])): # Calculate the tic bin xtic_bin = int(tics[i][j] / run.delta_t * xbins) # Calculate the toc bin xtoc_bin = int(tocs[i][j] / run.delta_t * xbins) # Populate the bins for xbin in range(xtic_bin, xtoc_bin): grid[i, xbin] = 1 # Sort the rows of the grid if sort_threads: counts = [np.sum(grid[i, :]) for i in range(run.nthread)] sinds = np.argsort(counts)[::-1] grid = grid[sinds, :] # Set up the figure with a main plot for the grid and a histogram on top fig = plt.figure(figsize=(16, 8)) gs = gridspec.GridSpec( 2, 2, height_ratios=[2, 10], width_ratios=[20, 1], hspace=0.0 ) ax_grid = fig.add_subplot(gs[1, 0]) ax_hist = fig.add_subplot(gs[0, 0]) ax_hist.grid(True) ax_hist.set_axisbelow(True) cax = fig.add_subplot(gs[:, 1]) # Plot the grid im = ax_grid.pcolormesh( np.linspace(0, run.delta_t, xbins), np.linspace(0, run.nthread - 1, run.nthread), grid, cmap="coolwarm", ) # Start and end of time-step ax_grid.plot([0, 0], ax_grid.get_ylim(), "k--", linewidth=1) ax_grid.plot([run.delta_t, run.delta_t], ax_grid.get_ylim(), "k--", linewidth=1) ax_grid.set_xlabel("Wall clock time [ms]") ax_grid.set_ylabel("Thread ID") # Create the colorbar cbar = fig.colorbar(im, cax=cax, ticks=[0, 1]) cbar.set_ticklabels(["Inactive", "Task"]) # Ensure the histogram aligns perfectly with the grid plot ax_hist.set_xlim(ax_grid.get_xlim()) # Calculate the sum across each column for the histogram column_sums = grid.sum(axis=0) ax_hist.bar( np.linspace(0, run.delta_t, xbins), column_sums, width=np.diff(np.linspace(0, run.delta_t, xbins))[0], align="edge", zorder=1, ) ax_hist.set_ylabel("Tasks Running") ax_hist.set_ylim(0, run.nthread + 0.1 * run.nthread) # Turn off the x-axis ticks and tick labels on the histogram ax_hist.tick_params(axis="x", which="both", bottom=False, top=False) ax_hist.set_xticklabels([]) # Define the filename filename = "task_activity" if ci_type is not None and cj_type is not None: filename += f"_types{ci_type}-{cj_type}" if ci_subtype is not None and cj_subtype is not None: filename += f"_subtypes{ci_subtype}-{cj_subtype}" if ci_type is not None and cj_type is None: filename += f"_type{ci_type}" if ci_subtype is not None and cj_subtype is None: filename += f"_subtype{ci_subtype}" if ci_type is None and cj_type is not None: filename += f"_type{cj_type}" if ci_subtype is None and cj_subtype is not None: filename += f"_subtype{cj_subtype}" if depth is not None: filename += f"_depth{depth}" if run.mpimode: filename += f"_rank{run.rank}" if sort_threads: filename += "_sorted" fig.tight_layout() # Save the figure fig.savefig(f"{output}{run.name}_{filename}.png", bbox_inches="tight") plt.close(fig) if __name__ == "__main__": # Define the command line arguments parser = argparse.ArgumentParser( description="Produce task analysis plots for SWIFT zoom simulations" ) # Adding files argument parser.add_argument( "--files", nargs="+", help="List of files to combine on outputs", required=True ) # Adding labels argument parser.add_argument("--labels", nargs="+", help="List of labels", default=[]) # Adding output directory parser.add_argument("--outdir", help="Output directory", default=".") # Adding output base name parser.add_argument("--outbase", help="Output base name", default="") # Parse the arguments args = parser.parse_args() files = args.files labels = args.labels outdir = args.outdir outbase = args.outbase output = f"{outdir}/{outbase}" if output[-1] != "_" and output[-1] != "/": output += "_" print("Writing outputs to:", output) if len(labels) == 0: labels = files print("Using filenames as labels") if len(files) != len(labels): raise ValueError("Number of files and labels must match") # Parse all the task files runs = {} for f, lab in zip(files, labels): runs[lab] = TaskParser(f, lab) # Below we call the functions for all common useful combinations of filters # Detailed stacked histograms make_task_hist_split(runs, output=output) make_task_hist_time_split(runs, output=output) make_task_hist_time_split(runs, order_by_count=False, output=output) # Counts of tasks make_task_hist(runs, output=output) make_task_hist(runs, ci_type=1, cj_type=1, output=output) make_task_hist(runs, ci_type=2, cj_type=2, output=output) make_task_hist(runs, ci_type=3, cj_type=3, output=output) make_task_hist(runs, ci_type=1, cj_type=2, output=output) make_task_hist(runs, ci_type=1, cj_type=3, output=output) make_task_hist(runs, ci_type=2, cj_type=3, output=output) # Count useful subtypes tasks to be sure there are none make_task_hist(runs, ci_subtype=0, output=output) # regular make_task_hist(runs, ci_subtype=1, output=output) # neighbour make_task_hist(runs, ci_subtype=2, output=output) # void # Counts of tasks but only depth 0 make_task_hist(runs, depth=0, output=output) make_task_hist(runs, ci_type=1, cj_type=1, depth=0, output=output) make_task_hist(runs, ci_type=2, cj_type=2, depth=0, output=output) make_task_hist(runs, ci_type=3, cj_type=3, depth=0, output=output) make_task_hist(runs, ci_type=1, cj_type=3, depth=0, output=output) make_task_hist(runs, ci_type=1, cj_type=2, depth=0, output=output) make_task_hist(runs, ci_type=2, cj_type=3, depth=0, output=output) # Time weighted counts of tasks make_task_hist_time_weighted(runs, output=output) make_task_hist_time_weighted(runs, ci_type=1, cj_type=1, output=output) make_task_hist_time_weighted(runs, ci_type=2, cj_type=2, output=output) make_task_hist_time_weighted(runs, ci_type=3, cj_type=3, output=output) make_task_hist_time_weighted(runs, ci_type=1, cj_type=3, output=output) make_task_hist_time_weighted(runs, ci_type=1, cj_type=2, output=output) make_task_hist_time_weighted(runs, ci_type=2, cj_type=3, output=output) # Time weighted counts of tasks but only depth 0 make_task_hist_time_weighted(runs, depth=0, output=output) make_task_hist_time_weighted(runs, ci_type=1, cj_type=1, depth=0, output=output) make_task_hist_time_weighted(runs, ci_type=2, cj_type=2, depth=0, output=output) make_task_hist_time_weighted(runs, ci_type=3, cj_type=3, depth=0, output=output) make_task_hist_time_weighted(runs, ci_type=1, cj_type=3, depth=0, output=output) make_task_hist_time_weighted(runs, ci_type=1, cj_type=2, depth=0, output=output) make_task_hist_time_weighted(runs, ci_type=2, cj_type=3, depth=0, output=output) # Pair distance histograms make_pair_mindist_plot(runs, output=output) make_pair_mindist_plot(runs, ci_type=1, cj_type=1, output=output) make_pair_mindist_plot(runs, ci_type=2, cj_type=2, output=output) make_pair_mindist_plot(runs, ci_type=3, cj_type=3, output=output) make_pair_mindist_plot(runs, ci_type=1, cj_type=3, output=output) make_pair_mindist_plot(runs, ci_type=1, cj_type=2, output=output) make_pair_mindist_plot(runs, ci_type=2, cj_type=3, output=output) # Pair multipole distance histograms make_pair_mpoledist_plot(runs, output=output) make_pair_mpoledist_plot(runs, ci_type=1, cj_type=1, output=output) make_pair_mpoledist_plot(runs, ci_type=2, cj_type=2, output=output) make_pair_mpoledist_plot(runs, ci_type=3, cj_type=3, output=output) make_pair_mpoledist_plot(runs, ci_type=1, cj_type=3, output=output) make_pair_mpoledist_plot(runs, ci_type=1, cj_type=2, output=output) make_pair_mpoledist_plot(runs, ci_type=2, cj_type=3, output=output) # Make the plots which only plot a single run at a time for run in runs.values(): # Distance comparison make_mindist_mpoledist_comp(run, output=output) make_mindist_mpoledist_comp(run, ci_type=1, cj_type=1, output=output) make_mindist_mpoledist_comp(run, ci_type=2, cj_type=2, output=output) make_mindist_mpoledist_comp(run, ci_type=3, cj_type=3, output=output) make_mindist_mpoledist_comp(run, ci_type=1, cj_type=3, output=output) make_mindist_mpoledist_comp(run, ci_type=1, cj_type=2, output=output) make_mindist_mpoledist_comp(run, ci_type=2, cj_type=3, output=output) # Make the task plots showing tasks per thread as a function of time make_task_plot(run, output=output) make_task_plot(run, ci_type=1, output=output) make_task_plot(run, ci_type=2, output=output) make_task_plot(run, ci_type=3, output=output) make_task_plot(run, ci_type=1, cj_type=1, output=output) make_task_plot(run, ci_type=2, cj_type=2, output=output) make_task_plot(run, ci_type=3, cj_type=3, output=output) make_task_plot(run, ci_type=1, cj_type=3, output=output) make_task_plot(run, ci_type=1, cj_type=2, output=output) make_task_plot(run, ci_type=2, cj_type=3, output=output) # Make task activity plots make_task_activity_plot(run, output=output) make_task_activity_plot(run, ci_type=1, output=output) make_task_activity_plot(run, ci_type=2, output=output) make_task_activity_plot(run, ci_type=3, output=output) make_task_activity_plot(run, ci_type=1, cj_type=1, output=output) make_task_activity_plot(run, ci_type=2, cj_type=2, output=output) make_task_activity_plot(run, ci_type=3, cj_type=3, output=output) make_task_activity_plot(run, ci_type=1, cj_type=3, output=output) make_task_activity_plot(run, ci_type=1, cj_type=2, output=output) make_task_activity_plot(run, ci_type=2, cj_type=3, output=output) # Make task activity plots but don't sort the threads make_task_activity_plot(run, sort_threads=False, output=output) make_task_activity_plot(run, ci_type=1, sort_threads=False, output=output) make_task_activity_plot(run, ci_type=2, sort_threads=False, output=output) make_task_activity_plot(run, ci_type=3, sort_threads=False, output=output) make_task_activity_plot( run, ci_type=1, cj_type=1, sort_threads=False, output=output ) make_task_activity_plot( run, ci_type=2, cj_type=2, sort_threads=False, output=output ) make_task_activity_plot( run, ci_type=3, cj_type=3, sort_threads=False, output=output ) make_task_activity_plot( run, ci_type=1, cj_type=3, sort_threads=False, output=output ) make_task_activity_plot( run, ci_type=1, cj_type=2, sort_threads=False, output=output ) make_task_activity_plot( run, ci_type=2, cj_type=3, sort_threads=False, output=output )