Skip to content

Commit

Permalink
Merge pull request #78 from j-f-gonzalez/main
Browse files Browse the repository at this point in the history 
disc/surface_density.py: added option to bin in log scale
  • Loading branch information
@ttricco
ttricco authored Mar 22, 2024
2 parents f0a75e5 + fb7c002 commit ad7c241
Showing 1 changed file with 38 additions and 12 deletions.
50 changes: 38 additions & 12 deletions sarracen/disc/surface_density.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ def _bin_particles_by_radius(data: 'SarracenDataFrame',
r_in: float = None,
r_out: float = None,
bins: int = 300,
log: bool = False,
geometry: str = 'cylindrical',
origin: list = None):
"""
Expand All @@ -40,6 +41,8 @@ def _bin_particles_by_radius(data: 'SarracenDataFrame',
bins : int, optional
Defines the number of equal-width bins in the range [r_in, r_out].
Default is 300.
log : bool, optional
Whether to bin in log scale or not. Defaults to False.
geometry : str, optional
Coordinate system to use to calculate the particle radii. Can be
either *spherical* or *cylindrical*. Defaults to *cylindrical*.
Expand Down Expand Up @@ -71,24 +74,36 @@ def _bin_particles_by_radius(data: 'SarracenDataFrame',
if r_out is None:
r_out = r.max() + sys.float_info.epsilon

bin_edges = np.linspace(r_in, r_out, bins+1)
if log:
bin_edges = np.logspace(np.log10(r_in), np.log10(r_out), bins+1)
else:
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:
def _get_bin_midpoints(bin_edges: np.ndarray,
log: bool = False) -> np.ndarray:
"""
Calculate the midpoint of bins given their edges.
bin_edges: ndarray
Locations of the bin edges.
log : bool, optional
Whether to bin in log scale or not. Defaults to False.
"""

return 0.5 * (bin_edges[1:] - bin_edges[:-1]) + bin_edges[:-1]
if log:
return np.sqrt(bin_edges[:-1] * bin_edges[1:])
else:
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,
log: bool = False,
geometry: str = 'cylindrical',
origin: list = None,
retbins: bool = False):
Expand All @@ -110,6 +125,8 @@ def surface_density(data: 'SarracenDataFrame',
bins : int, optional
Defines the number of equal-width bins in the range [r_in, r_out].
Default is 300.
log : bool, optional
Whether to bin in log scale or not. Defaults to False.
geometry : str, optional
Coordinate system to use to calculate the particle radii. Can be
either *spherical* or *cylindrical*. Defaults to *cylindrical*.
Expand Down Expand Up @@ -139,7 +156,7 @@ def surface_density(data: 'SarracenDataFrame',
"""

origin = _get_origin(origin)
rbins, bin_edges = _bin_particles_by_radius(data, r_in, r_out, bins,
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)
Expand All @@ -151,7 +168,7 @@ def surface_density(data: 'SarracenDataFrame',
sigma = data.groupby(rbins).count().iloc[:, 0] * mass

if retbins:
return (sigma / areas).to_numpy(), _get_bin_midpoints(bin_edges)
return (sigma / areas).to_numpy(), _get_bin_midpoints(bin_edges, log)
else:
return (sigma / areas).to_numpy()

Expand Down Expand Up @@ -214,6 +231,7 @@ def angular_momentum(data: 'SarracenDataFrame',
r_in: float = None,
r_out: float = None,
bins: int = 300,
log: bool = False,
geometry: str = 'cylindrical',
origin: list = None,
retbins: bool = False,
Expand All @@ -236,6 +254,8 @@ def angular_momentum(data: 'SarracenDataFrame',
bins : int, optional
Defines the number of equal-width bins in the range [r_in, r_out].
Default is 300.
log : bool, optional
Whether to bin in log scale or not. Defaults to False.
geometry : str, optional
Coordinate system to use to calculate the particle radii. Can be
either *spherical* or *cylindrical*. Defaults to *cylindrical*.
Expand All @@ -259,14 +279,14 @@ def angular_momentum(data: 'SarracenDataFrame',
"""

origin = _get_origin(origin)
rbins, bin_edges = _bin_particles_by_radius(data, r_in, r_out, bins,
rbins, bin_edges = _bin_particles_by_radius(data, r_in, r_out, bins, log,
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)
return Lx, Ly, Lz, _get_bin_midpoints(bin_edges, log)
else:
return Lx, Ly, Lz

Expand Down Expand Up @@ -305,6 +325,7 @@ def scale_height(data: 'SarracenDataFrame',
r_in: float = None,
r_out: float = None,
bins: int = 300,
log: bool = False,
geometry: str = 'cylindrical',
origin: list = None,
retbins: bool = False):
Expand All @@ -329,6 +350,8 @@ def scale_height(data: 'SarracenDataFrame',
bins : int, optional
Defines the number of equal-width bins in the range [r_in, r_out].
Default is 300.
log : bool, optional
Whether to bin in log scale or not. Defaults to False.
geometry : str, optional
Coordinate system to use to calculate the particle radii. Can be
either *spherical* or *cylindrical*. Defaults to *cylindrical*.
Expand Down Expand Up @@ -357,10 +380,10 @@ def scale_height(data: 'SarracenDataFrame',
"""

origin = _get_origin(origin)
rbins, bin_edges = _bin_particles_by_radius(data, r_in, r_out, bins,
rbins, bin_edges = _bin_particles_by_radius(data, r_in, r_out, bins, log,
geometry, origin)

midpoints = _get_bin_midpoints(bin_edges)
midpoints = _get_bin_midpoints(bin_edges, log)
H = _calc_scale_height(data, rbins, origin).to_numpy() / midpoints

if retbins:
Expand All @@ -373,6 +396,7 @@ def honH(data: 'SarracenDataFrame',
r_in: float = None,
r_out: float = None,
bins: int = 300,
log: bool = False,
geometry: str = 'cylindrical',
origin: list = None,
retbins: bool = False):
Expand All @@ -394,6 +418,8 @@ def honH(data: 'SarracenDataFrame',
bins : int, optional
Defines the number of equal-width bins in the range [r_in, r_out].
Default is 300.
log : bool, optional
Whether to bin in log scale or not. Defaults to False.
geometry : str, optional
Coordinate system to use to calculate the particle radii. Can be
either *spherical* or *cylindrical*. Defaults to *cylindrical*.
Expand Down Expand Up @@ -421,14 +447,14 @@ def honH(data: 'SarracenDataFrame',
"""

origin = _get_origin(origin)
rbins, bin_edges = _bin_particles_by_radius(data, r_in, r_out, bins,
rbins, bin_edges = _bin_particles_by_radius(data, r_in, r_out, bins, log,
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)
return mean_h / H, _get_bin_midpoints(bin_edges, log)
else:
return mean_h / H
return mean_h / H

0 comments on commit ad7c241

Please sign in to comment.