From 933ddb3f175ea07c04b86bc87482ea614eec3d9e Mon Sep 17 00:00:00 2001
From: Bhavin Patel <15210802+bpatel2107@users.noreply.github.com>
Date: Mon, 28 Oct 2024 11:01:39 +0000
Subject: [PATCH] Add species dependency to stress output

---
 cgyro/src/cgyro_flux.f90           | 3 ++-
 cgyro/src/cgyro_globals.F90        | 4 ++--
 cgyro/src/cgyro_init_manager.F90   | 4 ++--
 cgyro/src/cgyro_write_timedata.f90 | 4 ++--
 4 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/cgyro/src/cgyro_flux.f90 b/cgyro/src/cgyro_flux.f90
index 4da5f6765..5d61efae2 100644
--- a/cgyro/src/cgyro_flux.f90
+++ b/cgyro/src/cgyro_flux.f90
@@ -237,6 +237,7 @@ subroutine cgyro_flux
            do iv=nv1,nv2
 
               iv_loc = iv_loc+1
+              is = is_v(iv)
               ix = ix_v(iv)
               ie = ie_v(iv)
 
@@ -246,7 +247,7 @@ subroutine cgyro_flux
               do ic=1,nc
                  ir = ir_c(ic)
                  it = it_c(ic)
-                 stress_integrated_loc(ir, it, itor, :) = stress_integrated_loc(ir, it, itor, :) + stress(ic, iv_loc, itor, :) * dv
+                 stress_integrated_loc(ir, it, is, itor, :) = stress_integrated_loc(ir, it, is, itor, :) + stress(ic, iv_loc, itor, :) * dv
               end do
            enddo
         end if
diff --git a/cgyro/src/cgyro_globals.F90 b/cgyro/src/cgyro_globals.F90
index 40639d3c5..5ba51e664 100644
--- a/cgyro/src/cgyro_globals.F90
+++ b/cgyro/src/cgyro_globals.F90
@@ -345,8 +345,8 @@ module cgyro_globals
   ! Distributions
   complex, dimension(:,:,:,:), allocatable :: rhs
   complex, dimension(:,:,:,:), allocatable :: stress
-  complex, dimension(:,:,:,:), allocatable :: stress_integrated
-  complex, dimension(:,:,:,:), allocatable :: stress_integrated_loc
+  complex, dimension(:,:,:,:,:), allocatable :: stress_integrated
+  complex, dimension(:,:,:,:,:), allocatable :: stress_integrated_loc
   complex, dimension(:,:,:), allocatable :: h_x
   complex, dimension(:,:,:), allocatable :: g_x
   complex, dimension(:,:,:), allocatable :: h0_x
diff --git a/cgyro/src/cgyro_init_manager.F90 b/cgyro/src/cgyro_init_manager.F90
index f7e78631e..d42c0dc9a 100644
--- a/cgyro/src/cgyro_init_manager.F90
+++ b/cgyro/src/cgyro_init_manager.F90
@@ -309,8 +309,8 @@ subroutine cgyro_init_manager
 
        if (stress_print_flag == 1) then
           allocate(stress(nc,nv_loc,nt1:nt2,n_field))
-          allocate(stress_integrated_loc(n_radial,n_theta,nt1:nt2,n_field))
-          allocate(stress_integrated(n_radial,n_theta,nt1:nt2,n_field))
+          allocate(stress_integrated_loc(n_radial,n_theta,n_species,nt1:nt2,n_field))
+          allocate(stress_integrated(n_radial,n_theta,n_species,nt1:nt2,n_field))
 #if defined(OMPGPU)
 !$omp target enter data map(alloc:stress,stress_integrated,stress_integrated_loc)
 #elif defined(_OPENACC)
diff --git a/cgyro/src/cgyro_write_timedata.f90 b/cgyro/src/cgyro_write_timedata.f90
index 1183fa30b..7741afd40 100644
--- a/cgyro/src/cgyro_write_timedata.f90
+++ b/cgyro/src/cgyro_write_timedata.f90
@@ -98,8 +98,8 @@ subroutine cgyro_write_timedata
      ! Sort out subset of theta values for plotting
         call cgyro_write_distributed_bcomplex(&
              trim(path)//binfile_stress(i_field),&
-             size(stress_integrated(:,:,:,i_field)),&
-             stress_integrated(:,:,:,i_field))
+             size(stress_integrated(:,:,:,:,i_field)),&
+             stress_integrated(:,:,:,:,i_field))
      enddo
 
   end if