def uniform_allocation(N,rank):
""" Uniformly allocates N of tensors of total rank r=(p+q) into
T(k,r-k) for k=0,1,...,r. For orthogonal representations there is no
distinction between p and q, so this op is equivalent to N*T(rank)."""
if N==0: return 0
even_split = sum((N//(rank+1))*T(k,rank-k) for k in range(rank+1))
ragged = sum(random.sample([T(k,rank-k) for k in range(rank+1)],N%(rank+1)))
return even_split+ragged
def binomial_allocation(N,rank,G):
""" Allocates N of tensors of total rank r=(p+q) into
T(k,r-k) for k=0,1,...,r to match the binomial distribution.
For orthogonal representations there is no
distinction between p and q, so this op is equivalent to N*T(rank)."""
if N==0: return 0
n_binoms = N//(2**rank)
n_leftover = N%(2**rank)
even_split = sum([n_binoms*int(binom(rank,k))*T(k,rank-k,G) for k in range(rank+1)])
ps = np.random.binomial(rank,.5,n_leftover)
ragged = sum([T(int(p),rank-int(p),G) for p in ps])
out = even_split+ragged
return out
def __init__(self,N=1024,k=5):
super().__init__()
self.dim = (1+3)*k
self.X = torch.randn(N,self.dim)
self.X[:,:k] = F.softplus(self.X[:,:k]) # Masses
mi = self.X[:,:k]
ri = self.X[:,k:].reshape(-1,k,3)
I = torch.eye(3)
r2 = (ri**2).sum(-1)[...,None,None]
inertia = (mi[:,:,None,None]*(r2*I - ri[...,None]*ri[...,None,:])).sum(1)
self.Y = inertia.reshape(-1,9)
self.rep_in = k*Scalar+k*Vector
self.rep_out = T(2)
self.symmetry = O(3)
self.X = self.X.numpy()
self.Y = self.Y.numpy()
# One has to be careful computing offset and scale in a way so that standardizing
# does not violate equivariance
Xmean = self.X.mean(0)
Xmean[k:] = 0
Xstd = np.zeros_like(Xmean)
Xstd[:k] = np.abs(self.X[:,:k]).mean(0)#.std(0)
#Xstd[k:] = (np.sqrt((self.X[:,k:].reshape(N,k,3)**2).mean((0,2))[:,None]) + np.zeros((k,3))).reshape(k*3)
Xstd[k:] = (np.abs(self.X[:,k:].reshape(N,k,3)).mean((0,2))[:,None] + np.zeros((k,3))).reshape(k*3)
Ymean = 0*self.Y.mean(0)
#Ystd = np.sqrt(((self.Y-Ymean)**2).mean((0,1)))+ np.zeros_like(Ymean)
Ystd = np.abs(self.Y-Ymean).mean((0,1)) + np.zeros_like(Ymean)
self.stats =0,1,0,1#Xmean,Xstd,Ymean,Ystd
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.rep_in = 4 * T(1) #Vector
self.rep_out = T(0) #Scalar
self.symmetry = O2eR3()
self.stats = (0, 1, 0, 1)