complex data type

[toaster-robotics]

Are there any plans to incorporate the complex data type into warp? It would be very helpful for electro-magnetic use cases.

[toaster-robotics]

I just noticed that there's a note in the Limitations of the documentation stating that complex numbers are not currently supported. I've spent the weekend porting some code over to warp and it's veeeeery fast! I implemented what I need so far for complex numbers. It's just not very clean at the moment. Here is a portion of what I have:

@wp.struct
class ComplexFloat:
    real: wp.float32
    imag: wp.float32

    def __init__(self, real: wp.float32, imag: wp.float32) -> None:
        self.real = real
        self.imag = imag

@wp.struct
class ComplexVec3:
    real: wp.vec3
    imag: wp.vec3

    def __init__(self, real: wp.vec3, imag: wp.vec3) -> None:
        self.real = real
        self.imag = imag

@wp.func
def complex_add(c1: ComplexFloat, c2: ComplexFloat) -> ComplexFloat:
    return ComplexFloat(c1.real + c2.real, c1.imag + c2.imag)

@wp.func
def complex_add(c1: ComplexFloat, c2: wp.float32) -> ComplexFloat:
    return ComplexFloat(c1.real + c2, c1.imag)

@wp.func
def complex_add(c1: ComplexVec3, c2: ComplexVec3) -> ComplexVec3:
    return ComplexVec3(c1.real + c2.real, c1.imag + c2.imag)

I tried doing add, multiply etc inside the object (def __mul__(self, other)) but warp doesn't see the function. For example, here is what I get for __neg__:

@wp.struct
class ComplexFloat:
    real: wp.float32
    imag: wp.float32

    def __init__(self, real: wp.float32, imag: wp.float32) -> None:
        self.real = real
        self.imag = imag

    def __neg__(self):
        return ComplexFloat(-self.real, -self.imag)

    raise WarpCodegenError(
warp.codegen.WarpCodegenError: Error while parsing function "xxxx" at /home/jason/workspace/xxxxx.py:120:
            E = -ComplexFloat(0.0, distance)
;Couldn't find function overload for 'neg' that matched inputs with types: [MathLib.ComplexFloat]

I tried adding @wp.func to __neg__ but I get a different error:

warp.codegen.WarpCodegenError: Incomplete argument annotations on function __neg__

Any suggestions on how to get this working?

[daedalus5]

@wp.struct objects in Warp are just meant to be containers storing a mixture of variables of various data types, for convenience. So any complex number manipulations should be abstracted to standalone @wp.func functions or kernels. For instance:

import warp as wp


@wp.struct
class ComplexFloat:
    real: wp.float32
    imag: wp.float32

@wp.func
def complex_multiply(c1: ComplexFloat, c2: ComplexFloat):
    a = c1.real
    b = c1.imag
    c = c2.real
    d = c2.imag
    return ComplexFloat(a*c - b*d, a*d + b*c)

@wp.func
def complex_conjugate(c: ComplexFloat):
    return ComplexFloat(c.real, -c.imag)

@wp.func
def complex_length(c: ComplexFloat):
    c_conjugate = complex_conjugate(c)
    return wp.sqrt(complex_multiply(c, c_conjugate).real)

@wp.kernel
def complex_kernel(complex_nums: wp.array(dtype=ComplexFloat), lengths: wp.array(dtype=float)):
    tid = wp.tid()
    lengths[tid] = complex_length(complex_nums[tid])

N = 1

# create structs
inputs = []
for _ in range(N):
    inputs.append(ComplexFloat())

# initalize Data
for i in range(N):
    inputs[i].real = 0.1
    inputs[i].imag = 0.1

inputs = wp.array(inputs, dtype=ComplexFloat)
lengths = wp.zeros(N, dtype=float)

wp.launch(complex_kernel, N, inputs=[inputs], outputs=[lengths])

print(lengths)

You are free to create ComplexFloats on the fly in warp kernels and functions, but if you intend to store data and access it after launching kernels, you'll need to pre-allocate warp input/output arrays.