complex float* compute_psf(unsigned int N, const long img2_dims[N], const long trj_dims[N], const complex float* traj, const complex float* weights)
{
long ksp_dims1[N];
md_select_dims(N, ~MD_BIT(0), ksp_dims1, trj_dims);
struct linop_s* op2 = nufft_create(N, ksp_dims1, img2_dims, trj_dims, traj, NULL,
nufft_conf_defaults, false);
complex float* ones = md_alloc(N, ksp_dims1, CFL_SIZE);
md_zfill(N, ksp_dims1, ones, 1.);
if (NULL != weights) {
md_zmul(N, ksp_dims1, ones, ones, weights);
md_zmulc(N, ksp_dims1, ones, ones, weights);
}
complex float* psft = md_alloc(N, img2_dims, CFL_SIZE);
linop_adjoint_unchecked(op2, psft, ones);
md_free(ones);
linop_free(op2);
return psft;
}
void noir_recon(const struct noir_conf_s* conf, const long dims[DIMS], complex float* outbuf, complex float* sensout, const complex float* psf, const complex float* mask, const complex float* kspace)
{
long imgs_dims[DIMS];
long coil_dims[DIMS];
long data_dims[DIMS];
long img1_dims[DIMS];
md_select_dims(DIMS, FFT_FLAGS|MAPS_FLAG|CSHIFT_FLAG, imgs_dims, dims);
md_select_dims(DIMS, FFT_FLAGS|COIL_FLAG|MAPS_FLAG, coil_dims, dims);
md_select_dims(DIMS, FFT_FLAGS|COIL_FLAG, data_dims, dims);
md_select_dims(DIMS, FFT_FLAGS, img1_dims, dims);
long skip = md_calc_size(DIMS, imgs_dims);
long size = skip + md_calc_size(DIMS, coil_dims);
long data_size = md_calc_size(DIMS, data_dims);
long d1[1] = { size };
complex float* img = md_alloc_sameplace(1, d1, CFL_SIZE, kspace);
complex float* imgH = md_alloc_sameplace(1, d1, CFL_SIZE, kspace);
md_clear(DIMS, imgs_dims, img, CFL_SIZE);
md_zfill(DIMS, img1_dims, outbuf, 1.); // initial only first image
md_copy(DIMS, img1_dims, img, outbuf, CFL_SIZE);
md_clear(DIMS, coil_dims, img + skip, CFL_SIZE);
md_clear(DIMS, imgs_dims, imgH, CFL_SIZE);
md_clear(DIMS, coil_dims, imgH + skip, CFL_SIZE);
struct noir_data* ndata = noir_init(dims, mask, psf, conf->rvc, conf->usegpu);
struct data data = { ndata };
struct iter3_irgnm_conf irgnm_conf = { .iter = conf->iter, .alpha = conf->alpha, .redu = conf->redu };
iter3_irgnm(&irgnm_conf.base, frw, der, adj, &data, size * 2, (float*)img, data_size * 2, (const float*)kspace);
md_copy(DIMS, imgs_dims, outbuf, img, CFL_SIZE);
if (NULL != sensout) {
assert(!conf->usegpu);
noir_forw_coils(ndata, sensout, img + skip);
}
noir_free(ndata);
md_free(img);
md_free(imgH);
}
int main_nlinv(int argc, char* argv[])
{
int iter = 8;
float l1 = -1.;
bool waterfat = false;
bool rvc = false;
bool normalize = true;
float restrict_fov = -1.;
float csh[3] = { 0., 0., 0. };
bool usegpu = false;
const char* psf = NULL;
const struct opt_s opts[] = {
{ 'l', true, opt_float, &l1, NULL },
{ 'i', true, opt_int, &iter, NULL },
{ 'c', false, opt_set, &rvc, NULL },
{ 'N', false, opt_clear, &normalize, NULL },
{ 'f', true, opt_float, &restrict_fov, NULL },
{ 'p', true, opt_string, &psf, NULL },
{ 'g', false, opt_set, &usegpu, NULL },
};
cmdline(&argc, argv, 2, 3, usage_str, help_str, ARRAY_SIZE(opts), opts);
num_init();
assert(iter > 0);
long ksp_dims[DIMS];
complex float* kspace_data = load_cfl(argv[1], DIMS, ksp_dims);
long dims[DIMS];
md_copy_dims(DIMS, dims, ksp_dims);
if (waterfat)
dims[CSHIFT_DIM] = 2;
long img_dims[DIMS];
md_select_dims(DIMS, FFT_FLAGS|CSHIFT_FLAG, img_dims, dims);
long img_strs[DIMS];
md_calc_strides(DIMS, img_strs, img_dims, CFL_SIZE);
complex float* image = create_cfl(argv[2], DIMS, img_dims);
long msk_dims[DIMS];
md_select_dims(DIMS, FFT_FLAGS, msk_dims, dims);
long msk_strs[DIMS];
md_calc_strides(DIMS, msk_strs, msk_dims, CFL_SIZE);
complex float* mask;
complex float* norm = md_alloc(DIMS, msk_dims, CFL_SIZE);
complex float* sens;
if (4 == argc) {
sens = create_cfl(argv[3], DIMS, ksp_dims);
} else {
sens = md_alloc(DIMS, ksp_dims, CFL_SIZE);
}
complex float* pattern = NULL;
long pat_dims[DIMS];
if (NULL != psf) {
pattern = load_cfl(psf, DIMS, pat_dims);
// FIXME: check compatibility
} else {
pattern = md_alloc(DIMS, img_dims, CFL_SIZE);
estimate_pattern(DIMS, ksp_dims, COIL_DIM, pattern, kspace_data);
}
if (waterfat) {
size_t size = md_calc_size(DIMS, msk_dims);
md_copy(DIMS, msk_dims, pattern + size, pattern, CFL_SIZE);
long shift_dims[DIMS];
md_select_dims(DIMS, FFT_FLAGS, shift_dims, msk_dims);
long shift_strs[DIMS];
md_calc_strides(DIMS, shift_strs, shift_dims, CFL_SIZE);
complex float* shift = md_alloc(DIMS, shift_dims, CFL_SIZE);
unsigned int X = shift_dims[READ_DIM];
unsigned int Y = shift_dims[PHS1_DIM];
unsigned int Z = shift_dims[PHS2_DIM];
for (unsigned int x = 0; x < X; x++)
for (unsigned int y = 0; y < Y; y++)
for (unsigned int z = 0; z < Z; z++)
shift[(z * Z + y) * Y + x] = cexp(2.i * M_PI * ((csh[0] * x) / X + (csh[1] * y) / Y + (csh[2] * z) / Z));
md_zmul2(DIMS, msk_dims, msk_strs, pattern + size, msk_strs, pattern + size, shift_strs, shift);
md_free(shift);
}
#if 0
float scaling = 1. / estimate_scaling(ksp_dims, NULL, kspace_data);
#else
float scaling = 100. / md_znorm(DIMS, ksp_dims, kspace_data);
#endif
debug_printf(DP_INFO, "Scaling: %f\n", scaling);
md_zsmul(DIMS, ksp_dims, kspace_data, kspace_data, scaling);
if (-1. == restrict_fov) {
mask = md_alloc(DIMS, msk_dims, CFL_SIZE);
md_zfill(DIMS, msk_dims, mask, 1.);
} else {
float restrict_dims[DIMS] = { [0 ... DIMS - 1] = 1. };
restrict_dims[0] = restrict_fov;
restrict_dims[1] = restrict_fov;
restrict_dims[2] = restrict_fov;
mask = compute_mask(DIMS, msk_dims, restrict_dims);
}
#ifdef USE_CUDA
if (usegpu) {
complex float* kspace_gpu = md_alloc_gpu(DIMS, ksp_dims, CFL_SIZE);
md_copy(DIMS, ksp_dims, kspace_gpu, kspace_data, CFL_SIZE);
noir_recon(dims, iter, l1, image, NULL, pattern, mask, kspace_gpu, rvc, usegpu);
md_free(kspace_gpu);
md_zfill(DIMS, ksp_dims, sens, 1.);
} else
#endif
noir_recon(dims, iter, l1, image, sens, pattern, mask, kspace_data, rvc, usegpu);
if (normalize) {
md_zrss(DIMS, ksp_dims, COIL_FLAG, norm, sens);
md_zmul2(DIMS, img_dims, img_strs, image, img_strs, image, msk_strs, norm);
}
if (4 == argc) {
long strs[DIMS];
md_calc_strides(DIMS, strs, ksp_dims, CFL_SIZE);
if (norm)
md_zdiv2(DIMS, ksp_dims, strs, sens, strs, sens, img_strs, norm);
fftmod(DIMS, ksp_dims, FFT_FLAGS, sens, sens);
unmap_cfl(DIMS, ksp_dims, sens);
} else {
md_free(sens);
}
md_free(norm);
md_free(mask);
if (NULL != psf)
unmap_cfl(DIMS, pat_dims, pattern);
else
md_free(pattern);
unmap_cfl(DIMS, img_dims, image);
unmap_cfl(DIMS, ksp_dims, kspace_data);
exit(0);
}
/**
*
* NUFFT operator initialization
*
* @param N - number of dimensions
* @param ksp_dims - kspace dimension
* @param cim_dims - coil images dimension
* @param traj - trajectory
* @param conf - configuration options
* @param use_gpu - use gpu boolean
*
*/
struct linop_s* nufft_create(unsigned int N, const long ksp_dims[N], const long cim_dims[N], const long traj_dims[N], const complex float* traj, const complex float* weights, struct nufft_conf_s conf, bool use_gpu)
{
struct nufft_data* data = (struct nufft_data*)xmalloc(sizeof(struct nufft_data));
data->N = N;
data->use_gpu = use_gpu;
data->traj = traj;
data->conf = conf;
data->width = 3.;
data->beta = calc_beta(2., data->width);
// get dims
assert(md_check_compat(N - 3, 0, ksp_dims + 3, cim_dims + 3));
unsigned int ND = N + 3;
data->ksp_dims = xmalloc(ND * sizeof(long));
data->cim_dims = xmalloc(ND * sizeof(long));
data->cml_dims = xmalloc(ND * sizeof(long));
data->img_dims = xmalloc(ND * sizeof(long));
data->trj_dims = xmalloc(ND * sizeof(long));
data->lph_dims = xmalloc(ND * sizeof(long));
data->psf_dims = xmalloc(ND * sizeof(long));
data->wgh_dims = xmalloc(ND * sizeof(long));
data->ksp_strs = xmalloc(ND * sizeof(long));
data->cim_strs = xmalloc(ND * sizeof(long));
data->cml_strs = xmalloc(ND * sizeof(long));
data->img_strs = xmalloc(ND * sizeof(long));
data->trj_strs = xmalloc(ND * sizeof(long));
data->lph_strs = xmalloc(ND * sizeof(long));
data->psf_strs = xmalloc(ND * sizeof(long));
data->wgh_strs = xmalloc(ND * sizeof(long));
md_singleton_dims(ND, data->cim_dims);
md_singleton_dims(ND, data->ksp_dims);
md_copy_dims(N, data->cim_dims, cim_dims);
md_copy_dims(N, data->ksp_dims, ksp_dims);
md_select_dims(ND, FFT_FLAGS, data->img_dims, data->cim_dims);
assert(3 == traj_dims[0]);
assert(traj_dims[1] == ksp_dims[1]);
assert(traj_dims[2] == ksp_dims[2]);
assert(md_check_compat(N - 3, ~0, traj_dims + 3, ksp_dims + 3));
assert(md_check_bounds(N - 3, ~0, traj_dims + 3, ksp_dims + 3));
md_singleton_dims(ND, data->trj_dims);
md_copy_dims(N, data->trj_dims, traj_dims);
// get strides
md_calc_strides(ND, data->cim_strs, data->cim_dims, CFL_SIZE);
md_calc_strides(ND, data->img_strs, data->img_dims, CFL_SIZE);
md_calc_strides(ND, data->trj_strs, data->trj_dims, CFL_SIZE);
md_calc_strides(ND, data->ksp_strs, data->ksp_dims, CFL_SIZE);
data->weights = NULL;
if (NULL != weights) {
md_singleton_dims(ND, data->wgh_dims);
md_select_dims(N, ~MD_BIT(0), data->wgh_dims, data->trj_dims);
md_calc_strides(ND, data->wgh_strs, data->wgh_dims, CFL_SIZE);
complex float* tmp = md_alloc(ND, data->wgh_dims, CFL_SIZE);
md_copy(ND, data->wgh_dims, tmp, weights, CFL_SIZE);
data->weights = tmp;
}
complex float* roll = md_alloc(ND, data->img_dims, CFL_SIZE);
rolloff_correction(2., data->width, data->beta, data->img_dims, roll);
data->roll = roll;
complex float* linphase = compute_linphases(N, data->lph_dims, data->img_dims);
md_calc_strides(ND, data->lph_strs, data->lph_dims, CFL_SIZE);
if (!conf.toeplitz)
md_zmul2(ND, data->lph_dims, data->lph_strs, linphase, data->lph_strs, linphase, data->img_strs, data->roll);
fftmod(ND, data->lph_dims, FFT_FLAGS, linphase, linphase);
fftscale(ND, data->lph_dims, FFT_FLAGS, linphase, linphase);
// md_zsmul(ND, data->lph_dims, linphase, linphase, 1. / (float)(data->trj_dims[1] * data->trj_dims[2]));
complex float* fftm = md_alloc(ND, data->img_dims, CFL_SIZE);
md_zfill(ND, data->img_dims, fftm, 1.);
fftmod(ND, data->img_dims, FFT_FLAGS, fftm, fftm);
data->fftmod = fftm;
data->linphase = linphase;
data->psf = NULL;
if (conf.toeplitz) {
#if 0
md_copy_dims(ND, data->psf_dims, data->lph_dims);
#else
md_copy_dims(3, data->psf_dims, data->lph_dims);
md_copy_dims(ND - 3, data->psf_dims + 3, data->trj_dims + 3);
data->psf_dims[N] = data->lph_dims[N];
#endif
md_calc_strides(ND, data->psf_strs, data->psf_dims, CFL_SIZE);
data->psf = compute_psf2(N, data->psf_dims, data->trj_dims, data->traj, data->weights);
}
md_copy_dims(ND, data->cml_dims, data->cim_dims);
data->cml_dims[N + 0] = data->lph_dims[N + 0];
md_calc_strides(ND, data->cml_strs, data->cml_dims, CFL_SIZE);
data->cm2_dims = xmalloc(ND * sizeof(long));
// !
md_copy_dims(ND, data->cm2_dims, data->cim_dims);
for (int i = 0; i < 3; i++)
data->cm2_dims[i] = (1 == cim_dims[i]) ? 1 : (2 * cim_dims[i]);
data->grid = md_alloc(ND, data->cml_dims, CFL_SIZE);
data->fft_op = linop_fft_create(ND, data->cml_dims, FFT_FLAGS, use_gpu);
return linop_create(N, ksp_dims, N, cim_dims,
data, nufft_apply, nufft_apply_adjoint, nufft_apply_normal, NULL, nufft_free_data);
}
struct noir_data* noir_init(const long dims[DIMS], const complex float* mask, const complex float* psf, bool rvc, bool use_gpu)
{
#ifdef USE_CUDA
md_alloc_fun_t my_alloc = use_gpu ? md_alloc_gpu : md_alloc;
#else
assert(!use_gpu);
md_alloc_fun_t my_alloc = md_alloc;
#endif
PTR_ALLOC(struct noir_data, data);
data->rvc = rvc;
md_copy_dims(DIMS, data->dims, dims);
md_select_dims(DIMS, FFT_FLAGS|COIL_FLAG|CSHIFT_FLAG, data->sign_dims, dims);
md_calc_strides(DIMS, data->sign_strs, data->sign_dims, CFL_SIZE);
md_select_dims(DIMS, FFT_FLAGS|COIL_FLAG|MAPS_FLAG, data->coil_dims, dims);
md_calc_strides(DIMS, data->coil_strs, data->coil_dims, CFL_SIZE);
md_select_dims(DIMS, FFT_FLAGS|MAPS_FLAG|CSHIFT_FLAG, data->imgs_dims, dims);
md_calc_strides(DIMS, data->imgs_strs, data->imgs_dims, CFL_SIZE);
md_select_dims(DIMS, FFT_FLAGS|COIL_FLAG, data->data_dims, dims);
md_calc_strides(DIMS, data->data_strs, data->data_dims, CFL_SIZE);
md_select_dims(DIMS, FFT_FLAGS, data->mask_dims, dims);
md_calc_strides(DIMS, data->mask_strs, data->mask_dims, CFL_SIZE);
md_select_dims(DIMS, FFT_FLAGS, data->wght_dims, dims);
md_calc_strides(DIMS, data->wght_strs, data->wght_dims, CFL_SIZE);
md_select_dims(DIMS, FFT_FLAGS|CSHIFT_FLAG, data->ptrn_dims, dims);
md_calc_strides(DIMS, data->ptrn_strs, data->ptrn_dims, CFL_SIZE);
complex float* weights = md_alloc(DIMS, data->wght_dims, CFL_SIZE);
noir_calc_weights(dims, weights);
fftmod(DIMS, data->wght_dims, FFT_FLAGS, weights, weights);
fftscale(DIMS, data->wght_dims, FFT_FLAGS, weights, weights);
data->weights = weights;
#ifdef USE_CUDA
if (use_gpu) {
data->weights = md_gpu_move(DIMS, data->wght_dims, weights, CFL_SIZE);
}
#endif
complex float* ptr = my_alloc(DIMS, data->ptrn_dims, CFL_SIZE);
md_copy(DIMS, data->ptrn_dims, ptr, psf, CFL_SIZE);
fftmod(DIMS, data->ptrn_dims, FFT_FLAGS, ptr, ptr);
data->pattern = ptr;
complex float* msk = my_alloc(DIMS, data->mask_dims, CFL_SIZE);
if (NULL == mask) {
assert(!use_gpu);
md_zfill(DIMS, data->mask_dims, msk, 1.);
} else {
md_copy(DIMS, data->mask_dims, msk, mask, CFL_SIZE);
}
// fftmod(DIMS, data->mask_dims, 7, msk, msk);
fftscale(DIMS, data->mask_dims, FFT_FLAGS, msk, msk);
data->mask = msk;
data->sens = my_alloc(DIMS, data->coil_dims, CFL_SIZE);
data->xn = my_alloc(DIMS, data->imgs_dims, CFL_SIZE);
data->tmp = my_alloc(DIMS, data->sign_dims, CFL_SIZE);
return data;
}
