modified surface_density function to compute surface densities of gas…

sarracen/disc/surface_density.py

@@ -1,5 +1,6 @@
 import numpy as np
 import pandas as pd
+import warnings
 from ..sarracen_dataframe import SarracenDataFrame
 from .utils import _get_mass, _get_origin
 from .utils import _bin_particles_by_radius, _get_bin_midpoints
@@ -83,6 +84,13 @@ def surface_density(data: 'SarracenDataFrame',
                     origin: Union[list, None] = None,
                     retbins: bool = False) -> Union[np.ndarray,
                                                     Tuple[np.ndarray,
+                                                          np.ndarray],
+                                                    Tuple[np.ndarray,
+                                                          np.ndarray,
+                                                          np.ndarray],
+                                                    Tuple[np.ndarray,
+                                                          np.ndarray,
+                                                          np.ndarray,
                                                           np.ndarray]]:
     """
     Calculates the 1D azimuthally-averaged surface density profile.
@@ -139,24 +147,76 @@ def surface_density(data: 'SarracenDataFrame',
 
     """
 
+    if 'itype' in data.columns and data['itype'].nunique() > 1:
+        msg = ('Surface density being calculated from multiple particle types.'
+               ' Ensure that this is what you intend.')
+        warnings.warn(msg, UserWarning, stacklevel=2)
+
     origin = _get_origin(origin)
     rbins, bin_edges = _bin_particles_by_radius(data, r_in, r_out, bins, log,
                                                 geometry, origin)
 
     areas = np.pi * (bin_edges[1:] ** 2 - bin_edges[:-1] ** 2)
 
-    mass = _get_mass(data)
-    if isinstance(mass, pd.Series):
-        sigma = mass.groupby(rbins, observed=False).sum()
-    else:
-        sigma = data.groupby(rbins, observed=False).size() * mass
+    ndustsmall = data.params.get('ndustsmall') or len(data.dustfracscol)
 
-    sigma = sigma.reindex(rbins.cat.categories, fill_value=0)
-
-    if retbins:
-        return (sigma / areas).to_numpy(), _get_bin_midpoints(bin_edges, log)
-    else:
-        return (sigma / areas).to_numpy()
+    mass = _get_mass(data)
+    if ndustsmall == 0:  # gas-only dump
+        if isinstance(mass, pd.Series):
+            sigma = mass.groupby(rbins, observed=False).sum()
+        else:
+            counts = data.groupby(rbins, observed=False).size().astype(float)
+            sigma = counts * mass
+        sigma = sigma.reindex(rbins.cat.categories, fill_value=0)
+        sigma = (sigma / areas).to_numpy()
+
+        if retbins:
+            return sigma, _get_bin_midpoints(bin_edges, log)
+        else:
+            return sigma
+
+    elif ndustsmall == 1:  # 'dust-as-mixture' dump
+        if data.dustfracscol[0] not in data.columns:
+            raise KeyError('dustfrac not found in data')
+
+        mass_gas = mass * (1. - data[data.dustfracscol[0]])
+        sigma_gas = mass_gas.groupby(rbins, observed=False).sum()
+        sigma_gas = sigma_gas.reindex(rbins.cat.categories, fill_value=0)
+        sigma_gas = (sigma_gas / areas).to_numpy()
+
+        mass_dust = mass * data[data.dustfracscol[0]]
+        sigma_dust = mass_dust.groupby(rbins, observed=False).sum()
+        sigma_dust = sigma_dust.reindex(rbins.cat.categories, fill_value=0)
+        sigma_dust = (sigma_dust / areas).to_numpy()
+
+        if retbins:
+            return sigma_gas, sigma_dust, _get_bin_midpoints(bin_edges, log)
+        else:
+            return sigma_gas, sigma_dust
+
+    elif ndustsmall > 1:  # multiple grain sizes
+        if 'dustfrac_total' not in data.columns:
+            data.calc_one_fluid_quantities()
+
+        mass_gas = mass * (1. - data['dustfrac_total'])
+        sigma_gas = mass_gas.groupby(rbins, observed=False).sum()
+        sigma_gas = sigma_gas.reindex(rbins.cat.categories, fill_value=0)
+        sigma_gas = (sigma_gas / areas).to_numpy()
+
+        mass_dtot = mass * data['dustfrac_total']
+        sigma_dtot = mass_dtot.groupby(rbins, observed=False).sum()
+        sigma_dtot = sigma_dtot.reindex(rbins.cat.categories, fill_value=0)
+        sigma_dtot = (sigma_dtot / areas).to_numpy()
+
+        sigma_d = np.empty([ndustsmall, bins])
+        for i in range(ndustsmall):
+            mass_d = mass * data[data.dustfracscol[i]]
+            sigma_d[i] = mass_d.groupby(rbins, observed=False).sum() / areas
+        if retbins:
+            return sigma_gas, sigma_dtot, sigma_d, \
+                    _get_bin_midpoints(bin_edges, log)
+        else:
+            return sigma_gas, sigma_dtot, sigma_d
 
 
 def _calc_angular_momentum(data: 'SarracenDataFrame',

sarracen/tests/disc/test_surface_density.py

@@ -79,3 +79,107 @@ def test_parts_vs_whole() -> None:
 
     assert_allclose(sigma_in, sigma_all[:100], atol=1e-15, rtol=0.0)
     assert_allclose(sigma_out, sigma_all[100:], atol=1e-15, rtol=0.0)
+
+
+@pytest.mark.filterwarnings('ignore:Surface density being*:UserWarning')
+def test_itype_combining() -> None:
+    """ Adding profiles from different itypes should equal the total. """
+
+    # randomly place particles
+    rng = np.random.default_rng(seed=5)
+    x = rng.random(200)
+    y = rng.random(200)
+    z = rng.random(200)
+    vx = rng.random(200)
+    vy = rng.random(200)
+    vz = rng.random(200)
+    mass = [3.2e-4] * 100 + [5.22e-6] * 100
+    itype = [1] * 100 + [7] * 100
+
+    sdf = SarracenDataFrame(data={'x': x, 'y': y, 'z': z,
+                                  'vx': vx, 'vy': vy, 'vz': vz,
+                                  'mass': mass, 'itype': itype})
+    sigma_1 = surface_density(sdf[sdf.itype == 1], r_in=0, r_out=1, bins=30)
+    sigma_7 = surface_density(sdf[sdf.itype == 7], r_in=0, r_out=1, bins=30)
+    sigma_total = surface_density(sdf, r_in=0, r_out=1, bins=30)
+
+    assert_allclose(sigma_1 + sigma_7, sigma_total, atol=1e-15, rtol=0.0)
+
+
+def test_one_fluid_single_grain() -> None:
+    """ Parts vs whole test for single-grain one-fluid dust. """
+
+    # randomly place particles
+    rng = np.random.default_rng(seed=5)
+    x = rng.random(100)
+    y = rng.random(100)
+    z = rng.random(100)
+    vx = rng.random(100)
+    vy = rng.random(100)
+    vz = rng.random(100)
+    dustfrac = rng.random(100)
+    mass = [3.2e-4] * 100
+
+    sdf = SarracenDataFrame(data={'x': x, 'y': y, 'z': z,
+                                  'vx': vx, 'vy': vy, 'vz': vz,
+                                  'dustfrac': dustfrac, 'mass': mass})
+
+    sigma_g_in, sigma_d_in = surface_density(sdf, r_in=0, r_out=0.5, bins=50)
+    sigma_g_out, sigma_d_out = surface_density(sdf, r_in=0.5, r_out=1, bins=50)
+    sigma_g_all, sigma_d_all = surface_density(sdf, r_in=0, r_out=1, bins=100)
+
+    assert_allclose(sigma_g_in, sigma_g_all[:50], atol=1e-15, rtol=0.0)
+    assert_allclose(sigma_g_out, sigma_g_all[50:], atol=1e-15, rtol=0.0)
+
+    assert_allclose(sigma_d_in, sigma_d_all[:50], atol=1e-15, rtol=0.0)
+    assert_allclose(sigma_d_out, sigma_d_all[50:], atol=1e-15, rtol=0.0)
+
+
+def test_one_fluid_multi_grain() -> None:
+    """ Parts vs whole test for multi-grain one-fluid dust. """
+
+    # randomly place particles
+    rng = np.random.default_rng(seed=5)
+    x = rng.random(100)
+    y = rng.random(100)
+    z = rng.random(100)
+    vx = rng.random(100)
+    vy = rng.random(100)
+    vz = rng.random(100)
+    dustfrac_0 = rng.random(100)
+    dustfrac_1 = rng.random(100)
+    dustfrac_2 = rng.random(100)
+    dustfrac_total = dustfrac_0 + dustfrac_1 + dustfrac_2
+    mass = [3.2e-4] * 100
+
+    sdf = SarracenDataFrame(data={'x': x, 'y': y, 'z': z,
+                                  'vx': vx, 'vy': vy, 'vz': vz,
+                                  'dustfrac_0': dustfrac_0,
+                                  'dustfrac_1': dustfrac_1,
+                                  'dustfrac_2': dustfrac_2,
+                                  'dustfrac_total': dustfrac_total,
+                                  'mass': mass})
+
+    sdf.params['ndustsmall'] = 3
+
+    sigmas = surface_density(sdf, r_in=0, r_out=0.5, bins=50)
+    sigma_g_in, sigma_d_tot_in, sigma_d_in = sigmas
+
+    sigmas = surface_density(sdf, r_in=0.5, r_out=1, bins=50)
+    sigma_g_out, sigma_d_tot_out, sigma_d_out = sigmas
+
+    sigmas = surface_density(sdf, r_in=0, r_out=1, bins=100)
+    sigma_g_all, sigma_d_tot_all, sigma_d_all = sigmas
+
+    assert_allclose(sigma_g_in, sigma_g_all[:50], atol=1e-15, rtol=0)
+    assert_allclose(sigma_g_out, sigma_g_all[50:], atol=1e-15, rtol=0)
+
+    assert_allclose(sigma_d_tot_in, sigma_d_tot_all[:50], atol=1e-15, rtol=0)
+    assert_allclose(sigma_d_tot_out, sigma_d_tot_all[50:], atol=1e-15, rtol=0)
+
+    assert_allclose(sigma_d_in[0], sigma_d_all[0][:50], atol=1e-15, rtol=0)
+    assert_allclose(sigma_d_out[0], sigma_d_all[0][50:], atol=1e-15, rtol=0)
+    assert_allclose(sigma_d_in[1], sigma_d_all[1][:50], atol=1e-15, rtol=0)
+    assert_allclose(sigma_d_out[1], sigma_d_all[1][50:], atol=1e-15, rtol=0)
+    assert_allclose(sigma_d_in[2], sigma_d_all[2][:50], atol=1e-15, rtol=0)
+    assert_allclose(sigma_d_out[2], sigma_d_all[2][50:], atol=1e-15, rtol=0)