Generates Bernoulli-distributed spike trains based on input intensity. Inputs must be non-negative. Spikes correspond to successful Bernoulli trials, with success probability equal to (normalized in [0, 1]) input value. :param datum: Tensor of shape ``[n_1, ..., n_k]``. :param
(
datum: torch.Tensor,
time: Optional[int] = None,
dt: float = 1.0,
device="cpu",
**kwargs,
)
| 48 | |
| 49 | |
| 50 | def bernoulli( |
| 51 | datum: torch.Tensor, |
| 52 | time: Optional[int] = None, |
| 53 | dt: float = 1.0, |
| 54 | device="cpu", |
| 55 | **kwargs, |
| 56 | ) -> torch.Tensor: |
| 57 | # language=rst |
| 58 | """ |
| 59 | Generates Bernoulli-distributed spike trains based on input intensity. Inputs must |
| 60 | be non-negative. Spikes correspond to successful Bernoulli trials, with success |
| 61 | probability equal to (normalized in [0, 1]) input value. |
| 62 | |
| 63 | :param datum: Tensor of shape ``[n_1, ..., n_k]``. |
| 64 | :param time: Length of Bernoulli spike train per input variable. |
| 65 | :param dt: Simulation time step. |
| 66 | :return: Tensor of shape ``[time, n_1, ..., n_k]`` of Bernoulli-distributed spikes. |
| 67 | |
| 68 | Keyword arguments: |
| 69 | |
| 70 | :param float max_prob: Maximum probability of spike per Bernoulli trial. |
| 71 | """ |
| 72 | # Setting kwargs. |
| 73 | max_prob = kwargs.get("max_prob", 1.0) |
| 74 | |
| 75 | assert 0 <= max_prob <= 1, "Maximum firing probability must be in range [0, 1]" |
| 76 | assert (datum >= 0).all(), "Inputs must be non-negative" |
| 77 | |
| 78 | shape, size = datum.shape, datum.numel() |
| 79 | datum = datum.flatten().to(device) |
| 80 | |
| 81 | if time is not None: |
| 82 | time = int(time / dt) |
| 83 | |
| 84 | # Normalize inputs and rescale (spike probability proportional to input intensity). |
| 85 | if datum.max() > 1.0: |
| 86 | datum /= datum.max() |
| 87 | |
| 88 | # Make spike data from Bernoulli sampling. |
| 89 | if time is None: |
| 90 | spikes = torch.bernoulli(max_prob * datum).to(device) |
| 91 | spikes = spikes.view(*shape) |
| 92 | else: |
| 93 | spikes = torch.bernoulli(max_prob * datum.repeat([time, 1])) |
| 94 | spikes = spikes.view(time, *shape) |
| 95 | |
| 96 | return spikes.byte() |
| 97 | |
| 98 | |
| 99 | def poisson( |