From 68e171f4511fabb8d0795fa83866ab2dddb7a770 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 27 Feb 2026 14:34:34 -0500 Subject: [PATCH 001/161] Initial integration of NUG into abinitio class. --- src/aspire/abinitio/__init__.py | 1 + src/aspire/abinitio/commonline_nug.py | 789 ++++++++++++++++++++++++++ 2 files changed, 790 insertions(+) create mode 100644 src/aspire/abinitio/commonline_nug.py diff --git a/src/aspire/abinitio/__init__.py b/src/aspire/abinitio/__init__.py index d8e2cfddcd..4b6248c4c7 100644 --- a/src/aspire/abinitio/__init__.py +++ b/src/aspire/abinitio/__init__.py @@ -6,6 +6,7 @@ ) from .commonline_base import Orient3D from .commonline_matrix import CLOrient3D +from .commonline_nug import CommonlineNUG from .commonline_sdp import CommonlineSDP from .commonline_lud import CommonlineLUD from .commonline_irls import CommonlineIRLS diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py new file mode 100644 index 0000000000..eb092e8604 --- /dev/null +++ b/src/aspire/abinitio/commonline_nug.py @@ -0,0 +1,789 @@ +import logging +import time + +import numpy as np +import cupy as cp +from scipy.special import sph_harm, binom, jacobi, factorial +from scipy.io import loadmat +from scipy.spatial.transform import Rotation as spr + +from aspire.abinitio import CLOrient3D +from aspire.nufft import nufft +from aspire.numeric import fft, xp +from aspire.operators import wemd_embed + +logger = logging.getLogger(__name__) + + +class CommonlineNUG(CLOrient3D): + """ + Class to estimate 3D orientations using non-uqique games. + """ + + def __init__( + self, + src, + symmetry=None, + n_rad=None, + n_theta=360, + max_shift=0.15, + shift_step=1, + mask=True, + Lmax=12, + loss='l1', + T=36, + max_iter=501, + rho=0.05, + ratio=1, + factor=1.0, + mult=1.5, + Ngrid=16317, + Nstep_yI=10, + **kwargs, + ): + """ + Initialize object for estimating 3D orientations for molecules with C3 and C4 symmetry. + + :param src: The source object of 2D denoised or class-averaged images with metadata + :param symmetry: A string, ie. 'C3', indicating the symmetry type. + :param n_rad: The number of points in the radial direction + :param n_theta: The number of points in the theta direction. Default = 360. + """ + + super().__init__( + src, + n_rad=n_rad, + n_theta=n_theta, + max_shift=max_shift, + shift_step=shift_step, + mask=mask, + **kwargs, + ) + + self.Lmax = Lmax + self.loss = loss + self.T = T + self.max_iter = max_iter + self.rho = rho + self.ratio = ratio + self.factor = factor + self.mult = mult + self.Ngrid = Ngrid + self.Nstep_yI = Nstep_yI + self.sym = symmetry + + def estimate_rotations(self): + sym_euler,S=self.Symmetry_Euler(self.sym) + imgs = self.src.images[:] + C = self.compute_coeff(imgs, self.loss, self.Lmax, T=self.T) + X_est = self.admm_sym_J( + self.sym, + C, + self.Lmax, + self.n_img, + self.Ngrid, + self.max_iter, + self.rho, + self.ratio, + self.factor, + self.mult, + self.Nstep_yI, + ) + + R_est,Euler_est = self.euler_est(X_est[0],X_est[S-1],self.sym,self.n_img) + self.rotations = R_est + + return R_est + + ####################### + # Compute Coeffs Step # + ####################### + + def compute_coeff(self, Img, loss, Lmax, T): + # compute the coefficient matrix + N,L,_=Img.shape; n_theta=360 + angular_sampling = np.arange(0, 360, 1) + line_proj=np.zeros((L,n_theta,N)) + Img_pft=np.zeros((L,n_theta,N),dtype=complex) + Img = Img.asnumpy() + for n in range(N): + line_proj[:,:,n],Img_pft[:,:,n]=self.fast_radon_transform(Img[n], angular_sampling) + + dim_wave=len(wemd_embed(line_proj[:,0,0])) + WE=np.zeros((dim_wave,n_theta,N)) + for i in range(N): + for theta in range(n_theta): + WE[:,theta,i]=wemd_embed(line_proj[:,theta,i]) + + + def fij(alpha,gamma,i,j,loss): + if loss=='l1': + Ii_hat=Img_pft[:,:,i]; Ij_hat=Img_pft[:,:,j] + idxi=np.round((alpha-np.pi/2)*n_theta/2/np.pi)% n_theta + idxj=np.round((-gamma-np.pi/2)*n_theta/2/np.pi)% n_theta + + Si=Ii_hat[:,int(idxi)]; Sj=Ij_hat[:,int(idxj)] + return np.linalg.norm(Si-Sj,1) + + if loss=='wemd': + idxi=np.round((alpha-np.pi/2)*n_theta/2/np.pi)% n_theta + idxj=np.round((-gamma-np.pi/2)*n_theta/2/np.pi)% n_theta + + Si=WE[:,int(idxi),i]; Sj=WE[:,int(idxj),j] + return np.linalg.norm(Si-Sj,1) + + alpha_grid=np.arange(2*T)*np.pi/T + beta_grid=(2*np.arange(2*T)+1)*np.pi/4/T + gamma_grid=np.arange(2*T)*np.pi/T + + bT=np.zeros(2*T); + for l in range(2*T): + ss=0 + for m in range(T): + ss=ss+np.sin(beta_grid[l]*(2*m+1))/(2*m+1) + bT[l]=2/T*np.sin(beta_grid[l])*ss + + BTK=[]; + for k in range(1,Lmax+1): + dk=2*k+1 + btk=np.zeros((dk,dk)) + for l in range(2*T): btk=btk+bT[l]*self.Wd(k,beta_grid[l]) + BTK.append(btk.T) + + + def fijhat_k(k,F): + dk=2*k+1; + + exp_alpha_grid=np.zeros((2*T,dk),dtype=complex) + for m in range(-k,k+1): + exp_alpha_grid[:,m+k]=np.exp(1j*m*alpha_grid) + + exp_gamma_grid=np.zeros((2*T,dk),dtype=complex) + for m in range(-k,k+1): + exp_gamma_grid[:,m+k]=np.exp(1j*m*gamma_grid) + + S=(exp_alpha_grid.T@F@exp_gamma_grid).T + fhat=BTK[k-1]*S/4/T**2 + return fhat + + C=[] + for k in range(1,Lmax+1): + dk=2*k+1; C.append(np.zeros((N*dk,N*dk),dtype=complex)) + for i in range(N): + for j in range(i+1,N): + Fij=np.zeros((2*T,2*T)) + for j1 in range(2*T): + for j2 in range(2*T): + Fij[j1,j2]=fij(alpha_grid[j1],gamma_grid[j2],i,j,loss) + for k in range(1,Lmax+1): + dk=2*k+1; C[k-1][j*dk:(j+1)*dk,i*dk:(i+1)*dk]=fijhat_k(k,Fij)#*dk + for k in range(1,Lmax+1): + C[k-1]=C[k-1]+C[k-1].conj().T + + for i in range(N): + Fii=np.zeros((2*T,2*T)) + for j1 in range(2*T): + for j2 in range(2*T): + Fii[j1,j2]=fij(alpha_grid[j1],gamma_grid[j2],i,i,loss) + for k in range(1,Lmax+1): + dk=2*k+1 + C[k-1][i*dk:(i+1)*dk,i*dk:(i+1)*dk]=fijhat_k(k,Fii)#*dk + + + for k in range(1,Lmax+1): + [T,Tinv] = self.complex2real(k) + C[k-1] = np.real(np.kron(np.eye(N),Tinv) @ C[k-1]@np.kron(np.eye(N),T)) + C[k-1]=np.round(C[k-1],10) + + return C + + @staticmethod + def fast_radon_transform(array, angles, use_ramp=False): + + angles = np.array(angles).flatten() + img_size = array.shape[1] + rads = angles / 180 * np.pi + y_idx = np.arange(-img_size / 2, img_size / 2) / img_size * 2 + x_theta = y_idx[:, np.newaxis] * np.sin(rads)[np.newaxis, :] + y_theta = y_idx[:, np.newaxis] * np.cos(rads)[np.newaxis, :] + + pts = np.pi * np.vstack( + [ + x_theta.flatten(), + y_theta.flatten(), + ] + ) + pts = pts.astype(array.dtype) + + #array = array.astype(np.float32) + lines_f = nufft(array, pts).reshape((img_size, -1)) + + if img_size % 2 == 0: + lines_f[0, :] = 0 + + if use_ramp: + freqs = np.abs(np.pi * y_idx) + lines_f *= freqs[:, np.newaxis] + + projections = np.real(xp.asnumpy(fft.centered_ifft(xp.asarray(lines_f), axis=0))) + + return projections, lines_f + + + @staticmethod + def Wd(J,beta): + # compute Wigner small d matrix + d=np.zeros((2*J+1,2*J+1)); + + for m in range(-J,J+1): + for n in range(-J,J+1): + smin=max(0,m-n); smax=min(J+m,J-n); + for s in range(smin,smax+1): + mul=np.sqrt(factorial(J+m))/factorial(J+m-s)*np.sqrt(factorial(J+n))/factorial(s)*np.sqrt(factorial(J-m))/factorial(n-m+s)\ + *np.sqrt(factorial(J-n))/factorial(J-n-s) + d[n+J,m+J]+=mul*(-1)**(n-m+s)*(np.cos(beta/2))**(2*J+m-n-2*s)*(np.sin(beta/2))**(n-m+2*s); + + return d + + @staticmethod + def complex2real(ell): + # compute transformation matrices that convert complex representations to real ones + diml=2*ell+1 + Tinv=np.zeros((diml,diml),dtype=complex) + for i in range(diml): + if iell: + Tinv[i,i]=(-1)**(i-ell)/np.sqrt(2); Tinv[i,diml-1-i]=1/np.sqrt(2) + + T=np.zeros((diml,diml),dtype=complex) + for i in range(diml): + if iell: + T[i,i]=(-1)**(i-ell)/np.sqrt(2); T[i,diml-1-i]=1j*(-1)**(i-ell)/np.sqrt(2) + return T,Tinv + + + ############# + # ADMM Step # + ############# + + def admm_sym_J( + self, + sym, + C, + Lmax, + N, + Ngrid, + max_iter, + rho, + ratio, + factor, + mult=1, + Nstep_yI=20, + verbose=True, + GPU=True, + ): + # admm for symmetric case + if GPU: + import cupy as cp + else: + import numpy as cp + + C0,C1,normC,AEq,bEq,AEqAEqtinv,AI_mat_diag,AI_mat_offdiag,bI,Lambda,d0,d1,D0,D1,idx_diag,idx_offdiag,IDX_upper,IDX_lower,X0,X1,Xq,S0,S1,Sq=self.ADMM_preprocessing(C,Lmax,N,Ngrid,GPU) + + # S=int(sym[1]); sym_euler=Symmetry_Euler(sym) + rank_Ak,_=self.compute_rank(sym,Lmax); print(rank_Ak) + + # rank_Ak=cp.zeros(Lmax) + # for ell in range(Lmax): rank_Ak[ell]=np.linalg.matrix_rank(Ak(ell+1,sym_euler)) + + AE=[]; AEAETinv=[] + for k in range(1,Lmax+1): + s0=k**2; s1=(k+1)**2 + AEk=cp.zeros((1+s0+s1,2*(s0+s1))) + AEk[0,:s0]=cp.eye(k).T.reshape(-1); AEk[0,s0:s0+s1]=cp.eye(k+1).T.reshape(-1) + for count in range(1,1+s0): + AEk[count,count-1]=1; AEk[count,count-1+s0+s1]=1 + for count in range(1+s0,1+s0+s1): + AEk[count,count-1]=1; AEk[count,count-1+s1+s0]=1 + AE.append(AEk) + AEAETinv.append(np.linalg.pinv(AEk@AEk.T)) + bE=cp.zeros((Lmax+D0+D1)) + for k in range(Lmax): + bE[k+d0[k]+d1[k]:]=rank_Ak[k] + bE[k+1+d0[k]+d1[k]:k+1+d0[k+1]+d1[k]]=cp.eye(k+1).T.reshape(-1) + bE[k+1+d0[k+1]+d1[k]:k+1+d0[k+1]+d1[k+1]]=cp.eye(k+2).T.reshape(-1) + bE=np.repeat(bE[:,np.newaxis],N,axis=1) + P=[]; + for k in range(1,Lmax+1): + dk=2*k+1; Pk=cp.eye(dk) + for m in range(k): + for l in range(k-m): + Pk[(m+2*l,m+2*l+1),:]=Pk[(m+2*l+1,m+2*l),:] + P.append(Pk) + + + def fun_AE(X0,X1,Xd0,Xd1,Xq): + z=cp.zeros((Lmax+D0+D1,N)) + for k in range(Lmax): + z[k+d0[k]+d1[k]:k+1+d0[k+1]+d1[k+1]]=AE[k]@ cp.concatenate((X0[d0[k]:d0[k+1],idx_diag],X1[d1[k]:d1[k+1],idx_diag],Xd0[d0[k]:d0[k+1]],Xd1[d1[k]:d1[k+1]]),axis=0) + return z, AEq@cp.concatenate((Xq,X0[:1,idx_offdiag],X1[:4,idx_offdiag]),axis=0) + + def fun_AET(yE,yEq): + Z0=cp.zeros((D0,N*(N+1)//2)); Z1=cp.zeros((D1,N*(N+1)//2)); + Zd0=cp.zeros((D0,N)); Zd1=cp.zeros((D1,N)) + for k in range(Lmax): + s0=(k+1)**2; s1=(k+2)**2 + Ztmp=AE[k].T@yE[k+d0[k]+d1[k]:k+1+d0[k+1]+d1[k+1]] + Z0[d0[k]:d0[k+1],idx_diag]=Ztmp[:s0] + Z1[d1[k]:d1[k+1],idx_diag]=Ztmp[s0:s0+s1] + Zd0[d0[k]:d0[k+1]]=Ztmp[s0+s1:2*s0+s1] + Zd1[d1[k]:d1[k+1]]=Ztmp[2*s0+s1:2*s0+2*s1] + Zq=AEq.T@yEq + Z0[:1,idx_offdiag]=Zq[16:17]; Z1[:4,idx_offdiag]=Zq[17:] + return Z0,Z1,Zd0,Zd1,Zq[:16] + + def fun_AI(X0,X1): + z=cp.zeros((Ngrid,N*(N+1)//2)); tmp=cp.concatenate((X0,X1),axis=0); + z[:,idx_diag]=AI_mat_diag@tmp[:,idx_diag]; z[:,idx_offdiag]=AI_mat_offdiag@tmp[:,idx_offdiag] + # return AI_mat@cp.concatenate((X0,X1),axis=0) + return z + + def fun_AIT(yI): + # Z=AI_mat.T@yI + # return Z[:D0,:], Z[D0:,:] + Z=cp.zeros((D0+D1,N*(N+1)//2)); Z[:,idx_diag]=AI_mat_diag.T@yI[:,idx_diag]; Z[:,idx_offdiag]=AI_mat_offdiag.T@yI[:,idx_offdiag]; + return Z[:D0,:], Z[D0:,:] + + def update_S(C0,C1,yE,yEq,yI,X0,X1,Xd0,Xd1,Xq,rho,Lmax,N): + Z0,Z1,Zd0,Zd1,Zq=fun_AET(yE,yEq); AIT_yI0,AIT_yI1=fun_AIT(yI) + S0=C0-Z0-AIT_yI0-X0/rho; S1=C1-Z1-AIT_yI1-X1/rho; + Sd0=-Zd0-Xd0/rho; Sd1=-Zd1-Xd1/rho; Sq=-Zq-Xq/rho + tic0=time.perf_counter(); + for k in range(1,Lmax+1): + tmp=self.mat_block(S0[d0[k-1]:d0[k],:],N,k,IDX_upper,IDX_lower,idx_offdiag) + tmp=self.psd_projection(tmp) + S0[d0[k-1]:d0[k],:]=self.vec_block(tmp,N,k,IDX_upper) + + tmp=self.mat_block(S1[d1[k-1]:d1[k],:],N,k+1,IDX_upper,IDX_lower,idx_offdiag) + tmp=self.psd_projection(tmp) + S1[d1[k-1]:d1[k],:]=self.vec_block(tmp,N,k+1,IDX_upper) + toc0=time.perf_counter(); Time[0]+=toc0-tic0 + tic1=time.perf_counter(); + for k in range(1,Lmax+1): + for n in range(N): + tmp=self.transform_back_block(Sd0[d0[k-1]:d0[k],n],Sd1[d1[k-1]:d1[k],n],k,P[k-1]) + tmp=self.psd_projection(tmp) + Sd0[d0[k-1]:d0[k],n],Sd1[d1[k-1]:d1[k],n]=self.transform_block(tmp,k,P[k-1]) + toc1=time.perf_counter(); Time[1]+=toc1-tic1 + tic2=time.perf_counter(); + for count in range(N*(N-1)//2): + tmp=self.psd_projection(Sq[:,count].reshape(4,4).T) + Sq[:,count]=tmp.T.reshape(16) + toc2=time.perf_counter(); Time[2]+=toc2-tic2 + return S0,S1,Sd0,Sd1,Sq + + + def update_yE(C0,C1,X0,X1,Xd0,Xd1,Xq,S0,S1,Sd0,Sd1,Sq,yI,rho): + AIT_yI0,AIT_yI1=fun_AIT(yI) + z,zq=fun_AE(-X0/rho+C0-S0-AIT_yI0,-X1/rho+C1-S1-AIT_yI1,-Xd0/rho-Sd0,-Xd1/rho-Sd1,-Xq/rho-Sq) + yE=cp.zeros((Lmax+D0+D1,N)) + for k in range(Lmax): + yE[k+d0[k]+d1[k]:k+1+d0[k+1]+d1[k+1]]=AEAETinv[k]@(bE[k+d0[k]+d1[k]:k+1+d0[k+1]+d1[k+1]]/rho+z[k+d0[k]+d1[k]:k+1+d0[k+1]+d1[k+1]]) + yEq=AEqAEqtinv@(bEq/rho+zq) + return yE,yEq + + + def update_yI(C0,C1,X0,X1,S0,S1,yE,yEq,yI,rho,Lambda): + Z0,Z1,_,_,_=fun_AET(yE,yEq); AIT_yI0,AIT_yI1=fun_AIT(yI) + tmp=fun_AI(-X0/rho+C0-S0-Z0-AIT_yI0,-X1/rho+C1-S1-Z1-AIT_yI1) + yI=yI+bI/rho/Lambda+tmp/Lambda + yI=np.maximum(yI,0) + return yI + + def update_X(C0,C1,X0,X1,Xd0,Xd1,Xq,yE,yEq,yI,S0,S1,Sd0,Sd1,Sq,rho): + Z0,Z1,Zd0,Zd1,Zq=fun_AET(yE,yEq); AIT_yI0,AIT_yI1=fun_AIT(yI) + tmp=S0+Z0+AIT_yI0-C0; X0=X0+mult*rho*tmp; resX0=np.linalg.norm(tmp) + tmp=S1+Z1+AIT_yI1-C1; X1=X1+mult*rho*tmp; resX1=np.linalg.norm(tmp) + tmp=Sd0+Zd0; Xd0=Xd0+mult*rho*tmp; resXd0=np.linalg.norm(tmp) + tmp=Sd1+Zd1; Xd1=Xd1+mult*rho*tmp; resXd1=np.linalg.norm(tmp) + tmp=Sq+Zq; Xq=Xq+mult*rho*tmp; resXq=np.linalg.norm(tmp) + return X0,X1,Xd0,Xd1,Xq,np.sqrt(resX0**2+resX1**2+resXd0**2+resXd1**2+resXq**2) + + def update_rho(X0,X1,Xd0,Xd1,Xq,bE,bEq,bI,res_X,rho,factor,normC): + z,zq=fun_AE(X0,X1,Xd0,Xd1,Xq) + res_eq=np.linalg.norm(z-bE)/(1+np.linalg.norm(bE))+np.linalg.norm(zq-bEq)/(1+np.linalg.norm(bEq)); + res_inq=np.linalg.norm(np.maximum(bI-fun_AI(X0,X1),0))/(1+abs(bI)*np.sqrt(Ngrid*N**2)); + p_resnorm=res_eq+res_inq + d_resnorm=res_X/(1+normC) + if d_resnorm>ratio*p_resnorm: rho=rho*factor + if d_resnorm=i: IDX_upper.append(i*N+j); + for j in range(N): + for i in range(N): + if j500: + Lambda=cp.linalg.norm(z); z=z/cp.linalg.norm(z); z=AI@(AI.T@z) + Lambda+=2000; print('Largest eigenvalue of AIAIT is approximately %1.2f'%Lambda) + # Lambda=cp.linalg.eigvalsh(AI@AI.T)[-1]; print(Lambda) + return Lambda + + def compute_rank(self, sym,Lmax): + sym_euler,S=self.Symmetry_Euler(sym); rk=cp.zeros(Lmax); A=[] + for k in range(1,Lmax+1): + dk=2*k+1; Ak=np.zeros((dk,dk),dtype=np.complex128) + for s in range(S): Ak+=self.WD(k,sym_euler[s]) + Ak=np.round(Ak/S,6); A.append(Ak); rk[k-1]=np.linalg.matrix_rank(Ak) + return rk,A + + @staticmethod + def Symmetry_Euler(sym): + if sym[0]=='C': + order=int(sym[1:]); sym_euler=np.zeros((order,3)) + for i in range(order): sym_euler[i]=spr.from_euler('zyx',[2*np.pi/order*i,0,0]).as_euler('zyz') + + if sym[0]=='D': + order=int(sym[1:]); sym_euler=np.zeros((2*order,3)) + for i in range(order): + sym_euler[i]=spr.from_euler('zyx',[2*np.pi/order*i,0,0]).as_euler('zyz') + sym_euler[i+order]=spr.from_euler('zyx',[2*np.pi/order*i,0,np.pi]).as_euler('zyz') + + if sym=='T12': + sym_euler=np.zeros((12,3)) + sym_euler[1]=spr.from_rotvec(2*np.pi/3*np.array([0,0,1])).as_euler('zyz') + sym_euler[2]=spr.from_rotvec(4*np.pi/3*np.array([0,0,1])).as_euler('zyz') + + + return sym_euler,sym_euler.shape[0] + + def WD(self, J,euler): + # compute Wigner D matrix + alpha=euler[0]; beta=euler[1]; gamma=euler[2]; d=self.Wd(J,beta); + D=np.diag(np.exp(-1j*alpha*np.arange(-J,J+1)))@ d @np.diag(np.exp(-1j*gamma*np.arange(-J,J+1))) + return D + + @staticmethod + def Wd(J,beta): + # compute Wigner small d matrix + d=np.zeros((2*J+1,2*J+1)); + + for m in range(-J,J+1): + for n in range(-J,J+1): + smin=max(0,m-n); smax=min(J+m,J-n); + for s in range(smin,smax+1): + mul=np.sqrt(factorial(J+m))/factorial(J+m-s)*np.sqrt(factorial(J+n))/factorial(s)*np.sqrt(factorial(J-m))/factorial(n-m+s)\ + *np.sqrt(factorial(J-n))/factorial(J-n-s) + d[n+J,m+J]+=mul*(-1)**(n-m+s)*(np.cos(beta/2))**(2*J+m-n-2*s)*(np.sin(beta/2))**(n-m+2*s); + + return d + + @staticmethod + def mat_block(vecA,N,sz,IDX_upper,IDX_lower,idx_offdiag): + tmp=vecA.T.reshape(N*(N+1)//2,sz,sz).transpose(0,2,1) + AA=cp.zeros((N**2,sz,sz)); AA[IDX_upper]=tmp; AA[IDX_lower]=tmp[idx_offdiag].transpose(0,2,1) + return (AA.reshape(N,N,sz,sz).transpose(0,2,1,3)).reshape(N*sz,N*sz) + + @staticmethod + def psd_projection(B): + # compute the PSD part of a symmstric matrix + evals,evecs=cp.linalg.eigh((B+B.T)/2); evals=cp.maximum(evals,0) + return (evecs*evals) @ evecs.T + + @staticmethod + def transform_block(A,k,Pk=None): + if Pk is None: + dk=2*k+1; Pk=cp.eye(dk) + for m in range(k): + for l in range(k-m): + Pk[(m+2*l,m+2*l+1),:]=Pk[(m+2*l+1,m+2*l),:] + AT=Pk@A@Pk.T + return AT[:k,:k].T.reshape(-1), AT[k:,k:].T.reshape(-1) + + @staticmethod + def transform_back_block(A0,A1,k,Pk=None): + dk=2*k+1; A=cp.zeros((dk,dk)); A[:k,:k]=A0.reshape(k,k).T; A[k:,k:]=A1.reshape(k+1,k+1).T + if Pk is None: + Pk=cp.eye(dk) + for m in range(k): + for l in range(k-m): + Pk[(m+2*l,m+2*l+1),:]=Pk[(m+2*l+1,m+2*l),:] + return Pk.T@A@Pk From f206b403f93a0917cd3907bbc49aa0de31bb9d63 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 27 Feb 2026 14:37:49 -0500 Subject: [PATCH 002/161] isort, black --- src/aspire/abinitio/commonline_nug.py | 1441 ++++++++++++++++--------- 1 file changed, 939 insertions(+), 502 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index eb092e8604..97a7420d64 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1,11 +1,11 @@ import logging import time -import numpy as np import cupy as cp -from scipy.special import sph_harm, binom, jacobi, factorial +import numpy as np from scipy.io import loadmat from scipy.spatial.transform import Rotation as spr +from scipy.special import binom, factorial, jacobi, sph_harm from aspire.abinitio import CLOrient3D from aspire.nufft import nufft @@ -27,10 +27,10 @@ def __init__( n_rad=None, n_theta=360, max_shift=0.15, - shift_step=1, + shift_step=1, mask=True, Lmax=12, - loss='l1', + loss="l1", T=36, max_iter=501, rho=0.05, @@ -71,9 +71,9 @@ def __init__( self.Ngrid = Ngrid self.Nstep_yI = Nstep_yI self.sym = symmetry - + def estimate_rotations(self): - sym_euler,S=self.Symmetry_Euler(self.sym) + sym_euler, S = self.Symmetry_Euler(self.sym) imgs = self.src.images[:] C = self.compute_coeff(imgs, self.loss, self.Lmax, T=self.T) X_est = self.admm_sym_J( @@ -90,7 +90,7 @@ def estimate_rotations(self): self.Nstep_yI, ) - R_est,Euler_est = self.euler_est(X_est[0],X_est[S-1],self.sym,self.n_img) + R_est, Euler_est = self.euler_est(X_est[0], X_est[S - 1], self.sym, self.n_img) self.rotations = R_est return R_est @@ -98,104 +98,116 @@ def estimate_rotations(self): ####################### # Compute Coeffs Step # ####################### - + def compute_coeff(self, Img, loss, Lmax, T): - # compute the coefficient matrix - N,L,_=Img.shape; n_theta=360 - angular_sampling = np.arange(0, 360, 1) - line_proj=np.zeros((L,n_theta,N)) - Img_pft=np.zeros((L,n_theta,N),dtype=complex) + # compute the coefficient matrix + N, L, _ = Img.shape + n_theta = 360 + angular_sampling = np.arange(0, 360, 1) + line_proj = np.zeros((L, n_theta, N)) + Img_pft = np.zeros((L, n_theta, N), dtype=complex) Img = Img.asnumpy() for n in range(N): - line_proj[:,:,n],Img_pft[:,:,n]=self.fast_radon_transform(Img[n], angular_sampling) + line_proj[:, :, n], Img_pft[:, :, n] = self.fast_radon_transform( + Img[n], angular_sampling + ) - dim_wave=len(wemd_embed(line_proj[:,0,0])) - WE=np.zeros((dim_wave,n_theta,N)) + dim_wave = len(wemd_embed(line_proj[:, 0, 0])) + WE = np.zeros((dim_wave, n_theta, N)) for i in range(N): for theta in range(n_theta): - WE[:,theta,i]=wemd_embed(line_proj[:,theta,i]) - - - def fij(alpha,gamma,i,j,loss): - if loss=='l1': - Ii_hat=Img_pft[:,:,i]; Ij_hat=Img_pft[:,:,j] - idxi=np.round((alpha-np.pi/2)*n_theta/2/np.pi)% n_theta - idxj=np.round((-gamma-np.pi/2)*n_theta/2/np.pi)% n_theta - - Si=Ii_hat[:,int(idxi)]; Sj=Ij_hat[:,int(idxj)] - return np.linalg.norm(Si-Sj,1) - - if loss=='wemd': - idxi=np.round((alpha-np.pi/2)*n_theta/2/np.pi)% n_theta - idxj=np.round((-gamma-np.pi/2)*n_theta/2/np.pi)% n_theta - - Si=WE[:,int(idxi),i]; Sj=WE[:,int(idxj),j] - return np.linalg.norm(Si-Sj,1) - - alpha_grid=np.arange(2*T)*np.pi/T - beta_grid=(2*np.arange(2*T)+1)*np.pi/4/T - gamma_grid=np.arange(2*T)*np.pi/T - - bT=np.zeros(2*T); - for l in range(2*T): - ss=0 + WE[:, theta, i] = wemd_embed(line_proj[:, theta, i]) + + def fij(alpha, gamma, i, j, loss): + if loss == "l1": + Ii_hat = Img_pft[:, :, i] + Ij_hat = Img_pft[:, :, j] + idxi = np.round((alpha - np.pi / 2) * n_theta / 2 / np.pi) % n_theta + idxj = np.round((-gamma - np.pi / 2) * n_theta / 2 / np.pi) % n_theta + + Si = Ii_hat[:, int(idxi)] + Sj = Ij_hat[:, int(idxj)] + return np.linalg.norm(Si - Sj, 1) + + if loss == "wemd": + idxi = np.round((alpha - np.pi / 2) * n_theta / 2 / np.pi) % n_theta + idxj = np.round((-gamma - np.pi / 2) * n_theta / 2 / np.pi) % n_theta + + Si = WE[:, int(idxi), i] + Sj = WE[:, int(idxj), j] + return np.linalg.norm(Si - Sj, 1) + + alpha_grid = np.arange(2 * T) * np.pi / T + beta_grid = (2 * np.arange(2 * T) + 1) * np.pi / 4 / T + gamma_grid = np.arange(2 * T) * np.pi / T + + bT = np.zeros(2 * T) + for l in range(2 * T): + ss = 0 for m in range(T): - ss=ss+np.sin(beta_grid[l]*(2*m+1))/(2*m+1) - bT[l]=2/T*np.sin(beta_grid[l])*ss - - BTK=[]; - for k in range(1,Lmax+1): - dk=2*k+1 - btk=np.zeros((dk,dk)) - for l in range(2*T): btk=btk+bT[l]*self.Wd(k,beta_grid[l]) + ss = ss + np.sin(beta_grid[l] * (2 * m + 1)) / (2 * m + 1) + bT[l] = 2 / T * np.sin(beta_grid[l]) * ss + + BTK = [] + for k in range(1, Lmax + 1): + dk = 2 * k + 1 + btk = np.zeros((dk, dk)) + for l in range(2 * T): + btk = btk + bT[l] * self.Wd(k, beta_grid[l]) BTK.append(btk.T) + def fijhat_k(k, F): + dk = 2 * k + 1 - def fijhat_k(k,F): - dk=2*k+1; - - exp_alpha_grid=np.zeros((2*T,dk),dtype=complex) - for m in range(-k,k+1): - exp_alpha_grid[:,m+k]=np.exp(1j*m*alpha_grid) + exp_alpha_grid = np.zeros((2 * T, dk), dtype=complex) + for m in range(-k, k + 1): + exp_alpha_grid[:, m + k] = np.exp(1j * m * alpha_grid) - exp_gamma_grid=np.zeros((2*T,dk),dtype=complex) - for m in range(-k,k+1): - exp_gamma_grid[:,m+k]=np.exp(1j*m*gamma_grid) + exp_gamma_grid = np.zeros((2 * T, dk), dtype=complex) + for m in range(-k, k + 1): + exp_gamma_grid[:, m + k] = np.exp(1j * m * gamma_grid) - S=(exp_alpha_grid.T@F@exp_gamma_grid).T - fhat=BTK[k-1]*S/4/T**2 + S = (exp_alpha_grid.T @ F @ exp_gamma_grid).T + fhat = BTK[k - 1] * S / 4 / T**2 return fhat - C=[] - for k in range(1,Lmax+1): - dk=2*k+1; C.append(np.zeros((N*dk,N*dk),dtype=complex)) + C = [] + for k in range(1, Lmax + 1): + dk = 2 * k + 1 + C.append(np.zeros((N * dk, N * dk), dtype=complex)) for i in range(N): - for j in range(i+1,N): - Fij=np.zeros((2*T,2*T)) - for j1 in range(2*T): - for j2 in range(2*T): - Fij[j1,j2]=fij(alpha_grid[j1],gamma_grid[j2],i,j,loss) - for k in range(1,Lmax+1): - dk=2*k+1; C[k-1][j*dk:(j+1)*dk,i*dk:(i+1)*dk]=fijhat_k(k,Fij)#*dk - for k in range(1,Lmax+1): - C[k-1]=C[k-1]+C[k-1].conj().T + for j in range(i + 1, N): + Fij = np.zeros((2 * T, 2 * T)) + for j1 in range(2 * T): + for j2 in range(2 * T): + Fij[j1, j2] = fij(alpha_grid[j1], gamma_grid[j2], i, j, loss) + for k in range(1, Lmax + 1): + dk = 2 * k + 1 + C[k - 1][j * dk : (j + 1) * dk, i * dk : (i + 1) * dk] = fijhat_k( + k, Fij + ) # *dk + for k in range(1, Lmax + 1): + C[k - 1] = C[k - 1] + C[k - 1].conj().T for i in range(N): - Fii=np.zeros((2*T,2*T)) - for j1 in range(2*T): - for j2 in range(2*T): - Fii[j1,j2]=fij(alpha_grid[j1],gamma_grid[j2],i,i,loss) - for k in range(1,Lmax+1): - dk=2*k+1 - C[k-1][i*dk:(i+1)*dk,i*dk:(i+1)*dk]=fijhat_k(k,Fii)#*dk - - - for k in range(1,Lmax+1): - [T,Tinv] = self.complex2real(k) - C[k-1] = np.real(np.kron(np.eye(N),Tinv) @ C[k-1]@np.kron(np.eye(N),T)) - C[k-1]=np.round(C[k-1],10) - - return C + Fii = np.zeros((2 * T, 2 * T)) + for j1 in range(2 * T): + for j2 in range(2 * T): + Fii[j1, j2] = fij(alpha_grid[j1], gamma_grid[j2], i, i, loss) + for k in range(1, Lmax + 1): + dk = 2 * k + 1 + C[k - 1][i * dk : (i + 1) * dk, i * dk : (i + 1) * dk] = fijhat_k( + k, Fii + ) # *dk + + for k in range(1, Lmax + 1): + [T, Tinv] = self.complex2real(k) + C[k - 1] = np.real( + np.kron(np.eye(N), Tinv) @ C[k - 1] @ np.kron(np.eye(N), T) + ) + C[k - 1] = np.round(C[k - 1], 10) + + return C @staticmethod def fast_radon_transform(array, angles, use_ramp=False): @@ -215,7 +227,7 @@ def fast_radon_transform(array, angles, use_ramp=False): ) pts = pts.astype(array.dtype) - #array = array.astype(np.float32) + # array = array.astype(np.float32) lines_f = nufft(array, pts).reshape((img_size, -1)) if img_size % 2 == 0: @@ -225,69 +237,87 @@ def fast_radon_transform(array, angles, use_ramp=False): freqs = np.abs(np.pi * y_idx) lines_f *= freqs[:, np.newaxis] - projections = np.real(xp.asnumpy(fft.centered_ifft(xp.asarray(lines_f), axis=0))) + projections = np.real( + xp.asnumpy(fft.centered_ifft(xp.asarray(lines_f), axis=0)) + ) return projections, lines_f - @staticmethod - def Wd(J,beta): + def Wd(J, beta): # compute Wigner small d matrix - d=np.zeros((2*J+1,2*J+1)); - - for m in range(-J,J+1): - for n in range(-J,J+1): - smin=max(0,m-n); smax=min(J+m,J-n); - for s in range(smin,smax+1): - mul=np.sqrt(factorial(J+m))/factorial(J+m-s)*np.sqrt(factorial(J+n))/factorial(s)*np.sqrt(factorial(J-m))/factorial(n-m+s)\ - *np.sqrt(factorial(J-n))/factorial(J-n-s) - d[n+J,m+J]+=mul*(-1)**(n-m+s)*(np.cos(beta/2))**(2*J+m-n-2*s)*(np.sin(beta/2))**(n-m+2*s); - - return d + d = np.zeros((2 * J + 1, 2 * J + 1)) + + for m in range(-J, J + 1): + for n in range(-J, J + 1): + smin = max(0, m - n) + smax = min(J + m, J - n) + for s in range(smin, smax + 1): + mul = ( + np.sqrt(factorial(J + m)) + / factorial(J + m - s) + * np.sqrt(factorial(J + n)) + / factorial(s) + * np.sqrt(factorial(J - m)) + / factorial(n - m + s) + * np.sqrt(factorial(J - n)) + / factorial(J - n - s) + ) + d[n + J, m + J] += ( + mul + * (-1) ** (n - m + s) + * (np.cos(beta / 2)) ** (2 * J + m - n - 2 * s) + * (np.sin(beta / 2)) ** (n - m + 2 * s) + ) + + return d @staticmethod def complex2real(ell): # compute transformation matrices that convert complex representations to real ones - diml=2*ell+1 - Tinv=np.zeros((diml,diml),dtype=complex) + diml = 2 * ell + 1 + Tinv = np.zeros((diml, diml), dtype=complex) for i in range(diml): - if iell: - Tinv[i,i]=(-1)**(i-ell)/np.sqrt(2); Tinv[i,diml-1-i]=1/np.sqrt(2) - - T=np.zeros((diml,diml),dtype=complex) + if i < ell: + Tinv[i, i] = 1j / np.sqrt(2) + Tinv[i, diml - 1 - i] = -1j * (-1) ** (i - ell) / np.sqrt(2) + if i == ell: + Tinv[i, i] = 1 + if i > ell: + Tinv[i, i] = (-1) ** (i - ell) / np.sqrt(2) + Tinv[i, diml - 1 - i] = 1 / np.sqrt(2) + + T = np.zeros((diml, diml), dtype=complex) for i in range(diml): - if iell: - T[i,i]=(-1)**(i-ell)/np.sqrt(2); T[i,diml-1-i]=1j*(-1)**(i-ell)/np.sqrt(2) - return T,Tinv - + if i < ell: + T[i, i] = -1j / np.sqrt(2) + T[i, diml - 1 - i] = 1 / np.sqrt(2) + if i == ell: + T[i, i] = 1 + if i > ell: + T[i, i] = (-1) ** (i - ell) / np.sqrt(2) + T[i, diml - 1 - i] = 1j * (-1) ** (i - ell) / np.sqrt(2) + return T, Tinv ############# # ADMM Step # ############# def admm_sym_J( - self, - sym, - C, - Lmax, - N, - Ngrid, - max_iter, - rho, - ratio, - factor, - mult=1, - Nstep_yI=20, - verbose=True, - GPU=True, + self, + sym, + C, + Lmax, + N, + Ngrid, + max_iter, + rho, + ratio, + factor, + mult=1, + Nstep_yI=20, + verbose=True, + GPU=True, ): # admm for symmetric case if GPU: @@ -295,236 +325,463 @@ def admm_sym_J( else: import numpy as cp - C0,C1,normC,AEq,bEq,AEqAEqtinv,AI_mat_diag,AI_mat_offdiag,bI,Lambda,d0,d1,D0,D1,idx_diag,idx_offdiag,IDX_upper,IDX_lower,X0,X1,Xq,S0,S1,Sq=self.ADMM_preprocessing(C,Lmax,N,Ngrid,GPU) + ( + C0, + C1, + normC, + AEq, + bEq, + AEqAEqtinv, + AI_mat_diag, + AI_mat_offdiag, + bI, + Lambda, + d0, + d1, + D0, + D1, + idx_diag, + idx_offdiag, + IDX_upper, + IDX_lower, + X0, + X1, + Xq, + S0, + S1, + Sq, + ) = self.ADMM_preprocessing(C, Lmax, N, Ngrid, GPU) # S=int(sym[1]); sym_euler=Symmetry_Euler(sym) - rank_Ak,_=self.compute_rank(sym,Lmax); print(rank_Ak) + rank_Ak, _ = self.compute_rank(sym, Lmax) + print(rank_Ak) # rank_Ak=cp.zeros(Lmax) # for ell in range(Lmax): rank_Ak[ell]=np.linalg.matrix_rank(Ak(ell+1,sym_euler)) - AE=[]; AEAETinv=[] - for k in range(1,Lmax+1): - s0=k**2; s1=(k+1)**2 - AEk=cp.zeros((1+s0+s1,2*(s0+s1))) - AEk[0,:s0]=cp.eye(k).T.reshape(-1); AEk[0,s0:s0+s1]=cp.eye(k+1).T.reshape(-1) - for count in range(1,1+s0): - AEk[count,count-1]=1; AEk[count,count-1+s0+s1]=1 - for count in range(1+s0,1+s0+s1): - AEk[count,count-1]=1; AEk[count,count-1+s1+s0]=1 + AE = [] + AEAETinv = [] + for k in range(1, Lmax + 1): + s0 = k**2 + s1 = (k + 1) ** 2 + AEk = cp.zeros((1 + s0 + s1, 2 * (s0 + s1))) + AEk[0, :s0] = cp.eye(k).T.reshape(-1) + AEk[0, s0 : s0 + s1] = cp.eye(k + 1).T.reshape(-1) + for count in range(1, 1 + s0): + AEk[count, count - 1] = 1 + AEk[count, count - 1 + s0 + s1] = 1 + for count in range(1 + s0, 1 + s0 + s1): + AEk[count, count - 1] = 1 + AEk[count, count - 1 + s1 + s0] = 1 AE.append(AEk) - AEAETinv.append(np.linalg.pinv(AEk@AEk.T)) - bE=cp.zeros((Lmax+D0+D1)) + AEAETinv.append(np.linalg.pinv(AEk @ AEk.T)) + bE = cp.zeros((Lmax + D0 + D1)) for k in range(Lmax): - bE[k+d0[k]+d1[k]:]=rank_Ak[k] - bE[k+1+d0[k]+d1[k]:k+1+d0[k+1]+d1[k]]=cp.eye(k+1).T.reshape(-1) - bE[k+1+d0[k+1]+d1[k]:k+1+d0[k+1]+d1[k+1]]=cp.eye(k+2).T.reshape(-1) - bE=np.repeat(bE[:,np.newaxis],N,axis=1) - P=[]; - for k in range(1,Lmax+1): - dk=2*k+1; Pk=cp.eye(dk) + bE[k + d0[k] + d1[k] :] = rank_Ak[k] + bE[k + 1 + d0[k] + d1[k] : k + 1 + d0[k + 1] + d1[k]] = cp.eye( + k + 1 + ).T.reshape(-1) + bE[k + 1 + d0[k + 1] + d1[k] : k + 1 + d0[k + 1] + d1[k + 1]] = cp.eye( + k + 2 + ).T.reshape(-1) + bE = np.repeat(bE[:, np.newaxis], N, axis=1) + P = [] + for k in range(1, Lmax + 1): + dk = 2 * k + 1 + Pk = cp.eye(dk) for m in range(k): - for l in range(k-m): - Pk[(m+2*l,m+2*l+1),:]=Pk[(m+2*l+1,m+2*l),:] - P.append(Pk) - - - def fun_AE(X0,X1,Xd0,Xd1,Xq): - z=cp.zeros((Lmax+D0+D1,N)) + for l in range(k - m): + Pk[(m + 2 * l, m + 2 * l + 1), :] = Pk[ + (m + 2 * l + 1, m + 2 * l), : + ] + P.append(Pk) + + def fun_AE(X0, X1, Xd0, Xd1, Xq): + z = cp.zeros((Lmax + D0 + D1, N)) for k in range(Lmax): - z[k+d0[k]+d1[k]:k+1+d0[k+1]+d1[k+1]]=AE[k]@ cp.concatenate((X0[d0[k]:d0[k+1],idx_diag],X1[d1[k]:d1[k+1],idx_diag],Xd0[d0[k]:d0[k+1]],Xd1[d1[k]:d1[k+1]]),axis=0) - return z, AEq@cp.concatenate((Xq,X0[:1,idx_offdiag],X1[:4,idx_offdiag]),axis=0) - - def fun_AET(yE,yEq): - Z0=cp.zeros((D0,N*(N+1)//2)); Z1=cp.zeros((D1,N*(N+1)//2)); - Zd0=cp.zeros((D0,N)); Zd1=cp.zeros((D1,N)) + z[k + d0[k] + d1[k] : k + 1 + d0[k + 1] + d1[k + 1]] = AE[ + k + ] @ cp.concatenate( + ( + X0[d0[k] : d0[k + 1], idx_diag], + X1[d1[k] : d1[k + 1], idx_diag], + Xd0[d0[k] : d0[k + 1]], + Xd1[d1[k] : d1[k + 1]], + ), + axis=0, + ) + return z, AEq @ cp.concatenate( + (Xq, X0[:1, idx_offdiag], X1[:4, idx_offdiag]), axis=0 + ) + + def fun_AET(yE, yEq): + Z0 = cp.zeros((D0, N * (N + 1) // 2)) + Z1 = cp.zeros((D1, N * (N + 1) // 2)) + Zd0 = cp.zeros((D0, N)) + Zd1 = cp.zeros((D1, N)) for k in range(Lmax): - s0=(k+1)**2; s1=(k+2)**2 - Ztmp=AE[k].T@yE[k+d0[k]+d1[k]:k+1+d0[k+1]+d1[k+1]] - Z0[d0[k]:d0[k+1],idx_diag]=Ztmp[:s0] - Z1[d1[k]:d1[k+1],idx_diag]=Ztmp[s0:s0+s1] - Zd0[d0[k]:d0[k+1]]=Ztmp[s0+s1:2*s0+s1] - Zd1[d1[k]:d1[k+1]]=Ztmp[2*s0+s1:2*s0+2*s1] - Zq=AEq.T@yEq - Z0[:1,idx_offdiag]=Zq[16:17]; Z1[:4,idx_offdiag]=Zq[17:] - return Z0,Z1,Zd0,Zd1,Zq[:16] - - def fun_AI(X0,X1): - z=cp.zeros((Ngrid,N*(N+1)//2)); tmp=cp.concatenate((X0,X1),axis=0); - z[:,idx_diag]=AI_mat_diag@tmp[:,idx_diag]; z[:,idx_offdiag]=AI_mat_offdiag@tmp[:,idx_offdiag] + s0 = (k + 1) ** 2 + s1 = (k + 2) ** 2 + Ztmp = AE[k].T @ yE[k + d0[k] + d1[k] : k + 1 + d0[k + 1] + d1[k + 1]] + Z0[d0[k] : d0[k + 1], idx_diag] = Ztmp[:s0] + Z1[d1[k] : d1[k + 1], idx_diag] = Ztmp[s0 : s0 + s1] + Zd0[d0[k] : d0[k + 1]] = Ztmp[s0 + s1 : 2 * s0 + s1] + Zd1[d1[k] : d1[k + 1]] = Ztmp[2 * s0 + s1 : 2 * s0 + 2 * s1] + Zq = AEq.T @ yEq + Z0[:1, idx_offdiag] = Zq[16:17] + Z1[:4, idx_offdiag] = Zq[17:] + return Z0, Z1, Zd0, Zd1, Zq[:16] + + def fun_AI(X0, X1): + z = cp.zeros((Ngrid, N * (N + 1) // 2)) + tmp = cp.concatenate((X0, X1), axis=0) + z[:, idx_diag] = AI_mat_diag @ tmp[:, idx_diag] + z[:, idx_offdiag] = AI_mat_offdiag @ tmp[:, idx_offdiag] # return AI_mat@cp.concatenate((X0,X1),axis=0) return z def fun_AIT(yI): # Z=AI_mat.T@yI - # return Z[:D0,:], Z[D0:,:] - Z=cp.zeros((D0+D1,N*(N+1)//2)); Z[:,idx_diag]=AI_mat_diag.T@yI[:,idx_diag]; Z[:,idx_offdiag]=AI_mat_offdiag.T@yI[:,idx_offdiag]; - return Z[:D0,:], Z[D0:,:] - - def update_S(C0,C1,yE,yEq,yI,X0,X1,Xd0,Xd1,Xq,rho,Lmax,N): - Z0,Z1,Zd0,Zd1,Zq=fun_AET(yE,yEq); AIT_yI0,AIT_yI1=fun_AIT(yI) - S0=C0-Z0-AIT_yI0-X0/rho; S1=C1-Z1-AIT_yI1-X1/rho; - Sd0=-Zd0-Xd0/rho; Sd1=-Zd1-Xd1/rho; Sq=-Zq-Xq/rho - tic0=time.perf_counter(); - for k in range(1,Lmax+1): - tmp=self.mat_block(S0[d0[k-1]:d0[k],:],N,k,IDX_upper,IDX_lower,idx_offdiag) - tmp=self.psd_projection(tmp) - S0[d0[k-1]:d0[k],:]=self.vec_block(tmp,N,k,IDX_upper) - - tmp=self.mat_block(S1[d1[k-1]:d1[k],:],N,k+1,IDX_upper,IDX_lower,idx_offdiag) - tmp=self.psd_projection(tmp) - S1[d1[k-1]:d1[k],:]=self.vec_block(tmp,N,k+1,IDX_upper) - toc0=time.perf_counter(); Time[0]+=toc0-tic0 - tic1=time.perf_counter(); - for k in range(1,Lmax+1): + # return Z[:D0,:], Z[D0:,:] + Z = cp.zeros((D0 + D1, N * (N + 1) // 2)) + Z[:, idx_diag] = AI_mat_diag.T @ yI[:, idx_diag] + Z[:, idx_offdiag] = AI_mat_offdiag.T @ yI[:, idx_offdiag] + return Z[:D0, :], Z[D0:, :] + + def update_S(C0, C1, yE, yEq, yI, X0, X1, Xd0, Xd1, Xq, rho, Lmax, N): + Z0, Z1, Zd0, Zd1, Zq = fun_AET(yE, yEq) + AIT_yI0, AIT_yI1 = fun_AIT(yI) + S0 = C0 - Z0 - AIT_yI0 - X0 / rho + S1 = C1 - Z1 - AIT_yI1 - X1 / rho + Sd0 = -Zd0 - Xd0 / rho + Sd1 = -Zd1 - Xd1 / rho + Sq = -Zq - Xq / rho + tic0 = time.perf_counter() + for k in range(1, Lmax + 1): + tmp = self.mat_block( + S0[d0[k - 1] : d0[k], :], N, k, IDX_upper, IDX_lower, idx_offdiag + ) + tmp = self.psd_projection(tmp) + S0[d0[k - 1] : d0[k], :] = self.vec_block(tmp, N, k, IDX_upper) + + tmp = self.mat_block( + S1[d1[k - 1] : d1[k], :], + N, + k + 1, + IDX_upper, + IDX_lower, + idx_offdiag, + ) + tmp = self.psd_projection(tmp) + S1[d1[k - 1] : d1[k], :] = self.vec_block(tmp, N, k + 1, IDX_upper) + toc0 = time.perf_counter() + Time[0] += toc0 - tic0 + tic1 = time.perf_counter() + for k in range(1, Lmax + 1): for n in range(N): - tmp=self.transform_back_block(Sd0[d0[k-1]:d0[k],n],Sd1[d1[k-1]:d1[k],n],k,P[k-1]) - tmp=self.psd_projection(tmp) - Sd0[d0[k-1]:d0[k],n],Sd1[d1[k-1]:d1[k],n]=self.transform_block(tmp,k,P[k-1]) - toc1=time.perf_counter(); Time[1]+=toc1-tic1 - tic2=time.perf_counter(); - for count in range(N*(N-1)//2): - tmp=self.psd_projection(Sq[:,count].reshape(4,4).T) - Sq[:,count]=tmp.T.reshape(16) - toc2=time.perf_counter(); Time[2]+=toc2-tic2 - return S0,S1,Sd0,Sd1,Sq - - - def update_yE(C0,C1,X0,X1,Xd0,Xd1,Xq,S0,S1,Sd0,Sd1,Sq,yI,rho): - AIT_yI0,AIT_yI1=fun_AIT(yI) - z,zq=fun_AE(-X0/rho+C0-S0-AIT_yI0,-X1/rho+C1-S1-AIT_yI1,-Xd0/rho-Sd0,-Xd1/rho-Sd1,-Xq/rho-Sq) - yE=cp.zeros((Lmax+D0+D1,N)) + tmp = self.transform_back_block( + Sd0[d0[k - 1] : d0[k], n], + Sd1[d1[k - 1] : d1[k], n], + k, + P[k - 1], + ) + tmp = self.psd_projection(tmp) + Sd0[d0[k - 1] : d0[k], n], Sd1[d1[k - 1] : d1[k], n] = ( + self.transform_block(tmp, k, P[k - 1]) + ) + toc1 = time.perf_counter() + Time[1] += toc1 - tic1 + tic2 = time.perf_counter() + for count in range(N * (N - 1) // 2): + tmp = self.psd_projection(Sq[:, count].reshape(4, 4).T) + Sq[:, count] = tmp.T.reshape(16) + toc2 = time.perf_counter() + Time[2] += toc2 - tic2 + return S0, S1, Sd0, Sd1, Sq + + def update_yE(C0, C1, X0, X1, Xd0, Xd1, Xq, S0, S1, Sd0, Sd1, Sq, yI, rho): + AIT_yI0, AIT_yI1 = fun_AIT(yI) + z, zq = fun_AE( + -X0 / rho + C0 - S0 - AIT_yI0, + -X1 / rho + C1 - S1 - AIT_yI1, + -Xd0 / rho - Sd0, + -Xd1 / rho - Sd1, + -Xq / rho - Sq, + ) + yE = cp.zeros((Lmax + D0 + D1, N)) for k in range(Lmax): - yE[k+d0[k]+d1[k]:k+1+d0[k+1]+d1[k+1]]=AEAETinv[k]@(bE[k+d0[k]+d1[k]:k+1+d0[k+1]+d1[k+1]]/rho+z[k+d0[k]+d1[k]:k+1+d0[k+1]+d1[k+1]]) - yEq=AEqAEqtinv@(bEq/rho+zq) - return yE,yEq - - - def update_yI(C0,C1,X0,X1,S0,S1,yE,yEq,yI,rho,Lambda): - Z0,Z1,_,_,_=fun_AET(yE,yEq); AIT_yI0,AIT_yI1=fun_AIT(yI) - tmp=fun_AI(-X0/rho+C0-S0-Z0-AIT_yI0,-X1/rho+C1-S1-Z1-AIT_yI1) - yI=yI+bI/rho/Lambda+tmp/Lambda - yI=np.maximum(yI,0) + yE[k + d0[k] + d1[k] : k + 1 + d0[k + 1] + d1[k + 1]] = AEAETinv[k] @ ( + bE[k + d0[k] + d1[k] : k + 1 + d0[k + 1] + d1[k + 1]] / rho + + z[k + d0[k] + d1[k] : k + 1 + d0[k + 1] + d1[k + 1]] + ) + yEq = AEqAEqtinv @ (bEq / rho + zq) + return yE, yEq + + def update_yI(C0, C1, X0, X1, S0, S1, yE, yEq, yI, rho, Lambda): + Z0, Z1, _, _, _ = fun_AET(yE, yEq) + AIT_yI0, AIT_yI1 = fun_AIT(yI) + tmp = fun_AI( + -X0 / rho + C0 - S0 - Z0 - AIT_yI0, -X1 / rho + C1 - S1 - Z1 - AIT_yI1 + ) + yI = yI + bI / rho / Lambda + tmp / Lambda + yI = np.maximum(yI, 0) return yI - def update_X(C0,C1,X0,X1,Xd0,Xd1,Xq,yE,yEq,yI,S0,S1,Sd0,Sd1,Sq,rho): - Z0,Z1,Zd0,Zd1,Zq=fun_AET(yE,yEq); AIT_yI0,AIT_yI1=fun_AIT(yI) - tmp=S0+Z0+AIT_yI0-C0; X0=X0+mult*rho*tmp; resX0=np.linalg.norm(tmp) - tmp=S1+Z1+AIT_yI1-C1; X1=X1+mult*rho*tmp; resX1=np.linalg.norm(tmp) - tmp=Sd0+Zd0; Xd0=Xd0+mult*rho*tmp; resXd0=np.linalg.norm(tmp) - tmp=Sd1+Zd1; Xd1=Xd1+mult*rho*tmp; resXd1=np.linalg.norm(tmp) - tmp=Sq+Zq; Xq=Xq+mult*rho*tmp; resXq=np.linalg.norm(tmp) - return X0,X1,Xd0,Xd1,Xq,np.sqrt(resX0**2+resX1**2+resXd0**2+resXd1**2+resXq**2) - - def update_rho(X0,X1,Xd0,Xd1,Xq,bE,bEq,bI,res_X,rho,factor,normC): - z,zq=fun_AE(X0,X1,Xd0,Xd1,Xq) - res_eq=np.linalg.norm(z-bE)/(1+np.linalg.norm(bE))+np.linalg.norm(zq-bEq)/(1+np.linalg.norm(bEq)); - res_inq=np.linalg.norm(np.maximum(bI-fun_AI(X0,X1),0))/(1+abs(bI)*np.sqrt(Ngrid*N**2)); - p_resnorm=res_eq+res_inq - d_resnorm=res_X/(1+normC) - if d_resnorm>ratio*p_resnorm: rho=rho*factor - if d_resnorm ratio * p_resnorm: + rho = rho * factor + if d_resnorm < ratio * p_resnorm: + rho = rho / factor + return rho, p_resnorm, d_resnorm def print_updates(verbose=True): # X_admm=transform_coeff_back(X0,X1,Lmax,N); obj_p=0 # for k in range(Lmax): obj_p+=cp.trace(C[k]@X_admm[k]) - if verbose: - obj_p=cp.vdot(C0[:,idx_diag],X0[:,idx_diag])+cp.vdot(C1[:,idx_diag],X1[:,idx_diag])+2*cp.vdot(C0[:,idx_offdiag],X0[:,idx_offdiag])+2*cp.vdot(C1[:,idx_offdiag],X1[:,idx_offdiag]) - obj_d=cp.vdot(yE,bE)+2*cp.vdot(yEq,bEq)+cp.vdot(yI[:,idx_diag],bI*cp.ones((Ngrid,N)))+2*cp.vdot(yI[:,idx_offdiag],bI*cp.ones((Ngrid,N*(N-1)//2))); - - - z,zq=fun_AE(X0,X1,Xd0,Xd1,Xq); - res_eq=np.linalg.norm(z-bE)/(1+np.linalg.norm(bE))+np.linalg.norm(zq-bEq)/(1+np.linalg.norm(bEq)); - res_inq=np.linalg.norm(np.maximum(bI-fun_AI(X0,X1),0))/(1+abs(bI)*np.sqrt(Ngrid*N*(N+1)/2)); - res_psdX=0 - for k in range(1,Lmax+1): - tmp=self.mat_block(X0[d0[k-1]:d0[k],:],N,k,IDX_upper,IDX_lower,idx_offdiag); res_psdX+=np.linalg.norm(self.psd_projection(-tmp)); #res_psdX+=norm(self.psd_projection(-tmp))/(1+norm(tmp)) - tmp=self.mat_block(X1[d1[k-1]:d1[k],:],N,k+1,IDX_upper,IDX_lower,idx_offdiag); res_psdX+=np.linalg.norm(self.psd_projection(-tmp)); #res_psdX+=norm(self.psd_projection(-tmp))/(1+norm(tmp)) - res_psdX=res_psdX/(1+np.linalg.norm(X0)+np.linalg.norm(X1)) - res_psdD=0 - for k in range(1,Lmax+1): + if verbose: + obj_p = ( + cp.vdot(C0[:, idx_diag], X0[:, idx_diag]) + + cp.vdot(C1[:, idx_diag], X1[:, idx_diag]) + + 2 * cp.vdot(C0[:, idx_offdiag], X0[:, idx_offdiag]) + + 2 * cp.vdot(C1[:, idx_offdiag], X1[:, idx_offdiag]) + ) + obj_d = ( + cp.vdot(yE, bE) + + 2 * cp.vdot(yEq, bEq) + + cp.vdot(yI[:, idx_diag], bI * cp.ones((Ngrid, N))) + + 2 + * cp.vdot( + yI[:, idx_offdiag], bI * cp.ones((Ngrid, N * (N - 1) // 2)) + ) + ) + + z, zq = fun_AE(X0, X1, Xd0, Xd1, Xq) + res_eq = np.linalg.norm(z - bE) / ( + 1 + np.linalg.norm(bE) + ) + np.linalg.norm(zq - bEq) / (1 + np.linalg.norm(bEq)) + res_inq = np.linalg.norm(np.maximum(bI - fun_AI(X0, X1), 0)) / ( + 1 + abs(bI) * np.sqrt(Ngrid * N * (N + 1) / 2) + ) + res_psdX = 0 + for k in range(1, Lmax + 1): + tmp = self.mat_block( + X0[d0[k - 1] : d0[k], :], + N, + k, + IDX_upper, + IDX_lower, + idx_offdiag, + ) + res_psdX += np.linalg.norm(self.psd_projection(-tmp)) + # res_psdX+=norm(self.psd_projection(-tmp))/(1+norm(tmp)) + tmp = self.mat_block( + X1[d1[k - 1] : d1[k], :], + N, + k + 1, + IDX_upper, + IDX_lower, + idx_offdiag, + ) + res_psdX += np.linalg.norm(self.psd_projection(-tmp)) + # res_psdX+=norm(self.psd_projection(-tmp))/(1+norm(tmp)) + res_psdX = res_psdX / (1 + np.linalg.norm(X0) + np.linalg.norm(X1)) + res_psdD = 0 + for k in range(1, Lmax + 1): for n in range(N): - tmp=self.transform_back_block(Xd0[d0[k-1]:d0[k],n],Xd1[d1[k-1]:d1[k],n],k,P[k-1]) - res_psdD+=np.linalg.norm(self.psd_projection(-tmp)); #res_psdD+=norm(self.psd_projection(-tmp))/(1+norm(tmp)) - res_psdD=res_psdD/(1+np.linalg.norm(Xd0)+np.linalg.norm(Xd1)) - res_psdQ=0 - for count in range(N*(N-1)//2): - tmp=Xq[:,count].reshape(4,4).T; res_psdQ+=np.linalg.norm(self.psd_projection(-tmp)); #res_psdQ+=norm(self.psd_projection(-tmp))/(1+norm(tmp)) - res_psdQ=res_psdQ/(1+np.linalg.norm(Xq)) - - - normS=np.sqrt(np.linalg.norm(S0)**2+np.linalg.norm(S1)**2+np.linalg.norm(Sd0)**2+np.linalg.norm(Sd1)**2+np.linalg.norm(Sq)**2) - normX=np.sqrt(np.linalg.norm(X0)**2+np.linalg.norm(X1)**2+np.linalg.norm(Xd0)**2+np.linalg.norm(Xd1)**2+np.linalg.norm(Xq)**2) - p_res=res_eq+res_inq+res_psdX+res_psdD+res_psdQ - d_res=res_X/(1+normC) - print('Iter %i'%t+': p_res=%1.5f'%p_res+', d_res=%1.5f'%d_res+', obj_primal=%1.2f'%obj_p+', obj_dual=%1.2f'%obj_d+', duality gap=%1.2f'%(obj_p-obj_d)+ - '\n eq_res=%1.5f'%res_eq+', inq_res=%1.5f'%res_inq+', psd_res=%1.5f'%(res_psdX+res_psdD+res_psdQ)+', |S|=%1.2f'%normS+', |X|=%1.2f'%normX) - - - Xd0=cp.zeros((D0,N)); Xd1=cp.zeros((D1,N)); Sd0=cp.zeros(Xd0.shape); Sd1=cp.zeros(Xd1.shape); - yI=cp.zeros((Ngrid,N*(N+1)//2)); yE=cp.zeros(bE.shape); yEq=cp.zeros(bEq.shape) - - IDX=np.arange(3); Time=np.zeros(4) + tmp = self.transform_back_block( + Xd0[d0[k - 1] : d0[k], n], + Xd1[d1[k - 1] : d1[k], n], + k, + P[k - 1], + ) + res_psdD += np.linalg.norm(self.psd_projection(-tmp)) + # res_psdD+=norm(self.psd_projection(-tmp))/(1+norm(tmp)) + res_psdD = res_psdD / (1 + np.linalg.norm(Xd0) + np.linalg.norm(Xd1)) + res_psdQ = 0 + for count in range(N * (N - 1) // 2): + tmp = Xq[:, count].reshape(4, 4).T + res_psdQ += np.linalg.norm(self.psd_projection(-tmp)) + # res_psdQ+=norm(self.psd_projection(-tmp))/(1+norm(tmp)) + res_psdQ = res_psdQ / (1 + np.linalg.norm(Xq)) + + normS = np.sqrt( + np.linalg.norm(S0) ** 2 + + np.linalg.norm(S1) ** 2 + + np.linalg.norm(Sd0) ** 2 + + np.linalg.norm(Sd1) ** 2 + + np.linalg.norm(Sq) ** 2 + ) + normX = np.sqrt( + np.linalg.norm(X0) ** 2 + + np.linalg.norm(X1) ** 2 + + np.linalg.norm(Xd0) ** 2 + + np.linalg.norm(Xd1) ** 2 + + np.linalg.norm(Xq) ** 2 + ) + p_res = res_eq + res_inq + res_psdX + res_psdD + res_psdQ + d_res = res_X / (1 + normC) + print( + "Iter %i" % t + + ": p_res=%1.5f" % p_res + + ", d_res=%1.5f" % d_res + + ", obj_primal=%1.2f" % obj_p + + ", obj_dual=%1.2f" % obj_d + + ", duality gap=%1.2f" % (obj_p - obj_d) + + "\n eq_res=%1.5f" % res_eq + + ", inq_res=%1.5f" % res_inq + + ", psd_res=%1.5f" % (res_psdX + res_psdD + res_psdQ) + + ", |S|=%1.2f" % normS + + ", |X|=%1.2f" % normX + ) + + Xd0 = cp.zeros((D0, N)) + Xd1 = cp.zeros((D1, N)) + Sd0 = cp.zeros(Xd0.shape) + Sd1 = cp.zeros(Xd1.shape) + yI = cp.zeros((Ngrid, N * (N + 1) // 2)) + yE = cp.zeros(bE.shape) + yEq = cp.zeros(bEq.shape) + + IDX = np.arange(3) + Time = np.zeros(4) for t in range(max_iter): np.random.shuffle(IDX) for idx in IDX: - if idx==0: S0,S1,Sd0,Sd1,Sq=update_S(C0,C1,yE,yEq,yI,X0,X1,Xd0,Xd1,Xq,rho,Lmax,N); - if idx==1: yE,yEq=update_yE(C0,C1,X0,X1,Xd0,Xd1,Xq,S0,S1,Sd0,Sd1,Sq,yI,rho); - if idx==2: + if idx == 0: + S0, S1, Sd0, Sd1, Sq = update_S( + C0, C1, yE, yEq, yI, X0, X1, Xd0, Xd1, Xq, rho, Lmax, N + ) + if idx == 1: + yE, yEq = update_yE( + C0, C1, X0, X1, Xd0, Xd1, Xq, S0, S1, Sd0, Sd1, Sq, yI, rho + ) + if idx == 2: for kk in range(Nstep_yI): - yI=update_yI(C0,C1,X0,X1,S0,S1,yE,yEq,yI,rho,Lambda) - X0,X1,Xd0,Xd1,Xq,res_X=update_X(C0,C1,X0,X1,Xd0,Xd1,Xq,yE,yEq,yI,S0,S1,Sd0,Sd1,Sq,rho); - if t%100==0: print_updates(verbose) - rho,p_resnorm,d_resnorm=update_rho(X0,X1,Xd0,Xd1,Xq,bE,bEq,bI,res_X,rho,factor,normC) - - X_admm=self.transform_coeff_back(X0,X1,Lmax,N,IDX_upper,IDX_lower,idx_offdiag) + yI = update_yI(C0, C1, X0, X1, S0, S1, yE, yEq, yI, rho, Lambda) + X0, X1, Xd0, Xd1, Xq, res_X = update_X( + C0, C1, X0, X1, Xd0, Xd1, Xq, yE, yEq, yI, S0, S1, Sd0, Sd1, Sq, rho + ) + if t % 100 == 0: + print_updates(verbose) + rho, p_resnorm, d_resnorm = update_rho( + X0, X1, Xd0, Xd1, Xq, bE, bEq, bI, res_X, rho, factor, normC + ) + + X_admm = self.transform_coeff_back( + X0, X1, Lmax, N, IDX_upper, IDX_lower, idx_offdiag + ) if GPU: - for k in range(Lmax): X_admm[k]=X_admm[k].get() - return X_admm - + for k in range(Lmax): + X_admm[k] = X_admm[k].get() + return X_admm - def ADMM_preprocessing(self, C,Lmax,N,Ngrid,GPU=True): + def ADMM_preprocessing(self, C, Lmax, N, Ngrid, GPU=True): # compute necessary quantities for ADMM - if GPU: import cupy as cp - else: import numpy as cp + if GPU: + import cupy as cp + else: + import numpy as cp - # compute some useful index sets - count=0; idx_diag=[]; idx_offdiag=[]; + # compute some useful index sets + count = 0 + idx_diag = [] + idx_offdiag = [] for i in range(N): - for j in range(i,N): - if j==i: idx_diag.append(count) - else: idx_offdiag.append(count) - count+=1 - IDX_upper=[]; IDX_lower=[]; + for j in range(i, N): + if j == i: + idx_diag.append(count) + else: + idx_offdiag.append(count) + count += 1 + IDX_upper = [] + IDX_lower = [] for i in range(N): for j in range(N): - if j>=i: IDX_upper.append(i*N+j); + if j >= i: + IDX_upper.append(i * N + j) for j in range(N): for i in range(N): - if j500: - Lambda=cp.linalg.norm(z); z=z/cp.linalg.norm(z); z=AI@(AI.T@z) - Lambda+=2000; print('Largest eigenvalue of AIAIT is approximately %1.2f'%Lambda) + z = cp.random.normal(0, 1, (Ngrid, N**2)) + Lambda = 0 + + while abs(Lambda - cp.linalg.norm(z)) > 500: + Lambda = cp.linalg.norm(z) + z = z / cp.linalg.norm(z) + z = AI @ (AI.T @ z) + Lambda += 2000 + print("Largest eigenvalue of AIAIT is approximately %1.2f" % Lambda) # Lambda=cp.linalg.eigvalsh(AI@AI.T)[-1]; print(Lambda) - return Lambda - - def compute_rank(self, sym,Lmax): - sym_euler,S=self.Symmetry_Euler(sym); rk=cp.zeros(Lmax); A=[] - for k in range(1,Lmax+1): - dk=2*k+1; Ak=np.zeros((dk,dk),dtype=np.complex128) - for s in range(S): Ak+=self.WD(k,sym_euler[s]) - Ak=np.round(Ak/S,6); A.append(Ak); rk[k-1]=np.linalg.matrix_rank(Ak) - return rk,A + return Lambda + + def compute_rank(self, sym, Lmax): + sym_euler, S = self.Symmetry_Euler(sym) + rk = cp.zeros(Lmax) + A = [] + for k in range(1, Lmax + 1): + dk = 2 * k + 1 + Ak = np.zeros((dk, dk), dtype=np.complex128) + for s in range(S): + Ak += self.WD(k, sym_euler[s]) + Ak = np.round(Ak / S, 6) + A.append(Ak) + rk[k - 1] = np.linalg.matrix_rank(Ak) + return rk, A @staticmethod def Symmetry_Euler(sym): - if sym[0]=='C': - order=int(sym[1:]); sym_euler=np.zeros((order,3)) - for i in range(order): sym_euler[i]=spr.from_euler('zyx',[2*np.pi/order*i,0,0]).as_euler('zyz') - - if sym[0]=='D': - order=int(sym[1:]); sym_euler=np.zeros((2*order,3)) - for i in range(order): - sym_euler[i]=spr.from_euler('zyx',[2*np.pi/order*i,0,0]).as_euler('zyz') - sym_euler[i+order]=spr.from_euler('zyx',[2*np.pi/order*i,0,np.pi]).as_euler('zyz') - - if sym=='T12': - sym_euler=np.zeros((12,3)) - sym_euler[1]=spr.from_rotvec(2*np.pi/3*np.array([0,0,1])).as_euler('zyz') - sym_euler[2]=spr.from_rotvec(4*np.pi/3*np.array([0,0,1])).as_euler('zyz') - - - return sym_euler,sym_euler.shape[0] - - def WD(self, J,euler): + if sym[0] == "C": + order = int(sym[1:]) + sym_euler = np.zeros((order, 3)) + for i in range(order): + sym_euler[i] = spr.from_euler( + "zyx", [2 * np.pi / order * i, 0, 0] + ).as_euler("zyz") + + if sym[0] == "D": + order = int(sym[1:]) + sym_euler = np.zeros((2 * order, 3)) + for i in range(order): + sym_euler[i] = spr.from_euler( + "zyx", [2 * np.pi / order * i, 0, 0] + ).as_euler("zyz") + sym_euler[i + order] = spr.from_euler( + "zyx", [2 * np.pi / order * i, 0, np.pi] + ).as_euler("zyz") + + if sym == "T12": + sym_euler = np.zeros((12, 3)) + sym_euler[1] = spr.from_rotvec( + 2 * np.pi / 3 * np.array([0, 0, 1]) + ).as_euler("zyz") + sym_euler[2] = spr.from_rotvec( + 4 * np.pi / 3 * np.array([0, 0, 1]) + ).as_euler("zyz") + + return sym_euler, sym_euler.shape[0] + + def WD(self, J, euler): # compute Wigner D matrix - alpha=euler[0]; beta=euler[1]; gamma=euler[2]; d=self.Wd(J,beta); - D=np.diag(np.exp(-1j*alpha*np.arange(-J,J+1)))@ d @np.diag(np.exp(-1j*gamma*np.arange(-J,J+1))) + alpha = euler[0] + beta = euler[1] + gamma = euler[2] + d = self.Wd(J, beta) + D = ( + np.diag(np.exp(-1j * alpha * np.arange(-J, J + 1))) + @ d + @ np.diag(np.exp(-1j * gamma * np.arange(-J, J + 1))) + ) return D @staticmethod - def Wd(J,beta): + def Wd(J, beta): # compute Wigner small d matrix - d=np.zeros((2*J+1,2*J+1)); - - for m in range(-J,J+1): - for n in range(-J,J+1): - smin=max(0,m-n); smax=min(J+m,J-n); - for s in range(smin,smax+1): - mul=np.sqrt(factorial(J+m))/factorial(J+m-s)*np.sqrt(factorial(J+n))/factorial(s)*np.sqrt(factorial(J-m))/factorial(n-m+s)\ - *np.sqrt(factorial(J-n))/factorial(J-n-s) - d[n+J,m+J]+=mul*(-1)**(n-m+s)*(np.cos(beta/2))**(2*J+m-n-2*s)*(np.sin(beta/2))**(n-m+2*s); - - return d + d = np.zeros((2 * J + 1, 2 * J + 1)) + + for m in range(-J, J + 1): + for n in range(-J, J + 1): + smin = max(0, m - n) + smax = min(J + m, J - n) + for s in range(smin, smax + 1): + mul = ( + np.sqrt(factorial(J + m)) + / factorial(J + m - s) + * np.sqrt(factorial(J + n)) + / factorial(s) + * np.sqrt(factorial(J - m)) + / factorial(n - m + s) + * np.sqrt(factorial(J - n)) + / factorial(J - n - s) + ) + d[n + J, m + J] += ( + mul + * (-1) ** (n - m + s) + * (np.cos(beta / 2)) ** (2 * J + m - n - 2 * s) + * (np.sin(beta / 2)) ** (n - m + 2 * s) + ) + + return d @staticmethod - def mat_block(vecA,N,sz,IDX_upper,IDX_lower,idx_offdiag): - tmp=vecA.T.reshape(N*(N+1)//2,sz,sz).transpose(0,2,1) - AA=cp.zeros((N**2,sz,sz)); AA[IDX_upper]=tmp; AA[IDX_lower]=tmp[idx_offdiag].transpose(0,2,1) - return (AA.reshape(N,N,sz,sz).transpose(0,2,1,3)).reshape(N*sz,N*sz) + def mat_block(vecA, N, sz, IDX_upper, IDX_lower, idx_offdiag): + tmp = vecA.T.reshape(N * (N + 1) // 2, sz, sz).transpose(0, 2, 1) + AA = cp.zeros((N**2, sz, sz)) + AA[IDX_upper] = tmp + AA[IDX_lower] = tmp[idx_offdiag].transpose(0, 2, 1) + return (AA.reshape(N, N, sz, sz).transpose(0, 2, 1, 3)).reshape(N * sz, N * sz) @staticmethod def psd_projection(B): # compute the PSD part of a symmstric matrix - evals,evecs=cp.linalg.eigh((B+B.T)/2); evals=cp.maximum(evals,0) - return (evecs*evals) @ evecs.T + evals, evecs = cp.linalg.eigh((B + B.T) / 2) + evals = cp.maximum(evals, 0) + return (evecs * evals) @ evecs.T @staticmethod - def transform_block(A,k,Pk=None): + def transform_block(A, k, Pk=None): if Pk is None: - dk=2*k+1; Pk=cp.eye(dk) + dk = 2 * k + 1 + Pk = cp.eye(dk) for m in range(k): - for l in range(k-m): - Pk[(m+2*l,m+2*l+1),:]=Pk[(m+2*l+1,m+2*l),:] - AT=Pk@A@Pk.T - return AT[:k,:k].T.reshape(-1), AT[k:,k:].T.reshape(-1) - + for l in range(k - m): + Pk[(m + 2 * l, m + 2 * l + 1), :] = Pk[ + (m + 2 * l + 1, m + 2 * l), : + ] + AT = Pk @ A @ Pk.T + return AT[:k, :k].T.reshape(-1), AT[k:, k:].T.reshape(-1) + @staticmethod - def transform_back_block(A0,A1,k,Pk=None): - dk=2*k+1; A=cp.zeros((dk,dk)); A[:k,:k]=A0.reshape(k,k).T; A[k:,k:]=A1.reshape(k+1,k+1).T + def transform_back_block(A0, A1, k, Pk=None): + dk = 2 * k + 1 + A = cp.zeros((dk, dk)) + A[:k, :k] = A0.reshape(k, k).T + A[k:, k:] = A1.reshape(k + 1, k + 1).T if Pk is None: - Pk=cp.eye(dk) + Pk = cp.eye(dk) for m in range(k): - for l in range(k-m): - Pk[(m+2*l,m+2*l+1),:]=Pk[(m+2*l+1,m+2*l),:] - return Pk.T@A@Pk + for l in range(k - m): + Pk[(m + 2 * l, m + 2 * l + 1), :] = Pk[ + (m + 2 * l + 1, m + 2 * l), : + ] + return Pk.T @ A @ Pk From dfd701cbeeb3edcda14becd6d14b29b51d68e3e3 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 6 Mar 2026 11:26:41 -0500 Subject: [PATCH 003/161] Use numeric to handle cupy/numpy. Remove duplicate function. --- src/aspire/abinitio/commonline_nug.py | 255 +++++++++++--------------- 1 file changed, 110 insertions(+), 145 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 97a7420d64..c2f9aa249c 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -5,7 +5,7 @@ import numpy as np from scipy.io import loadmat from scipy.spatial.transform import Rotation as spr -from scipy.special import binom, factorial, jacobi, sph_harm +from scipy.special import factorial from aspire.abinitio import CLOrient3D from aspire.nufft import nufft @@ -73,6 +73,7 @@ def __init__( self.sym = symmetry def estimate_rotations(self): + breakpoint() sym_euler, S = self.Symmetry_Euler(self.sym) imgs = self.src.images[:] C = self.compute_coeff(imgs, self.loss, self.Lmax, T=self.T) @@ -142,18 +143,18 @@ def fij(alpha, gamma, i, j, loss): gamma_grid = np.arange(2 * T) * np.pi / T bT = np.zeros(2 * T) - for l in range(2 * T): + for n in range(2 * T): ss = 0 for m in range(T): - ss = ss + np.sin(beta_grid[l] * (2 * m + 1)) / (2 * m + 1) - bT[l] = 2 / T * np.sin(beta_grid[l]) * ss + ss = ss + np.sin(beta_grid[n] * (2 * m + 1)) / (2 * m + 1) + bT[n] = 2 / T * np.sin(beta_grid[n]) * ss BTK = [] for k in range(1, Lmax + 1): dk = 2 * k + 1 btk = np.zeros((dk, dk)) - for l in range(2 * T): - btk = btk + bT[l] * self.Wd(k, beta_grid[l]) + for n in range(2 * T): + btk = btk + bT[n] * self.Wd(k, beta_grid[n]) BTK.append(btk.T) def fijhat_k(k, F): @@ -227,7 +228,7 @@ def fast_radon_transform(array, angles, use_ramp=False): ) pts = pts.astype(array.dtype) - # array = array.astype(np.float32) + #array = array.astype(np.float32) lines_f = nufft(array, pts).reshape((img_size, -1)) if img_size % 2 == 0: @@ -243,35 +244,6 @@ def fast_radon_transform(array, angles, use_ramp=False): return projections, lines_f - @staticmethod - def Wd(J, beta): - # compute Wigner small d matrix - d = np.zeros((2 * J + 1, 2 * J + 1)) - - for m in range(-J, J + 1): - for n in range(-J, J + 1): - smin = max(0, m - n) - smax = min(J + m, J - n) - for s in range(smin, smax + 1): - mul = ( - np.sqrt(factorial(J + m)) - / factorial(J + m - s) - * np.sqrt(factorial(J + n)) - / factorial(s) - * np.sqrt(factorial(J - m)) - / factorial(n - m + s) - * np.sqrt(factorial(J - n)) - / factorial(J - n - s) - ) - d[n + J, m + J] += ( - mul - * (-1) ** (n - m + s) - * (np.cos(beta / 2)) ** (2 * J + m - n - 2 * s) - * (np.sin(beta / 2)) ** (n - m + 2 * s) - ) - - return d - @staticmethod def complex2real(ell): # compute transformation matrices that convert complex representations to real ones @@ -317,14 +289,8 @@ def admm_sym_J( mult=1, Nstep_yI=20, verbose=True, - GPU=True, ): # admm for symmetric case - if GPU: - import cupy as cp - else: - import numpy as cp - ( C0, C1, @@ -350,11 +316,11 @@ def admm_sym_J( S0, S1, Sq, - ) = self.ADMM_preprocessing(C, Lmax, N, Ngrid, GPU) + ) = self.ADMM_preprocessing(C, Lmax, N, Ngrid) # S=int(sym[1]); sym_euler=Symmetry_Euler(sym) rank_Ak, _ = self.compute_rank(sym, Lmax) - print(rank_Ak) + logger.info(f"Rank of Ak: {rank_Ak}") # rank_Ak=cp.zeros(Lmax) # for ell in range(Lmax): rank_Ak[ell]=np.linalg.matrix_rank(Ak(ell+1,sym_euler)) @@ -364,9 +330,9 @@ def admm_sym_J( for k in range(1, Lmax + 1): s0 = k**2 s1 = (k + 1) ** 2 - AEk = cp.zeros((1 + s0 + s1, 2 * (s0 + s1))) - AEk[0, :s0] = cp.eye(k).T.reshape(-1) - AEk[0, s0 : s0 + s1] = cp.eye(k + 1).T.reshape(-1) + AEk = xp.zeros((1 + s0 + s1, 2 * (s0 + s1))) + AEk[0, :s0] = xp.eye(k).T.reshape(-1) + AEk[0, s0 : s0 + s1] = xp.eye(k + 1).T.reshape(-1) for count in range(1, 1 + s0): AEk[count, count - 1] = 1 AEk[count, count - 1 + s0 + s1] = 1 @@ -375,20 +341,20 @@ def admm_sym_J( AEk[count, count - 1 + s1 + s0] = 1 AE.append(AEk) AEAETinv.append(np.linalg.pinv(AEk @ AEk.T)) - bE = cp.zeros((Lmax + D0 + D1)) + bE = xp.zeros((Lmax + D0 + D1)) for k in range(Lmax): bE[k + d0[k] + d1[k] :] = rank_Ak[k] - bE[k + 1 + d0[k] + d1[k] : k + 1 + d0[k + 1] + d1[k]] = cp.eye( + bE[k + 1 + d0[k] + d1[k] : k + 1 + d0[k + 1] + d1[k]] = xp.eye( k + 1 ).T.reshape(-1) - bE[k + 1 + d0[k + 1] + d1[k] : k + 1 + d0[k + 1] + d1[k + 1]] = cp.eye( + bE[k + 1 + d0[k + 1] + d1[k] : k + 1 + d0[k + 1] + d1[k + 1]] = xp.eye( k + 2 ).T.reshape(-1) bE = np.repeat(bE[:, np.newaxis], N, axis=1) P = [] for k in range(1, Lmax + 1): dk = 2 * k + 1 - Pk = cp.eye(dk) + Pk = xp.eye(dk) for m in range(k): for l in range(k - m): Pk[(m + 2 * l, m + 2 * l + 1), :] = Pk[ @@ -397,11 +363,11 @@ def admm_sym_J( P.append(Pk) def fun_AE(X0, X1, Xd0, Xd1, Xq): - z = cp.zeros((Lmax + D0 + D1, N)) + z = xp.zeros((Lmax + D0 + D1, N)) for k in range(Lmax): z[k + d0[k] + d1[k] : k + 1 + d0[k + 1] + d1[k + 1]] = AE[ k - ] @ cp.concatenate( + ] @ xp.concatenate( ( X0[d0[k] : d0[k + 1], idx_diag], X1[d1[k] : d1[k + 1], idx_diag], @@ -410,15 +376,15 @@ def fun_AE(X0, X1, Xd0, Xd1, Xq): ), axis=0, ) - return z, AEq @ cp.concatenate( + return z, AEq @ xp.concatenate( (Xq, X0[:1, idx_offdiag], X1[:4, idx_offdiag]), axis=0 ) def fun_AET(yE, yEq): - Z0 = cp.zeros((D0, N * (N + 1) // 2)) - Z1 = cp.zeros((D1, N * (N + 1) // 2)) - Zd0 = cp.zeros((D0, N)) - Zd1 = cp.zeros((D1, N)) + Z0 = xp.zeros((D0, N * (N + 1) // 2)) + Z1 = xp.zeros((D1, N * (N + 1) // 2)) + Zd0 = xp.zeros((D0, N)) + Zd1 = xp.zeros((D1, N)) for k in range(Lmax): s0 = (k + 1) ** 2 s1 = (k + 2) ** 2 @@ -433,17 +399,17 @@ def fun_AET(yE, yEq): return Z0, Z1, Zd0, Zd1, Zq[:16] def fun_AI(X0, X1): - z = cp.zeros((Ngrid, N * (N + 1) // 2)) - tmp = cp.concatenate((X0, X1), axis=0) + z = xp.zeros((Ngrid, N * (N + 1) // 2)) + tmp = xp.concatenate((X0, X1), axis=0) z[:, idx_diag] = AI_mat_diag @ tmp[:, idx_diag] z[:, idx_offdiag] = AI_mat_offdiag @ tmp[:, idx_offdiag] - # return AI_mat@cp.concatenate((X0,X1),axis=0) + # return AI_mat@xp.concatenate((X0,X1),axis=0) return z def fun_AIT(yI): # Z=AI_mat.T@yI # return Z[:D0,:], Z[D0:,:] - Z = cp.zeros((D0 + D1, N * (N + 1) // 2)) + Z = xp.zeros((D0 + D1, N * (N + 1) // 2)) Z[:, idx_diag] = AI_mat_diag.T @ yI[:, idx_diag] Z[:, idx_offdiag] = AI_mat_offdiag.T @ yI[:, idx_offdiag] return Z[:D0, :], Z[D0:, :] @@ -508,7 +474,7 @@ def update_yE(C0, C1, X0, X1, Xd0, Xd1, Xq, S0, S1, Sd0, Sd1, Sq, yI, rho): -Xd1 / rho - Sd1, -Xq / rho - Sq, ) - yE = cp.zeros((Lmax + D0 + D1, N)) + yE = xp.zeros((Lmax + D0 + D1, N)) for k in range(Lmax): yE[k + d0[k] + d1[k] : k + 1 + d0[k + 1] + d1[k + 1]] = AEAETinv[k] @ ( bE[k + d0[k] + d1[k] : k + 1 + d0[k + 1] + d1[k + 1]] / rho @@ -574,21 +540,21 @@ def update_rho(X0, X1, Xd0, Xd1, Xq, bE, bEq, bI, res_X, rho, factor, normC): def print_updates(verbose=True): # X_admm=transform_coeff_back(X0,X1,Lmax,N); obj_p=0 - # for k in range(Lmax): obj_p+=cp.trace(C[k]@X_admm[k]) + # for k in range(Lmax): obj_p+=xp.trace(C[k]@X_admm[k]) if verbose: obj_p = ( - cp.vdot(C0[:, idx_diag], X0[:, idx_diag]) - + cp.vdot(C1[:, idx_diag], X1[:, idx_diag]) - + 2 * cp.vdot(C0[:, idx_offdiag], X0[:, idx_offdiag]) - + 2 * cp.vdot(C1[:, idx_offdiag], X1[:, idx_offdiag]) + xp.vdot(C0[:, idx_diag], X0[:, idx_diag]) + + xp.vdot(C1[:, idx_diag], X1[:, idx_diag]) + + 2 * xp.vdot(C0[:, idx_offdiag], X0[:, idx_offdiag]) + + 2 * xp.vdot(C1[:, idx_offdiag], X1[:, idx_offdiag]) ) obj_d = ( - cp.vdot(yE, bE) - + 2 * cp.vdot(yEq, bEq) - + cp.vdot(yI[:, idx_diag], bI * cp.ones((Ngrid, N))) + xp.vdot(yE, bE) + + 2 * xp.vdot(yEq, bEq) + + xp.vdot(yI[:, idx_diag], bI * xp.ones((Ngrid, N))) + 2 - * cp.vdot( - yI[:, idx_offdiag], bI * cp.ones((Ngrid, N * (N - 1) // 2)) + * xp.vdot( + yI[:, idx_offdiag], bI * xp.ones((Ngrid, N * (N - 1) // 2)) ) ) @@ -657,7 +623,7 @@ def print_updates(verbose=True): ) p_res = res_eq + res_inq + res_psdX + res_psdD + res_psdQ d_res = res_X / (1 + normC) - print( + logger.info( "Iter %i" % t + ": p_res=%1.5f" % p_res + ", d_res=%1.5f" % d_res @@ -671,13 +637,13 @@ def print_updates(verbose=True): + ", |X|=%1.2f" % normX ) - Xd0 = cp.zeros((D0, N)) - Xd1 = cp.zeros((D1, N)) - Sd0 = cp.zeros(Xd0.shape) - Sd1 = cp.zeros(Xd1.shape) - yI = cp.zeros((Ngrid, N * (N + 1) // 2)) - yE = cp.zeros(bE.shape) - yEq = cp.zeros(bEq.shape) + Xd0 = xp.zeros((D0, N)) + Xd1 = xp.zeros((D1, N)) + Sd0 = xp.zeros(Xd0.shape) + Sd1 = xp.zeros(Xd1.shape) + yI = xp.zeros((Ngrid, N * (N + 1) // 2)) + yE = xp.zeros(bE.shape) + yEq = xp.zeros(bEq.shape) IDX = np.arange(3) Time = np.zeros(4) @@ -693,7 +659,7 @@ def print_updates(verbose=True): C0, C1, X0, X1, Xd0, Xd1, Xq, S0, S1, Sd0, Sd1, Sq, yI, rho ) if idx == 2: - for kk in range(Nstep_yI): + for _ in range(Nstep_yI): yI = update_yI(C0, C1, X0, X1, S0, S1, yE, yEq, yI, rho, Lambda) X0, X1, Xd0, Xd1, Xq, res_X = update_X( C0, C1, X0, X1, Xd0, Xd1, Xq, yE, yEq, yI, S0, S1, Sd0, Sd1, Sq, rho @@ -707,18 +673,16 @@ def print_updates(verbose=True): X_admm = self.transform_coeff_back( X0, X1, Lmax, N, IDX_upper, IDX_lower, idx_offdiag ) - if GPU: - for k in range(Lmax): - X_admm[k] = X_admm[k].get() + # if self.GPU: + # for k in range(Lmax): + # X_admm[k] = X_admm[k].get() + for k in range(Lmax): + X_admm[k] = xp.asnumpy(X_admm[k]) return X_admm - def ADMM_preprocessing(self, C, Lmax, N, Ngrid, GPU=True): + def ADMM_preprocessing(self, C, Lmax, N, Ngrid): # compute necessary quantities for ADMM - if GPU: - import cupy as cp - else: - import numpy as cp - + # compute some useful index sets count = 0 idx_diag = [] @@ -751,7 +715,7 @@ def ADMM_preprocessing(self, C, Lmax, N, Ngrid, GPU=True): Cnorm = np.sqrt(Cnorm) Xnorm = np.sqrt(Xnorm) for k in range(Lmax): - C[k] = cp.asarray(Xnorm / Cnorm * C[k]) + C[k] = xp.asarray(Xnorm / Cnorm * C[k]) C0, C1 = self.transform_coeff(C, Lmax, N, IDX_upper) normC = np.sqrt(np.linalg.norm(C0) ** 2 + np.linalg.norm(C1) ** 2) del C @@ -766,12 +730,12 @@ def ADMM_preprocessing(self, C, Lmax, N, Ngrid, GPU=True): D1 = d1[-1] # AE and bE for quaternion constraints - AEq = cp.asarray(loadmat("data/Eq_constraints/AEqJ.mat")["AEq"]) - AEqAEqtinv = cp.asarray(loadmat("data/Eq_constraints/AEqJ.mat")["AEqAEqtinv"]) - bEq = cp.zeros(17) - bEq[:16] = cp.eye(4).reshape(-1) / 4 + AEq = xp.asarray(loadmat("data/Eq_constraints/AEqJ.mat")["AEq"]) + AEqAEqtinv = xp.asarray(loadmat("data/Eq_constraints/AEqJ.mat")["AEqAEqtinv"]) + bEq = xp.zeros(17) + bEq[:16] = xp.eye(4).reshape(-1) / 4 bEq[-1] = 1 - bEq = cp.repeat(bEq[:, cp.newaxis], N * (N - 1) // 2, axis=1) + bEq = xp.repeat(bEq[:, xp.newaxis], N * (N - 1) // 2, axis=1) # AI and bI W0 = [ @@ -790,7 +754,7 @@ def ADMM_preprocessing(self, C, Lmax, N, Ngrid, GPU=True): # w1[d1[k-1]:d1[k]]=(Lmax-k+2)*(Lmax-k+1)*(k+0.5)*W1[k-1][p].T.reshape(-1) # # this needs double checking # AI_mat[p,:d0[-1]]=w0; AI_mat[p,d0[-1]:]=w1 - # AI_mat=cp.asarray(AI_mat) / 10; + # AI_mat=xp.asarray(AI_mat) / 10; # bI=-(Lmax+2)*(Lmax+1)/2 / 10 AI_mat_offdiag = np.zeros((Ngrid, D0 + D1)) @@ -833,10 +797,10 @@ def ADMM_preprocessing(self, C, Lmax, N, Ngrid, GPU=True): # this needs double checking AI_mat_diag[p, : d0[-1]] = w0 AI_mat_diag[p, d0[-1] :] = w1 - AI_mat_diag = cp.asarray(AI_mat_diag) / 1 - AI_mat_offdiag = cp.asarray(AI_mat_offdiag) / 1 + AI_mat_diag = xp.asarray(AI_mat_diag) / 1 + AI_mat_offdiag = xp.asarray(AI_mat_offdiag) / 1 bI = -(Lmax + 2) * (Lmax + 1) / 2 / 1 - + np.save("AI_mat_offdiag_aspire.npy", AI_mat_offdiag) # largest eigenvalue for AIAIT Lambda = self.largest_eigenvalue(AI_mat_offdiag, Ngrid, N) @@ -844,16 +808,16 @@ def ADMM_preprocessing(self, C, Lmax, N, Ngrid, GPU=True): II = [] for k in range(1, Lmax + 1): dk = 2 * k + 1 - II.append(cp.eye(N * dk)) + II.append(xp.eye(N * dk)) I0, I1 = self.transform_coeff(II, Lmax, N, IDX_upper) - X0 = cp.zeros((D0, N * (N + 1) // 2)) - X1 = cp.zeros((D1, N * (N + 1) // 2)) - Xq = cp.zeros((16, N * (N - 1) // 2)) - # X0,X1=transform_coeff(II,Lmax,N); Xq=cp.zeros((16,N*(N-1))); Xq[0,:]=1; - S0 = cp.copy(I0) - S1 = cp.copy(I1) - Sq = cp.zeros(Xq.shape) - # S0=cp.zeros(X0.shape); S1=cp.zeros(X1.shape); Sq=cp.zeros(Xq.shape) + X0 = xp.zeros((D0, N * (N + 1) // 2)) + X1 = xp.zeros((D1, N * (N + 1) // 2)) + Xq = xp.zeros((16, N * (N - 1) // 2)) + # X0,X1=transform_coeff(II,Lmax,N); Xq=xp.zeros((16,N*(N-1))); Xq[0,:]=1; + S0 = xp.copy(I0) + S1 = xp.copy(I1) + Sq = xp.zeros(Xq.shape) + # S0=xp.zeros(X0.shape); S1=xp.zeros(X1.shape); Sq=xp.zeros(Xq.shape) # return C0,C1,normC,AEq,bEq,AEqAEqtinv,AI_mat,bI,Lambda,d0,d1,D0,D1,idx_diag,idx_offdiag,IDX_upper,IDX_lower,X0,X1,Xq,S0,S1,Sq return ( @@ -1007,8 +971,8 @@ def transform_coeff(self, A, Lmax, N, IDX_upper): for k in range(1, Lmax + 1): d0.append(d0[-1] + k**2) d1.append(d1[-1] + (k + 1) ** 2) - A0 = cp.zeros((d0[-1], N * (N + 1) // 2)) - A1 = cp.zeros((d1[-1], N * (N + 1) // 2)) + A0 = xp.zeros((d0[-1], N * (N + 1) // 2)) + A1 = xp.zeros((d1[-1], N * (N + 1) // 2)) for k in range(1, Lmax + 1): a0, a1 = self.permutek(A[k - 1], k, N) A0[d0[k - 1] : d0[k], :] = self.vec_block(a0, N, k, IDX_upper) @@ -1024,7 +988,7 @@ def transform_coeff_back(self, A0, A1, Lmax, N, IDX_upper, IDX_lower, idx_offdia A = [] for k in range(1, Lmax + 1): dk = 2 * k + 1 - Ak = cp.zeros((N * dk, N * dk)) + Ak = xp.zeros((N * dk, N * dk)) Ak[: N * k, : N * k] = self.mat_block( A0[d0[k - 1] : d0[k], :], N, k, IDX_upper, IDX_lower, idx_offdiag ) @@ -1037,20 +1001,20 @@ def transform_coeff_back(self, A0, A1, Lmax, N, IDX_upper, IDX_lower, idx_offdia @staticmethod def permutek(Ak, k, N): - AkP = cp.copy(Ak) + AkP = xp.copy(Ak) dk = 2 * k + 1 - Pk = cp.eye(dk) + Pk = xp.eye(dk) for m in range(k): - for l in range(k - m): - Pk[(m + 2 * l, m + 2 * l + 1), :] = Pk[(m + 2 * l + 1, m + 2 * l), :] - AkP = cp.kron(cp.eye(N), Pk) @ Ak @ cp.kron(cp.eye(N), Pk.T) + for n in range(k - m): + Pk[(m + 2 * n, m + 2 * n + 1), :] = Pk[(m + 2 * n + 1, m + 2 * n), :] + AkP = xp.kron(xp.eye(N), Pk) @ Ak @ xp.kron(xp.eye(N), Pk.T) - Pk = cp.eye(N * dk) - idx = cp.concatenate((cp.arange(dk - k, dk), cp.arange(k + 1))) + Pk = xp.eye(N * dk) + idx = xp.concatenate((xp.arange(dk - k, dk), xp.arange(k + 1))) for m in range(N - 1): - for l in range(N - 1 - m): - Pk[k * (m + 1) + l * dk : k * (m + 1) + (l + 1) * dk] = Pk[ - k * (m + 1) + l * dk : k * (m + 1) + (l + 1) * dk + for n in range(N - 1 - m): + Pk[k * (m + 1) + n * dk : k * (m + 1) + (n + 1) * dk] = Pk[ + k * (m + 1) + n * dk : k * (m + 1) + (n + 1) * dk ][idx, :] AkP = Pk @ AkP @ Pk.T return AkP[: N * k, : N * k], AkP[N * k :, N * k :] @@ -1058,20 +1022,20 @@ def permutek(Ak, k, N): @staticmethod def permutek_back(Ak, k, N): dk = 2 * k + 1 - Pk = cp.eye(N * dk) - idx = cp.concatenate((cp.arange(dk - k, dk), cp.arange(k + 1))) + Pk = xp.eye(N * dk) + idx = xp.concatenate((xp.arange(dk - k, dk), xp.arange(k + 1))) for m in range(N - 1): - for l in range(N - 1 - m): - Pk[k * (m + 1) + l * dk : k * (m + 1) + (l + 1) * dk] = Pk[ - k * (m + 1) + l * dk : k * (m + 1) + (l + 1) * dk + for n in range(N - 1 - m): + Pk[k * (m + 1) + n * dk : k * (m + 1) + (n + 1) * dk] = Pk[ + k * (m + 1) + n * dk : k * (m + 1) + (n + 1) * dk ][idx, :] AkB = Pk.T @ Ak @ Pk dk = 2 * k + 1 - Pk = cp.eye(dk) + Pk = xp.eye(dk) for m in range(k): - for l in range(k - m): - Pk[(m + 2 * l, m + 2 * l + 1), :] = Pk[(m + 2 * l + 1, m + 2 * l), :] - AkB = cp.kron(cp.eye(N), Pk.T) @ AkB @ cp.kron(cp.eye(N), Pk) + for n in range(k - m): + Pk[(m + 2 * n, m + 2 * n + 1), :] = Pk[(m + 2 * n + 1, m + 2 * n), :] + AkB = xp.kron(xp.eye(N), Pk.T) @ AkB @ xp.kron(xp.eye(N), Pk) return AkB @staticmethod @@ -1082,21 +1046,22 @@ def vec_block(A, N, sz, IDX_upper): @staticmethod def largest_eigenvalue(AI, Ngrid, N): # find the largest eigenvalue of the operator AI - z = cp.random.normal(0, 1, (Ngrid, N**2)) + np.random.seed(0) + z = xp.random.normal(0, 1, (Ngrid, N**2)) Lambda = 0 - while abs(Lambda - cp.linalg.norm(z)) > 500: - Lambda = cp.linalg.norm(z) - z = z / cp.linalg.norm(z) + while abs(Lambda - xp.linalg.norm(z)) > 500: + Lambda = xp.linalg.norm(z) + z = z / xp.linalg.norm(z) z = AI @ (AI.T @ z) Lambda += 2000 - print("Largest eigenvalue of AIAIT is approximately %1.2f" % Lambda) - # Lambda=cp.linalg.eigvalsh(AI@AI.T)[-1]; print(Lambda) + logger.info("Largest eigenvalue of AIAIT is approximately %1.2f" % Lambda) + # Lambda=xp.linalg.eigvalsh(AI@AI.T)[-1]; print(Lambda) return Lambda def compute_rank(self, sym, Lmax): sym_euler, S = self.Symmetry_Euler(sym) - rk = cp.zeros(Lmax) + rk = xp.zeros(Lmax) A = [] for k in range(1, Lmax + 1): dk = 2 * k + 1 @@ -1185,7 +1150,7 @@ def Wd(J, beta): @staticmethod def mat_block(vecA, N, sz, IDX_upper, IDX_lower, idx_offdiag): tmp = vecA.T.reshape(N * (N + 1) // 2, sz, sz).transpose(0, 2, 1) - AA = cp.zeros((N**2, sz, sz)) + AA = xp.zeros((N**2, sz, sz)) AA[IDX_upper] = tmp AA[IDX_lower] = tmp[idx_offdiag].transpose(0, 2, 1) return (AA.reshape(N, N, sz, sz).transpose(0, 2, 1, 3)).reshape(N * sz, N * sz) @@ -1193,15 +1158,15 @@ def mat_block(vecA, N, sz, IDX_upper, IDX_lower, idx_offdiag): @staticmethod def psd_projection(B): # compute the PSD part of a symmstric matrix - evals, evecs = cp.linalg.eigh((B + B.T) / 2) - evals = cp.maximum(evals, 0) + evals, evecs = xp.linalg.eigh((B + B.T) / 2) + evals = xp.maximum(evals, 0) return (evecs * evals) @ evecs.T @staticmethod def transform_block(A, k, Pk=None): if Pk is None: dk = 2 * k + 1 - Pk = cp.eye(dk) + Pk = xp.eye(dk) for m in range(k): for l in range(k - m): Pk[(m + 2 * l, m + 2 * l + 1), :] = Pk[ @@ -1213,11 +1178,11 @@ def transform_block(A, k, Pk=None): @staticmethod def transform_back_block(A0, A1, k, Pk=None): dk = 2 * k + 1 - A = cp.zeros((dk, dk)) + A = xp.zeros((dk, dk)) A[:k, :k] = A0.reshape(k, k).T A[k:, k:] = A1.reshape(k + 1, k + 1).T if Pk is None: - Pk = cp.eye(dk) + Pk = xp.eye(dk) for m in range(k): for l in range(k - m): Pk[(m + 2 * l, m + 2 * l + 1), :] = Pk[ From 153c3f83df4dbf6216a3785ce69bbb6f36b3dd12 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 9 Mar 2026 10:22:27 -0400 Subject: [PATCH 004/161] tox --- src/aspire/abinitio/commonline_nug.py | 23 ++++++++++------------- 1 file changed, 10 insertions(+), 13 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index c2f9aa249c..410ebb6fc8 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1,7 +1,6 @@ import logging import time -import cupy as cp import numpy as np from scipy.io import loadmat from scipy.spatial.transform import Rotation as spr @@ -73,7 +72,6 @@ def __init__( self.sym = symmetry def estimate_rotations(self): - breakpoint() sym_euler, S = self.Symmetry_Euler(self.sym) imgs = self.src.images[:] C = self.compute_coeff(imgs, self.loss, self.Lmax, T=self.T) @@ -228,7 +226,7 @@ def fast_radon_transform(array, angles, use_ramp=False): ) pts = pts.astype(array.dtype) - #array = array.astype(np.float32) + # array = array.astype(np.float32) lines_f = nufft(array, pts).reshape((img_size, -1)) if img_size % 2 == 0: @@ -356,9 +354,9 @@ def admm_sym_J( dk = 2 * k + 1 Pk = xp.eye(dk) for m in range(k): - for l in range(k - m): - Pk[(m + 2 * l, m + 2 * l + 1), :] = Pk[ - (m + 2 * l + 1, m + 2 * l), : + for el in range(k - m): + Pk[(m + 2 * el, m + 2 * el + 1), :] = Pk[ + (m + 2 * el + 1, m + 2 * el), : ] P.append(Pk) @@ -682,7 +680,6 @@ def print_updates(verbose=True): def ADMM_preprocessing(self, C, Lmax, N, Ngrid): # compute necessary quantities for ADMM - # compute some useful index sets count = 0 idx_diag = [] @@ -1168,9 +1165,9 @@ def transform_block(A, k, Pk=None): dk = 2 * k + 1 Pk = xp.eye(dk) for m in range(k): - for l in range(k - m): - Pk[(m + 2 * l, m + 2 * l + 1), :] = Pk[ - (m + 2 * l + 1, m + 2 * l), : + for el in range(k - m): + Pk[(m + 2 * el, m + 2 * el + 1), :] = Pk[ + (m + 2 * el + 1, m + 2 * el), : ] AT = Pk @ A @ Pk.T return AT[:k, :k].T.reshape(-1), AT[k:, k:].T.reshape(-1) @@ -1184,8 +1181,8 @@ def transform_back_block(A0, A1, k, Pk=None): if Pk is None: Pk = xp.eye(dk) for m in range(k): - for l in range(k - m): - Pk[(m + 2 * l, m + 2 * l + 1), :] = Pk[ - (m + 2 * l + 1, m + 2 * l), : + for el in range(k - m): + Pk[(m + 2 * el, m + 2 * el + 1), :] = Pk[ + (m + 2 * el + 1, m + 2 * el), : ] return Pk.T @ A @ Pk From 6097ffcde5146584041060e165fd0aa1b9be0e05 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 12 Mar 2026 14:52:11 -0400 Subject: [PATCH 005/161] build full pf --- src/aspire/abinitio/commonline_nug.py | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 410ebb6fc8..27d2e134b3 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -9,7 +9,7 @@ from aspire.abinitio import CLOrient3D from aspire.nufft import nufft from aspire.numeric import fft, xp -from aspire.operators import wemd_embed +from aspire.operators import PolarFT, wemd_embed logger = logging.getLogger(__name__) @@ -71,6 +71,12 @@ def __init__( self.Nstep_yI = Nstep_yI self.sym = symmetry + self._build_full_pft() + + def _build_full_pft(self): + pf = self.pf + self.pf_full = PolarFT.half_to_full(pf) + def estimate_rotations(self): sym_euler, S = self.Symmetry_Euler(self.sym) imgs = self.src.images[:] @@ -105,6 +111,8 @@ def compute_coeff(self, Img, loss, Lmax, T): angular_sampling = np.arange(0, 360, 1) line_proj = np.zeros((L, n_theta, N)) Img_pft = np.zeros((L, n_theta, N), dtype=complex) + + # Replace with Image.project() later Img = Img.asnumpy() for n in range(N): line_proj[:, :, n], Img_pft[:, :, n] = self.fast_radon_transform( @@ -126,6 +134,17 @@ def fij(alpha, gamma, i, j, loss): Si = Ii_hat[:, int(idxi)] Sj = Ij_hat[:, int(idxj)] + + # Using aspire PolarFT. Replace later + # Ii_hat = self.pf_full[i] + # Ij_hat = self.pf_full[j] + # idxi = np.round((alpha - np.pi / 2) * n_theta / 2 / np.pi) % n_theta + # idxj = np.round((-gamma - np.pi / 2) * n_theta / 2 / np.pi) % n_theta + + # Si = Ii_hat[int(idxi)] + # Sj = Ij_hat[int(idxj)] + # norm_new = np.linalg.norm(Si - Sj, 1) + return np.linalg.norm(Si - Sj, 1) if loss == "wemd": @@ -729,6 +748,7 @@ def ADMM_preprocessing(self, C, Lmax, N, Ngrid): # AE and bE for quaternion constraints AEq = xp.asarray(loadmat("data/Eq_constraints/AEqJ.mat")["AEq"]) AEqAEqtinv = xp.asarray(loadmat("data/Eq_constraints/AEqJ.mat")["AEqAEqtinv"]) + bEq = xp.zeros(17) bEq[:16] = xp.eye(4).reshape(-1) / 4 bEq[-1] = 1 From 32ec1340bfa29cf7ee4fb7f748ec7baf5fb1b9cd Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 7 Apr 2026 09:48:52 -0400 Subject: [PATCH 006/161] compute_fejer_weights method. Still need to optimize (very slow) --- src/aspire/abinitio/commonline_nug.py | 46 ++++++++++++++++++++++++++- 1 file changed, 45 insertions(+), 1 deletion(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 27d2e134b3..41f48f35f2 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -41,7 +41,7 @@ def __init__( **kwargs, ): """ - Initialize object for estimating 3D orientations for molecules with C3 and C4 symmetry. + Initialize object for estimating 3D orientations for symmetric molecules. :param src: The source object of 2D denoised or class-averaged images with metadata :param symmetry: A string, ie. 'C3', indicating the symmetry type. @@ -763,6 +763,9 @@ def ADMM_preprocessing(self, C, Lmax, N, Ngrid): loadmat("data/Fejer/Ngrid=%i" % Ngrid + "/k=%i" % k + ".mat")["W1"] for k in range(1, Lmax + 1) ] + + W0k, W1k = self.compute_fejer_weights() + # AI_mat=np.zeros((Ngrid,D0+D1)) # for p in range(Ngrid): # w0=np.zeros(D0); w1=np.zeros(D1) @@ -864,6 +867,47 @@ def ADMM_preprocessing(self, C, Lmax, N, Ngrid): Sq, ) + def compute_fejer_weights(self): + SO3_grid = loadmat("data/SO3_grid.mat")["SO3"] + Ngrid = SO3_grid.shape[0] + start = 1 + + TT = [] + TTI = [] + for ell in range(start, self.Lmax + 1): + T, Tinv = self.complex2real(ell) + TT.append(T) + TTI.append(Tinv) + + def permutek_block(Ak, k): + dk = 2 * k + 1 + Pk = np.eye(dk) + for m in range(k): + for l in range(k - m): + Pk[(m + 2 * l, m + 2 * l + 1), :] = Pk[ + (m + 2 * l + 1, m + 2 * l), : + ] + AkP = Pk @ Ak @ Pk.T + return AkP[:k, :k], AkP[k:, k:] + + W0 = [] + W1 = [] + for k in range(start, self.Lmax + 1): + print(k) + W0k = np.zeros((Ngrid, k, k)) + W1k = np.zeros((Ngrid, k + 1, k + 1)) + + TkT = TT[k - start].T + TinvkT = TTI[k - start].T + + for p in range(Ngrid): + w = np.real(TkT @ self.WD(k, SO3_grid[p]).conj() @ TinvkT) + W0k[p], W1k[p] = permutek_block(w, k) + + W0.append(W0k) + W1.append(W1k) + return W0, W1 + ######################### # Euler Estimation Step # ######################### From 5b0250c1bd6d3fcf594bb74a02fe5641749bb646 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 8 Apr 2026 15:48:19 -0400 Subject: [PATCH 007/161] vectorize Wigner matrix comps. Vectorize fejer_weights comp. --- src/aspire/abinitio/commonline_nug.py | 56 ++++++++++----------------- 1 file changed, 20 insertions(+), 36 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 41f48f35f2..ff0a590ae0 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -168,10 +168,7 @@ def fij(alpha, gamma, i, j, loss): BTK = [] for k in range(1, Lmax + 1): - dk = 2 * k + 1 - btk = np.zeros((dk, dk)) - for n in range(2 * T): - btk = btk + bT[n] * self.Wd(k, beta_grid[n]) + btk = np.sum(bT[:, None, None] * self.Wd(k, beta_grid), axis=0) BTK.append(btk.T) def fijhat_k(k, F): @@ -755,16 +752,7 @@ def ADMM_preprocessing(self, C, Lmax, N, Ngrid): bEq = xp.repeat(bEq[:, xp.newaxis], N * (N - 1) // 2, axis=1) # AI and bI - W0 = [ - loadmat("data/Fejer/Ngrid=%i" % Ngrid + "/k=%i" % k + ".mat")["W0"] - for k in range(1, Lmax + 1) - ] - W1 = [ - loadmat("data/Fejer/Ngrid=%i" % Ngrid + "/k=%i" % k + ".mat")["W1"] - for k in range(1, Lmax + 1) - ] - - W0k, W1k = self.compute_fejer_weights() + W0, W1 = self.compute_fejer_weights() # AI_mat=np.zeros((Ngrid,D0+D1)) # for p in range(Ngrid): @@ -888,7 +876,7 @@ def permutek_block(Ak, k): (m + 2 * l + 1, m + 2 * l), : ] AkP = Pk @ Ak @ Pk.T - return AkP[:k, :k], AkP[k:, k:] + return AkP[..., :k, :k], AkP[..., k:, k:] W0 = [] W1 = [] @@ -900,9 +888,8 @@ def permutek_block(Ak, k): TkT = TT[k - start].T TinvkT = TTI[k - start].T - for p in range(Ngrid): - w = np.real(TkT @ self.WD(k, SO3_grid[p]).conj() @ TinvkT) - W0k[p], W1k[p] = permutek_block(w, k) + w = np.real(TkT @ self.WD(k, SO3_grid).conj() @ TinvkT) + W0k, W1k = permutek_block(w, k) W0.append(W0k) W1.append(W1k) @@ -979,14 +966,15 @@ def find_gamma(Xm, beta, alpha): Jk[S + 1 :: 2] = -1 Jk[S - 1 :: -2] = -1 Jk = np.diag(Jk) + ws = self.Wd(S, beta) for i in range(N): + wi = ws[i] for j in range(i + 1, N): Di = np.exp(-1j * np.arange(-S, S + 1) * alpha[i]) Dj = np.exp(-1j * np.arange(-S, S + 1) * alpha[j]) Xijm = Xm[dk * i : dk * (i + 1), dk * j : dk * (j + 1)] DXijmD = np.diag(Di.conj()) @ Xijm @ np.diag(Dj) - wi = self.Wd(S, beta[i]) - wj = self.Wd(S, beta[j]) + wj = ws[j] C1 = ( wi[:, 0][:, None] @ (wj[:, 0][:, None].T) + Jk @ wi[:, 0][:, None] @ (wj[:, 0][:, None].T) @ Jk @@ -1125,10 +1113,7 @@ def compute_rank(self, sym, Lmax): rk = xp.zeros(Lmax) A = [] for k in range(1, Lmax + 1): - dk = 2 * k + 1 - Ak = np.zeros((dk, dk), dtype=np.complex128) - for s in range(S): - Ak += self.WD(k, sym_euler[s]) + Ak = np.sum(self.WD(k, sym_euler), axis=0) Ak = np.round(Ak / S, 6) A.append(Ak) rk[k - 1] = np.linalg.matrix_rank(Ak) @@ -1168,22 +1153,22 @@ def Symmetry_Euler(sym): def WD(self, J, euler): # compute Wigner D matrix - alpha = euler[0] - beta = euler[1] - gamma = euler[2] + alpha = euler[:, 0] + beta = euler[:, 1] + gamma = euler[:, 2] d = self.Wd(J, beta) - D = ( - np.diag(np.exp(-1j * alpha * np.arange(-J, J + 1))) - @ d - @ np.diag(np.exp(-1j * gamma * np.arange(-J, J + 1))) - ) + + m = np.arange(-J, J + 1) + left = np.exp(-1j * alpha[:, None] * m[None, :]) + right = np.exp(-1j * gamma[:, None] * m[None, :]) + D = left[:, :, None] * d * right[:, None, :] + return D @staticmethod def Wd(J, beta): # compute Wigner small d matrix - d = np.zeros((2 * J + 1, 2 * J + 1)) - + d = np.zeros((len(beta), 2 * J + 1, 2 * J + 1)) for m in range(-J, J + 1): for n in range(-J, J + 1): smin = max(0, m - n) @@ -1199,13 +1184,12 @@ def Wd(J, beta): * np.sqrt(factorial(J - n)) / factorial(J - n - s) ) - d[n + J, m + J] += ( + d[:, n + J, m + J] += ( mul * (-1) ** (n - m + s) * (np.cos(beta / 2)) ** (2 * J + m - n - 2 * s) * (np.sin(beta / 2)) ** (n - m + 2 * s) ) - return d @staticmethod From 4558b19b3e513a3ac2023678ac88e77181a19bad Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 16 Apr 2026 11:45:56 -0400 Subject: [PATCH 008/161] Add SO3 grid generation --- src/aspire/abinitio/commonline_nug.py | 38 ++++++++++++++++++++++++++- src/aspire/utils/__init__.py | 1 + 2 files changed, 38 insertions(+), 1 deletion(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index ff0a590ae0..01dc3fe995 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -10,6 +10,7 @@ from aspire.nufft import nufft from aspire.numeric import fft, xp from aspire.operators import PolarFT, wemd_embed +from aspire.utils import cart2sph logger = logging.getLogger(__name__) @@ -785,6 +786,18 @@ def ADMM_preprocessing(self, C, Lmax, N, Ngrid): # this needs double checking AI_mat_offdiag[p, : d0[-1]] = w0 AI_mat_offdiag[p, d0[-1] :] = w1 + + # Vectorized version, added by Josh + # AI_mat_offdiag_new = np.zeros((Ngrid, D0 + D1)) + # for k in range(1, Lmax + 1): + # scale = (Lmax - k + 2) * (Lmax - k + 1) * (k + 0.5) + + # block0 = scale * W0[k - 1].transpose(0, 2, 1).reshape(Ngrid, -1) + # block1 = scale * W1[k - 1].transpose(0, 2, 1).reshape(Ngrid, -1) + + # AI_mat_offdiag_new[:, d0[k - 1]:d0[k]] = block0 + # AI_mat_offdiag_new[:, d0[-1] + d1[k - 1]:d0[-1] + d1[k]] = block1 + AI_mat_diag = np.zeros((Ngrid, D0 + D1)) for p in range(Ngrid): w0 = np.zeros(D0) @@ -856,7 +869,7 @@ def ADMM_preprocessing(self, C, Lmax, N, Ngrid): ) def compute_fejer_weights(self): - SO3_grid = loadmat("data/SO3_grid.mat")["SO3"] + SO3_grid = self.discretize_SO3() Ngrid = SO3_grid.shape[0] start = 1 @@ -895,6 +908,29 @@ def permutek_block(Ak, k): W1.append(W1k) return W0, W1 + def discretize_SO3(self): + S2 = loadmat("design20.mat")["design"] + S2_size = S2.shape[0] + + # discretize S1 + S1_size = round(np.sqrt(np.pi * S2_size)) + alpha = np.linspace(0, 2 * np.pi, S1_size) + + # discretize S2 + gamma, beta, _ = cart2sph(S2[:, 0], S2[:, 1], S2[:, 2]) + beta = np.pi / 2 - beta + gamma = gamma + np.pi + + # SO(3) in Euler ZYZ + SO3 = np.zeros((S2_size * S1_size, 3)) + count = 0 + for i in range(S1_size): + for j in range(S2_size): + SO3[count] = [alpha[i], beta[j], gamma[j]] + count += 1 + + return SO3 + ######################### # Euler Estimation Step # ######################### diff --git a/src/aspire/utils/__init__.py b/src/aspire/utils/__init__.py index ae896ebeb8..b7bd2d1947 100644 --- a/src/aspire/utils/__init__.py +++ b/src/aspire/utils/__init__.py @@ -2,6 +2,7 @@ from .coor_trans import ( # isort:skip mean_aligned_angular_distance, cart2pol, + cart2sph, crop_pad_2d, crop_pad_3d, grid_1d, From d32fcb92f4729c67d15d1e8feb4355d4d417f701 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 17 Apr 2026 10:46:35 -0400 Subject: [PATCH 009/161] tox --- src/aspire/abinitio/commonline_nug.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 01dc3fe995..14d8cded0b 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -884,9 +884,9 @@ def permutek_block(Ak, k): dk = 2 * k + 1 Pk = np.eye(dk) for m in range(k): - for l in range(k - m): - Pk[(m + 2 * l, m + 2 * l + 1), :] = Pk[ - (m + 2 * l + 1, m + 2 * l), : + for ell in range(k - m): + Pk[(m + 2 * ell, m + 2 * ell + 1), :] = Pk[ + (m + 2 * ell + 1, m + 2 * ell), : ] AkP = Pk @ Ak @ Pk.T return AkP[..., :k, :k], AkP[..., k:, k:] @@ -894,7 +894,6 @@ def permutek_block(Ak, k): W0 = [] W1 = [] for k in range(start, self.Lmax + 1): - print(k) W0k = np.zeros((Ngrid, k, k)) W1k = np.zeros((Ngrid, k + 1, k + 1)) From 6f0006c18eaa86c6a862c88ea1c98ff60c552c6b Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 20 Apr 2026 11:14:38 -0400 Subject: [PATCH 010/161] use aspire symmetry parsing --- src/aspire/abinitio/commonline_nug.py | 59 +++++++-------------------- 1 file changed, 14 insertions(+), 45 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 14d8cded0b..abedcd5431 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -11,6 +11,7 @@ from aspire.numeric import fft, xp from aspire.operators import PolarFT, wemd_embed from aspire.utils import cart2sph +from aspire.volume import SymmetryGroup logger = logging.getLogger(__name__) @@ -70,7 +71,11 @@ def __init__( self.mult = mult self.Ngrid = Ngrid self.Nstep_yI = Nstep_yI - self.sym = symmetry + + # Handle symmetry + self.sym_grp = SymmetryGroup.parse(symmetry) + self.sym_euler = self.sym_grp.rotations.angles + self.n_sym = len(self.sym_euler) self._build_full_pft() @@ -79,11 +84,9 @@ def _build_full_pft(self): self.pf_full = PolarFT.half_to_full(pf) def estimate_rotations(self): - sym_euler, S = self.Symmetry_Euler(self.sym) imgs = self.src.images[:] C = self.compute_coeff(imgs, self.loss, self.Lmax, T=self.T) X_est = self.admm_sym_J( - self.sym, C, self.Lmax, self.n_img, @@ -96,7 +99,7 @@ def estimate_rotations(self): self.Nstep_yI, ) - R_est, Euler_est = self.euler_est(X_est[0], X_est[S - 1], self.sym, self.n_img) + R_est, Euler_est = self.euler_est(X_est[0], X_est[self.n_sym - 1]) self.rotations = R_est return R_est @@ -292,7 +295,6 @@ def complex2real(ell): def admm_sym_J( self, - sym, C, Lmax, N, @@ -333,8 +335,7 @@ def admm_sym_J( Sq, ) = self.ADMM_preprocessing(C, Lmax, N, Ngrid) - # S=int(sym[1]); sym_euler=Symmetry_Euler(sym) - rank_Ak, _ = self.compute_rank(sym, Lmax) + rank_Ak, _ = self.compute_rank(Lmax) logger.info(f"Rank of Ak: {rank_Ak}") # rank_Ak=cp.zeros(Lmax) @@ -934,8 +935,9 @@ def discretize_SO3(self): # Euler Estimation Step # ######################### - def euler_est(self, X1, XS, sym, N): - S = int(sym[1]) + def euler_est(self, X1, XS): + S = self.n_sym + N = self.n_img sym_euler = np.zeros((S, 3)) for s in range(S): sym_euler[s] = [2 * np.pi * s / S, 0, 0] @@ -1143,49 +1145,16 @@ def largest_eigenvalue(AI, Ngrid, N): # Lambda=xp.linalg.eigvalsh(AI@AI.T)[-1]; print(Lambda) return Lambda - def compute_rank(self, sym, Lmax): - sym_euler, S = self.Symmetry_Euler(sym) + def compute_rank(self, Lmax): rk = xp.zeros(Lmax) A = [] for k in range(1, Lmax + 1): - Ak = np.sum(self.WD(k, sym_euler), axis=0) - Ak = np.round(Ak / S, 6) + Ak = np.sum(self.WD(k, self.sym_euler), axis=0) + Ak = np.round(Ak / self.n_sym, 6) A.append(Ak) rk[k - 1] = np.linalg.matrix_rank(Ak) return rk, A - @staticmethod - def Symmetry_Euler(sym): - if sym[0] == "C": - order = int(sym[1:]) - sym_euler = np.zeros((order, 3)) - for i in range(order): - sym_euler[i] = spr.from_euler( - "zyx", [2 * np.pi / order * i, 0, 0] - ).as_euler("zyz") - - if sym[0] == "D": - order = int(sym[1:]) - sym_euler = np.zeros((2 * order, 3)) - for i in range(order): - sym_euler[i] = spr.from_euler( - "zyx", [2 * np.pi / order * i, 0, 0] - ).as_euler("zyz") - sym_euler[i + order] = spr.from_euler( - "zyx", [2 * np.pi / order * i, 0, np.pi] - ).as_euler("zyz") - - if sym == "T12": - sym_euler = np.zeros((12, 3)) - sym_euler[1] = spr.from_rotvec( - 2 * np.pi / 3 * np.array([0, 0, 1]) - ).as_euler("zyz") - sym_euler[2] = spr.from_rotvec( - 4 * np.pi / 3 * np.array([0, 0, 1]) - ).as_euler("zyz") - - return sym_euler, sym_euler.shape[0] - def WD(self, J, euler): # compute Wigner D matrix alpha = euler[:, 0] From 3c630e007e868ebbe89d81c4ca69041894894250 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 20 Apr 2026 13:31:19 -0400 Subject: [PATCH 011/161] Use aspire ZYZ rotation convention --- src/aspire/abinitio/commonline_nug.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index abedcd5431..2f9d8a9c19 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -10,7 +10,7 @@ from aspire.nufft import nufft from aspire.numeric import fft, xp from aspire.operators import PolarFT, wemd_embed -from aspire.utils import cart2sph +from aspire.utils import Rotation, cart2sph from aspire.volume import SymmetryGroup logger = logging.getLogger(__name__) @@ -1041,11 +1041,7 @@ def find_gamma(Xm, beta, alpha): Euler_est[:, 0] = find_alpha(X1) Euler_est[:, 1] = find_beta(X1) Euler_est[:, 2] = find_gamma(XS, Euler_est[:, 1], Euler_est[:, 0]) - R_est = np.zeros((N, 3, 3)) - for n in range(N): - R_est[n] = spr.from_euler("zyz", np.flip(Euler_est[n])).as_matrix().T - # note that the order of alpha and gamma are swapped due to convention - + R_est = Rotation.from_euler(Euler_est).matrices.transpose(0, 2, 1) return R_est, Euler_est #################### From f5caebda8dce5129ba6ffd67216b0c50cbc423c7 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 21 Apr 2026 16:06:49 -0400 Subject: [PATCH 012/161] vectorize psd_projection, transform_block, transform_back_block --- src/aspire/abinitio/commonline_nug.py | 42 +++++++++++++++++++-------- 1 file changed, 30 insertions(+), 12 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 2f9d8a9c19..cf8647d064 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -3,7 +3,6 @@ import numpy as np from scipy.io import loadmat -from scipy.spatial.transform import Rotation as spr from scipy.special import factorial from aspire.abinitio import CLOrient3D @@ -335,6 +334,7 @@ def admm_sym_J( Sq, ) = self.ADMM_preprocessing(C, Lmax, N, Ngrid) + n_pairs = N * (N - 1) // 2 rank_Ak, _ = self.compute_rank(Lmax) logger.info(f"Rank of Ak: {rank_Ak}") @@ -473,10 +473,10 @@ def update_S(C0, C1, yE, yEq, yI, X0, X1, Xd0, Xd1, Xq, rho, Lmax, N): ) toc1 = time.perf_counter() Time[1] += toc1 - tic1 + tic2 = time.perf_counter() - for count in range(N * (N - 1) // 2): - tmp = self.psd_projection(Sq[:, count].reshape(4, 4).T) - Sq[:, count] = tmp.T.reshape(16) + Sq = self.psd_projection(Sq.T.reshape(n_pairs, 4, 4)) + Sq = Sq.T.reshape(-1, n_pairs) toc2 = time.perf_counter() Time[2] += toc2 - tic2 return S0, S1, Sd0, Sd1, Sq @@ -822,7 +822,7 @@ def ADMM_preprocessing(self, C, Lmax, N, Ngrid): AI_mat_diag = xp.asarray(AI_mat_diag) / 1 AI_mat_offdiag = xp.asarray(AI_mat_offdiag) / 1 bI = -(Lmax + 2) * (Lmax + 1) / 2 / 1 - np.save("AI_mat_offdiag_aspire.npy", AI_mat_offdiag) + # largest eigenvalue for AIAIT Lambda = self.largest_eigenvalue(AI_mat_offdiag, Ngrid, N) @@ -1203,12 +1203,17 @@ def mat_block(vecA, N, sz, IDX_upper, IDX_lower, idx_offdiag): @staticmethod def psd_projection(B): # compute the PSD part of a symmstric matrix - evals, evecs = xp.linalg.eigh((B + B.T) / 2) + B_sym = (B + B.swapaxes(-1, -2)) / 2 + evals, evecs = xp.linalg.eigh(B_sym) evals = xp.maximum(evals, 0) - return (evecs * evals) @ evecs.T + return (evecs * evals[..., None, :]) @ evecs.swapaxes(-1, -2) @staticmethod def transform_block(A, k, Pk=None): + single = A.ndim == 2 + if single: + A = A[None, :, :] + if Pk is None: dk = 2 * k + 1 Pk = xp.eye(dk) @@ -1218,14 +1223,26 @@ def transform_block(A, k, Pk=None): (m + 2 * el + 1, m + 2 * el), : ] AT = Pk @ A @ Pk.T - return AT[:k, :k].T.reshape(-1), AT[k:, k:].T.reshape(-1) + A0 = AT[:, :k, :k].swapaxes(-1, -2).reshape(A.shape[0], -1) + A1 = AT[:, k:, k:].swapaxes(-1, -2).reshape(A.shape[0], -1) + + if single: + return A0[0], A1[0] + + return A0, A1 @staticmethod def transform_back_block(A0, A1, k, Pk=None): dk = 2 * k + 1 - A = xp.zeros((dk, dk)) - A[:k, :k] = A0.reshape(k, k).T - A[k:, k:] = A1.reshape(k + 1, k + 1).T + single = A0.ndim == 1 + + if single: + A0 = A0[None, :] + A1 = A1[None, :] + + A = xp.zeros((A0.shape[0], dk, dk)) + A[:, :k, :k] = A0.reshape(-1, k, k).swapaxes(-1, -2) + A[:, k:, k:] = A1.reshape(-1, k + 1, k + 1).swapaxes(-1, -2) if Pk is None: Pk = xp.eye(dk) for m in range(k): @@ -1233,4 +1250,5 @@ def transform_back_block(A0, A1, k, Pk=None): Pk[(m + 2 * el, m + 2 * el + 1), :] = Pk[ (m + 2 * el + 1, m + 2 * el), : ] - return Pk.T @ A @ Pk + out = Pk.T @ A @ Pk + return out[0] if single else out From f134b01f757af1caebf3b2ee1a825bf392dd4b08 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 22 Apr 2026 11:12:27 -0400 Subject: [PATCH 013/161] Vectorize update_S. 2x speedup! --- src/aspire/abinitio/commonline_nug.py | 27 ++++++++++++++++----------- 1 file changed, 16 insertions(+), 11 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index cf8647d064..27f7f48912 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -458,19 +458,24 @@ def update_S(C0, C1, yE, yEq, yI, X0, X1, Xd0, Xd1, Xq, rho, Lmax, N): S1[d1[k - 1] : d1[k], :] = self.vec_block(tmp, N, k + 1, IDX_upper) toc0 = time.perf_counter() Time[0] += toc0 - tic0 + + # See if we can switch everything to C order before here. + Sd0 = Sd0.T + Sd1 = Sd1.T tic1 = time.perf_counter() for k in range(1, Lmax + 1): - for n in range(N): - tmp = self.transform_back_block( - Sd0[d0[k - 1] : d0[k], n], - Sd1[d1[k - 1] : d1[k], n], - k, - P[k - 1], - ) - tmp = self.psd_projection(tmp) - Sd0[d0[k - 1] : d0[k], n], Sd1[d1[k - 1] : d1[k], n] = ( - self.transform_block(tmp, k, P[k - 1]) - ) + tmp = self.transform_back_block( + Sd0[:, d0[k - 1] : d0[k]], + Sd1[:, d1[k - 1] : d1[k]], + k, + P[k - 1], + ) + tmp = self.psd_projection(tmp) + Sd0[:, d0[k - 1] : d0[k]], Sd1[:, d1[k - 1] : d1[k]] = ( + self.transform_block(tmp, k, P[k - 1]) + ) + Sd0 = Sd0.T + Sd1 = Sd1.T toc1 = time.perf_counter() Time[1] += toc1 - tic1 From 864a2bae87e2aeae9ca8c5fe244b5439bafd1d29 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 22 Apr 2026 12:02:32 -0400 Subject: [PATCH 014/161] vectorize print update component. --- src/aspire/abinitio/commonline_nug.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 27f7f48912..45df6798b1 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -611,15 +611,16 @@ def print_updates(verbose=True): res_psdX = res_psdX / (1 + np.linalg.norm(X0) + np.linalg.norm(X1)) res_psdD = 0 for k in range(1, Lmax + 1): - for n in range(N): - tmp = self.transform_back_block( - Xd0[d0[k - 1] : d0[k], n], - Xd1[d1[k - 1] : d1[k], n], - k, - P[k - 1], - ) - res_psdD += np.linalg.norm(self.psd_projection(-tmp)) - # res_psdD+=norm(self.psd_projection(-tmp))/(1+norm(tmp)) + tmp = self.transform_back_block( + Xd0[d0[k - 1] : d0[k]].T, + Xd1[d1[k - 1] : d1[k]].T, + k, + P[k - 1], + ) + res_psdD += np.linalg.norm( + self.psd_projection(-tmp), axis=(-2, -1) + ).sum() + # res_psdD+=norm(self.psd_projection(-tmp))/(1+norm(tmp)) res_psdD = res_psdD / (1 + np.linalg.norm(Xd0) + np.linalg.norm(Xd1)) res_psdQ = 0 for count in range(N * (N - 1) // 2): From a1022b8f12dedbacf5710bc30f443bd07e84eabf Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 24 Apr 2026 15:36:19 -0400 Subject: [PATCH 015/161] intial add of proximal_refine --- src/aspire/abinitio/commonline_nug.py | 97 +++++++++++++++++++++++++++ 1 file changed, 97 insertions(+) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 45df6798b1..ccd0f33226 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -39,6 +39,7 @@ def __init__( mult=1.5, Ngrid=16317, Nstep_yI=10, + perform_pr=False, **kwargs, ): """ @@ -70,6 +71,7 @@ def __init__( self.mult = mult self.Ngrid = Ngrid self.Nstep_yI = Nstep_yI + self.perform_pr = perform_pr # Handle symmetry self.sym_grp = SymmetryGroup.parse(symmetry) @@ -98,6 +100,26 @@ def estimate_rotations(self): self.Nstep_yI, ) + if self.perform_pr: + weight = 1 / (1 + np.arange(self.Lmax)) + Penalty = [1, 1, 1, 1] + r = [3, 2, 1, 0] + X_est = self.proximal_refine( + X_est, + C, + self.n_img, + weight, + Penalty, + r, + self.Ngrid, + self.max_iter, + self.rho, + self.ratio, + self.factor, + self.mult, + self.Nstep_yI, + ) + R_est, Euler_est = self.euler_est(X_est[0], X_est[self.n_sym - 1]) self.rotations = R_est @@ -937,6 +959,81 @@ def discretize_SO3(self): return SO3 + ############################ + # Proximal Refinement Step # + ############################ + + def proximal_refine( + self, + X_admm, + C, + N, + weight, + Penalty, + r, + Ngrid, + max_iter, + rho, + ratio, + factor, + mult, + Nstep_yI, + verbose=True, + ): + + def Ak(J, Euler): + # compute Ak matrix + order = Euler.shape[0] + # A = np.zeros((2 * J + 1, 2 * J + 1), dtype=np.complex128) + # for i in range(order): + # A += self.WD(J, Euler[i]) + A = self.WD(J, Euler).sum(axis=0) + return np.round(A / order, 10) + + rank_Ak = xp.zeros(self.Lmax) + for k in range(self.Lmax): + C[k] = xp.asnumpy(C[k]) + rank_Ak[k] = np.linalg.matrix_rank( + Ak(k + 1, self.sym_euler), tol=1e-6, hermitian=True + ) + + def low_rank_proj(X, r): + Xproj = [] + for k in range(self.Lmax): + dk = 2 * k + 3 + rk = int(rank_Ak[k] * 2) + r + tmp = np.copy(X[k]) + for i in range(N): + u, s, v = np.linalg.svd( + tmp[i * dk : (i + 1) * dk, i * dk : (i + 1) * dk] + ) + tmp[i * dk : (i + 1) * dk, i * dk : (i + 1) * dk] = ( + u[:, :rk] @ np.diag(s[:rk]) @ v[:rk] + ) + Xproj.append(tmp) + return Xproj + + Niter = len(r) + CC = [None] * self.Lmax + for iter in range(Niter): + X_prox = low_rank_proj(X_admm, r[iter]) + for k in range(self.Lmax): + CC[k] = C[k] - Penalty[iter] * weight[k] * (X_prox[k] + X_prox[k].T) / 2 + X_prox = self.admm_sym_J( + CC, + self.Lmax, + N, + Ngrid, + max_iter, + rho, + ratio, + factor, + mult, + Nstep_yI, + verbose, + ) + return X_prox + ######################### # Euler Estimation Step # ######################### From b680ff9ec8e5b378302b4bbde9438920a43a5ddd Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 30 Apr 2026 14:31:03 -0400 Subject: [PATCH 016/161] Update base class. Fix Proximal Refinement bug. Add PR update diagnostic log. --- src/aspire/abinitio/commonline_nug.py | 55 +++++++++++++++++++-------- 1 file changed, 39 insertions(+), 16 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index ccd0f33226..3998dfc90a 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -5,7 +5,7 @@ from scipy.io import loadmat from scipy.special import factorial -from aspire.abinitio import CLOrient3D +from aspire.abinitio import Orient3D from aspire.nufft import nufft from aspire.numeric import fft, xp from aspire.operators import PolarFT, wemd_embed @@ -15,7 +15,7 @@ logger = logging.getLogger(__name__) -class CommonlineNUG(CLOrient3D): +class CommonlineNUG(Orient3D): """ Class to estimate 3D orientations using non-uqique games. """ @@ -980,28 +980,33 @@ def proximal_refine( Nstep_yI, verbose=True, ): - def Ak(J, Euler): # compute Ak matrix order = Euler.shape[0] - # A = np.zeros((2 * J + 1, 2 * J + 1), dtype=np.complex128) - # for i in range(order): - # A += self.WD(J, Euler[i]) A = self.WD(J, Euler).sum(axis=0) return np.round(A / order, 10) - rank_Ak = xp.zeros(self.Lmax) + def rel_change(A, B, eps=1e-12): + num = 0.0 + den = 0.0 + for k in range(len(A)): + num += np.linalg.norm(A[k] - B[k]) ** 2 + den += np.linalg.norm(B[k]) ** 2 + return np.sqrt(num) / max(np.sqrt(den), eps) + + rank_Ak = np.zeros(self.Lmax) + C_base = [None] * self.Lmax for k in range(self.Lmax): - C[k] = xp.asnumpy(C[k]) + C_base[k] = xp.asnumpy(C[k]).copy() rank_Ak[k] = np.linalg.matrix_rank( Ak(k + 1, self.sym_euler), tol=1e-6, hermitian=True ) - def low_rank_proj(X, r): + def low_rank_proj(X, r_step): Xproj = [] for k in range(self.Lmax): dk = 2 * k + 3 - rk = int(rank_Ak[k] * 2) + r + rk = min(int(rank_Ak[k] * 2) + r_step, dk) tmp = np.copy(X[k]) for i in range(N): u, s, v = np.linalg.svd( @@ -1015,11 +1020,18 @@ def low_rank_proj(X, r): Niter = len(r) CC = [None] * self.Lmax - for iter in range(Niter): - X_prox = low_rank_proj(X_admm, r[iter]) + current = [np.copy(Xk) for Xk in X_admm] + + for step in range(Niter): + X_proj = low_rank_proj(current, r[step]) + for k in range(self.Lmax): - CC[k] = C[k] - Penalty[iter] * weight[k] * (X_prox[k] + X_prox[k].T) / 2 - X_prox = self.admm_sym_J( + CC[k] = ( + C_base[k] + - Penalty[step] * weight[k] * (X_proj[k] + X_proj[k].T) / 2 + ) + + X_next = self.admm_sym_J( CC, self.Lmax, N, @@ -1030,9 +1042,20 @@ def low_rank_proj(X, r): factor, mult, Nstep_yI, - verbose, + verbose=False, ) - return X_prox + + if verbose: + logger.info( + "Proximal refine step %d/%d: relative update %.3e", + step + 1, + Niter, + rel_change(X_next, current), + ) + + current = [np.copy(Xk) for Xk in X_next] + + return current ######################### # Euler Estimation Step # From 22831d6abf66818c2c600ccd26c4dda9b5ac5612 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 5 May 2026 14:23:49 -0400 Subject: [PATCH 017/161] Add euler_est_Dm --- src/aspire/abinitio/commonline_nug.py | 151 +++++++++++++++++++++++++- 1 file changed, 149 insertions(+), 2 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 3998dfc90a..5635aa6edf 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -10,7 +10,7 @@ from aspire.numeric import fft, xp from aspire.operators import PolarFT, wemd_embed from aspire.utils import Rotation, cart2sph -from aspire.volume import SymmetryGroup +from aspire.volume import CnSymmetryGroup, DnSymmetryGroup, SymmetryGroup logger = logging.getLogger(__name__) @@ -120,7 +120,11 @@ def estimate_rotations(self): self.Nstep_yI, ) - R_est, Euler_est = self.euler_est(X_est[0], X_est[self.n_sym - 1]) + if isinstance(self.sym_grp, CnSymmetryGroup): + R_est, Euler_est = self.euler_est(X_est[0], X_est[self.n_sym - 1]) + elif isinstance(self.sym_grp, DnSymmetryGroup): + R_est, Euler_est = self.euler_est_Dm(X_est) + self.rotations = R_est return R_est @@ -1170,6 +1174,149 @@ def find_gamma(Xm, beta, alpha): R_est = Rotation.from_euler(Euler_est).matrices.transpose(0, 2, 1) return R_est, Euler_est + def euler_est_Dm(self, X_est): + + X2 = X_est[1] + S = self.sym_grp.order + N = self.n_img + XS = X_est[S - 1] + dk = 2 * S + 1 + + def find_alpha_beta(X2): + def find_phase(A, B): + # find a number c that minimizes ||cA-B||_F + Ar = np.real(A) + Ai = np.imag(A) + Br = np.real(B) + Bi = np.imag(B) + c = (np.vdot(Ar, Br) + np.vdot(Ai, Bi)) / ( + np.vdot(Ar, Ar) + np.vdot(Ai, Ai) + ) + 1j * (np.vdot(Ar, Bi) - np.vdot(Ai, Br)) / ( + np.vdot(Ar, Ar) + np.vdot(Ai, Ai) + ) + return c / abs(c) + + T, Tinv = self.complex2real(2) + X2 = np.kron(np.eye(N), T) @ X2 @ np.kron(np.eye(N), Tinv) + + B1 = np.zeros((N, N)) + for i in range(N): + for j in range(N): + Xij = X2[5 * i : 5 * (i + 1), 5 * j : 5 * (j + 1)] + B1[i, j] = np.real((Xij[2, 2] - 2 * abs(Xij[0, 0]) + 0.5) / 3 * 2) + e1, v1 = np.linalg.eigh(B1) + idx = np.argmax(e1) + b1 = v1[:, idx] * np.sqrt(e1[idx]) * np.sign(v1[0, idx]) + beta_est = np.arccos(np.clip(np.sqrt(b1), -1, 1)) % np.pi + + Aminus = np.zeros((N, N), dtype=complex) + Aplus = np.zeros((N, N), dtype=complex) + for i in range(N): + for j in range(N): + Xij = X2[5 * i : 5 * (i + 1), 5 * j : 5 * (j + 1)] + if abs(beta_est[i]) < 1e-6: + beta_est[i] = 1e-6 + if abs(beta_est[j]) < 1e-6: + beta_est[j] = 1e-6 + Aminus[i, j] = ( + Xij[1, 1] + / np.sin(2 * beta_est[i]) + / np.sin(2 * beta_est[j]) + * 8 + / 3 + ) + Aplus[i, j] = ( + -Xij[1, 3] + / np.sin(2 * beta_est[i]) + / np.sin(2 * beta_est[j]) + * 8 + / 3 + ) + + evals, evecs = np.linalg.eigh(Aminus) + idx = np.argmax(abs(evals)) + Z = evecs[:, idx] * np.sqrt(abs(evals[idx])) + c = find_phase(Z[:, None] @ Z[:, None].T, Aplus) + Z = np.sqrt(c) * Z + alpha_est = (np.angle(Z)) % (2 * np.pi) + + return alpha_est, beta_est + + def find_gamma(Xm, alpha, beta): + def LS_D(W1, W2, W3, W4, Br, Bi): + A = np.array( + [ + [np.vdot(W1 + W4, W1 + W4), np.vdot(W1 + W4, W2 + W3)], + [np.vdot(W2 + W3, W1 + W4), np.vdot(W2 + W3, W2 + W3)], + ] + ) + B = np.array([np.vdot(W1 + W4, Br), np.vdot(W2 + W3, Br)]) + a, c = np.linalg.lstsq(A, B)[0] + + A = np.array( + [ + [np.vdot(W1 - W4, W1 - W4), np.vdot(W1 - W4, W3 - W2)], + [np.vdot(W1 - W4, W3 - W2), np.vdot(W3 - W2, W3 - W2)], + ] + ) + B = np.array([np.vdot(W1 - W4, Bi), np.vdot(W3 - W2, Bi)]) + b, d = np.linalg.lstsq(A, B)[0] + return a + 1j * b + + [T, Tinv] = self.complex2real(S) + Xm = np.kron(np.eye(N), T) @ Xm @ np.kron(np.eye(N), Tinv) + C = np.zeros((N, N), dtype=complex) + Jk = np.ones(dk) + Jk[S + 1 :: 2] = -1 + Jk[S - 1 :: -2] = -1 + Jk = np.diag(Jk) + ws = self.Wd(S, beta) + for i in range(N): + wi = ws[i] + for j in range(i + 1, N): + Di = np.exp(-1j * np.arange(-S, S + 1) * alpha[i]) + Dj = np.exp(-1j * np.arange(-S, S + 1) * alpha[j]) + Xijm = Xm[dk * i : dk * (i + 1), dk * j : dk * (j + 1)] + DXijmD = np.diag(Di.conj()) @ Xijm @ np.diag(Dj) + wj = ws[j] + W1 = ( + wi[:, 0][:, np.newaxis] @ wj[:, 0][:, np.newaxis].T + + Jk @ wi[:, 0][:, np.newaxis] @ wj[:, 0][:, np.newaxis].T @ Jk + ) + W2 = ( + wi[:, -1][:, np.newaxis] @ wj[:, 0][:, np.newaxis].T + + Jk @ wi[:, -1][:, np.newaxis] @ wj[:, 0][:, np.newaxis].T @ Jk + ) + W3 = ( + wi[:, 0][:, np.newaxis] @ wj[:, -1][:, np.newaxis].T + + Jk @ wi[:, 0][:, np.newaxis] @ wj[:, -1][:, np.newaxis].T @ Jk + ) + W4 = ( + wi[:, -1][:, np.newaxis] @ wj[:, -1][:, np.newaxis].T + + Jk + @ wi[:, -1][:, np.newaxis] + @ wj[:, -1][:, np.newaxis].T + @ Jk + ) + Br = np.real(4 * DXijmD) + Bi = np.imag(4 * DXijmD) + C[i, j] = LS_D(W1, W2, W3, W4, Br, Bi) + C += C.T.conj() + np.eye(N) + evals, evecs = np.linalg.eigh(C) + idx = np.argmax(evals) + c = evecs[:, idx] * np.sqrt(evals[idx]) + return (np.angle(c) / S) % (2 * np.pi) + + alpha_est, beta_est = find_alpha_beta(X2) + gamma_est = find_gamma(XS, alpha_est, beta_est) + Euler_est = np.zeros((N, 3)) + Euler_est[:, 0] = alpha_est + Euler_est[:, 1] = beta_est + Euler_est[:, 2] = gamma_est + R_est = Rotation.from_euler(Euler_est).matrices.transpose(0, 2, 1) + + return R_est, Euler_est + #################### # Helper Functions # #################### From d700d69236e1f34e037ea12b81d57743ad951805 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 13 May 2026 15:32:07 -0400 Subject: [PATCH 018/161] Use in-house polarFT and Sinogram. Add symmetry handling and logs --- src/aspire/abinitio/commonline_nug.py | 48 +++++++++++++++++---------- 1 file changed, 30 insertions(+), 18 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 5635aa6edf..ece406f0bb 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -74,7 +74,18 @@ def __init__( self.perform_pr = perform_pr # Handle symmetry - self.sym_grp = SymmetryGroup.parse(symmetry) + if symmetry is None: + logger.info( + f"Symmetry not provided. Using Source symmetry: {str(self.src.symmetry_group)}" + ) + self.sym_grp = self.src.symmetry_group + else: + if symmetry != str(self.src.symmetry_group): + logger.info( + f"Provided symmetry, {symmetry}, does not match source, {str(self.src.symmetry_group)}" + ) + logger.info(f"Using provided symmetry: {symmetry}") + self.sym_grp = SymmetryGroup.parse(symmetry) self.sym_euler = self.sym_grp.rotations.angles self.n_sym = len(self.sym_euler) @@ -142,11 +153,12 @@ def compute_coeff(self, Img, loss, Lmax, T): Img_pft = np.zeros((L, n_theta, N), dtype=complex) # Replace with Image.project() later - Img = Img.asnumpy() - for n in range(N): - line_proj[:, :, n], Img_pft[:, :, n] = self.fast_radon_transform( - Img[n], angular_sampling - ) + line_proj = Img.project(angular_sampling).asnumpy().T + # Img = Img.asnumpy() + # for n in range(N): + # line_proj[:, :, n], Img_pft[:, :, n] = self.fast_radon_transform( + # Img[n], angular_sampling + # ) dim_wave = len(wemd_embed(line_proj[:, 0, 0])) WE = np.zeros((dim_wave, n_theta, N)) @@ -156,22 +168,22 @@ def compute_coeff(self, Img, loss, Lmax, T): def fij(alpha, gamma, i, j, loss): if loss == "l1": - Ii_hat = Img_pft[:, :, i] - Ij_hat = Img_pft[:, :, j] - idxi = np.round((alpha - np.pi / 2) * n_theta / 2 / np.pi) % n_theta - idxj = np.round((-gamma - np.pi / 2) * n_theta / 2 / np.pi) % n_theta + # Ii_hat = Img_pft[:, :, i] + # Ij_hat = Img_pft[:, :, j] + # idxi = np.round((alpha - np.pi / 2) * n_theta / 2 / np.pi) % n_theta + # idxj = np.round((-gamma - np.pi / 2) * n_theta / 2 / np.pi) % n_theta - Si = Ii_hat[:, int(idxi)] - Sj = Ij_hat[:, int(idxj)] + # Si = Ii_hat[:, int(idxi)] + # Sj = Ij_hat[:, int(idxj)] # Using aspire PolarFT. Replace later - # Ii_hat = self.pf_full[i] - # Ij_hat = self.pf_full[j] - # idxi = np.round((alpha - np.pi / 2) * n_theta / 2 / np.pi) % n_theta - # idxj = np.round((-gamma - np.pi / 2) * n_theta / 2 / np.pi) % n_theta + Ii_hat = self.pf_full[i] + Ij_hat = self.pf_full[j] + idxi = np.round((alpha - np.pi / 2) * n_theta / 2 / np.pi) % n_theta + idxj = np.round((-gamma - np.pi / 2) * n_theta / 2 / np.pi) % n_theta - # Si = Ii_hat[int(idxi)] - # Sj = Ij_hat[int(idxj)] + Si = Ii_hat[int(idxi)] + Sj = Ij_hat[int(idxj)] # norm_new = np.linalg.norm(Si - Sj, 1) return np.linalg.norm(Si - Sj, 1) From e1b4f37d7843b18d4f13937c980bd8f07d4ade32 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 13 May 2026 15:46:44 -0400 Subject: [PATCH 019/161] cleanup --- src/aspire/abinitio/commonline_nug.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index ece406f0bb..23624aae1f 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -149,11 +149,12 @@ def compute_coeff(self, Img, loss, Lmax, T): N, L, _ = Img.shape n_theta = 360 angular_sampling = np.arange(0, 360, 1) - line_proj = np.zeros((L, n_theta, N)) - Img_pft = np.zeros((L, n_theta, N), dtype=complex) - # Replace with Image.project() later + # Using ASPIRE Image.project(). Leaving original method in comments for now. line_proj = Img.project(angular_sampling).asnumpy().T + + # line_proj = np.zeros((L, n_theta, N)) + # Img_pft = np.zeros((L, n_theta, N), dtype=complex) # Img = Img.asnumpy() # for n in range(N): # line_proj[:, :, n], Img_pft[:, :, n] = self.fast_radon_transform( From 1f5b6b8dbcf1c5c6694eca82bd82b070ab545688 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 14 May 2026 14:26:39 -0400 Subject: [PATCH 020/161] Use saff-kuijlaars method for S2 grid --- src/aspire/abinitio/__init__.py | 1 + src/aspire/abinitio/commonline_d2.py | 35 ++----------------------- src/aspire/abinitio/commonline_nug.py | 26 +++++++++--------- src/aspire/abinitio/commonline_utils.py | 31 ++++++++++++++++++++++ 4 files changed, 46 insertions(+), 47 deletions(-) diff --git a/src/aspire/abinitio/__init__.py b/src/aspire/abinitio/__init__.py index 4b6248c4c7..cb68c8f544 100644 --- a/src/aspire/abinitio/__init__.py +++ b/src/aspire/abinitio/__init__.py @@ -3,6 +3,7 @@ from .commonline_utils import ( build_outer_products, g_sync, + saff_kuijlaars, ) from .commonline_base import Orient3D from .commonline_matrix import CLOrient3D diff --git a/src/aspire/abinitio/commonline_d2.py b/src/aspire/abinitio/commonline_d2.py index 34f28ee775..13ec282346 100644 --- a/src/aspire/abinitio/commonline_d2.py +++ b/src/aspire/abinitio/commonline_d2.py @@ -10,7 +10,7 @@ from aspire.utils.random import randn from aspire.volume import DnSymmetryGroup -from .commonline_utils import _generate_shift_phase_and_filter +from .commonline_utils import _generate_shift_phase_and_filter, saff_kuijlaars logger = logging.getLogger(__name__) @@ -160,7 +160,7 @@ def _generate_lookup_data(self): logger.info("Generating commonline lookup data.") # Generate uniform grid on sphere with Saff-Kuijlaars and take one quarter # of sphere because of D2 symmetry redundancy. - sphere_grid = self._saff_kuijlaars(self.grid_res) + sphere_grid = saff_kuijlaars(self.grid_res) octant1_mask = np.all(sphere_grid > 0, axis=1) octant2_mask = ( (sphere_grid[:, 0] > 0) & (sphere_grid[:, 1] > 0) & (sphere_grid[:, 2] < 0) @@ -1822,37 +1822,6 @@ def _circ_seq(n1, n2, L): return seq - @staticmethod - def _saff_kuijlaars(N): - """ - Generates N vertices on the unit sphere that are approximately evenly distributed. - - This implements the recommended algorithm in spherical coordinates - (theta, phi) according to "Distributing many points on a sphere" - by E.B. Saff and A.B.J. Kuijlaars, Mathematical Intelligencer 19.1 - (1997) 5--11. - - :param N: Number of vertices to generate. - - :return: Nx3 array of vertices in cartesian coordinates. - """ - k = np.arange(1, N + 1) - h = -1 + 2 * (k - 1) / (N - 1) - theta = np.arccos(h) - phi = np.zeros(N) - - for i in range(1, N - 1): - phi[i] = (phi[i - 1] + 3.6 / (np.sqrt(N * (1 - h[i] ** 2)))) % (2 * np.pi) - - # Spherical coordinates - x = np.sin(theta) * np.cos(phi) - y = np.sin(theta) * np.sin(phi) - z = np.cos(theta) - - mesh = np.column_stack((x, y, z)) - - return mesh - @staticmethod def _mark_equators(sphere_grid, eq_filter_angle): """ diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 23624aae1f..70cee9ccb0 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -12,6 +12,8 @@ from aspire.utils import Rotation, cart2sph from aspire.volume import CnSymmetryGroup, DnSymmetryGroup, SymmetryGroup +from .commonline_utils import saff_kuijlaars + logger = logging.getLogger(__name__) @@ -37,7 +39,7 @@ def __init__( ratio=1, factor=1.0, mult=1.5, - Ngrid=16317, + S2_grid=441, Nstep_yI=10, perform_pr=False, **kwargs, @@ -69,7 +71,7 @@ def __init__( self.ratio = ratio self.factor = factor self.mult = mult - self.Ngrid = Ngrid + self.S2_grid = S2_grid self.Nstep_yI = Nstep_yI self.perform_pr = perform_pr @@ -102,7 +104,6 @@ def estimate_rotations(self): C, self.Lmax, self.n_img, - self.Ngrid, self.max_iter, self.rho, self.ratio, @@ -122,7 +123,6 @@ def estimate_rotations(self): weight, Penalty, r, - self.Ngrid, self.max_iter, self.rho, self.ratio, @@ -336,7 +336,6 @@ def admm_sym_J( C, Lmax, N, - Ngrid, max_iter, rho, ratio, @@ -371,9 +370,10 @@ def admm_sym_J( S0, S1, Sq, - ) = self.ADMM_preprocessing(C, Lmax, N, Ngrid) + ) = self.ADMM_preprocessing(C, Lmax, N) n_pairs = N * (N - 1) // 2 + Ngrid = self.Ngrid rank_Ak, _ = self.compute_rank(Lmax) logger.info(f"Rank of Ak: {rank_Ak}") @@ -741,7 +741,7 @@ def print_updates(verbose=True): X_admm[k] = xp.asnumpy(X_admm[k]) return X_admm - def ADMM_preprocessing(self, C, Lmax, N, Ngrid): + def ADMM_preprocessing(self, C, Lmax, N): # compute necessary quantities for ADMM # compute some useful index sets count = 0 @@ -799,8 +799,7 @@ def ADMM_preprocessing(self, C, Lmax, N, Ngrid): bEq = xp.repeat(bEq[:, xp.newaxis], N * (N - 1) // 2, axis=1) # AI and bI - W0, W1 = self.compute_fejer_weights() - + W0, W1, Ngrid = self.compute_fejer_weights() # AI_mat=np.zeros((Ngrid,D0+D1)) # for p in range(Ngrid): # w0=np.zeros(D0); w1=np.zeros(D1) @@ -811,7 +810,6 @@ def ADMM_preprocessing(self, C, Lmax, N, Ngrid): # AI_mat[p,:d0[-1]]=w0; AI_mat[p,d0[-1]:]=w1 # AI_mat=xp.asarray(AI_mat) / 10; # bI=-(Lmax+2)*(Lmax+1)/2 / 10 - AI_mat_offdiag = np.zeros((Ngrid, D0 + D1)) for p in range(Ngrid): w0 = np.zeros(D0) @@ -886,6 +884,7 @@ def ADMM_preprocessing(self, C, Lmax, N, Ngrid): Sq = xp.zeros(Xq.shape) # S0=xp.zeros(X0.shape); S1=xp.zeros(X1.shape); Sq=xp.zeros(Xq.shape) + self.Ngrid = Ngrid # return C0,C1,normC,AEq,bEq,AEqAEqtinv,AI_mat,bI,Lambda,d0,d1,D0,D1,idx_diag,idx_offdiag,IDX_upper,IDX_lower,X0,X1,Xq,S0,S1,Sq return ( C0, @@ -951,10 +950,11 @@ def permutek_block(Ak, k): W0.append(W0k) W1.append(W1k) - return W0, W1 + return W0, W1, Ngrid def discretize_SO3(self): - S2 = loadmat("design20.mat")["design"] + # S2 = loadmat("design20.mat")["design"] + S2 = saff_kuijlaars(self.S2_grid) S2_size = S2.shape[0] # discretize S1 @@ -988,7 +988,6 @@ def proximal_refine( weight, Penalty, r, - Ngrid, max_iter, rho, ratio, @@ -1052,7 +1051,6 @@ def low_rank_proj(X, r_step): CC, self.Lmax, N, - Ngrid, max_iter, rho, ratio, diff --git a/src/aspire/abinitio/commonline_utils.py b/src/aspire/abinitio/commonline_utils.py index 46a3a932c0..5fe88a21b1 100644 --- a/src/aspire/abinitio/commonline_utils.py +++ b/src/aspire/abinitio/commonline_utils.py @@ -288,6 +288,37 @@ def _cl_angles_to_ind(cl_angles, n_theta): return ind +def saff_kuijlaars(N): + """ + Generates N vertices on the unit sphere that are approximately evenly distributed. + + This implements the recommended algorithm in spherical coordinates + (theta, phi) according to "Distributing many points on a sphere" + by E.B. Saff and A.B.J. Kuijlaars, Mathematical Intelligencer 19.1 + (1997) 5--11. + + :param N: Number of vertices to generate. + + :return: Nx3 array of vertices in cartesian coordinates. + """ + k = np.arange(1, N + 1) + h = -1 + 2 * (k - 1) / (N - 1) + theta = np.arccos(h) + phi = np.zeros(N) + + for i in range(1, N - 1): + phi[i] = (phi[i - 1] + 3.6 / (np.sqrt(N * (1 - h[i] ** 2)))) % (2 * np.pi) + + # Spherical coordinates + x = np.sin(theta) * np.cos(phi) + y = np.sin(theta) * np.sin(phi) + z = np.cos(theta) + + mesh = np.column_stack((x, y, z)) + + return mesh + + def g_sync(rots, order, rots_gt): """ Given ground truth rotations, synchronize estimated rotations over From fef88c80c3a8db7e363a55d817869984805afbe3 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 15 May 2026 10:41:14 -0400 Subject: [PATCH 021/161] initial test file --- src/aspire/abinitio/commonline_nug.py | 2 +- tests/test_nug.py | 97 +++++++++++++++++++++++++++ 2 files changed, 98 insertions(+), 1 deletion(-) create mode 100644 tests/test_nug.py diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 70cee9ccb0..c3bae03b70 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -28,7 +28,7 @@ def __init__( symmetry=None, n_rad=None, n_theta=360, - max_shift=0.15, + max_shift=0, shift_step=1, mask=True, Lmax=12, diff --git a/tests/test_nug.py b/tests/test_nug.py new file mode 100644 index 0000000000..4affb492e7 --- /dev/null +++ b/tests/test_nug.py @@ -0,0 +1,97 @@ +import numpy as np +import pytest + +from aspire.abinitio import CommonlineNUG +from aspire.source import Simulation +from aspire.utils import ( + J_conjugate, + Random, + Rotation, + all_pairs, + mean_aligned_angular_distance, + utest_tolerance, +) +from aspire.volume import CnSymmetricVolume, CnSymmetryGroup + +DTYPE = [np.float32] +RESOLUTION = [48, 49] +N_IMG = [5] +OFFSETS = [0] +ORDER = [3, 4] +PR = [False] +SEED = 1980 + + +@pytest.fixture(params=DTYPE, ids=lambda x: f"dtype={x}", scope="module") +def dtype(request): + return request.param + + +@pytest.fixture(params=RESOLUTION, ids=lambda x: f"resolution={x}", scope="module") +def resolution(request): + return request.param + + +@pytest.fixture(params=N_IMG, ids=lambda x: f"n images={x}", scope="module") +def n_img(request): + return request.param + + +@pytest.fixture(params=OFFSETS, ids=lambda x: f"offsets={x}", scope="module") +def offsets(request): + return request.param + + +@pytest.fixture(params=ORDER, ids=lambda x: f"order={x}", scope="module") +def order(request): + return request.param + + +@pytest.fixture(params=PR, ids=lambda x: f"proximal_refine={x}", scope="module") +def proximal_refine(request): + return request.param + + +############ +# Fixtures # +############ + +@pytest.fixture(scope="module") +def source(n_img, resolution, dtype, offsets, order): + vol = CnSymmetricVolume( + L=resolution, order=order, C=1, K=100, dtype=dtype, seed=SEED + ).generate() + + src = Simulation( + n=n_img, + L=resolution, + vols=vol, + offsets=offsets, + amplitudes=1, + seed=SEED, + ) + src = src.cache() # Precompute image stack + + return src + +@pytest.fixture(scope="module") +def orient_est(src, proximal_refine): + orient_est = CommonlineNUG( + src, + perform_pr=proximal_refine, + ) + + return orient_est + + +######### +# Tests # +######### + + +def test_dtypes(orient_est): + """ + Check dtypes for each major step of the algorithm. + """ + pass + From e04864e22192def7df9cb09cb7bde885394f4e12 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 15 May 2026 14:30:08 -0400 Subject: [PATCH 022/161] clean up estimate_rotations --- src/aspire/abinitio/commonline_nug.py | 160 +++++++++++--------------- tests/test_nug.py | 12 +- 2 files changed, 71 insertions(+), 101 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index c3bae03b70..f64a4f2689 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -42,6 +42,7 @@ def __init__( S2_grid=441, Nstep_yI=10, perform_pr=False, + verbose=True, **kwargs, ): """ @@ -74,6 +75,7 @@ def __init__( self.S2_grid = S2_grid self.Nstep_yI = Nstep_yI self.perform_pr = perform_pr + self.verbose = verbose # Handle symmetry if symmetry is None: @@ -98,56 +100,24 @@ def _build_full_pft(self): self.pf_full = PolarFT.half_to_full(pf) def estimate_rotations(self): - imgs = self.src.images[:] - C = self.compute_coeff(imgs, self.loss, self.Lmax, T=self.T) - X_est = self.admm_sym_J( - C, - self.Lmax, - self.n_img, - self.max_iter, - self.rho, - self.ratio, - self.factor, - self.mult, - self.Nstep_yI, - ) - - if self.perform_pr: - weight = 1 / (1 + np.arange(self.Lmax)) - Penalty = [1, 1, 1, 1] - r = [3, 2, 1, 0] - X_est = self.proximal_refine( - X_est, - C, - self.n_img, - weight, - Penalty, - r, - self.max_iter, - self.rho, - self.ratio, - self.factor, - self.mult, - self.Nstep_yI, - ) - - if isinstance(self.sym_grp, CnSymmetryGroup): - R_est, Euler_est = self.euler_est(X_est[0], X_est[self.n_sym - 1]) - elif isinstance(self.sym_grp, DnSymmetryGroup): - R_est, Euler_est = self.euler_est_Dm(X_est) - - self.rotations = R_est - - return R_est + self.compute_coeff() + self.perform_admm() + self.euler_est() + return self.rotations ####################### # Compute Coeffs Step # ####################### - def compute_coeff(self, Img, loss, Lmax, T): + def compute_coeff(self): # compute the coefficient matrix - N, L, _ = Img.shape - n_theta = 360 + Img = self.src.images[:] + N = self.n_img + L = self.src.L + n_theta = self.n_theta + loss = self.loss + Lmax = self.Lmax + T = self.T angular_sampling = np.arange(0, 360, 1) # Using ASPIRE Image.project(). Leaving original method in comments for now. @@ -264,7 +234,7 @@ def fijhat_k(k, F): ) C[k - 1] = np.round(C[k - 1], 10) - return C + self.C = C @staticmethod def fast_radon_transform(array, angles, use_ramp=False): @@ -331,19 +301,31 @@ def complex2real(ell): # ADMM Step # ############# - def admm_sym_J( - self, - C, - Lmax, - N, - max_iter, - rho, - ratio, - factor, - mult=1, - Nstep_yI=20, - verbose=True, - ): + def perform_admm(self): + X_est = self.admm_sym_J(self.C, self.verbose) + + if self.perform_pr: + weight = 1 / (1 + np.arange(self.Lmax)) + Penalty = [1, 1, 1, 1] + r = [3, 2, 1, 0] + X_est = self.proximal_refine( + X_est, + weight, + Penalty, + r, + ) + self.X_est = X_est + + def admm_sym_J(self, C, verbose): + Lmax = self.Lmax + N = self.n_img + max_iter = self.max_iter + rho = self.rho + ratio = self.ratio + factor = self.factor + mult = self.mult + Nstep_yI = self.Nstep_yI + # admm for symmetric case ( C0, @@ -370,7 +352,7 @@ def admm_sym_J( S0, S1, Sq, - ) = self.ADMM_preprocessing(C, Lmax, N) + ) = self.ADMM_preprocessing(C) n_pairs = N * (N - 1) // 2 Ngrid = self.Ngrid @@ -598,7 +580,7 @@ def update_rho(X0, X1, Xd0, Xd1, Xq, bE, bEq, bI, res_X, rho, factor, normC): rho = rho / factor return rho, p_resnorm, d_resnorm - def print_updates(verbose=True): + def print_updates(verbose): # X_admm=transform_coeff_back(X0,X1,Lmax,N); obj_p=0 # for k in range(Lmax): obj_p+=xp.trace(C[k]@X_admm[k]) if verbose: @@ -741,9 +723,11 @@ def print_updates(verbose=True): X_admm[k] = xp.asnumpy(X_admm[k]) return X_admm - def ADMM_preprocessing(self, C, Lmax, N): + def ADMM_preprocessing(self, C): # compute necessary quantities for ADMM # compute some useful index sets + Lmax = self.Lmax + N = self.n_img count = 0 idx_diag = [] idx_offdiag = [] @@ -980,22 +964,17 @@ def discretize_SO3(self): # Proximal Refinement Step # ############################ - def proximal_refine( - self, - X_admm, - C, - N, - weight, - Penalty, - r, - max_iter, - rho, - ratio, - factor, - mult, - Nstep_yI, - verbose=True, - ): + def proximal_refine(self, X_admm, weight, Penalty, r): + Lmax = self.Lmax + N = self.n_img + max_iter = self.max_iter + rho = self.rho + ratio = self.ratio + factor = self.factor + mult = self.mult + Nstep_yI = self.Nstep_yI + C = self.C + def Ak(J, Euler): # compute Ak matrix order = Euler.shape[0] @@ -1047,20 +1026,9 @@ def low_rank_proj(X, r_step): - Penalty[step] * weight[k] * (X_proj[k] + X_proj[k].T) / 2 ) - X_next = self.admm_sym_J( - CC, - self.Lmax, - N, - max_iter, - rho, - ratio, - factor, - mult, - Nstep_yI, - verbose=False, - ) + X_next = self.admm_sym_J(CC, verbose=False) - if verbose: + if self.verbose: logger.info( "Proximal refine step %d/%d: relative update %.3e", step + 1, @@ -1076,7 +1044,17 @@ def low_rank_proj(X, r_step): # Euler Estimation Step # ######################### - def euler_est(self, X1, XS): + def euler_est(self): + X_est = self.X_est + if isinstance(self.sym_grp, CnSymmetryGroup): + R_est, Euler_est = self.euler_est_Cm(X_est[0], X_est[self.n_sym - 1]) + elif isinstance(self.sym_grp, DnSymmetryGroup): + R_est, Euler_est = self.euler_est_Dm(X_est) + + self.Euler_est = Euler_est + self.rotations = R_est + + def euler_est_Cm(self, X1, XS): S = self.n_sym N = self.n_img sym_euler = np.zeros((S, 3)) diff --git a/tests/test_nug.py b/tests/test_nug.py index 4affb492e7..a221f36bd7 100644 --- a/tests/test_nug.py +++ b/tests/test_nug.py @@ -3,15 +3,6 @@ from aspire.abinitio import CommonlineNUG from aspire.source import Simulation -from aspire.utils import ( - J_conjugate, - Random, - Rotation, - all_pairs, - mean_aligned_angular_distance, - utest_tolerance, -) -from aspire.volume import CnSymmetricVolume, CnSymmetryGroup DTYPE = [np.float32] RESOLUTION = [48, 49] @@ -56,6 +47,7 @@ def proximal_refine(request): # Fixtures # ############ + @pytest.fixture(scope="module") def source(n_img, resolution, dtype, offsets, order): vol = CnSymmetricVolume( @@ -74,6 +66,7 @@ def source(n_img, resolution, dtype, offsets, order): return src + @pytest.fixture(scope="module") def orient_est(src, proximal_refine): orient_est = CommonlineNUG( @@ -94,4 +87,3 @@ def test_dtypes(orient_est): Check dtypes for each major step of the algorithm. """ pass - From a8b8f7f8aacbd953e0631b6032ce2c869c944a26 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 19 May 2026 16:23:45 -0400 Subject: [PATCH 023/161] Explicit dtypes. Add dtype test. Add estimate_rotations test. --- src/aspire/abinitio/commonline_nug.py | 271 ++++++++++++++++---------- tests/test_nug.py | 66 ++++++- 2 files changed, 222 insertions(+), 115 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index f64a4f2689..944394c5a4 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -2,14 +2,13 @@ import time import numpy as np -from scipy.io import loadmat from scipy.special import factorial from aspire.abinitio import Orient3D from aspire.nufft import nufft from aspire.numeric import fft, xp from aspire.operators import PolarFT, wemd_embed -from aspire.utils import Rotation, cart2sph +from aspire.utils import Rotation, cart2sph, complex_type from aspire.volume import CnSymmetryGroup, DnSymmetryGroup, SymmetryGroup from .commonline_utils import saff_kuijlaars @@ -113,12 +112,11 @@ def compute_coeff(self): # compute the coefficient matrix Img = self.src.images[:] N = self.n_img - L = self.src.L n_theta = self.n_theta loss = self.loss Lmax = self.Lmax T = self.T - angular_sampling = np.arange(0, 360, 1) + angular_sampling = np.arange(0, 360, 1, dtype=np.float64) # Using ASPIRE Image.project(). Leaving original method in comments for now. line_proj = Img.project(angular_sampling).asnumpy().T @@ -130,13 +128,6 @@ def compute_coeff(self): # line_proj[:, :, n], Img_pft[:, :, n] = self.fast_radon_transform( # Img[n], angular_sampling # ) - - dim_wave = len(wemd_embed(line_proj[:, 0, 0])) - WE = np.zeros((dim_wave, n_theta, N)) - for i in range(N): - for theta in range(n_theta): - WE[:, theta, i] = wemd_embed(line_proj[:, theta, i]) - def fij(alpha, gamma, i, j, loss): if loss == "l1": # Ii_hat = Img_pft[:, :, i] @@ -160,6 +151,12 @@ def fij(alpha, gamma, i, j, loss): return np.linalg.norm(Si - Sj, 1) if loss == "wemd": + dim_wave = len(wemd_embed(line_proj[:, 0, 0])) + WE = np.zeros((dim_wave, n_theta, N), dtype=np.float64) + for i in range(N): + for theta in range(n_theta): + WE[:, theta, i] = wemd_embed(line_proj[:, theta, i]) + idxi = np.round((alpha - np.pi / 2) * n_theta / 2 / np.pi) % n_theta idxj = np.round((-gamma - np.pi / 2) * n_theta / 2 / np.pi) % n_theta @@ -167,11 +164,11 @@ def fij(alpha, gamma, i, j, loss): Sj = WE[:, int(idxj), j] return np.linalg.norm(Si - Sj, 1) - alpha_grid = np.arange(2 * T) * np.pi / T - beta_grid = (2 * np.arange(2 * T) + 1) * np.pi / 4 / T - gamma_grid = np.arange(2 * T) * np.pi / T + alpha_grid = np.arange(2 * T, dtype=np.float64) * np.pi / T + beta_grid = (2 * np.arange(2 * T, dtype=np.float64) + 1) * np.pi / 4 / T + gamma_grid = np.arange(2 * T, dtype=np.float64) * np.pi / T - bT = np.zeros(2 * T) + bT = np.zeros(2 * T, dtype=np.float64) for n in range(2 * T): ss = 0 for m in range(T): @@ -186,11 +183,11 @@ def fij(alpha, gamma, i, j, loss): def fijhat_k(k, F): dk = 2 * k + 1 - exp_alpha_grid = np.zeros((2 * T, dk), dtype=complex) + exp_alpha_grid = np.zeros((2 * T, dk), dtype=complex_type(np.float64)) for m in range(-k, k + 1): exp_alpha_grid[:, m + k] = np.exp(1j * m * alpha_grid) - exp_gamma_grid = np.zeros((2 * T, dk), dtype=complex) + exp_gamma_grid = np.zeros((2 * T, dk), dtype=complex_type(np.float64)) for m in range(-k, k + 1): exp_gamma_grid[:, m + k] = np.exp(1j * m * gamma_grid) @@ -201,10 +198,10 @@ def fijhat_k(k, F): C = [] for k in range(1, Lmax + 1): dk = 2 * k + 1 - C.append(np.zeros((N * dk, N * dk), dtype=complex)) + C.append(np.zeros((N * dk, N * dk), dtype=complex_type(np.float64))) for i in range(N): for j in range(i + 1, N): - Fij = np.zeros((2 * T, 2 * T)) + Fij = np.zeros((2 * T, 2 * T), dtype=np.float64) for j1 in range(2 * T): for j2 in range(2 * T): Fij[j1, j2] = fij(alpha_grid[j1], gamma_grid[j2], i, j, loss) @@ -217,7 +214,7 @@ def fijhat_k(k, F): C[k - 1] = C[k - 1] + C[k - 1].conj().T for i in range(N): - Fii = np.zeros((2 * T, 2 * T)) + Fii = np.zeros((2 * T, 2 * T), dtype=np.float64) for j1 in range(2 * T): for j2 in range(2 * T): Fii[j1, j2] = fij(alpha_grid[j1], gamma_grid[j2], i, i, loss) @@ -230,9 +227,10 @@ def fijhat_k(k, F): for k in range(1, Lmax + 1): [T, Tinv] = self.complex2real(k) C[k - 1] = np.real( - np.kron(np.eye(N), Tinv) @ C[k - 1] @ np.kron(np.eye(N), T) + np.kron(np.eye(N, dtype=np.float64), Tinv) + @ C[k - 1] + @ np.kron(np.eye(N, dtype=np.float64), T) ) - C[k - 1] = np.round(C[k - 1], 10) self.C = C @@ -270,11 +268,10 @@ def fast_radon_transform(array, angles, use_ramp=False): return projections, lines_f - @staticmethod - def complex2real(ell): + def complex2real(self, ell): # compute transformation matrices that convert complex representations to real ones diml = 2 * ell + 1 - Tinv = np.zeros((diml, diml), dtype=complex) + Tinv = np.zeros((diml, diml), dtype=complex_type(np.float64)) for i in range(diml): if i < ell: Tinv[i, i] = 1j / np.sqrt(2) @@ -285,7 +282,7 @@ def complex2real(ell): Tinv[i, i] = (-1) ** (i - ell) / np.sqrt(2) Tinv[i, diml - 1 - i] = 1 / np.sqrt(2) - T = np.zeros((diml, diml), dtype=complex) + T = np.zeros((diml, diml), dtype=complex_type(np.float64)) for i in range(diml): if i < ell: T[i, i] = -1j / np.sqrt(2) @@ -367,9 +364,9 @@ def admm_sym_J(self, C, verbose): for k in range(1, Lmax + 1): s0 = k**2 s1 = (k + 1) ** 2 - AEk = xp.zeros((1 + s0 + s1, 2 * (s0 + s1))) - AEk[0, :s0] = xp.eye(k).T.reshape(-1) - AEk[0, s0 : s0 + s1] = xp.eye(k + 1).T.reshape(-1) + AEk = xp.zeros((1 + s0 + s1, 2 * (s0 + s1)), dtype=np.float64) + AEk[0, :s0] = xp.eye(k, dtype=np.float64).T.reshape(-1) + AEk[0, s0 : s0 + s1] = xp.eye(k + 1, dtype=np.float64).T.reshape(-1) for count in range(1, 1 + s0): AEk[count, count - 1] = 1 AEk[count, count - 1 + s0 + s1] = 1 @@ -378,7 +375,7 @@ def admm_sym_J(self, C, verbose): AEk[count, count - 1 + s1 + s0] = 1 AE.append(AEk) AEAETinv.append(np.linalg.pinv(AEk @ AEk.T)) - bE = xp.zeros((Lmax + D0 + D1)) + bE = xp.zeros((Lmax + D0 + D1), dtype=np.float64) for k in range(Lmax): bE[k + d0[k] + d1[k] :] = rank_Ak[k] bE[k + 1 + d0[k] + d1[k] : k + 1 + d0[k + 1] + d1[k]] = xp.eye( @@ -391,7 +388,7 @@ def admm_sym_J(self, C, verbose): P = [] for k in range(1, Lmax + 1): dk = 2 * k + 1 - Pk = xp.eye(dk) + Pk = xp.eye(dk, dtype=np.float64) for m in range(k): for el in range(k - m): Pk[(m + 2 * el, m + 2 * el + 1), :] = Pk[ @@ -400,7 +397,7 @@ def admm_sym_J(self, C, verbose): P.append(Pk) def fun_AE(X0, X1, Xd0, Xd1, Xq): - z = xp.zeros((Lmax + D0 + D1, N)) + z = xp.zeros((Lmax + D0 + D1, N), dtype=np.float64) for k in range(Lmax): z[k + d0[k] + d1[k] : k + 1 + d0[k + 1] + d1[k + 1]] = AE[ k @@ -418,10 +415,10 @@ def fun_AE(X0, X1, Xd0, Xd1, Xq): ) def fun_AET(yE, yEq): - Z0 = xp.zeros((D0, N * (N + 1) // 2)) - Z1 = xp.zeros((D1, N * (N + 1) // 2)) - Zd0 = xp.zeros((D0, N)) - Zd1 = xp.zeros((D1, N)) + Z0 = xp.zeros((D0, N * (N + 1) // 2), dtype=np.float64) + Z1 = xp.zeros((D1, N * (N + 1) // 2), dtype=np.float64) + Zd0 = xp.zeros((D0, N), dtype=np.float64) + Zd1 = xp.zeros((D1, N), dtype=np.float64) for k in range(Lmax): s0 = (k + 1) ** 2 s1 = (k + 2) ** 2 @@ -436,7 +433,7 @@ def fun_AET(yE, yEq): return Z0, Z1, Zd0, Zd1, Zq[:16] def fun_AI(X0, X1): - z = xp.zeros((Ngrid, N * (N + 1) // 2)) + z = xp.zeros((Ngrid, N * (N + 1) // 2), dtype=np.float64) tmp = xp.concatenate((X0, X1), axis=0) z[:, idx_diag] = AI_mat_diag @ tmp[:, idx_diag] z[:, idx_offdiag] = AI_mat_offdiag @ tmp[:, idx_offdiag] @@ -446,7 +443,7 @@ def fun_AI(X0, X1): def fun_AIT(yI): # Z=AI_mat.T@yI # return Z[:D0,:], Z[D0:,:] - Z = xp.zeros((D0 + D1, N * (N + 1) // 2)) + Z = xp.zeros((D0 + D1, N * (N + 1) // 2), dtype=np.float64) Z[:, idx_diag] = AI_mat_diag.T @ yI[:, idx_diag] Z[:, idx_offdiag] = AI_mat_offdiag.T @ yI[:, idx_offdiag] return Z[:D0, :], Z[D0:, :] @@ -516,7 +513,7 @@ def update_yE(C0, C1, X0, X1, Xd0, Xd1, Xq, S0, S1, Sd0, Sd1, Sq, yI, rho): -Xd1 / rho - Sd1, -Xq / rho - Sq, ) - yE = xp.zeros((Lmax + D0 + D1, N)) + yE = xp.zeros((Lmax + D0 + D1, N), dtype=np.float64) for k in range(Lmax): yE[k + d0[k] + d1[k] : k + 1 + d0[k + 1] + d1[k + 1]] = AEAETinv[k] @ ( bE[k + d0[k] + d1[k] : k + 1 + d0[k + 1] + d1[k + 1]] / rho @@ -680,13 +677,13 @@ def print_updates(verbose): + ", |X|=%1.2f" % normX ) - Xd0 = xp.zeros((D0, N)) - Xd1 = xp.zeros((D1, N)) - Sd0 = xp.zeros(Xd0.shape) - Sd1 = xp.zeros(Xd1.shape) - yI = xp.zeros((Ngrid, N * (N + 1) // 2)) - yE = xp.zeros(bE.shape) - yEq = xp.zeros(bEq.shape) + Xd0 = xp.zeros((D0, N), dtype=np.float64) + Xd1 = xp.zeros((D1, N), dtype=np.float64) + Sd0 = xp.zeros(Xd0.shape, dtype=np.float64) + Sd1 = xp.zeros(Xd1.shape, dtype=np.float64) + yI = xp.zeros((Ngrid, N * (N + 1) // 2), dtype=np.float64) + yE = xp.zeros(bE.shape, dtype=np.float64) + yEq = xp.zeros(bEq.shape, dtype=np.float64) IDX = np.arange(3) Time = np.zeros(4) @@ -774,11 +771,13 @@ def ADMM_preprocessing(self, C): D1 = d1[-1] # AE and bE for quaternion constraints - AEq = xp.asarray(loadmat("data/Eq_constraints/AEqJ.mat")["AEq"]) - AEqAEqtinv = xp.asarray(loadmat("data/Eq_constraints/AEqJ.mat")["AEqAEqtinv"]) + # AEq = xp.asarray(loadmat("data/Eq_constraints/AEqJ.mat")["AEq"]) + # AEqAEqtinv = xp.asarray(loadmat("data/Eq_constraints/AEqJ.mat")["AEqAEqtinv"]) + AEq = xp.asarray(self.construct_AEq()) + AEqAEqtinv = xp.linalg.pinv(AEq @ AEq.T) - bEq = xp.zeros(17) - bEq[:16] = xp.eye(4).reshape(-1) / 4 + bEq = xp.zeros(17, dtype=np.float64) + bEq[:16] = xp.eye(4, dtype=np.float64).reshape(-1) / 4 bEq[-1] = 1 bEq = xp.repeat(bEq[:, xp.newaxis], N * (N - 1) // 2, axis=1) @@ -794,10 +793,10 @@ def ADMM_preprocessing(self, C): # AI_mat[p,:d0[-1]]=w0; AI_mat[p,d0[-1]:]=w1 # AI_mat=xp.asarray(AI_mat) / 10; # bI=-(Lmax+2)*(Lmax+1)/2 / 10 - AI_mat_offdiag = np.zeros((Ngrid, D0 + D1)) + AI_mat_offdiag = np.zeros((Ngrid, D0 + D1), dtype=np.float64) for p in range(Ngrid): - w0 = np.zeros(D0) - w1 = np.zeros(D1) + w0 = np.zeros(D0, dtype=np.float64) + w1 = np.zeros(D1, dtype=np.float64) for k in range(1, Lmax + 1): w0[d0[k - 1] : d0[k]] = ( (Lmax - k + 2) @@ -826,10 +825,10 @@ def ADMM_preprocessing(self, C): # AI_mat_offdiag_new[:, d0[k - 1]:d0[k]] = block0 # AI_mat_offdiag_new[:, d0[-1] + d1[k - 1]:d0[-1] + d1[k]] = block1 - AI_mat_diag = np.zeros((Ngrid, D0 + D1)) + AI_mat_diag = np.zeros((Ngrid, D0 + D1), dtype=np.float64) for p in range(Ngrid): - w0 = np.zeros(D0) - w1 = np.zeros(D1) + w0 = np.zeros(D0, dtype=np.float64) + w1 = np.zeros(D1, dtype=np.float64) for k in range(1, Lmax + 1): w0[d0[k - 1] : d0[k]] = ( (Lmax - k + 2) @@ -857,15 +856,15 @@ def ADMM_preprocessing(self, C): II = [] for k in range(1, Lmax + 1): dk = 2 * k + 1 - II.append(xp.eye(N * dk)) + II.append(xp.eye(N * dk, dtype=np.float64)) I0, I1 = self.transform_coeff(II, Lmax, N, IDX_upper) - X0 = xp.zeros((D0, N * (N + 1) // 2)) - X1 = xp.zeros((D1, N * (N + 1) // 2)) - Xq = xp.zeros((16, N * (N - 1) // 2)) + X0 = xp.zeros((D0, N * (N + 1) // 2), dtype=np.float64) + X1 = xp.zeros((D1, N * (N + 1) // 2), dtype=np.float64) + Xq = xp.zeros((16, N * (N - 1) // 2), dtype=np.float64) # X0,X1=transform_coeff(II,Lmax,N); Xq=xp.zeros((16,N*(N-1))); Xq[0,:]=1; S0 = xp.copy(I0) S1 = xp.copy(I1) - Sq = xp.zeros(Xq.shape) + Sq = xp.zeros(Xq.shape, dtype=np.float64) # S0=xp.zeros(X0.shape); S1=xp.zeros(X1.shape); Sq=xp.zeros(Xq.shape) self.Ngrid = Ngrid @@ -911,7 +910,7 @@ def compute_fejer_weights(self): def permutek_block(Ak, k): dk = 2 * k + 1 - Pk = np.eye(dk) + Pk = np.eye(dk, dtype=np.float64) for m in range(k): for ell in range(k - m): Pk[(m + 2 * ell, m + 2 * ell + 1), :] = Pk[ @@ -923,8 +922,8 @@ def permutek_block(Ak, k): W0 = [] W1 = [] for k in range(start, self.Lmax + 1): - W0k = np.zeros((Ngrid, k, k)) - W1k = np.zeros((Ngrid, k + 1, k + 1)) + W0k = np.zeros((Ngrid, k, k), dtype=np.float64) + W1k = np.zeros((Ngrid, k + 1, k + 1), dtype=np.float64) TkT = TT[k - start].T TinvkT = TTI[k - start].T @@ -951,7 +950,7 @@ def discretize_SO3(self): gamma = gamma + np.pi # SO(3) in Euler ZYZ - SO3 = np.zeros((S2_size * S1_size, 3)) + SO3 = np.zeros((S2_size * S1_size, 3), dtype=np.float64) count = 0 for i in range(S1_size): for j in range(S2_size): @@ -965,14 +964,7 @@ def discretize_SO3(self): ############################ def proximal_refine(self, X_admm, weight, Penalty, r): - Lmax = self.Lmax N = self.n_img - max_iter = self.max_iter - rho = self.rho - ratio = self.ratio - factor = self.factor - mult = self.mult - Nstep_yI = self.Nstep_yI C = self.C def Ak(J, Euler): @@ -989,7 +981,7 @@ def rel_change(A, B, eps=1e-12): den += np.linalg.norm(B[k]) ** 2 return np.sqrt(num) / max(np.sqrt(den), eps) - rank_Ak = np.zeros(self.Lmax) + rank_Ak = np.zeros(self.Lmax, dtype=np.float64) C_base = [None] * self.Lmax for k in range(self.Lmax): C_base[k] = xp.asnumpy(C[k]).copy() @@ -1052,18 +1044,26 @@ def euler_est(self): R_est, Euler_est = self.euler_est_Dm(X_est) self.Euler_est = Euler_est - self.rotations = R_est + self.rotations = R_est.astype(self.dtype) def euler_est_Cm(self, X1, XS): S = self.n_sym N = self.n_img - sym_euler = np.zeros((S, 3)) + sym_euler = np.zeros((S, 3), dtype=np.float64) for s in range(S): sym_euler[s] = [2 * np.pi * s / S, 0, 0] [T, Tinv] = self.complex2real(1) - X1 = np.kron(np.eye(N), T) @ X1 @ np.kron(np.eye(N), Tinv) + X1 = ( + np.kron(np.eye(N, dtype=np.float64), T) + @ X1 + @ np.kron(np.eye(N, dtype=np.float64), Tinv) + ) [T, Tinv] = self.complex2real(S) - XS = np.kron(np.eye(N), T) @ XS @ np.kron(np.eye(N), Tinv) + XS = ( + np.kron(np.eye(N, dtype=np.float64), T) + @ XS + @ np.kron(np.eye(N, dtype=np.float64), Tinv) + ) def find_phase(A, B): # find a number c that minimizes ||cA-B||_F @@ -1079,8 +1079,8 @@ def find_phase(A, B): return c / abs(c) def find_beta(X1): - B1 = np.zeros((N, N)) - B2 = np.zeros((N, N)) + B1 = np.zeros((N, N), dtype=np.float64) + B2 = np.zeros((N, N), dtype=np.float64) for i in range(N): for j in range(N): Xij = X1[3 * i : 3 * (i + 1), 3 * j : 3 * (j + 1)] @@ -1096,8 +1096,8 @@ def find_beta(X1): return beta def find_alpha(X1): - ZZbar = np.zeros((N, N), dtype=complex) - ZZ = np.zeros((N, N), dtype=complex) + ZZbar = np.zeros((N, N), dtype=complex_type(np.float64)) + ZZ = np.zeros((N, N), dtype=complex_type(np.float64)) for i in range(N): for j in range(N): z = X1[3 * i : 3 * (i + 1), 3 * j : 3 * (j + 1)][0, 0] @@ -1112,12 +1112,12 @@ def find_alpha(X1): c = find_phase(Z[:, None] @ Z[:, None].T, ZZ) Z = np.sqrt(c) * Z - return np.angle(Z) + return np.angle(Z).astype(np.float64) dk = 2 * S + 1 def find_gamma(Xm, beta, alpha): - C = np.zeros((N, N), dtype=complex) + C = np.zeros((N, N), dtype=complex_type(np.float64)) Jk = np.ones(dk) Jk[S + 1 :: 2] = -1 Jk[S - 1 :: -2] = -1 @@ -1150,13 +1150,13 @@ def find_gamma(Xm, beta, alpha): C[i, j] = np.vdot(C1 + C2, np.real(C3)) / np.vdot( C1 + C2, C1 + C2 ) + 1j * np.vdot(C1 - C2, np.imag(C3)) / np.vdot(C1 - C2, C1 - C2) - C += C.T.conj() + np.eye(N) + C += C.T.conj() + np.eye(N, dtype=np.float64) evals, evecs = np.linalg.eigh(C) idx = np.argmax(evals) c = evecs[:, idx] * np.sqrt(evals[idx]) return np.angle(c) / S - Euler_est = np.zeros((N, 3)) + Euler_est = np.zeros((N, 3), dtype=np.float64) Euler_est[:, 0] = find_alpha(X1) Euler_est[:, 1] = find_beta(X1) Euler_est[:, 2] = find_gamma(XS, Euler_est[:, 1], Euler_est[:, 0]) @@ -1186,9 +1186,13 @@ def find_phase(A, B): return c / abs(c) T, Tinv = self.complex2real(2) - X2 = np.kron(np.eye(N), T) @ X2 @ np.kron(np.eye(N), Tinv) + X2 = ( + np.kron(np.eye(N, dtype=np.float64), T) + @ X2 + @ np.kron(np.eye(N, dtype=np.float64), Tinv) + ) - B1 = np.zeros((N, N)) + B1 = np.zeros((N, N), dtype=np.float64) for i in range(N): for j in range(N): Xij = X2[5 * i : 5 * (i + 1), 5 * j : 5 * (j + 1)] @@ -1198,8 +1202,8 @@ def find_phase(A, B): b1 = v1[:, idx] * np.sqrt(e1[idx]) * np.sign(v1[0, idx]) beta_est = np.arccos(np.clip(np.sqrt(b1), -1, 1)) % np.pi - Aminus = np.zeros((N, N), dtype=complex) - Aplus = np.zeros((N, N), dtype=complex) + Aminus = np.zeros((N, N), dtype=complex_type(np.float64)) + Aplus = np.zeros((N, N), dtype=complex_type(np.float64)) for i in range(N): for j in range(N): Xij = X2[5 * i : 5 * (i + 1), 5 * j : 5 * (j + 1)] @@ -1253,8 +1257,12 @@ def LS_D(W1, W2, W3, W4, Br, Bi): return a + 1j * b [T, Tinv] = self.complex2real(S) - Xm = np.kron(np.eye(N), T) @ Xm @ np.kron(np.eye(N), Tinv) - C = np.zeros((N, N), dtype=complex) + Xm = ( + np.kron(np.eye(N, dtype=np.float64), T) + @ Xm + @ np.kron(np.eye(N, dtype=np.float64), Tinv) + ) + C = np.zeros((N, N), dtype=complex_type(np.float64)) Jk = np.ones(dk) Jk[S + 1 :: 2] = -1 Jk[S - 1 :: -2] = -1 @@ -1290,7 +1298,7 @@ def LS_D(W1, W2, W3, W4, Br, Bi): Br = np.real(4 * DXijmD) Bi = np.imag(4 * DXijmD) C[i, j] = LS_D(W1, W2, W3, W4, Br, Bi) - C += C.T.conj() + np.eye(N) + C += C.T.conj() + np.eye(N, dtype=np.float64) evals, evecs = np.linalg.eigh(C) idx = np.argmax(evals) c = evecs[:, idx] * np.sqrt(evals[idx]) @@ -1298,7 +1306,7 @@ def LS_D(W1, W2, W3, W4, Br, Bi): alpha_est, beta_est = find_alpha_beta(X2) gamma_est = find_gamma(XS, alpha_est, beta_est) - Euler_est = np.zeros((N, 3)) + Euler_est = np.zeros((N, 3), dtype=np.float64) Euler_est[:, 0] = alpha_est Euler_est[:, 1] = beta_est Euler_est[:, 2] = gamma_est @@ -1315,8 +1323,8 @@ def transform_coeff(self, A, Lmax, N, IDX_upper): for k in range(1, Lmax + 1): d0.append(d0[-1] + k**2) d1.append(d1[-1] + (k + 1) ** 2) - A0 = xp.zeros((d0[-1], N * (N + 1) // 2)) - A1 = xp.zeros((d1[-1], N * (N + 1) // 2)) + A0 = xp.zeros((d0[-1], N * (N + 1) // 2), dtype=np.float64) + A1 = xp.zeros((d1[-1], N * (N + 1) // 2), dtype=np.float64) for k in range(1, Lmax + 1): a0, a1 = self.permutek(A[k - 1], k, N) A0[d0[k - 1] : d0[k], :] = self.vec_block(a0, N, k, IDX_upper) @@ -1332,7 +1340,7 @@ def transform_coeff_back(self, A0, A1, Lmax, N, IDX_upper, IDX_lower, idx_offdia A = [] for k in range(1, Lmax + 1): dk = 2 * k + 1 - Ak = xp.zeros((N * dk, N * dk)) + Ak = xp.zeros((N * dk, N * dk), dtype=np.float64) Ak[: N * k, : N * k] = self.mat_block( A0[d0[k - 1] : d0[k], :], N, k, IDX_upper, IDX_lower, idx_offdiag ) @@ -1347,13 +1355,17 @@ def transform_coeff_back(self, A0, A1, Lmax, N, IDX_upper, IDX_lower, idx_offdia def permutek(Ak, k, N): AkP = xp.copy(Ak) dk = 2 * k + 1 - Pk = xp.eye(dk) + Pk = xp.eye(dk, dtype=AkP.dtype) for m in range(k): for n in range(k - m): Pk[(m + 2 * n, m + 2 * n + 1), :] = Pk[(m + 2 * n + 1, m + 2 * n), :] - AkP = xp.kron(xp.eye(N), Pk) @ Ak @ xp.kron(xp.eye(N), Pk.T) + AkP = ( + xp.kron(xp.eye(N, dtype=AkP.dtype), Pk) + @ Ak + @ xp.kron(xp.eye(N, dtype=AkP.dtype), Pk.T) + ) - Pk = xp.eye(N * dk) + Pk = xp.eye(N * dk, dtype=AkP.dtype) idx = xp.concatenate((xp.arange(dk - k, dk), xp.arange(k + 1))) for m in range(N - 1): for n in range(N - 1 - m): @@ -1366,7 +1378,7 @@ def permutek(Ak, k, N): @staticmethod def permutek_back(Ak, k, N): dk = 2 * k + 1 - Pk = xp.eye(N * dk) + Pk = xp.eye(N * dk, dtype=Ak.dtype) idx = xp.concatenate((xp.arange(dk - k, dk), xp.arange(k + 1))) for m in range(N - 1): for n in range(N - 1 - m): @@ -1375,11 +1387,15 @@ def permutek_back(Ak, k, N): ][idx, :] AkB = Pk.T @ Ak @ Pk dk = 2 * k + 1 - Pk = xp.eye(dk) + Pk = xp.eye(dk, dtype=Ak.dtype) for m in range(k): for n in range(k - m): Pk[(m + 2 * n, m + 2 * n + 1), :] = Pk[(m + 2 * n + 1, m + 2 * n), :] - AkB = xp.kron(xp.eye(N), Pk.T) @ AkB @ xp.kron(xp.eye(N), Pk) + AkB = ( + xp.kron(xp.eye(N, dtype=Ak.dtype), Pk.T) + @ AkB + @ xp.kron(xp.eye(N, dtype=Ak.dtype), Pk) + ) return AkB @staticmethod @@ -1404,7 +1420,7 @@ def largest_eigenvalue(AI, Ngrid, N): return Lambda def compute_rank(self, Lmax): - rk = xp.zeros(Lmax) + rk = xp.zeros(Lmax, dtype=np.float64) A = [] for k in range(1, Lmax + 1): Ak = np.sum(self.WD(k, self.sym_euler), axis=0) @@ -1430,7 +1446,7 @@ def WD(self, J, euler): @staticmethod def Wd(J, beta): # compute Wigner small d matrix - d = np.zeros((len(beta), 2 * J + 1, 2 * J + 1)) + d = np.zeros((len(beta), 2 * J + 1, 2 * J + 1), dtype=beta.dtype) for m in range(-J, J + 1): for n in range(-J, J + 1): smin = max(0, m - n) @@ -1457,7 +1473,7 @@ def Wd(J, beta): @staticmethod def mat_block(vecA, N, sz, IDX_upper, IDX_lower, idx_offdiag): tmp = vecA.T.reshape(N * (N + 1) // 2, sz, sz).transpose(0, 2, 1) - AA = xp.zeros((N**2, sz, sz)) + AA = xp.zeros((N**2, sz, sz), dtype=vecA.dtype) AA[IDX_upper] = tmp AA[IDX_lower] = tmp[idx_offdiag].transpose(0, 2, 1) return (AA.reshape(N, N, sz, sz).transpose(0, 2, 1, 3)).reshape(N * sz, N * sz) @@ -1478,7 +1494,7 @@ def transform_block(A, k, Pk=None): if Pk is None: dk = 2 * k + 1 - Pk = xp.eye(dk) + Pk = xp.eye(dk, dtype=A.dtype) for m in range(k): for el in range(k - m): Pk[(m + 2 * el, m + 2 * el + 1), :] = Pk[ @@ -1502,11 +1518,11 @@ def transform_back_block(A0, A1, k, Pk=None): A0 = A0[None, :] A1 = A1[None, :] - A = xp.zeros((A0.shape[0], dk, dk)) + A = xp.zeros((A0.shape[0], dk, dk), dtype=A0.dtype) A[:, :k, :k] = A0.reshape(-1, k, k).swapaxes(-1, -2) A[:, k:, k:] = A1.reshape(-1, k + 1, k + 1).swapaxes(-1, -2) if Pk is None: - Pk = xp.eye(dk) + Pk = xp.eye(dk, dtype=A0.dtype) for m in range(k): for el in range(k - m): Pk[(m + 2 * el, m + 2 * el + 1), :] = Pk[ @@ -1514,3 +1530,42 @@ def transform_back_block(A0, A1, k, Pk=None): ] out = Pk.T @ A @ Pk return out[0] if single else out + + def construct_AEq(self): + """ + Construct the linear equality matrix for quaternion constraints. + """ + AEq = np.zeros((17, 21), np.float64) + + # First 16 rows: identity constraints on first 16 variables + AEq[:16, :16] = np.eye(16, dtype=np.float64) + + # Extra columns 16:21 + extra = 0.25 * np.array( + [ + [-1, 1, 0, 0, 1], + [0, 0, 0, 0, 0], + [0, 0, 1, -1, 0], + [0, 0, 0, 0, 0], + [0, 0, 0, 0, 0], + [1, 1, 0, 0, -1], + [0, 0, 0, 0, 0], + [0, 0, 1, 1, 0], + [0, 0, 1, -1, 0], + [0, 0, 0, 0, 0], + [-1, -1, 0, 0, -1], + [0, 0, 0, 0, 0], + [0, 0, 0, 0, 0], + [0, 0, 1, 1, 0], + [0, 0, 0, 0, 0], + [1, -1, 0, 0, 1], + ], + dtype=np.float64, + ) + + AEq[:16, 16:] = extra + + # Last row: redundant trace/sum constraint + AEq[16, [0, 5, 10, 15]] = 1 + + return AEq diff --git a/tests/test_nug.py b/tests/test_nug.py index a221f36bd7..1eb65a71d6 100644 --- a/tests/test_nug.py +++ b/tests/test_nug.py @@ -3,12 +3,13 @@ from aspire.abinitio import CommonlineNUG from aspire.source import Simulation +from aspire.volume import CnSymmetricVolume, SymmetryGroup -DTYPE = [np.float32] -RESOLUTION = [48, 49] +DTYPE = [np.float32, pytest.param(np.float64, marks=pytest.mark.expensive)] +RESOLUTION = [48, pytest.param(48, marks=pytest.mark.expensive)] N_IMG = [5] OFFSETS = [0] -ORDER = [3, 4] +ORDER = [3, pytest.param(4, marks=pytest.mark.expensive)] PR = [False] SEED = 1980 @@ -68,12 +69,12 @@ def source(n_img, resolution, dtype, offsets, order): @pytest.fixture(scope="module") -def orient_est(src, proximal_refine): +def orient_est(source, proximal_refine): orient_est = CommonlineNUG( - src, + source, perform_pr=proximal_refine, ) - + orient_est.estimate_rotations() return orient_est @@ -86,4 +87,55 @@ def test_dtypes(orient_est): """ Check dtypes for each major step of the algorithm. """ - pass + assert orient_est.dtype == orient_est.src.dtype + + # Intermediate steps use doubles + for Ci in orient_est.C: + assert Ci.dtype == np.float64 + + for Xi in orient_est.X_est: + assert Xi.dtype == np.float64 + + assert orient_est.rotations.dtype == orient_est.dtype + + +def test_estimate_rotations_pairwise(orient_est): + """ """ + MSE = compare_rots_sym( + orient_est.rotations, orient_est.src.rotations, orient_est.sym_grp + ) + np.testing.assert_array_less(MSE, 0.3) + + +########### +# Helpers # +########### + + +def compare_rots_sym(R_est, R_true, sym): + N = R_true.shape[0] + sym_euler = SymmetryGroup.parse(sym).matrices + order = sym_euler.shape[0] + J = np.array([[-1, 0, 0], [0, -1, 0], [0, 0, 1]]) + error = np.zeros((N, N)) + errorJ = np.zeros((N, N)) + for i in range(N): + for j in range(N): + e = np.zeros(order) + eJ = np.zeros(order) + for s in range(order): + Rs = sym_euler[s] + e[s] = ( + np.linalg.norm(R_est[i].T @ R_est[j] - R_true[i].T @ Rs @ R_true[j]) + ** 2 + ) + eJ[s] = ( + np.linalg.norm( + R_est[i].T @ R_est[j] - J @ R_true[i].T @ Rs @ R_true[j] @ J + ) + ** 2 + ) + error[i, j] = e.min() + errorJ[i, j] = eJ.min() + E = min(error.sum(), errorJ.sum()) / N**2 + return E From 7e73f7f5ae475b437ae9eb2e4e08f812369b7343 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 22 May 2026 14:16:54 -0400 Subject: [PATCH 024/161] Support shifts --- src/aspire/abinitio/commonline_nug.py | 22 +++++++++++++++++----- tests/test_nug.py | 1 + 2 files changed, 18 insertions(+), 5 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 944394c5a4..0bb1ebf8bd 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -11,7 +11,7 @@ from aspire.utils import Rotation, cart2sph, complex_type from aspire.volume import CnSymmetryGroup, DnSymmetryGroup, SymmetryGroup -from .commonline_utils import saff_kuijlaars +from .commonline_utils import _generate_shift_phase_and_filter, saff_kuijlaars logger = logging.getLogger(__name__) @@ -27,7 +27,7 @@ def __init__( symmetry=None, n_rad=None, n_theta=360, - max_shift=0, + max_shift=0.15, shift_step=1, mask=True, Lmax=12, @@ -98,6 +98,16 @@ def _build_full_pft(self): pf = self.pf self.pf_full = PolarFT.half_to_full(pf) + # Prepare the shift phases to try and generate filter for common-line detection + r_max = self.pf_full.shape[2] + self.shifts, self.shift_phases, h = _generate_shift_phase_and_filter( + r_max, self.max_shift, self.shift_step, self.dtype + ) + + # Apply bandpass filter, normalize each ray of each image + # Note that only use half of each ray + # self.pf_full = self._apply_filter_and_norm("ijk, k -> ijk", pf_full, r_max, h) + def estimate_rotations(self): self.compute_coeff() self.perform_admm() @@ -138,7 +148,6 @@ def fij(alpha, gamma, i, j, loss): # Si = Ii_hat[:, int(idxi)] # Sj = Ij_hat[:, int(idxj)] - # Using aspire PolarFT. Replace later Ii_hat = self.pf_full[i] Ij_hat = self.pf_full[j] idxi = np.round((alpha - np.pi / 2) * n_theta / 2 / np.pi) % n_theta @@ -146,9 +155,11 @@ def fij(alpha, gamma, i, j, loss): Si = Ii_hat[int(idxi)] Sj = Ij_hat[int(idxj)] - # norm_new = np.linalg.norm(Si - Sj, 1) - return np.linalg.norm(Si - Sj, 1) + # Apply shifts + Sj_shifted = self.shift_phases * Sj + norms = np.linalg.norm(Si[None] - Sj_shifted, 1, axis=1) + return norms.min() if loss == "wemd": dim_wave = len(wemd_embed(line_proj[:, 0, 0])) @@ -162,6 +173,7 @@ def fij(alpha, gamma, i, j, loss): Si = WE[:, int(idxi), i] Sj = WE[:, int(idxj), j] + return np.linalg.norm(Si - Sj, 1) alpha_grid = np.arange(2 * T, dtype=np.float64) * np.pi / T diff --git a/tests/test_nug.py b/tests/test_nug.py index 1eb65a71d6..6f1b96654f 100644 --- a/tests/test_nug.py +++ b/tests/test_nug.py @@ -72,6 +72,7 @@ def source(n_img, resolution, dtype, offsets, order): def orient_est(source, proximal_refine): orient_est = CommonlineNUG( source, + max_shift=0, perform_pr=proximal_refine, ) orient_est.estimate_rotations() From 48dfa14cc7c82f910f3300b7a5cde359a7d5739b Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 26 May 2026 11:35:26 -0400 Subject: [PATCH 025/161] test shifts --- tests/test_nug.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/tests/test_nug.py b/tests/test_nug.py index 6f1b96654f..3b6ebb264b 100644 --- a/tests/test_nug.py +++ b/tests/test_nug.py @@ -6,9 +6,9 @@ from aspire.volume import CnSymmetricVolume, SymmetryGroup DTYPE = [np.float32, pytest.param(np.float64, marks=pytest.mark.expensive)] -RESOLUTION = [48, pytest.param(48, marks=pytest.mark.expensive)] -N_IMG = [5] -OFFSETS = [0] +RESOLUTION = [48, pytest.param(49, marks=pytest.mark.expensive)] +N_IMG = [15] +OFFSETS = [0, pytest.param(None, marks=pytest.mark.expensive)] ORDER = [3, pytest.param(4, marks=pytest.mark.expensive)] PR = [False] SEED = 1980 @@ -70,9 +70,15 @@ def source(n_img, resolution, dtype, offsets, order): @pytest.fixture(scope="module") def orient_est(source, proximal_refine): + max_shift = 0 + shift_step = 1 + if source.offsets.all() != 0: + max_shift = 0.20 + shift_step = 0.25 orient_est = CommonlineNUG( source, - max_shift=0, + max_shift=max_shift, + shift_step=shift_step, perform_pr=proximal_refine, ) orient_est.estimate_rotations() @@ -105,7 +111,7 @@ def test_estimate_rotations_pairwise(orient_est): MSE = compare_rots_sym( orient_est.rotations, orient_est.src.rotations, orient_est.sym_grp ) - np.testing.assert_array_less(MSE, 0.3) + np.testing.assert_array_less(MSE, 0.1) ########### From 063350a27a90c8ac63d595ff18c346a287b8bd59 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 26 May 2026 13:17:04 -0400 Subject: [PATCH 026/161] Test Dn symmetry --- tests/test_nug.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/tests/test_nug.py b/tests/test_nug.py index 3b6ebb264b..5d5be97ffb 100644 --- a/tests/test_nug.py +++ b/tests/test_nug.py @@ -3,7 +3,7 @@ from aspire.abinitio import CommonlineNUG from aspire.source import Simulation -from aspire.volume import CnSymmetricVolume, SymmetryGroup +from aspire.volume import CnSymmetricVolume, DnSymmetricVolume, SymmetryGroup DTYPE = [np.float32, pytest.param(np.float64, marks=pytest.mark.expensive)] RESOLUTION = [48, pytest.param(49, marks=pytest.mark.expensive)] @@ -12,6 +12,10 @@ ORDER = [3, pytest.param(4, marks=pytest.mark.expensive)] PR = [False] SEED = 1980 +VOLUME = [ + CnSymmetricVolume, + pytest.param(DnSymmetricVolume, marks=pytest.mark.expensive), +] @pytest.fixture(params=DTYPE, ids=lambda x: f"dtype={x}", scope="module") @@ -44,14 +48,19 @@ def proximal_refine(request): return request.param +@pytest.fixture(params=VOLUME, ids=lambda x: f"Volume={x}", scope="module") +def Volume(request): + return request.param + + ############ # Fixtures # ############ @pytest.fixture(scope="module") -def source(n_img, resolution, dtype, offsets, order): - vol = CnSymmetricVolume( +def source(n_img, resolution, dtype, offsets, order, Volume): + vol = Volume( L=resolution, order=order, C=1, K=100, dtype=dtype, seed=SEED ).generate() From ecdef9a79a9d9f189e0bdf22bf64e838a0e38cbf Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 29 May 2026 10:43:18 -0400 Subject: [PATCH 027/161] move compare_rots_sym to utils --- src/aspire/abinitio/__init__.py | 1 + src/aspire/abinitio/commonline_utils.py | 30 +++++++++++++++++++++ tests/test_nug.py | 36 +------------------------ 3 files changed, 32 insertions(+), 35 deletions(-) diff --git a/src/aspire/abinitio/__init__.py b/src/aspire/abinitio/__init__.py index cb68c8f544..42f40f5955 100644 --- a/src/aspire/abinitio/__init__.py +++ b/src/aspire/abinitio/__init__.py @@ -2,6 +2,7 @@ from .J_sync import JSync from .commonline_utils import ( build_outer_products, + compare_rots_sym, g_sync, saff_kuijlaars, ) diff --git a/src/aspire/abinitio/commonline_utils.py b/src/aspire/abinitio/commonline_utils.py index 5fe88a21b1..bdf955bd4d 100644 --- a/src/aspire/abinitio/commonline_utils.py +++ b/src/aspire/abinitio/commonline_utils.py @@ -5,6 +5,7 @@ from aspire.operators import PolarFT from aspire.utils import J_conjugate, Rotation, all_pairs, anorm, cyclic_rotations, tqdm +from aspire.volume import SymmetryGroup logger = logging.getLogger(__name__) @@ -418,3 +419,32 @@ def build_outer_products(n, dtype): viis[i] = np.outer(gt_vis[i], gt_vis[i]) return vijs, viis, gt_vis + + +def compare_rots_sym(R_est, R_true, sym): + N = R_true.shape[0] + sym_euler = SymmetryGroup.parse(sym).matrices + order = sym_euler.shape[0] + J = np.array([[-1, 0, 0], [0, -1, 0], [0, 0, 1]]) + error = np.zeros((N, N)) + errorJ = np.zeros((N, N)) + for i in range(N): + for j in range(N): + e = np.zeros(order) + eJ = np.zeros(order) + for s in range(order): + Rs = sym_euler[s] + e[s] = ( + np.linalg.norm(R_est[i].T @ R_est[j] - R_true[i].T @ Rs @ R_true[j]) + ** 2 + ) + eJ[s] = ( + np.linalg.norm( + R_est[i].T @ R_est[j] - J @ R_true[i].T @ Rs @ R_true[j] @ J + ) + ** 2 + ) + error[i, j] = e.min() + errorJ[i, j] = eJ.min() + E = min(error.sum(), errorJ.sum()) / N**2 + return E diff --git a/tests/test_nug.py b/tests/test_nug.py index 5d5be97ffb..62c86834a0 100644 --- a/tests/test_nug.py +++ b/tests/test_nug.py @@ -1,7 +1,7 @@ import numpy as np import pytest -from aspire.abinitio import CommonlineNUG +from aspire.abinitio import CommonlineNUG, compare_rots_sym from aspire.source import Simulation from aspire.volume import CnSymmetricVolume, DnSymmetricVolume, SymmetryGroup @@ -121,37 +121,3 @@ def test_estimate_rotations_pairwise(orient_est): orient_est.rotations, orient_est.src.rotations, orient_est.sym_grp ) np.testing.assert_array_less(MSE, 0.1) - - -########### -# Helpers # -########### - - -def compare_rots_sym(R_est, R_true, sym): - N = R_true.shape[0] - sym_euler = SymmetryGroup.parse(sym).matrices - order = sym_euler.shape[0] - J = np.array([[-1, 0, 0], [0, -1, 0], [0, 0, 1]]) - error = np.zeros((N, N)) - errorJ = np.zeros((N, N)) - for i in range(N): - for j in range(N): - e = np.zeros(order) - eJ = np.zeros(order) - for s in range(order): - Rs = sym_euler[s] - e[s] = ( - np.linalg.norm(R_est[i].T @ R_est[j] - R_true[i].T @ Rs @ R_true[j]) - ** 2 - ) - eJ[s] = ( - np.linalg.norm( - R_est[i].T @ R_est[j] - J @ R_true[i].T @ Rs @ R_true[j] @ J - ) - ** 2 - ) - error[i, j] = e.min() - errorJ[i, j] = eJ.min() - E = min(error.sum(), errorJ.sum()) / N**2 - return E From d2fe08efa08c8893e8450a19fbaa9c6b92676975 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 4 Jun 2026 14:34:51 -0400 Subject: [PATCH 028/161] Add generalized g_sync --- src/aspire/abinitio/__init__.py | 1 + src/aspire/abinitio/commonline_utils.py | 181 +++++++++++++++++++++++- tests/test_nug.py | 10 +- tests/test_orient_symmetric.py | 2 +- 4 files changed, 185 insertions(+), 9 deletions(-) diff --git a/src/aspire/abinitio/__init__.py b/src/aspire/abinitio/__init__.py index 42f40f5955..c5b2a23d0e 100644 --- a/src/aspire/abinitio/__init__.py +++ b/src/aspire/abinitio/__init__.py @@ -4,6 +4,7 @@ build_outer_products, compare_rots_sym, g_sync, + g_sync_finite_group, saff_kuijlaars, ) from .commonline_base import Orient3D diff --git a/src/aspire/abinitio/commonline_utils.py b/src/aspire/abinitio/commonline_utils.py index bdf955bd4d..1d73b8fa16 100644 --- a/src/aspire/abinitio/commonline_utils.py +++ b/src/aspire/abinitio/commonline_utils.py @@ -5,7 +5,7 @@ from aspire.operators import PolarFT from aspire.utils import J_conjugate, Rotation, all_pairs, anorm, cyclic_rotations, tqdm -from aspire.volume import SymmetryGroup +from aspire.volume import CnSymmetryGroup, SymmetryGroup logger = logging.getLogger(__name__) @@ -320,10 +320,31 @@ def saff_kuijlaars(N): return mesh -def g_sync(rots, order, rots_gt): +def g_sync(rots, rots_gt, symmetry): """ Given ground truth rotations, synchronize estimated rotations over - symmetry group elements. + symmetry group elements. This method dispatches to either the faster + cyclic implementation for Cn symmetry or the generalized version for + all other symmetries. + + :param rots: Estimated rotation matrices + :param rots_gt: Ground truth rotation matrices + :param symmetry: The symmetry of the underlying molecule. + + :return: g-synchronized ground truth rotations. + """ + sym_grp = SymmetryGroup.parse(symmetry) + + if isinstance(sym_grp, CnSymmetryGroup): + return g_sync_cyclic(rots, rots_gt, symmetry) + # else + return g_sync_finite_group(rots, rots_gt, symmetry) + + +def g_sync_cyclic(rots, rots_gt, symmetry): + """ + Given ground truth rotations, synchronize estimated rotations over + cyclic symmetry group elements. Every estimated rotation might be a version of the ground truth rotation rotated by g^{s_i}, where s_i = 0, 1, ..., order. This method synchronizes the @@ -331,8 +352,8 @@ def g_sync(rots, order, rots_gt): to all estimates for error analysis. :param rots: Estimated rotation matrices - :param order: The cyclic order asssociated with the symmetry of the underlying molecule. - :param rots_gt: Ground truth rotation matrices. + :param rots_gt: Ground truth rotation matrices + :param symmetry: The symmetry of the underlying molecule. :return: g-synchronized ground truth rotations. """ @@ -342,8 +363,8 @@ def g_sync(rots, order, rots_gt): n_img = len(rots) dtype = rots.dtype - rots_symm = cyclic_rotations(order, dtype).matrices - + rots_symm = SymmetryGroup.parse(symmetry).matrices + order = len(rots_symm) A_g = np.zeros((n_img, n_img), dtype=complex) pairs = all_pairs(n_img) @@ -386,6 +407,152 @@ def g_sync(rots, order, rots_gt): return rots_gt_sync +def g_sync_finite_group(rots, rots_gt, symmetry): + """ + Synchronize ground-truth rotations over a finite symmetry group. + + This is a finite-group generalization of cyclic synchronization. The + pairwise matching step estimates relative symmetry elements between image + pairs. The spectral synchronization step then recovers one symmetry + element per image that is globally consistent with those pairwise estimates. + + Unlike the cyclic case, the relative symmetry cannot generally be encoded + as a scalar complex phase. For non-commutative groups such as D_n, we instead + represent each group element by its left-regular permutation matrix. + + :param rots: Estimated rotation matrices + :param rots_gt: Ground truth rotation matrices + :param symmetry: The symmetry of the underlying molecule. + + :return: g-synchronized ground truth rotations. + """ + assert len(rots) == len( + rots_gt + ), "Number of estimates not equal to number of references." + + n_img = len(rots) + dtype = rots.dtype + + # All matrices in the symmetry group. + G = SymmetryGroup.parse(symmetry).matrices + n_group = len(G) + + def find_group_index(A): + """Return the index of the group matrix closest to A.""" + dists = np.linalg.norm(G - A, axis=(1, 2)) + return np.argmin(dists) + + # Build multiplication and inverse tables for the symmetry group. + # mult[a, b] is the index c such that: + # G[a] @ G[b] == G[c] + # inv[a] is the index b such that: + # G[a] @ G[b] == identity + mult = np.empty((n_group, n_group), dtype=int) + inv = np.empty(n_group, dtype=int) + + for a in range(n_group): + inv[a] = find_group_index(G[a].T) + + for b in range(n_group): + mult[a, b] = find_group_index(G[a] @ G[b]) + + # Build the left-regular representation of the group. + # + # reps[a] is an n_group x n_group permutation matrix representing left + # multiplication by group element a: + # + # reps[a] @ e_b = e_{a b} + # + # This lets us store arbitrary finite-group relative elements in a block + # synchronization matrix. This is the key generalization beyond cyclic + # scalar phases in g_sync_Cn. + reps = np.zeros((n_group, n_group, n_group), dtype=float) + + for a in range(n_group): + for b in range(n_group): + reps[a, mult[a, b], b] = 1.0 + + # Block synchronization matrix. + # A is made of n_img x n_img blocks, each of size n_group x n_group. + # Block (i, j) stores the representation of the estimated relative + # symmetry element between images i and j. + A = np.zeros((n_img * n_group, n_img * n_group), dtype=float) + + # The relative symmetry from an image to itself is the identity. + for i in range(n_img): + sl_i = slice(i * n_group, (i + 1) * n_group) + A[sl_i, sl_i] = np.eye(n_group) + + for i, j in all_pairs(n_img): + # Estimated relative rotation. + Ri = rots[i] + Rj = rots[j] + Rij = Ri.T @ Rj + + # Ground-truth rotations for this pair. + Ri_gt = rots_gt[i] + Rj_gt = rots_gt[j] + + # Try every symmetry element and find which one makes the + # ground-truth relative rotation most closely match the estimated + # relative rotation. + diffs = np.zeros(n_group, dtype=float) + + for s, g_s in enumerate(G): + Rij_gt = Ri_gt.T @ g_s @ Rj_gt + diffs[s] = min( + np.linalg.norm(Rij - Rij_gt), + np.linalg.norm(Rij - J_conjugate(Rij_gt)), + ) + + # Estimates relative group element h_i^{-1} h_j. + idx_ij = np.argmin(diffs) + + sl_i = slice(i * n_group, (i + 1) * n_group) + sl_j = slice(j * n_group, (j + 1) * n_group) + + # Store the pairwise relative group elements in blocks (i, j)/(j, i). + A[sl_i, sl_j] = reps[idx_ij] + A[sl_j, sl_i] = reps[idx_ij].T + + # Spectral synchronization: + # In the noiseless case, this block matrix has a top eigenspace of + # dimension n_group. That eigenspace encodes the unknown per-image group + # elements, up to one global group action. + _, eig_vecs = np.linalg.eigh(A) + V = eig_vecs[:, -n_group:] + + # Fix the global gauge using image 0 as a reference: + # The recovered group elements are only determined up to a common global + # symmetry. This is fine for error analysis, because a single remaining + # global rotation/symmetry can still be applied later. + V0 = V[:n_group, :] + V0_pinv = np.linalg.pinv(V0) + + rots_gt_sync = np.zeros_like(rots_gt, dtype=dtype) + + for i, rot_gt in enumerate(rots_gt): + sl_i = slice(i * n_group, (i + 1) * n_group) + Vi = V[sl_i, :] + + # Compare image i's eigenspace block to the reference block. + # Ideally, this matrix is close to the representation of one group + # element. Noise makes it approximate, so we round it to the nearest + # valid group representation below. + M_i = Vi @ V0_pinv + + # Round to nearest group representation. + dists = np.linalg.norm(M_i - reps, axis=(1, 2)) + q_i = np.argmin(dists) + + # Apply the synchronized group element to the ground-truth rotation. + h_i_sync = inv[q_i] + + rots_gt_sync[i] = G[h_i_sync] @ rot_gt + + return rots_gt_sync + + def build_outer_products(n, dtype): """ Builds sets of outer products of 3rd rows of rotation matrices. diff --git a/tests/test_nug.py b/tests/test_nug.py index 62c86834a0..c83eb57303 100644 --- a/tests/test_nug.py +++ b/tests/test_nug.py @@ -1,8 +1,9 @@ import numpy as np import pytest -from aspire.abinitio import CommonlineNUG, compare_rots_sym +from aspire.abinitio import CommonlineNUG, compare_rots_sym, g_sync from aspire.source import Simulation +from aspire.utils import mean_aligned_angular_distance from aspire.volume import CnSymmetricVolume, DnSymmetricVolume, SymmetryGroup DTYPE = [np.float32, pytest.param(np.float64, marks=pytest.mark.expensive)] @@ -121,3 +122,10 @@ def test_estimate_rotations_pairwise(orient_est): orient_est.rotations, orient_est.src.rotations, orient_est.sym_grp ) np.testing.assert_array_less(MSE, 0.1) + + +def test_estimate_rotations(orient_est): + gt_rots_synced = g_sync( + orient_est.rotations, orient_est.src.rotations, orient_est.sym_grp + ) + mean_aligned_angular_distance(orient_est.rotations, gt_rots_synced, 8.0) diff --git a/tests/test_orient_symmetric.py b/tests/test_orient_symmetric.py index ed9c5a6904..df86ba020a 100644 --- a/tests/test_orient_symmetric.py +++ b/tests/test_orient_symmetric.py @@ -123,7 +123,7 @@ def test_estimate_rotations(n_img, L, order, dtype): rots_gt = src.rotations # g-synchronize ground truth rotations. - rots_gt_sync = g_sync(rots_est, order, rots_gt) + rots_gt_sync = g_sync(rots_est, rots_gt, src.symmetry_group) # Register estimates to ground truth rotations and check that the # mean angular distance between them is less than 3 degrees. From a0fceff12dad393ce698e0ef9c3a3b60dffe8fff Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 4 Jun 2026 14:37:29 -0400 Subject: [PATCH 029/161] tox --- src/aspire/abinitio/commonline_utils.py | 2 +- tests/test_nug.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/aspire/abinitio/commonline_utils.py b/src/aspire/abinitio/commonline_utils.py index 1d73b8fa16..3d2745b656 100644 --- a/src/aspire/abinitio/commonline_utils.py +++ b/src/aspire/abinitio/commonline_utils.py @@ -4,7 +4,7 @@ from numpy.linalg import eigh, norm from aspire.operators import PolarFT -from aspire.utils import J_conjugate, Rotation, all_pairs, anorm, cyclic_rotations, tqdm +from aspire.utils import J_conjugate, Rotation, all_pairs, anorm, tqdm from aspire.volume import CnSymmetryGroup, SymmetryGroup logger = logging.getLogger(__name__) diff --git a/tests/test_nug.py b/tests/test_nug.py index c83eb57303..a2f5adad1a 100644 --- a/tests/test_nug.py +++ b/tests/test_nug.py @@ -4,7 +4,7 @@ from aspire.abinitio import CommonlineNUG, compare_rots_sym, g_sync from aspire.source import Simulation from aspire.utils import mean_aligned_angular_distance -from aspire.volume import CnSymmetricVolume, DnSymmetricVolume, SymmetryGroup +from aspire.volume import CnSymmetricVolume, DnSymmetricVolume DTYPE = [np.float32, pytest.param(np.float64, marks=pytest.mark.expensive)] RESOLUTION = [48, pytest.param(49, marks=pytest.mark.expensive)] From 7f469840fbf30576984347857ad4424c3814f085 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 8 Jun 2026 14:34:55 -0400 Subject: [PATCH 030/161] Add g_sync test --- src/aspire/abinitio/commonline_utils.py | 3 ++ tests/test_commonline_utils.py | 53 ++++++++++++++++++++++++- tests/test_orient_d2.py | 2 +- 3 files changed, 55 insertions(+), 3 deletions(-) diff --git a/src/aspire/abinitio/commonline_utils.py b/src/aspire/abinitio/commonline_utils.py index 3d2745b656..ec4c0cfbf7 100644 --- a/src/aspire/abinitio/commonline_utils.py +++ b/src/aspire/abinitio/commonline_utils.py @@ -394,6 +394,9 @@ def g_sync_cyclic(rots, rots_gt, symmetry): _, eig_vecs = eigh(A_g) leading_eig_vec = eig_vecs[:, -1] + # Remove arbitrary global phase (eigh returns +-1 eigenvector) + leading_eig_vec *= np.exp(-1j * np.angle(leading_eig_vec[0])) + angles = np.exp(1j * 2 * np.pi / order * np.arange(order)) rots_gt_sync = np.zeros((n_img, 3, 3), dtype=dtype) diff --git a/tests/test_commonline_utils.py b/tests/test_commonline_utils.py index 64d015d1dd..818d829c6a 100644 --- a/tests/test_commonline_utils.py +++ b/tests/test_commonline_utils.py @@ -1,13 +1,20 @@ import numpy as np import pytest -from aspire.abinitio import JSync +from aspire.abinitio import JSync, g_sync from aspire.abinitio.commonline_utils import ( _complete_third_row_to_rot, _estimate_third_rows, build_outer_products, ) -from aspire.utils import J_conjugate, Rotation, randn, utest_tolerance +from aspire.utils import ( + J_conjugate, + Rotation, + mean_aligned_angular_distance, + randn, + utest_tolerance, +) +from aspire.volume import SymmetryGroup DTYPES = [np.float32, np.float64] @@ -116,3 +123,45 @@ def test_J_sync(dtype): np.testing.assert_allclose(Rijs_sync, Rijs_gt) assert Rijs_sync.dtype == dtype + + +@pytest.mark.parametrize("symmetry", ["C3", "C4", "D3", "D4", "T", "O"]) +def test_g_sync(symmetry): + n = 100 + dtype = np.float64 + + # Get symmetry group matrices + gs = SymmetryGroup.parse(symmetry).matrices + + # Build set of ground truth rotations + gt_rots = Rotation.generate_random_rotations(n, dtype=dtype) + + # Build set of estimates which are close to ground truth + # by generating set of small perturbation rotations to apply + # to ground truth rotations. + target_mean_deg = 2.0 + axes = np.random.normal(size=(n, 3)).astype(dtype) + axes /= np.linalg.norm(axes, axis=1, keepdims=True) + angles = np.random.uniform(0, 2 * np.deg2rad(target_mean_deg), n).astype(dtype) + delta_rots = Rotation.from_rotvec(axes * angles[:, None], dtype=dtype) + perturbed_rots = Rotation(delta_rots.matrices @ gt_rots.matrices) + + # Get mean ang dist for aligned estimates + # and check we're close to target. + og_maad = mean_aligned_angular_distance(perturbed_rots, gt_rots) + np.testing.assert_array_less(abs(og_maad - target_mean_deg), 0.2) + + # Simulate symmetry desynchronization + g_idx = np.random.randint(len(gs), size=n) + desynced_rots = Rotation(gs[g_idx] @ perturbed_rots.matrices) + + # Mean aligned angular distance of unsynced rots should be bad + np.testing.assert_array_less( + 10 * og_maad, mean_aligned_angular_distance(desynced_rots, gt_rots) + ) + + # Perform g_sync and check that mean aligned angular distance + # matches ground truth MAAD to within .1 degrees. + rots_gt_synced = g_sync(desynced_rots, gt_rots, symmetry) + est_maad = mean_aligned_angular_distance(desynced_rots, rots_gt_synced) + np.testing.assert_array_less(abs(og_maad - est_maad), 0.1) diff --git a/tests/test_orient_d2.py b/tests/test_orient_d2.py index cc2013d07a..2d1bb73d2b 100644 --- a/tests/test_orient_d2.py +++ b/tests/test_orient_d2.py @@ -1,7 +1,7 @@ import numpy as np import pytest -from aspire.abinitio import CLSymmetryD2 +from aspire.abinitio import CLSymmetryD2, g_sync from aspire.source import Simulation from aspire.utils import ( J_conjugate, From 5046b3f0023f39e9657ec55fe05a5c9ea843eba9 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 8 Jun 2026 15:20:11 -0400 Subject: [PATCH 031/161] Use new g_sync in D2 test. Remove unused method --- tests/test_orient_d2.py | 74 +---------------------------------------- 1 file changed, 1 insertion(+), 73 deletions(-) diff --git a/tests/test_orient_d2.py b/tests/test_orient_d2.py index 2d1bb73d2b..1c05691c94 100644 --- a/tests/test_orient_d2.py +++ b/tests/test_orient_d2.py @@ -103,7 +103,7 @@ def test_estimate_rotations(orient_est): rots_gt = orient_est.src.rotations # g-sync ground truth rotations. - rots_gt_sync = g_sync_d2(rots_est, rots_gt) + rots_gt_sync = g_sync(rots_est, rots_gt, "D2") # Register estimates to ground truth rotations and check that the mean angular # distance between them is less than 5 degrees. @@ -382,78 +382,6 @@ def test_sync_signs(orient_est): assert rots_est.dtype == orient_est.dtype -#################### -# Helper Functions # -#################### - - -def g_sync_d2(rots, rots_gt): - """ - Every estimated rotation might be a version of the ground truth rotation - rotated by g^{s_i}, where s_i = 0, 1, ..., order. This method synchronizes the - ground truth rotations so that only a single global rotation need be applied - to all estimates for error analysis. - - :param rots: Estimated rotation matrices - :param rots_gt: Ground truth rotation matrices. - - :return: g-synchronized ground truth rotations. - """ - assert len(rots) == len( - rots_gt - ), "Number of estimates not equal to number of references." - n_img = len(rots) - dtype = rots.dtype - - rots_symm = DnSymmetryGroup(2).matrices.astype(dtype, copy=False)[[0, 2, 1, 3]] - order = len(rots_symm) - - A_g = np.zeros((n_img, n_img), dtype=complex) - - pairs = all_pairs(n_img) - - for i, j in pairs: - Ri = rots[i] - Rj = rots[j] - Rij = Ri.T @ Rj - - Ri_gt = rots_gt[i] - Rj_gt = rots_gt[j] - - diffs = np.zeros(order) - for s, g_s in enumerate(rots_symm): - Rij_gt = Ri_gt.T @ g_s @ Rj_gt - diffs[s] = min( - [ - np.linalg.norm(Rij - Rij_gt), - np.linalg.norm(Rij - J_conjugate(Rij_gt)), - ] - ) - - idx = np.argmin(diffs) - - A_g[i, j] = np.exp(-1j * 2 * np.pi / order * idx) - - # A_g(k,l) is exp(-j(-theta_k+theta_l)) - # Diagonal elements correspond to exp(-i*0) so put 1. - # This is important only for verification purposes that spectrum is (K,0,0,0...,0). - A_g += np.conj(A_g).T + np.eye(n_img) - _, eig_vecs = np.linalg.eigh(A_g) - leading_eig_vec = eig_vecs[:, -1] - - angles = np.exp(1j * 2 * np.pi / order * np.arange(order)) - rots_gt_sync = np.zeros((n_img, 3, 3), dtype=dtype) - - for i, rot_gt in enumerate(rots_gt): - # Since the closest ccw or cw rotation are just as good, - # we take the absolute value of the angle differences. - angle_dists = np.abs(np.angle(leading_eig_vec[i] / angles)) - power_g_Ri = np.argmin(angle_dists) - rots_gt_sync[i] = rots_symm[power_g_Ri] @ rot_gt - - return rots_gt_sync - - def build_cl_from_source(source): # Search for common lines over less shifts for 0 offsets. max_shift = 0 From 7ea89f3a8364b1c796e38e4a2bbc11eb96374172 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 8 Jun 2026 15:22:36 -0400 Subject: [PATCH 032/161] remove unused imports --- tests/test_orient_d2.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/test_orient_d2.py b/tests/test_orient_d2.py index 1c05691c94..fba062ab9a 100644 --- a/tests/test_orient_d2.py +++ b/tests/test_orient_d2.py @@ -7,11 +7,10 @@ J_conjugate, Random, Rotation, - all_pairs, mean_aligned_angular_distance, utest_tolerance, ) -from aspire.volume import DnSymmetricVolume, DnSymmetryGroup +from aspire.volume import DnSymmetricVolume ############## # Parameters # From c63fe4a7bb9ab3a67df7cf5ddb78f7f3205104dd Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 9 Jun 2026 13:47:41 -0400 Subject: [PATCH 033/161] remove fast_radon_transform and wemd code path --- src/aspire/abinitio/commonline_nug.py | 104 ++++---------------------- tests/test_nug.py | 2 +- 2 files changed, 16 insertions(+), 90 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 0bb1ebf8bd..3114160fce 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -5,9 +5,8 @@ from scipy.special import factorial from aspire.abinitio import Orient3D -from aspire.nufft import nufft -from aspire.numeric import fft, xp -from aspire.operators import PolarFT, wemd_embed +from aspire.numeric import xp +from aspire.operators import PolarFT from aspire.utils import Rotation, cart2sph, complex_type from aspire.volume import CnSymmetryGroup, DnSymmetryGroup, SymmetryGroup @@ -31,7 +30,6 @@ def __init__( shift_step=1, mask=True, Lmax=12, - loss="l1", T=36, max_iter=501, rho=0.05, @@ -64,7 +62,6 @@ def __init__( ) self.Lmax = Lmax - self.loss = loss self.T = T self.max_iter = max_iter self.rho = rho @@ -120,61 +117,24 @@ def estimate_rotations(self): def compute_coeff(self): # compute the coefficient matrix - Img = self.src.images[:] N = self.n_img n_theta = self.n_theta - loss = self.loss Lmax = self.Lmax T = self.T - angular_sampling = np.arange(0, 360, 1, dtype=np.float64) - - # Using ASPIRE Image.project(). Leaving original method in comments for now. - line_proj = Img.project(angular_sampling).asnumpy().T - - # line_proj = np.zeros((L, n_theta, N)) - # Img_pft = np.zeros((L, n_theta, N), dtype=complex) - # Img = Img.asnumpy() - # for n in range(N): - # line_proj[:, :, n], Img_pft[:, :, n] = self.fast_radon_transform( - # Img[n], angular_sampling - # ) - def fij(alpha, gamma, i, j, loss): - if loss == "l1": - # Ii_hat = Img_pft[:, :, i] - # Ij_hat = Img_pft[:, :, j] - # idxi = np.round((alpha - np.pi / 2) * n_theta / 2 / np.pi) % n_theta - # idxj = np.round((-gamma - np.pi / 2) * n_theta / 2 / np.pi) % n_theta - - # Si = Ii_hat[:, int(idxi)] - # Sj = Ij_hat[:, int(idxj)] - - Ii_hat = self.pf_full[i] - Ij_hat = self.pf_full[j] - idxi = np.round((alpha - np.pi / 2) * n_theta / 2 / np.pi) % n_theta - idxj = np.round((-gamma - np.pi / 2) * n_theta / 2 / np.pi) % n_theta - - Si = Ii_hat[int(idxi)] - Sj = Ij_hat[int(idxj)] - - # Apply shifts - Sj_shifted = self.shift_phases * Sj - norms = np.linalg.norm(Si[None] - Sj_shifted, 1, axis=1) - return norms.min() - - if loss == "wemd": - dim_wave = len(wemd_embed(line_proj[:, 0, 0])) - WE = np.zeros((dim_wave, n_theta, N), dtype=np.float64) - for i in range(N): - for theta in range(n_theta): - WE[:, theta, i] = wemd_embed(line_proj[:, theta, i]) - idxi = np.round((alpha - np.pi / 2) * n_theta / 2 / np.pi) % n_theta - idxj = np.round((-gamma - np.pi / 2) * n_theta / 2 / np.pi) % n_theta + def fij(alpha, gamma, i, j): + Ii_hat = self.pf_full[i] + Ij_hat = self.pf_full[j] + idxi = np.round((alpha - np.pi / 2) * n_theta / 2 / np.pi) % n_theta + idxj = np.round((-gamma - np.pi / 2) * n_theta / 2 / np.pi) % n_theta - Si = WE[:, int(idxi), i] - Sj = WE[:, int(idxj), j] + Si = Ii_hat[int(idxi)] + Sj = Ij_hat[int(idxj)] - return np.linalg.norm(Si - Sj, 1) + # Apply shifts + Sj_shifted = self.shift_phases * Sj + norms = np.linalg.norm(Si[None] - Sj_shifted, 1, axis=1) + return norms.min() alpha_grid = np.arange(2 * T, dtype=np.float64) * np.pi / T beta_grid = (2 * np.arange(2 * T, dtype=np.float64) + 1) * np.pi / 4 / T @@ -216,7 +176,7 @@ def fijhat_k(k, F): Fij = np.zeros((2 * T, 2 * T), dtype=np.float64) for j1 in range(2 * T): for j2 in range(2 * T): - Fij[j1, j2] = fij(alpha_grid[j1], gamma_grid[j2], i, j, loss) + Fij[j1, j2] = fij(alpha_grid[j1], gamma_grid[j2], i, j) for k in range(1, Lmax + 1): dk = 2 * k + 1 C[k - 1][j * dk : (j + 1) * dk, i * dk : (i + 1) * dk] = fijhat_k( @@ -229,7 +189,7 @@ def fijhat_k(k, F): Fii = np.zeros((2 * T, 2 * T), dtype=np.float64) for j1 in range(2 * T): for j2 in range(2 * T): - Fii[j1, j2] = fij(alpha_grid[j1], gamma_grid[j2], i, i, loss) + Fii[j1, j2] = fij(alpha_grid[j1], gamma_grid[j2], i, i) for k in range(1, Lmax + 1): dk = 2 * k + 1 C[k - 1][i * dk : (i + 1) * dk, i * dk : (i + 1) * dk] = fijhat_k( @@ -246,40 +206,6 @@ def fijhat_k(k, F): self.C = C - @staticmethod - def fast_radon_transform(array, angles, use_ramp=False): - - angles = np.array(angles).flatten() - img_size = array.shape[1] - rads = angles / 180 * np.pi - y_idx = np.arange(-img_size / 2, img_size / 2) / img_size * 2 - x_theta = y_idx[:, np.newaxis] * np.sin(rads)[np.newaxis, :] - y_theta = y_idx[:, np.newaxis] * np.cos(rads)[np.newaxis, :] - - pts = np.pi * np.vstack( - [ - x_theta.flatten(), - y_theta.flatten(), - ] - ) - pts = pts.astype(array.dtype) - - # array = array.astype(np.float32) - lines_f = nufft(array, pts).reshape((img_size, -1)) - - if img_size % 2 == 0: - lines_f[0, :] = 0 - - if use_ramp: - freqs = np.abs(np.pi * y_idx) - lines_f *= freqs[:, np.newaxis] - - projections = np.real( - xp.asnumpy(fft.centered_ifft(xp.asarray(lines_f), axis=0)) - ) - - return projections, lines_f - def complex2real(self, ell): # compute transformation matrices that convert complex representations to real ones diml = 2 * ell + 1 diff --git a/tests/test_nug.py b/tests/test_nug.py index a2f5adad1a..b9699e7f66 100644 --- a/tests/test_nug.py +++ b/tests/test_nug.py @@ -6,7 +6,7 @@ from aspire.utils import mean_aligned_angular_distance from aspire.volume import CnSymmetricVolume, DnSymmetricVolume -DTYPE = [np.float32, pytest.param(np.float64, marks=pytest.mark.expensive)] +DTYPE = [np.float64, pytest.param(np.float32, marks=pytest.mark.expensive)] RESOLUTION = [48, pytest.param(49, marks=pytest.mark.expensive)] N_IMG = [15] OFFSETS = [0, pytest.param(None, marks=pytest.mark.expensive)] From 7fede11579609e29c144958d1579067e84d52c4c Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 9 Jun 2026 15:12:24 -0400 Subject: [PATCH 034/161] cleanup --- src/aspire/abinitio/commonline_nug.py | 25 ++++--------------------- 1 file changed, 4 insertions(+), 21 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 3114160fce..a236cac83f 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -17,7 +17,8 @@ class CommonlineNUG(Orient3D): """ - Class to estimate 3D orientations using non-uqique games. + Class to estimate 3D orientations using non-uqique games for molecules with cyclic + or dihedral symmetry. """ def __init__( @@ -95,16 +96,12 @@ def _build_full_pft(self): pf = self.pf self.pf_full = PolarFT.half_to_full(pf) - # Prepare the shift phases to try and generate filter for common-line detection + # Prepare the shift phases for common-line detection r_max = self.pf_full.shape[2] - self.shifts, self.shift_phases, h = _generate_shift_phase_and_filter( + self.shifts, self.shift_phases, _ = _generate_shift_phase_and_filter( r_max, self.max_shift, self.shift_step, self.dtype ) - # Apply bandpass filter, normalize each ray of each image - # Note that only use half of each ray - # self.pf_full = self._apply_filter_and_norm("ijk, k -> ijk", pf_full, r_max, h) - def estimate_rotations(self): self.compute_coeff() self.perform_admm() @@ -709,8 +706,6 @@ def ADMM_preprocessing(self, C): D1 = d1[-1] # AE and bE for quaternion constraints - # AEq = xp.asarray(loadmat("data/Eq_constraints/AEqJ.mat")["AEq"]) - # AEqAEqtinv = xp.asarray(loadmat("data/Eq_constraints/AEqJ.mat")["AEqAEqtinv"]) AEq = xp.asarray(self.construct_AEq()) AEqAEqtinv = xp.linalg.pinv(AEq @ AEq.T) @@ -721,16 +716,6 @@ def ADMM_preprocessing(self, C): # AI and bI W0, W1, Ngrid = self.compute_fejer_weights() - # AI_mat=np.zeros((Ngrid,D0+D1)) - # for p in range(Ngrid): - # w0=np.zeros(D0); w1=np.zeros(D1) - # for k in range(1,Lmax+1): - # w0[d0[k-1]:d0[k]]=(Lmax-k+2)*(Lmax-k+1)*(k+0.5)*W0[k-1][p].T.reshape(-1) - # w1[d1[k-1]:d1[k]]=(Lmax-k+2)*(Lmax-k+1)*(k+0.5)*W1[k-1][p].T.reshape(-1) - # # this needs double checking - # AI_mat[p,:d0[-1]]=w0; AI_mat[p,d0[-1]:]=w1 - # AI_mat=xp.asarray(AI_mat) / 10; - # bI=-(Lmax+2)*(Lmax+1)/2 / 10 AI_mat_offdiag = np.zeros((Ngrid, D0 + D1), dtype=np.float64) for p in range(Ngrid): w0 = np.zeros(D0, dtype=np.float64) @@ -874,7 +859,6 @@ def permutek_block(Ak, k): return W0, W1, Ngrid def discretize_SO3(self): - # S2 = loadmat("design20.mat")["design"] S2 = saff_kuijlaars(self.S2_grid) S2_size = S2.shape[0] @@ -1354,7 +1338,6 @@ def largest_eigenvalue(AI, Ngrid, N): z = AI @ (AI.T @ z) Lambda += 2000 logger.info("Largest eigenvalue of AIAIT is approximately %1.2f" % Lambda) - # Lambda=xp.linalg.eigvalsh(AI@AI.T)[-1]; print(Lambda) return Lambda def compute_rank(self, Lmax): From 50f1d4d01ad20d59a21287f6bfee2a5bdd798cb6 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 10 Jun 2026 11:13:24 -0400 Subject: [PATCH 035/161] Add function docstrings --- src/aspire/abinitio/commonline_nug.py | 89 ++++++++++++++++++++++++--- 1 file changed, 81 insertions(+), 8 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index a236cac83f..3fe18e2d84 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -17,8 +17,7 @@ class CommonlineNUG(Orient3D): """ - Class to estimate 3D orientations using non-uqique games for molecules with cyclic - or dihedral symmetry. + Estimate orientations of cyclically or dihedrally symmetric molecules using the non-unique games framework. """ def __init__( @@ -44,7 +43,7 @@ def __init__( **kwargs, ): """ - Initialize object for estimating 3D orientations for symmetric molecules. + Initialize the symmetric NUG orientation estimator. :param src: The source object of 2D denoised or class-averaged images with metadata :param symmetry: A string, ie. 'C3', indicating the symmetry type. @@ -93,6 +92,9 @@ def __init__( self._build_full_pft() def _build_full_pft(self): + """ + Construct the full polar Fourier transforms and candidate shift phases. + """ pf = self.pf self.pf_full = PolarFT.half_to_full(pf) @@ -103,6 +105,9 @@ def _build_full_pft(self): ) def estimate_rotations(self): + """ + Estimate rotations by computing NUG coefficients, solving the SDP relaxation, and recovering Euler angles. + """ self.compute_coeff() self.perform_admm() self.euler_est() @@ -113,6 +118,9 @@ def estimate_rotations(self): ####################### def compute_coeff(self): + """ + Compute the truncated Fourier coefficient matrices of the pairwise common-line losses. + """ # compute the coefficient matrix N = self.n_img n_theta = self.n_theta @@ -204,7 +212,9 @@ def fijhat_k(k, F): self.C = C def complex2real(self, ell): - # compute transformation matrices that convert complex representations to real ones + """ + Construct the transformation matrices between complex and real degree ell representations. + """ diml = 2 * ell + 1 Tinv = np.zeros((diml, diml), dtype=complex_type(np.float64)) for i in range(diml): @@ -234,6 +244,9 @@ def complex2real(self, ell): ############# def perform_admm(self): + """ + Solve the symmetric NUG relaxation and optionally apply proximal refinement. + """ X_est = self.admm_sym_J(self.C, self.verbose) if self.perform_pr: @@ -249,6 +262,9 @@ def perform_admm(self): self.X_est = X_est def admm_sym_J(self, C, verbose): + """ + Solve the symmetry-constrained NUG semidefinite relaxation using ADMM. + """ Lmax = self.Lmax N = self.n_img max_iter = self.max_iter @@ -656,8 +672,9 @@ def print_updates(verbose): return X_admm def ADMM_preprocessing(self, C): - # compute necessary quantities for ADMM - # compute some useful index sets + """ + Construct the transformed coefficients, constraints, indices, and initial variables used by ADMM. + """ Lmax = self.Lmax N = self.n_img count = 0 @@ -820,6 +837,9 @@ def ADMM_preprocessing(self, C): ) def compute_fejer_weights(self): + """ + Evaluate the real Wigner representation blocks used by the discretized Fejer inequality constraints. + """ SO3_grid = self.discretize_SO3() Ngrid = SO3_grid.shape[0] start = 1 @@ -859,6 +879,9 @@ def permutek_block(Ak, k): return W0, W1, Ngrid def discretize_SO3(self): + """ + Construct an approximately uniform Euler-angle grid over SO(3). + """ S2 = saff_kuijlaars(self.S2_grid) S2_size = S2.shape[0] @@ -886,6 +909,9 @@ def discretize_SO3(self): ############################ def proximal_refine(self, X_admm, weight, Penalty, r): + """ + Refine the relaxed solution by iteratively encouraging lower-rank representation matrices. + """ N = self.n_img C = self.C @@ -959,6 +985,9 @@ def low_rank_proj(X, r_step): ######################### def euler_est(self): + """ + Recover rotations and Euler angles using the estimator for the configured symmetry group. + """ X_est = self.X_est if isinstance(self.sym_grp, CnSymmetryGroup): R_est, Euler_est = self.euler_est_Cm(X_est[0], X_est[self.n_sym - 1]) @@ -969,6 +998,9 @@ def euler_est(self): self.rotations = R_est.astype(self.dtype) def euler_est_Cm(self, X1, XS): + """ + Recover Euler angles from the degree-one and degree-m solutions for cyclic symmetry. + """ S = self.n_sym N = self.n_img sym_euler = np.zeros((S, 3), dtype=np.float64) @@ -1086,7 +1118,9 @@ def find_gamma(Xm, beta, alpha): return R_est, Euler_est def euler_est_Dm(self, X_est): - + """ + Recover Euler angles from the degree-two and degree-m solutions for dihedral symmetry. + """ X2 = X_est[1] S = self.sym_grp.order N = self.n_img @@ -1240,6 +1274,9 @@ def LS_D(W1, W2, W3, W4, Br, Bi): # Helper Functions # #################### def transform_coeff(self, A, Lmax, N, IDX_upper): + """ + Transform representation matrices into the two block-vector forms used by ADMM. + """ d0 = [0] d1 = [0] for k in range(1, Lmax + 1): @@ -1254,6 +1291,9 @@ def transform_coeff(self, A, Lmax, N, IDX_upper): return A0, A1 def transform_coeff_back(self, A0, A1, Lmax, N, IDX_upper, IDX_lower, idx_offdiag): + """ + Reconstruct representation matrices from the ADMM block-vector forms. + """ d0 = [0] d1 = [0] for k in range(1, Lmax + 1): @@ -1275,6 +1315,9 @@ def transform_coeff_back(self, A0, A1, Lmax, N, IDX_upper, IDX_lower, idx_offdia @staticmethod def permutek(Ak, k, N): + """ + Permute and split a degree-k matrix into blocks of sizes k and k + 1. + """ AkP = xp.copy(Ak) dk = 2 * k + 1 Pk = xp.eye(dk, dtype=AkP.dtype) @@ -1299,6 +1342,9 @@ def permutek(Ak, k, N): @staticmethod def permutek_back(Ak, k, N): + """ + Undo the degree-k block permutation and reconstruct the full matrix. + """ dk = 2 * k + 1 Pk = xp.eye(N * dk, dtype=Ak.dtype) idx = xp.concatenate((xp.arange(dk - k, dk), xp.arange(k + 1))) @@ -1322,11 +1368,17 @@ def permutek_back(Ak, k, N): @staticmethod def vec_block(A, N, sz, IDX_upper): + """ + Vectorize the upper-triangular image-pair blocks of a block matrix. + """ vecA = (A.reshape(N, sz, N, sz).transpose(0, 2, 3, 1)).reshape(N**2, sz**2).T return vecA[:, IDX_upper] @staticmethod def largest_eigenvalue(AI, Ngrid, N): + """ + Estimate the largest eigenvalue of the Fejér constraint operator. + """ # find the largest eigenvalue of the operator AI np.random.seed(0) z = xp.random.normal(0, 1, (Ngrid, N**2)) @@ -1341,6 +1393,9 @@ def largest_eigenvalue(AI, Ngrid, N): return Lambda def compute_rank(self, Lmax): + """ + Compute the ranks and matrices of the symmetry-averaging projectors at each degree. + """ rk = xp.zeros(Lmax, dtype=np.float64) A = [] for k in range(1, Lmax + 1): @@ -1351,6 +1406,9 @@ def compute_rank(self, Lmax): return rk, A def WD(self, J, euler): + """ + Evaluate degree-J Wigner D matrices at the supplied ZYZ Euler angles. + """ # compute Wigner D matrix alpha = euler[:, 0] beta = euler[:, 1] @@ -1366,6 +1424,9 @@ def WD(self, J, euler): @staticmethod def Wd(J, beta): + """ + Evaluate degree-J Wigner small-d matrices at the supplied polar angles. + """ # compute Wigner small d matrix d = np.zeros((len(beta), 2 * J + 1, 2 * J + 1), dtype=beta.dtype) for m in range(-J, J + 1): @@ -1393,6 +1454,9 @@ def Wd(J, beta): @staticmethod def mat_block(vecA, N, sz, IDX_upper, IDX_lower, idx_offdiag): + """ + Reconstruct a symmetric block matrix from its vectorized upper-triangular blocks. + """ tmp = vecA.T.reshape(N * (N + 1) // 2, sz, sz).transpose(0, 2, 1) AA = xp.zeros((N**2, sz, sz), dtype=vecA.dtype) AA[IDX_upper] = tmp @@ -1401,6 +1465,9 @@ def mat_block(vecA, N, sz, IDX_upper, IDX_lower, idx_offdiag): @staticmethod def psd_projection(B): + """ + Project one or more symmetric matrices onto the positive semidefinite cone. + """ # compute the PSD part of a symmstric matrix B_sym = (B + B.swapaxes(-1, -2)) / 2 evals, evecs = xp.linalg.eigh(B_sym) @@ -1409,6 +1476,9 @@ def psd_projection(B): @staticmethod def transform_block(A, k, Pk=None): + """ + Permute and vectorize the two invariant blocks of a degree-k matrix. + """ single = A.ndim == 2 if single: A = A[None, :, :] @@ -1432,6 +1502,9 @@ def transform_block(A, k, Pk=None): @staticmethod def transform_back_block(A0, A1, k, Pk=None): + """ + Reconstruct a degree-k matrix from its two invariant block vectors. + """ dk = 2 * k + 1 single = A0.ndim == 1 @@ -1454,7 +1527,7 @@ def transform_back_block(A0, A1, k, Pk=None): def construct_AEq(self): """ - Construct the linear equality matrix for quaternion constraints. + Construct the linear equality operator encoding the quaternion constraints. """ AEq = np.zeros((17, 21), np.float64) From 1ce900ec4de98359cefe4ebe013cfd6daba5afa7 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 15 Jun 2026 11:53:29 -0400 Subject: [PATCH 036/161] Add param descrips for init --- src/aspire/abinitio/commonline_nug.py | 23 +++++++++++++++++++---- src/aspire/abinitio/commonline_utils.py | 2 +- 2 files changed, 20 insertions(+), 5 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 3fe18e2d84..94504f236b 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -45,10 +45,25 @@ def __init__( """ Initialize the symmetric NUG orientation estimator. - :param src: The source object of 2D denoised or class-averaged images with metadata - :param symmetry: A string, ie. 'C3', indicating the symmetry type. - :param n_rad: The number of points in the radial direction - :param n_theta: The number of points in the theta direction. Default = 360. + :param src: Source containing the input projection images. + :param symmetry: Cyclic or dihedral symmetry specification, such as 'C3' or + 'D4'. If omitted, uses the source symmetry. + :param n_rad: Number of radial samples in the polar Fourier transform. + :param n_theta: Number of angular samples in the polar Fourier transform. + :param max_shift: Maximum shift considered when comparing common lines. + :param shift_step: Sampling interval for candidate shifts. + :param mask: Whether to apply a circular mask to the input images. + :param Lmax: Maximum Wigner representation degree used in the relaxation. + :param T: Quadrature resolution used to compute the Fourier coefficients. + :param max_iter: Number of ADMM iterations. + :param rho: Initial ADMM penalty parameter. + :param ratio: Residual ratio used when updating the ADMM penalty. + :param factor: Scaling factor used when updating the ADMM penalty. + :param mult: Step-size multiplier for the ADMM primal update. + :param S2_grid: Number of sphere samples used to discretize SO(3). + :param Nstep_yI: Number of inequality-multiplier updates per ADMM iteration. + :param perform_pr: Whether to apply proximal refinement after ADMM. + :param verbose: Whether to log ADMM progress. """ super().__init__( diff --git a/src/aspire/abinitio/commonline_utils.py b/src/aspire/abinitio/commonline_utils.py index ec4c0cfbf7..283dfaddc8 100644 --- a/src/aspire/abinitio/commonline_utils.py +++ b/src/aspire/abinitio/commonline_utils.py @@ -468,7 +468,7 @@ def find_group_index(A): # # This lets us store arbitrary finite-group relative elements in a block # synchronization matrix. This is the key generalization beyond cyclic - # scalar phases in g_sync_Cn. + # scalar phases in g_sync_cyclic. reps = np.zeros((n_group, n_group, n_group), dtype=float) for a in range(n_group): From 7ec02079b2c61d635bb69c3cb5d9eef3883350d5 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 15 Jun 2026 15:11:42 -0400 Subject: [PATCH 037/161] more params. cleanup. --- src/aspire/abinitio/commonline_nug.py | 120 +++++++++++++++++++------- 1 file changed, 88 insertions(+), 32 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 94504f236b..61bcaabdab 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -122,6 +122,8 @@ def _build_full_pft(self): def estimate_rotations(self): """ Estimate rotations by computing NUG coefficients, solving the SDP relaxation, and recovering Euler angles. + + :return: Estimated rotation matrices. """ self.compute_coeff() self.perform_admm() @@ -229,6 +231,10 @@ def fijhat_k(k, F): def complex2real(self, ell): """ Construct the transformation matrices between complex and real degree ell representations. + + :param ell: Wigner representation degree. + + :return: Forward and inverse change-of-basis matrices. """ diml = 2 * ell + 1 Tinv = np.zeros((diml, diml), dtype=complex_type(np.float64)) @@ -279,6 +285,11 @@ def perform_admm(self): def admm_sym_J(self, C, verbose): """ Solve the symmetry-constrained NUG semidefinite relaxation using ADMM. + + :param C: Fourier coefficient matrices of the NUG objective. + :param verbose: Whether to log ADMM progress. + + :return: Relaxed representation matrices. """ Lmax = self.Lmax N = self.n_img @@ -319,7 +330,7 @@ def admm_sym_J(self, C, verbose): n_pairs = N * (N - 1) // 2 Ngrid = self.Ngrid - rank_Ak, _ = self.compute_rank(Lmax) + rank_Ak, _ = self.compute_rank() logger.info(f"Rank of Ak: {rank_Ak}") # rank_Ak=cp.zeros(Lmax) @@ -544,8 +555,6 @@ def update_rho(X0, X1, Xd0, Xd1, Xq, bE, bEq, bI, res_X, rho, factor, normC): return rho, p_resnorm, d_resnorm def print_updates(verbose): - # X_admm=transform_coeff_back(X0,X1,Lmax,N); obj_p=0 - # for k in range(Lmax): obj_p+=xp.trace(C[k]@X_admm[k]) if verbose: obj_p = ( xp.vdot(C0[:, idx_diag], X0[:, idx_diag]) @@ -676,12 +685,7 @@ def print_updates(verbose): X0, X1, Xd0, Xd1, Xq, bE, bEq, bI, res_X, rho, factor, normC ) - X_admm = self.transform_coeff_back( - X0, X1, Lmax, N, IDX_upper, IDX_lower, idx_offdiag - ) - # if self.GPU: - # for k in range(Lmax): - # X_admm[k] = X_admm[k].get() + X_admm = self.transform_coeff_back(X0, X1, IDX_upper, IDX_lower, idx_offdiag) for k in range(Lmax): X_admm[k] = xp.asnumpy(X_admm[k]) return X_admm @@ -689,6 +693,10 @@ def print_updates(verbose): def ADMM_preprocessing(self, C): """ Construct the transformed coefficients, constraints, indices, and initial variables used by ADMM. + + :param C: Fourier coefficient matrices of the NUG objective. + + :return: Quantities required by the ADMM solver. """ Lmax = self.Lmax N = self.n_img @@ -724,7 +732,7 @@ def ADMM_preprocessing(self, C): Xnorm = np.sqrt(Xnorm) for k in range(Lmax): C[k] = xp.asarray(Xnorm / Cnorm * C[k]) - C0, C1 = self.transform_coeff(C, Lmax, N, IDX_upper) + C0, C1 = self.transform_coeff(C, IDX_upper) normC = np.sqrt(np.linalg.norm(C0) ** 2 + np.linalg.norm(C1) ** 2) del C @@ -812,18 +820,16 @@ def ADMM_preprocessing(self, C): for k in range(1, Lmax + 1): dk = 2 * k + 1 II.append(xp.eye(N * dk, dtype=np.float64)) - I0, I1 = self.transform_coeff(II, Lmax, N, IDX_upper) + I0, I1 = self.transform_coeff(II, IDX_upper) X0 = xp.zeros((D0, N * (N + 1) // 2), dtype=np.float64) X1 = xp.zeros((D1, N * (N + 1) // 2), dtype=np.float64) Xq = xp.zeros((16, N * (N - 1) // 2), dtype=np.float64) - # X0,X1=transform_coeff(II,Lmax,N); Xq=xp.zeros((16,N*(N-1))); Xq[0,:]=1; S0 = xp.copy(I0) S1 = xp.copy(I1) Sq = xp.zeros(Xq.shape, dtype=np.float64) - # S0=xp.zeros(X0.shape); S1=xp.zeros(X1.shape); Sq=xp.zeros(Xq.shape) self.Ngrid = Ngrid - # return C0,C1,normC,AEq,bEq,AEqAEqtinv,AI_mat,bI,Lambda,d0,d1,D0,D1,idx_diag,idx_offdiag,IDX_upper,IDX_lower,X0,X1,Xq,S0,S1,Sq + return ( C0, C1, @@ -854,6 +860,8 @@ def ADMM_preprocessing(self, C): def compute_fejer_weights(self): """ Evaluate the real Wigner representation blocks used by the discretized Fejer inequality constraints. + + :return: Two sets of block weights and the number of SO(3) grid points. """ SO3_grid = self.discretize_SO3() Ngrid = SO3_grid.shape[0] @@ -896,6 +904,8 @@ def permutek_block(Ak, k): def discretize_SO3(self): """ Construct an approximately uniform Euler-angle grid over SO(3). + + :return: Array of ZYZ Euler angles. """ S2 = saff_kuijlaars(self.S2_grid) S2_size = S2.shape[0] @@ -926,6 +936,13 @@ def discretize_SO3(self): def proximal_refine(self, X_admm, weight, Penalty, r): """ Refine the relaxed solution by iteratively encouraging lower-rank representation matrices. + + :param X_admm: Initial relaxed representation matrices. + :param weight: Degree-dependent refinement weights. + :param Penalty: Penalty value for each refinement step. + :param r: Rank offset for each refinement step. + + :return: Refined representation matrices. """ N = self.n_img C = self.C @@ -1014,7 +1031,12 @@ def euler_est(self): def euler_est_Cm(self, X1, XS): """ - Recover Euler angles from the degree-one and degree-m solutions for cyclic symmetry. + Recover Euler angles for cyclic symmetry. + + :param X1: Relaxed degree-one representation matrix. + :param XS: Relaxed representation matrix at the symmetry order. + + :return: Estimated rotation matrices and Euler angles. """ S = self.n_sym N = self.n_img @@ -1134,7 +1156,11 @@ def find_gamma(Xm, beta, alpha): def euler_est_Dm(self, X_est): """ - Recover Euler angles from the degree-two and degree-m solutions for dihedral symmetry. + Recover Euler angles for dihedral symmetry. + + :param X_est: Relaxed representation matrices. + + :return: Estimated rotation matrices and Euler angles. """ X2 = X_est[1] S = self.sym_grp.order @@ -1288,34 +1314,48 @@ def LS_D(W1, W2, W3, W4, Br, Bi): #################### # Helper Functions # #################### - def transform_coeff(self, A, Lmax, N, IDX_upper): + def transform_coeff(self, A, IDX_upper): """ - Transform representation matrices into the two block-vector forms used by ADMM. + Convert representation matrices to the block-vector form used by ADMM. + + :param A: Representation matrices. + :param IDX_upper: Indices of upper-triangular image pair blocks. + + :return: The two block-vector coefficient arrays. """ d0 = [0] d1 = [0] - for k in range(1, Lmax + 1): + for k in range(1, self.Lmax + 1): d0.append(d0[-1] + k**2) d1.append(d1[-1] + (k + 1) ** 2) - A0 = xp.zeros((d0[-1], N * (N + 1) // 2), dtype=np.float64) - A1 = xp.zeros((d1[-1], N * (N + 1) // 2), dtype=np.float64) - for k in range(1, Lmax + 1): - a0, a1 = self.permutek(A[k - 1], k, N) - A0[d0[k - 1] : d0[k], :] = self.vec_block(a0, N, k, IDX_upper) - A1[d1[k - 1] : d1[k], :] = self.vec_block(a1, N, k + 1, IDX_upper) + A0 = xp.zeros((d0[-1], self.n_img * (self.n_img + 1) // 2), dtype=np.float64) + A1 = xp.zeros((d1[-1], self.n_img * (self.n_img + 1) // 2), dtype=np.float64) + for k in range(1, self.Lmax + 1): + a0, a1 = self.permutek(A[k - 1], k, self.n_img) + A0[d0[k - 1] : d0[k], :] = self.vec_block(a0, self.n_img, k, IDX_upper) + A1[d1[k - 1] : d1[k], :] = self.vec_block(a1, self.n_img, k + 1, IDX_upper) return A0, A1 - def transform_coeff_back(self, A0, A1, Lmax, N, IDX_upper, IDX_lower, idx_offdiag): + def transform_coeff_back(self, A0, A1, IDX_upper, IDX_lower, idx_offdiag): """ - Reconstruct representation matrices from the ADMM block-vector forms. + Reconstruct representation matrices from their ADMM block-vector form. + + :param A0: First block-vector array. + :param A1: Second block-vector array. + :param IDX_upper: Indices of upper-triangular image-pair blocks. + :param IDX_lower: Indices of lower-triangular image-pair blocks. + :param idx_offdiag: Indices of off-diagonal image pairs. + + :return: Reconstructed representation matrices. """ d0 = [0] d1 = [0] - for k in range(1, Lmax + 1): + N = self.n_img + for k in range(1, self.Lmax + 1): d0.append(d0[-1] + k**2) d1.append(d1[-1] + (k + 1) ** 2) A = [] - for k in range(1, Lmax + 1): + for k in range(1, self.Lmax + 1): dk = 2 * k + 1 Ak = xp.zeros((N * dk, N * dk), dtype=np.float64) Ak[: N * k, : N * k] = self.mat_block( @@ -1332,6 +1372,12 @@ def transform_coeff_back(self, A0, A1, Lmax, N, IDX_upper, IDX_lower, idx_offdia def permutek(Ak, k, N): """ Permute and split a degree-k matrix into blocks of sizes k and k + 1. + + :param Ak: Degree-k block matrix. + :param k: Representation degree. + :param N: Number of images. + + :return: The two permuted matrix blocks. """ AkP = xp.copy(Ak) dk = 2 * k + 1 @@ -1359,6 +1405,12 @@ def permutek(Ak, k, N): def permutek_back(Ak, k, N): """ Undo the degree-k block permutation and reconstruct the full matrix. + + :param Ak: Permuted degree-k matrix. + :param k: Representation degree. + :param N: Number of images. + + :return: Matrix in the original block ordering. """ dk = 2 * k + 1 Pk = xp.eye(N * dk, dtype=Ak.dtype) @@ -1407,13 +1459,17 @@ def largest_eigenvalue(AI, Ngrid, N): logger.info("Largest eigenvalue of AIAIT is approximately %1.2f" % Lambda) return Lambda - def compute_rank(self, Lmax): + def compute_rank(self): """ Compute the ranks and matrices of the symmetry-averaging projectors at each degree. + + :param Lmax: Maximum representation degree. + + :return: Ranks and symmetry-averaging matrices for each degree. """ - rk = xp.zeros(Lmax, dtype=np.float64) + rk = xp.zeros(self.Lmax, dtype=np.float64) A = [] - for k in range(1, Lmax + 1): + for k in range(1, self.Lmax + 1): Ak = np.sum(self.WD(k, self.sym_euler), axis=0) Ak = np.round(Ak / self.n_sym, 6) A.append(Ak) From 0692245c214811df4a0ea28ce791327ddac11264 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 15 Jun 2026 15:21:34 -0400 Subject: [PATCH 038/161] remove commented out code --- src/aspire/abinitio/commonline_nug.py | 19 ++++--------------- 1 file changed, 4 insertions(+), 15 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 61bcaabdab..5d01f8cc77 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -333,9 +333,6 @@ def admm_sym_J(self, C, verbose): rank_Ak, _ = self.compute_rank() logger.info(f"Rank of Ak: {rank_Ak}") - # rank_Ak=cp.zeros(Lmax) - # for ell in range(Lmax): rank_Ak[ell]=np.linalg.matrix_rank(Ak(ell+1,sym_euler)) - AE = [] AEAETinv = [] for k in range(1, Lmax + 1): @@ -414,12 +411,9 @@ def fun_AI(X0, X1): tmp = xp.concatenate((X0, X1), axis=0) z[:, idx_diag] = AI_mat_diag @ tmp[:, idx_diag] z[:, idx_offdiag] = AI_mat_offdiag @ tmp[:, idx_offdiag] - # return AI_mat@xp.concatenate((X0,X1),axis=0) return z def fun_AIT(yI): - # Z=AI_mat.T@yI - # return Z[:D0,:], Z[D0:,:] Z = xp.zeros((D0 + D1, N * (N + 1) // 2), dtype=np.float64) Z[:, idx_diag] = AI_mat_diag.T @ yI[:, idx_diag] Z[:, idx_offdiag] = AI_mat_offdiag.T @ yI[:, idx_offdiag] @@ -454,7 +448,6 @@ def update_S(C0, C1, yE, yEq, yI, X0, X1, Xd0, Xd1, Xq, rho, Lmax, N): toc0 = time.perf_counter() Time[0] += toc0 - tic0 - # See if we can switch everything to C order before here. Sd0 = Sd0.T Sd1 = Sd1.T tic1 = time.perf_counter() @@ -590,7 +583,6 @@ def print_updates(verbose): idx_offdiag, ) res_psdX += np.linalg.norm(self.psd_projection(-tmp)) - # res_psdX+=norm(self.psd_projection(-tmp))/(1+norm(tmp)) tmp = self.mat_block( X1[d1[k - 1] : d1[k], :], N, @@ -600,7 +592,6 @@ def print_updates(verbose): idx_offdiag, ) res_psdX += np.linalg.norm(self.psd_projection(-tmp)) - # res_psdX+=norm(self.psd_projection(-tmp))/(1+norm(tmp)) res_psdX = res_psdX / (1 + np.linalg.norm(X0) + np.linalg.norm(X1)) res_psdD = 0 for k in range(1, Lmax + 1): @@ -613,13 +604,11 @@ def print_updates(verbose): res_psdD += np.linalg.norm( self.psd_projection(-tmp), axis=(-2, -1) ).sum() - # res_psdD+=norm(self.psd_projection(-tmp))/(1+norm(tmp)) res_psdD = res_psdD / (1 + np.linalg.norm(Xd0) + np.linalg.norm(Xd1)) res_psdQ = 0 for count in range(N * (N - 1) // 2): tmp = Xq[:, count].reshape(4, 4).T res_psdQ += np.linalg.norm(self.psd_projection(-tmp)) - # res_psdQ+=norm(self.psd_projection(-tmp))/(1+norm(tmp)) res_psdQ = res_psdQ / (1 + np.linalg.norm(Xq)) normS = np.sqrt( @@ -773,7 +762,7 @@ def ADMM_preprocessing(self, C): * (k + 0.5) * W1[k - 1][p].T.reshape(-1) ) - # this needs double checking + # this needs double checking (Ruiyi) AI_mat_offdiag[p, : d0[-1]] = w0 AI_mat_offdiag[p, d0[-1] :] = w1 @@ -805,7 +794,7 @@ def ADMM_preprocessing(self, C): * (k + 0.5) * (0.5 * W1[k - 1][p] + 0.5 * W1[k - 1][p].T).T.reshape(-1) ) - # this needs double checking + # this needs double checking (Ruiyi) AI_mat_diag[p, : d0[-1]] = w0 AI_mat_diag[p, d0[-1] :] = w1 AI_mat_diag = xp.asarray(AI_mat_diag) / 1 @@ -1093,10 +1082,10 @@ def find_alpha(X1): for j in range(N): z = X1[3 * i : 3 * (i + 1), 3 * j : 3 * (j + 1)][0, 0] ZZbar[i, j] = z / abs(z) - # ZZbar[i,j]=z*2/sin(beta[i])/sin(beta[j]) + z = X1[3 * i : 3 * (i + 1), 3 * j : 3 * (j + 1)][0, 2] ZZ[i, j] = -z / abs(z) - # ZZ[i,j]=-z*2/sin(beta[i])/sin(beta[j]) + evals, evecs = np.linalg.eigh(ZZbar) idx = np.argmax(abs(evals)) Z = evecs[:, idx] * np.sqrt(abs(evals[idx])) From b9472289f7298a4af58cb0d1149efb8ba3b77326 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 17 Jun 2026 14:46:34 -0400 Subject: [PATCH 039/161] Test g_sync for clean case --- tests/test_commonline_utils.py | 35 +++++++++++++++++++++++++++------- 1 file changed, 28 insertions(+), 7 deletions(-) diff --git a/tests/test_commonline_utils.py b/tests/test_commonline_utils.py index 818d829c6a..a511ab4d18 100644 --- a/tests/test_commonline_utils.py +++ b/tests/test_commonline_utils.py @@ -144,24 +144,45 @@ def test_g_sync(symmetry): axes /= np.linalg.norm(axes, axis=1, keepdims=True) angles = np.random.uniform(0, 2 * np.deg2rad(target_mean_deg), n).astype(dtype) delta_rots = Rotation.from_rotvec(axes * angles[:, None], dtype=dtype) - perturbed_rots = Rotation(delta_rots.matrices @ gt_rots.matrices) + noisy_rots = Rotation(delta_rots.matrices @ gt_rots.matrices) # Get mean ang dist for aligned estimates # and check we're close to target. - og_maad = mean_aligned_angular_distance(perturbed_rots, gt_rots) + og_maad = mean_aligned_angular_distance(noisy_rots, gt_rots) np.testing.assert_array_less(abs(og_maad - target_mean_deg), 0.2) - # Simulate symmetry desynchronization + # Simulate symmetry desynchronization for clean and noisy case. g_idx = np.random.randint(len(gs), size=n) - desynced_rots = Rotation(gs[g_idx] @ perturbed_rots.matrices) + desynced_noisy_rots = Rotation(gs[g_idx] @ noisy_rots) + desynced_clean_rots = Rotation(gs[g_idx] @ gt_rots.matrices) + + # Apply a global rotation to the noisy rotations to + # to simulate a set of estimated rotations + desynced_noisy_rots = ( + Rotation.generate_random_rotations(1, dtype=dtype).matrices + @ desynced_noisy_rots + ) # Mean aligned angular distance of unsynced rots should be bad np.testing.assert_array_less( - 10 * og_maad, mean_aligned_angular_distance(desynced_rots, gt_rots) + 10 * og_maad, mean_aligned_angular_distance(desynced_noisy_rots, gt_rots) ) # Perform g_sync and check that mean aligned angular distance # matches ground truth MAAD to within .1 degrees. - rots_gt_synced = g_sync(desynced_rots, gt_rots, symmetry) - est_maad = mean_aligned_angular_distance(desynced_rots, rots_gt_synced) + rots_gt_synced_to_noisy = g_sync(desynced_noisy_rots, gt_rots, symmetry) + est_maad = mean_aligned_angular_distance( + desynced_noisy_rots, rots_gt_synced_to_noisy + ) np.testing.assert_array_less(abs(og_maad - est_maad), 0.1) + + # For the clean case the synced rotations should match allclose up to + # a global multiplication by one of the symmetry group elements. + gt_rots_synced_to_clean = g_sync(desynced_clean_rots, gt_rots, symmetry) + errs = np.linalg.norm( + gs @ gt_rots_synced_to_clean[0] - desynced_clean_rots[0], axis=(-2, -1) + ) + best_g = np.argmin(errs) + np.testing.assert_allclose( + gs[best_g] @ gt_rots_synced_to_clean, desynced_clean_rots + ) From 0d813ab04f88225c5fb24ab543a6733f360d5875 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 23 Jun 2026 12:11:39 -0400 Subject: [PATCH 040/161] xp in admm --- src/aspire/abinitio/commonline_nug.py | 26 +++++++++++++------------- tests/test_nug.py | 1 + 2 files changed, 14 insertions(+), 13 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 5d01f8cc77..6f25c6aaef 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -348,7 +348,7 @@ def admm_sym_J(self, C, verbose): AEk[count, count - 1] = 1 AEk[count, count - 1 + s1 + s0] = 1 AE.append(AEk) - AEAETinv.append(np.linalg.pinv(AEk @ AEk.T)) + AEAETinv.append(xp.linalg.pinv(AEk @ AEk.T)) bE = xp.zeros((Lmax + D0 + D1), dtype=np.float64) for k in range(Lmax): bE[k + d0[k] + d1[k] :] = rank_Ak[k] @@ -358,7 +358,7 @@ def admm_sym_J(self, C, verbose): bE[k + 1 + d0[k + 1] + d1[k] : k + 1 + d0[k + 1] + d1[k + 1]] = xp.eye( k + 2 ).T.reshape(-1) - bE = np.repeat(bE[:, np.newaxis], N, axis=1) + bE = xp.repeat(bE[:, np.newaxis], N, axis=1) P = [] for k in range(1, Lmax + 1): dk = 2 * k + 1 @@ -499,7 +499,7 @@ def update_yI(C0, C1, X0, X1, S0, S1, yE, yEq, yI, rho, Lambda): -X0 / rho + C0 - S0 - Z0 - AIT_yI0, -X1 / rho + C1 - S1 - Z1 - AIT_yI1 ) yI = yI + bI / rho / Lambda + tmp / Lambda - yI = np.maximum(yI, 0) + yI = xp.maximum(yI, 0) return yI def update_X( @@ -509,35 +509,35 @@ def update_X( AIT_yI0, AIT_yI1 = fun_AIT(yI) tmp = S0 + Z0 + AIT_yI0 - C0 X0 = X0 + mult * rho * tmp - resX0 = np.linalg.norm(tmp) + resX0 = xp.linalg.norm(tmp) tmp = S1 + Z1 + AIT_yI1 - C1 X1 = X1 + mult * rho * tmp - resX1 = np.linalg.norm(tmp) + resX1 = xp.linalg.norm(tmp) tmp = Sd0 + Zd0 Xd0 = Xd0 + mult * rho * tmp - resXd0 = np.linalg.norm(tmp) + resXd0 = xp.linalg.norm(tmp) tmp = Sd1 + Zd1 Xd1 = Xd1 + mult * rho * tmp - resXd1 = np.linalg.norm(tmp) + resXd1 = xp.linalg.norm(tmp) tmp = Sq + Zq Xq = Xq + mult * rho * tmp - resXq = np.linalg.norm(tmp) + resXq = xp.linalg.norm(tmp) return ( X0, X1, Xd0, Xd1, Xq, - np.sqrt(resX0**2 + resX1**2 + resXd0**2 + resXd1**2 + resXq**2), + xp.sqrt(resX0**2 + resX1**2 + resXd0**2 + resXd1**2 + resXq**2), ) def update_rho(X0, X1, Xd0, Xd1, Xq, bE, bEq, bI, res_X, rho, factor, normC): z, zq = fun_AE(X0, X1, Xd0, Xd1, Xq) - res_eq = np.linalg.norm(z - bE) / (1 + np.linalg.norm(bE)) + np.linalg.norm( + res_eq = xp.linalg.norm(z - bE) / (1 + xp.linalg.norm(bE)) + xp.linalg.norm( zq - bEq - ) / (1 + np.linalg.norm(bEq)) - res_inq = np.linalg.norm(np.maximum(bI - fun_AI(X0, X1), 0)) / ( - 1 + abs(bI) * np.sqrt(Ngrid * N**2) + ) / (1 + xp.linalg.norm(bEq)) + res_inq = xp.linalg.norm(xp.maximum(bI - fun_AI(X0, X1), 0)) / ( + 1 + abs(bI) * xp.sqrt(Ngrid * N**2) ) p_resnorm = res_eq + res_inq d_resnorm = res_X / (1 + normC) diff --git a/tests/test_nug.py b/tests/test_nug.py index b9699e7f66..2297f7a362 100644 --- a/tests/test_nug.py +++ b/tests/test_nug.py @@ -90,6 +90,7 @@ def orient_est(source, proximal_refine): max_shift=max_shift, shift_step=shift_step, perform_pr=proximal_refine, + verbose=False, ) orient_est.estimate_rotations() return orient_est From 77b23d6115c1bdb9e6e601dac9040a21d385dee3 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 23 Jun 2026 15:26:46 -0400 Subject: [PATCH 041/161] raise on unsupported symmetry --- src/aspire/abinitio/commonline_nug.py | 8 +++++++- tests/test_nug.py | 10 +++++++++- 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 6f25c6aaef..8120929d0a 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -101,6 +101,12 @@ def __init__( ) logger.info(f"Using provided symmetry: {symmetry}") self.sym_grp = SymmetryGroup.parse(symmetry) + + if not isinstance(self.sym_grp, (CnSymmetryGroup, DnSymmetryGroup)): + raise ValueError( + f"This algorithm supports cyclic or dihedral symmetry. Found {str(self.sym_grp)}." + ) + self.sym_euler = self.sym_grp.rotations.angles self.n_sym = len(self.sym_euler) @@ -722,7 +728,7 @@ def ADMM_preprocessing(self, C): for k in range(Lmax): C[k] = xp.asarray(Xnorm / Cnorm * C[k]) C0, C1 = self.transform_coeff(C, IDX_upper) - normC = np.sqrt(np.linalg.norm(C0) ** 2 + np.linalg.norm(C1) ** 2) + normC = xp.sqrt(xp.linalg.norm(C0) ** 2 + xp.linalg.norm(C1) ** 2) del C # compute the block sizes for X diff --git a/tests/test_nug.py b/tests/test_nug.py index 2297f7a362..b5f1cc0abf 100644 --- a/tests/test_nug.py +++ b/tests/test_nug.py @@ -4,7 +4,7 @@ from aspire.abinitio import CommonlineNUG, compare_rots_sym, g_sync from aspire.source import Simulation from aspire.utils import mean_aligned_angular_distance -from aspire.volume import CnSymmetricVolume, DnSymmetricVolume +from aspire.volume import CnSymmetricVolume, DnSymmetricVolume, TSymmetricVolume DTYPE = [np.float64, pytest.param(np.float32, marks=pytest.mark.expensive)] RESOLUTION = [48, pytest.param(49, marks=pytest.mark.expensive)] @@ -130,3 +130,11 @@ def test_estimate_rotations(orient_est): orient_est.rotations, orient_est.src.rotations, orient_est.sym_grp ) mean_aligned_angular_distance(orient_est.rotations, gt_rots_synced, 8.0) + + +def test_unspupported_symmetry_raises(dtype): + vol = TSymmetricVolume(L=16, C=1, K=10, dtype=dtype).generate() + src = Simulation(n=3, vols=vol) + + with pytest.raises(ValueError, match="supports cyclic or dihedral symmetry"): + _ = CommonlineNUG(src) From 9e3c27ceda2da30fe94527bc458ce78247d45e68 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 23 Jun 2026 16:10:38 -0400 Subject: [PATCH 042/161] mark tests expensive. Add smoke test. --- tests/test_nug.py | 44 +++++++++++++++++++++++++++++++++++++++----- 1 file changed, 39 insertions(+), 5 deletions(-) diff --git a/tests/test_nug.py b/tests/test_nug.py index b5f1cc0abf..5c27969c53 100644 --- a/tests/test_nug.py +++ b/tests/test_nug.py @@ -6,16 +6,16 @@ from aspire.utils import mean_aligned_angular_distance from aspire.volume import CnSymmetricVolume, DnSymmetricVolume, TSymmetricVolume -DTYPE = [np.float64, pytest.param(np.float32, marks=pytest.mark.expensive)] -RESOLUTION = [48, pytest.param(49, marks=pytest.mark.expensive)] +DTYPE = [np.float64, np.float32] +RESOLUTION = [48, 49] N_IMG = [15] -OFFSETS = [0, pytest.param(None, marks=pytest.mark.expensive)] -ORDER = [3, pytest.param(4, marks=pytest.mark.expensive)] +OFFSETS = [0, None] +ORDER = [3, 4] PR = [False] SEED = 1980 VOLUME = [ CnSymmetricVolume, - pytest.param(DnSymmetricVolume, marks=pytest.mark.expensive), + DnSymmetricVolume, ] @@ -101,6 +101,38 @@ def orient_est(source, proximal_refine): ######### +def test_smoke_nug(dtype, Volume): + """ + Perform quick smoke test since other tests are long running. + """ + vol = Volume(L=32, order=3, C=1, dtype=dtype, seed=SEED).generate() + + src = Simulation( + n=20, + vols=vol, + offsets=0, + amplitudes=1, + seed=SEED, + ).cache() + + orient_est = CommonlineNUG( + src, + Lmax=4, + T=5, + max_iter=10, + S2_grid=50, + max_shift=0, + mask=False, + verbose=False, + ) + + rots = orient_est.estimate_rotations() + + assert rots.shape == (src.n, 3, 3) + assert rots.dtype == dtype + + +@pytest.mark.expensive def test_dtypes(orient_est): """ Check dtypes for each major step of the algorithm. @@ -117,6 +149,7 @@ def test_dtypes(orient_est): assert orient_est.rotations.dtype == orient_est.dtype +@pytest.mark.expensive def test_estimate_rotations_pairwise(orient_est): """ """ MSE = compare_rots_sym( @@ -125,6 +158,7 @@ def test_estimate_rotations_pairwise(orient_est): np.testing.assert_array_less(MSE, 0.1) +@pytest.mark.expensive def test_estimate_rotations(orient_est): gt_rots_synced = g_sync( orient_est.rotations, orient_est.src.rotations, orient_est.sym_grp From 36a5dbba56f047cacff1024761c14b76bb8a7ba6 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 24 Jun 2026 15:48:19 -0400 Subject: [PATCH 043/161] tunable proximal refinement iterations --- src/aspire/abinitio/commonline_nug.py | 29 +++++++++++++++------------ tests/test_nug.py | 24 ++++++++++++++++------ 2 files changed, 34 insertions(+), 19 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 8120929d0a..1132fafd42 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -38,7 +38,7 @@ def __init__( mult=1.5, S2_grid=441, Nstep_yI=10, - perform_pr=False, + pr_iters=None, verbose=True, **kwargs, ): @@ -62,7 +62,8 @@ def __init__( :param mult: Step-size multiplier for the ADMM primal update. :param S2_grid: Number of sphere samples used to discretize SO(3). :param Nstep_yI: Number of inequality-multiplier updates per ADMM iteration. - :param perform_pr: Whether to apply proximal refinement after ADMM. + :param pr_iters: Number of proximal refinement iterations. Default of None + does not perform proximal refinement. Recommended value is 4. :param verbose: Whether to log ADMM progress. """ @@ -85,7 +86,6 @@ def __init__( self.mult = mult self.S2_grid = S2_grid self.Nstep_yI = Nstep_yI - self.perform_pr = perform_pr self.verbose = verbose # Handle symmetry @@ -110,6 +110,13 @@ def __init__( self.sym_euler = self.sym_grp.rotations.angles self.n_sym = len(self.sym_euler) + # Set up proximal refinement terms + if pr_iters is not None: + self.pr_weights = 1 / (1 + np.arange(self.Lmax)) + self.pr_penalty = [1] * pr_iters + self.pr_rank = list(range(pr_iters - 1, -1, -1)) # [pr_iters - 1,..., 0] + self.pr_iters = pr_iters + self._build_full_pft() def _build_full_pft(self): @@ -276,15 +283,12 @@ def perform_admm(self): """ X_est = self.admm_sym_J(self.C, self.verbose) - if self.perform_pr: - weight = 1 / (1 + np.arange(self.Lmax)) - Penalty = [1, 1, 1, 1] - r = [3, 2, 1, 0] + if self.pr_iters is not None: X_est = self.proximal_refine( X_est, - weight, - Penalty, - r, + self.pr_weights, + self.pr_penalty, + self.pr_rank, ) self.X_est = X_est @@ -980,11 +984,10 @@ def low_rank_proj(X, r_step): Xproj.append(tmp) return Xproj - Niter = len(r) CC = [None] * self.Lmax current = [np.copy(Xk) for Xk in X_admm] - for step in range(Niter): + for step in range(self.pr_iters): X_proj = low_rank_proj(current, r[step]) for k in range(self.Lmax): @@ -999,7 +1002,7 @@ def low_rank_proj(X, r_step): logger.info( "Proximal refine step %d/%d: relative update %.3e", step + 1, - Niter, + self.pr_iters, rel_change(X_next, current), ) diff --git a/tests/test_nug.py b/tests/test_nug.py index 5c27969c53..c68f70fbe3 100644 --- a/tests/test_nug.py +++ b/tests/test_nug.py @@ -11,7 +11,7 @@ N_IMG = [15] OFFSETS = [0, None] ORDER = [3, 4] -PR = [False] +PR = [None] SEED = 1980 VOLUME = [ CnSymmetricVolume, @@ -89,7 +89,8 @@ def orient_est(source, proximal_refine): source, max_shift=max_shift, shift_step=shift_step, - perform_pr=proximal_refine, + max_iter=201, + pr_iters=proximal_refine, verbose=False, ) orient_est.estimate_rotations() @@ -101,7 +102,7 @@ def orient_est(source, proximal_refine): ######### -def test_smoke_nug(dtype, Volume): +def test_smoke(dtype, Volume): """ Perform quick smoke test since other tests are long running. """ @@ -122,7 +123,6 @@ def test_smoke_nug(dtype, Volume): max_iter=10, S2_grid=50, max_shift=0, - mask=False, verbose=False, ) @@ -151,7 +151,11 @@ def test_dtypes(orient_est): @pytest.mark.expensive def test_estimate_rotations_pairwise(orient_est): - """ """ + """ + Check mean squared error between estimates and ground truth + pairwise rotations, Rij. This serves as a reference to the error + metric used by the researcher in the related publication. + """ MSE = compare_rots_sym( orient_est.rotations, orient_est.src.rotations, orient_est.sym_grp ) @@ -160,13 +164,21 @@ def test_estimate_rotations_pairwise(orient_est): @pytest.mark.expensive def test_estimate_rotations(orient_est): + """ + Check that the mean angular distance between estimates and ground + truth, after symmetry synchronization and global alignment, are + within 10 degrees. + """ gt_rots_synced = g_sync( orient_est.rotations, orient_est.src.rotations, orient_est.sym_grp ) - mean_aligned_angular_distance(orient_est.rotations, gt_rots_synced, 8.0) + mean_aligned_angular_distance(orient_est.rotations, gt_rots_synced, 10.0) def test_unspupported_symmetry_raises(dtype): + """ + Check that we raise for symmetries other than Cn/Dn. + """ vol = TSymmetricVolume(L=16, C=1, K=10, dtype=dtype).generate() src = Simulation(n=3, vols=vol) From d1de07ff69ba0581ab0cb36130467b553d5e6e6b Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 25 Jun 2026 09:11:38 -0400 Subject: [PATCH 044/161] code coverage for proximal refinement and print_update computations. --- tests/test_nug.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_nug.py b/tests/test_nug.py index c68f70fbe3..d3fcb6de83 100644 --- a/tests/test_nug.py +++ b/tests/test_nug.py @@ -121,9 +121,10 @@ def test_smoke(dtype, Volume): Lmax=4, T=5, max_iter=10, + pr_iters=1, S2_grid=50, max_shift=0, - verbose=False, + verbose=True, ) rots = orient_est.estimate_rotations() From ecfbfb95dfc163d718daaec015176c36d79bedba Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 25 Jun 2026 09:58:23 -0400 Subject: [PATCH 045/161] test compare_rots_sym --- tests/test_commonline_utils.py | 29 ++++++++++++++++++++++++++++- 1 file changed, 28 insertions(+), 1 deletion(-) diff --git a/tests/test_commonline_utils.py b/tests/test_commonline_utils.py index a511ab4d18..df8edfcad2 100644 --- a/tests/test_commonline_utils.py +++ b/tests/test_commonline_utils.py @@ -1,7 +1,7 @@ import numpy as np import pytest -from aspire.abinitio import JSync, g_sync +from aspire.abinitio import JSync, compare_rots_sym, g_sync from aspire.abinitio.commonline_utils import ( _complete_third_row_to_rot, _estimate_third_rows, @@ -186,3 +186,30 @@ def test_g_sync(symmetry): np.testing.assert_allclose( gs[best_g] @ gt_rots_synced_to_clean, desynced_clean_rots ) + + +@pytest.mark.parametrize("symmetry", ["C3", "C4", "D3", "D4", "T", "O"]) +def test_compare_rots_sym(symmetry): + """ + This method find the mean squared error between all pairs of rotations + taking into account each rotation being multiplied by an arbitrary + symmetry group element. In this test we check that a set of rotations + multiplied by random symmetry group elements gives a zero MSE when compared + to the original set. + """ + n = 100 + dtype = np.float64 + + # Get symmetry group matrices + gs = SymmetryGroup.parse(symmetry).matrices + + # Build set of ground truth rotations + gt_rots = Rotation.generate_random_rotations(n, dtype=dtype).matrices + + # Multiply by random group elements + g_idx = np.random.randint(len(gs), size=n) + rots_with_sym = gs[g_idx] @ gt_rots + + # Check MSE is zero + MSE = compare_rots_sym(rots_with_sym, gt_rots, symmetry) + np.testing.assert_allclose(MSE, 0, atol=np.finfo(dtype).eps) From 5c2a915add21299970945932deafae7de5f2701f Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 25 Jun 2026 11:10:25 -0400 Subject: [PATCH 046/161] test symmetry handling log messages --- tests/test_nug.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/tests/test_nug.py b/tests/test_nug.py index d3fcb6de83..5ae3fa4025 100644 --- a/tests/test_nug.py +++ b/tests/test_nug.py @@ -185,3 +185,28 @@ def test_unspupported_symmetry_raises(dtype): with pytest.raises(ValueError, match="supports cyclic or dihedral symmetry"): _ = CommonlineNUG(src) + + +def test_symmetry_logging(caplog): + """ + Check expected log messages for mismatch or unprovided symmetry. + """ + vol = CnSymmetricVolume(L=16, order=3, C=1, dtype=np.float64, seed=SEED).generate() + src = Simulation(n=3, vols=vol).cache() + + # Check unprovided symmetry populates with source symmetry and logs messsage + caplog.clear() + with caplog.at_level("INFO"): + orient_est = CommonlineNUG(src) + + assert str(orient_est.sym_grp) == "C3" + assert "Symmetry not provided. Using Source symmetry: C3" in caplog.text + + # Check we use provided symmetry on mismatch, with log + caplog.clear() + with caplog.at_level("INFO"): + orient_est = CommonlineNUG(src, symmetry="D3") + + assert str(orient_est.sym_grp) == "D3" + assert "Provided symmetry, D3, does not match source, C3" in caplog.text + assert "Using provided symmetry: D3" in caplog.text From 552ee514d9e9cb57c5a2ce1bc33f6b6a4a139cc4 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 25 Jun 2026 13:30:36 -0400 Subject: [PATCH 047/161] remove singleton handling for block transformations --- src/aspire/abinitio/commonline_nug.py | 36 +++------------------------ 1 file changed, 3 insertions(+), 33 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 1132fafd42..e39c82dedf 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1544,55 +1544,25 @@ def psd_projection(B): return (evecs * evals[..., None, :]) @ evecs.swapaxes(-1, -2) @staticmethod - def transform_block(A, k, Pk=None): + def transform_block(A, k, Pk): """ Permute and vectorize the two invariant blocks of a degree-k matrix. """ - single = A.ndim == 2 - if single: - A = A[None, :, :] - - if Pk is None: - dk = 2 * k + 1 - Pk = xp.eye(dk, dtype=A.dtype) - for m in range(k): - for el in range(k - m): - Pk[(m + 2 * el, m + 2 * el + 1), :] = Pk[ - (m + 2 * el + 1, m + 2 * el), : - ] AT = Pk @ A @ Pk.T A0 = AT[:, :k, :k].swapaxes(-1, -2).reshape(A.shape[0], -1) A1 = AT[:, k:, k:].swapaxes(-1, -2).reshape(A.shape[0], -1) - - if single: - return A0[0], A1[0] - return A0, A1 @staticmethod - def transform_back_block(A0, A1, k, Pk=None): + def transform_back_block(A0, A1, k, Pk): """ Reconstruct a degree-k matrix from its two invariant block vectors. """ dk = 2 * k + 1 - single = A0.ndim == 1 - - if single: - A0 = A0[None, :] - A1 = A1[None, :] - A = xp.zeros((A0.shape[0], dk, dk), dtype=A0.dtype) A[:, :k, :k] = A0.reshape(-1, k, k).swapaxes(-1, -2) A[:, k:, k:] = A1.reshape(-1, k + 1, k + 1).swapaxes(-1, -2) - if Pk is None: - Pk = xp.eye(dk, dtype=A0.dtype) - for m in range(k): - for el in range(k - m): - Pk[(m + 2 * el, m + 2 * el + 1), :] = Pk[ - (m + 2 * el + 1, m + 2 * el), : - ] - out = Pk.T @ A @ Pk - return out[0] if single else out + return Pk.T @ A @ Pk def construct_AEq(self): """ From 87bcbff3fa3fbaf5e40cb7514974cfd7a8b1b7f5 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 25 Jun 2026 15:40:16 -0400 Subject: [PATCH 048/161] fij loss funtion docstring --- src/aspire/abinitio/commonline_nug.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index e39c82dedf..a9de66bc61 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -158,6 +158,23 @@ def compute_coeff(self): T = self.T def fij(alpha, gamma, i, j): + """ + Evaluate the pairwise common-line loss used to approximate NUG Fourier coefficients. + + For a relative orientation parameterized by ZYZ Euler angles, the common-line + loss f_ij depends only on the first and third angles, alpha and gamma. + This function samples the corresponding polar Fourier rays from images i and j, + compares them over all candidate 1D shifts, and returns the minimum shifted + L1 mismatch. + + :param alpha: First ZYZ Euler angle of the relative orientation, in radians. + :param gamma: Third ZYZ Euler angle of the relative orientation, in radians. + :param i: Index of the first image. + :param j: Index of the second image. + + :return: Minimum shifted L1 mismatch between the induced common-line rays. + """ + Ii_hat = self.pf_full[i] Ij_hat = self.pf_full[j] idxi = np.round((alpha - np.pi / 2) * n_theta / 2 / np.pi) % n_theta From 30c307aeea213b26f9c3045045594bf34b4d9f7f Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 29 Jun 2026 10:02:20 -0400 Subject: [PATCH 049/161] Add Defaults to param descriptions --- src/aspire/abinitio/commonline_nug.py | 41 ++++++++++++++++++--------- 1 file changed, 27 insertions(+), 14 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index a9de66bc61..6f74d9550f 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -43,28 +43,41 @@ def __init__( **kwargs, ): """ - Initialize the symmetric NUG orientation estimator. + Initialize the symmetric NUG orientation estimator. All default values match those used + for publication results, with exception of `pr_iters`, which controls how many iterations + of proximal refinement to be used. The default of None runs the algorithm without proximal + refinement. Set to `pr_iters` to 4 to match the proximal refinement workflow in the related + publication. :param src: Source containing the input projection images. :param symmetry: Cyclic or dihedral symmetry specification, such as 'C3' or 'D4'. If omitted, uses the source symmetry. - :param n_rad: Number of radial samples in the polar Fourier transform. - :param n_theta: Number of angular samples in the polar Fourier transform. - :param max_shift: Maximum shift considered when comparing common lines. - :param shift_step: Sampling interval for candidate shifts. - :param mask: Whether to apply a circular mask to the input images. + :param n_rad: Number of radial samples in the polar Fourier transform. If None, + n_rad will default to the ceiling of half the resolution of the source. + :param n_theta: Number of angular samples in the polar Fourier transform. This + value must be even. Default is 360. + :param max_shift: Determines maximum range for shifts for common-line detection + as a proportion of the resolution. Default is 0.15. + :param shift_step:Resolution of shift estimation common-line detection in pixels. + Default is 1 pixel. + :param mask: Option to mask `src.images` with a fuzzy mask (boolean). + Default, `True`, applies a mask. :param Lmax: Maximum Wigner representation degree used in the relaxation. + Default is 12. :param T: Quadrature resolution used to compute the Fourier coefficients. - :param max_iter: Number of ADMM iterations. - :param rho: Initial ADMM penalty parameter. - :param ratio: Residual ratio used when updating the ADMM penalty. - :param factor: Scaling factor used when updating the ADMM penalty. - :param mult: Step-size multiplier for the ADMM primal update. - :param S2_grid: Number of sphere samples used to discretize SO(3). + Default is 36. + :param max_iter: Number of ADMM iterations. Default is 501. + :param rho: Initial ADMM penalty parameter. Default 0.05. + :param ratio: Residual ratio used when updating the ADMM penalty. Default is 1. + :param factor: Scaling factor used when updating the ADMM penalty. Default is 1.0. + :param mult: Step-size multiplier for the ADMM primal update. Default is 1.5. + :param S2_grid: Number of sphere samples used to discretize SO(3). Default is 441. :param Nstep_yI: Number of inequality-multiplier updates per ADMM iteration. + Default is 10. :param pr_iters: Number of proximal refinement iterations. Default of None - does not perform proximal refinement. Recommended value is 4. - :param verbose: Whether to log ADMM progress. + does not perform proximal refinement. Recommended value when using is 4. + :param verbose: Whether to log ADMM and proximal refinement progress. + Default is True. """ super().__init__( From 8e4990bc628e314e7263cf745a92e92aa6130998 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 29 Jun 2026 11:22:30 -0400 Subject: [PATCH 050/161] add vectorized computation. some cleanup. --- src/aspire/abinitio/commonline_nug.py | 106 ++++++++++---------------- 1 file changed, 39 insertions(+), 67 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 6f74d9550f..789802d78e 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -101,7 +101,8 @@ def __init__( self.Nstep_yI = Nstep_yI self.verbose = verbose - # Handle symmetry + # Handle symmetry. + # Get from source if not provided. Warn on mismatch if provided. if symmetry is None: logger.info( f"Symmetry not provided. Using Source symmetry: {str(self.src.symmetry_group)}" @@ -174,8 +175,6 @@ def fij(alpha, gamma, i, j): """ Evaluate the pairwise common-line loss used to approximate NUG Fourier coefficients. - For a relative orientation parameterized by ZYZ Euler angles, the common-line - loss f_ij depends only on the first and third angles, alpha and gamma. This function samples the corresponding polar Fourier rays from images i and j, compares them over all candidate 1D shifts, and returns the minimum shifted L1 mismatch. @@ -398,7 +397,7 @@ def admm_sym_J(self, C, verbose): bE[k + 1 + d0[k + 1] + d1[k] : k + 1 + d0[k + 1] + d1[k + 1]] = xp.eye( k + 2 ).T.reshape(-1) - bE = xp.repeat(bE[:, np.newaxis], N, axis=1) + bE = xp.repeat(bE[:, None], N, axis=1) P = [] for k in range(1, Lmax + 1): dk = 2 * k + 1 @@ -781,65 +780,41 @@ def ADMM_preprocessing(self, C): bEq = xp.zeros(17, dtype=np.float64) bEq[:16] = xp.eye(4, dtype=np.float64).reshape(-1) / 4 bEq[-1] = 1 - bEq = xp.repeat(bEq[:, xp.newaxis], N * (N - 1) // 2, axis=1) + bEq = xp.repeat(bEq[:, None], N * (N - 1) // 2, axis=1) - # AI and bI + # Compute AI and bI W0, W1, Ngrid = self.compute_fejer_weights() - AI_mat_offdiag = np.zeros((Ngrid, D0 + D1), dtype=np.float64) - for p in range(Ngrid): - w0 = np.zeros(D0, dtype=np.float64) - w1 = np.zeros(D1, dtype=np.float64) - for k in range(1, Lmax + 1): - w0[d0[k - 1] : d0[k]] = ( - (Lmax - k + 2) - * (Lmax - k + 1) - * (k + 0.5) - * W0[k - 1][p].T.reshape(-1) - ) - w1[d1[k - 1] : d1[k]] = ( - (Lmax - k + 2) - * (Lmax - k + 1) - * (k + 0.5) - * W1[k - 1][p].T.reshape(-1) - ) - # this needs double checking (Ruiyi) - AI_mat_offdiag[p, : d0[-1]] = w0 - AI_mat_offdiag[p, d0[-1] :] = w1 - # Vectorized version, added by Josh - # AI_mat_offdiag_new = np.zeros((Ngrid, D0 + D1)) - # for k in range(1, Lmax + 1): - # scale = (Lmax - k + 2) * (Lmax - k + 1) * (k + 0.5) + # AI_mat_offdiag computation has been vectorized, but originally + # had the note by block1: this needs double checking (Ruiyi) + AI_mat_offdiag = np.zeros((Ngrid, D0 + D1), dtype=np.float64) + for k in range(1, Lmax + 1): + scale = (Lmax - k + 2) * (Lmax - k + 1) * (k + 0.5) - # block0 = scale * W0[k - 1].transpose(0, 2, 1).reshape(Ngrid, -1) - # block1 = scale * W1[k - 1].transpose(0, 2, 1).reshape(Ngrid, -1) + block0 = scale * W0[k - 1].transpose(0, 2, 1).reshape(Ngrid, -1) + block1 = scale * W1[k - 1].transpose(0, 2, 1).reshape(Ngrid, -1) - # AI_mat_offdiag_new[:, d0[k - 1]:d0[k]] = block0 - # AI_mat_offdiag_new[:, d0[-1] + d1[k - 1]:d0[-1] + d1[k]] = block1 + AI_mat_offdiag[:, d0[k - 1] : d0[k]] = block0 + AI_mat_offdiag[:, d0[-1] + d1[k - 1] : d0[-1] + d1[k]] = block1 + # AI_mat_diag computation has been vectorized, but originally + # had the note by block1: this needs double checking (Ruiyi) AI_mat_diag = np.zeros((Ngrid, D0 + D1), dtype=np.float64) - for p in range(Ngrid): - w0 = np.zeros(D0, dtype=np.float64) - w1 = np.zeros(D1, dtype=np.float64) - for k in range(1, Lmax + 1): - w0[d0[k - 1] : d0[k]] = ( - (Lmax - k + 2) - * (Lmax - k + 1) - * (k + 0.5) - * (0.5 * W0[k - 1][p] + 0.5 * W0[k - 1][p].T).T.reshape(-1) - ) - w1[d1[k - 1] : d1[k]] = ( - (Lmax - k + 2) - * (Lmax - k + 1) - * (k + 0.5) - * (0.5 * W1[k - 1][p] + 0.5 * W1[k - 1][p].T).T.reshape(-1) - ) - # this needs double checking (Ruiyi) - AI_mat_diag[p, : d0[-1]] = w0 - AI_mat_diag[p, d0[-1] :] = w1 - AI_mat_diag = xp.asarray(AI_mat_diag) / 1 - AI_mat_offdiag = xp.asarray(AI_mat_offdiag) / 1 - bI = -(Lmax + 2) * (Lmax + 1) / 2 / 1 + for k in range(1, Lmax + 1): + scale = (Lmax - k + 2) * (Lmax - k + 1) * (k + 0.5) + + W0_sym = 0.5 * (W0[k - 1] + W0[k - 1].transpose(0, 2, 1)) + W1_sym = 0.5 * (W1[k - 1] + W1[k - 1].transpose(0, 2, 1)) + + block0 = scale * W0_sym.transpose(0, 2, 1).reshape(Ngrid, -1) + block1 = scale * W1_sym.transpose(0, 2, 1).reshape(Ngrid, -1) + + AI_mat_diag[:, d0[k - 1] : d0[k]] = block0 + AI_mat_diag[:, d0[-1] + d1[k - 1] : d0[-1] + d1[k]] = block1 + + AI_mat_diag = xp.asarray(AI_mat_diag) + AI_mat_offdiag = xp.asarray(AI_mat_offdiag) + bI = -(Lmax + 2) * (Lmax + 1) / 2 # largest eigenvalue for AIAIT Lambda = self.largest_eigenvalue(AI_mat_offdiag, Ngrid, N) @@ -1302,23 +1277,20 @@ def LS_D(W1, W2, W3, W4, Br, Bi): DXijmD = np.diag(Di.conj()) @ Xijm @ np.diag(Dj) wj = ws[j] W1 = ( - wi[:, 0][:, np.newaxis] @ wj[:, 0][:, np.newaxis].T - + Jk @ wi[:, 0][:, np.newaxis] @ wj[:, 0][:, np.newaxis].T @ Jk + wi[:, 0][:, None] @ wj[:, 0][:, None].T + + Jk @ wi[:, 0][:, None] @ wj[:, 0][:, None].T @ Jk ) W2 = ( - wi[:, -1][:, np.newaxis] @ wj[:, 0][:, np.newaxis].T - + Jk @ wi[:, -1][:, np.newaxis] @ wj[:, 0][:, np.newaxis].T @ Jk + wi[:, -1][:, None] @ wj[:, 0][:, None].T + + Jk @ wi[:, -1][:, None] @ wj[:, 0][:, None].T @ Jk ) W3 = ( - wi[:, 0][:, np.newaxis] @ wj[:, -1][:, np.newaxis].T - + Jk @ wi[:, 0][:, np.newaxis] @ wj[:, -1][:, np.newaxis].T @ Jk + wi[:, 0][:, None] @ wj[:, -1][:, None].T + + Jk @ wi[:, 0][:, None] @ wj[:, -1][:, None].T @ Jk ) W4 = ( - wi[:, -1][:, np.newaxis] @ wj[:, -1][:, np.newaxis].T - + Jk - @ wi[:, -1][:, np.newaxis] - @ wj[:, -1][:, np.newaxis].T - @ Jk + wi[:, -1][:, None] @ wj[:, -1][:, None].T + + Jk @ wi[:, -1][:, None] @ wj[:, -1][:, None].T @ Jk ) Br = np.real(4 * DXijmD) Bi = np.imag(4 * DXijmD) From 4de983bd724b8f7720acd740ebc85dd83907900f Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 29 Jun 2026 11:42:47 -0400 Subject: [PATCH 051/161] fix test docstring --- tests/test_commonline_utils.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/test_commonline_utils.py b/tests/test_commonline_utils.py index df8edfcad2..a460d6cefb 100644 --- a/tests/test_commonline_utils.py +++ b/tests/test_commonline_utils.py @@ -191,11 +191,11 @@ def test_g_sync(symmetry): @pytest.mark.parametrize("symmetry", ["C3", "C4", "D3", "D4", "T", "O"]) def test_compare_rots_sym(symmetry): """ - This method find the mean squared error between all pairs of rotations - taking into account each rotation being multiplied by an arbitrary - symmetry group element. In this test we check that a set of rotations - multiplied by random symmetry group elements gives a zero MSE when compared - to the original set. + Thwe compare_rots_sym method finds the mean squared error between all pairs + of relative rotations, Ri.T @ Rj, taking into account each rotation being + multiplied by an arbitrary symmetry group element, ie. g @ Ri. In this test + we check that a set of rotations multiplied by random symmetry group elements + gives a zero MSE when compared to the original set. """ n = 100 dtype = np.float64 From efa559915bf392ed2e23d6469c3784c8bb6cd2a2 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 29 Jun 2026 14:09:52 -0400 Subject: [PATCH 052/161] remove time --- src/aspire/abinitio/commonline_nug.py | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 789802d78e..ce536f5953 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1,5 +1,4 @@ import logging -import time import numpy as np from scipy.special import factorial @@ -466,7 +465,6 @@ def update_S(C0, C1, yE, yEq, yI, X0, X1, Xd0, Xd1, Xq, rho, Lmax, N): Sd0 = -Zd0 - Xd0 / rho Sd1 = -Zd1 - Xd1 / rho Sq = -Zq - Xq / rho - tic0 = time.perf_counter() for k in range(1, Lmax + 1): tmp = self.mat_block( S0[d0[k - 1] : d0[k], :], N, k, IDX_upper, IDX_lower, idx_offdiag @@ -484,12 +482,9 @@ def update_S(C0, C1, yE, yEq, yI, X0, X1, Xd0, Xd1, Xq, rho, Lmax, N): ) tmp = self.psd_projection(tmp) S1[d1[k - 1] : d1[k], :] = self.vec_block(tmp, N, k + 1, IDX_upper) - toc0 = time.perf_counter() - Time[0] += toc0 - tic0 Sd0 = Sd0.T Sd1 = Sd1.T - tic1 = time.perf_counter() for k in range(1, Lmax + 1): tmp = self.transform_back_block( Sd0[:, d0[k - 1] : d0[k]], @@ -503,14 +498,8 @@ def update_S(C0, C1, yE, yEq, yI, X0, X1, Xd0, Xd1, Xq, rho, Lmax, N): ) Sd0 = Sd0.T Sd1 = Sd1.T - toc1 = time.perf_counter() - Time[1] += toc1 - tic1 - - tic2 = time.perf_counter() Sq = self.psd_projection(Sq.T.reshape(n_pairs, 4, 4)) Sq = Sq.T.reshape(-1, n_pairs) - toc2 = time.perf_counter() - Time[2] += toc2 - tic2 return S0, S1, Sd0, Sd1, Sq def update_yE(C0, C1, X0, X1, Xd0, Xd1, Xq, S0, S1, Sd0, Sd1, Sq, yI, rho): @@ -689,7 +678,6 @@ def print_updates(verbose): yEq = xp.zeros(bEq.shape, dtype=np.float64) IDX = np.arange(3) - Time = np.zeros(4) for t in range(max_iter): np.random.shuffle(IDX) for idx in IDX: From 1aea6c3d88afe46c1e31d9c26e4cef0ac56e8d0e Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 29 Jun 2026 14:58:43 -0400 Subject: [PATCH 053/161] cleanup index creation --- src/aspire/abinitio/commonline_nug.py | 33 +++++++++++---------------- 1 file changed, 13 insertions(+), 20 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index ce536f5953..a41952b0e3 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -716,26 +716,19 @@ def ADMM_preprocessing(self, C): """ Lmax = self.Lmax N = self.n_img - count = 0 - idx_diag = [] - idx_offdiag = [] - for i in range(N): - for j in range(i, N): - if j == i: - idx_diag.append(count) - else: - idx_offdiag.append(count) - count += 1 - IDX_upper = [] - IDX_lower = [] - for i in range(N): - for j in range(N): - if j >= i: - IDX_upper.append(i * N + j) - for j in range(N): - for i in range(N): - if j < i: - IDX_lower.append(i * N + j) + + # Indices into the packed upper-triangular image-pair columns: + # diagonal self-pairs (i, i) and off-diagonal pairs (i, j), i < j. + iu, ju = np.triu_indices(N) + idx_diag = xp.asarray(np.flatnonzero(iu == ju)) + idx_offdiag = xp.asarray(np.flatnonzero(iu != ju)) + + # Linear indices into the full N x N image-pair grid. IDX_upper selects + # upper-triangular pairs (i, j), i <= j; IDX_lower selects the matching + # lower-triangular pairs (j, i), i < j, in the same order as idx_offdiag. + iu_off, ju_off = np.triu_indices(N, k=1) + IDX_upper = xp.asarray(iu * N + ju) + IDX_lower = xp.asarray(ju_off * N + iu_off) # normalize C matrix Cnorm = 0 From 4b020903d6717de75a5ac01c0510f389a6ce04f4 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 30 Jun 2026 11:25:00 -0400 Subject: [PATCH 054/161] docstrings and code comments for main ADMM step --- src/aspire/abinitio/commonline_nug.py | 89 ++++++++++++++++++++++++++- 1 file changed, 87 insertions(+), 2 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index a41952b0e3..2ce85a1db0 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -338,7 +338,8 @@ def admm_sym_J(self, C, verbose): mult = self.mult Nstep_yI = self.Nstep_yI - # admm for symmetric case + # Solve the symmetry-constrained SDP relaxation in the real, packed + # representation basis prepared by ADMM_preprocessing. ( C0, C1, @@ -371,6 +372,8 @@ def admm_sym_J(self, C, verbose): rank_Ak, _ = self.compute_rank() logger.info(f"Rank of Ak: {rank_Ak}") + # Build per-degree equality operators coupling the k and k+1 invariant + # blocks. These encode the representation constraints imposed by symmetry. AE = [] AEAETinv = [] for k in range(1, Lmax + 1): @@ -387,6 +390,9 @@ def admm_sym_J(self, C, verbose): AEk[count, count - 1 + s1 + s0] = 1 AE.append(AEk) AEAETinv.append(xp.linalg.pinv(AEk @ AEk.T)) + + # Right-hand sides for the equality constraints: symmetry projector ranks + # and identity constraints on diagonal representation blocks. bE = xp.zeros((Lmax + D0 + D1), dtype=np.float64) for k in range(Lmax): bE[k + d0[k] + d1[k] :] = rank_Ak[k] @@ -397,6 +403,9 @@ def admm_sym_J(self, C, verbose): k + 2 ).T.reshape(-1) bE = xp.repeat(bE[:, None], N, axis=1) + + # Degree-wise permutations used to move between full Wigner blocks and the + # two invariant block variables used by the relaxation. P = [] for k in range(1, Lmax + 1): dk = 2 * k + 1 @@ -409,6 +418,15 @@ def admm_sym_J(self, C, verbose): P.append(Pk) def fun_AE(X0, X1, Xd0, Xd1, Xq): + """ + Apply the equality constraint operator to the current primal variables. + + z contains the per-image, per-degree equality constraints that couple + diagonal X0/X1 blocks to auxiliary diagonal variables Xd0/Xd1. + + zq contains the off-diagonal quaternion constraints built from Xq and + the lowest-degree off-diagonal representation blocks. + """ z = xp.zeros((Lmax + D0 + D1, N), dtype=np.float64) for k in range(Lmax): z[k + d0[k] + d1[k] : k + 1 + d0[k + 1] + d1[k + 1]] = AE[ @@ -422,11 +440,24 @@ def fun_AE(X0, X1, Xd0, Xd1, Xq): ), axis=0, ) - return z, AEq @ xp.concatenate( + zq = AEq @ xp.concatenate( (Xq, X0[:1, idx_offdiag], X1[:4, idx_offdiag]), axis=0 ) + return z, zq def fun_AET(yE, yEq): + """ + Apply the adjoint of the equality constraint operator. + + yE contains multipliers for the per-image, per-degree equality + constraints produced as z by fun_AE. + + yEq contains multipliers for the off-diagonal quaternion constraints + produced as zq by fun_AE. + + Returns the contribution of these equality multipliers back into each + packed primal variable block: X0, X1, Xd0, Xd1, and Xq. + """ Z0 = xp.zeros((D0, N * (N + 1) // 2), dtype=np.float64) Z1 = xp.zeros((D1, N * (N + 1) // 2), dtype=np.float64) Zd0 = xp.zeros((D0, N), dtype=np.float64) @@ -445,6 +476,15 @@ def fun_AET(yE, yEq): return Z0, Z1, Zd0, Zd1, Zq[:16] def fun_AI(X0, X1): + """ + Evaluate the Fejer inequality operator on the packed representation blocks. + + Each output column corresponds to one packed image pair. Each row + corresponds to one sampled SO(3) grid point. + + The result is compared against bI to enforce the discretized + nonnegativity constraints of the truncated SDP relaxation. + """ z = xp.zeros((Ngrid, N * (N + 1) // 2), dtype=np.float64) tmp = xp.concatenate((X0, X1), axis=0) z[:, idx_diag] = AI_mat_diag @ tmp[:, idx_diag] @@ -452,19 +492,36 @@ def fun_AI(X0, X1): return z def fun_AIT(yI): + """ + Apply the adjoint of the Fejer inequality operator. + """ Z = xp.zeros((D0 + D1, N * (N + 1) // 2), dtype=np.float64) Z[:, idx_diag] = AI_mat_diag.T @ yI[:, idx_diag] Z[:, idx_offdiag] = AI_mat_offdiag.T @ yI[:, idx_offdiag] return Z[:D0, :], Z[D0:, :] def update_S(C0, C1, yE, yEq, yI, X0, X1, Xd0, Xd1, Xq, rho, Lmax, N): + """ + Update PSD slack variables for the current ADMM iterate. + + The slack variables S0/S1/Sd0/Sd1/Sq carry the semidefinite constraints. + This step forms the unconstrained slack minimizers and projects each + matrix block onto the positive semidefinite cone. + """ + # Equality multipliers mapped back to each primal variable block. Z0, Z1, Zd0, Zd1, Zq = fun_AET(yE, yEq) + + # Inequality multipliers mapped back to the X0/X1 variable blocks. AIT_yI0, AIT_yI1 = fun_AIT(yI) + + # Unconstrained slack updates before PSD projection. S0 = C0 - Z0 - AIT_yI0 - X0 / rho S1 = C1 - Z1 - AIT_yI1 - X1 / rho Sd0 = -Zd0 - Xd0 / rho Sd1 = -Zd1 - Xd1 / rho Sq = -Zq - Xq / rho + + # Project packed X0/X1 blocks degree by degree. for k in range(1, Lmax + 1): tmp = self.mat_block( S0[d0[k - 1] : d0[k], :], N, k, IDX_upper, IDX_lower, idx_offdiag @@ -483,6 +540,7 @@ def update_S(C0, C1, yE, yEq, yI, X0, X1, Xd0, Xd1, Xq, rho, Lmax, N): tmp = self.psd_projection(tmp) S1[d1[k - 1] : d1[k], :] = self.vec_block(tmp, N, k + 1, IDX_upper) + # Project the diagonal coupling blocks. Sd0 = Sd0.T Sd1 = Sd1.T for k in range(1, Lmax + 1): @@ -498,11 +556,19 @@ def update_S(C0, C1, yE, yEq, yI, X0, X1, Xd0, Xd1, Xq, rho, Lmax, N): ) Sd0 = Sd0.T Sd1 = Sd1.T + + # Project each 4x4 quaternion slack block for off-diagonal image pairs. Sq = self.psd_projection(Sq.T.reshape(n_pairs, 4, 4)) Sq = Sq.T.reshape(-1, n_pairs) return S0, S1, Sd0, Sd1, Sq def update_yE(C0, C1, X0, X1, Xd0, Xd1, Xq, S0, S1, Sd0, Sd1, Sq, yI, rho): + """ + Update equality multipliers for the ADMM iterate. + + yE enforces the per-degree symmetry/diagonal constraints. + yEq enforces the off-diagonal quaternion constraints. + """ AIT_yI0, AIT_yI1 = fun_AIT(yI) z, zq = fun_AE( -X0 / rho + C0 - S0 - AIT_yI0, @@ -521,6 +587,12 @@ def update_yE(C0, C1, X0, X1, Xd0, Xd1, Xq, S0, S1, Sd0, Sd1, Sq, yI, rho): return yE, yEq def update_yI(C0, C1, X0, X1, S0, S1, yE, yEq, yI, rho, Lambda): + """ + Update nonnegative multipliers for the Fejer inequality constraints. + + yI corresponds to the discretized constraints fun_AI(X0, X1) >= bI, + one multiplier per SO(3) grid point and packed image-pair column. + """ Z0, Z1, _, _, _ = fun_AET(yE, yEq) AIT_yI0, AIT_yI1 = fun_AIT(yI) tmp = fun_AI( @@ -533,6 +605,12 @@ def update_yI(C0, C1, X0, X1, S0, S1, yE, yEq, yI, rho, Lambda): def update_X( C0, C1, X0, X1, Xd0, Xd1, Xq, yE, yEq, yI, S0, S1, Sd0, Sd1, Sq, rho ): + """ + Update the primal ADMM multipliers and compute their combined residual. + + Each tmp below is the current residual for one block family. + The multiplier is advanced by mult * rho * tmp. + """ Z0, Z1, Zd0, Zd1, Zq = fun_AET(yE, yEq) AIT_yI0, AIT_yI1 = fun_AIT(yI) tmp = S0 + Z0 + AIT_yI0 - C0 @@ -560,6 +638,9 @@ def update_X( ) def update_rho(X0, X1, Xd0, Xd1, Xq, bE, bEq, bI, res_X, rho, factor, normC): + """ + Update the ADMM penalty parameter rho. + """ z, zq = fun_AE(X0, X1, Xd0, Xd1, Xq) res_eq = xp.linalg.norm(z - bE) / (1 + xp.linalg.norm(bE)) + xp.linalg.norm( zq - bEq @@ -677,6 +758,8 @@ def print_updates(verbose): yE = xp.zeros(bE.shape, dtype=np.float64) yEq = xp.zeros(bEq.shape, dtype=np.float64) + # Run ADMM iterations, randomly ordering the block updates before each primal + # multiplier update and penalty adjustment. IDX = np.arange(3) for t in range(max_iter): np.random.shuffle(IDX) @@ -701,6 +784,8 @@ def print_updates(verbose): X0, X1, Xd0, Xd1, Xq, bE, bEq, bI, res_X, rho, factor, normC ) + # Convert the optimized packed block variables back to full degree-wise + # representation matrices for Euler-angle recovery. X_admm = self.transform_coeff_back(X0, X1, IDX_upper, IDX_lower, idx_offdiag) for k in range(Lmax): X_admm[k] = xp.asnumpy(X_admm[k]) From d69e204b601bf7f29e02059401112015c536c76d Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 30 Jun 2026 14:21:40 -0400 Subject: [PATCH 055/161] docstring for compare_rots_sym --- src/aspire/abinitio/commonline_utils.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/src/aspire/abinitio/commonline_utils.py b/src/aspire/abinitio/commonline_utils.py index 283dfaddc8..0976450491 100644 --- a/src/aspire/abinitio/commonline_utils.py +++ b/src/aspire/abinitio/commonline_utils.py @@ -592,6 +592,18 @@ def build_outer_products(n, dtype): def compare_rots_sym(R_est, R_true, sym): + """ + Compute the mean squared error between all pairs of estimate relative rotations, + R_est(i).T @ R_est(j), and a reference set, R_true(i).T @ R_true(j), taking into + account each estimate rotation being multiplied by an arbitrary symmetry group + element, ie. g @ R_est(i), where g is an element of the symmetry group 'sym'. + + :param R_est: Nx3x3 array of estimate rotations. + :param R_true: Nx3x3 array of reference rotations. + :param sym: A string designating the symmetry of the underlying molecule, ie 'D3'. + + :return: The pairwise mean squared error. + """ N = R_true.shape[0] sym_euler = SymmetryGroup.parse(sym).matrices order = sym_euler.shape[0] From fe5911be63fbf135a0e7affb7af146b373cd6db3 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 1 Jul 2026 10:44:33 -0400 Subject: [PATCH 056/161] compute_coeff code comments --- src/aspire/abinitio/commonline_nug.py | 23 +++++++++++++++++++---- 1 file changed, 19 insertions(+), 4 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 2ce85a1db0..39f3f0751b 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -45,7 +45,7 @@ def __init__( Initialize the symmetric NUG orientation estimator. All default values match those used for publication results, with exception of `pr_iters`, which controls how many iterations of proximal refinement to be used. The default of None runs the algorithm without proximal - refinement. Set to `pr_iters` to 4 to match the proximal refinement workflow in the related + refinement. Set `pr_iters=4` to match t3he proximal refinement workflow in the related publication. :param src: Source containing the input projection images. @@ -57,7 +57,7 @@ def __init__( value must be even. Default is 360. :param max_shift: Determines maximum range for shifts for common-line detection as a proportion of the resolution. Default is 0.15. - :param shift_step:Resolution of shift estimation common-line detection in pixels. + :param shift_step: Resolution of shift estimation common-line detection in pixels. Default is 1 pixel. :param mask: Option to mask `src.images` with a fuzzy mask (boolean). Default, `True`, applies a mask. @@ -164,7 +164,7 @@ def compute_coeff(self): """ Compute the truncated Fourier coefficient matrices of the pairwise common-line losses. """ - # compute the coefficient matrix + # Build degree-wise coefficient matrices C[k] for the linear SDP objective. N = self.n_img n_theta = self.n_theta Lmax = self.Lmax @@ -199,10 +199,12 @@ def fij(alpha, gamma, i, j): norms = np.linalg.norm(Si[None] - Sj_shifted, 1, axis=1) return norms.min() + # Quadrature grid used to sample the pairwise loss over SO(3) Euler angles. alpha_grid = np.arange(2 * T, dtype=np.float64) * np.pi / T beta_grid = (2 * np.arange(2 * T, dtype=np.float64) + 1) * np.pi / 4 / T gamma_grid = np.arange(2 * T, dtype=np.float64) * np.pi / T + # Beta quadrature weights for the truncated Wigner-D Fourier expansion. bT = np.zeros(2 * T, dtype=np.float64) for n in range(2 * T): ss = 0 @@ -210,12 +212,16 @@ def fij(alpha, gamma, i, j): ss = ss + np.sin(beta_grid[n] * (2 * m + 1)) / (2 * m + 1) bT[n] = 2 / T * np.sin(beta_grid[n]) * ss + # Precompute degree-wise beta/Wigner weight matrices used in each coefficient transform. BTK = [] for k in range(1, Lmax + 1): btk = np.sum(bT[:, None, None] * self.Wd(k, beta_grid), axis=0) BTK.append(btk.T) def fijhat_k(k, F): + """ + Approximate the degree-k Fourier coefficient block of a sampled pairwise loss. + """ dk = 2 * k + 1 exp_alpha_grid = np.zeros((2 * T, dk), dtype=complex_type(np.float64)) @@ -230,10 +236,14 @@ def fijhat_k(k, F): fhat = BTK[k - 1] * S / 4 / T**2 return fhat + # Allocate one coefficient matrix per Wigner degree; + # each is block-indexed by image pair. C = [] for k in range(1, Lmax + 1): dk = 2 * k + 1 C.append(np.zeros((N * dk, N * dk), dtype=complex_type(np.float64))) + + # Compute off-diagonal image-pair losses and insert their degree-wise coefficients. for i in range(N): for j in range(i + 1, N): Fij = np.zeros((2 * T, 2 * T), dtype=np.float64) @@ -245,9 +255,12 @@ def fijhat_k(k, F): C[k - 1][j * dk : (j + 1) * dk, i * dk : (i + 1) * dk] = fijhat_k( k, Fij ) # *dk + + # Fill the conjugate transpose blocks so each coefficient matrix is Hermitian. for k in range(1, Lmax + 1): C[k - 1] = C[k - 1] + C[k - 1].conj().T + # Diagonal blocks encode self-pair losses and are handled separately. for i in range(N): Fii = np.zeros((2 * T, 2 * T), dtype=np.float64) for j1 in range(2 * T): @@ -259,6 +272,7 @@ def fijhat_k(k, F): k, Fii ) # *dk + # Convert complex Wigner coefficients to the real representation basis used by ADMM. for k in range(1, Lmax + 1): [T, Tinv] = self.complex2real(k) C[k - 1] = np.real( @@ -1641,7 +1655,8 @@ def construct_AEq(self): # First 16 rows: identity constraints on first 16 variables AEq[:16, :16] = np.eye(16, dtype=np.float64) - # Extra columns 16:21 + # Columns 16:21 map the low-degree X0/X1 entries into the quaternion + # convex-hull constraint Xq = I/4 - linear(X^(1)). extra = 0.25 * np.array( [ [-1, 1, 0, 0, 1], From 1e526819dac761f910162b66ca2710f738d80219 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 1 Jul 2026 11:34:33 -0400 Subject: [PATCH 057/161] proximal_refine code commentary --- src/aspire/abinitio/commonline_nug.py | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 39f3f0751b..3af452cf99 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1032,12 +1032,18 @@ def proximal_refine(self, X_admm, weight, Penalty, r): C = self.C def Ak(J, Euler): - # compute Ak matrix + """ + Average the degree-J representation over the symmetry + group to form the symmetry projector Ak. + """ order = Euler.shape[0] A = self.WD(J, Euler).sum(axis=0) return np.round(A / order, 10) def rel_change(A, B, eps=1e-12): + """ + Measure relative change across all degree-wise relaxed matrices. + """ num = 0.0 den = 0.0 for k in range(len(A)): @@ -1045,7 +1051,12 @@ def rel_change(A, B, eps=1e-12): den += np.linalg.norm(B[k]) ** 2 return np.sqrt(num) / max(np.sqrt(den), eps) + # The symmetry projector ranks give the expected low-rank structure + # for each degree of a valid orbit solution. rank_Ak = np.zeros(self.Lmax, dtype=np.float64) + + # Keep an unmodified copy of the original objective coefficients; each + # refinement step adds a temporary proximal term to these. C_base = [None] * self.Lmax for k in range(self.Lmax): C_base[k] = xp.asnumpy(C[k]).copy() @@ -1054,6 +1065,9 @@ def rel_change(A, B, eps=1e-12): ) def low_rank_proj(X, r_step): + """ + Project diagonal image blocks toward the rank expected from symmetry. + """ Xproj = [] for k in range(self.Lmax): dk = 2 * k + 3 @@ -1069,9 +1083,12 @@ def low_rank_proj(X, r_step): Xproj.append(tmp) return Xproj + # CC stores the modified objective coefficients for the next proximal ADMM solve. CC = [None] * self.Lmax current = [np.copy(Xk) for Xk in X_admm] + # Proximal refinement step. + # Repeatedly bias the SDP objective toward the low-rank projection and resolve. for step in range(self.pr_iters): X_proj = low_rank_proj(current, r[step]) @@ -1091,6 +1108,7 @@ def low_rank_proj(X, r_step): rel_change(X_next, current), ) + # Use the refined SDP solution as the starting point for the next projection step current = [np.copy(Xk) for Xk in X_next] return current From 2db60f458fc30ec93c6c148947c393d572008077 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 1 Jul 2026 14:27:21 -0400 Subject: [PATCH 058/161] some cleanup --- src/aspire/abinitio/commonline_utils.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/aspire/abinitio/commonline_utils.py b/src/aspire/abinitio/commonline_utils.py index 0976450491..2f7da1cd24 100644 --- a/src/aspire/abinitio/commonline_utils.py +++ b/src/aspire/abinitio/commonline_utils.py @@ -414,13 +414,13 @@ def g_sync_finite_group(rots, rots_gt, symmetry): """ Synchronize ground-truth rotations over a finite symmetry group. - This is a finite-group generalization of cyclic synchronization. The + This is a finite-group generalization of cyclic synchronization. The pairwise matching step estimates relative symmetry elements between image - pairs. The spectral synchronization step then recovers one symmetry + pairs. The spectral synchronization step then recovers one symmetry element per image that is globally consistent with those pairwise estimates. Unlike the cyclic case, the relative symmetry cannot generally be encoded - as a scalar complex phase. For non-commutative groups such as D_n, we instead + as a scalar complex phase. For non-commutative groups such as D_n, we instead represent each group element by its left-regular permutation matrix. :param rots: Estimated rotation matrices @@ -467,7 +467,7 @@ def find_group_index(A): # reps[a] @ e_b = e_{a b} # # This lets us store arbitrary finite-group relative elements in a block - # synchronization matrix. This is the key generalization beyond cyclic + # synchronization matrix. This is the key generalization beyond cyclic # scalar phases in g_sync_cyclic. reps = np.zeros((n_group, n_group, n_group), dtype=float) From bf1030a0b7e9668cca676390b251ba05b0fc7292 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 28 Jul 2026 14:04:11 -0400 Subject: [PATCH 059/161] missing xp's in print_update --- src/aspire/abinitio/commonline_nug.py | 64 +++++++++++++-------------- 1 file changed, 32 insertions(+), 32 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 3af452cf99..2650d942da 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -689,11 +689,11 @@ def print_updates(verbose): ) z, zq = fun_AE(X0, X1, Xd0, Xd1, Xq) - res_eq = np.linalg.norm(z - bE) / ( - 1 + np.linalg.norm(bE) - ) + np.linalg.norm(zq - bEq) / (1 + np.linalg.norm(bEq)) - res_inq = np.linalg.norm(np.maximum(bI - fun_AI(X0, X1), 0)) / ( - 1 + abs(bI) * np.sqrt(Ngrid * N * (N + 1) / 2) + res_eq = xp.linalg.norm(z - bE) / ( + 1 + xp.linalg.norm(bE) + ) + xp.linalg.norm(zq - bEq) / (1 + xp.linalg.norm(bEq)) + res_inq = xp.linalg.norm(xp.maximum(bI - fun_AI(X0, X1), 0)) / ( + 1 + abs(bI) * xp.sqrt(Ngrid * N * (N + 1) / 2) ) res_psdX = 0 for k in range(1, Lmax + 1): @@ -705,7 +705,7 @@ def print_updates(verbose): IDX_lower, idx_offdiag, ) - res_psdX += np.linalg.norm(self.psd_projection(-tmp)) + res_psdX += xp.linalg.norm(self.psd_projection(-tmp)) tmp = self.mat_block( X1[d1[k - 1] : d1[k], :], N, @@ -714,8 +714,8 @@ def print_updates(verbose): IDX_lower, idx_offdiag, ) - res_psdX += np.linalg.norm(self.psd_projection(-tmp)) - res_psdX = res_psdX / (1 + np.linalg.norm(X0) + np.linalg.norm(X1)) + res_psdX += xp.linalg.norm(self.psd_projection(-tmp)) + res_psdX = res_psdX / (1 + xp.linalg.norm(X0) + xp.linalg.norm(X1)) res_psdD = 0 for k in range(1, Lmax + 1): tmp = self.transform_back_block( @@ -724,29 +724,29 @@ def print_updates(verbose): k, P[k - 1], ) - res_psdD += np.linalg.norm( + res_psdD += xp.linalg.norm( self.psd_projection(-tmp), axis=(-2, -1) ).sum() - res_psdD = res_psdD / (1 + np.linalg.norm(Xd0) + np.linalg.norm(Xd1)) + res_psdD = res_psdD / (1 + xp.linalg.norm(Xd0) + xp.linalg.norm(Xd1)) res_psdQ = 0 for count in range(N * (N - 1) // 2): tmp = Xq[:, count].reshape(4, 4).T - res_psdQ += np.linalg.norm(self.psd_projection(-tmp)) - res_psdQ = res_psdQ / (1 + np.linalg.norm(Xq)) - - normS = np.sqrt( - np.linalg.norm(S0) ** 2 - + np.linalg.norm(S1) ** 2 - + np.linalg.norm(Sd0) ** 2 - + np.linalg.norm(Sd1) ** 2 - + np.linalg.norm(Sq) ** 2 + res_psdQ += xp.linalg.norm(self.psd_projection(-tmp)) + res_psdQ = res_psdQ / (1 + xp.linalg.norm(Xq)) + + normS = xp.sqrt( + xp.linalg.norm(S0) ** 2 + + xp.linalg.norm(S1) ** 2 + + xp.linalg.norm(Sd0) ** 2 + + xp.linalg.norm(Sd1) ** 2 + + xp.linalg.norm(Sq) ** 2 ) - normX = np.sqrt( - np.linalg.norm(X0) ** 2 - + np.linalg.norm(X1) ** 2 - + np.linalg.norm(Xd0) ** 2 - + np.linalg.norm(Xd1) ** 2 - + np.linalg.norm(Xq) ** 2 + normX = xp.sqrt( + xp.linalg.norm(X0) ** 2 + + xp.linalg.norm(X1) ** 2 + + xp.linalg.norm(Xd0) ** 2 + + xp.linalg.norm(Xd1) ** 2 + + xp.linalg.norm(Xq) ** 2 ) p_res = res_eq + res_inq + res_psdX + res_psdD + res_psdQ d_res = res_X / (1 + normC) @@ -764,13 +764,13 @@ def print_updates(verbose): + ", |X|=%1.2f" % normX ) - Xd0 = xp.zeros((D0, N), dtype=np.float64) - Xd1 = xp.zeros((D1, N), dtype=np.float64) - Sd0 = xp.zeros(Xd0.shape, dtype=np.float64) - Sd1 = xp.zeros(Xd1.shape, dtype=np.float64) - yI = xp.zeros((Ngrid, N * (N + 1) // 2), dtype=np.float64) - yE = xp.zeros(bE.shape, dtype=np.float64) - yEq = xp.zeros(bEq.shape, dtype=np.float64) + Xd0 = xp.zeros((D0, N), dtype=xp.float64) + Xd1 = xp.zeros((D1, N), dtype=xp.float64) + Sd0 = xp.zeros(Xd0.shape, dtype=xp.float64) + Sd1 = xp.zeros(Xd1.shape, dtype=xp.float64) + yI = xp.zeros((Ngrid, N * (N + 1) // 2), dtype=xp.float64) + yE = xp.zeros(bE.shape, dtype=xp.float64) + yEq = xp.zeros(bEq.shape, dtype=xp.float64) # Run ADMM iterations, randomly ordering the block updates before each primal # multiplier update and penalty adjustment. From eac5698c8c44bacdb847d0597aa6233ab6574718 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 29 Jul 2026 15:01:11 -0400 Subject: [PATCH 060/161] add reference for g_sync_finite_group method --- src/aspire/abinitio/commonline_utils.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/aspire/abinitio/commonline_utils.py b/src/aspire/abinitio/commonline_utils.py index 2f7da1cd24..4cdf29a0bf 100644 --- a/src/aspire/abinitio/commonline_utils.py +++ b/src/aspire/abinitio/commonline_utils.py @@ -412,7 +412,7 @@ def g_sync_cyclic(rots, rots_gt, symmetry): def g_sync_finite_group(rots, rots_gt, symmetry): """ - Synchronize ground-truth rotations over a finite symmetry group. + Synchronize ground truth rotations over a finite symmetry group. This is a finite-group generalization of cyclic synchronization. The pairwise matching step estimates relative symmetry elements between image @@ -423,6 +423,12 @@ def g_sync_finite_group(rots, rots_gt, symmetry): as a scalar complex phase. For non-commutative groups such as D_n, we instead represent each group element by its left-regular permutation matrix. + This follows the standard spectral group-synchronization template: build a + block matrix of pairwise group measurements, compute the leading group-sized + eigenspace, and round the recovered blocks back to group elements. See Ling, + "Near-Optimal Performance Bounds for Orthogonal and Permutation Group + Synchronization via Spectral Methods", Appl. Comput. Harmon. Anal. 2022. + :param rots: Estimated rotation matrices :param rots_gt: Ground truth rotation matrices :param symmetry: The symmetry of the underlying molecule. From 6b5f16ef06c5235311cd007da2140c9376b122ad Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 13 Aug 2026 15:08:57 -0400 Subject: [PATCH 061/161] Initial port of C1 admm --- src/aspire/abinitio/commonline_nug.py | 406 +++++++++++++++++++++++++- 1 file changed, 400 insertions(+), 6 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 2650d942da..8474c7b428 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -7,7 +7,7 @@ from aspire.numeric import xp from aspire.operators import PolarFT from aspire.utils import Rotation, cart2sph, complex_type -from aspire.volume import CnSymmetryGroup, DnSymmetryGroup, SymmetryGroup +from aspire.volume import CnSymmetryGroup, DnSymmetryGroup, IdentitySymmetryGroup, SymmetryGroup from .commonline_utils import _generate_shift_phase_and_filter, saff_kuijlaars @@ -115,9 +115,9 @@ def __init__( logger.info(f"Using provided symmetry: {symmetry}") self.sym_grp = SymmetryGroup.parse(symmetry) - if not isinstance(self.sym_grp, (CnSymmetryGroup, DnSymmetryGroup)): + if not isinstance(self.sym_grp, (IdentitySymmetryGroup, CnSymmetryGroup, DnSymmetryGroup)): raise ValueError( - f"This algorithm supports cyclic or dihedral symmetry. Found {str(self.sym_grp)}." + f"This algorithm supports Cn, Dn, and asymmetric molecules. Found {str(self.sym_grp)}." ) self.sym_euler = self.sym_grp.rotations.angles @@ -321,19 +321,414 @@ def complex2real(self, ell): def perform_admm(self): """ - Solve the symmetric NUG relaxation and optionally apply proximal refinement. + Solve the NUG relaxation and optionally apply proximal refinement. + + The trivial symmetry group C1 uses the asymmetric formulation. + Nontrivial cyclic and dihedral symmetry groups use the + symmetry-constrained formulation. """ - X_est = self.admm_sym_J(self.C, self.verbose) + is_asymmetric = self.n_sym == 1 + + if is_asymmetric: + X_est = self.admm_plain_J(self.C, self.verbose) + else: + X_est = self.admm_sym_J(self.C, self.verbose) if self.pr_iters is not None: + if is_asymmetric: + raise NotImplementedError( + "Proximal refinement is not yet implemented for asymmetric NUG." + ) + X_est = self.proximal_refine( X_est, self.pr_weights, self.pr_penalty, self.pr_rank, ) + self.X_est = X_est + def admm_plain_J(self, C, verbose): + """ + Solve the asymmetric NUG semidefinite relaxation using ADMM. + + :param C: Fourier coefficient matrices of the NUG objective. + :param verbose: Whether to log ADMM progress. + + :return: Relaxed representation matrices. + """ + Lmax = self.Lmax + N = self.n_img + max_iter = self.max_iter + rho = self.rho + ratio = self.ratio + factor = self.factor + mult = self.mult + Nstep_yI = self.Nstep_yI + + ( + C0, + C1, + normC, + AEq, + bEq, + AEqAEqtinv, + AI_mat_diag, + AI_mat_offdiag, + bI, + Lambda, + d0, + d1, + D0, + D1, + idx_diag, + idx_offdiag, + IDX_upper, + IDX_lower, + X0, + X1, + Xq, + S0, + S1, + Sq, + ) = self.ADMM_preprocessing(C) + + Ngrid = self.Ngrid + n_pairs = N * (N - 1) // 2 + n_packed = N * (N + 1) // 2 + + # In the asymmetric formulation, diagonal representation blocks are identity. + bE0 = xp.zeros(D0, dtype=np.float64) + bE1 = xp.zeros(D1, dtype=np.float64) + + for k in range(1, Lmax + 1): + bE0[d0[k - 1] : d0[k]] = xp.eye(k, dtype=np.float64).reshape(-1) + bE1[d1[k - 1] : d1[k]] = xp.eye(k + 1, dtype=np.float64).reshape(-1) + + bE0 = xp.repeat(bE0[:, None], N, axis=1) + bE1 = xp.repeat(bE1[:, None], N, axis=1) + + def fun_AE(X0, X1, Xq): + """ + Apply the asymmetric equality constraints. + """ + z0 = X0[:, idx_diag] + z1 = X1[:, idx_diag] + zq = AEq @ xp.concatenate( + (Xq, X0[:1, idx_offdiag], X1[:4, idx_offdiag]), axis=0 + ) + return z0, z1, zq + + def fun_AET(yE0, yE1, yEq): + """ + Apply the adjoint of the asymmetric equality operator. + """ + Z0 = xp.zeros((D0, n_packed), dtype=np.float64) + Z1 = xp.zeros((D1, n_packed), dtype=np.float64) + Z0[:, idx_diag] = yE0 + Z1[:, idx_diag] = yE1 + + Zq = AEq.T @ yEq + Z0[:1, idx_offdiag] = Zq[16:17] + Z1[:4, idx_offdiag] = Zq[17:] + return Z0, Z1, Zq[:16] + + def fun_AI(X0, X1): + """ + Apply the Fejer inequality operator to off-diagonal image pairs. + + The plain formulation does not apply these inequalities to the fixed + diagonal representation blocks. + """ + tmp = xp.concatenate( + (X0[:, idx_offdiag], X1[:, idx_offdiag]), axis=0 + ) + return AI_mat_offdiag @ tmp + + def fun_AIT(yI): + """ + Apply the adjoint of the asymmetric Fejer operator. + """ + Z = xp.zeros((D0 + D1, n_packed), dtype=np.float64) + Z[:, idx_offdiag] = AI_mat_offdiag.T @ yI + return Z[:D0], Z[D0:] + + def update_S(C0, C1, yE0, yE1, yEq, yI, X0, X1, Xq, rho): + """ + Update and project the asymmetric PSD slack variables. + """ + Z0, Z1, Zq = fun_AET(yE0, yE1, yEq) + AIT_yI0, AIT_yI1 = fun_AIT(yI) + + S0 = C0 - Z0 - AIT_yI0 - X0 / rho + S1 = C1 - Z1 - AIT_yI1 - X1 / rho + Sq = -Zq - Xq / rho + + for k in range(1, Lmax + 1): + tmp = self.mat_block( + S0[d0[k - 1] : d0[k]], + N, + k, + IDX_upper, + IDX_lower, + idx_offdiag, + ) + tmp = self.psd_projection(tmp) + S0[d0[k - 1] : d0[k]] = self.vec_block(tmp, N, k, IDX_upper) + + tmp = self.mat_block( + S1[d1[k - 1] : d1[k]], + N, + k + 1, + IDX_upper, + IDX_lower, + idx_offdiag, + ) + tmp = self.psd_projection(tmp) + S1[d1[k - 1] : d1[k]] = self.vec_block(tmp, N, k + 1, IDX_upper) + + # This is the established batched equivalent of projecting each 4x4 + # quaternion block and storing tmp.T.reshape(16) in each column. + Sq = self.psd_projection(Sq.T.reshape(n_pairs, 4, 4)) + Sq = Sq.T.reshape(16, n_pairs) + return S0, S1, Sq + + def update_yE(C0, C1, S0, S1, Sq, yI, X0, X1, Xq, rho): + """ + Update the asymmetric equality multipliers. + """ + AIT_yI0, AIT_yI1 = fun_AIT(yI) + z0, z1, zq = fun_AE( + -X0 / rho + C0 - S0 - AIT_yI0, + -X1 / rho + C1 - S1 - AIT_yI1, + -Xq / rho - Sq, + ) + yE0 = bE0 / rho + z0 + yE1 = bE1 / rho + z1 + yEq = AEqAEqtinv @ (bEq / rho + zq) + return yE0, yE1, yEq + + def update_yI( + C0, + C1, + S0, + S1, + yE0, + yE1, + yEq, + yI, + X0, + X1, + rho, + Lambda, + ): + """ + Update the nonnegative Fejer inequality multipliers. + """ + AET_yE0, AET_yE1, _ = fun_AET(yE0, yE1, yEq) + AIT_yI0, AIT_yI1 = fun_AIT(yI) + tmp = fun_AI( + -X0 / rho + C0 - S0 - AIT_yI0 - AET_yE0, + -X1 / rho + C1 - S1 - AIT_yI1 - AET_yE1, + ) + yI = yI + bI / rho / Lambda + tmp / Lambda + yI = xp.maximum(yI, 0) + return yI + + def update_X( + C0, + C1, + S0, + S1, + Sq, + yE0, + yE1, + yEq, + yI, + X0, + X1, + Xq, + rho, + ): + """ + Update the asymmetric primal multipliers. + """ + Z0, Z1, Zq = fun_AET(yE0, yE1, yEq) + AIT_yI0, AIT_yI1 = fun_AIT(yI) + + tmp = S0 + AIT_yI0 + Z0 - C0 + X0 = X0 + mult * rho * tmp + resX0 = xp.linalg.norm(tmp) + + tmp = S1 + AIT_yI1 + Z1 - C1 + X1 = X1 + mult * rho * tmp + resX1 = xp.linalg.norm(tmp) + + tmp = Sq + Zq + Xq = Xq + mult * rho * tmp + resXq = xp.linalg.norm(tmp) + + res_X = xp.sqrt(resX0**2 + resX1**2 + resXq**2) + return X0, X1, Xq, res_X + + def update_rho(X0, X1, Xq, res_X, rho): + """ + Update the ADMM penalty parameter. + """ + z0, z1, zq = fun_AE(X0, X1, Xq) + res_eq = ( + xp.linalg.norm(bE0 - z0) / (1 + xp.linalg.norm(bE0)) + + xp.linalg.norm(bE1 - z1) / (1 + xp.linalg.norm(bE1)) + + xp.linalg.norm(bEq - zq) / (1 + xp.linalg.norm(bEq)) + ) + res_inq = xp.linalg.norm( + xp.maximum(bI - fun_AI(X0, X1), 0) + ) / (1 + abs(bI) * xp.sqrt(Ngrid * n_pairs)) + + p_resnorm = res_eq + res_inq + d_resnorm = res_X / (1 + normC) + if d_resnorm > ratio * p_resnorm: + rho = rho * factor + if d_resnorm < ratio * p_resnorm: + rho = rho / factor + return rho, p_resnorm, d_resnorm + + def print_updates(verbose): + """ + Log diagnostics for the current asymmetric ADMM iterate. + """ + if not verbose: + return + + obj_p = ( + xp.vdot(C0[:, idx_diag], X0[:, idx_diag]) + + xp.vdot(C1[:, idx_diag], X1[:, idx_diag]) + + 2 * xp.vdot(C0[:, idx_offdiag], X0[:, idx_offdiag]) + + 2 * xp.vdot(C1[:, idx_offdiag], X1[:, idx_offdiag]) + ) + obj_d = ( + xp.vdot(yE0, bE0) + + xp.vdot(yE1, bE1) + + 2 * xp.vdot(yEq, bEq) + + 2 * bI * xp.sum(yI) + ) + + z0, z1, zq = fun_AE(X0, X1, Xq) + res_eq = ( + xp.linalg.norm(bE0 - z0) / (1 + xp.linalg.norm(bE0)) + + xp.linalg.norm(bE1 - z1) / (1 + xp.linalg.norm(bE1)) + + xp.linalg.norm(bEq - zq) / (1 + xp.linalg.norm(bEq)) + ) + res_inq = xp.linalg.norm( + xp.maximum(bI - fun_AI(X0, X1), 0) + ) / (1 + abs(bI) * xp.sqrt(Ngrid * n_pairs)) + + res_psdX = 0 + for k in range(1, Lmax + 1): + tmp = self.mat_block( + X0[d0[k - 1] : d0[k]], + N, + k, + IDX_upper, + IDX_lower, + idx_offdiag, + ) + res_psdX += xp.linalg.norm(self.psd_projection(-tmp)) + + tmp = self.mat_block( + X1[d1[k - 1] : d1[k]], + N, + k + 1, + IDX_upper, + IDX_lower, + idx_offdiag, + ) + res_psdX += xp.linalg.norm(self.psd_projection(-tmp)) + + res_psdX /= 1 + xp.linalg.norm(X0) + xp.linalg.norm(X1) + + Xq_blocks = Xq.T.reshape(n_pairs, 4, 4) + res_psdQ = xp.linalg.norm( + self.psd_projection(-Xq_blocks), axis=(-2, -1) + ).sum() + res_psdQ /= 1 + xp.linalg.norm(Xq) + + normS = xp.sqrt( + xp.linalg.norm(S0) ** 2 + + xp.linalg.norm(S1) ** 2 + + xp.linalg.norm(Sq) ** 2 + ) + normX = xp.sqrt( + xp.linalg.norm(X0) ** 2 + + xp.linalg.norm(X1) ** 2 + + xp.linalg.norm(Xq) ** 2 + ) + p_res = res_eq + res_inq + res_psdX + res_psdQ + d_res = res_X / (1 + normC) + + logger.info( + "Iter %i" % t + + ": p_res=%1.5f" % p_res + + ", d_res=%1.5f" % d_res + + ", obj_primal=%1.2f" % obj_p + + ", obj_dual=%1.2f" % obj_d + + ", duality gap=%1.2f" % (obj_p - obj_d) + + "\n eq_res=%1.5f" % res_eq + + ", inq_res=%1.5f" % res_inq + + ", psd_res=%1.5f" % (res_psdX + res_psdQ) + + ", |S|=%1.2f" % normS + + ", |X|=%1.2f" % normX + ) + + yI = xp.zeros((Ngrid, n_pairs), dtype=np.float64) + yE0 = xp.zeros(bE0.shape, dtype=np.float64) + yE1 = xp.zeros(bE1.shape, dtype=np.float64) + yEq = xp.zeros(bEq.shape, dtype=np.float64) + + IDX = np.arange(3) + for t in range(max_iter): + #np.random.shuffle(IDX) + for idx in IDX: + if idx == 0: + S0, S1, Sq = update_S( + C0, C1, yE0, yE1, yEq, yI, X0, X1, Xq, rho + ) + elif idx == 1: + yE0, yE1, yEq = update_yE( + C0, C1, S0, S1, Sq, yI, X0, X1, Xq, rho + ) + else: + for _ in range(Nstep_yI): + yI = update_yI( + C0, + C1, + S0, + S1, + yE0, + yE1, + yEq, + yI, + X0, + X1, + rho, + Lambda, + ) + + X0, X1, Xq, res_X = update_X( + C0, C1, S0, S1, Sq, yE0, yE1, yEq, yI, X0, X1, Xq, rho + ) + if t % 100 == 0: + print_updates(verbose) + rho, p_resnorm, d_resnorm = update_rho(X0, X1, Xq, res_X, rho) + + X_admm = self.transform_coeff_back( + X0, X1, IDX_upper, IDX_lower, idx_offdiag + ) + for k in range(Lmax): + X_admm[k] = xp.asnumpy(X_admm[k]) + return X_admm + def admm_sym_J(self, C, verbose): """ Solve the symmetry-constrained NUG semidefinite relaxation using ADMM. @@ -1545,7 +1940,6 @@ def largest_eigenvalue(AI, Ngrid, N): Estimate the largest eigenvalue of the Fejér constraint operator. """ # find the largest eigenvalue of the operator AI - np.random.seed(0) z = xp.random.normal(0, 1, (Ngrid, N**2)) Lambda = 0 From 897c2139376f69df8d0b78aa0721f117c4528bb2 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 14 Aug 2026 11:38:08 -0400 Subject: [PATCH 062/161] build clmatrix from X_est for C1 euler est --- src/aspire/abinitio/commonline_nug.py | 91 ++++++++++++++++++++------- 1 file changed, 69 insertions(+), 22 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 8474c7b428..c51004da31 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -7,7 +7,12 @@ from aspire.numeric import xp from aspire.operators import PolarFT from aspire.utils import Rotation, cart2sph, complex_type -from aspire.volume import CnSymmetryGroup, DnSymmetryGroup, IdentitySymmetryGroup, SymmetryGroup +from aspire.volume import ( + CnSymmetryGroup, + DnSymmetryGroup, + IdentitySymmetryGroup, + SymmetryGroup, +) from .commonline_utils import _generate_shift_phase_and_filter, saff_kuijlaars @@ -115,7 +120,9 @@ def __init__( logger.info(f"Using provided symmetry: {symmetry}") self.sym_grp = SymmetryGroup.parse(symmetry) - if not isinstance(self.sym_grp, (IdentitySymmetryGroup, CnSymmetryGroup, DnSymmetryGroup)): + if not isinstance( + self.sym_grp, (IdentitySymmetryGroup, CnSymmetryGroup, DnSymmetryGroup) + ): raise ValueError( f"This algorithm supports Cn, Dn, and asymmetric molecules. Found {str(self.sym_grp)}." ) @@ -441,9 +448,7 @@ def fun_AI(X0, X1): The plain formulation does not apply these inequalities to the fixed diagonal representation blocks. """ - tmp = xp.concatenate( - (X0[:, idx_offdiag], X1[:, idx_offdiag]), axis=0 - ) + tmp = xp.concatenate((X0[:, idx_offdiag], X1[:, idx_offdiag]), axis=0) return AI_mat_offdiag @ tmp def fun_AIT(yI): @@ -582,9 +587,9 @@ def update_rho(X0, X1, Xq, res_X, rho): + xp.linalg.norm(bE1 - z1) / (1 + xp.linalg.norm(bE1)) + xp.linalg.norm(bEq - zq) / (1 + xp.linalg.norm(bEq)) ) - res_inq = xp.linalg.norm( - xp.maximum(bI - fun_AI(X0, X1), 0) - ) / (1 + abs(bI) * xp.sqrt(Ngrid * n_pairs)) + res_inq = xp.linalg.norm(xp.maximum(bI - fun_AI(X0, X1), 0)) / ( + 1 + abs(bI) * xp.sqrt(Ngrid * n_pairs) + ) p_resnorm = res_eq + res_inq d_resnorm = res_X / (1 + normC) @@ -620,9 +625,9 @@ def print_updates(verbose): + xp.linalg.norm(bE1 - z1) / (1 + xp.linalg.norm(bE1)) + xp.linalg.norm(bEq - zq) / (1 + xp.linalg.norm(bEq)) ) - res_inq = xp.linalg.norm( - xp.maximum(bI - fun_AI(X0, X1), 0) - ) / (1 + abs(bI) * xp.sqrt(Ngrid * n_pairs)) + res_inq = xp.linalg.norm(xp.maximum(bI - fun_AI(X0, X1), 0)) / ( + 1 + abs(bI) * xp.sqrt(Ngrid * n_pairs) + ) res_psdX = 0 for k in range(1, Lmax + 1): @@ -688,16 +693,12 @@ def print_updates(verbose): IDX = np.arange(3) for t in range(max_iter): - #np.random.shuffle(IDX) + # np.random.shuffle(IDX) for idx in IDX: if idx == 0: - S0, S1, Sq = update_S( - C0, C1, yE0, yE1, yEq, yI, X0, X1, Xq, rho - ) + S0, S1, Sq = update_S(C0, C1, yE0, yE1, yEq, yI, X0, X1, Xq, rho) elif idx == 1: - yE0, yE1, yEq = update_yE( - C0, C1, S0, S1, Sq, yI, X0, X1, Xq, rho - ) + yE0, yE1, yEq = update_yE(C0, C1, S0, S1, Sq, yI, X0, X1, Xq, rho) else: for _ in range(Nstep_yI): yI = update_yI( @@ -722,9 +723,7 @@ def print_updates(verbose): print_updates(verbose) rho, p_resnorm, d_resnorm = update_rho(X0, X1, Xq, res_X, rho) - X_admm = self.transform_coeff_back( - X0, X1, IDX_upper, IDX_lower, idx_offdiag - ) + X_admm = self.transform_coeff_back(X0, X1, IDX_upper, IDX_lower, idx_offdiag) for k in range(Lmax): X_admm[k] = xp.asnumpy(X_admm[k]) return X_admm @@ -1517,7 +1516,9 @@ def euler_est(self): Recover rotations and Euler angles using the estimator for the configured symmetry group. """ X_est = self.X_est - if isinstance(self.sym_grp, CnSymmetryGroup): + if isinstance(self.sym_grp, IdentitySymmetryGroup): + R_est, Euler_est = self.euler_est_C1(X_est[0]) + elif isinstance(self.sym_grp, CnSymmetryGroup): R_est, Euler_est = self.euler_est_Cm(X_est[0], X_est[self.n_sym - 1]) elif isinstance(self.sym_grp, DnSymmetryGroup): R_est, Euler_est = self.euler_est_Dm(X_est) @@ -1525,6 +1526,52 @@ def euler_est(self): self.Euler_est = Euler_est self.rotations = R_est.astype(self.dtype) + def euler_est_C1(self, X1): + """ + Recover Euler angles for asymmetric molecules. + + :param X1: Relaxed degree-one representation matrix. + + :return: Estimated rotation matrices and Euler angles. + """ + clmatrix = self.build_commonline_matrix_from_X1(X1) + + def build_commonline_matrix_from_X1(self, X1): + clmatrix = -np.ones((self.n_img, self.n_img), dtype=self.dtype) + + [T, Tinv] = self.complex2real(1) + + for i in range(1, self.n_img): + for j in range(i): + Rij = X1[ + 3 * i : 3 * (i + 1), + 3 * j : 3 * (j + 1), + ] + + Wij = T @ Rij @ Tinv + + gamma_ij = (np.angle(Wij[0, 2]) + np.angle(Wij[2, 2])) / 2 + + alpha_ij = (np.angle(Wij[2, 0]) + np.angle(Wij[2, 2])) / 2 + + theta_i = -alpha_ij - np.pi / 2 + theta_j = gamma_ij - np.pi / 2 + + idx_i = self.n_theta * theta_i / (2 * np.pi) + idx_j = self.n_theta * theta_j / (2 * np.pi) + + clmatrix[i, j] = np.round(idx_i) % self.n_theta + clmatrix[j, i] = np.round(idx_j) % self.n_theta + + half = self.n_theta // 2 + + if clmatrix[j, i] >= half: + clmatrix[j, i] -= half + clmatrix[i, j] += half + clmatrix[i, j] %= self.n_theta + + return clmatrix + def euler_est_Cm(self, X1, XS): """ Recover Euler angles for cyclic symmetry. From ed19b7630653069e50b1f186ac063ed921914381 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 18 Aug 2026 15:24:07 -0400 Subject: [PATCH 063/161] populate upper tri with clmatrix loop. Port scalar rotratio_eulerangle. --- src/aspire/abinitio/commonline_nug.py | 12 ++++---- src/aspire/abinitio/sync_voting.py | 43 +++++++++++++++++++++++++++ 2 files changed, 49 insertions(+), 6 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index c51004da31..bded99938a 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1541,8 +1541,8 @@ def build_commonline_matrix_from_X1(self, X1): [T, Tinv] = self.complex2real(1) - for i in range(1, self.n_img): - for j in range(i): + for i in range(self.n_img - 1): + for j in range(i + 1, self.n_img): Rij = X1[ 3 * i : 3 * (i + 1), 3 * j : 3 * (j + 1), @@ -1565,10 +1565,10 @@ def build_commonline_matrix_from_X1(self, X1): half = self.n_theta // 2 - if clmatrix[j, i] >= half: - clmatrix[j, i] -= half - clmatrix[i, j] += half - clmatrix[i, j] %= self.n_theta + if clmatrix[i, j] >= half: + clmatrix[i, j] -= half + clmatrix[j, i] += half + clmatrix[j, i] %= self.n_theta return clmatrix diff --git a/src/aspire/abinitio/sync_voting.py b/src/aspire/abinitio/sync_voting.py index b7d9620840..b90d22dba5 100644 --- a/src/aspire/abinitio/sync_voting.py +++ b/src/aspire/abinitio/sync_voting.py @@ -59,6 +59,49 @@ def _syncmatrix_ij_vote_3n( return rot +def _rotratio_eulerangle(clmatrix, n_theta): + """ + Given a 3x3 common lines matrix, where the index of each common line is + between 1 and n_theta, compute the rotation that takes image 1 to image 2. + + :param clmatrix: A 3-by-3 common lines matrix. + :param n_theta: Number of angular samples used to discretize the common + lines. + :return: A 3-by-3 rotation matrix taking the first image to the second, + or ``None`` if the common lines form a degenerate triangle. + """ + + # Prepare the theta values from the differences of common-line indices. + # + # These correspond to idx3, idx2, and idx1, respectively, in the + # original MATLAB implementation. + cl_diff1 = np.asarray([clmatrix[0, 2] - clmatrix[0, 1]]) + cl_diff2 = np.asarray([clmatrix[1, 2] - clmatrix[1, 0]]) + cl_diff3 = np.asarray([clmatrix[2, 1] - clmatrix[2, 0]]) + + # Calculate the cos values of rotation angles between the two images. + c_alpha, _ = _get_cos_phis( + cl_diff1, + cl_diff2, + cl_diff3, + n_theta, + sync=False, + ) + + if len(c_alpha) == 0: + return None + + alpha = np.arccos(c_alpha[0]) + + # Convert the Euler angles with ZYZ conversion to rotation matrices + angles = np.zeros(3) + angles[0] = clmatrix[0, 1] * 2 * np.pi / n_theta + np.pi / 2 + angles[1] = alpha + angles[2] = -np.pi / 2 - clmatrix[1, 0] * 2 * np.pi / n_theta + + return Rotation.from_euler(angles).matrices + + def _rotratio_eulerangle_vec(clmatrix, i, j, good_k, n_theta): """ Compute the rotation that takes image i to image j From 1d41afc52176b0a622b83b88b5b7fd58494fdc32 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 18 Aug 2026 15:52:49 -0400 Subject: [PATCH 064/161] tox --- src/aspire/abinitio/commonline_nug.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index bded99938a..031cb4d469 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1536,6 +1536,8 @@ def euler_est_C1(self, X1): """ clmatrix = self.build_commonline_matrix_from_X1(X1) + return clmatrix + def build_commonline_matrix_from_X1(self, X1): clmatrix = -np.ones((self.n_img, self.n_img), dtype=self.dtype) From eabd1d6a0d5d738134ba68d82f76ec5511085cf3 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 19 Aug 2026 10:08:42 -0400 Subject: [PATCH 065/161] fix symmetry log message/test --- src/aspire/abinitio/commonline_nug.py | 2 +- tests/test_nug.py | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 031cb4d469..fe715da30e 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -124,7 +124,7 @@ def __init__( self.sym_grp, (IdentitySymmetryGroup, CnSymmetryGroup, DnSymmetryGroup) ): raise ValueError( - f"This algorithm supports Cn, Dn, and asymmetric molecules. Found {str(self.sym_grp)}." + f"This algorithm supports cyclic, dihedral, and asymmetric molecules. Found {str(self.sym_grp)}." ) self.sym_euler = self.sym_grp.rotations.angles diff --git a/tests/test_nug.py b/tests/test_nug.py index 5ae3fa4025..136cb85e1c 100644 --- a/tests/test_nug.py +++ b/tests/test_nug.py @@ -183,7 +183,9 @@ def test_unspupported_symmetry_raises(dtype): vol = TSymmetricVolume(L=16, C=1, K=10, dtype=dtype).generate() src = Simulation(n=3, vols=vol) - with pytest.raises(ValueError, match="supports cyclic or dihedral symmetry"): + with pytest.raises( + ValueError, match="supports cyclic, dihedral, and asymmetric molecules" + ): _ = CommonlineNUG(src) From 984a51fbbdc6b97f293677b88f72c6dfb4261675 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 19 Aug 2026 14:20:00 -0400 Subject: [PATCH 066/161] port of matlab cryo_syncmatrix --- src/aspire/abinitio/sync_voting.py | 114 ++++++++++++++++++++++++++++- 1 file changed, 112 insertions(+), 2 deletions(-) diff --git a/src/aspire/abinitio/sync_voting.py b/src/aspire/abinitio/sync_voting.py index b90d22dba5..aa926c8368 100644 --- a/src/aspire/abinitio/sync_voting.py +++ b/src/aspire/abinitio/sync_voting.py @@ -65,8 +65,8 @@ def _rotratio_eulerangle(clmatrix, n_theta): between 1 and n_theta, compute the rotation that takes image 1 to image 2. :param clmatrix: A 3-by-3 common lines matrix. - :param n_theta: Number of angular samples used to discretize the common - lines. + :param n_theta: The number of points in the theta direction (common lines). + :return: A 3-by-3 rotation matrix taking the first image to the second, or ``None`` if the common lines form a degenerate triangle. """ @@ -102,6 +102,116 @@ def _rotratio_eulerangle(clmatrix, n_theta): return Rotation.from_euler(angles).matrices +def _syncmatrix(clmatrix, n_theta, dtype): + """ + Construct the CryoEM synchronization matrix, given a common lines matrix + clmatrix, that was constructed using angular resolution of n_theta radial + lines per image. + + :param clmatrix: A commonline matrix + :param n_theta: The number of points in the theta direction (common lines). + :param dtype: dtype of the synchronization matrix. + + :return: The 2N-by-2N synchronization matrix, where N is the number of + images. + """ + if clmatrix.ndim != 2 or clmatrix.shape[0] != clmatrix.shape[1]: + raise ValueError("clmatrix must be a square matrix.") + + n_img = clmatrix.shape[0] + + # This tolerance is relevant only if n_theta is very large (1.0e15), + # to check for accuracy. Otherwise, ignore it and the resulting output. + tol = 1.0e-12 + + # S is initialized with identity 2-by-2 diagonal blocks. + # + # Use an n_img-by-2-by-n_img-by-2 view so that synchronization blocks + # can be addressed directly as S[i, :, j, :]. + syncmatrix = np.eye( + 2 * n_img, + dtype=dtype, + ).reshape(n_img, 2, n_img, 2) + + for i in range(n_img - 1): + for j in range(i + 1, n_img): + # Each triplet (i, j, k) defines some rotation matrix from + # image i to image j. Not all images k give rise to such a + # rotation since, for example, the resulting triangle can be + # too small. + # + # Store all valid rotation matrices. This is not strictly + # necessary, as the matrices could be processed on the fly. + rotations = [] + + for k in range(n_img): + if k == i or k == j: + continue + + indices = [i, j, k] + cl_triplet = clmatrix[indices][:, indices] + + rot = _rotratio_eulerangle(cl_triplet, n_theta) + + # _rotratio_eulerangle returns None when the common lines + # form a degenerate triangle. + if rot is not None: + # Average over the handedness ambiguity. + rot = (rot + J_conjugate(rot)) / 2 + rotations.append(rot) + + # Merge all rotations computed by all triplets (i, j, k) into + # a single rotation that takes the frame of image i into the + # frame of image j. + if rotations: + rotations = np.stack(rotations) + rot_mean = np.mean(rotations, axis=0) + + # Preserve the consistency diagnostic from the MATLAB code. + diff = rotations - rot_mean + denominator = np.linalg.norm(rotations) + + err = np.linalg.norm(diff) / denominator + if err > tol: + logger.debug( + f"Inconsistent rotations: [i={i}, j={j}] " + f"err={err:e}, tol={tol:e}" + ) + else: + # Images i and j may correspond to the same viewing + # direction and differ only by an in-plane rotation. No + # triangle can then be formed using any third image k. + # + # The MATLAB implementation optionally uses reference + # rotations here. Reference rotations are not available to + # the estimator, so put zero. + rot_mean = np.zeros((3, 3), dtype=dtype) + + # S is a 2N-by-2N matrix containing N-by-N blocks of size 2-by-2. + # The (i, j) block is + # + # [r11 r12] + # [r21 r22], + # + # where r_kl is the inner product of column k of rotation i and + # column l of rotation j. + # + # To extract the required block, note that the handedness- + # averaged relative rotation has the form + # + # [r11 r12 0] + # [r21 r22 0] + # [ 0 0 r33]. + rot_block = rot_mean[:2, :2] + + # Put the block in the correct location in S and its transpose + # in the symmetric location. + syncmatrix[i, :, j, :] = rot_block + syncmatrix[j, :, i, :] = rot_block.T + + return syncmatrix.reshape(2 * n_img, 2 * n_img) + + def _rotratio_eulerangle_vec(clmatrix, i, j, good_k, n_theta): """ Compute the rotation that takes image i to image j From 26874596061d5e069ec6de3ee87484068b5ef0c2 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 19 Aug 2026 14:26:34 -0400 Subject: [PATCH 067/161] tox --- src/aspire/abinitio/sync_voting.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/aspire/abinitio/sync_voting.py b/src/aspire/abinitio/sync_voting.py index aa926c8368..c171a4a606 100644 --- a/src/aspire/abinitio/sync_voting.py +++ b/src/aspire/abinitio/sync_voting.py @@ -2,7 +2,7 @@ import numpy as np -from aspire.utils import Rotation +from aspire.utils import J_conjugate, Rotation logger = logging.getLogger(__name__) From 6a7323d0e48e3bb27a4f57139a9329dff71bbdf0 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 20 Aug 2026 09:30:52 -0400 Subject: [PATCH 068/161] move syncrotations to sync_voting.py --- src/aspire/abinitio/commonline_sync.py | 93 ++----------------------- src/aspire/abinitio/sync_voting.py | 96 +++++++++++++++++++++++++- 2 files changed, 101 insertions(+), 88 deletions(-) diff --git a/src/aspire/abinitio/commonline_sync.py b/src/aspire/abinitio/commonline_sync.py index c40751cdbb..93b1930bd2 100644 --- a/src/aspire/abinitio/commonline_sync.py +++ b/src/aspire/abinitio/commonline_sync.py @@ -3,9 +3,11 @@ import numpy as np from aspire.abinitio import CLOrient3D -from aspire.abinitio.sync_voting import _rotratio_eulerangle_vec, _vote_ij -from aspire.utils import nearest_rotations -from aspire.utils.matlab_compat import stable_eigsh +from aspire.abinitio.sync_voting import ( + _rotratio_eulerangle_vec, + _syncrotations, + _vote_ij, +) logger = logging.getLogger(__name__) @@ -78,90 +80,7 @@ def estimate_rotations(self): if self.syncmatrix is None: self.syncmatrix_vote() - S = self.syncmatrix - sz = S.shape - assert sz[0] == sz[1], "syncmatrix must be a square matrix." - assert sz[0] % 2 == 0, "syncmatrix must be a square matrix of size 2Kx2K." - - n_img = sz[0] // 2 - - # S is a 2Kx2K matrix (K=n_img), containing KxK blocks of size 2x2. - # The [i,j] block is given by [r11 r12; r12 r22], where - # r_{kl}=, k,l=1,2, namely, the dot product of - # column k of R_{i} and columns l of R_{j}. Thus, given the true - # rotations R_{1},...,R_{K}, S is decomposed as S=W^{T}W where - # W=(R_{1}^{1},R_{1}^{2},...,R_{K}^{1},R_{K}^{2}), where R_{j}^{k} - # the k column of R_{j}. Therefore, S is a rank-3 matrix, and thus, it - # three eigenvectors that correspond to non-zero eigenvalues, are linear - # combinations of the column space of S, namely, W^{T}. - - # Extract three eigenvectors corresponding to non-zero eigenvalues. - d, v = stable_eigsh(S, 10) - sort_idx = np.argsort(-d) - logger.info( - f"Top 10 eigenvalues from synchronization voting matrix: {d[sort_idx]}" - ) - - # Only need the top 3 eigen-vectors. - v = v[:, sort_idx[:3]] - # According to the structure of W^{T} above, the odd rows of V, denoted V1, - # are a linear combination of the vectors R_{i}^{1}, i=1,...,K, that is of - # column 1 of all rotation matrices. Similarly, the even rows of V, - # denoted, V2, are linear combinations of R_{i}^{1}, i=1,...,K. - v1 = v[: 2 * n_img : 2].T.copy() - v2 = v[1 : 2 * n_img : 2].T.copy() - - # We look for a linear transformation (3 x 3 matrix) A such that - # A*V1'=R1 and A*V2=R2 are the columns of the rotations matrices. - # Therefore: - # V1 * A'*A V1' = 1 - # V2 * A'*A V2' = 1 - # V1 * A'*A V2' = 0 - # These are 3*K linear equations for 9 matrix entries of A'*A - # Actually, there are only 6 unknown variables, because A'*A is symmetric. - # So we will truncate from 9 variables to 6 variables corresponding - # to the upper half of the matrix A'*A - truncated_equations = np.zeros((3 * n_img, 6), dtype=self.dtype) - k = 0 - for i in range(3): - for j in range(i, 3): - truncated_equations[0::3, k] = v1[i] * v1[j] - truncated_equations[1::3, k] = v2[i] * v2[j] - truncated_equations[2::3, k] = v1[i] * v2[j] - k += 1 - - # b = [1 1 0 1 1 0 ...]' is the right hand side vector - b = np.ones(3 * n_img) - b[2::3] = 0 - - # Find the least squares approximation of A'*A in vector form - ATA_vec = np.linalg.lstsq(truncated_equations, b, rcond=None)[0] - - # Construct the matrix A'*A from the vectorized matrix. - ATA = np.zeros((3, 3), dtype=self.dtype) - upper_mask = np.triu_indices(3) - ATA[upper_mask] = ATA_vec - lower_mask = np.tril_indices(3) - ATA[lower_mask] = ATA.T[lower_mask] - - # The Cholesky decomposition of A'*A gives A - # numpy returns lower, matlab upper - a = np.linalg.cholesky(ATA) - - # Recover the rotations. The first two columns of all rotation - # matrices are given by unmixing V1 and V2 using A. The third - # column is the cross product of the first two. - r1 = np.dot(a, v1) - r2 = np.dot(a, v2) - r3 = np.cross(r1, r2, axis=0) - - rotations = np.empty((n_img, 3, 3), dtype=self.dtype) - rotations[:, :, 0] = r1.T - rotations[:, :, 1] = r2.T - rotations[:, :, 2] = r3.T - # Make sure that we got rotations by enforcing R to be - # a rotation (in case the error is large) - self.rotations = nearest_rotations(rotations) + self.rotations = _syncrotations(self.syncmatrix) def syncmatrix_vote(self): """ diff --git a/src/aspire/abinitio/sync_voting.py b/src/aspire/abinitio/sync_voting.py index c171a4a606..a6f362fce4 100644 --- a/src/aspire/abinitio/sync_voting.py +++ b/src/aspire/abinitio/sync_voting.py @@ -2,7 +2,8 @@ import numpy as np -from aspire.utils import J_conjugate, Rotation +from aspire.utils import J_conjugate, Rotation, nearest_rotations +from aspire.utils.matlab_compat import stable_eigsh logger = logging.getLogger(__name__) @@ -59,6 +60,99 @@ def _syncmatrix_ij_vote_3n( return rot +def _syncrotations(S): + """ + Compute the rotations from the syncronization matrix S. + + :param S: A 2Kx2K synchronization matrix. + + :return: Kx3x3 rotations. + """ + sz = S.shape + dtype = S.dtype + assert sz[0] == sz[1], "syncmatrix must be a square matrix." + assert sz[0] % 2 == 0, "syncmatrix must be a square matrix of size 2Kx2K." + + n_img = sz[0] // 2 + + # S is a 2Kx2K matrix (K=n_img), containing KxK blocks of size 2x2. + # The [i,j] block is given by [r11 r12; r12 r22], where + # r_{kl}=, k,l=1,2, namely, the dot product of + # column k of R_{i} and columns l of R_{j}. Thus, given the true + # rotations R_{1},...,R_{K}, S is decomposed as S=W^{T}W where + # W=(R_{1}^{1},R_{1}^{2},...,R_{K}^{1},R_{K}^{2}), where R_{j}^{k} + # the k column of R_{j}. Therefore, S is a rank-3 matrix, and thus, it + # three eigenvectors that correspond to non-zero eigenvalues, are linear + # combinations of the column space of S, namely, W^{T}. + + # Extract three eigenvectors corresponding to non-zero eigenvalues. + d, v = stable_eigsh(S, 10) + sort_idx = np.argsort(-d) + logger.info(f"Top 10 eigenvalues from synchronization voting matrix: {d[sort_idx]}") + + # Only need the top 3 eigen-vectors. + v = v[:, sort_idx[:3]] + # According to the structure of W^{T} above, the odd rows of V, denoted V1, + # are a linear combination of the vectors R_{i}^{1}, i=1,...,K, that is of + # column 1 of all rotation matrices. Similarly, the even rows of V, + # denoted, V2, are linear combinations of R_{i}^{1}, i=1,...,K. + v1 = v[: 2 * n_img : 2].T.copy() + v2 = v[1 : 2 * n_img : 2].T.copy() + + # We look for a linear transformation (3 x 3 matrix) A such that + # A*V1'=R1 and A*V2=R2 are the columns of the rotations matrices. + # Therefore: + # V1 * A'*A V1' = 1 + # V2 * A'*A V2' = 1 + # V1 * A'*A V2' = 0 + # These are 3*K linear equations for 9 matrix entries of A'*A + # Actually, there are only 6 unknown variables, because A'*A is symmetric. + # So we will truncate from 9 variables to 6 variables corresponding + # to the upper half of the matrix A'*A + truncated_equations = np.zeros((3 * n_img, 6), dtype=dtype) + k = 0 + for i in range(3): + for j in range(i, 3): + truncated_equations[0::3, k] = v1[i] * v1[j] + truncated_equations[1::3, k] = v2[i] * v2[j] + truncated_equations[2::3, k] = v1[i] * v2[j] + k += 1 + + # b = [1 1 0 1 1 0 ...]' is the right hand side vector + b = np.ones(3 * n_img) + b[2::3] = 0 + + # Find the least squares approximation of A'*A in vector form + ATA_vec = np.linalg.lstsq(truncated_equations, b, rcond=None)[0] + + # Construct the matrix A'*A from the vectorized matrix. + ATA = np.zeros((3, 3), dtype=dtype) + upper_mask = np.triu_indices(3) + ATA[upper_mask] = ATA_vec + lower_mask = np.tril_indices(3) + ATA[lower_mask] = ATA.T[lower_mask] + + # The Cholesky decomposition of A'*A gives A + # numpy returns lower, matlab upper + a = np.linalg.cholesky(ATA) + + # Recover the rotations. The first two columns of all rotation + # matrices are given by unmixing V1 and V2 using A. The third + # column is the cross product of the first two. + r1 = np.dot(a, v1) + r2 = np.dot(a, v2) + r3 = np.cross(r1, r2, axis=0) + + rotations = np.empty((n_img, 3, 3), dtype=dtype) + rotations[:, :, 0] = r1.T + rotations[:, :, 1] = r2.T + rotations[:, :, 2] = r3.T + + # Make sure that we got rotations by enforcing R to be + # a rotation (in case the error is large) + return nearest_rotations(rotations) + + def _rotratio_eulerangle(clmatrix, n_theta): """ Given a 3x3 common lines matrix, where the index of each common line is From 8b80350f0f5b87fa84232168eee31316d4bd2ffb Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 24 Aug 2026 10:28:45 -0400 Subject: [PATCH 069/161] add syncmatrix_from_X1 --- src/aspire/abinitio/commonline_nug.py | 63 +++++++++++++++++++++++++-- src/aspire/abinitio/sync_voting.py | 2 +- 2 files changed, 60 insertions(+), 5 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index fe715da30e..e7b926e613 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -4,6 +4,7 @@ from scipy.special import factorial from aspire.abinitio import Orient3D +from aspire.abinitio.sync_voting import _syncmatrix, _syncrotations from aspire.numeric import xp from aspire.operators import PolarFT from aspire.utils import Rotation, cart2sph, complex_type @@ -1517,13 +1518,13 @@ def euler_est(self): """ X_est = self.X_est if isinstance(self.sym_grp, IdentitySymmetryGroup): - R_est, Euler_est = self.euler_est_C1(X_est[0]) + R_est = self.euler_est_C1(X_est[0]) elif isinstance(self.sym_grp, CnSymmetryGroup): R_est, Euler_est = self.euler_est_Cm(X_est[0], X_est[self.n_sym - 1]) elif isinstance(self.sym_grp, DnSymmetryGroup): R_est, Euler_est = self.euler_est_Dm(X_est) - self.Euler_est = Euler_est + # self.Euler_est = Euler_est self.rotations = R_est.astype(self.dtype) def euler_est_C1(self, X1): @@ -1534,9 +1535,63 @@ def euler_est_C1(self, X1): :return: Estimated rotation matrices and Euler angles. """ - clmatrix = self.build_commonline_matrix_from_X1(X1) + ### clmatrix = self.build_commonline_matrix_from_X1(X1) + ### S = _syncmatrix(clmatrix, self.n_theta, self.dtype) + S = self.syncmatrix_from_X1(X1) + rots = _syncrotations(S) + return rots - return clmatrix + def syncmatrix_from_X1(self, X1): + """ + Construct the synchronization matrix directly from the degree-one + NUG solution. + + :param X1: Degree-one NUG solution of shape (3 * n_img, 3 * n_img). + + :return: Synchronization matrix of shape (2 * n_img, 2 * n_img). + """ + # View X1 as an n_img-by-n_img array of 3-by-3 blocks. + X1_blocks = X1.reshape(self.n_img, 3, self.n_img, 3) + + # The NUG block decomposition gives each degree-one block X1_ij the + # structure + # + # [a 0 b] + # [0 c 0] + # [d 0 e]. + # + # The one-dimensional component occupies index 1, while the + # two-dimensional component used to construct the synchronization + # matrix occupies indices 0 and 2. Extract that 2-by-2 component from + # every image-pair block. + corner_blocks = X1_blocks[:, [0, 2], :, :] + corner_blocks = corner_blocks[:, :, :, [0, 2]] + + # The extracted blocks use the degree-one real spherical-harmonic + # coordinate ordering (y, -x), while the synchronization matrix uses + # Cartesian ordering (x, y). If + # + # P = [[0, 1], + # [-1, 0]], + # + # maps Cartesian coordinates to the real-basis coordinates, then each + # block must be converted according to + # + # S_ij = P.T @ corner_ij @ P. + # + # The assignments below apply this change of basis to every image-pair + # block simultaneously. + syncmatrix = np.empty_like(corner_blocks) + syncmatrix[:, 0, :, 0] = corner_blocks[:, 1, :, 1] + syncmatrix[:, 0, :, 1] = -corner_blocks[:, 1, :, 0] + syncmatrix[:, 1, :, 0] = -corner_blocks[:, 0, :, 1] + syncmatrix[:, 1, :, 1] = corner_blocks[:, 0, :, 0] + + # Collapse the image and Cartesian-coordinate axes to obtain the + # conventional 2*n_img-by-2*n_img synchronization matrix. + syncmatrix = syncmatrix.reshape(2 * self.n_img, 2 * self.n_img) + + return syncmatrix def build_commonline_matrix_from_X1(self, X1): clmatrix = -np.ones((self.n_img, self.n_img), dtype=self.dtype) diff --git a/src/aspire/abinitio/sync_voting.py b/src/aspire/abinitio/sync_voting.py index a6f362fce4..2a1b294d2c 100644 --- a/src/aspire/abinitio/sync_voting.py +++ b/src/aspire/abinitio/sync_voting.py @@ -193,7 +193,7 @@ def _rotratio_eulerangle(clmatrix, n_theta): angles[1] = alpha angles[2] = -np.pi / 2 - clmatrix[1, 0] * 2 * np.pi / n_theta - return Rotation.from_euler(angles).matrices + return Rotation.from_euler(angles).matrices[0] def _syncmatrix(clmatrix, n_theta, dtype): From 6d2f2168e6164cb7527c03de362134d206be2e82 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 24 Aug 2026 14:59:48 -0400 Subject: [PATCH 070/161] Add form_ground_truth_X --- src/aspire/abinitio/commonline_nug.py | 53 +++++++++++++++++++++++++++ 1 file changed, 53 insertions(+) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index e7b926e613..43b8890063 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -2201,3 +2201,56 @@ def construct_AEq(self): AEq[16, [0, 5, 10, 15]] = 1 return AEq + + def form_ground_truth_X(self, euler_angles): + """ + Construct handedness-averaged ground-truth NUG representation matrices. + + This helper is intended for validating the relaxed ADMM solution when + ground-truth rotations are available. + + :param euler_angles: Ground-truth ZYZ Euler angles of shape (n_img, 3). + :return: List containing one ground-truth representation matrix for + each degree from 1 through Lmax. + """ + X_gt = [] + + # Internally, euler_angles correspond to R.T, so we adjust them + # here so X_gt corresponds to the ground truth rotations. + euler_angles = -euler_angles[:, ::-1] + + for k in range(1, self.Lmax + 1): + dk = 2 * k + 1 + + # Evaluate the degree-k Wigner representations. + wigner = self.WD(k, euler_angles) + + # Average the degree-k representation over the molecular symmetry + # group. For the asymmetric case, self.sym_euler contains only the + # identity and Ak is therefore the identity matrix. + Ak = np.mean(self.WD(k, self.sym_euler), axis=0) + + # Construct the globally handedness-conjugated solution. + Jk = np.ones(dk) + Jk[k + 1 :: 2] = -1 + Jk[k - 1 :: -2] = -1 + Jk = np.diag(Jk) + + wigner_J = Jk @ wigner @ Jk + + # Convert each representation to the real basis used by ADMM. + _, Tinv = self.complex2real(k) + wigner_real = Tinv @ wigner + wigner_J_real = Tinv @ wigner_J + + # Stack the image representations and construct their Gram matrices. + wigner_real = wigner_real.reshape(self.n_img * dk, dk) + wigner_J_real = wigner_J_real.reshape(self.n_img * dk, dk) + + Xk = wigner_real @ Ak @ wigner_real.conj().T + XJk = wigner_J_real @ Ak @ wigner_J_real.conj().T + + # Average the two globally indistinguishable handedness choices. + X_gt.append(np.real((Xk + XJk) / 2)) + + return X_gt From d68e8b875d680c21bb116b6a887b273cf9b952e8 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 25 Aug 2026 16:05:55 -0400 Subject: [PATCH 071/161] Add C1 NUG to testing --- tests/test_nug.py | 78 +++++++++++++++++++++++++---------------------- 1 file changed, 41 insertions(+), 37 deletions(-) diff --git a/tests/test_nug.py b/tests/test_nug.py index 136cb85e1c..755bbd21aa 100644 --- a/tests/test_nug.py +++ b/tests/test_nug.py @@ -2,21 +2,17 @@ import pytest from aspire.abinitio import CommonlineNUG, compare_rots_sym, g_sync +from aspire.downloader import emdb_2660 from aspire.source import Simulation from aspire.utils import mean_aligned_angular_distance from aspire.volume import CnSymmetricVolume, DnSymmetricVolume, TSymmetricVolume DTYPE = [np.float64, np.float32] RESOLUTION = [48, 49] -N_IMG = [15] +N_IMG = [30] OFFSETS = [0, None] -ORDER = [3, 4] -PR = [None] +SYMMETRY = ["C1", "C3", "C4", "D3", "D4"] SEED = 1980 -VOLUME = [ - CnSymmetricVolume, - DnSymmetricVolume, -] @pytest.fixture(params=DTYPE, ids=lambda x: f"dtype={x}", scope="module") @@ -39,18 +35,8 @@ def offsets(request): return request.param -@pytest.fixture(params=ORDER, ids=lambda x: f"order={x}", scope="module") -def order(request): - return request.param - - -@pytest.fixture(params=PR, ids=lambda x: f"proximal_refine={x}", scope="module") -def proximal_refine(request): - return request.param - - -@pytest.fixture(params=VOLUME, ids=lambda x: f"Volume={x}", scope="module") -def Volume(request): +@pytest.fixture(params=SYMMETRY, ids=lambda x: f"symmetry={x}", scope="module") +def symmetry(request): return request.param @@ -60,37 +46,53 @@ def Volume(request): @pytest.fixture(scope="module") -def source(n_img, resolution, dtype, offsets, order, Volume): - vol = Volume( - L=resolution, order=order, C=1, K=100, dtype=dtype, seed=SEED - ).generate() +def volume(symmetry, resolution, dtype): + if symmetry == "C1": + vol = emdb_2660().astype(dtype).downsample(resolution) + if symmetry == "C3": + vol = CnSymmetricVolume( + L=resolution, order=3, C=1, K=100, dtype=dtype, seed=SEED + ).generate() + if symmetry == "C4": + vol = CnSymmetricVolume( + L=resolution, order=4, C=1, K=100, dtype=dtype, seed=SEED + ).generate() + if symmetry == "D3": + vol = DnSymmetricVolume( + L=resolution, order=3, C=1, K=100, dtype=dtype, seed=SEED + ).generate() + if symmetry == "D4": + vol = DnSymmetricVolume( + L=resolution, order=4, C=1, K=100, dtype=dtype, seed=SEED + ).generate() + return vol + +@pytest.fixture(scope="module") +def source(n_img, offsets, volume): src = Simulation( n=n_img, - L=resolution, - vols=vol, + vols=volume, offsets=offsets, amplitudes=1, seed=SEED, - ) - src = src.cache() # Precompute image stack - + ).cache() # Precompute image stack return src @pytest.fixture(scope="module") -def orient_est(source, proximal_refine): +def orient_est(source): max_shift = 0 shift_step = 1 if source.offsets.all() != 0: - max_shift = 0.20 - shift_step = 0.25 + max_shift = 0.30 + shift_step = 0.5 orient_est = CommonlineNUG( source, max_shift=max_shift, shift_step=shift_step, max_iter=201, - pr_iters=proximal_refine, + pr_iters=None, verbose=False, ) orient_est.estimate_rotations() @@ -160,7 +162,7 @@ def test_estimate_rotations_pairwise(orient_est): MSE = compare_rots_sym( orient_est.rotations, orient_est.src.rotations, orient_est.sym_grp ) - np.testing.assert_array_less(MSE, 0.1) + np.testing.assert_array_less(MSE, 0.15) @pytest.mark.expensive @@ -170,10 +172,12 @@ def test_estimate_rotations(orient_est): truth, after symmetry synchronization and global alignment, are within 10 degrees. """ - gt_rots_synced = g_sync( - orient_est.rotations, orient_est.src.rotations, orient_est.sym_grp - ) - mean_aligned_angular_distance(orient_est.rotations, gt_rots_synced, 10.0) + gt_rots = orient_est.src.rotations + if orient_est.sym_grp.order > 1: + gt_rots = g_sync( + orient_est.rotations, orient_est.src.rotations, orient_est.sym_grp + ) + mean_aligned_angular_distance(orient_est.rotations, gt_rots, 10.0) def test_unspupported_symmetry_raises(dtype): From 92a7e3c04b7233af6a6b0a7ba3491e635e615a1e Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 25 Aug 2026 16:09:20 -0400 Subject: [PATCH 072/161] remove yutong implementation --- src/aspire/abinitio/commonline_nug.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 43b8890063..c3181df338 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -4,7 +4,7 @@ from scipy.special import factorial from aspire.abinitio import Orient3D -from aspire.abinitio.sync_voting import _syncmatrix, _syncrotations +from aspire.abinitio.sync_voting import _syncrotations from aspire.numeric import xp from aspire.operators import PolarFT from aspire.utils import Rotation, cart2sph, complex_type @@ -1535,8 +1535,6 @@ def euler_est_C1(self, X1): :return: Estimated rotation matrices and Euler angles. """ - ### clmatrix = self.build_commonline_matrix_from_X1(X1) - ### S = _syncmatrix(clmatrix, self.n_theta, self.dtype) S = self.syncmatrix_from_X1(X1) rots = _syncrotations(S) return rots From 97683ad7e98494fee5fcca58a803f4a6561940e7 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 26 Aug 2026 09:14:36 -0400 Subject: [PATCH 073/161] fix smoke test --- tests/test_nug.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/test_nug.py b/tests/test_nug.py index 755bbd21aa..c2cd8d0847 100644 --- a/tests/test_nug.py +++ b/tests/test_nug.py @@ -104,9 +104,12 @@ def orient_est(source): ######### +@pytest.mark.parametrize("Volume", [CnSymmetricVolume, DnSymmetricVolume]) def test_smoke(dtype, Volume): """ - Perform quick smoke test since other tests are long running. + Perform quick smoke test since other tests are expensive. + This runs through both symmetry cases and the proximal refine + code path. """ vol = Volume(L=32, order=3, C=1, dtype=dtype, seed=SEED).generate() From e7edb080f99d81182429f85cf8f93fbf0dbc4127 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 26 Aug 2026 09:21:19 -0400 Subject: [PATCH 074/161] supress expected gimbal-lock warning --- src/aspire/abinitio/commonline_nug.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index c3181df338..8a9296571b 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -128,7 +128,13 @@ def __init__( f"This algorithm supports cyclic, dihedral, and asymmetric molecules. Found {str(self.sym_grp)}." ) - self.sym_euler = self.sym_grp.rotations.angles + # Get Euler angles for the symmetry group. + # Supress expected gimbal-lock warning (due to identity matrix). + sym_rotations = Rotation( + self.sym_grp.matrices, + gimble_lock_warnings=False, + ) + self.sym_euler = sym_rotations.angles self.n_sym = len(self.sym_euler) # Set up proximal refinement terms From 2984df3d5d4bf36b6992ac9ad6c3b9d86a651f02 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 26 Aug 2026 09:38:22 -0400 Subject: [PATCH 075/161] Remove yutong rounding procedure and associated methods --- src/aspire/abinitio/commonline_nug.py | 36 ------ src/aspire/abinitio/sync_voting.py | 155 +------------------------- 2 files changed, 1 insertion(+), 190 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 8a9296571b..4aad3b3e4e 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1597,42 +1597,6 @@ def syncmatrix_from_X1(self, X1): return syncmatrix - def build_commonline_matrix_from_X1(self, X1): - clmatrix = -np.ones((self.n_img, self.n_img), dtype=self.dtype) - - [T, Tinv] = self.complex2real(1) - - for i in range(self.n_img - 1): - for j in range(i + 1, self.n_img): - Rij = X1[ - 3 * i : 3 * (i + 1), - 3 * j : 3 * (j + 1), - ] - - Wij = T @ Rij @ Tinv - - gamma_ij = (np.angle(Wij[0, 2]) + np.angle(Wij[2, 2])) / 2 - - alpha_ij = (np.angle(Wij[2, 0]) + np.angle(Wij[2, 2])) / 2 - - theta_i = -alpha_ij - np.pi / 2 - theta_j = gamma_ij - np.pi / 2 - - idx_i = self.n_theta * theta_i / (2 * np.pi) - idx_j = self.n_theta * theta_j / (2 * np.pi) - - clmatrix[i, j] = np.round(idx_i) % self.n_theta - clmatrix[j, i] = np.round(idx_j) % self.n_theta - - half = self.n_theta // 2 - - if clmatrix[i, j] >= half: - clmatrix[i, j] -= half - clmatrix[j, i] += half - clmatrix[j, i] %= self.n_theta - - return clmatrix - def euler_est_Cm(self, X1, XS): """ Recover Euler angles for cyclic symmetry. diff --git a/src/aspire/abinitio/sync_voting.py b/src/aspire/abinitio/sync_voting.py index 2a1b294d2c..2fdf3a8a3a 100644 --- a/src/aspire/abinitio/sync_voting.py +++ b/src/aspire/abinitio/sync_voting.py @@ -2,7 +2,7 @@ import numpy as np -from aspire.utils import J_conjugate, Rotation, nearest_rotations +from aspire.utils import Rotation, nearest_rotations from aspire.utils.matlab_compat import stable_eigsh logger = logging.getLogger(__name__) @@ -153,159 +153,6 @@ def _syncrotations(S): return nearest_rotations(rotations) -def _rotratio_eulerangle(clmatrix, n_theta): - """ - Given a 3x3 common lines matrix, where the index of each common line is - between 1 and n_theta, compute the rotation that takes image 1 to image 2. - - :param clmatrix: A 3-by-3 common lines matrix. - :param n_theta: The number of points in the theta direction (common lines). - - :return: A 3-by-3 rotation matrix taking the first image to the second, - or ``None`` if the common lines form a degenerate triangle. - """ - - # Prepare the theta values from the differences of common-line indices. - # - # These correspond to idx3, idx2, and idx1, respectively, in the - # original MATLAB implementation. - cl_diff1 = np.asarray([clmatrix[0, 2] - clmatrix[0, 1]]) - cl_diff2 = np.asarray([clmatrix[1, 2] - clmatrix[1, 0]]) - cl_diff3 = np.asarray([clmatrix[2, 1] - clmatrix[2, 0]]) - - # Calculate the cos values of rotation angles between the two images. - c_alpha, _ = _get_cos_phis( - cl_diff1, - cl_diff2, - cl_diff3, - n_theta, - sync=False, - ) - - if len(c_alpha) == 0: - return None - - alpha = np.arccos(c_alpha[0]) - - # Convert the Euler angles with ZYZ conversion to rotation matrices - angles = np.zeros(3) - angles[0] = clmatrix[0, 1] * 2 * np.pi / n_theta + np.pi / 2 - angles[1] = alpha - angles[2] = -np.pi / 2 - clmatrix[1, 0] * 2 * np.pi / n_theta - - return Rotation.from_euler(angles).matrices[0] - - -def _syncmatrix(clmatrix, n_theta, dtype): - """ - Construct the CryoEM synchronization matrix, given a common lines matrix - clmatrix, that was constructed using angular resolution of n_theta radial - lines per image. - - :param clmatrix: A commonline matrix - :param n_theta: The number of points in the theta direction (common lines). - :param dtype: dtype of the synchronization matrix. - - :return: The 2N-by-2N synchronization matrix, where N is the number of - images. - """ - if clmatrix.ndim != 2 or clmatrix.shape[0] != clmatrix.shape[1]: - raise ValueError("clmatrix must be a square matrix.") - - n_img = clmatrix.shape[0] - - # This tolerance is relevant only if n_theta is very large (1.0e15), - # to check for accuracy. Otherwise, ignore it and the resulting output. - tol = 1.0e-12 - - # S is initialized with identity 2-by-2 diagonal blocks. - # - # Use an n_img-by-2-by-n_img-by-2 view so that synchronization blocks - # can be addressed directly as S[i, :, j, :]. - syncmatrix = np.eye( - 2 * n_img, - dtype=dtype, - ).reshape(n_img, 2, n_img, 2) - - for i in range(n_img - 1): - for j in range(i + 1, n_img): - # Each triplet (i, j, k) defines some rotation matrix from - # image i to image j. Not all images k give rise to such a - # rotation since, for example, the resulting triangle can be - # too small. - # - # Store all valid rotation matrices. This is not strictly - # necessary, as the matrices could be processed on the fly. - rotations = [] - - for k in range(n_img): - if k == i or k == j: - continue - - indices = [i, j, k] - cl_triplet = clmatrix[indices][:, indices] - - rot = _rotratio_eulerangle(cl_triplet, n_theta) - - # _rotratio_eulerangle returns None when the common lines - # form a degenerate triangle. - if rot is not None: - # Average over the handedness ambiguity. - rot = (rot + J_conjugate(rot)) / 2 - rotations.append(rot) - - # Merge all rotations computed by all triplets (i, j, k) into - # a single rotation that takes the frame of image i into the - # frame of image j. - if rotations: - rotations = np.stack(rotations) - rot_mean = np.mean(rotations, axis=0) - - # Preserve the consistency diagnostic from the MATLAB code. - diff = rotations - rot_mean - denominator = np.linalg.norm(rotations) - - err = np.linalg.norm(diff) / denominator - if err > tol: - logger.debug( - f"Inconsistent rotations: [i={i}, j={j}] " - f"err={err:e}, tol={tol:e}" - ) - else: - # Images i and j may correspond to the same viewing - # direction and differ only by an in-plane rotation. No - # triangle can then be formed using any third image k. - # - # The MATLAB implementation optionally uses reference - # rotations here. Reference rotations are not available to - # the estimator, so put zero. - rot_mean = np.zeros((3, 3), dtype=dtype) - - # S is a 2N-by-2N matrix containing N-by-N blocks of size 2-by-2. - # The (i, j) block is - # - # [r11 r12] - # [r21 r22], - # - # where r_kl is the inner product of column k of rotation i and - # column l of rotation j. - # - # To extract the required block, note that the handedness- - # averaged relative rotation has the form - # - # [r11 r12 0] - # [r21 r22 0] - # [ 0 0 r33]. - rot_block = rot_mean[:2, :2] - - # Put the block in the correct location in S and its transpose - # in the symmetric location. - syncmatrix[i, :, j, :] = rot_block - syncmatrix[j, :, i, :] = rot_block.T - - return syncmatrix.reshape(2 * n_img, 2 * n_img) - - def _rotratio_eulerangle_vec(clmatrix, i, j, good_k, n_theta): """ Compute the rotation that takes image i to image j From d0ec0f4c591d401504e10cf5c0c01fa102b03438 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 26 Aug 2026 10:27:13 -0400 Subject: [PATCH 076/161] cleanup/rename rotation recovery procedure --- src/aspire/abinitio/commonline_nug.py | 28 ++++++++++++++------------- 1 file changed, 15 insertions(+), 13 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 4aad3b3e4e..d77eda8a5d 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -161,13 +161,13 @@ def _build_full_pft(self): def estimate_rotations(self): """ - Estimate rotations by computing NUG coefficients, solving the SDP relaxation, and recovering Euler angles. + Estimate rotations by computing NUG coefficients, solving the SDP relaxation, and recovering rotations. :return: Estimated rotation matrices. """ self.compute_coeff() self.perform_admm() - self.euler_est() + self.recover_rotations() return self.rotations ####################### @@ -1514,26 +1514,28 @@ def low_rank_proj(X, r_step): return current - ######################### - # Euler Estimation Step # - ######################### + ########################## + # Rotation Recovery Step # + ########################## - def euler_est(self): + def recover_rotations(self): """ - Recover rotations and Euler angles using the estimator for the configured symmetry group. + Recover rotations from the NUG representation matrices, using synchronization-based + rounding for asymmetric molecules and symmetry-specific Euler estimation for cyclic and + dihedral molecules. """ X_est = self.X_est if isinstance(self.sym_grp, IdentitySymmetryGroup): - R_est = self.euler_est_C1(X_est[0]) + R_est = self.rounding_C1(X_est[0]) elif isinstance(self.sym_grp, CnSymmetryGroup): - R_est, Euler_est = self.euler_est_Cm(X_est[0], X_est[self.n_sym - 1]) + R_est = self.euler_est_Cm(X_est[0], X_est[self.n_sym - 1]) elif isinstance(self.sym_grp, DnSymmetryGroup): - R_est, Euler_est = self.euler_est_Dm(X_est) + R_est = self.euler_est_Dm(X_est) # self.Euler_est = Euler_est self.rotations = R_est.astype(self.dtype) - def euler_est_C1(self, X1): + def rounding_C1(self, X1): """ Recover Euler angles for asymmetric molecules. @@ -1720,7 +1722,7 @@ def find_gamma(Xm, beta, alpha): Euler_est[:, 1] = find_beta(X1) Euler_est[:, 2] = find_gamma(XS, Euler_est[:, 1], Euler_est[:, 0]) R_est = Rotation.from_euler(Euler_est).matrices.transpose(0, 2, 1) - return R_est, Euler_est + return R_est def euler_est_Dm(self, X_est): """ @@ -1874,7 +1876,7 @@ def LS_D(W1, W2, W3, W4, Br, Bi): Euler_est[:, 2] = gamma_est R_est = Rotation.from_euler(Euler_est).matrices.transpose(0, 2, 1) - return R_est, Euler_est + return R_est #################### # Helper Functions # From a5c7d46ea2f0ca5f69c6001f064d75e90febbd92 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 26 Aug 2026 13:39:19 -0400 Subject: [PATCH 077/161] de-duplicate euler estimation code with helper functions --- src/aspire/abinitio/commonline_nug.py | 406 +++++++++++++++----------- 1 file changed, 237 insertions(+), 169 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index d77eda8a5d..953a823df7 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1610,34 +1610,9 @@ def euler_est_Cm(self, X1, XS): """ S = self.n_sym N = self.n_img - sym_euler = np.zeros((S, 3), dtype=np.float64) - for s in range(S): - sym_euler[s] = [2 * np.pi * s / S, 0, 0] - [T, Tinv] = self.complex2real(1) - X1 = ( - np.kron(np.eye(N, dtype=np.float64), T) - @ X1 - @ np.kron(np.eye(N, dtype=np.float64), Tinv) - ) - [T, Tinv] = self.complex2real(S) - XS = ( - np.kron(np.eye(N, dtype=np.float64), T) - @ XS - @ np.kron(np.eye(N, dtype=np.float64), Tinv) - ) - def find_phase(A, B): - # find a number c that minimizes ||cA-B||_F - Ar = np.real(A) - Ai = np.imag(A) - Br = np.real(B) - Bi = np.imag(B) - c = (np.vdot(Ar, Br) + np.vdot(Ai, Bi)) / ( - np.vdot(Ar, Ar) + np.vdot(Ai, Ai) - ) + 1j * (np.vdot(Ar, Bi) - np.vdot(Ai, Br)) / ( - np.vdot(Ar, Ar) + np.vdot(Ai, Ai) - ) - return c / abs(c) + X1 = self._real_to_complex_representation(X1, degree=1) + XS = self._real_to_complex_representation(XS, degree=S) def find_beta(X1): B1 = np.zeros((N, N), dtype=np.float64) @@ -1671,57 +1646,26 @@ def find_alpha(X1): idx = np.argmax(abs(evals)) Z = evecs[:, idx] * np.sqrt(abs(evals[idx])) - c = find_phase(Z[:, None] @ Z[:, None].T, ZZ) + c = self._find_phase(Z[:, None] @ Z[:, None].T, ZZ) Z = np.sqrt(c) * Z return np.angle(Z).astype(np.float64) - dk = 2 * S + 1 + alpha_est = find_alpha(X1) + beta_est = find_beta(X1) + gamma_est = self._estimate_gamma( + XS, + alpha_est, + beta_est, + degree=S, + pair_estimator=self._cyclic_gamma_pair, + ) + + R_est = self._assemble_rotation_estimates( + alpha_est, + beta_est, + gamma_est, + ) - def find_gamma(Xm, beta, alpha): - C = np.zeros((N, N), dtype=complex_type(np.float64)) - Jk = np.ones(dk) - Jk[S + 1 :: 2] = -1 - Jk[S - 1 :: -2] = -1 - Jk = np.diag(Jk) - ws = self.Wd(S, beta) - for i in range(N): - wi = ws[i] - for j in range(i + 1, N): - Di = np.exp(-1j * np.arange(-S, S + 1) * alpha[i]) - Dj = np.exp(-1j * np.arange(-S, S + 1) * alpha[j]) - Xijm = Xm[dk * i : dk * (i + 1), dk * j : dk * (j + 1)] - DXijmD = np.diag(Di.conj()) @ Xijm @ np.diag(Dj) - wj = ws[j] - C1 = ( - wi[:, 0][:, None] @ (wj[:, 0][:, None].T) - + Jk @ wi[:, 0][:, None] @ (wj[:, 0][:, None].T) @ Jk - ) / 2 - C2 = ( - wi[:, -1][:, None] @ (wj[:, -1][:, None].T) - + Jk @ wi[:, -1][:, None] @ (wj[:, -1][:, None].T) @ Jk - ) / 2 - C3 = ( - DXijmD - - ( - wi[:, S][:, None] @ (wj[:, S][:, None].T) - + Jk @ wi[:, S][:, None] @ (wj[:, S][:, None].T) @ Jk - ) - / 2 - ) - C[i, j] = np.vdot(C1 + C2, np.real(C3)) / np.vdot( - C1 + C2, C1 + C2 - ) + 1j * np.vdot(C1 - C2, np.imag(C3)) / np.vdot(C1 - C2, C1 - C2) - C += C.T.conj() + np.eye(N, dtype=np.float64) - evals, evecs = np.linalg.eigh(C) - idx = np.argmax(evals) - c = evecs[:, idx] * np.sqrt(evals[idx]) - return np.angle(c) / S - - Euler_est = np.zeros((N, 3), dtype=np.float64) - Euler_est[:, 0] = find_alpha(X1) - Euler_est[:, 1] = find_beta(X1) - Euler_est[:, 2] = find_gamma(XS, Euler_est[:, 1], Euler_est[:, 0]) - R_est = Rotation.from_euler(Euler_est).matrices.transpose(0, 2, 1) return R_est def euler_est_Dm(self, X_est): @@ -1732,33 +1676,13 @@ def euler_est_Dm(self, X_est): :return: Estimated rotation matrices and Euler angles. """ - X2 = X_est[1] S = self.sym_grp.order N = self.n_img - XS = X_est[S - 1] - dk = 2 * S + 1 - def find_alpha_beta(X2): - def find_phase(A, B): - # find a number c that minimizes ||cA-B||_F - Ar = np.real(A) - Ai = np.imag(A) - Br = np.real(B) - Bi = np.imag(B) - c = (np.vdot(Ar, Br) + np.vdot(Ai, Bi)) / ( - np.vdot(Ar, Ar) + np.vdot(Ai, Ai) - ) + 1j * (np.vdot(Ar, Bi) - np.vdot(Ai, Br)) / ( - np.vdot(Ar, Ar) + np.vdot(Ai, Ai) - ) - return c / abs(c) - - T, Tinv = self.complex2real(2) - X2 = ( - np.kron(np.eye(N, dtype=np.float64), T) - @ X2 - @ np.kron(np.eye(N, dtype=np.float64), Tinv) - ) + X2 = self._real_to_complex_representation(X_est[1], degree=2) + XS = self._real_to_complex_representation(X_est[S - 1], degree=S) + def find_alpha_beta(X2): B1 = np.zeros((N, N), dtype=np.float64) for i in range(N): for j in range(N): @@ -1796,87 +1720,231 @@ def find_phase(A, B): evals, evecs = np.linalg.eigh(Aminus) idx = np.argmax(abs(evals)) Z = evecs[:, idx] * np.sqrt(abs(evals[idx])) - c = find_phase(Z[:, None] @ Z[:, None].T, Aplus) + c = self._find_phase(Z[:, None] @ Z[:, None].T, Aplus) Z = np.sqrt(c) * Z alpha_est = (np.angle(Z)) % (2 * np.pi) return alpha_est, beta_est - def find_gamma(Xm, alpha, beta): - def LS_D(W1, W2, W3, W4, Br, Bi): - A = np.array( - [ - [np.vdot(W1 + W4, W1 + W4), np.vdot(W1 + W4, W2 + W3)], - [np.vdot(W2 + W3, W1 + W4), np.vdot(W2 + W3, W2 + W3)], - ] - ) - B = np.array([np.vdot(W1 + W4, Br), np.vdot(W2 + W3, Br)]) - a, c = np.linalg.lstsq(A, B)[0] + alpha_est, beta_est = find_alpha_beta(X2) + gamma_est = self._estimate_gamma( + XS, + alpha_est, + beta_est, + degree=S, + pair_estimator=self._dihedral_gamma_pair, + wrap=True, + ) - A = np.array( - [ - [np.vdot(W1 - W4, W1 - W4), np.vdot(W1 - W4, W3 - W2)], - [np.vdot(W1 - W4, W3 - W2), np.vdot(W3 - W2, W3 - W2)], - ] + R_est = self._assemble_rotation_estimates( + alpha_est, + beta_est, + gamma_est, + ) + + return R_est + + ############################ + # Euler Estimation Helpers # + ############################ + def _real_to_complex_representation(self, X, degree): + """ + Convert a degree-wise representation matrix from the real basis used + by ADMM to the complex Wigner basis. + + :param X: Representation matrix of shape + (n_img * (2 * degree + 1), n_img * (2 * degree + 1)). + :param degree: Wigner representation degree. + + :return: Representation matrix in the complex Wigner basis. + """ + T, Tinv = self.complex2real(degree) + identity = np.eye(self.n_img, dtype=np.float64) + + return np.kron(identity, T) @ X @ np.kron(identity, Tinv) + + @staticmethod + def _find_phase(A, B): + """ + Find the unit-modulus scalar c that minimizes ||cA - B||_F. + + :param A: First complex-valued matrix. + :param B: Second complex-valued matrix. + + :return: Unit-modulus complex phase. + """ + Ar = np.real(A) + Ai = np.imag(A) + Br = np.real(B) + Bi = np.imag(B) + + denominator = np.vdot(Ar, Ar) + np.vdot(Ai, Ai) + + c = (np.vdot(Ar, Br) + np.vdot(Ai, Bi)) / denominator + 1j * ( + np.vdot(Ar, Bi) - np.vdot(Ai, Br) + ) / denominator + + return c / abs(c) + + @staticmethod + def _handedness_matrix(degree): + """ + Construct the handedness-conjugation matrix at a Wigner degree. + + :param degree: Wigner representation degree. + + :return: Diagonal matrix of shape (2 * degree + 1, 2 * degree + 1). + """ + signs = np.ones(2 * degree + 1) + signs[degree + 1 :: 2] = -1 + signs[degree - 1 :: -2] = -1 + + return np.diag(signs) + + @staticmethod + def _assemble_rotation_estimates(alpha, beta, gamma): + """ + Assemble internal NUG Euler estimates and convert rotation matrices. + + :param alpha: Estimated alpha angles. + :param beta: Estimated beta angles. + :param gamma: Estimated gamma angles. + + :return: rotation matrices. + """ + euler_nug = np.column_stack((alpha, beta, gamma)).astype( + np.float64, + copy=False, + ) + + rotations = Rotation.from_euler(euler_nug).matrices.swapaxes(-1, -2) + + return rotations + + def _estimate_gamma( + self, + Xm, + alpha, + beta, + degree, + pair_estimator, + wrap=False, + ): + """ + Estimate gamma angles by synchronizing pairwise phase measurements. + + Xm must already be expressed in the complex Wigner basis. + + :param Xm: Complex-basis representation matrix at `degree`. + :param alpha: Estimated alpha angles. + :param beta: Estimated beta angles. + :param degree: Representation degree used for gamma recovery. + :param pair_estimator: Callable that computes one pairwise phase + measurement. + :param wrap: Whether to wrap the results to [0, 2*pi). + + :return: Estimated gamma angles. + """ + N = self.n_img + dk = 2 * degree + 1 + modes = np.arange(-degree, degree + 1) + + C = np.zeros((N, N), dtype=complex_type(np.float64)) + Jk = self._handedness_matrix(degree) + ws = self.Wd(degree, beta) + + for i in range(N): + wi = ws[i] + Di = np.exp(-1j * modes * alpha[i]) + + for j in range(i + 1, N): + wj = ws[j] + Dj = np.exp(-1j * modes * alpha[j]) + + Xij = Xm[ + dk * i : dk * (i + 1), + dk * j : dk * (j + 1), + ] + + DXijD = np.diag(Di.conj()) @ Xij @ np.diag(Dj) + + C[i, j] = pair_estimator( + wi, + wj, + DXijD, + Jk, + degree, ) - B = np.array([np.vdot(W1 - W4, Bi), np.vdot(W3 - W2, Bi)]) - b, d = np.linalg.lstsq(A, B)[0] - return a + 1j * b - - [T, Tinv] = self.complex2real(S) - Xm = ( - np.kron(np.eye(N, dtype=np.float64), T) - @ Xm - @ np.kron(np.eye(N, dtype=np.float64), Tinv) - ) - C = np.zeros((N, N), dtype=complex_type(np.float64)) - Jk = np.ones(dk) - Jk[S + 1 :: 2] = -1 - Jk[S - 1 :: -2] = -1 - Jk = np.diag(Jk) - ws = self.Wd(S, beta) - for i in range(N): - wi = ws[i] - for j in range(i + 1, N): - Di = np.exp(-1j * np.arange(-S, S + 1) * alpha[i]) - Dj = np.exp(-1j * np.arange(-S, S + 1) * alpha[j]) - Xijm = Xm[dk * i : dk * (i + 1), dk * j : dk * (j + 1)] - DXijmD = np.diag(Di.conj()) @ Xijm @ np.diag(Dj) - wj = ws[j] - W1 = ( - wi[:, 0][:, None] @ wj[:, 0][:, None].T - + Jk @ wi[:, 0][:, None] @ wj[:, 0][:, None].T @ Jk - ) - W2 = ( - wi[:, -1][:, None] @ wj[:, 0][:, None].T - + Jk @ wi[:, -1][:, None] @ wj[:, 0][:, None].T @ Jk - ) - W3 = ( - wi[:, 0][:, None] @ wj[:, -1][:, None].T - + Jk @ wi[:, 0][:, None] @ wj[:, -1][:, None].T @ Jk - ) - W4 = ( - wi[:, -1][:, None] @ wj[:, -1][:, None].T - + Jk @ wi[:, -1][:, None] @ wj[:, -1][:, None].T @ Jk - ) - Br = np.real(4 * DXijmD) - Bi = np.imag(4 * DXijmD) - C[i, j] = LS_D(W1, W2, W3, W4, Br, Bi) - C += C.T.conj() + np.eye(N, dtype=np.float64) - evals, evecs = np.linalg.eigh(C) - idx = np.argmax(evals) - c = evecs[:, idx] * np.sqrt(evals[idx]) - return (np.angle(c) / S) % (2 * np.pi) - alpha_est, beta_est = find_alpha_beta(X2) - gamma_est = find_gamma(XS, alpha_est, beta_est) - Euler_est = np.zeros((N, 3), dtype=np.float64) - Euler_est[:, 0] = alpha_est - Euler_est[:, 1] = beta_est - Euler_est[:, 2] = gamma_est - R_est = Rotation.from_euler(Euler_est).matrices.transpose(0, 2, 1) + C += C.T.conj() + np.eye(N, dtype=np.float64) - return R_est + eigenvalues, eigenvectors = np.linalg.eigh(C) + leading_idx = np.argmax(eigenvalues) + phase_vector = eigenvectors[:, leading_idx] * np.sqrt(eigenvalues[leading_idx]) + + gamma = np.angle(phase_vector) / degree + + if wrap: + gamma %= 2 * np.pi + + return gamma + + @staticmethod + def _cyclic_gamma_pair(wi, wj, DXijD, Jk, degree): + outer_first = np.outer(wi[:, 0], wj[:, 0]) + C1 = (outer_first + Jk @ outer_first @ Jk) / 2 + + outer_last = np.outer(wi[:, -1], wj[:, -1]) + C2 = (outer_last + Jk @ outer_last @ Jk) / 2 + + outer_middle = np.outer(wi[:, degree], wj[:, degree]) + C3 = DXijD - (outer_middle + Jk @ outer_middle @ Jk) / 2 + + return np.vdot(C1 + C2, np.real(C3)) / np.vdot(C1 + C2, C1 + C2) + 1j * np.vdot( + C1 - C2, np.imag(C3) + ) / np.vdot(C1 - C2, C1 - C2) + + @staticmethod + def _solve_dihedral_gamma_pair(W1, W2, W3, W4, Br, Bi): + A = np.array( + [ + [np.vdot(W1 + W4, W1 + W4), np.vdot(W1 + W4, W2 + W3)], + [np.vdot(W2 + W3, W1 + W4), np.vdot(W2 + W3, W2 + W3)], + ] + ) + B = np.array([np.vdot(W1 + W4, Br), np.vdot(W2 + W3, Br)]) + a, _ = np.linalg.lstsq(A, B, rcond=None)[0] + + A = np.array( + [ + [np.vdot(W1 - W4, W1 - W4), np.vdot(W1 - W4, W3 - W2)], + [np.vdot(W1 - W4, W3 - W2), np.vdot(W3 - W2, W3 - W2)], + ] + ) + B = np.array([np.vdot(W1 - W4, Bi), np.vdot(W3 - W2, Bi)]) + b, _ = np.linalg.lstsq(A, B, rcond=None)[0] + + return a + 1j * b + + def _dihedral_gamma_pair(self, wi, wj, DXijD, Jk, degree): + outer_11 = np.outer(wi[:, 0], wj[:, 0]) + outer_21 = np.outer(wi[:, -1], wj[:, 0]) + outer_12 = np.outer(wi[:, 0], wj[:, -1]) + outer_22 = np.outer(wi[:, -1], wj[:, -1]) + + W1 = outer_11 + Jk @ outer_11 @ Jk + W2 = outer_21 + Jk @ outer_21 @ Jk + W3 = outer_12 + Jk @ outer_12 @ Jk + W4 = outer_22 + Jk @ outer_22 @ Jk + + return self._solve_dihedral_gamma_pair( + W1, + W2, + W3, + W4, + np.real(4 * DXijD), + np.imag(4 * DXijD), + ) #################### # Helper Functions # From 1a241cd2ac3328257525be7ee52e249f46a02512 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 26 Aug 2026 14:50:32 -0400 Subject: [PATCH 078/161] dosctrings for helper functions --- src/aspire/abinitio/commonline_nug.py | 73 ++++++++++++++++++++++++++- 1 file changed, 71 insertions(+), 2 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 953a823df7..25e72bccee 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1812,12 +1812,12 @@ def _assemble_rotation_estimates(alpha, beta, gamma): :return: rotation matrices. """ - euler_nug = np.column_stack((alpha, beta, gamma)).astype( + euler_angles = np.column_stack((alpha, beta, gamma)).astype( np.float64, copy=False, ) - rotations = Rotation.from_euler(euler_nug).matrices.swapaxes(-1, -2) + rotations = Rotation.from_euler(euler_angles).matrices.swapaxes(-1, -2) return rotations @@ -1891,6 +1891,28 @@ def _estimate_gamma( @staticmethod def _cyclic_gamma_pair(wi, wj, DXijD, Jk, degree): + """ + Estimate the pairwise gamma phase for cyclic symmetry. + + Construct the handedness-invariant Wigner terms induced by images + i and j, subtract the contribution of the central Wigner component + from the corresponding representation block, and fit the remaining + real and imaginary components. The resulting complex value is inserted + into the phase-synchronization matrix used to recover the gamma angles. + + :param wi: Degree-k Wigner small-d matrix evaluated at the + estimated beta angle of image i. Shape (2 * k + 1, 2 * k + 1). + :param wj: Degree-k Wigner small-d matrix evaluated at the + estimated beta angle of image j. Shape (2 * k + 1, 2 * k + 1). + :param DXijD: Complex representation block for images i and j after + removing their estimated alpha phases. Shape (2 * k + 1, 2 * k + 1). + :param Jk: Degree-k handedness-conjugation matrix. Shape + (2 * k + 1, 2 * k + 1). + :param degree: Wigner representation degree (k). + + :return: Complex pairwise phase measurement used for gamma + synchronization. + """ outer_first = np.outer(wi[:, 0], wj[:, 0]) C1 = (outer_first + Jk @ outer_first @ Jk) / 2 @@ -1906,6 +1928,33 @@ def _cyclic_gamma_pair(wi, wj, DXijD, Jk, degree): @staticmethod def _solve_dihedral_gamma_pair(W1, W2, W3, W4, Br, Bi): + """ + Fit one dihedral pairwise gamma-phase measurement. + + Solve separate two-variable least-squares systems for the real and + imaginary components of the pairwise phase. The systems express the + real and imaginary parts of the alpha-corrected representation block + in terms of handedness-invariant combinations of the four extreme + Wigner outer products. + + :param W1: Handedness-invariant outer product formed from the first + Wigner columns of both images. + :param W2: Handedness-invariant outer product formed from the last + Wigner column of the first image and the first Wigner column of + the second image. + :param W3: Handedness-invariant outer product formed from the first + Wigner column of the first image and the last Wigner column of + the second image. + :param W4: Handedness-invariant outer product formed from the last + Wigner columns of both images. + :param Br: Scaled real part of the alpha-corrected representation + block. + :param Bi: Scaled imaginary part of the alpha-corrected representation + block. + + :return: Complex pairwise phase measurement whose real and imaginary + components are obtained from the two least-squares fits. + """ A = np.array( [ [np.vdot(W1 + W4, W1 + W4), np.vdot(W1 + W4, W2 + W3)], @@ -1927,6 +1976,26 @@ def _solve_dihedral_gamma_pair(W1, W2, W3, W4, Br, Bi): return a + 1j * b def _dihedral_gamma_pair(self, wi, wj, DXijD, Jk, degree): + """ + Estimate the pairwise gamma phase for dihedral symmetry. + + Form the four handedness-invariant outer products associated with the + outer Wigner columns for images i and j. Fit these terms to the real + and imaginary parts of the alpha-corrected representation block to + obtain the complex measurement used for gamma synchronization. + + :param wi: Degree-k Wigner small-d matrix evaluated at the estimated + beta angle of image i. Shape (2 * k + 1, 2 * k + 1). + :param wj: Degree-k Wigner small-d matrix evaluated at the estimated + beta angle of image j. Shape (2 * k + 1, 2 * k + 1). + :param DXijD: Complex representation block for images i and j after + removing their estimated alpha phases. Shape (2 * k + 1, 2 * k + 1). + :param Jk: Degree-k handedness-conjugation matrix. Shape (2 * k + 1, 2 * k + 1). + :param degree: Wigner representation degree. This parameter is included + to provide the common pair-estimator interface used by `_estimate_gamma`. + + :return: Complex pairwise phase measurement used for gamma synchronization. + """ outer_11 = np.outer(wi[:, 0], wj[:, 0]) outer_21 = np.outer(wi[:, -1], wj[:, 0]) outer_12 = np.outer(wi[:, 0], wj[:, -1]) From 987a6871eef1997e1e36158d8582ee266d30267c Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 26 Aug 2026 15:27:58 -0400 Subject: [PATCH 079/161] Add code comments to euler_est_Cm --- src/aspire/abinitio/commonline_nug.py | 46 +++++++++++++++++++++++++-- 1 file changed, 43 insertions(+), 3 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 25e72bccee..9112b3604c 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1606,52 +1606,92 @@ def euler_est_Cm(self, X1, XS): :param X1: Relaxed degree-one representation matrix. :param XS: Relaxed representation matrix at the symmetry order. - :return: Estimated rotation matrices and Euler angles. + :return: Estimated rotation matrices. """ S = self.n_sym N = self.n_img + # Convert the relaxed NUG matrices from the real representation used by + # the SDP solver to the complex Wigner-D basis used in the recovery formulas. X1 = self._real_to_complex_representation(X1, degree=1) XS = self._real_to_complex_representation(XS, degree=S) def find_beta(X1): B1 = np.zeros((N, N), dtype=np.float64) B2 = np.zeros((N, N), dtype=np.float64) + + # From equation (36) in the paper, we have + # + # 2 * abs(X1_ij[0, 0]) = sin(beta_i) * sin(beta_j) + # real(X1_ij[1, 1]) = cos(beta_i) * cos(beta_j). + # + # Therefore, B1 and B2 are approximately rank-one outer-product + # matrices containing the sine and cosine factors, respectively. for i in range(N): for j in range(N): Xij = X1[3 * i : 3 * (i + 1), 3 * j : 3 * (j + 1)] B1[i, j] = abs(Xij[0, 0]) * 2 B2[i, j] = np.real(Xij[1, 1]) + + # Recover vectors approximating sin(beta_i) and cos(beta_i) from + # the leading eigenpairs of their rank-one Gram matrices. e1, v1 = np.linalg.eigh(B1) idx = np.argmax(e1) b1 = -v1[:, idx] * np.sqrt(e1[idx]) + e2, v2 = np.linalg.eigh(B2) idx = np.argmax(e2) b2 = v2[:, idx] * np.sqrt(e2[idx]) + + # Combine the synchronized sine and cosine factors to recover beta. + # The eigenvector sign ambiguity corresponds to the global + # beta <-> pi - beta ambiguity described in the paper. beta = np.arctan(b1 / b2) % np.pi return beta def find_alpha(X1): ZZbar = np.zeros((N, N), dtype=complex_type(np.float64)) ZZ = np.zeros((N, N), dtype=complex_type(np.float64)) + + # Equation (36) also gives the pairwise phase measurements + # + # X1_ij[0, 0] / abs(X1_ij[0, 0]) + # = exp(i * (alpha_i - alpha_j)) + # + # -X1_ij[0, 2] / abs(X1_ij[0, 2]) + # = exp(i * (alpha_i + alpha_j)). + # + # Store these as the difference-phase and sum-phase matrices. for i in range(N): for j in range(N): - z = X1[3 * i : 3 * (i + 1), 3 * j : 3 * (j + 1)][0, 0] + Xij = X1[ + 3 * i : 3 * (i + 1), + 3 * j : 3 * (j + 1), + ] + + z = Xij[0, 0] ZZbar[i, j] = z / abs(z) - z = X1[3 * i : 3 * (i + 1), 3 * j : 3 * (j + 1)][0, 2] + z = Xij[0, 2] ZZ[i, j] = -z / abs(z) + # Angular synchronization of the difference phases recovers + # exp(i * alpha_i) up to a common phase. evals, evecs = np.linalg.eigh(ZZbar) idx = np.argmax(abs(evals)) Z = evecs[:, idx] * np.sqrt(abs(evals[idx])) + # Use the sum phases to resolve the remaining common phase. c = self._find_phase(Z[:, None] @ Z[:, None].T, ZZ) Z = np.sqrt(c) * Z return np.angle(Z).astype(np.float64) alpha_est = find_alpha(X1) beta_est = find_beta(X1) + + # Estimate the remaining Euler angle from the degree-S solution. + # The cyclic pair estimator extracts measurements of + # exp(i * S * (gamma_i - gamma_j)), which are then synchronized. gamma_est = self._estimate_gamma( XS, alpha_est, From b15e4528bfc06bd49e4d9a9b985bd9f46db80860 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 26 Aug 2026 15:53:02 -0400 Subject: [PATCH 080/161] code comments for euler_est_Dm --- src/aspire/abinitio/commonline_nug.py | 35 +++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 9112b3604c..ab8eb03314 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1719,25 +1719,51 @@ def euler_est_Dm(self, X_est): S = self.sym_grp.order N = self.n_img + # Convert the degree-two and degree-S solutions from the real + # representation used by the SDP solver to the complex Wigner-D basis. X2 = self._real_to_complex_representation(X_est[1], degree=2) XS = self._real_to_complex_representation(X_est[S - 1], degree=S) def find_alpha_beta(X2): B1 = np.zeros((N, N), dtype=np.float64) + + # Equations (59)-(63) combine two entries of each degree-two block + # to isolate + # + # B1[i, j] = cos(beta_i)^2 * cos(beta_j)^2. + # + # Thus, B1 is approximately a rank-one outer-product matrix. for i in range(N): for j in range(N): Xij = X2[5 * i : 5 * (i + 1), 5 * j : 5 * (j + 1)] B1[i, j] = np.real((Xij[2, 2] - 2 * abs(Xij[0, 0]) + 0.5) / 3 * 2) + + # Recover the vector cos(beta_i)^2 from the leading eigenpair. + # Its entries should be nonnegative, so orient the eigenvector + # using the sign of its first entry. e1, v1 = np.linalg.eigh(B1) idx = np.argmax(e1) b1 = v1[:, idx] * np.sqrt(e1[idx]) * np.sign(v1[0, idx]) + + # Taking the square root recovers abs(cos(beta_i)). Consequently, + # each image retains an independent beta_i <-> pi - beta_i + # ambiguity, as described in the paper. beta_est = np.arccos(np.clip(np.sqrt(b1), -1, 1)) % np.pi Aminus = np.zeros((N, N), dtype=complex_type(np.float64)) Aplus = np.zeros((N, N), dtype=complex_type(np.float64)) + + # Equations (64)-(66) remove the known sin(2*beta) factors from + # two entries of X2_ij, producing the pairwise phases + # + # Aminus[i, j] = exp(i * (alpha_i - alpha_j)) + # Aplus[i, j] = exp(i * (alpha_i + alpha_j)). for i in range(N): for j in range(N): Xij = X2[5 * i : 5 * (i + 1), 5 * j : 5 * (j + 1)] + + # Avoid division by a numerically vanishing + # sin(2*beta) factor. if abs(beta_est[i]) < 1e-6: beta_est[i] = 1e-6 if abs(beta_est[j]) < 1e-6: @@ -1757,9 +1783,14 @@ def find_alpha_beta(X2): / 3 ) + # Angular synchronization of the difference phases recovers + # exp(i * alpha_i) up to a common phase. evals, evecs = np.linalg.eigh(Aminus) idx = np.argmax(abs(evals)) Z = evecs[:, idx] * np.sqrt(abs(evals[idx])) + + # Use the sum phases to resolve the remaining common phase. + # The square root leaves a global pi ambiguity. c = self._find_phase(Z[:, None] @ Z[:, None].T, Aplus) Z = np.sqrt(c) * Z alpha_est = (np.angle(Z)) % (2 * np.pi) @@ -1767,6 +1798,10 @@ def find_alpha_beta(X2): return alpha_est, beta_est alpha_est, beta_est = find_alpha_beta(X2) + + # Estimate gamma from the degree-S solution. The dihedral pair + # estimator separates the gamma-difference and gamma-sum terms, + # retaining the difference phases needed for synchronization. gamma_est = self._estimate_gamma( XS, alpha_est, From 7a429eba11a392c5890ed00e665b2a2f1f384e0c Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 27 Aug 2026 11:31:29 -0400 Subject: [PATCH 081/161] code comments for _estimate_gamma --- src/aspire/abinitio/commonline_nug.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index ab8eb03314..7521b534b1 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1924,8 +1924,13 @@ def _estimate_gamma( dk = 2 * degree + 1 modes = np.arange(-degree, degree + 1) + # Store pairwise estimates of exp(i * degree * (gamma_i - gamma_j)). C = np.zeros((N, N), dtype=complex_type(np.float64)) + + # Construct the handedness-conjugation matrix. Jk = self._handedness_matrix(degree) + + # Evaluate the Wigner small-d matrices at the recovered beta angles. ws = self.Wd(degree, beta) for i in range(N): @@ -1936,13 +1941,23 @@ def _estimate_gamma( wj = ws[j] Dj = np.exp(-1j * modes * alpha[j]) + # Each degree-m image-pair block factors as + # + # X_ij = D(alpha_i) @ M_ij @ D(alpha_j)^*, + # + # where M_ij contains the beta- and gamma-dependent terms. Xij = Xm[ dk * i : dk * (i + 1), dk * j : dk * (j + 1), ] + # Cancel the known alpha-dependent factors to isolate M_ij: + # + # D(alpha_i)^* @ X_ij @ D(alpha_j). DXijD = np.diag(Di.conj()) @ Xij @ np.diag(Dj) + # Since beta is already estimated, the remaining unknown coefficients + # are pairwise phases involving gamma_i and gamma_j. C[i, j] = pair_estimator( wi, wj, @@ -1951,12 +1966,16 @@ def _estimate_gamma( degree, ) + # Complete the Hermitian synchronization matrix. C += C.T.conj() + np.eye(N, dtype=np.float64) + # Angular synchronization recovers exp(i * degree * gamma_i) up to one common phase. eigenvalues, eigenvectors = np.linalg.eigh(C) leading_idx = np.argmax(eigenvalues) phase_vector = eigenvectors[:, leading_idx] * np.sqrt(eigenvalues[leading_idx]) + # Divide by the representation degree to recover gamma. Each estimate + # remains defined modulo 2*pi/degree, as expected from the symmetry. gamma = np.angle(phase_vector) / degree if wrap: From d420689b73ab7f4055e40bb2f05d5df7204a33fd Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 31 Aug 2026 09:35:44 -0400 Subject: [PATCH 082/161] use k for degree everywhere --- src/aspire/abinitio/commonline_nug.py | 42 +++++++++++++-------------- 1 file changed, 21 insertions(+), 21 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 7521b534b1..66ffc9fd93 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1432,13 +1432,13 @@ def proximal_refine(self, X_admm, weight, Penalty, r): N = self.n_img C = self.C - def Ak(J, Euler): + def Ak(k, Euler): """ - Average the degree-J representation over the symmetry + Average the degree-k representation over the symmetry group to form the symmetry projector Ak. """ order = Euler.shape[0] - A = self.WD(J, Euler).sum(axis=0) + A = self.WD(k, Euler).sum(axis=0) return np.round(A / order, 10) def rel_change(A, B, eps=1e-12): @@ -2273,17 +2273,17 @@ def compute_rank(self): rk[k - 1] = np.linalg.matrix_rank(Ak) return rk, A - def WD(self, J, euler): + def WD(self, k, euler): """ - Evaluate degree-J Wigner D matrices at the supplied ZYZ Euler angles. + Evaluate degree-k Wigner D matrices at the supplied ZYZ Euler angles. """ # compute Wigner D matrix alpha = euler[:, 0] beta = euler[:, 1] gamma = euler[:, 2] - d = self.Wd(J, beta) + d = self.Wd(k, beta) - m = np.arange(-J, J + 1) + m = np.arange(-k, k + 1) left = np.exp(-1j * alpha[:, None] * m[None, :]) right = np.exp(-1j * gamma[:, None] * m[None, :]) D = left[:, :, None] * d * right[:, None, :] @@ -2291,31 +2291,31 @@ def WD(self, J, euler): return D @staticmethod - def Wd(J, beta): + def Wd(k, beta): """ - Evaluate degree-J Wigner small-d matrices at the supplied polar angles. + Evaluate degree-k Wigner small-d matrices at the supplied polar angles. """ # compute Wigner small d matrix - d = np.zeros((len(beta), 2 * J + 1, 2 * J + 1), dtype=beta.dtype) - for m in range(-J, J + 1): - for n in range(-J, J + 1): + d = np.zeros((len(beta), 2 * k + 1, 2 * k + 1), dtype=beta.dtype) + for m in range(-k, k + 1): + for n in range(-k, k + 1): smin = max(0, m - n) - smax = min(J + m, J - n) + smax = min(k + m, k - n) for s in range(smin, smax + 1): mul = ( - np.sqrt(factorial(J + m)) - / factorial(J + m - s) - * np.sqrt(factorial(J + n)) + np.sqrt(factorial(k + m)) + / factorial(k + m - s) + * np.sqrt(factorial(k + n)) / factorial(s) - * np.sqrt(factorial(J - m)) + * np.sqrt(factorial(k - m)) / factorial(n - m + s) - * np.sqrt(factorial(J - n)) - / factorial(J - n - s) + * np.sqrt(factorial(k - n)) + / factorial(k - n - s) ) - d[:, n + J, m + J] += ( + d[:, n + k, m + k] += ( mul * (-1) ** (n - m + s) - * (np.cos(beta / 2)) ** (2 * J + m - n - 2 * s) + * (np.cos(beta / 2)) ** (2 * k + m - n - 2 * s) * (np.sin(beta / 2)) ** (n - m + 2 * s) ) return d From 81500287408b9e3f110f799a77c035edaad79765 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 31 Aug 2026 11:02:11 -0400 Subject: [PATCH 083/161] helper function _project_representation_blocks --- src/aspire/abinitio/commonline_nug.py | 117 +++++++++++++++++--------- 1 file changed, 78 insertions(+), 39 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 66ffc9fd93..2ae92be974 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -477,28 +477,15 @@ def update_S(C0, C1, yE0, yE1, yEq, yI, X0, X1, Xq, rho): S1 = C1 - Z1 - AIT_yI1 - X1 / rho Sq = -Zq - Xq / rho - for k in range(1, Lmax + 1): - tmp = self.mat_block( - S0[d0[k - 1] : d0[k]], - N, - k, - IDX_upper, - IDX_lower, - idx_offdiag, - ) - tmp = self.psd_projection(tmp) - S0[d0[k - 1] : d0[k]] = self.vec_block(tmp, N, k, IDX_upper) - - tmp = self.mat_block( - S1[d1[k - 1] : d1[k]], - N, - k + 1, - IDX_upper, - IDX_lower, - idx_offdiag, - ) - tmp = self.psd_projection(tmp) - S1[d1[k - 1] : d1[k]] = self.vec_block(tmp, N, k + 1, IDX_upper) + S0, S1 = self._project_representation_blocks( + S0, + S1, + d0, + d1, + IDX_upper, + IDX_lower, + idx_offdiag, + ) # This is the established batched equivalent of projecting each 4x4 # quaternion block and storing tmp.T.reshape(16) in each column. @@ -937,23 +924,15 @@ def update_S(C0, C1, yE, yEq, yI, X0, X1, Xd0, Xd1, Xq, rho, Lmax, N): Sq = -Zq - Xq / rho # Project packed X0/X1 blocks degree by degree. - for k in range(1, Lmax + 1): - tmp = self.mat_block( - S0[d0[k - 1] : d0[k], :], N, k, IDX_upper, IDX_lower, idx_offdiag - ) - tmp = self.psd_projection(tmp) - S0[d0[k - 1] : d0[k], :] = self.vec_block(tmp, N, k, IDX_upper) - - tmp = self.mat_block( - S1[d1[k - 1] : d1[k], :], - N, - k + 1, - IDX_upper, - IDX_lower, - idx_offdiag, - ) - tmp = self.psd_projection(tmp) - S1[d1[k - 1] : d1[k], :] = self.vec_block(tmp, N, k + 1, IDX_upper) + S0, S1 = self._project_representation_blocks( + S0, + S1, + d0, + d1, + IDX_upper, + IDX_lower, + idx_offdiag, + ) # Project the diagonal coupling blocks. Sd0 = Sd0.T @@ -2112,6 +2091,66 @@ def _dihedral_gamma_pair(self, wi, wj, DXijD, Jk, degree): #################### # Helper Functions # #################### + + def _project_representation_blocks( + self, + Z0, + Z1, + d0, + d1, + IDX_upper, + IDX_lower, + idx_offdiag, + ): + """ + Project packed degree-wise representation blocks onto the PSD cone. + + :param Z0: Packed blocks of sizes 1 through `Lmax`. + :param Z1: Packed blocks of sizes 2 through `Lmax + 1`. + :param d0: Cumulative row offsets for the blocks in `Z0`. + :param d1: Cumulative row offsets for the blocks in `Z1`. + :param IDX_upper: Linear indices of upper-triangular image pairs. + :param IDX_lower: Linear indices of lower-triangular image pairs. + :param idx_offdiag: Linear indices of off-diagonal image pairs in the + packed upper-triangular representation. + + :return: The projected packed arrays `Z0` and `Z1`. + """ + for k in range(1, self.Lmax + 1): + block0 = self.mat_block( + Z0[d0[k - 1] : d0[k]], + self.n_img, + k, + IDX_upper, + IDX_lower, + idx_offdiag, + ) + block0 = self.psd_projection(block0) + Z0[d0[k - 1] : d0[k]] = self.vec_block( + block0, + self.n_img, + k, + IDX_upper, + ) + + block1 = self.mat_block( + Z1[d1[k - 1] : d1[k]], + self.n_img, + k + 1, + IDX_upper, + IDX_lower, + idx_offdiag, + ) + block1 = self.psd_projection(block1) + Z1[d1[k - 1] : d1[k]] = self.vec_block( + block1, + self.n_img, + k + 1, + IDX_upper, + ) + + return Z0, Z1 + def transform_coeff(self, A, IDX_upper): """ Convert representation matrices to the block-vector form used by ADMM. From 7825d95a4562a7f097b7e34f20071b6aaf8614e7 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 31 Aug 2026 15:07:13 -0400 Subject: [PATCH 084/161] organize helper function sections --- src/aspire/abinitio/commonline_nug.py | 1025 +++++++++++++------------ 1 file changed, 517 insertions(+), 508 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 2ae92be974..91b0575342 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -297,38 +297,6 @@ def fijhat_k(k, F): self.C = C - def complex2real(self, ell): - """ - Construct the transformation matrices between complex and real degree ell representations. - - :param ell: Wigner representation degree. - - :return: Forward and inverse change-of-basis matrices. - """ - diml = 2 * ell + 1 - Tinv = np.zeros((diml, diml), dtype=complex_type(np.float64)) - for i in range(diml): - if i < ell: - Tinv[i, i] = 1j / np.sqrt(2) - Tinv[i, diml - 1 - i] = -1j * (-1) ** (i - ell) / np.sqrt(2) - if i == ell: - Tinv[i, i] = 1 - if i > ell: - Tinv[i, i] = (-1) ** (i - ell) / np.sqrt(2) - Tinv[i, diml - 1 - i] = 1 / np.sqrt(2) - - T = np.zeros((diml, diml), dtype=complex_type(np.float64)) - for i in range(diml): - if i < ell: - T[i, i] = -1j / np.sqrt(2) - T[i, diml - 1 - i] = 1 / np.sqrt(2) - if i == ell: - T[i, i] = 1 - if i > ell: - T[i, i] = (-1) ** (i - ell) / np.sqrt(2) - T[i, diml - 1 - i] = 1j * (-1) ** (i - ell) / np.sqrt(2) - return T, Tinv - ############# # ADMM Step # ############# @@ -1185,6 +1153,10 @@ def print_updates(verbose): X_admm[k] = xp.asnumpy(X_admm[k]) return X_admm + ################ + # ADMM Helpers # + ################ + def ADMM_preprocessing(self, C): """ Construct the transformed coefficients, constraints, indices, and initial variables used by ADMM. @@ -1393,126 +1365,429 @@ def discretize_SO3(self): return SO3 - ############################ - # Proximal Refinement Step # - ############################ - - def proximal_refine(self, X_admm, weight, Penalty, r): + @staticmethod + def largest_eigenvalue(AI, Ngrid, N): """ - Refine the relaxed solution by iteratively encouraging lower-rank representation matrices. + Estimate the largest eigenvalue of the Fejér constraint operator. + """ + # find the largest eigenvalue of the operator AI + z = xp.random.normal(0, 1, (Ngrid, N**2)) + Lambda = 0 - :param X_admm: Initial relaxed representation matrices. - :param weight: Degree-dependent refinement weights. - :param Penalty: Penalty value for each refinement step. - :param r: Rank offset for each refinement step. + while abs(Lambda - xp.linalg.norm(z)) > 500: + Lambda = xp.linalg.norm(z) + z = z / xp.linalg.norm(z) + z = AI @ (AI.T @ z) + Lambda += 2000 + logger.info("Largest eigenvalue of AIAIT is approximately %1.2f" % Lambda) + return Lambda - :return: Refined representation matrices. + def compute_rank(self): """ - N = self.n_img - C = self.C + Compute the ranks and matrices of the symmetry-averaging projectors at each degree. - def Ak(k, Euler): - """ - Average the degree-k representation over the symmetry - group to form the symmetry projector Ak. - """ - order = Euler.shape[0] - A = self.WD(k, Euler).sum(axis=0) - return np.round(A / order, 10) + :param Lmax: Maximum representation degree. - def rel_change(A, B, eps=1e-12): - """ - Measure relative change across all degree-wise relaxed matrices. - """ - num = 0.0 - den = 0.0 - for k in range(len(A)): - num += np.linalg.norm(A[k] - B[k]) ** 2 - den += np.linalg.norm(B[k]) ** 2 - return np.sqrt(num) / max(np.sqrt(den), eps) + :return: Ranks and symmetry-averaging matrices for each degree. + """ + rk = xp.zeros(self.Lmax, dtype=np.float64) + A = [] + for k in range(1, self.Lmax + 1): + Ak = np.sum(self.WD(k, self.sym_euler), axis=0) + Ak = np.round(Ak / self.n_sym, 6) + A.append(Ak) + rk[k - 1] = np.linalg.matrix_rank(Ak) + return rk, A - # The symmetry projector ranks give the expected low-rank structure - # for each degree of a valid orbit solution. - rank_Ak = np.zeros(self.Lmax, dtype=np.float64) + def construct_AEq(self): + """ + Construct the linear equality operator encoding the quaternion constraints. + """ + AEq = np.zeros((17, 21), np.float64) - # Keep an unmodified copy of the original objective coefficients; each - # refinement step adds a temporary proximal term to these. - C_base = [None] * self.Lmax - for k in range(self.Lmax): - C_base[k] = xp.asnumpy(C[k]).copy() - rank_Ak[k] = np.linalg.matrix_rank( - Ak(k + 1, self.sym_euler), tol=1e-6, hermitian=True - ) + # First 16 rows: identity constraints on first 16 variables + AEq[:16, :16] = np.eye(16, dtype=np.float64) - def low_rank_proj(X, r_step): - """ - Project diagonal image blocks toward the rank expected from symmetry. - """ - Xproj = [] - for k in range(self.Lmax): - dk = 2 * k + 3 - rk = min(int(rank_Ak[k] * 2) + r_step, dk) - tmp = np.copy(X[k]) - for i in range(N): - u, s, v = np.linalg.svd( - tmp[i * dk : (i + 1) * dk, i * dk : (i + 1) * dk] - ) - tmp[i * dk : (i + 1) * dk, i * dk : (i + 1) * dk] = ( - u[:, :rk] @ np.diag(s[:rk]) @ v[:rk] - ) - Xproj.append(tmp) - return Xproj + # Columns 16:21 map the low-degree X0/X1 entries into the quaternion + # convex-hull constraint Xq = I/4 - linear(X^(1)). + extra = 0.25 * np.array( + [ + [-1, 1, 0, 0, 1], + [0, 0, 0, 0, 0], + [0, 0, 1, -1, 0], + [0, 0, 0, 0, 0], + [0, 0, 0, 0, 0], + [1, 1, 0, 0, -1], + [0, 0, 0, 0, 0], + [0, 0, 1, 1, 0], + [0, 0, 1, -1, 0], + [0, 0, 0, 0, 0], + [-1, -1, 0, 0, -1], + [0, 0, 0, 0, 0], + [0, 0, 0, 0, 0], + [0, 0, 1, 1, 0], + [0, 0, 0, 0, 0], + [1, -1, 0, 0, 1], + ], + dtype=np.float64, + ) - # CC stores the modified objective coefficients for the next proximal ADMM solve. - CC = [None] * self.Lmax - current = [np.copy(Xk) for Xk in X_admm] + AEq[:16, 16:] = extra - # Proximal refinement step. - # Repeatedly bias the SDP objective toward the low-rank projection and resolve. - for step in range(self.pr_iters): - X_proj = low_rank_proj(current, r[step]) + # Last row: redundant trace/sum constraint + AEq[16, [0, 5, 10, 15]] = 1 - for k in range(self.Lmax): - CC[k] = ( - C_base[k] - - Penalty[step] * weight[k] * (X_proj[k] + X_proj[k].T) / 2 - ) + return AEq - X_next = self.admm_sym_J(CC, verbose=False) + def _project_representation_blocks( + self, + Z0, + Z1, + d0, + d1, + IDX_upper, + IDX_lower, + idx_offdiag, + ): + """ + Project packed degree-wise representation blocks onto the PSD cone. - if self.verbose: - logger.info( - "Proximal refine step %d/%d: relative update %.3e", - step + 1, - self.pr_iters, - rel_change(X_next, current), - ) + :param Z0: Packed blocks of sizes 1 through `Lmax`. + :param Z1: Packed blocks of sizes 2 through `Lmax + 1`. + :param d0: Cumulative row offsets for the blocks in `Z0`. + :param d1: Cumulative row offsets for the blocks in `Z1`. + :param IDX_upper: Linear indices of upper-triangular image pairs. + :param IDX_lower: Linear indices of lower-triangular image pairs. + :param idx_offdiag: Linear indices of off-diagonal image pairs in the + packed upper-triangular representation. - # Use the refined SDP solution as the starting point for the next projection step - current = [np.copy(Xk) for Xk in X_next] + :return: The projected packed arrays `Z0` and `Z1`. + """ + for k in range(1, self.Lmax + 1): + block0 = self.mat_block( + Z0[d0[k - 1] : d0[k]], + self.n_img, + k, + IDX_upper, + IDX_lower, + idx_offdiag, + ) + block0 = self.psd_projection(block0) + Z0[d0[k - 1] : d0[k]] = self.vec_block( + block0, + self.n_img, + k, + IDX_upper, + ) - return current + block1 = self.mat_block( + Z1[d1[k - 1] : d1[k]], + self.n_img, + k + 1, + IDX_upper, + IDX_lower, + idx_offdiag, + ) + block1 = self.psd_projection(block1) + Z1[d1[k - 1] : d1[k]] = self.vec_block( + block1, + self.n_img, + k + 1, + IDX_upper, + ) - ########################## - # Rotation Recovery Step # - ########################## + return Z0, Z1 - def recover_rotations(self): + @staticmethod + def psd_projection(B): """ - Recover rotations from the NUG representation matrices, using synchronization-based - rounding for asymmetric molecules and symmetry-specific Euler estimation for cyclic and - dihedral molecules. + Project one or more symmetric matrices onto the positive semidefinite cone. """ - X_est = self.X_est - if isinstance(self.sym_grp, IdentitySymmetryGroup): - R_est = self.rounding_C1(X_est[0]) - elif isinstance(self.sym_grp, CnSymmetryGroup): - R_est = self.euler_est_Cm(X_est[0], X_est[self.n_sym - 1]) - elif isinstance(self.sym_grp, DnSymmetryGroup): - R_est = self.euler_est_Dm(X_est) - - # self.Euler_est = Euler_est - self.rotations = R_est.astype(self.dtype) + # compute the PSD part of a symmstric matrix + B_sym = (B + B.swapaxes(-1, -2)) / 2 + evals, evecs = xp.linalg.eigh(B_sym) + evals = xp.maximum(evals, 0) + return (evecs * evals[..., None, :]) @ evecs.swapaxes(-1, -2) + + def transform_coeff(self, A, IDX_upper): + """ + Convert representation matrices to the block-vector form used by ADMM. + + :param A: Representation matrices. + :param IDX_upper: Indices of upper-triangular image pair blocks. + + :return: The two block-vector coefficient arrays. + """ + d0 = [0] + d1 = [0] + for k in range(1, self.Lmax + 1): + d0.append(d0[-1] + k**2) + d1.append(d1[-1] + (k + 1) ** 2) + A0 = xp.zeros((d0[-1], self.n_img * (self.n_img + 1) // 2), dtype=np.float64) + A1 = xp.zeros((d1[-1], self.n_img * (self.n_img + 1) // 2), dtype=np.float64) + for k in range(1, self.Lmax + 1): + a0, a1 = self.permutek(A[k - 1], k, self.n_img) + A0[d0[k - 1] : d0[k], :] = self.vec_block(a0, self.n_img, k, IDX_upper) + A1[d1[k - 1] : d1[k], :] = self.vec_block(a1, self.n_img, k + 1, IDX_upper) + return A0, A1 + + def transform_coeff_back(self, A0, A1, IDX_upper, IDX_lower, idx_offdiag): + """ + Reconstruct representation matrices from their ADMM block-vector form. + + :param A0: First block-vector array. + :param A1: Second block-vector array. + :param IDX_upper: Indices of upper-triangular image-pair blocks. + :param IDX_lower: Indices of lower-triangular image-pair blocks. + :param idx_offdiag: Indices of off-diagonal image pairs. + + :return: Reconstructed representation matrices. + """ + d0 = [0] + d1 = [0] + N = self.n_img + for k in range(1, self.Lmax + 1): + d0.append(d0[-1] + k**2) + d1.append(d1[-1] + (k + 1) ** 2) + A = [] + for k in range(1, self.Lmax + 1): + dk = 2 * k + 1 + Ak = xp.zeros((N * dk, N * dk), dtype=np.float64) + Ak[: N * k, : N * k] = self.mat_block( + A0[d0[k - 1] : d0[k], :], N, k, IDX_upper, IDX_lower, idx_offdiag + ) + Ak[N * k :, N * k :] = self.mat_block( + A1[d1[k - 1] : d1[k], :], N, k + 1, IDX_upper, IDX_lower, idx_offdiag + ) + Ak = self.permutek_back(Ak, k, N) + A.append(Ak) + return A + + @staticmethod + def permutek(Ak, k, N): + """ + Permute and split a degree-k matrix into blocks of sizes k and k + 1. + + :param Ak: Degree-k block matrix. + :param k: Representation degree. + :param N: Number of images. + + :return: The two permuted matrix blocks. + """ + AkP = xp.copy(Ak) + dk = 2 * k + 1 + Pk = xp.eye(dk, dtype=AkP.dtype) + for m in range(k): + for n in range(k - m): + Pk[(m + 2 * n, m + 2 * n + 1), :] = Pk[(m + 2 * n + 1, m + 2 * n), :] + AkP = ( + xp.kron(xp.eye(N, dtype=AkP.dtype), Pk) + @ Ak + @ xp.kron(xp.eye(N, dtype=AkP.dtype), Pk.T) + ) + + Pk = xp.eye(N * dk, dtype=AkP.dtype) + idx = xp.concatenate((xp.arange(dk - k, dk), xp.arange(k + 1))) + for m in range(N - 1): + for n in range(N - 1 - m): + Pk[k * (m + 1) + n * dk : k * (m + 1) + (n + 1) * dk] = Pk[ + k * (m + 1) + n * dk : k * (m + 1) + (n + 1) * dk + ][idx, :] + AkP = Pk @ AkP @ Pk.T + return AkP[: N * k, : N * k], AkP[N * k :, N * k :] + + @staticmethod + def permutek_back(Ak, k, N): + """ + Undo the degree-k block permutation and reconstruct the full matrix. + + :param Ak: Permuted degree-k matrix. + :param k: Representation degree. + :param N: Number of images. + + :return: Matrix in the original block ordering. + """ + dk = 2 * k + 1 + Pk = xp.eye(N * dk, dtype=Ak.dtype) + idx = xp.concatenate((xp.arange(dk - k, dk), xp.arange(k + 1))) + for m in range(N - 1): + for n in range(N - 1 - m): + Pk[k * (m + 1) + n * dk : k * (m + 1) + (n + 1) * dk] = Pk[ + k * (m + 1) + n * dk : k * (m + 1) + (n + 1) * dk + ][idx, :] + AkB = Pk.T @ Ak @ Pk + dk = 2 * k + 1 + Pk = xp.eye(dk, dtype=Ak.dtype) + for m in range(k): + for n in range(k - m): + Pk[(m + 2 * n, m + 2 * n + 1), :] = Pk[(m + 2 * n + 1, m + 2 * n), :] + AkB = ( + xp.kron(xp.eye(N, dtype=Ak.dtype), Pk.T) + @ AkB + @ xp.kron(xp.eye(N, dtype=Ak.dtype), Pk) + ) + return AkB + + @staticmethod + def vec_block(A, N, sz, IDX_upper): + """ + Vectorize the upper-triangular image-pair blocks of a block matrix. + """ + vecA = (A.reshape(N, sz, N, sz).transpose(0, 2, 3, 1)).reshape(N**2, sz**2).T + return vecA[:, IDX_upper] + + @staticmethod + def mat_block(vecA, N, sz, IDX_upper, IDX_lower, idx_offdiag): + """ + Reconstruct a symmetric block matrix from its vectorized upper-triangular blocks. + """ + tmp = vecA.T.reshape(N * (N + 1) // 2, sz, sz).transpose(0, 2, 1) + AA = xp.zeros((N**2, sz, sz), dtype=vecA.dtype) + AA[IDX_upper] = tmp + AA[IDX_lower] = tmp[idx_offdiag].transpose(0, 2, 1) + return (AA.reshape(N, N, sz, sz).transpose(0, 2, 1, 3)).reshape(N * sz, N * sz) + + @staticmethod + def transform_block(A, k, Pk): + """ + Permute and vectorize the two invariant blocks of a degree-k matrix. + """ + AT = Pk @ A @ Pk.T + A0 = AT[:, :k, :k].swapaxes(-1, -2).reshape(A.shape[0], -1) + A1 = AT[:, k:, k:].swapaxes(-1, -2).reshape(A.shape[0], -1) + return A0, A1 + + @staticmethod + def transform_back_block(A0, A1, k, Pk): + """ + Reconstruct a degree-k matrix from its two invariant block vectors. + """ + dk = 2 * k + 1 + A = xp.zeros((A0.shape[0], dk, dk), dtype=A0.dtype) + A[:, :k, :k] = A0.reshape(-1, k, k).swapaxes(-1, -2) + A[:, k:, k:] = A1.reshape(-1, k + 1, k + 1).swapaxes(-1, -2) + return Pk.T @ A @ Pk + + ############################ + # Proximal Refinement Step # + ############################ + + def proximal_refine(self, X_admm, weight, Penalty, r): + """ + Refine the relaxed solution by iteratively encouraging lower-rank representation matrices. + + :param X_admm: Initial relaxed representation matrices. + :param weight: Degree-dependent refinement weights. + :param Penalty: Penalty value for each refinement step. + :param r: Rank offset for each refinement step. + + :return: Refined representation matrices. + """ + N = self.n_img + C = self.C + + def Ak(k, Euler): + """ + Average the degree-k representation over the symmetry + group to form the symmetry projector Ak. + """ + order = Euler.shape[0] + A = self.WD(k, Euler).sum(axis=0) + return np.round(A / order, 10) + + def rel_change(A, B, eps=1e-12): + """ + Measure relative change across all degree-wise relaxed matrices. + """ + num = 0.0 + den = 0.0 + for k in range(len(A)): + num += np.linalg.norm(A[k] - B[k]) ** 2 + den += np.linalg.norm(B[k]) ** 2 + return np.sqrt(num) / max(np.sqrt(den), eps) + + # The symmetry projector ranks give the expected low-rank structure + # for each degree of a valid orbit solution. + rank_Ak = np.zeros(self.Lmax, dtype=np.float64) + + # Keep an unmodified copy of the original objective coefficients; each + # refinement step adds a temporary proximal term to these. + C_base = [None] * self.Lmax + for k in range(self.Lmax): + C_base[k] = xp.asnumpy(C[k]).copy() + rank_Ak[k] = np.linalg.matrix_rank( + Ak(k + 1, self.sym_euler), tol=1e-6, hermitian=True + ) + + def low_rank_proj(X, r_step): + """ + Project diagonal image blocks toward the rank expected from symmetry. + """ + Xproj = [] + for k in range(self.Lmax): + dk = 2 * k + 3 + rk = min(int(rank_Ak[k] * 2) + r_step, dk) + tmp = np.copy(X[k]) + for i in range(N): + u, s, v = np.linalg.svd( + tmp[i * dk : (i + 1) * dk, i * dk : (i + 1) * dk] + ) + tmp[i * dk : (i + 1) * dk, i * dk : (i + 1) * dk] = ( + u[:, :rk] @ np.diag(s[:rk]) @ v[:rk] + ) + Xproj.append(tmp) + return Xproj + + # CC stores the modified objective coefficients for the next proximal ADMM solve. + CC = [None] * self.Lmax + current = [np.copy(Xk) for Xk in X_admm] + + # Proximal refinement step. + # Repeatedly bias the SDP objective toward the low-rank projection and resolve. + for step in range(self.pr_iters): + X_proj = low_rank_proj(current, r[step]) + + for k in range(self.Lmax): + CC[k] = ( + C_base[k] + - Penalty[step] * weight[k] * (X_proj[k] + X_proj[k].T) / 2 + ) + + X_next = self.admm_sym_J(CC, verbose=False) + + if self.verbose: + logger.info( + "Proximal refine step %d/%d: relative update %.3e", + step + 1, + self.pr_iters, + rel_change(X_next, current), + ) + + # Use the refined SDP solution as the starting point for the next projection step + current = [np.copy(Xk) for Xk in X_next] + + return current + + ########################## + # Rotation Recovery Step # + ########################## + + def recover_rotations(self): + """ + Recover rotations from the NUG representation matrices, using synchronization-based + rounding for asymmetric molecules and symmetry-specific Euler estimation for cyclic and + dihedral molecules. + """ + X_est = self.X_est + if isinstance(self.sym_grp, IdentitySymmetryGroup): + R_est = self.rounding_C1(X_est[0]) + elif isinstance(self.sym_grp, CnSymmetryGroup): + R_est = self.euler_est_Cm(X_est[0], X_est[self.n_sym - 1]) + elif isinstance(self.sym_grp, DnSymmetryGroup): + R_est = self.euler_est_Dm(X_est) + + # self.Euler_est = Euler_est + self.rotations = R_est.astype(self.dtype) def rounding_C1(self, X1): """ @@ -1798,24 +2073,9 @@ def find_alpha_beta(X2): return R_est - ############################ - # Euler Estimation Helpers # - ############################ - def _real_to_complex_representation(self, X, degree): - """ - Convert a degree-wise representation matrix from the real basis used - by ADMM to the complex Wigner basis. - - :param X: Representation matrix of shape - (n_img * (2 * degree + 1), n_img * (2 * degree + 1)). - :param degree: Wigner representation degree. - - :return: Representation matrix in the complex Wigner basis. - """ - T, Tinv = self.complex2real(degree) - identity = np.eye(self.n_img, dtype=np.float64) - - return np.kron(identity, T) @ X @ np.kron(identity, Tinv) + ############################# + # Rotation Recovery Helpers # + ############################# @staticmethod def _find_phase(A, B): @@ -1840,21 +2100,6 @@ def _find_phase(A, B): return c / abs(c) - @staticmethod - def _handedness_matrix(degree): - """ - Construct the handedness-conjugation matrix at a Wigner degree. - - :param degree: Wigner representation degree. - - :return: Diagonal matrix of shape (2 * degree + 1, 2 * degree + 1). - """ - signs = np.ones(2 * degree + 1) - signs[degree + 1 :: 2] = -1 - signs[degree - 1 :: -2] = -1 - - return np.diag(signs) - @staticmethod def _assemble_rotation_estimates(alpha, beta, gamma): """ @@ -2023,294 +2268,105 @@ def _solve_dihedral_gamma_pair(W1, W2, W3, W4, Br, Bi): :param Br: Scaled real part of the alpha-corrected representation block. :param Bi: Scaled imaginary part of the alpha-corrected representation - block. - - :return: Complex pairwise phase measurement whose real and imaginary - components are obtained from the two least-squares fits. - """ - A = np.array( - [ - [np.vdot(W1 + W4, W1 + W4), np.vdot(W1 + W4, W2 + W3)], - [np.vdot(W2 + W3, W1 + W4), np.vdot(W2 + W3, W2 + W3)], - ] - ) - B = np.array([np.vdot(W1 + W4, Br), np.vdot(W2 + W3, Br)]) - a, _ = np.linalg.lstsq(A, B, rcond=None)[0] - - A = np.array( - [ - [np.vdot(W1 - W4, W1 - W4), np.vdot(W1 - W4, W3 - W2)], - [np.vdot(W1 - W4, W3 - W2), np.vdot(W3 - W2, W3 - W2)], - ] - ) - B = np.array([np.vdot(W1 - W4, Bi), np.vdot(W3 - W2, Bi)]) - b, _ = np.linalg.lstsq(A, B, rcond=None)[0] - - return a + 1j * b - - def _dihedral_gamma_pair(self, wi, wj, DXijD, Jk, degree): - """ - Estimate the pairwise gamma phase for dihedral symmetry. - - Form the four handedness-invariant outer products associated with the - outer Wigner columns for images i and j. Fit these terms to the real - and imaginary parts of the alpha-corrected representation block to - obtain the complex measurement used for gamma synchronization. - - :param wi: Degree-k Wigner small-d matrix evaluated at the estimated - beta angle of image i. Shape (2 * k + 1, 2 * k + 1). - :param wj: Degree-k Wigner small-d matrix evaluated at the estimated - beta angle of image j. Shape (2 * k + 1, 2 * k + 1). - :param DXijD: Complex representation block for images i and j after - removing their estimated alpha phases. Shape (2 * k + 1, 2 * k + 1). - :param Jk: Degree-k handedness-conjugation matrix. Shape (2 * k + 1, 2 * k + 1). - :param degree: Wigner representation degree. This parameter is included - to provide the common pair-estimator interface used by `_estimate_gamma`. - - :return: Complex pairwise phase measurement used for gamma synchronization. - """ - outer_11 = np.outer(wi[:, 0], wj[:, 0]) - outer_21 = np.outer(wi[:, -1], wj[:, 0]) - outer_12 = np.outer(wi[:, 0], wj[:, -1]) - outer_22 = np.outer(wi[:, -1], wj[:, -1]) - - W1 = outer_11 + Jk @ outer_11 @ Jk - W2 = outer_21 + Jk @ outer_21 @ Jk - W3 = outer_12 + Jk @ outer_12 @ Jk - W4 = outer_22 + Jk @ outer_22 @ Jk - - return self._solve_dihedral_gamma_pair( - W1, - W2, - W3, - W4, - np.real(4 * DXijD), - np.imag(4 * DXijD), - ) - - #################### - # Helper Functions # - #################### - - def _project_representation_blocks( - self, - Z0, - Z1, - d0, - d1, - IDX_upper, - IDX_lower, - idx_offdiag, - ): - """ - Project packed degree-wise representation blocks onto the PSD cone. - - :param Z0: Packed blocks of sizes 1 through `Lmax`. - :param Z1: Packed blocks of sizes 2 through `Lmax + 1`. - :param d0: Cumulative row offsets for the blocks in `Z0`. - :param d1: Cumulative row offsets for the blocks in `Z1`. - :param IDX_upper: Linear indices of upper-triangular image pairs. - :param IDX_lower: Linear indices of lower-triangular image pairs. - :param idx_offdiag: Linear indices of off-diagonal image pairs in the - packed upper-triangular representation. - - :return: The projected packed arrays `Z0` and `Z1`. - """ - for k in range(1, self.Lmax + 1): - block0 = self.mat_block( - Z0[d0[k - 1] : d0[k]], - self.n_img, - k, - IDX_upper, - IDX_lower, - idx_offdiag, - ) - block0 = self.psd_projection(block0) - Z0[d0[k - 1] : d0[k]] = self.vec_block( - block0, - self.n_img, - k, - IDX_upper, - ) - - block1 = self.mat_block( - Z1[d1[k - 1] : d1[k]], - self.n_img, - k + 1, - IDX_upper, - IDX_lower, - idx_offdiag, - ) - block1 = self.psd_projection(block1) - Z1[d1[k - 1] : d1[k]] = self.vec_block( - block1, - self.n_img, - k + 1, - IDX_upper, - ) - - return Z0, Z1 - - def transform_coeff(self, A, IDX_upper): - """ - Convert representation matrices to the block-vector form used by ADMM. - - :param A: Representation matrices. - :param IDX_upper: Indices of upper-triangular image pair blocks. - - :return: The two block-vector coefficient arrays. - """ - d0 = [0] - d1 = [0] - for k in range(1, self.Lmax + 1): - d0.append(d0[-1] + k**2) - d1.append(d1[-1] + (k + 1) ** 2) - A0 = xp.zeros((d0[-1], self.n_img * (self.n_img + 1) // 2), dtype=np.float64) - A1 = xp.zeros((d1[-1], self.n_img * (self.n_img + 1) // 2), dtype=np.float64) - for k in range(1, self.Lmax + 1): - a0, a1 = self.permutek(A[k - 1], k, self.n_img) - A0[d0[k - 1] : d0[k], :] = self.vec_block(a0, self.n_img, k, IDX_upper) - A1[d1[k - 1] : d1[k], :] = self.vec_block(a1, self.n_img, k + 1, IDX_upper) - return A0, A1 - - def transform_coeff_back(self, A0, A1, IDX_upper, IDX_lower, idx_offdiag): - """ - Reconstruct representation matrices from their ADMM block-vector form. - - :param A0: First block-vector array. - :param A1: Second block-vector array. - :param IDX_upper: Indices of upper-triangular image-pair blocks. - :param IDX_lower: Indices of lower-triangular image-pair blocks. - :param idx_offdiag: Indices of off-diagonal image pairs. - - :return: Reconstructed representation matrices. - """ - d0 = [0] - d1 = [0] - N = self.n_img - for k in range(1, self.Lmax + 1): - d0.append(d0[-1] + k**2) - d1.append(d1[-1] + (k + 1) ** 2) - A = [] - for k in range(1, self.Lmax + 1): - dk = 2 * k + 1 - Ak = xp.zeros((N * dk, N * dk), dtype=np.float64) - Ak[: N * k, : N * k] = self.mat_block( - A0[d0[k - 1] : d0[k], :], N, k, IDX_upper, IDX_lower, idx_offdiag - ) - Ak[N * k :, N * k :] = self.mat_block( - A1[d1[k - 1] : d1[k], :], N, k + 1, IDX_upper, IDX_lower, idx_offdiag - ) - Ak = self.permutek_back(Ak, k, N) - A.append(Ak) - return A - - @staticmethod - def permutek(Ak, k, N): - """ - Permute and split a degree-k matrix into blocks of sizes k and k + 1. - - :param Ak: Degree-k block matrix. - :param k: Representation degree. - :param N: Number of images. + block. - :return: The two permuted matrix blocks. + :return: Complex pairwise phase measurement whose real and imaginary + components are obtained from the two least-squares fits. """ - AkP = xp.copy(Ak) - dk = 2 * k + 1 - Pk = xp.eye(dk, dtype=AkP.dtype) - for m in range(k): - for n in range(k - m): - Pk[(m + 2 * n, m + 2 * n + 1), :] = Pk[(m + 2 * n + 1, m + 2 * n), :] - AkP = ( - xp.kron(xp.eye(N, dtype=AkP.dtype), Pk) - @ Ak - @ xp.kron(xp.eye(N, dtype=AkP.dtype), Pk.T) + A = np.array( + [ + [np.vdot(W1 + W4, W1 + W4), np.vdot(W1 + W4, W2 + W3)], + [np.vdot(W2 + W3, W1 + W4), np.vdot(W2 + W3, W2 + W3)], + ] ) + B = np.array([np.vdot(W1 + W4, Br), np.vdot(W2 + W3, Br)]) + a, _ = np.linalg.lstsq(A, B, rcond=None)[0] - Pk = xp.eye(N * dk, dtype=AkP.dtype) - idx = xp.concatenate((xp.arange(dk - k, dk), xp.arange(k + 1))) - for m in range(N - 1): - for n in range(N - 1 - m): - Pk[k * (m + 1) + n * dk : k * (m + 1) + (n + 1) * dk] = Pk[ - k * (m + 1) + n * dk : k * (m + 1) + (n + 1) * dk - ][idx, :] - AkP = Pk @ AkP @ Pk.T - return AkP[: N * k, : N * k], AkP[N * k :, N * k :] + A = np.array( + [ + [np.vdot(W1 - W4, W1 - W4), np.vdot(W1 - W4, W3 - W2)], + [np.vdot(W1 - W4, W3 - W2), np.vdot(W3 - W2, W3 - W2)], + ] + ) + B = np.array([np.vdot(W1 - W4, Bi), np.vdot(W3 - W2, Bi)]) + b, _ = np.linalg.lstsq(A, B, rcond=None)[0] - @staticmethod - def permutek_back(Ak, k, N): + return a + 1j * b + + def _dihedral_gamma_pair(self, wi, wj, DXijD, Jk, degree): """ - Undo the degree-k block permutation and reconstruct the full matrix. + Estimate the pairwise gamma phase for dihedral symmetry. - :param Ak: Permuted degree-k matrix. - :param k: Representation degree. - :param N: Number of images. + Form the four handedness-invariant outer products associated with the + outer Wigner columns for images i and j. Fit these terms to the real + and imaginary parts of the alpha-corrected representation block to + obtain the complex measurement used for gamma synchronization. - :return: Matrix in the original block ordering. + :param wi: Degree-k Wigner small-d matrix evaluated at the estimated + beta angle of image i. Shape (2 * k + 1, 2 * k + 1). + :param wj: Degree-k Wigner small-d matrix evaluated at the estimated + beta angle of image j. Shape (2 * k + 1, 2 * k + 1). + :param DXijD: Complex representation block for images i and j after + removing their estimated alpha phases. Shape (2 * k + 1, 2 * k + 1). + :param Jk: Degree-k handedness-conjugation matrix. Shape (2 * k + 1, 2 * k + 1). + :param degree: Wigner representation degree. This parameter is included + to provide the common pair-estimator interface used by `_estimate_gamma`. + + :return: Complex pairwise phase measurement used for gamma synchronization. """ - dk = 2 * k + 1 - Pk = xp.eye(N * dk, dtype=Ak.dtype) - idx = xp.concatenate((xp.arange(dk - k, dk), xp.arange(k + 1))) - for m in range(N - 1): - for n in range(N - 1 - m): - Pk[k * (m + 1) + n * dk : k * (m + 1) + (n + 1) * dk] = Pk[ - k * (m + 1) + n * dk : k * (m + 1) + (n + 1) * dk - ][idx, :] - AkB = Pk.T @ Ak @ Pk - dk = 2 * k + 1 - Pk = xp.eye(dk, dtype=Ak.dtype) - for m in range(k): - for n in range(k - m): - Pk[(m + 2 * n, m + 2 * n + 1), :] = Pk[(m + 2 * n + 1, m + 2 * n), :] - AkB = ( - xp.kron(xp.eye(N, dtype=Ak.dtype), Pk.T) - @ AkB - @ xp.kron(xp.eye(N, dtype=Ak.dtype), Pk) + outer_11 = np.outer(wi[:, 0], wj[:, 0]) + outer_21 = np.outer(wi[:, -1], wj[:, 0]) + outer_12 = np.outer(wi[:, 0], wj[:, -1]) + outer_22 = np.outer(wi[:, -1], wj[:, -1]) + + W1 = outer_11 + Jk @ outer_11 @ Jk + W2 = outer_21 + Jk @ outer_21 @ Jk + W3 = outer_12 + Jk @ outer_12 @ Jk + W4 = outer_22 + Jk @ outer_22 @ Jk + + return self._solve_dihedral_gamma_pair( + W1, + W2, + W3, + W4, + np.real(4 * DXijD), + np.imag(4 * DXijD), ) - return AkB - @staticmethod - def vec_block(A, N, sz, IDX_upper): + def _real_to_complex_representation(self, X, degree): """ - Vectorize the upper-triangular image-pair blocks of a block matrix. + Convert a degree-wise representation matrix from the real basis used + by ADMM to the complex Wigner basis. + + :param X: Representation matrix of shape + (n_img * (2 * degree + 1), n_img * (2 * degree + 1)). + :param degree: Wigner representation degree. + + :return: Representation matrix in the complex Wigner basis. """ - vecA = (A.reshape(N, sz, N, sz).transpose(0, 2, 3, 1)).reshape(N**2, sz**2).T - return vecA[:, IDX_upper] + T, Tinv = self.complex2real(degree) + identity = np.eye(self.n_img, dtype=np.float64) + + return np.kron(identity, T) @ X @ np.kron(identity, Tinv) @staticmethod - def largest_eigenvalue(AI, Ngrid, N): - """ - Estimate the largest eigenvalue of the Fejér constraint operator. + def _handedness_matrix(degree): """ - # find the largest eigenvalue of the operator AI - z = xp.random.normal(0, 1, (Ngrid, N**2)) - Lambda = 0 + Construct the handedness-conjugation matrix at a Wigner degree. - while abs(Lambda - xp.linalg.norm(z)) > 500: - Lambda = xp.linalg.norm(z) - z = z / xp.linalg.norm(z) - z = AI @ (AI.T @ z) - Lambda += 2000 - logger.info("Largest eigenvalue of AIAIT is approximately %1.2f" % Lambda) - return Lambda + :param degree: Wigner representation degree. - def compute_rank(self): + :return: Diagonal matrix of shape (2 * degree + 1, 2 * degree + 1). """ - Compute the ranks and matrices of the symmetry-averaging projectors at each degree. + signs = np.ones(2 * degree + 1) + signs[degree + 1 :: 2] = -1 + signs[degree - 1 :: -2] = -1 - :param Lmax: Maximum representation degree. + return np.diag(signs) - :return: Ranks and symmetry-averaging matrices for each degree. - """ - rk = xp.zeros(self.Lmax, dtype=np.float64) - A = [] - for k in range(1, self.Lmax + 1): - Ak = np.sum(self.WD(k, self.sym_euler), axis=0) - Ak = np.round(Ak / self.n_sym, 6) - A.append(Ak) - rk[k - 1] = np.linalg.matrix_rank(Ak) - return rk, A + ########################## + # Representation Helpers # + ########################## def WD(self, k, euler): """ @@ -2359,88 +2415,41 @@ def Wd(k, beta): ) return d - @staticmethod - def mat_block(vecA, N, sz, IDX_upper, IDX_lower, idx_offdiag): - """ - Reconstruct a symmetric block matrix from its vectorized upper-triangular blocks. - """ - tmp = vecA.T.reshape(N * (N + 1) // 2, sz, sz).transpose(0, 2, 1) - AA = xp.zeros((N**2, sz, sz), dtype=vecA.dtype) - AA[IDX_upper] = tmp - AA[IDX_lower] = tmp[idx_offdiag].transpose(0, 2, 1) - return (AA.reshape(N, N, sz, sz).transpose(0, 2, 1, 3)).reshape(N * sz, N * sz) - - @staticmethod - def psd_projection(B): - """ - Project one or more symmetric matrices onto the positive semidefinite cone. - """ - # compute the PSD part of a symmstric matrix - B_sym = (B + B.swapaxes(-1, -2)) / 2 - evals, evecs = xp.linalg.eigh(B_sym) - evals = xp.maximum(evals, 0) - return (evecs * evals[..., None, :]) @ evecs.swapaxes(-1, -2) - - @staticmethod - def transform_block(A, k, Pk): - """ - Permute and vectorize the two invariant blocks of a degree-k matrix. + def complex2real(self, ell): """ - AT = Pk @ A @ Pk.T - A0 = AT[:, :k, :k].swapaxes(-1, -2).reshape(A.shape[0], -1) - A1 = AT[:, k:, k:].swapaxes(-1, -2).reshape(A.shape[0], -1) - return A0, A1 + Construct the transformation matrices between complex and real degree ell representations. - @staticmethod - def transform_back_block(A0, A1, k, Pk): - """ - Reconstruct a degree-k matrix from its two invariant block vectors. - """ - dk = 2 * k + 1 - A = xp.zeros((A0.shape[0], dk, dk), dtype=A0.dtype) - A[:, :k, :k] = A0.reshape(-1, k, k).swapaxes(-1, -2) - A[:, k:, k:] = A1.reshape(-1, k + 1, k + 1).swapaxes(-1, -2) - return Pk.T @ A @ Pk + :param ell: Wigner representation degree. - def construct_AEq(self): - """ - Construct the linear equality operator encoding the quaternion constraints. + :return: Forward and inverse change-of-basis matrices. """ - AEq = np.zeros((17, 21), np.float64) - - # First 16 rows: identity constraints on first 16 variables - AEq[:16, :16] = np.eye(16, dtype=np.float64) - - # Columns 16:21 map the low-degree X0/X1 entries into the quaternion - # convex-hull constraint Xq = I/4 - linear(X^(1)). - extra = 0.25 * np.array( - [ - [-1, 1, 0, 0, 1], - [0, 0, 0, 0, 0], - [0, 0, 1, -1, 0], - [0, 0, 0, 0, 0], - [0, 0, 0, 0, 0], - [1, 1, 0, 0, -1], - [0, 0, 0, 0, 0], - [0, 0, 1, 1, 0], - [0, 0, 1, -1, 0], - [0, 0, 0, 0, 0], - [-1, -1, 0, 0, -1], - [0, 0, 0, 0, 0], - [0, 0, 0, 0, 0], - [0, 0, 1, 1, 0], - [0, 0, 0, 0, 0], - [1, -1, 0, 0, 1], - ], - dtype=np.float64, - ) - - AEq[:16, 16:] = extra + diml = 2 * ell + 1 + Tinv = np.zeros((diml, diml), dtype=complex_type(np.float64)) + for i in range(diml): + if i < ell: + Tinv[i, i] = 1j / np.sqrt(2) + Tinv[i, diml - 1 - i] = -1j * (-1) ** (i - ell) / np.sqrt(2) + if i == ell: + Tinv[i, i] = 1 + if i > ell: + Tinv[i, i] = (-1) ** (i - ell) / np.sqrt(2) + Tinv[i, diml - 1 - i] = 1 / np.sqrt(2) - # Last row: redundant trace/sum constraint - AEq[16, [0, 5, 10, 15]] = 1 + T = np.zeros((diml, diml), dtype=complex_type(np.float64)) + for i in range(diml): + if i < ell: + T[i, i] = -1j / np.sqrt(2) + T[i, diml - 1 - i] = 1 / np.sqrt(2) + if i == ell: + T[i, i] = 1 + if i > ell: + T[i, i] = (-1) ** (i - ell) / np.sqrt(2) + T[i, diml - 1 - i] = 1j * (-1) ** (i - ell) / np.sqrt(2) + return T, Tinv - return AEq + ###################### + # Validation Helpers # + ###################### def form_ground_truth_X(self, euler_angles): """ From 502cf1850ab9df696e184116f4f2a105457e4597 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 9 Sep 2026 10:02:30 -0400 Subject: [PATCH 085/161] Add NUG dataset to downloader. Add gallery experiment reproducing paper results --- .../experimental_nug_pipeline_10081.py | 113 ++++++++++++++++++ src/aspire/downloader/__init__.py | 1 + src/aspire/downloader/data_fetcher.py | 14 +++ src/aspire/downloader/registry.py | 3 + 4 files changed, 131 insertions(+) create mode 100644 gallery/experiments/experimental_nug_pipeline_10081.py diff --git a/gallery/experiments/experimental_nug_pipeline_10081.py b/gallery/experiments/experimental_nug_pipeline_10081.py new file mode 100644 index 0000000000..7d07c90240 --- /dev/null +++ b/gallery/experiments/experimental_nug_pipeline_10081.py @@ -0,0 +1,113 @@ +""" +Symmetric NUG Pipeline +====================== + +This notebook demonstrates reproducing simliar results to the +experiment found in: + +.. admonition:: Publication + + | Shi, Y., Singer, A., Wang, J., & Yang, R. + | Orientation estimation of cryo-EM images for molecules with + | arbitrary symmetry by semidefinite programming + +This pipeline begins with 17 class averages that were generated (as described +in the publication) from the EMPIAR 10081 picked particle data, available here: + +https://www.ebi.ac.uk/empiar/EMPIAR-10081 + +First, orientation estimation is performed using the Non-Unique Games +algorithm adapted for handling symmetric particles. A volume is then +reconstructed from the estimated orientations and saved to an output +file for inspection. The reconstructed volume can be compared with the +EMPIAR-10081 associated reconstruction found here: + +https://www.ebi.ac.uk/emdb/EMD-8511 +""" + +# %% +# Imports +# ------- +# First import the necessary utilities and ASPIRE +# classes used throughout the notebook. +import logging +import numpy as np + +from aspire.abinitio import CommonlineNUG +from aspire.downloader import nug_10081 +from aspire.image import Image +from aspire.reconstruction import MeanEstimator +from aspire.source import ArrayImageSource, OrientedSource + +logger = logging.getLogger(__name__) + + +# %% +# Setup +# ----- + +# Note: Use of GPU is recommended for the NUG algorithm. +# Below we setup some output files to save off orientation +# and reconstruction results. +oriented_fn = "nug_10081_oriented.star" +volume_output_fn = "nug_10081_abinitio.mrc" + + +# %% +# Load the Dataset +# ---------------- +# +# The 17 class averages can be accessed as an ASPIRE ``Image`` object +# via the built-in downloader utility. + +logger.info("Load Precomputed Class Averages") +ims = nug_10081() +avgs = ArrayImageSource(ims, pixel_size=1.3, symmetry_group="C4") + + +# %% +# Orientation Estimation +# ---------------------- +# +# Create an orientation estimation object for the ``avgs``. +# The ``CommonlineNUG`` algorithm will detect the symmetry from +# the ``avgs`` metadata and use the approriate NUG solver to estimate +# orientations. This experiment uses 15 iterations of the optional +# proximal refinement method to improve the initial estimation. + +logger.info("Begin Orientation Estimation") +orient_est = CommonlineNUG(avgs, pr_iters=15) + +# Create an ``OrientedSource`` class instance that performs orientation +# estimation in a lazy fashion upon request of images or rotations. +oriented_src = OrientedSource(avgs, orient_est) + +# Estimate orientations and override the estimated shifts. +# In the case of insufficient number of images, estimating shifts fails. +oriented_src.rotations +oriented_src = oriented_src.update(offsets=0) + +# Confirm offsets have been set to zero. +logger.info(f"Offsets: {oriented_src.offsets}") + +# Save oriented source. +oriented_src.save(oriented_fn) + +# %% +# Volume Reconstruction +# ---------------------- +# +# Using the oriented source, attempt to reconstruct a volume. Since +# this is a C4 symmetric molecule the ``symmetry_group`` source attribute +# will flow through the pipeline to ``oriented_src``. Then each image will be +# repeated for each of the 3 additional rotations during +# back-projection. In general, this boosts the effective number of images used in +# the reconstruction from ``n_imgs`` to ``4 * n_imgs``, or from 17 to 68 in this case. + +# Setup an estimator to perform the back projection. +logger.info("Begin Volume reconstruction") +estimator = MeanEstimator(oriented_src) + +# Perform the estimation and save the reconstructed volume. +est_vol = estimator.estimate() +est_vol.save(volume_output_fn) diff --git a/src/aspire/downloader/__init__.py b/src/aspire/downloader/__init__.py index 714bba0e53..516047857f 100644 --- a/src/aspire/downloader/__init__.py +++ b/src/aspire/downloader/__init__.py @@ -20,4 +20,5 @@ emdb_14621, remove_downloads, simulated_channelspin, + nug_10081, ) diff --git a/src/aspire/downloader/data_fetcher.py b/src/aspire/downloader/data_fetcher.py index efaadfe5a0..1f5e1610c5 100644 --- a/src/aspire/downloader/data_fetcher.py +++ b/src/aspire/downloader/data_fetcher.py @@ -408,3 +408,17 @@ def simulated_channelspin(): data["rots"] = Rotation(_LegacySimulation.rots_zyx_to_legacy_aspire(data["rots"])) return data + + +def nug_10081(): + """ + Download class averages used in the symmetric NUG experiment. + + The dataset contains 17 manually selected, 256-by-256 RELION class + averages derived from EMPIAR-10081. The class averages were used for + the C4 experimental example in the symmetric NUG manuscript. + + :return: An `Image` instance containing the 17 class averages. + """ + file_path = fetch_data("nug_10081.npy") + return Image(np.load(file_path).astype(np.float64), pixel_size=1.3) diff --git a/src/aspire/downloader/registry.py b/src/aspire/downloader/registry.py index 7baac9038f..af25099568 100644 --- a/src/aspire/downloader/registry.py +++ b/src/aspire/downloader/registry.py @@ -21,6 +21,7 @@ "emdb_9571.map": "sha256:e399b5a53c23e06e2384b15150ad1f9e6120a4799ea1fdfcaf213948d959bb65", "emdb_10920.map": "sha256:8e08a511903d2f088b5ed69446843dad8214fa34128e513b90a3fc75cf3fa795", "simulated_channelspin.npz": "c0752674acb85417f6a77a28ac55280c1926c73fda9e25ce0a9940728b1dfcc8", + "nug_10081.npy": "8167ba8d802cda80eb485dff16365d523ea45aa3131915bc1a06b73ca3f27ffa", } registry_urls = { @@ -45,6 +46,7 @@ "emdb_9571.map": "https://ftp.ebi.ac.uk/pub/databases/emdb/structures/EMD-9571/map/emd_9571.map.gz", "emdb_10920.map": "https://ftp.ebi.ac.uk/pub/databases/emdb/structures/EMD-10920/map/emd_10920.map.gz", "simulated_channelspin.npz": "https://zenodo.org/records/8186548/files/example_FakeKV_dataset.npz", + "nug_10081.npy": "https://zenodo.org/records/22233860/files/empiar_10081_nug_17_class_averages.npy", } file_to_method_map = { @@ -69,4 +71,5 @@ "emdb_9571.map": "emdb_9571", "emdb_10920.map": "emdb_10920", "simulated_channelspin.npz": "simulated_channelspin", + "nug_10081.npy": "nug_10081", } From 7658c6df7b99a13180d76e4cc9e4a98db75e8b5c Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 9 Sep 2026 10:28:57 -0400 Subject: [PATCH 086/161] uncomment admm random update order --- src/aspire/abinitio/commonline_nug.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 91b0575342..3d126e914c 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -655,7 +655,7 @@ def print_updates(verbose): IDX = np.arange(3) for t in range(max_iter): - # np.random.shuffle(IDX) + np.random.shuffle(IDX) for idx in IDX: if idx == 0: S0, S1, Sq = update_S(C0, C1, yE0, yE1, yEq, yI, X0, X1, Xq, rho) From d298cabe49ab5149e3a57a723e9811b6ec16a016 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 10 Sep 2026 10:39:39 -0400 Subject: [PATCH 087/161] disable shift search space in gallery --- .../experimental_nug_pipeline_10081.py | 16 +++++++++------- src/aspire/downloader/__init__.py | 2 +- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/gallery/experiments/experimental_nug_pipeline_10081.py b/gallery/experiments/experimental_nug_pipeline_10081.py index 7d07c90240..c43c19cd7c 100644 --- a/gallery/experiments/experimental_nug_pipeline_10081.py +++ b/gallery/experiments/experimental_nug_pipeline_10081.py @@ -31,6 +31,7 @@ # First import the necessary utilities and ASPIRE # classes used throughout the notebook. import logging + import numpy as np from aspire.abinitio import CommonlineNUG @@ -58,9 +59,9 @@ # ---------------- # # The 17 class averages can be accessed as an ASPIRE ``Image`` object -# via the built-in downloader utility. +# via the built-in downloader utility. -logger.info("Load Precomputed Class Averages") +logger.info("Load Precomputed Class Averages") ims = nug_10081() avgs = ArrayImageSource(ims, pixel_size=1.3, symmetry_group="C4") @@ -71,12 +72,13 @@ # # Create an orientation estimation object for the ``avgs``. # The ``CommonlineNUG`` algorithm will detect the symmetry from -# the ``avgs`` metadata and use the approriate NUG solver to estimate -# orientations. This experiment uses 15 iterations of the optional -# proximal refinement method to improve the initial estimation. +# the ``avgs`` metadata and use the approriate NUG ADMM-solver to estimate +# orientations. To replicate the original experiment as closely as possible, +# we use the 15 iterations of proximal refinement and disable searching over +# the commonline shift space. logger.info("Begin Orientation Estimation") -orient_est = CommonlineNUG(avgs, pr_iters=15) +orient_est = CommonlineNUG(avgs, max_shift=0, pr_iters=15) # Create an ``OrientedSource`` class instance that performs orientation # estimation in a lazy fashion upon request of images or rotations. @@ -88,7 +90,7 @@ oriented_src = oriented_src.update(offsets=0) # Confirm offsets have been set to zero. -logger.info(f"Offsets: {oriented_src.offsets}") +logger.info(f"Offsets: {oriented_src.offsets}") # Save oriented source. oriented_src.save(oriented_fn) diff --git a/src/aspire/downloader/__init__.py b/src/aspire/downloader/__init__.py index 516047857f..eba986696e 100644 --- a/src/aspire/downloader/__init__.py +++ b/src/aspire/downloader/__init__.py @@ -18,7 +18,7 @@ emdb_8511, emdb_10835, emdb_14621, + nug_10081, remove_downloads, simulated_channelspin, - nug_10081, ) From e1bc5dc9d5b21df94d40d6091b26d0e8d8a5fb00 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 10 Sep 2026 15:49:21 -0400 Subject: [PATCH 088/161] tox --- gallery/experiments/experimental_nug_pipeline_10081.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/gallery/experiments/experimental_nug_pipeline_10081.py b/gallery/experiments/experimental_nug_pipeline_10081.py index c43c19cd7c..ba76aa45bd 100644 --- a/gallery/experiments/experimental_nug_pipeline_10081.py +++ b/gallery/experiments/experimental_nug_pipeline_10081.py @@ -32,11 +32,8 @@ # classes used throughout the notebook. import logging -import numpy as np - from aspire.abinitio import CommonlineNUG from aspire.downloader import nug_10081 -from aspire.image import Image from aspire.reconstruction import MeanEstimator from aspire.source import ArrayImageSource, OrientedSource From bc2285abc21e6977c77b255798dd5cff1c14fd91 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 16 Sep 2026 15:58:55 -0400 Subject: [PATCH 089/161] Add gallery for D3 experiment. Add D3 dataset to downloader. --- .../experimental_nug_pipeline_12036.py | 112 ++++++++++++++++++ src/aspire/downloader/__init__.py | 1 + src/aspire/downloader/data_fetcher.py | 20 +++- src/aspire/downloader/registry.py | 3 + 4 files changed, 134 insertions(+), 2 deletions(-) create mode 100644 gallery/experiments/experimental_nug_pipeline_12036.py diff --git a/gallery/experiments/experimental_nug_pipeline_12036.py b/gallery/experiments/experimental_nug_pipeline_12036.py new file mode 100644 index 0000000000..4db3e88722 --- /dev/null +++ b/gallery/experiments/experimental_nug_pipeline_12036.py @@ -0,0 +1,112 @@ +""" +Dn-Symmetric NUG Pipeline +========================= + +This notebook demonstrates reproducing simliar results to the +experiment found in: + +.. admonition:: Publication + + | Shi, Y., Singer, A., Wang, J., & Yang, R. + | Orientation estimation of cryo-EM images for molecules with + | arbitrary symmetry by semidefinite programming + +This pipeline begins with 93 class averages that were generated (as described +in the publication) from the EMPIAR 12036 picked particle data, available here: + +https://www.ebi.ac.uk/empiar/EMPIAR-12036 + +First, orientation estimation is performed using the Non-Unique Games +algorithm adapted for handling symmetric particles. A volume is then +reconstructed from the estimated orientations and saved to an output +file for inspection. The reconstructed volume can be compared with the +EMPIAR-12036 associated reconstruction found here: + +https://www.ebi.ac.uk/emdb/EMD-14803 +""" + +# %% +# Imports +# ------- +# First import the necessary utilities and ASPIRE +# classes used throughout the notebook. +import logging + +from aspire.abinitio import CommonlineNUG +from aspire.downloader import nug_12036 +from aspire.image import Image +from aspire.reconstruction import MeanEstimator +from aspire.source import ArrayImageSource, OrientedSource + +logger = logging.getLogger(__name__) + + +# %% +# Setup +# ----- + +# Note: Use of GPU is recommended for the NUG algorithm. +# Below we setup some output files to save off orientation +# and reconstruction results. +oriented_fn = "nug_12036_oriented.star" +volume_output_fn = "nug_12036_abinitio.mrc" + + +# %% +# Load the Dataset +# ---------------- +# +# The 93 class averages can be accessed as an ASPIRE ``Image`` object +# via the built-in downloader utility. +logger.info("Load Precomputed Class Averages") +ims = nug_12036() +avgs = ArrayImageSource(ims, symmetry_group="D3") + + +# %% +# Orientation Estimation +# ---------------------- +# +# Create an orientation estimation object for the ``avgs``. +# The ``CommonlineNUG`` algorithm will detect the symmetry from +# the ``avgs`` metadata and use the approriate NUG ADMM-solver to estimate +# orientations. To replicate the original experiment as closely as possible, +# we use the 5 iterations of proximal refinement and disable searching over +# the commonline shift space. + +logger.info("Begin Orientation Estimation") +orient_est = CommonlineNUG(avgs, max_shift=0, pr_iters=5) + +# Create an ``OrientedSource`` class instance that performs orientation +# estimation in a lazy fashion upon request of images or rotations. +oriented_src = OrientedSource(avgs, orient_est) + +# Estimate orientations and override the estimated shifts. +# In the case of insufficient number of images, estimating shifts fails. +oriented_src.rotations +oriented_src = oriented_src.update(offsets=0) + +# Confirm offsets have been set to zero. +logger.info(f"Offsets: {oriented_src.offsets}") + +# Save oriented source. +oriented_src.save(oriented_fn) + +# %% +# Volume Reconstruction +# ---------------------- +# +# Using the oriented source, attempt to reconstruct a volume. Since +# this is a D3 symmetric molecule the ``symmetry_group`` source attribute +# will flow through the pipeline to ``oriented_src``. Then each image will be +# repeated for each of the 5 additional symmetry induced rotations during +# back-projection. In general, this boosts the effective number of images used in +# the reconstruction from ``n_imgs`` to ``6 * n_imgs``, or from 93 to 558 in this case. + +# Setup an estimator to perform the back projection. +logger.info("Begin Volume reconstruction") +estimator = MeanEstimator(oriented_src) + +# Perform the estimation and save the reconstructed volume. +est_vol = estimator.estimate() +est_vol.save(volume_output_fn) diff --git a/src/aspire/downloader/__init__.py b/src/aspire/downloader/__init__.py index eba986696e..9151e7a2d1 100644 --- a/src/aspire/downloader/__init__.py +++ b/src/aspire/downloader/__init__.py @@ -19,6 +19,7 @@ emdb_10835, emdb_14621, nug_10081, + nug_12036, remove_downloads, simulated_channelspin, ) diff --git a/src/aspire/downloader/data_fetcher.py b/src/aspire/downloader/data_fetcher.py index 1f5e1610c5..9a53320979 100644 --- a/src/aspire/downloader/data_fetcher.py +++ b/src/aspire/downloader/data_fetcher.py @@ -415,10 +415,26 @@ def nug_10081(): Download class averages used in the symmetric NUG experiment. The dataset contains 17 manually selected, 256-by-256 RELION class - averages derived from EMPIAR-10081. The class averages were used for - the C4 experimental example in the symmetric NUG manuscript. + averages with 1.3 Angstrom pixel size, derived from EMPIAR-10081. + The class averages were used for the C4 experimental example in the + symmetric NUG manuscript. :return: An `Image` instance containing the 17 class averages. """ file_path = fetch_data("nug_10081.npy") return Image(np.load(file_path).astype(np.float64), pixel_size=1.3) + + +def nug_12036(): + """ + Download class averages used in the symmetric NUG experiment. + + The dataset contains 93 manually selected, 320-by-320 RELION class + averages with 0.945 Angstrom pixel size, derived from EMPIAR-12036. + The class averages were used for the D3 experimental example in the + symmetric NUG manuscript. + + :return: An `Image` instance containing the 93 class averages. + """ + file_path = fetch_data("nug_12036.mrcs") + return Image.load(file_path) diff --git a/src/aspire/downloader/registry.py b/src/aspire/downloader/registry.py index af25099568..0d1a56e264 100644 --- a/src/aspire/downloader/registry.py +++ b/src/aspire/downloader/registry.py @@ -22,6 +22,7 @@ "emdb_10920.map": "sha256:8e08a511903d2f088b5ed69446843dad8214fa34128e513b90a3fc75cf3fa795", "simulated_channelspin.npz": "c0752674acb85417f6a77a28ac55280c1926c73fda9e25ce0a9940728b1dfcc8", "nug_10081.npy": "8167ba8d802cda80eb485dff16365d523ea45aa3131915bc1a06b73ca3f27ffa", + "nug_12036.mrcs": "7981a9a1e4343cbb6b1323630eb08926b38b5058074d81ef1663f719444f3569", } registry_urls = { @@ -47,6 +48,7 @@ "emdb_10920.map": "https://ftp.ebi.ac.uk/pub/databases/emdb/structures/EMD-10920/map/emd_10920.map.gz", "simulated_channelspin.npz": "https://zenodo.org/records/8186548/files/example_FakeKV_dataset.npz", "nug_10081.npy": "https://zenodo.org/records/22233860/files/empiar_10081_nug_17_class_averages.npy", + "nug_12036.mrcs": "https://zenodo.org/records/22801412/files/empiar_12036_nug_93_class_averages.mrcs", } file_to_method_map = { @@ -72,4 +74,5 @@ "emdb_10920.map": "emdb_10920", "simulated_channelspin.npz": "simulated_channelspin", "nug_10081.npy": "nug_10081", + "nug_12036.mrcs": "nug_12036", } From eecfaa326b6a3788aff5a057947ac86a2d069b67 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 16 Sep 2026 16:03:05 -0400 Subject: [PATCH 090/161] remove unused import --- gallery/experiments/experimental_nug_pipeline_12036.py | 1 - 1 file changed, 1 deletion(-) diff --git a/gallery/experiments/experimental_nug_pipeline_12036.py b/gallery/experiments/experimental_nug_pipeline_12036.py index 4db3e88722..86d19cd7ec 100644 --- a/gallery/experiments/experimental_nug_pipeline_12036.py +++ b/gallery/experiments/experimental_nug_pipeline_12036.py @@ -34,7 +34,6 @@ from aspire.abinitio import CommonlineNUG from aspire.downloader import nug_12036 -from aspire.image import Image from aspire.reconstruction import MeanEstimator from aspire.source import ArrayImageSource, OrientedSource From 7073929af95c5403a1e74b3ecde23f85001b608b Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 17 Sep 2026 11:44:01 -0400 Subject: [PATCH 091/161] handle no proximal refine for C1 in init. Add test. --- src/aspire/abinitio/commonline_nug.py | 20 ++++++++++++-------- tests/test_nug.py | 17 +++++++++++++++++ 2 files changed, 29 insertions(+), 8 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 3d126e914c..0382712150 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -79,8 +79,10 @@ def __init__( :param S2_grid: Number of sphere samples used to discretize SO(3). Default is 441. :param Nstep_yI: Number of inequality-multiplier updates per ADMM iteration. Default is 10. - :param pr_iters: Number of proximal refinement iterations. Default of None - does not perform proximal refinement. Recommended value when using is 4. + :param pr_iters: Number of proximal-refinement iterations for nontrivial + cyclic and dihedral symmetry groups. Default is None, which disables + proximal refinement. Supplying a value for an asymmetric (C1) + problem raises a ValueError. :param verbose: Whether to log ADMM and proximal refinement progress. Default is True. """ @@ -137,7 +139,14 @@ def __init__( self.sym_euler = sym_rotations.angles self.n_sym = len(self.sym_euler) - # Set up proximal refinement terms + # Proximal refinement is implemented only for cyclic and dihedral symmetry groups. + if self.n_sym == 1 and pr_iters is not None: + raise ValueError( + "Proximal refinement is not supported for asymmetric (C1) NUG. " + "Set pr_iters=None." + ) + + # Set up proximal refinement terms. if pr_iters is not None: self.pr_weights = 1 / (1 + np.arange(self.Lmax)) self.pr_penalty = [1] * pr_iters @@ -317,11 +326,6 @@ def perform_admm(self): X_est = self.admm_sym_J(self.C, self.verbose) if self.pr_iters is not None: - if is_asymmetric: - raise NotImplementedError( - "Proximal refinement is not yet implemented for asymmetric NUG." - ) - X_est = self.proximal_refine( X_est, self.pr_weights, diff --git a/tests/test_nug.py b/tests/test_nug.py index c2cd8d0847..8d10c59c2d 100644 --- a/tests/test_nug.py +++ b/tests/test_nug.py @@ -219,3 +219,20 @@ def test_symmetry_logging(caplog): assert str(orient_est.sym_grp) == "D3" assert "Provided symmetry, D3, does not match source, C3" in caplog.text assert "Using provided symmetry: D3" in caplog.text + + +def test_c1_proximal_refinement_raises(dtype): + vol = emdb_2660().astype(dtype).downsample(16) + src = Simulation( + n=3, + vols=vol, + offsets=0, + amplitudes=1, + seed=SEED, + ) + + with pytest.raises( + ValueError, + match=r"Proximal refinement is not supported for asymmetric", + ): + CommonlineNUG(src, pr_iters=1) From fa5596a81450088cc8f5918da977838b264f98cc Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 17 Sep 2026 13:33:53 -0400 Subject: [PATCH 092/161] test docstring --- tests/test_nug.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/tests/test_nug.py b/tests/test_nug.py index 8d10c59c2d..2149b0edeb 100644 --- a/tests/test_nug.py +++ b/tests/test_nug.py @@ -222,14 +222,12 @@ def test_symmetry_logging(caplog): def test_c1_proximal_refinement_raises(dtype): + """ + Proximal refinement is only supported for cyclic and dihedral problems. + Check we raise when pr_iters are provided. + """ vol = emdb_2660().astype(dtype).downsample(16) - src = Simulation( - n=3, - vols=vol, - offsets=0, - amplitudes=1, - seed=SEED, - ) + src = Simulation(n=3, vols=vol) with pytest.raises( ValueError, From dfb4ffcd1a01b3d87a9f22ea6fa2d29dae8c7e40 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 17 Sep 2026 14:02:53 -0400 Subject: [PATCH 093/161] update init docstring --- src/aspire/abinitio/commonline_nug.py | 24 ++++++++++++++---------- 1 file changed, 14 insertions(+), 10 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 0382712150..d155efc545 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -48,15 +48,19 @@ def __init__( **kwargs, ): """ - Initialize the symmetric NUG orientation estimator. All default values match those used - for publication results, with exception of `pr_iters`, which controls how many iterations - of proximal refinement to be used. The default of None runs the algorithm without proximal - refinement. Set `pr_iters=4` to match t3he proximal refinement workflow in the related - publication. + Initialize the NUG orientation estimator. + + The estimator supports asymmetric molecules (C1) and molecules with + cyclic (Cn) or dihedral (Dn) symmetry. Most defaults match those used + for the experiments in the associated publication. Proximal refinement + is available only for nontrivial cyclic and dihedral symmetry groups. + + For complete experimental workflows, see the NUG examples in + ``gallery/experiments``. :param src: Source containing the input projection images. - :param symmetry: Cyclic or dihedral symmetry specification, such as 'C3' or - 'D4'. If omitted, uses the source symmetry. + :param symmetry: Cyclic or dihedral symmetry specification, such as 'C1', 'C3', + or 'D4'. If None, use the symmetry specified by `src`. Default is None. :param n_rad: Number of radial samples in the polar Fourier transform. If None, n_rad will default to the ceiling of half the resolution of the source. :param n_theta: Number of angular samples in the polar Fourier transform. This @@ -80,9 +84,9 @@ def __init__( :param Nstep_yI: Number of inequality-multiplier updates per ADMM iteration. Default is 10. :param pr_iters: Number of proximal-refinement iterations for nontrivial - cyclic and dihedral symmetry groups. Default is None, which disables - proximal refinement. Supplying a value for an asymmetric (C1) - problem raises a ValueError. + cyclic and dihedral symmetry groups. If None, do not perform proximal + refinement. Supplying a value for an asymmetric C1 problem raises a + `ValueError`. Default is None. :param verbose: Whether to log ADMM and proximal refinement progress. Default is True. """ From b1ea62fdef0fcda48433a38e8a7d008ab9c35810 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 17 Sep 2026 14:41:24 -0400 Subject: [PATCH 094/161] _largest_eigenvalue docs --- src/aspire/abinitio/commonline_nug.py | 20 ++++++++++++++++---- 1 file changed, 16 insertions(+), 4 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index d155efc545..08f8488206 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1257,7 +1257,7 @@ def ADMM_preprocessing(self, C): bI = -(Lmax + 2) * (Lmax + 1) / 2 # largest eigenvalue for AIAIT - Lambda = self.largest_eigenvalue(AI_mat_offdiag, Ngrid, N) + Lambda = self._largest_eigenvalue(AI_mat_offdiag, N) # initialization II = [] @@ -1374,12 +1374,24 @@ def discretize_SO3(self): return SO3 @staticmethod - def largest_eigenvalue(AI, Ngrid, N): + def _largest_eigenvalue(AI, N): """ - Estimate the largest eigenvalue of the Fejér constraint operator. + Estimate the largest eigenvalue of the Fejér constraint operator, + AI @ AI.T, using power iteration. + + A fixed empirical margin is added to the estimate, following the + researcher's implementation, to obtain a more conservative step + size for the inequality-multiplier update. + + :param AI: Matrix representation of the Fejer inequality operator, + with shape (Ngrid, n_coefficients). + :param N: Number of images. + + :return: Estimated largest eigenvalue of AI @ AI.T. """ # find the largest eigenvalue of the operator AI - z = xp.random.normal(0, 1, (Ngrid, N**2)) + n_grid = AI.shape[0] + z = xp.random.normal(0, 1, (n_grid, N**2)) Lambda = 0 while abs(Lambda - xp.linalg.norm(z)) > 500: From 53e8fd54959eceda47ec3aea0ef8524cac97bd79 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 17 Sep 2026 14:46:35 -0400 Subject: [PATCH 095/161] psd_projection docs --- src/aspire/abinitio/commonline_nug.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 08f8488206..a689f2b603 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1521,7 +1521,15 @@ def _project_representation_blocks( @staticmethod def psd_projection(B): """ - Project one or more symmetric matrices onto the positive semidefinite cone. + Project real sqaure matrices onto the positive semidefinite cone. + + Each matrix is first symmetrized. Negative eigenvalues are then + replaced with zero before the matrix is reconstructed. Leading + dimensions are treated as stack dimensions. + + :param B: Matrix or stack of matrices with shape (..., n, n). + + :return: Positive semidefinite projection with same shape as B. """ # compute the PSD part of a symmstric matrix B_sym = (B + B.swapaxes(-1, -2)) / 2 From bcd70cfc1a6fd241fc7f7848398039983fd949bb Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 17 Sep 2026 15:18:51 -0400 Subject: [PATCH 096/161] vec/mat_block docs --- src/aspire/abinitio/commonline_nug.py | 35 +++++++++++++++++++++++++-- 1 file changed, 33 insertions(+), 2 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index a689f2b603..a27878bfaf 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1659,7 +1659,20 @@ def permutek_back(Ak, k, N): @staticmethod def vec_block(A, N, sz, IDX_upper): """ - Vectorize the upper-triangular image-pair blocks of a block matrix. + Pack the upper-triangular image-pair blocks of a block matrix. + + The input is interpreted as an N-by-N array of sz-by-sz blocks. + Blocks from the upper triangle, including the diagonal, are vectorized + column-wise and stored as columns of the returned array. + + :param A: Block matrix with shape (N * sz, N * sz). + :param N: Number of block rows and columns, corresponding to the + number of images. + :param sz: Row and column size of each square block. + :param IDX_upper: Linear indices of the upper-triangular blocks in + the flattened N-by-N block grid. + + :return: Packed block array with shape (sz**2, N * (N + 1) // 2). """ vecA = (A.reshape(N, sz, N, sz).transpose(0, 2, 3, 1)).reshape(N**2, sz**2).T return vecA[:, IDX_upper] @@ -1667,7 +1680,25 @@ def vec_block(A, N, sz, IDX_upper): @staticmethod def mat_block(vecA, N, sz, IDX_upper, IDX_lower, idx_offdiag): """ - Reconstruct a symmetric block matrix from its vectorized upper-triangular blocks. + Unpack upper-triangular blocks into a symmetric block matrix. + + Each column of `vecA` is reshaped column-wise into an sz-by-sz block. + The blocks are placed in the upper triangle of an N-by-N block grid, + including the diagonal. Transposes of the off-diagonal blocks are placed + in the corresponding lower-triangular positions. + + :param vecA: Packed block array with shape (sz**2, N * (N + 1) // 2). + :param N: Number of block rows and columns, corresponding to the + number of images. + :param sz: Row and column size of each square block. + :param IDX_upper: Linear indices of the upper-triangular blocks in + the flattened N-by-N block grid. + :param IDX_lower: Linear indices of the corresponding + lower-triangular off-diagonal blocks in the flattened block grid. + :param idx_offdiag: Column indices selecting the off-diagonal blocks + from vecA. + + :return: Symmetric block matrix with shape (N * sz, N * sz). """ tmp = vecA.T.reshape(N * (N + 1) // 2, sz, sz).transpose(0, 2, 1) AA = xp.zeros((N**2, sz, sz), dtype=vecA.dtype) From 7662262030aee9be5ea562ed50966a948f18815e Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 17 Sep 2026 15:31:44 -0400 Subject: [PATCH 097/161] transform_block/transform_back_block docs --- src/aspire/abinitio/commonline_nug.py | 28 +++++++++++++++++++++++++-- 1 file changed, 26 insertions(+), 2 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index a27878bfaf..706ead1243 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1709,7 +1709,18 @@ def mat_block(vecA, N, sz, IDX_upper, IDX_lower, idx_offdiag): @staticmethod def transform_block(A, k, Pk): """ - Permute and vectorize the two invariant blocks of a degree-k matrix. + Permute and pack the two components of degree-k matrices. + + Each input matrix is transformed as Pk @ A @ Pk.T. The leading k-by-k + and trailing (k + 1)-by-(k + 1) diagonal blocks are then vectorized + column-wise. Entries in the two off-diagonal blocks are not retained. + + :param A: Batch of degree-k matrices with shape (n_blocks, 2 * k + 1, 2 * k + 1). + :param k: Wigner representation degree. + :param Pk: Permutation matrix with shape (2 * k + 1, 2 * k + 1). + + :return: Packed leading and trailing block arrays with shapes + (n_blocks, k**2) and (n_blocks, (k + 1)**2), respectively. """ AT = Pk @ A @ Pk.T A0 = AT[:, :k, :k].swapaxes(-1, -2).reshape(A.shape[0], -1) @@ -1719,7 +1730,20 @@ def transform_block(A, k, Pk): @staticmethod def transform_back_block(A0, A1, k, Pk): """ - Reconstruct a degree-k matrix from its two invariant block vectors. + Reconstruct degree-k matrices from two packed components. + + Each row of A0 and A1 is interpreted as a column-wise vectorization + of a k-by-k or (k + 1)-by-(k + 1) block. The two blocks are placed + on the diagonal of a (2 * k + 1)-by-(2 * k + 1) matrix, with zeros in the + off-diagonal blocks. The permutation represented by Pk is then reversed. + + :param A0: Packed leading blocks with shape (n_blocks, k**2). + :param A1: Packed trailing blocks with shape (n_blocks, (k + 1)**2). + :param k: Wigner representation degree. + :param Pk: Permutation matrix with shape (2 * k + 1, 2 * k + 1). + + :return: Reconstructed degree-k matrices with shape + (n_blocks, 2 * k + 1, 2 * k + 1). """ dk = 2 * k + 1 A = xp.zeros((A0.shape[0], dk, dk), dtype=A0.dtype) From 58731135c43e0451ae4423efdb4244110cc1bafe Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 18 Sep 2026 09:21:04 -0400 Subject: [PATCH 098/161] docs construct_AEq, WD, Wd --- src/aspire/abinitio/commonline_nug.py | 28 ++++++++++++++++++++++++--- 1 file changed, 25 insertions(+), 3 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 706ead1243..5ece177bed 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1406,8 +1406,6 @@ def compute_rank(self): """ Compute the ranks and matrices of the symmetry-averaging projectors at each degree. - :param Lmax: Maximum representation degree. - :return: Ranks and symmetry-averaging matrices for each degree. """ rk = xp.zeros(self.Lmax, dtype=np.float64) @@ -1422,6 +1420,14 @@ def compute_rank(self): def construct_AEq(self): """ Construct the linear equality operator encoding the quaternion constraints. + + The operator acts on 21 variables consisting of the 16 vectorized + entries of an auxiliary 4-by-4 quaternion matrix followed by five + entries from the lowest-degree representation blocks. The first + 16 rows couple the quaternion and representation variables. The + final row constrains the trace of the quaternion matrix. + + :return: Equality-constraint matrix with shape (17, 21). """ AEq = np.zeros((17, 21), np.float64) @@ -1521,7 +1527,7 @@ def _project_representation_blocks( @staticmethod def psd_projection(B): """ - Project real sqaure matrices onto the positive semidefinite cone. + Project real square matrices onto the positive semidefinite cone. Each matrix is first symmetrized. Negative eigenvalues are then replaced with zero before the matrix is reconstructed. Leading @@ -2454,6 +2460,11 @@ def _handedness_matrix(degree): def WD(self, k, euler): """ Evaluate degree-k Wigner D matrices at the supplied ZYZ Euler angles. + + :param k: Wigner representation degree. + :param euler: ZYZ Euler angles with shape (n_rotations, 3). + + :return: Complex Wigner D matrices with shape (n_rotations, 2 * k + 1, 2 * k + 1). """ # compute Wigner D matrix alpha = euler[:, 0] @@ -2472,6 +2483,17 @@ def WD(self, k, euler): def Wd(k, beta): """ Evaluate degree-k Wigner small-d matrices at the supplied polar angles. + + Matrix rows and columns correspond to orders from -k through k. These + real, beta-dependent matrices are used to construct the corresponding + Wigner D matrices. + + :param k: Wigner representation degree. + :param beta: One-dimensional array of polar angles in radians, with + shape (n_rotations,). + + :return: Wigner small-d matrices with shape + (n_rotations, 2 * k + 1, 2 * k + 1). """ # compute Wigner small d matrix d = np.zeros((len(beta), 2 * k + 1, 2 * k + 1), dtype=beta.dtype) From bf5738710fe5cde56d5e4c4cb98110b57c5c98a9 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 18 Sep 2026 11:04:30 -0400 Subject: [PATCH 099/161] Fix indexing and variable naming for AE and bE construction --- src/aspire/abinitio/commonline_nug.py | 56 ++++++++++++++++++--------- 1 file changed, 37 insertions(+), 19 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 5ece177bed..8a2a8d9203 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -754,33 +754,51 @@ def admm_sym_J(self, C, verbose): # blocks. These encode the representation constraints imposed by symmetry. AE = [] AEAETinv = [] - for k in range(1, Lmax + 1): - s0 = k**2 - s1 = (k + 1) ** 2 - AEk = xp.zeros((1 + s0 + s1, 2 * (s0 + s1)), dtype=np.float64) - AEk[0, :s0] = xp.eye(k, dtype=np.float64).T.reshape(-1) - AEk[0, s0 : s0 + s1] = xp.eye(k + 1, dtype=np.float64).T.reshape(-1) - for count in range(1, 1 + s0): - AEk[count, count - 1] = 1 - AEk[count, count - 1 + s0 + s1] = 1 - for count in range(1 + s0, 1 + s0 + s1): - AEk[count, count - 1] = 1 - AEk[count, count - 1 + s1 + s0] = 1 + for degree_idx in range(Lmax): + degree = degree_idx + 1 + size0 = degree**2 + size1 = (degree + 1) ** 2 + n_components = size0 + size1 + + AEk = xp.zeros((1 + n_components, 2 * n_components), dtype=np.float64) + + # Row zero contains the trace constraint for the two components. + AEk[0, :size0] = xp.eye(degree, dtype=np.float64).T.reshape(-1) + AEk[0, size0:n_components] = xp.eye(degree + 1, dtype=np.float64).T.reshape( + -1 + ) + + # The remaining rows couple every representation entry to its + # corresponding auxiliary diagonal variable. + for component_idx in range(n_components): + row_idx = component_idx + 1 + AEk[row_idx, component_idx] = 1 + AEk[row_idx, component_idx + n_components] = 1 + AE.append(AEk) AEAETinv.append(xp.linalg.pinv(AEk @ AEk.T)) # Right-hand sides for the equality constraints: symmetry projector ranks # and identity constraints on diagonal representation blocks. bE = xp.zeros((Lmax + D0 + D1), dtype=np.float64) - for k in range(Lmax): - bE[k + d0[k] + d1[k] :] = rank_Ak[k] - bE[k + 1 + d0[k] + d1[k] : k + 1 + d0[k + 1] + d1[k]] = xp.eye( - k + 1 + for degree_idx in range(Lmax): + degree = degree_idx + 1 + + start = degree_idx + d0[degree_idx] + d1[degree_idx] + split = degree_idx + 1 + d0[degree_idx + 1] + d1[degree_idx] + stop = degree_idx + 1 + d0[degree_idx + 1] + d1[degree_idx + 1] + + bE[start:] = rank_Ak[degree_idx] + bE[start + 1 : split] = xp.eye( + degree, + dtype=np.float64, ).T.reshape(-1) - bE[k + 1 + d0[k + 1] + d1[k] : k + 1 + d0[k + 1] + d1[k + 1]] = xp.eye( - k + 2 + bE_new[split:stop] = xp.eye( + degree + 1, + dtype=np.float64, ).T.reshape(-1) - bE = xp.repeat(bE[:, None], N, axis=1) + + bE_new = xp.repeat(bE_new[:, None], N, axis=1) # Degree-wise permutations used to move between full Wigner blocks and the # two invariant block variables used by the relaxation. From 2a1bdc8b026dea9a62798e4e30db2db6e0cbebb6 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 18 Sep 2026 15:32:06 -0400 Subject: [PATCH 100/161] fix more indexing. Use k for wigner degree and k_idx for python indexing --- src/aspire/abinitio/commonline_nug.py | 100 +++++++++++++++++--------- 1 file changed, 66 insertions(+), 34 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 8a2a8d9203..b02656605f 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -717,7 +717,7 @@ def admm_sym_J(self, C, verbose): Nstep_yI = self.Nstep_yI # Solve the symmetry-constrained SDP relaxation in the real, packed - # representation basis prepared by ADMM_preprocessing. + # ion basis prepared by ADMM_preprocessing. ( C0, C1, @@ -754,19 +754,17 @@ def admm_sym_J(self, C, verbose): # blocks. These encode the representation constraints imposed by symmetry. AE = [] AEAETinv = [] - for degree_idx in range(Lmax): - degree = degree_idx + 1 - size0 = degree**2 - size1 = (degree + 1) ** 2 + for k_idx in range(Lmax): + k = k_idx + 1 + size0 = k**2 + size1 = (k + 1) ** 2 n_components = size0 + size1 AEk = xp.zeros((1 + n_components, 2 * n_components), dtype=np.float64) # Row zero contains the trace constraint for the two components. - AEk[0, :size0] = xp.eye(degree, dtype=np.float64).T.reshape(-1) - AEk[0, size0:n_components] = xp.eye(degree + 1, dtype=np.float64).T.reshape( - -1 - ) + AEk[0, :size0] = xp.eye(k, dtype=np.float64).T.reshape(-1) + AEk[0, size0:n_components] = xp.eye(k + 1, dtype=np.float64).T.reshape(-1) # The remaining rows couple every representation entry to its # corresponding auxiliary diagonal variable. @@ -781,24 +779,24 @@ def admm_sym_J(self, C, verbose): # Right-hand sides for the equality constraints: symmetry projector ranks # and identity constraints on diagonal representation blocks. bE = xp.zeros((Lmax + D0 + D1), dtype=np.float64) - for degree_idx in range(Lmax): - degree = degree_idx + 1 + for k_idx in range(Lmax): + k = k_idx + 1 - start = degree_idx + d0[degree_idx] + d1[degree_idx] - split = degree_idx + 1 + d0[degree_idx + 1] + d1[degree_idx] - stop = degree_idx + 1 + d0[degree_idx + 1] + d1[degree_idx + 1] + start = k_idx + d0[k_idx] + d1[k_idx] + split = k_idx + 1 + d0[k_idx + 1] + d1[k_idx] + stop = k_idx + 1 + d0[k_idx + 1] + d1[k_idx + 1] - bE[start:] = rank_Ak[degree_idx] + bE[start:] = rank_Ak[k_idx] bE[start + 1 : split] = xp.eye( - degree, + k, dtype=np.float64, ).T.reshape(-1) - bE_new[split:stop] = xp.eye( - degree + 1, + bE[split:stop] = xp.eye( + k + 1, dtype=np.float64, ).T.reshape(-1) - bE_new = xp.repeat(bE_new[:, None], N, axis=1) + bE = xp.repeat(bE[:, None], N, axis=1) # Degree-wise permutations used to move between full Wigner blocks and the # two invariant block variables used by the relaxation. @@ -1507,9 +1505,20 @@ def _project_representation_blocks( :return: The projected packed arrays `Z0` and `Z1`. """ - for k in range(1, self.Lmax + 1): + for k_idx in range(self.Lmax): + k = k_idx + 1 + + block0_rows = slice( + d0[k_idx], + d0[k_idx + 1], + ) + block1_rows = slice( + d1[k_idx], + d1[k_idx + 1], + ) + block0 = self.mat_block( - Z0[d0[k - 1] : d0[k]], + Z0[block0_rows], self.n_img, k, IDX_upper, @@ -1517,7 +1526,7 @@ def _project_representation_blocks( idx_offdiag, ) block0 = self.psd_projection(block0) - Z0[d0[k - 1] : d0[k]] = self.vec_block( + Z0[block0_rows] = self.vec_block( block0, self.n_img, k, @@ -1525,7 +1534,7 @@ def _project_representation_blocks( ) block1 = self.mat_block( - Z1[d1[k - 1] : d1[k]], + Z1[block1_rows], self.n_img, k + 1, IDX_upper, @@ -1533,13 +1542,12 @@ def _project_representation_blocks( idx_offdiag, ) block1 = self.psd_projection(block1) - Z1[d1[k - 1] : d1[k]] = self.vec_block( + Z1[block1_rows] = self.vec_block( block1, self.n_img, k + 1, IDX_upper, ) - return Z0, Z1 @staticmethod @@ -1575,12 +1583,20 @@ def transform_coeff(self, A, IDX_upper): for k in range(1, self.Lmax + 1): d0.append(d0[-1] + k**2) d1.append(d1[-1] + (k + 1) ** 2) + A0 = xp.zeros((d0[-1], self.n_img * (self.n_img + 1) // 2), dtype=np.float64) A1 = xp.zeros((d1[-1], self.n_img * (self.n_img + 1) // 2), dtype=np.float64) - for k in range(1, self.Lmax + 1): - a0, a1 = self.permutek(A[k - 1], k, self.n_img) - A0[d0[k - 1] : d0[k], :] = self.vec_block(a0, self.n_img, k, IDX_upper) - A1[d1[k - 1] : d1[k], :] = self.vec_block(a1, self.n_img, k + 1, IDX_upper) + for k_idx in range(self.Lmax): + k = k_idx + 1 + + block0_rows = slice(d0[k_idx], d0[k_idx + 1]) + block1_rows = slice(d1[k_idx], d1[k_idx + 1]) + + block0, block1 = self.permutek(A[k_idx], k, self.n_img) + + A0[block0_rows] = self.vec_block(block0, self.n_img, k, IDX_upper) + A1[block1_rows] = self.vec_block(block1, self.n_img, k + 1, IDX_upper) + return A0, A1 def transform_coeff_back(self, A0, A1, IDX_upper, IDX_lower, idx_offdiag): @@ -1601,15 +1617,31 @@ def transform_coeff_back(self, A0, A1, IDX_upper, IDX_lower, idx_offdiag): for k in range(1, self.Lmax + 1): d0.append(d0[-1] + k**2) d1.append(d1[-1] + (k + 1) ** 2) + A = [] - for k in range(1, self.Lmax + 1): - dk = 2 * k + 1 - Ak = xp.zeros((N * dk, N * dk), dtype=np.float64) + for k_idx in range(self.Lmax): + k = k_idx + 1 + dimension = 2 * k + 1 + + block0_rows = slice(d0[k_idx], d0[k_idx + 1]) + block1_rows = slice(d1[k_idx], d1[k_idx + 1]) + + Ak = xp.zeros((N * dimension, N * dimension), dtype=np.float64) Ak[: N * k, : N * k] = self.mat_block( - A0[d0[k - 1] : d0[k], :], N, k, IDX_upper, IDX_lower, idx_offdiag + A0[block0_rows], + N, + k, + IDX_upper, + IDX_lower, + idx_offdiag, ) Ak[N * k :, N * k :] = self.mat_block( - A1[d1[k - 1] : d1[k], :], N, k + 1, IDX_upper, IDX_lower, idx_offdiag + A1[block1_rows], + N, + k + 1, + IDX_upper, + IDX_lower, + idx_offdiag, ) Ak = self.permutek_back(Ak, k, N) A.append(Ak) From 8ebe2023d610910ed4d4a01a3734d1b6897aab4f Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 21 Sep 2026 11:17:13 -0400 Subject: [PATCH 101/161] cleanup indexing fun_AE/AET --- src/aspire/abinitio/commonline_nug.py | 51 ++++++++++++++++++--------- 1 file changed, 34 insertions(+), 17 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index b02656605f..c7d65ea15b 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -717,7 +717,7 @@ def admm_sym_J(self, C, verbose): Nstep_yI = self.Nstep_yI # Solve the symmetry-constrained SDP relaxation in the real, packed - # ion basis prepared by ADMM_preprocessing. + # representation basis prepared by ADMM_preprocessing. ( C0, C1, @@ -822,15 +822,20 @@ def fun_AE(X0, X1, Xd0, Xd1, Xq): the lowest-degree off-diagonal representation blocks. """ z = xp.zeros((Lmax + D0 + D1, N), dtype=np.float64) - for k in range(Lmax): - z[k + d0[k] + d1[k] : k + 1 + d0[k + 1] + d1[k + 1]] = AE[ - k - ] @ xp.concatenate( + for k_idx in range(Lmax): + block0_rows = slice(d0[k_idx], d0[k_idx + 1]) + block1_rows = slice(d1[k_idx], d1[k_idx + 1]) + + constraint_start = k_idx + d0[k_idx] + d1[k_idx] + constraint_stop = k_idx + 1 + d0[k_idx + 1] + d1[k_idx + 1] + constraint_rows = slice(constraint_start, constraint_stop) + + z[constraint_rows] = AE[k_idx] @ xp.concatenate( ( - X0[d0[k] : d0[k + 1], idx_diag], - X1[d1[k] : d1[k + 1], idx_diag], - Xd0[d0[k] : d0[k + 1]], - Xd1[d1[k] : d1[k + 1]], + X0[block0_rows, idx_diag], + X1[block1_rows, idx_diag], + Xd0[block0_rows], + Xd1[block1_rows], ), axis=0, ) @@ -856,14 +861,26 @@ def fun_AET(yE, yEq): Z1 = xp.zeros((D1, N * (N + 1) // 2), dtype=np.float64) Zd0 = xp.zeros((D0, N), dtype=np.float64) Zd1 = xp.zeros((D1, N), dtype=np.float64) - for k in range(Lmax): - s0 = (k + 1) ** 2 - s1 = (k + 2) ** 2 - Ztmp = AE[k].T @ yE[k + d0[k] + d1[k] : k + 1 + d0[k + 1] + d1[k + 1]] - Z0[d0[k] : d0[k + 1], idx_diag] = Ztmp[:s0] - Z1[d1[k] : d1[k + 1], idx_diag] = Ztmp[s0 : s0 + s1] - Zd0[d0[k] : d0[k + 1]] = Ztmp[s0 + s1 : 2 * s0 + s1] - Zd1[d1[k] : d1[k + 1]] = Ztmp[2 * s0 + s1 : 2 * s0 + 2 * s1] + + for k_idx in range(Lmax): + k = k_idx + 1 # degree + size0 = k**2 + size1 = (k + 1) ** 2 + + block0_rows = slice(d0[k_idx], d0[k_idx + 1]) + block1_rows = slice(d1[k_idx], d1[k_idx + 1]) + + constraint_start = k_idx + d0[k_idx] + d1[k_idx] + constraint_stop = k_idx + 1 + d0[k_idx + 1] + d1[k_idx + 1] + constraint_rows = slice(constraint_start, constraint_stop) + + Ztmp = AE[k_idx].T @ yE[constraint_rows] + + Z0[block0_rows, idx_diag] = Ztmp[:size0] + Z1[block1_rows, idx_diag] = Ztmp[size0 : size0 + size1] + Zd0[block0_rows] = Ztmp[size0 + size1 : 2 * size0 + size1] + Zd1[block1_rows] = Ztmp[2 * size0 + size1 : 2 * size0 + 2 * size1] + Zq = AEq.T @ yEq Z0[:1, idx_offdiag] = Zq[16:17] Z1[:4, idx_offdiag] = Zq[17:] From 6e7ce0f3d3d1ab0eb9ea77c39096abfba37d8a69 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 21 Sep 2026 14:00:56 -0400 Subject: [PATCH 102/161] cleanup indexing update_S --- src/aspire/abinitio/commonline_nug.py | 21 ++++++++++++++------- 1 file changed, 14 insertions(+), 7 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index c7d65ea15b..73ba509cb0 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -156,7 +156,7 @@ def __init__( self.pr_penalty = [1] * pr_iters self.pr_rank = list(range(pr_iters - 1, -1, -1)) # [pr_iters - 1,..., 0] self.pr_iters = pr_iters - + self.count_S_iter = 0 self._build_full_pft() def _build_full_pft(self): @@ -946,17 +946,24 @@ def update_S(C0, C1, yE, yEq, yI, X0, X1, Xd0, Xd1, Xq, rho, Lmax, N): # Project the diagonal coupling blocks. Sd0 = Sd0.T Sd1 = Sd1.T - for k in range(1, Lmax + 1): + for k_idx in range(Lmax): + k = k_idx + 1 # degreee + block0_rows = slice(d0[k_idx], d0[k_idx + 1]) + block1_rows = slice(d1[k_idx], d1[k_idx + 1]) + permutation = P[k_idx] + tmp = self.transform_back_block( - Sd0[:, d0[k - 1] : d0[k]], - Sd1[:, d1[k - 1] : d1[k]], + Sd0[:, block0_rows], + Sd1[:, block1_rows], k, - P[k - 1], + permutation, ) tmp = self.psd_projection(tmp) - Sd0[:, d0[k - 1] : d0[k]], Sd1[:, d1[k - 1] : d1[k]] = ( - self.transform_block(tmp, k, P[k - 1]) + + Sd0[:, block0_rows], Sd1[:, block1_rows] = self.transform_block( + tmp, k, permutation ) + Sd0 = Sd0.T Sd1 = Sd1.T From 155640f1c6fcc1718dc04bf69bdd738357386eaa Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 21 Sep 2026 14:13:39 -0400 Subject: [PATCH 103/161] cleanup indexing update_yE --- src/aspire/abinitio/commonline_nug.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 73ba509cb0..f7ea61c7d1 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -988,10 +988,13 @@ def update_yE(C0, C1, X0, X1, Xd0, Xd1, Xq, S0, S1, Sd0, Sd1, Sq, yI, rho): -Xq / rho - Sq, ) yE = xp.zeros((Lmax + D0 + D1, N), dtype=np.float64) - for k in range(Lmax): - yE[k + d0[k] + d1[k] : k + 1 + d0[k + 1] + d1[k + 1]] = AEAETinv[k] @ ( - bE[k + d0[k] + d1[k] : k + 1 + d0[k + 1] + d1[k + 1]] / rho - + z[k + d0[k] + d1[k] : k + 1 + d0[k + 1] + d1[k + 1]] + for k_idx in range(Lmax): + constraint_start = k_idx + d0[k_idx] + d1[k_idx] + constraint_stop = k_idx + 1 + d0[k_idx + 1] + d1[k_idx + 1] + constraint_rows = slice(constraint_start, constraint_stop) + + yE[constraint_rows] = AEAETinv[k_idx] @ ( + bE[constraint_rows] / rho + z[constraint_rows] ) yEq = AEqAEqtinv @ (bEq / rho + zq) return yE, yEq From fd962adafe024cc349dd77e87f1602155af5c5d4 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 21 Sep 2026 14:30:15 -0400 Subject: [PATCH 104/161] cleanup indexing in print_updates --- src/aspire/abinitio/commonline_nug.py | 30 +++++++++++++++++++-------- 1 file changed, 21 insertions(+), 9 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index f7ea61c7d1..8b15d08bd5 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1094,10 +1094,15 @@ def print_updates(verbose): res_inq = xp.linalg.norm(xp.maximum(bI - fun_AI(X0, X1), 0)) / ( 1 + abs(bI) * xp.sqrt(Ngrid * N * (N + 1) / 2) ) + res_psdX = 0 - for k in range(1, Lmax + 1): + for k_idx in range(Lmax): + k = k_idx + 1 # degree + block0_rows = slice(d0[k_idx], d0[k_idx + 1]) + block1_rows = slice(d1[k_idx], d1[k_idx + 1]) + tmp = self.mat_block( - X0[d0[k - 1] : d0[k], :], + X0[block0_rows], N, k, IDX_upper, @@ -1105,8 +1110,9 @@ def print_updates(verbose): idx_offdiag, ) res_psdX += xp.linalg.norm(self.psd_projection(-tmp)) + tmp = self.mat_block( - X1[d1[k - 1] : d1[k], :], + X1[block1_rows], N, k + 1, IDX_upper, @@ -1115,21 +1121,27 @@ def print_updates(verbose): ) res_psdX += xp.linalg.norm(self.psd_projection(-tmp)) res_psdX = res_psdX / (1 + xp.linalg.norm(X0) + xp.linalg.norm(X1)) + res_psdD = 0 - for k in range(1, Lmax + 1): + for k_idx in range(Lmax): + k = k_idx + 1 # degree + block0_rows = slice(d0[k_idx], d0[k_idx + 1]) + block1_rows = slice(d1[k_idx], d1[k_idx + 1]) + tmp = self.transform_back_block( - Xd0[d0[k - 1] : d0[k]].T, - Xd1[d1[k - 1] : d1[k]].T, + Xd0[block0_rows].T, + Xd1[block1_rows].T, k, - P[k - 1], + P[k_idx], ) res_psdD += xp.linalg.norm( self.psd_projection(-tmp), axis=(-2, -1) ).sum() res_psdD = res_psdD / (1 + xp.linalg.norm(Xd0) + xp.linalg.norm(Xd1)) + res_psdQ = 0 - for count in range(N * (N - 1) // 2): - tmp = Xq[:, count].reshape(4, 4).T + for pair in range(n_pairs): + tmp = Xq[:, pair].reshape(4, 4).T res_psdQ += xp.linalg.norm(self.psd_projection(-tmp)) res_psdQ = res_psdQ / (1 + xp.linalg.norm(Xq)) From 885c6de8452a0fb8804f290de080d8d0dccc35e3 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 21 Sep 2026 14:36:26 -0400 Subject: [PATCH 105/161] typos, cleanup --- src/aspire/abinitio/commonline_nug.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 8b15d08bd5..ff33cf6764 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -156,7 +156,7 @@ def __init__( self.pr_penalty = [1] * pr_iters self.pr_rank = list(range(pr_iters - 1, -1, -1)) # [pr_iters - 1,..., 0] self.pr_iters = pr_iters - self.count_S_iter = 0 + self._build_full_pft() def _build_full_pft(self): @@ -947,7 +947,7 @@ def update_S(C0, C1, yE, yEq, yI, X0, X1, Xd0, Xd1, Xq, rho, Lmax, N): Sd0 = Sd0.T Sd1 = Sd1.T for k_idx in range(Lmax): - k = k_idx + 1 # degreee + k = k_idx + 1 # degree block0_rows = slice(d0[k_idx], d0[k_idx + 1]) block1_rows = slice(d1[k_idx], d1[k_idx + 1]) permutation = P[k_idx] @@ -1140,8 +1140,8 @@ def print_updates(verbose): res_psdD = res_psdD / (1 + xp.linalg.norm(Xd0) + xp.linalg.norm(Xd1)) res_psdQ = 0 - for pair in range(n_pairs): - tmp = Xq[:, pair].reshape(4, 4).T + for pair_idx in range(n_pairs): + tmp = Xq[:, pair_idx].reshape(4, 4).T res_psdQ += xp.linalg.norm(self.psd_projection(-tmp)) res_psdQ = res_psdQ / (1 + xp.linalg.norm(Xq)) @@ -1602,7 +1602,7 @@ def psd_projection(B): :return: Positive semidefinite projection with same shape as B. """ - # compute the PSD part of a symmstric matrix + # compute the PSD part of a symmetric matrix B_sym = (B + B.swapaxes(-1, -2)) / 2 evals, evecs = xp.linalg.eigh(B_sym) evals = xp.maximum(evals, 0) From 31e9380c06dce5d8793e9c8d7e40321a97acac89 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 21 Sep 2026 15:30:25 -0400 Subject: [PATCH 106/161] cleanup indexing asymmetric admm --- src/aspire/abinitio/commonline_nug.py | 48 +++++++++++++++------------ 1 file changed, 26 insertions(+), 22 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index ff33cf6764..4073db13aa 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -391,10 +391,13 @@ def admm_plain_J(self, C, verbose): # In the asymmetric formulation, diagonal representation blocks are identity. bE0 = xp.zeros(D0, dtype=np.float64) bE1 = xp.zeros(D1, dtype=np.float64) + for k_idx in range(Lmax): + k = k_idx + 1 + block0_rows = slice(d0[k_idx], d0[k_idx + 1]) + block1_rows = slice(d1[k_idx], d1[k_idx + 1]) - for k in range(1, Lmax + 1): - bE0[d0[k - 1] : d0[k]] = xp.eye(k, dtype=np.float64).reshape(-1) - bE1[d1[k - 1] : d1[k]] = xp.eye(k + 1, dtype=np.float64).reshape(-1) + bE0[block0_rows] = xp.eye(k, dtype=np.float64).reshape(-1) + bE1[block1_rows] = xp.eye(k + 1, dtype=np.float64).reshape(-1) bE0 = xp.repeat(bE0[:, None], N, axis=1) bE1 = xp.repeat(bE1[:, None], N, axis=1) @@ -600,9 +603,13 @@ def print_updates(verbose): ) res_psdX = 0 - for k in range(1, Lmax + 1): + for k_idx in range(Lmax): + k = k_idx + 1 + block0_rows = slice(d0[k_idx], d0[k_idx + 1]) + block1_rows = slice(d1[k_idx], d1[k_idx + 1]) + tmp = self.mat_block( - X0[d0[k - 1] : d0[k]], + X0[block0_rows], N, k, IDX_upper, @@ -612,7 +619,7 @@ def print_updates(verbose): res_psdX += xp.linalg.norm(self.psd_projection(-tmp)) tmp = self.mat_block( - X1[d1[k - 1] : d1[k]], + X1[block1_rows], N, k + 1, IDX_upper, @@ -620,7 +627,6 @@ def print_updates(verbose): idx_offdiag, ) res_psdX += xp.linalg.norm(self.psd_projection(-tmp)) - res_psdX /= 1 + xp.linalg.norm(X0) + xp.linalg.norm(X1) Xq_blocks = Xq.T.reshape(n_pairs, 4, 4) @@ -661,13 +667,13 @@ def print_updates(verbose): yE1 = xp.zeros(bE1.shape, dtype=np.float64) yEq = xp.zeros(bEq.shape, dtype=np.float64) - IDX = np.arange(3) + update_order = np.arange(3) for t in range(max_iter): - np.random.shuffle(IDX) - for idx in IDX: - if idx == 0: + np.random.shuffle(update_order) + for update_idx in update_order: + if update_idx == 0: S0, S1, Sq = update_S(C0, C1, yE0, yE1, yEq, yI, X0, X1, Xq, rho) - elif idx == 1: + elif update_idx == 1: yE0, yE1, yEq = update_yE(C0, C1, S0, S1, Sq, yI, X0, X1, Xq, rho) else: for _ in range(Nstep_yI): @@ -694,8 +700,7 @@ def print_updates(verbose): rho, p_resnorm, d_resnorm = update_rho(X0, X1, Xq, res_X, rho) X_admm = self.transform_coeff_back(X0, X1, IDX_upper, IDX_lower, idx_offdiag) - for k in range(Lmax): - X_admm[k] = xp.asnumpy(X_admm[k]) + X_admm = [xp.asnumpy(Xk) for Xk in X_admm] return X_admm def admm_sym_J(self, C, verbose): @@ -1185,19 +1190,19 @@ def print_updates(verbose): # Run ADMM iterations, randomly ordering the block updates before each primal # multiplier update and penalty adjustment. - IDX = np.arange(3) + update_order = np.arange(3) for t in range(max_iter): - np.random.shuffle(IDX) - for idx in IDX: - if idx == 0: + np.random.shuffle(update_order) + for update_idx in update_order: + if update_idx == 0: S0, S1, Sd0, Sd1, Sq = update_S( C0, C1, yE, yEq, yI, X0, X1, Xd0, Xd1, Xq, rho, Lmax, N ) - if idx == 1: + if update_idx == 1: yE, yEq = update_yE( C0, C1, X0, X1, Xd0, Xd1, Xq, S0, S1, Sd0, Sd1, Sq, yI, rho ) - if idx == 2: + if update_idx == 2: for _ in range(Nstep_yI): yI = update_yI(C0, C1, X0, X1, S0, S1, yE, yEq, yI, rho, Lambda) X0, X1, Xd0, Xd1, Xq, res_X = update_X( @@ -1212,8 +1217,7 @@ def print_updates(verbose): # Convert the optimized packed block variables back to full degree-wise # representation matrices for Euler-angle recovery. X_admm = self.transform_coeff_back(X0, X1, IDX_upper, IDX_lower, idx_offdiag) - for k in range(Lmax): - X_admm[k] = xp.asnumpy(X_admm[k]) + X_admm = [xp.asnumpy(Xk) for Xk in X_admm] return X_admm ################ From 9e02ccf0423730e7bac196ae4d0e61624200e539 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 29 Sep 2026 09:53:47 -0400 Subject: [PATCH 107/161] Fix indexing bug in Xnorm. Expose missing volumes in downloader init (used for testing Xnorm bug). --- src/aspire/abinitio/commonline_nug.py | 13 +++++++++---- src/aspire/downloader/__init__.py | 2 ++ 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 4073db13aa..f9a50fccbb 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1251,14 +1251,19 @@ def ADMM_preprocessing(self, C): # normalize C matrix Cnorm = 0 Xnorm = 0 - for k in range(Lmax): + for k_idx in range(Lmax): + k = k_idx + 1 # degree dk = 2 * k + 1 - Cnorm += np.linalg.norm(C[k]) ** 2 + + Cnorm += np.linalg.norm(C[k_idx]) ** 2 Xnorm += dk * N**2 + Cnorm = np.sqrt(Cnorm) Xnorm = np.sqrt(Xnorm) - for k in range(Lmax): - C[k] = xp.asarray(Xnorm / Cnorm * C[k]) + + for k_idx in range(Lmax): + C[k_idx] = xp.asarray(Xnorm / Cnorm * C[k_idx]) + C0, C1 = self.transform_coeff(C, IDX_upper) normC = xp.sqrt(xp.linalg.norm(C0) ** 2 + xp.linalg.norm(C1) ** 2) del C diff --git a/src/aspire/downloader/__init__.py b/src/aspire/downloader/__init__.py index 9151e7a2d1..d95667225b 100644 --- a/src/aspire/downloader/__init__.py +++ b/src/aspire/downloader/__init__.py @@ -18,8 +18,10 @@ emdb_8511, emdb_10835, emdb_14621, + emdb_22308, nug_10081, nug_12036, + emdb_14803, remove_downloads, simulated_channelspin, ) From 22ac4cd9f2dd99cefff45ec2edb67a1008ed913f Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 29 Sep 2026 10:36:49 -0400 Subject: [PATCH 108/161] Indexing for AI_mat_(off)diag computation loops --- src/aspire/abinitio/commonline_nug.py | 26 ++++++++++++++++---------- src/aspire/downloader/__init__.py | 2 +- 2 files changed, 17 insertions(+), 11 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index f9a50fccbb..446764aa9d 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1292,29 +1292,35 @@ def ADMM_preprocessing(self, C): # AI_mat_offdiag computation has been vectorized, but originally # had the note by block1: this needs double checking (Ruiyi) AI_mat_offdiag = np.zeros((Ngrid, D0 + D1), dtype=np.float64) - for k in range(1, Lmax + 1): + for k_idx in range(Lmax): + k = k_idx + 1 scale = (Lmax - k + 2) * (Lmax - k + 1) * (k + 0.5) + block0_cols = slice(d0[k_idx], d0[k_idx + 1]) + block1_cols = slice(D0 + d1[k_idx], D0 + d1[k_idx + 1]) - block0 = scale * W0[k - 1].transpose(0, 2, 1).reshape(Ngrid, -1) - block1 = scale * W1[k - 1].transpose(0, 2, 1).reshape(Ngrid, -1) + block0 = scale * W0[k_idx].transpose(0, 2, 1).reshape(Ngrid, -1) + block1 = scale * W1[k_idx].transpose(0, 2, 1).reshape(Ngrid, -1) - AI_mat_offdiag[:, d0[k - 1] : d0[k]] = block0 - AI_mat_offdiag[:, d0[-1] + d1[k - 1] : d0[-1] + d1[k]] = block1 + AI_mat_offdiag[:, block0_cols] = block0 + AI_mat_offdiag[:, block1_cols] = block1 # AI_mat_diag computation has been vectorized, but originally # had the note by block1: this needs double checking (Ruiyi) AI_mat_diag = np.zeros((Ngrid, D0 + D1), dtype=np.float64) - for k in range(1, Lmax + 1): + for k_idx in range(Lmax): + k = k_idx + 1 scale = (Lmax - k + 2) * (Lmax - k + 1) * (k + 0.5) + block0_cols = slice(d0[k_idx], d0[k_idx + 1]) + block1_cols = slice(D0 + d1[k_idx], D0 + d1[k_idx + 1]) - W0_sym = 0.5 * (W0[k - 1] + W0[k - 1].transpose(0, 2, 1)) - W1_sym = 0.5 * (W1[k - 1] + W1[k - 1].transpose(0, 2, 1)) + W0_sym = 0.5 * (W0[k_idx] + W0[k_idx].transpose(0, 2, 1)) + W1_sym = 0.5 * (W1[k_idx] + W1[k_idx].transpose(0, 2, 1)) block0 = scale * W0_sym.transpose(0, 2, 1).reshape(Ngrid, -1) block1 = scale * W1_sym.transpose(0, 2, 1).reshape(Ngrid, -1) - AI_mat_diag[:, d0[k - 1] : d0[k]] = block0 - AI_mat_diag[:, d0[-1] + d1[k - 1] : d0[-1] + d1[k]] = block1 + AI_mat_diag[:, block0_cols] = block0 + AI_mat_diag[:, block1_cols] = block1 AI_mat_diag = xp.asarray(AI_mat_diag) AI_mat_offdiag = xp.asarray(AI_mat_offdiag) diff --git a/src/aspire/downloader/__init__.py b/src/aspire/downloader/__init__.py index d95667225b..5bc2818362 100644 --- a/src/aspire/downloader/__init__.py +++ b/src/aspire/downloader/__init__.py @@ -18,10 +18,10 @@ emdb_8511, emdb_10835, emdb_14621, + emdb_14803, emdb_22308, nug_10081, nug_12036, - emdb_14803, remove_downloads, simulated_channelspin, ) From 9098a82b783a2c751d399744572fdf1bf84077cc Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 29 Sep 2026 10:56:32 -0400 Subject: [PATCH 109/161] Indexing for compute_rank --- src/aspire/abinitio/commonline_nug.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 446764aa9d..faec9fdb13 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1480,11 +1480,12 @@ def compute_rank(self): """ rk = xp.zeros(self.Lmax, dtype=np.float64) A = [] - for k in range(1, self.Lmax + 1): + for k_idx in range(self.Lmax): + k = k_idx + 1 # degree Ak = np.sum(self.WD(k, self.sym_euler), axis=0) Ak = np.round(Ak / self.n_sym, 6) A.append(Ak) - rk[k - 1] = np.linalg.matrix_rank(Ak) + rk[k_idx] = np.linalg.matrix_rank(Ak) return rk, A def construct_AEq(self): From c95a9a301aac07e65253a7b1f74cbd3ad2d343eb Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 29 Sep 2026 11:29:59 -0400 Subject: [PATCH 110/161] update indexing in compute_coeff --- src/aspire/abinitio/commonline_nug.py | 32 +++++++++++++++------------ 1 file changed, 18 insertions(+), 14 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index faec9fdb13..5c7e737990 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -245,10 +245,11 @@ def fij(alpha, gamma, i, j): btk = np.sum(bT[:, None, None] * self.Wd(k, beta_grid), axis=0) BTK.append(btk.T) - def fijhat_k(k, F): + def fijhat_k(k_idx, F): """ Approximate the degree-k Fourier coefficient block of a sampled pairwise loss. """ + k = k_idx + 1 # degree dk = 2 * k + 1 exp_alpha_grid = np.zeros((2 * T, dk), dtype=complex_type(np.float64)) @@ -260,7 +261,7 @@ def fijhat_k(k, F): exp_gamma_grid[:, m + k] = np.exp(1j * m * gamma_grid) S = (exp_alpha_grid.T @ F @ exp_gamma_grid).T - fhat = BTK[k - 1] * S / 4 / T**2 + fhat = BTK[k_idx] * S / 4 / T**2 return fhat # Allocate one coefficient matrix per Wigner degree; @@ -277,15 +278,16 @@ def fijhat_k(k, F): for j1 in range(2 * T): for j2 in range(2 * T): Fij[j1, j2] = fij(alpha_grid[j1], gamma_grid[j2], i, j) - for k in range(1, Lmax + 1): + for k_idx in range(Lmax): + k = k_idx + 1 dk = 2 * k + 1 - C[k - 1][j * dk : (j + 1) * dk, i * dk : (i + 1) * dk] = fijhat_k( - k, Fij + C[k_idx][j * dk : (j + 1) * dk, i * dk : (i + 1) * dk] = fijhat_k( + k_idx, Fij ) # *dk # Fill the conjugate transpose blocks so each coefficient matrix is Hermitian. - for k in range(1, Lmax + 1): - C[k - 1] = C[k - 1] + C[k - 1].conj().T + for k_idx in range(Lmax): + C[k_idx] = C[k_idx] + C[k_idx].conj().T # Diagonal blocks encode self-pair losses and are handled separately. for i in range(N): @@ -293,18 +295,20 @@ def fijhat_k(k, F): for j1 in range(2 * T): for j2 in range(2 * T): Fii[j1, j2] = fij(alpha_grid[j1], gamma_grid[j2], i, i) - for k in range(1, Lmax + 1): + for k_idx in range(Lmax): + k = k_idx + 1 dk = 2 * k + 1 - C[k - 1][i * dk : (i + 1) * dk, i * dk : (i + 1) * dk] = fijhat_k( - k, Fii + C[k_idx][i * dk : (i + 1) * dk, i * dk : (i + 1) * dk] = fijhat_k( + k_idx, Fii ) # *dk # Convert complex Wigner coefficients to the real representation basis used by ADMM. - for k in range(1, Lmax + 1): - [T, Tinv] = self.complex2real(k) - C[k - 1] = np.real( + for k_idx in range(Lmax): + k = k_idx + 1 + T, Tinv = self.complex2real(k) + C[k_idx] = np.real( np.kron(np.eye(N, dtype=np.float64), Tinv) - @ C[k - 1] + @ C[k_idx] @ np.kron(np.eye(N, dtype=np.float64), T) ) From 596c8740321b336e313309abadccc8c4744be682 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 29 Sep 2026 11:47:02 -0400 Subject: [PATCH 111/161] Indexing in compute_fejer_weights --- src/aspire/abinitio/commonline_nug.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 5c7e737990..5dda6589c7 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1383,12 +1383,12 @@ def compute_fejer_weights(self): """ SO3_grid = self.discretize_SO3() Ngrid = SO3_grid.shape[0] - start = 1 TT = [] TTI = [] - for ell in range(start, self.Lmax + 1): - T, Tinv = self.complex2real(ell) + for k_idx in range(self.Lmax): + k = k_idx + 1 + T, Tinv = self.complex2real(k) TT.append(T) TTI.append(Tinv) @@ -1405,12 +1405,13 @@ def permutek_block(Ak, k): W0 = [] W1 = [] - for k in range(start, self.Lmax + 1): + for k_idx in range(self.Lmax): + k = k_idx + 1 W0k = np.zeros((Ngrid, k, k), dtype=np.float64) W1k = np.zeros((Ngrid, k + 1, k + 1), dtype=np.float64) - TkT = TT[k - start].T - TinvkT = TTI[k - start].T + TkT = TT[k_idx].T + TinvkT = TTI[k_idx].T w = np.real(TkT @ self.WD(k, SO3_grid).conj() @ TinvkT) W0k, W1k = permutek_block(w, k) From 443112f1ddba14d3c7223896fd5bda9351c23256 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 29 Sep 2026 11:55:43 -0400 Subject: [PATCH 112/161] indexing in proximal_refine --- src/aspire/abinitio/commonline_nug.py | 29 ++++++++++++++++----------- 1 file changed, 17 insertions(+), 12 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 5dda6589c7..15197a5912 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1912,10 +1912,11 @@ def rel_change(A, B, eps=1e-12): # Keep an unmodified copy of the original objective coefficients; each # refinement step adds a temporary proximal term to these. C_base = [None] * self.Lmax - for k in range(self.Lmax): - C_base[k] = xp.asnumpy(C[k]).copy() - rank_Ak[k] = np.linalg.matrix_rank( - Ak(k + 1, self.sym_euler), tol=1e-6, hermitian=True + for k_idx in range(self.Lmax): + k = k_idx + 1 + C_base[k_idx] = xp.asnumpy(C[k_idx]).copy() + rank_Ak[k_idx] = np.linalg.matrix_rank( + Ak(k, self.sym_euler), tol=1e-6, hermitian=True ) def low_rank_proj(X, r_step): @@ -1923,10 +1924,11 @@ def low_rank_proj(X, r_step): Project diagonal image blocks toward the rank expected from symmetry. """ Xproj = [] - for k in range(self.Lmax): - dk = 2 * k + 3 - rk = min(int(rank_Ak[k] * 2) + r_step, dk) - tmp = np.copy(X[k]) + for k_idx in range(self.Lmax): + k = k_idx + 1 + dk = 2 * k + 1 + rk = min(int(rank_Ak[k_idx] * 2) + r_step, dk) + tmp = np.copy(X[k_idx]) for i in range(N): u, s, v = np.linalg.svd( tmp[i * dk : (i + 1) * dk, i * dk : (i + 1) * dk] @@ -1946,10 +1948,13 @@ def low_rank_proj(X, r_step): for step in range(self.pr_iters): X_proj = low_rank_proj(current, r[step]) - for k in range(self.Lmax): - CC[k] = ( - C_base[k] - - Penalty[step] * weight[k] * (X_proj[k] + X_proj[k].T) / 2 + for k_idx in range(self.Lmax): + CC[k_idx] = ( + C_base[k_idx] + - Penalty[step] + * weight[k_idx] + * (X_proj[k_idx] + X_proj[k_idx].T) + / 2 ) X_next = self.admm_sym_J(CC, verbose=False) From fe8eac04d23d6496e94c0110a6248ebff1dcaf10 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 1 Oct 2026 10:55:38 -0400 Subject: [PATCH 113/161] Roundtrip test for mat_block/vec_block --- tests/test_nug.py | 34 ++++++++++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/tests/test_nug.py b/tests/test_nug.py index 2149b0edeb..fe2716c028 100644 --- a/tests/test_nug.py +++ b/tests/test_nug.py @@ -3,6 +3,7 @@ from aspire.abinitio import CommonlineNUG, compare_rots_sym, g_sync from aspire.downloader import emdb_2660 +from aspire.numeric import xp from aspire.source import Simulation from aspire.utils import mean_aligned_angular_distance from aspire.volume import CnSymmetricVolume, DnSymmetricVolume, TSymmetricVolume @@ -183,6 +184,39 @@ def test_estimate_rotations(orient_est): mean_aligned_angular_distance(orient_est.rotations, gt_rots, 10.0) +@pytest.mark.parametrize("block_size", [2, 3, 12, 13]) +def test_mat_block_vec_block_roundtrip(block_size): + """ + Packing and unpacking preserve nonsymmetric image-pair blocks. + """ + # Build random symmetric image pair block matrix + n_img = 5 + rng = np.random.default_rng(SEED) + rand_matrix = rng.normal(size=(n_img * block_size, n_img * block_size)) + sym_mat = rand_matrix + rand_matrix.T + + # The full matrix is symmetric, but an off-diagonal image-pair + # block is not generally symmetric. Check at least one block is + # not symmetric as a safegaurd to expose mistaken transposes. + pair_block = sym_mat[:block_size, block_size : 2 * block_size] + assert not np.allclose(pair_block, pair_block.T) + + # Construct indices as in ADMM_preprocessing + i_upper, j_upper = np.triu_indices(n_img) + idx_upper = xp.asarray(i_upper * n_img + j_upper) + idx_offdiag = xp.asarray(np.flatnonzero(i_upper != j_upper)) + + i_off, j_off = np.triu_indices(n_img, k=1) + idx_lower = xp.asarray(j_off * n_img + i_off) + + # Perform round trip and check mat_block(vec_block(A)) == A + packed = CommonlineNUG.vec_block(xp.asarray(sym_mat), n_img, block_size, idx_upper) + restored = CommonlineNUG.mat_block( + packed, n_img, block_size, idx_upper, idx_lower, idx_offdiag + ) + np.testing.assert_allclose(xp.asnumpy(restored), sym_mat) + + def test_unspupported_symmetry_raises(dtype): """ Check that we raise for symmetries other than Cn/Dn. From 5c3cf4f2ad0847eb20f4c71f91265ee32527e758 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 1 Oct 2026 12:16:47 -0400 Subject: [PATCH 114/161] Add transform_block roundtrip test. --- tests/test_nug.py | 33 ++++++++++++++++++++++++++++++++- 1 file changed, 32 insertions(+), 1 deletion(-) diff --git a/tests/test_nug.py b/tests/test_nug.py index fe2716c028..04a37271fc 100644 --- a/tests/test_nug.py +++ b/tests/test_nug.py @@ -187,7 +187,7 @@ def test_estimate_rotations(orient_est): @pytest.mark.parametrize("block_size", [2, 3, 12, 13]) def test_mat_block_vec_block_roundtrip(block_size): """ - Packing and unpacking preserve nonsymmetric image-pair blocks. + Check that mat_block(vec_block(A)) == A, for symmetric block matrix A. """ # Build random symmetric image pair block matrix n_img = 5 @@ -217,6 +217,37 @@ def test_mat_block_vec_block_roundtrip(block_size): np.testing.assert_allclose(xp.asnumpy(restored), sym_mat) +def test_transform_block_round_trip(): + """ + Reconstruct degree-k matrices from their two packed components. + """ + k = 2 + dk = 2 * k + 1 + rng = np.random.default_rng(SEED) + + # In the permuted basis, a degree-k matrix has components of sizes + # k-by-k and (k+1)-by-(k+1). Give both components arbitrary entries + # so a row/column packing mismatch is visible. + components = np.zeros((2, dk, dk)) + components[:, :k, :k] = rng.normal(size=(2, k, k)) + components[:, k:, k:] = rng.normal(size=(2, k + 1, k + 1)) + + # Keep the cross-component blocks zero: transform_block intentionally + # discards them, so they cannot be recovered by a round trip. + + # Reproduce NUG's permutation, which groups the odd-indexed coordinates + # before the even-indexed coordinates. + row_order = np.r_[np.arange(1, dk, 2), np.arange(0, dk, 2)] + permutation = xp.asarray(np.eye(dk)[row_order]) + matrices = permutation.T @ xp.asarray(components) @ permutation + + # Pack the components as done in ADMM, then reconstruct the full matrices. + packed0, packed1 = CommonlineNUG.transform_block(matrices, k, permutation) + restored = CommonlineNUG.transform_back_block(packed0, packed1, k, permutation) + + np.testing.assert_allclose(xp.asnumpy(restored), xp.asnumpy(matrices)) + + def test_unspupported_symmetry_raises(dtype): """ Check that we raise for symmetries other than Cn/Dn. From 4d9717617e4ce21518a08e943dac3fc462465a69 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 1 Oct 2026 15:07:57 -0400 Subject: [PATCH 115/161] Row-wise ADMM block packing --- src/aspire/abinitio/commonline_nug.py | 25 ++++++++++++++----------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 15197a5912..4997af33ec 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1302,8 +1302,8 @@ def ADMM_preprocessing(self, C): block0_cols = slice(d0[k_idx], d0[k_idx + 1]) block1_cols = slice(D0 + d1[k_idx], D0 + d1[k_idx + 1]) - block0 = scale * W0[k_idx].transpose(0, 2, 1).reshape(Ngrid, -1) - block1 = scale * W1[k_idx].transpose(0, 2, 1).reshape(Ngrid, -1) + block0 = scale * W0[k_idx].reshape(Ngrid, -1) + block1 = scale * W1[k_idx].reshape(Ngrid, -1) AI_mat_offdiag[:, block0_cols] = block0 AI_mat_offdiag[:, block1_cols] = block1 @@ -1320,8 +1320,8 @@ def ADMM_preprocessing(self, C): W0_sym = 0.5 * (W0[k_idx] + W0[k_idx].transpose(0, 2, 1)) W1_sym = 0.5 * (W1[k_idx] + W1[k_idx].transpose(0, 2, 1)) - block0 = scale * W0_sym.transpose(0, 2, 1).reshape(Ngrid, -1) - block1 = scale * W1_sym.transpose(0, 2, 1).reshape(Ngrid, -1) + block0 = scale * W0_sym.reshape(Ngrid, -1) + block1 = scale * W1_sym.reshape(Ngrid, -1) AI_mat_diag[:, block0_cols] = block0 AI_mat_diag[:, block1_cols] = block1 @@ -1534,7 +1534,10 @@ def construct_AEq(self): dtype=np.float64, ) - AEq[:16, 16:] = extra + # `extra` uses the original column-wise order + # [X0_00, X1_00, X1_10, X1_01, X1_11]. + # Swap the two off-diagonal columns to match row-wise packing of X1. + AEq[:16, 16:] = extra[:, [0, 1, 3, 2, 4]] # Last row: redundant trace/sum constraint AEq[16, [0, 5, 10, 15]] = 1 @@ -1790,7 +1793,7 @@ def vec_block(A, N, sz, IDX_upper): :return: Packed block array with shape (sz**2, N * (N + 1) // 2). """ - vecA = (A.reshape(N, sz, N, sz).transpose(0, 2, 3, 1)).reshape(N**2, sz**2).T + vecA = (A.reshape(N, sz, N, sz).transpose(0, 2, 1, 3)).reshape(N**2, sz**2).T return vecA[:, IDX_upper] @staticmethod @@ -1816,7 +1819,7 @@ def mat_block(vecA, N, sz, IDX_upper, IDX_lower, idx_offdiag): :return: Symmetric block matrix with shape (N * sz, N * sz). """ - tmp = vecA.T.reshape(N * (N + 1) // 2, sz, sz).transpose(0, 2, 1) + tmp = vecA.T.reshape(N * (N + 1) // 2, sz, sz) AA = xp.zeros((N**2, sz, sz), dtype=vecA.dtype) AA[IDX_upper] = tmp AA[IDX_lower] = tmp[idx_offdiag].transpose(0, 2, 1) @@ -1839,8 +1842,8 @@ def transform_block(A, k, Pk): (n_blocks, k**2) and (n_blocks, (k + 1)**2), respectively. """ AT = Pk @ A @ Pk.T - A0 = AT[:, :k, :k].swapaxes(-1, -2).reshape(A.shape[0], -1) - A1 = AT[:, k:, k:].swapaxes(-1, -2).reshape(A.shape[0], -1) + A0 = AT[:, :k, :k].reshape(A.shape[0], -1) + A1 = AT[:, k:, k:].reshape(A.shape[0], -1) return A0, A1 @staticmethod @@ -1863,8 +1866,8 @@ def transform_back_block(A0, A1, k, Pk): """ dk = 2 * k + 1 A = xp.zeros((A0.shape[0], dk, dk), dtype=A0.dtype) - A[:, :k, :k] = A0.reshape(-1, k, k).swapaxes(-1, -2) - A[:, k:, k:] = A1.reshape(-1, k + 1, k + 1).swapaxes(-1, -2) + A[:, :k, :k] = A0.reshape(-1, k, k) + A[:, k:, k:] = A1.reshape(-1, k + 1, k + 1) return Pk.T @ A @ Pk ############################ From 56a6d50ab8ca2744b2f52814b232cdc89c6856e5 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 1 Oct 2026 15:13:59 -0400 Subject: [PATCH 116/161] update docstrings column-wise -> row-wise --- src/aspire/abinitio/commonline_nug.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 4997af33ec..5c7a1229c4 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1782,7 +1782,7 @@ def vec_block(A, N, sz, IDX_upper): The input is interpreted as an N-by-N array of sz-by-sz blocks. Blocks from the upper triangle, including the diagonal, are vectorized - column-wise and stored as columns of the returned array. + row-wise and stored as columns of the returned array. :param A: Block matrix with shape (N * sz, N * sz). :param N: Number of block rows and columns, corresponding to the @@ -1801,7 +1801,7 @@ def mat_block(vecA, N, sz, IDX_upper, IDX_lower, idx_offdiag): """ Unpack upper-triangular blocks into a symmetric block matrix. - Each column of `vecA` is reshaped column-wise into an sz-by-sz block. + Each column of `vecA` is reshaped row-wise into an sz-by-sz block. The blocks are placed in the upper triangle of an N-by-N block grid, including the diagonal. Transposes of the off-diagonal blocks are placed in the corresponding lower-triangular positions. @@ -1832,7 +1832,7 @@ def transform_block(A, k, Pk): Each input matrix is transformed as Pk @ A @ Pk.T. The leading k-by-k and trailing (k + 1)-by-(k + 1) diagonal blocks are then vectorized - column-wise. Entries in the two off-diagonal blocks are not retained. + row-wise. Entries in the two off-diagonal blocks are not retained. :param A: Batch of degree-k matrices with shape (n_blocks, 2 * k + 1, 2 * k + 1). :param k: Wigner representation degree. @@ -1851,7 +1851,7 @@ def transform_back_block(A0, A1, k, Pk): """ Reconstruct degree-k matrices from two packed components. - Each row of A0 and A1 is interpreted as a column-wise vectorization + Each row of A0 and A1 is interpreted as a row-wise vectorization of a k-by-k or (k + 1)-by-(k + 1) block. The two blocks are placed on the diagonal of a (2 * k + 1)-by-(2 * k + 1) matrix, with zeros in the off-diagonal blocks. The permutation represented by Pk is then reversed. From fb7d81cd80e2d9130347e06df52f34fb0d167294 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 1 Oct 2026 15:35:29 -0400 Subject: [PATCH 117/161] remove unnecessary transposes in eye().T --- src/aspire/abinitio/commonline_nug.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 5c7a1229c4..688210f86a 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -772,8 +772,8 @@ def admm_sym_J(self, C, verbose): AEk = xp.zeros((1 + n_components, 2 * n_components), dtype=np.float64) # Row zero contains the trace constraint for the two components. - AEk[0, :size0] = xp.eye(k, dtype=np.float64).T.reshape(-1) - AEk[0, size0:n_components] = xp.eye(k + 1, dtype=np.float64).T.reshape(-1) + AEk[0, :size0] = xp.eye(k, dtype=np.float64).reshape(-1) + AEk[0, size0:n_components] = xp.eye(k + 1, dtype=np.float64).reshape(-1) # The remaining rows couple every representation entry to its # corresponding auxiliary diagonal variable. @@ -799,11 +799,11 @@ def admm_sym_J(self, C, verbose): bE[start + 1 : split] = xp.eye( k, dtype=np.float64, - ).T.reshape(-1) + ).reshape(-1) bE[split:stop] = xp.eye( k + 1, dtype=np.float64, - ).T.reshape(-1) + ).reshape(-1) bE = xp.repeat(bE[:, None], N, axis=1) From d49fd41c2977bdc074ed4b47be2bfc5b629593aa Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 6 Oct 2026 11:23:15 -0400 Subject: [PATCH 118/161] remove unecessary xp.vdot and xp.ones call. replace with xp.sum --- src/aspire/abinitio/commonline_nug.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 688210f86a..eec3d51fee 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1086,14 +1086,12 @@ def print_updates(verbose): + 2 * xp.vdot(C0[:, idx_offdiag], X0[:, idx_offdiag]) + 2 * xp.vdot(C1[:, idx_offdiag], X1[:, idx_offdiag]) ) + obj_d = ( xp.vdot(yE, bE) + 2 * xp.vdot(yEq, bEq) - + xp.vdot(yI[:, idx_diag], bI * xp.ones((Ngrid, N))) - + 2 - * xp.vdot( - yI[:, idx_offdiag], bI * xp.ones((Ngrid, N * (N - 1) // 2)) - ) + + bI * xp.sum(yI[:, idx_diag]) + + 2 * bI * xp.sum(yI[:, idx_offdiag]) ) z, zq = fun_AE(X0, X1, Xd0, Xd1, Xq) From b7cb2763009e21ea76cda6a01973759306222cae Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 6 Oct 2026 11:36:25 -0400 Subject: [PATCH 119/161] host-scalar cleanup --- src/aspire/abinitio/commonline_nug.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index eec3d51fee..492c52359b 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -565,7 +565,7 @@ def update_rho(X0, X1, Xq, res_X, rho): + xp.linalg.norm(bEq - zq) / (1 + xp.linalg.norm(bEq)) ) res_inq = xp.linalg.norm(xp.maximum(bI - fun_AI(X0, X1), 0)) / ( - 1 + abs(bI) * xp.sqrt(Ngrid * n_pairs) + 1 + abs(bI) * np.sqrt(Ngrid * n_pairs) ) p_resnorm = res_eq + res_inq @@ -603,7 +603,7 @@ def print_updates(verbose): + xp.linalg.norm(bEq - zq) / (1 + xp.linalg.norm(bEq)) ) res_inq = xp.linalg.norm(xp.maximum(bI - fun_AI(X0, X1), 0)) / ( - 1 + abs(bI) * xp.sqrt(Ngrid * n_pairs) + 1 + abs(bI) * np.sqrt(Ngrid * n_pairs) ) res_psdX = 0 @@ -1068,7 +1068,7 @@ def update_rho(X0, X1, Xd0, Xd1, Xq, bE, bEq, bI, res_X, rho, factor, normC): zq - bEq ) / (1 + xp.linalg.norm(bEq)) res_inq = xp.linalg.norm(xp.maximum(bI - fun_AI(X0, X1), 0)) / ( - 1 + abs(bI) * xp.sqrt(Ngrid * N**2) + 1 + abs(bI) * np.sqrt(Ngrid * N**2) ) p_resnorm = res_eq + res_inq d_resnorm = res_X / (1 + normC) @@ -1099,7 +1099,7 @@ def print_updates(verbose): 1 + xp.linalg.norm(bE) ) + xp.linalg.norm(zq - bEq) / (1 + xp.linalg.norm(bEq)) res_inq = xp.linalg.norm(xp.maximum(bI - fun_AI(X0, X1), 0)) / ( - 1 + abs(bI) * xp.sqrt(Ngrid * N * (N + 1) / 2) + 1 + abs(bI) * np.sqrt(Ngrid * N * (N + 1) / 2) ) res_psdX = 0 From 94166c8624bcdb9d082609a4962b6aa78bc246e0 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 6 Oct 2026 11:43:39 -0400 Subject: [PATCH 120/161] keep compute_rank on CPU --- src/aspire/abinitio/commonline_nug.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 492c52359b..c559baac21 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1481,7 +1481,7 @@ def compute_rank(self): :return: Ranks and symmetry-averaging matrices for each degree. """ - rk = xp.zeros(self.Lmax, dtype=np.float64) + rk = np.zeros(self.Lmax, dtype=np.float64) A = [] for k_idx in range(self.Lmax): k = k_idx + 1 # degree From fdc1cb7228a1a7faf0f458c1574ac3ee6730a3b6 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 6 Oct 2026 13:28:53 -0400 Subject: [PATCH 121/161] cleanup right-hand side, bE, to remove small xp.eye calls --- src/aspire/abinitio/commonline_nug.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index c559baac21..483d0cc412 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -795,15 +795,9 @@ def admm_sym_J(self, C, verbose): split = k_idx + 1 + d0[k_idx + 1] + d1[k_idx] stop = k_idx + 1 + d0[k_idx + 1] + d1[k_idx + 1] - bE[start:] = rank_Ak[k_idx] - bE[start + 1 : split] = xp.eye( - k, - dtype=np.float64, - ).reshape(-1) - bE[split:stop] = xp.eye( - k + 1, - dtype=np.float64, - ).reshape(-1) + bE[start] = rank_Ak[k_idx] + bE[start + 1 : split : k + 1] = 1 + bE[split:stop:k + 2] = 1 bE = xp.repeat(bE[:, None], N, axis=1) From 726f62b36f4edab0270b9b7c691c4c802a0c6152 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 6 Oct 2026 13:48:08 -0400 Subject: [PATCH 122/161] cleanup asymmetric bE0/1 construction. --- src/aspire/abinitio/commonline_nug.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 483d0cc412..aaacd7de97 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -396,12 +396,12 @@ def admm_plain_J(self, C, verbose): bE0 = xp.zeros(D0, dtype=np.float64) bE1 = xp.zeros(D1, dtype=np.float64) for k_idx in range(Lmax): - k = k_idx + 1 - block0_rows = slice(d0[k_idx], d0[k_idx + 1]) - block1_rows = slice(d1[k_idx], d1[k_idx + 1]) + k = k_idx + 1 # degree - bE0[block0_rows] = xp.eye(k, dtype=np.float64).reshape(-1) - bE1[block1_rows] = xp.eye(k + 1, dtype=np.float64).reshape(-1) + # Encode the k-by-k and (k+1)-by-(k+1) identity blocks. A flattened + # m-by-m identity has diagonal entries every m+1 positions. + bE0[d0[k_idx] : d0[k_idx + 1] : k + 1] = 1 + bE1[d1[k_idx] : d1[k_idx + 1] : k + 2] = 1 bE0 = xp.repeat(bE0[:, None], N, axis=1) bE1 = xp.repeat(bE1[:, None], N, axis=1) From 2c98d3c2cc35da87a9bafcc5fbf225a08e082cad Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 6 Oct 2026 14:10:59 -0400 Subject: [PATCH 123/161] Remove unecessary xp.copy() --- src/aspire/abinitio/commonline_nug.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index aaacd7de97..8384aab419 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1713,16 +1713,15 @@ def permutek(Ak, k, N): :return: The two permuted matrix blocks. """ - AkP = xp.copy(Ak) dk = 2 * k + 1 - Pk = xp.eye(dk, dtype=AkP.dtype) + Pk = xp.eye(dk, dtype=Ak.dtype) for m in range(k): for n in range(k - m): Pk[(m + 2 * n, m + 2 * n + 1), :] = Pk[(m + 2 * n + 1, m + 2 * n), :] AkP = ( - xp.kron(xp.eye(N, dtype=AkP.dtype), Pk) + xp.kron(xp.eye(N, dtype=Ak.dtype), Pk) @ Ak - @ xp.kron(xp.eye(N, dtype=AkP.dtype), Pk.T) + @ xp.kron(xp.eye(N, dtype=Ak.dtype), Pk.T) ) Pk = xp.eye(N * dk, dtype=AkP.dtype) From 892f0c6335e9819f177e11681a92c40b19f219b2 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 6 Oct 2026 14:21:27 -0400 Subject: [PATCH 124/161] tox --- src/aspire/abinitio/commonline_nug.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 8384aab419..2f407d7882 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -797,7 +797,7 @@ def admm_sym_J(self, C, verbose): bE[start] = rank_Ak[k_idx] bE[start + 1 : split : k + 1] = 1 - bE[split:stop:k + 2] = 1 + bE[split : stop : k + 2] = 1 bE = xp.repeat(bE[:, None], N, axis=1) From 039c8e267ad37b3b89cdee399a8a84691c0b77b5 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 6 Oct 2026 15:48:55 -0400 Subject: [PATCH 125/161] cleanup AEk construction. Remove xp.eye() --- src/aspire/abinitio/commonline_nug.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 2f407d7882..880b6938fd 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -772,8 +772,8 @@ def admm_sym_J(self, C, verbose): AEk = xp.zeros((1 + n_components, 2 * n_components), dtype=np.float64) # Row zero contains the trace constraint for the two components. - AEk[0, :size0] = xp.eye(k, dtype=np.float64).reshape(-1) - AEk[0, size0:n_components] = xp.eye(k + 1, dtype=np.float64).reshape(-1) + AEk[0, : size0 : k + 1] = 1 + AEk[0, size0 : n_components : k + 2] = 1 # The remaining rows couple every representation entry to its # corresponding auxiliary diagonal variable. From 9d457182c4599a733949bdd540d658cbc72033d6 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 6 Oct 2026 15:54:54 -0400 Subject: [PATCH 126/161] cleanup bEq construction. remove xp.eye(). --- src/aspire/abinitio/commonline_nug.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 880b6938fd..0fb9410d1b 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1278,7 +1278,7 @@ def ADMM_preprocessing(self, C): AEqAEqtinv = xp.linalg.pinv(AEq @ AEq.T) bEq = xp.zeros(17, dtype=np.float64) - bEq[:16] = xp.eye(4, dtype=np.float64).reshape(-1) / 4 + bEq[:16:5] = 0.25 # Flattened np.eye(4)/4 bEq[-1] = 1 bEq = xp.repeat(bEq[:, None], N * (N - 1) // 2, axis=1) From 68a74b8fa1b77dc2754ff92d1bd2c25e807b388b Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 6 Oct 2026 16:05:47 -0400 Subject: [PATCH 127/161] keep small array operations on cpu in permutek --- src/aspire/abinitio/commonline_nug.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 0fb9410d1b..3b575448a0 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1714,10 +1714,11 @@ def permutek(Ak, k, N): :return: The two permuted matrix blocks. """ dk = 2 * k + 1 - Pk = xp.eye(dk, dtype=Ak.dtype) + Pk = np.eye(dk, dtype=Ak.dtype) for m in range(k): for n in range(k - m): Pk[(m + 2 * n, m + 2 * n + 1), :] = Pk[(m + 2 * n + 1, m + 2 * n), :] + Pk = xp.asarray(Pk) AkP = ( xp.kron(xp.eye(N, dtype=Ak.dtype), Pk) @ Ak From baf69d4d581a38899984b9f598a5bf86acef22a7 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 7 Oct 2026 10:36:37 -0400 Subject: [PATCH 128/161] permutek_back, keep small array operations on cpu. --- src/aspire/abinitio/commonline_nug.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 3b575448a0..0813825e72 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1756,10 +1756,11 @@ def permutek_back(Ak, k, N): ][idx, :] AkB = Pk.T @ Ak @ Pk dk = 2 * k + 1 - Pk = xp.eye(dk, dtype=Ak.dtype) + Pk = np.eye(dk, dtype=Ak.dtype) for m in range(k): for n in range(k - m): Pk[(m + 2 * n, m + 2 * n + 1), :] = Pk[(m + 2 * n + 1, m + 2 * n), :] + Pk = xp.asarray(Pk) AkB = ( xp.kron(xp.eye(N, dtype=Ak.dtype), Pk.T) @ AkB From 6a800aaadbf03fe1d45c64137daf3c6d257f1d4f Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 7 Oct 2026 11:59:20 -0400 Subject: [PATCH 129/161] remove unecessary use of transform_coeff on identity matrices --- src/aspire/abinitio/commonline_nug.py | 24 ++++++++++++++---------- 1 file changed, 14 insertions(+), 10 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 0813825e72..87b9bf7c98 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1325,17 +1325,21 @@ def ADMM_preprocessing(self, C): # largest eigenvalue for AIAIT Lambda = self._largest_eigenvalue(AI_mat_offdiag, N) - # initialization - II = [] - for k in range(1, Lmax + 1): - dk = 2 * k + 1 - II.append(xp.eye(N * dk, dtype=np.float64)) - I0, I1 = self.transform_coeff(II, IDX_upper) - X0 = xp.zeros((D0, N * (N + 1) // 2), dtype=np.float64) - X1 = xp.zeros((D1, N * (N + 1) // 2), dtype=np.float64) + # Initialize the primal variables in packed image-pair form. + n_packed = N * (N + 1) // 2 + X0 = xp.zeros((D0, n_packed), dtype=np.float64) + X1 = xp.zeros((D1, n_packed), dtype=np.float64) Xq = xp.zeros((16, N * (N - 1) // 2), dtype=np.float64) - S0 = xp.copy(I0) - S1 = xp.copy(I1) + + # The initial slack matrices are identity blocks for self-pairs and zero + # for distinct image pairs. Build them directly in packed form. + S0 = xp.zeros_like(X0) + S1 = xp.zeros_like(X1) + for k_idx in range(Lmax): + k = k_idx + 1 + S0[d0[k_idx] : d0[k_idx + 1] : k + 1, idx_diag] = 1 + S1[d1[k_idx] : d1[k_idx + 1] : k + 2, idx_diag] = 1 + Sq = xp.zeros(Xq.shape, dtype=np.float64) self.Ngrid = Ngrid From 3c8c43f33f2a5848ef0978b6bd0d94f1ac12462b Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 8 Oct 2026 11:46:06 -0400 Subject: [PATCH 130/161] Add trange to ADMM loop. Use f-strings for all logging. --- src/aspire/abinitio/commonline_nug.py | 48 ++++++++++----------------- 1 file changed, 18 insertions(+), 30 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 87b9bf7c98..dc18256300 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -7,7 +7,7 @@ from aspire.abinitio.sync_voting import _syncrotations from aspire.numeric import xp from aspire.operators import PolarFT -from aspire.utils import Rotation, cart2sph, complex_type +from aspire.utils import Rotation, cart2sph, complex_type, trange from aspire.volume import ( CnSymmetryGroup, DnSymmetryGroup, @@ -653,17 +653,12 @@ def print_updates(verbose): d_res = res_X / (1 + normC) logger.info( - "Iter %i" % t - + ": p_res=%1.5f" % p_res - + ", d_res=%1.5f" % d_res - + ", obj_primal=%1.2f" % obj_p - + ", obj_dual=%1.2f" % obj_d - + ", duality gap=%1.2f" % (obj_p - obj_d) - + "\n eq_res=%1.5f" % res_eq - + ", inq_res=%1.5f" % res_inq - + ", psd_res=%1.5f" % (res_psdX + res_psdQ) - + ", |S|=%1.2f" % normS - + ", |X|=%1.2f" % normX + f"Iter {t}: p_res={p_res:.5f}, d_res={d_res:.5f}, " + f"obj_primal={obj_p:.2f}, obj_dual={obj_d:.2f}, " + f"duality gap={obj_p - obj_d:.2f}\n" + f" eq_res={res_eq:.5f}, inq_res={res_inq:.5f}, " + f"psd_res={res_psdX + res_psdQ:.5f}, " + f"|S|={normS:.2f}, |X|={normX:.2f}" ) yI = xp.zeros((Ngrid, n_pairs), dtype=np.float64) @@ -672,7 +667,7 @@ def print_updates(verbose): yEq = xp.zeros(bEq.shape, dtype=np.float64) update_order = np.arange(3) - for t in range(max_iter): + for t in trange(max_iter, desc="C1 NUG ADMM"): np.random.shuffle(update_order) for update_idx in update_order: if update_idx == 0: @@ -1163,17 +1158,12 @@ def print_updates(verbose): p_res = res_eq + res_inq + res_psdX + res_psdD + res_psdQ d_res = res_X / (1 + normC) logger.info( - "Iter %i" % t - + ": p_res=%1.5f" % p_res - + ", d_res=%1.5f" % d_res - + ", obj_primal=%1.2f" % obj_p - + ", obj_dual=%1.2f" % obj_d - + ", duality gap=%1.2f" % (obj_p - obj_d) - + "\n eq_res=%1.5f" % res_eq - + ", inq_res=%1.5f" % res_inq - + ", psd_res=%1.5f" % (res_psdX + res_psdD + res_psdQ) - + ", |S|=%1.2f" % normS - + ", |X|=%1.2f" % normX + f"Iter {t}: p_res={p_res:.5f}, d_res={d_res:.5f}, " + f"obj_primal={obj_p:.2f}, obj_dual={obj_d:.2f}, " + f"duality gap={obj_p - obj_d:.2f}\n" + f" eq_res={res_eq:.5f}, inq_res={res_inq:.5f}, " + f"psd_res={res_psdX + res_psdD + res_psdQ:.5f}, " + f"|S|={normS:.2f}, |X|={normX:.2f}" ) Xd0 = xp.zeros((D0, N), dtype=xp.float64) @@ -1187,7 +1177,7 @@ def print_updates(verbose): # Run ADMM iterations, randomly ordering the block updates before each primal # multiplier update and penalty adjustment. update_order = np.arange(3) - for t in range(max_iter): + for t in trange(max_iter, desc="Symmetric NUG ADMM"): np.random.shuffle(update_order) for update_idx in update_order: if update_idx == 0: @@ -1470,7 +1460,7 @@ def _largest_eigenvalue(AI, N): z = z / xp.linalg.norm(z) z = AI @ (AI.T @ z) Lambda += 2000 - logger.info("Largest eigenvalue of AIAIT is approximately %1.2f" % Lambda) + logger.info(f"Largest eigenvalue of AIAIT is approximately {Lambda:.2f}") return Lambda def compute_rank(self): @@ -1961,10 +1951,8 @@ def low_rank_proj(X, r_step): if self.verbose: logger.info( - "Proximal refine step %d/%d: relative update %.3e", - step + 1, - self.pr_iters, - rel_change(X_next, current), + f"Proximal refine step {step + 1}/{self.pr_iters}: " + f"relative update {rel_change(X_next, current):.3e}" ) # Use the refined SDP solution as the starting point for the next projection step From ca7befcf1758b7ac1ed1a8ee3ade9789409d4359 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 8 Oct 2026 13:55:42 -0400 Subject: [PATCH 131/161] add progress bars to coefficient computation --- src/aspire/abinitio/commonline_nug.py | 33 +++++++++++++++------------ 1 file changed, 19 insertions(+), 14 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index dc18256300..63cdc3a9e4 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -7,7 +7,7 @@ from aspire.abinitio.sync_voting import _syncrotations from aspire.numeric import xp from aspire.operators import PolarFT -from aspire.utils import Rotation, cart2sph, complex_type, trange +from aspire.utils import Rotation, cart2sph, complex_type, tqdm, trange from aspire.volume import ( CnSymmetryGroup, DnSymmetryGroup, @@ -272,25 +272,30 @@ def fijhat_k(k_idx, F): C.append(np.zeros((N * dk, N * dk), dtype=complex_type(np.float64))) # Compute off-diagonal image-pair losses and insert their degree-wise coefficients. - for i in range(N): - for j in range(i + 1, N): - Fij = np.zeros((2 * T, 2 * T), dtype=np.float64) - for j1 in range(2 * T): - for j2 in range(2 * T): - Fij[j1, j2] = fij(alpha_grid[j1], gamma_grid[j2], i, j) - for k_idx in range(Lmax): - k = k_idx + 1 - dk = 2 * k + 1 - C[k_idx][j * dk : (j + 1) * dk, i * dk : (i + 1) * dk] = fijhat_k( - k_idx, Fij - ) # *dk + with tqdm( + total=N * (N - 1) // 2, + desc="Computing off-diagonal coefficients", + ) as pbar: + for i in range(N): + for j in range(i + 1, N): + Fij = np.zeros((2 * T, 2 * T), dtype=np.float64) + for j1 in range(2 * T): + for j2 in range(2 * T): + Fij[j1, j2] = fij(alpha_grid[j1], gamma_grid[j2], i, j) + for k_idx in range(Lmax): + k = k_idx + 1 + dk = 2 * k + 1 + C[k_idx][j * dk : (j + 1) * dk, i * dk : (i + 1) * dk] = ( + fijhat_k(k_idx, Fij) + ) # *dk + pbar.update(1) # Fill the conjugate transpose blocks so each coefficient matrix is Hermitian. for k_idx in range(Lmax): C[k_idx] = C[k_idx] + C[k_idx].conj().T # Diagonal blocks encode self-pair losses and are handled separately. - for i in range(N): + for i in trange(N, desc="Computing diagonal coefficients"): Fii = np.zeros((2 * T, 2 * T), dtype=np.float64) for j1 in range(2 * T): for j2 in range(2 * T): From cf3d19b10ec308a7531d4429c5e16ec0b2078c60 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 8 Oct 2026 14:07:26 -0400 Subject: [PATCH 132/161] Gallery example for C1 experiment --- ...perimental_nug_simulated_pipeline_10028.py | 115 ++++++++++++++++++ 1 file changed, 115 insertions(+) create mode 100644 gallery/experiments/experimental_nug_simulated_pipeline_10028.py diff --git a/gallery/experiments/experimental_nug_simulated_pipeline_10028.py b/gallery/experiments/experimental_nug_simulated_pipeline_10028.py new file mode 100644 index 0000000000..e56d7101f6 --- /dev/null +++ b/gallery/experiments/experimental_nug_simulated_pipeline_10028.py @@ -0,0 +1,115 @@ +""" +C1 NUG Rotation Accuracy on Simulated EMDB-2660 Images +====================================================== + +Estimate orientations of noisy, simulated projections of EMDB-2660 using +the asymmetric (C1) NUG method. Compare the estimated rotations with the +known simulation rotations at several signal-to-noise ratios (SNRs). A +volume is reconstructed and saved off for comparison against the generating +volume. + +This experiment is based on: + +.. admonition:: Reference + + A. S. Bandeira, Y. Chen, R. R. Lederman, and A. Singer, + "Non-unique games over compact groups and orientation estimation + in cryo-EM," Inverse Problems 36(6), 064002 (2020). + https://doi.org/10.1088/1361-6420/ab7d2c + +The ASPIRE-Python implementation of NUG differs from the original MATLAB +implementation, so this example is intended to reproduce the experimental +workflow rather than its exact numerical results. +""" + +# %% +# Imports +# ------- +import logging + +import numpy as np + +from aspire.abinitio import CommonlineNUG +from aspire.downloader import emdb_2660 +from aspire.noise import WhiteNoiseAdder +from aspire.reconstruction import MeanEstimator +from aspire.source import OrientedSource, Simulation +from aspire.utils import Rotation + +logger = logging.getLogger(__name__) + + +# %% +# Experiment settings +# ------------------- +# Start with 100 images. The paper also reports results for 500 images and +# SNR ranging from 1 to 1/128, which requires substantially more computation. +SEED = 1980 +RESOLUTION = 129 +N_IMAGES = (100,) +SNR_VALUES = (1, 1 / 2, 1 / 4, 1 / 8) + +# Download and downsample once; every SNR uses the same underlying volume. +volume = emdb_2660().astype(np.float64).downsample(RESOLUTION) + + +# %% +# Estimate rotations at each SNR +# ------------------------------ +results = [] + +for n_images in N_IMAGES: + for snr in SNR_VALUES: + logger.info(f"Estimating C1 rotations: n={n_images}, SNR={snr}") + + # Keeping the simulation seed fixed gives corresponding ground-truth + # orientations across SNR conditions. The noise seed is also fixed, + # so the noise realization is scaled as the requested SNR changes. + source = Simulation( + n=n_images, + vols=volume, + offsets=0, + amplitudes=1, + seed=SEED, + noise_adder=WhiteNoiseAdder.from_snr(snr=snr, seed=SEED), + ).cache() + + # The ADMM update order uses NumPy randomness. Reset it for a + # reproducible run at each SNR. + np.random.seed(SEED) + + estimator = CommonlineNUG( + source, + symmetry="C1", + max_shift=0, + max_iter=501, + verbose=False, + ) + + oriented_source = OrientedSource(source, estimator) + estimated_rotations = oriented_source.rotations + + # Rotation.mse registers the estimates to ground truth before + # computing the mean squared error. + mse = Rotation(estimated_rotations).mse(Rotation(source.rotations)) + results.append((n_images, snr, mse)) + + print(f"n={n_images} SNR={snr} rotation MSE={mse}", flush=True) + + # The simulation has zero offsets. Use the estimated orientations + # with zero shifts to reconstruct from the cached noisy images. + oriented_source = oriented_source.update(offsets=0) + + snr_tag = f"{snr:g}".replace(".", "p") + volume_filename = f"nug_2660_c1_n{n_images}_snr{snr_tag}_recon.mrc" + + logger.info(f"Reconstructing and saving {volume_filename}") + estimated_volume = MeanEstimator(oriented_source).estimate() + estimated_volume.save(volume_filename, overwrite=True) +# %% +# Summary +# ------- +print("\nC1 NUG rotation accuracy") +print(f"{'Images':>6} {'SNR':>10} {'MSE':>14}") +for n_images, snr, mse in results: + print(f"{n_images:6d} {snr:10g} {mse:14.6g}") From 7b1230a7f3200f05ae055ee5ad17b80cccc67095 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 8 Oct 2026 14:11:15 -0400 Subject: [PATCH 133/161] add gallery print error to tox.ini --- tox.ini | 1 + 1 file changed, 1 insertion(+) diff --git a/tox.ini b/tox.ini index 37b5bcfcbe..f278dd6af5 100644 --- a/tox.ini +++ b/tox.ini @@ -67,6 +67,7 @@ max-line-length = 88 extend-ignore = E203, E501 per-file-ignores = __init__.py: F401 + gallery/experiments/experimental_nug_simulated_pipeline_10028.py: T201 gallery/tutorials/tutorials/*.py: T201, E402 gallery/tutorials/aspire_introduction.py: T201, F401, E402 gallery/tutorials/configuration.py: T201, E402 From 5aba7b9f8e2af0e5af36128158bc4cc2fd34d68a Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 12 Oct 2026 14:15:18 -0400 Subject: [PATCH 134/161] Clean up gallery. Use full size simulation with noise, then downsample. --- ...perimental_nug_simulated_pipeline_10028.py | 27 ++++++++----------- 1 file changed, 11 insertions(+), 16 deletions(-) diff --git a/gallery/experiments/experimental_nug_simulated_pipeline_10028.py b/gallery/experiments/experimental_nug_simulated_pipeline_10028.py index e56d7101f6..e87d79285e 100644 --- a/gallery/experiments/experimental_nug_simulated_pipeline_10028.py +++ b/gallery/experiments/experimental_nug_simulated_pipeline_10028.py @@ -47,11 +47,10 @@ SEED = 1980 RESOLUTION = 129 N_IMAGES = (100,) -SNR_VALUES = (1, 1 / 2, 1 / 4, 1 / 8) - -# Download and downsample once; every SNR uses the same underlying volume. -volume = emdb_2660().astype(np.float64).downsample(RESOLUTION) +SNR_VALUES = (1, 1 / 2, 1 / 4, 1 / 8, 1 / 16) +# Download volume once; every SNR uses the same underlying volume. +volume = emdb_2660().astype(np.float64) # %% # Estimate rotations at each SNR @@ -60,29 +59,24 @@ for n_images in N_IMAGES: for snr in SNR_VALUES: - logger.info(f"Estimating C1 rotations: n={n_images}, SNR={snr}") - + logger.info(f"Estimating rotations: n={n_images}, SNR={snr}") # Keeping the simulation seed fixed gives corresponding ground-truth - # orientations across SNR conditions. The noise seed is also fixed, - # so the noise realization is scaled as the requested SNR changes. + # orientations across SNR conditions. source = Simulation( n=n_images, vols=volume, offsets=0, amplitudes=1, seed=SEED, - noise_adder=WhiteNoiseAdder.from_snr(snr=snr, seed=SEED), - ).cache() - - # The ADMM update order uses NumPy randomness. Reset it for a - # reproducible run at each SNR. - np.random.seed(SEED) + noise_adder=WhiteNoiseAdder.from_snr(snr=snr), + ).downsample(RESOLUTION).cache() + # Initialize the NUG orientation estimator, matching paramters + # to the referenced experiment. estimator = CommonlineNUG( source, symmetry="C1", max_shift=0, - max_iter=501, verbose=False, ) @@ -100,12 +94,13 @@ # with zero shifts to reconstruct from the cached noisy images. oriented_source = oriented_source.update(offsets=0) - snr_tag = f"{snr:g}".replace(".", "p") + snr_tag = f"{snr:g}" volume_filename = f"nug_2660_c1_n{n_images}_snr{snr_tag}_recon.mrc" logger.info(f"Reconstructing and saving {volume_filename}") estimated_volume = MeanEstimator(oriented_source).estimate() estimated_volume.save(volume_filename, overwrite=True) + # %% # Summary # ------- From c3798852b7699260c084e2c61efd17ac5ff30407 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 12 Oct 2026 14:20:40 -0400 Subject: [PATCH 135/161] black --- ...perimental_nug_simulated_pipeline_10028.py | 20 +++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/gallery/experiments/experimental_nug_simulated_pipeline_10028.py b/gallery/experiments/experimental_nug_simulated_pipeline_10028.py index e87d79285e..f78a7f287e 100644 --- a/gallery/experiments/experimental_nug_simulated_pipeline_10028.py +++ b/gallery/experiments/experimental_nug_simulated_pipeline_10028.py @@ -62,14 +62,18 @@ logger.info(f"Estimating rotations: n={n_images}, SNR={snr}") # Keeping the simulation seed fixed gives corresponding ground-truth # orientations across SNR conditions. - source = Simulation( - n=n_images, - vols=volume, - offsets=0, - amplitudes=1, - seed=SEED, - noise_adder=WhiteNoiseAdder.from_snr(snr=snr), - ).downsample(RESOLUTION).cache() + source = ( + Simulation( + n=n_images, + vols=volume, + offsets=0, + amplitudes=1, + seed=SEED, + noise_adder=WhiteNoiseAdder.from_snr(snr=snr), + ) + .downsample(RESOLUTION) + .cache() + ) # Initialize the NUG orientation estimator, matching paramters # to the referenced experiment. From 405c78d93a5c4aaad2ab7ea7b68216045dafaed1 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 12 Oct 2026 14:46:59 -0400 Subject: [PATCH 136/161] Add dataset references --- .../experimental_nug_simulated_pipeline_10028.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/gallery/experiments/experimental_nug_simulated_pipeline_10028.py b/gallery/experiments/experimental_nug_simulated_pipeline_10028.py index f78a7f287e..204f867013 100644 --- a/gallery/experiments/experimental_nug_simulated_pipeline_10028.py +++ b/gallery/experiments/experimental_nug_simulated_pipeline_10028.py @@ -8,6 +8,12 @@ volume is reconstructed and saved off for comparison against the generating volume. +EMDB-2660 is the volume entry reconstructed from the experimental dataset +EMPIAR 10028, both available here: + +https://www.ebi.ac.uk/emdb/EMD-2660 +https://www.ebi.ac.uk/empiar/EMPIAR-10028 + This experiment is based on: .. admonition:: Reference From c5b4f5e3673e45dd7ac8c38c5bec8efd9cbb55f4 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 13 Oct 2026 13:39:48 -0400 Subject: [PATCH 137/161] staticmethods --- src/aspire/abinitio/commonline_nug.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 63cdc3a9e4..80a1fa4ad1 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1484,7 +1484,8 @@ def compute_rank(self): rk[k_idx] = np.linalg.matrix_rank(Ak) return rk, A - def construct_AEq(self): + @staticmethod + def construct_AEq(): """ Construct the linear equality operator encoding the quaternion constraints. @@ -2628,7 +2629,8 @@ def Wd(k, beta): ) return d - def complex2real(self, ell): + @staticmethod + def complex2real(ell): """ Construct the transformation matrices between complex and real degree ell representations. From 78b7dcc2a9f322d9f7c35bdf195c684384ee494a Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 13 Oct 2026 13:59:44 -0400 Subject: [PATCH 138/161] Make helpers private --- src/aspire/abinitio/commonline_nug.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 80a1fa4ad1..53c8260cee 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -756,7 +756,7 @@ def admm_sym_J(self, C, verbose): n_pairs = N * (N - 1) // 2 Ngrid = self.Ngrid - rank_Ak, _ = self.compute_rank() + rank_Ak, _ = self._compute_rank() logger.info(f"Rank of Ak: {rank_Ak}") # Build per-degree equality operators coupling the k and k+1 invariant @@ -1269,7 +1269,7 @@ def ADMM_preprocessing(self, C): D1 = d1[-1] # AE and bE for quaternion constraints - AEq = xp.asarray(self.construct_AEq()) + AEq = xp.asarray(self._construct_AEq()) AEqAEqtinv = xp.linalg.pinv(AEq @ AEq.T) bEq = xp.zeros(17, dtype=np.float64) @@ -1278,7 +1278,7 @@ def ADMM_preprocessing(self, C): bEq = xp.repeat(bEq[:, None], N * (N - 1) // 2, axis=1) # Compute AI and bI - W0, W1, Ngrid = self.compute_fejer_weights() + W0, W1, Ngrid = self._compute_fejer_weights() # AI_mat_offdiag computation has been vectorized, but originally # had the note by block1: this needs double checking (Ruiyi) @@ -1366,13 +1366,13 @@ def ADMM_preprocessing(self, C): Sq, ) - def compute_fejer_weights(self): + def _compute_fejer_weights(self): """ Evaluate the real Wigner representation blocks used by the discretized Fejer inequality constraints. :return: Two sets of block weights and the number of SO(3) grid points. """ - SO3_grid = self.discretize_SO3() + SO3_grid = self._discretize_SO3() Ngrid = SO3_grid.shape[0] TT = [] @@ -1411,7 +1411,7 @@ def permutek_block(Ak, k): W1.append(W1k) return W0, W1, Ngrid - def discretize_SO3(self): + def _discretize_SO3(self): """ Construct an approximately uniform Euler-angle grid over SO(3). @@ -1468,7 +1468,7 @@ def _largest_eigenvalue(AI, N): logger.info(f"Largest eigenvalue of AIAIT is approximately {Lambda:.2f}") return Lambda - def compute_rank(self): + def _compute_rank(self): """ Compute the ranks and matrices of the symmetry-averaging projectors at each degree. @@ -1485,7 +1485,7 @@ def compute_rank(self): return rk, A @staticmethod - def construct_AEq(): + def _construct_AEq(): """ Construct the linear equality operator encoding the quaternion constraints. From 60f2d5140090d0077c6d18b4707939012932110d Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 13 Oct 2026 14:22:42 -0400 Subject: [PATCH 139/161] Rename remaining helpers as private --- src/aspire/abinitio/commonline_nug.py | 74 +++++++++++++-------------- 1 file changed, 37 insertions(+), 37 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 53c8260cee..52ba09acc4 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -242,7 +242,7 @@ def fij(alpha, gamma, i, j): # Precompute degree-wise beta/Wigner weight matrices used in each coefficient transform. BTK = [] for k in range(1, Lmax + 1): - btk = np.sum(bT[:, None, None] * self.Wd(k, beta_grid), axis=0) + btk = np.sum(bT[:, None, None] * self._Wd(k, beta_grid), axis=0) BTK.append(btk.T) def fijhat_k(k_idx, F): @@ -310,7 +310,7 @@ def fijhat_k(k_idx, F): # Convert complex Wigner coefficients to the real representation basis used by ADMM. for k_idx in range(Lmax): k = k_idx + 1 - T, Tinv = self.complex2real(k) + T, Tinv = self._complex2real(k) C[k_idx] = np.real( np.kron(np.eye(N, dtype=np.float64), Tinv) @ C[k_idx] @@ -477,7 +477,7 @@ def update_S(C0, C1, yE0, yE1, yEq, yI, X0, X1, Xq, rho): # This is the established batched equivalent of projecting each 4x4 # quaternion block and storing tmp.T.reshape(16) in each column. - Sq = self.psd_projection(Sq.T.reshape(n_pairs, 4, 4)) + Sq = self._psd_projection(Sq.T.reshape(n_pairs, 4, 4)) Sq = Sq.T.reshape(16, n_pairs) return S0, S1, Sq @@ -625,7 +625,7 @@ def print_updates(verbose): IDX_lower, idx_offdiag, ) - res_psdX += xp.linalg.norm(self.psd_projection(-tmp)) + res_psdX += xp.linalg.norm(self._psd_projection(-tmp)) tmp = self.mat_block( X1[block1_rows], @@ -635,12 +635,12 @@ def print_updates(verbose): IDX_lower, idx_offdiag, ) - res_psdX += xp.linalg.norm(self.psd_projection(-tmp)) + res_psdX += xp.linalg.norm(self._psd_projection(-tmp)) res_psdX /= 1 + xp.linalg.norm(X0) + xp.linalg.norm(X1) Xq_blocks = Xq.T.reshape(n_pairs, 4, 4) res_psdQ = xp.linalg.norm( - self.psd_projection(-Xq_blocks), axis=(-2, -1) + self._psd_projection(-Xq_blocks), axis=(-2, -1) ).sum() res_psdQ /= 1 + xp.linalg.norm(Xq) @@ -703,7 +703,7 @@ def print_updates(verbose): print_updates(verbose) rho, p_resnorm, d_resnorm = update_rho(X0, X1, Xq, res_X, rho) - X_admm = self.transform_coeff_back(X0, X1, IDX_upper, IDX_lower, idx_offdiag) + X_admm = self._transform_coeff_back(X0, X1, IDX_upper, IDX_lower, idx_offdiag) X_admm = [xp.asnumpy(Xk) for Xk in X_admm] return X_admm @@ -961,7 +961,7 @@ def update_S(C0, C1, yE, yEq, yI, X0, X1, Xd0, Xd1, Xq, rho, Lmax, N): k, permutation, ) - tmp = self.psd_projection(tmp) + tmp = self._psd_projection(tmp) Sd0[:, block0_rows], Sd1[:, block1_rows] = self.transform_block( tmp, k, permutation @@ -971,7 +971,7 @@ def update_S(C0, C1, yE, yEq, yI, X0, X1, Xd0, Xd1, Xq, rho, Lmax, N): Sd1 = Sd1.T # Project each 4x4 quaternion slack block for off-diagonal image pairs. - Sq = self.psd_projection(Sq.T.reshape(n_pairs, 4, 4)) + Sq = self._psd_projection(Sq.T.reshape(n_pairs, 4, 4)) Sq = Sq.T.reshape(-1, n_pairs) return S0, S1, Sd0, Sd1, Sq @@ -1110,7 +1110,7 @@ def print_updates(verbose): IDX_lower, idx_offdiag, ) - res_psdX += xp.linalg.norm(self.psd_projection(-tmp)) + res_psdX += xp.linalg.norm(self._psd_projection(-tmp)) tmp = self.mat_block( X1[block1_rows], @@ -1120,7 +1120,7 @@ def print_updates(verbose): IDX_lower, idx_offdiag, ) - res_psdX += xp.linalg.norm(self.psd_projection(-tmp)) + res_psdX += xp.linalg.norm(self._psd_projection(-tmp)) res_psdX = res_psdX / (1 + xp.linalg.norm(X0) + xp.linalg.norm(X1)) res_psdD = 0 @@ -1136,14 +1136,14 @@ def print_updates(verbose): P[k_idx], ) res_psdD += xp.linalg.norm( - self.psd_projection(-tmp), axis=(-2, -1) + self._psd_projection(-tmp), axis=(-2, -1) ).sum() res_psdD = res_psdD / (1 + xp.linalg.norm(Xd0) + xp.linalg.norm(Xd1)) res_psdQ = 0 for pair_idx in range(n_pairs): tmp = Xq[:, pair_idx].reshape(4, 4).T - res_psdQ += xp.linalg.norm(self.psd_projection(-tmp)) + res_psdQ += xp.linalg.norm(self._psd_projection(-tmp)) res_psdQ = res_psdQ / (1 + xp.linalg.norm(Xq)) normS = xp.sqrt( @@ -1207,7 +1207,7 @@ def print_updates(verbose): # Convert the optimized packed block variables back to full degree-wise # representation matrices for Euler-angle recovery. - X_admm = self.transform_coeff_back(X0, X1, IDX_upper, IDX_lower, idx_offdiag) + X_admm = self._transform_coeff_back(X0, X1, IDX_upper, IDX_lower, idx_offdiag) X_admm = [xp.asnumpy(Xk) for Xk in X_admm] return X_admm @@ -1255,7 +1255,7 @@ def ADMM_preprocessing(self, C): for k_idx in range(Lmax): C[k_idx] = xp.asarray(Xnorm / Cnorm * C[k_idx]) - C0, C1 = self.transform_coeff(C, IDX_upper) + C0, C1 = self._transform_coeff(C, IDX_upper) normC = xp.sqrt(xp.linalg.norm(C0) ** 2 + xp.linalg.norm(C1) ** 2) del C @@ -1379,7 +1379,7 @@ def _compute_fejer_weights(self): TTI = [] for k_idx in range(self.Lmax): k = k_idx + 1 - T, Tinv = self.complex2real(k) + T, Tinv = self._complex2real(k) TT.append(T) TTI.append(Tinv) @@ -1404,7 +1404,7 @@ def permutek_block(Ak, k): TkT = TT[k_idx].T TinvkT = TTI[k_idx].T - w = np.real(TkT @ self.WD(k, SO3_grid).conj() @ TinvkT) + w = np.real(TkT @ self._WD(k, SO3_grid).conj() @ TinvkT) W0k, W1k = permutek_block(w, k) W0.append(W0k) @@ -1478,7 +1478,7 @@ def _compute_rank(self): A = [] for k_idx in range(self.Lmax): k = k_idx + 1 # degree - Ak = np.sum(self.WD(k, self.sym_euler), axis=0) + Ak = np.sum(self._WD(k, self.sym_euler), axis=0) Ak = np.round(Ak / self.n_sym, 6) A.append(Ak) rk[k_idx] = np.linalg.matrix_rank(Ak) @@ -1580,7 +1580,7 @@ def _project_representation_blocks( IDX_lower, idx_offdiag, ) - block0 = self.psd_projection(block0) + block0 = self._psd_projection(block0) Z0[block0_rows] = self.vec_block( block0, self.n_img, @@ -1596,7 +1596,7 @@ def _project_representation_blocks( IDX_lower, idx_offdiag, ) - block1 = self.psd_projection(block1) + block1 = self._psd_projection(block1) Z1[block1_rows] = self.vec_block( block1, self.n_img, @@ -1606,7 +1606,7 @@ def _project_representation_blocks( return Z0, Z1 @staticmethod - def psd_projection(B): + def _psd_projection(B): """ Project real square matrices onto the positive semidefinite cone. @@ -1624,7 +1624,7 @@ def psd_projection(B): evals = xp.maximum(evals, 0) return (evecs * evals[..., None, :]) @ evecs.swapaxes(-1, -2) - def transform_coeff(self, A, IDX_upper): + def _transform_coeff(self, A, IDX_upper): """ Convert representation matrices to the block-vector form used by ADMM. @@ -1647,14 +1647,14 @@ def transform_coeff(self, A, IDX_upper): block0_rows = slice(d0[k_idx], d0[k_idx + 1]) block1_rows = slice(d1[k_idx], d1[k_idx + 1]) - block0, block1 = self.permutek(A[k_idx], k, self.n_img) + block0, block1 = self._permutek(A[k_idx], k, self.n_img) A0[block0_rows] = self.vec_block(block0, self.n_img, k, IDX_upper) A1[block1_rows] = self.vec_block(block1, self.n_img, k + 1, IDX_upper) return A0, A1 - def transform_coeff_back(self, A0, A1, IDX_upper, IDX_lower, idx_offdiag): + def _transform_coeff_back(self, A0, A1, IDX_upper, IDX_lower, idx_offdiag): """ Reconstruct representation matrices from their ADMM block-vector form. @@ -1698,12 +1698,12 @@ def transform_coeff_back(self, A0, A1, IDX_upper, IDX_lower, idx_offdiag): IDX_lower, idx_offdiag, ) - Ak = self.permutek_back(Ak, k, N) + Ak = self._permutek_back(Ak, k, N) A.append(Ak) return A @staticmethod - def permutek(Ak, k, N): + def _permutek(Ak, k, N): """ Permute and split a degree-k matrix into blocks of sizes k and k + 1. @@ -1736,7 +1736,7 @@ def permutek(Ak, k, N): return AkP[: N * k, : N * k], AkP[N * k :, N * k :] @staticmethod - def permutek_back(Ak, k, N): + def _permutek_back(Ak, k, N): """ Undo the degree-k block permutation and reconstruct the full matrix. @@ -1887,7 +1887,7 @@ def Ak(k, Euler): group to form the symmetry projector Ak. """ order = Euler.shape[0] - A = self.WD(k, Euler).sum(axis=0) + A = self._WD(k, Euler).sum(axis=0) return np.round(A / order, 10) def rel_change(A, B, eps=1e-12): @@ -2353,7 +2353,7 @@ def _estimate_gamma( Jk = self._handedness_matrix(degree) # Evaluate the Wigner small-d matrices at the recovered beta angles. - ws = self.Wd(degree, beta) + ws = self._Wd(degree, beta) for i in range(N): wi = ws[i] @@ -2542,7 +2542,7 @@ def _real_to_complex_representation(self, X, degree): :return: Representation matrix in the complex Wigner basis. """ - T, Tinv = self.complex2real(degree) + T, Tinv = self._complex2real(degree) identity = np.eye(self.n_img, dtype=np.float64) return np.kron(identity, T) @ X @ np.kron(identity, Tinv) @@ -2566,7 +2566,7 @@ def _handedness_matrix(degree): # Representation Helpers # ########################## - def WD(self, k, euler): + def _WD(self, k, euler): """ Evaluate degree-k Wigner D matrices at the supplied ZYZ Euler angles. @@ -2579,7 +2579,7 @@ def WD(self, k, euler): alpha = euler[:, 0] beta = euler[:, 1] gamma = euler[:, 2] - d = self.Wd(k, beta) + d = self._Wd(k, beta) m = np.arange(-k, k + 1) left = np.exp(-1j * alpha[:, None] * m[None, :]) @@ -2589,7 +2589,7 @@ def WD(self, k, euler): return D @staticmethod - def Wd(k, beta): + def _Wd(k, beta): """ Evaluate degree-k Wigner small-d matrices at the supplied polar angles. @@ -2630,7 +2630,7 @@ def Wd(k, beta): return d @staticmethod - def complex2real(ell): + def _complex2real(ell): """ Construct the transformation matrices between complex and real degree ell representations. @@ -2687,12 +2687,12 @@ def form_ground_truth_X(self, euler_angles): dk = 2 * k + 1 # Evaluate the degree-k Wigner representations. - wigner = self.WD(k, euler_angles) + wigner = self._WD(k, euler_angles) # Average the degree-k representation over the molecular symmetry # group. For the asymmetric case, self.sym_euler contains only the # identity and Ak is therefore the identity matrix. - Ak = np.mean(self.WD(k, self.sym_euler), axis=0) + Ak = np.mean(self._WD(k, self.sym_euler), axis=0) # Construct the globally handedness-conjugated solution. Jk = np.ones(dk) @@ -2703,7 +2703,7 @@ def form_ground_truth_X(self, euler_angles): wigner_J = Jk @ wigner @ Jk # Convert each representation to the real basis used by ADMM. - _, Tinv = self.complex2real(k) + _, Tinv = self._complex2real(k) wigner_real = Tinv @ wigner wigner_J_real = Tinv @ wigner_J From de1fbe334414ee5fac98f394241562742b6d8f8c Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 14 Oct 2026 15:43:39 -0400 Subject: [PATCH 140/161] return preprocessing variables as dict and unpack with descriptions --- src/aspire/abinitio/commonline_nug.py | 160 +++++++++++++------------- 1 file changed, 79 insertions(+), 81 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 52ba09acc4..b8a6bd70a8 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -366,32 +366,33 @@ def admm_plain_J(self, C, verbose): mult = self.mult Nstep_yI = self.Nstep_yI - ( - C0, - C1, - normC, - AEq, - bEq, - AEqAEqtinv, - AI_mat_diag, - AI_mat_offdiag, - bI, - Lambda, - d0, - d1, - D0, - D1, - idx_diag, - idx_offdiag, - IDX_upper, - IDX_lower, - X0, - X1, - Xq, - S0, - S1, - Sq, - ) = self.ADMM_preprocessing(C) + # Prepare the objective, constraints, block layout, and initial ADMM variables. + pre = self.ADMM_preprocessing(C) + + # Packed objective coefficients and their norm. + C0, C1, normC = pre["C0"], pre["C1"], pre["normC"] + + # Quaternion equality constraints. + AEq, bEq = pre["AEq"], pre["bEq"] + AEqAEqtinv = pre["AEqAEqtinv"] + + # Fejér inequality constraints and spectral scale. + AI_mat_diag = pre["AI_mat_diag"] + AI_mat_offdiag = pre["AI_mat_offdiag"] + bI, Lambda = pre["bI"], pre["Lambda"] + + # Row boundaries by degree: d0 applies to C0/X0/S0, and d1 to C1/X1/S1. + # Their final entries give the total row counts. + d0, d1 = pre["d0"], pre["d1"] + D0, D1 = d0[-1], d1[-1] + + # Indices into packed columns and the full image-pair grid, respectively. + idx_diag, idx_offdiag = pre["idx_diag"], pre["idx_offdiag"] + IDX_upper, IDX_lower = pre["IDX_upper"], pre["IDX_lower"] + + # Initial primal and slack variables. + X0, X1, Xq = pre["X0"], pre["X1"], pre["Xq"] + S0, S1, Sq = pre["S0"], pre["S1"], pre["Sq"] Ngrid = self.Ngrid n_pairs = N * (N - 1) // 2 @@ -725,34 +726,33 @@ def admm_sym_J(self, C, verbose): mult = self.mult Nstep_yI = self.Nstep_yI - # Solve the symmetry-constrained SDP relaxation in the real, packed - # representation basis prepared by ADMM_preprocessing. - ( - C0, - C1, - normC, - AEq, - bEq, - AEqAEqtinv, - AI_mat_diag, - AI_mat_offdiag, - bI, - Lambda, - d0, - d1, - D0, - D1, - idx_diag, - idx_offdiag, - IDX_upper, - IDX_lower, - X0, - X1, - Xq, - S0, - S1, - Sq, - ) = self.ADMM_preprocessing(C) + # Prepare the objective, constraints, block layout, and initial ADMM variables. + pre = self.ADMM_preprocessing(C) + + # Packed objective coefficients and their norm. + C0, C1, normC = pre["C0"], pre["C1"], pre["normC"] + + # Quaternion equality constraints. + AEq, bEq = pre["AEq"], pre["bEq"] + AEqAEqtinv = pre["AEqAEqtinv"] + + # Fejér inequality constraints and spectral scale. + AI_mat_diag = pre["AI_mat_diag"] + AI_mat_offdiag = pre["AI_mat_offdiag"] + bI, Lambda = pre["bI"], pre["Lambda"] + + # Row boundaries by degree: d0 applies to C0/X0/S0, and d1 to C1/X1/S1. + # Their final entries give the total row counts. + d0, d1 = pre["d0"], pre["d1"] + D0, D1 = d0[-1], d1[-1] + + # Indices into packed columns and the full image-pair grid, respectively. + idx_diag, idx_offdiag = pre["idx_diag"], pre["idx_offdiag"] + IDX_upper, IDX_lower = pre["IDX_upper"], pre["IDX_lower"] + + # Initial primal and slack variables. + X0, X1, Xq = pre["X0"], pre["X1"], pre["Xq"] + S0, S1, Sq = pre["S0"], pre["S1"], pre["Sq"] n_pairs = N * (N - 1) // 2 Ngrid = self.Ngrid @@ -1221,7 +1221,7 @@ def ADMM_preprocessing(self, C): :param C: Fourier coefficient matrices of the NUG objective. - :return: Quantities required by the ADMM solver. + :return: Dictionary of quantities required by the ADMM solver. """ Lmax = self.Lmax N = self.n_img @@ -1339,32 +1339,30 @@ def ADMM_preprocessing(self, C): self.Ngrid = Ngrid - return ( - C0, - C1, - normC, - AEq, - bEq, - AEqAEqtinv, - AI_mat_diag, - AI_mat_offdiag, - bI, - Lambda, - d0, - d1, - D0, - D1, - idx_diag, - idx_offdiag, - IDX_upper, - IDX_lower, - X0, - X1, - Xq, - S0, - S1, - Sq, - ) + return { + "C0": C0, + "C1": C1, + "normC": normC, + "AEq": AEq, + "bEq": bEq, + "AEqAEqtinv": AEqAEqtinv, + "AI_mat_diag": AI_mat_diag, + "AI_mat_offdiag": AI_mat_offdiag, + "bI": bI, + "Lambda": Lambda, + "d0": d0, + "d1": d1, + "idx_diag": idx_diag, + "idx_offdiag": idx_offdiag, + "IDX_upper": IDX_upper, + "IDX_lower": IDX_lower, + "X0": X0, + "X1": X1, + "Xq": Xq, + "S0": S0, + "S1": S1, + "Sq": Sq, + } def _compute_fejer_weights(self): """ From 1900546f24d9ce4a5b45534e893cbb01d12a7ac3 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 14 Oct 2026 15:48:26 -0400 Subject: [PATCH 141/161] remove explicit reference to manuscript equation numbers --- src/aspire/abinitio/commonline_nug.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index b8a6bd70a8..70f7732a14 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -2070,7 +2070,7 @@ def find_beta(X1): B1 = np.zeros((N, N), dtype=np.float64) B2 = np.zeros((N, N), dtype=np.float64) - # From equation (36) in the paper, we have + # The cyclic-symmetry degree-one block gives # # 2 * abs(X1_ij[0, 0]) = sin(beta_i) * sin(beta_j) # real(X1_ij[1, 1]) = cos(beta_i) * cos(beta_j). @@ -2103,7 +2103,7 @@ def find_alpha(X1): ZZbar = np.zeros((N, N), dtype=complex_type(np.float64)) ZZ = np.zeros((N, N), dtype=complex_type(np.float64)) - # Equation (36) also gives the pairwise phase measurements + # The same degree-one block also gives the pairwise phase measurements # # X1_ij[0, 0] / abs(X1_ij[0, 0]) # = exp(i * (alpha_i - alpha_j)) @@ -2177,8 +2177,7 @@ def euler_est_Dm(self, X_est): def find_alpha_beta(X2): B1 = np.zeros((N, N), dtype=np.float64) - # Equations (59)-(63) combine two entries of each degree-two block - # to isolate + # Combining two entries of the dihedral degree-two block isolates # # B1[i, j] = cos(beta_i)^2 * cos(beta_j)^2. # @@ -2203,8 +2202,8 @@ def find_alpha_beta(X2): Aminus = np.zeros((N, N), dtype=complex_type(np.float64)) Aplus = np.zeros((N, N), dtype=complex_type(np.float64)) - # Equations (64)-(66) remove the known sin(2*beta) factors from - # two entries of X2_ij, producing the pairwise phases + # Divide two degree-two block entries by their known sin(2*beta) + # factors to recover the pairwise phases: # # Aminus[i, j] = exp(i * (alpha_i - alpha_j)) # Aplus[i, j] = exp(i * (alpha_i + alpha_j)). From cc58640a311d1235a5c9827b955bac27881e1f96 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 15 Oct 2026 11:38:53 -0400 Subject: [PATCH 142/161] remove unused variable --- src/aspire/abinitio/commonline_nug.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 70f7732a14..8d6e0ae2b2 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -377,7 +377,6 @@ def admm_plain_J(self, C, verbose): AEqAEqtinv = pre["AEqAEqtinv"] # Fejér inequality constraints and spectral scale. - AI_mat_diag = pre["AI_mat_diag"] AI_mat_offdiag = pre["AI_mat_offdiag"] bI, Lambda = pre["bI"], pre["Lambda"] From 4527bc56513344b5d0380bbfa8be925cfa16c410 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 15 Oct 2026 14:38:54 -0400 Subject: [PATCH 143/161] build ADMM permutation matrics on CPU --- src/aspire/abinitio/commonline_nug.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 8d6e0ae2b2..47a160aae6 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -805,13 +805,13 @@ def admm_sym_J(self, C, verbose): P = [] for k in range(1, Lmax + 1): dk = 2 * k + 1 - Pk = xp.eye(dk, dtype=np.float64) + Pk = np.eye(dk, dtype=np.float64) for m in range(k): for el in range(k - m): Pk[(m + 2 * el, m + 2 * el + 1), :] = Pk[ (m + 2 * el + 1, m + 2 * el), : ] - P.append(Pk) + P.append(xp.asarray(Pk)) def fun_AE(X0, X1, Xd0, Xd1, Xq): """ From 5f64cab646fb63a76cb3c5e5bd22194f93b8d072 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 15 Oct 2026 15:59:10 -0400 Subject: [PATCH 144/161] Replace dense ADMM permutation matrices with index-based grouping --- src/aspire/abinitio/commonline_nug.py | 39 +++++++++++++++------------ 1 file changed, 22 insertions(+), 17 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 47a160aae6..c0eb93afba 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1711,25 +1711,30 @@ def _permutek(Ak, k, N): :return: The two permuted matrix blocks. """ dk = 2 * k + 1 - Pk = np.eye(dk, dtype=Ak.dtype) - for m in range(k): - for n in range(k - m): - Pk[(m + 2 * n, m + 2 * n + 1), :] = Pk[(m + 2 * n + 1, m + 2 * n), :] - Pk = xp.asarray(Pk) - AkP = ( - xp.kron(xp.eye(N, dtype=Ak.dtype), Pk) - @ Ak - @ xp.kron(xp.eye(N, dtype=Ak.dtype), Pk.T) + + # The real degree-k basis interleaves its k- and (k+1)-coordinate + # components. Put the odd-indexed basis coordinates first. + coordinate_order = xp.asarray( + np.concatenate((np.arange(1, dk, 2), np.arange(0, dk, 2))) ) - Pk = xp.eye(N * dk, dtype=AkP.dtype) - idx = xp.concatenate((xp.arange(dk - k, dk), xp.arange(k + 1))) - for m in range(N - 1): - for n in range(N - 1 - m): - Pk[k * (m + 1) + n * dk : k * (m + 1) + (n + 1) * dk] = Pk[ - k * (m + 1) + n * dk : k * (m + 1) + (n + 1) * dk - ][idx, :] - AkP = Pk @ AkP @ Pk.T + # Reshape into blocks (image i, coordinate i, image j, coordinate j) + # and reorder the coordinate axes. + blocks = Ak.reshape(N, dk, N, dk) + blocks = blocks[:, coordinate_order, :, :] + blocks = blocks[:, :, :, coordinate_order] + AkP = blocks.reshape(N * dk, N * dk) + + # Group the k-coordinate component of every image first, followed by + # the (k + 1)-coordinate component of every image. + image_offsets = dk * np.arange(N)[:, None] + first_component = (image_offsets + np.arange(k)).reshape(-1) + second_component = (image_offsets + np.arange(k, dk)).reshape(-1) + group_order = xp.asarray(np.concatenate((first_component, second_component))) + + # Apply the same permutation to rows and columns. + AkP = AkP[group_order[:, None], group_order[None, :]] + return AkP[: N * k, : N * k], AkP[N * k :, N * k :] @staticmethod From 1e8b668d67df56584b0d6e90cc46b871bb3a8c62 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 16 Oct 2026 09:29:09 -0400 Subject: [PATCH 145/161] Replace dense inverse ADMM permutation matrices with index-based reordering --- src/aspire/abinitio/commonline_nug.py | 38 +++++++++++++-------------- 1 file changed, 18 insertions(+), 20 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index c0eb93afba..1fe3f08803 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1749,26 +1749,24 @@ def _permutek_back(Ak, k, N): :return: Matrix in the original block ordering. """ dk = 2 * k + 1 - Pk = xp.eye(N * dk, dtype=Ak.dtype) - idx = xp.concatenate((xp.arange(dk - k, dk), xp.arange(k + 1))) - for m in range(N - 1): - for n in range(N - 1 - m): - Pk[k * (m + 1) + n * dk : k * (m + 1) + (n + 1) * dk] = Pk[ - k * (m + 1) + n * dk : k * (m + 1) + (n + 1) * dk - ][idx, :] - AkB = Pk.T @ Ak @ Pk - dk = 2 * k + 1 - Pk = np.eye(dk, dtype=Ak.dtype) - for m in range(k): - for n in range(k - m): - Pk[(m + 2 * n, m + 2 * n + 1), :] = Pk[(m + 2 * n + 1, m + 2 * n), :] - Pk = xp.asarray(Pk) - AkB = ( - xp.kron(xp.eye(N, dtype=Ak.dtype), Pk.T) - @ AkB - @ xp.kron(xp.eye(N, dtype=Ak.dtype), Pk) - ) - return AkB + + # Undo the grouping of all k-coordinate components before all + # (k + 1)-coordinate components. + image_offsets = dk * np.arange(N)[:, None] + first_component = (image_offsets + np.arange(k)).reshape(-1) + second_component = (image_offsets + np.arange(k, dk)).reshape(-1) + group_order = np.concatenate((first_component, second_component)) + inverse_group_order = xp.asarray(np.argsort(group_order)) + AkB = Ak[inverse_group_order[:, None], inverse_group_order[None, :]] + + # Undo the within-image ordering of the degree-k basis coordinates. + coordinate_order = np.concatenate((np.arange(1, dk, 2), np.arange(0, dk, 2))) + inverse_coordinate_order = xp.asarray(np.argsort(coordinate_order)) + + blocks = AkB.reshape(N, dk, N, dk) + blocks = blocks[:, inverse_coordinate_order, :, :] + blocks = blocks[:, :, :, inverse_coordinate_order] + return blocks.reshape(N * dk, N * dk) @staticmethod def vec_block(A, N, sz, IDX_upper): From 01da9dd89a54f538876b551293e7555ef77e92ac Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 16 Oct 2026 10:39:35 -0400 Subject: [PATCH 146/161] Improve docstrings for _permutek/_permutek_back --- src/aspire/abinitio/commonline_nug.py | 29 +++++++++++++++++++++------ 1 file changed, 23 insertions(+), 6 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 1fe3f08803..8d8d581bbf 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1702,13 +1702,22 @@ def _transform_coeff_back(self, A0, A1, IDX_upper, IDX_lower, idx_offdiag): @staticmethod def _permutek(Ak, k, N): """ - Permute and split a degree-k matrix into blocks of sizes k and k + 1. + Separate the two coordinate groups of a degree-k matrix. - :param Ak: Degree-k block matrix. + Each image contributes 2k + 1 coordinates. Within an image, the + k-coordinate group occupies odd indices and the k + 1 group occupies + even indices. Reorder both axes within each image, then place the + k-coordinate group from every image before the (k + 1)-coordinate + group. Return the two diagonal blocks so they can be processed + separately by ADMM. Cross-group blocks are not returned. + + :param Ak: Matrix in image-wise coordinate order, with shape + (N * (2k + 1), N * (2k + 1)). :param k: Representation degree. :param N: Number of images. - :return: The two permuted matrix blocks. + :return: The k-component and (k + 1)-component blocks, with shapes + (N*k, N*k) and (N*(k + 1), N*(k + 1)). """ dk = 2 * k + 1 @@ -1740,13 +1749,21 @@ def _permutek(Ak, k, N): @staticmethod def _permutek_back(Ak, k, N): """ - Undo the degree-k block permutation and reconstruct the full matrix. + Restore image-wise coordinate order after ADMM block processing. + + The input is in component-grouped order: all N*k coordinates of + the first group precede all N*(k + 1) coordinates of the second. + `_transform_coeff_back` places the two processed blocks on the + diagonal of this matrix, with zeros elsewhere. This method then + restores the original image-wise coordinate order. - :param Ak: Permuted degree-k matrix. + :param Ak: Full matrix in component-grouped order, with shape + (N * (2k + 1), N * (2k + 1)). :param k: Representation degree. :param N: Number of images. - :return: Matrix in the original block ordering. + :return: Matrix in image-wise coordinate order, with the same shape + as Ak. """ dk = 2 * k + 1 From ff808d17727228d852f666383fb2ae2f08c9b704 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 16 Oct 2026 15:37:41 -0400 Subject: [PATCH 147/161] Replace _largest_eigenvalue power method with eigvalsh. Use 3.58GiB less memory. --- src/aspire/abinitio/commonline_nug.py | 32 ++++++++------------------- 1 file changed, 9 insertions(+), 23 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 8d8d581bbf..04c382ce0d 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1317,7 +1317,7 @@ def ADMM_preprocessing(self, C): bI = -(Lmax + 2) * (Lmax + 1) / 2 # largest eigenvalue for AIAIT - Lambda = self._largest_eigenvalue(AI_mat_offdiag, N) + Lambda = self._largest_eigenvalue_gram(AI_mat_offdiag) # Initialize the primal variables in packed image-pair form. n_packed = N * (N + 1) // 2 @@ -1437,33 +1437,19 @@ def _discretize_SO3(self): return SO3 @staticmethod - def _largest_eigenvalue(AI, N): + def _largest_eigenvalue_gram(AI): """ - Estimate the largest eigenvalue of the Fejér constraint operator, - AI @ AI.T, using power iteration. + Compute the denominator for the Fejér inequality-multiplier update. - A fixed empirical margin is added to the estimate, following the - researcher's implementation, to obtain a more conservative step - size for the inequality-multiplier update. + The largest eigenvalue of AI @ AI.T is computed from the smaller + AI.T @ AI Gram matrix. Retain the fixed 2000 offset used by the + original implementation. - :param AI: Matrix representation of the Fejer inequality operator, - with shape (Ngrid, n_coefficients). - :param N: Number of images. + :param AI: Real Fejér inequality operator of shape (Ngrid, D0 + D1). - :return: Estimated largest eigenvalue of AI @ AI.T. + :return: Largest eigenvalue of AI @ AI.T, plus 2000. """ - # find the largest eigenvalue of the operator AI - n_grid = AI.shape[0] - z = xp.random.normal(0, 1, (n_grid, N**2)) - Lambda = 0 - - while abs(Lambda - xp.linalg.norm(z)) > 500: - Lambda = xp.linalg.norm(z) - z = z / xp.linalg.norm(z) - z = AI @ (AI.T @ z) - Lambda += 2000 - logger.info(f"Largest eigenvalue of AIAIT is approximately {Lambda:.2f}") - return Lambda + return xp.linalg.eigvalsh(AI.T @ AI)[-1] + 2000 def _compute_rank(self): """ From bb3bf5d8712af1323272eb3b9c2aff09c9a0545c Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 19 Oct 2026 11:21:16 -0400 Subject: [PATCH 148/161] vectorize pairwise loss function (fij) --- src/aspire/abinitio/commonline_nug.py | 64 +++++++++++++-------------- 1 file changed, 32 insertions(+), 32 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 04c382ce0d..a3bccc916f 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -197,39 +197,45 @@ def compute_coeff(self): Lmax = self.Lmax T = self.T - def fij(alpha, gamma, i, j): + # Quadrature grid used to sample the pairwise loss over SO(3) Euler angles. + alpha_grid = np.arange(2 * T, dtype=np.float64) * np.pi / T + beta_grid = (2 * np.arange(2 * T, dtype=np.float64) + 1) * np.pi / 4 / T + gamma_grid = np.arange(2 * T, dtype=np.float64) * np.pi / T + + # Map the sampled angles to polar Fourier ray indices once for all image pairs. + alpha_idx = ( + np.round((alpha_grid - np.pi / 2) * n_theta / 2 / np.pi) % n_theta + ).astype(int) + gamma_idx = ( + np.round((-gamma_grid - np.pi / 2) * n_theta / 2 / np.pi) % n_theta + ).astype(int) + + def fij(i, j): """ - Evaluate the pairwise common-line loss used to approximate NUG Fourier coefficients. + Evaluate the pairwise common-line loss on the alpha–gamma grid. - This function samples the corresponding polar Fourier rays from images i and j, - compares them over all candidate 1D shifts, and returns the minimum shifted - L1 mismatch. + For each sampled pair of angles, select the corresponding polar Fourier + rays from images i and j. Compare the rays using the L1 norm over radial + samples, then retain the smallest mismatch over the candidate 1D shifts. - :param alpha: First ZYZ Euler angle of the relative orientation, in radians. - :param gamma: Third ZYZ Euler angle of the relative orientation, in radians. :param i: Index of the first image. :param j: Index of the second image. - :return: Minimum shifted L1 mismatch between the induced common-line rays. + :return: Loss array of shape (2*T, 2*T). Rows correspond to alpha_grid + and columns to gamma_grid. """ + rays_i = self.pf_full[i, alpha_idx] + rays_j = self.pf_full[j, gamma_idx] + loss_grid = np.empty((2 * T, 2 * T), dtype=np.float64) - Ii_hat = self.pf_full[i] - Ij_hat = self.pf_full[j] - idxi = np.round((alpha - np.pi / 2) * n_theta / 2 / np.pi) % n_theta - idxj = np.round((-gamma - np.pi / 2) * n_theta / 2 / np.pi) % n_theta + for gamma_col, ray_j in enumerate(rays_j): + shifted_j = self.shift_phases * ray_j + differences = rays_i[:, None, :] - shifted_j[None, :, :] + loss_grid[:, gamma_col] = np.linalg.norm( + differences, ord=1, axis=-1 + ).min(axis=-1) - Si = Ii_hat[int(idxi)] - Sj = Ij_hat[int(idxj)] - - # Apply shifts - Sj_shifted = self.shift_phases * Sj - norms = np.linalg.norm(Si[None] - Sj_shifted, 1, axis=1) - return norms.min() - - # Quadrature grid used to sample the pairwise loss over SO(3) Euler angles. - alpha_grid = np.arange(2 * T, dtype=np.float64) * np.pi / T - beta_grid = (2 * np.arange(2 * T, dtype=np.float64) + 1) * np.pi / 4 / T - gamma_grid = np.arange(2 * T, dtype=np.float64) * np.pi / T + return loss_grid # Beta quadrature weights for the truncated Wigner-D Fourier expansion. bT = np.zeros(2 * T, dtype=np.float64) @@ -278,10 +284,7 @@ def fijhat_k(k_idx, F): ) as pbar: for i in range(N): for j in range(i + 1, N): - Fij = np.zeros((2 * T, 2 * T), dtype=np.float64) - for j1 in range(2 * T): - for j2 in range(2 * T): - Fij[j1, j2] = fij(alpha_grid[j1], gamma_grid[j2], i, j) + Fij = fij(i, j) for k_idx in range(Lmax): k = k_idx + 1 dk = 2 * k + 1 @@ -296,10 +299,7 @@ def fijhat_k(k_idx, F): # Diagonal blocks encode self-pair losses and are handled separately. for i in trange(N, desc="Computing diagonal coefficients"): - Fii = np.zeros((2 * T, 2 * T), dtype=np.float64) - for j1 in range(2 * T): - for j2 in range(2 * T): - Fii[j1, j2] = fij(alpha_grid[j1], gamma_grid[j2], i, i) + Fii = fij(i, i) for k_idx in range(Lmax): k = k_idx + 1 dk = 2 * k + 1 From e73f7bbf64765feff545073f36485418fa2e5e0d Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 19 Oct 2026 13:46:36 -0400 Subject: [PATCH 149/161] extract nested functions from rotation recovery main methods. --- src/aspire/abinitio/commonline_nug.py | 333 +++++++++++++++----------- 1 file changed, 188 insertions(+), 145 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index a3bccc916f..007637e044 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -2071,78 +2071,9 @@ def euler_est_Cm(self, X1, XS): X1 = self._real_to_complex_representation(X1, degree=1) XS = self._real_to_complex_representation(XS, degree=S) - def find_beta(X1): - B1 = np.zeros((N, N), dtype=np.float64) - B2 = np.zeros((N, N), dtype=np.float64) - - # The cyclic-symmetry degree-one block gives - # - # 2 * abs(X1_ij[0, 0]) = sin(beta_i) * sin(beta_j) - # real(X1_ij[1, 1]) = cos(beta_i) * cos(beta_j). - # - # Therefore, B1 and B2 are approximately rank-one outer-product - # matrices containing the sine and cosine factors, respectively. - for i in range(N): - for j in range(N): - Xij = X1[3 * i : 3 * (i + 1), 3 * j : 3 * (j + 1)] - B1[i, j] = abs(Xij[0, 0]) * 2 - B2[i, j] = np.real(Xij[1, 1]) - - # Recover vectors approximating sin(beta_i) and cos(beta_i) from - # the leading eigenpairs of their rank-one Gram matrices. - e1, v1 = np.linalg.eigh(B1) - idx = np.argmax(e1) - b1 = -v1[:, idx] * np.sqrt(e1[idx]) - - e2, v2 = np.linalg.eigh(B2) - idx = np.argmax(e2) - b2 = v2[:, idx] * np.sqrt(e2[idx]) - - # Combine the synchronized sine and cosine factors to recover beta. - # The eigenvector sign ambiguity corresponds to the global - # beta <-> pi - beta ambiguity described in the paper. - beta = np.arctan(b1 / b2) % np.pi - return beta - - def find_alpha(X1): - ZZbar = np.zeros((N, N), dtype=complex_type(np.float64)) - ZZ = np.zeros((N, N), dtype=complex_type(np.float64)) - - # The same degree-one block also gives the pairwise phase measurements - # - # X1_ij[0, 0] / abs(X1_ij[0, 0]) - # = exp(i * (alpha_i - alpha_j)) - # - # -X1_ij[0, 2] / abs(X1_ij[0, 2]) - # = exp(i * (alpha_i + alpha_j)). - # - # Store these as the difference-phase and sum-phase matrices. - for i in range(N): - for j in range(N): - Xij = X1[ - 3 * i : 3 * (i + 1), - 3 * j : 3 * (j + 1), - ] - - z = Xij[0, 0] - ZZbar[i, j] = z / abs(z) - - z = Xij[0, 2] - ZZ[i, j] = -z / abs(z) - - # Angular synchronization of the difference phases recovers - # exp(i * alpha_i) up to a common phase. - evals, evecs = np.linalg.eigh(ZZbar) - idx = np.argmax(abs(evals)) - Z = evecs[:, idx] * np.sqrt(abs(evals[idx])) - - # Use the sum phases to resolve the remaining common phase. - c = self._find_phase(Z[:, None] @ Z[:, None].T, ZZ) - Z = np.sqrt(c) * Z - return np.angle(Z).astype(np.float64) - - alpha_est = find_alpha(X1) - beta_est = find_beta(X1) + # Estimate alpha/beta Euler angles. + alpha_est = self._find_alpha(X1) + beta_est = self._find_beta(X1) # Estimate the remaining Euler angle from the degree-S solution. # The cyclic pair estimator extracts measurements of @@ -2179,79 +2110,8 @@ def euler_est_Dm(self, X_est): X2 = self._real_to_complex_representation(X_est[1], degree=2) XS = self._real_to_complex_representation(X_est[S - 1], degree=S) - def find_alpha_beta(X2): - B1 = np.zeros((N, N), dtype=np.float64) - - # Combining two entries of the dihedral degree-two block isolates - # - # B1[i, j] = cos(beta_i)^2 * cos(beta_j)^2. - # - # Thus, B1 is approximately a rank-one outer-product matrix. - for i in range(N): - for j in range(N): - Xij = X2[5 * i : 5 * (i + 1), 5 * j : 5 * (j + 1)] - B1[i, j] = np.real((Xij[2, 2] - 2 * abs(Xij[0, 0]) + 0.5) / 3 * 2) - - # Recover the vector cos(beta_i)^2 from the leading eigenpair. - # Its entries should be nonnegative, so orient the eigenvector - # using the sign of its first entry. - e1, v1 = np.linalg.eigh(B1) - idx = np.argmax(e1) - b1 = v1[:, idx] * np.sqrt(e1[idx]) * np.sign(v1[0, idx]) - - # Taking the square root recovers abs(cos(beta_i)). Consequently, - # each image retains an independent beta_i <-> pi - beta_i - # ambiguity, as described in the paper. - beta_est = np.arccos(np.clip(np.sqrt(b1), -1, 1)) % np.pi - - Aminus = np.zeros((N, N), dtype=complex_type(np.float64)) - Aplus = np.zeros((N, N), dtype=complex_type(np.float64)) - - # Divide two degree-two block entries by their known sin(2*beta) - # factors to recover the pairwise phases: - # - # Aminus[i, j] = exp(i * (alpha_i - alpha_j)) - # Aplus[i, j] = exp(i * (alpha_i + alpha_j)). - for i in range(N): - for j in range(N): - Xij = X2[5 * i : 5 * (i + 1), 5 * j : 5 * (j + 1)] - - # Avoid division by a numerically vanishing - # sin(2*beta) factor. - if abs(beta_est[i]) < 1e-6: - beta_est[i] = 1e-6 - if abs(beta_est[j]) < 1e-6: - beta_est[j] = 1e-6 - Aminus[i, j] = ( - Xij[1, 1] - / np.sin(2 * beta_est[i]) - / np.sin(2 * beta_est[j]) - * 8 - / 3 - ) - Aplus[i, j] = ( - -Xij[1, 3] - / np.sin(2 * beta_est[i]) - / np.sin(2 * beta_est[j]) - * 8 - / 3 - ) - - # Angular synchronization of the difference phases recovers - # exp(i * alpha_i) up to a common phase. - evals, evecs = np.linalg.eigh(Aminus) - idx = np.argmax(abs(evals)) - Z = evecs[:, idx] * np.sqrt(abs(evals[idx])) - - # Use the sum phases to resolve the remaining common phase. - # The square root leaves a global pi ambiguity. - c = self._find_phase(Z[:, None] @ Z[:, None].T, Aplus) - Z = np.sqrt(c) * Z - alpha_est = (np.angle(Z)) % (2 * np.pi) - - return alpha_est, beta_est - - alpha_est, beta_est = find_alpha_beta(X2) + # Estimate alpha/beta Euler angles. + alpha_est, beta_est = self._find_alpha_beta(X2) # Estimate gamma from the degree-S solution. The dihedral pair # estimator separates the gamma-difference and gamma-sum terms, @@ -2300,6 +2160,189 @@ def _find_phase(A, B): return c / abs(c) + def _find_beta(self, X1): + """ + Estimate beta angles for cyclic symmetry from the degree-one solution. + + Entries of each image pair block provide approximate outer products of + sin(beta) and cos(beta). Recover those factors from their leading + eigenvectors, then combine them to estimate one beta angle per image. + + :param X1: Degree-one representation matrix in the complex Wigner basis, + with shape (3 * n_img, 3 * n_img). + + :return: Estimated beta angles in radians, with shape (n_img,). + """ + N = self.n_img + B1 = np.zeros((N, N), dtype=np.float64) + B2 = np.zeros((N, N), dtype=np.float64) + + # The cyclic symmetry degree-one block gives + # + # 2 * abs(X1_ij[0, 0]) = sin(beta_i) * sin(beta_j) + # real(X1_ij[1, 1]) = cos(beta_i) * cos(beta_j). + # + # Therefore, B1 and B2 are approximately rank-one outer-product + # matrices containing the sine and cosine factors, respectively. + for i in range(N): + for j in range(N): + Xij = X1[3 * i : 3 * (i + 1), 3 * j : 3 * (j + 1)] + B1[i, j] = abs(Xij[0, 0]) * 2 + B2[i, j] = np.real(Xij[1, 1]) + + # Recover vectors approximating sin(beta_i) and cos(beta_i) from + # the leading eigenpairs of their rank-one Gram matrices. + e1, v1 = np.linalg.eigh(B1) + idx = np.argmax(e1) + b1 = -v1[:, idx] * np.sqrt(e1[idx]) + + e2, v2 = np.linalg.eigh(B2) + idx = np.argmax(e2) + b2 = v2[:, idx] * np.sqrt(e2[idx]) + + # Combine the synchronized sine and cosine factors to recover beta. + # The eigenvector sign ambiguity corresponds to the global + # beta <-> pi - beta ambiguity described in the paper. + beta = np.arctan(b1 / b2) % np.pi + return beta + + def _find_alpha(self, X1): + """ + Estimate alpha angles for cyclic symmetry from the degree-one solution. + + Extract pairwise alpha difference and sum phases from the image pair + blocks. Synchronize the difference phases and use the sum phases to + resolve the remaining common phase. + + :param X1: Degree-one representation matrix in the complex Wigner basis, + with shape (3 * n_img, 3 * n_img). + + :return: Estimated alpha angles in radians, with shape (n_img,). + """ + N = self.n_img + ZZbar = np.zeros((N, N), dtype=complex_type(np.float64)) + ZZ = np.zeros((N, N), dtype=complex_type(np.float64)) + + # The same degree-one block also gives the pairwise phase measurements + # + # X1_ij[0, 0] / abs(X1_ij[0, 0]) + # = exp(i * (alpha_i - alpha_j)) + # + # -X1_ij[0, 2] / abs(X1_ij[0, 2]) + # = exp(i * (alpha_i + alpha_j)). + # + # Store these as the difference-phase and sum-phase matrices. + for i in range(N): + for j in range(N): + Xij = X1[ + 3 * i : 3 * (i + 1), + 3 * j : 3 * (j + 1), + ] + + z = Xij[0, 0] + ZZbar[i, j] = z / abs(z) + + z = Xij[0, 2] + ZZ[i, j] = -z / abs(z) + + # Angular synchronization of the difference phases recovers + # exp(i * alpha_i) up to a common phase. + evals, evecs = np.linalg.eigh(ZZbar) + idx = np.argmax(abs(evals)) + Z = evecs[:, idx] * np.sqrt(abs(evals[idx])) + + # Use the sum phases to resolve the remaining common phase. + c = self._find_phase(Z[:, None] @ Z[:, None].T, ZZ) + Z = np.sqrt(c) * Z + return np.angle(Z).astype(np.float64) + + def _find_alpha_beta(self, X2): + """ + Estimate alpha and beta angles for dihedral symmetry. + + Use the degree-two image pair blocks to estimate beta from an approximate + rank-one matrix. Then extract pairwise alpha difference and sum phases, + synchronize them, and resolve the remaining common phase. The beta + estimates retain the ambiguity inherent in the dihedral data. + + :param X2: Degree-two representation matrix in the complex Wigner basis, + with shape (5 * n_img, 5 * n_img). + + :return: Tuple of estimated alpha and beta angle arrays in radians; + each has shape (n_img,). + """ + N = self.n_img + B1 = np.zeros((N, N), dtype=np.float64) + + # Combining two entries of the dihedral degree-two block isolates + # + # B1[i, j] = cos(beta_i)^2 * cos(beta_j)^2. + # + # Thus, B1 is approximately a rank-one outer-product matrix. + for i in range(N): + for j in range(N): + Xij = X2[5 * i : 5 * (i + 1), 5 * j : 5 * (j + 1)] + B1[i, j] = np.real((Xij[2, 2] - 2 * abs(Xij[0, 0]) + 0.5) / 3 * 2) + + # Recover the vector cos(beta_i)^2 from the leading eigenpair. + # Its entries should be nonnegative, so orient the eigenvector + # using the sign of its first entry. + e1, v1 = np.linalg.eigh(B1) + idx = np.argmax(e1) + b1 = v1[:, idx] * np.sqrt(e1[idx]) * np.sign(v1[0, idx]) + + # Taking the square root recovers abs(cos(beta_i)). Consequently, + # each image retains an independent beta_i <-> pi - beta_i + # ambiguity, as described in the paper. + beta_est = np.arccos(np.clip(np.sqrt(b1), -1, 1)) % np.pi + + Aminus = np.zeros((N, N), dtype=complex_type(np.float64)) + Aplus = np.zeros((N, N), dtype=complex_type(np.float64)) + + # Divide two degree-two block entries by their known sin(2*beta) + # factors to recover the pairwise phases: + # + # Aminus[i, j] = exp(i * (alpha_i - alpha_j)) + # Aplus[i, j] = exp(i * (alpha_i + alpha_j)). + for i in range(N): + for j in range(N): + Xij = X2[5 * i : 5 * (i + 1), 5 * j : 5 * (j + 1)] + + # Avoid division by a numerically vanishing + # sin(2*beta) factor. + if abs(beta_est[i]) < 1e-6: + beta_est[i] = 1e-6 + if abs(beta_est[j]) < 1e-6: + beta_est[j] = 1e-6 + Aminus[i, j] = ( + Xij[1, 1] + / np.sin(2 * beta_est[i]) + / np.sin(2 * beta_est[j]) + * 8 + / 3 + ) + Aplus[i, j] = ( + -Xij[1, 3] + / np.sin(2 * beta_est[i]) + / np.sin(2 * beta_est[j]) + * 8 + / 3 + ) + + # Angular synchronization of the difference phases recovers + # exp(i * alpha_i) up to a common phase. + evals, evecs = np.linalg.eigh(Aminus) + idx = np.argmax(abs(evals)) + Z = evecs[:, idx] * np.sqrt(abs(evals[idx])) + + # Use the sum phases to resolve the remaining common phase. + # The square root leaves a global pi ambiguity. + c = self._find_phase(Z[:, None] @ Z[:, None].T, Aplus) + Z = np.sqrt(c) * Z + alpha_est = (np.angle(Z)) % (2 * np.pi) + + return alpha_est, beta_est + @staticmethod def _assemble_rotation_estimates(alpha, beta, gamma): """ From 386362aee593c36a35c3c3f3df9c32a35dfc40d1 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 19 Oct 2026 14:35:12 -0400 Subject: [PATCH 150/161] replace complex_type() with np.complex128 --- src/aspire/abinitio/commonline_nug.py | 23 +++++++++++------------ 1 file changed, 11 insertions(+), 12 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 007637e044..b5af85442e 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -258,11 +258,11 @@ def fijhat_k(k_idx, F): k = k_idx + 1 # degree dk = 2 * k + 1 - exp_alpha_grid = np.zeros((2 * T, dk), dtype=complex_type(np.float64)) + exp_alpha_grid = np.zeros((2 * T, dk), dtype=np.complex128) for m in range(-k, k + 1): exp_alpha_grid[:, m + k] = np.exp(1j * m * alpha_grid) - exp_gamma_grid = np.zeros((2 * T, dk), dtype=complex_type(np.float64)) + exp_gamma_grid = np.zeros((2 * T, dk), dtype=np.complex128) for m in range(-k, k + 1): exp_gamma_grid[:, m + k] = np.exp(1j * m * gamma_grid) @@ -275,7 +275,7 @@ def fijhat_k(k_idx, F): C = [] for k in range(1, Lmax + 1): dk = 2 * k + 1 - C.append(np.zeros((N * dk, N * dk), dtype=complex_type(np.float64))) + C.append(np.zeros((N * dk, N * dk), dtype=np.complex128)) # Compute off-diagonal image-pair losses and insert their degree-wise coefficients. with tqdm( @@ -1033,7 +1033,7 @@ def update_X( resX0 = xp.linalg.norm(tmp) tmp = S1 + Z1 + AIT_yI1 - C1 X1 = X1 + mult * rho * tmp - resX1 = xp.linalg.norm(tmp) + resX1 = Xp.linalg.norm(tmp) tmp = Sd0 + Zd0 Xd0 = Xd0 + mult * rho * tmp resXd0 = xp.linalg.norm(tmp) @@ -1256,7 +1256,6 @@ def ADMM_preprocessing(self, C): C0, C1 = self._transform_coeff(C, IDX_upper) normC = xp.sqrt(xp.linalg.norm(C0) ** 2 + xp.linalg.norm(C1) ** 2) - del C # compute the block sizes for X d0 = [0] @@ -2220,8 +2219,8 @@ def _find_alpha(self, X1): :return: Estimated alpha angles in radians, with shape (n_img,). """ N = self.n_img - ZZbar = np.zeros((N, N), dtype=complex_type(np.float64)) - ZZ = np.zeros((N, N), dtype=complex_type(np.float64)) + ZZbar = np.zeros((N, N), dtype=np.complex128) + ZZ = np.zeros((N, N), dtype=np.complex128) # The same degree-one block also gives the pairwise phase measurements # @@ -2296,8 +2295,8 @@ def _find_alpha_beta(self, X2): # ambiguity, as described in the paper. beta_est = np.arccos(np.clip(np.sqrt(b1), -1, 1)) % np.pi - Aminus = np.zeros((N, N), dtype=complex_type(np.float64)) - Aplus = np.zeros((N, N), dtype=complex_type(np.float64)) + Aminus = np.zeros((N, N), dtype=np.complex128) + Aplus = np.zeros((N, N), dtype=np.complex128) # Divide two degree-two block entries by their known sin(2*beta) # factors to recover the pairwise phases: @@ -2392,7 +2391,7 @@ def _estimate_gamma( modes = np.arange(-degree, degree + 1) # Store pairwise estimates of exp(i * degree * (gamma_i - gamma_j)). - C = np.zeros((N, N), dtype=complex_type(np.float64)) + C = np.zeros((N, N), dtype=np.complex128) # Construct the handedness-conjugation matrix. Jk = self._handedness_matrix(degree) @@ -2684,7 +2683,7 @@ def _complex2real(ell): :return: Forward and inverse change-of-basis matrices. """ diml = 2 * ell + 1 - Tinv = np.zeros((diml, diml), dtype=complex_type(np.float64)) + Tinv = np.zeros((diml, diml), dtype=np.complex128) for i in range(diml): if i < ell: Tinv[i, i] = 1j / np.sqrt(2) @@ -2695,7 +2694,7 @@ def _complex2real(ell): Tinv[i, i] = (-1) ** (i - ell) / np.sqrt(2) Tinv[i, diml - 1 - i] = 1 / np.sqrt(2) - T = np.zeros((diml, diml), dtype=complex_type(np.float64)) + T = np.zeros((diml, diml), dtype=np.complex128) for i in range(diml): if i < ell: T[i, i] = -1j / np.sqrt(2) From 258cf4493e20e752b172baf4cf57999ffdafd603 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Mon, 19 Oct 2026 14:38:59 -0400 Subject: [PATCH 151/161] typo --- src/aspire/abinitio/commonline_nug.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index b5af85442e..0132d6db3d 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1033,7 +1033,7 @@ def update_X( resX0 = xp.linalg.norm(tmp) tmp = S1 + Z1 + AIT_yI1 - C1 X1 = X1 + mult * rho * tmp - resX1 = Xp.linalg.norm(tmp) + resX1 = xp.linalg.norm(tmp) tmp = Sd0 + Zd0 Xd0 = Xd0 + mult * rho * tmp resXd0 = xp.linalg.norm(tmp) From 68ce9713ae603dcfae41d40ee342c2e95341105d Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Thu, 22 Oct 2026 14:58:14 -0400 Subject: [PATCH 152/161] remove unused import and variables --- src/aspire/abinitio/commonline_nug.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 0132d6db3d..056c189ed4 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -7,7 +7,7 @@ from aspire.abinitio.sync_voting import _syncrotations from aspire.numeric import xp from aspire.operators import PolarFT -from aspire.utils import Rotation, cart2sph, complex_type, tqdm, trange +from aspire.utils import Rotation, cart2sph, tqdm, trange from aspire.volume import ( CnSymmetryGroup, DnSymmetryGroup, @@ -2063,7 +2063,6 @@ def euler_est_Cm(self, X1, XS): :return: Estimated rotation matrices. """ S = self.n_sym - N = self.n_img # Convert the relaxed NUG matrices from the real representation used by # the SDP solver to the complex Wigner-D basis used in the recovery formulas. @@ -2102,7 +2101,6 @@ def euler_est_Dm(self, X_est): :return: Estimated rotation matrices and Euler angles. """ S = self.sym_grp.order - N = self.n_img # Convert the degree-two and degree-S solutions from the real # representation used by the SDP solver to the complex Wigner-D basis. From 16a915fe283da6ffbc76abbace6862f3b6580ede Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 27 Oct 2026 10:11:34 -0400 Subject: [PATCH 153/161] Batch C1 Fejer inequality operations to reduce GPU memory use --- src/aspire/abinitio/commonline_nug.py | 143 +++++++++++++++++++++----- 1 file changed, 117 insertions(+), 26 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 056c189ed4..648932f04b 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -181,6 +181,11 @@ def estimate_rotations(self): self.compute_coeff() self.perform_admm() self.recover_rotations() + + # Return unused CuPy allocations before downstream reconstruction. + if hasattr(xp, "get_default_memory_pool"): + xp.get_default_memory_pool().free_all_blocks() + return self.rotations ####################### @@ -352,6 +357,10 @@ def admm_plain_J(self, C, verbose): """ Solve the asymmetric NUG semidefinite relaxation using ADMM. + To limit peak GPU memory use, Fejer inequality operations process + off-diagonal image pairs in chunks, and large temporary arrays are + released before memory intensive projections. + :param C: Fourier coefficient matrices of the NUG objective. :param verbose: Whether to log ADMM progress. @@ -376,7 +385,7 @@ def admm_plain_J(self, C, verbose): AEq, bEq = pre["AEq"], pre["bEq"] AEqAEqtinv = pre["AEqAEqtinv"] - # Fejér inequality constraints and spectral scale. + # Fejer inequality constraints and spectral scale. AI_mat_offdiag = pre["AI_mat_offdiag"] bI, Lambda = pre["bI"], pre["Lambda"] @@ -393,10 +402,24 @@ def admm_plain_J(self, C, verbose): X0, X1, Xq = pre["X0"], pre["X1"], pre["Xq"] S0, S1, Sq = pre["S0"], pre["S1"], pre["Sq"] + # Safe to remove pre now that it is unpacked. + del pre + Ngrid = self.Ngrid n_pairs = N * (N - 1) // 2 n_packed = N * (N + 1) // 2 + # Process off-diagonal image pairs in chunks to limit temporary GPU memory + # used by the Fejer inequality operators. + pair_chunk_size = 1024 + + def pair_chunks(): + """ + Yield bounded slices of the off-diagonal image-pair columns. + """ + for start in range(0, n_pairs, pair_chunk_size): + yield slice(start, min(start + pair_chunk_size, n_pairs)) + # In the asymmetric formulation, diagonal representation blocks are identity. bE0 = xp.zeros(D0, dtype=np.float64) bE1 = xp.zeros(D1, dtype=np.float64) @@ -436,24 +459,52 @@ def fun_AET(yE0, yE1, yEq): Z1[:4, idx_offdiag] = Zq[17:] return Z0, Z1, Zq[:16] - def fun_AI(X0, X1): + def fun_AI(X0_chunk, X1_chunk): + """ + Apply the Fejer operator to one chunk of off-diagonal pairs. """ - Apply the Fejer inequality operator to off-diagonal image pairs. + coefficients = xp.concatenate((X0_chunk, X1_chunk), axis=0) + return AI_mat_offdiag @ coefficients - The plain formulation does not apply these inequalities to the fixed - diagonal representation blocks. + def fun_AIT_chunk(yI_chunk): """ - tmp = xp.concatenate((X0[:, idx_offdiag], X1[:, idx_offdiag]), axis=0) - return AI_mat_offdiag @ tmp + Apply the Fejer adjoint to one chunk of off-diagonal pairs. + """ + return AI_mat_offdiag.T @ yI_chunk def fun_AIT(yI): """ - Apply the adjoint of the asymmetric Fejer operator. + Assemble the Fejer adjoint in packed image-pair order. """ Z = xp.zeros((D0 + D1, n_packed), dtype=np.float64) - Z[:, idx_offdiag] = AI_mat_offdiag.T @ yI + + for columns in pair_chunks(): + packed_columns = idx_offdiag[columns] + Z[:, packed_columns] = fun_AIT_chunk(yI[:, columns]) + return Z[:D0], Z[D0:] + def fejer_residual_norm(X0, X1): + """ + Compute the norm of the positive Fejer inequality residual in chunks. + This avoids storing a residual over the full Fejer grid. + """ + squared_norm = xp.zeros((), dtype=np.float64) + + for columns in pair_chunks(): + packed_columns = idx_offdiag[columns] + residual = xp.maximum( + bI + - fun_AI( + X0[:, packed_columns], + X1[:, packed_columns], + ), + 0, + ) + squared_norm += xp.sum(residual * residual) + + return xp.sqrt(squared_norm) + def update_S(C0, C1, yE0, yE1, yEq, yI, X0, X1, Xq, rho): """ Update and project the asymmetric PSD slack variables. @@ -461,10 +512,26 @@ def update_S(C0, C1, yE0, yE1, yEq, yI, X0, X1, Xq, rho): Z0, Z1, Zq = fun_AET(yE0, yE1, yEq) AIT_yI0, AIT_yI1 = fun_AIT(yI) - S0 = C0 - Z0 - AIT_yI0 - X0 / rho - S1 = C1 - Z1 - AIT_yI1 - X1 / rho + S0 = C0.copy() + S0 -= Z0 + S0 -= AIT_yI0 + xp.divide(X0, rho, out=Z0) # Z0 is no longer needed. + S0 -= Z0 + + S1 = C1.copy() + S1 -= Z1 + S1 -= AIT_yI1 + xp.divide(X1, rho, out=Z1) # Z1 is no longer needed. + S1 -= Z1 + + # S0 = C0 - Z0 - AIT_yI0 - X0 / rho + # S1 = C1 - Z1 - AIT_yI1 - X1 / rho Sq = -Zq - Xq / rho + # These full size adjoint and equality arrays are no longer needed during + # the PSD eigendecompositions. Release them for the eigensolver workspace. + del Z0, Z1, AIT_yI0, AIT_yI1 + S0, S1 = self._project_representation_blocks( S0, S1, @@ -475,8 +542,8 @@ def update_S(C0, C1, yE0, yE1, yEq, yI, X0, X1, Xq, rho): idx_offdiag, ) - # This is the established batched equivalent of projecting each 4x4 - # quaternion block and storing tmp.T.reshape(16) in each column. + # This is the batched equivalent of projecting each 4x4 quaternion block + # and storing tmp.T.reshape(16) in each column. Sq = self._psd_projection(Sq.T.reshape(n_pairs, 4, 4)) Sq = Sq.T.reshape(16, n_pairs) return S0, S1, Sq @@ -511,16 +578,35 @@ def update_yI( Lambda, ): """ - Update the nonnegative Fejer inequality multipliers. + Update the nonnegative Fejer multipliers in pair-column chunks. + Each chunk is updated in place to limit temporary GPU allocations. """ AET_yE0, AET_yE1, _ = fun_AET(yE0, yE1, yEq) - AIT_yI0, AIT_yI1 = fun_AIT(yI) - tmp = fun_AI( - -X0 / rho + C0 - S0 - AIT_yI0 - AET_yE0, - -X1 / rho + C1 - S1 - AIT_yI1 - AET_yE1, - ) - yI = yI + bI / rho / Lambda + tmp / Lambda - yI = xp.maximum(yI, 0) + + for columns in pair_chunks(): + packed_columns = idx_offdiag[columns] + AIT_yI = fun_AIT_chunk(yI[:, columns]) + + tmp0 = ( + -X0[:, packed_columns] / rho + + C0[:, packed_columns] + - S0[:, packed_columns] + - AIT_yI[:D0] + - AET_yE0[:, packed_columns] + ) + tmp1 = ( + -X1[:, packed_columns] / rho + + C1[:, packed_columns] + - S1[:, packed_columns] + - AIT_yI[D0:] + - AET_yE1[:, packed_columns] + ) + + yI_chunk = yI[:, columns] + yI_chunk += bI / rho / Lambda + yI_chunk += fun_AI(tmp0, tmp1) / Lambda + xp.maximum(yI_chunk, 0, out=yI_chunk) + return yI def update_X( @@ -569,7 +655,7 @@ def update_rho(X0, X1, Xq, res_X, rho): + xp.linalg.norm(bE1 - z1) / (1 + xp.linalg.norm(bE1)) + xp.linalg.norm(bEq - zq) / (1 + xp.linalg.norm(bEq)) ) - res_inq = xp.linalg.norm(xp.maximum(bI - fun_AI(X0, X1), 0)) / ( + res_inq = fejer_residual_norm(X0, X1) / ( 1 + abs(bI) * np.sqrt(Ngrid * n_pairs) ) @@ -607,7 +693,8 @@ def print_updates(verbose): + xp.linalg.norm(bE1 - z1) / (1 + xp.linalg.norm(bE1)) + xp.linalg.norm(bEq - zq) / (1 + xp.linalg.norm(bEq)) ) - res_inq = xp.linalg.norm(xp.maximum(bI - fun_AI(X0, X1), 0)) / ( + + res_inq = fejer_residual_norm(X0, X1) / ( 1 + abs(bI) * np.sqrt(Ngrid * n_pairs) ) @@ -703,6 +790,10 @@ def print_updates(verbose): print_updates(verbose) rho, p_resnorm, d_resnorm = update_rho(X0, X1, Xq, res_X, rho) + # The inequality multipliers and other ADMM-only arrays are no longer needed + # while reconstructing the full degree matrices. + del yI, C0, C1, S0, S1, Sq, Xq, AI_mat_offdiag + X_admm = self._transform_coeff_back(X0, X1, IDX_upper, IDX_lower, idx_offdiag) X_admm = [xp.asnumpy(Xk) for Xk in X_admm] return X_admm @@ -735,7 +826,7 @@ def admm_sym_J(self, C, verbose): AEq, bEq = pre["AEq"], pre["bEq"] AEqAEqtinv = pre["AEqAEqtinv"] - # Fejér inequality constraints and spectral scale. + # Fejer inequality constraints and spectral scale. AI_mat_diag = pre["AI_mat_diag"] AI_mat_offdiag = pre["AI_mat_offdiag"] bI, Lambda = pre["bI"], pre["Lambda"] @@ -1438,13 +1529,13 @@ def _discretize_SO3(self): @staticmethod def _largest_eigenvalue_gram(AI): """ - Compute the denominator for the Fejér inequality-multiplier update. + Compute the denominator for the Fejer inequality-multiplier update. The largest eigenvalue of AI @ AI.T is computed from the smaller AI.T @ AI Gram matrix. Retain the fixed 2000 offset used by the original implementation. - :param AI: Real Fejér inequality operator of shape (Ngrid, D0 + D1). + :param AI: Real Fejer inequality operator of shape (Ngrid, D0 + D1). :return: Largest eigenvalue of AI @ AI.T, plus 2000. """ From df45e4b7bf4c749d25a01e74113ec5752bafe24f Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 27 Oct 2026 10:25:01 -0400 Subject: [PATCH 154/161] Release ADMM arrays by setting to None --- src/aspire/abinitio/commonline_nug.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 648932f04b..88477e0485 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -792,7 +792,11 @@ def print_updates(verbose): # The inequality multipliers and other ADMM-only arrays are no longer needed # while reconstructing the full degree matrices. - del yI, C0, C1, S0, S1, Sq, Xq, AI_mat_offdiag + yI = None + C0 = C1 = None + S0 = S1 = Sq = None + Xq = None + AI_mat_offdiag = None X_admm = self._transform_coeff_back(X0, X1, IDX_upper, IDX_lower, idx_offdiag) X_admm = [xp.asnumpy(Xk) for Xk in X_admm] From b7a6195d1172809d2b9b250088d15a619fd9d193 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 27 Oct 2026 10:57:35 -0400 Subject: [PATCH 155/161] ADMM section titles for readability --- src/aspire/abinitio/commonline_nug.py | 49 ++++++++++++++++++++------- 1 file changed, 37 insertions(+), 12 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 88477e0485..52c4b7c5bd 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -375,6 +375,8 @@ def admm_plain_J(self, C, verbose): mult = self.mult Nstep_yI = self.Nstep_yI + # --------------------------- Preprocessing and problem data ---------------------------- + # Prepare the objective, constraints, block layout, and initial ADMM variables. pre = self.ADMM_preprocessing(C) @@ -409,17 +411,6 @@ def admm_plain_J(self, C, verbose): n_pairs = N * (N - 1) // 2 n_packed = N * (N + 1) // 2 - # Process off-diagonal image pairs in chunks to limit temporary GPU memory - # used by the Fejer inequality operators. - pair_chunk_size = 1024 - - def pair_chunks(): - """ - Yield bounded slices of the off-diagonal image-pair columns. - """ - for start in range(0, n_pairs, pair_chunk_size): - yield slice(start, min(start + pair_chunk_size, n_pairs)) - # In the asymmetric formulation, diagonal representation blocks are identity. bE0 = xp.zeros(D0, dtype=np.float64) bE1 = xp.zeros(D1, dtype=np.float64) @@ -434,6 +425,8 @@ def pair_chunks(): bE0 = xp.repeat(bE0[:, None], N, axis=1) bE1 = xp.repeat(bE1[:, None], N, axis=1) + # -------------------------------- Constraint operators --------------------------------- + def fun_AE(X0, X1, Xq): """ Apply the asymmetric equality constraints. @@ -459,6 +452,14 @@ def fun_AET(yE0, yE1, yEq): Z1[:4, idx_offdiag] = Zq[17:] return Z0, Z1, Zq[:16] + def pair_chunks(): + """ + Yield off-diagonal image-pair slices that bound GPU temporaries. + """ + chunk_size = 1024 + for start in range(0, n_pairs, chunk_size): + yield slice(start, min(start + chunk_size, n_pairs)) + def fun_AI(X0_chunk, X1_chunk): """ Apply the Fejer operator to one chunk of off-diagonal pairs. @@ -505,6 +506,8 @@ def fejer_residual_norm(X0, X1): return xp.sqrt(squared_norm) + # ---------------------------------- ADMM update rules ---------------------------------- + def update_S(C0, C1, yE0, yE1, yEq, yI, X0, X1, Xq, rho): """ Update and project the asymmetric PSD slack variables. @@ -667,6 +670,8 @@ def update_rho(X0, X1, Xq, res_X, rho): rho = rho / factor return rho, p_resnorm, d_resnorm + # ----------------------------- Optional iteration logging ------------------------------ + def print_updates(verbose): """ Log diagnostics for the current asymmetric ADMM iterate. @@ -753,6 +758,8 @@ def print_updates(verbose): f"|S|={normS:.2f}, |X|={normX:.2f}" ) + # ----------------------------------- ADMM iterations ----------------------------------- + yI = xp.zeros((Ngrid, n_pairs), dtype=np.float64) yE0 = xp.zeros(bE0.shape, dtype=np.float64) yE1 = xp.zeros(bE1.shape, dtype=np.float64) @@ -798,6 +805,8 @@ def print_updates(verbose): Xq = None AI_mat_offdiag = None + # ---------------------------- Assemble degree-wise results ----------------------------- + X_admm = self._transform_coeff_back(X0, X1, IDX_upper, IDX_lower, idx_offdiag) X_admm = [xp.asnumpy(Xk) for Xk in X_admm] return X_admm @@ -820,6 +829,8 @@ def admm_sym_J(self, C, verbose): mult = self.mult Nstep_yI = self.Nstep_yI + # --------------------------- Preprocessing and problem data ---------------------------- + # Prepare the objective, constraints, block layout, and initial ADMM variables. pre = self.ADMM_preprocessing(C) @@ -850,6 +861,9 @@ def admm_sym_J(self, C, verbose): n_pairs = N * (N - 1) // 2 Ngrid = self.Ngrid + + # ---------------------------- Symmetry-specific constraints ---------------------------- + rank_Ak, _ = self._compute_rank() logger.info(f"Rank of Ak: {rank_Ak}") @@ -908,6 +922,8 @@ def admm_sym_J(self, C, verbose): ] P.append(xp.asarray(Pk)) + # -------------------------------- Constraint operators --------------------------------- + def fun_AE(X0, X1, Xd0, Xd1, Xq): """ Apply the equality constraint operator to the current primal variables. @@ -1008,6 +1024,8 @@ def fun_AIT(yI): Z[:, idx_offdiag] = AI_mat_offdiag.T @ yI[:, idx_offdiag] return Z[:D0, :], Z[D0:, :] + # ---------------------------------- ADMM update rules ---------------------------------- + def update_S(C0, C1, yE, yEq, yI, X0, X1, Xd0, Xd1, Xq, rho, Lmax, N): """ Update PSD slack variables for the current ADMM iterate. @@ -1166,6 +1184,8 @@ def update_rho(X0, X1, Xd0, Xd1, Xq, bE, bEq, bI, res_X, rho, factor, normC): rho = rho / factor return rho, p_resnorm, d_resnorm + # ----------------------------- Optional iteration logging ------------------------------ + def print_updates(verbose): if verbose: obj_p = ( @@ -1265,6 +1285,8 @@ def print_updates(verbose): f"|S|={normS:.2f}, |X|={normX:.2f}" ) + # ----------------------------------- ADMM iterations ----------------------------------- + Xd0 = xp.zeros((D0, N), dtype=xp.float64) Xd1 = xp.zeros((D1, N), dtype=xp.float64) Sd0 = xp.zeros(Xd0.shape, dtype=xp.float64) @@ -1299,6 +1321,8 @@ def print_updates(verbose): X0, X1, Xd0, Xd1, Xq, bE, bEq, bI, res_X, rho, factor, normC ) + # ---------------------------- Assemble degree-wise results ----------------------------- + # Convert the optimized packed block variables back to full degree-wise # representation matrices for Euler-angle recovery. X_admm = self._transform_coeff_back(X0, X1, IDX_upper, IDX_lower, idx_offdiag) @@ -1311,7 +1335,8 @@ def print_updates(verbose): def ADMM_preprocessing(self, C): """ - Construct the transformed coefficients, constraints, indices, and initial variables used by ADMM. + Construct the transformed coefficients, constraints, indices, and + initial variables used by ADMM. :param C: Fourier coefficient matrices of the NUG objective. From 7998dcb53ba92b5039eab215af5a413b5b266c87 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 27 Oct 2026 12:19:41 -0400 Subject: [PATCH 156/161] make dtype test faster --- tests/test_nug.py | 40 +++++++++++++++++++++++++++++----------- 1 file changed, 29 insertions(+), 11 deletions(-) diff --git a/tests/test_nug.py b/tests/test_nug.py index 04a37271fc..1dd3a5f7e9 100644 --- a/tests/test_nug.py +++ b/tests/test_nug.py @@ -139,21 +139,39 @@ def test_smoke(dtype, Volume): assert rots.dtype == dtype -@pytest.mark.expensive -def test_dtypes(orient_est): +# Narrow the symmetry parametrization to cover each path only once. +@pytest.mark.parametrize("symmetry", ["C1", "C3", "D3"], indirect=True) +def test_dtypes(volume): """ - Check dtypes for each major step of the algorithm. + Check input/output dtype and internal precision for asymmetric, cyclic, + and dihedral NUG paths. """ - assert orient_est.dtype == orient_est.src.dtype - - # Intermediate steps use doubles - for Ci in orient_est.C: - assert Ci.dtype == np.float64 + src = Simulation( + n=20, + vols=volume, + offsets=0, + amplitudes=1, + seed=SEED, + ).cache() - for Xi in orient_est.X_est: - assert Xi.dtype == np.float64 + # Detuned estimator for quick test + orient_est = CommonlineNUG( + src, + max_shift=0, + Lmax=4, + T=5, + S2_grid=50, + max_iter=2, + Nstep_yI=1, + pr_iters=None, + verbose=False, + ) + rots = orient_est.estimate_rotations() - assert orient_est.rotations.dtype == orient_est.dtype + assert orient_est.dtype == src.dtype + assert all(Ci.dtype == np.float64 for Ci in orient_est.C) + assert all(Xi.dtype == np.float64 for Xi in orient_est.X_est) + assert rots.dtype == src.dtype @pytest.mark.expensive From 7076d70f2932a4831bd4be7cabef5038285bf65f Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 27 Oct 2026 12:26:19 -0400 Subject: [PATCH 157/161] C1 pair batching smoke test --- tests/test_nug.py | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/tests/test_nug.py b/tests/test_nug.py index 1dd3a5f7e9..26db92a7fc 100644 --- a/tests/test_nug.py +++ b/tests/test_nug.py @@ -202,6 +202,34 @@ def test_estimate_rotations(orient_est): mean_aligned_angular_distance(orient_est.rotations, gt_rots, 10.0) +def test_c1_admm_pair_batching(): + """ + Run C1 ADMM with enough image pairs to require two chunks. + This is just a smoke test. + """ + n_img = 46 + + # assert we have more pairs than the 1024 batch size. + assert n_img * (n_img - 1) // 2 > 1024 + + vol = emdb_2660().astype(np.float64).downsample(16) + src = Simulation(n=n_img, vols=vol, offsets=0, seed=SEED).cache() + estimator = CommonlineNUG( + src, + Lmax=4, + T=5, + S2_grid=50, + max_shift=0, + max_iter=2, + Nstep_yI=1, + pr_iters=None, + verbose=False, + ) + + estimator.compute_coeff() + estimator.perform_admm() + + @pytest.mark.parametrize("block_size", [2, 3, 12, 13]) def test_mat_block_vec_block_roundtrip(block_size): """ From 5bbd4385a7d4ff070a65c20d6432964d165895cb Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 27 Oct 2026 13:53:14 -0400 Subject: [PATCH 158/161] typos --- gallery/experiments/experimental_nug_pipeline_10081.py | 4 ++-- gallery/experiments/experimental_nug_pipeline_12036.py | 4 ++-- .../experiments/experimental_nug_simulated_pipeline_10028.py | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/gallery/experiments/experimental_nug_pipeline_10081.py b/gallery/experiments/experimental_nug_pipeline_10081.py index ba76aa45bd..d0b68c4fb9 100644 --- a/gallery/experiments/experimental_nug_pipeline_10081.py +++ b/gallery/experiments/experimental_nug_pipeline_10081.py @@ -2,7 +2,7 @@ Symmetric NUG Pipeline ====================== -This notebook demonstrates reproducing simliar results to the +This notebook demonstrates reproducing similar results to the experiment found in: .. admonition:: Publication @@ -69,7 +69,7 @@ # # Create an orientation estimation object for the ``avgs``. # The ``CommonlineNUG`` algorithm will detect the symmetry from -# the ``avgs`` metadata and use the approriate NUG ADMM-solver to estimate +# the ``avgs`` metadata and use the appropriate NUG ADMM-solver to estimate # orientations. To replicate the original experiment as closely as possible, # we use the 15 iterations of proximal refinement and disable searching over # the commonline shift space. diff --git a/gallery/experiments/experimental_nug_pipeline_12036.py b/gallery/experiments/experimental_nug_pipeline_12036.py index 86d19cd7ec..1c0904690d 100644 --- a/gallery/experiments/experimental_nug_pipeline_12036.py +++ b/gallery/experiments/experimental_nug_pipeline_12036.py @@ -2,7 +2,7 @@ Dn-Symmetric NUG Pipeline ========================= -This notebook demonstrates reproducing simliar results to the +This notebook demonstrates reproducing similar results to the experiment found in: .. admonition:: Publication @@ -68,7 +68,7 @@ # # Create an orientation estimation object for the ``avgs``. # The ``CommonlineNUG`` algorithm will detect the symmetry from -# the ``avgs`` metadata and use the approriate NUG ADMM-solver to estimate +# the ``avgs`` metadata and use the appropriate NUG ADMM-solver to estimate # orientations. To replicate the original experiment as closely as possible, # we use the 5 iterations of proximal refinement and disable searching over # the commonline shift space. diff --git a/gallery/experiments/experimental_nug_simulated_pipeline_10028.py b/gallery/experiments/experimental_nug_simulated_pipeline_10028.py index 204f867013..623ee3eb0c 100644 --- a/gallery/experiments/experimental_nug_simulated_pipeline_10028.py +++ b/gallery/experiments/experimental_nug_simulated_pipeline_10028.py @@ -81,7 +81,7 @@ .cache() ) - # Initialize the NUG orientation estimator, matching paramters + # Initialize the NUG orientation estimator, matching parameters # to the referenced experiment. estimator = CommonlineNUG( source, From 63996ad70399d1b0a1de68460e5808f04bddfb0a Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Tue, 27 Oct 2026 15:13:48 -0400 Subject: [PATCH 159/161] Flake8 fix also contained in PR #1411 --- tests/test_orient_symmetric.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_orient_symmetric.py b/tests/test_orient_symmetric.py index df86ba020a..ff07c96657 100644 --- a/tests/test_orient_symmetric.py +++ b/tests/test_orient_symmetric.py @@ -192,8 +192,8 @@ def test_self_relative_rotations(n_img, L, order, dtype): Rii_gt = rot_gt.T @ g @ rot_gt Rii = Riis[i] cases = np.array([Rii, Rii.T, J_conjugate(Rii), J_conjugate(Rii.T)]) - for i, estimate in enumerate(cases): - dist[i] = Rotation.angle_dist(estimate, Rii_gt) + for case_idx, estimate in enumerate(cases): + dist[case_idx] = Rotation.angle_dist(estimate, Rii_gt) angular_distance[i] = min(dist) mean_angular_distance = np.mean(angular_distance) * 180 / np.pi From bcbeac0d95ad04fba9dd9b9b379485b9b6bd139a Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Wed, 28 Oct 2026 15:09:47 -0400 Subject: [PATCH 160/161] remove commented out code --- src/aspire/abinitio/commonline_nug.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 52c4b7c5bd..64df94ce20 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -295,7 +295,7 @@ def fijhat_k(k_idx, F): dk = 2 * k + 1 C[k_idx][j * dk : (j + 1) * dk, i * dk : (i + 1) * dk] = ( fijhat_k(k_idx, Fij) - ) # *dk + ) pbar.update(1) # Fill the conjugate transpose blocks so each coefficient matrix is Hermitian. @@ -310,7 +310,7 @@ def fijhat_k(k_idx, F): dk = 2 * k + 1 C[k_idx][i * dk : (i + 1) * dk, i * dk : (i + 1) * dk] = fijhat_k( k_idx, Fii - ) # *dk + ) # Convert complex Wigner coefficients to the real representation basis used by ADMM. for k_idx in range(Lmax): From eada3db21861c7d2f375177aff7fcd89005c8372 Mon Sep 17 00:00:00 2001 From: Josh Carmichael Date: Fri, 30 Oct 2026 12:41:03 -0400 Subject: [PATCH 161/161] Replace ADMM_preprocessing dictionary with data class --- src/aspire/abinitio/commonline_nug.py | 576 ++++++++++++++++---------- 1 file changed, 346 insertions(+), 230 deletions(-) diff --git a/src/aspire/abinitio/commonline_nug.py b/src/aspire/abinitio/commonline_nug.py index 64df94ce20..c293810d75 100644 --- a/src/aspire/abinitio/commonline_nug.py +++ b/src/aspire/abinitio/commonline_nug.py @@ -1,4 +1,5 @@ import logging +from dataclasses import dataclass import numpy as np from scipy.special import factorial @@ -20,6 +21,73 @@ logger = logging.getLogger(__name__) +@dataclass +class _ADMMData: + """ + Data constructed for the C1 and symmetric ADMM solvers. + + Array fields use the active NumPy or CuPy backend. Some array fields are + set to None after their last use to release GPU memory. All fields are + populated when preprocessing returns. + + Objective: + C0, C1: Coefficients packed into the two representation blocks. + normC: Combined norm of those coefficients. + + Quaternion equalities: + AEq: Equality-constraint matrix. + bEq: Equality right-hand side. + AEqAEqtinv: Pseudoinverse used when updating equality multipliers. + + Fejer inequalities: + AI_mat_diag: Operator for self-pairs (used by the symmetric solver). + AI_mat_offdiag: Operator for distinct image pairs. + bI: Inequality bound. + Lambda: Scale used for inequality-multiplier updates. + + Block layout: + d0, d1: Degree-block row boundaries for the two packed arrays. + idx_diag, idx_offdiag: Column indices in packed image-pair order. + IDX_upper, IDX_lower: Indices in the full image-pair grid. + + ADMM variables: + X0, X1, Xq: Primal variables for the two blocks and quaternion terms. + S0, S1, Sq: Corresponding slack variables. + """ + + # Objective + C0: xp.ndarray | None + C1: xp.ndarray | None + normC: float | xp.ndarray + + # Quaternion equalities + AEq: xp.ndarray + bEq: xp.ndarray + AEqAEqtinv: xp.ndarray + + # Fejer inequalities + AI_mat_diag: xp.ndarray | None + AI_mat_offdiag: xp.ndarray | None + bI: float + Lambda: float | xp.ndarray + + # Block layout + d0: list[int] + d1: list[int] + idx_diag: xp.ndarray + idx_offdiag: xp.ndarray + IDX_upper: xp.ndarray + IDX_lower: xp.ndarray + + # Initial and subsequently updated ADMM variables + X0: xp.ndarray + X1: xp.ndarray + Xq: xp.ndarray | None + S0: xp.ndarray | None + S1: xp.ndarray | None + Sq: xp.ndarray | None + + class CommonlineNUG(Orient3D): """ Estimate orientations of cyclically or dihedrally symmetric molecules using the non-unique games framework. @@ -379,33 +447,8 @@ def admm_plain_J(self, C, verbose): # Prepare the objective, constraints, block layout, and initial ADMM variables. pre = self.ADMM_preprocessing(C) - - # Packed objective coefficients and their norm. - C0, C1, normC = pre["C0"], pre["C1"], pre["normC"] - - # Quaternion equality constraints. - AEq, bEq = pre["AEq"], pre["bEq"] - AEqAEqtinv = pre["AEqAEqtinv"] - - # Fejer inequality constraints and spectral scale. - AI_mat_offdiag = pre["AI_mat_offdiag"] - bI, Lambda = pre["bI"], pre["Lambda"] - - # Row boundaries by degree: d0 applies to C0/X0/S0, and d1 to C1/X1/S1. - # Their final entries give the total row counts. - d0, d1 = pre["d0"], pre["d1"] - D0, D1 = d0[-1], d1[-1] - - # Indices into packed columns and the full image-pair grid, respectively. - idx_diag, idx_offdiag = pre["idx_diag"], pre["idx_offdiag"] - IDX_upper, IDX_lower = pre["IDX_upper"], pre["IDX_lower"] - - # Initial primal and slack variables. - X0, X1, Xq = pre["X0"], pre["X1"], pre["Xq"] - S0, S1, Sq = pre["S0"], pre["S1"], pre["Sq"] - - # Safe to remove pre now that it is unpacked. - del pre + pre.AI_mat_diag = None # C1 does not use this operator. + D0, D1 = pre.d0[-1], pre.d1[-1] Ngrid = self.Ngrid n_pairs = N * (N - 1) // 2 @@ -419,8 +462,8 @@ def admm_plain_J(self, C, verbose): # Encode the k-by-k and (k+1)-by-(k+1) identity blocks. A flattened # m-by-m identity has diagonal entries every m+1 positions. - bE0[d0[k_idx] : d0[k_idx + 1] : k + 1] = 1 - bE1[d1[k_idx] : d1[k_idx + 1] : k + 2] = 1 + bE0[pre.d0[k_idx] : pre.d0[k_idx + 1] : k + 1] = 1 + bE1[pre.d1[k_idx] : pre.d1[k_idx + 1] : k + 2] = 1 bE0 = xp.repeat(bE0[:, None], N, axis=1) bE1 = xp.repeat(bE1[:, None], N, axis=1) @@ -431,10 +474,10 @@ def fun_AE(X0, X1, Xq): """ Apply the asymmetric equality constraints. """ - z0 = X0[:, idx_diag] - z1 = X1[:, idx_diag] - zq = AEq @ xp.concatenate( - (Xq, X0[:1, idx_offdiag], X1[:4, idx_offdiag]), axis=0 + z0 = X0[:, pre.idx_diag] + z1 = X1[:, pre.idx_diag] + zq = pre.AEq @ xp.concatenate( + (Xq, X0[:1, pre.idx_offdiag], X1[:4, pre.idx_offdiag]), axis=0 ) return z0, z1, zq @@ -444,12 +487,12 @@ def fun_AET(yE0, yE1, yEq): """ Z0 = xp.zeros((D0, n_packed), dtype=np.float64) Z1 = xp.zeros((D1, n_packed), dtype=np.float64) - Z0[:, idx_diag] = yE0 - Z1[:, idx_diag] = yE1 + Z0[:, pre.idx_diag] = yE0 + Z1[:, pre.idx_diag] = yE1 - Zq = AEq.T @ yEq - Z0[:1, idx_offdiag] = Zq[16:17] - Z1[:4, idx_offdiag] = Zq[17:] + Zq = pre.AEq.T @ yEq + Z0[:1, pre.idx_offdiag] = Zq[16:17] + Z1[:4, pre.idx_offdiag] = Zq[17:] return Z0, Z1, Zq[:16] def pair_chunks(): @@ -465,13 +508,13 @@ def fun_AI(X0_chunk, X1_chunk): Apply the Fejer operator to one chunk of off-diagonal pairs. """ coefficients = xp.concatenate((X0_chunk, X1_chunk), axis=0) - return AI_mat_offdiag @ coefficients + return pre.AI_mat_offdiag @ coefficients def fun_AIT_chunk(yI_chunk): """ Apply the Fejer adjoint to one chunk of off-diagonal pairs. """ - return AI_mat_offdiag.T @ yI_chunk + return pre.AI_mat_offdiag.T @ yI_chunk def fun_AIT(yI): """ @@ -480,7 +523,7 @@ def fun_AIT(yI): Z = xp.zeros((D0 + D1, n_packed), dtype=np.float64) for columns in pair_chunks(): - packed_columns = idx_offdiag[columns] + packed_columns = pre.idx_offdiag[columns] Z[:, packed_columns] = fun_AIT_chunk(yI[:, columns]) return Z[:D0], Z[D0:] @@ -493,9 +536,9 @@ def fejer_residual_norm(X0, X1): squared_norm = xp.zeros((), dtype=np.float64) for columns in pair_chunks(): - packed_columns = idx_offdiag[columns] + packed_columns = pre.idx_offdiag[columns] residual = xp.maximum( - bI + pre.bI - fun_AI( X0[:, packed_columns], X1[:, packed_columns], @@ -538,11 +581,11 @@ def update_S(C0, C1, yE0, yE1, yEq, yI, X0, X1, Xq, rho): S0, S1 = self._project_representation_blocks( S0, S1, - d0, - d1, - IDX_upper, - IDX_lower, - idx_offdiag, + pre.d0, + pre.d1, + pre.IDX_upper, + pre.IDX_lower, + pre.idx_offdiag, ) # This is the batched equivalent of projecting each 4x4 quaternion block @@ -563,7 +606,7 @@ def update_yE(C0, C1, S0, S1, Sq, yI, X0, X1, Xq, rho): ) yE0 = bE0 / rho + z0 yE1 = bE1 / rho + z1 - yEq = AEqAEqtinv @ (bEq / rho + zq) + yEq = pre.AEqAEqtinv @ (pre.bEq / rho + zq) return yE0, yE1, yEq def update_yI( @@ -587,7 +630,7 @@ def update_yI( AET_yE0, AET_yE1, _ = fun_AET(yE0, yE1, yEq) for columns in pair_chunks(): - packed_columns = idx_offdiag[columns] + packed_columns = pre.idx_offdiag[columns] AIT_yI = fun_AIT_chunk(yI[:, columns]) tmp0 = ( @@ -606,7 +649,7 @@ def update_yI( ) yI_chunk = yI[:, columns] - yI_chunk += bI / rho / Lambda + yI_chunk += pre.bI / rho / Lambda yI_chunk += fun_AI(tmp0, tmp1) / Lambda xp.maximum(yI_chunk, 0, out=yI_chunk) @@ -656,14 +699,14 @@ def update_rho(X0, X1, Xq, res_X, rho): res_eq = ( xp.linalg.norm(bE0 - z0) / (1 + xp.linalg.norm(bE0)) + xp.linalg.norm(bE1 - z1) / (1 + xp.linalg.norm(bE1)) - + xp.linalg.norm(bEq - zq) / (1 + xp.linalg.norm(bEq)) + + xp.linalg.norm(pre.bEq - zq) / (1 + xp.linalg.norm(pre.bEq)) ) res_inq = fejer_residual_norm(X0, X1) / ( - 1 + abs(bI) * np.sqrt(Ngrid * n_pairs) + 1 + abs(pre.bI) * np.sqrt(Ngrid * n_pairs) ) p_resnorm = res_eq + res_inq - d_resnorm = res_X / (1 + normC) + d_resnorm = res_X / (1 + pre.normC) if d_resnorm > ratio * p_resnorm: rho = rho * factor if d_resnorm < ratio * p_resnorm: @@ -680,74 +723,74 @@ def print_updates(verbose): return obj_p = ( - xp.vdot(C0[:, idx_diag], X0[:, idx_diag]) - + xp.vdot(C1[:, idx_diag], X1[:, idx_diag]) - + 2 * xp.vdot(C0[:, idx_offdiag], X0[:, idx_offdiag]) - + 2 * xp.vdot(C1[:, idx_offdiag], X1[:, idx_offdiag]) + xp.vdot(pre.C0[:, pre.idx_diag], pre.X0[:, pre.idx_diag]) + + xp.vdot(pre.C1[:, pre.idx_diag], pre.X1[:, pre.idx_diag]) + + 2 * xp.vdot(pre.C0[:, pre.idx_offdiag], pre.X0[:, pre.idx_offdiag]) + + 2 * xp.vdot(pre.C1[:, pre.idx_offdiag], pre.X1[:, pre.idx_offdiag]) ) obj_d = ( xp.vdot(yE0, bE0) + xp.vdot(yE1, bE1) - + 2 * xp.vdot(yEq, bEq) - + 2 * bI * xp.sum(yI) + + 2 * xp.vdot(yEq, pre.bEq) + + 2 * pre.bI * xp.sum(yI) ) - z0, z1, zq = fun_AE(X0, X1, Xq) + z0, z1, zq = fun_AE(pre.X0, pre.X1, pre.Xq) res_eq = ( xp.linalg.norm(bE0 - z0) / (1 + xp.linalg.norm(bE0)) + xp.linalg.norm(bE1 - z1) / (1 + xp.linalg.norm(bE1)) - + xp.linalg.norm(bEq - zq) / (1 + xp.linalg.norm(bEq)) + + xp.linalg.norm(pre.bEq - zq) / (1 + xp.linalg.norm(pre.bEq)) ) - res_inq = fejer_residual_norm(X0, X1) / ( - 1 + abs(bI) * np.sqrt(Ngrid * n_pairs) + res_inq = fejer_residual_norm(pre.X0, pre.X1) / ( + 1 + abs(pre.bI) * np.sqrt(Ngrid * n_pairs) ) res_psdX = 0 for k_idx in range(Lmax): k = k_idx + 1 - block0_rows = slice(d0[k_idx], d0[k_idx + 1]) - block1_rows = slice(d1[k_idx], d1[k_idx + 1]) + block0_rows = slice(pre.d0[k_idx], pre.d0[k_idx + 1]) + block1_rows = slice(pre.d1[k_idx], pre.d1[k_idx + 1]) tmp = self.mat_block( - X0[block0_rows], + pre.X0[block0_rows], N, k, - IDX_upper, - IDX_lower, - idx_offdiag, + pre.IDX_upper, + pre.IDX_lower, + pre.idx_offdiag, ) res_psdX += xp.linalg.norm(self._psd_projection(-tmp)) tmp = self.mat_block( - X1[block1_rows], + pre.X1[block1_rows], N, k + 1, - IDX_upper, - IDX_lower, - idx_offdiag, + pre.IDX_upper, + pre.IDX_lower, + pre.idx_offdiag, ) res_psdX += xp.linalg.norm(self._psd_projection(-tmp)) - res_psdX /= 1 + xp.linalg.norm(X0) + xp.linalg.norm(X1) + res_psdX /= 1 + xp.linalg.norm(pre.X0) + xp.linalg.norm(pre.X1) - Xq_blocks = Xq.T.reshape(n_pairs, 4, 4) + Xq_blocks = pre.Xq.T.reshape(n_pairs, 4, 4) res_psdQ = xp.linalg.norm( self._psd_projection(-Xq_blocks), axis=(-2, -1) ).sum() - res_psdQ /= 1 + xp.linalg.norm(Xq) + res_psdQ /= 1 + xp.linalg.norm(pre.Xq) normS = xp.sqrt( - xp.linalg.norm(S0) ** 2 - + xp.linalg.norm(S1) ** 2 - + xp.linalg.norm(Sq) ** 2 + xp.linalg.norm(pre.S0) ** 2 + + xp.linalg.norm(pre.S1) ** 2 + + xp.linalg.norm(pre.Sq) ** 2 ) normX = xp.sqrt( - xp.linalg.norm(X0) ** 2 - + xp.linalg.norm(X1) ** 2 - + xp.linalg.norm(Xq) ** 2 + xp.linalg.norm(pre.X0) ** 2 + + xp.linalg.norm(pre.X1) ** 2 + + xp.linalg.norm(pre.Xq) ** 2 ) p_res = res_eq + res_inq + res_psdX + res_psdQ - d_res = res_X / (1 + normC) + d_res = res_X / (1 + pre.normC) logger.info( f"Iter {t}: p_res={p_res:.5f}, d_res={d_res:.5f}, " @@ -763,51 +806,78 @@ def print_updates(verbose): yI = xp.zeros((Ngrid, n_pairs), dtype=np.float64) yE0 = xp.zeros(bE0.shape, dtype=np.float64) yE1 = xp.zeros(bE1.shape, dtype=np.float64) - yEq = xp.zeros(bEq.shape, dtype=np.float64) + yEq = xp.zeros(pre.bEq.shape, dtype=np.float64) update_order = np.arange(3) for t in trange(max_iter, desc="C1 NUG ADMM"): np.random.shuffle(update_order) for update_idx in update_order: if update_idx == 0: - S0, S1, Sq = update_S(C0, C1, yE0, yE1, yEq, yI, X0, X1, Xq, rho) + pre.S0, pre.S1, pre.Sq = update_S( + pre.C0, pre.C1, yE0, yE1, yEq, yI, pre.X0, pre.X1, pre.Xq, rho + ) elif update_idx == 1: - yE0, yE1, yEq = update_yE(C0, C1, S0, S1, Sq, yI, X0, X1, Xq, rho) + yE0, yE1, yEq = update_yE( + pre.C0, + pre.C1, + pre.S0, + pre.S1, + pre.Sq, + yI, + pre.X0, + pre.X1, + pre.Xq, + rho, + ) else: for _ in range(Nstep_yI): yI = update_yI( - C0, - C1, - S0, - S1, + pre.C0, + pre.C1, + pre.S0, + pre.S1, yE0, yE1, yEq, yI, - X0, - X1, + pre.X0, + pre.X1, rho, - Lambda, + pre.Lambda, ) - X0, X1, Xq, res_X = update_X( - C0, C1, S0, S1, Sq, yE0, yE1, yEq, yI, X0, X1, Xq, rho + pre.X0, pre.X1, pre.Xq, res_X = update_X( + pre.C0, + pre.C1, + pre.S0, + pre.S1, + pre.Sq, + yE0, + yE1, + yEq, + yI, + pre.X0, + pre.X1, + pre.Xq, + rho, ) if t % 100 == 0: print_updates(verbose) - rho, p_resnorm, d_resnorm = update_rho(X0, X1, Xq, res_X, rho) + rho, p_resnorm, d_resnorm = update_rho(pre.X0, pre.X1, pre.Xq, res_X, rho) # The inequality multipliers and other ADMM-only arrays are no longer needed # while reconstructing the full degree matrices. yI = None - C0 = C1 = None - S0 = S1 = Sq = None - Xq = None - AI_mat_offdiag = None + pre.C0 = pre.C1 = None + pre.S0 = pre.S1 = pre.Sq = None + pre.Xq = None + pre.AI_mat_offdiag = None # ---------------------------- Assemble degree-wise results ----------------------------- - X_admm = self._transform_coeff_back(X0, X1, IDX_upper, IDX_lower, idx_offdiag) + X_admm = self._transform_coeff_back( + pre.X0, pre.X1, pre.IDX_upper, pre.IDX_lower, pre.idx_offdiag + ) X_admm = [xp.asnumpy(Xk) for Xk in X_admm] return X_admm @@ -833,31 +903,7 @@ def admm_sym_J(self, C, verbose): # Prepare the objective, constraints, block layout, and initial ADMM variables. pre = self.ADMM_preprocessing(C) - - # Packed objective coefficients and their norm. - C0, C1, normC = pre["C0"], pre["C1"], pre["normC"] - - # Quaternion equality constraints. - AEq, bEq = pre["AEq"], pre["bEq"] - AEqAEqtinv = pre["AEqAEqtinv"] - - # Fejer inequality constraints and spectral scale. - AI_mat_diag = pre["AI_mat_diag"] - AI_mat_offdiag = pre["AI_mat_offdiag"] - bI, Lambda = pre["bI"], pre["Lambda"] - - # Row boundaries by degree: d0 applies to C0/X0/S0, and d1 to C1/X1/S1. - # Their final entries give the total row counts. - d0, d1 = pre["d0"], pre["d1"] - D0, D1 = d0[-1], d1[-1] - - # Indices into packed columns and the full image-pair grid, respectively. - idx_diag, idx_offdiag = pre["idx_diag"], pre["idx_offdiag"] - IDX_upper, IDX_lower = pre["IDX_upper"], pre["IDX_lower"] - - # Initial primal and slack variables. - X0, X1, Xq = pre["X0"], pre["X1"], pre["Xq"] - S0, S1, Sq = pre["S0"], pre["S1"], pre["Sq"] + D0, D1 = pre.d0[-1], pre.d1[-1] n_pairs = N * (N - 1) // 2 Ngrid = self.Ngrid @@ -899,9 +945,9 @@ def admm_sym_J(self, C, verbose): for k_idx in range(Lmax): k = k_idx + 1 - start = k_idx + d0[k_idx] + d1[k_idx] - split = k_idx + 1 + d0[k_idx + 1] + d1[k_idx] - stop = k_idx + 1 + d0[k_idx + 1] + d1[k_idx + 1] + start = k_idx + pre.d0[k_idx] + pre.d1[k_idx] + split = k_idx + 1 + pre.d0[k_idx + 1] + pre.d1[k_idx] + stop = k_idx + 1 + pre.d0[k_idx + 1] + pre.d1[k_idx + 1] bE[start] = rank_Ak[k_idx] bE[start + 1 : split : k + 1] = 1 @@ -936,24 +982,24 @@ def fun_AE(X0, X1, Xd0, Xd1, Xq): """ z = xp.zeros((Lmax + D0 + D1, N), dtype=np.float64) for k_idx in range(Lmax): - block0_rows = slice(d0[k_idx], d0[k_idx + 1]) - block1_rows = slice(d1[k_idx], d1[k_idx + 1]) + block0_rows = slice(pre.d0[k_idx], pre.d0[k_idx + 1]) + block1_rows = slice(pre.d1[k_idx], pre.d1[k_idx + 1]) - constraint_start = k_idx + d0[k_idx] + d1[k_idx] - constraint_stop = k_idx + 1 + d0[k_idx + 1] + d1[k_idx + 1] + constraint_start = k_idx + pre.d0[k_idx] + pre.d1[k_idx] + constraint_stop = k_idx + 1 + pre.d0[k_idx + 1] + pre.d1[k_idx + 1] constraint_rows = slice(constraint_start, constraint_stop) z[constraint_rows] = AE[k_idx] @ xp.concatenate( ( - X0[block0_rows, idx_diag], - X1[block1_rows, idx_diag], + X0[block0_rows, pre.idx_diag], + X1[block1_rows, pre.idx_diag], Xd0[block0_rows], Xd1[block1_rows], ), axis=0, ) - zq = AEq @ xp.concatenate( - (Xq, X0[:1, idx_offdiag], X1[:4, idx_offdiag]), axis=0 + zq = pre.AEq @ xp.concatenate( + (Xq, X0[:1, pre.idx_offdiag], X1[:4, pre.idx_offdiag]), axis=0 ) return z, zq @@ -980,23 +1026,23 @@ def fun_AET(yE, yEq): size0 = k**2 size1 = (k + 1) ** 2 - block0_rows = slice(d0[k_idx], d0[k_idx + 1]) - block1_rows = slice(d1[k_idx], d1[k_idx + 1]) + block0_rows = slice(pre.d0[k_idx], pre.d0[k_idx + 1]) + block1_rows = slice(pre.d1[k_idx], pre.d1[k_idx + 1]) - constraint_start = k_idx + d0[k_idx] + d1[k_idx] - constraint_stop = k_idx + 1 + d0[k_idx + 1] + d1[k_idx + 1] + constraint_start = k_idx + pre.d0[k_idx] + pre.d1[k_idx] + constraint_stop = k_idx + 1 + pre.d0[k_idx + 1] + pre.d1[k_idx + 1] constraint_rows = slice(constraint_start, constraint_stop) Ztmp = AE[k_idx].T @ yE[constraint_rows] - Z0[block0_rows, idx_diag] = Ztmp[:size0] - Z1[block1_rows, idx_diag] = Ztmp[size0 : size0 + size1] + Z0[block0_rows, pre.idx_diag] = Ztmp[:size0] + Z1[block1_rows, pre.idx_diag] = Ztmp[size0 : size0 + size1] Zd0[block0_rows] = Ztmp[size0 + size1 : 2 * size0 + size1] Zd1[block1_rows] = Ztmp[2 * size0 + size1 : 2 * size0 + 2 * size1] - Zq = AEq.T @ yEq - Z0[:1, idx_offdiag] = Zq[16:17] - Z1[:4, idx_offdiag] = Zq[17:] + Zq = pre.AEq.T @ yEq + Z0[:1, pre.idx_offdiag] = Zq[16:17] + Z1[:4, pre.idx_offdiag] = Zq[17:] return Z0, Z1, Zd0, Zd1, Zq[:16] def fun_AI(X0, X1): @@ -1011,8 +1057,8 @@ def fun_AI(X0, X1): """ z = xp.zeros((Ngrid, N * (N + 1) // 2), dtype=np.float64) tmp = xp.concatenate((X0, X1), axis=0) - z[:, idx_diag] = AI_mat_diag @ tmp[:, idx_diag] - z[:, idx_offdiag] = AI_mat_offdiag @ tmp[:, idx_offdiag] + z[:, pre.idx_diag] = pre.AI_mat_diag @ tmp[:, pre.idx_diag] + z[:, pre.idx_offdiag] = pre.AI_mat_offdiag @ tmp[:, pre.idx_offdiag] return z def fun_AIT(yI): @@ -1020,8 +1066,8 @@ def fun_AIT(yI): Apply the adjoint of the Fejer inequality operator. """ Z = xp.zeros((D0 + D1, N * (N + 1) // 2), dtype=np.float64) - Z[:, idx_diag] = AI_mat_diag.T @ yI[:, idx_diag] - Z[:, idx_offdiag] = AI_mat_offdiag.T @ yI[:, idx_offdiag] + Z[:, pre.idx_diag] = pre.AI_mat_diag.T @ yI[:, pre.idx_diag] + Z[:, pre.idx_offdiag] = pre.AI_mat_offdiag.T @ yI[:, pre.idx_offdiag] return Z[:D0, :], Z[D0:, :] # ---------------------------------- ADMM update rules ---------------------------------- @@ -1051,11 +1097,11 @@ def update_S(C0, C1, yE, yEq, yI, X0, X1, Xd0, Xd1, Xq, rho, Lmax, N): S0, S1 = self._project_representation_blocks( S0, S1, - d0, - d1, - IDX_upper, - IDX_lower, - idx_offdiag, + pre.d0, + pre.d1, + pre.IDX_upper, + pre.IDX_lower, + pre.idx_offdiag, ) # Project the diagonal coupling blocks. @@ -1063,8 +1109,8 @@ def update_S(C0, C1, yE, yEq, yI, X0, X1, Xd0, Xd1, Xq, rho, Lmax, N): Sd1 = Sd1.T for k_idx in range(Lmax): k = k_idx + 1 # degree - block0_rows = slice(d0[k_idx], d0[k_idx + 1]) - block1_rows = slice(d1[k_idx], d1[k_idx + 1]) + block0_rows = slice(pre.d0[k_idx], pre.d0[k_idx + 1]) + block1_rows = slice(pre.d1[k_idx], pre.d1[k_idx + 1]) permutation = P[k_idx] tmp = self.transform_back_block( @@ -1104,14 +1150,14 @@ def update_yE(C0, C1, X0, X1, Xd0, Xd1, Xq, S0, S1, Sd0, Sd1, Sq, yI, rho): ) yE = xp.zeros((Lmax + D0 + D1, N), dtype=np.float64) for k_idx in range(Lmax): - constraint_start = k_idx + d0[k_idx] + d1[k_idx] - constraint_stop = k_idx + 1 + d0[k_idx + 1] + d1[k_idx + 1] + constraint_start = k_idx + pre.d0[k_idx] + pre.d1[k_idx] + constraint_stop = k_idx + 1 + pre.d0[k_idx + 1] + pre.d1[k_idx + 1] constraint_rows = slice(constraint_start, constraint_stop) yE[constraint_rows] = AEAETinv[k_idx] @ ( bE[constraint_rows] / rho + z[constraint_rows] ) - yEq = AEqAEqtinv @ (bEq / rho + zq) + yEq = pre.AEqAEqtinv @ (pre.bEq / rho + zq) return yE, yEq def update_yI(C0, C1, X0, X1, S0, S1, yE, yEq, yI, rho, Lambda): @@ -1126,7 +1172,7 @@ def update_yI(C0, C1, X0, X1, S0, S1, yE, yEq, yI, rho, Lambda): tmp = fun_AI( -X0 / rho + C0 - S0 - Z0 - AIT_yI0, -X1 / rho + C1 - S1 - Z1 - AIT_yI1 ) - yI = yI + bI / rho / Lambda + tmp / Lambda + yI = yI + pre.bI / rho / Lambda + tmp / Lambda yI = xp.maximum(yI, 0) return yI @@ -1189,59 +1235,63 @@ def update_rho(X0, X1, Xd0, Xd1, Xq, bE, bEq, bI, res_X, rho, factor, normC): def print_updates(verbose): if verbose: obj_p = ( - xp.vdot(C0[:, idx_diag], X0[:, idx_diag]) - + xp.vdot(C1[:, idx_diag], X1[:, idx_diag]) - + 2 * xp.vdot(C0[:, idx_offdiag], X0[:, idx_offdiag]) - + 2 * xp.vdot(C1[:, idx_offdiag], X1[:, idx_offdiag]) + xp.vdot(pre.C0[:, pre.idx_diag], pre.X0[:, pre.idx_diag]) + + xp.vdot(pre.C1[:, pre.idx_diag], pre.X1[:, pre.idx_diag]) + + 2 + * xp.vdot(pre.C0[:, pre.idx_offdiag], pre.X0[:, pre.idx_offdiag]) + + 2 + * xp.vdot(pre.C1[:, pre.idx_offdiag], pre.X1[:, pre.idx_offdiag]) ) obj_d = ( xp.vdot(yE, bE) - + 2 * xp.vdot(yEq, bEq) - + bI * xp.sum(yI[:, idx_diag]) - + 2 * bI * xp.sum(yI[:, idx_offdiag]) + + 2 * xp.vdot(yEq, pre.bEq) + + pre.bI * xp.sum(yI[:, pre.idx_diag]) + + 2 * pre.bI * xp.sum(yI[:, pre.idx_offdiag]) ) - z, zq = fun_AE(X0, X1, Xd0, Xd1, Xq) + z, zq = fun_AE(pre.X0, pre.X1, Xd0, Xd1, pre.Xq) res_eq = xp.linalg.norm(z - bE) / ( 1 + xp.linalg.norm(bE) - ) + xp.linalg.norm(zq - bEq) / (1 + xp.linalg.norm(bEq)) - res_inq = xp.linalg.norm(xp.maximum(bI - fun_AI(X0, X1), 0)) / ( - 1 + abs(bI) * np.sqrt(Ngrid * N * (N + 1) / 2) - ) + ) + xp.linalg.norm(zq - pre.bEq) / (1 + xp.linalg.norm(pre.bEq)) + res_inq = xp.linalg.norm( + xp.maximum(pre.bI - fun_AI(pre.X0, pre.X1), 0) + ) / (1 + abs(pre.bI) * np.sqrt(Ngrid * N * (N + 1) / 2)) res_psdX = 0 for k_idx in range(Lmax): k = k_idx + 1 # degree - block0_rows = slice(d0[k_idx], d0[k_idx + 1]) - block1_rows = slice(d1[k_idx], d1[k_idx + 1]) + block0_rows = slice(pre.d0[k_idx], pre.d0[k_idx + 1]) + block1_rows = slice(pre.d1[k_idx], pre.d1[k_idx + 1]) tmp = self.mat_block( - X0[block0_rows], + pre.X0[block0_rows], N, k, - IDX_upper, - IDX_lower, - idx_offdiag, + pre.IDX_upper, + pre.IDX_lower, + pre.idx_offdiag, ) res_psdX += xp.linalg.norm(self._psd_projection(-tmp)) tmp = self.mat_block( - X1[block1_rows], + pre.X1[block1_rows], N, k + 1, - IDX_upper, - IDX_lower, - idx_offdiag, + pre.IDX_upper, + pre.IDX_lower, + pre.idx_offdiag, ) res_psdX += xp.linalg.norm(self._psd_projection(-tmp)) - res_psdX = res_psdX / (1 + xp.linalg.norm(X0) + xp.linalg.norm(X1)) + res_psdX = res_psdX / ( + 1 + xp.linalg.norm(pre.X0) + xp.linalg.norm(pre.X1) + ) res_psdD = 0 for k_idx in range(Lmax): k = k_idx + 1 # degree - block0_rows = slice(d0[k_idx], d0[k_idx + 1]) - block1_rows = slice(d1[k_idx], d1[k_idx + 1]) + block0_rows = slice(pre.d0[k_idx], pre.d0[k_idx + 1]) + block1_rows = slice(pre.d1[k_idx], pre.d1[k_idx + 1]) tmp = self.transform_back_block( Xd0[block0_rows].T, @@ -1256,26 +1306,26 @@ def print_updates(verbose): res_psdQ = 0 for pair_idx in range(n_pairs): - tmp = Xq[:, pair_idx].reshape(4, 4).T + tmp = pre.Xq[:, pair_idx].reshape(4, 4).T res_psdQ += xp.linalg.norm(self._psd_projection(-tmp)) - res_psdQ = res_psdQ / (1 + xp.linalg.norm(Xq)) + res_psdQ = res_psdQ / (1 + xp.linalg.norm(pre.Xq)) normS = xp.sqrt( - xp.linalg.norm(S0) ** 2 - + xp.linalg.norm(S1) ** 2 + xp.linalg.norm(pre.S0) ** 2 + + xp.linalg.norm(pre.S1) ** 2 + xp.linalg.norm(Sd0) ** 2 + xp.linalg.norm(Sd1) ** 2 - + xp.linalg.norm(Sq) ** 2 + + xp.linalg.norm(pre.Sq) ** 2 ) normX = xp.sqrt( - xp.linalg.norm(X0) ** 2 - + xp.linalg.norm(X1) ** 2 + xp.linalg.norm(pre.X0) ** 2 + + xp.linalg.norm(pre.X1) ** 2 + xp.linalg.norm(Xd0) ** 2 + xp.linalg.norm(Xd1) ** 2 - + xp.linalg.norm(Xq) ** 2 + + xp.linalg.norm(pre.Xq) ** 2 ) p_res = res_eq + res_inq + res_psdX + res_psdD + res_psdQ - d_res = res_X / (1 + normC) + d_res = res_X / (1 + pre.normC) logger.info( f"Iter {t}: p_res={p_res:.5f}, d_res={d_res:.5f}, " f"obj_primal={obj_p:.2f}, obj_dual={obj_d:.2f}, " @@ -1293,7 +1343,7 @@ def print_updates(verbose): Sd1 = xp.zeros(Xd1.shape, dtype=xp.float64) yI = xp.zeros((Ngrid, N * (N + 1) // 2), dtype=xp.float64) yE = xp.zeros(bE.shape, dtype=xp.float64) - yEq = xp.zeros(bEq.shape, dtype=xp.float64) + yEq = xp.zeros(pre.bEq.shape, dtype=xp.float64) # Run ADMM iterations, randomly ordering the block updates before each primal # multiplier update and penalty adjustment. @@ -1302,30 +1352,95 @@ def print_updates(verbose): np.random.shuffle(update_order) for update_idx in update_order: if update_idx == 0: - S0, S1, Sd0, Sd1, Sq = update_S( - C0, C1, yE, yEq, yI, X0, X1, Xd0, Xd1, Xq, rho, Lmax, N + pre.S0, pre.S1, Sd0, Sd1, pre.Sq = update_S( + pre.C0, + pre.C1, + yE, + yEq, + yI, + pre.X0, + pre.X1, + Xd0, + Xd1, + pre.Xq, + rho, + Lmax, + N, ) if update_idx == 1: yE, yEq = update_yE( - C0, C1, X0, X1, Xd0, Xd1, Xq, S0, S1, Sd0, Sd1, Sq, yI, rho + pre.C0, + pre.C1, + pre.X0, + pre.X1, + Xd0, + Xd1, + pre.Xq, + pre.S0, + pre.S1, + Sd0, + Sd1, + pre.Sq, + yI, + rho, ) if update_idx == 2: for _ in range(Nstep_yI): - yI = update_yI(C0, C1, X0, X1, S0, S1, yE, yEq, yI, rho, Lambda) - X0, X1, Xd0, Xd1, Xq, res_X = update_X( - C0, C1, X0, X1, Xd0, Xd1, Xq, yE, yEq, yI, S0, S1, Sd0, Sd1, Sq, rho + yI = update_yI( + pre.C0, + pre.C1, + pre.X0, + pre.X1, + pre.S0, + pre.S1, + yE, + yEq, + yI, + rho, + pre.Lambda, + ) + pre.X0, pre.X1, Xd0, Xd1, pre.Xq, res_X = update_X( + pre.C0, + pre.C1, + pre.X0, + pre.X1, + Xd0, + Xd1, + pre.Xq, + yE, + yEq, + yI, + pre.S0, + pre.S1, + Sd0, + Sd1, + pre.Sq, + rho, ) if t % 100 == 0: print_updates(verbose) rho, p_resnorm, d_resnorm = update_rho( - X0, X1, Xd0, Xd1, Xq, bE, bEq, bI, res_X, rho, factor, normC + pre.X0, + pre.X1, + Xd0, + Xd1, + pre.Xq, + bE, + pre.bEq, + pre.bI, + res_X, + rho, + factor, + pre.normC, ) # ---------------------------- Assemble degree-wise results ----------------------------- # Convert the optimized packed block variables back to full degree-wise # representation matrices for Euler-angle recovery. - X_admm = self._transform_coeff_back(X0, X1, IDX_upper, IDX_lower, idx_offdiag) + X_admm = self._transform_coeff_back( + pre.X0, pre.X1, pre.IDX_upper, pre.IDX_lower, pre.idx_offdiag + ) X_admm = [xp.asnumpy(Xk) for Xk in X_admm] return X_admm @@ -1340,7 +1455,8 @@ def ADMM_preprocessing(self, C): :param C: Fourier coefficient matrices of the NUG objective. - :return: Dictionary of quantities required by the ADMM solver. + :return: An _ADMMData instance containing the constraints, layout, + and ADMM variables. """ Lmax = self.Lmax N = self.n_img @@ -1457,30 +1573,30 @@ def ADMM_preprocessing(self, C): self.Ngrid = Ngrid - return { - "C0": C0, - "C1": C1, - "normC": normC, - "AEq": AEq, - "bEq": bEq, - "AEqAEqtinv": AEqAEqtinv, - "AI_mat_diag": AI_mat_diag, - "AI_mat_offdiag": AI_mat_offdiag, - "bI": bI, - "Lambda": Lambda, - "d0": d0, - "d1": d1, - "idx_diag": idx_diag, - "idx_offdiag": idx_offdiag, - "IDX_upper": IDX_upper, - "IDX_lower": IDX_lower, - "X0": X0, - "X1": X1, - "Xq": Xq, - "S0": S0, - "S1": S1, - "Sq": Sq, - } + return _ADMMData( + C0=C0, + C1=C1, + normC=normC, + AEq=AEq, + bEq=bEq, + AEqAEqtinv=AEqAEqtinv, + AI_mat_diag=AI_mat_diag, + AI_mat_offdiag=AI_mat_offdiag, + bI=bI, + Lambda=Lambda, + d0=d0, + d1=d1, + idx_diag=idx_diag, + idx_offdiag=idx_offdiag, + IDX_upper=IDX_upper, + IDX_lower=IDX_lower, + X0=X0, + X1=X1, + Xq=Xq, + S0=S0, + S1=S1, + Sq=Sq, + ) def _compute_fejer_weights(self): """