#optimizer=torch.optim.AdamW([{'params':G.parameters(),'lr':1e-3},{'params':GV.parameters(),'lr':1e-3}])
scheduler = ReduceLROnPlateau(optimizer,
mode='min',
factor=0.7,
patience=6,
verbose=True,
min_lr=5e-5)
trnFiles = os.listdir('/Shared/lss_jcb/abdul/prashant_cardiac_data/Data/d2/')
sz = len(trnFiles)
#data2=np.zeros((sz,1,n_select,N,N)).astype(np.complex64)
rndm = random.sample(range(sz), sz)
#%%
nuf_ob = KbNufft(im_size=(nx, nx), norm='ortho').to(dtype)
nuf_ob = nuf_ob.to(gpu)
adjnuf_ob = AdjKbNufft(im_size=(nx, nx), norm='ortho').to(dtype)
adjnuf_ob = adjnuf_ob.to(gpu)
smapT = torch.ones((1, 1, 2, nx, nx)).cuda()
nuf_ob = MriSenseNufft(im_size=(nx, nx), smap=smapT, norm='ortho').to(dtype)
nuf_ob = nuf_ob.to(gpu)
adjnuf_ob = AdjMriSenseNufft(im_size=(nx, nx), smap=smapT,
norm='ortho').to(dtype)
adjnuf_ob = adjnuf_ob.to(gpu)
d2 = sio.loadmat('/Shared/lss_jcb/abdul/prashant_cardiac_data/Data/ktraj.mat')
#%% take them to gpu
atbT = atbT.to(gpu)
ktrajT1 = ktrajT1.to(gpu)
smapT = smapT.to(gpu)
ktrajT = ktrajT.to(gpu)
#dcfT=dcfT.to(gpu)
LT = LT.to(gpu)
VT = VT.to(gpu)
sbasis = sbasis.to(gpu)
#%% generate atb
sigma = 0.0
lam = 1e3
cgIter = 1
cgTol = 1e-15
nuf_ob = KbNufft(im_size=im_size).to(dtype)
nuf_ob = nuf_ob.to(gpu)
adjnuf_ob = AdjKbNufft(im_size=im_size).to(dtype)
adjnuf_ob = adjnuf_ob.to(gpu)
#real_mat, imag_mat = precomp_sparse_mats(ktrajT, adjnuf_ob)
#interp_mats = {'real_interp_mats': real_mat, 'imag_interp_mats': imag_mat}
nufft_ob1 = MriSenseNufft(im_size=im_size, smap=smapT2).to(dtype)
nufft_ob1 = nufft_ob1.to(gpu)
nufft_ob = MriSenseNufft(im_size=im_size, smap=smapT).to(dtype)
nufft_ob = nufft_ob.to(gpu)
adjnufft_ob = AdjMriSenseNufft(im_size=im_size, smap=smapT).to(dtype)
adjnufft_ob = adjnufft_ob.to(gpu)
#nufft_ob = MriSenseNufft(smap=smapT,im_size=im_size).to(dtype)
#adjnufft_ob = AdjMriSenseNufft(smap=smapT,im_size=im_size ).to(dtype)
#%% x = shepp_logan_phantom().astype(np.complex) im_size = x.shape x = np.stack((np.real(x), np.imag(x))) # convert to tensor, unsqueeze batch and coil dimension # output size: (1, 1, 2, ny, nx) x = torch.tensor(x).unsqueeze(0).unsqueeze(0) klength = 64 ktraj = np.stack((np.zeros(64), np.linspace(-np.pi, np.pi, klength))) # convert to tensor, unsqueeze batch dimension # output size: (1, 2, klength) ktraj = torch.tensor(ktraj).unsqueeze(0) nufft_ob = KbNufft(im_size=im_size) # outputs a (1, 1, 2, klength) vector of k-space data kdata = nufft_ob(x, ktraj) #%% generate atb sigma = 0.0 lam = 1e-5 cgIter = 10 cgTol = 1e-15 A = lambda x: pt_A(x, csmT, maskT) At = lambda x: pt_At(x, csmT, maskT) bT = At(orgT) noiseT = torch.randn(bT.shape) * sigma noiseT = noiseT.to(gpu) bT = bT + noiseT
smap = np.stack((np.real(csmTrn), np.imag(csmTrn)), axis=1)
smap=np.tile(smap,(NF,1,1,1,1))
smapT = torch.tensor(smap).to(dtype)
# convert to tensor, unsqueeze batch dimension
# output size: (1, 2, klength)
ktraj = np.stack((np.real(ktraj), np.imag(ktraj)), axis=1)
ktrajT = torch.tensor(ktraj).to(dtype)
#%% take them to gpu
kdataT=kdataT.cuda()
smapT=smapT.cuda()
ktrajT=ktrajT.cuda()
dcfT=dcfT.cuda()
#%% generate atb (Initial reconstruction)
nuf_ob = KbNufft(im_size=im_size).to(dtype)
nuf_ob=nuf_ob.cuda()
adjnuf_ob = AdjKbNufft(im_size=im_size).to(dtype)
adjnuf_ob=adjnuf_ob.cuda()
real_mat, imag_mat = precomp_sparse_mats(ktrajT, adjnuf_ob)
interp_mats = {'real_interp_mats': real_mat, 'imag_interp_mats': imag_mat}
nufft_ob = MriSenseNufft(im_size=im_size, smap=smapT).to(dtype)
nufft_ob=nufft_ob.cuda()
adjnufft_ob = AdjMriSenseNufft(im_size=im_size, smap=smapT).to(dtype)
adjnufft_ob=adjnufft_ob.cuda()
At=lambda x: adjnufft_ob(x*dcfT,ktrajT,interp_mats)
A=lambda x: nufft_ob(x,ktrajT,interp_mats)
#inverse NUFFT transform -->Initial reconstruction
#GV.load_state_dict(torch.load('wtsVB-400.pt'))
#GV.load_state_dict(torch.load('./PTmodels/27Oct_112451am_500ep_27oct/wts-500.pt'))
#optimizer=torch.optim.SGD([{'params':G.parameters(),'lr':5e-3,'momentum':0.9}])
#optimizer=torch.optim.AdamW([{'params':G.parameters(),'lr':1e-4}])
optimizer=torch.optim.AdamW([{'params':G.parameters(),'lr':1e-4},{'params':GV.parameters(),'lr':1e-4}])
scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.7, patience=6, verbose=True, min_lr=1e-5)
trnFiles=os.listdir('/Shared/lss_jcb/abdul/prashant_cardiac_data/Data/d2/')
sz=len(trnFiles)
#data2=np.zeros((sz,1,n_select,N,N)).astype(np.complex64)
rndm=random.sample(range(sz),sz)
#%%
nuf_ob = KbNufft(im_size=(nx,nx)).to(dtype)
nuf_ob=nuf_ob.to(gpu)
adjnuf_ob = AdjKbNufft(im_size=(nx,nx)).to(dtype)
adjnuf_ob=adjnuf_ob.to(gpu)
smapT=torch.ones((1,1,2,nx,nx)).cuda()
nuf_ob = MriSenseNufft(im_size=(nx,nx),smap=smapT).to(dtype)
nuf_ob=nuf_ob.to(gpu)
adjnuf_ob = AdjMriSenseNufft(im_size=(nx,nx), smap=smapT).to(dtype)
adjnuf_ob=adjnuf_ob.to(gpu)