def condition(state):
# rolling average
x_cumsum_window = np.convolve(np.abs(state.info_loss), window,
"valid")
n_zeros = int(np.sum(np.where(x_cumsum_window > 0.0, 0, 1)))
return jax.lax.ne(n_zeros, 1) or state.ilayer > state.max_layers
def moving_average(x, o): """ x: array to take moving average og o: # of days (entries) on either side of the current value to average over """ w=o*2+1 # width of window to average over, current day in center of window y=np.convolve(x, np.ones(w), 'full') den=np.concatenate((np.arange(o+1,w),w*np.ones(len(x)-w+1),np.arange(w-1,o,step=-1))) z=y[o:-o]/den return z
def get_minimum_zeroth_element(x: Array, window_size: int = 10) -> int:
# window for the convolution
window = np.ones(window_size) / window_size
# rolling average
x_cumsum_window = np.convolve(np.abs(x), window, "valid")
# get minimum zeroth element
min_idx = int(np.min(np.argwhere(x_cumsum_window == 0.0)[0]))
return min_idx
def info_loss_condition(state):
# rolling average of absolute value
x_cumsum_window = jnp.convolve(jnp.abs(state.info_loss), window,
"valid")
# count number of zeros in moving window
n_zeros = jnp.sum(jnp.where(x_cumsum_window > 0.0, 0,
1)).astype(jnp.int32)
return jnp.logical_and(
jax.lax.ne(n_zeros, jnp.array([1.0], dtype=jnp.int32)),
state.ilayer < state.max_layers,
)
def conv2d(data, filt):
# data_s=np.empty(np.shape(data))
# nr=data.shape[1]
with loops.Scope() as s:
s.data_s = np.empty(np.shape(data))
for r in s.range(data.shape[1]):
# data_s[:,r] = np.convolve(data[:,r], filt, 'same')
#s.data_s = jax.ops.index_update(s.data_s, jax.ops.index[:,r], np.convolve(data[:,r], filt, 'same'))
s.data_s = s.data_s.at[:,
r].set(np.convolve(data[:, r], filt,
'same'))
return s.data_s
def ipgauss2(nus, F0, varr_kernel, beta):
"""Apply the Gaussian IP response to a spectrum F.
Args:
nus: input wavenumber, evenly log-spaced
F0: original spectrum (F0)
varr_kernel: velocity array for the rotational kernel
beta: STD of a Gaussian broadening (IP+microturbulence)
Return:
response-applied spectrum (F)
"""
x = varr_kernel / beta
kernel = jnp.exp(-x * x / 2.0)
kernel = kernel / jnp.sum(kernel, axis=0)
F = jnp.convolve(F0, kernel, mode='same')
return F
def rigidrot2(nus, F0, varr_kernel, vsini, u1=0.0, u2=0.0):
"""Apply the Rotation response to a spectrum F using jax.lax.scan.
Args:
nus: wavenumber, evenly log-spaced
F0: original spectrum (F0)
varr_kernel: velocity array for the rotational kernel
vsini: V sini for rotation
beta: STD of a Gaussian broadening (IP+microturbulence)
RV: radial velocity
u1: Limb-darkening coefficient 1
u2: Limb-darkening coefficient 2
Return:
response-applied spectrum (F)
"""
x = varr_kernel / vsini
x2 = x * x
kernel = rotkernel(x, u1, u2)
kernel = kernel / jnp.sum(kernel, axis=0)
F = jnp.convolve(F0, kernel, mode='same')
return F
def test_cuDNN():
# convolution requires cuDNN
v = jnp.array([1, 9, 1])
s = jnp.array([0, 0, 0, 0, 1, 2, 0, 0, 0, 2, 0, 0, 0, 0, 0, 1])
c = jnp.convolve(s, v, mode='same')
res = c - jnp.array([
0.,
0.,
0.,
1.,
11.,
19.,
2.,
0.,
2.,
18.,
2.,
0.,
0.,
0.,
1.,
9.,
])
assert jnp.sum(res**2) == 0.0
def convolve(a, v, mode='full'): if isinstance(a, JaxArray): a = a.value if isinstance(v, JaxArray): v = v.value return JaxArray(jnp.convolve(a, v, mode))
def tick_buffer(carry, X):
params, _ = carry
B = params["B"]
return carry, jnp.convolve(X, B)[:-(B.size - 1)]