+ t_to_keep = (
+ torch.arange(self.nb_iterations, device=result.device) * self.speed
+ )
+
+ for l in range(self.nb_iterations * self.speed):
+ fine[l] = collision_okay()
+ for n in range(self.nb_birds):
+ c = col[n]
+ result[l, i[n], j[n]] = c
+ result[l, i[n] - vi[n], j[n]] = c
+ result[l, i[n], j[n] - vj[n]] = c
+
+ if (i[n] == 0 and vi[n] == -1) or (
+ i[n] == self.height - 1 and vi[n] == 1
+ ):
+ vi[n] = -vi[n]
+
+ if (j[n] == 0 and vj[n] == -1) or (
+ j[n] == self.width - 1 and vj[n] == 1
+ ):
+ vj[n] = -vj[n]
+
+ i[n] += vi[n]
+ j[n] += vj[n]