octopi 1.4.0__py3-none-any.whl

octopi/utils/submit_slurm.py

@@ -0,0 +1,95 @@
+def create_shellsubmit(
+    job_name, 
+    output_file,
+    shell_name,
+    conda_path,
+    command,
+    num_gpus = 1, 
+    gpu_constraint = 'h100'):
+
+    if num_gpus > 0:
+        slurm_gpus = f'#SBATCH --partition=gpu\n#SBATCH --gpus={gpu_constraint}:{num_gpus}'
+    else:
+        slurm_gpus = f'#SBATCH --partition=cpu'
+
+    shell_script_content = f"""#!/bin/bash
+
+{slurm_gpus}
+#SBATCH --time=18:00:00
+#SBATCH --cpus-per-task=4
+#SBATCH --mem-per-cpu=16G
+#SBATCH --job-name={job_name}
+#SBATCH --output={output_file}
+
+ml anaconda 
+conda activate {conda_path}
+{command}
+"""
+
+    # Save to file
+    with open(shell_name, "w") as file:
+        file.write(shell_script_content)
+
+    print(f"\nShell script has been created successfully as {shell_name}\n")
+
+def create_shellsubmit_array(
+    job_name, 
+    output_file,
+    shell_name,
+    conda_path,
+    command,
+    job_array = [min, max]):
+
+    shell_script_content = f"""#!/bin/bash
+
+#SBATCH --time=18:00:00
+#SBATCH --cpus-per-task=4
+#SBATCH --mem-per-cpu=16G
+#SBATCH --job-name={job_name}
+#SBATCH --output={output_file}
+#SBATCH --array={job_array[0]}-{job_array[1]}
+
+ml anaconda 
+conda activate {conda_path}
+{command}
+"""
+
+    # Save to file
+    with open(shell_name, "w") as file:
+        file.write(shell_script_content)
+
+    print(f"\nShell script has been created successfully as {shell_name}\n")
+
+def create_multiconfig_shellsubmit(
+    job_name, 
+    output_file,
+    shell_name,
+    conda_path,
+    base_inputs,
+    config_inputs,
+    command):
+
+    multiconfig = f"""#! /bin/bash
+
+#SBATCH --job-name={job_name}
+#SBATCH --time=24:00:00
+#SBATCH --cpus-per-task=4
+#SBATCH --mem-per-cpu=12G
+#SBATCH --partition=cpu
+#SBATCH --output={output_file}
+
+ml anaconda 
+{conda_path}
+
+{base_inputs}
+
+{config_inputs}
+
+{command}
+"""
+
+    # Save to file
+    with open(shell_name, "w") as file:
+        file.write(multiconfig)
+
+    print(f"\nShell script has been created successfully as {shell_name}\n")

octopi/utils/visualization_tools.py

@@ -0,0 +1,290 @@
+from ipywidgets import interact, IntSlider, fixed
+from copick_utils.io import readers
+import matplotlib.colors as mcolors
+from typing import Optional, List
+import matplotlib.pyplot as plt
+import numpy as np
+import copick
+
+# Define the interactive function
+def interact_3d_seg(vol, seg):
+    """
+    Interactively show the segmentation on a tomogram.
+    
+    Args:
+        vol (numpy.ndarray): The tomogram to show the segmentation on.
+        seg (numpy.ndarray): The segmentation to show on the tomogram.
+    """
+
+    # Get the number of slices for the slider range
+    max_slices = vol.shape[0] - 1
+    middle_slice = int(max_slices // 2)
+
+    # Launch the Interactive Widget
+    interact(
+        show_tomo_segmentation,
+        tomo=fixed(vol), seg=fixed(seg),
+        vol_slice=IntSlider(min=0, max=max_slices, step=1, value=middle_slice)
+    )
+
+# Define the plotting function
+def show_tomo_segmentation(tomo, seg, vol_slice):
+    """
+    Show Segmentation on a Tomogram Slice.
+
+    Args:
+        tomo (numpy.ndarray): The tomogram to show the segmentation on.
+        seg (numpy.ndarray): The segmentation to show on the tomogram.
+        vol_slice (int): The slice index to show.
+    """
+    
+    plt.figure(figsize=(20, 10))
+    
+    # Tomogram
+    plt.subplot(1, 3, 1)
+    plt.title('Tomogram')
+    plt.imshow(tomo[vol_slice], cmap='gray')
+    plt.axis('off')
+    
+    # Painted Segmentation
+    plt.subplot(1, 3, 2)
+    plt.title('Painted Segmentation from Picks')
+    plt.imshow(seg[vol_slice], cmap='viridis')
+    plt.axis('off')
+
+    plt.subplot(1, 3, 3)
+    plt.title('Overlay')
+    plt.imshow(tomo[vol_slice], cmap='gray')
+    plt.imshow(seg[vol_slice], cmap='viridis', alpha=0.5)  # Add alpha=0.5 for 50% transparency
+    plt.axis('off')
+    
+    plt.tight_layout()
+    plt.show()
+
+def show_labeled_tomo_segmentation(tomo, seg, seg_labels, unique_values, vol_slice):
+
+        # # Check unique values in segmentation to ensure correct mapping
+        # unique_values = np.unique(seg)
+
+    plt.figure(figsize=(20, 10))
+    num_classes = len(seg_labels)        
+
+    # Dynamically update the labels and colormap based on unique values
+    seg_labels_filtered = {k: v for k, v in seg_labels.items() if k in unique_values}
+    num_classes = len(seg_labels_filtered)
+
+    # Create a discrete colormap
+    colors = plt.cm.tab20b(np.linspace(0, 1, num_classes))  # You can use other colormaps like 'Set3', 'tab20', etc.
+    cmap = mcolors.ListedColormap(colors)
+    bounds = list(seg_labels_filtered.keys()) + [max(seg_labels_filtered.keys())]
+    # norm = mcolors.BoundaryNorm(bounds, cmap.N)
+
+    # Tomogram plot
+    plt.subplot(1, 2, 1)
+    plt.title('Tomogram')
+    plt.imshow(tomo[vol_slice], cmap='gray')
+    plt.axis('off')
+
+    # Prediction segmentation plot
+    plt.subplot(1, 2, 2)
+    plt.title('Prediction Segmentation')
+    im = plt.imshow(seg[vol_slice], cmap=cmap)  # Use norm and cmap for segmentation
+    plt.axis('off')
+
+    # Add the labeled color bar
+    cbar = plt.colorbar(im, ticks=list(seg_labels_filtered.keys()))
+    cbar.ax.set_yticklabels([seg_labels_filtered[i] for i in seg_labels_filtered.keys()])  # Set custom labels
+
+    plt.tight_layout()
+    plt.show()    
+
+def interact_points(
+    tomo, config, run_id, user_id='octopi', 
+    session_id = None, pt_size = 15,
+    slice_proximity_threshold = 3
+    ):
+    """
+    Interactively show the points on a tomogram.
+
+    Args:
+        tomo (numpy.ndarray): The tomogram to show the points on.
+        run_id (str): The ID of the run to show the points on.
+        user_id (str): The ID of the user to show the points on.
+        session_id (str): The ID of the session to show the points on.
+        slice_proximity_threshold (int): The threshold for the proximity of the points to the slice.
+        pt_size (int): The size of the points to show.
+    """
+
+    # Load Copick Project and Run
+    root = copick.from_file(config)
+    run = root.get_run(run_id)
+
+    # Get objects that can be Picked
+    objects = [(obj.name, obj.label, obj.radius) for obj in root.pickable_objects if obj.is_particle]
+
+    # Get the number of slices for the slider range
+    max_slices = tomo.shape[0] - 1
+    middle_slice = int(max_slices // 2)
+
+    # Launch the Interactive Widget
+    interact(
+        show_tomo_points,
+        tomo=fixed(tomo), run=fixed(run), objects=fixed(objects), 
+        user_id=fixed(user_id), session_id=fixed(session_id), 
+        slice_proximity_threshold=fixed(slice_proximity_threshold),
+        pt_size=fixed(pt_size),
+        vol_slice=IntSlider(min=0, max=max_slices, step=1, value=middle_slice)
+    )
+
+def show_tomo_points(
+        tomo, run, objects, user_id, 
+        vol_slice, session_id = None, 
+        slice_proximity_threshold = 3, pt_size = 15
+    ):
+    """
+    Show Coordinates on a Tomogram Slice.
+
+    Args:
+        tomo (numpy.ndarray): The tomogram to show the points on.
+        run (copick.Run): The Copick Run to show the points from.
+        objects (list): List of pickable objects.
+        user_id (str): The ID of the user to show the points from.
+        vol_slice (int): The slice index to show.
+        session_id (str): The ID of the session to show the points from.
+        slice_proximity_threshold (int): The threshold for the proximity of the points to the slice
+    """
+
+
+    plt.figure(figsize=(20, 10))
+
+    plt.imshow(tomo[vol_slice],cmap='gray')
+    plt.axis('off')
+
+    for name,_,_ in objects:
+        try:    
+            coordinates = readers.coordinates(run, name=name, user_id=user_id, session_id=session_id)
+            close_points = coordinates[np.abs(coordinates[:, 0] - vol_slice) <= slice_proximity_threshold]
+            plt.scatter(close_points[:, 2], close_points[:, 1], label=name, s=pt_size)
+        except:
+            pass
+
+    plt.show()
+
+def compare_tomo_points(tomo, run, objects, vol_slice, user_id1, user_id2, 
+                        session_id1 = None, session_id2 = None, slice_proximity_threshold = 3):
+    plt.figure(figsize=(20, 10))
+
+    plt.subplot(1, 2, 1)
+    plt.imshow(tomo[vol_slice],cmap='gray')
+    plt.title(f'{user_id1} Picks')
+
+    for name,_,_ in objects:
+        try:
+            coordinates = readers.coordinates(run, name=name, user_id=user_id1, session_id=session_id1)
+            close_points = coordinates[np.abs(coordinates[:, 0] - vol_slice) <= slice_proximity_threshold]
+            plt.scatter(close_points[:, 2], close_points[:, 1], label=name, s=15)    
+        except: 
+            pass
+
+    plt.subplot(1, 2, 2)
+    plt.imshow(tomo[vol_slice],cmap='gray')
+    plt.title(f'{user_id2} Picks')
+    
+    for name,_,_ in objects:
+        try:
+            coordinates = readers.coordinates(run, name=name, user_id=user_id2, session_id=session_id2)
+            close_points = coordinates[np.abs(coordinates[:, 0] - vol_slice) <= slice_proximity_threshold]
+            plt.scatter(close_points[:, 2], close_points[:, 1], label=name, s=15)
+        except:
+            pass
+
+    plt.axis('off')
+    plt.show()
+
+def plot_training_results(
+    results, 
+    class_names: Optional[List[str]] = None,
+    save_plot: str = None,
+    fig = None, axs = None):
+    """
+    Plot Training Results including Loss, Recall, Precision, and F1 Score.
+
+    Args:
+        results (dict): A dictionary containing training metrics.
+        class_names (list, optional): List of class names for labeling. Defaults to None.
+        save_plot (str): Path to save the plot image.
+        fig (matplotlib.figure.Figure, optional): Existing figure to plot on. Defaults to None.
+        axs (numpy.ndarray, optional): Existing axes to plot on. Defaults to None.
+    """
+
+    # Create a 2x2 subplot layout
+    if fig is None:
+        fig, axs = plt.subplots(2, 2, figsize=(12, 10))
+    else:
+        # Clear previos plots
+        for ax in axs.flatten():
+            ax.clear()
+
+    fig.suptitle("Metrics Over Epochs", fontsize=16)
+
+    # Unpack the data for loss (logged every epoch)
+    epochs_loss = [epoch for epoch, _ in results['loss']]
+    loss = [value for _, value in results['loss']]
+    val_epochs_loss = [epoch for epoch, _ in results['val_loss']]
+    val_loss = [value for _,value in results['val_loss']]
+
+    # Plot Training Loss in the top-left
+    axs[0, 0].plot(epochs_loss, loss, label="Training Loss")
+    axs[0, 0].plot(val_epochs_loss, val_loss, label='Validation Loss')
+    axs[0, 0].set_xlabel("Epochs")
+    axs[0, 0].set_ylabel("Loss")
+    axs[0, 0].set_title("Training Loss")
+    axs[0, 0].legend()
+    axs[0, 0].tick_params(axis='both', direction='in', top=True, right=True, length=6, width=1)
+
+    # For metrics that are logged every `val_interval` epochs
+    epochs_metrics = [epoch for epoch, _ in results['avg_recall']]
+    
+    # Determine the number of classes and names
+    num_classes = len([key for key in results.keys() if key.startswith('recall_class')])
+
+    if class_names is None or len(class_names) != num_classes:
+        class_names = [f"Class {i+1}" for i in range(num_classes)]
+
+    # Plot Recall in the top-right
+    for class_idx in range(num_classes):
+        recall_class = [value for _, value in results[f'recall_class{class_idx+1}']]
+        axs[0, 1].plot(epochs_metrics, recall_class, label=f"{class_names[class_idx]}")
+    axs[0, 1].set_xlabel("Epochs")
+    axs[0, 1].set_ylabel("Recall")
+    axs[0, 1].set_title("Recall per Class")
+    # axs[0, 1].legend()
+    axs[0, 1].tick_params(axis='both', direction='in', top=True, right=True, length=6, width=1)
+
+    # Plot Precision in the bottom-left
+    for class_idx in range(num_classes):
+        precision_class = [value for _, value in results[f'precision_class{class_idx+1}']]
+        axs[1, 0].plot(epochs_metrics, precision_class, label=f"{class_names[class_idx]}")
+    axs[1, 0].set_xlabel("Epochs")
+    axs[1, 0].set_ylabel("Precision")
+    axs[1, 0].set_title("Precision per Class")
+    axs[1, 0].legend()
+    axs[1, 0].tick_params(axis='both', direction='in', top=True, right=True, length=6, width=1)
+
+    # Plot F1 Score in the bottom-right
+    for class_idx in range(num_classes):
+        f1_class = [value for _, value in results[f'f1_class{class_idx+1}']]
+        axs[1, 1].plot(epochs_metrics, f1_class, label=f"{class_names[class_idx]}")
+    axs[1, 1].set_xlabel("Epochs")
+    axs[1, 1].set_ylabel("F1 Score")
+    axs[1, 1].set_title("F1 Score per Class")
+    # axs[1, 1].legend()
+    axs[1, 1].tick_params(axis='both', direction='in', top=True, right=True, length=6, width=1)
+
+    # Adjust layout and show plot
+    plt.tight_layout(rect=[0, 0, 1, 0.96])  # Leave space for the main title
+
+    fig.savefig(save_plot)
+    fig.canvas.draw()
+
+    return fig, axs

octopi/workflows.py

@@ -0,0 +1,262 @@
+from octopi.extract.localize import process_localization
+import octopi.processing.evaluate as octopi_evaluate
+from monai.metrics import ConfusionMatrixMetric
+from octopi.models import common as builder
+from octopi.pytorch import segmentation
+from octopi.pytorch import trainer 
+from octopi.utils import io
+import multiprocess as mp
+import copick, torch, os
+from tqdm import tqdm
+    
+def train(data_generator, loss_function, num_crops = 16,
+          model_config = None, model_weights = None, lr0 = 1e-3,
+          model_save_path = 'results', best_metric = 'fBeta2', 
+          num_epochs = 1000, use_ema = True, val_interval = 10,
+          sw_bs = 4, overlap = 0.5, ):
+    """
+    Train a UNet Model for Segmentation
+
+    Args:
+        config (str): Path to the Copick Config File
+        target_info (list): List containing the target user ID, target session ID, and target algorithm
+        tomo_algorithm (str): The tomographic algorithm to use for segmentation
+        voxel_size (float): The voxel size of the data
+        loss_function (str): The loss function to use for training
+        model_config (dict): The model configuration
+        model_weights (str): The path to the model weights
+        trainRunIDs (list): The list of run IDs to use for training
+        validateRunIDs (list): The list of run IDs to use for validation
+        model_save_path (str): The path to save the model
+        best_metric (str): The metric to use for early stopping
+        num_epochs (int): The number of epochs to train for
+        val_interval (int): The number of epoch intervals for validation during training
+        sw_bs (int): The sliding window batch size for validation
+        overlap (float): The overlap for sliding window inference during validation
+    """
+
+    # If No Model Configuration is Provided, Use the Default Configuration
+    if model_config is None:
+        model_config = {
+            'architecture': 'Unet',
+            'num_classes': data_generator.Nclasses,
+            'dim_in': 80,
+            'strides': [2, 2, 1],
+            'channels': [48, 64, 80, 80],
+            'dropout': 0.0, 'num_res_units': 1,
+        }
+        print('No Model Configuration Provided, Using Default Configuration')
+        print(model_config)
+    # extract the model config from full config dict
+    elif isinstance(model_config, dict) and 'model' in model_config: 
+        model_config = model_config['model']
+
+    # Monai Functions
+    metrics_function = ConfusionMatrixMetric(include_background=False, metric_name=["recall",'precision','f1 score'], reduction="none")
+
+    # Build the Model
+    model_builder = builder.get_model(model_config['architecture'])
+    model = model_builder.build_model(model_config)
+
+    # Load the Model Weights if Provided 
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    if model_weights: 
+        print(f'Loading Model Weights from: {model_weights}\n')
+        state_dict = torch.load(model_weights, map_location=device, weights_only=True)
+        model.load_state_dict(state_dict)     
+    model.to(device) 
+
+    # Optimizer
+    optimizer = torch.optim.AdamW(
+        model.parameters(), lr=lr0, weight_decay=1e-4
+    )
+
+    # Create UNet-Trainer
+    model_trainer = trainer.ModelTrainer(
+        model, device, loss_function, metrics_function, 
+        optimizer, use_ema = use_ema
+    )
+    model_trainer.sw_bs = sw_bs
+    model_trainer.overlap = overlap
+
+    # Train the Model
+    print(f'🔃 Starting Training...\nSaving Training Results to: {model_save_path}/\n')
+    results = model_trainer.train(
+        data_generator, model_save_path, max_epochs=num_epochs,
+        crop_size=model_config['dim_in'], my_num_samples=num_crops,
+        val_interval=val_interval, best_metric=best_metric, verbose=True
+    )
+    print('✅ Training Complete!')
+
+    # Save parameters and results
+    print(f'💾 Saving Training Parameters and Results to: {model_save_path}/\n')
+    parameters_save_name = os.path.join(model_save_path, "model_config.yaml")
+    io.save_parameters_to_yaml(model_builder, model_trainer, data_generator, parameters_save_name)
+
+    # TODO: Write Results to CSV...
+    results_save_name = os.path.join(model_save_path, "results.csv")
+    io.save_results_to_csv(results, results_save_name)
+
+def segment(config, tomo_algorithm, voxel_size, model_weights, model_config, 
+            seg_info = ['predict', 'octopi', '1'], use_tta = False, run_ids = None):
+    """
+    Segment a Dataset using a Trained Model or Ensemble of Models
+
+    Args:
+        config (str): Path to the Copick Config File
+        tomo_algorithm (str): The tomographic algorithm to use for segmentation
+        voxel_size (float): The voxel size of the data
+        model_weights (str, list): The path to the model weights or a list of paths to the model weights
+        model_config (str, list): The model configuration or a list of model configurations
+        seg_info (list): The segmentation information
+        use_tta (bool): Whether to use test time augmentation
+        run_ids (list): The list of run IDs to use for segmentation
+    """
+
+    # Initialize the Predictor
+    predict = segmentation.Predictor(
+        config, model_config, model_weights,
+        apply_tta = use_tta
+    )
+
+    # Run batch prediction
+    predict.batch_predict(
+        runIDs=run_ids,
+        num_tomos_per_batch=15,
+        tomo_algorithm=tomo_algorithm,
+        voxel_spacing=voxel_size,
+        segmentation_name=seg_info[0],
+        segmentation_user_id=seg_info[1],
+        segmentation_session_id=seg_info[2]
+    )
+
+    print('✅ Segmentation Complete!')
+
+def localize(config, voxel_size, seg_info, pick_user_id, pick_session_id, n_procs = 16,
+            method = 'watershed', filter_size = 10, radius_min_scale = 0.4, radius_max_scale = 1.0,
+            run_ids = None, pick_objects = None):
+    """
+    Extract 3D Coordinates from the Segmentation Maps
+
+    Args:
+        config (str): Path to the Copick Config File
+        voxel_size (float): The voxel size of the data
+        seg_info (list): The segmentation information
+        pick_user_id (str): The user ID of the pick
+        pick_session_id (str): The session ID of the pick
+        n_procs (int): The number of processes to use for parallelization
+        method (str): The method to use for localization
+        filter_size (int): The filter size to use for localization
+        radius_min_scale (float): The minimum radius scale to use for localization
+        radius_max_scale (float): The maximum radius scale to use for localization
+        run_ids (list): The list of run IDs to use for localization
+    """
+    
+    # Load the Copick Config
+    root = copick.from_file(config) 
+
+    # Get objects that can be Picked build into mutable rows
+    objects = [[obj.name, int(obj.label), float(obj.radius)]
+            for obj in root.pickable_objects if obj.is_particle]
+
+    # Verify each object has the required attributes
+    for obj in objects:
+        if len(obj) < 3 or not isinstance(obj[2], (float, int)):
+            raise ValueError(f"Invalid object format: {obj}. Expected a tuple with (name, label, radius).")
+  
+    # Load the Model Output Configuration
+    seg_config = io.get_config(config, seg_info[0], 'segment', seg_info[1], seg_info[2])
+
+    # sync labels from the model config and remove objects not in model labels
+    label_map = seg_config.get('labels', {})
+    for row in objects.copy():  # avoid modifying the list while iterating
+        name, label, radius = row
+        if name in label_map and label != label_map[name]:
+            row[1] = int(label_map[name])  # mutate in place
+        elif name not in label_map: # remove this entry from objects 
+            objects.remove(row)
+
+    # Filter objects based on the provided list
+    if pick_objects is not None:
+        objects0 = objects.copy()  # avoid modifying the list while iterating for error tracking
+        objects = [obj for obj in objects if obj[0] in pick_objects]
+        if len(objects) == 0:
+            raise ValueError(f"No valid objects found for localization after filtering. Mismatched names: {pick_objects} and {[obj[0] for obj in objects0]}")
+
+    # Get all RunIDs
+    if run_ids is None:
+        run_ids = [run.name for run in root.runs]
+    n_run_ids = len(run_ids)
+
+    # Exit if No Runs are Available
+    if n_run_ids == 0:
+        print(f"No runs available for localization with the specified voxel size - {voxel_size}.\nExiting...")
+        return
+
+     # Run Localization - Main Parallelization Loop
+    print(f"Using {n_procs} processes to parallelize across {n_run_ids} run IDs.")
+    with mp.Pool(processes=n_procs) as pool:
+        with tqdm(total=n_run_ids, desc="Localization", unit="run") as pbar:
+            worker_func = lambda run_id: process_localization(
+                root.get_run(run_id),  
+                objects, 
+                seg_info,
+                method, 
+                voxel_size,
+                filter_size,
+                radius_min_scale, 
+                radius_max_scale,
+                pick_session_id,
+                pick_user_id
+            )
+
+            for _ in pool.imap_unordered(worker_func, run_ids, chunksize=1):
+                pbar.update(1)
+
+    print('✅ Localization Complete!')
+    
+
+def evaluate(config, 
+             gt_user_id, gt_session_id,
+             pred_user_id, pred_session_id,
+             run_ids = None, distance_threshold = 0.5, save_path = None):
+    """
+    Evaluate the Localization on a Dataset
+
+    Args:
+        config (str): Path to the Copick Config File
+        gt_user_id (str): The user ID of the ground truth
+        gt_session_id (str): The session ID of the ground truth
+        pred_user_id (str): The user ID of the predicted coordinates
+        pred_session_id (str): The session ID of the predicted coordinates
+        run_ids (list): The list of run IDs to use for evaluation
+        distance_threshold (float): The distance threshold to use for evaluation
+        save_path (str): The path to save the evaluation results
+    """
+             
+    print('Running Evaluation on the Following Query:')
+    print(f'Distance Threshold: {distance_threshold}')
+    print(f'GT User ID: {gt_user_id}, GT Session ID: {gt_session_id}')
+    print(f'Pred User ID: {pred_user_id}, Pred Session ID: {pred_session_id}')
+    print(f'Run IDs: {run_ids}')
+    
+    # Load the Copick Config
+    root = copick.from_file(config) 
+
+    # For Now Lets Assume Object Names are None..
+    object_names = None
+
+    # Run Evaluation
+    eval = octopi_evaluate.evaluator(
+        config,
+        gt_user_id,
+        gt_session_id,
+        pred_user_id,
+        pred_session_id, 
+        object_names=object_names
+    )
+
+    eval.run(
+        distance_threshold_scale=distance_threshold, 
+        runIDs=run_ids, save_path=save_path
+    )