Coverage for src/algolib/core/complex.py: 100%

1# src/algolib/maths/complex/complex.py 

2""" 

3A lightweight Complex number implementation (without using Python's built-in ``complex``). 

4 

5This module provides a small, well-documented ``Complex`` class suitable for learning 

6and algorithmic implementations. It supports algebraic form (a + bi) and polar form 

7(r·e^{iθ}), with common operations. 

8 

9Notes 

10----- 

11- This class uses plain floats and is immutable. 

12- We intentionally avoid Python's built-in :class:`complex` to practice fundamentals. 

13 

14Examples 

15-------- 

16>>> z1 = Complex(3, 4) 

17>>> z1.modulus() 

185.0 

19>>> z2 = Complex.from_polar(2, math.pi/2) 

20>>> (z1 + z2).re # doctest: +ELLIPSIS 

213.0 

22>>> (z1.conjugate()).im 

23-4.0 

24""" 

25 

26from __future__ import annotations 

27 

28import math 

29from dataclasses import dataclass 

30from typing import Iterable, Tuple 

31 

32from algolib.exceptions import InvalidTypeError, InvalidValueError 

33 

34 

35Number = float # for readability 

36 

37 

38@dataclass(frozen=True) 

39class Complex: 

40 """ 

41 Complex number in algebraic form :math:`a + b \\mathrm{i}`. 

42 

43 Parameters 

44 ---------- 

45 re : float 

46 Real part :math:`a`. 

47 im : float 

48 Imaginary part :math:`b`. 

49 

50 Raises 

51 ------ 

52 InvalidTypeError 

53 If either part is not a real number (int/float). 

54 """ 

55 

56 re: Number 

57 im: Number 

58 

59 # ------------------------------- constructors ------------------------------- 

60 

61 def __post_init__(self) -> None: 

62 """ 

63 Post-initialization validation and normalization. 

64 

65 Ensures that the real and imaginary parts are valid numbers (int or float) 

66 and converts them to floats for consistency. 

67 

68 Raises 

69 ------ 

70 InvalidTypeError 

71 If `re` or `im` is not a real number. 

72 """ 

73 if not isinstance(self.re, (int, float)) or not isinstance(self.im, (int, float)): 

74 raise InvalidTypeError("re and im must be real numbers (int or float).") 

75 

76 # freeze as floats (even if ints given) 

77 object.__setattr__(self, "re", float(self.re)) 

78 object.__setattr__(self, "im", float(self.im)) 

79 

80 @staticmethod 

81 def from_polar(r: Number, theta: Number) -> "Complex": 

82 r""" 

83 Construct a complex number from polar coordinates. 

84 

85 Parameters 

86 ---------- 

87 r : float 

88 The modulus (radius) of the complex number. Must be non-negative. 

89 theta : float 

90 The argument (angle in radians) of the complex number. 

91 

92 Returns 

93 ------- 

94 Complex 

95 The complex number corresponding to the polar coordinates. 

96 

97 Raises 

98 ------ 

99 InvalidTypeError 

100 If `r` or `theta` is not a real number. 

101 InvalidValueError 

102 If `r` is negative. 

103 

104 Examples 

105 -------- 

106 >>> z = Complex.from_polar(2, math.pi / 2) 

107 >>> z 

108 Complex(re=1.2246467991473532e-16, im=2.0) 

109 """ 

110 if not isinstance(r, (int, float)) or not isinstance(theta, (int, float)): 

111 raise InvalidTypeError("r and theta must be real numbers (int or float).") 

112 if r < 0: 

113 raise InvalidValueError(f"radius r must be non-negative, got {r}") 

114 return Complex(r * math.cos(theta), r * math.sin(theta)) 

115 

116 @staticmethod 

117 def from_cartesian(re: Number, im: Number) -> "Complex": 

118 """ 

119 Construct a complex number from Cartesian coordinates. 

120 

121 Parameters 

122 ---------- 

123 re : float 

124 The real part of the complex number. 

125 im : float 

126 The imaginary part of the complex number. 

127 

128 Returns 

129 ------- 

130 Complex 

131 The complex number corresponding to the Cartesian coordinates. 

132 

133 Examples 

134 -------- 

135 >>> z = Complex.from_cartesian(3, 4) 

136 >>> z 

137 Complex(re=3.0, im=4.0) 

138 """ 

139 return Complex(re, im) 

140 

141 @staticmethod 

142 def from_iterable(pair: Iterable[Number]) -> "Complex": 

143 r""" 

144 Construct a complex number from an iterable of two numbers. 

145 

146 Parameters 

147 ---------- 

148 pair : Iterable[float] 

149 An iterable containing exactly two elements: the real and imaginary parts. 

150 

151 Returns 

152 ------- 

153 Complex 

154 The complex number constructed from the iterable. 

155 

156 Raises 

157 ------ 

158 InvalidTypeError 

159 If the iterable does not contain exactly two numeric elements. 

160 

161 Examples 

162 -------- 

163 >>> z = Complex.from_iterable([3, 4]) 

164 >>> z 

165 Complex(re=3.0, im=4.0) 

166 """ 

167 try: 

168 re, im = pair # type: ignore[misc] 

169 except Exception as e: # noqa: BLE001 

170 raise InvalidTypeError("pair must be an iterable of length 2 (re, im).") from e 

171 return Complex(re, im) 

172 

173 # ------------------------------- properties --------------------------------- 

174 

175 def to_tuple(self) -> Tuple[Number, Number]: 

176 r""" 

177 Convert the complex number to a tuple of its real and imaginary parts. 

178 

179 Returns 

180 ------- 

181 tuple of float 

182 A tuple `(re, im)` representing the real and imaginary parts. 

183 

184 Examples 

185 -------- 

186 >>> z = Complex(3, 4) 

187 >>> z.to_tuple() 

188 (3.0, 4.0) 

189 """ 

190 return (self.re, self.im) 

191 

192 # --------------------------------- queries ---------------------------------- 

193 

194 def modulus(self) -> float: 

195 r""" 

196 Compute the modulus (absolute value) of the complex number. 

197 

198 Returns 

199 ------- 

200 float 

201 The modulus :math:`\sqrt{a^2 + b^2}`. 

202 

203 Examples 

204 -------- 

205 >>> z = Complex(3, 4) 

206 >>> z.modulus() 

207 5.0 

208 """ 

209 x, y = abs(self.re), abs(self.im) 

210 if x < y: 

211 x, y = y, x 

212 if x == 0.0: # 避免 0/0 

213 return 0.0 

214 r = y / x 

215 return x * (1.0 + r * r) ** 0.5 

216 

217 def argument(self) -> float: 

218 r""" 

219 Compute the principal argument of the complex number. 

220 

221 Returns 

222 ------- 

223 float 

224 The argument (angle in radians) in the range :math:`(-\pi, \pi]`. 

225 

226 Examples 

227 -------- 

228 >>> z = Complex(0, 1) 

229 >>> z.argument() 

230 1.5707963267948966 

231 """ 

232 return math.atan2(self.im, self.re) 

233 

234 def conjugate(self) -> "Complex": 

235 r""" 

236 Compute the complex conjugate of the number. 

237 

238 Returns 

239 ------- 

240 Complex 

241 The conjugate :math:`a - b \mathrm{i}`. 

242 

243 Examples 

244 -------- 

245 >>> z = Complex(3, 4) 

246 >>> z.conjugate() 

247 Complex(re=3.0, im=-4.0) 

248 """ 

249 return Complex(self.re, -self.im) 

250 

251 def normalized(self) -> "Complex": 

252 r""" 

253 Normalize the complex number by :math:`z/|z|` to have a modulus of 1. 

254 

255 Returns 

256 ------- 

257 Complex 

258 The normalized complex number. 

259 

260 Raises 

261 ------ 

262 InvalidValueError 

263 If the complex number is zero. 

264 

265 Examples 

266 -------- 

267 >>> z = Complex(3, 4) 

268 >>> z.normalized() 

269 Complex(re=0.6, im=0.8) 

270 """ 

271 r = math.hypot(self.re, self.im) 

272 if r == 0.0: 

273 raise InvalidValueError("cannot normalize 0+0i.") 

274 # 先标准化 

275 x, y = self.re / r, self.im / r 

276 # 计算（标准化后）的模长平方，理论应为 1，但会有舍入 

277 m2 = x * x + y * y 

278 # 用一次性的缩放把模长“微校准”到 1 

279 # 注意：sqrt 和乘法的舍入会再带来极小误差，但会压到 ~1e-15 量级 

280 if m2 != 1.0 and m2 > 0.0 and math.isfinite(m2): 

281 adj = 1.0 / math.sqrt(m2) 

282 x, y = x * adj, y * adj 

283 return Complex(x, y) 

284 

285 # ---------------------------------- algebra --------------------------------- 

286 

287 def __add__(self, other: "Complex") -> "Complex": 

288 """ 

289 Add two complex numbers. 

290 

291 Parameters 

292 ---------- 

293 other : :class:`Complex` 

294 The complex number to add. 

295 

296 Returns 

297 ------- 

298 :class:`Complex` 

299 The sum of the two complex numbers. 

300 

301 Raises 

302 ------ 

303 InvalidTypeError 

304 If ``other`` is not :class:`Complex`. 

305 

306 Examples 

307 -------- 

308 >>> z1 = Complex(3, 4) 

309 >>> z2 = Complex(1, -1) 

310 >>> z1 + z2 

311 Complex(re=4.0, im=3.0) 

312 """ 

313 if not isinstance(other, Complex): 

314 raise InvalidTypeError("other must be Complex.") 

315 return Complex(self.re + other.re, self.im + other.im) 

316 

317 def __sub__(self, other: "Complex") -> "Complex": 

318 """ 

319 Subtract two complex numbers. 

320 """ 

321 if not isinstance(other, Complex): 

322 raise InvalidTypeError("other must be Complex.") 

323 return Complex(self.re - other.re, self.im - other.im) 

324 

325 def __mul__(self, other: "Complex") -> "Complex": 

326 r""" 

327 Multiply two complex numbers. 

328 

329 Formula 

330 ------- 

331 :math:`(a+bi)(c+di) = (ac - bd) + (ad + bc)i`. 

332 """ 

333 if not isinstance(other, Complex): 

334 raise InvalidTypeError("other must be Complex.") 

335 a, b, c, d = self.re, self.im, other.re, other.im 

336 return Complex(a * c - b * d, a * d + b * c) 

337 

338 def __truediv__(self, other: "Complex") -> "Complex": 

339 r""" 

340 Robust complex division using scaling + Smith's algorithm to avoid overflow/underflow. 

341 

342 For z = a+bi, w = c+di: 

343 1) scale = max(|a|,|b|,|c|,|d|) ; divide all by scale (if scale>0) 

344 2) use Smith's branch on the scaled values. 

345 """ 

346 if not isinstance(other, Complex): 

347 raise InvalidTypeError("other must be Complex.") 

348 a, b = self.re, self.im 

349 c, d = other.re, other.im 

350 if c == 0.0 and d == 0.0: 

351 raise InvalidValueError("division by zero complex.") 

352 

353 # ---- 统一缩放，避免中间量溢出/下溢 ---- 

354 scale = max(abs(a), abs(b), abs(c), abs(d)) 

355 if scale > 0.0: 

356 a /= scale; 

357 b /= scale; 

358 c /= scale; 

359 d /= scale 

360 # 若 scale==0，这里意味着 self==0 且 other!=0，直接走下面公式也安全 

361 

362 ac = abs(c) 

363 ad = abs(d) 

364 if ac >= ad: 

365 # |c| >= |d| 

366 t = d / c # |t| <= 1 

367 denom = c + d * t # 数值上 ~ (c^2 + d^2)/c 

368 re = (a + b * t) / denom 

369 im = (b - a * t) / denom 

370 else: 

371 t = c / d # |t| <= 1 

372 denom = d + c * t # 数值上 ~ (c^2 + d^2)/d 

373 re = (a * t + b) / denom 

374 im = (b * t - a) / denom 

375 

376 return Complex(re, im) 

377 

378 def __neg__(self) -> "Complex": 

379 """Unary minus.""" 

380 return Complex(-self.re, -self.im) 

381 

382 def __abs__(self) -> float: 

383 """Builtin ``abs(z)`` -> modulus.""" 

384 return self.modulus() 

385 

386 # ------------------------------- comparisons -------------------------------- 

387 

388 def almost_equal(self, other: "Complex", tol: float = 1e-12) -> bool: 

389 """ 

390 Return True if each component differs by at most ``tol`` (absolute). 

391 

392 This is safer than exact float equality. 

393 """ 

394 if not isinstance(other, Complex): 

395 return False 

396 return (abs(self.re - other.re) <= tol) and (abs(self.im - other.im) <= tol) 

397 

398 def __eq__(self, other: object) -> bool: # exact equality 

399 if not isinstance(other, Complex): 

400 return NotImplemented 

401 return self.re == other.re and self.im == other.im 

402 

403 # --------------------------------- polar ------------------------------------ 

404 

405 def to_polar(self) -> Tuple[float, float]: 

406 r""" 

407 Convert the complex number to polar coordinates. 

408 

409 Returns 

410 ------- 

411 tuple of float 

412 A tuple ``(r, theta)`` where ``r`` is the modulus and ``theta`` is the argument. 

413 

414 Examples 

415 -------- 

416 >>> z = Complex(3, 4) 

417 >>> z.to_polar() 

418 (5.0, 0.9272952180016122) 

419 """ 

420 return (self.modulus(), self.argument()) 

421 

422 # --------------------------------- display ---------------------------------- 

423 

424 def __repr__(self) -> str: 

425 return f"Complex(re={self.re:.12g}, im={self.im:.12g})" 

426 

427 def __str__(self) -> str: 

428 sign = "+" if self.im >= 0 else "-" 

429 return f"{self.re:.12g} {sign} {abs(self.im):.12g}i"