File size: 4,891 Bytes
feab7f5 | 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 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | import math
import pytest
import torch
import kernels
surf = kernels.get_kernel("phanerozoic/surf", version=1, trust_remote_code=True)
requires_cuda = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
G = 9.81
def _bumpy_bed(N, L):
xs = torch.linspace(0, L, N, device="cuda")
X, Y = torch.meshgrid(xs, xs, indexing="xy")
b = (-3.0 + 0.8 * torch.exp(-((X - L / 2) ** 2 + (Y - L / 2) ** 2) / 200)
+ 0.5 * torch.sin(X * 0.3) * torch.cos(Y * 0.2))
return b.float().contiguous(), X, Y
@requires_cuda
def test_nsw_lake_at_rest():
"""Still water over an uneven bed must stay still: the well-balancing test.
A scheme that is not well-balanced spins up spurious currents from the bed
slope and never settles.
"""
N, L = 192, 100.0
dx = dy = L / N
b, _, _ = _bumpy_bed(N, L)
h = (0.0 - b).clamp(min=0)
hu = torch.zeros_like(h)
hv = torch.zeros_like(h)
dt = 0.3 * dx / math.sqrt(G * float(h.max()))
for _ in range(200):
h, hu, hv = surf.nsw_step(h, hu, hv, b, dx, dy, dt)
u = hu / h.clamp(min=1e-3)
v = hv / h.clamp(min=1e-3)
assert float(torch.sqrt(u * u + v * v).max()) < 1e-4
assert float((h + b).abs().max()) < 1e-4
@requires_cuda
def test_nsw_dam_break_bounded_and_finite():
"""A radial dam break stays finite and bounded (shock capturing works)."""
N, L = 192, 100.0
dx = dy = L / N
xs = torch.linspace(0, L, N, device="cuda")
X, Y = torch.meshgrid(xs, xs, indexing="xy")
b = torch.full((N, N), -5.0, device="cuda")
r = torch.sqrt((X - L / 2) ** 2 + (Y - L / 2) ** 2)
eta = torch.where(r < 15, torch.ones_like(r), torch.zeros_like(r))
h = (eta - b).clamp(min=0).contiguous()
hu = torch.zeros_like(h)
hv = torch.zeros_like(h)
dt = 0.3 * dx / math.sqrt(G * float(h.max()))
for _ in range(150):
h, hu, hv = surf.nsw_step(h, hu, hv, b, dx, dy, dt)
eta = h + b
assert torch.isfinite(eta).all()
assert -1.0 < float(eta.min()) and float(eta.max()) < 1.0
@requires_cuda
def test_nsw_mass_conservation_closed_bump():
"""With no in/outflow the total water volume is conserved."""
N, L = 128, 60.0
dx = dy = L / N
b = torch.full((N, N), -4.0, device="cuda")
xs = torch.linspace(0, L, N, device="cuda")
X, Y = torch.meshgrid(xs, xs, indexing="xy")
eta = 0.3 * torch.exp(-((X - L / 2) ** 2 + (Y - L / 2) ** 2) / 40)
h = (eta - b).clamp(min=0).contiguous()
hu = torch.zeros_like(h)
hv = torch.zeros_like(h)
m0 = float(h.sum())
dt = 0.3 * dx / math.sqrt(G * float(h.max()))
for _ in range(120):
h, hu, hv = surf.nsw_step(h, hu, hv, b, dx, dy, dt)
assert abs(float(h.sum()) - m0) / m0 < 1e-4
@requires_cuda
def test_sph_isolated_particles_freefall():
"""SPH particles with no neighbours inside 2h feel gravity and nothing else."""
pos = torch.tensor([[0.0, 0.0, 0.0], [0.5, 0.0, 0.0]], device="cuda")
vel = torch.zeros_like(pos)
rho = torch.full((2,), 1000.0, device="cuda")
ptype = torch.zeros(2, dtype=torch.int32, device="cuda")
dp = 0.01
h = 1.5 * dp
csz = 2 * h
origin = (-0.1, -0.1, -0.1)
grid = (int(0.8 / csz) + 2,) * 3
perm, cs, ce = surf.build_cells(pos, origin, csz, grid)
pos, vel, rho, ptype = pos[perm], vel[perm], rho[perm], ptype[perm]
drho, acc, xsph = surf.sph_forces(pos, vel, rho, ptype, cs, ce, grid, origin, csz,
h, 1000.0 * dp ** 3, c0=20.0)
assert torch.allclose(acc[:, 2], torch.full((2,), -G, device="cuda"), atol=1e-4)
assert float(acc[:, :2].abs().max()) < 1e-5
assert float(drho.abs().max()) < 1e-6
@requires_cuda
def test_sph_block_is_finite_and_repels():
"""A compressed block pushes apart: pressure forces are finite and outward."""
dp = 0.02
n = 12
g1 = torch.arange(n, device="cuda", dtype=torch.float32) * dp
X, Y, Z = torch.meshgrid(g1, g1, g1, indexing="ij")
pos = torch.stack([X.reshape(-1), Y.reshape(-1), Z.reshape(-1)], 1).contiguous()
vel = torch.zeros_like(pos)
rho = torch.full((pos.shape[0],), 1050.0, device="cuda") # compressed
ptype = torch.zeros(pos.shape[0], dtype=torch.int32, device="cuda")
h = 1.5 * dp
csz = 2 * h
origin = (-0.1, -0.1, -0.1)
grid = (int((n * dp + 0.3) / csz) + 2,) * 3
perm, cs, ce = surf.build_cells(pos, origin, csz, grid)
pos, vel, rho, ptype = pos[perm], vel[perm], rho[perm], ptype[perm]
drho, acc, xsph = surf.sph_forces(pos, vel, rho, ptype, cs, ce, grid, origin, csz,
h, 1000.0 * dp ** 3, c0=30.0)
assert torch.isfinite(acc).all() and torch.isfinite(drho).all()
ctr = pos.mean(0)
outward = ((pos - ctr) * (acc - torch.tensor([0.0, 0.0, -G], device="cuda"))).sum(1)
assert float(outward.mean()) > 0 # compressed block pushes outward
|