Merge pull request #1199 from ErwanH29/nearest_neighbour_fix

Resolving issue #978: To ensure particles do not consider themselves when calculating nearest neighbours

Author: elkogerville

examples/textbook/relax_gas_and_stars.py

@@ -1,11 +1,20 @@
 import numpy
 import pickle
-from amuse.lab import *
+
+from amuse.lab import (
+    units, nbody_system, 
+    new_plummer_model, write_set_to_file,
+    new_salpeter_mass_distribution, 
+)
 from amuse.community.fastkick.interface import FastKick
-from amuse.ext.relax_sph import relax
+from amuse.community.fi.interface import Fi
+from amuse.ext.relax_sph import relax, monitor_energy
 from amuse.ext.spherical_model import new_gas_plummer_distribution
 from amuse.community.fractalcluster.interface import new_fractal_cluster_model
 
+
+from prepare_figure import *
+
 ###BOOKLISTSTART1###
 def check_energy_conservation(system, i_step, time, n_steps):
     unit = units.J
@@ -163,7 +172,6 @@ def make_map(sph,N=100,L=1):
 def plot_hydro_and_stars(hydro, stars):
     x_label = "x [pc]"
     y_label = "y [pc]"
-    from prepare_figure import *
     fig = single_frame(x_label, y_label, logx=False, logy=False,
                        xsize=12, ysize=12)
 

src/amuse/datamodel/particle_attributes.py

@@ -923,14 +923,15 @@ def distances_squared(particles, other_particles):
     return (dxdydz**2).sum(-1)
 
 
-def nearest_neighbour(particles, neighbours=None, max_array_length=10000000):
+def nearest_neighbour(particles, neighbours=None, self_search=False, max_array_length=10000000):
     """
     Returns the nearest neighbour of each particle in this set. If the 'neighbours'
     particle set is supplied, the search is performed on the neighbours set, for
     each particle in the orignal set. Otherwise the nearest neighbour in the same
     set is searched.
 
     :argument neighbours: the particle set in which to search for the nearest neighbour (optional)
+    :argument self_search: if True, the nearest neighbour can be the particle itself (default False)
 
     >>> from amuse.datamodel import Particles
     >>> particles = Particles(3)
@@ -964,7 +965,7 @@ def nearest_neighbour(particles, neighbours=None, max_array_length=10000000):
         )
         for indices in indices_in_each_batch:
             distances_squared = particles[indices].distances_squared(other_particles)
-            if neighbours is None:
+            if not self_search and neighbours is None:
                 diagonal_indices = (numpy.arange(len(indices)), indices)
                 distances_squared.number[
                     diagonal_indices
@@ -973,11 +974,11 @@ def nearest_neighbour(particles, neighbours=None, max_array_length=10000000):
         return other_particles[numpy.concatenate(neighbour_indices)]
 
     distances_squared = particles.distances_squared(other_particles)
-    if neighbours is None:
+    if not self_search and neighbours is None:  # can't be your own neighbour
         diagonal_indices = numpy.diag_indices(len(particles))
         distances_squared.number[
             diagonal_indices
-        ] = numpy.inf  # can't be your own neighbour
+        ] = numpy.inf  
     return other_particles[distances_squared.argmin(axis=1)]
 
 

src/tests/core_tests/test_particle_attributes.py

@@ -173,6 +173,7 @@ def test11(self):
         particles.z = 0.0 | units.m
         self.assertEqual(particles.nearest_neighbour()[0], particles[1])
         self.assertEqual(particles.nearest_neighbour()[1:].key, particles[:-1].key)
+        self.assertEqual(particles.nearest_neighbour(self_search=True).key, particles.key)
 
         neighbours = Particles(3)
         neighbours.x = [1.0, 10.0, 100.0] | units.m