<jupyter_start><jupyter_text>Create data<jupyter_code>rng = RandomState(20190401)
# i freq
# amp_level
inputs =  []

for i_class, freq_amps  in enumerate((((2,3), (4,1)), ((2,1), (4,3)))):
    fft_coefs = np.zeros((300, 33), dtype=np.complex64)
    for i_freq, amp in freq_amps:
        amps = np.abs(rng.randn(300) + amp)
        phases = rng.rand(300) * 2 * np.pi - np.pi
        this_fft_coefs_a = amps * np.exp(phases * 1j)
        fft_coefs[:,i_freq] = this_fft_coefs_a
        
        signals = np.fft.irfft(fft_coefs)
        inputs.append(np_to_var(signals, dtype=np.float32))
a = th.rfft(inputs[0].view(300,-1), signal_ndim=1)

net = AmplitudePhaseFFT()

b = net(a)

c = invert(net,b)
th.mean(th.abs(c - a))
def th_angle(x,y):
    # yes reverse oder to (y,x)...
    return th.atan2(y, x)


def th_amp(x,y):
    return th.sqrt((x * x) + (y * y))

class AmplitudePhaseFFT(nn.Module):
    def __init__(self):
        super(AmplitudePhaseFFT, self).__init__()

    def forward(self, inputs):
        expanded = th.cat(
            (inputs[:, :1], th.zeros(inputs.shape[0], 1, device=inputs.device),
             inputs[:, 1:], th.zeros(inputs.shape[0], 1, device=inputs.device)),
            dim=1)
        unfolded = expanded.view(expanded.shape[0], expanded.shape[1] // 2, 2)
        
        x = unfolded.transpose(len(unfolded.shape)-1,0)[0].transpose(
            len(unfolded.shape) - 2,0)
        y = unfolded.transpose(len(unfolded.shape)-1,0)[1].transpose(
            len(unfolded.shape) - 2,0)
        
        amps = th_amp(x, y)
        phases = th_angle(x, y)
        
        return th.stack((amps, phases), dim=-1)
    
    def invert(self, out):
        amps = out.transpose(len(out.shape) - 1,0)[0].transpose(
            len(out.shape) - 2,0)
        phases = out.transpose(len(out.shape) - 1,0)[1].transpose(
            len(out.shape) - 2,0)
        x = th.cos(phases)
        y = th.sin(phases)
        x_y = th.stack((x,y), dim=-1)
        x_y = x_y * amps.unsqueeze(-1)
        flattened = x_y.view(x_y.shape[0], -1)
        reduced = th.cat((flattened[:, :1], flattened[:, 2:-1]), dim=1)
        return reduced

from rfft import RFFT, Interleave
from discriminator import ProjectionDiscriminator
from reversible.revnet import SubsampleSplitter, ViewAs
from reversible.util import set_random_seeds
from reversible.revnet import init_model_params
from torch.nn import ConstantPad2d
import torch as th
from conv_spectral_norm import conv_spectral_norm
from disttransform import DistTransformResNet

feature_model = nn.Sequential(RFFT(),
                              AmplitudePhaseFFT())

outs_2 = feature_model(inputs[1])
print(np.mean(np.abs(var_to_np(outs_2[:,:,0]) - np.abs(np.fft.rfft(var_to_np(inputs[1]),
                                                            norm='ortho')))))
np.mean(np.abs(var_to_np(outs_2[:,:,1]) - np.angle(np.fft.rfft(var_to_np(inputs[1]),
                                                            norm='ortho'))))
np_angles = np.angle(np.fft.rfft(var_to_np(inputs[1]),norm='ortho'))
th_angles = outs_2[:,:,1]
plt.plot((var_to_np(th_angles) - np_angles).T);
fft_th = th.rfft(inputs[1], signal_ndim=1, normalized=True)
fft_th_2 = feature_model[0](inputs[1])
fft_th_2 - fft_th
fft_np = np.fft.rfft(var_to_np(inputs[1]), norm='ortho')
fft_th = th.rfft(inputs[1], signal_ndim=1, normalized=True)
fft_np = np.stack((np.real(fft_np), np.imag(fft_np)), axis=-1)
np.mean(np.abs(var_to_np(fft_th) -fft_np))
outs = feature_model(inputs[0])

outs.shape
out_fft = invert(list(feature_model.children())[1], outs)
out_fft.shape
ins = invert(feature_model, outs)
th.mean(np.abs(ins - inputs[0]))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1602721088435374, 'fraction_numerical': 0.02857142857142857, 'mean_word_length': 2.9783549783549783, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '277937', 'license': 'MIT', 'n_tokens_mistral': 1484, 'n_tokens_neox': 1434, 'n_words': 237, 'path': 'notebooks/splitted-hierarchical/.ipynb_checkpoints/TwoClassToyDataAmpFFT-checkpoint.ipynb', 'repo_name': 'robintibor/reversible2', 'size': '3675'}