prophet/python/fbprophet/make_holidays.py
ziye666 95fa4460a9 Fix bug about holiday year different from input year and others (#670)
* Allow both both hoidays and append holidays

Match holidays in predict and fit

Add test for append_holiday features; minor fixes

Add column name validation for append_holidays names; allow only one country

Fix bug about holiday year different from input year and also other bugs

Change function description

Add append holiday feature for R

* Add test for R/Pyrhon; fix bugs
2018-10-18 17:40:17 -07:00

69 lines
2.1 KiB
Python

# Copyright (c) 2017-present, Facebook, Inc.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree. An additional grant
# of patent rights can be found in the PATENTS file in the same directory.
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals
import pandas as pd
import numpy as np
import warnings
import holidays as hdays_part1
import fbprophet.hdays as hdays_part2
def get_holiday_names(country):
"""Return all possible holiday names of given country
Parameters
----------
country: country name
Returns
-------
A set of all possible holiday names of given country
"""
years = np.arange(1995, 2045)
try:
with warnings.catch_warnings():
warnings.simplefilter("ignore")
holiday_names = getattr(hdays_part2, country)(years=years).values()
except AttributeError:
try:
holiday_names = getattr(hdays_part1, country)(years=years).values()
except AttributeError:
raise AttributeError(
"Holidays in {} are not currently supported!".format(country))
return set(holiday_names)
def make_holidays_df(year_list, country):
"""Make dataframe of holidays for given years and countries
Parameters
----------
year_list: a list of years
country: country name
Returns
-------
Dataframe with 'ds' and 'holiday', which can directly feed
to 'holidays' params in Prophet
"""
try:
holidays = getattr(hdays_part2, country)(years=year_list)
except AttributeError:
try:
holidays = getattr(hdays_part1, country)(years=year_list)
except AttributeError:
raise AttributeError(
"Holidays in {} are not currently supported!".format(country))
holidays_df = pd.DataFrame(list(holidays.items()), columns=['ds', 'holiday'])
holidays_df.reset_index(inplace=True, drop=True)
holidays_df['ds'] = pd.to_datetime(holidays_df['ds'])
return (holidays_df)