Added utility function for generating smoothing lengths (#822)

docs/source/galaxy/line_of_sight.ipynb

@@ -19,7 +19,6 @@
    "source": [
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
-    "from scipy.spatial import cKDTree\n",
     "from unyt import Mpc, Msun, Myr, yr\n",
     "\n",
     "from synthesizer.kernel_functions import Kernel\n",
@@ -29,22 +28,7 @@
     "from synthesizer.particle.gas import Gas\n",
     "from synthesizer.particle.particles import CoordinateGenerator\n",
     "from synthesizer.particle.stars import sample_sfzh\n",
-    "\n",
-    "\n",
-    "def calculate_smoothing_lengths(positions, num_neighbors=56):\n",
-    "    \"\"\"Calculate the SPH smoothing lengths for a set of coordinates.\"\"\"\n",
-    "    tree = cKDTree(positions)\n",
-    "    distances, _ = tree.query(positions, k=num_neighbors + 1)\n",
-    "\n",
-    "    # The k-th nearest neighbor distance (k = num_neighbors)\n",
-    "    kth_distances = distances[:, num_neighbors]\n",
-    "\n",
-    "    # Set the smoothing length to the k-th nearest neighbor\n",
-    "    # distance divided by 2.0\n",
-    "    smoothing_lengths = kth_distances / 2.0\n",
-    "\n",
-    "    return smoothing_lengths\n",
-    "\n",
+    "from synthesizer.particle.utils import calculate_smoothing_lengths\n",
     "\n",
     "# Define the grid (normally this would be defined by an SPS grid)\n",
     "log10ages = np.arange(6.0, 10.5, 0.1)\n",
@@ -73,7 +57,7 @@
     "coords = CoordinateGenerator.generate_3D_gaussian(nstars) * Mpc\n",
     "\n",
     "# Calculate the smoothing lengths\n",
-    "smls = calculate_smoothing_lengths(coords) * Mpc\n",
+    "smls = calculate_smoothing_lengths(coords)\n",
     "\n",
     "# Sample the parametric SFZH, producing a particle Stars object\n",
     "# we will also pass some keyword arguments for some example attributes\n",
@@ -107,7 +91,7 @@
     "coords = CoordinateGenerator.generate_3D_gaussian(ngas) * Mpc\n",
     "\n",
     "# Calculate the smoothing lengths\n",
-    "smls = calculate_smoothing_lengths(coords) * Mpc\n",
+    "smls = calculate_smoothing_lengths(coords)\n",
     "\n",
     "gas = Gas(\n",
     "    masses=np.random.uniform(10**6, 10**6.5, ngas) * Msun,\n",

examples/particle/plot_get_los_attenutation.py

@@ -10,7 +10,6 @@
 
 import matplotlib.pyplot as plt
 import numpy as np
-from scipy.spatial import cKDTree
 from unyt import Mpc, Msun, Myr
 
 from synthesizer.grid import Grid
@@ -21,22 +20,7 @@
 from synthesizer.particle.gas import Gas
 from synthesizer.particle.particles import CoordinateGenerator
 from synthesizer.particle.stars import sample_sfzh
-
-
-def calculate_smoothing_lengths(positions, num_neighbors=56):
-    """Calculate the SPH smoothing lengths for a set of coordinates."""
-    tree = cKDTree(positions)
-    distances, _ = tree.query(positions, k=num_neighbors + 1)
-
-    # The k-th nearest neighbor distance (k = num_neighbors)
-    kth_distances = distances[:, num_neighbors]
-
-    # Set the smoothing length to the k-th nearest neighbor
-    # distance divided by 2.0
-    smoothing_lengths = kth_distances / 2.0
-
-    return smoothing_lengths
-
+from synthesizer.particle.utils import calculate_smoothing_lengths
 
 plt.rcParams["font.family"] = "DeJavu Serif"
 plt.rcParams["font.serif"] = ["Times New Roman"]
@@ -92,7 +76,7 @@ def calculate_smoothing_lengths(positions, num_neighbors=56):
 )
 
 # Calculate the smoothing lengths
-smls = calculate_smoothing_lengths(coords, num_neighbors=56) * Mpc
+smls = calculate_smoothing_lengths(coords, num_neighbours=56)
 
 # Sample the SFZH, producing a Stars object
 # we will also pass some keyword arguments for attributes
@@ -122,7 +106,7 @@ def calculate_smoothing_lengths(positions, num_neighbors=56):
 )
 
 # Calculate the smoothing lengths
-smls = calculate_smoothing_lengths(coords, num_neighbors=56) * Mpc
+smls = calculate_smoothing_lengths(coords, num_neighbours=56)
 
 gas = Gas(
     masses=np.random.uniform(10**6, 10**6.5, ngas) * Msun,

examples/particle/plot_los_spectra.py

@@ -10,7 +10,6 @@
 
 import matplotlib.pyplot as plt
 import numpy as np
-from scipy.spatial import cKDTree
 from unyt import Mpc, Msun, Myr
 
 from synthesizer.emission_models import TotalEmission
@@ -23,26 +22,12 @@
 from synthesizer.particle.gas import Gas
 from synthesizer.particle.particles import CoordinateGenerator
 from synthesizer.particle.stars import sample_sfzh
+from synthesizer.particle.utils import calculate_smoothing_lengths
 
 plt.rcParams["font.family"] = "DeJavu Serif"
 plt.rcParams["font.serif"] = ["Times New Roman"]
 
 
-def calculate_smoothing_lengths(positions, num_neighbors=56):
-    """Calculate the SPH smoothing lengths for a set of coordinates."""
-    tree = cKDTree(positions)
-    distances, _ = tree.query(positions, k=num_neighbors + 1)
-
-    # The k-th nearest neighbor distance (k = num_neighbors)
-    kth_distances = distances[:, num_neighbors]
-
-    # Set the smoothing length to the k-th nearest neighbor
-    # distance divided by 2.0
-    smoothing_lengths = kth_distances / 2.0
-
-    return smoothing_lengths
-
-
 # Set the seed
 np.random.seed(42)
 
@@ -89,7 +74,7 @@ def calculate_smoothing_lengths(positions, num_neighbors=56):
 ngas = 1000
 
 # Generate some random coordinates
-coords = CoordinateGenerator.generate_3D_gaussian(nstars)
+coords = CoordinateGenerator.generate_3D_gaussian(nstars) * Mpc
 
 # Calculate smoothing lengths
 smls = calculate_smoothing_lengths(coords)
@@ -102,25 +87,25 @@ def calculate_smoothing_lengths(positions, num_neighbors=56):
     param_stars.log10ages,
     param_stars.log10metallicities,
     nstars,
-    coordinates=coords * Mpc,
+    coordinates=coords,
     current_masses=np.full(nstars, 10**8.7 / nstars) * Msun,
-    smoothing_lengths=smls * Mpc,
+    smoothing_lengths=smls,
     redshift=1,
 )
 
 # Now make the gas
 
 # Generate some random coordinates
-coords = CoordinateGenerator.generate_3D_gaussian(ngas)
+coords = CoordinateGenerator.generate_3D_gaussian(ngas) * Mpc
 
 # Calculate the smoothing lengths
 smls = calculate_smoothing_lengths(coords)
 
 gas = Gas(
     masses=np.random.uniform(10**6, 10**6.5, ngas) * Msun,
     metallicities=np.random.uniform(0.01, 0.05, ngas),
-    coordinates=coords * Mpc,
-    smoothing_lengths=smls * Mpc,
+    coordinates=coords,
+    smoothing_lengths=smls,
     dust_to_metal_ratio=0.2,
 )
 

examples/particle/plot_rotated_galaxy.py

@@ -15,27 +15,11 @@
 
 import matplotlib.pyplot as plt
 import numpy as np
-from scipy.spatial import cKDTree
 from unyt import Msun, Myr, degree, km, kpc, s
 
 from synthesizer.kernel_functions import Kernel
 from synthesizer.particle import CoordinateGenerator, Galaxy, Gas, Stars
-
-
-def calculate_smoothing_lengths(positions, num_neighbors=56):
-    """Calculate the SPH smoothing lengths for a set of coordinates."""
-    tree = cKDTree(positions)
-    distances, _ = tree.query(positions, k=num_neighbors + 1)
-
-    # The k-th nearest neighbor distance (k = num_neighbors)
-    kth_distances = distances[:, num_neighbors]
-
-    # Set the smoothing length to the k-th nearest neighbor
-    # distance divided by 2.0
-    smoothing_lengths = kth_distances / 2.0
-
-    return smoothing_lengths
-
+from synthesizer.particle.utils import calculate_smoothing_lengths
 
 # Set the seed
 np.random.seed(42)
@@ -79,7 +63,7 @@ def calculate_smoothing_lengths(positions, num_neighbors=56):
 initial_masses = masses.copy()
 redshift = 0.0
 centre = np.array([0.0, 0.0, 0.0]) * kpc
-smoothing_lengths = calculate_smoothing_lengths(coords) * kpc
+smoothing_lengths = calculate_smoothing_lengths(coords)
 
 # We'll start by simply using some stars
 stars = Stars(
@@ -151,7 +135,7 @@ def calculate_smoothing_lengths(positions, num_neighbors=56):
     redshift=redshift,
     centre=centre,
     dust_to_metal_ratio=0.3,
-    smoothing_lengths=calculate_smoothing_lengths(gas_coords) * kpc,
+    smoothing_lengths=calculate_smoothing_lengths(gas_coords),
 )
 
 # Make the galaxy

profiling/compare_los_loop_tree.py

@@ -13,7 +13,6 @@
 import matplotlib.pyplot as plt
 import numpy as np
 from matplotlib.lines import Line2D
-from scipy.spatial import cKDTree
 from unyt import Myr
 
 from synthesizer.grid import Grid
@@ -24,22 +23,7 @@
 from synthesizer.particle.gas import Gas
 from synthesizer.particle.particles import CoordinateGenerator
 from synthesizer.particle.stars import sample_sfzh
-
-
-def calculate_smoothing_lengths(positions, num_neighbors=56):
-    """Calculate the SPH smoothing lengths for a set of coordinates."""
-    tree = cKDTree(positions)
-    distances, _ = tree.query(positions, k=num_neighbors + 1)
-
-    # The k-th nearest neighbor distance (k = num_neighbors)
-    kth_distances = distances[:, num_neighbors]
-
-    # Set the smoothing length to the k-th nearest neighbor
-    # distance divided by 2.0
-    smoothing_lengths = kth_distances / 2.0
-
-    return smoothing_lengths
-
+from synthesizer.particle.utils import calculate_smoothing_lengths
 
 plt.rcParams["font.family"] = "DeJavu Serif"
 plt.rcParams["font.serif"] = ["Times New Roman"]

profiling/image_strong_scaling.py

@@ -15,7 +15,6 @@
 import matplotlib.pyplot as plt
 import numpy as np
 from astropy.cosmology import Planck15 as cosmo
-from scipy.spatial import cKDTree
 from unyt import Myr, kpc
 
 from synthesizer.emission_models import IncidentEmission
@@ -27,6 +26,7 @@
 from synthesizer.particle.galaxy import Galaxy
 from synthesizer.particle.particles import CoordinateGenerator
 from synthesizer.particle.stars import sample_sfzh
+from synthesizer.particle.utils import calculate_smoothing_lengths
 
 plt.rcParams["font.family"] = "DeJavu Serif"
 plt.rcParams["font.serif"] = ["Times New Roman"]
@@ -35,21 +35,6 @@
 np.random.seed(42)
 
 
-def calculate_smoothing_lengths(positions, num_neighbors=56):
-    """Calculate the SPH smoothing lengths for a set of coordinates."""
-    tree = cKDTree(positions)
-    distances, _ = tree.query(positions, k=num_neighbors + 1)
-
-    # The k-th nearest neighbor distance (k = num_neighbors)
-    kth_distances = distances[:, num_neighbors]
-
-    # Set the smoothing length to the k-th nearest neighbor
-    # distance divided by 2.0
-    smoothing_lengths = kth_distances / 2.0
-
-    return smoothing_lengths
-
-
 def image_strong_scaling(
     basename,
     max_threads=8,

profiling/los_surf_strong_scaling.py

@@ -14,7 +14,6 @@
 import matplotlib.gridspec as gridspec
 import matplotlib.pyplot as plt
 import numpy as np
-from scipy.spatial import cKDTree
 from unyt import Myr
 
 from synthesizer.kernel_functions import Kernel
@@ -24,6 +23,7 @@
 from synthesizer.particle.gas import Gas
 from synthesizer.particle.particles import CoordinateGenerator
 from synthesizer.particle.stars import sample_sfzh
+from synthesizer.particle.utils import calculate_smoothing_lengths
 
 plt.rcParams["font.family"] = "DeJavu Serif"
 plt.rcParams["font.serif"] = ["Times New Roman"]
@@ -32,21 +32,6 @@
 np.random.seed(42)
 
 
-def calculate_smoothing_lengths(positions, num_neighbors=56):
-    """Calculate the SPH smoothing lengths for a set of coordinates."""
-    tree = cKDTree(positions)
-    distances, _ = tree.query(positions, k=num_neighbors + 1)
-
-    # The k-th nearest neighbor distance (k = num_neighbors)
-    kth_distances = distances[:, num_neighbors]
-
-    # Set the smoothing length to the k-th nearest neighbor
-    # distance divided by 2.0
-    smoothing_lengths = kth_distances / 2.0
-
-    return smoothing_lengths
-
-
 def los_surface_density_strong_scaling(
     basename,
     max_threads=8,

profiling/spectral_cube_strong_scaling.py

@@ -15,7 +15,6 @@
 import matplotlib.pyplot as plt
 import numpy as np
 from astropy.cosmology import Planck15 as cosmo
-from scipy.spatial import cKDTree
 from unyt import Myr, kpc
 
 from synthesizer.emission_models import IncidentEmission
@@ -27,6 +26,7 @@
 from synthesizer.particle.galaxy import Galaxy
 from synthesizer.particle.particles import CoordinateGenerator
 from synthesizer.particle.stars import sample_sfzh
+from synthesizer.particle.utils import calculate_smoothing_lengths
 
 plt.rcParams["font.family"] = "DeJavu Serif"
 plt.rcParams["font.serif"] = ["Times New Roman"]
@@ -35,21 +35,6 @@
 np.random.seed(42)
 
 
-def calculate_smoothing_lengths(positions, num_neighbors=56):
-    """Calculate the SPH smoothing lengths for a set of coordinates."""
-    tree = cKDTree(positions)
-    distances, _ = tree.query(positions, k=num_neighbors + 1)
-
-    # The k-th nearest neighbor distance (k = num_neighbors)
-    kth_distances = distances[:, num_neighbors]
-
-    # Set the smoothing length to the k-th nearest neighbor
-    # distance divided by 2.0
-    smoothing_lengths = kth_distances / 2.0
-
-    return smoothing_lengths
-
-
 def cube_strong_scaling(
     basename,
     max_threads=8,