File size: 1,815 Bytes
626ab93 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | from __future__ import annotations
import torch
from torch.nn import functional as F
LAPLACIAN = torch.tensor(
[
[0.05, 0.20, 0.05],
[0.20, -1.00, 0.20],
[0.05, 0.20, 0.05],
],
dtype=torch.float32,
).reshape(1, 1, 3, 3)
def initial_state(
batch: int,
size: int,
seed: int,
) -> torch.Tensor:
generator = torch.Generator().manual_seed(seed)
u = torch.ones(batch, 1, size, size)
v = torch.zeros(batch, 1, size, size)
radius = max(2, size // 10)
center = size // 2
u[
:,
:,
center - radius : center + radius,
center - radius : center + radius,
] = 0.50
v[
:,
:,
center - radius : center + radius,
center - radius : center + radius,
] = 0.25
noise = torch.rand((batch, 1, size, size), generator=generator) * 0.04 - 0.02
u = torch.clamp(u + noise, 0, 1)
v = torch.clamp(v - noise, 0, 1)
return torch.cat([u, v], dim=1)
def gray_scott_step(
state: torch.Tensor,
feed: torch.Tensor,
kill: torch.Tensor,
diffusion_u: float = 0.16,
diffusion_v: float = 0.08,
) -> torch.Tensor:
kernel = LAPLACIAN.to(state)
u = state[:, 0:1]
v = state[:, 1:2]
padded_u = F.pad(u, (1, 1, 1, 1), mode="circular")
padded_v = F.pad(v, (1, 1, 1, 1), mode="circular")
laplacian_u = F.conv2d(padded_u, kernel)
laplacian_v = F.conv2d(padded_v, kernel)
reaction = u * v.square()
feed_field = feed.reshape(-1, 1, 1, 1)
kill_field = kill.reshape(-1, 1, 1, 1)
delta_u = diffusion_u * laplacian_u - reaction + feed_field * (1 - u)
delta_v = diffusion_v * laplacian_v + reaction - (feed_field + kill_field) * v
return torch.clamp(
torch.cat([u + delta_u, v + delta_v], dim=1),
0,
1,
)
|