Nightly bin a dataset of radial velocities
| import numpy as np | |
| ############################################################################### | |
| # the following is mostly a copy of the scipy implementation of | |
| # binned_statistic and binned_statistic_dd | |
| # but allowing for a weights parameter | |
| from scipy._lib.six import callable, xrange | |
| from scipy._lib._numpy_compat import suppress_warnings | |
| ## careful here! | |
| import warnings | |
| warnings.simplefilter(action='ignore', category=FutureWarning) | |
| def binned_statistic(x, values, statistic='mean', bins=10, range=None, | |
| weights=None): | |
| try: | |
| N = len(bins) | |
| except TypeError: | |
| N = 1 | |
| if N != 1: | |
| bins = [np.asarray(bins, float)] | |
| if range is not None: | |
| if len(range) == 2: | |
| range = [range] | |
| medians, edges, binnumbers = binned_statistic_dd( | |
| [x], values, statistic, bins, range, weights) | |
| return medians, edges[0], binnumbers | |
| def binned_statistic_dd(sample, values, statistic='mean', bins=10, range=None, | |
| weights=None, expand_binnumbers=False): | |
| known_stats = ['mean', 'median', 'count', 'sum', 'std', 'min', 'max'] | |
| if not callable(statistic) and statistic not in known_stats: | |
| raise ValueError('invalid statistic %r' % (statistic, )) | |
| # `Ndim` is the number of dimensions (e.g. `2` for `binned_statistic_2d`) | |
| # `Dlen` is the length of elements along each dimension. | |
| # This code is based on np.histogramdd | |
| try: | |
| # `sample` is an ND-array. | |
| Dlen, Ndim = sample.shape | |
| except (AttributeError, ValueError): | |
| # `sample` is a sequence of 1D arrays. | |
| sample = np.atleast_2d(sample).T | |
| Dlen, Ndim = sample.shape | |
| # Store initial shape of `values` to preserve it in the output | |
| values = np.asarray(values) | |
| input_shape = list(values.shape) | |
| # Make sure that `values` is 2D to iterate over rows | |
| values = np.atleast_2d(values) | |
| if weights is not None: | |
| weights = np.atleast_2d(weights) | |
| Vdim, Vlen = values.shape | |
| # Make sure `values` match `sample` | |
| if (statistic != 'count' and Vlen != Dlen): | |
| raise AttributeError('The number of `values` elements must match the ' | |
| 'length of each `sample` dimension.') | |
| nbin = np.empty(Ndim, int) # Number of bins in each dimension | |
| edges = Ndim * [None] # Bin edges for each dim (will be 2D array) | |
| dedges = Ndim * [None] # Spacing between edges (will be 2D array) | |
| try: | |
| M = len(bins) | |
| if M != Ndim: | |
| raise AttributeError('The dimension of bins must be equal ' | |
| 'to the dimension of the sample x.') | |
| except TypeError: | |
| bins = Ndim * [bins] | |
| # Select range for each dimension | |
| # Used only if number of bins is given. | |
| if range is None: | |
| smin = np.atleast_1d(np.array(sample.min(axis=0), float)) | |
| smax = np.atleast_1d(np.array(sample.max(axis=0), float)) | |
| else: | |
| smin = np.zeros(Ndim) | |
| smax = np.zeros(Ndim) | |
| for i in xrange(Ndim): | |
| smin[i], smax[i] = range[i] | |
| # Make sure the bins have a finite width. | |
| for i in xrange(len(smin)): | |
| if smin[i] == smax[i]: | |
| smin[i] = smin[i] - .5 | |
| smax[i] = smax[i] + .5 | |
| # Create edge arrays | |
| for i in xrange(Ndim): | |
| if np.isscalar(bins[i]): | |
| nbin[i] = bins[i] + 2 # +2 for outlier bins | |
| edges[i] = np.linspace(smin[i], smax[i], nbin[i] - 1) | |
| else: | |
| edges[i] = np.asarray(bins[i], float) | |
| nbin[i] = len(edges[i]) + 1 # +1 for outlier bins | |
| dedges[i] = np.diff(edges[i]) | |
| nbin = np.asarray(nbin) | |
| # Compute the bin number each sample falls into, in each dimension | |
| sampBin = [np.digitize(sample[:, i], edges[i]) for i in xrange(Ndim)] | |
| # Using `digitize`, values that fall on an edge are put in the right bin. | |
| # For the rightmost bin, we want values equal to the right | |
| # edge to be counted in the last bin, and not as an outlier. | |
| for i in xrange(Ndim): | |
| # Find the rounding precision | |
| decimal = int(-np.log10(dedges[i].min())) + 6 | |
| # Find which points are on the rightmost edge. | |
| on_edge = np.where( | |
| np.around(sample[:, i], decimal) == np.around( | |
| edges[i][-1], decimal))[0] | |
| # Shift these points one bin to the left. | |
| sampBin[i][on_edge] -= 1 | |
| # Compute the sample indices in the flattened statistic matrix. | |
| binnumbers = np.ravel_multi_index(sampBin, nbin) | |
| result = np.empty([Vdim, nbin.prod()], float) | |
| if statistic == 'mean': | |
| result.fill(np.nan) | |
| flatcount = np.bincount(binnumbers, None) | |
| a = flatcount.nonzero() | |
| for vv in xrange(Vdim): | |
| flatsum = np.bincount(binnumbers, values[vv]) | |
| result[vv, a] = flatsum[a] / flatcount[a] | |
| elif statistic == 'std': | |
| result.fill(0) | |
| flatcount = np.bincount(binnumbers, None) | |
| a = flatcount.nonzero() | |
| for vv in xrange(Vdim): | |
| flatsum = np.bincount(binnumbers, values[vv]) | |
| flatsum2 = np.bincount(binnumbers, values[vv]**2) | |
| result[vv, a] = np.sqrt(flatsum2[a] / flatcount[a] - | |
| (flatsum[a] / flatcount[a])**2) | |
| elif statistic == 'count': | |
| result.fill(0) | |
| flatcount = np.bincount(binnumbers, None) | |
| a = np.arange(len(flatcount)) | |
| result[:, a] = flatcount[np.newaxis, :] | |
| elif statistic == 'sum': | |
| result.fill(0) | |
| for vv in xrange(Vdim): | |
| flatsum = np.bincount(binnumbers, values[vv]) | |
| a = np.arange(len(flatsum)) | |
| result[vv, a] = flatsum | |
| elif statistic == 'median': | |
| result.fill(np.nan) | |
| for i in np.unique(binnumbers): | |
| for vv in xrange(Vdim): | |
| result[vv, i] = np.median(values[vv, binnumbers == i]) | |
| elif statistic == 'min': | |
| result.fill(np.nan) | |
| for i in np.unique(binnumbers): | |
| for vv in xrange(Vdim): | |
| result[vv, i] = np.min(values[vv, binnumbers == i]) | |
| elif statistic == 'max': | |
| result.fill(np.nan) | |
| for i in np.unique(binnumbers): | |
| for vv in xrange(Vdim): | |
| result[vv, i] = np.max(values[vv, binnumbers == i]) | |
| elif callable(statistic): | |
| with np.errstate(invalid='ignore'), suppress_warnings() as sup: | |
| sup.filter(RuntimeWarning) | |
| try: | |
| null = statistic([], weights=[]) | |
| except: | |
| null = np.nan | |
| result.fill(null) | |
| for i in np.unique(binnumbers): | |
| for vv in xrange(Vdim): | |
| if weights is None: | |
| result[vv, i] = statistic(values[vv, binnumbers == i]) | |
| else: | |
| result[vv, i] = statistic(values[vv, binnumbers == i], | |
| weights[vv, binnumbers == i]) | |
| # Shape into a proper matrix | |
| result = result.reshape(np.append(Vdim, nbin)) | |
| # Remove outliers (indices 0 and -1 for each bin-dimension). | |
| core = [slice(None)] + Ndim * [slice(1, -1)] | |
| result = result[core] | |
| # Unravel binnumbers into an ndarray, each row the bins for each dimension | |
| if (expand_binnumbers and Ndim > 1): | |
| binnumbers = np.asarray(np.unravel_index(binnumbers, nbin)) | |
| if np.any(result.shape[1:] != nbin - 2): | |
| raise RuntimeError('Internal Shape Error') | |
| # Reshape to have output (`reulst`) match input (`values`) shape | |
| result = result.reshape(input_shape[:-1] + list(nbin - 2)) | |
| return result, edges, binnumbers | |
| # put back the documentation | |
| from scipy.stats import binned_statistic as old_binned_statistic,\ | |
| binned_statistic_dd as old_binned_statistic_dd | |
| doc1 = old_binned_statistic.__doc__ | |
| doc2 = old_binned_statistic_dd.__doc__ | |
| binned_statistic.__doc__ = doc1 | |
| binned_statistic_dd.__doc__ = doc2 | |
| ############################################################################### | |
| def binRV(time, rv, err=None, stat='wmean', tstat='wmean', estat='addquad'): | |
| """ Bin a dataset of radial-velocity observations. | |
| Parameters | |
| ---------- | |
| time : array | |
| The array of times where the radial velocity is measured. This function | |
| does nightly binning, based on the integer part of this array. | |
| rv : array | |
| The radial-velocity values. | |
| err : array (optional) | |
| The uncertainty on the radial-velocity values. | |
| stat : string or callable (optional) | |
| The statistic to compute for the radial velocities. By default, this is | |
| 'mean' if the `err` array is not provided and weighted mean ('wmean') | |
| if `err` is provided. In that case, the routine does inverse variance | |
| averaging, using 1/`err`**2 as weights. See other available statistics | |
| on the `binned_statistic` function. | |
| tstat : string or callable (optional) | |
| The statistic to compute for the times. The default is the same as for | |
| the radial velocities. See other available statistics on the | |
| `binned_statistic` function. | |
| estat : string or callable (optional) | |
| The statistic to compute for the errors. The default is to add them | |
| quadratically and divide by the number. See other available statistics | |
| on the `binned_statistic` function. | |
| Notes | |
| ----- | |
| Arguably, the most justified way to perform binning is to use the defaults, | |
| with times, radial velocities and errors as input. This will lead to a | |
| weighted average of the times, a weighted average of the radial velocities | |
| and quadratically added errors. | |
| """ | |
| # "nightly" binning, based on the integer part of the time array | |
| intt = time.astype(int) | |
| _, indices = np.unique(intt, return_index=True) | |
| # include the last time in the array of bins, so we don't miss a point | |
| # but be careful: if it is alone, perturb the last time by a small ammount | |
| # so that the difference time[-1] - bins[-1] is not zero | |
| if indices[-1] == time.size - 1: | |
| bins = np.r_[time[indices], time[-1] + 1e-10] | |
| else: | |
| bins = np.r_[time[indices], time[-1]] | |
| # weighted mean | |
| wmean = lambda a, e: np.average(a, weights=1 / e**2) | |
| # bin the RVs | |
| if (err is not None) and (stat == 'wmean'): | |
| stat = wmean # default is weighted mean | |
| brv = binned_statistic(time, rv, statistic=stat, bins=bins, range=None, | |
| weights=err) | |
| # bin the times | |
| if (err is not None) and (tstat == 'wmean'): | |
| tstat = wmean # default is weighted mean | |
| times = binned_statistic(time, time, statistic=tstat, bins=bins, | |
| range=None, weights=err) | |
| # if there are errors, bin them too | |
| if err is not None: | |
| if estat == 'addquad': | |
| # this function adds the elements of `a` quadratically | |
| add_quadratically = lambda a: np.sqrt(np.add.reduce(a**2)) | |
| # this function then divides by the number of elements | |
| mean_quadratically = lambda a: add_quadratically(a) / a.size | |
| # the two previous functions are only separated for clarity | |
| # bin | |
| estat = mean_quadratically | |
| errors = binned_statistic(time, err, statistic=estat, bins=bins) | |
| else: | |
| errors = binned_statistic(time, err, statistic=estat, bins=bins) | |
| return times[0], brv[0], errors[0] | |
| else: | |
| return times[0], brv[0] |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment