functions.py

"""
Functions for training the cosmoGAN algorithm and to analyze the results.
"""


###################################
##    Importing the packages:    ##
###################################

import os
import tensorflow as tf
import sys
import time
import numpy as np
import pprint
import matplotlib.pyplot as plt


###################################
##   Defining the functions:     ##
###################################


def linear(input_, output_size, scope=None, stddev=0.02, bias_start=0.0, transpose_b=False):

    shape = input_.get_shape().as_list()
    if not transpose_b:
        w_shape = [shape[1], output_size]
    else:
        w_shape = [output_size, shape[1]]

    with tf.variable_scope(scope or "linear"):
        matrix = tf.get_variable('w', w_shape, tf.float32,
                                 tf.random_normal_initializer(stddev=stddev))
        bias = tf.get_variable('b', [output_size],
            initializer=tf.constant_initializer(bias_start))

        return tf.matmul(input_, matrix, transpose_b=transpose_b) + bias

def conv2d(input_, out_channels, data_format, kernel=5, stride=2, stddev=0.02, name="conv2d"):

    if data_format == "NHWC":
        in_channels = input_.get_shape()[-1]
        strides = [1, stride, stride, 1]
    else: # NCHW
        in_channels = input_.get_shape()[1]
        strides = [1, 1, stride, stride]

    with tf.variable_scope(name):
        w = tf.get_variable('w', [kernel, kernel, in_channels, out_channels],
                            initializer=tf.truncated_normal_initializer(stddev=stddev))
        conv = tf.nn.conv2d(input_, w, strides=strides, padding='SAME', data_format=data_format)

        biases = tf.get_variable('biases', [out_channels], initializer=tf.constant_initializer(0.0))
        conv = tf.reshape(tf.nn.bias_add(conv, biases, data_format=data_format), conv.get_shape())

        return conv

def conv2d_transpose(input_, output_shape, data_format, kernel=5, stride=2, stddev=0.02,
                     name="conv2d_transpose"):

    if data_format == "NHWC":
        in_channels = input_.get_shape()[-1]
        out_channels = output_shape[-1]
        strides = [1, stride, stride, 1]
    else:
        in_channels = input_.get_shape()[1]
        out_channels = output_shape[1]
        strides = [1, 1, stride, stride]

    with tf.variable_scope(name):
        w = tf.get_variable('w', [kernel, kernel, out_channels, in_channels],
                            initializer=tf.random_normal_initializer(stddev=stddev))

        deconv = tf.nn.conv2d_transpose(input_, w, output_shape=output_shape, strides=strides, data_format=data_format)

        biases = tf.get_variable('biases', [out_channels], initializer=tf.constant_initializer(0.0))
        deconv = tf.reshape(tf.nn.bias_add(deconv, biases, data_format=data_format), deconv.get_shape())

        return deconv


def lrelu(x, alpha=0.2, name="lrelu"):
    with tf.name_scope(name):
      return tf.maximum(x, alpha*x)

class dcgan(object):
    def __init__(self, output_size=64, batch_size=64,
                 nd_layers=4, ng_layers=4, df_dim=128, gf_dim=128,
                 c_dim=1, z_dim=100, flip_labels=0.01, data_format="NHWC",
                 gen_prior=tf.random_normal, transpose_b=False):

        self.output_size = output_size
        self.batch_size = batch_size
        self.nd_layers = nd_layers
        self.ng_layers = ng_layers
        self.df_dim = df_dim
        self.gf_dim = gf_dim
        self.c_dim = c_dim
        self.z_dim = z_dim
        self.flip_labels = flip_labels
        self.data_format = data_format
        self.gen_prior = gen_prior
        self.transpose_b = transpose_b # transpose weight matrix in linear layers for (possible) better performance when running on HSW/KNL
        self.stride = 2 # this is fixed for this architecture

        self._check_architecture_consistency()


        self.batchnorm_kwargs = {'epsilon' : 1e-5, 'decay': 0.9,
                                 'updates_collections': None, 'scale': True,
                                 'fused': True, 'data_format': self.data_format}

    def training_graph(self):

        if self.data_format == "NHWC":
            self.images = tf.placeholder(tf.float32, [self.batch_size, self.output_size, self.output_size, self.c_dim], name='real_images')
        else:
            self.images = tf.placeholder(tf.float32, [self.batch_size, self.c_dim, self.output_size, self.output_size], name='real_images')

        self.z = self.gen_prior(shape=[self.batch_size, self.z_dim])

        with tf.variable_scope("discriminator") as d_scope:
            d_prob_real, d_logits_real = self.discriminator(self.images, is_training=True)

        with tf.variable_scope("generator") as g_scope:
            g_images = self.generator(self.z, is_training=True)

        with tf.variable_scope("discriminator") as d_scope:
            d_scope.reuse_variables()
            d_prob_fake, d_logits_fake = self.discriminator(g_images, is_training=True)

        with tf.name_scope("losses"):
            with tf.name_scope("d"):
                d_label_real, d_label_fake = self._labels()
                self.d_loss_real = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=d_logits_real, labels=d_label_real, name="real"))
                self.d_loss_fake = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=d_logits_fake, labels=d_label_fake, name="fake"))
                self.d_loss = self.d_loss_real + self.d_loss_fake
            with tf.name_scope("g"):
                self.g_loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=d_logits_fake, labels=tf.ones_like(d_logits_fake)))

        self.d_summary = tf.summary.merge([tf.summary.histogram("prob/real", d_prob_real),
                                           tf.summary.histogram("prob/fake", d_prob_fake),
                                           tf.summary.scalar("loss/real", self.d_loss_real),
                                           tf.summary.scalar("loss/fake", self.d_loss_fake),
                                           tf.summary.scalar("loss/d", self.d_loss)])

        g_sum = [tf.summary.scalar("loss/g", self.g_loss)]
        if self.data_format == "NHWC": # tf.summary.image is not implemented for NCHW
            g_sum.append(tf.summary.image("G", g_images, max_outputs=4))
        self.g_summary = tf.summary.merge(g_sum)

        t_vars = tf.trainable_variables()
        self.d_vars = [var for var in t_vars if 'discriminator/' in var.name]
        self.g_vars = [var for var in t_vars if 'generator/' in var.name]

        with tf.variable_scope("counters") as counters_scope:
            self.epoch = tf.Variable(-1, name='epoch', trainable=False)
            self.increment_epoch = tf.assign(self.epoch, self.epoch+1)
            self.global_step = tf.train.get_or_create_global_step()

        self.saver = tf.train.Saver(max_to_keep=8000)

    def inference_graph(self):

        if self.data_format == "NHWC":
            self.images = tf.placeholder(tf.float32, [self.batch_size, self.output_size, self.output_size, self.c_dim], name='real_images')
        else:
            self.images = tf.placeholder(tf.float32, [self.batch_size, self.c_dim, self.output_size, self.output_size], name='real_images')

        self.z = tf.placeholder(tf.float32, [None, self.z_dim], name='z')

        with tf.variable_scope("discriminator") as d_scope:
            self.D,_ = self.discriminator(self.images, is_training=False)

        with tf.variable_scope("generator") as g_scope:
            self.G = self.generator(self.z, is_training=False)

        with tf.variable_scope("counters") as counters_scope:
            self.epoch = tf.Variable(-1, name='epoch', trainable=False)
            self.increment_epoch = tf.assign(self.epoch, self.epoch+1)
            self.global_step = tf.train.get_or_create_global_step()

        self.saver = tf.train.Saver(max_to_keep=8000)

    def optimizer(self, learning_rate, beta1):

        d_optim = tf.train.AdamOptimizer(learning_rate, beta1=beta1) \
                                         .minimize(self.d_loss, var_list=self.d_vars, global_step=self.global_step)

        g_optim = tf.train.AdamOptimizer(learning_rate, beta1=beta1) \
                                         .minimize(self.g_loss, var_list=self.g_vars)

        return tf.group(d_optim, g_optim, name="all_optims")


    def generator(self, z, is_training):

        map_size = self.output_size/int(2**self.ng_layers)
        num_channels = self.gf_dim * int(2**(self.ng_layers -1))

        # h0 = relu(BN(reshape(FC(z))))
        z_ = linear(z, num_channels*map_size*map_size, 'h0_lin', transpose_b=self.transpose_b)
        h0 = tf.reshape(z_, self._tensor_data_format(-1, map_size, map_size, num_channels))
        bn0 = tf.contrib.layers.batch_norm(h0, is_training=is_training, scope='bn0', **self.batchnorm_kwargs)
        h0 = tf.nn.relu(bn0)

        chain = h0
        for h in range(1, self.ng_layers):
            # h1 = relu(BN(conv2d_transpose(h0)))
            map_size *= self.stride
            num_channels /= 2
            chain = conv2d_transpose(chain,
                                     self._tensor_data_format(self.batch_size, map_size, map_size, num_channels),
                                     stride=self.stride, data_format=self.data_format, name='h%i_conv2d_T'%h)
            chain = tf.contrib.layers.batch_norm(chain, is_training=is_training, scope='bn%i'%h, **self.batchnorm_kwargs)
            chain = tf.nn.relu(chain)

        # h1 = conv2d_transpose(h0)
        map_size *= self.stride
        hn = conv2d_transpose(chain,
                              self._tensor_data_format(self.batch_size, map_size, map_size, self.c_dim),
                              stride=self.stride, data_format=self.data_format, name='h%i_conv2d_T'%(self.ng_layers))

        return tf.nn.tanh(hn)


    def discriminator(self, image, is_training):

        # h0 = lrelu(conv2d(image))
        h0 = lrelu(conv2d(image, self.df_dim, self.data_format, name='h0_conv'))

        chain = h0
        for h in range(1, self.nd_layers):
            # h1 = lrelu(BN(conv2d(h0)))
            chain = conv2d(chain, self.df_dim*(2**h), self.data_format, name='h%i_conv'%h)
            chain = tf.contrib.layers.batch_norm(chain, is_training=is_training, scope='bn%i'%h, **self.batchnorm_kwargs)
            chain = lrelu(chain)

        # h1 = linear(reshape(h0))
        hn = linear(tf.reshape(chain, [self.batch_size, -1]), 1, 'h%i_lin'%self.nd_layers, transpose_b=self.transpose_b)

        return tf.nn.sigmoid(hn), hn

    def _labels(self):
        with tf.name_scope("labels"):
            ones = tf.ones([self.batch_size, 1])
            zeros = tf.zeros([self.batch_size, 1])
            flip_labels = tf.constant(self.flip_labels)

            if self.flip_labels > 0:
                prob = tf.random_uniform([], 0, 1)

                d_label_real = tf.cond(tf.less(prob, flip_labels), lambda: zeros, lambda: ones)
                d_label_fake = tf.cond(tf.less(prob, flip_labels), lambda: ones, lambda: zeros)
            else:
                d_label_real = ones
                d_label_fake = zeros

        return d_label_real, d_label_fake

    def _tensor_data_format(self, N, H, W, C):
        if self.data_format == "NHWC":
            return [int(N), int(H), int(W), int(C)]
        else:
            return [int(N), int(C), int(H), int(W)]

    def _check_architecture_consistency(self):

        if self.output_size/2**self.nd_layers < 1:
            print("Error: Number of discriminator conv. layers are larger than the output_size for this architecture")
            exit(0)

        if self.output_size/2**self.ng_layers < 1:
            print("Error: Number of generator conv_transpose layers are larger than the output_size for this architecture")
            exit(0)


def save_checkpoint(sess, saver, tag, checkpoint_dir, counter, step=False):

    model_name = tag + '.model-'
    if step:
        model_name += 'step'
    else:
        model_name += 'epoch'

    if not os.path.exists(checkpoint_dir):
        os.makedirs(checkpoint_dir)

    saver.save(sess, os.path.join(checkpoint_dir, model_name), global_step=counter)

def load_checkpoint(sess, saver, tag, checkpoint_dir, counter=None, step=False):
    print(" [*] Reading checkpoints...")

    if step:
        counter_name = 'step'
    else:
        counter_name = 'epoch'

    ckpt = tf.train.get_checkpoint_state(checkpoint_dir)
    if ckpt and ckpt.model_checkpoint_path:
        ckpt_name = os.path.basename(ckpt.model_checkpoint_path)

        if not counter==None:
            ckpt_name_epoch = ckpt_name[:ckpt_name.find(counter_name)] + counter_name + '-%i'%counter
            if os.path.exists(os.path.join(checkpoint_dir, ckpt_name_epoch+'.index')):
                ckpt_name = ckpt_name_epoch
            else:
                print("Checkpoint for ", counter_name , counter_name, "doesn't exist. Using latest checkpoint instead!")

        saver.restore(sess, os.path.join(checkpoint_dir, ckpt_name))
        print(" [*] Success to read {}".format(ckpt_name))
        return True
    else:
        print(" [*] Failed to find a checkpoint")
        return False

def train_dcgan(data, config):

    training_graph = tf.Graph()

    with training_graph.as_default():

        gan = dcgan(output_size=config.output_size,
                          batch_size=config.batch_size,
                          nd_layers=config.nd_layers,
                          ng_layers=config.ng_layers,
                          df_dim=config.df_dim,
                          gf_dim=config.gf_dim,
                          c_dim=config.c_dim,
                          z_dim=config.z_dim,
                          flip_labels=config.flip_labels,
                          data_format=config.data_format,
                          transpose_b=config.transpose_matmul_b)

        save_every_step = True if config.save_every_step == 'True' else False

        gan.training_graph()
        update_op = gan.optimizer(config.learning_rate, config.beta1)

        checkpoint_dir = os.path.join(config.checkpoint_dir, config.experiment)

        with tf.Session() as sess:
            writer = tf.summary.FileWriter('./logs/'+config.experiment+'/train', sess.graph)

            init_op = tf.global_variables_initializer()
            sess.run(init_op)

            load_checkpoint(sess, gan.saver, 'dcgan', checkpoint_dir, step=save_every_step)

            epoch = sess.run(gan.increment_epoch)
            start_time = time.time()
            for epoch in range(epoch, epoch + config.epoch):

                perm = np.random.permutation(data.shape[0])
                num_batches = data.shape[0] // config.batch_size

                for idx in range(0, num_batches):
                    batch_images = data[perm[idx*config.batch_size:(idx+1)*config.batch_size]]

                    _, g_sum, d_sum = sess.run([update_op, gan.g_summary, gan.d_summary],
                                               feed_dict={gan.images: batch_images})

                    global_step = gan.global_step.eval()
                    writer.add_summary(g_sum, global_step)
                    writer.add_summary(d_sum, global_step)

                    ## Uncomment this (and comment out the lines below) to save every step:
                    #save_checkpoint(sess, gan.saver, 'dcgan', checkpoint_dir, global_step, step=True)

                    if save_every_step:
                        save_checkpoint(sess, gan.saver, 'dcgan', checkpoint_dir, global_step, step=True)

                    if config.verbose:
                        errD_fake = gan.d_loss_fake.eval()
                        errD_real = gan.d_loss_real.eval({gan.images: batch_images})
                        errG = gan.g_loss.eval()

                        print("Epoch: [%2d] Step: [%4d/%4d] time: %4.4f, d_loss: %.8f, g_loss: %.8f" \
                                % (epoch, idx, num_batches, time.time() - start_time, errD_fake+errD_real, errG))

                    elif global_step%100 == 0:
                        print("Epoch: [%2d] Step: [%4d/%4d] time: %4.4f"%(epoch, idx, num_batches, time.time() - start_time))

                # save a checkpoint every epoch
        ## Save only every 100th epoch (uncomment if it breaks):
                if epoch % 100 == 0:
                    save_checkpoint(sess, gan.saver, 'dcgan', checkpoint_dir, epoch, step=True)

                sess.run(gan.increment_epoch)

### A function to load the checkpoint:
def Loader(checkpoint_dir_pt ,checkpoint_epoch, output_size = 256, z = 256):
    '''
    Initiates a tensorflow session and loads the input checkpoint at a given epoch.
    Also produces a batch of samples corresponding to the newest.
    '''
    with tf.Graph().as_default() as g:
        with tf.Session(graph=g) as sess:

            gan = dcgan(output_size=output_size,
                          nd_layers=4,
                          ng_layers=4,
                          df_dim=64,
                          gf_dim=64,
                          z_dim=z,
                          data_format="NHWC")

            gan.inference_graph()

            load_checkpoint(sess, gan.saver, 'dcgan', checkpoint_dir_pt, counter=checkpoint_epoch)

            t_vars = tf.trainable_variables()
            d_vars = [var for var in t_vars if 'discriminator/' in var.name]
            g_vars = [var for var in t_vars if 'generator/' in var.name]

            g_vars = [(sess.run(var), var.name) for var in g_vars]
            d_vars = [(sess.run(var), var.name) for var in d_vars]


    with tf.Graph().as_default() as g:
        with tf.Session(graph=g) as sess:

            gan = dcgan(output_size=output_size,
                          nd_layers=4,
                          ng_layers=4,
                          df_dim=64,
                          gf_dim=64,
                          z_dim=z,
                          data_format="NHWC")

            gan.inference_graph()

            load_checkpoint(sess, gan.saver, 'dcgan', checkpoint_dir_pt, counter=checkpoint_epoch)

            z_sample = np.random.normal(size=(gan.batch_size, gan.z_dim))
            #z_sample = np.random.normal(size=(batch_size, gan.z_dim))
            samples = sess.run(gan.G, feed_dict={gan.z: z_sample})

            t_vars = tf.trainable_variables()
            d_vars = [var for var in t_vars if 'discriminator/' in var.name]
            g_vars = [var for var in t_vars if 'generator/' in var.name]

            g_vars = [(sess.run(var), var.name) for var in g_vars]# if 'g_h' in var.name]
            d_vars = [(sess.run(var), var.name) for var in d_vars]#if 'd_h' in var.name]

    return z_sample, samples


def get_data():
    data = np.load(config.datafile, mmap_mode='r')

    if config.data_format == 'NHWC':
        data = np.expand_dims(data, axis=-1)
    else: # 'NCHW'
        data = np.expand_dims(data, axis=1)

    return data


def binavg(pk,k_min, k_max, kgrid,nkbins):
    '''
    Bin averaging for the powerspectrum calculations
    '''
    kgrid[0,0] = 1.0
    ikbin = np.digitize(kgrid,np.linspace(k_min,k_max,nkbins+1),right=False)

    nmodes,pkavg,kmean = np.zeros(nkbins,dtype=int),np.full(nkbins,0.),np.full(nkbins,0.)
    for ik in range(nkbins):
        nmodes[ik] = int(np.sum(np.array([ikbin == ik+1])))
        if (nmodes[ik] > 0):
            pkavg[ik] = np.mean(pk[ikbin == ik+1])
            kmean[ik] = np.mean(kgrid[ikbin == ik+1])

    return pkavg, nmodes, kmean


def powerspectrum_ensemble(data1, data2, k_min=0, k_max=1,nkbins=11, h = 0.7 ):
    '''
    A function to calculate the powerspectrum for a 2-D overdensity field.
    k_max = 256*pi/512 ##the Nyquist frequency calculated using the number of pixes * pi/physical length of the array
    to find and estimate for this calculate the fundamental frequency kf = 2*pi/Vbox^1/3
    '''
    kx = 2. * np.pi * np.fft.fftfreq(256, d=1.0)
    ky = 2. * np.pi * np.fft.fftfreq(256, d=1.0)
    kgrid = np.sqrt(kx[:,np.newaxis]**2.0 + kx[np.newaxis,:]**2.0)

    plt.figure(1)

    for i in range(len(data1)):

        delta_r = data1[i]
        delta_k = np.fft.fftn(delta_r)
        pk = np.real(delta_k * np.conj(delta_k))
        pkavg, nmodes, kmean = binavg(pk,k_min, k_max, kgrid, nkbins)

        delta_r2 = data2[i]
        delta_k2 = np.fft.fftn(delta_r2)
        pk2 = np.real(delta_k2 * np.conj(delta_k2))
        pkavg2, nmodes2, kmean2 = binavg(pk2,k_min, k_max, kgrid, nkbins)

        plt.plot(kmean*h,pkavg, label ='P(k) for data1', color = 'red', alpha = 0.4)
        plt.plot(kmean2*h, pkavg2, label = 'P(k) for data2', color = 'blue', alpha = 0.4)
        plt.xscale('log')
        plt.yscale('log')
        plt.xlabel(r'$k$ $[h$ $Mpc^{-1}]$', fontsize = 13)
        plt.ylabel('P(k)', fontsize = 13)
        plt.tick_params(right=True, top=True, direction='in', which = 'both')
        plt.tick_params(which='major', length=4)
        plt.tight_layout()

    plt.legend(['Training data', 'Gan-produced'], fontsize = 11)
    plt.tight_layout()
    plt.show()


def powerspectrum_mean_std(data, k_min=0, k_max=1,nkbins=11, h = 0.7):
    '''
    Calculates the mean + standard deviation of an ensemble of cw slices/WL maps
    '''
    ## Defining the grid in 2-D:
    kx = 2. * np.pi * np.fft.fftfreq(256, d=1.0)
    ky = 2. * np.pi * np.fft.fftfreq(256, d=1.0)

    kgrid = np.sqrt(kx[:,np.newaxis]**2.0 + kx[np.newaxis,:]**2.0)

    kmeans= [] ## arrays to store each powerspectrum
    pkavgs = []

    for i in data:
        delta_r = i
        delta_k = np.fft.fftn(delta_r)
        pk = np.real(delta_k * np.conj(delta_k))
        pkavg, nmodes, kmean = binavg(pk,k_min, k_max, kgrid, nkbins)

        ## Storing the results:
        kmeans.append(kmean)
        pkavgs.append(pkavg)

    return kmeans,pkavgs



def Load_samples(z_sample, checkpoint_dir_pt, checkpoint_epoch, output_size=256, z = 64):
    with tf.Graph().as_default() as g:
        with tf.Session(graph=g) as sess:

            gan = dcgan(output_size=output_size,
                          nd_layers=4,
                          ng_layers=4,
                          df_dim=64,
                          gf_dim=64,
                          z_dim=z,
                          data_format="NHWC")

            gan.inference_graph()

            load_checkpoint(sess, gan.saver, 'dcgan', checkpoint_dir_pt, counter=checkpoint_epoch)

            t_vars = tf.trainable_variables()
            d_vars = [var for var in t_vars if 'discriminator/' in var.name]
            g_vars = [var for var in t_vars if 'generator/' in var.name]

            g_vars = [(sess.run(var), var.name) for var in g_vars]
            d_vars = [(sess.run(var), var.name) for var in d_vars]

    with tf.Graph().as_default() as g:
        with tf.Session(graph=g) as sess:

            gan = dcgan(output_size=output_size,
                          nd_layers=4,
                          ng_layers=4,
                          df_dim=64,
                          gf_dim=64,
                          z_dim=z,
                          data_format="NHWC")

            gan.inference_graph()

            load_checkpoint(sess, gan.saver, 'dcgan', checkpoint_dir_pt, counter=checkpoint_epoch)

    #        z_sample = np.random.normal(size=(gan.batch_size, gan.z_dim))
            samples = sess.run(gan.G, feed_dict={gan.z: z_sample})

            t_vars = tf.trainable_variables()
            d_vars = [var for var in t_vars if 'discriminator/' in var.name]
            g_vars = [var for var in t_vars if 'generator/' in var.name]

            g_vars = [(sess.run(var), var.name) for var in g_vars]# if 'g_h' in var.name]
            d_vars = [(sess.run(var), var.name) for var in d_vars]#if 'd_h' in var.name]

            return samples



def Z_interpolate(Z1,Z2, batch_size, vector_size):
    '''
    A function that reads in two points in N-dimensional space (default = 256)
    Z1 and Z2 and finds a number of intermediate points on the line connecting
    the two points.
    '''
    Z_int =  np.ones((batch_size,vector_size))
    for i in range(0,vector_size):
        Z_delta = np.linspace(Z1[i], Z2[i],num = batch_size)
        for j in range(0,batch_size):
            Z_int[j][i] = Z_delta[j]
    return Z_int


def powerspectrum_i(data, color, alpha = 1.0, label = None):
	## Defining the grid in 2-D:
	kx = 2. * np.pi * np.fft.fftfreq(256, d=1.0)
	ky = 2. * np.pi * np.fft.fftfreq(256, d=1.0)

	kgrid = np.sqrt(kx[:,np.newaxis]**2.0 + kx[np.newaxis,:]**2.0)


	k_min = 0.0
	k_max = 1.0   ## k_max = 256*pi/512 ##the Nyquist frequency calculated using the number of pixes * pi/physical length of the array
	nkbins = 11   ## to find and estimate for this calculate the fundamental frequency kf = 2*pi/Vbox^1/3
	              ## --> in our case kf = 2*pi/Area_box^{1/2}. And then kmax-kmin/kf ~ nkbins

	## For data1:

	delta_r = data
	delta_k = np.fft.fftn(delta_r)
	pk = np.real(delta_k * np.conj(delta_k))
	pkavg, nmodes, kmean = binavg(pk,k_min, k_max, kgrid,nkbins)

	h = 0.7

	## Plotting the results:

	plt.figure(1,figsize=(5.5,4.5))
	plt.plot(kmean,pkavg, color = color, alpha = alpha, label = label) ## Need to figure out the units here
	plt.xscale('log')
	plt.yscale('log')
	plt.xlabel(r'$k$ [h $\mathrm{Mpc^{-1}}$]', fontsize = 13)
	plt.ylabel(r'$P(k)$', fontsize = 13)
	plt.tick_params(right=True, top=True, direction='in', which = 'both')
	plt.tick_params(which='major', length=4)
	plt.tight_layout()