1import numpy as np
 2
 3META = {
 4    "name": "complex_array_channel",
 5    "domain": "array-valued complex low-rank tensors",
 6    "description": "Regression from noisy complex array channels to clean channels represented by continuous rank-one steering atoms."
 7}
 8
 9def steering(n, u):
10    x = np.arange(n, dtype=np.float32) - (n - 1) / 2.0
11    return np.exp(1j * np.pi * x * u) / np.sqrt(n)
12
13def make_channel(rng, n=8, k=2):
14    h = np.zeros((n, n), dtype=np.complex64)
15    for _ in range(k):
16        ur, ut = rng.uniform(-0.8, 0.8, 2)
17        g = (rng.normal() + 1j * rng.normal()) / np.sqrt(2 * k)
18        h += g * np.outer(steering(n, ur), steering(n, ut).conj())
19    return h
20
21def pack(h):
22    return np.concatenate([h.real.reshape(-1), h.imag.reshape(-1)]).astype(np.float32)
23
24def get_dataset(seed, n_train, n_test):
25    rng = np.random.RandomState(seed)
26    def sample(num):
27        xs, ys = [], []
28        for _ in range(num):
29            clean = make_channel(rng)
30            noise = (rng.normal(size=clean.shape) + 1j * rng.normal(size=clean.shape)).astype(np.complex64) * 0.035
31            xs.append(pack(clean + noise))
32            ys.append(pack(clean))
33        return np.asarray(xs, dtype=np.float32), np.asarray(ys, dtype=np.float32)
34    xtr, ytr = sample(n_train)
35    xte, yte = sample(n_test)
36    return {"xtr": xtr, "ytr": ytr, "xte": xte, "yte": yte,
37            "task": "regression", "metric": "mse"}