Merge pull request astropy#897 from aphearin/bijective_dist_matching

Bijective distribution matching

CHANGES.rst

@@ -17,6 +17,8 @@
 
 - Added new bin-free implementation of `conditional_abunmatch`.
 
+- Added new utils function `bijective_distribution_matching`
+
 
 0.6 (2017-12-15)
 ----------------

docs/function_usage/utility_functions.rst

@@ -39,13 +39,15 @@ Matching one distribution to another
 
     distribution_matching_indices
     resample_x_to_match_y
+    bijective_distribution_matching
 
 Rotations, dot products, and other operations in 3d space
 ===============================================================
 
 .. autosummary::
 
     elementwise_dot
+    elementwise_norm
     angles_between_list_of_vectors
     vectors_between_list_of_vectors
     rotation_matrices_from_angles

halotools/utils/distribution_matching.py

@@ -5,7 +5,8 @@
 from astropy.utils.misc import NumpyRNGContext
 
 
-__all__ = ('distribution_matching_indices', 'resample_x_to_match_y')
+__all__ = ('distribution_matching_indices', 'resample_x_to_match_y',
+        'bijective_distribution_matching')
 
 
 def distribution_matching_indices(input_distribution, output_distribution,
@@ -123,3 +124,43 @@ def resample_x_to_match_y(x, y, bins, seed=None):
     indices = np.empty_like(x).astype(int)
     indices[idx_sorted_x] = idx[idx_sorted_xnew]
     return indices
+
+
+def bijective_distribution_matching(x_in, x_desired):
+    """ Replace the values in ``x_in`` with ``x_desired``, preserving the rank-order of ``x_in``
+
+    Parameters
+    ----------
+    x_in : ndarray
+        Numpy array of shape (npts, )
+
+    x_desired : ndarray
+        Numpy array of shape (npts, )
+
+    Returns
+    -------
+    x_out : ndarray
+        Numpy array of shape (npts, )
+
+    Examples
+    --------
+    >>> npts = int(1e5)
+    >>> x_in = np.random.normal(loc=0, scale=0.5, size=npts)
+    >>> x_desired = np.random.normal(loc=2, scale=1, size=npts)
+    >>> x_out = bijective_distribution_matching(x_in, x_desired)
+
+    In the figure below, the left hand panel shows that the output distribution
+    is in exact agreement with the desired distribution. The right hand panel
+    shows that the rank-order of the input distribution is preserved.
+
+    .. image:: /_static/bijective_distribution_matching_demo.png
+    """
+    x_in = np.atleast_1d(x_in)
+    x_desired = np.atleast_1d(x_desired)
+    x_out = np.zeros_like(x_in)
+    idx_sorted = np.argsort(x_in)
+    x_out[idx_sorted] = np.sort(x_desired)
+    return x_out
+
+
+

halotools/utils/tests/test_distribution_matching.py

@@ -4,6 +4,8 @@
 from astropy.utils.misc import NumpyRNGContext
 
 from ..distribution_matching import distribution_matching_indices, resample_x_to_match_y
+from ..distribution_matching import bijective_distribution_matching
+
 
 __all__ = ('test_distribution_matching_indices1', )
 
@@ -50,3 +52,16 @@ def test_resample_x_to_match_y():
         assert np.allclose(result, correct_result, atol=0.02)
     except TypeError:
         pass
+
+
+def test_bijective_distribution_matching():
+    npts = int(1e5)
+    with NumpyRNGContext(fixed_seed):
+        x_in = np.random.normal(loc=0, scale=0.5, size=npts)
+        x_desired = np.random.normal(loc=2, scale=1, size=npts)
+
+    x_out = bijective_distribution_matching(x_in, x_desired)
+    assert np.allclose(np.sort(x_out), np.sort(x_desired))
+
+    idx_x_in_sorted = np.argsort(x_in)
+    assert np.all(np.diff(x_out[idx_x_in_sorted])>=0)