From ac995e6524b479a2c9d0b8c05438c23cc82962e0 Mon Sep 17 00:00:00 2001 From: Jean-Francois Gonzalez Date: Thu, 24 Jul 2025 17:35:27 +0200 Subject: [PATCH 1/5] modified surface_density function to compute surface densities of gas and dust for one-fluid dumps --- sarracen/disc/surface_density.py | 52 +++++++++++++++++++++++++++----- 1 file changed, 44 insertions(+), 8 deletions(-) diff --git a/sarracen/disc/surface_density.py b/sarracen/disc/surface_density.py index 45e693a..f59b030 100644 --- a/sarracen/disc/surface_density.py +++ b/sarracen/disc/surface_density.py @@ -118,6 +118,8 @@ def surface_density(data: 'SarracenDataFrame', ------ ValueError If the *geometry* is not *cylindrical* or *spherical*. + If the dump is 'dust-as-particles' or + 'dust-as-mixture' AND 'dust-as-particles', for now. Notes ----- @@ -134,15 +136,49 @@ def surface_density(data: 'SarracenDataFrame', areas = np.pi * (bin_edges[1:] ** 2 - bin_edges[:-1] ** 2) mass = _get_mass(data) - if isinstance(mass, pd.Series): - sigma = mass.groupby(rbins).sum() + if (int(data.params['ndustsmall']) == 0 and + int(data.params['ndustlarge']) == 0): # gas-only dump + if isinstance(mass, pd.Series): + sigma = mass.groupby(rbins).sum() + else: + sigma = data.groupby(rbins).count().iloc[:, 0] * mass + sigma = (sigma / areas).to_numpy() + if retbins: + return sigma, _get_bin_midpoints(bin_edges, log) + else: + return sigma + elif (int(data.params['ndustsmall']) != 0 and + int(data.params['ndustlarge']) == 0): # 'dust-as-mixture' dump + if int(data.params['ndustsmall']) == 1: # only one grain size + mass_gas = mass * (1. - data['dustfrac']) + sigma_gas = (mass_gas.groupby(rbins).sum() / areas).to_numpy() + mass_dust = mass * data['dustfrac'] + sigma_dust = (mass_dust.groupby(rbins).sum() / areas).to_numpy() + if retbins: + return sigma_gas, sigma_dust, \ + _get_bin_midpoints(bin_edges, log) + else: + return sigma_gas, sigma_dust + else: # 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).sum() / areas).to_numpy() + mass_dust_total = mass * data['dustfrac_total'] + sigma_dust_total = (mass_dust_total.groupby(rbins).sum() / + areas).to_numpy() + sigma_dust = np.empty([int(data.params['ndustsmall']), bins]) + for i in range(int(data.params['ndustsmall'])): + sigma_dust[i] = ((mass * data[data.dustfracscol[i]]). + groupby(rbins).sum() / areas).to_numpy() + if retbins: + return sigma_gas, sigma_dust_total, sigma_dust, \ + _get_bin_midpoints(bin_edges, log) + else: + return sigma_gas, sigma_dust_total, sigma_dust else: - sigma = data.groupby(rbins).count().iloc[:, 0] * mass - - if retbins: - return (sigma / areas).to_numpy(), _get_bin_midpoints(bin_edges, log) - else: - return (sigma / areas).to_numpy() + raise ValueError("'dust-as-mixture' AND 'dust-as-particles' dump or" + + "'dust-as-particles' dump: not available yet.") def _calc_angular_momentum(data: 'SarracenDataFrame', From 17f6330a1f9150e482d9c21c47f7ef49cd9e4052 Mon Sep 17 00:00:00 2001 From: Terrence Tricco Date: Tue, 28 Jul 2026 06:40:03 -0230 Subject: [PATCH 2/5] tidy reading of ndust numbers from params --- sarracen/disc/surface_density.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/sarracen/disc/surface_density.py b/sarracen/disc/surface_density.py index 0aee3a8..5b7e84d 100644 --- a/sarracen/disc/surface_density.py +++ b/sarracen/disc/surface_density.py @@ -145,22 +145,25 @@ def surface_density(data: 'SarracenDataFrame', areas = np.pi * (bin_edges[1:] ** 2 - bin_edges[:-1] ** 2) + ndustsmall = data.params.get('ndustsmall', 0) + ndustlarge = data.params.get('ndustlarge', 0) + mass = _get_mass(data) - if (int(data.params['ndustsmall']) == 0 and - int(data.params['ndustlarge']) == 0): # gas-only dump + if ndustsmall == 0 and ndustlarge == 0: # gas-only dump if isinstance(mass, pd.Series): sigma = mass.groupby(rbins, observed=False).sum() else: - sigma = data.groupby(rbins, observed=False).count().iloc[:, 0] * mass + sigma = data.groupby(rbins, observed=False).size() * 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 (int(data.params['ndustsmall']) != 0 and - int(data.params['ndustlarge']) == 0): # 'dust-as-mixture' dump - if int(data.params['ndustsmall']) == 1: # only one grain size + + elif ndustsmall > 0 and ndustlarge == 0: # 'dust-as-mixture' dump + if ndustsmall == 1: # only one grain size mass_gas = mass * (1. - data['dustfrac']) sigma_gas = mass_gas.groupby(rbins, observed=False).sum() sigma_gas = sigma_gas.reindex(rbins.cat.categories, fill_value=0) @@ -189,8 +192,8 @@ def surface_density(data: 'SarracenDataFrame', sigma_dtot = sigma_dtot.reindex(rbins.cat.categories, fill_value=0) sigma_dtot = (sigma_dtot / areas).to_numpy() - sigma_d = np.empty([int(data.params['ndustsmall']), bins]) - for i in range(int(data.params['ndustsmall'])): + sigma_d = np.empty([ndustsmall, bins]) + for i in range(ndustsmall): sigma_d[i] = mass * data[data.dustfracscol[i]].groupby(rbins, observed=False).sum() sigma_d[i] = sigma_d[i].reindex(rbins.cat.categories, fill_value=0) sigma_d[i] = (sigma_d[i] / areas).to_numpy() From eb27a645f4c4c0a8341d49f5afa6d5d282090586 Mon Sep 17 00:00:00 2001 From: Terrence Tricco Date: Tue, 28 Jul 2026 07:15:30 -0230 Subject: [PATCH 3/5] warn if surface density is computed from multiple particle types --- sarracen/disc/surface_density.py | 91 ++++++++++++++++---------------- 1 file changed, 46 insertions(+), 45 deletions(-) diff --git a/sarracen/disc/surface_density.py b/sarracen/disc/surface_density.py index 5b7e84d..059c966 100644 --- a/sarracen/disc/surface_density.py +++ b/sarracen/disc/surface_density.py @@ -139,6 +139,10 @@ def surface_density(data: 'SarracenDataFrame', """ + if 'itype' in data.columns and data['itype'].nunique() > 1: + raise UserWarning('Surface density being calculated from multiple ' + 'particle types. This is not recommended.') + origin = _get_origin(origin) rbins, bin_edges = _bin_particles_by_radius(data, r_in, r_out, bins, log, geometry, origin) @@ -146,10 +150,9 @@ def surface_density(data: 'SarracenDataFrame', areas = np.pi * (bin_edges[1:] ** 2 - bin_edges[:-1] ** 2) ndustsmall = data.params.get('ndustsmall', 0) - ndustlarge = data.params.get('ndustlarge', 0) mass = _get_mass(data) - if ndustsmall == 0 and ndustlarge == 0: # gas-only dump + if ndustsmall == 0: # gas-only dump if isinstance(mass, pd.Series): sigma = mass.groupby(rbins, observed=False).sum() else: @@ -162,49 +165,47 @@ def surface_density(data: 'SarracenDataFrame', else: return sigma - elif ndustsmall > 0 and ndustlarge == 0: # 'dust-as-mixture' dump - if ndustsmall == 1: # only one grain size - mass_gas = mass * (1. - data['dustfrac']) - 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['dustfrac'] - 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 - else: # 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): - sigma_d[i] = mass * data[data.dustfracscol[i]].groupby(rbins, observed=False).sum() - sigma_d[i] = sigma_d[i].reindex(rbins.cat.categories, fill_value=0) - sigma_d[i] = (sigma_d[i] / areas).to_numpy() - if retbins: - return sigma_gas, sigma_dtot, sigma_d, \ - _get_bin_midpoints(bin_edges, log) - else: - return sigma_gas, sigma_dtot, sigma_d - else: - raise ValueError("'dust-as-mixture' AND 'dust-as-particles' dump or" + - "'dust-as-particles' dump: not available yet.") + elif ndustsmall == 1: # 'dust-as-mixture' dump + mass_gas = mass * (1. - data['dustfrac']) + 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['dustfrac'] + 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() + sigma_d[i] = sigma_d[i].reindex(rbins.cat.categories, fill_value=0) + sigma_d[i] = (sigma_d[i] / areas).to_numpy() + 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', From bc7c5a925871d28abc51000e92f02466dcd89b98 Mon Sep 17 00:00:00 2001 From: Terrence Tricco Date: Tue, 28 Jul 2026 08:15:08 -0230 Subject: [PATCH 4/5] unit test of surface density with multiple itypes --- sarracen/disc/surface_density.py | 6 +++-- sarracen/tests/disc/test_surface_density.py | 27 +++++++++++++++++++++ 2 files changed, 31 insertions(+), 2 deletions(-) diff --git a/sarracen/disc/surface_density.py b/sarracen/disc/surface_density.py index 059c966..607f060 100644 --- a/sarracen/disc/surface_density.py +++ b/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 @@ -140,8 +141,9 @@ def surface_density(data: 'SarracenDataFrame', """ if 'itype' in data.columns and data['itype'].nunique() > 1: - raise UserWarning('Surface density being calculated from multiple ' - 'particle types. This is not recommended.') + msg = ('Surface density being calculated from multiple particle types. ' + 'This is not recommended.') + warnings.warn(msg, UserWarning, stacklevel=2) origin = _get_origin(origin) rbins, bin_edges = _bin_particles_by_radius(data, r_in, r_out, bins, log, diff --git a/sarracen/tests/disc/test_surface_density.py b/sarracen/tests/disc/test_surface_density.py index 5ca40ac..5317f80 100644 --- a/sarracen/tests/disc/test_surface_density.py +++ b/sarracen/tests/disc/test_surface_density.py @@ -79,3 +79,30 @@ 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 calculated*: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) From 69da412ac6a46bc229877b08fce716d4912640c0 Mon Sep 17 00:00:00 2001 From: Terrence Tricco Date: Tue, 28 Jul 2026 09:12:45 -0230 Subject: [PATCH 5/5] unit tests of one fluid surface density --- sarracen/disc/surface_density.py | 31 +++++--- sarracen/tests/disc/test_surface_density.py | 83 ++++++++++++++++++++- 2 files changed, 100 insertions(+), 14 deletions(-) diff --git a/sarracen/disc/surface_density.py b/sarracen/disc/surface_density.py index 607f060..4070837 100644 --- a/sarracen/disc/surface_density.py +++ b/sarracen/disc/surface_density.py @@ -84,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. @@ -141,8 +148,8 @@ def surface_density(data: 'SarracenDataFrame', """ if 'itype' in data.columns and data['itype'].nunique() > 1: - msg = ('Surface density being calculated from multiple particle types. ' - 'This is not recommended.') + 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) @@ -151,14 +158,15 @@ def surface_density(data: 'SarracenDataFrame', areas = np.pi * (bin_edges[1:] ** 2 - bin_edges[:-1] ** 2) - ndustsmall = data.params.get('ndustsmall', 0) + ndustsmall = data.params.get('ndustsmall') or len(data.dustfracscol) mass = _get_mass(data) if ndustsmall == 0: # gas-only dump if isinstance(mass, pd.Series): sigma = mass.groupby(rbins, observed=False).sum() else: - sigma = data.groupby(rbins, observed=False).size() * mass + 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() @@ -168,25 +176,28 @@ def surface_density(data: 'SarracenDataFrame', return sigma elif ndustsmall == 1: # 'dust-as-mixture' dump - mass_gas = mass * (1. - data['dustfrac']) + 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['dustfrac'] + 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) + 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) @@ -200,9 +211,7 @@ def surface_density(data: 'SarracenDataFrame', 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() - sigma_d[i] = sigma_d[i].reindex(rbins.cat.categories, fill_value=0) - sigma_d[i] = (sigma_d[i] / areas).to_numpy() + 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) diff --git a/sarracen/tests/disc/test_surface_density.py b/sarracen/tests/disc/test_surface_density.py index 5317f80..e4507de 100644 --- a/sarracen/tests/disc/test_surface_density.py +++ b/sarracen/tests/disc/test_surface_density.py @@ -81,9 +81,7 @@ def test_parts_vs_whole() -> None: assert_allclose(sigma_out, sigma_all[100:], atol=1e-15, rtol=0.0) -@pytest.mark.filterwarnings( - 'ignore:Surface density being calculated*:UserWarning' -) +@pytest.mark.filterwarnings('ignore:Surface density being*:UserWarning') def test_itype_combining() -> None: """ Adding profiles from different itypes should equal the total. """ @@ -106,3 +104,82 @@ def test_itype_combining() -> None: 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)