Skip to content

Commit a1c6353

Browse files
jvdp1jvdp1
committed
modification to have the same behaviour as Fortran sum
1 parent 72500e1 commit a1c6353

File tree

5 files changed

+98
-70
lines changed

5 files changed

+98
-70
lines changed

‎src/stdlib_experimental_stat.f90

Lines changed: 20 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -22,20 +22,34 @@ module function mean_1_qp_qp(mat) result(res)
2222
real(qp) ::res
2323
end function mean_1_qp_qp
2424

25+
module function mean_2_all_sp_sp(mat) result(res)
26+
real(sp), intent(in) :: mat(:,:)
27+
real(sp) ::res
28+
end function mean_2_all_sp_sp
29+
module function mean_2_all_dp_dp(mat) result(res)
30+
real(dp), intent(in) :: mat(:,:)
31+
real(dp) ::res
32+
end function mean_2_all_dp_dp
33+
module function mean_2_all_qp_qp(mat) result(res)
34+
real(qp), intent(in) :: mat(:,:)
35+
real(qp) ::res
36+
end function mean_2_all_qp_qp
37+
38+
2539
module function mean_2_sp_sp(mat, dim) result(res)
2640
real(sp), intent(in) :: mat(:,:)
27-
integer, intent(in), optional :: dim
28-
real(sp), allocatable ::res(:)
41+
integer, intent(in) :: dim
42+
real(sp) :: res(size(mat)/size(mat, dim))
2943
end function mean_2_sp_sp
3044
module function mean_2_dp_dp(mat, dim) result(res)
3145
real(dp), intent(in) :: mat(:,:)
32-
integer, intent(in), optional :: dim
33-
real(dp), allocatable ::res(:)
46+
integer, intent(in) :: dim
47+
real(dp) :: res(size(mat)/size(mat, dim))
3448
end function mean_2_dp_dp
3549
module function mean_2_qp_qp(mat, dim) result(res)
3650
real(qp), intent(in) :: mat(:,:)
37-
integer, intent(in), optional :: dim
38-
real(qp), allocatable ::res(:)
51+
integer, intent(in) :: dim
52+
real(qp) :: res(size(mat)/size(mat, dim))
3953
end function mean_2_qp_qp
4054
end interface
4155

‎src/stdlib_experimental_stat.fypp.f90

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,11 +20,19 @@ module function mean_1_${k1}$_${k1}$(mat) result(res)
2020
end function mean_1_${k1}$_${k1}$
2121
#:endfor
2222

23+
#:for i1, k1, t1 in ikt
24+
module function mean_2_all_${k1}$_${k1}$(mat) result(res)
25+
${t1}$, intent(in) :: mat(:,:)
26+
${t1}$ ::res
27+
end function mean_2_all_${k1}$_${k1}$
28+
#:endfor
29+
30+
2331
#:for i1, k1, t1 in ikt
2432
module function mean_2_${k1}$_${k1}$(mat, dim) result(res)
2533
${t1}$, intent(in) :: mat(:,:)
26-
integer, intent(in), optional :: dim
27-
${t1}$, allocatable ::res(:)
34+
integer, intent(in) :: dim
35+
${t1}$ :: res(size(mat)/size(mat, dim))
2836
end function mean_2_${k1}$_${k1}$
2937
#:endfor
3038
end interface

‎src/stdlib_experimental_stat_mean.f90

Lines changed: 45 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -29,78 +29,79 @@ module function mean_1_qp_qp(mat) result(res)
2929

3030
end function mean_1_qp_qp
3131

32-
module function mean_2_sp_sp(mat, dim) result(res)
32+
module function mean_2_all_sp_sp(mat) result(res)
3333
real(sp), intent(in) :: mat(:,:)
34-
integer, intent(in), optional :: dim
35-
real(sp), allocatable ::res(:)
34+
real(sp) ::res
3635

37-
integer :: i
38-
integer :: dim_
36+
res = sum(mat) / real(size(mat), sp)
37+
38+
end function mean_2_all_sp_sp
39+
module function mean_2_all_dp_dp(mat) result(res)
40+
real(dp), intent(in) :: mat(:,:)
41+
real(dp) ::res
42+
43+
res = sum(mat) / real(size(mat), dp)
3944

40-
dim_ = optval(dim, 1)
45+
end function mean_2_all_dp_dp
46+
module function mean_2_all_qp_qp(mat) result(res)
47+
real(qp), intent(in) :: mat(:,:)
48+
real(qp) ::res
4149

42-
if (dim_ < 0 .or. dim_ > 2 ) call error_stop("ERROR (mean): invalid argument (dim) ")
50+
res = sum(mat) / real(size(mat), qp)
4351

44-
allocate(res(size(mat, dim_)))
52+
end function mean_2_all_qp_qp
4553

46-
if (dim_ == 1) then
47-
do i=1, size(mat, dim_)
48-
res(i) = mean_1_sp_sp(mat(i,:))
49-
end do
50-
else if (dim_ == 2) then
51-
do i=1, size(mat, dim_)
54+
module function mean_2_sp_sp(mat, dim) result(res)
55+
real(sp), intent(in) :: mat(:,:)
56+
integer, intent(in) :: dim
57+
real(sp) :: res(size(mat)/size(mat, dim))
58+
59+
integer :: i
60+
61+
if (dim == 1) then
62+
do i=1, size(mat)/size(mat, dim)
5263
res(i) = mean_1_sp_sp(mat(:,i))
5364
end do
65+
else if (dim == 2) then
66+
do i=1, size(mat)/size(mat, dim)
67+
res(i) = mean_1_sp_sp(mat(i,:))
68+
end do
5469
end if
5570

5671
end function mean_2_sp_sp
5772
module function mean_2_dp_dp(mat, dim) result(res)
5873
real(dp), intent(in) :: mat(:,:)
59-
integer, intent(in), optional :: dim
60-
real(dp), allocatable ::res(:)
74+
integer, intent(in) :: dim
75+
real(dp) :: res(size(mat)/size(mat, dim))
6176

6277
integer :: i
63-
integer :: dim_
64-
65-
dim_ = optval(dim, 1)
66-
67-
if (dim_ < 0 .or. dim_ > 2 ) call error_stop("ERROR (mean): invalid argument (dim) ")
6878

69-
allocate(res(size(mat, dim_)))
70-
71-
if (dim_ == 1) then
72-
do i=1, size(mat, dim_)
73-
res(i) = mean_1_dp_dp(mat(i,:))
74-
end do
75-
else if (dim_ == 2) then
76-
do i=1, size(mat, dim_)
79+
if (dim == 1) then
80+
do i=1, size(mat)/size(mat, dim)
7781
res(i) = mean_1_dp_dp(mat(:,i))
7882
end do
83+
else if (dim == 2) then
84+
do i=1, size(mat)/size(mat, dim)
85+
res(i) = mean_1_dp_dp(mat(i,:))
86+
end do
7987
end if
8088

8189
end function mean_2_dp_dp
8290
module function mean_2_qp_qp(mat, dim) result(res)
8391
real(qp), intent(in) :: mat(:,:)
84-
integer, intent(in), optional :: dim
85-
real(qp), allocatable ::res(:)
92+
integer, intent(in) :: dim
93+
real(qp) :: res(size(mat)/size(mat, dim))
8694

8795
integer :: i
88-
integer :: dim_
89-
90-
dim_ = optval(dim, 1)
91-
92-
if (dim_ < 0 .or. dim_ > 2 ) call error_stop("ERROR (mean): invalid argument (dim) ")
93-
94-
allocate(res(size(mat, dim_)))
9596

96-
if (dim_ == 1) then
97-
do i=1, size(mat, dim_)
98-
res(i) = mean_1_qp_qp(mat(i,:))
99-
end do
100-
else if (dim_ == 2) then
101-
do i=1, size(mat, dim_)
97+
if (dim == 1) then
98+
do i=1, size(mat)/size(mat, dim)
10299
res(i) = mean_1_qp_qp(mat(:,i))
103100
end do
101+
else if (dim == 2) then
102+
do i=1, size(mat)/size(mat, dim)
103+
res(i) = mean_1_qp_qp(mat(i,:))
104+
end do
104105
end if
105106

106107
end function mean_2_qp_qp

‎src/stdlib_experimental_stat_mean.fypp.f90

Lines changed: 17 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -22,28 +22,31 @@ end function mean_1_${k1}$_${k1}$
2222
#:endfor
2323

2424
#:for i1, k1, t1 in ikt
25-
module function mean_2_${k1}$_${k1}$(mat, dim) result(res)
25+
module function mean_2_all_${k1}$_${k1}$(mat) result(res)
2626
${t1}$, intent(in) :: mat(:,:)
27-
integer, intent(in), optional :: dim
28-
${t1}$, allocatable ::res(:)
27+
${t1}$ ::res
2928

30-
integer :: i
31-
integer :: dim_
29+
res = sum(mat) / real(size(mat), ${k1}$)
3230

33-
dim_ = optval(dim, 1)
31+
end function mean_2_all_${k1}$_${k1}$
32+
#:endfor
3433

35-
if (dim_ < 0 .or. dim_ > 2 ) call error_stop("ERROR (mean): invalid argument (dim) ")
34+
#:for i1, k1, t1 in ikt
35+
module function mean_2_${k1}$_${k1}$(mat, dim) result(res)
36+
${t1}$, intent(in) :: mat(:,:)
37+
integer, intent(in) :: dim
38+
${t1}$ :: res(size(mat)/size(mat, dim))
3639

37-
allocate(res(size(mat, dim_)))
40+
integer :: i
3841

39-
if (dim_ == 1) then
40-
do i=1, size(mat, dim_)
41-
res(i) = mean_1_${k1}$_${k1}$(mat(i,:))
42-
end do
43-
else if (dim_ == 2) then
44-
do i=1, size(mat, dim_)
42+
if (dim == 1) then
43+
do i=1, size(mat)/size(mat, dim)
4544
res(i) = mean_1_${k1}$_${k1}$(mat(:,i))
4645
end do
46+
else if (dim == 2) then
47+
do i=1, size(mat)/size(mat, dim)
48+
res(i) = mean_1_${k1}$_${k1}$(mat(i,:))
49+
end do
4750
end if
4851

4952
end function mean_2_${k1}$_${k1}$

‎src/tests/stat/test_mean.f90

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,14 +11,16 @@ program test_mean
1111
!sp
1212
call loadtxt("array1.dat", s)
1313

14-
call assert(sum( mean(s) - [1.5_sp, 3.5_sp, 5.5_sp, 7.5_sp] ) == 0.0_sp)
15-
call assert(sum( mean(s, dim = 2) - [4.0_sp, 5.0_sp] ) == 0.0_sp)
14+
call assert( mean(s) - 4.5_sp == 0.0_sp)
15+
call assert(sum( mean(s, dim = 1) - [4.0_sp, 5.0_sp] ) == 0.0_sp)
16+
call assert(sum( mean(s, dim = 2) - [1.5_dp, 3.5_dp, 5.5_dp, 7.5_dp] ) == 0.0_sp)
1617

1718
!dp
1819
call loadtxt("array1.dat", d)
1920

20-
call assert(sum( mean(d) - [1.5_dp, 3.5_dp, 5.5_dp, 7.5_dp] ) == 0.0_dp)
21-
call assert(sum( mean(d, dim = 2) - [4.0_dp, 5.0_dp] ) == 0.0_dp)
21+
call assert(mean(d) - 4.5_dp == 0.0_dp)
22+
call assert(sum( mean(d, dim = 1) - [4.0_dp, 5.0_dp] ) == 0.0_dp)
23+
call assert(sum( mean(d, dim = 2) - [1.5_dp, 3.5_dp, 5.5_dp, 7.5_dp] ) == 0.0_dp)
2224

2325

2426
contains

0 commit comments

Comments
 (0)