Merge pull request statsmodels#1483 from jseabold/tsa-tools

ENH: Some utility functions for working with dates

statsmodels/tsa/base/datetools.py

@@ -20,6 +20,18 @@ def __getitem__(self, key):
                        'A' : pandas_datetools.yearEnd,
                        'Q' : pandas_datetools.quarterEnd}
 
+try:
+    from pandas import Period
+except ImportError:  # not sure when this was added in pandas
+    Period = datetime.datetime  # HACK
+
+
+def _is_datetime_index(dates):
+    if isinstance(dates[0], (datetime.datetime, Period)):
+        return True  # TimeStamp is a datetime subclass
+    else:
+        return False
+
 def _index_date(date, dates):
     """
     Gets the index number of a date in a date index.

statsmodels/tsa/base/tsa_model.py

@@ -6,10 +6,6 @@
 from pandas import Index
 from pandas import datetools as pandas_datetools
 import datetime
-try:
-    from pandas import Period
-except ImportError: # not sure what version this was added
-    Period = datetime.datetime # HACK
 
 _freq_to_pandas = datetools._freq_to_pandas
 
@@ -49,7 +45,7 @@ def _init_dates(self, dates, freq):
             dates = self.data.row_labels
 
         if dates is not None:
-            if (not isinstance(dates[0], (datetime.datetime, Period)) and
+            if (not datetools._is_datetime_index(dates) and
                     isinstance(self.data, data.PandasData)):
                 raise ValueError("Given a pandas object and the index does "
                                  "not contain dates")

statsmodels/tsa/tests/test_tsa_tools.py

@@ -196,6 +196,13 @@ def test_add_lag_drop_noinsert():
     lag_data = sm.tsa.add_lag(data, 'realgdp', 3, insert=False, drop=True)
     assert_equal(lag_data.view((float,len(lag_data.dtype.names))), results)
 
+def test_freq_to_period():
+    from pandas.tseries.frequencies import to_offset
+    freqs = ['A', 'AS-MAR', 'Q', 'QS', 'QS-APR', 'W', 'W-MON', 'B']
+    expected = [1, 1, 4, 4, 4, 52, 52, 52]
+    for i, j in zip(freqs, expected):
+        assert_equal(tools.freq_to_period(i), j)
+        assert_equal(tools.freq_to_period(to_offset(i)), j)
 
 if __name__ == '__main__':
     #running them directly

statsmodels/tsa/tsatools.py

@@ -1,6 +1,9 @@
 import numpy as np
 import numpy.lib.recfunctions as nprf
 from statsmodels.tools.tools import add_constant
+from pandas.tseries import offsets
+from pandas.tseries.frequencies import to_offset
+
 
 def add_trend(X, trend="c", prepend=False):
     """
@@ -570,6 +573,42 @@ def unintegrate(x, levels):
     x0 = levels[0]
     return np.cumsum(np.r_[x0, x])
 
+
+def freq_to_period(freq):
+    """
+    Convert a pandas frequency to a periodicity
+
+    Parameters
+    ----------
+    freq : str or offset
+        Frequency to convert
+
+    Returns
+    -------
+    period : int
+        Periodicity of freq
+
+    Notes
+    -----
+    Annual maps to 1, quarterly maps to 4, monthly to 12, weekly to 52.
+    """
+    if not isinstance(freq, offsets.DateOffset):
+        freq = to_offset(freq)  # go ahead and standardize
+    freq = freq.rule_code.upper()
+
+    if freq == 'A' or freq.startswith(('A-', 'AS-')):
+        return 1
+    elif freq == 'Q' or freq.startswith(('Q-', 'QS-')):
+        return 4
+    elif freq == 'M' or freq.startswith(('M-', 'MS-')):
+        return 12
+    elif freq == 'B' or freq == 'W' or freq.startswith('W-'):
+        return 52
+    else:  # pragma : no cover
+        raise ValueError("freq {} not understood. Please report if you "
+                         "think this in error.".format(freq))
+
+
 __all__ = ['lagmat', 'lagmat2ds','add_trend', 'duplication_matrix',
            'elimination_matrix', 'commutation_matrix',
            'vec', 'vech', 'unvec', 'unvech']