Pipweights

CHANGES.rst

@@ -20,6 +20,10 @@ Changes
 - Migrate doctests to Python 3.8 [#261]
 - Migrate Python tests to pytest [#265]
 
+Fixes
+-----
+- Add additional check to tell if it's safe to redirect stdout/err [#270]
+
 
 2.4.0 (2021-09-30)
 ==================

Corrfunc/mocks/DDrppi_mocks.py

@@ -14,7 +14,7 @@
 __all__ = ('DDrppi_mocks', )
 
 
-def DDrppi_mocks(autocorr, cosmology, nthreads, pimax, binfile,
+def DDrppi_mocks(autocorr, cosmology, nthreads, binfile, pimax, npibins,
                  RA1, DEC1, CZ1, weights1=None,
                  RA2=None, DEC2=None, CZ2=None, weights2=None,
                  is_comoving_dist=False,
@@ -24,10 +24,11 @@ def DDrppi_mocks(autocorr, cosmology, nthreads, pimax, binfile,
                  zbin_refine_factor=1, max_cells_per_dim=100,
                  copy_particles=True, enable_min_sep_opt=True,
                  c_api_timer=False, isa='fastest',
-                 weight_type=None, bin_type='custom'):
+                 weight_type=None, bin_type='custom',
+                 pair_weights=None, sep_pair_weights=None, attrs_pair_weights=None):
     """
     Calculate the 2-D pair-counts corresponding to the projected correlation
-    function, :math:`\\xi(r_p, \pi)`. Pairs which are separated by less
+    function, :math:`\\xi(r_p, \\pi)`. Pairs which are separated by less
     than the ``rp`` bins (specified in ``binfile``) in the
     X-Y plane, and less than ``pimax`` in the Z-dimension are
     counted. The input positions are expected to be on-sky co-ordinates.
@@ -42,16 +43,15 @@ def DDrppi_mocks(autocorr, cosmology, nthreads, pimax, binfile,
 
 
     .. note:: that this module only returns pair counts and not the actual
-       correlation function :math:`\\xi(r_p, \pi)` or :math:`wp(r_p)`. See the
+       correlation function :math:`\\xi(r_p, \\pi)` or :math:`wp(r_p)`. See the
        utilities :py:mod:`Corrfunc.utils.convert_3d_counts_to_cf` and
        :py:mod:`Corrfunc.utils.convert_rp_pi_counts_to_wp` for computing
-       :math:`\\xi(r_p, \pi)` and :math:`wp(r_p)` respectively from the
+       :math:`\\xi(r_p, \\pi)` and :math:`wp(r_p)` respectively from the
        pair counts.
 
 
     Parameters
     ----------
-
     autocorr : boolean, required
         Boolean flag for auto/cross-correlation. If autocorr is set to 1,
         then the second set of particle positions are not required.
@@ -75,15 +75,6 @@ def DDrppi_mocks(autocorr, cosmology, nthreads, pimax, binfile,
         The number of OpenMP threads to use. Has no effect if OpenMP was not
         enabled during library compilation.
 
-    pimax : double
-        A double-precision value for the maximum separation along
-        the Z-dimension.
-
-        Distances along the :math:`\\pi` direction are binned with unit
-        depth. For instance, if ``pimax=40``, then 40 bins will be created
-        along the ``pi`` direction. Only pairs with ``0 <= dz < pimax``
-        are counted (no equality).
-
     binfile : string or an list/array of floats
         For string input: filename specifying the ``rp`` bins for
         ``DDrppi_mocks``. The file should contain white-space separated values
@@ -96,6 +87,19 @@ def DDrppi_mocks(autocorr, cosmology, nthreads, pimax, binfile,
         input specifying **14** (logarithmic) bins between 0.1 and 10.0. This
         array does not need to be sorted.
 
+    pimax : double
+        A double-precision value for the maximum separation along
+        the Z-dimension.
+
+        Distances along the :math:`\\pi` direction are binned with unit
+        depth. For instance, if ``pimax=40``, then 40 bins will be created
+        along the ``pi`` direction. Only pairs with ``0 <= dz < pimax``
+        are counted (no equality).
+
+    npibins : int
+        The number of linear ``pi`` bins, with the bins ranging from
+        from (0, :math:`\\pi_{max}`)
+
     RA1 : array-like, real (float/double)
         The array of Right Ascensions for the first set of points. RA's
         are expected to be in [0.0, 360.0], but the code will try to fix cases
@@ -172,7 +176,7 @@ def DDrppi_mocks(autocorr, cosmology, nthreads, pimax, binfile,
         you need accurate ``rpavg`` values, then pass in double precision
         arrays for the particle positions.
 
-    fast_divide_and_NR_steps: integer (default 0)
+    fast_divide_and_NR_steps : integer (default 0)
         Replaces the division in ``AVX`` implementation with an approximate
         reciprocal, followed by ``fast_divide_and_NR_steps`` of Newton-Raphson.
         Can improve runtime by ~15-20% on older computers. Value of 0 uses
@@ -182,18 +186,18 @@ def DDrppi_mocks(autocorr, cosmology, nthreads, pimax, binfile,
         Controls the refinement on the cell sizes. Can have up to a 20% impact
         on runtime.
 
-    max_cells_per_dim: integer, default is 100, typical values in [50-300]
+    max_cells_per_dim : integer, default is 100, typical values in [50-300]
         Controls the maximum number of cells per dimension. Total number of
         cells can be up to (max_cells_per_dim)^3. Only increase if ``rpmax`` is
         too small relative to the boxsize (and increasing helps the runtime).
 
-    copy_particles: boolean (default True)
+    copy_particles : boolean (default True)
         Boolean flag to make a copy of the particle positions
         If set to False, the particles will be re-ordered in-place
 
         .. versionadded:: 2.3.0
 
-    enable_min_sep_opt: boolean (default true)
+    enable_min_sep_opt : boolean (default true)
        Boolean flag to allow optimizations based on min. separation between
        pairs of cells. Here to allow for comparison studies.
 
@@ -236,26 +240,35 @@ def DDrppi_mocks(autocorr, cosmology, nthreads, pimax, binfile,
         ``rtol = 1e-05`` *and* ``atol = 1e-08`` (relative and absolute tolerance)
         of ``np.linspace(binfile[0], binfile[-1], len(binfile))``.
 
-    Returns
-    --------
+    pair_weights : array-like, optional. Default: None.
+        Array of pair weights.
 
+    sep_pair_weights : array-like, optional. Default: None.
+        Array of separations corresponding to ``pair_weights``.
+
+    attrs_pair_weights : tuple. Default: None.
+        Attributes for pair weights; in case ``weight_type`` is "inverse_bitwise",
+        the tuple of (offset to be added to the bitwise counts,
+        default weight value if denominator is zero).
+
+    Returns
+    -------
     results : Numpy structured array
 
         A numpy structured array containing [rpmin, rpmax, rpavg, pimax,
         npairs, weightavg] for each radial bin specified in the ``binfile``.
         If ``output_ravg`` is not set, then ``rpavg`` will be set to 0.0 for
         all bins; similarly for ``weightavg``. ``npairs`` contains the number
         of pairs in that bin and can be used to compute the actual
-        :math:`\\xi(r_p, \pi)` or :math:`wp(rp)` by combining with
+        :math:`\\xi(r_p, \\pi)` or :math:`wp(rp)` by combining with
         (DR, RR) counts.
 
     api_time : float, optional
         Only returned if ``c_api_timer`` is set.  ``api_time`` measures only
         the time spent within the C library and ignores all python overhead.
 
     Example
-    --------
-
+    -------
     >>> from __future__ import print_function
     >>> import numpy as np
     >>> from os.path import dirname, abspath, join as pjoin
@@ -344,7 +357,7 @@ def DDrppi_mocks(autocorr, cosmology, nthreads, pimax, binfile,
 
     import numpy as np
     from Corrfunc.utils import translate_isa_string_to_enum, translate_bin_type_string_to_enum,\
-        fix_ra_dec, return_file_with_rbins, convert_to_native_endian,\
+        fix_ra_dec, get_edges, convert_to_native_endian,\
         sys_pipes, process_weights
     from future.utils import bytes_to_native_str
 
@@ -361,27 +374,37 @@ def DDrppi_mocks(autocorr, cosmology, nthreads, pimax, binfile,
     weights1, weights2 = process_weights(weights1, weights2, RA1, RA2, weight_type, autocorr)
 
     # Ensure all input arrays are native endian
-    RA1, DEC1, CZ1, weights1, RA2, DEC2, CZ2, weights2 = [
+    RA1, DEC1, CZ1, RA2, DEC2, CZ2 = [
             convert_to_native_endian(arr, warn=True) for arr in
-            [RA1, DEC1, CZ1, weights1, RA2, DEC2, CZ2, weights2]]
+            [RA1, DEC1, CZ1, RA2, DEC2, CZ2]]
 
     fix_ra_dec(RA1, DEC1)
     if autocorr == 0:
         fix_ra_dec(RA2, DEC2)
 
+    if weights1 is not None:
+        weights1 = [convert_to_native_endian(arr, warn=True) for arr in weights1]
+    if weights2 is not None:
+        weights2 = [convert_to_native_endian(arr, warn=True) for arr in weights2]
+
+    if pair_weights is not None:
+        pair_weights = convert_to_native_endian(pair_weights, warn=True)
+        sep_pair_weights = convert_to_native_endian(sep_pair_weights, warn=True)
+
     # Passing None parameters breaks the parsing code, so avoid this
     kwargs = {}
-    for k in ['weights1', 'weights2', 'weight_type', 'RA2', 'DEC2', 'CZ2']:
+    for k in ['weights1', 'weights2', 'weight_type', 'RA2', 'DEC2', 'CZ2',
+              'pair_weights', 'sep_pair_weights', 'attrs_pair_weights']:
         v = locals()[k]
         if v is not None:
             kwargs[k] = v
 
     integer_isa = translate_isa_string_to_enum(isa)
     integer_bin_type = translate_bin_type_string_to_enum(bin_type)
-    rbinfile, delete_after_use = return_file_with_rbins(binfile)
+    rbinfile = get_edges(binfile)
     with sys_pipes():
         extn_results = DDrppi_extn(autocorr, cosmology, nthreads,
-                                   pimax, rbinfile,
+                                   rbinfile, pimax, npibins,
                                    RA1, DEC1, CZ1,
                                    is_comoving_dist=is_comoving_dist,
                                    verbose=verbose,
@@ -402,10 +425,6 @@ def DDrppi_mocks(autocorr, cosmology, nthreads, pimax, binfile,
     else:
         extn_results, api_time = extn_results
 
-    if delete_after_use:
-        import os
-        os.remove(rbinfile)
-
     results_dtype = np.dtype([(bytes_to_native_str(b'rmin'), np.float64),
                               (bytes_to_native_str(b'rmax'), np.float64),
                               (bytes_to_native_str(b'rpavg'), np.float64),