Merge pull request #786 from galacticusorg/fixTreeReadBeginAt

Reimplement the `beginAt` functionality in the `mergerTreeConstructorRead` class
galacticusorg · Feb 10, 2025 · a3161eb · a3161eb
2 parents 6c5e27c + f2c7617
commit a3161eb
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 14 deletions.
diff --git a/source/merger_trees.construct.read.F90 b/source/merger_trees.construct.read.F90
@@ -235,6 +235,7 @@
      integer                                                     , allocatable, dimension(:) :: indexNamedReals                                 , indexNamedIntegers
      logical                                                                                 :: importerOpen                          =  .false.
      integer         (kind_int8                                 )                            :: beginAt
+     logical                                                                                 :: foundBeginAt                          =  .false.
      double precision                                                                        :: treeWeightCurrent
      logical                                                                                 :: allowBranchJumps
      logical                                                                                 :: allowSubhaloPromotions                          , alwaysPromoteMostMassive
@@ -744,6 +745,8 @@ function readConstructorInternal(fileNames,outputTimeSnapTolerance,forestSizeMax
     ! Initialize statuses.
     self%warningNestedHierarchyIssued           =.false.
     self%warningSplitForestNestedHierarchyIssued=.false.
+    ! Set initial state indicating if the first tree to process has been found.
+    self%foundBeginAt                           =self%beginAt == -1_kind_int8
     ! Initialize split forests counter.
     splitForestUniqueID=mpiCounter()
     ! Get array of output times.
@@ -995,18 +998,34 @@ function readConstruct(self,treeNumber,finished) result(tree)
     treeNumberInternal=treeStateStoreSequence
     ! Determine if we have any split forests to return.
     returnSplitForest=allocated(self%splitForestTreeSize)
-    ! Find the maximum tree number in the current file.
-    treeNumberMaximum=int(self%mergerTreeImporter_%treeCount(),kind=c_size_t)
-    ! Check if we need to move to a new file.
-    do while (treeNumber-self%treeNumberOffset > treeNumberMaximum .and. self%fileCurrent < size(self%fileNames))
-       self%fileCurrent     =self%fileCurrent     +1
-       self%treeNumberOffset=self%treeNumberOffset+treeNumberMaximum
-       call self%mergerTreeImporter_%close(                                                        )
-       call self%mergerTreeImporter_%open (File_Name_Expand(char(self%fileNames(self%fileCurrent))))
+    ! Scan trees until we find one to process. This allows us to skip trees until the tree index specified by `beginAt` is found.
+    do while (.true.)
+       ! Find the maximum tree number in the current file.
        treeNumberMaximum=int(self%mergerTreeImporter_%treeCount(),kind=c_size_t)
-    end do
-    treeNumberOffset=treeNumber-self%treeNumberOffset
-    if (treeNumberOffset <= treeNumberMaximum) then
+       ! Check if we need to move to a new file.
+       do while (treeNumber-self%treeNumberOffset > treeNumberMaximum .and. self%fileCurrent < size(self%fileNames))
+          self%fileCurrent     =self%fileCurrent     +1
+          self%treeNumberOffset=self%treeNumberOffset+treeNumberMaximum
+          call self%mergerTreeImporter_%close(                                                        )
+          call self%mergerTreeImporter_%open (File_Name_Expand(char(self%fileNames(self%fileCurrent))))
+          treeNumberMaximum=int(self%mergerTreeImporter_%treeCount(),kind=c_size_t)
+       end do
+       treeNumberOffset=treeNumber-self%treeNumberOffset
+       ! If all trees are used up, we're done.
+       if (treeNumberOffset > treeNumberMaximum) then
+          nullify(tree)
+          finished=.true.
+          exit
+       end if
+       ! Test if this is the first forest that we are to process, or if we have already found that forest.
+       if (.not.self%foundBeginAt .and. self%mergerTreeImporter_%treeIndex (int(treeNumberOffset)) /= self%beginAt) then
+          ! Decrement our internal offset, and try again - this will move us to trying the next forest in the file.
+          self%treeNumberOffset=self%treeNumberOffset-1_c_size_t
+          cycle
+       else
+          ! The first forest to process is found - record this.
+          self%foundBeginAt=.true.
+       end if
        ! Set tree properties.
        allocate(tree)
        ! treeIndex
@@ -1392,9 +1411,9 @@ function readConstruct(self,treeNumber,finished) result(tree)
        end select
        ! Deallocate nodes.
        deallocate(nodes)
-    else
-       nullify(tree)
-    end if
+       ! Indicate that we are not finished.
+       exit
+    end do
     finished=.not.associated(tree)
     return
   end function readConstruct

diff --git a/testSuite/test-methods.xml b/testSuite/test-methods.xml
@@ -238,6 +238,8 @@
       </componentHotHalo>
       <presetOrbits value="true"/>
       <presetOrbitsBoundOnly value="true"/>
+      <beginAt value="-1"/>
+      <beginAt value="7000124000000"/>
       <componentSatellite value="standard" parameterLevel="top">
     	<presetMergerTimes value="false" parameterLevel="-1"/>
     	<presetMergerNodes value="false" parameterLevel="-1"/>