Merge pull request #7 from starkiller-astro/explicit_loops

Loop restructuring
starkiller-astro · Jul 15, 2024 · d8a33d3 · d8a33d3
2 parents 94020c1 + 2f0b92c
commit d8a33d3
Show file tree

Hide file tree

Showing 11 changed files with 1,533 additions and 1,295 deletions.
diff --git a/source/xnet_conditions.f90 b/source/xnet_conditions.f90
@@ -132,35 +132,8 @@ Subroutine t9rhofind_vector(kstep,tf,nf,t9f,rhof,mask_in)
 
     Do izb = zb_lo, zb_hi
       If ( mask(izb) ) Then
-
-        ! For constant conditions (nh = 1), set temperature and density
-        If ( nh(izb) == 1 ) Then
-          t9f(izb) = t9h(1,izb)
-          rhof(izb) = rhoh(1,izb)
-          nf(izb) = 1
-
-        ! Otherwise, calculate T9 and rho by interpolation
-        Else
-          Do n = 1, nh(izb)
-            If ( tf(izb) <= th(n,izb) ) Exit
-          EndDo
-          nf(izb) = n
-          If ( n > 1 .and. n <= nh(izb) ) Then
-            rdt = 1.0 / (th(n,izb)-th(n-1,izb))
-            dt = tf(izb) - th(n-1,izb)
-            dt9 = t9h(n,izb) - t9h(n-1,izb)
-            drho = rhoh(n,izb) - rhoh(n-1,izb)
-            t9f(izb) = dt*rdt*dt9 + t9h(n-1,izb)
-            rhof(izb) = dt*rdt*drho + rhoh(n-1,izb)
-          ElseIf ( n == 1 ) Then
-            t9f(izb) = t9h(1,izb)
-            rhof(izb) = rhoh(1,izb)
-          Else
-            t9f(izb) = t9h(nh(izb),izb)
-            rhof(izb) = rhoh(nh(izb),izb)
-            Write(lun_stdout,*) 'Time beyond thermodynamic range',tf(izb),' >',th(nh(izb),izb)
-          EndIf
-        EndIf
+        Call t9rhofind_scalar(kstep,tf(izb),nf(izb),t9f(izb),rhof(izb), &
+          & nh(izb),th(:,izb),t9h(:,izb),rhoh(:,izb))
       EndIf
     EndDo
 

diff --git a/source/xnet_controls.f90 b/source/xnet_controls.f90
@@ -205,7 +205,7 @@ Subroutine read_controls(data_dir)
     Allocate (zone_id(3,nzevolve))
     Allocate (lzactive(nzevolve))
     Allocate (iweak(nzevolve),lun_ev(nzevolve),lun_ts(nzevolve))
-    Allocate (kmon(2,nzevolve),ktot(5,nzevolve))
+    Allocate (kmon(5,nzevolve),ktot(5,nzevolve))
     !$omp parallel default(shared)
     zb_offset = (tid-1) * nzbatchmx
     zb_lo = zb_offset + 1

diff --git a/source/xnet_data.f90 b/source/xnet_data.f90
@@ -108,19 +108,19 @@ Subroutine partf(t9,mask_in)
     !-----------------------------------------------------------------------------------------------
     ! This routine calculates the nuclear partition functions as a function of temperature.
     !-----------------------------------------------------------------------------------------------
-    Use xnet_controls, Only: idiag, iheat, lun_diag, nzevolve, zb_lo, zb_hi, lzactive
+    Use xnet_controls, Only: idiag, iheat, lun_diag, zb_lo, zb_hi, lzactive
     Use xnet_types, Only: dp
     Use xnet_util, Only: safe_exp
     Implicit None
 
     ! Input variables
-    Real(dp), Intent(in) :: t9(nzevolve)
+    Real(dp), Intent(in) :: t9(zb_lo:zb_hi)
 
     ! Optional variables
     Logical, Optional, Target, Intent(in) :: mask_in(zb_lo:zb_hi)
 
     ! Local variables
-    Integer :: i, ii, izb
+    Integer :: i, ii, k, izb
     Real(dp) :: rdt9
     Logical, Pointer :: mask(:)
 
@@ -139,27 +139,39 @@ Subroutine partf(t9,mask_in)
         ii = i
 
         ! Linear interpolation in log-space
+        gg(0,izb) = 1.0 ! placeholder for non-nuclei, gamma-rays, etc.
         Select Case (ii)
         Case (1)
-          gg(1:ny,izb) = g(1,1:ny)
+          Do k = 1, ny
+            gg(k,izb) = g(1,k)
+          EndDo
         Case (ng+1)
-          gg(1:ny,izb) = g(ng,1:ny)
+          Do k = 1, ny
+            gg(k,izb) = g(ng,k)
+          EndDo
         Case Default
           rdt9 = (t9(izb)-t9i(ii-1)) / (t9i(ii)-t9i(ii-1))
-          gg(1:ny,izb) = safe_exp( rdt9*log(g(ii,1:ny)) + (1.0-rdt9)*log(g(ii-1,1:ny)) )
+          Do k = 1, ny
+            gg(k,izb) = safe_exp( rdt9*log(g(ii,k)) + (1.0-rdt9)*log(g(ii-1,k)) )
+          EndDo
         End Select
-        gg(0,izb) = 1.0 ! placeholder for non-nuclei, gamma-rays, etc.
 
         If ( iheat > 0 ) Then
+          dlngdt9(0,izb) = 0.0
           Select Case (ii)
           Case (1)
-            dlngdt9(1:ny,izb) = log(g(2,1:ny)/g(1,1:ny)) / (t9i(2)-t9i(1))
+            Do k = 1, ny
+              dlngdt9(k,izb) = log(g(2,k)/g(1,k)) / (t9i(2)-t9i(1))
+            EndDo
           Case (ng+1)
-            dlngdt9(1:ny,izb) = log(g(ng,1:ny)/g(ng-1,1:ny)) / (t9i(ng)-t9i(ng-1))
+            Do k = 1, ny
+              dlngdt9(k,izb) = log(g(ng,k)/g(ng-1,k)) / (t9i(ng)-t9i(ng-1))
+            EndDo
           Case Default
-            dlngdt9(1:ny,izb) = log(g(ii,1:ny)/g(ii-1,1:ny)) / (t9i(ii)-t9i(ii-1))
+            Do k = 1, ny
+              dlngdt9(k,izb) = log(g(ii,k)/g(ii-1,k)) / (t9i(ii)-t9i(ii-1))
+            EndDo
           End Select
-          dlngdt9(0,izb) = 0.0
         EndIf
       EndIf
     EndDo

diff --git a/source/xnet_evolve.f90 b/source/xnet_evolve.f90
@@ -54,12 +54,6 @@ Subroutine full_net(kstep)
     start_timer = xnet_wtime()
     timer_xnet = timer_xnet - start_timer
 
-    ! Initialize counters
-    kstep = 0
-    Do izb = zb_lo, zb_hi
-      kmon(:,izb) = 0
-      ktot(:,izb) = 0
-    EndDo
 
     ! Set reaction controls not read in from control
     idiag0 = idiag
@@ -76,10 +70,16 @@ Subroutine full_net(kstep)
       t9t(izb) = t9(izb)
       rhoo(izb) = rho(izb)
       rhot(izb) = rho(izb)
-      yet(izb) = ye(izb)
       yeo(izb) = ye(izb)
-      yo(:,izb) = y(:,izb)
-      yt(:,izb) = y(:,izb)
+      yet(izb) = ye(izb)
+      Do k = 1, ny
+        yo(k,izb) = y(k,izb)
+        yt(k,izb) = y(k,izb)
+      EndDo
+      Do k = 1, 5
+        kmon(k,izb) = 0
+        ktot(k,izb) = 0
+      EndDo
     EndDo
 
     en0 = 0.0
@@ -100,16 +100,20 @@ Subroutine full_net(kstep)
     ! Output initial abundances and conditions
     Call ts_output(0,delta_en,edot)
 
-    ! Start evolution
+    ! Initialize timestep loop flags
+    kstep = 0
+    mykstep = 0
+    lzsolve = lzactive(zb_lo:zb_hi)
+    lzoutput = lzsolve
     Do izb = zb_lo, zb_hi
-      If ( lzactive(izb) ) Then
+      If ( lzsolve(izb) ) Then
         its(izb) = 0
       Else
         its(izb) = -1
       EndIf
-      lzsolve(izb) = lzactive(izb)
-      mykstep(izb) = 0
     EndDo
+
+    ! Start evolution
     Do kstep = 1, kstmx
 
       ! Determine if this is an output step
@@ -129,47 +133,55 @@ Subroutine full_net(kstep)
         Call solve_be(kstep,its)
       End Select
 
-      Do izb = zb_lo, zb_hi
-        izone = izb + szbatch - zb_lo
+      ! If convergence is successful, output timestep results
+      If ( idiag >= 1 ) Then
+        Do izb = zb_lo, zb_hi
+          izone = izb + szbatch - zb_lo
+          If ( its(izb) == 0 .and. idiag >= 1 ) Then
+            Write(lun_diag,"(2(a,i5),5es14.7)") &
+              & 'KStep ',kstep,' Zone ',izone,t(izb),tdel(izb),t9(izb),rho(izb),ye(izb)
+            If ( idiag >= 3 ) Then
+              Write(lun_diag,"(a8)") 'delta Y'
+              Write(lun_diag,"(8x,a5,4es12.4)") &
+                & (nname(k),y(k,izb),yo(k,izb),(y(k,izb)-yo(k,izb)),(tdel(izb)*ydot(k,izb)),k=1,ny)
+              If ( iheat > 0 ) Write(lun_diag,"(8x,a5,4es12.4)") &
+                & 'T9',t9(izb),t9o(izb),t9(izb)-t9o(izb),tdel(izb)*t9dot(izb)
+            EndIf
+          ElseIf ( its(izb) == 1 ) Then
+            Write(lun_diag,"(a,2(i5,a))") 'KStep ',kstep,' Zone ',izone,' Inter!!!'
+          EndIf
+        EndDo
+      EndIf
 
-        ! If convergence is successful, output timestep results
+      Do izb = zb_lo, zb_hi
         If ( its(izb) == 0 ) Then
-          If ( idiag >= 1 ) Write(lun_diag,"(2(a,i5),5es14.7)") &
-            & 'KStep ',kstep,' Zone ',izone,t(izb),tdel(izb),t9(izb),rho(izb),ye(izb)
-          If ( idiag >= 3 ) Then
-            Write(lun_diag,"(a8)") 'delta Y'
-            Write(lun_diag,"(8x,a5,4es12.4)") &
-              & (nname(k),y(k,izb),yo(k,izb),(y(k,izb)-yo(k,izb)),(tdel(izb)*ydot(k,izb)),k=1,ny)
-            If ( iheat > 0 ) Write(lun_diag,"(8x,a5,4es12.4)") &
-              & 'T9',t9(izb),t9o(izb),t9(izb)-t9o(izb),tdel(izb)*t9dot(izb)
-          EndIf
+
           enold(izb) = enm(izb)
           Call benuc(yt(:,izb),enb(izb),enm(izb))
           delta_en(izb) = enm(izb) - en0(izb)
           edot(izb) = -(enm(izb)-enold(izb)) / tdel(izb)
 
-        ! If reduced timesteps fail to successfully integrate, warn and flag to remove from loop
-        ElseIf ( its(izb) == 1 ) Then
-          If ( idiag >= 0 ) Write(lun_diag,"(a,2(i5,a))") 'KStep ',kstep,' Zone ',izone,' Inter!!!'
-          its(izb) = 2
-        EndIf
-      EndDo
-
-      lzoutput = ( lzsolve .and. its <= 0 )
-      Call ts_output(kstep,delta_en,edot,mask_in = lzoutput)
-
-      Do izb = zb_lo, zb_hi
-        izone = izb + szbatch - zb_lo
-
           ! If this zone reaches the stop time, flag it to remove from loop
           If ( t(izb) >= tstop(izb) ) Then
             mykstep(izb) = kstep
             its(izb) = -1
+            lzsolve(izb) = .false.
           EndIf
+
+        ! If reduced timesteps fail to successfully integrate, flag it to remove from loop
+        ElseIf ( its(izb) == 1 ) Then
+          its(izb) = 2
+          lzsolve(izb) = .false.
+          lzoutput(izb) = .false.
+        ElseIf ( its(izb) < 0 ) Then
+          lzsolve(izb) = .false.
+          lzoutput(izb) = .false.
+        EndIf
       EndDo
 
+      Call ts_output(kstep,delta_en,edot,mask_in = lzoutput)
+
       ! Test if all zones have stopped
-      lzsolve = ( its == 0 )
       If ( .not. any( lzsolve ) ) Exit
     EndDo