File size: 1,238 Bytes
eff9b8d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 | from __future__ import annotations
import numpy as np
import torch
def irregular_batch(
batch: int,
length: int,
seed: int,
*,
large_gaps: bool,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
rng = np.random.default_rng(seed)
low, high = ((0.12, 0.40) if large_gaps else (0.02, 0.12))
delta_time = rng.uniform(low, high, size=(batch, length))
time = np.cumsum(delta_time, axis=1)
amplitudes = rng.uniform(0.5, 1.5, size=(batch, 3, 1))
frequencies = rng.uniform(
np.array([0.7, 1.7, 3.5])[None, :, None],
np.array([1.1, 2.3, 5.0])[None, :, None],
size=(batch, 3, 1),
)
phases = rng.uniform(0, 2 * np.pi, size=(batch, 3, 1))
values = (
amplitudes
* np.sin(frequencies * time[:, None, :] + phases)
* np.array([1.0, 0.5, 0.2])[None, :, None]
).sum(axis=1)
values += rng.normal(0, 0.01, values.shape)
current = np.concatenate([np.zeros((batch, 1)), values[:, :-1]], axis=1)
inputs = np.stack([current, delta_time], axis=2).astype(np.float32)
return (
torch.from_numpy(inputs),
torch.from_numpy(values.astype(np.float32)).unsqueeze(2),
torch.from_numpy(time.astype(np.float32)),
)
|