Coverage for src/invariant/ops/poly.py: 90.00%

1"""Polynomial operations for testing Invariant's DAG capabilities."""

2

3 from typing import Any

4

5 from invariant.types import Polynomial

6

7

8 def poly_from_coefficients(coefficients: list[int]) -> Polynomial:

9 """Create a polynomial from a list of coefficients.

10

11 Args:

12 coefficients: List of integer coefficients. Index i represents

13 the coefficient of x^i. Trailing zeros are automatically

14 stripped for canonical form.

15

16 Returns:

17 Polynomial with the given coefficients (trailing zeros stripped).

18

19 Raises:

20 TypeError: If coefficients are not integers.

21 """

22 # Ensure all coefficients are integers

23 int_coeffs = []

24 for i, coeff in enumerate(coefficients):

25 if not isinstance(coeff, int):

26 raise TypeError(f"Coefficient at index {i} must be int, got {type(coeff)}")

27 int_coeffs.append(coeff)

28

29 return Polynomial(int_coeffs)

30

31

32 def poly_add(a: Polynomial, b: Polynomial) -> Polynomial:

33 """Add two polynomials.

34

35 Args:

36 a: First polynomial.

37 b: Second polynomial.

38

39 Returns:

40 Polynomial representing the sum.

41 """

42 max_len = max(len(a.coefficients), len(b.coefficients))

43 result_coeffs = []

44 for i in range(max_len):

45 coeff_a = a.coefficients[i] if i < len(a.coefficients) else 0

46 coeff_b = b.coefficients[i] if i < len(b.coefficients) else 0

47 result_coeffs.append(coeff_a + coeff_b)

48 return Polynomial(tuple(result_coeffs))

49

50

51 def poly_multiply(a: Polynomial, b: Polynomial) -> Polynomial:

52 """Multiply two polynomials (convolution of coefficient lists).

53

54 Args:

55 a: First polynomial.

56 b: Second polynomial.

57

58 Returns:

59 Polynomial representing the product.

60 """

61 # Result degree is len(a) + len(b) - 2 (0-indexed)

62 result_len = len(a.coefficients) + len(b.coefficients) - 1

63 result_coeffs = [0] * result_len

64

65 # Convolution: result[i+j] += a[i] * b[j]

66 for i, coeff_a in enumerate(a.coefficients):

67 for j, coeff_b in enumerate(b.coefficients):

68 result_coeffs[i + j] += coeff_a * coeff_b

69

70 return Polynomial(tuple(result_coeffs))

71

72

73 def poly_scale(poly: Polynomial, scalar: int) -> Polynomial:

74 """Scale a polynomial by a scalar.

75

76 Args:

77 poly: Polynomial to scale.

78 scalar: int scalar value.

79

80 Returns:

81 Polynomial with all coefficients multiplied by scalar.

82 """

83 # Multiply each coefficient by scalar

84 new_coeffs = tuple(c * scalar for c in poly.coefficients)

85 return Polynomial(new_coeffs)

86

87

88 def poly_derivative(poly: Polynomial) -> Polynomial:

89 """Compute the derivative of a polynomial.

90

91 The derivative of c[i] * x^i is c[i] * i * x^(i-1), which means

92 we multiply each coefficient by its index and shift down one degree.

93

94 Args:

95 poly: Polynomial to differentiate.

96

97 Returns:

98 Polynomial representing the derivative.

99 """

100 # Derivative: c[i] * i becomes the coefficient of x^(i-1)

101 # For i=0, the derivative is 0 (constant term disappears)

102 if len(poly.coefficients) <= 1:

103 # Constant or zero polynomial -> derivative is zero

104 return Polynomial((0,))

105

106 new_coeffs = []

107 for i in range(1, len(poly.coefficients)):

108 # Coefficient at index i becomes i * coeff[i] at index i-1

109 new_coeffs.append(poly.coefficients[i] * i)

110

111 return Polynomial(tuple(new_coeffs))

112

113

114 def poly_evaluate(poly: Polynomial, x: int) -> int:

115 """Evaluate a polynomial at a point using Horner's method.

116

117 Args:

118 poly: Polynomial to evaluate.

119 x: int value to evaluate at.

120

121 Returns:

122 int result of evaluation.

123 """

124 # Horner's method: evaluate from highest degree down

125 # For polynomial c[0] + c[1]*x + c[2]*x^2 + ... + c[n]*x^n

126 # Start with result = c[n], then result = result * x + c[n-1], etc.

127 result = 0

128 for i in range(len(poly.coefficients) - 1, -1, -1):

129 result = result * x + poly.coefficients[i]

130

131 return result

132

133

134# Package of polynomial operations

135 OPS: dict[str, Any] = {

136 "from_coefficients": poly_from_coefficients,

137 "add": poly_add,

138 "multiply": poly_multiply,

139 "evaluate": poly_evaluate,

140 "scale": poly_scale,

141 "derivative": poly_derivative,

142}