/*******************************************************

* Copyright (c) 2014, ArrayFire

* All rights reserved.

*

* This file is distributed under 3-clause BSD license.

* The complete license agreement can be obtained at:

* http://arrayfire.com/licenses/BSD-3-Clause

********************************************************/

#include <arrayfire.h>

#include <gtest/gtest.h>

#include <half.hpp>

#include <testHelpers.hpp>

#include <af/dim4.hpp>

#include <af/traits.hpp>

#include <algorithm>

#include <ctime>

#include <iostream>

#include <string>

#include <vector>

using af::array;

using af::cdouble;

using af::cfloat;

using af::constant;

using af::dim4;

using af::randu;

using half_float::half;

using std::endl;

using std::string;

using std::vector;

template<typename T>

class Mean : public ::testing::Test {

public:

virtual void SetUp() {}

};

// create a list of types to be tested

// This list does not allow to cleanly add the af_half/half_float type : at the

// moment half tested in some special unittests

typedef ::testing::Types<cdouble, cfloat, float, double, int, uint, intl, uintl,

char, uchar, short, ushort, half_float::half>

TestTypes;

// register the type list

TYPED_TEST_CASE(Mean, TestTypes);

template<typename T>

struct f32HelperType {

typedef

typename cond_type<is_same_type<T, double>::value, double, float>::type

type;

};

template<typename T>

struct c32HelperType {

typedef typename cond_type<is_same_type<T, cfloat>::value, cfloat,

typename f32HelperType<T>::type>::type type;

};

template<typename T>

struct elseType {

typedef typename cond_type<is_same_type<T, uintl>::value ||

is_same_type<T, intl>::value,

double, T>::type type;

};

template<typename T>

struct meanOutType {

typedef typename cond_type<

is_same_type<T, float>::value || is_same_type<T, int>::value ||

is_same_type<T, uint>::value || is_same_type<T, uchar>::value ||

is_same_type<T, short>::value || is_same_type<T, ushort>::value ||

is_same_type<T, char>::value,

float, typename elseType<T>::type>::type type;

};

template<typename T>

void meanDimTest(string pFileName, dim_t dim, bool isWeighted = false) {

typedef typename meanOutType<T>::type outType;

SUPPORTED_TYPE_CHECK(T);

SUPPORTED_TYPE_CHECK(outType);

double tol = 1.0e-3;

if ((af_dtype)af::dtype_traits<T>::af_type == f16) tol = 4.e-3;

vector<dim4> numDims;

vector<vector<int> > in;

vector<vector<float> > tests;

readTestsFromFile<int, float>(pFileName, numDims, in, tests);

dim4 goldDims = numDims[0];

goldDims[dim] = 1;

if (!isWeighted) {

dim4 dims = numDims[0];

vector<T> input(in[0].begin(), in[0].end());

array inArray(dims, &(input.front()));

array outArray = mean(inArray, dim);

vector<outType> outData(dims.elements());

outArray.host((void*)outData.data());

vector<outType> currGoldBar(tests[0].begin(), tests[0].end());

dim4 goldDims = dims;

goldDims[dim] = 1;

ASSERT_VEC_ARRAY_NEAR(currGoldBar, goldDims, outArray, tol);

} else {

dim4 dims = numDims[0];

dim4 wdims = numDims[1];

vector<T> input(in[0].begin(), in[0].end());

vector<float> weights(in[1].size());

transform(in[1].begin(), in[1].end(), weights.begin(),

convert_to<float, int>);

array inArray(dims, &(input.front()));

array wtsArray(wdims, &(weights.front()));

array outArray = mean(inArray, wtsArray, dim);

vector<outType> outData(dims.elements());

outArray.host((void*)outData.data());

vector<outType> currGoldBar(tests[0].begin(), tests[0].end());

ASSERT_VEC_ARRAY_NEAR(currGoldBar, goldDims, outArray, tol);

}

}

TYPED_TEST(Mean, Dim0Matrix) {

meanDimTest<TypeParam>(string(TEST_DIR "/mean/mean_dim0_matrix.test"), 0);

}

TYPED_TEST(Mean, Dim1Cube) {

meanDimTest<TypeParam>(string(TEST_DIR "/mean/mean_dim1_cube.test"), 1);

}

TYPED_TEST(Mean, Dim0HyperCube) {

meanDimTest<TypeParam>(string(TEST_DIR "/mean/mean_dim0_hypercube.test"),

0);

}

TYPED_TEST(Mean, Dim2Matrix) {

meanDimTest<TypeParam>(string(TEST_DIR "/mean/mean_dim2_matrix.test"), 2);

}

TYPED_TEST(Mean, Dim2Cube) {

meanDimTest<TypeParam>(string(TEST_DIR "/mean/mean_dim2_cube.test"), 2);

}

TYPED_TEST(Mean, Dim2HyperCube) {

meanDimTest<TypeParam>(string(TEST_DIR "/mean/mean_dim2_hypercube.test"),

2);

}

TYPED_TEST(Mean, Wtd_Dim0Matrix) {

meanDimTest<TypeParam>(string(TEST_DIR "/mean/wtd_mean_dim0_mat.test"), 0,

true);

}

TYPED_TEST(Mean, Wtd_Dim1Matrix) {

meanDimTest<TypeParam>(string(TEST_DIR "/mean/wtd_mean_dim1_mat.test"), 1,

true);

}

template<typename T>

void meanAllTest(T const_value, dim4 dims) {

typedef typename meanOutType<T>::type outType;

SUPPORTED_TYPE_CHECK(T);

SUPPORTED_TYPE_CHECK(outType);

using af::array;

using af::mean;

vector<T> hundred(dims.elements(), const_value);

outType gold = outType(0);

// for(auto i:hundred) gold += i;

for (int i = 0; i < (int)hundred.size(); i++) { gold = gold + hundred[i]; }

gold = gold / dims.elements();

array a(dims, &(hundred.front()));

outType output = mean<outType>(a);

ASSERT_NEAR(::real(output), ::real(gold), 1.0e-3);

ASSERT_NEAR(::imag(output), ::imag(gold), 1.0e-3);

}

template<>

void meanAllTest(half_float::half const_value, dim4 dims) {

SUPPORTED_TYPE_CHECK(half_float::half);

using af::array;

using af::mean;

vector<float> hundred(dims.elements(), const_value);

float gold = float(0);

for (int i = 0; i < (int)hundred.size(); i++) { gold = gold + hundred[i]; }

gold = gold / dims.elements();

array a = array(dims, &(hundred.front())).as(f16);

half output = mean<half>(a);

af_half output2 = mean<af_half>(a);

// make sure output2 and output are binary equals. This is necessary

// because af_half is not a complete type

half output2_copy;

memcpy(static_cast<void*>(&output2_copy), &output2, sizeof(af_half));

ASSERT_EQ(output, output2_copy);

ASSERT_NEAR(output, gold, 1.0e-3);

}

TEST(MeanAll, f64) { meanAllTest<double>(2.1, dim4(10, 10, 1, 1)); }

TEST(MeanAll, f32) { meanAllTest<float>(2.1f, dim4(10, 5, 2, 1)); }

TEST(MeanAll, f16) { meanAllTest<half>((half)0.3f, dim4(10, 5, 2, 1)); }

TEST(MeanAll, s32) { meanAllTest<int>(2, dim4(5, 5, 2, 2)); }

TEST(MeanAll, u32) { meanAllTest<unsigned>(2, dim4(100, 1, 1, 1)); }

TEST(MeanAll, s8) { meanAllTest<char>(2, dim4(5, 5, 2, 2)); }

TEST(MeanAll, u8) { meanAllTest<uchar>(2, dim4(100, 1, 1, 1)); }

TEST(MeanAll, c32) { meanAllTest<cfloat>(cfloat(2.1f), dim4(10, 5, 2, 1)); }

TEST(MeanAll, s16) { meanAllTest<short>(2, dim4(5, 5, 2, 2)); }

TEST(MeanAll, u16) { meanAllTest<ushort>(2, dim4(100, 1, 1, 1)); }

TEST(MeanAll, c64) { meanAllTest<cdouble>(cdouble(2.1), dim4(10, 10, 1, 1)); }

template<typename T>

T random() {

return T(std::rand() % 10);

}

template<>

half random<half>() {

// create values from -0.5 to 0.5 to ensure sum does not deviate

// too far out of half's useful range

float r = static_cast<float>(rand()) / static_cast<float>(RAND_MAX) - 0.5f;

return half(r);

}

template<>

cfloat random<cfloat>() {

return cfloat(float(std::rand() % 10), float(std::rand() % 10));

}

template<>

cdouble random<cdouble>() {

return cdouble(double(std::rand() % 10), double(std::rand() % 10));

}

template<typename T>

class WeightedMean : public ::testing::Test {

public:

virtual void SetUp() {}

};

// register the type list

TYPED_TEST_CASE(WeightedMean, TestTypes);

template<typename T, typename wtsType>

void weightedMeanAllTest(dim4 dims) {

typedef typename meanOutType<T>::type outType;

SUPPORTED_TYPE_CHECK(T);

SUPPORTED_TYPE_CHECK(outType);

SUPPORTED_TYPE_CHECK(wtsType);

using af::array;

using af::mean;

std::srand(std::time(0));

vector<T> data(dims.elements());

vector<wtsType> wts(dims.elements());

std::generate(data.begin(), data.end(), random<T>);

std::generate(wts.begin(), wts.end(), random<wtsType>);

outType wtdSum = outType(0);

wtsType wtsSum = wtsType(0);

for (int i = 0; i < (int)data.size(); i++) {

wtdSum = wtdSum + data[i] * wts[i];

wtsSum = wtsSum + wts[i];

}

outType gold = wtdSum / outType(wtsSum);

array a(dims, &(data.front()));

array w(dims, &(wts.front()));

outType output = mean<outType>(a, w);

ASSERT_NEAR(::real(output), ::real(gold), 1.0e-2);

ASSERT_NEAR(::imag(output), ::imag(gold), 1.0e-2);

}

TYPED_TEST(WeightedMean, Basic) {

weightedMeanAllTest<TypeParam, float>(dim4(32, 30, 33, 17));

}

TEST(WeightedMean, Broadacst) {

float val = 0.5f;

array a = randu(4096, 32);

array w = constant(val, a.dims());

array c = mean(a);

array d = mean(a, w);

vector<float> hc(c.elements());

vector<float> hd(d.elements());

c.host(hc.data());

d.host(hd.data());

for (size_t i = 0; i < hc.size(); i++) {

// C and D are the same because they are normalized by the sum of the

// weights.

ASSERT_NEAR(hc[i], hd[i], 1E-5);

}

}

TEST(Mean, Issue2093) {

const int NELEMS = 512;

array data = randu(1, NELEMS);

array wts = constant(1.0f, 1, NELEMS);

vector<float> hdata(NELEMS);

data.host(hdata.data());

array out = mean(data, wts, 1);

float outVal;

out.host(&outVal);

float expected = 0.0;

for (size_t i = 0; i < NELEMS; ++i) expected += hdata[i];

expected /= NELEMS;

ASSERT_NEAR(outVal, expected, 0.001);

}

TEST(MeanAll, SubArray) {

// Fixes Issue 2636

using af::mean;

using af::span;

using af::sum;

const dim4 inDims(10, 10, 10, 10);

array in = randu(inDims);

array sub = in(0, span, span, span);

size_t nElems = sub.elements();

ASSERT_FLOAT_EQ(mean<float>(sub), sum<float>(sub) / nElems);

}

TEST(MeanHalf, dim0) {

SUPPORTED_TYPE_CHECK(half_float::half);

// Keeping N low to be able to run on 6GB GPUs

int N = 1024;

const dim4 inDims(N, N, 1, 1);

array in = randu(inDims, f16);

array m16 = af::mean(in, 0);

array m32 = af::mean(in.as(f32), 0);

// Some diffs appears at 0.0001 max diff : example: float: 0.507014 vs half:

// 0.506836

ASSERT_ARRAYS_NEAR(m16.as(f32), m32, 0.001f);

}