def main():
    Log.set_loglevel(logging.DEBUG)
    
    prior = Gaussian(Sigma=eye(2) * 100)
    
    posterior = OzonePosterior(prior, logdet_alg="scikits",
                               solve_method="scikits")
    
    proposal_cov = diag([ 4.000000000000000e-05, 1.072091680000000e+02])
    mcmc_sampler = StandardMetropolis(posterior, scale=1.0, cov=proposal_cov)
    
    start = asarray([-11.35, -13.1])
    mcmc_params = MCMCParams(start=start, num_iterations=5000)
    chain = MCMCChain(mcmc_sampler, mcmc_params)
    
    chain.append_mcmc_output(StatisticsOutput(print_from=1, lag=1))
    
    home = expanduser("~")
    folder = os.sep.join([home, "sample_ozone_posterior_average_serial"])
    store_chain_output = StoreChainOutput(folder)
    chain.append_mcmc_output(store_chain_output)
    
    loaded = store_chain_output.load_last_stored_chain()
    if loaded is None:
        logging.info("Running chain from scratch")
    else:
        logging.info("Running chain from iteration %d" % loaded.iteration)
        chain = loaded
        
    chain.run()
    
    f = open(folder + os.sep + "final_chain", "w")
    dump(chain, f)
    f.close()
Esempio n. 2
0
 def low_rank_approx(K, d):
     """
     Returns a low rank approximation factor L of the given psd matrix such that
     LL^T \approx K with a given number of principal components to use
 
     K - psd matrix to compute low-rank approximation of
     d - number of principal components to use
     
     returns (L, s, V) where
     L - LL^T \approx K
     s - 1D vector of Eigenvalues
     V - matrix containing Eigen-row-vectors
     """
     # perform SVD and only use first d components. Note that U^T=V if K psd and
     # rows of V are Eigenvectors of K
     U, s, V = svd(K)
     U = U[:, 0:d]
     V = V[0:d, :]
     s = s[0:d]
     S = diag(s)
     
     # K \approx=U.dot(S.dot(V))
     L = sqrt(S).dot(V)
     
     # LL^T \approx K
     return (L, s, V)
Esempio n. 3
0
 def fit(self, func, params_dist, pre_X = None, pre_y = None):
     time_start = time.time()
     if self.random_state is not None: np.random.seed(self.random_state)
     grid, grid_scaled = self.get_grid(params_dist)
     mu       = np.zeros(self.n_grid) + self.mu_prior
     sigma    = np.ones(self.n_grid)*self.sigma_prior
     X        = np.zeros((self.max_iter, self.n_params))
     X_scaled = np.matrix(np.zeros((self.max_iter, self.n_params)))
     y        = np.zeros(self.max_iter)
     if (pre_X is not None) and (pre_y is not None):
         pre_X_mat, pre_X_scaled = self.scale(pre_X, pre_y, params_dist)
         X        = np.vstack([pre_X_mat, X]) 
         X_scaled = np.vstack([pre_X_scaled, X_scaled])
         y        = np.concatenate([pre_y, y])
         pre_len  = len(pre_y)
     else: pre_len = 0
     if self.verbose:
         params_name = [i[:9] for i in self.params_name]
         logger.info('%4s|%9s|%9s|%9s', 'Iter','Func','Max',
                 '|'.join(['{:9s}'.format(i) for i in params_name]))
     for i in xrange(pre_len, pre_len + self.max_iter):
         #beta        = (i + 1)**2
         if self.beta_mode == 'log':
             d       = len(self.params_name)
             beta    = 2*np.log(2 *(i + 1)**2 * np.pi**2 /.3) + \
                       2*d*np.log( (i+1)**2 * d * 2)
         elif self.beta_mode == 'linear':
             beta    = i + 1
         elif self.beta_mode == 'square':
             beta    = (i + 1)**2
         else:
             logger.error("What The Hell. Change Beta Parameter")
         idx         = np.argmax(mu + np.sqrt(beta)*sigma)
         X[i,:]      = grid[idx]
         X_scaled[i] = grid_scaled[idx]
         y[i]        = func(**dict(zip(self.params_name, X[i])))
         KT          = self.kernel(X_scaled[:(i + 1)], X_scaled[:(i + 1)])*\
                         self.sigma_prior
         invKT       = inv(KT + self.sig**2*identity(i + 1))
         grid, grid_scaled = self.get_grid(params_dist)
         kT          = self.kernel(X_scaled[:(i + 1)], grid_scaled)*\
                       self.sigma_prior**2
         mu          = self.mu_prior + \
                       kT.T.dot(invKT).dot(y[:(i + 1)] - self.mu_prior)
         sigma2      = np.ones(self.n_grid)*self.sigma_prior**2 - \
                         diag(kT.T.dot(invKT).dot(kT))
         sigma       = np.sqrt(sigma2)
         ### Save Data
         if self.verbose:
             logger.info('%4d|%9.4g|%9.4g|%s', i, y[i], np.max(y[:(i + 1)]),
                     '|'.join(['{:9.4g}'.format(ii) for ii in X[i]]))
         if time.time() - time_start > self.time_budget:
             break
     self.X      = X[:(i + 1)]
     self.y      = y[:(i + 1)]
     self.mu     = mu
     self.beta   = beta
     self.sigma  = sigma
     self.grid   = grid
def corrcoef(x, y=None, rowvar=1, bias=0):
    """The correlation coefficients
    """
    c = cov(x, y, rowvar, bias)
    try:
        d = diag(c)
    except ValueError: # scalar covariance
        return 1
    return c/sqrt(multiply.outer(d,d))
def main():
    Log.set_loglevel(logging.DEBUG)
    
    prior = Gaussian(Sigma=eye(2) * 100)
    num_estimates = 1000
    
    home = expanduser("~")
    folder = os.sep.join([home, "sample_ozone_posterior_rr_sge"])
    
    # cluster admin set project jump for me to exclusively allocate nodes
    parameter_prefix = ""  # #$ -P jump"
    
    cluster_parameters = BatchClusterParameters(foldername=folder,
                                            memory=7.8,
                                            loglevel=logging.DEBUG,
                                            parameter_prefix=parameter_prefix,
                                            max_walltime=60 * 60 * 24 - 1)
        
    computation_engine = SGEComputationEngine(cluster_parameters, check_interval=10)
    
    rr_instance = RussianRoulette(1e-3, block_size=400)
    
    posterior = OzonePosteriorRREngine(rr_instance=rr_instance,
                                       computation_engine=computation_engine,
                                       num_estimates=num_estimates,
                                       prior=prior)
    
    posterior.logdet_method = "shogun_estimate"
    
    proposal_cov = diag([ 4.000000000000000e-05, 1.072091680000000e+02])
    mcmc_sampler = StandardMetropolis(posterior, scale=1.0, cov=proposal_cov)
    
    start = asarray([-11.55, -10.1])
    mcmc_params = MCMCParams(start=start, num_iterations=5000)
    chain = MCMCChain(mcmc_sampler, mcmc_params)
    
#    chain.append_mcmc_output(PlottingOutput(None, plot_from=1, lag=1))
    chain.append_mcmc_output(StatisticsOutput(print_from=1, lag=1))
    
    store_chain_output = StoreChainOutput(folder, lag=1)
    chain.append_mcmc_output(store_chain_output)
    
    loaded = store_chain_output.load_last_stored_chain()
    if loaded is None:
        logging.info("Running chain from scratch")
    else:
        logging.info("Running chain from iteration %d" % loaded.iteration)
        chain = loaded
        
    chain.run()
    
    f = open(folder + os.sep + "final_chain", "w")
    dump(chain, f)
    f.close()
Esempio n. 6
0
 def log_pdf_at_quantile(self, alphas):
     """
     Computes the log-pdf at a given 1d-vector of quantiles
     """
     chi2_instance = chi2(self.dimension)
     cuttoffs = chi2_instance.isf(1 - alphas)
     
     log_determinant_part = -sum(log(diag(self.L)))
     quadratic_part = -0.5 * cuttoffs
     const_part = -0.5 * len(self.L) * log(2 * pi)
     
     return const_part + log_determinant_part + quadratic_part
Esempio n. 7
0
        def calc_error(xi, sef):
            Xk = self._get_X(xi).T

            covarM = matrix(self.var_covar)
            varY_hat = (Xk.T * covarM * Xk)
            varY_hat = sum(diag(varY_hat))
            if error_calc == 'sem':
                se = sef * sqrt(varY_hat)
            else:
                se = sqrt(sef ** 2 + sef ** 2 * varY_hat)

            return se
def main():
    Log.set_loglevel(logging.DEBUG)
    
    modulename = "sample_ozone_posterior_average_slurm"
    if not FileSystem.cmd_exists("sbatch"):
        engine = SerialComputationEngine()
    else:
        johns_slurm_hack = "#SBATCH --partition=intel-ivy,wrkstn,compute"
        johns_slurm_hack = "#SBATCH --partition=intel-ivy,compute"
        
        folder = os.sep + os.sep.join(["nfs", "data3", "ucabhst", modulename])
        batch_parameters = BatchClusterParameters(foldername=folder, max_walltime=24 * 60 * 60,
                                                  resubmit_on_timeout=False, memory=3,
                                                  parameter_prefix=johns_slurm_hack)
        engine = SlurmComputationEngine(batch_parameters, check_interval=1,
                                do_clean_up=True)
    
    
    prior = Gaussian(Sigma=eye(2) * 100)
    num_estimates = 100
    
    posterior = OzonePosteriorAverageEngine(computation_engine=engine,
                                        num_estimates=num_estimates,
                                        prior=prior)
    posterior.logdet_method = "shogun_estimate"
    
    proposal_cov = diag([ 4.000000000000000e-05, 1.072091680000000e+02])
    mcmc_sampler = StandardMetropolis(posterior, scale=1.0, cov=proposal_cov)
    
    start = asarray([-11.35, -13.1])
    mcmc_params = MCMCParams(start=start, num_iterations=2000)
    chain = MCMCChain(mcmc_sampler, mcmc_params)
    
    chain.append_mcmc_output(StatisticsOutput(print_from=1, lag=1))
    
    home = expanduser("~")
    folder = os.sep.join([home, modulename])
    store_chain_output = StoreChainOutput(folder)
    chain.append_mcmc_output(store_chain_output)
    
    loaded = store_chain_output.load_last_stored_chain()
    if loaded is None:
        logging.info("Running chain from scratch")
    else:
        logging.info("Running chain from iteration %d" % loaded.iteration)
        chain = loaded
        
    chain.run()
    
    f = open(folder + os.sep + "final_chain", "w")
    dump(chain, f)
    f.close()
Esempio n. 9
0
 def emp_quantiles(self, X, quantiles=arange(0.1, 1, 0.1)):
     # need inverse chi2 cdf with self.dimension degrees of freedom
     chi2_instance = chi2(self.dimension)
     cutoffs = chi2_instance.isf(1 - quantiles)
     # whitening
     D, U = eig(self.L.dot(self.L.T))
     D = D ** (-0.5)
     W = (diag(D).dot(U.T).dot((X - self.mu).T)).T
     norms_squared = array([norm(w) ** 2 for w in W])
     results = zeros([len(quantiles)])
     for jj in range(0, len(quantiles)):
         results[jj] = mean(norms_squared < cutoffs[jj])
     return results
Esempio n. 10
0
 def log_pdf(self, X):
     assert(len(shape(X)) == 2)
     assert(shape(X)[1] == self.dimension)
     
     log_determinant_part = -sum(log(diag(self.L)))
     
     quadratic_parts = zeros(len(X))
     for i in range(len(X)):
         x = X[i] - self.mu
         
         # solve y=K^(-1)x = L^(-T)L^(-1)x
         y = solve_triangular(self.L, x.T, lower=True)
         y = solve_triangular(self.L.T, y, lower=False)
         quadratic_parts[i] = -0.5 * x.dot(y)
         
     const_part = -0.5 * len(self.L) * log(2 * pi)
     
     return const_part + log_determinant_part + quadratic_parts
def main():
    Log.set_loglevel(logging.DEBUG)

    prior = Gaussian(Sigma=eye(2) * 100)
    num_estimates = 2

    home = expanduser("~")
    folder = os.sep.join([home, "sample_ozone_posterior_rr_sge"])

    computation_engine = SerialComputationEngine()

    rr_instance = RussianRoulette(1e-3, block_size=10)

    posterior = OzonePosteriorRREngine(
        rr_instance=rr_instance, computation_engine=computation_engine, num_estimates=num_estimates, prior=prior
    )

    posterior.logdet_method = "shogun_estimate"

    proposal_cov = diag([4.000000000000000e-05, 1.072091680000000e02])
    mcmc_sampler = StandardMetropolis(posterior, scale=1.0, cov=proposal_cov)

    start = asarray([-11.35, -13.1])
    mcmc_params = MCMCParams(start=start, num_iterations=200)
    chain = MCMCChain(mcmc_sampler, mcmc_params)

    #    chain.append_mcmc_output(PlottingOutput(None, plot_from=1, lag=1))
    chain.append_mcmc_output(StatisticsOutput(print_from=1, lag=1))

    store_chain_output = StoreChainOutput(folder, lag=50)
    chain.append_mcmc_output(store_chain_output)

    loaded = store_chain_output.load_last_stored_chain()
    if loaded is None:
        logging.info("Running chain from scratch")
    else:
        logging.info("Running chain from iteration %d" % loaded.iteration)
        chain = loaded

    chain.run()

    f = open(folder + os.sep + "final_chain", "w")
    dump(chain, f)
    f.close()
Esempio n. 12
0
def roots(p):
    """ Return the roots of the polynomial coefficients in p.

        The values in the rank-1 array p are coefficients of a polynomial.
        If the length of p is n+1 then the polynomial is
        p[0] * x**n + p[1] * x**(n-1) + ... + p[n-1]*x + p[n]
    """
    # If input is scalar, this makes it an array
    p = atleast_1d(p)
    if len(p.shape) != 1:
        raise ValueError,"Input must be a rank-1 array."

    # find non-zero array entries
    non_zero = NX.nonzero(NX.ravel(p))[0]

    # Return an empty array if polynomial is all zeros
    if len(non_zero) == 0:
        return NX.array([])

    # find the number of trailing zeros -- this is the number of roots at 0.
    trailing_zeros = len(p) - non_zero[-1] - 1

    # strip leading and trailing zeros
    p = p[int(non_zero[0]):int(non_zero[-1])+1]

    # casting: if incoming array isn't floating point, make it floating point.
    if not issubclass(p.dtype.type, (NX.floating, NX.complexfloating)):
        p = p.astype(float)

    N = len(p)
    if N > 1:
        # build companion matrix and find its eigenvalues (the roots)
        A = diag(NX.ones((N-2,), p.dtype), -1)
        A[0, :] = -p[1:] / p[0]
        roots = _eigvals(A)
    else:
        roots = NX.array([])

    # tack any zeros onto the back of the array
    roots = hstack((roots, NX.zeros(trailing_zeros, roots.dtype)))
    return roots
Esempio n. 13
0
def roots(p):
    """
    Return the roots of a polynomial with coefficients given in p.

    The values in the rank-1 array `p` are coefficients of a polynomial.
    If the length of `p` is n+1 then the polynomial is described by::

      p[0] * x**n + p[1] * x**(n-1) + ... + p[n-1]*x + p[n]

    Parameters
    ----------
    p : array_like
        Rank-1 array of polynomial coefficients.

    Returns
    -------
    out : ndarray
        An array containing the complex roots of the polynomial.

    Raises
    ------
    ValueError :
        When `p` cannot be converted to a rank-1 array.

    See also
    --------
    poly : Find the coefficients of a polynomial with a given sequence
           of roots.
    polyval : Evaluate a polynomial at a point.
    polyfit : Least squares polynomial fit.
    poly1d : A one-dimensional polynomial class.

    Notes
    -----
    The algorithm relies on computing the eigenvalues of the
    companion matrix [1]_.

    References
    ----------
    .. [1] R. A. Horn & C. R. Johnson, *Matrix Analysis*.  Cambridge, UK:
        Cambridge University Press, 1999, pp. 146-7.

    Examples
    --------
    >>> coeff = [3.2, 2, 1]
    >>> np.roots(coeff)
    array([-0.3125+0.46351241j, -0.3125-0.46351241j])

    """
    # If input is scalar, this makes it an array
    p = atleast_1d(p)
    if len(p.shape) != 1:
        raise ValueError("Input must be a rank-1 array.")

    # find non-zero array entries
    non_zero = NX.nonzero(NX.ravel(p))[0]

    # Return an empty array if polynomial is all zeros
    if len(non_zero) == 0:
        return NX.array([])

    # find the number of trailing zeros -- this is the number of roots at 0.
    trailing_zeros = len(p) - non_zero[-1] - 1

    # strip leading and trailing zeros
    p = p[int(non_zero[0]):int(non_zero[-1])+1]

    # casting: if incoming array isn't floating point, make it floating point.
    if not issubclass(p.dtype.type, (NX.floating, NX.complexfloating)):
        p = p.astype(float)

    N = len(p)
    if N > 1:
        # build companion matrix and find its eigenvalues (the roots)
        A = diag(NX.ones((N-2,), p.dtype), -1)
        A[0, :] = -p[1:] / p[0]
        roots = eigvals(A)
    else:
        roots = NX.array([])

    # tack any zeros onto the back of the array
    roots = hstack((roots, NX.zeros(trailing_zeros, roots.dtype)))
    return roots
Esempio n. 14
0
def roots(p):
    """
    Return the roots of a polynomial with coefficients given in p.

    The values in the rank-1 array `p` are coefficients of a polynomial.
    If the length of `p` is n+1 then the polynomial is described by
    p[0] * x**n + p[1] * x**(n-1) + ... + p[n-1]*x + p[n]

    Parameters
    ----------
    p : array_like of shape(M,)
        Rank-1 array of polynomial co-efficients.

    Returns
    -------
    out : ndarray
        An array containing the complex roots of the polynomial.

    Raises
    ------
    ValueError:
        When `p` cannot be converted to a rank-1 array.

    Examples
    --------

    >>> coeff = [3.2, 2, 1]
    >>> print np.roots(coeff)
    [-0.3125+0.46351241j -0.3125-0.46351241j]

    """
    # If input is scalar, this makes it an array
    p = atleast_1d(p)
    if len(p.shape) != 1:
        raise ValueError, "Input must be a rank-1 array."

    # find non-zero array entries
    non_zero = NX.nonzero(NX.ravel(p))[0]

    # Return an empty array if polynomial is all zeros
    if len(non_zero) == 0:
        return NX.array([])

    # find the number of trailing zeros -- this is the number of roots at 0.
    trailing_zeros = len(p) - non_zero[-1] - 1

    # strip leading and trailing zeros
    p = p[int(non_zero[0]):int(non_zero[-1]) + 1]

    # casting: if incoming array isn't floating point, make it floating point.
    if not issubclass(p.dtype.type, (NX.floating, NX.complexfloating)):
        p = p.astype(float)

    N = len(p)
    if N > 1:
        # build companion matrix and find its eigenvalues (the roots)
        A = diag(NX.ones((N - 2, ), p.dtype), -1)
        A[0, :] = -p[1:] / p[0]
        roots = eigvals(A)
    else:
        roots = NX.array([])

    # tack any zeros onto the back of the array
    roots = hstack((roots, NX.zeros(trailing_zeros, roots.dtype)))
    return roots
Esempio n. 15
0
def GPUCB(func = f, kernel = DoubleExponential,
        params_dist = {'x': Uniform(start = 0, end = 5)},
        prev_X = None, prev_y = None,
          sig = .1, mu_prior = 0, sigma_prior = 1, n_grid = 100, n_iter = 10,
          seed = 2, time_budget = 36000):
    time_start = time.time()
    np.random.seed(seed)
    n_params          = len(params_dist)
    params_name       = params_dist.keys()
    grid, grid_scaled = GetRandGrid(n_grid, params_dist)
    mu = np.zeros(n_grid) + mu_prior
    sigma = np.ones(n_grid)*sigma_prior
    X = np.empty((n_iter, n_params))
    X_scaled = np.matrix(np.empty((n_iter, n_params)))
    y = np.empty(n_iter)
    logger.info("%4s |%9s |%9s |%s", "Iter", "Func", "Max",
            '|'.join(['{:6s}'.format(i) for i in params_name]))
    for i in xrange(n_iter):
        #beta = 2*np.log((i+1)**2*2*np.pi**2/3/.1) + \
        #       2*n_params*np.log((i+1)**2*n_params)
        beta = (i+1)**2
        #ipdb.set_trace()
        idx = np.argmax(mu + np.sqrt(beta)*sigma)
        X[i,:] = grid[idx]
        X_scaled[i] = grid_scaled[idx]
        y[i] = func(**dict(zip(params_name, X[i])))
        invKT = inv(kernel(X_scaled[:i+1], X_scaled[:i+1])*sigma_prior**2 + 
                    sig**2*identity(i + 1))
        grid, grid_scaled = GetRandGrid(n_grid, params_dist)
        kT = kernel(X_scaled[:i+1], grid_scaled)*sigma_prior**2
        mu = mu_prior + kT.T.dot(invKT).dot(y[:i+1] - mu_prior)
        sigma2 = np.ones(n_grid)*sigma_prior**2 - diag(kT.T.dot(invKT).dot(kT))
        sigma = np.sqrt(sigma2)

        logger.info("%4d |%9.4g |%9.4g |%s" , i, y[i], np.max(y[:i+1]),
                '|'.join(['{:6.2g}'.format(i) for i in X[i]]))
        if time.time() - time_start > time_budget:
            break
        ipdb.set_trace()
        if True:
            figure(1); plt.clf(); xlim((0,5)); ylim(-4,10);
            index = np.argsort(grid[:,0])
            gr = grid[:,0]
            plot(gr[index], mu[index], color = 'red', label = "Mean")
            plot(gr[index], mu[index] + sigma[index], color = 'blue', 
                    label = "Mean + Sigma")
            plot(gr[index], mu[index] - sigma[index], color = 'blue',
                    label = "Mean - Sigma")
            plot(X[:i+1,0], y[:i+1], 'o', color = 'green', label = "Eval Points")
            plot(np.linspace(0,5, num = 500),func(np.linspace(0,5, num = 500)),
                    color = 'green', label = "True Func")
            plot(gr[index], mu[index] + sqrt(beta)*sigma[index], 
                    color = 'yellow', label = "Mean + sqrt(B)*Sigma")
            plt.grid()
            legend(loc = 2)
            show()
    return {'X': X,
            'y': y,
            'mu': mu,
            'beta': beta,
            'sigma': sigma,
            'grid': grid}
Esempio n. 16
0
def roots(p):
    """
    Return the roots of a polynomial with coefficients given in p.

    The values in the rank-1 array `p` are coefficients of a polynomial.
    If the length of `p` is n+1 then the polynomial is described by
    p[0] * x**n + p[1] * x**(n-1) + ... + p[n-1]*x + p[n]

    Parameters
    ----------
    p : array_like of shape(M,)
        Rank-1 array of polynomial co-efficients.

    Returns
    -------
    out : ndarray
        An array containing the complex roots of the polynomial.

    Raises
    ------
    ValueError:
        When `p` cannot be converted to a rank-1 array.

    Examples
    --------

    >>> coeff = [3.2, 2, 1]
    >>> print np.roots(coeff)
    [-0.3125+0.46351241j -0.3125-0.46351241j]

    """
    # If input is scalar, this makes it an array
    p = atleast_1d(p)
    if len(p.shape) != 1:
        raise ValueError,"Input must be a rank-1 array."

    # find non-zero array entries
    non_zero = NX.nonzero(NX.ravel(p))[0]

    # Return an empty array if polynomial is all zeros
    if len(non_zero) == 0:
        return NX.array([])

    # find the number of trailing zeros -- this is the number of roots at 0.
    trailing_zeros = len(p) - non_zero[-1] - 1

    # strip leading and trailing zeros
    p = p[int(non_zero[0]):int(non_zero[-1])+1]

    # casting: if incoming array isn't floating point, make it floating point.
    if not issubclass(p.dtype.type, (NX.floating, NX.complexfloating)):
        p = p.astype(float)

    N = len(p)
    if N > 1:
        # build companion matrix and find its eigenvalues (the roots)
        A = diag(NX.ones((N-2,), p.dtype), -1)
        A[0, :] = -p[1:] / p[0]
        roots = eigvals(A)
    else:
        roots = NX.array([])

    # tack any zeros onto the back of the array
    roots = hstack((roots, NX.zeros(trailing_zeros, roots.dtype)))
    return roots
def roots(p):
    """
    Return the roots of a polynomial with coefficients given in p.

    The values in the rank-1 array `p` are coefficients of a polynomial.
    If the length of `p` is n+1 then the polynomial is described by::

      p[0] * x**n + p[1] * x**(n-1) + ... + p[n-1]*x + p[n]

    Parameters
    ----------
    p : array_like
        Rank-1 array of polynomial coefficients.

    Returns
    -------
    out : ndarray
        An array containing the complex roots of the polynomial.

    Raises
    ------
    ValueError :
        When `p` cannot be converted to a rank-1 array.

    See also
    --------
    poly : Find the coefficients of a polynomial with a given sequence
           of roots.
    polyval : Evaluate a polynomial at a point.
    polyfit : Least squares polynomial fit.
    poly1d : A one-dimensional polynomial class.

    Notes
    -----
    The algorithm relies on computing the eigenvalues of the
    companion matrix [1]_.

    References
    ----------
    .. [1] R. A. Horn & C. R. Johnson, *Matrix Analysis*.  Cambridge, UK:
        Cambridge University Press, 1999, pp. 146-7.

    Examples
    --------
    >>> coeff = [3.2, 2, 1]
    >>> np.roots(coeff)
    array([-0.3125+0.46351241j, -0.3125-0.46351241j])

    """
    # If input is scalar, this makes it an array
    p = atleast_1d(p)
    if len(p.shape) != 1:
        raise ValueError, "Input must be a rank-1 array."

    # find non-zero array entries
    non_zero = NX.nonzero(NX.ravel(p))[0]

    # Return an empty array if polynomial is all zeros
    if len(non_zero) == 0:
        return NX.array([])

    # find the number of trailing zeros -- this is the number of roots at 0.
    trailing_zeros = len(p) - non_zero[-1] - 1

    # strip leading and trailing zeros
    p = p[int(non_zero[0]):int(non_zero[-1]) + 1]

    # casting: if incoming array isn't floating point, make it floating point.
    if not issubclass(p.dtype.type, (NX.floating, NX.complexfloating)):
        p = p.astype(float)

    N = len(p)
    if N > 1:
        # build companion matrix and find its eigenvalues (the roots)
        A = diag(NX.ones((N - 2, ), p.dtype), -1)
        A[0, :] = -p[1:] / p[0]
        roots = eigvals(A)
    else:
        roots = NX.array([])

    # tack any zeros onto the back of the array
    roots = hstack((roots, NX.zeros(trailing_zeros, roots.dtype)))
    return roots