Runs a single simulation step. :param x: Inputs to the layer.
(self, x: torch.Tensor)
| 1263 | self.register_buffer("u", self.b * self.v) # Neuron recovery. |
| 1264 | |
| 1265 | def forward(self, x: torch.Tensor) -> None: |
| 1266 | # language=rst |
| 1267 | """ |
| 1268 | Runs a single simulation step. |
| 1269 | |
| 1270 | :param x: Inputs to the layer. |
| 1271 | """ |
| 1272 | |
| 1273 | # Voltage and recovery reset. |
| 1274 | self.v = torch.where(self.s, self.c, self.v) |
| 1275 | self.u = torch.where(self.s, self.u + self.d, self.u) |
| 1276 | |
| 1277 | # Add inter-columnar input. |
| 1278 | if self.s.any(): |
| 1279 | x += torch.cat( |
| 1280 | [self.S[:, self.s[i]].sum(dim=1)[None] for i in range(self.s.shape[0])], |
| 1281 | dim=0, |
| 1282 | ) |
| 1283 | |
| 1284 | # Apply v and u updates. |
| 1285 | self.v += self.dt * 0.5 * (0.04 * self.v**2 + 5 * self.v + 140 - self.u + x) |
| 1286 | self.v += self.dt * 0.5 * (0.04 * self.v**2 + 5 * self.v + 140 - self.u + x) |
| 1287 | self.u += self.dt * self.a * (self.b * self.v - self.u) |
| 1288 | |
| 1289 | # Voltage clipping to lower bound. |
| 1290 | if self.lbound is not None: |
| 1291 | self.v.masked_fill_(self.v < self.lbound, self.lbound) |
| 1292 | |
| 1293 | # Check for spiking neurons. |
| 1294 | self.s = self.v >= self.thresh |
| 1295 | |
| 1296 | super().forward(x) |
| 1297 | |
| 1298 | def reset_state_variables(self) -> None: |
| 1299 | # language=rst |