| 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") |
| 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 |
|
|
