Two-process model of sleep

[raphaelvallat]

A crude translation R --> Python from https://github.com/PabRod/sleepR/blob/master/R/borbely.R I think the code could be optimized. Worth adding to a new yasa.modeling module?

def circ(t):
    """Circadian process."""
    return np.sin(params["w"] * t - params["alpha"])

def borbely(ts, y0, first_awake=False):
    """Two-process model."""    
    # Helper functions
    def Hs(t, H0):
        """Sleep"""
        return H0 * np.exp(-t / params['Xis'])

    def Ha(t, H0):
        """Awake"""
        return params['mu'] + (H0 - params['mu']) * np.exp(-t / params['Xiw'])

    def Hc(t, H0, awake):
        """Sleep and awake"""
        asleep = not awake
        return awake * Ha(t, H0) + asleep * Hs(t, H0)

    upper = params["Hu0"] + params['a'] * circ(ts)
    lower = params['Hl0'] + params['a'] * circ(ts)

    # Initialize variables
    n_hours = len(ts)
    H = np.empty(n_hours)
    awake = np.empty(n_hours, dtype=bool)
    H[0] = y0
    awake[0] = first_awake

    # Loop over hours
    for i in np.arange(1, n_hours):
        # Update homeostatic pressure
        H[i] = Hc(ts[i] - ts[i-1], H[i-1], awake[i-1])
        # Set triggers
        asleep = not awake[i-1]
        wake_trigger = asleep and (H[i] <= lower[i])
        sleep_trigger = (awake[i-1]) and (H[i] >= upper[i])

        # Update awake/sleeping status
        if wake_trigger:
            awake[i] = True
        elif sleep_trigger:
            awake[i] = False
        else:
            awake[i] = awake[i-1]

    # Convert to dataframe
    sol = pd.DataFrame({"time": ts, "H": H, "awake": awake, "lower": lower})
    return sol

# Define parameters
params = dict(
    Xis = 4, # hours
    Xiw = 18.2, # hours
    mu = 1, # 1
    a = 0.1, # 0.1
    w = 2 * np.pi / 24, # h^-1
    alpha = 0, # h
    Hu0 = 0.6, # 1
    Hl0 = 0.17) # 1

# Compute process C and S
hours = np.arange(48)
df_out = borbely(hours, y0=params["Hl0"], first_awake=True)
df_out.head()

which roughly gives the following plot (plotting H (S) and lower (C) against time):

[raphaelvallat]

Closing this, as I feel that this is out of the scope of YASA. Feel free to reopen!