diff --git a/docs/api.rst b/docs/api.rst index d19a61c..397c264 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -80,3 +80,14 @@ The default smoothing kernel is the cubic spline. Additional smoothing kernels a kernels.CubicSplineKernel kernels.QuarticSplineKernel kernels.QuinticSplineKernel + + +Disc +---- + +Accretion disc analysis routines are in the disc module. + +.. autosummary:: + :toctree: api/ + + disc.surface_density diff --git a/sarracen/__init__.py b/sarracen/__init__.py index 69aba3a..1c38c30 100644 --- a/sarracen/__init__.py +++ b/sarracen/__init__.py @@ -8,5 +8,6 @@ from .interpolate import interpolate_2d, interpolate_2d_line, interpolate_3d_proj, interpolate_3d_cross from .render import render, streamlines, arrowplot +import sarracen.disc __version__ = "1.2.3" diff --git a/sarracen/disc/__init__.py b/sarracen/disc/__init__.py new file mode 100644 index 0000000..c35d96f --- /dev/null +++ b/sarracen/disc/__init__.py @@ -0,0 +1 @@ +from ..disc.surface_density import surface_density, angular_momentum, scale_height, honH diff --git a/sarracen/disc/surface_density.py b/sarracen/disc/surface_density.py new file mode 100644 index 0000000..735146e --- /dev/null +++ b/sarracen/disc/surface_density.py @@ -0,0 +1,434 @@ +import numpy as np +import pandas as pd +import sys +from ..sarracen_dataframe import SarracenDataFrame + + +def _get_mass(data: 'SarracenDataFrame'): + if data.mcol == None: + if 'mass' not in data.params: + raise KeyError("'mass' column does not exist in this SarracenDataFrame.") + return data.params['mass'] + + return data[data.mcol] + + +def _get_origin(origin: list) -> list: + if origin is None: + return [0.0, 0.0, 0.0] + else: + return origin + + +def _bin_particles_by_radius(data: 'SarracenDataFrame', + r_in: float = None, + r_out: float = None, + bins: int = 300, + geometry: str = 'cylindrical', + origin: list = None): + """ + Utility function to bin particles in discrete intervals by radius. + + Parameters + ---------- + data: SarracenDataFrame + The particle dataset. + r_in : float, optional + Inner radius of the disc. Defaults to the minimum r value. + r_out : float, optional + Outer radius of the disc. Defaults to the maximum r value. + bins : int, optional + Defines the number of equal-width bins in the range [r_in, r_out]. + Default is 300. + geometry : str, optional + Coordinate system to use to calculate the particle radii. Can be + either *spherical* or *cylindrical*. Defaults to *cylindrical*. + origin : array-like, optional + The x, y and z centre point around which to compute radii. Defaults to + [0, 0, 0]. + + Returns + ------- + rbins: Series + The radial bin to which each particle belongs. + bin_edges: ndarray + Locations of the bin edges. + """ + + if geometry == 'spherical': + r = np.sqrt((data[data.xcol] - origin[0]) ** 2 + + (data[data.ycol] - origin[1]) ** 2 + + (data[data.zcol] - origin[2]) ** 2) + elif geometry == 'cylindrical': + r = np.sqrt((data[data.xcol] - origin[0]) ** 2 + + (data[data.ycol] - origin[1]) ** 2) + else: + raise ValueError("geometry should be either 'cylindrical' or 'spherical'") + + # should we add epsilon here? + if r_in is None: + r_in = r.min() - sys.float_info.epsilon + if r_out is None: + r_out = r.max() + sys.float_info.epsilon + + bin_edges = np.linspace(r_in, r_out, bins+1) + rbins = pd.cut(r, bin_edges) + + return rbins, bin_edges + + +def _get_bin_midpoints(bin_edges: np.ndarray) -> np.ndarray: + """ + Calculate the midpoint of bins given their edges. + """ + + return 0.5 * (bin_edges[1:] - bin_edges[:-1]) + bin_edges[:-1] + + +def surface_density(data: 'SarracenDataFrame', + r_in: float = None, + r_out: float = None, + bins: int = 300, + geometry: str = 'cylindrical', + origin: list = None, + retbins: bool = False): + """ + Calculates the 1D azimuthally-averaged surface density profile. + + The surface density profile is computed by segmenting the particles into + radial bins (rings) and dividing the total mass contained within each bin + by the area of its respective ring. + + Parameters + ---------- + data : SarracenDataFrame + Particle data, in a SarracenDataFrame. + r_in : float, optional + Inner radius of the disc. Defaults to the minimum r value. + r_out : float, optional + Outer radius of the disc. Defaults to the maximum r value. + bins : int, optional + Defines the number of equal-width bins in the range [r_in, r_out]. + Default is 300. + geometry : str, optional + Coordinate system to use to calculate the particle radii. Can be + either *spherical* or *cylindrical*. Defaults to *cylindrical*. + origin : array-like, optional + The x, y and z centre point around which to compute radii. Defaults to + [0, 0, 0]. + retbins : bool, optional + Whether to return the midpoints of the bins or not. Defaults to False. + + Returns + ------- + array + A NumPy array of length bins containing the surface density profile. + array, optional + The midpoint values of each bin. Only returned if *retbins=True*. + + Raises + ------ + ValueError + If the *geometry* is not *cylindrical* or *spherical*. + + See Also + -------- + The surface density averaging procedure for SPH is described in section + 3.2.6 of Lodato & Price, MNRAS (2010), `doi:10.1111/j.1365-2966.2010.16526.x + `_. + """ + + origin = _get_origin(origin) + rbins, bin_edges = _bin_particles_by_radius(data, r_in, r_out, bins, + 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).sum() + else: + sigma = data.groupby(rbins).count().iloc[:, 0] * mass + + if retbins: + return (sigma / areas).to_numpy(), _get_bin_midpoints(bin_edges) + else: + return (sigma / areas).to_numpy() + + +def _calc_angular_momentum(data: 'SarracenDataFrame', + rbins: pd.Series, + origin: list, + unit_vector: bool): + """ + Utility function to calculate angular momentum of the disc. + + Parameters + ---------- + data: SarracenDataFrame + Particle data, in a SarracenDataFrame. + rbins: Series + The radial bin to which each particle belongs. + origin: list + The x, y and z centre point around which to compute radii. + unit_vector: bool + Whether to convert the angular momentum to unit vectors. + Default is True. + + Returns + ------- + Lx, Ly, Lz: Series + The x, y and z components of the angular momentum per bin. + """ + + mass = _get_mass(data) + + x_data = data[data.xcol].to_numpy() - origin[0] + y_data = data[data.ycol].to_numpy() - origin[1] + z_data = data[data.zcol].to_numpy() - origin[2] + + Lx = y_data * data[data.vzcol] - z_data * data[data.vycol] + Ly = z_data * data[data.vxcol] - x_data * data[data.vzcol] + Lz = x_data * data[data.vycol] - y_data * data[data.vxcol] + + if isinstance(mass, float): + Lx = (mass * Lx).groupby(rbins).sum() + Ly = (mass * Ly).groupby(rbins).sum() + Lz = (mass * Lz).groupby(rbins).sum() + else: + Lx = (data[data.mcol] * Lx).groupby(rbins).sum() + Ly = (data[data.mcol] * Ly).groupby(rbins).sum() + Lz = (data[data.mcol] * Lz).groupby(rbins).sum() + + if unit_vector: + Lmag = 1.0 / np.sqrt(Lx ** 2 + Ly ** 2 + Lz ** 2) + + Lx = Lx * Lmag + Ly = Ly * Lmag + Lz = Lz * Lmag + + return Lx, Ly, Lz + + +def angular_momentum(data: 'SarracenDataFrame', + r_in: float = None, + r_out: float = None, + bins: int = 300, + geometry: str = 'cylindrical', + origin: list = None, + retbins: bool = False, + unit_vector: bool = True): + """ + Calculates the angular momentum profile of the disc. + + The profile is computed by segmenting the particles into radial bins + (rings) and summing the angular momentum of the particles within each + bin. + + Parameters + ---------- + data : SarracenDataFrame + Particle data, in a SarracenDataFrame. + r_in : float, optional + Inner radius of the disc. Defaults to the minimum r value. + r_out : float, optional + Outer radius of the disc. Defaults to the maximum r value. + bins : int, optional + Defines the number of equal-width bins in the range [r_in, r_out]. + Default is 300. + geometry : str, optional + Coordinate system to use to calculate the particle radii. Can be + either *spherical* or *cylindrical*. Defaults to *cylindrical*. + origin : array-like, optional + The x, y and z centre point around which to compute radii. Defaults to + [0, 0, 0]. + retbins : bool, optional + Whether to return the midpoints of the bins or not. Defaults to False. + + Returns + ------- + array + A NumPy array of length bins containing the angular momentum profile. + array, optional + The midpoint values of each bin. Only returned if *retbins=True*. + + Raises + ------ + ValueError + If the *geometry* is not *cylindrical* or *spherical*. + """ + + origin = _get_origin(origin) + rbins, bin_edges = _bin_particles_by_radius(data, r_in, r_out, bins, + geometry, origin) + + Lx, Ly, Lz = _calc_angular_momentum(data, rbins, origin, unit_vector) + Lx, Ly, Lz = Lx.to_numpy(), Ly.to_numpy(), Lz.to_numpy() + + if retbins: + return Lx, Ly, Lz, _get_bin_midpoints(bin_edges) + else: + return Lx, Ly, Lz + + +def _calc_scale_height(data: 'SarracenDataFrame', + rbins: pd.Series, + origin: list = None): + """ + Utility function to calculate the scale height of the disc. + + Parameters + ---------- + data: SarracenDataFrame + Particle data, in a SarracenDataFrame. + rbins: Series + The radial bin to which each particle belongs. + origin : array-like, optional + The x, y and z centre point around which to compute radii. Defaults to + [0, 0, 0]. + + Returns + ------- + H: Series + The scale height of the disc. + """ + Lx, Ly, Lz = _calc_angular_momentum(data, rbins, origin, unit_vector=True) + + zdash = rbins.map(Lx).to_numpy() * data[data.xcol] \ + + rbins.map(Ly).to_numpy() * data[data.ycol] \ + + rbins.map(Lz).to_numpy() * data[data.zcol] + + return zdash.groupby(rbins).std() + + +def scale_height(data: 'SarracenDataFrame', + r_in: float = None, + r_out: float = None, + bins: int = 300, + geometry: str = 'cylindrical', + origin: list = None, + retbins: bool = False): + """ + Calculates the scale height, H/R, of the disc. + + The scale height, H/R, is computed by segmenting the particles into radial + bins (rings) and calculating the angular momentum profile of the disc. + Each particle takes the dot product of its position vector with the + angular momentum vector of its corresponding bin. The standard deviation + of this result per bin yields the scale height profile of the disc, which + is divided by the midpoint radius of each bin. + + Parameters + ---------- + data : SarracenDataFrame + Particle data, in a SarracenDataFrame. + r_in : float, optional + Inner radius of the disc. Defaults to the minimum r value. + r_out : float, optional + Outer radius of the disc. Defaults to the maximum r value. + bins : int, optional + Defines the number of equal-width bins in the range [r_in, r_out]. + Default is 300. + geometry : str, optional + Coordinate system to use to calculate the particle radii. Can be + either *spherical* or *cylindrical*. Defaults to *cylindrical*. + origin : array-like, optional + The x, y and z centre point around which to compute radii. Defaults to + [0, 0, 0]. + retbins : bool, optional + Whether to return the midpoints of the bins or not. Defaults to False. + + Returns + ------- + array + A NumPy array of length bins scale height, H, profile. + array, optional + The midpoint values of each bin. Only returned if *retbins=True*. + + Raises + ------ + ValueError + If the *geometry* is not *cylindrical* or *spherical*. + + See Also + -------- + :func:`angular_momentum` : Calculate the angular momentum profile of a + disc. + """ + + origin = _get_origin(origin) + rbins, bin_edges = _bin_particles_by_radius(data, r_in, r_out, bins, + geometry, origin) + + midpoints = _get_bin_midpoints(bin_edges) + H = _calc_scale_height(data, rbins, origin).to_numpy() / midpoints + + if retbins: + return H, midpoints + else: + return H + + +def honH(data: 'SarracenDataFrame', + r_in: float = None, + r_out: float = None, + bins: int = 300, + geometry: str = 'cylindrical', + origin: list = None, + retbins: bool = False): + """ + Calculates /H, the averaged smoothing length divided by the scale height. + + The profile is computed by segmenting the particles into radial bins + (rings). The average smoothing length in each bin is divided by the scale + height as calculated for that bin. + + Parameters + ---------- + data : SarracenDataFrame + Particle data, in a SarracenDataFrame. + r_in : float, optional + Inner radius of the disc. Defaults to the minimum r value. + r_out : float, optional + Outer radius of the disc. Defaults to the maximum r value. + bins : int, optional + Defines the number of equal-width bins in the range [r_in, r_out]. + Default is 300. + geometry : str, optional + Coordinate system to use to calculate the particle radii. Can be + either *spherical* or *cylindrical*. Defaults to *cylindrical*. + origin : array-like, optional + The x, y and z centre point around which to compute radii. Defaults to + [0, 0, 0]. + retbins : bool, optional + Whether to return the midpoints of the bins or not. Defaults to False. + + Returns + ------- + array + A NumPy array of length bins containing the /H profile. + array, optional + The midpoint values of each bin. Only returned if *retbins=True*. + + Raises + ------ + ValueError + If the *geometry* is not *cylindrical* or *spherical*. + + See Also + -------- + :func:`scale_height` : Calculate the scale height of a disc. + """ + + origin = _get_origin(origin) + rbins, bin_edges = _bin_particles_by_radius(data, r_in, r_out, bins, + geometry, origin) + + H = _calc_scale_height(data, rbins, origin).to_numpy() + + mean_h = data.groupby(rbins)[data.hcol].mean().to_numpy() + + if retbins: + return mean_h / H, _get_bin_midpoints(bin_edges) + else: + return mean_h / H \ No newline at end of file diff --git a/sarracen/sarracen_dataframe.py b/sarracen/sarracen_dataframe.py index 59ff3a8..1c36f04 100644 --- a/sarracen/sarracen_dataframe.py +++ b/sarracen/sarracen_dataframe.py @@ -36,7 +36,8 @@ class SarracenDataFrame(DataFrame): """ _internal_names = pd.DataFrame._internal_names + ['_xcol', '_ycol', '_zcol', - '_mcol', '_rhocol', '_hcol'] + '_mcol', '_rhocol', '_hcol', + '_vxcol', '_vycol', '_vzcol'] _internal_names_set = set(_internal_names) _metadata = ['_params', '_units', '_kernel'] @@ -106,7 +107,10 @@ def __init__(self, data=None, params=None, *args, **kwargs): self._units = None self.units = Series([np.nan for _ in range(len(self.columns))]) - self._xcol, self._ycol, self._zcol, self._hcol, self._mcol, self._rhocol = None, None, None, None, None, None + self._xcol, self._ycol, self._zcol = None, None, None + self._hcol, self._mcol, self._rhocol = None, None, None + self._vxcol, self._vycol, self._vzcol = None, None, None + self._identify_special_columns() self._kernel = CubicSplineKernel() @@ -167,6 +171,15 @@ def _identify_special_columns(self): elif 'density' in self.columns: self.rhocol = 'density' + # Look for the keyword 'rho' or 'density' in the data. + if 'vx' in self.columns: + self.vxcol = 'vx' + if 'vy' in self.columns: + self.vycol = 'vy' + if 'vz' in self.columns: + self.vzcol = 'vz' + + def create_mass_column(self): """ Create a new column 'm', copied from the 'massoftype' dataset parameter. @@ -473,6 +486,54 @@ def rhocol(self, new_col: str): if new_col in self or new_col is None: self._rhocol = new_col + + @property + def vxcol(self): + """ + str : Label of the column which contains the x-component of the velocity. + + If this is set to a column which does not exist in the dataset, the column + label will remain set to the old value. + """ + return self._vxcol + + @vxcol.setter + def vxcol(self, new_col: str): + if new_col in self or new_col is None: + self._vxcol = new_col + + @property + def vycol(self): + """ + str : Label of the column which contains the y-component of the velocity. + + If this is set to a column which does not exist in the dataset, the column + label will remain set to the old value. + """ + return self._vycol + + @vycol.setter + def vycol(self, new_col: str): + if new_col in self or new_col is None: + self._vycol = new_col + + + @property + def vzcol(self): + """ + str : Label of the column which contains the z-component of the velocity. + + If this is set to a column which does not exist in the dataset, the column + label will remain set to the old value. + """ + return self._vzcol + + @vzcol.setter + def vzcol(self, new_col: str): + if new_col in self or new_col is None: + self._vzcol = new_col + + @property def kernel(self): """ diff --git a/sarracen/tests/disc/__init__.py b/sarracen/tests/disc/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/sarracen/tests/disc/test_angular_momentum.py b/sarracen/tests/disc/test_angular_momentum.py new file mode 100644 index 0000000..6d3cd45 --- /dev/null +++ b/sarracen/tests/disc/test_angular_momentum.py @@ -0,0 +1,62 @@ +import pandas as pd +import numpy as np +from numpy.testing import assert_array_equal, assert_allclose +from sarracen import SarracenDataFrame +from sarracen.disc import angular_momentum +import pytest + + +def test_mass_equivalency(): + """ Column mass result should equal global params mass result. """ + + # 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) + mass = [3.2e-4] * 100 + + sdf = SarracenDataFrame(data={'x': x, 'y': y, 'z': z, + 'vx': vx, 'vy': vy, 'vz': vz, + 'mass': mass}) + Lx1, Ly1, Lz1 = angular_momentum(sdf) + + sdf = SarracenDataFrame(data={'x': x, 'y': y, 'z': z, + 'vx': vx, 'vy': vy, 'vz': vz}, + params={'mass': 3.2e-4}) + Lx2, Ly2, Lz2 = angular_momentum(sdf) + + assert_array_equal(Lx1, Lx2) + assert_array_equal(Ly1, Ly2) + assert_array_equal(Lz1, Lz2) + + +def test_parts_vs_whole(): + """ Profiles should be the same for matching bins. """ + + # 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) + mass = [3.2e-4] * 100 + + sdf = SarracenDataFrame(data={'x': x, 'y': y, 'z': z, + 'vx': vx, 'vy': vy, 'vz': vz, + 'mass': mass}) + Lx_in, Ly_in, Lz_in = angular_momentum(sdf, r_in=0.0, r_out=0.5, bins=100) + Lx_out, Ly_out, Lz_out = angular_momentum(sdf, r_in=0.5, r_out=1.0, bins=100) + Lx_all, Ly_all, Lz_all = angular_momentum(sdf, r_in=0.0, r_out=1.0, bins=200) + + assert_array_equal(Lx_in, Lx_all[:100]) + assert_array_equal(Lx_out, Lx_all[100:]) + assert_array_equal(Ly_in, Ly_all[:100]) + assert_array_equal(Ly_out, Ly_all[100:]) + assert_array_equal(Lz_in, Lz_all[:100]) + assert_array_equal(Lz_out, Lz_all[100:]) \ No newline at end of file diff --git a/sarracen/tests/disc/test_surface_density.py b/sarracen/tests/disc/test_surface_density.py new file mode 100644 index 0000000..2ff02a8 --- /dev/null +++ b/sarracen/tests/disc/test_surface_density.py @@ -0,0 +1,82 @@ +import pandas as pd +import numpy as np +from numpy.testing import assert_array_equal, assert_allclose +from sarracen import SarracenDataFrame +from sarracen.disc import surface_density +import pytest + + +def test_mass_equivalency(): + """ Column mass result should equal global params mass result. """ + + # randomly place particles + rng = np.random.default_rng(seed=5) + x = rng.random(100) + y = rng.random(100) + z = rng.random(100) + mass = [3.2e-4] * 100 + + sdf = SarracenDataFrame(data={'x': x, 'y': y, 'z': z, 'mass': mass}) + sigma1 = surface_density(sdf) + + sdf = SarracenDataFrame(data={'x': x, 'y': y, 'z': z}, + params={'mass': 3.2e-4}) + sigma2 = surface_density(sdf) + + assert_array_equal(sigma1, sigma2) + + +@pytest.mark.parametrize("geometry", ['cylindrical', 'spherical']) +def test_origin(geometry): + """ Expect same profile regardless of origin. """ + + rng = np.random.default_rng(seed=5) + x = rng.random(100) + y = rng.random(100) + z = rng.random(100) + + sdf = SarracenDataFrame(data={'x': x, 'y': y, 'z': z}, + params={'mass': 3.2e-4}) + sigma_zero = surface_density(sdf, origin=[0.0, 0.0, 0.0], + bins=30, geometry=geometry) + + sdf = SarracenDataFrame(data={'x': x - 2e-8, + 'y': y + 2e-8, + 'z': z + 3e-9}, + params={'mass': 3.2e-4}) + sigma_eps = surface_density(sdf, origin=[2e-8, -2e-8, 3e-9], + bins=30, geometry=geometry) + + sdf = SarracenDataFrame(data={'x': x + 5.6e4, + 'y': y - 8.7e3, + 'z': z + 5.4e6}, + params={'mass': 3.2e-4}) + sigma_large = surface_density(sdf, origin=[5.6e4, -8.7e3, 5.4e6], + bins=30, geometry=geometry) + + assert_allclose(sigma_zero, sigma_eps, atol=2e-8, rtol=0.0) + assert_allclose(sigma_zero, sigma_large, atol=2e-8, rtol=0.0) + + +def test_parts_vs_whole(): + """ Profiles should be the same for matching bins. """ + + # 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) + mass = [3.2e-4] * 100 + + sdf = SarracenDataFrame(data={'x': x, 'y': y, 'z': z, + 'vx': vx, 'vy': vy, 'vz': vz, + 'mass': mass}) + sigma_in = surface_density(sdf, r_in=0.0, r_out=0.5, bins=100) + sigma_out = surface_density(sdf, r_in=0.5, r_out=1.0, bins=100) + sigma_all = surface_density(sdf, r_in=0.0, r_out=1.0, bins=200) + + 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)