FEAT: Added support for float coefficients

src/polysolve/__init__.py

@@ -15,7 +15,7 @@ except ImportError:
 # The CUDA kernel for the fitness function
 _FITNESS_KERNEL = """
 extern "C" __global__ void fitness_kernel(
-    const long long* coefficients, 
+    const double* coefficients, 
     int num_coefficients, 
     const double* x_vals, 
     double* ranks, 
@@ -76,13 +76,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
-                                   must be largest_exponent + 1.
+            coefficients (List[Union[int, float]]): A list of integer or float
+                                                   coefficients. The list size
+                                                   must be largest_exponent + 1.
 
         Raises:
             ValueError: If the input is invalid.
@@ -96,7 +97,7 @@ 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)
+        self.coefficients = np.array(coefficients, dtype=np.float64)
         self._initialized = True
 
     def _check_initialized(self):
@@ -279,7 +280,7 @@ class Function:
         fitness_gpu = cupy.RawKernel(_FITNESS_KERNEL, 'fitness_kernel')
         
         # Move coefficients to GPU
-        d_coefficients = cupy.array(self.coefficients, dtype=cupy.int64)
+        d_coefficients = cupy.array(self.coefficients, dtype=cupy.float64)
         
         # Create initial random solutions on the GPU
         d_solutions = cupy.random.uniform(
@@ -337,12 +338,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 +361,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