Added infrastructure to allow users to turn on or off coarrays

python/swiftest/swiftest/simulation_class.py

@@ -299,6 +299,11 @@ def __init__(self,read_param: bool = False,
             Parameter input file equivalent: `ENCOUNTER_CHECK`
         dask : bool, default False
             Use Dask to lazily load data (useful for very large datasets)
+        coarray : bool, default False
+            If true, will employ Coarrays on test particle structures to run in single program/multiple data parallel mode. 
+            *Note" In order to use this capability, Swiftest must be compiled for Coarray support. Only certain integrators
+            can use Coarrays: RMVS, WHM, Helio are all compatible, but SyMBA is not, due to the way tp-pl close encounters 
+            are handeled.
         verbose : bool, default True
             If set to True, then more information is printed by Simulation methods as they are executed. Setting to
             False suppresses most messages other than errors.
@@ -824,6 +829,7 @@ def set_parameter(self, verbose: bool = True, **kwargs):
             "ephemeris_date": "MBCL",
             "restart": False,
             "encounter_save" : "NONE",
+            "coarray" : False,
             "simdir" : self.simdir
         }
         param_file = kwargs.pop("param_file",None)
@@ -1065,6 +1071,7 @@ def set_feature(self,
                     interaction_loops: Literal["TRIANGULAR", "FLAT"] | None = None,
                     encounter_check_loops: Literal["TRIANGULAR", "SORTSWEEP"] | None = None,
                     encounter_save: Literal["NONE", "TRAJECTORY", "CLOSEST", "BOTH"] | None = None,
+                    coarray: bool | None = None,
                     verbose: bool | None = None,
                     simdir: str | os.PathLike = None, 
                     **kwargs: Any
@@ -1130,6 +1137,11 @@ def set_feature(self,
             * "SORTSWEEP" : A Sort-Sweep algorithm is used to reduce the population of potential close encounter bodies.
               This algorithm is still in development, and does not necessarily speed up the encounter checking.
               Use with caution.
+        coarray : bool, default False
+            If true, will employ Coarrays on test particle structures to run in single program/multiple data parallel mode. 
+            *Note" In order to use this capability, Swiftest must be compiled for Coarray support. Only certain integrators
+            can use Coarrays: RMVS, WHM, Helio are all compatible, but SyMBA is not, due to the way tp-pl close encounters 
+            are handeled.           
         tides: bool, optional
             Turns on tidal model (IN DEVELOPMENT - IGNORED)
         Yarkovsky: bool, optional
@@ -1279,7 +1291,14 @@ def set_feature(self,
             self.driver_executable = self.binary_path / "swiftest_driver"
             self.param_file = Path(kwargs.pop("param_file","param.in"))
 
+        if coarray is not None:
+            if self.codename == "Swiftest":
+                self.param["COARRAY"] = coarray
+                update_list.append("coarray")
+
+
         self.param["TIDES"] = False
+
 
         feature_dict = self.get_feature(update_list, verbose)
         return feature_dict
@@ -1321,6 +1340,7 @@ def get_feature(self, arg_list: str | List[str] | None = None, verbose: bool | N
                      "big_discard": "BIG_DISCARD",
                      "interaction_loops": "INTERACTION_LOOPS",
                      "encounter_check_loops": "ENCOUNTER_CHECK",
+                     "coarray" : "COARRAY",
                      "restart": "RESTART"
                      }
 

src/base/base_module.f90

@@ -69,6 +69,7 @@ module base
       character(NAMELEN)                      :: interaction_loops    = "ADAPTIVE"      !! Method used to compute interaction loops. Options are "TRIANGULAR", "FLAT", or "ADAPTIVE" 
       character(NAMELEN)                      :: encounter_check_plpl = "ADAPTIVE"      !! Method used to compute pl-pl encounter checks. Options are "TRIANGULAR", "SORTSWEEP", or "ADAPTIVE" 
       character(NAMELEN)                      :: encounter_check_pltp = "ADAPTIVE"      !! Method used to compute pl-tp encounter checks. Options are "TRIANGULAR", "SORTSWEEP", or "ADAPTIVE" 
+      logical                                 :: lcoarray             = .false.         !! Use Coarrays for test particle parallelization.
 
       ! The following are used internally, and are not set by the user, but instead are determined by the input value of INTERACTION_LOOPS
       logical :: lflatten_interactions     = .false. !! Use the flattened upper triangular matrix for pl-pl interaction loops
@@ -2131,7 +2132,6 @@ subroutine base_coclone_param(self)
          ! Arguments
          class(base_parameters),intent(inout),codimension[*]  :: self  !! Collection of parameters
          ! Internals
-         integer(I4B) :: i
 
          call coclone(self%integrator)
          call coclone(self%param_file_name)
@@ -2192,12 +2192,10 @@ subroutine base_coclone_param(self)
          call coclone(self%ltides        )
          call coclone(self%E_orbit_orig )
          call coclone(self%GMtot_orig  )
-         do i = 1, NDIM
-            call coclone(self%L_total_orig(i))
-            call coclone(self%L_orbit_orig(i))
-            call coclone(self%L_spin_orig(i))
-            call coclone(self%L_escape(i))
-         end do
+         call coclonevec(self%L_total_orig)
+         call coclonevec(self%L_orbit_orig)
+         call coclonevec(self%L_spin_orig)
+         call coclonevec(self%L_escape)
          call coclone(self%GMescape    )
          call coclone(self%E_collisions )
          call coclone(self%E_untracked  )
@@ -2211,6 +2209,7 @@ subroutine base_coclone_param(self)
          call coclone(self%lyarkovsky)
          call coclone(self%lyorp     )
          call coclone(self%seed)
+         call coclone(self%lcoarray)
 
          return
       end subroutine base_coclone_param 

src/swiftest/swiftest_coarray.f90

@@ -455,8 +455,8 @@ module subroutine swiftest_coarray_collect_system(nbody_system, param)
         class(swiftest_tp), allocatable, codimension[:] :: cotp
         character(len=NAMELEN) :: image_num_char
 
-        sync all
-        if (allocated(nbody_system%tp%id)) write(*,*) "Image: ",this_image(), "before collecting ids: ",nbody_system%tp%id
+        if (.not.param%lcoarray) return
+
         sync all
         if (this_image() == 1) then
             write(image_num_char,*) num_images()
@@ -470,8 +470,6 @@ module subroutine swiftest_coarray_collect_system(nbody_system, param)
 
         deallocate(cotp)
 
-        if (this_image() == 1) write(*,*) "Image: ",this_image(), "After collecting ids: ",nbody_system%tp%id
-
         return
     end subroutine swiftest_coarray_collect_system
 
@@ -493,6 +491,7 @@ module subroutine swiftest_coarray_distribute_system(nbody_system, param)
         class(swiftest_tp), allocatable, codimension[:] :: cotp
         class(swiftest_tp), allocatable :: tmp
 
+        if (.not.param%lcoarray) return
         sync all
         if (this_image() == 1) then
             write(image_num_char,*) num_images()
@@ -543,6 +542,8 @@ module subroutine swiftest_coarray_initialize_system(nbody_system, param)
         class(swiftest_nbody_system), allocatable, codimension[:] :: tmp_system
         character(len=NAMELEN) :: image_num_char
 
+        if (.not.param%lcoarray) return
+
         sync all
         if (this_image() == 1) then
             write(image_num_char,*) num_images()

src/swiftest/swiftest_io.f90

@@ -2071,6 +2071,9 @@ module subroutine swiftest_io_param_reader(self, unit, iotype, v_list, iostat, i
                case ("ENCOUNTER_SAVE")
                   call swiftest_io_toupper(param_value)
                   read(param_value, *) param%encounter_save
+               case ("COARRAY")
+                  call swiftest_io_toupper(param_value)
+                  if (param_value == "YES" .or. param_value == 'T') param%lcoarray = .true. 
                case("SEED")
                   read(param_value, *) nseeds_from_file
                   ! Because the number of seeds can vary between compilers/systems, we need to make sure we can handle cases in which the input file has a different
@@ -2317,6 +2320,26 @@ module subroutine swiftest_io_param_reader(self, unit, iotype, v_list, iostat, i
             param%lencounter_sas_pltp = .false.
          end select
 
+
+         if (param%lcoarray) then
+#ifdef COARRAY
+            if (num_images() == 1) then
+               write(iomsg, *) "Only one Coarray image detected. Coarrays will not be used."
+               param%lcoarray = .false.
+            end if
+
+            select case(param%integrator)
+            case(INT_WHM, INT_RMVS, INT_HELIO)
+            case default   
+               write(iomsg, *) "Coarray-based parallelization of test particles are not compatible with this integrator. This parameter will be ignored."
+               param%lcoarray = .false.
+            end select
+#else
+            write(iomsg,*) "Coarray capability not detected. Swiftest must be compiled with Coarrays enabled. to use this feature."
+            param%lcoarray = .false.
+#endif
+         end if
+
          iostat = 0
 
          call param%set_display(param%display_style)