97 lines
3.2 KiB
Python
97 lines
3.2 KiB
Python
from sympy.abc import x, y
|
|
import numpy as np
|
|
|
|
from .carry import Carry
|
|
|
|
|
|
def add(
|
|
expansion: np.ndarray, carry: Carry, val: int = 1, center: tuple[int, int] = (0, 0)
|
|
):
|
|
"""
|
|
Add `val` to a position in an `expansion` designated the `center`, then apply `carry`.
|
|
"""
|
|
expansion[center] += val
|
|
carry.apply(expansion)
|
|
return expansion
|
|
|
|
|
|
def times(expansion: np.ndarray, carry: Carry, val: int):
|
|
"""
|
|
Multiply all digits in an `expansion` by `val`, then apply `carry`.
|
|
"""
|
|
expansion *= val
|
|
carry.apply(expansion)
|
|
return expansion
|
|
|
|
|
|
def poly_from_array(
|
|
array: np.ndarray, center: tuple[int, int] = (0, 0), x_var=x, y_var=y
|
|
):
|
|
"""
|
|
Convert a numpy array to a polynomial in two variables.
|
|
Note that indices are row-major in Numpy, so center is in the form (y, x).
|
|
"""
|
|
ret = 0
|
|
for i, row in enumerate(array):
|
|
for j, val in enumerate(row):
|
|
ret += val * (x_var ** (i - center[1]) * y_var ** (j - center[0]))
|
|
return ret
|
|
|
|
|
|
def _first_nonzero(arr: np.ndarray, axis: int):
|
|
"""
|
|
Find the index of the first nonzero entry of `arr` along `axis`.
|
|
Modified from https://stackoverflow.com/questions/47269390
|
|
"""
|
|
mask = arr != 0
|
|
return min(
|
|
np.where(mask.any(axis=axis), mask.argmax(axis=axis), arr.shape[axis] + 1)
|
|
)
|
|
|
|
|
|
def latex_polynumber(
|
|
arr: np.ndarray, center: tuple[int, int] | None = None, show_zero: bool = False
|
|
):
|
|
"""
|
|
Convert a 2D Numpy array into a LaTeX array representing a polynomial in two variables.
|
|
Vertical lines and `\\hline`s are placed to indicate the transition between negative
|
|
and nonnegative powers, similar to a typical fractional separator (decimal point).
|
|
|
|
By default, this assumes the constant term is in position (0, 0) of the array.
|
|
It may be alternatively supplied by `center`, which designates `arr[center]` as the constant term.
|
|
Note that indices are row-major in Numpy, so `center` is in the form (y, x).
|
|
|
|
`show_zero` defaults to False, which causes entries in `arr` which are 0 to not be drawn.
|
|
When true, all 0 entries will be drawn.
|
|
"""
|
|
upper_left = _first_nonzero(arr, 0), _first_nonzero(arr, 1)
|
|
lower_right = (
|
|
len(arr) - _first_nonzero(np.flip(arr, 0), 0) - 1,
|
|
len(arr[1]) - _first_nonzero(np.flip(arr, 1), 1) - 1,
|
|
)
|
|
|
|
center_left, center_top = center or (0, 0)
|
|
if center is not None:
|
|
# this does not need offset, since we iterate over this range
|
|
center_top = center[0]
|
|
# but this does for the array environment argument
|
|
center_left = center[1] - upper_left[1]
|
|
|
|
num_columns = lower_right[1] - upper_left[1]
|
|
column_layout = ("c" * center_left) + "|" + ("c" * (num_columns - center_left))
|
|
|
|
# build array output
|
|
ret = "\\begin{array}{" + column_layout + "}"
|
|
for i in range(upper_left[0], lower_right[0] + 1):
|
|
if i == center_top:
|
|
ret += " \\hline "
|
|
ret += " & ".join(
|
|
[
|
|
str(arr[i, j] if arr[i, j] != 0 or show_zero else "")
|
|
for j in range(upper_left[1], lower_right[1] + 1)
|
|
]
|
|
)
|
|
# next row
|
|
ret += " \\\\ "
|
|
return ret + "\\end{array}"
|