example.py

#!/usr/bin/env python
# example.py - rcampbel 2024-11-01
# Notional code to act as example for ensemble project

# Comes with Python
import sys
import argparse
import logging

# Added by environment
import numpy as np
import matplotlib.pyplot as plt

# List of selectable parameter sets
ALL_PARAMS = [[23, 45, 67],
              [98, 102, 7],
              [1, 2, 3]]

def run(task_id):
    """Perform calculations, create plot."""
    logging.info('Run')

    # Extract parameters using task_id as index
    # NOTE Another option: use task_id to create random values
    params = ALL_PARAMS[task_id]

    # Calculate
    output = np.array(params) * 2 + 1

    # Create bar plot
    plt.figure(figsize=(6, 4))
    plt.bar(params, output, color=['skyblue', 'salmon', 'lightgreen'])
    plt.xlabel('Parameters')
    plt.ylabel('Output')
    plt.title('Example Plot with Three Values')
    plt.savefig(f'plot-{task_id}.png', dpi=300)


if __name__ == "__main__":

    # Command line arugments
    parser = argparse.ArgumentParser(description='Example')
    parser.add_argument('--task_id', type=int, default=None, help="Current job's index value in job array (usually $SLURM_ARRAY_TASK_ID)")
    args = parser.parse_args()

    # Set up logging
    logging.basicConfig(format=f'%(levelname)s: %(message)s (%(filename)s:%(lineno)d, %(asctime)s, task:{args.task_id})',
                        stream=sys.stdout,
                        level=logging.INFO),

    # Suppress most log items from plot module
    logging.getLogger('matplotlib.font_manager').setLevel(logging.WARNING)  # Suppress debug entries for "...findfont..."

    logging.info('Start')

    # Run our logic
    run(int(args.task_id))

    logging.info('End')