v0.7.0 - Complex Number Support (#25)

Reviewed-on: #25
Co-authored-by: Jonathan Rampersad <rampersad.jonathan@gmail.com>
Co-committed-by: Jonathan Rampersad <rampersad.jonathan@gmail.com>

README.md

@@ -12,6 +12,7 @@ A Python library for representing, manipulating, and solving polynomial equation
 ## Key Features
 
 * **Numerically Stable Solver**: Makes complex calculations **practical**. Leverage your GPU to power the robust genetic algorithm, solving high-degree polynomials accurately in a reasonable timeframe.
+* **Complex Number Support**: Fully supports complex coefficients and finding roots in the complex plane (e.g., $x^2 + 1 = 0 \to \pm i$).
 * **Numba Accelerated CPU Solver**: The default genetic algorithm is JIT-compiled with Numba for high-speed CPU performance, right out of the box.
 * **CUDA Accelerated**: Leverage NVIDIA GPUs for a massive performance boost when finding roots in large solution spaces.
 * **Create and Manipulate Polynomials**: Easily define polynomials of any degree using integer or float coefficients, and perform arithmetic operations like addition, subtraction, multiplication, and scaling.
@@ -75,12 +76,19 @@ roots_analytic = f1.quadratic_solve()
 print(f"Analytic roots: {sorted(roots_analytic)}")
 # > Analytic roots: [-1.0, 2.5]
 
-# 6. Find roots with the genetic algorithm (Numba CPU)
+# 6. Find REAL roots with the genetic algorithm (Numba CPU)
-#    This is the default, JIT-compiled CPU solver.
+#    This is the default, JIT-compiled CPU solver.
 ga_opts = GA_Options(num_of_generations=20)
 roots_ga = f1.get_real_roots(ga_opts, use_cuda=False)
-print(f"Approximate roots from GA: {roots_ga[:2]}")
+print(f"Approximate real roots: {roots_ga[:2]}")
-# > Approximate roots from GA: [-1.000..., 2.500...]
+# > Approximate real roots: [-1.000..., 2.500...]
+
+# 7. Find ALL roots (Real + Complex)
+#    Use get_roots() to search the complex plane.
+f_complex = Function(2, [1, 0, 1]) # x^2 + 1
+roots_all = f_complex.get_roots(ga_opts)
+print(f"Approximate complex roots: {roots_all}")
+# > Approximate complex roots: [-1.00...j, 1.00...j]
 
 # If you installed a CUDA extra, you can run it on the GPU:
 # roots_ga_gpu = f1.get_real_roots(ga_opts, use_cuda=True)
@@ -114,7 +122,10 @@ ga_robust_search = GA_Options(
     # Increase the crossover blend factor to 0.75.
     # This allows new solutions to be created further
     # away from their parents, increasing exploration.
-    blend_alpha=0.75
+    blend_alpha=0.75,
+
+    # Enable complex root finding (default is True)
+    find_complex=True
 )
 
 # Pass the custom options to the solver

pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 [project]
 # --- Core Metadata ---
 name = "polysolve"
-version = "0.6.1"
+version = "0.7.0"
 authors = [
   { name="Jonathan Rampersad", email="jonathan@jono-rams.work" },
 ]

tests/test_polysolve.py

@@ -37,6 +37,17 @@ def m_func_2() -> Function:
     f.set_coeffs([5, -4])
     return f
 
+@pytest.fixture
+def base_func():
+    f = Function(2)
+    f.set_coeffs([1, 2, 3])
+    return f
+
+@pytest.fixture
+def complex_func():
+    f = Function(2, [1, 2, 2])
+    return f
+
 # --- Core Functionality Tests ---
 
 def test_solve_y(quadratic_func):
@@ -95,6 +106,32 @@ def test_function_multiplication(m_func_1, m_func_2):
     assert result.largest_exponent == 3
     assert np.array_equal(result.coefficients, [10, 7, -7, -4])
 
+def test_equality(base_func):
+    """Tests the __eq__ method for the Function class."""
+    
+    # 1. Test for equality with a new, identical object
+    f_identical = Function(2)
+    f_identical.set_coeffs([1, 2, 3])
+    assert base_func == f_identical
+
+    # 2. Test for inequality (different coefficients)
+    f_different = Function(2)
+    f_different.set_coeffs([1, 9, 3])
+    assert base_func != f_different
+
+    # 3. Test for inequality (different degree)
+    f_diff_degree = Function(1)
+    f_diff_degree.set_coeffs([1, 2])
+    assert base_func != f_diff_degree
+
+    # 4. Test against a different type
+    assert base_func != "some_string"
+    assert base_func != 123
+
+    # 5. Test against an uninitialized Function
+    f_uninitialized = Function(2)
+    assert base_func != f_uninitialized
+
 # --- Genetic Algorithm Root-Finding Tests ---
 
 def test_get_real_roots_numpy(quadratic_func):
@@ -130,3 +167,36 @@ def test_get_real_roots_cuda(quadratic_func):
     # Verify that the CUDA implementation also finds the correct roots within tolerance.
     npt.assert_allclose(np.sort(roots), np.sort(expected_roots), atol=1e-2)
 
+def test_get_roots_numpy(complex_func):
+    """
+    Tests that the NumPy-based genetic algorithm approximates the roots correctly.
+    """
+    # Using more generations for higher accuracy in testing
+    ga_opts = GA_Options(num_of_generations=50, data_size=200000, selection_percentile=0.66, root_precision=3)
+    
+    roots = complex_func.get_roots(ga_opts, use_cuda=False)
+    
+    # Check if the algorithm found values close to the two known roots.
+    # We don't know which order they'll be in, so we check for presence.
+    expected_roots = np.array([-1.0-1.j, -1.0+1.j])
+    
+    npt.assert_allclose(np.sort(roots), np.sort(expected_roots), atol=1e-2)
+
+
+@pytest.mark.skipif(not _CUPY_AVAILABLE, reason="CuPy is not installed, skipping CUDA test.")
+def test_get_roots_cuda(complex_func):
+    """
+    Tests that the CUDA-based genetic algorithm approximates the roots correctly.
+    This test implicitly verifies that the CUDA kernel is functioning.
+    It will be skipped automatically if CuPy is not available.
+    """
+    
+    ga_opts = GA_Options(num_of_generations=50, data_size=200000, selection_percentile=0.66, root_precision=3)
+    
+    roots = complex_func.get_roots(ga_opts, use_cuda=True)
+    
+    expected_roots = np.array([-1.0-1.j, -1+1.j])
+    
+    # Verify that the CUDA implementation also finds the correct roots within tolerance.
+    npt.assert_allclose(np.sort(roots), np.sort(expected_roots), atol=1e-2)
+