chore: update sigma working with complex number

.pre-commit-config.yaml

@@ -6,3 +6,8 @@ repos:
         entry: bash ./build.sh clean
         language: system
         pass_filenames: false
+
+  - repo: https://github.com/psf/black-pre-commit-mirror
+    rev: 25.1.0
+    hooks:
+      - id: black

build.sh

@@ -2,7 +2,7 @@
 
 make_libs () {
   echo "make project"
-  uv run setup.py build_ext --inplace;
+  uv run setup.py build_ext --inplace
 }
 
 make_as_package() {

mega/__init__.py

@@ -1,5 +1,5 @@
 from .op.function import Gamma, Haversine, JordanTotient
-from .op.arithmethic import SigmaZ,  EulerPhi
+from .op.arithmethic import SigmaZ, EulerPhi
 from .op.tensor import Tensor
 
 # utils folders

mega/op/arithmethic.pyx

@@ -46,9 +46,9 @@ cdef class SigmaZ:
     4
     """
     cdef int n
-    cdef int z
+    cdef object z
 
-    def __cinit__(self, int n, int z):
+    def __cinit__(self, int n, object z):
         """
         constructor that validating input before storing value
 
@@ -60,56 +60,84 @@ cdef class SigmaZ:
         """
         if n <= 0:
             raise ValueError("only acc positive integer")
-        if z < 0:
-            raise ValueError("exponent must be non-negative")
+        if isinstance(z, (float, int)):
+            self.z = <double>z
+        elif isinstance(z, complex):
+            self.z = z
+        else:
+            raise TypeError("exponent must be numeric")
 
         self.n = n
         self.z = z
 
     def __dealloc__(self):
         pass
 
-    cpdef long compute(self):
+    cpdef object compute(self):
         """
         compute σ_z(n), the sum of the z-th powers of all positive divisors of n
 
         Return:
             (long): computed value of σ_z(n)
         """
-        cdef long total = 0
         cdef int i = 1
-        cdef int limit = <int>sqrt(self.n)
         cdef int oth_div
-        cdef long power_i, power_oth
-
-        # σ₀(n) count the number of positive divisor of n
+        cdef double power_i, power_oth
+        cdef double z_real
+        cdef complex z_complex
+        cdef long total_real = 0
+
+        # when z == 0, just counting number of divisor
         if self.z == 0:
+            total_real = 0
+            while i * i <= self.n:
+                if self.n % i == 0:
+                    oth_div = self.n // i
+                    if i == oth_div:
+                        total_real += 1 # perfect square
+                    else:
+                        total_real += 2 # two distinct divisor
+                i += 1
+            return total_real
+
+        # try convert z to double for optimal compute
+        try:
+            z_real = <double>self.z
+        except TypeError:
+            z_real = -1.0 # mark as invalid for fallback logic
+
+        # if z is numeric and >= 0
+        if z_real >= 0:
+            total_real = 0
             while i * i <= self.n:
                 if self.n % i == 0:
                     oth_div = self.n // i
+                    power_i_real = <long>(pow(i, z_real))
+                    power_oth_real = <long>(pow(oth_div, z_real))
+
                     if i == oth_div:
-                        total += 1 # perfect square: count once
+                        total_real += power_i_real # add once for perfect square
                     else:
-                        total += 2 # two distrint divisor: i and oth_div
+                        total_real += power_i_real + power_oth_real
                 i += 1
-            return total
+            return total_real
+
+
+        import numpy as np
+        z_complex = self.z
+        total_complex = 0j
 
-        # σ_z(n) = Σ d^z over all d dividing n
+        # same loop structure, using numpy for complex support
         while i * i <= self.n:
             if self.n % i == 0:
                 oth_div = self.n // i
-
-                # compute power using pow
-                power_i = <long>(pow(i, self.z))
-                power_oth = <long>(pow(oth_div, self.z))
-
-                if i == oth_div:
-                    total += power_i
-                else:
-                    total += power_i + power_oth
+                total_complex += np.power(i, z_complex)
+                if i != oth_div:
+                    total_complex += np.power(oth_div, z_complex)
             i += 1
-        return total
 
+        return total_complex
+
     def __repr__(self):
         """
         Return string representation of the object

mega/op/function.pyx

@@ -26,14 +26,15 @@
 from mega.utils.constant cimport PI_NUMBER, SQRT_PI
 from libc.math cimport sqrt, exp, sin, pow, cos
 
+
 def prime_factors(int n, bint unique=False) -> list[int]:
     """"
     return prime factor of positive integer n
 
     Parameter:
         n (int): integer to be factorized, must greater than zero
         unique (bool): if True, return only distinct prime factors
-        
+
     Return:
         list[int]: list of prime factor of `n`, sorted in increasing order
 
@@ -47,7 +48,7 @@ def prime_factors(int n, bint unique=False) -> list[int]:
         raise ValueError("only positive integer are acc")
     cdef int i = 2
     cdef list factors = []
-    
+
     # trie division algorithm
     while i * i <= n:
         while n % i == 0:
@@ -69,6 +70,7 @@ def prime_factors(int n, bint unique=False) -> list[int]:
     else:
         return factors
 
+
 cdef class Haversine:
     """
     compute the haversine function of an angle θ
@@ -116,7 +118,7 @@ cdef class Haversine:
     cpdef double compute(self):
         """
         compute haversine function
-        
+
         Return:
             (double): value of haversine(theta) = (1 - cos(theta)) / 2
         """
@@ -127,7 +129,7 @@ cdef class Haversine:
         """
         get the current angle stored in the haversine instance
         usefull for debug or chaining operations
-        
+
         Return:
             (double): current theta (in radians)
         """
@@ -195,7 +197,7 @@ cdef class Gamma:
         if z < 0.5:
             y = 1.0 - z
             neg = 1
-        
+
         for i in range(8):
             x += self.LANCZOS_COEFF[i] / (y + i)
 
@@ -205,7 +207,7 @@ cdef class Gamma:
             return PI_NUMBER / (sin(PI_NUMBER * z) * tmp)
         else:
             return tmp
-    
+
     def __repr__(self) -> str:
         return f"Gamma({self.point}) = {Gamma(self.point).compute()}"
 

mega/op/tensor.pyx

@@ -48,7 +48,7 @@ cdef class Tensor:
     cdef double *double_data
     cdef float *float_data
 
-    def __cinit__(self, tuple shape, list data=None, str dtype='long'):
+    def __cinit__(self, tuple shape, list data=None, str dtype="long"):
         """
         initialize tensor with given shape, optional data, and specified data type
 
@@ -59,7 +59,7 @@ cdef class Tensor:
                                     long, int, float, double, default was long
         """
         self._dtype = dtype.lower()
-        if self._dtype not in ('int', 'long', 'double', 'float'):
+        if self._dtype not in ("int", "long", "double", "float"):
             raise ValueError(f"unsupported dtype: {dtype}")
 
         self.ndim = len(shape)
@@ -76,13 +76,13 @@ cdef class Tensor:
             self.shape[i] = dim
             self.size *= dim
 
-        if self._dtype == 'int':
+        if self._dtype == "int":
             self.int_data = [0] * self.size
             if data is not None:
                 for i in range(min(self.size, len(data))):
                     self.int_data[i] = data[i]
 
-        elif self._dtype == 'long':
+        elif self._dtype == "long":
             self.long_data = <long*>malloc(self.size * sizeof(long))
             if not self.long_data:
                 raise MemoryError("failed to allocate long data")
@@ -93,7 +93,7 @@ cdef class Tensor:
                 for i in range(self.size):
                     self.long_data[i] = data[i]
 
-        elif self._dtype == 'double':
+        elif self._dtype == "double":
             self.double_data = <double*>malloc(self.size * sizeof(double))
             if not self.double_data:
                 raise MemoryError("failed to allocate double data")
@@ -112,11 +112,11 @@ cdef class Tensor:
         """
         if self.shape is not NULL:
             free(self.shape)
-        if hasattr(self, 'long_data') and self.long_data is not NULL:
+        if hasattr(self, "long_data") and self.long_data is not NULL:
             free(self.long_data)
-        if hasattr(self, 'double_data') and self.double_data is not NULL:
+        if hasattr(self, "double_data") and self.double_data is not NULL:
             free(self.double_data)
-        if hasattr(self, 'float_data') and self.float_data is not NULL:
+        if hasattr(self, "float_data") and self.float_data is not NULL:
             free(self.float_data)
 
     cpdef str dtype(self):
@@ -158,13 +158,13 @@ cdef class Tensor:
                 raise IndexError(f"Index out of bounds at axis {i}: {pos}")
             offset += pos * self._compute_stride(i)
 
-        if self._dtype == 'int':
+        if self._dtype == "int":
             return self.int_data[offset]
-        elif self._dtype == 'long':
+        elif self._dtype == "long":
             return self.long_data[offset]
-        elif self._dtype == 'double':
+        elif self._dtype == "double":
             return self.double_data[offset]
-        elif self._dtype == 'float':
+        elif self._dtype == "float":
             return self.float_data[offset]
 
     def __setitem__(self, object idxs, long value):
@@ -195,13 +195,13 @@ cdef class Tensor:
                 raise IndexError(f"Index out of bounds at axis {i}: {pos}")
             offset += pos * self._compute_stride(i)
 
-        if self._dtype == 'int':
+        if self._dtype == "int":
             self.int_data[offset] = value
-        elif self._dtype == 'long':
+        elif self._dtype == "long":
             self.long_data[offset] = value
-        elif self._dtype == 'double':
+        elif self._dtype == "double":
             self.double_data[offset] = value
-        elif self._dtype == 'float':
+        elif self._dtype == "float":
             self.float_data[offset] = value
 
     cdef long _compute_stride(self, int axis):
@@ -233,21 +233,21 @@ cdef class Tensor:
         shape_str = ", ".join(shape_list)
 
         data_preview = []
-        if self._dtype == 'int':
+        if self._dtype == "int":
             for i in range(min(5, self.size)):
                 data_preview.append(str(self.int_data[i]))
-        elif self._dtype == 'long':
+        elif self._dtype == "long":
             for i in range(min(5, self.size)):
                 data_preview.append(str(self.long_data[i]))
-        elif self._dtype == 'double':
+        elif self._dtype == "double":
             for i in range(min(5, self.size)):
                 data_preview.append(f"{self.double_data[i]:.4f}")
-        elif self._dtype == 'float':
+        elif self._dtype == "float":
             for i in range(min(5, self.size)):
                 data_preview.append(f"{self.float_data[i]:.4f}")
 
         data_str = ", ".join(data_preview)
         if self.size > 5:
-            data_str += ', ...'
+            data_str += ", ..."
 
         return f"Tensor(shape=({shape_str}), dtype={self._dtype}, data=[{data_str}])"

pyproject.toml

@@ -10,13 +10,13 @@ requires-python = ">=3.13"
 dependencies = [
     "cython>=3.1.0",
     "numpy>=2.2.6",
-    "pre-commit>=4.2.0",
     "setuptools>=80.7.1",
 ]
 
 [dependency-groups]
 dev = [
     "black>=25.1.0",
     "matplotlib>=3.10.3",
+    "pre-commit>=4.2.0",
     "pytest>=8.3.5",
 ]

setup.py

@@ -19,9 +19,7 @@
         include_dirs=[np.get_include()],
     ),
     Extension(
-        "mega.op.tensor",
-        ["mega/op/tensor.pyx"],
-        include_dirs=[np.get_include()]
+        "mega.op.tensor", ["mega/op/tensor.pyx"], include_dirs=[np.get_include()]
     ),
 ]