Fixes a bug in coupled decomposition where membership was not rotatin

import numpy as np

from scipy.stats import pearsonr

RADIAL_SIZE = 720

RADIAL_STEP = 2 * np.pi / RADIAL_SIZE

THETAS = np.round(np.arange(0, 2 * np.pi, RADIAL_STEP), 5)

def cart2polar(x, y):

    theta = np.arctan2(y, x)

    return theta

    return -theta

def index_coords(data, origin=None):

    """Creates x & y coords for the indicies in a numpy array "data".

    if origin is None:

        origin = (nx//2, ny//2)

    # Determine that the theta coords will be

    # Determine that the min and max r and theta coords will be...

    x, y = index_coords(data, origin=origin)

    theta = cart2polar(x, y)

    # -180 to 180 conversion to 0 to 360

    theta[theta < 0] += 2 * np.pi

    return theta

def coupled_decomposition(sdata, ddata, sorigin=(), dorigin=(), M=4, sub_skp=None):

def coupled_decomposition(sdata, ddata, sorigin=(), dorigin=(), M=4, theta_steps=720, theta=None):

    """

    Apply coupled decomposition to two 2d images.

    sdata : ndarray

    stheta = reproject_image_into_polar(sdata, origin=(int(soriginx), int(soriginy)))

    dtheta = reproject_image_into_polar(ddata, origin=(int(doriginx), int(doriginy)))

    if theta == None:

        # Compute the mean profiles for each radial slice

    smean = np.empty(RADIAL_SIZE)

    dmean = np.empty(RADIAL_SIZE)

    for i, t in enumerate(THETAS):

        # The way this method words, it is possible to get nan values in some of the steps as this is discrete

        smean[i] = np.mean(sdata[(t <= stheta) & (stheta <= t + RADIAL_STEP)])

        dmean[i] = np.mean(ddata[(t <= dtheta) & (dtheta <= t + RADIAL_STEP)])

        smean = np.empty(theta_steps)

        dmean = np.empty(theta_steps)

        radial_step = 2 * np.pi / theta_steps # 0.5 deg

        thetas = np.arange(0, 2 * np.pi, radial_step)

        for i, t in enumerate(thetas):

            smean[i] = np.nanmean(sdata[(stheta >= t) & (stheta <= t + radial_step)])

            dmean[i] = np.nanmean(ddata[(dtheta >= t) & (dtheta <= t + radial_step)])

        # Rotate the second image around the origin and compute the correlation coeff. for each 0.5 degree rotation.

        maxp = -1

        maxidx = 0

    for j in range(RADIAL_SIZE):

        dsearch=np.empty(theta_steps)

        for j in range(theta_steps):

            dsearch = np.concatenate((dmean[j:], dmean[:j]))

        r, p = pearsonr(smean, dsearch)

            r, _ = pearsonr(smean, dsearch)

            if r >= maxp:

                maxp = r

                maxidx = j

        # Maximum correlation (theta) defines the angle of rotation for the destination image

    theta = THETAS[maxidx]

    if theta <= np.pi:

        lam = theta

    else:

        lam = 2 * np.pi - theta

        theta = thetas[maxidx]

    # Classify the sub-images based on the decomposition size (M) and theta

    breaks = np.linspace(0, 2 * np.pi, M + 1)

    for i, t in enumerate(breaks[:-1]):

        smembership[(t <= stheta) & ( stheta <= breaks[i+1])] = i

    for i, t in enumerate(breaks[:-1]):

        # Handle the boundary crossers

        start_theta = t + theta

        stop_theta = breaks[i + 1] + theta

        if stop_theta > 2 * np.pi:

            stop_theta -= 2 * np.pi

        if start_theta > 2 * np.pi:

            start_theta -= 2 * np.pi

        if start_theta > stop_theta:

            # Handles the case where theta is a negative rotation

            dmembership[(start_theta <= dtheta) & (dtheta <= 2 * np.pi)] = i

            dmembership[(0 <= dtheta) * dtheta <= stop_theta + lam] = i

            dmembership[(start_theta <= dtheta) & (dtheta <= stop_theta)] = i

    smembership[(stheta >= 0) & (stheta <= np.pi/2)] = 0

    smembership[(stheta >= np.pi/2) & (stheta <= np.pi)] = 1

    smembership[(stheta >=np.pi) & (stheta <= 3 * np.pi/2)] = 2

    smembership[(stheta >= 3 * np.pi/2) & (stheta <= 2 * np.pi)] = 3

    if 0 <= theta <= np.pi / 2:

        order = [0,1,2,3]

    elif np.pi/2 <= theta <= np.pi:

        order = [1,2,3,0]

    elif np.pi <= theta <= 3*np.pi/2:

        order = [2,3,0,1]

    elif 3*np.pi/2 <= theta <= 2*np.pi:

        order = [3,0,1,2]

    def wrap(v):

        return v % (2 * np.pi)

    def classify(start, stop, classid):

        start = wrap(start)

        stop = wrap(stop)

        if start > stop:

            dmembership[(dtheta >=start) & (dtheta <= 2*np.pi)] = classid

            dmembership[(dtheta >=0) & (dtheta <= stop)] = classid

        else:

            # Handles the standard case without boundary crossers

            dmembership[(start_theta <= dtheta) & (dtheta <= stop_theta)] = i

            dmembership[(dtheta >= start) & (dtheta <= stop)] = classid

    classify(theta, theta + np.pi/2, 0)

    classify(theta + np.pi/2, theta + np.pi, 1)

    classify(theta + np.pi, theta + 3*np.pi/2,2)

    classify(theta + 3*np.pi/2, theta + 2*np.pi, 3)

    return smembership, dmembership