Skip to content

Commit

Permalink
modification to have the same behaviour as Fortran sum
Browse files Browse the repository at this point in the history
  • Loading branch information
@jvdp1
jvdp1 committed Jan 19, 2020
1 parent 72500e1 commit a1c6353
Show file tree
Hide file tree
Showing 5 changed files with 98 additions and 70 deletions.
26 changes: 20 additions & 6 deletions src/stdlib_experimental_stat.f90
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,20 +22,34 @@ module function mean_1_qp_qp(mat) result(res)
real(qp) ::res
end function mean_1_qp_qp

module function mean_2_all_sp_sp(mat) result(res)
real(sp), intent(in) :: mat(:,:)
real(sp) ::res
end function mean_2_all_sp_sp
module function mean_2_all_dp_dp(mat) result(res)
real(dp), intent(in) :: mat(:,:)
real(dp) ::res
end function mean_2_all_dp_dp
module function mean_2_all_qp_qp(mat) result(res)
real(qp), intent(in) :: mat(:,:)
real(qp) ::res
end function mean_2_all_qp_qp


module function mean_2_sp_sp(mat, dim) result(res)
real(sp), intent(in) :: mat(:,:)
integer, intent(in), optional :: dim
real(sp), allocatable ::res(:)
integer, intent(in) :: dim
real(sp) :: res(size(mat)/size(mat, dim))
end function mean_2_sp_sp
module function mean_2_dp_dp(mat, dim) result(res)
real(dp), intent(in) :: mat(:,:)
integer, intent(in), optional :: dim
real(dp), allocatable ::res(:)
integer, intent(in) :: dim
real(dp) :: res(size(mat)/size(mat, dim))
end function mean_2_dp_dp
module function mean_2_qp_qp(mat, dim) result(res)
real(qp), intent(in) :: mat(:,:)
integer, intent(in), optional :: dim
real(qp), allocatable ::res(:)
integer, intent(in) :: dim
real(qp) :: res(size(mat)/size(mat, dim))
end function mean_2_qp_qp
end interface

Expand Down
12 changes: 10 additions & 2 deletions src/stdlib_experimental_stat.fypp.f90
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,11 +20,19 @@ module function mean_1_${k1}$_${k1}$(mat) result(res)
end function mean_1_${k1}$_${k1}$
#:endfor

#:for i1, k1, t1 in ikt
module function mean_2_all_${k1}$_${k1}$(mat) result(res)
${t1}$, intent(in) :: mat(:,:)
${t1}$ ::res
end function mean_2_all_${k1}$_${k1}$
#:endfor


#:for i1, k1, t1 in ikt
module function mean_2_${k1}$_${k1}$(mat, dim) result(res)
${t1}$, intent(in) :: mat(:,:)
integer, intent(in), optional :: dim
${t1}$, allocatable ::res(:)
integer, intent(in) :: dim
${t1}$ :: res(size(mat)/size(mat, dim))
end function mean_2_${k1}$_${k1}$
#:endfor
end interface
Expand Down
89 changes: 45 additions & 44 deletions src/stdlib_experimental_stat_mean.f90
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,78 +29,79 @@ module function mean_1_qp_qp(mat) result(res)

end function mean_1_qp_qp

module function mean_2_sp_sp(mat, dim) result(res)
module function mean_2_all_sp_sp(mat) result(res)
real(sp), intent(in) :: mat(:,:)
integer, intent(in), optional :: dim
real(sp), allocatable ::res(:)
real(sp) ::res

integer :: i
integer :: dim_
res = sum(mat) / real(size(mat), sp)

end function mean_2_all_sp_sp
module function mean_2_all_dp_dp(mat) result(res)
real(dp), intent(in) :: mat(:,:)
real(dp) ::res

res = sum(mat) / real(size(mat), dp)

dim_ = optval(dim, 1)
end function mean_2_all_dp_dp
module function mean_2_all_qp_qp(mat) result(res)
real(qp), intent(in) :: mat(:,:)
real(qp) ::res

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

allocate(res(size(mat, dim_)))
end function mean_2_all_qp_qp

if (dim_ == 1) then
do i=1, size(mat, dim_)
res(i) = mean_1_sp_sp(mat(i,:))
end do
else if (dim_ == 2) then
do i=1, size(mat, dim_)
module function mean_2_sp_sp(mat, dim) result(res)
real(sp), intent(in) :: mat(:,:)
integer, intent(in) :: dim
real(sp) :: res(size(mat)/size(mat, dim))

integer :: i

if (dim == 1) then
do i=1, size(mat)/size(mat, dim)
res(i) = mean_1_sp_sp(mat(:,i))
end do
else if (dim == 2) then
do i=1, size(mat)/size(mat, dim)
res(i) = mean_1_sp_sp(mat(i,:))
end do
end if

end function mean_2_sp_sp
module function mean_2_dp_dp(mat, dim) result(res)
real(dp), intent(in) :: mat(:,:)
integer, intent(in), optional :: dim
real(dp), allocatable ::res(:)
integer, intent(in) :: dim
real(dp) :: res(size(mat)/size(mat, dim))

integer :: i
integer :: dim_

dim_ = optval(dim, 1)

if (dim_ < 0 .or. dim_ > 2 ) call error_stop("ERROR (mean): invalid argument (dim) ")

allocate(res(size(mat, dim_)))

if (dim_ == 1) then
do i=1, size(mat, dim_)
res(i) = mean_1_dp_dp(mat(i,:))
end do
else if (dim_ == 2) then
do i=1, size(mat, dim_)
if (dim == 1) then
do i=1, size(mat)/size(mat, dim)
res(i) = mean_1_dp_dp(mat(:,i))
end do
else if (dim == 2) then
do i=1, size(mat)/size(mat, dim)
res(i) = mean_1_dp_dp(mat(i,:))
end do
end if

end function mean_2_dp_dp
module function mean_2_qp_qp(mat, dim) result(res)
real(qp), intent(in) :: mat(:,:)
integer, intent(in), optional :: dim
real(qp), allocatable ::res(:)
integer, intent(in) :: dim
real(qp) :: res(size(mat)/size(mat, dim))

integer :: i
integer :: dim_

dim_ = optval(dim, 1)

if (dim_ < 0 .or. dim_ > 2 ) call error_stop("ERROR (mean): invalid argument (dim) ")

allocate(res(size(mat, dim_)))

if (dim_ == 1) then
do i=1, size(mat, dim_)
res(i) = mean_1_qp_qp(mat(i,:))
end do
else if (dim_ == 2) then
do i=1, size(mat, dim_)
if (dim == 1) then
do i=1, size(mat)/size(mat, dim)
res(i) = mean_1_qp_qp(mat(:,i))
end do
else if (dim == 2) then
do i=1, size(mat)/size(mat, dim)
res(i) = mean_1_qp_qp(mat(i,:))
end do
end if

end function mean_2_qp_qp
Expand Down
31 changes: 17 additions & 14 deletions src/stdlib_experimental_stat_mean.fypp.f90
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,28 +22,31 @@ end function mean_1_${k1}$_${k1}$
#:endfor

#:for i1, k1, t1 in ikt
module function mean_2_${k1}$_${k1}$(mat, dim) result(res)
module function mean_2_all_${k1}$_${k1}$(mat) result(res)
${t1}$, intent(in) :: mat(:,:)
integer, intent(in), optional :: dim
${t1}$, allocatable ::res(:)
${t1}$ ::res

integer :: i
integer :: dim_
res = sum(mat) / real(size(mat), ${k1}$)

dim_ = optval(dim, 1)
end function mean_2_all_${k1}$_${k1}$
#:endfor

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

allocate(res(size(mat, dim_)))
integer :: i

if (dim_ == 1) then
do i=1, size(mat, dim_)
res(i) = mean_1_${k1}$_${k1}$(mat(i,:))
end do
else if (dim_ == 2) then
do i=1, size(mat, dim_)
if (dim == 1) then
do i=1, size(mat)/size(mat, dim)
res(i) = mean_1_${k1}$_${k1}$(mat(:,i))
end do
else if (dim == 2) then
do i=1, size(mat)/size(mat, dim)
res(i) = mean_1_${k1}$_${k1}$(mat(i,:))
end do
end if

end function mean_2_${k1}$_${k1}$
Expand Down
10 changes: 6 additions & 4 deletions src/tests/stat/test_mean.f90
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,14 +11,16 @@ program test_mean
!sp
call loadtxt("array1.dat", s)

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

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

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


contains
Expand Down

0 comments on commit a1c6353

Please sign in to comment.