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

Commit

Permalink
Vectorized drift
Browse files Browse the repository at this point in the history
  • Loading branch information
@daminton
daminton committed Sep 21, 2021
1 parent 3b9999c commit de11649
Show file tree
Hide file tree
Showing 9 changed files with 158 additions and 97 deletions.
29 changes: 18 additions & 11 deletions Makefile
View file
Original file line number Diff line number Diff line change
Expand Up @@ -93,11 +93,6 @@ lib:
ln -s $(SWIFTEST_HOME)/Makefile.Defines .; \
ln -s $(SWIFTEST_HOME)/Makefile .; \
make libdir
cd $(SWIFTEST_HOME)/src/drift; \
rm -f Makefile.Defines Makefile; \
ln -s $(SWIFTEST_HOME)/Makefile.Defines .; \
ln -s $(SWIFTEST_HOME)/Makefile .; \
make libdir
cd $(SWIFTEST_HOME)/src/gr; \
rm -f Makefile.Defines Makefile; \
ln -s $(SWIFTEST_HOME)/Makefile.Defines .; \
Expand Down Expand Up @@ -128,11 +123,6 @@ lib:
ln -s $(SWIFTEST_HOME)/Makefile.Defines .; \
ln -s $(SWIFTEST_HOME)/Makefile .; \
make libdir
cd $(SWIFTEST_HOME)/src/orbel; \
rm -f Makefile.Defines Makefile; \
ln -s $(SWIFTEST_HOME)/Makefile.Defines .; \
ln -s $(SWIFTEST_HOME)/Makefile .; \
make libdir
cd $(SWIFTEST_HOME)/src/setup; \
rm -f Makefile.Defines Makefile; \
ln -s $(SWIFTEST_HOME)/Makefile.Defines .; \
Expand Down Expand Up @@ -187,6 +177,24 @@ fast:
ln -s $(SWIFTEST_HOME)/Makefile .; \
make fastdir

cd $(SWIFTEST_HOME)/src/orbel; \
rm -f Makefile.Defines Makefile; \
ln -s $(SWIFTEST_HOME)/Makefile.Defines .; \
ln -s $(SWIFTEST_HOME)/Makefile .; \
make fastdir

cd $(SWIFTEST_HOME)/src/drift; \
rm -f Makefile.Defines Makefile; \
ln -s $(SWIFTEST_HOME)/Makefile.Defines .; \
ln -s $(SWIFTEST_HOME)/Makefile .; \
make fastdir

cd $(SWIFTEST_HOME)/src/helio; \
$(FORTRAN) $(FFASTFLAGS) -I$(SWIFTEST_HOME)/include -I$(NETCDF_FORTRAN_HOME)/include $(IMKL) -c helio_drift.f90; \
$(AR) rv $(SWIFTEST_HOME)/lib/libswiftest.a *.o *.smod; \
$(INSTALL_DATA) *.smod $(SWIFTEST_HOME)/include; \
rm -f *.o *.smod

cd $(SWIFTEST_HOME)/src/rmvs; \
$(FORTRAN) $(FFASTFLAGS) -I$(SWIFTEST_HOME)/include -I$(NETCDF_FORTRAN_HOME)/include $(IMKL) -c rmvs_encounter_check.f90; \
$(AR) rv $(SWIFTEST_HOME)/lib/libswiftest.a *.o *.smod; \
Expand All @@ -211,7 +219,6 @@ fastdir:
$(INSTALL_DATA) *.smod $(SWIFTEST_HOME)/include; \
rm -f *.o *.smod


drivers:
cd $(SWIFTEST_HOME)/src/main; \
rm -f Makefile.Defines Makefile; \
Expand Down
43 changes: 24 additions & 19 deletions src/drift/drift.f90
View file
Original file line number Diff line number Diff line change
Expand Up @@ -77,18 +77,18 @@ module subroutine drift_all(mu, x, v, n, param, dt, lmask, iflag)
else
where(lmask(1:n)) dtp(1:n) = dt
end if
!$omp parallel do default(private) &
!$omp shared(n, lmask, mu, x, v, dtp, iflag)
do i = 1,n
!$omp simd
do i = 1, n
if (lmask(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
!$omp end parallel do
!$omp end simd

return
end subroutine drift_all


module pure elemental subroutine drift_one(mu, px, py, pz, vx, vy, vz, dt, iflag)
!$omp declare simd(drift_one)
!! author: The Purdue Swiftest Team - David A. Minton, Carlisle A. Wishard, Jennifer L.L. Pouplin, and Jacob R. Elliott
!!
!! Perform Danby drift for one body, redoing drift with smaller substeps if original accuracy is insufficient
Expand All @@ -104,27 +104,22 @@ module pure elemental subroutine drift_one(mu, px, py, pz, vx, vy, vz, dt, iflag
! Internals
integer(I4B) :: i
real(DP) :: dttmp
real(DP), dimension(NDIM) :: x, v

x = [px, py, pz]
v = [vx, vy, vz]

call drift_dan(mu, x(:), v(:), dt, iflag)
call drift_dan(mu, px, py, pz, vx, vy, vz, dt, iflag)
if (iflag /= 0) then
dttmp = 0.1_DP * dt
do i = 1, 10
call drift_dan(mu, x(:), v(:), dttmp, iflag)
call drift_dan(mu, px, py, pz, vx, vy, vz, dttmp, iflag)
if (iflag /= 0) exit
end do
end if
px = x(1); py = x(2); pz = x(3)
vx = v(1); vy = v(2); vz = v(3)

return
end subroutine drift_one


pure subroutine drift_dan(mu, x0, v0, dt0, iflag)
pure subroutine drift_dan(mu, px0, py0, pz0, vx0, vy0, vz0, dt0, iflag)
!$omp declare simd(drift_dan)
!! author: David A. Minton
!!
!! Perform Kepler drift, solving Kepler's equation in appropriate variables
Expand All @@ -134,14 +129,16 @@ pure subroutine drift_dan(mu, x0, v0, dt0, iflag)
implicit none
integer(I4B), intent(out) :: iflag
real(DP), intent(in) :: mu, dt0
real(DP), dimension(:), intent(inout) :: x0, v0
real(DP), intent(inout) :: px0, py0, pz0, vx0, vy0, vz0
real(DP) :: dt, f, g, fdot, gdot, c1, c2, c3, u, alpha, fp, r0
real(DP) :: v0s, a, asq, en, dm, ec, es, esq, xkep, fchk, s, c
real(DP), dimension(NDIM) :: x, v
real(DP), dimension(NDIM) :: x, v, x0, v0

! Executable code
iflag = 0
dt = dt0
x0 = [px0, py0, pz0]
v0 = [vx0, vy0, vz0]
r0 = sqrt(dot_product(x0(:), x0(:)))
v0s = dot_product(v0(:), v0(:))
u = dot_product(x0(:), v0(:))
Expand Down Expand Up @@ -172,8 +169,8 @@ pure subroutine drift_dan(mu, x0, v0, dt0, iflag)
gdot = (c - 1.0_DP) / fp + 1.0_DP
x(:) = x0(:) * f + v0(:) * g
v(:) = x0(:) * fdot + v0(:) * gdot
x0(:) = x(:)
v0(:) = v(:)
px0 = x(1); py0 = x(2); pz0 = x(3)
vx0 = v(1); vy0 = v(2); vz0 = v(3)
iflag = 0
return
end if
Expand All @@ -187,15 +184,16 @@ pure subroutine drift_dan(mu, x0, v0, dt0, iflag)
gdot = 1.0_DP - mu / fp * c2
x(:) = x0(:) * f + v0(:) * g
v(:) = x0(:) * fdot + v0(:) * gdot
x0(:) = x(:)
v0(:) = v(:)
px0 = x(1); py0 = x(2); pz0 = x(3)
vx0 = v(1); vy0 = v(2); vz0 = v(3)
end if

return
end subroutine drift_dan


pure subroutine drift_kepmd(dm, es, ec, x, s, c)
!$omp declare simd(drift_kepmd)
!! author: David A. Minton
!!
!! Solve Kepler's equation in difference form for an ellipse for small input dm and eccentricity
Expand Down Expand Up @@ -235,6 +233,7 @@ end subroutine drift_kepmd


pure subroutine drift_kepu(dt,r0,mu,alpha,u,fp,c1,c2,c3,iflag)
!$omp declare simd(drift_kepu)
!! author: David A. Minton
!!
!! Solve Kepler's equation in universal variables
Expand Down Expand Up @@ -263,6 +262,7 @@ end subroutine drift_kepu


pure subroutine drift_kepu_fchk(dt, r0, mu, alpha, u, s, f)
!$omp declare simd(drift_kepu_fchk)
!! author: David A. Minton
!!
!! Computes the value of f, the function whose root we are trying to find in universal variables
Expand All @@ -286,6 +286,7 @@ end subroutine drift_kepu_fchk


pure subroutine drift_kepu_guess(dt, r0, mu, alpha, u, s)
!$omp declare simd(drift_kepu_guess)
!! author: David A. Minton
!!
!! Compute initial guess for solving Kepler's equation using universal variables
Expand Down Expand Up @@ -324,6 +325,7 @@ end subroutine drift_kepu_guess


pure subroutine drift_kepu_lag(s, dt, r0, mu, alpha, u, fp, c1, c2, c3, iflag)
!$omp declare simd(drift_kepu_lag)
!! author: David A. Minton
!!
!! Solve Kepler's equation in universal variables using Laguerre's method
Expand Down Expand Up @@ -369,6 +371,7 @@ end subroutine drift_kepu_lag


pure subroutine drift_kepu_new(s, dt, r0, mu, alpha, u, fp, c1, c2, c3, iflag)
!$omp declare simd(drift_kepu_new)
!! author: David A. Minton
!!
!! Solve Kepler's equation in universal variables using Newton's method
Expand Down Expand Up @@ -411,6 +414,7 @@ end subroutine drift_kepu_new


pure subroutine drift_kepu_p3solve(dt, r0, mu, alpha, u, s, iflag)
!$omp declare simd(drift_kepu_p3solve)
!! author: David A. Minton
!!
!! Computes real root of cubic involved in setting initial guess for solving Kepler's equation in universal variables
Expand Down Expand Up @@ -455,6 +459,7 @@ end subroutine drift_kepu_p3solve


pure subroutine drift_kepu_stumpff(x, c0, c1, c2, c3)
!$omp declare simd(drift_kepu_stumpff)
!! author: David A. Minton
!!
!! Compute Stumpff functions needed for Kepler drift in universal variables
Expand Down
39 changes: 36 additions & 3 deletions src/helio/helio_drift.f90
View file
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,41 @@ module subroutine helio_drift_tp(self, system, param, dt)
return
end subroutine helio_drift_tp


pure elemental subroutine helio_drift_linear_one(xhx, xhy, xhz, ptx, pty, ptz, dt)
!$omp declare simd(helio_drift_linear_one)
implicit none
real(DP), intent(inout) :: xhx, xhy, xhz
real(DP), intent(in) :: ptx, pty, ptz, dt

xhx = xhx + ptx * dt
xhy = xhy + pty * dt
xhz = xhz + ptz * dt

return
end subroutine helio_drift_linear_one


subroutine helio_drift_linear_all(xh, pt, dt, n, lmask)
implicit none
! Arguments
real(DP), dimension(:,:), intent(inout) :: xh
real(DP), dimension(:), intent(in) :: pt
real(DP), intent(in) :: dt
integer(I4B), intent(in) :: n
logical, dimension(:), intent(in) :: lmask
! Internals
integer(I4B) :: i

!$omp parallel do simd default(shared) schedule(static)
do i = 1, n
if (lmask(i)) call helio_drift_linear_one(xh(1,i), xh(2,i), xh(3,i), pt(1), pt(2), pt(3), dt)
end do
!$omp end parallel do simd

return
end subroutine helio_drift_linear_all


module subroutine helio_drift_linear_pl(self, cb, dt, lbeg)
!! author: David A. Minton
Expand All @@ -100,9 +135,7 @@ module subroutine helio_drift_linear_pl(self, cb, dt, lbeg)
pt(2) = sum(pl%Gmass(1:npl) * pl%vb(2,1:npl), self%lmask(1:npl))
pt(3) = sum(pl%Gmass(1:npl) * pl%vb(3,1:npl), self%lmask(1:npl))
pt(:) = pt(:) / cb%Gmass
do concurrent(i = 1:npl, self%lmask(i))
pl%xh(:,i) = pl%xh(:,i) + pt(:) * dt
end do
call helio_drift_linear_all(pl%xh(:,:), pt(:), dt, npl, pl%lmask(:))

if (lbeg) then
cb%ptbeg = pt(:)
Expand Down
47 changes: 19 additions & 28 deletions src/modules/swiftest_classes.f90
View file
Original file line number Diff line number Diff line change
Expand Up @@ -558,6 +558,7 @@ module subroutine drift_body(self, system, param, dt)
end subroutine drift_body

module pure elemental subroutine drift_one(mu, px, py, pz, vx, vy, vz, dt, iflag)
!$omp declare simd(drift_one)
implicit none
real(DP), intent(in) :: mu !! G * (Mcb + m), G = gravitational constant, Mcb = mass of central body, m = mass of body to drift
real(DP), intent(inout) :: px, py, pz, vx, vy, vz !! Position and velocity of body to drift
Expand Down Expand Up @@ -1059,45 +1060,35 @@ module subroutine orbel_el2xv_vec(self, cb)
end subroutine orbel_el2xv_vec

module pure subroutine orbel_scget(angle, sx, cx)
!$omp declare simd(orbel_scget)
implicit none
real(DP), intent(in) :: angle
real(DP), intent(out) :: sx, cx
end subroutine orbel_scget

module pure subroutine orbel_xv2aeq(mu, x, v, a, e, q)
module pure subroutine orbel_xv2aeq(mu, px, py, pz, vx, vy, vz, a, e, q)
!$omp declare simd(orbel_xv2aeq)
implicit none
real(DP), intent(in) :: mu !! Gravitational constant
real(DP), dimension(:), intent(in) :: x !! Position vector
real(DP), dimension(:), intent(in) :: v !! Velocity vector
real(DP), intent(out) :: a !! semimajor axis
real(DP), intent(out) :: e !! eccentricity
real(DP), intent(out) :: q !! periapsis
real(DP), intent(in) :: mu !! Gravitational constant
real(DP), intent(in) :: px,py,pz !! Position vector
real(DP), intent(in) :: vx,vy,vz !! Velocity vector
real(DP), intent(out) :: a !! semimajor axis
real(DP), intent(out) :: e !! eccentricity
real(DP), intent(out) :: q !! periapsis
end subroutine orbel_xv2aeq

module pure subroutine orbel_xv2aqt(mu, x, v, a, q, capm, tperi)
module pure subroutine orbel_xv2aqt(mu, px, py, pz, vx, vy, vz, a, q, capm, tperi)
!$omp declare simd(orbel_xv2aqt)
implicit none
real(DP), intent(in) :: mu !! Gravitational constant
real(DP), dimension(:), intent(in) :: x !! Position vector
real(DP), dimension(:), intent(in) :: v !! Velocity vector
real(DP), intent(out) :: a !! semimajor axis
real(DP), intent(out) :: q !! periapsis
real(DP), intent(out) :: capm !! mean anomaly
real(DP), intent(out) :: tperi !! time of pericenter passage
real(DP), intent(in) :: mu !! Gravitational constant
real(DP), intent(in) :: px,py,pz !! Position vector
real(DP), intent(in) :: vx,vy,vz !! Velocity vector
real(DP), intent(out) :: a !! semimajor axis
real(DP), intent(out) :: q !! periapsis
real(DP), intent(out) :: capm !! mean anomaly
real(DP), intent(out) :: tperi !! time of pericenter passage
end subroutine orbel_xv2aqt

module pure subroutine orbel_xv2el(mu, x, v, a, e, inc, capom, omega, capm)
implicit none
real(DP), intent(in) :: mu !! Gravitational constant
real(DP), dimension(:), intent(in) :: x !! Position vector
real(DP), dimension(:), intent(in) :: v !! Velocity vector
real(DP), intent(out) :: a !! semimajor axis
real(DP), intent(out) :: e !! eccentricity
real(DP), intent(out) :: inc !! inclination
real(DP), intent(out) :: capom !! longitude of ascending node
real(DP), intent(out) :: omega !! argument of periapsis
real(DP), intent(out) :: capm !! mean anomaly
end subroutine orbel_xv2el

module subroutine orbel_xv2el_vec(self, cb)
implicit none
class(swiftest_body), intent(inout) :: self !! Swiftest body object
Expand Down
Loading

0 comments on commit de11649

Please sign in to comment.