1import numpy as np
 2import torch
 3
 4META = {
 5    'name': 'piola_surface_operator',
 6    'domain': 'pde',
 7    'description': 'Manufactured vector-valued surface operator on a fixed reference chart with affine deformations and contravariant Piola transport.'
 8}
 9
10N = 6
11q = np.linspace(0.08, 0.92, N).astype(np.float32)
12xx, yy = np.meshgrid(q, q, indexing='ij')
13X = np.stack([xx.ravel(), yy.ravel()], axis=1)
14M = len(X)
15DSX = np.pi*np.cos(np.pi*X[:,0])*np.sin(np.pi*X[:,1])
16DSY = np.pi*np.sin(np.pi*X[:,0])*np.cos(np.pi*X[:,1])
17
18
19def geometry(y):
20    a, b, c = [float(v) for v in y]
21    F = np.zeros((M, 3, 2), dtype=np.float32)
22    F[:,0,0] = 1 + a*DSX; F[:,0,1] = a*DSY
23    F[:,1,0] = b*DSX; F[:,1,1] = 1 + b*DSY
24    F[:,2,0] = c*DSX; F[:,2,1] = c*DSY
25    G = np.einsum('nki,nkj->nij', F, F)
26    J = np.sqrt(np.maximum(np.linalg.det(G), 1e-8)).astype(np.float32)
27    return F, J
28
29
30def piola(F, J, u):
31    return np.einsum('nki,ni->nk', F, u) / J[:, None]
32
33
34def inverse_piola(F, J, v):
35    G = np.einsum('nki,nkj->nij', F, F)
36    rhs = np.einsum('nki,nk->ni', F, J[:, None] * v)
37    return np.einsum('nij,nj->ni', np.linalg.inv(G), rhs)
38
39
40def _sample(rng, n, physical):
41    ys = rng.uniform(-0.35, 0.35, (n, 3)).astype(np.float32)
42    out = []
43    for y in ys:
44        F, J = geometry(y)
45        u = rng.normal(size=(M, 2)).astype(np.float32)
46        u += 0.45*np.stack([np.sin(2*np.pi*X[:,0]), np.cos(2*np.pi*X[:,1])], axis=1)
47        target = 0.7*u + 0.3*u.mean(axis=0, keepdims=True)
48        inpvec = piola(F, J, u) if physical else np.concatenate([u, np.zeros((M, 1), dtype=np.float32)], axis=1)
49        # Same architecture and tensor shape; baseline sees physical components,
50        # idea sees the inverse-Piola fixed-reference components.
51        geomfeat = np.broadcast_to(y, (M, 3))
52        meanfeat = np.broadcast_to(inpvec.mean(axis=0), (M, 3))
53        feats = np.concatenate([X, geomfeat, inpvec, meanfeat], axis=1)
54        out.append((feats.reshape(-1), target.reshape(-1)))
55    return np.stack([z[0] for z in out]).astype(np.float32), np.stack([z[1] for z in out]).astype(np.float32)
56
57
58
59def math_check():
60    y=np.array([0.18,-0.13,0.16], dtype=np.float32)
61    F,J=geometry(y)
62    u=np.stack([X[:,0]**2+X[:,1], X[:,0]-X[:,1]**2], axis=1).astype(np.float32)
63    v=np.stack([np.sin(2*np.pi*X[:,0]), np.cos(2*np.pi*X[:,1])], axis=1).astype(np.float32)
64    lhs=float(np.mean(np.sum(piola(F,J,u)*piola(F,J,v),axis=1)*J))
65    pu, pv = np.einsum('nki,ni->nk',F,u), np.einsum('nki,ni->nk',F,v)
66    rhs=float(np.mean(np.sum(pu*pv,axis=1)/J))
67    return {'area_cancellation_relative_error': abs(lhs-rhs)/(abs(rhs)+1e-12),
68            'J_min': float(J.min()), 'J_max': float(J.max())}
69
70def get_dataset(seed, n_train, n_test, physical=True):
71    xtr, ytr = _sample(np.random.RandomState(seed), n_train, physical)
72    xte, yte = _sample(np.random.RandomState(seed + 5000), n_test, physical)
73    return {'xtr': torch.from_numpy(xtr), 'ytr': torch.from_numpy(ytr),
74            'xte': torch.from_numpy(xte), 'yte': torch.from_numpy(yte),
75            'task': 'regression', 'metric': 'mse', 'input_shape': (xtr.shape[1],),
76            'out_dim': ytr.shape[1]}