1import numpy as np
 2
 3META = {
 4    "name": "orientation_phase_quotient",
 5    "domain": "molecular_orientation_symmetry",
 6    "description": "Cubic apolar molecular orientations with isotropic/nematic phase labels and nuisance global rotations/sign flips.",
 7}
 8
 9
10def _p2(x):
11    return 0.5 * (3.0 * x * x - 1.0)
12
13
14def _rot(rng):
15    q, r = np.linalg.qr(rng.normal(size=(3, 3)))
16    q = q @ np.diag(np.where(np.diag(r) >= 0, 1.0, -1.0))
17    if np.linalg.det(q) < 0:
18        q[:, 0] *= -1
19    return q
20
21
22def _corr_features(u, L=4, shells=2):
23    a = u.reshape(L, L, L, 3)
24    channels = []
25    for radius in range(1, shells + 1):
26        offsets = [(dx, dy, dz)
27                   for dx in range(-radius, radius + 1)
28                   for dy in range(-radius, radius + 1)
29                   for dz in range(-radius, radius + 1)
30                   if (dx, dy, dz) != (0, 0, 0)
31                   and dx * dx + dy * dy + dz * dz == radius * radius]
32        c = np.zeros((L, L, L), dtype=np.float64)
33        for dx, dy, dz in offsets:
34            b = np.roll(a, (dx, dy, dz), axis=(0, 1, 2))
35            c += _p2(np.sum(a * b, axis=-1))
36        channels.append(c / len(offsets))
37    return np.stack(channels, axis=0).reshape(-1).astype(np.float32)
38
39
40def _sample(rng, phase, L=4):
41    n = L ** 3
42    if phase == 0:
43        u = rng.normal(size=(n, 3))
44    else:
45        d = rng.normal(size=3)
46        d /= np.linalg.norm(d)
47        u = d + rng.normal(scale=0.28, size=(n, 3))
48    u /= np.linalg.norm(u, axis=1, keepdims=True)
49    u *= rng.choice([-1.0, 1.0], size=(n, 1))
50    # Global frame is nuisance; independent signs are the apolar gauge.
51    return u @ _rot(rng).T
52
53
54def get_dataset(seed, n_train, n_test):
55    rng = np.random.default_rng(int(seed))
56    total = int(n_train) + int(n_test)
57    x, y = [], []
58    for k in range(total):
59        phase = k % 2
60        x.append(_sample(rng, phase))
61        y.append(phase)
62    x = np.asarray(x, dtype=np.float32)
63    y = np.asarray(y, dtype=np.int64)
64    perm = rng.permutation(total)
65    x, y = x[perm], y[perm]
66    return {
67        "xtr": x[:n_train].reshape(n_train, -1),
68        "ytr": y[:n_train],
69        "xte": x[n_train:].reshape(n_test, -1),
70        "yte": y[n_train:],
71        "task": "classification",
72        "metric": "err",
73        "input_shape": (4 * 4 * 4 * 3,),
74        "out_dim": 2,
75    }
76
77
78def invariant_dataset(ds):
79    def conv(x):
80        return np.asarray([_corr_features(v.reshape(4 ** 3, 3)) for v in x], dtype=np.float32)
81    out = dict(ds)
82    out["xtr"] = conv(ds["xtr"])
83    out["xte"] = conv(ds["xte"])
84    out["input_shape"] = (2 * 4 ** 3,)
85    return out
86
87
88def math_check(seed=71):
89    rng = np.random.default_rng(seed)
90    u = _sample(rng, 1)
91    v = u @ _rot(rng).T
92    v *= rng.choice([-1.0, 1.0], size=(len(v), 1))
93    return float(np.max(np.abs(_corr_features(u) - _corr_features(v))))