WeekDayConstrainedDist for HCiw

WeekDayConstrainedDist.py

class WeekDayConstrainedDist(ciw.dists.Distribution):

    def __init__(self, dist, offset=0, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.dist = dist
        self.offset = offset

    def sample(self, t=None, ind=None):
        tau = (np.floor(t) - self.offset + 7) % 7
        if tau < 5:
            return self.dist.sample(t, ind)
        elif tau == 5:
            return 2
        else:
            return 1