<commit_before>from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

from time import time
import os.path
from math import exp, log
import sys


import numpy as np
from scipy.special import gammaln
from numpy.linalg import inv
from scipy.ndimage import median_filter, gaussian_filter

from ._version import __version__
from .tools import calculate_speeds, center_image, center_image_asymmetric

#############################################################################
# This is adapted from the BASEX Matlab code provided by the Reisler group.
#
# Please cite: "The Gaussian basis-set expansion Abel transform method"
# <NAME>, <NAME>, <NAME>, and <NAME>,
# Review of Scientific Instruments 73, 2634 (2002).
#
# Version 0.6 - 2015-12-01
#   Another major code reformatting
# Version 0.5 - 2015-11-16
#   Code cleanup
# Version 0.4 - 2015-05-0
#   Major code update see pull request
#   https://github.com/DanHickstein/pyBASEX/pull/3
# Version 0.3 - 2015-02-01
#   Added documentation
# Version 0.2 - 2014-10-09
#   Adding a "center_and_transform" function to make things easier
# Versions 0.1 - 2012
#   First port to Python
#
#############################################################################


def BASEX(data, center, n, 
        nbf='auto',  basis_dir='./', calc_speeds=False, vertical_symmetry=False, dr=1.0, verbose=True):

    """ This function that centers the image, performs the BASEX transform (loads or generates basis sets), 
        and (optionally) calculates the radial integration of the image (calc_speeds)

    Parameters:
    -----------
      - data:  a NxM numpy array
            If data is smaller than the size of the basis set, zeros will be padded on the edges.
      - n:  * odd integer -
                Abel inverse transform will be performed on a `n x n` area of the image
            * list in format [n_vert, n_horz] - 
                Abel inverse transform will be performed on a `n[0] x n[1]` area of the image
      - nbf: * integer - 
                number of basis functions. If nbf='auto', it is set to n//2 + 1.
             * list in format [nbf_vert, nbf_horz] -
                If nbf='auto', it is set to [n_vert, n_horz//2 + 1]
      - center: * integer - 
                    the center column of the image
                * tuple (x,y) -
                    the center of the image in (x,y) format
      - basis_dir: string
            path to the directory for saving / loading the basis set coefficients.
            If None, the basis set will not be saved to disk. 
      - dr: float
            size of one pixel in the radial direction
      - calc_speeds: True/False
            determines if the speed distribution should be calculated
      - vertical_symmetry: True/False
            determines if the data has up/down symmetry 
      - verbose: True/False
            Set to True to see more output for debugging

    Returns:
       if calc_speeds=False: the processed image
       if calc_speeds=True:  the processes images, arrays with the calculated speeds

    """
    # make dimension-of-rawdata into list to account for rectangular n
    if type(n) is not list: n = [ n ] 

    # duplicate elements of n if n is single-valued (rawdata is square)
    if len(n) < 2: n = n*2

    # make sure n_horz is odd
    n[1] = 2 * (n[1] // 2) + 1 

    # make sure that the data is the right shape (1D must be converted to 2D)
    data = np.atleast_2d(data) # if passed a 1D array convert it to 2D
    if data.shape[0] == 1:
        data_ndim = 1
    elif data.shape[1] == 1:
        raise ValueError('Wrong input shape for data {0}, should be  (N1, N2) or (1, N), not (N, 1)'.format(data.shape))
    else:
        data_ndim = 2

    if vertical_symmetry:
        image = center_image(data, center=center, n=n[1], ndim=data_ndim)

        if verbose:
            t1 = time()

        M, Mc, M_left, M_right = get_basis_sets_cached(n[1], nbf, basis_dir, verbose)

        if verbose:
            print('{:.2f} seconds'.format((time() - t1)))

        # Do the actual transform
        if verbose:
            print('Reconstructing image...         ')
            t1 = time()

        recon = basex_transform(image, M, Mc, M_left, M_right, dr)

        if verbose:
            print('%.2f seconds' % (time() - t1))

        if data_ndim == 1:
            recon = recon[0, :] # taking one row, since they are all the same anyway

        if calc_speeds:
            if verbose:
                print('Generating speed distribution...')
                t1 = time()

            speeds = calculate_speeds(recon)

            if verbose:
                print('%.2f seconds' % (time() - t1))
            return recon, speeds
        else:
            return recon



def basex_transform(rawdata, M, Mc, M_left, M_right, dr=1.0):
    """ This is the internal function that does the actual BASEX transform

     Parameters
     ----------
      - rawdata: a NxN ndarray of the raw image.
      - M, Mc, M_left, M_right: 2D arrays given given by the basis set calculation function
      - dr: float: pixel size

     Returns
     -------
      IM: The abel-transformed image, a slice of the 3D distribution
    """

    Ci = M_left.dot(rawdata).dot(M_right)
    # P = dot(dot(Mc,Ci),M.T) # This calculates the projection, which should recreate the original image
    IM = Mc.dot(Ci).dot(Mc.T)

    # use an heuristic scaling factor to match the analytical abel transform
    # see https://github.com/PyAbel/PyAbel/issues/4
    # When nbf = n//2, we need the MAGIC_NUMBER = 8.053
    # When nbf = n//2 + 1, see below,
    MAGIC_NUMBER = 1.1122 
    IM *= MAGIC_NUMBER/dr

    return IM

def basex_transform_asymmetric(rawdata, M_vert, M_horz, Mc_vert, Mc_horz, vert_left, horz_right, dr=1.0):<commit_msg>Main implementation of asymmetric case.<commit_after>from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

from time import time
import os.path
from math import exp, log
import sys


import numpy as np
from scipy.special import gammaln
from numpy.linalg import inv
from scipy.ndimage import median_filter, gaussian_filter

from ._version import __version__
from .tools import calculate_speeds, center_image, center_image_asymmetric

#############################################################################
# This is adapted from the BASEX Matlab code provided by the Reisler group.
#
# Please cite: "The Gaussian basis-set expansion Abel transform method"
# <NAME>, <NAME>, <NAME>, and <NAME>,
# Review of Scientific Instruments 73, 2634 (2002).
#
# Version 0.6 - 2015-12-01
#   Another major code reformatting
# Version 0.5 - 2015-11-16
#   Code cleanup
# Version 0.4 - 2015-05-0
#   Major code update see pull request
#   https://github.com/DanHickstein/pyBASEX/pull/3
# Version 0.3 - 2015-02-01
#   Added documentation
# Version 0.2 - 2014-10-09
#   Adding a "center_and_transform" function to make things easier
# Versions 0.1 - 2012
#   First port to Python
#
#############################################################################


def BASEX(data, center, n, 
        nbf='auto',  basis_dir='./', calc_speeds=False, vertical_symmetry=False, dr=1.0, verbose=True):

    """ This function that centers the image, performs the BASEX transform (loads or generates basis sets), 
        and (optionally) calculates the radial integration of the image (calc_speeds)

    Parameters:
    -----------
      - data:  a NxM numpy array
            If data is smaller than the size of the basis set, zeros will be padded on the edges.
      - n:  * odd integer -
                Abel inverse transform will be performed on a `n x n` area of the image
            * list in format [n_vert, n_horz] - 
                Abel inverse transform will be performed on a `n[0] x n[1]` area of the image
      - nbf: * integer - 
                number of basis functions. If nbf='auto', it is set to n//2 + 1.
             * list in format [nbf_vert, nbf_horz] -
                If nbf='auto', it is set to [n_vert, n_horz//2 + 1]
      - center: * integer - 
                    the center column of the image
                * tuple (x,y) -
                    the center of the image in (x,y) format
      - basis_dir: string
            path to the directory for saving / loading the basis set coefficients.
            If None, the basis set will not be saved to disk. 
      - dr: float
            size of one pixel in the radial direction
      - calc_speeds: True/False
            determines if the speed distribution should be calculated
      - vertical_symmetry: True/False
            determines if the data has up/down symmetry 
      - verbose: True/False
            Set to True to see more output for debugging

    Returns:
       if calc_speeds=False: the processed image
       if calc_speeds=True:  the processes images, arrays with the calculated speeds

    """
    # make dimension-of-rawdata into list to account for rectangular n
    if type(n) is not list: n = [ n ] 

    # duplicate elements of n if n is single-valued (rawdata is square)
    if len(n) < 2: n = n*2

    # make sure n_horz is odd
    n[1] = 2 * (n[1] // 2) + 1 

    # make sure that the data is the right shape (1D must be converted to 2D)
    data = np.atleast_2d(data) # if passed a 1D array convert it to 2D
    if data.shape[0] == 1:
        data_ndim = 1
    elif data.shape[1] == 1:
        raise ValueError('Wrong input shape for data {0}, should be  (N1, N2) or (1, N), not (N, 1)'.format(data.shape))
    else:
        data_ndim = 2

    if vertical_symmetry:
        image = center_image(data, center=center, n=n[1], ndim=data_ndim)

        if verbose:
            t1 = time()

        M, Mc, M_left, M_right = get_basis_sets_cached(n[1], nbf, basis_dir, verbose)

        if verbose:
            print('{:.2f} seconds'.format((time() - t1)))

        # Do the actual transform
        if verbose:
            print('Reconstructing image...         ')
            t1 = time()

        recon = basex_transform(image, M, Mc, M_left, M_right, dr)

        if verbose:
            print('%.2f seconds' % (time() - t1))

        if data_ndim == 1:
            recon = recon[0, :] # taking one row, since they are all the same anyway

        if calc_speeds:
            if verbose:
                print('Generating speed distribution...')
                t1 = time()

            speeds = calculate_speeds(recon)

            if verbose:
                print('%.2f seconds' % (time() - t1))
            return recon, speeds
        else:
            return recon

    else: # No vertical symmetry
        if type(center) == tuple: cx, cy = center
        elif type(center) == int: cx = int(center)
        else: raise ValueError("Center specified incorrectly. Must be tuple (x,y) or integer (column)")

        image = center_image_asymmetric(data, center_column = cx, n_vert = n[0], n_horz = n[1], verbose)

        if verbose:
            t1 = time()

        M_vert, M_horz, Mc_vert, Mc_horz, vert_left, horz_right = get_basis_sets_cached_asymmetric(n_vert = n[0], n_horz = n[1], nbf, basis_dir, verbose)

        if verbose:
            print('{:.2f} seconds'.format((time() - t1)))

        #Do the actual transform
        if verbose:
            print('Reconstructing image...         ')
            t1 = time()

        recon = basex_transform_asymmetric(image, M_vert, M_horz, Mc_vert, Mc_horz, vert_left, horz_right, dr)

        if verbose:
            print('%.2f seconds' % (time() - t1))

        if data_ndim == 1:
            recon = recon[0, :] # taking one row, since they are all the same anyway

        # -------------------------------------------------
        # asymmetric speeds calculation not implemented yet
        # -------------------------------------------------
        # if calc_speeds:
        #     if verbose:
        #         print('Generating speed distribution...')
        #         t1 = time()

        #     speeds = calculate_speeds(recon, n)

        #     if verbose:
        #         print('%.2f seconds' % (time() - t1))
        #     return recon, speeds
        # else:
        #     return recon

        return recon




def basex_transform(rawdata, M, Mc, M_left, M_right, dr=1.0):
    """ This is the internal function that does the actual BASEX transform

     Parameters
     ----------
      - rawdata: a NxN ndarray of the raw image.
      - M, Mc, M_left, M_right: 2D arrays given given by the basis set calculation function
      - dr: float: pixel size

     Returns
     -------
      IM: The abel-transformed image, a slice of the 3D distribution
    """

    Ci = M_left.dot(rawdata).dot(M_right)
    # P = dot(dot(Mc,Ci),M.T) # This calculates the projection, which should recreate the original image
    IM = Mc.dot(Ci).dot(Mc.T)

    # use an heuristic scaling factor to match the analytical abel transform
    # see https://github.com/PyAbel/PyAbel/issues/4
    # When nbf = n//2, we need the MAGIC_NUMBER = 8.053
    # When nbf = n//2 + 1, see below,
    MAGIC_NUMBER = 1.1122 
    IM *= MAGIC_NUMBER/dr

    return IM

def basex_transform_asymmetric(rawdata, M_vert, M_horz, Mc_vert, Mc_horz, vert_left, horz_right, dr=1.0):
{'fraction_non_alphanumeric': 0.12209911663422668, 'fraction_numerical': 0.02043719119628687, 'mean_word_length': 2.2268115942028985, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 11, 'https://': 4, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6364463', 'n_tokens_mistral': 4226, 'n_tokens_neox': 3882, 'n_words': 1439}