Fit shm on reference data

approxSignal.py

from util import *
from rish import rish
N_shm=2

def approxSignal(imgPath, maskPath):

    directory = os.path.dirname(imgPath)
    inPrefix = imgPath.split('.nii')[0]
    prefix = os.path.split(inPrefix)[-1]
    outPrefix = os.path.join(directory, 'harm', prefix)

    b0, shm_coeff, qb_model = rish(imgPath, maskPath, inPrefix, outPrefix, N_shm)
    B = qb_model.B

    img= load(imgPath)
    hdr= img.header
    affine= img.affine

    mapped_cs= []
    shs_same_level= [[0, 1], [1, 6], [6, 15], [15, 28], [28, 45]]
    for i in range(0, N_shm+1, 2):

        # load data and mask
        fileName= os.path.join(directory, 'harm', f'{prefix}_L{i}.nii.gz')
        img= load(fileName).get_data()

        ind= int(i/2)
        for level in range(shs_same_level[ind][0], shs_same_level[ind][1]):
            mapped_cs.append(img * shm_coeff[ :,:,:,level])


    S_hat= np.dot(np.moveaxis(mapped_cs, 0, -1), B.T)
    # keep only upper half of the reconstructed signal
    S_hat= S_hat[..., :int(S_hat.shape[3]/2)]
    np.nan_to_num(S_hat).clip(min= 0., max= 1., out= S_hat)

    # affine= templateAffine for all Scale_L{i}
    mappedFile= os.path.join(directory, f'{prefix}_mapped_cs.nii.gz')
    save_nifti(mappedFile, S_hat, affine, hdr)

    # un-normalize approximated data
    S_hat_dwi= applymask(S_hat, b0) # overriding applymask function with a nonbinary mask b0

    # place b0s in proper indices
    S_hat_final= stack_b0(qb_model.gtab.b0s_mask, S_hat_dwi, b0)


    # save approximated data
    harmImg= os.path.join(directory, f'Reconstructed_{prefix}.nii.gz')
    save_nifti(harmImg, S_hat_final, affine, hdr)

    shutil.copyfile(inPrefix + '.bvec', harmImg.split('.nii')[0] + '.bvec')
    shutil.copyfile(inPrefix + '.bval', harmImg.split('.nii')[0] + '.bval')
    

if __name__ == '__main__':
    imgPath= 'test_data_release/Site2_Y_DWI_b1100.nii.gz'
    maskPath= 'test_data_release/Site2_Y_mask.nii.gz'
    approxSignal(imgPath, maskPath)