Chris@87: from __future__ import division, absolute_import, print_function Chris@87: Chris@87: ''' Chris@87: >>> p = np.poly1d([1.,2,3]) Chris@87: >>> p Chris@87: poly1d([ 1., 2., 3.]) Chris@87: >>> print(p) Chris@87: 2 Chris@87: 1 x + 2 x + 3 Chris@87: >>> q = np.poly1d([3.,2,1]) Chris@87: >>> q Chris@87: poly1d([ 3., 2., 1.]) Chris@87: >>> print(q) Chris@87: 2 Chris@87: 3 x + 2 x + 1 Chris@87: >>> print(np.poly1d([1.89999+2j, -3j, -5.12345678, 2+1j])) Chris@87: 3 2 Chris@87: (1.9 + 2j) x - 3j x - 5.123 x + (2 + 1j) Chris@87: >>> print(np.poly1d([-3, -2, -1])) Chris@87: 2 Chris@87: -3 x - 2 x - 1 Chris@87: Chris@87: >>> p(0) Chris@87: 3.0 Chris@87: >>> p(5) Chris@87: 38.0 Chris@87: >>> q(0) Chris@87: 1.0 Chris@87: >>> q(5) Chris@87: 86.0 Chris@87: Chris@87: >>> p * q Chris@87: poly1d([ 3., 8., 14., 8., 3.]) Chris@87: >>> p / q Chris@87: (poly1d([ 0.33333333]), poly1d([ 1.33333333, 2.66666667])) Chris@87: >>> p + q Chris@87: poly1d([ 4., 4., 4.]) Chris@87: >>> p - q Chris@87: poly1d([-2., 0., 2.]) Chris@87: >>> p ** 4 Chris@87: poly1d([ 1., 8., 36., 104., 214., 312., 324., 216., 81.]) Chris@87: Chris@87: >>> p(q) Chris@87: poly1d([ 9., 12., 16., 8., 6.]) Chris@87: >>> q(p) Chris@87: poly1d([ 3., 12., 32., 40., 34.]) Chris@87: Chris@87: >>> np.asarray(p) Chris@87: array([ 1., 2., 3.]) Chris@87: >>> len(p) Chris@87: 2 Chris@87: Chris@87: >>> p[0], p[1], p[2], p[3] Chris@87: (3.0, 2.0, 1.0, 0) Chris@87: Chris@87: >>> p.integ() Chris@87: poly1d([ 0.33333333, 1. , 3. , 0. ]) Chris@87: >>> p.integ(1) Chris@87: poly1d([ 0.33333333, 1. , 3. , 0. ]) Chris@87: >>> p.integ(5) Chris@87: poly1d([ 0.00039683, 0.00277778, 0.025 , 0. , 0. , Chris@87: 0. , 0. , 0. ]) Chris@87: >>> p.deriv() Chris@87: poly1d([ 2., 2.]) Chris@87: >>> p.deriv(2) Chris@87: poly1d([ 2.]) Chris@87: Chris@87: >>> q = np.poly1d([1.,2,3], variable='y') Chris@87: >>> print(q) Chris@87: 2 Chris@87: 1 y + 2 y + 3 Chris@87: >>> q = np.poly1d([1.,2,3], variable='lambda') Chris@87: >>> print(q) Chris@87: 2 Chris@87: 1 lambda + 2 lambda + 3 Chris@87: Chris@87: >>> np.polydiv(np.poly1d([1,0,-1]), np.poly1d([1,1])) Chris@87: (poly1d([ 1., -1.]), poly1d([ 0.])) Chris@87: Chris@87: ''' Chris@87: import numpy as np Chris@87: from numpy.testing import ( Chris@87: run_module_suite, TestCase, assert_, assert_equal, assert_array_equal, Chris@87: assert_almost_equal, rundocs Chris@87: ) Chris@87: Chris@87: Chris@87: class TestDocs(TestCase): Chris@87: def test_doctests(self): Chris@87: return rundocs() Chris@87: Chris@87: def test_roots(self): Chris@87: assert_array_equal(np.roots([1, 0, 0]), [0, 0]) Chris@87: Chris@87: def test_str_leading_zeros(self): Chris@87: p = np.poly1d([4, 3, 2, 1]) Chris@87: p[3] = 0 Chris@87: assert_equal(str(p), Chris@87: " 2\n" Chris@87: "3 x + 2 x + 1") Chris@87: Chris@87: p = np.poly1d([1, 2]) Chris@87: p[0] = 0 Chris@87: p[1] = 0 Chris@87: assert_equal(str(p), " \n0") Chris@87: Chris@87: def test_polyfit(self): Chris@87: c = np.array([3., 2., 1.]) Chris@87: x = np.linspace(0, 2, 7) Chris@87: y = np.polyval(c, x) Chris@87: err = [1, -1, 1, -1, 1, -1, 1] Chris@87: weights = np.arange(8, 1, -1)**2/7.0 Chris@87: Chris@87: # check 1D case Chris@87: m, cov = np.polyfit(x, y+err, 2, cov=True) Chris@87: est = [3.8571, 0.2857, 1.619] Chris@87: assert_almost_equal(est, m, decimal=4) Chris@87: val0 = [[2.9388, -5.8776, 1.6327], Chris@87: [-5.8776, 12.7347, -4.2449], Chris@87: [1.6327, -4.2449, 2.3220]] Chris@87: assert_almost_equal(val0, cov, decimal=4) Chris@87: Chris@87: m2, cov2 = np.polyfit(x, y+err, 2, w=weights, cov=True) Chris@87: assert_almost_equal([4.8927, -1.0177, 1.7768], m2, decimal=4) Chris@87: val = [[8.7929, -10.0103, 0.9756], Chris@87: [-10.0103, 13.6134, -1.8178], Chris@87: [0.9756, -1.8178, 0.6674]] Chris@87: assert_almost_equal(val, cov2, decimal=4) Chris@87: Chris@87: # check 2D (n,1) case Chris@87: y = y[:, np.newaxis] Chris@87: c = c[:, np.newaxis] Chris@87: assert_almost_equal(c, np.polyfit(x, y, 2)) Chris@87: # check 2D (n,2) case Chris@87: yy = np.concatenate((y, y), axis=1) Chris@87: cc = np.concatenate((c, c), axis=1) Chris@87: assert_almost_equal(cc, np.polyfit(x, yy, 2)) Chris@87: Chris@87: m, cov = np.polyfit(x, yy + np.array(err)[:, np.newaxis], 2, cov=True) Chris@87: assert_almost_equal(est, m[:, 0], decimal=4) Chris@87: assert_almost_equal(est, m[:, 1], decimal=4) Chris@87: assert_almost_equal(val0, cov[:, :, 0], decimal=4) Chris@87: assert_almost_equal(val0, cov[:, :, 1], decimal=4) Chris@87: Chris@87: def test_objects(self): Chris@87: from decimal import Decimal Chris@87: p = np.poly1d([Decimal('4.0'), Decimal('3.0'), Decimal('2.0')]) Chris@87: p2 = p * Decimal('1.333333333333333') Chris@87: assert_(p2[1] == Decimal("3.9999999999999990")) Chris@87: p2 = p.deriv() Chris@87: assert_(p2[1] == Decimal('8.0')) Chris@87: p2 = p.integ() Chris@87: assert_(p2[3] == Decimal("1.333333333333333333333333333")) Chris@87: assert_(p2[2] == Decimal('1.5')) Chris@87: assert_(np.issubdtype(p2.coeffs.dtype, np.object_)) Chris@87: Chris@87: def test_complex(self): Chris@87: p = np.poly1d([3j, 2j, 1j]) Chris@87: p2 = p.integ() Chris@87: assert_((p2.coeffs == [1j, 1j, 1j, 0]).all()) Chris@87: p2 = p.deriv() Chris@87: assert_((p2.coeffs == [6j, 2j]).all()) Chris@87: Chris@87: def test_integ_coeffs(self): Chris@87: p = np.poly1d([3, 2, 1]) Chris@87: p2 = p.integ(3, k=[9, 7, 6]) Chris@87: assert_( Chris@87: (p2.coeffs == [1/4./5., 1/3./4., 1/2./3., 9/1./2., 7, 6]).all()) Chris@87: Chris@87: def test_zero_dims(self): Chris@87: try: Chris@87: np.poly(np.zeros((0, 0))) Chris@87: except ValueError: Chris@87: pass Chris@87: Chris@87: if __name__ == "__main__": Chris@87: run_module_suite()