Is there a better way to create a water year datetime column in Python?

Viewed 198

For one of my projects I need to change from a calendar year to water year and have it in the YYYY-MM-DD format. A common definition of a water year is October 1 - September 30th, which is what I am going for. I have drafted some code that works, but I feel like there is a better way to approach this and I wanted to see if anyone has any suggestions.

I start out by creating the monthly date range I am interested in:

import pandas as pd
date = pd.DataFrame(pd.date_range('1915-01-01', '2011-12-31', freq='MS'))
date.columns = ['Date']
Date = date['Date']
date_set = date.set_index(Date)
print(date_set)

Which prints the following:

    Date
0    1915-01-01
1    1915-02-01
2    1915-03-01
3    1915-04-01
4    1915-05-01
        ...
1159 2011-08-01
1160 2011-09-01
1161 2011-10-01
1162 2011-11-01
1163 2011-12-01

Now use a function to set the WY, where the months October through December get +1 added to their year:

def convert_to_WY(row):
    if row['Date'].month>=10:
        return(pd.datetime(row['Date'].year+1,1,1).year)
    else:
        return(pd.datetime(row['Date'].year,1,1).year)

date_set['WaterYear'] = date_set.apply(lambda x: convert_to_WY(x), axis=1)

This works and gives me a column called "WaterYear" where it lists the year based on the conditions specified in the function. I want the WaterYear column to have the format YYYY-MM-DD and the way I went about this is shown below:

date_set['Month'] = pd.DatetimeIndex(date_set['Date']).month
date_set['Day'] = pd.DatetimeIndex(date_set['Date']).day
date_set['WaterYear_F'] = date_set.apply(lambda row: datetime(row['WaterYear'], row['Month'], row['Day']), axis=1)
date_set.drop(['Month','Day','WaterYear','Date'], axis=1, inplace = True)
print(date_set.head(15))

After printing the result I got was what I wanted:

           WaterYear_F
Date                  
1915-01-01  1915-01-01
1915-02-01  1915-02-01
1915-03-01  1915-03-01
1915-04-01  1915-04-01
1915-05-01  1915-05-01
1915-06-01  1915-06-01
1915-07-01  1915-07-01
1915-08-01  1915-08-01
1915-09-01  1915-09-01
1915-10-01  1916-10-01
1915-11-01  1916-11-01
1915-12-01  1916-12-01
1916-01-01  1916-01-01
1916-02-01  1916-02-01
1916-03-01  1916-03-01

I just feel like this is a very roundabout way to do this and am looking for any tips or suggestions to clean this up some.

0 Answers
Related