work in progress, started on fft function

b_asic/sfg_generators.py

@@ -10,6 +10,7 @@ import numpy as np
 
 from b_asic.core_operations import (
     Addition,
+    Butterfly,
     ConstantMultiplication,
     Name,
     SymmetricTwoportAdaptor,
@@ -371,3 +372,82 @@ def direct_form_2_iir(
     output = Output()
     output <<= add
     return SFG([input_op], [output], name=Name(name))
+
+
+def radix_2_fft(
+    points: int,
+    mult_properties: Optional[Union[Dict[str, int], Dict[str, Dict[str, int]]]] = None,
+    add_properties: Optional[Union[Dict[str, int], Dict[str, Dict[str, int]]]] = None,
+) -> SFG:
+    if points % 2 != 0:
+        raise ValueError("Points must be a power of two.")
+
+    # add all inputs so that x0 gets rendered on top
+    inputs = []
+    for i in range(points - 1, -1, -1):
+        inputs.append(Input(name=f"Input: {i}"))
+
+    # bits = int(np.log2(points))
+    # bit_reversed_indices = [_get_bit_reversed(i, bits) for i in list(range(points))]
+    # for stage in range(bits):
+
+    # create and connect all stages in the FFT
+    proceeding_stage = inputs
+    for stage in range(int(np.log2(points))):
+        split_size = stage // (2**stage)
+        # create and connect butterflies
+        butterflies = []
+        for i in range(split_size // 2):
+            op_a = proceeding_stage[i]
+            op_b = proceeding_stage[i + points // 2]
+            butterflies.append(Butterfly(op_a, op_b))
+        # proceeding_stage = butterflies
+
+    outputs = []
+    for i in range(points - 1, -1, -1):
+        outputs.append(Output(proceeding_stage[i].output(0)))
+    # print("Points:", points, "Bits:", bits, "Rev:", bit_reversed_indices)
+    # chunks = [bit_reversed_indices[i:i+2] for i in range(0, len(bit_reversed_indices), 2)]
+    # print("Chunks:", chunks)
+
+    # inputs = []
+    # for chunk in reversed(chunks):
+    #     for i in chunk:
+    #         inputs.append(Input(name=f"Input: {i}"))
+
+    # # for stage in range(bits):
+    # butterflies = []
+    # for i in range(points//2):
+    #     butterflies.append(Butterfly(inputs[2*i], inputs[2*i+1]))
+
+    # outputs = []
+    # for bf in butterflies:
+    #     outputs.append(Output(bf.output(0)))
+    #     outputs.append(Output(bf.output(1)))
+
+    return SFG(inputs=inputs, outputs=outputs)
+
+
+def _get_butterfly_index(index, number_of_bits, stage) -> int:
+    shifted_index = _circular_right_shift(index, number_of_bits, stage)
+    bit_reversed_index = _get_bit_reversed(shifted_index, number_of_bits)
+    return int(bit_reversed_index, 2)
+
+
+def _get_bit_reversed(number, number_of_bits) -> int:
+    reversed_number = 0
+    for i in range(number_of_bits):
+        # mask out the current bit
+        current_bit = (number >> i) & 1
+        # compute the position of the current bit in the reversed string
+        reversed_pos = number_of_bits - 1 - i
+        # place the current bit in that position
+        reversed_number |= current_bit << reversed_pos
+    return reversed_number
+
+
+def _circular_right_shift(number: int, number_of_bits: int, shift: int) -> int:
+    bits = bin(number)[2:]
+    shift = shift % number_of_bits
+    shifted_bits = bits[-shift:] + bits[:-shift]
+    return int(shifted_bits, 2)