diff --git a/docs/example.rst b/docs/example.rst index 907aa68..0be64f4 100644 --- a/docs/example.rst +++ b/docs/example.rst @@ -1,3 +1,5 @@ +Example usage +============= .. code:: ipython3 %matplotlib inline @@ -20,7 +22,7 @@ That is, the ensemble average of the squared-difference in RM for sources with angular seperation :math:`\delta\theta`. We also need to correct for the impact of errors by: -.. math:: SF_{\text{RM}}(\delta\theta) = SF_{\text{RM},\text{obs}}(\delta\theta) - SF_{\sigma_\text{RM}}(\delta\theta) +.. math:: SF_{\text{RM}}(\delta\theta) = SF_{\text{RM},\text{obs}}(\delta\theta) - SF_{\sigma_\text{RM}}(\delta\theta) See Haverkorn et al. 2004 (2004ApJ…609..776H) for details. @@ -41,15 +43,15 @@ the paper’s plots using a web plot digitiser. 2.2444782900152482, 2.2476963207124476, 2.2837806390213578,]) * (u.rad / u.m**2)**2 - mao_sep = 10**np.array([-0.7729091483767441, - -0.5386163683663935, - -0.2730532911440767, - -0.02550632317850443, - 0.21819567988496358, - 0.47213008276920787, - 0.7173429798998987, - 0.9643533199726302, - 1.18882007856649, + mao_sep = 10**np.array([-0.7729091483767441, + -0.5386163683663935, + -0.2730532911440767, + -0.02550632317850443, + 0.21819567988496358, + 0.47213008276920787, + 0.7173429798998987, + 0.9643533199726302, + 1.18882007856649, 1.3453070240944185,]) * u.deg .. code:: ipython3 @@ -104,14 +106,14 @@ astropy table for convenience rms.append(rm) e_rms.append(e_rm) flags.append(flag) - + mao_rm_tab = Table() mao_rm_tab.add_column(coords, name='coordinates') mao_rm_tab.add_column(rms, name='RM') mao_rm_tab.add_column(e_rms, name='e_RM') mao_rm_tab.add_column(incs, name='included') mao_rm_tab.add_column(flags, name='flag') - + mao_rm_tab @@ -240,7 +242,7 @@ broken power law. Here we’re using ``nestle`` to do the sampling. All ln_noise_evidence: nan ln_evidence: -114.499 +/- 0.130 ln_bayes_factor: nan +/- 0.130 - + 2022-11-07 14:04:01.465 INFO structurefunction - fit_data: Fitting results: 2022-11-07 14:04:01.467 INFO structurefunction - fit_data: amplitude: 180 ± 10 2022-11-07 14:04:01.468 INFO structurefunction - fit_data: x_break: 22 ± 4 @@ -312,7 +314,7 @@ broken power law. Here we’re using ``nestle`` to do the sampling. All ln_noise_evidence: nan ln_evidence: -113.771 +/- 0.124 ln_bayes_factor: nan +/- 0.124 - + 2022-11-07 14:04:47.401 INFO structurefunction - fit_data: Fitting results: 2022-11-07 14:04:47.402 INFO structurefunction - fit_data: amplitude: 180 ± 20 2022-11-07 14:04:47.403 INFO structurefunction - fit_data: x_break: 15 ± 9 @@ -437,7 +439,7 @@ the triple-point structure function is implemented. ln_noise_evidence: nan ln_evidence: -107.875 +/- 0.118 ln_bayes_factor: nan +/- 0.118 - + 2022-11-07 14:15:22.691 INFO structurefunction - fit_data: Fitting results: 2022-11-07 14:15:22.692 INFO structurefunction - fit_data: amplitude: 590 ± 70 2022-11-07 14:15:22.693 INFO structurefunction - fit_data: x_break: 16 ± 9 diff --git a/structurefunction.py b/structurefunction.py index bfcd62f..4916d98 100644 --- a/structurefunction.py +++ b/structurefunction.py @@ -15,7 +15,7 @@ import numpy as np import pandas as pd import xarray as xr -from astropy.coordinates import SkyCoord +from astropy.coordinates import SkyCoord, UnitSphericalRepresentation from astropy.visualization import quantity_support from scipy.optimize import curve_fit from sigfig import round @@ -342,6 +342,7 @@ def sf_two_point( rm_err_2: np.ndarray, dtheta: u.Quantity, bins: u.Quantity, + bin_unit: u.Quantity, ) -> SFResult: """Compute the two-point structure function @@ -351,7 +352,8 @@ def sf_two_point( rm_err_1 (np.ndarray): RM errors of source 1 in pair (n_samples, n_pairs) rm_err_2 (np.ndarray): RM errors of source 2 in pair (n_samples, n_pairs) dtheta (u.Quantity): Separation between sources in pair (n_pairs) - bins (u.Quantity): Angular separation bins + bins (u.Quantity): Angular (or 3D Euclidean) separation bins + bin_unit (u.Quantity): Unit to set bins to when resolving units Returns: SFResult: Structure function results @@ -366,13 +368,13 @@ def sf_two_point( rm_err_2=(["sample", "source_pair"], rm_err_2), ), coords=dict( - seps=("source_pair", dtheta.to(u.deg)), + seps=("source_pair", dtheta.to(bin_unit)), sample=("sample", np.arange(samples)), ), ) # Groupby separation - grp = data_xr.groupby_bins("seps", bins.to(u.deg).value) + grp = data_xr.groupby_bins("seps", bins.to(bin_unit).value) # Compute Structure Function sf_xr = grp.apply(lambda x: ((x.rm_1 - x.rm_2) ** 2).mean(dim="source_pair")) @@ -392,7 +394,7 @@ def sf_two_point( count = grp.count(dim="source_pair").rm_1[:, 0] # Get bin centers - c_bins = np.array([i.mid for i in sf_corr_xr.seps_bins.values]) * u.deg + c_bins = np.array([i.mid for i in sf_corr_xr.seps_bins.values]) * bin_unit return SFResult( med.values, @@ -412,6 +414,7 @@ def sf_three_point( src_2: np.ndarray, dtheta: u.Quantity, bins: u.Quantity, + bin_unit: u.Quantity, ) -> SFResult: """Compute the three-point structure function @@ -423,7 +426,8 @@ def sf_three_point( src_1 (np.ndarray): Source 1 in pair (n_pairs) src_2 (np.ndarray): Source 2 in pair (n_pairs) dtheta (u.Quantity): Separation between sources in pair (n_pairs) - bins (u.Quantity): Angular separation bins + bins (u.Quantity): Angular (or 3D Euclidean) separation bins + bin_unit (u.Quantity): Unit to set bins to when resolving units Returns: SFResult: Structure function result @@ -438,7 +442,7 @@ def sf_three_point( rm_err_2=(["sample", "source_pair"], rm_err_2), ), coords=dict( - seps=("source_pair", dtheta.to(u.deg)), + seps=("source_pair", dtheta.to(bin_unit)), sample=("sample", np.arange(samples)), src_1=("source_pair", src_1), src_2=("source_pair", src_2), @@ -446,7 +450,7 @@ def sf_three_point( ) # Groupby separation - grp = data_xr.groupby_bins("seps", bins.to(u.deg).value) + grp = data_xr.groupby_bins("seps", bins.to(bin_unit).value) rm_1s = [] rm_2s = [] @@ -526,7 +530,7 @@ def sf_three_point( count = triple_grp.count(dim="source_triplet").rm_1[:, 0] # Get bin centers - c_bins = np.array([i for i in sf_t_xr_corr.seps.values]) * u.deg + c_bins = np.array([i for i in sf_t_xr_corr.seps.values]) * bin_unit return SFResult( med.values, @@ -831,22 +835,45 @@ def structure_function( d_rm_1, d_rm_2 = combinate(d_rm_dist) # Get the angular separation of the source pairs + # Check if coords have distance + has_distance = ~issubclass( + coords.data.__class__, + UnitSphericalRepresentation, + ) + logger.info("Getting angular separations...") ra_1, ra_2 = combinate(coords.ra) dec_1, dec2 = combinate(coords.dec) + if has_distance: + dist_1, dist_2 = combinate( + coords.distance + ) - coords_1 = SkyCoord(ra_1, dec_1) - coords_2 = SkyCoord(ra_2, dec2) - dtheta = coords_1.separation(coords_2) + coords_1 = SkyCoord( + ra_1, + dec_1, + distance=dist_1 if has_distance else None, + ) + coords_2 = SkyCoord( + ra_2, + dec2, + distance=dist_2 if has_distance else None, + ) + if has_distance: + dtheta = coords_1.separation_3d(coords_2) + bin_unit = u.pc + else: + dtheta = coords_1.separation(coords_2) + bin_unit = u.deg # Auto compute bins if type(bins) is int: logger.info("Auto-computing bins...") nbins = bins - start = np.log10(np.min(dtheta).to(u.deg).value) - stop = np.log10(np.max(dtheta).to(u.deg).value) - bins = np.logspace(start, stop, nbins, endpoint=True) * u.deg + start = np.log10(np.min(dtheta).to(bin_unit).value) + stop = np.log10(np.max(dtheta).to(bin_unit).value) + bins = np.logspace(start, stop, nbins, endpoint=True) * bin_unit logger.info(f"Maximal angular separation: {np.max(dtheta)}") logger.info(f"Minimal angular separation: {np.min(dtheta)}") else: @@ -863,6 +890,7 @@ def structure_function( rm_err_2=d_rm_2.T, dtheta=dtheta, bins=bins, + bin_unit=bin_unit, ) saturate = np.nanvar(data) * 2 elif n_point == 3: @@ -877,6 +905,7 @@ def structure_function( src_2=src_2, dtheta=dtheta, bins=bins, + bin_unit=bin_unit, ) saturate = np.nanvar(data) * 6 else: