| 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)), |
| ) |
|
|
|
|
