Trendlevel: TypeError 'Axes' object is not subscriptable

[Tanvi-Jain01]

@nipunbatra @patel-zeel
This dataset mydata.csv is already present in the doc folder of vayu.

Code:

file_path = r'C:\..\mydata.csv'

# Read the CSV file into a DataFrame
df = pd.read_csv(file_path)

#df.reset_index(inplace=True)
df['date'] = pd.to_datetime(df['date'])
print(df)

df_2003 = df[df['date'].dt.year == 2003]
print(df_2003)

from vayu.trendLevel import trendLevel
trendLevel(df_2003, 'pm25')

Error:

TypeError                                 Traceback (most recent call last)
Cell In[59], line 2
      1 from vayu.trendLevel import trendLevel
----> 2 trendLevel(df_2003, 'pm25')

File ~\anaconda3\lib\site-packages\vayu\trendLevel.py:45, in trendLevel(df, pollutant, **kwargs)
     34 t = pollutant_series_year.groupby(
     35     [pollutant_series_year.index.month, pollutant_series_year.index.hour]
     36 ).mean()
     37 two_d_array = t.values.reshape(12, 24).T
     38 sns.heatmap(
     39     two_d_array,
     40     cbar=True,
     41     linewidth=0,
     42     cmap="Spectral_r",
     43     vmin=0,
     44     vmax=400,
---> 45     ax=ax[i],
     46 )
     47 ax[i].set_title(year_string)
     48 ax[i].invert_yaxis()

TypeError: 'Axes' object is not subscriptable

Source Code:

vayu/vayu/trendLevel.py

Lines 48 to 53 in ef99aef

	ax=ax[i],
	)
	ax[i].set_xticklabels(months)
	ax[i].set_ylabel("Hour of the Day")
	ax[i].set_title(year_string)
	ax[i].invert_yaxis()

Explaination:

The error occurred when there was only one unique year because of the way the ax object was handled in the code.

In the original code, when there was only one unique year, ax was a single Axes object, not an array of Axes objects. However, the code was still trying to access ax[i] inside the loop, assuming it was an array. This caused the TypeError because a single Axes object is not subscriptable like an array.

Solution:

By using the condition ax[i] if num_unique_years > 1 else ax to assign the ax object to heatmap_ax, we ensure that heatmap_ax is always a valid Axes object, whether it's a single object or an array. This avoids the TypeError and allows the code to work correctly for both cases.

Improved Code:

def trendLevel(df, pollutant, **kwargs):
    """
    Plot that shows the overall pollutant trend for every year in the 
    df. It takes the average hour value of each month and plots a heatmap
    showing what times of the year there is a high concentration of the 
    pollutant.

    Parameters
    ----------
    df: pd.DataFrame
        Data frame of complete data
    pollutant: str
        Name of the data series in df to produce plot

    Returns
    -------
    None
    """
    import matplotlib.pyplot as plt
    import seaborn as sns

    df.index = pd.to_datetime(df.date)
    pollutant_series = df[pollutant]
    unique_years = np.unique(df.index.year)
    num_unique_years = len(unique_years)
    fig, ax = plt.subplots(nrows=num_unique_years, figsize=(20, 20))

    months = ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul',
              'Aug', 'Sep', 'Oct', 'Nov', 'Dec']


    for i, year in enumerate(unique_years):
       
        year_string = str(year)
        pollutant_series_year = pollutant_series[year_string]
        t = pollutant_series_year.groupby(
            [pollutant_series_year.index.month, pollutant_series_year.index.hour]
        ).mean()
        two_d_array = t.values.reshape(12, 24).T
        heatmap_ax = ax[i] if num_unique_years > 1 else ax
        sns.heatmap(
            two_d_array,
            cbar=True,
            linewidth=0,
            cmap="Spectral_r",
            vmin=0,
            vmax=400,
            ax=heatmap_ax
        )
        heatmap_ax.set_xticklabels(months)
        heatmap_ax.set_ylabel("Hour of the Day")
        heatmap_ax.set_title(year_string)
        heatmap_ax.invert_yaxis()
        
    plt.savefig("Trendlevelplot.png", bbox_inches="tight",dpi=300)
    print("Your plots has also been saved")
    plt.show()  # Display the plot

trendLevel(df, 'pm25')

Output: