| import numpy as np |
| import torch |
|
|
|
|
| def generate_spikes( |
| size: int, |
| spikes_type: str = "choose_randomly", |
| spike_intervals: int | None = None, |
| n_spikes: int | None = None, |
| to_keep_rate: float = 0.4, |
| ): |
| spikes = np.zeros(size) |
| if size < 120: |
| build_up_points = 1 |
| elif size < 250: |
| build_up_points = np.random.choice([2, 1], p=[0.3, 0.7]) |
| else: |
| build_up_points = np.random.choice([3, 2, 1], p=[0.15, 0.45, 0.4]) |
|
|
| spike_duration = build_up_points * 2 |
|
|
| if spikes_type == "choose_randomly": |
| spikes_type = np.random.choice(["regular", "patchy", "random"], p=[0.4, 0.5, 0.1]) |
|
|
| if spikes_type == "patchy" and size < 64: |
| spikes_type = "regular" |
|
|
| if spikes_type in ["regular", "patchy"]: |
| if spike_intervals is None: |
| upper_bound = np.ceil(spike_duration / 0.05) |
| lower_bound = np.ceil(spike_duration / 0.15) |
| spike_intervals = np.random.randint(lower_bound, upper_bound) |
| n_spikes = np.ceil(size / spike_intervals) |
| spike_intervals = np.arange(spike_intervals, size, spike_intervals) |
| if spikes_type == "patchy": |
| patch_size = np.random.randint(2, max(n_spikes * 0.7, 3)) |
| to_keep = np.random.randint(np.ceil(patch_size * to_keep_rate), patch_size) |
| else: |
| n_spikes = ( |
| n_spikes if n_spikes is not None else np.random.randint(4, min(max(size // (spike_duration * 3), 6), 20)) |
| ) |
| spike_intervals = np.sort(np.random.choice(np.arange(spike_duration, size), size=n_spikes, replace=False)) |
|
|
| constant_build_rate = False |
| if spikes_type in ["regular", "patchy"]: |
| random_ = np.random.random() |
| constant_build_rate = True |
|
|
| patch_count = 0 |
| spike_intervals -= 1 |
| for interval in spike_intervals: |
| interval = np.round(interval).astype(int) |
| if spikes_type == "patchy": |
| if patch_count >= patch_size: |
| patch_count = 0 |
| if patch_count < to_keep: |
| patch_count += 1 |
| else: |
| patch_count += 1 |
| continue |
| if not constant_build_rate: |
| random_ = np.random.random() |
| build_up_rate = np.random.uniform(0.5, 2) if random_ < 0.7 else np.random.uniform(2.5, 5) |
|
|
| spike_start = interval - build_up_points + 1 |
| for i in range(build_up_points): |
| if 0 <= spike_start + i < len(spikes): |
| spikes[spike_start + i] = build_up_rate * (i + 1) |
|
|
| for i in range(1, build_up_points): |
| if (interval + i) < len(spikes): |
| spikes[interval + i] = spikes[interval - i] |
|
|
| |
| spikes += 1 |
| spikes = spikes * np.random.choice([1, -1], 1, p=[0.7, 0.3]) |
|
|
| return torch.Tensor(spikes) |
|
|
|
|
| def generate_peak_spikes(ts_size, peak_period, spikes_type="regular"): |
| return generate_spikes(ts_size, spikes_type=spikes_type, spike_intervals=peak_period) |
|
|
