Returns functions to map from ordinal to biz day and back.
(weekend_mask=None, holidays=None)
| 21 | |
| 22 | |
| 23 | def business_day_mappers(weekend_mask=None, holidays=None): |
| 24 | """Returns functions to map from ordinal to biz day and back.""" |
| 25 | if weekend_mask is None and holidays is None: |
| 26 | return (lambda x: (x, tf.ones_like(x, dtype=tf.bool))), (lambda x: x) |
| 27 | |
| 28 | weekday_fwd, weekday_back = _week_day_mappers(weekend_mask) |
| 29 | |
| 30 | if holidays is None: |
| 31 | return weekday_fwd, weekday_back |
| 32 | |
| 33 | # Apply the weekend adjustment to the holidays as well |
| 34 | holidays_raw = tf.convert_to_tensor(holidays, dtype=tf.int32) |
| 35 | holidays, is_weekday = weekday_fwd(holidays_raw) |
| 36 | |
| 37 | # Keep only the holidays that are not on weekends |
| 38 | holidays = holidays[is_weekday] |
| 39 | |
| 40 | # The above step can lead to an empty holidays set which causes problems. |
| 41 | # To mitigate this, we add a safe fake holiday. |
| 42 | holidays = tf.concat([[0], holidays], axis=0) |
| 43 | reverse_holidays = tf.reverse(-holidays, axis=[0]) |
| 44 | num_holidays = tf.size(holidays) - 1 |
| 45 | |
| 46 | def bizday_fwd(x): |
| 47 | """Calculates business day ordinal and whether it is a business day.""" |
| 48 | left = tf.searchsorted(holidays, x, side='left') |
| 49 | right = num_holidays - tf.searchsorted(reverse_holidays, -x, side='left') |
| 50 | is_bizday = tf.not_equal(left, right) |
| 51 | bizday_ordinal = x - right |
| 52 | return bizday_ordinal, is_bizday |
| 53 | |
| 54 | cum_holidays = tf.range(num_holidays + 1, dtype=holidays.dtype) |
| 55 | bizday_at_holidays = holidays - cum_holidays |
| 56 | |
| 57 | def bizday_back(x): |
| 58 | left = tf.searchsorted(bizday_at_holidays, x, side='left') |
| 59 | ordinal = x + left - 1 |
| 60 | return ordinal |
| 61 | |
| 62 | def from_ordinal(ordinals): |
| 63 | """Maps ordinals to business day and whether it is a work day.""" |
| 64 | ordinals = tf.convert_to_tensor(ordinals, dtype=tf.int32) |
| 65 | weekday_values, is_weekday = weekday_fwd(ordinals) |
| 66 | biz_ordinal, is_bizday = bizday_fwd(weekday_values) |
| 67 | return biz_ordinal, (is_weekday & is_bizday) |
| 68 | |
| 69 | def to_ordinal(biz_values): |
| 70 | """Maps from business day count to ordinals.""" |
| 71 | return weekday_back(bizday_back(biz_values)) |
| 72 | |
| 73 | return from_ordinal, to_ordinal |
| 74 | |
| 75 | |
| 76 | def _week_day_mappers(weekend_mask): |
nothing calls this directly
no test coverage detected