1import json, math, random, time
  2import numpy as np
  3import torch
  4from torch import nn
  5
  6SEED = 1254
  7random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
  8torch.set_num_threads(min(8, torch.get_num_threads()))
  9DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 10try:
 11    if DEVICE == "cuda":
 12        torch.cuda.empty_cache()
 13        torch.zeros(1, device=DEVICE)
 14except Exception:
 15    DEVICE = "cpu"
 16
 17def exact_incidence_check():
 18    rng = np.random.default_rng(SEED)
 19    A = rng.integers(0, 2, size=(17, 23)).astype(np.float32)
 20    A[A.sum(1) == 0, 0] = 1
 21    s = rng.normal(size=23).astype(np.float32)
 22    dense = A @ s
 23    # Same operation expressed as row-wise sparse entries.
 24    sparse = np.array([s[row.astype(bool)].sum() for row in A])
 25    return {"max_abs_error": float(np.max(np.abs(dense-sparse))),
 26            "all_nonempty_rows": bool(np.all(A.sum(1) > 0))}
 27
 28def normalization_sweep():
 29    # For s_m=c, a k-atom action has logit c*k^(1-alpha).
 30    c, eps = 2.0, 1e-12
 31    ks = np.arange(1, 33, dtype=float)
 32    rows = []
 33    for alpha in (0.0, 0.5, 1.0):
 34        observed = c * ks / (ks + eps) ** alpha
 35        predicted = c * ks ** (1-alpha)
 36        slope = float(np.polyfit(np.log(ks), np.log(observed), 1)[0])
 37        rel = float(np.max(np.abs(observed-predicted) / np.maximum(1e-9, np.abs(predicted))))
 38        rows.append({"alpha": alpha, "predicted_loglog_slope": 1-alpha,
 39                     "observed_loglog_slope": slope, "max_relative_error": rel,
 40                     "logit_k1": float(observed[0]), "logit_k32": float(observed[-1])})
 41    return rows
 42
 43def softmax_bias_sweep():
 44    # Equal atom scores: raw sum's odds between k and 1 atom are exp((k-1)c).
 45    c = 0.25
 46    out=[]
 47    for alpha in (0., .5, 1.):
 48        vals=[]
 49        for k in (1,2,4,8,16):
 50            l1=c
 51            lk=c*k/(k**alpha)
 52            vals.append(math.exp(lk-l1))
 53        pred_last=math.exp(c*(16**(1-alpha)-1))
 54        out.append({"alpha":alpha, "odds_k16_vs_k1_observed":vals[-1],
 55                    "odds_k16_vs_k1_predicted":pred_last,
 56                    "odds_all_sizes":vals})
 57    return out
 58
 59def make_batch(n, max_atoms, d, max_actions, rng):
 60    # Each target action is the additive utility of its distinct atomic entities.
 61    xs=[]; inc=[]; labels=[]
 62    true_w=np.array([1.1,-.8,.6,.35], dtype=np.float32)[:d]
 63    for _ in range(n):
 64        na=int(rng.integers(3,max_atoms+1)); pa=int(rng.integers(3,max_actions+1))
 65        x=rng.normal(size=(na,d)).astype(np.float32)
 66        A=np.zeros((pa,na),np.float32)
 67        for p in range(pa):
 68            cnt=int(rng.integers(1,min(na,5)+1)); A[p,rng.choice(na,cnt,replace=False)]=1
 69        util=A @ (x@true_w) + .08*rng.normal(size=pa)
 70        xs.append(x); inc.append(A); labels.append(int(np.argmax(util)))
 71    return xs,inc,labels
 72
 73class Structured(nn.Module):
 74    def __init__(self,d):
 75        super().__init__(); self.scorer=nn.Sequential(nn.Linear(d,16),nn.ReLU(),nn.Linear(16,1))
 76    def forward(self,x,A): return A @ self.scorer(x).squeeze(-1)
 77
 78class PaddedMLP(nn.Module):
 79    def __init__(self,d,max_atoms,max_actions):
 80        super().__init__(); self.na=max_atoms; self.pa=max_actions
 81        self.net=nn.Sequential(nn.Linear(max_atoms*d,32),nn.ReLU(),nn.Linear(32,max_actions))
 82    def forward(self,x,A):
 83        # Standard padded action head receives a fixed-size flattened graph and
 84        # emits a fixed global action vector; absent actions are masked.
 85        z=torch.zeros(self.na,x.shape[1],device=x.device); z[:x.shape[0]]=x
 86        logits=self.net(z.reshape(-1)).unsqueeze(0).expand(A.shape[0],-1)
 87        # This baseline cannot represent instance-specific action composition;
 88        # use one shared graph score plus a fixed action index bias and mask.
 89        return logits[:, :A.shape[0]].squeeze(0)
 90
 91def mini_experiment():
 92    rng=np.random.default_rng(SEED+1); d=4; train_n=700; test_n=250
 93    tr=make_batch(train_n,8,d,12,rng); te=make_batch(test_n,16,d,32,rng)
 94    # Baseline is intentionally a conventional padded fixed-action MLP.
 95    models=[("structured",Structured(d).to(DEVICE)),("padded_mlp",PaddedMLP(d,16,32).to(DEVICE))]
 96    results={}
 97    for name,model in models:
 98        opt=torch.optim.Adam(model.parameters(),lr=3e-3)
 99        t0=time.time()
100        for epoch in range(35):
101            order=rng.permutation(train_n)
102            for ii in order:
103                x=torch.tensor(tr[0][ii],device=DEVICE); A=torch.tensor(tr[1][ii],device=DEVICE)
104                y=torch.tensor(tr[2][ii],device=DEVICE)
105                opt.zero_grad(); logits=model(x,A).reshape(1,-1)
106                loss=nn.functional.cross_entropy(logits,y.reshape(1)); loss.backward(); opt.step()
107        correct=0; total=0; losses=[]
108        with torch.no_grad():
109            for x0,A0,y0 in zip(*te):
110                x=torch.tensor(x0,device=DEVICE); A=torch.tensor(A0,device=DEVICE)
111                logits=model(x,A); y=torch.tensor(y0,device=DEVICE)
112                losses.append(float(nn.functional.cross_entropy(logits.reshape(1,-1),y.reshape(1))))
113                correct += int(logits.argmax().item()==y0); total += 1
114        results[name]={"unseen_schema_accuracy":correct/total,
115                       "unseen_schema_ce":float(np.mean(losses)),
116                       "parameters":sum(p.numel() for p in model.parameters()),
117                       "seconds":time.time()-t0}
118    # Activation storage comparison for a batch of heterogeneous schemas.
119    results["padding_elements_per_instance"] = 16*32*d
120    results["incidence_elements_per_instance_example"] = int(np.mean([a.size for a in te[1]]))
121    return results
122
123def main():
124    result={"device":DEVICE,"exact_incidence":exact_incidence_check(),
125            "prediction_1_exact_additivity": "l_p equals sum of incident s_m; numerical max error reported above",
126            "prediction_2_normalization_scaling":normalization_sweep(),
127            "prediction_3_equal-score_size_bias":softmax_bias_sweep(),
128            "mini_experiment":mini_experiment()}
129    with open("results.json","w") as f: json.dump(result,f,indent=2)
130    print(json.dumps(result,indent=2))
131if __name__ == "__main__": main()