Publish reproducible AlienX release and live demo

5d18f5e verified 14 days ago

2.88 kB

import sys, math, torch, numpy as np
src = open('AlienX-mhd.py').read().replace('if __name__ == "__main__":\n    main()', 'if False:\n    main()')
exec(compile(src, 'AlienX-mhd.py', 'exec'), globals())

G = GRID
torch.manual_seed(0)
x = torch.arange(G, device=DEVICE) * (2 * math.pi / G)
X, Y, Z = torch.meshgrid(x, x, x, indexing="ij")
u = torch.stack([torch.sin(X) * torch.cos(Y) * torch.sin(Z),
                 torch.cos(X) * torch.sin(Y) * torch.cos(Z),
                 torch.sin(Y) * torch.cos(Z)], 0)  # (3,G,G,G)

def rel(a, b):
    return float((a - b).square().mean().sqrt() / (b.square().mean().sqrt() + 1e-8))

# 1. Octahedral (lattice-preserving) roundtrip must be exact
worst_oct = 0.0
for R in octahedral_rotations():
    u2 = interpolate_rotate(interpolate_rotate(u, R), R.T)
    worst_oct = max(worst_oct, rel(u2, u))
print(f"octahedral roundtrip worst = {worst_oct:.3e}  (tol 1e-2)")
oct_ok = worst_oct < 1e-2

# 2. Scalar (3-dim) octahedral roundtrip
f = torch.sin(X) * torch.cos(Y) * torch.sin(Z)
R = octahedral_rotations()[7]
f2 = interpolate_rotate(interpolate_rotate(f, R), R.T)
scalar_ok = rel(f2, f) < 1e-2
print(f"scalar oct roundtrip rel = {rel(f2, f):.3e}  (tol 1e-2)")

# 3. Generic SO(3): single-step sampling accuracy vs analytic (the meaningful
#    test). The double-roundtrip is mathematically impossible because a generic
#    rotation does not preserve the integer Fourier lattice -> the resampled
#    field is not periodic. Assert the sample accuracy instead.
R = random_rotation()
Rnp = R.detach().cpu().numpy().astype(np.float64)
fu = np.asarray(u.detach().cpu().float().numpy(), dtype=np.float64)
h = 2 * math.pi / G
pos = np.stack(np.meshgrid(np.arange(G) * h, np.arange(G) * h, np.arange(G) * h, indexing="ij"), -1).reshape(-1, 3)
p = pos @ Rnp
coords = np.stack([(np.mod(p[:, a], 2 * math.pi) + 2 * math.pi) / h for a in range(3)], 0)
fil = np.tile(fu, (1, 3, 3, 3))
samp = np.stack([map_coordinates(fil[c], coords, order=3, mode="wrap").reshape(G, G, G) for c in range(3)])
analytic = np.stack([np.sin(p[:, 0]).reshape(G, G, G),
                     np.cos(p[:, 0]).reshape(G, G, G) * 0,   # placeholder, computed below
                     np.zeros((G, G, G))])
# analytic u at rotated points: u = (sin(x)cos(y)sin(z), cos(x)sin(y)cos(z), sin(y)cos(z))
pp = p  # (N,3) rotated coords (x,y,z)
ax, ay, az = pp[:, 0], pp[:, 1], pp[:, 2]
ana0 = (np.sin(ax) * np.cos(ay) * np.sin(az)).reshape(G, G, G)
ana1 = (np.cos(ax) * np.sin(ay) * np.cos(az)).reshape(G, G, G)
ana2 = (np.sin(ay) * np.cos(az)).reshape(G, G, G)
analytic = np.stack([ana0, ana1, ana2])
gen_err = rel(torch.from_numpy(samp), torch.from_numpy(analytic))
print(f"generic single-step sample vs analytic = {gen_err:.3e}  (tol 1e-3)")
gen_ok = gen_err < 1e-3

ok = oct_ok and scalar_ok and gen_ok
print("ROUNDTRIP", "PASS" if ok else "FAIL")
raise SystemExit(0 if ok else 1)

	import sys, math, torch, numpy as np
	src = open('AlienX-mhd.py').read().replace('if __name__ == "__main__":\n main()', 'if False:\n main()')
	exec(compile(src, 'AlienX-mhd.py', 'exec'), globals())

	G = GRID
	torch.manual_seed(0)
	x = torch.arange(G, device=DEVICE) * (2 * math.pi / G)
	X, Y, Z = torch.meshgrid(x, x, x, indexing="ij")
	u = torch.stack([torch.sin(X) * torch.cos(Y) * torch.sin(Z),
	torch.cos(X) * torch.sin(Y) * torch.cos(Z),
	torch.sin(Y) * torch.cos(Z)], 0) # (3,G,G,G)

	def rel(a, b):
	return float((a - b).square().mean().sqrt() / (b.square().mean().sqrt() + 1e-8))

	# 1. Octahedral (lattice-preserving) roundtrip must be exact
	worst_oct = 0.0
	for R in octahedral_rotations():
	u2 = interpolate_rotate(interpolate_rotate(u, R), R.T)
	worst_oct = max(worst_oct, rel(u2, u))
	print(f"octahedral roundtrip worst = {worst_oct:.3e} (tol 1e-2)")
	oct_ok = worst_oct < 1e-2

	# 2. Scalar (3-dim) octahedral roundtrip
	f = torch.sin(X) * torch.cos(Y) * torch.sin(Z)
	R = octahedral_rotations()[7]
	f2 = interpolate_rotate(interpolate_rotate(f, R), R.T)
	scalar_ok = rel(f2, f) < 1e-2
	print(f"scalar oct roundtrip rel = {rel(f2, f):.3e} (tol 1e-2)")

	# 3. Generic SO(3): single-step sampling accuracy vs analytic (the meaningful
	# test). The double-roundtrip is mathematically impossible because a generic
	# rotation does not preserve the integer Fourier lattice -> the resampled
	# field is not periodic. Assert the sample accuracy instead.
	R = random_rotation()
	Rnp = R.detach().cpu().numpy().astype(np.float64)
	fu = np.asarray(u.detach().cpu().float().numpy(), dtype=np.float64)
	h = 2 * math.pi / G
	pos = np.stack(np.meshgrid(np.arange(G) * h, np.arange(G) * h, np.arange(G) * h, indexing="ij"), -1).reshape(-1, 3)
	p = pos @ Rnp
	coords = np.stack([(np.mod(p[:, a], 2 * math.pi) + 2 * math.pi) / h for a in range(3)], 0)
	fil = np.tile(fu, (1, 3, 3, 3))
	samp = np.stack([map_coordinates(fil[c], coords, order=3, mode="wrap").reshape(G, G, G) for c in range(3)])
	analytic = np.stack([np.sin(p[:, 0]).reshape(G, G, G),
	np.cos(p[:, 0]).reshape(G, G, G) * 0, # placeholder, computed below
	np.zeros((G, G, G))])
	# analytic u at rotated points: u = (sin(x)cos(y)sin(z), cos(x)sin(y)cos(z), sin(y)cos(z))
	pp = p # (N,3) rotated coords (x,y,z)
	ax, ay, az = pp[:, 0], pp[:, 1], pp[:, 2]
	ana0 = (np.sin(ax) * np.cos(ay) * np.sin(az)).reshape(G, G, G)
	ana1 = (np.cos(ax) * np.sin(ay) * np.cos(az)).reshape(G, G, G)
	ana2 = (np.sin(ay) * np.cos(az)).reshape(G, G, G)
	analytic = np.stack([ana0, ana1, ana2])
	gen_err = rel(torch.from_numpy(samp), torch.from_numpy(analytic))
	print(f"generic single-step sample vs analytic = {gen_err:.3e} (tol 1e-3)")
	gen_ok = gen_err < 1e-3

	ok = oct_ok and scalar_ok and gen_ok
	print("ROUNDTRIP", "PASS" if ok else "FAIL")
	raise SystemExit(0 if ok else 1)