pull_geothermal_data.py

import eia
import pandas as pd
import matplotlib.pyplot as plt

def retrieve_time_series(api, series_ID):
    """
    Return the time series dataframe, based on API and unique Series ID
    """
    #Retrieve Data By Series ID 
    series_search = api.data_by_series(series=series_ID)
    ##Create a pandas dataframe from the retrieved time series
    df = pd.DataFrame(series_search)
    return df

def plot_data(df, x_variable, y_variable, title):
    """
    Plot the x- and y- variables against each other, where the variables are columns in
    a pandas dataframe
    Args:
        df: Pandas dataframe. 
        x_variable: String. Name of x-variable column
        y_variable: String. Name of y-variable column
        title: String. Desired title name
    """
    fig, ax = plt.subplots()
    ax.plot_date(df[x_variable], 
                 df[y_variable], marker='', linestyle='-', label=y_variable)
    fig.autofmt_xdate()
    plt.title(title)
    plt.show()

#### EXECUTE IN MAIN FUNCTION ####
#Create EIA API using your specific API key
api_key = "YOUR EIA API KEY HERE" ###SIGN UP ON THE EIA WEBSITE TO GET A FREE API KEY
api = eia.API(api_key)
#Declare desired series ID
series_ID='TOTAL.GEEGPUS.M'
df = retrieve_time_series(api, series_ID)
df.reset_index(level=0, inplace=True)
df.rename(columns={'index':'Date',
                   df.columns[1]:'Geothermal_net_generation'}, inplace=True)
#Convert the Date column into a date object
df['Date'] = df['Date'].str.rstrip()
df['Date'] = df['Date'].str.replace(' ', '-')
df['Date']=pd.to_datetime(df['Date'], format='%Y-%m')
#Plot the time series
plot_data(df, 'Date', 
          'Geothermal_net_generation', 
          'Net Generation for Geothermal over Time')