/**
* Proximal operator for l1-norm with unitary transform: f(x) = lambda || T x ||_1
*
* @param D number of dimensions
* @param dim dimensions of x
* @param lambda threshold parameter
* @param unitary_op unitary linear operator
* @param flags bitmask for joint soft-thresholding
*/
extern const struct operator_p_s* prox_unithresh_create(unsigned int D, const struct linop_s* unitary_op, const float lambda, const unsigned long flags)
{
PTR_ALLOC(struct thresh_s, data);
SET_TYPEID(thresh_s, data);
data->lambda = lambda;
data->D = D;
data->flags = flags;
data->unitary_op = unitary_op;
const long* dims = linop_domain(unitary_op)->dims;
PTR_ALLOC(long[D], ndim);
md_copy_dims(D, *ndim, dims);
data->dim = *PTR_PASS(ndim);
PTR_ALLOC(long[D], nstr);
md_calc_strides(D, *nstr, data->dim, CFL_SIZE);
data->str = *PTR_PASS(nstr);
// norm dimensions are the flagged transform dimensions
// FIXME should use linop_codomain(unitary_op)->N
PTR_ALLOC(long[D], norm_dim);
md_select_dims(D, ~flags, *norm_dim, linop_codomain(unitary_op)->dims);
data->norm_dim = *PTR_PASS(norm_dim);
return operator_p_create(D, dims, D, dims, CAST_UP(PTR_PASS(data)), unisoftthresh_apply, thresh_del);
}
const struct operator_p_s* prox_lineq_create(const struct linop_s* op, const complex float* y)
{
PTR_ALLOC(struct prox_lineq_data, pdata);
unsigned int N = linop_domain(op)->N;
const long* dims = linop_domain(op)->dims;
pdata->op = op;
pdata->adj = md_alloc_sameplace(N, dims, CFL_SIZE, y);
linop_adjoint(op, N, dims, pdata->adj, N, linop_codomain(op)->dims, y);
pdata->tmp = md_alloc_sameplace(N, dims, CFL_SIZE, y);
return operator_p_create(N, dims, N, dims, CAST_UP(PTR_PASS(pdata)), prox_lineq_apply, prox_lineq_del);
}
/**
* Efficiently chain two matrix linops by multiplying the actual matrices together.
* Stores a copy of the new matrix.
* Returns: C = B A
*
* @param a first matrix (applied to input)
* @param b second matrix (applied to output of first matrix)
*/
struct linop_s* linop_matrix_chain(const struct linop_s* a, const struct linop_s* b)
{
const struct operator_matrix_s* a_data = linop_get_data(a);
const struct operator_matrix_s* b_data = linop_get_data(b);
// check compatibility
assert(linop_codomain(a)->N == linop_domain(b)->N);
assert(md_calc_size(linop_codomain(a)->N, linop_codomain(a)->dims) == md_calc_size(linop_domain(b)->N, linop_domain(b)->dims));
assert(a_data->K_dim != b_data->T_dim); // FIXME error for now -- need to deal with this specially.
assert((a_data->T_dim == b_data->K_dim) && (a_data->T == b_data->K));
unsigned int N = linop_domain(a)->N;
long max_dims[N];
md_singleton_dims(N, max_dims);
max_dims[a_data->T_dim] = a_data->T;
max_dims[a_data->K_dim] = a_data->K;
max_dims[b_data->T_dim] = b_data->T;
long matrix_dims[N];
long matrix_strs[N];
md_select_dims(N, ~MD_BIT(a_data->T_dim), matrix_dims, max_dims);
md_calc_strides(N, matrix_strs, matrix_dims, CFL_SIZE);
complex float* matrix = md_alloc_sameplace(N, matrix_dims, CFL_SIZE, a_data->mat);
md_clear(N, matrix_dims, matrix, CFL_SIZE);
md_zfmac2(N, max_dims, matrix_strs, matrix, a_data->mat_iovec->strs, a_data->mat, b_data->mat_iovec->strs, b_data->mat);
struct linop_s* c = linop_matrix_create(N, linop_codomain(b)->dims, linop_domain(a)->dims, matrix_dims, matrix);
md_free(matrix);
return c;
}
static void prox_lineq_apply(const operator_data_t* _data, float mu, complex float* dst, const complex float* src)
{
UNUSED(mu);
struct prox_lineq_data* pdata = CAST_DOWN(prox_lineq_data, _data);
const struct linop_s* op = pdata->op;
linop_normal(op, linop_domain(op)->N, linop_domain(op)->dims, pdata->tmp, src);
md_zsub(linop_domain(op)->N, linop_domain(op)->dims, dst, src, pdata->tmp);
md_zadd(linop_domain(op)->N, linop_domain(op)->dims, dst, dst, pdata->adj);
}
/* Chambolle Pock Primal Dual algorithm. Solves G(x) + F(Ax)
* Assumes that G is in prox_ops[0], F is in prox_ops[1], A is in ops[1]
*/
void iter2_chambolle_pock(iter_conf* _conf,
const struct operator_s* normaleq_op,
unsigned int D,
const struct operator_p_s* prox_ops[D],
const struct linop_s* ops[D],
const float* biases[D],
const struct operator_p_s* xupdate_op,
long size, float* image, const float* image_adj,
struct iter_monitor_s* monitor)
{
assert(D == 2);
assert(NULL == biases);
assert(NULL == normaleq_op);
UNUSED(xupdate_op);
UNUSED(image_adj);
auto conf = CAST_DOWN(iter_chambolle_pock_conf, _conf);
const struct iovec_s* iv = linop_domain(ops[1]);
const struct iovec_s* ov = linop_codomain(ops[1]);
assert((long)md_calc_size(iv->N, iv->dims) * 2 == size);
// FIXME: sensible way to check for corrupt data?
#if 0
float eps = md_norm(1, MD_DIMS(size), image_adj);
if (checkeps(eps))
goto cleanup;
#else
float eps = 1.;
#endif
chambolle_pock(conf->maxiter, eps * conf->tol, conf->tau, conf->sigma, conf->theta, conf->decay, 2 * md_calc_size(iv->N, iv->dims), 2 * md_calc_size(ov->N, ov->dims), select_vecops(image),
OPERATOR2ITOP(ops[1]->forward), OPERATOR2ITOP(ops[1]->adjoint), OPERATOR_P2ITOP(prox_ops[1]), OPERATOR_P2ITOP(prox_ops[0]),
image, monitor);
//cleanup:
//;
}
const struct operator_p_s* prox_normaleq_create(const struct linop_s* op, const complex float* y)
{
PTR_ALLOC(struct prox_normaleq_data, pdata);
SET_TYPEID(prox_normaleq_data, pdata);
PTR_ALLOC(struct iter_conjgrad_conf, cgconf);
*cgconf = iter_conjgrad_defaults;
cgconf->maxiter = 10;
cgconf->l2lambda = 0;
pdata->cgconf = PTR_PASS(cgconf);
pdata->op = op;
pdata->size = 2 * md_calc_size(linop_domain(op)->N, linop_domain(op)->dims);
pdata->adj = md_alloc_sameplace(1, &(pdata->size), FL_SIZE, y);
linop_adjoint_iter((struct linop_s*)op, pdata->adj, (const float*)y);
return operator_p_create(linop_domain(op)->N, linop_domain(op)->dims,
linop_domain(op)->N, linop_domain(op)->dims,
CAST_UP(PTR_PASS(pdata)), prox_normaleq_apply, prox_normaleq_del);
}
/**
* Efficiently chain two matrix linops by multiplying the actual matrices together.
* Stores a copy of the new matrix.
* Returns: C = B A
*
* @param a first matrix (applied to input)
* @param b second matrix (applied to output of first matrix)
*/
struct linop_s* linop_matrix_chain(const struct linop_s* a, const struct linop_s* b)
{
const struct operator_matrix_s* a_data = CAST_DOWN(operator_matrix_s, linop_get_data(a));
const struct operator_matrix_s* b_data = CAST_DOWN(operator_matrix_s, linop_get_data(b));
// check compatibility
assert(linop_codomain(a)->N == linop_domain(b)->N);
assert(md_check_compat(linop_codomain(a)->N, 0u, linop_codomain(a)->dims, linop_domain(b)->dims));
unsigned int D = linop_domain(a)->N;
unsigned long outB_flags = md_nontriv_dims(D, linop_codomain(b)->dims);
unsigned long inB_flags = md_nontriv_dims(D, linop_domain(b)->dims);
unsigned long delB_flags = inB_flags & ~outB_flags;
unsigned int N = a_data->N;
assert(N == 2 * D);
long in_dims[N];
md_copy_dims(N, in_dims, a_data->in_dims);
long matA_dims[N];
md_copy_dims(N, matA_dims, a_data->mat_dims);
long matB_dims[N];
md_copy_dims(N, matB_dims, b_data->mat_dims);
long out_dims[N];
md_copy_dims(N, out_dims, b_data->out_dims);
for (unsigned int i = 0; i < D; i++) {
if (MD_IS_SET(delB_flags, i)) {
matA_dims[2 * i + 0] = a_data->mat_dims[2 * i + 1];
matA_dims[2 * i + 1] = a_data->mat_dims[2 * i + 0];
in_dims[2 * i + 0] = a_data->in_dims[2 * i + 1];
in_dims[2 * i + 1] = a_data->in_dims[2 * i + 0];
}
}
long matrix_dims[N];
md_singleton_dims(N, matrix_dims);
unsigned long iflags = md_nontriv_dims(N, in_dims);
unsigned long oflags = md_nontriv_dims(N, out_dims);
unsigned long flags = iflags | oflags;
// we combine a and b and sum over dims not in input or output
md_max_dims(N, flags, matrix_dims, matA_dims, matB_dims);
debug_printf(DP_DEBUG1, "tensor chain: %ld x %ld -> %ld\n",
md_calc_size(N, matA_dims), md_calc_size(N, matB_dims), md_calc_size(N, matrix_dims));
complex float* matrix = md_alloc(N, matrix_dims, CFL_SIZE);
debug_print_dims(DP_DEBUG2, N, matrix_dims);
debug_print_dims(DP_DEBUG2, N, in_dims);
debug_print_dims(DP_DEBUG2, N, matA_dims);
debug_print_dims(DP_DEBUG2, N, matB_dims);
debug_print_dims(DP_DEBUG2, N, out_dims);
md_ztenmul(N, matrix_dims, matrix, matA_dims, a_data->mat, matB_dims, b_data->mat);
// priv2 takes our doubled dimensions
struct operator_matrix_s* data = linop_matrix_priv2(N, out_dims, in_dims, matrix_dims, matrix);
/* although we internally use different dimensions we define the
* correct interface
*/
struct linop_s* c = linop_create(linop_codomain(b)->N, linop_codomain(b)->dims,
linop_domain(a)->N, linop_domain(a)->dims, CAST_UP(data),
linop_matrix_apply, linop_matrix_apply_adjoint,
linop_matrix_apply_normal, NULL, linop_matrix_del);
md_free(matrix);
return c;
}
