! ch13_particle_swarm.f95
! Companion code for Chapter 13: Optimization and Data Fitting.
!
! Bounded minimization by particle swarm optimization (PSO).
! Edit the objective, bounds, and search settings below. No derivatives needed.
! The routine accepts a vector of any dimension. Console output only.
! Repeat with other seeds and search budgets; global optimality is not guaranteed.
!
! Compile (Fortran 95):
! gfortran -std=f95 -Wall -Wextra -O2 ch13_particle_swarm.f95 -o ch13_particle_swarm
! Run: ./ch13_particle_swarm (Windows: ch13_particle_swarm.exe)

module particle_swarm_methods
    implicit none
    private
    integer, parameter :: dp=selected_real_kind(15,307)
    public :: dp, particle_swarm, initialize_random
contains
    subroutine initialize_random(seed)
        integer, intent(in) :: seed
        integer, allocatable :: seeds(:)
        integer :: n,i
        call random_seed(size=n)
        allocate(seeds(n))
        do i=1,n
            seeds(i)=1+modulo(modulo(seed,1000000)+37*modulo(i,1000000),1000000)
        end do
        call random_seed(put=seeds)
        deallocate(seeds)
        ! Repeatable on the same compiler/runtime; sequences depend on compiler.
    end subroutine initialize_random

    subroutine particle_swarm(f,lower,upper,N,K,inertia,c_personal,c_global,vmax,pbest,fbest,evaluations)
        interface
            function f(p) result(value)
                integer, parameter :: dp=selected_real_kind(15,307)
                real(dp), intent(in) :: p(:)
                real(dp) :: value
            end function f
        end interface
        real(dp), intent(in) :: lower(:),upper(:),inertia,c_personal,c_global,vmax(:)
        integer, intent(in) :: N,K
        real(dp), intent(out) :: pbest(:),fbest
        integer, intent(out) :: evaluations
        real(dp) :: X(size(lower),N),V(size(lower),N),personal(size(lower),N)
        real(dp) :: values(N),r1(size(lower),N),r2(size(lower),N),value
        integer :: d,i,j,iteration,idx(1)

        d=size(lower)
        if (d<1 .or. size(upper)/=d .or. size(vmax)/=d .or. size(pbest)/=d) &
            stop 'Bounds, velocities, and result must have matching dimensions.'
        if (.not. all(finite_value(lower)) .or. .not. all(finite_value(upper))) &
            stop 'Bounds must be finite.'
        if (.not. all(finite_value(upper-lower)) .or. any(upper<=lower)) stop 'Invalid bound widths.'
        if (.not. all(finite_value(vmax)) .or. any(vmax<=0.0_dp) .or. any(vmax>upper-lower)) &
            stop 'Require positive vmax no larger than bound widths.'
        if (N<2 .or. K<1) stop 'Require at least two particles and a positive iteration count.'
        if (.not. finite_value(inertia) .or. inertia<0.0_dp .or. inertia>=1.0_dp) stop 'Invalid inertia.'
        if (.not. finite_value(c_personal) .or. c_personal<0.0_dp) stop 'Invalid personal attraction.'
        if (.not. finite_value(c_global) .or. c_global<0.0_dp) stop 'Invalid global attraction.'

        call random_number(r1)
        call random_number(r2)
        do i=1,N
            X(:,i)=lower+(upper-lower)*r1(:,i)
            V(:,i)=vmax*(2.0_dp*r2(:,i)-1.0_dp)
            values(i)=checked_value(f,X(:,i))
        end do
        personal=X
        idx=minloc(values)
        fbest=values(idx(1))
        pbest=personal(:,idx(1))
        evaluations=N

        do iteration=1,K
            call random_number(r1)
            call random_number(r2)
            ! All particles use the global best from the previous update.
            ! Independent random factors for each coordinate and particle.
            do i=1,N
                V(:,i)=inertia*V(:,i)+c_personal*r1(:,i)*(personal(:,i)-X(:,i)) &
                    +c_global*r2(:,i)*(pbest-X(:,i))
                if (.not. all(finite_value(V(:,i)))) stop 'Velocity overflow: rescale settings.'
                V(:,i)=max(min(V(:,i),vmax),-vmax)
                X(:,i)=X(:,i)+V(:,i)
                ! Reflect overshoots and damp the reflected velocities.
                ! vmax <= interval width ensures one reflection suffices.
                do j=1,d
                    if (X(j,i)<lower(j)) then
                        X(j,i)=lower(j)+(lower(j)-X(j,i))
                        V(j,i)=-0.5_dp*V(j,i)
                    else if (X(j,i)>upper(j)) then
                        X(j,i)=upper(j)-(X(j,i)-upper(j))
                        V(j,i)=-0.5_dp*V(j,i)
                    end if
                end do
                X(:,i)=max(min(X(:,i),upper),lower)
                value=checked_value(f,X(:,i))
                evaluations=evaluations+1
                if (value<values(i)) then
                    personal(:,i)=X(:,i)
                    values(i)=value
                end if
            end do
            idx=minloc(values)
            fbest=values(idx(1))
            pbest=personal(:,idx(1))
        end do
    end subroutine particle_swarm

    function checked_value(f,p) result(value)
        interface
            function f(p) result(value)
                integer, parameter :: dp=selected_real_kind(15,307)
                real(dp), intent(in) :: p(:)
                real(dp) :: value
            end function f
        end interface
        real(dp), intent(in) :: p(:)
        real(dp) :: value
        value=f(p)
        if (.not. finite_value(value)) stop 'Objective returned a nonfinite value.'
    end function checked_value

    elemental logical function finite_value(x)
        real(dp), intent(in) :: x
        finite_value=abs(x)<=huge(x)
    end function finite_value
end module particle_swarm_methods

module swarm_example
    use particle_swarm_methods, only: dp
    implicit none
contains
    function objective(p) result(value)
        real(dp), intent(in) :: p(:)
        real(dp) :: value
        value=sin(3.0_dp*p(1))*sin(3.0_dp*p(2))+0.1_dp*(p(1)**2+p(2)**2)
    end function objective
end module swarm_example

program ch13_particle_swarm
    use particle_swarm_methods, only: dp, particle_swarm, initialize_random
    use swarm_example, only: objective
    implicit none
    ! Editable settings. Bounds and velocity limits contain one entry per variable.
    integer, parameter :: d=2,swarm_size=25,max_iter=150,seed=2
    real(dp), parameter :: lower(d)=(/-3.0_dp,-3.0_dp/),upper(d)=(/3.0_dp,3.0_dp/)
    real(dp), parameter :: vmax(d)=(/0.5_dp,0.5_dp/)
    real(dp), parameter :: inertia=0.72_dp,c_personal=1.4_dp,c_global=1.4_dp
    real(dp) :: pbest(d),fbest
    integer :: evaluations,j
    ! integer :: clock_seed  ! Uncomment for clock-based seeding.

    call initialize_random(seed)
    ! Alternative: comment out the line above, uncomment the clock_seed
    ! declaration above, and uncomment both lines below.
    ! call system_clock(count=clock_seed)
    ! call initialize_random(clock_seed)

    call particle_swarm(objective,lower,upper,swarm_size,max_iter, &
        inertia,c_personal,c_global,vmax,pbest,fbest,evaluations)
    write(*,'(/,a)') 'Particle swarm optimization'
    do j=1,d
        write(*,'(a,i0,a,es23.15)') 'Best p(',j,') = ',pbest(j)
    end do
    write(*,'(a,es23.15)') 'Best objective value: ',fbest
    write(*,'(a,i0)') 'Swarm updates:        ',max_iter
    write(*,'(a,i0)') 'Function calls:       ',evaluations
    write(*,'(a)') 'Status: Search budget completed.'
    write(*,'(a)') 'Repeat with other seeds; global optimality is not guaranteed.'
end program ch13_particle_swarm