added some complex functions

python/polars_ds/complex_ext.py

@@ -1,7 +1,9 @@
 """
 Tools for dealing with complex numbers columns inside Polars dataframe.
 
-Complex number columns are represented as a column of size-2 lists. An element will look like [re, im].
+Complex number columns are represented as a column of size-2 lists. By default, an element will look like [re, im],
+which is in coordinate form. All operations (except powi, which turns it into polar form internally) assume the number 
+is in coordinate form. There is a to_coord function provided for complex numbers in polar form [r, theta].
 """
 
 import polars as pl
@@ -62,6 +64,16 @@ def theta(self, degree: bool = False) -> pl.Expr:
                 .otherwise(pl.lit(math.nan, dtype=pl.Float64))
             )
 
+    def to_polar(self) -> pl.Expr:
+        """Turns a complex number in coordinate form into polar form."""
+        return pl.concat_list(self.modulus(), self.theta())
+
+    def to_coord(self) -> pl.Expr:
+        """Turns a complex number in polar form into coordinate form."""
+        r = self._expr.list.first()
+        theta = self._expr.list.last()
+        return pl.concat_list(r * theta.cos(), r * theta.sin())
+
     def conj(self) -> pl.Expr:
         """Returns complex conjugate."""
         return pl.concat_list(self._expr.list.first(), -self._expr.list.last())
@@ -150,3 +162,28 @@ def div(self, other: Union[float, complex, pl.Expr]) -> pl.Expr:
             new_real = x * x_inv - y * y_inv
             new_imag = x * y_inv + y * x_inv
             return pl.concat_list(new_real, new_imag)
+
+    def mul_by_i(self) -> pl.Expr:
+        """Multiplies self by i."""
+        x = self._expr.list.first()
+        y = self._expr.list.last()
+        return pl.concat_list(-y, x)
+
+    def pow(self, x: float) -> pl.Expr:
+        """Raises a complex number to the x power."""
+        if x == 0.0:
+            return pl.concat_list(
+                pl.when(self.modulus() == 0.0).then(math.nan).otherwise(1.0),
+                pl.lit(0.0, dtype=pl.Float64),
+            )
+        elif x == 1.0:
+            return self._expr
+        elif x == 2.0:
+            return self.mul(self._expr)
+        elif x == -1.0:
+            return self.inv()
+        else:
+            polar = self.to_polar()
+            r = polar.list.first()
+            theta = polar.list.last()
+            return pl.concat_list(r.pow(x) * (x * theta).cos(), r.pow(x) * (x * theta).sin())

tests/test.ipynb

@@ -23,36 +23,48 @@
     "from scipy import signal\n",
     "rng = np.random.default_rng()\n",
     "\n",
-    "npoints = 100_000\n",
+    "npoints = 10\n",
     "\n",
     "noise = rng.standard_normal(npoints)\n",
     "\n",
     "x = 3 + 2*np.linspace(0, 1, npoints) + noise\n",
     "\n",
     "df = pl.DataFrame({\n",
-    "    \"test\": x \n",
+    "    \"test\": x,\n",
+    "    \"entity_id\": [1] *  5 + [2] * 5\n",
     "})\n",
     "df.head()"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "b5263aa7",
+   "id": "a0a9731d",
    "metadata": {},
    "outputs": [],
    "source": [
-    "%timeit signal.detrend(x)"
+    "import polars as pl\n",
+    "import polars_ds\n",
+    "\n",
+    "df = pl.DataFrame({\n",
+    "    \"test\": x,\n",
+    "    \"entity_id\": [1] *  5 + [2] * 5\n",
+    "})\n",
+    "# And is 5x faster than Scipy.signal detrend on larger time series\n",
+    "df.select(\n",
+    "    pl.col(\"entity_id\"),\n",
+    "    pl.col(\"test\").num_ext.detrend().over(pl.col(\"entity_id\")).alias(\"test_detrended\") # linear detrend\n",
+    ")"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "7880e26d",
+   "id": "b5263aa7",
    "metadata": {},
    "outputs": [],
    "source": [
-    "%timeit df.select(pl.col(\"test\").num_ext.detrend2())"
+    "%timeit signal.detrend(x)"
    ]
   },
   {