Merge pull request #4 from AlecThomson/cartesian

Add support for cartesian coordinates
AlecThomson · Nov 25, 2022 · 7a7cd59 · 7a7cd59
2 parents 0369b72 + 9e2301b
commit 7a7cd59
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 30 deletions.
diff --git 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
@@ -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,15 +442,15 @@ 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),
         ),
     )
 
     # 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: