Last active
March 27, 2025 15:37
-
-
Save roguh/f497efc796efcf8ceefc7b671d29006c to your computer and use it in GitHub Desktop.
With this code, you can define a Polynomial and find its derivative, its root (real or complex), and print it out as LaTeX math notation like 5x^2 + 2x - 3
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
"""This program is for working with Polynomials. | |
With this code, you can define a Polynomial and: | |
- convert it to LaTeX math notation | |
something like 5x^2 + 2x - 3 | |
- find the derivative of a polynomial | |
- find a root of a one dimensional polynomial (floating point or complex/imaginary) | |
""" | |
from math import isclose | |
from typing import Sequence, Generic, TypeVar | |
import pytest | |
N = TypeVar("N") | |
# The type variable N depends on how the class is initialized | |
class Polynomial(Generic[N]): | |
# Using the type Sequence[float | int] means list[float], list[int], and list[int|float] are all treated as the same type | |
def __init__(self, coefficients: Sequence[N] = (0,)): | |
"""A Polynomial function. | |
Coefficients are represented with the highest order first. | |
For example: f(x) = 5 * x^2 + 2 * x - 3 | |
This Polynomial is simply represented by the list of coefficients 5, 2, -3. | |
With the as_tex() method, we get valid LaTeX math notation. | |
>>> Polynomial([5, 2, -3]).as_tex() | |
'5x^2 + 2x - 3' | |
The exponential term is implicit in the indices. | |
Index 0 has power 2, index 1 is power 1, index 2 is power 0 (a constant). | |
A Polynomial can be easily evaluated too: | |
>>> Polynomial([5, 2, -3])(2) | |
21 | |
Or in mathematical notation: | |
f(2) = 5 * 2**2 + 2 * 2 - 3 | |
= 5 * 4 + 4 - 3 | |
= 21 | |
Notice we just call the object as though it were a function (Python's magic __call__ method). | |
From here, we can find the rate of change of any Polynomial with the .diff() function: | |
f'(x) = 10x + 2 | |
>>> Polynomial([5, 2, -3]).diff().as_tex() | |
'10x + 2' | |
And finally, we can use this derivative to find a y-intercept or zero of the polynomial: | |
""" | |
# With no arguments, create a zero-function that always returns zero. | |
if len(coefficients) == 1 and coefficients[0] == 0: | |
self.coefficients = [] | |
else: | |
self.coefficients = list(coefficients) | |
def degree(self) -> int: | |
"""x^2 has degree 2. Degree of a zero-polynomial is defined as -1 here.""" | |
return len(self.coefficients) - 1 | |
def coeff_big_first(self) -> list[N]: | |
"""Return a copy of the list of coeffs with the higher orders first.""" | |
return list(self.coefficients) | |
def coeff_small_first(self) -> list[N]: | |
"""Return a copy of the list of coeffs with the smallest orders first.""" | |
return list(reversed(self.coefficients)) | |
def eval_simple(self, x: N) -> N: | |
# This is close to the mathematical definition | |
total = 0 | |
for pow, coeff in enumerate(self.coeff_small_first()): | |
total += coeff * x**pow | |
return total | |
def __call__(self, x: N) -> N: | |
"""Evaluate this polynomial at a specific value. | |
This is implemented with Horner's method, | |
which uses less multiplications and additions than a straightforward implementation like eval_simple() | |
Example for 10x^3 + 5x^2 + 9x + 1 | |
total = 0 | |
total = 10 + x * total = 10 | |
total = 5 + x * total = 5 + 10x | |
total = 9 + x * total = 9 + 5x + 10x^2 | |
total = 1 + x * total = 10x^3 + 5x^2 + 9x + 1 | |
""" | |
total = 0 | |
for coeff in self.coeff_big_first(): | |
# Multiply, then sum | |
total = coeff + x * total | |
return total | |
def diff(self) -> "Polynomial": | |
diff_coeff = [] | |
for pow, coeff in enumerate(self.coeff_small_first()): | |
if pow == 0: | |
continue | |
diff_coeff.insert(0, coeff * pow) | |
return Polynomial(diff_coeff) | |
def find_zero(self, initial_guess: N, max_iterations: int=25) -> N: | |
"""Finds an approximation to the zero-intercept of this polynomial. | |
Works with complex numbers as well: | |
>>> Polynomial([5, 2, 3]).find_zero(-1j) | |
(-0.19999999999999998-0.7483314773547882j) | |
>>> -1/5 - 1j * (14**.5/5) # This is the right answer | |
(-0.2-0.7483314773547882j) | |
Newton's method: | |
improved_guess = current guess + f(current_guess) / f'(current_guess) | |
Explanation: | |
the goal is to approximate or to guess the point where f(x) equals zero | |
we must start with a guess of the solution | |
f'(x) tells us the rate of change at the current_guess | |
we can use f(x) / f'(x) as a sort of arrow that points to an a more accurate guess | |
then we head in that direction and repeat until the guess is considered good enough | |
""" | |
r = initial_guess | |
iterations = 0 | |
while True: | |
y = self(r) | |
d = self.diff()(r) | |
# Can crash if d is too close to zero | |
if iterations > max_iterations: | |
return r | |
r = r - y / d | |
iterations += 1 | |
# Utility functions follow | |
def __eq__(self, other) -> bool: | |
"""Checks if two polynomials are equal.""" | |
if not isinstance(other, Polynomial): | |
return False | |
return self.coefficients == other.coefficients | |
def __str__(self) -> str: | |
"""Define __str__ for easier debugging, REPL usage, etc.""" | |
return f"Polynomial({list(self.coefficients)})" | |
__repr__ = __str__ | |
def as_tex(self) -> str: | |
"""Convert this polynomial into LaTeX mathematical notation. | |
>>> Polynomial([15, -8, 165, -3, 0, 0, 1, -10]).as_tex() | |
'15x^7 - 8x^6 + 165x^5 - 3x^4 + x - 10' | |
>>> Polynomial([-7, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 3, -1.3]).as_tex() | |
'- 7x^{12} + 3x - 1.3' | |
""" | |
assert all(not isinstance(n, complex) for n in self.coefficients), "Cannot print polynomials with complex coefficients (yet)" | |
mult_sep = "" # can also be " \\cdot " for a dot | |
coeff_strs = [] | |
if self.degree() == -1: | |
return "0" | |
for pow, signed_coeff in enumerate(self.coeff_small_first()): | |
if signed_coeff != 0: | |
# We'll add the sign +/- later | |
coeff = abs(signed_coeff) | |
if coeff == 1: | |
coeff_str = [] | |
else: | |
coeff_str = [f"{coeff}"] | |
if pow == 0: | |
this_str = "".join(coeff_str) | |
elif pow == 1: | |
this_str = mult_sep.join(coeff_str + ["x"]) | |
else: | |
pow_str = str(pow) | |
if len(pow_str) > 1: | |
pow_str = "{" + pow_str + "}" | |
this_str = mult_sep.join(coeff_str + [f"x^{pow_str}"]) | |
if signed_coeff < 0: | |
this_str = "- " + this_str | |
elif pow != self.degree(): | |
this_str = "+ " + this_str | |
coeff_strs.append(this_str) | |
pow -= 1 | |
return " ".join(reversed(coeff_strs)) | |
g = 9.82 | |
init_velocity, init_height = 4, 10 | |
motion = Polynomial([-g, init_velocity, init_height]) | |
test_xs = [0.0, 1, 2, 2.0, 3.14159265359, 100, 1000, 1e10] | |
test_xs = test_xs + [-x for x in test_xs] | |
@pytest.mark.parametrize("x", test_xs) | |
def test_zero_function(x): | |
assert Polynomial([0])(x) == 0 | |
assert Polynomial([])(x) == 0 | |
@pytest.mark.parametrize("x", test_xs) | |
@pytest.mark.parametrize("c", [1, 10, -123, 1e10]) | |
def test_constant_function(x, c): | |
assert Polynomial([c])(x) == c | |
@pytest.mark.parametrize("x", test_xs) | |
@pytest.mark.parametrize("c", [1, 10, -123, 1e10]) | |
def test_linear_function(x, c): | |
assert Polynomial([c, 0])(x) == c * x | |
@pytest.mark.parametrize("x", test_xs) | |
@pytest.mark.parametrize("c", [1, 10, -123, 1e10]) | |
def test_quadratic_function(x, c): | |
assert Polynomial([c, 0, 0])(x) == c * x**2 | |
@pytest.mark.parametrize("str_conv", [str, repr]) | |
def test_str(str_conv): | |
poly = Polynomial([5, 2, -3]) | |
expected = "Polynomial([5, 2, -3])" | |
assert str_conv(poly) == expected | |
assert str_conv(Polynomial([])) == "Polynomial([])" | |
assert str_conv(Polynomial([0])) == "Polynomial([])" | |
assert str_conv(Polynomial([5])) == "Polynomial([5])" | |
def test_simple(): | |
assert Polynomial([5, 0, 0]).as_tex() == "5x^2" | |
assert Polynomial([5, 2, -3]).as_tex() == "5x^2 + 2x - 3" | |
assert Polynomial([]).as_tex() == "0" | |
assert Polynomial([0]).as_tex() == "0" | |
assert Polynomial([5]).as_tex() == "5" | |
assert Polynomial([5, 10]).as_tex() == "5x + 10" | |
@pytest.mark.parametrize( | |
"poly, x, y", | |
( | |
(Polynomial([]), 0, 0), | |
(Polynomial([]), 10, 0), | |
(Polynomial([5]), 10, 5), | |
(Polynomial([2, 3]), 10, 23), | |
(Polynomial([5, 2, 3]), 10, 523), # 5 * 100 + 2 * 10 + 3 * 1 | |
(Polynomial([5, 2, -3]), 2, 21), | |
(Polynomial([5, 2, -3]), 4.5, 107.25), | |
), | |
) | |
def test_eval(poly, x, y): | |
assert poly.eval_simple(x) == y | |
assert poly(x) == poly.eval_simple(x) | |
assert poly(x) == y | |
def test_zero(): | |
assert isclose(Polynomial([5, 2, -3]).find_zero(-5), -1.0) | |
assert isclose(Polynomial([5, 2, -3]).find_zero(5), 0.6) | |
imag_zero = Polynomial([5, 2, 3]).find_zero(-1j) | |
assert isclose(imag_zero.real, -0.2) | |
assert isclose(imag_zero.imag, -0.7483314773547882) | |
def test_diff_tex(): | |
assert Polynomial([5, 0, 0]).diff().as_tex() == "10x" | |
assert Polynomial([5, 2, 3]).diff().as_tex() == "10x + 2" | |
def test_diff_eq(): | |
assert Polynomial([5, 0, 0]).diff() == Polynomial([10, 0]) | |
assert Polynomial([5, 2, 3]).diff() == Polynomial([10, 2]) | |
assert Polynomial([441142, 14, 14, 124, 124, 124, 124, 1, 41, 241, 24, 12415, 31, 37, 86, 84, 62, 51]).diff() == Polynomial( | |
[7499414, 224, 210, 1736, 1612, 1488, 1364, 10, 369, 1928, 168, 74490, 155, 148, 258, 168, 62] | |
) | |
def test_conways_constant(): | |
# Long ago, before computer screens existed, Dr. John Conway had to give a B.S. test to his students but he knew they'd all pass. | |
# He saw the weather was dreadful that time of the semester and decided to excuse all his students for the day, although he would show up to make sure the administrabots didn't fire him. | |
# Instead, his students threw a surprise party in the lecture room. | |
# They had a fun time, and among the many conversations that day, one student posed this question to Dr. John Conway: | |
# Look at this pattern: | |
# 1 | |
# 11 | |
# 21 | |
# What number comes next? | |
# ... | |
# The next number is 1211 :) | |
# Dr. John Conway couldn't predict the next numbers, despite all his mathematical knowledge. | |
# Eventually, the student explained the pattern and Dr. John Conway spent a lot of time studying this list of strings. | |
# | |
# | |
# | |
# This is the Look and Say sequence. | |
# | |
# The trick is it is not a numerical pattern, it's a combination of string manipulation and counting no more than 3 things. | |
# This pattern will only ever have 3 digits: 1, 2, and 3. | |
# | |
# We are counting how many digits appeared in the last string! | |
# 11 means we started with a single one digit. | |
# 21 means we saw two ones previously. | |
# 1211 means we saw a single two, and then a single one digit. | |
# 111221 and the next value means we saw a single one, a single two, and two ones. | |
# 312211 means we saw 3 ones, then 2 twos, then a 1 one. | |
# | |
# | |
# | |
# Here's another question: how fast does this string grow as we keep following the pattern? | |
# | |
# The strings grew from lengths 1, to 2, 2, 4, 6, 6, and so on. | |
# They can never shrink and there is a predictable pattern to their growth. | |
# Is the length of the string growing quadratically or exponentially, perhaps? | |
# | |
# If the 1st string is "1", the 64th string in this sequence is 36 million characters long (36,619,163). | |
# The 65th string is 47,734,823 characters long. | |
# If you want an approximation of how fast the strings will grow: | |
# 36619163 divided by 47734823 is about 1.3035... | |
# The rate of growth is exponential with a coefficient of about 1.3035... | |
# | |
# The 75th string is 676,406,711 characters long, or about 676 megabytes of uncompressed data. | |
# The ratio between the lengths of the 74th string and the 75th string is: | |
# 676406711/518891359 => 1.303561331804718 | |
# | |
# After many more iterations, this number settles down into a specific constant. | |
# This number is called Conway's Constant. | |
# | |
# To find the exact rate of growth, the rules for this string function can be represented as an 80+ dimensional matrix that shows how one string transitions into another. | |
# This matrix can be simplified to a 71-dimensional matrix, and the eigenfunction represents its growth rate. | |
# From here, that matrix is converted to a 71-degree polynomial. | |
# | |
# When you find the root (the zero) of this polynomial, you find how exactly how fast the string function's input grows as you infinitely apply it to the starting string "1" | |
# Remember if f(x) = 0, then x is the zero of the function. | |
# | |
# Because it's a polynomial, we can find its derivative and then use Newton's method to find the zero of this function. | |
# | |
# When we find the root of this 71-degree polynomial, we'll find an increasingly accurate value for Conway's constant. | |
conway_poly = Polynomial(conway_coeffs) | |
# The polynomial looks a bit like: | |
# f(x) = x^71 - x^69 - 2x^68 .... -6x^2 + 3x - 6 | |
# All of its coefficients are integers (because it represents string transitions, which only happen in one-character increments) | |
assert conway_poly.as_tex().startswith("x^{71} - x^{69} - 2x^{68}") | |
assert conway_poly.as_tex().endswith("- 6x^2 + 3x - 6") | |
# This polynomial has degree 71 | |
assert conway_poly.degree() == 71 | |
# Now we can find a 16 digit approximation for Conway's Constant | |
approx_zero = conway_poly.find_zero(initial_guess=1.3) | |
# The 64th term of the Look and Say sequence has a length of about 36 million, for the 65th it's about 47 million. | |
# Their ratio represents the rate of growth of the Look and Say sequence. | |
# This ratio is accurate to only 4 digits. | |
approx4 = 47_734_823 / 36_619_163 | |
assert isclose(approx_zero, approx4, rel_tol=1e-4) | |
# 16 digits of Conway's Constant from Wolfram MathWorld | |
conway_constant16 = 1.3035772690342966 | |
assert approx_zero == conway_constant16 | |
# Extra: We found 16 digits of this factor. Now how would one find 1000 digits? | |
conway_coeffs_raw = ( | |
"1 0 -1 -2 -1 2 2 1 -1 -1 -1 -1 -1 2 5 3 -2 -10 -3 -2 6 6 1 9 -3 -7 -8 -8 10 6 8 -5 -12 7 -7 7 1 -3 10 1 -6 -2 -10 -3 2 9 -3 14 -8 0 -7 9 3 -4 -10 -7 12 7 2 -12 -4 -2 5 0 1 -7 7 -4 12 -6 3 -6" | |
) | |
conway_coeffs = list(map(int, conway_coeffs_raw.split())) | |
if __name__ == "__main__": | |
pytest.main([__file__]) | |
print(Polynomial(conway_coeffs).find_zero(initial_guess=1.3)) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment