beta_coef(self, t) returns incorrect values

[arik-shurygin]

due to complex indexing logic it can be hard for a user to understand exactly which parameters will return which coefficients

Current Behavior

BETA_COEFICIENTS = jnp.array([-1.0, 0.0, 1.0])
BETA_TIMES = jnp.array([50, 125, 150])
def test_buggy(t):
    # copy pasted from AbstractParameters.beta_coef(self, t)
    return BETA_COEFICIENTS[jnp.maximum(0, jnp.searchsorted(BETA_TIMES, t) - 1)]
times_of_interest = jnp.array([0, 49, 50, 51, 124, 125, 126, 149, 150, 151])
print(">>> self.beta_coef(t=%s) \n %s"%(times_of_interest.tolist(), test_buggy(times_of_interest)))

>>> self.beta_coef(t=[0, 25, 49, 50, 51, 124, 125, 126, 149, 150, 151]) 
[-1.  -1. -1. -1. -1. -1.  0.  0.  0.  1.]

Proposed Solution

BETA_COEFICIENTS = jnp.array([-1.0, 0.0, 1.0])
BETA_TIMES = jnp.array([125, 150])
def test_proposed(t):
    return BETA_COEFICIENTS[jnp.maximum(0, jnp.searchsorted(BETA_TIMES, t, side='right'))]
times_of_interest = jnp.array([0, 49, 50, 51, 124, 125, 126, 149, 150, 151])
print(">>> self.beta_coef(t=%s) \n %s"%(times_of_interest.tolist(), test_proposed(times_of_interest)))

>>> self.beta_coef(t=[0, 49, 50, 51, 124, 125, 126, 149, 150, 151]) 
[-1. -1. -1. -1. -1.  0.  0.  0.  1.  1.]

This includes more sane boundaries on the times listed in BETA_TIMES, furthermore BETA_COEFICIENTS being 1 element longer than BETA_TIMES more clearly indicates that the first and last coefficient bound t = [-inf, BETA_TIMES[0] ) and t = [BETA_TIMES[-1], inf ) respectively.