# -*- coding: utf-8 -*-
"""confirming DCT implementation.ipynb

Automatically generated by Colaboratory.

"""
# run from a folder which has access to importing 'modules'

import numpy as np
import scipy as sp
from scipy.fftpack import dct
import torch
import torch.nn as nn

"""**Testing our DCTconv2d layer**: We want to confirm that our layer is doing the DCT of each patch.
To test this we implement our own convolution by hand, using Scipy's DCT, and compare this with the
output of a DCTconv2d on the same input. We see expected that they are equal. We repeate this for both
cases of needing to pad the input of the DCT and truncate the output of the DCT."""

class DCTconv2d(nn.Conv2d):

  def __init__(self, in_channels, out_channels, kernel_size, stride = 1, padding = 0, dilation = 1, groups = 1, bias = False):
    super(DCTconv2d, self).__init__(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride,
                                        padding=padding, dilation=dilation, groups=groups, bias=bias)
    patch_dim = int(in_channels/groups)*kernel_size*kernel_size
    f = sp.fftpack.dct(np.eye(max(patch_dim, out_channels)), norm='ortho')[:patch_dim, :out_channels]
    kernel_shape = (kernel_size,kernel_size,int(in_channels/groups),out_channels)
    self.weight.data = torch.from_numpy(np.transpose(np.reshape(f,kernel_shape), [3,2,0,1])).float()

    self.weight.requires_grad = False
    

# create  input
in_channels = 10
x = np.random.normal(size = (1,32,32,in_channels))

# By hand calculate DCT of each patch

num_filters = 20
k = 3
output = np.zeros((1, 32-k+1, 32-k+1, num_filters))
for i in range(32-k+1):
  for j in range(32-k+1):
    x_patch = x[:,i:i+k,j:j+k,:]
    x_patch_vec = np.reshape(x_patch, -1)
    if len(x_patch_vec)<num_filters:
      padded_x_patch = np.zeros(num_filters)
      padded_x_patch[:len(x_patch_vec)] = x_patch_vec
      output[:, i,j,:] = sp.fftpack.dct(padded_x_patch, norm='ortho')
    else:
      output[:, i,j,:] = sp.fftpack.dct(x_patch_vec, norm='ortho')[:num_filters]

# Create DCTconv2dlayer and reshape input for Pytorch
layer = DCTconv2d(10, 20, 3, 1, padding=0)
x_torch = np.transpose(x, (0,3,1,2))

# Pass input through the layer
layer_output = layer(torch.Tensor(x_torch))

# Check that they are equal (minus numerical errors)
output = np.transpose(output, (0,3,1,2))
if np.sum(np.isclose(layer_output.numpy(), output, atol=1e-06)) == np.product(output.shape):
    print("DCTconv2d gives correct output")

# create  input
in_channels = 3
x = np.random.normal(size = (1,32,32,in_channels))

# By hand calculate DCT of each patch

num_filters = 30
k = 3
output = np.zeros((1, 32-k+1, 32-k+1, num_filters))
for i in range(32-k+1):
  for j in range(32-k+1):
    x_patch = x[:,i:i+k,j:j+k,:]
    x_patch_vec = np.reshape(x_patch, -1)
    if len(x_patch_vec)<num_filters:
      padded_x_patch = np.zeros(num_filters)
      padded_x_patch[:len(x_patch_vec)] = x_patch_vec
      output[:, i,j,:] = sp.fftpack.dct(padded_x_patch, norm='ortho')
    else:
      output[:, i,j,:] = sp.fftpack.dct(x_patch_vec, norm='ortho')[:num_filters]

# Create DCTconv2dlayer and reshape input for Pytorch
layer = DCTconv2d(in_channels, num_filters, 3, 1, padding=0)
x_torch = np.transpose(x, (0,3,1,2))

# Pass input through the layer
layer_output = layer(torch.Tensor(x_torch))

# Check that they are equal (minus numerical errors)
output = np.transpose(output, (0,3,1,2))
if np.sum(np.isclose(layer_output.numpy(), output, atol=1e-6)) == np.product(output.shape):
    print("DCTconv2d gives correct output")

# -*- coding: utf-8 -*-
"""confirming DCT implementation.ipynb

Automatically generated by Colaboratory.

"""
# run from a folder which has access to importing 'modules'

import numpy as np
import scipy as sp
from scipy.fftpack import dct
import torch
import torch.nn as nn

"""**Testing our DCTconv2d layer**: We want to confirm that our layer is doing the DCT of each patch.
To test this we implement our own convolution by hand, using Scipy's DCT, and compare this with the
output of a DCTconv2d on the same input. We see expected that they are equal. We repeate this for both
cases of needing to pad the input of the DCT and truncate the output of the DCT."""

class DCTconv2d(nn.Conv2d):

  def __init__(self, in_channels, out_channels, kernel_size, stride = 1, padding = 0, dilation = 1, groups = 1, bias = False):
    super(DCTconv2d, self).__init__(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=stride,
                                        padding=padding, dilation=dilation, groups=groups, bias=bias)
    patch_dim = int(in_channels/groups)*kernel_size*kernel_size
    f = sp.fftpack.dct(np.eye(max(patch_dim, out_channels)), norm='ortho')[:patch_dim, :out_channels]
    kernel_shape = (kernel_size,kernel_size,int(in_channels/groups),out_channels)
    self.weight.data = torch.from_numpy(np.transpose(np.reshape(f,kernel_shape), [3,2,0,1])).float()

    self.weight.requires_grad = False
    

# create  input
in_channels = 10
x = np.random.normal(size = (1,32,32,in_channels))

# By hand calculate DCT of each patch

num_filters = 20
k = 3
output = np.zeros((1, 32-k+1, 32-k+1, num_filters))
for i in range(32-k+1):
  for j in range(32-k+1):
    x_patch = x[:,i:i+k,j:j+k,:]
    x_patch_vec = np.reshape(x_patch, -1)
    if len(x_patch_vec)<num_filters:
      padded_x_patch = np.zeros(num_filters)
      padded_x_patch[:len(x_patch_vec)] = x_patch_vec
      output[:, i,j,:] = sp.fftpack.dct(padded_x_patch, norm='ortho')
    else:
      output[:, i,j,:] = sp.fftpack.dct(x_patch_vec, norm='ortho')[:num_filters]

# Create DCTconv2dlayer and reshape input for Pytorch
layer = DCTconv2d(10, 20, 3, 1, padding=0)
x_torch = np.transpose(x, (0,3,1,2))

# Pass input through the layer
layer_output = layer(torch.Tensor(x_torch))

# Check that they are equal (minus numerical errors)
output = np.transpose(output, (0,3,1,2))
if np.sum(np.isclose(layer_output.numpy(), output, atol=1e-06)) == np.product(output.shape):
    print("DCTconv2d gives correct output")

# create  input
in_channels = 3
x = np.random.normal(size = (1,32,32,in_channels))

# By hand calculate DCT of each patch

num_filters = 30
k = 3
output = np.zeros((1, 32-k+1, 32-k+1, num_filters))
for i in range(32-k+1):
  for j in range(32-k+1):
    x_patch = x[:,i:i+k,j:j+k,:]
    x_patch_vec = np.reshape(x_patch, -1)
    if len(x_patch_vec)<num_filters:
      padded_x_patch = np.zeros(num_filters)
      padded_x_patch[:len(x_patch_vec)] = x_patch_vec
      output[:, i,j,:] = sp.fftpack.dct(padded_x_patch, norm='ortho')
    else:
      output[:, i,j,:] = sp.fftpack.dct(x_patch_vec, norm='ortho')[:num_filters]

# Create DCTconv2dlayer and reshape input for Pytorch
layer = DCTconv2d(in_channels, num_filters, 3, 1, padding=0)
x_torch = np.transpose(x, (0,3,1,2))

# Pass input through the layer
layer_output = layer(torch.Tensor(x_torch))

# Check that they are equal (minus numerical errors)
output = np.transpose(output, (0,3,1,2))
if np.sum(np.isclose(layer_output.numpy(), output, atol=1e-6)) == np.product(output.shape):
    print("DCTconv2d gives correct output")

# -*- coding: utf-8 -*- confirming DCT implementation.ipynb  Automatically generated by Colaboratory. # run from a folder which has access to importing 'modules' **Testing our DCTconv2d layer**: We want to confirm that our layer is doing the DCT of each patch. To test this we implement our own convolution by hand, using Scipy's DCT, and compare this with the output of a DCTconv2d on the same input. We see expected that they are equal. We repeate this for both cases of needing to pad the input of the DCT and truncate the output of the DCT. # create  input # By hand calculate DCT of each patch # Create DCTconv2dlayer and reshape input for Pytorch # Pass input through the layer # Check that they are equal (minus numerical errors) # create  input # By hand calculate DCT of each patch # Create DCTconv2dlayer and reshape input for Pytorch # Pass input through the layer # Check that they are equal (minus numerical errors)