wrapper to apply bottleneck moving window funcs on dask arrays
177798
def dask_rolling_wrapper(moving_func, a, window, min_count=None, axis=-1):
    ''''''
    dtype, fill_value = dtypes.maybe_promote(a.dtype)
    a = a.astype(dtype)
    # inputs for overlap
    if axis < 0:
        axis = a.ndim + axis
    depth = {d: 0 for d in range(a.ndim)}
    depth[axis] = (window + 1) // 2
    boundary = {d: fill_value for d in range(a.ndim)}
    # Create overlap array.
    ag = overlap(a, depth=depth, boundary=boundary)
    # apply rolling func
    out = ag.map_blocks(moving_func, window, min_count=min_count,
                        axis=axis, dtype=a.dtype)
    # trim array
    result = trim_internal(out, depth)
    return result
554152
def ffill(arr, dim=None, limit=None):
    ''''''
    import bottleneck as bn

    axis = arr.get_axis_num(dim)

    # work around for bottleneck 178
    _limit = limit if limit is not None else arr.shape[axis]

    return apply_ufunc(bn.push, arr,
                       dask='parallelized',
                       keep_attrs=True,
                       output_dtypes=[arr.dtype],
                       kwargs=dict(n=_limit, axis=axis)).transpose(*arr.dims)
864675
def gmdaZeros(shape, dtype, mask=None, numGhosts=1):
    
    res = GhostedMaskedDistArray(shape, dtype)
    res.mas = mask
    res.setNumberOfGhosts(numGhosts)
    res[:] = 0
    return res
668687
def _nanstd(array, axis=None, ddof=0):
    """"""

    if isinstance(axis, tuple):
        array = _move_tuple_axes_first(array, axis=axis)
        axis = 0
    return bottleneck.nanstd(array, axis=axis, ddof=ddof)
668968
def _nanmedian(array, axis=None):
    """"""

    if isinstance(axis, tuple):
        array = _move_tuple_axes_first(array, axis=axis)
        axis = 0
    return bottleneck.nanmedian(array, axis=axis)
934246
def _persist_inplace(self, **kwargs):
        
        # access .data to coerce everything to numpy or dask arrays
        lazy_data = {k: v._data for k, v in self.variables.items()
                     if isinstance(v._data, dask_array_type)}
        if lazy_data:
            import dask

            # evaluate all the dask arrays simultaneously
            evaluated_data = dask.persist(*lazy_data.values(), **kwargs)

            for k, data in zip(lazy_data, evaluated_data):
                self.variables[k].data = data

        return self
822350
def update_gradients_full(self, dL_dK, X, X2=None):
    #def dK_dtheta(self, dL_dK, X, X2, target):
        """"""
        X,slices = X[:,:-1],index_to_slices(X[:,-1])
        if X2 is None:
            X2,slices2 = X,slices
            K = np.zeros((X.shape[0], X.shape[0]))
        else:
            X2,slices2 = X2[:,:-1],index_to_slices(X2[:,-1])
        
        vyt = self.variance_Yt
        vyx = self.variance_Yx

        lyt = 1./(2*self.lengthscale_Yt)
        lyx = 1./(2*self.lengthscale_Yx)

        a = self.a
        b = self.b
        c = self.c

        tdist = (X[:,0][:,None] - X2[:,0][None,:])**2
        xdist = (X[:,1][:,None] - X2[:,1][None,:])**2
        #rdist = [tdist,xdist]
        ttdist = (X[:,0][:,None] - X2[:,0][None,:])
        
        rd=tdist.shape[0]

        dka = np.zeros([rd,rd])
        dkb = np.zeros([rd,rd])
        dkc = np.zeros([rd,rd])
        dkYdvart = np.zeros([rd,rd])
        dkYdvarx = np.zeros([rd,rd])
        dkYdlent = np.zeros([rd,rd])
        dkYdlenx = np.zeros([rd,rd])


        kyy = lambda tdist,xdist: np.exp(-lyt*(tdist) -lyx*(xdist))
        #k1 = lambda tdist: (lyt - lyt**2 * (tdist) )
        #k2 = lambda xdist: ( lyx**2 * (xdist)  - lyx )
        #k3 = lambda xdist: ( 3*lyx**2 - 6*xdist*lyx**3 + xdist**2*lyx**4 )
        #k4 = lambda tdist: -lyt*np.sqrt(tdist)

        k1 = lambda tdist: (2*lyt - 4*lyt**2 * (tdist) )

        k2 = lambda xdist: ( 4*lyx**2 * (xdist)  - 2*lyx )

        k3 = lambda xdist: ( 3*4*lyx**2 - 6*8*xdist*lyx**3 + 16*xdist**2*lyx**4 )

        k4 = lambda ttdist: 2*lyt*(ttdist)

        dkyydlyx = lambda tdist,xdist: kyy(tdist,xdist)*(-xdist)
        dkyydlyt = lambda tdist,xdist: kyy(tdist,xdist)*(-tdist)

        dk1dlyt = lambda tdist: 2. - 4*2.*lyt*tdist
        dk2dlyx = lambda xdist: (4.*2.*lyx*xdist -2.)
        dk3dlyx = lambda xdist: (6.*4.*lyx - 18.*8*xdist*lyx**2 + 4*16*xdist**2*lyx**3)

        dk4dlyt = lambda ttdist: 2*(ttdist)

        for i, s1 in enumerate(slices):
            for j, s2 in enumerate(slices2):
                for ss1 in s1:
                    for ss2 in s2:
                        if i==0 and j==0:
                            dka[ss1,ss2] = 0
                            dkb[ss1,ss2] = 0
                            dkc[ss1,ss2] = 0
                            dkYdvart[ss1,ss2] = vyx*kyy(tdist[ss1,ss2],xdist[ss1,ss2])
                            dkYdvarx[ss1,ss2] = vyt*kyy(tdist[ss1,ss2],xdist[ss1,ss2])
                            dkYdlenx[ss1,ss2] = vyt*vyx*dkyydlyx(tdist[ss1,ss2],xdist[ss1,ss2])
                            dkYdlent[ss1,ss2] = vyt*vyx*dkyydlyt(tdist[ss1,ss2],xdist[ss1,ss2])
                        elif i==0 and j==1:
                            dka[ss1,ss2] = -k2(xdist[ss1,ss2])*vyt*vyx*kyy(tdist[ss1,ss2],xdist[ss1,ss2])
                            dkb[ss1,ss2] = k4(ttdist[ss1,ss2])*vyt*vyx*kyy(tdist[ss1,ss2],xdist[ss1,ss2])
                            dkc[ss1,ss2] = vyt*vyx*kyy(tdist[ss1,ss2],xdist[ss1,ss2])
                            #dkYdvart[ss1,ss2] = 0
                            #dkYdvarx[ss1,ss2] = 0
                            #dkYdlent[ss1,ss2] = 0
                            #dkYdlenx[ss1,ss2] = 0
                            dkYdvart[ss1,ss2] = (-a*k2(xdist[ss1,ss2])+b*k4(ttdist[ss1,ss2])+c)*vyx*kyy(tdist[ss1,ss2],xdist[ss1,ss2])
                            dkYdvarx[ss1,ss2] = (-a*k2(xdist[ss1,ss2])+b*k4(ttdist[ss1,ss2])+c)*vyt*kyy(tdist[ss1,ss2],xdist[ss1,ss2])
                            dkYdlent[ss1,ss2] = vyt*vyx*dkyydlyt(tdist[ss1,ss2],xdist[ss1,ss2])* (-a*k2(xdist[ss1,ss2])+b*k4(ttdist[ss1,ss2])+c)+\
                            vyt*vyx*kyy(tdist[ss1,ss2],xdist[ss1,ss2])*b*dk4dlyt(ttdist[ss1,ss2])
                            dkYdlenx[ss1,ss2] = vyt*vyx*dkyydlyx(tdist[ss1,ss2],xdist[ss1,ss2])*(-a*k2(xdist[ss1,ss2])+b*k4(ttdist[ss1,ss2])+c)+\
                            vyt*vyx*kyy(tdist[ss1,ss2],xdist[ss1,ss2])*(-a*dk2dlyx(xdist[ss1,ss2]))
                        elif i==1 and j==1:
                            dka[ss1,ss2] = (2*a*k3(xdist[ss1,ss2]) - 2*c*k2(xdist[ss1,ss2]))*vyt*vyx* kyy(tdist[ss1,ss2],xdist[ss1,ss2])
                            dkb[ss1,ss2] = 2*b*k1(tdist[ss1,ss2])*vyt*vyx* kyy(tdist[ss1,ss2],xdist[ss1,ss2])
                            dkc[ss1,ss2] = (-2*a*k2(xdist[ss1,ss2]) + 2*c )*vyt*vyx* kyy(tdist[ss1,ss2],xdist[ss1,ss2])
                            dkYdvart[ss1,ss2] = ( b**2*k1(tdist[ss1,ss2]) - 2*a*c*k2(xdist[ss1,ss2]) + a**2*k3(xdist[ss1,ss2]) + c**2 )*vyx* kyy(tdist[ss1,ss2],xdist[ss1,ss2])
                            dkYdvarx[ss1,ss2] = ( b**2*k1(tdist[ss1,ss2]) - 2*a*c*k2(xdist[ss1,ss2]) + a**2*k3(xdist[ss1,ss2]) + c**2 )*vyt* kyy(tdist[ss1,ss2],xdist[ss1,ss2])
                            dkYdlent[ss1,ss2] = vyt*vyx*dkyydlyt(tdist[ss1,ss2],xdist[ss1,ss2])*( b**2*k1(tdist[ss1,ss2]) - 2*a*c*k2(xdist[ss1,ss2]) + a**2*k3(xdist[ss1,ss2]) + c**2 ) +\
                            vyx*vyt*kyy(tdist[ss1,ss2],xdist[ss1,ss2])*b**2*dk1dlyt(tdist[ss1,ss2])
                            dkYdlenx[ss1,ss2] = vyt*vyx*dkyydlyx(tdist[ss1,ss2],xdist[ss1,ss2])*( b**2*k1(tdist[ss1,ss2]) - 2*a*c*k2(xdist[ss1,ss2]) + a**2*k3(xdist[ss1,ss2]) + c**2 ) +\
                            vyx*vyt*kyy(tdist[ss1,ss2],xdist[ss1,ss2])* (-2*a*c*dk2dlyx(xdist[ss1,ss2]) + a**2*dk3dlyx(xdist[ss1,ss2]) )
                        else:
                            dka[ss1,ss2] = -k2(xdist[ss1,ss2])*vyt*vyx*kyy(tdist[ss1,ss2],xdist[ss1,ss2])
                            dkb[ss1,ss2] = -k4(ttdist[ss1,ss2])*vyt*vyx*kyy(tdist[ss1,ss2],xdist[ss1,ss2])
                            dkc[ss1,ss2] = vyt*vyx*kyy(tdist[ss1,ss2],xdist[ss1,ss2])
                            #dkYdvart[ss1,ss2] = 0
                            #dkYdvarx[ss1,ss2] = 0
                            #dkYdlent[ss1,ss2] = 0
                            #dkYdlenx[ss1,ss2] = 0
                            dkYdvart[ss1,ss2] = (-a*k2(xdist[ss1,ss2])-b*k4(ttdist[ss1,ss2])+c)*vyx*kyy(tdist[ss1,ss2],xdist[ss1,ss2])
                            dkYdvarx[ss1,ss2] = (-a*k2(xdist[ss1,ss2])-b*k4(ttdist[ss1,ss2])+c)*vyt*kyy(tdist[ss1,ss2],xdist[ss1,ss2])
                            dkYdlent[ss1,ss2] = vyt*vyx*dkyydlyt(tdist[ss1,ss2],xdist[ss1,ss2])* (-a*k2(xdist[ss1,ss2])-b*k4(ttdist[ss1,ss2])+c)+\
                            vyt*vyx*kyy(tdist[ss1,ss2],xdist[ss1,ss2])*(-1)*b*dk4dlyt(ttdist[ss1,ss2])
                            dkYdlenx[ss1,ss2] = vyt*vyx*dkyydlyx(tdist[ss1,ss2],xdist[ss1,ss2])*(-a*k2(xdist[ss1,ss2])-b*k4(ttdist[ss1,ss2])+c)+\
                            vyt*vyx*kyy(tdist[ss1,ss2],xdist[ss1,ss2])*(-a*dk2dlyx(xdist[ss1,ss2])) 

        self.a.gradient = np.sum(dka * dL_dK)  

        self.b.gradient = np.sum(dkb * dL_dK) 

        self.c.gradient = np.sum(dkc * dL_dK)


        self.variance_Yt.gradient = np.sum(dkYdvart * dL_dK)  # Vy

        self.variance_Yx.gradient = np.sum(dkYdvarx * dL_dK)

        self.lengthscale_Yt.gradient = np.sum(dkYdlent*(-0.5*self.lengthscale_Yt**(-2)) * dL_dK)    #ly np.sum(dktheta2*(-self.lengthscale_Y**(-2)) * dL_dK) 

        self.lengthscale_Yx.gradient =  np.sum(dkYdlenx*(-0.5*self.lengthscale_Yx**(-2)) * dL_dK)
1336737
def _bfill(arr, n=None, axis=-1):
    ''''''
    import bottleneck as bn

    arr = np.flip(arr, axis=axis)

    # fill
    arr = bn.push(arr, axis=axis, n=n)

    # reverse back to original
    return np.flip(arr, axis=axis)
89108
def vq_discrete_unbottleneck(x, hidden_size):
  """"""
  x_shape = common_layers.shape_list(x)
  x = tf.to_float(x)
  bottleneck_size = common_layers.shape_list(x)[-1]
  means, _, _ = get_vq_codebook(bottleneck_size, hidden_size)
  result = tf.matmul(tf.reshape(x, [-1, x_shape[-1]]), means)
  return tf.reshape(result, x_shape[:-1] + [hidden_size])
545658
def autostack(self):
    """"""
    num_slices = 0
    for v in self.graph.all_variables:
      num_slices += self.mesh_to_impl[v.mesh].size
    if num_slices >= 2 ** 16:
      # Startup times are slow with lots of variable slices.
      # Perform more aggressive stacking
      max_combined_slice_size = 2 ** 27
    else:
      # Stacking hurts memory utilization - only stack small variables.
      max_combined_slice_size = 2 ** 16
    self.graph.rewrite_stack_variables(
        mesh_to_impl=self.mesh_to_impl,
        max_combined_slice_size=max_combined_slice_size)
673454
def _arrayize_vectorized_indexer(indexer, shape):
    
    slices = [v for v in indexer.tuple if isinstance(v, slice)]
    if len(slices) == 0:
        return indexer

    arrays = [v for v in indexer.tuple if isinstance(v, np.ndarray)]
    n_dim = arrays[0].ndim if len(arrays) > 0 else 0
    i_dim = 0
    new_key = []
    for v, size in zip(indexer.tuple, shape):
        if isinstance(v, np.ndarray):
            new_key.append(np.reshape(v, v.shape + (1, ) * len(slices)))
        else:  # slice
            shape = ((1,) * (n_dim + i_dim) + (-1,) +
                     (1,) * (len(slices) - i_dim - 1))
            new_key.append(np.arange(*v.indices(size)).reshape(shape))
            i_dim += 1
    return VectorizedIndexer(tuple(new_key))
1363668
function window(data_array, windowing_function, alpha) {
	var datapoints = data_array.length;

	/* For each item in the array */
	for (var n=0; n<datapoints; ++n) {
		/* Apply the windowing function */
		data_array[n] *= windowing_function(n, datapoints, alpha);
	}

	return data_array;
}
220141
def multidispatch(*, nargs=None, nouts=None):
    
    def wrapper(f):
        return wraps(f)(MultiDispatchFunction(f, nargs=nargs, nouts=nouts))

    return wrapper
595917
def _exec_loop_moving_window(self, a_all, bd_all, mask, bd_idx):
        
        import scipy.linalg.lapack

        npt = bd_all.shape[0]
        n = bd_idx.shape[1]
        kvalues = np.zeros(npt)
        sigmasq = np.zeros(npt)

        for i in np.nonzero(~mask)[0]:
            b_selector = bd_idx[i]
            bd = bd_all[i]

            a_selector = np.concatenate((b_selector, np.array([a_all.shape[0] - 1])))
            a = a_all[a_selector[:, None], a_selector]

            if np.any(np.absolute(bd) <= self.eps):
                zero_value = True
                zero_index = np.where(np.absolute(bd) <= self.eps)
            else:
                zero_value = False
                zero_index = None
            b = np.zeros((n+1, 1))
            b[:n, 0] = - self.variogram_function(self.variogram_model_parameters, bd)
            if zero_value:
                b[zero_index[0], 0] = 0.0
            b[n, 0] = 1.0

            x = scipy.linalg.solve(a, b)

            kvalues[i] = x[:n, 0].dot(self.VALUES[b_selector])
            sigmasq[i] = - x[:, 0].dot(b[:, 0])

        return kvalues, sigmasq
267366
def precompute(self, cache_dir=None, swath_usage=0, **kwargs):
        """"""
        if kwargs.get('mask') is not None:
            LOG.warning("'mask' parameter has no affect during EWA "
                        "resampling")

        del kwargs
        source_geo_def = self.source_geo_def
        target_geo_def = self.target_geo_def

        if cache_dir:
            LOG.warning("'cache_dir' is not used by EWA resampling")

        # Satpy/PyResample don't support dynamic grids out of the box yet
        lons, lats = source_geo_def.get_lonlats()
        if isinstance(lons, xr.DataArray):
            # get dask arrays
            lons = lons.data
            lats = lats.data
        # we are remapping to a static unchanging grid/area with all of
        # its parameters specified
        chunks = (2,) + lons.chunks
        res = da.map_blocks(self._call_ll2cr, lons, lats,
                            target_geo_def, swath_usage,
                            dtype=lons.dtype, chunks=chunks, new_axis=[0])
        cols = res[0]
        rows = res[1]

        # save the dask arrays in the class instance cache
        # the on-disk cache will store the numpy arrays
        self.cache = {
            "rows": rows,
            "cols": cols,
        }

        return None
678100
def apply(self, func, *args, **kwds):
        
        # estimate time to pandas apply
        wrapped = self._wrapped_apply(func, *args, **kwds)
        n_repeats = 3
        timed = timeit.timeit(wrapped, number=n_repeats)
        samp_proc_est = timed / n_repeats
        est_apply_duration = samp_proc_est / self._SAMP_SIZE * self._nrows

        # if pandas apply takes too long, use dask
        if est_apply_duration > self._dask_threshold:
            return self._dask_apply(func, *args, **kwds)
        else:  # use pandas
            if self._progress_bar:
                tqdm.pandas(desc="Pandas Apply")
                return self._obj_pd.progress_apply(func, *args, **kwds)
            else:
                return self._obj_pd.apply(func, *args, **kwds)
1145520
def short_dask_repr(array, show_dtype=True):
    
    chunksize = tuple(c[0] for c in array.chunks)
    if show_dtype:
        return 'dask.array<shape={}, dtype={}, chunksize={}>'.format(
            array.shape, array.dtype, chunksize)
    else:
        return 'dask.array<shape={}, chunksize={}>'.format(
            array.shape, chunksize)