From e5e1cfce6c41421bb86ea6a70d0e58a19e136a69 Mon Sep 17 00:00:00 2001 From: Justin Smits Date: Sat, 9 Jul 2022 23:45:03 -0400 Subject: [PATCH] Added polynomial class and its tests --- AUTHORS | 1 + CHANGES | 6 + pint/polynomial.py | 201 +++++++++++++++++++++++++++ pint/testsuite/test_polynomial.py | 224 ++++++++++++++++++++++++++++++ 4 files changed, 432 insertions(+) create mode 100644 pint/polynomial.py create mode 100644 pint/testsuite/test_polynomial.py diff --git a/AUTHORS b/AUTHORS index e74dc6744..23dd221d8 100644 --- a/AUTHORS +++ b/AUTHORS @@ -33,6 +33,7 @@ Other contributors, listed alphabetically, are: * John David Reaver * Jonas Olson * Jules Chéron +* Justin Smits * Kaido Kert * Kenneth D. Mankoff * Kevin Davies diff --git a/CHANGES b/CHANGES index 41a0bbd10..ec7ff6ee7 100644 --- a/CHANGES +++ b/CHANGES @@ -1,6 +1,12 @@ Pint Changelog ============== +0.20.1 (unreleased) +------------------ + +- Added Polynomial class that inherits from numpy's polynomial + class and incorporates unit for the x and y variables. + 0.20 (unreleased) ------------------ diff --git a/pint/polynomial.py b/pint/polynomial.py new file mode 100644 index 000000000..0d501b58f --- /dev/null +++ b/pint/polynomial.py @@ -0,0 +1,201 @@ +""" + pint.polynomial + ~~~~~~~~~~~~~~ + + A polynomial class inheriting from numpy.polynomial.polynomial.Polynomial incorporating the pint Quantity unit values. + + :copyright: 2016 by Pint Authors, see AUTHORS for more details. + :license: BSD, see LICENSE for more details. +""" + +from __future__ import annotations + +from math import inf +from typing import Optional, Union + +from numpy.polynomial import polynomial as p + +from . import _DEFAULT_REGISTRY, Quantity, Unit + + +class Polynomial(p.Polynomial): + def __init__( + self, + coef: list[float], + x_unit: Unit = _DEFAULT_REGISTRY.dimensionless, + y_unit: Unit = _DEFAULT_REGISTRY.dimensionless, + ): + super(Polynomial, self).__init__(coef) + self.x_unit = x_unit + self.y_unit = y_unit + + @property + def y_intercept(self) -> Quantity: + return self.coef[0] * self.y_unit + + @staticmethod + def _get_roots_of_polynomial( + poly: p.Polynomial, + x_min: float = -inf, + x_max: float = inf, + real_only: bool = True, + ): + roots: set = set(poly.roots()) + if real_only: + roots = {float(root.real) for root in roots if root == root.real} + roots: list = list({root for root in roots if x_min <= float(root) <= x_max}) + return roots + + def _x_at_y( + self, + y_value: float, + x_min: float = -inf, + x_max: float = inf, + real_only: bool = True, + ) -> list[Union[float, complex]]: + solutions = self._get_roots_of_polynomial( + (self - y_value), x_min, x_max, real_only + ) + if len(solutions) == 1: + return solutions[0] + return solutions + + @property + def real_roots(self) -> list[float]: + return self._get_roots_of_polynomial(self, real_only=True) + + @property + def positive_real_roots(self) -> list[float]: + return self._get_roots_of_polynomial(self, x_min=0, real_only=True) + + @property + def x_intercept(self) -> Optional[Quantity]: + roots = self.real_roots + if len(roots) == 0: + return None + else: + root = min(roots) + return root * self.x_unit + + def solve(self, value: Quantity, min_value: float = -inf) -> Quantity: + x = value.m_as(self.x_unit) + return max(self(x), min_value) * self.y_unit + + def solve_for_x(self, y_value: Quantity, min_value: float = -inf) -> Quantity: + y = y_value.m_as(self.y_unit) + try: + return self._x_at_y(y, x_min=min_value, real_only=True) * self.x_unit + except TypeError: + raise ValueError( + "There are no values of {} that output {}!".format(self, y_value) + ) + + def integ(self, m=1, k=None, lbnd=None) -> "Polynomial": + k = k if k is not None else [] + return self.__class__( + super(Polynomial, self).integ(m, k, lbnd).coef, + self.x_unit, + self.y_unit * self.x_unit**m, + ) + + @property + def integral(self) -> "Polynomial": + return self.integ(1) + + def deriv(self, m=1) -> "Polynomial": + return self.__class__( + super(Polynomial, self).deriv(m).coef, + self.x_unit, + self.y_unit * self.x_unit**-m, + ) + + @property + def derivative(self) -> "Polynomial": + return self.deriv(1) + + def derivative_at(self, x: Quantity, derivative_order: int = 1) -> Quantity: + return self.deriv(derivative_order).solve(x) + + def __pow__(self, power, modulo=None): + x_unit, y_unit = self.x_unit, self.y_unit + if isinstance(power, self.__class__): + x_unit **= power.x_unit + y_unit **= power.y_unit + new_coefficients = p.polypow(self.coef, power) + else: + new_coefficients = p.polypow( + self.coef, super(Polynomial, self)._get_coefficients(power) + ) + + new_poly = sum(map(self.__class__, new_coefficients)) + return self.__class__(new_poly.coef, x_unit, y_unit) + + def __truediv__(self, other) -> "Polynomial": + x_unit, y_unit = self.x_unit, self.y_unit + if isinstance(other, self.__class__): + x_unit /= other.x_unit + y_unit /= other.y_unit + new_coefficients = p.polydiv(self.coef, other.coef) + else: + new_coefficients = p.polydiv( + self.coef, super(Polynomial, self)._get_coefficients(other) + ) + + new_poly = sum(map(self.__class__, new_coefficients)) + return self.__class__(new_poly.coef, x_unit, y_unit) + + def __mul__(self, other) -> "Polynomial": + if isinstance(other, self.__class__): + return self.__class__( + super(Polynomial, self).__mul__(other).coef, + self.x_unit * other.x_unit, + self.y_unit * other.y_unit, + ) + return self.__class__( + super(Polynomial, self).__mul__(other).coef, self.x_unit, self.y_unit + ) + + def __add__(self, other) -> "Polynomial": + self._polynomials_have_compatible_units(other) + return self.__class__( + super(Polynomial, self).__add__(other).coef, self.x_unit, self.y_unit + ) + + def __sub__(self, other) -> "Polynomial": + self._polynomials_have_compatible_units(other) + return self.__class__( + super(Polynomial, self).__sub__(other).coef, self.x_unit, self.y_unit + ) + + def __neg__(self) -> "Polynomial": + return self * -1 + + def __rtruediv__(self, other) -> "Polynomial": + if isinstance(other, self.__class__): + return other.__truediv__(self) + return super(Polynomial, self).__rtruediv__(other) + + def __rmul__(self, other) -> "Polynomial": + return self * other + + def __radd__(self, other) -> "Polynomial": + return self + other + + def __rsub__(self, other) -> "Polynomial": + return -self + other + + def _polynomials_have_compatible_units(self, other): + if not isinstance(other, self.__class__): + return + if self.x_unit == other.x_unit and self.y_unit == other.y_unit: + return + raise TypeError( + "Units between {} 1 ({}, {}) {} 2 ({}, {}) are not compatible".format( + self.__class__.__name__, + self.x_unit, + self.y_unit, + other.__class__.__name__, + other.x_unit, + other.y_unit, + ) + ) diff --git a/pint/testsuite/test_polynomial.py b/pint/testsuite/test_polynomial.py new file mode 100644 index 000000000..7c5355500 --- /dev/null +++ b/pint/testsuite/test_polynomial.py @@ -0,0 +1,224 @@ +import unittest +from math import sqrt + +from pint.testsuite import QuantityTestCase, helpers + +from ..polynomial import Polynomial + + +@helpers.requires_numpy +class TestPolynomial(QuantityTestCase, unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + + polys: list[Polynomial] = [ + Polynomial([0, 1], cls.ureg.ft, cls.ureg.s), + Polynomial([5, -5], cls.ureg.ft, cls.ureg.s), + Polynomial([0, 0, 2], cls.ureg.gal, cls.ureg.psi), + Polynomial([1, -3, 5], cls.ureg.ft, cls.ureg.s), + Polynomial([0, 0, 0, 1], cls.ureg.ft, cls.ureg.s), + Polynomial([3, -5, 9, 0.1], cls.ureg.ft, cls.ureg.s), + Polynomial([0, 0, 0, 0, 0, 1], cls.ureg.gal, cls.ureg.psi), + Polynomial([9, 2, 4, 1, 2], cls.ureg.ft, cls.ureg.s), + ] + cls.polys = polys + cls.linear_1, cls.linear_2 = polys[0], polys[1] + cls.parabolic_1, cls.parabolic_2 = polys[2], polys[3] + cls.cubic_1, cls.cubic_2 = polys[4], polys[5] + cls.polynomial_1, cls.polynomial_2 = polys[6], polys[7] + + def test_polynomial_addition(self): + + self.assertTrue( + all([a == b for a, b in zip((self.linear_1 + 5).coef, [5.0, 1.0])]) + ) + self.assertTrue( + all([a == b for a, b in zip((self.cubic_1 + 9.5).coef, [9.5, 0, 0, 1])]) + ) + self.assertTrue( + all([a == b for a, b in zip((9.5 + self.cubic_1).coef, [9.5, 0, 0, 1])]) + ) + self.assertRaises(TypeError, self.linear_1.__add__, self.parabolic_1) + self.assertTrue( + all( + [ + a == b + for a, b in zip( + (self.cubic_2 + self.polynomial_2.coef).coef, + [12, -3, 13, 1.1, 2], + ) + ] + ) + ) + self.assertTrue( + all( + [ + a == b + for a, b in zip( + (self.cubic_2 + self.polynomial_2).coef, [12, -3, 13, 1.1, 2] + ) + ] + ) + ) + + def test_polynomial_subtraction(self): + self.assertTrue( + all([a == b for a, b in zip((self.linear_1 - 5).coef, [-5, 1])]) + ) + self.assertTrue( + all([a == b for a, b in zip((self.cubic_1 - 9.5).coef, [-9.5, 0, 0, 1])]) + ) + self.assertTrue( + all([a == b for a, b in zip((9.5 - self.cubic_1).coef, [9.5, 0, 0, -1])]) + ) + + self.assertRaises(TypeError, self.linear_1.__sub__, self.parabolic_1) + self.assertTrue( + all( + [ + a == b + for a, b in zip( + (self.cubic_2 - self.polynomial_2.coef).coef, + [-6, -7, 5, -0.9, -2], + ) + ] + ) + ) + self.assertTrue( + all( + [ + a == b + for a, b in zip( + (self.cubic_2 - self.polynomial_2).coef, [-6, -7, 5, -0.9, -2] + ) + ] + ) + ) + + def test_correct_units_after_math(self): + self.assertEqual((self.linear_1 + 1).x_unit, self.linear_1.x_unit) + self.assertEqual((self.linear_1 - 1).x_unit, self.linear_1.x_unit) + self.assertEqual((self.linear_1 * 1).x_unit, self.linear_1.x_unit) + self.assertEqual((self.polynomial_1 / 2).x_unit, self.polynomial_1.x_unit) + + self.assertEqual((1 + self.linear_1).y_unit, self.linear_1.y_unit) + self.assertEqual((1 - self.linear_1).y_unit, self.linear_1.y_unit) + self.assertEqual((1 * self.linear_1).y_unit, self.linear_1.y_unit) + # self.assertEqual((1 / self.linear_1).y_unit, self.linear_1.y_unit ** -1) + + self.assertEqual( + (self.parabolic_1 + self.polynomial_1).x_unit, self.parabolic_1.x_unit + ) + self.assertEqual((self.linear_1 - self.linear_2).y_unit, self.linear_1.y_unit) + self.assertEqual( + (self.linear_1 * self.parabolic_1).x_unit, + self.linear_1.x_unit * self.parabolic_1.x_unit, + ) + self.assertEqual( + (self.linear_1 / self.parabolic_1).y_unit, + self.linear_1.y_unit / self.parabolic_1.y_unit, + ) + + def test_derivative_units(self): + for p in self.polys: + self.assertEqual(p.derivative.y_unit, p.y_unit / p.x_unit) + self.assertEqual(p.derivative.x_unit, p.x_unit) + + def test_integral_units(self): + for p in self.polys: + self.assertEqual(p.integral.y_unit, p.y_unit * p.x_unit) + self.assertEqual(p.integral.x_unit, p.x_unit) + + def test_poly_iter(self): + ys = [1, 1, 1, 1, 1, 5, 5, 10] + xs = [ + 1, + 0.8, + sqrt(0.5), + [0, 0.6], + 1, + 0.819323306231648, + 1.379729661461215, + 0.29934454191522236, + ] + for p, y, x in zip(self.polys, ys, xs): + try: + solution = p.solve_for_x(y * p.y_unit, min_value=0).m_as(p.x_unit) + if isinstance(solution, float): + self.assertAlmostEqual(solution, x) + else: + self.assertListEqual(list(solution), x) + except (AssertionError, ValueError) as e: + print(p, solution, x * p.x_unit) + raise e + + def test_solve_derivative(self): + # y_intercept of the derivative should always be the second coefficient + for p in self.polys: + self.assertEqual(p.derivative_at(0 * p.x_unit).m, p.coef[1]) + + def test_derivatives(self): + # y_intercept of the derivative should always be the second coefficient + for p in self.polys: + self.assertEqual(p.derivative.y_intercept.m, p.coef[1]) + + def test_y_intercept(self): + # y_intercept should always be the first coefficient + for p in self.polys: + self.assertEqual(p.y_intercept.m, p.coef[0]) + + def test_x_intercept(self): + intercepts = [0, 1, 0, None, 0, -90.55580410259788, 0, None] + for p, intercept in zip(self.polys, intercepts): + if intercept is None: + self.assertEqual(p.x_intercept, None) + else: + self.assertEqual(p.x_intercept, intercept * p.x_unit) + + def test_solve_linear(self): + self.assertEqual( + self.linear_1.solve(-3 * self.linear_1.x_unit), -3 * self.linear_1.y_unit + ) + self.assertEqual( + self.linear_1.solve(0 * self.linear_1.x_unit), self.linear_1.y_intercept + ) + self.assertEqual( + self.linear_1.solve(10 * self.linear_1.x_unit), 10 * self.linear_1.y_unit + ) + + self.assertEqual( + self.linear_2.solve(-3 * self.linear_2.x_unit), 20 * self.linear_2.y_unit + ) + self.assertEqual( + self.linear_2.solve(0.0 * self.linear_2.x_unit), self.linear_2.y_intercept + ) + self.assertEqual( + self.linear_2.solve(10 * self.linear_2.x_unit), -45 * self.linear_2.y_unit + ) + + def test_solve_parabolic(self): + self.assertEqual( + self.parabolic_1.solve(-3 * self.parabolic_1.x_unit), + 18 * self.parabolic_1.y_unit, + ) + self.assertEqual( + self.parabolic_1.solve(0 * self.parabolic_1.x_unit), + self.parabolic_1.y_intercept, + ) + self.assertEqual( + self.parabolic_1.solve(-10 * self.parabolic_1.x_unit), + 200 * self.parabolic_1.y_unit, + ) + + self.assertEqual( + self.parabolic_2.solve(-3 * self.parabolic_2.x_unit), + 55 * self.parabolic_2.y_unit, + ) + self.assertEqual( + self.parabolic_2.solve(0 * self.parabolic_2.x_unit), + self.parabolic_2.y_intercept, + ) + self.assertEqual( + self.parabolic_2.solve(-10 * self.parabolic_2.x_unit), + 531 * self.parabolic_2.y_unit, + )