1import json, math, random
 2from pathlib import Path
 3import numpy as np
 4from scipy.integrate import quad
 5import torch
 6from torch import nn
 7
 8SEED = 282
 9np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
10torch.set_num_threads(4)
11
12# Positive Mellin mixture h(a)=sum softplus(w)(1+softplus(s))^-a.
13class MellinGate(nn.Module):
14    def __init__(self, L=6, hidden=8):
15        super().__init__()
16        self.L = L
17        self.wraw = nn.Parameter(torch.zeros(L))
18        self.sraw = nn.Parameter(torch.linspace(-2.5, 1.5, L))
19        self.base = nn.Sequential(nn.Linear(1, hidden), nn.Tanh(), nn.Linear(hidden, 1))
20    def forward(self, z, a):
21        b = torch.nn.functional.softplus(self.base(z))
22        w = torch.nn.functional.softplus(self.wraw)
23        s = torch.nn.functional.softplus(self.sraw)
24        h = torch.exp(torch.log(w)[None, :] - a[:, None] * torch.log1p(s)[None, :]).sum(1, keepdim=True)
25        return b * h
26
27class FreeGate(nn.Module):
28    def __init__(self, hidden=12):
29        super().__init__()
30        self.net = nn.Sequential(nn.Linear(2, hidden), nn.Tanh(), nn.Linear(hidden, 1), nn.Softplus())
31    def forward(self, z, a):
32        return self.net(torch.cat([z, a[:, None]], 1))
33
34def gamma_identity_check():
35    # Adaptive quadrature verifies the gamma-normalized identity for an atomic measure.
36    mu = np.array([0.7, 1.2, 0.4]); rates = np.array([0.0, 0.8, 3.0])
37    aa = np.array([0.6, 1.0, 2.5, 5.0])
38    lhs = []
39    for a in aa:
40        f = lambda y: y**(a-1) * np.exp(-y) * np.sum(mu*np.exp(-rates*y))
41        lhs.append(quad(f, 0.0, np.inf, epsabs=1e-12, epsrel=1e-12)[0] / math.gamma(a))
42    lhs = np.array(lhs)
43    rhs = np.sum(mu[None,:] * (1+rates[None,:])**(-aa[:,None]), axis=1)
44    return float(np.max(np.abs(lhs-rhs))), lhs.tolist(), rhs.tolist()
45
46def structure_check():
47    w = np.array([0.4, 1.1, 0.8, 0.3]); s = np.array([0., .2, 1.5, 6.])
48    def h(a): return np.sum(w * (1+s)**(-a))
49    grid = np.linspace(.3, 8., 100)
50    # finite differences: (-1)^n h^(n) >= 0, and log convexity h h''-(h')^2 >=0.
51    deriv_min = []
52    for n in range(5):
53        exact = np.sum(w[None,:] * (-np.log1p(s)[None,:])**n * (1+s)[None,:]**(-grid[:,None]), axis=1)
54        deriv_min.append(float(np.min(((-1)**n)*exact)))
55    a0, d = .7, .45
56    H = np.array([[h(a0+(i+j)*d) for j in range(4)] for i in range(4)])
57    eig = np.linalg.eigvalsh(H)
58    hp = np.sum(w[None,:] * (-np.log1p(s)[None,:]) * (1+s)[None,:]**(-grid[:,None]), axis=1)
59    hpp = np.sum(w[None,:] * np.log1p(s)[None,:]**2 * (1+s)[None,:]**(-grid[:,None]), axis=1)
60    logconv_min = float(np.min(np.array([h(x) for x in grid])*hpp-hp**2))
61    return deriv_min, float(eig.min()), logconv_min
62
63def fit_model(model, ztr, atr, ytr, zte, ate, yte, steps=1800):
64    opt = torch.optim.Adam(model.parameters(), lr=.025)
65    for _ in range(steps):
66        pred = model(ztr, atr)
67        loss = ((pred-ytr)**2).mean()
68        opt.zero_grad(); loss.backward(); opt.step()
69    with torch.no_grad():
70        train = float(((model(ztr,atr)-ytr)**2).mean())
71        test = float(((model(zte,ate)-yte)**2).mean())
72    return train, test, sum(p.numel() for p in model.parameters())
73
74def toy_experiment():
75    # Feature is constant; target is a positive Mellin response, with observation noise.
76    rng = np.random.default_rng(SEED)
77    z = torch.zeros(180,1)
78    a = torch.linspace(0.5, 6.0, 180)
79    truew = torch.tensor([.25,.7,.5]); trues = torch.tensor([.05,.8,4.0])
80    target = sum(truew[i]*(1+trues[i])**(-a) for i in range(3))
81    noisy = target + torch.tensor(rng.normal(0,.012, len(a)), dtype=torch.float32)
82    trainmask = (a <= 4.5)
83    ztr,atr,ytr = z[trainmask],a[trainmask],noisy[trainmask,None]
84    zte,ate,yte = z[~trainmask],a[~trainmask],target[~trainmask,None]
85    # Equal-ish parameter count: free MLP has 49, Mellin has 6+6+ (1*8+8+8+1)=33.
86    # Use a scalar offset feature to make both models genuinely order-dependent.
87    torch.manual_seed(SEED)
88    free = fit_model(FreeGate(hidden=9), ztr,atr,ytr,zte,ate,yte)
89    torch.manual_seed(SEED)
90    mellin = fit_model(MellinGate(L=6,hidden=8), ztr,atr,ytr,zte,ate,yte)
91    return {"free_mlp": free, "positive_mellin": mellin, "train_orders": int(trainmask.sum()), "extrapolation_orders": int((~trainmask).sum())}
92
93def main():
94    ident = gamma_identity_check(); struct = structure_check(); exp = toy_experiment()
95    out = {"gamma_identity_max_abs_error": ident[0], "identity_lhs": ident[1], "identity_rhs": ident[2], "derivative_nonnegative_min_by_order": struct[0], "hankel_min_eigenvalue": struct[1], "log_convexity_min": struct[2], "toy": exp}
96    Path("results.json").write_text(json.dumps(out, indent=2))
97    print(json.dumps(out, indent=2))
98
99if __name__ == '__main__': main()