Skip to content
This repository was archived by the owner on Aug 28, 2024. It is now read-only.

Commit

Permalink
Consolidated redundant drift subroutines into one single shared algor…
Browse files Browse the repository at this point in the history 
…ithm. Each type-bound procedure now feeds the correct position and velocities to the universal drift_all function
  • Loading branch information
@daminton
daminton committed Jul 23, 2021
1 parent 205e666 commit 2ef9364
Show file tree
Hide file tree
Showing 5 changed files with 93 additions and 165 deletions.
99 changes: 16 additions & 83 deletions examples/helio_swifter_comparison/swiftest_vs_swifter.ipynb
View file

Large diffs are not rendered by default.

63 changes: 43 additions & 20 deletions src/drift/drift.f90
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,41 +26,64 @@ module subroutine drift_body(self, system, param, dt, mask)
logical, dimension(:), intent(in) :: mask !! Logical mask of size self%nbody that determines which bodies to drift.
! Internals
integer(I4B) :: i
real(DP) :: energy, vmag2, rmag !! Variables used in GR calculation
integer(I4B), dimension(:), allocatable :: iflag
real(DP), dimension(:), allocatable :: dtp

associate(n => self%nbody)
if (n == 0) return
allocate(iflag(n))
iflag(:) = 0
allocate(dtp(n))
if (param%lgr) then
do concurrent(i = 1:n, mask(i))
rmag = norm2(self%xh(:, i))
vmag2 = dot_product(self%vh(:, i), self%vh(:, i))
energy = 0.5_DP * vmag2 - self%mu(i) / rmag
dtp(i) = dt * (1.0_DP + 3 * param%inv_c2 * energy)
end do
else
where(mask(1:n)) dtp(1:n) = dt
end if
do concurrent(i = 1:n, mask(i))
call drift_one(self%mu(i), self%xh(1,i), self%xh(2,i), self%xh(3,i), &
self%vh(1,i), self%vh(2,i), self%vh(3,i), &
dtp(i), iflag(i))
end do
call drift_all(self%mu, self%xh, self%vh, self%nbody, param, dt, mask, iflag)
if (any(iflag(1:n) /= 0)) then
where(iflag(1:n) /= 0) self%status(1:n) = DISCARDED_DRIFTERR
do i = 1, n
if (iflag(i) /= 0) write(*, *) "Particle ", self%id(i), " lost due to error in Danby drift"
if (iflag(i) /= 0) write(*, *) " Body ", self%id(i), " lost due to error in Danby drift"
end do
end if
end associate

return
end subroutine drift_body

module pure subroutine drift_all(mu, x, v, n, param, dt, mask, iflag)
!! author: David A. Minton
!!
!! Loop bodies and call Danby drift routine on all bodies for the given position and velocity vector.
!!
!! Adapted from Hal Levison's Swift routine drift_tp.f
!! Adapted from David E. Kaufmann's Swifter routine whm_drift_tp.f9
implicit none
! Arguments
real(DP), dimension(:), intent(in) :: mu !! Vector of gravitational constants
real(DP), dimension(:,:), intent(inout) :: x, v !! Position and velocity vectors
integer(I4B), intent(in) :: n !! number of bodies
class(swiftest_parameters), intent(in) :: param !! Current run configuration parameters
real(DP), intent(in) :: dt !! Stepsize
logical, dimension(:), intent(in) :: mask !! Logical mask of size self%nbody that determines which bodies to drift.
integer(I4B), dimension(:), intent(out) :: iflag !! Vector of error flags. 0 means no problem
! Internals
integer(I4B) :: i
real(DP) :: energy, vmag2, rmag !! Variables used in GR calculation
real(DP), dimension(:), allocatable :: dtp

if (n == 0) return

allocate(dtp(n))
if (param%lgr) then
do concurrent(i = 1:n, mask(i))
rmag = norm2(x(:, i))
vmag2 = dot_product(v(:, i), v(:, i))
energy = 0.5_DP * vmag2 - mu(i) / rmag
dtp(i) = dt * (1.0_DP + 3 * param%inv_c2 * energy)
end do
else
where(mask(1:n)) dtp(1:n) = dt
end if
do concurrent(i = 1:n, mask(i))
call drift_one(mu(i), x(1,i), x(2,i), x(3,i), v(1,i), v(2,i), v(3,i), dtp(i), iflag(i))
end do

return
end subroutine drift_all

module pure elemental subroutine drift_one(mu, px, py, pz, vx, vy, vz, dt, iflag)
!! author: The Purdue Swiftest Team - David A. Minton, Carlisle A. Wishard, Jennifer L.L. Pouplin, and Jacob R. Elliott
!!
Expand Down
48 changes: 13 additions & 35 deletions src/helio/helio_drift.f90
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,41 +22,19 @@ module subroutine helio_drift_body(self, system, param, dt, mask)
integer(I4B), dimension(:),allocatable :: iflag !! Vectorized error code flag
real(DP), dimension(:), allocatable :: dtp, mu

if (self%nbody == 0) return

allocate(iflag(self%nbody))
iflag(:) = 0
allocate(dtp(self%nbody))
allocate(mu(self%nbody))
mu(:) = system%cb%Gmass

if (param%lgr) then
do concurrent(i = 1:self%nbody, mask(i))
rmag = norm2(self%xh(:, i))
vmag2 = dot_product(self%vb(:, i), self%vb(:, i))
energy = 0.5_DP * vmag2 - mu(i) / rmag
dtp(i) = dt * (1.0_DP + 3 * param%inv_c2 * energy)
end do
else
where(mask(1:self%nbody)) dtp(1:self%nbody) = dt
end if

do concurrent(i = 1:self%nbody, mask(i))
call drift_one(mu(i), self%xh(1,i), self%xh(2,i), self%xh(3,i), &
self%vb(1,i), self%vb(2,i), self%vb(3,i), &
dtp(i), iflag(i))
end do
if (any(iflag(1:self%nbody) /= 0)) then
do i = 1, self%nbody
if (iflag(i) /= 0) then
write(*, *) " Body", self%id(i), " is lost!!!!!!!!!!"
write(*, *) self%xh(:,i)
write(*, *) self%vb(:,i)
write(*, *) " stopping "
call util_exit(FAILURE)
end if
end do
end if
associate(n => self%nbody)
allocate(iflag(n))
iflag(:) = 0
allocate(mu(n))
mu(:) = system%cb%Gmass
call drift_all(mu, self%xh, self%vb, self%nbody, param, dt, mask, iflag)
if (any(iflag(1:n) /= 0)) then
where(iflag(1:n) /= 0) self%status(1:n) = DISCARDED_DRIFTERR
do i = 1, n
if (iflag(i) /= 0) write(*, *) " Body ", self%id(i), " lost due to error in Danby drift"
end do
end if
end associate

return
end subroutine helio_drift_body
Expand Down
13 changes: 12 additions & 1 deletion src/modules/swiftest_classes.f90
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ module swiftest_classes
implicit none
private
public :: discard_pl, discard_system, discard_tp
public :: drift_body, drift_one
public :: drift_all, drift_body, drift_one
public :: eucl_dist_index_plpl, eucl_dist_index_pltp
public :: gr_getaccb_ns_body, gr_p4_pos_kick, gr_pseudovel2vel, gr_vel2pseudovel
public :: io_dump_param, io_dump_swiftest, io_dump_system, io_get_args, io_get_token, io_param_reader, io_param_writer, io_read_body_in, &
Expand Down Expand Up @@ -383,6 +383,17 @@ module subroutine discard_tp(self, system, param)
class(swiftest_parameters), intent(in) :: param !! Current run configuration parameters
end subroutine discard_tp

module pure subroutine drift_all(mu, x, v, n, param, dt, mask, iflag)
implicit none
real(DP), dimension(:), intent(in) :: mu !! Vector of gravitational constants
real(DP), dimension(:,:), intent(inout) :: x, v !! Position and velocity vectors
integer(I4B), intent(in) :: n !! number of bodies
class(swiftest_parameters), intent(in) :: param !! Current run configuration parameters
real(DP), intent(in) :: dt !! Stepsize
logical, dimension(:), intent(in) :: mask !! Logical mask of size self%nbody that determines which bodies to drift.
integer(I4B), dimension(:), intent(out) :: iflag !! Vector of error flags. 0 means no problem
end subroutine drift_all

module subroutine drift_body(self, system, param, dt, mask)
implicit none
class(swiftest_body), intent(inout) :: self !! Swiftest particle data structure
Expand Down
35 changes: 9 additions & 26 deletions src/whm/whm_drift.f90
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,43 +17,26 @@ module subroutine whm_drift_pl(self, system, param, dt, mask)
logical, dimension(:), intent(in) :: mask !! Logical mask of size self%nbody that determines which bodies to drift
! Internals
integer(I4B) :: i
real(DP) :: energy, vmag2, rmag !! Variables used in GR calculation
integer(I4B), dimension(:), allocatable :: iflag
real(DP), dimension(:), allocatable :: dtp

associate(pl => self, npl => self%nbody)
if (npl == 0) return

allocate(iflag(npl))
iflag(:) = 0
allocate(dtp(npl))

if (param%lgr) then
do concurrent(i = 1:npl, mask(i))
rmag = norm2(pl%xj(:, i))
vmag2 = dot_product(pl%vj(:, i), pl%vj(:, i))
energy = 0.5_DP * vmag2 - pl%muj(i) / rmag
dtp(i) = dt * (1.0_DP + 3 * param%inv_c2 * energy)
end do
else
where(mask(1:npl)) dtp(1:npl) = dt
end if

do concurrent(i = 1:npl, mask(i))
call drift_one(pl%muj(i), pl%xj(1,i), pl%xj(2,i), pl%xj(3,i), &
pl%vj(1,i), pl%vj(2,i), pl%vj(3,i), &
dtp(i), iflag(i))
end do
call drift_all(pl%muj, pl%xj, pl%vj, npl, param, dt, mask, iflag)
if (any(iflag(1:npl) /= 0)) then
where(iflag(1:npl) /= 0) pl%status(1:npl) = DISCARDED_DRIFTERR
do i = 1, npl
if (iflag(i) /= 0) then
write(*, *) " Planet ", self%id(i), " is lost!!!!!!!!!!"
write(*, *) pl%xj(:,i)
write(*, *) pl%vj(:,i)
write(*, *) " stopping "
call util_exit(FAILURE)
if (iflag(i) /= 0) then
write(*, *) " Planet ", pl%id(i), " is lost!!!!!!!!!!!!"
WRITE(*, *) pl%muj(i), dt
WRITE(*, *) pl%xj(:,i)
WRITE(*, *) pl%vj(:,i)
WRITE(*, *) " STOPPING "
end if
end do
call util_exit(FAILURE)
end if
end associate

Expand Down

0 comments on commit 2ef9364

Please sign in to comment.