1import numpy as np
 2
 3META = {
 4    'name': 'masked_blur_inverse',
 5    'domain': 'pde_inverse',
 6    'description': 'Small linear inverse imaging task with spatially nonuniform detector coverage, analogous to a discretized boundary-value inverse problem.'
 7}
 8
 9def operator(n=16):
10    z = np.arange(n, dtype=np.float64)
11    g = np.exp(-((z[:, None] - z[None, :]) ** 2) / (2 * 1.15 ** 2))
12    g /= g.sum(axis=1, keepdims=True)
13    blur = np.kron(g, g)
14    yy, xx = np.mgrid[:n, :n]
15    coverage = np.where(xx < int(.58*n), 1.0, .18) + .15
16    A = np.sqrt(coverage.reshape(-1))[:, None] * blur
17    return A, coverage
18
19def get_dataset(seed, n_train, n_test):
20    rng = np.random.default_rng(seed)
21    n = 16
22    A, _ = operator(n)
23    def sample(count):
24        ys = np.zeros((count, n, n), dtype=np.float32)
25        yy, xx = np.mgrid[:n, :n]
26        for k in range(count):
27            cy1, cx1 = rng.integers(3, 13, size=2)
28            cy2, cx2 = rng.integers(3, 13, size=2)
29            ys[k] = (((yy-cy1)**2 + (xx-cx1)**2) <= 2.0**2).astype(np.float32)
30            ys[k] += .9 * (((yy-cy2)**2 + (xx-cx2)**2) <= 1.7**2)
31        flat = ys.reshape(count, -1).astype(np.float64)
32        clean = flat @ A.T
33        noise = .08 * np.std(clean, axis=1, keepdims=True) * rng.normal(size=clean.shape)
34        return (clean + noise).astype(np.float32), ys
35    xtr, ytr = sample(n_train)
36    xte, yte = sample(n_test)
37    return {'xtr': xtr, 'ytr': ytr, 'xte': xte, 'yte': yte,
38            'task': 'regression', 'metric': 'mse', 'input_shape': (n*n,),
39            'out_dim': n*n, 'operator': A.astype(np.float32), 'n': n}