# HG changeset patch # User Stephen Skory # Date 1247182559 25200 # Branch hopwrap # Node ID 64d88cef90c2cc542ac270c4375537773476072c # Parent 5b8ba3c49a70c83816ba5c530f31238bef2b845f Changed chainHOP to be more memory-efficient, which required adding a fortran subroutine. diff -r 5b8ba3c49a70c83816ba5c530f31238bef2b845f -r 64d88cef90c2cc542ac270c4375537773476072c yt/extensions/kdtree/Makefile --- a/yt/extensions/kdtree/Makefile Thu Jul 09 13:32:30 2009 -0700 +++ b/yt/extensions/kdtree/Makefile Thu Jul 09 16:35:59 2009 -0700 @@ -5,7 +5,7 @@ fKD: fKD.f90 fKD.v fKD_source.f90 # Forthon --compile_first fKD_source --no2underscores --with-numpy -g fKD fKD.f90 fKD_source.f90 - Forthon --compile_first fKD_source --no2underscores --with-numpy --fopt "-O3" fKD fKD_source.f90 + Forthon -F gfortran --compile_first fKD_source --no2underscores --with-numpy --fopt "-O3" fKD fKD_source.f90 clean: rm -rf build fKDpy.a fKDpy.so diff -r 5b8ba3c49a70c83816ba5c530f31238bef2b845f -r 64d88cef90c2cc542ac270c4375537773476072c yt/extensions/kdtree/fKD.f90 --- a/yt/extensions/kdtree/fKD.f90 Thu Jul 09 13:32:30 2009 -0700 +++ b/yt/extensions/kdtree/fKD.f90 Thu Jul 09 16:35:59 2009 -0700 @@ -68,9 +68,32 @@ end subroutine find_all_nn_nearest_neighbors +subroutine find_chunk_nearest_neighbors() + ! for a chunk of the full number of particles, find their nearest neighbors + use kdtree2_module + use fKD_module + use kdtree2module + use tree_nodemodule + use intervalmodule + + integer :: k + type(kdtree2_result),allocatable :: results(:) ! nearest neighbors + allocate(results(nn)) + do k=start,finish + qv(:) = pos(:,k) + call kdtree2_n_nearest(tp=tree2,qv=qv,nn=nn,results=results) + chunk_tags(:,k - start + 1) = results%idx + + end do + + deallocate(results) + return + +end subroutine find_chunk_nearest_neighbors + subroutine chainHOP_tags_dens() ! for all particles in pos, find their nearest neighbors, and calculate - ! their density. + ! their density. Return only nMerge nearest neighbors. use kdtree2_module use fKD_module use kdtree2module @@ -79,21 +102,25 @@ integer :: k, pj, i real :: ih2, fNorm, r2, rs + integer, allocatable :: temp_tags(:) + real, allocatable :: temp_dist(:) type(kdtree2_result),allocatable :: results(:) ! nearest neighbors allocate(results(nn)) + allocate(temp_tags(nn)) + allocate(temp_dist(nn)) do k=1,nparts qv(:) = pos(:,k) call kdtree2_n_nearest(tp=tree2,qv=qv,nn=nn,results=results) - nn_tags(:,k) = results%idx - nn_dist(:,k) = results%dis + temp_tags(:) = results%idx + temp_dist(:) = results%dis ! calculate the density for this particle - ih2 = 4.0/maxval(nn_dist(:,k)) + ih2 = 4.0/maxval(results%dis) fNorm = 0.5*sqrt(ih2)*ih2/3.1415926535897931 do i=1,nn - pj = nn_tags(i,k) - r2 = nn_dist(i,k)*ih2 + pj = temp_tags(i) + r2 = temp_dist(i) * ih2 rs = 2.0 - sqrt(r2) if (r2 < 1.0) then rs = (1.0 - 0.75*rs*r2) @@ -104,10 +131,15 @@ dens(k) = dens(k) + rs * mass(pj) dens(pj) = dens(pj) + rs * mass(k) end do + + ! record only nMerge nearest neighbors + nn_tags(:,k) = temp_tags(1:nMerge) end do deallocate(results) + deallocate(temp_dist) + deallocate(temp_tags) return end subroutine chainHOP_tags_dens diff -r 5b8ba3c49a70c83816ba5c530f31238bef2b845f -r 64d88cef90c2cc542ac270c4375537773476072c yt/extensions/kdtree/fKD.v --- a/yt/extensions/kdtree/fKD.v Thu Jul 09 13:32:30 2009 -0700 +++ b/yt/extensions/kdtree/fKD.v Thu Jul 09 16:35:59 2009 -0700 @@ -4,6 +4,7 @@ tags(:) _integer # particle ID tags dist(:) _real # interparticle spacings nn_tags(:,:) _integer # for all particles at once, [nth neighbor, index] +chunk_tags(:,:) _integer # for finding only a chunk of the nearest neighbors nn_dist(:,:) _real pos(3,:) _real dens(:) _real @@ -11,6 +12,9 @@ qv(3) real nparts integer nn integer +nMerge integer # number of nearest neighbors used in chain merging +start integer +finish integer tree2 _kdtree2 sort logical /.false./ rearrange logical /.true./ @@ -92,4 +96,5 @@ create_tree() subroutine free_tree() subroutine find_all_nn_nearest_neighbors subroutine +find_chunk_nearest_neighbors subroutine chainHOP_tags_dens subroutine diff -r 5b8ba3c49a70c83816ba5c530f31238bef2b845f -r 64d88cef90c2cc542ac270c4375537773476072c yt/lagos/chainHOP/chainHOP.py --- a/yt/lagos/chainHOP/chainHOP.py Thu Jul 09 13:32:30 2009 -0700 +++ b/yt/lagos/chainHOP/chainHOP.py Thu Jul 09 16:35:59 2009 -0700 @@ -18,6 +18,7 @@ self.mass = mass self.padded_particles = [] self.size = len(self.xpos) + self.nMerge = 4 self.chainID = na.ones(self.size,dtype='l') * -1 self._chain_hop() self.padded_particles = na.array(self.padded_particles) @@ -26,22 +27,22 @@ def _init_kd_tree(self): # Yes, we really do need to initialize this many arrays. # They're deleted in _chainHOP. - fKD.nn_tags = empty((self.num_neighbors,self.size), dtype='l') - fKD.nn_dist = empty((self.num_neighbors,self.size), dtype='d') + fKD.nn_tags = empty((self.nMerge + 2,self.size), dtype='l') fKD.dens = zeros(self.size, dtype='d') fKD.mass = empty(self.size, dtype='d') fKD.pos = empty((3,self.size), dtype='d') fKD.qv = empty(3, dtype='d') fKD.nn = self.num_neighbors + fKD.nMerge = self.nMerge + 2 fKD.nparts = self.size - fKD.sort = True # slower, but needed in _connect_chains - fKD.rearrange = True # faster, more memory - # this actually copies the data into the fortran space + fKD.sort = True # Slower, but needed in _connect_chains + fKD.rearrange = True # Faster, but uses more memory + # This actually copies the data into the fortran space fKD.pos[0, :] = self.xpos fKD.pos[1, :] = self.ypos fKD.pos[2, :] = self.zpos fKD.mass = self.mass - # now call the fortran + # Now call the fortran create_tree() def _is_inside(self): @@ -52,13 +53,32 @@ (points < RE).all(axis=1) ) def _densestNN(self): - # Find the densest nearest neighbor for all particles. - n_dens = na.take(self.density,self.NNtags) - max_loc = na.argmax(n_dens,axis=1) + """ + The first search of nearest neighbors did not return all + num_neighbor neighbors, so we need to do it again, but we're not + keeping the all of this data, just using it. + """ self.densestNN = na.ones(self.size,dtype='l') - # How do I do this better? - for i in xrange(self.size): - self.densestNN[i] = self.NNtags[i,max_loc[i]] + # We find nearest neighbors in chunks + chunksize = 10000 + + start = 1 # fortran counting! + finish = 0 + while finish < self.size: + finish = min(finish+chunksize,self.size) + fKD.chunk_tags = empty((self.num_neighbors,finish-start+1), dtype='l') + # call the fortran + fKD.start = start + fKD.finish = finish + find_chunk_nearest_neighbors() + chunk_NNtags = (fKD.chunk_tags - 1).transpose() + # find the densest nearest neighbors + n_dens = na.take(self.density,chunk_NNtags) + max_loc = na.argmax(n_dens,axis=1) + for i in xrange(finish - start + 1): # +1 for fortran counting + j = start + i - 1 # -1 for fortran counting + self.densestNN[j] = chunk_NNtags[i,max_loc[i]] + start = finish + 1 def _build_chains(self): """ @@ -257,7 +277,6 @@ self.densest_in_group[groupID] = max_dens def _chain_hop(self): - self.nMerge = 4 mylog.info('Building kd tree for %d particles...' % \ self.size) self._init_kd_tree() @@ -266,18 +285,19 @@ # Loop over the particles to find NN for each. mylog.info('Finding nearest neighbors/density...') chainHOP_tags_dens() - mylog.info('Copying results...') self.density = fKD.dens self.NNtags = (fKD.nn_tags - 1).transpose() - # When done with the tree free the memory, del these first. - del fKD.pos, fKD.dens, fKD.nn_dist, fKD.nn_tags, fKD.mass, fKD.dens - free_tree() # Frees the kdtree object. + # We can free these right now, the rest later. + del fKD.dens, fKD.nn_tags, fKD.mass, fKD.dens count = len(na.where(self.density >= self.threshold)[0]) print 'count above thresh', count # Now each particle has NNtags, and a local self density. # Let's find densest NN mylog.info('Finding densest nearest neighbors...') self._densestNN() + # Now we can free these + del fKD.pos, fKD.chunk_tags + free_tree() # Frees the kdtree object. chain_count = 0 for i in xrange(self.size): if i == self.densestNN[i]: chain_count += 1