116 lines
3.6 KiB
Python
116 lines
3.6 KiB
Python
|
"""Find the roots of a polynomial by *asking* the model for them.
|
||
|
|
|
||
|
|
Compare against numpy's actual numerical root finder. For nice polynomials
|
||
|
|
with integer roots, the AI often gets there. For uglier ones, it makes
|
||
|
|
plausible-looking guesses that don't quite check out. Both behaviors are
|
||
|
|
interesting — and both are what you should expect from a system that
|
||
|
|
"reasons" by predicting the most likely next characters.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
from ai_function import ask
|
||
|
|
|
||
|
|
|
||
|
|
def format_polynomial(coeffs: list[float]) -> str:
|
||
|
|
"""Render coefficients as a readable expression like '2x^3 - 5x + 1'."""
|
||
|
|
terms = []
|
||
|
|
degree = len(coeffs) - 1
|
||
|
|
for i, c in enumerate(coeffs):
|
||
|
|
if c == 0:
|
||
|
|
continue
|
||
|
|
power = degree - i
|
||
|
|
# Coefficient formatting: drop trailing zeros, keep sign in front.
|
||
|
|
abs_str = f"{abs(c):g}"
|
||
|
|
sign = "+" if c > 0 else "-"
|
||
|
|
if power == 0:
|
||
|
|
body = abs_str
|
||
|
|
else:
|
||
|
|
coef = "" if abs_str == "1" else abs_str
|
||
|
|
var = "x" if power == 1 else f"x^{power}"
|
||
|
|
body = f"{coef}{var}"
|
||
|
|
terms.append(f"{sign} {body}")
|
||
|
|
if not terms:
|
||
|
|
return "0"
|
||
|
|
expr = " ".join(terms)
|
||
|
|
# Drop the leading "+ " if positive.
|
||
|
|
if expr.startswith("+ "):
|
||
|
|
expr = expr[2:]
|
||
|
|
elif expr.startswith("- "):
|
||
|
|
expr = "-" + expr[2:]
|
||
|
|
return expr
|
||
|
|
|
||
|
|
|
||
|
|
def ai_polynomial_roots(coeffs: list[float]) -> str:
|
||
|
|
"""Same shape as numpy.roots: pass coefficients high-degree to low-degree."""
|
||
|
|
polynomial = format_polynomial(coeffs)
|
||
|
|
return ask(
|
||
|
|
"Find the real roots of the polynomial. Output only a Python list of "
|
||
|
|
"numbers rounded to 3 decimal places, like [1.0, -2.0]. No explanation.\n"
|
||
|
|
"Polynomial: x^2 - 5x + 6\n"
|
||
|
|
"Roots: [2.0, 3.0]\n"
|
||
|
|
"Polynomial: x^2 - 4\n"
|
||
|
|
"Roots: [-2.0, 2.0]\n"
|
||
|
|
"Polynomial: x^3 - 6x^2 + 11x - 6\n"
|
||
|
|
"Roots: [1.0, 2.0, 3.0]\n"
|
||
|
|
f"Polynomial: {polynomial}\n"
|
||
|
|
"Roots: "
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def py_polynomial_roots(coeffs: list[float]) -> list[float]:
|
||
|
|
return sorted(np.roots(coeffs).real.round(3).tolist())
|
||
|
|
|
||
|
|
|
||
|
|
def _run_one(coeffs: list[float]) -> None:
|
||
|
|
print(f" coefficients: {coeffs}")
|
||
|
|
print(f" polynomial: {format_polynomial(coeffs)}")
|
||
|
|
ai_out = ai_polynomial_roots(coeffs)
|
||
|
|
py_out = py_polynomial_roots(coeffs)
|
||
|
|
print(f" AI says: {ai_out}")
|
||
|
|
print(f" numpy says: {py_out}")
|
||
|
|
print()
|
||
|
|
|
||
|
|
|
||
|
|
def _canned_examples() -> None:
|
||
|
|
cases = [
|
||
|
|
[1, -7, 12], # roots 3, 4
|
||
|
|
[1, 0, -9], # roots -3, 3
|
||
|
|
[1, -6, 11, -6], # roots 1, 2, 3
|
||
|
|
[2, -3, -11, 6], # roots 3, -2, 0.5
|
||
|
|
]
|
||
|
|
for coeffs in cases:
|
||
|
|
_run_one(coeffs)
|
||
|
|
|
||
|
|
|
||
|
|
def _interactive() -> None:
|
||
|
|
print("--- interactive ---")
|
||
|
|
print(
|
||
|
|
"Enter polynomial coefficients high-degree to low-degree, separated by "
|
||
|
|
"commas or spaces. Example: `1, -7, 12` is x^2 - 7x + 12. 'q' to quit.\n"
|
||
|
|
)
|
||
|
|
while True:
|
||
|
|
try:
|
||
|
|
raw = input("coefficients > ").strip()
|
||
|
|
except (EOFError, KeyboardInterrupt):
|
||
|
|
print()
|
||
|
|
return
|
||
|
|
if raw.lower() in {"q", "quit", "exit"}:
|
||
|
|
return
|
||
|
|
try:
|
||
|
|
coeffs = [float(p) for p in raw.replace(",", " ").split() if p]
|
||
|
|
except ValueError:
|
||
|
|
print(" could not parse those as numbers, try again")
|
||
|
|
continue
|
||
|
|
if len(coeffs) < 2:
|
||
|
|
print(" need at least 2 coefficients (a linear polynomial)")
|
||
|
|
continue
|
||
|
|
_run_one(coeffs)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
_canned_examples()
|
||
|
|
_interactive()
|