Simple Gantt chart

gantt_chart.py

"""
Simple Gantt chart implementation
"""
from datetime import datetime as dt
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.dates import date2num
import matplotlib.dates as mdates
import numpy

matplotlib.rcParams.update({'font.size': 14})

default_color = 'IndianRed'


def create_gantt_chart(task_list):

    fig = plt.figure(figsize=(10, 4.5))
    ax = fig.add_subplot(111)

    ntasks = len(task_list)
    yticks = numpy.arange(ntasks)
    for y, task in zip(yticks, task_list):
        st = task['start_date']
        et = task['end_date']
        if st is None or et is None:
            continue
        st = date2num(st)
        et = date2num(et)
        color = task.get('color', default_color)
        duration = et - st
        ax.barh(y, duration, left=st, height=0.8,
                align='center', edgecolor='none', color=color)

    labels = [t['name'] for t in task_list]
    ax.set_yticks(yticks)
    ax.set_yticklabels(labels)
    ax.xaxis_date()
    ax.invert_yaxis()

    major_locator = mdates.YearLocator()
    minor_locator = mdates.MonthLocator([1, 4, 7, 10])

    ax.xaxis.set_major_locator(major_locator)
    ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y'))
    ax.xaxis.set_minor_locator(minor_locator)
    ax.grid()
    fig.autofmt_xdate()

    return fig, ax


task_list = [
    {
        'name': '1.1: 2D shallow water solver',
        'start_date': dt(2022, 9, 1),
        'end_date': dt(2023, 5, 30),
        'color': '#FFA07A',
    },
    {
        'name': '1.2: Hybridized version',
        'start_date': dt(2023, 6, 1),
        'end_date': dt(2024, 2, 28),
        'color': '#FFA07A',
    },
    {
        'name': '2.1: Energy conserving 2D solver',
        'start_date': dt(2024, 3, 1),
        'end_date': dt(2024, 12, 31),
        'color': '#87CEFA',
    },
    {
        'name': '2.2: Hybridized version',
        'start_date': dt(2025, 1, 1),
        'end_date': dt(2025, 12, 31),
        'color': '#87CEFA',
    },
    {
        'name': '2.3: Numerical analysis',
        'start_date': dt(2026, 1, 1),
        'end_date': dt(2026, 8, 31),
        'color': '#87CEFA',
    },
    {
        'name': '3.1: Explicit wetting-drying',
        'start_date': dt(2022, 10, 1),
        'end_date': dt(2023, 8, 31),
        'color': '#90EE90',
    },
    {
        'name': '3.2: Implicit wetting-drying',
        'start_date': dt(2023, 9, 1),
        'end_date': dt(2024, 7, 31),
        'color': '#90EE90',
    },
    {
        'name': '4.1: 3D formulation',
        'start_date': dt(2022, 9, 1),
        'end_date': dt(2024, 1, 31),
        'color': 'IndianRed',
    },
    {
        'name': '4.2: Regrid-remap',
        'start_date': dt(2023, 12, 1),
        'end_date': dt(2024, 11, 30),
        'color': 'IndianRed',
    },
    {
        'name': '4.3: Vertical diffusion',
        'start_date': dt(2024, 12, 1),
        'end_date': dt(2025, 8, 31),
        'color': 'IndianRed',
    },
    {
        'name': '4.2: Numerical analysis',
        'start_date': dt(2025, 9, 1),
        'end_date': dt(2026, 8, 31),
        'color': 'IndianRed',
    },
]

fig, ax = create_gantt_chart(task_list)

# project start/end dates
ax.axvline(dt(2022, 9, 1), color='k', linewidth=2.0)
ax.axvline(dt(2026, 8, 31), color='k', linewidth=2.0)
ax.set_xlim([dt(2022, 8, 1), dt(2026, 10, 1)])

imgfile = 'gantt_chart.png'
print('Saving image {:}'.format(imgfile))
plt.savefig(imgfile, dpi=200, bbox_inches='tight')