FEAT: Added support for float coefficients (#13)

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

README.md

@@ -9,7 +9,7 @@ A Python library for representing, manipulating, and solving polynomial equation
 
 ## Key Features
 
-* **Create and Manipulate Polynomials**: Easily define polynomials of any degree and perform arithmetic operations like addition, subtraction, multiplication, and scaling.
+* **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.
 * **Genetic Algorithm Solver**: Find approximate real roots for complex polynomials where analytical solutions are difficult or impossible.
 * **CUDA Accelerated**: Leverage NVIDIA GPUs for a massive performance boost when finding roots in large solution spaces.
 * **Analytical Solvers**: Includes standard, exact solvers for simple cases (e.g., `quadratic_solve`).
@@ -44,6 +44,7 @@ Here is a simple example of how to define a quadratic function, find its propert
 from polysolve import Function, GA_Options, quadratic_solve
 
 # 1. Define the function f(x) = 2x^2 - 3x - 5
+#    Coefficients can be integers or floats.
 f1 = Function(largest_exponent=2)
 f1.set_constants([2, -3, -5])
 

pyproject.toml

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

src/polysolve/__init__.py

@@ -12,8 +12,32 @@ try:
 except ImportError:
     _CUPY_AVAILABLE = False
 
-# The CUDA kernel for the fitness function
+# The CUDA kernels for the fitness function
-_FITNESS_KERNEL = """
+_FITNESS_KERNEL_FLOAT = """
+extern "C" __global__ void fitness_kernel(
+    const double* coefficients, 
+    int num_coefficients, 
+    const double* x_vals, 
+    double* ranks, 
+    int size, 
+    double y_val)
+{
+    int idx = threadIdx.x + blockIdx.x * blockDim.x;
+    if (idx < size)
+    {
+        double ans = 0;
+        int lrgst_expo = num_coefficients - 1;
+        for (int i = 0; i < num_coefficients; ++i)
+        {
+            ans += coefficients[i] * pow(x_vals[idx], (double)(lrgst_expo - i));
+        }
+
+        ans -= y_val;
+        ranks[idx] = (ans == 0) ? 1.7976931348623157e+308 : fabs(1.0 / ans);
+    }
+}
+"""
+_FITNESS_KERNEL_INT = """
 extern "C" __global__ void fitness_kernel(
     const long long* coefficients, 
     int num_coefficients, 
@@ -38,6 +62,7 @@ extern "C" __global__ void fitness_kernel(
 }
 """
 
+
 @dataclass
 class GA_Options:
     """
@@ -76,13 +101,14 @@ class Function:
         self.coefficients: Optional[np.ndarray] = None
         self._initialized = False
 
-    def set_coeffs(self, coefficients: List[int]):
+    def set_coeffs(self, coefficients: List[Union[int, float]]):
         """
         Sets the coefficients of the polynomial.
 
         Args:
-            coefficients (List[int]): A list of integer coefficients. The list size
+            coefficients (List[Union[int, float]]): A list of integer or float
-                                   must be largest_exponent + 1.
+                                                   coefficients. The list size
+                                                   must be largest_exponent + 1.
 
         Raises:
             ValueError: If the input is invalid.
@@ -96,7 +122,13 @@ class Function:
         if coefficients[0] == 0 and self._largest_exponent > 0:
             raise ValueError("The first constant (for the largest exponent) cannot be 0.")
         
-        self.coefficients = np.array(coefficients, dtype=np.int64)
+        # Check if any coefficient is a float
+        is_float = any(isinstance(c, float) for c in coefficients)
+
+        # Choose the dtype based on the input
+        target_dtype = np.float64 if is_float else np.int64
+
+        self.coefficients = np.array(coefficients, dtype=target_dtype)
         self._initialized = True
 
     def _check_initialized(self):
@@ -275,11 +307,16 @@ class Function:
 
     def _solve_x_cuda(self, y_val: float, options: GA_Options) -> np.ndarray:
         """Genetic algorithm implementation using CuPy (GPU/CUDA)."""
-        # Load the raw CUDA kernel
-        fitness_gpu = cupy.RawKernel(_FITNESS_KERNEL, 'fitness_kernel')
         
-        # Move coefficients to GPU
+        # Check the dtype of our coefficients array
-        d_coefficients = cupy.array(self.coefficients, dtype=cupy.int64)
+        if self.coefficients.dtype == np.float64:
+            fitness_gpu = cupy.RawKernel(_FITNESS_KERNEL_FLOAT, 'fitness_kernel')
+            d_coefficients = cupy.array(self.coefficients, dtype=cupy.float64)
+        elif self.coefficients.dtype == np.int64:
+            fitness_gpu = cupy.RawKernel(_FITNESS_KERNEL_INT, 'fitness_kernel')
+            d_coefficients = cupy.array(self.coefficients, dtype=cupy.int64)
+        else:
+            raise TypeError(f"Unsupported dtype for CUDA solver: {self.coefficients.dtype}")
         
         # Create initial random solutions on the GPU
         d_solutions = cupy.random.uniform(
@@ -337,12 +374,16 @@ class Function:
             power = self._largest_exponent - i
             
             # Coefficient part
-            if c == 1 and power != 0:
+            coeff_val = c
+            if c == int(c):
+                coeff_val = int(c)
+
+            if coeff_val == 1 and power != 0:
                 coeff = ""
-            elif c == -1 and power != 0:
+            elif coeff_val == -1 and power != 0:
                 coeff = "-"
             else:
-                coeff = str(c)
+                coeff = str(coeff_val)
 
             # Variable part
             if power == 0:
@@ -356,7 +397,7 @@ class Function:
             sign = ""
             if i > 0:
                 sign = " + " if c > 0 else " - "
-                coeff = str(abs(c))
+                coeff = str(abs(coeff_val))
                 if abs(c) == 1 and power != 0:
                     coeff = "" # Don't show 1 for non-constant terms