5c6c50274c11ecbf9be54052394922fabde80923,kymatio/scattering2d/core/scattering2d.py,,scattering2d,#,4

Before Change

            out_S_2.append(S_2_r)

    out_S = []
    out_S.extend(out_S_0)
    out_S.extend(out_S_1)
    out_S.extend(out_S_2)

    out_S = concatenate(out_S)
    return out_S


__all__ = ["scattering2d"]

After Change

    concatenate = backend.concatenate

    // S is simply a dictionary if we do not perform the averaging...
    S = {}

    U_r = pad(x)

    U_0_c = fft(U_r, "C2C")

    // First low pass filter
    U_1_c = cdgmm(U_0_c, phi[0])
    U_1_c = subsample_fourier(U_1_c, k=2 ** J)

    S_0 = fft(U_1_c, "C2R", inverse=True)
    S_0 = unpad(S_0)

    S[()] = S_0

    for n1 in range(len(psi)):
        j1 = psi[n1]["j"]
        U_1_c = cdgmm(U_0_c, psi[n1][0])
        if j1 > 0:
            U_1_c = subsample_fourier(U_1_c, k=2 ** j1)
        U_1_c = fft(U_1_c, "C2C", inverse=True)
        U_1_c = modulus(U_1_c)
        U_1_c = fft(U_1_c, "C2C")

        // Second low pass filter
        S_1_c = cdgmm(U_1_c, phi[j1])
        S_1_c = subsample_fourier(S_1_c, k=2 ** (J - j1))

        S_1_r = fft(S_1_c, "C2R", inverse=True)
        S_1_r = unpad(S_1_r)


        S[(n1,)] = S_1_r

        if max_order < 2:
            continue
        for n2 in range(len(psi)):
            j2 = psi[n2]["j"]
            if j2 <= j1:
                continue
            U_2_c = cdgmm(U_1_c, psi[n2][j1])
            U_2_c = subsample_fourier(U_2_c, k=2 ** (j2 - j1))
            U_2_c = fft(U_2_c, "C2C", inverse=True)
            U_2_c = modulus(U_2_c)
            U_2_c = fft(U_2_c, "C2C")

            // Third low pass filter
            S_2_c = cdgmm(U_2_c, phi[j2])
            S_2_c = subsample_fourier(S_2_c, k=2 ** (J - j2))

            S_2_r = fft(S_2_c, "C2R", inverse=True)
            S_2_r = unpad(S_2_r)

            S[(n1, n2)] = S_2_r

    if vectorize:
        sorted_keys = sorted(S.keys(), key=lambda x: (len(x),) + x)

        S = [S[key] for key in sorted_keys]

        S = concatenate(S)

    return S


__all__ = ["scattering2d"]

Instances