P_Sparsity.py

from numpy import zeros_like
import numpy as np
from .proxoperators import ProxOperator


class P_Sparsity(ProxOperator):
    def __init__(self, config):
        """
        Initialization

        Parameters
        ----------
        config : dict - Dictionary containing the problem configuration. It must contain the following mapping:
        'sparsity_parameter' : int
        """
        self.sparsity_parameter = config['sparsity_parameter']
        self.ny = config['Ny']

        if 'sparsity_support' in config:
            self.support = config['sparsity_support'].real.astype(np.uint)
        else:
            self.support = 1

        if self.sparsity_parameter > 30:
            def value_selection(original, indices, sparsity_parameter):
                idx_for_threshold = divmod(indices[-sparsity_parameter], self.ny)
                threshold_val = abs(original[idx_for_threshold].get())
                return (abs(original) >= threshold_val) * original
        else:
            def value_selection(original, indices, sparsity_parameter):
                out = zeros_like(original)
                hits = indices[-sparsity_parameter:].get()
                hit_idx = [divmod(hit, self.ny) for hit in hits]
                for _idx in hit_idx:
                    out[_idx[0], _idx[1]] = original[_idx[0], _idx[1]]
                return out

        self.value_selection = value_selection

    def work(self, u):
        """
        Parameters
        ----------
        u : array_like - Input

        Returns
        -------
        p_Sparsity : array_like
        """
        u *= self.support # apply support
        p_Sparsity = 0 * u
        sorting = np.argsort(abs(u), axis=None)  # gives indices of sorted array in ascending order
        indices = np.asarray([divmod(sorting[i], self.ny) for i in range(-1 * self.sparsity_parameter, 0)])
        p_Sparsity[indices[:, 0], indices[:, 1]] = u[indices[:, 0], indices[:, 1]]

        return p_Sparsity


class P_Sparsity_real(P_Sparsity):
    def __init__(self, config):
        super(P_Sparsity_real, self).__init__(config)

    def work(self, u):
        return super(P_Sparsity_real, self).work(u.real)


class P_Sparsity_Symmetric(P_Sparsity):
    def __init__(self, config):
        super(P_Sparsity_Symmetric, self).__init__(config=config)
        self.symmetry = config['symmetry_type'] # antisymmetric = -1, symmetric = 1
        self.symmetry_axis = config['symmetry_axis']  # -1 for last, 0 for first, 'both', 'all' or None for full flip

    def work(self, u):
        if self.symmetry_axis in ['both', 'all']:
            mirror = np.flip(u, axis=None)
        else:
            mirror = np.flip(u, axis=self.symmetry_axis)
        inp = (u + self.symmetry * mirror) / 2
        out = super(P_Sparsity_Symmetric, self).work(inp)
        return out


class P_Sparsity_Symmetric_real(P_Sparsity_Symmetric):
    def __init__(self, config):
        super(P_Sparsity_Symmetric_real, self).__init__(config=config)

    def work(self, u):
        return super(P_Sparsity_Symmetric_real, self).work(u.real)