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

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

2 

3from typing import Any 

4 

5from invariant.types import Polynomial 

6 

7 

8def 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 

32def 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 

51def 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 

73def 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 

88def 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 

114def 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 

135OPS: 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}