radix 2 dif fft work in progress, adding twiddle factors remaining

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,87 @@ def direct_form_2_iir(
     output = Output()
     output <<= add
     return SFG([input_op], [output], name=Name(name))
+
+
+def radix_2_dif_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 < 0:
+        raise ValueError("Points must be positive number.")
+    if points & (points - 1) != 0:
+        raise ValueError("Points must be a power of two.")
+
+    inputs = []
+    for i in range(points):
+        inputs.append(Input(name=f"Input: {i}"))
+
+    ports = inputs
+    number_of_stages = int(np.log2(points))
+
+    twiddles = _generate_twiddles(points, number_of_stages)
+    print(twiddles)
+
+    for stage in range(number_of_stages):
+        ports = _construct_dif_fft_stage(ports, number_of_stages, stage)
+
+    ports = _get_bit_reversed_ports(ports)
+    outputs = []
+    for i, port in enumerate(ports):
+        outputs.append(Output(port, name=f"Output: {i}"))
+
+    return SFG(inputs=inputs, outputs=outputs)
+
+
+def _construct_dif_fft_stage(ports_from_previous_stage, number_of_stages, stage):
+    ports = ports_from_previous_stage.copy()
+    number_of_butterflies = len(ports) // 2
+    number_of_groups = 2 ** (stage)
+    group_size = number_of_butterflies // number_of_groups
+
+    for group_index in range(number_of_groups):
+        for bf_index in range(group_size):
+            input1_index = group_index * 2 * group_size + bf_index
+            input2_index = input1_index + group_size
+
+            input1 = ports[input1_index]
+            input2 = ports[input2_index]
+
+            butterfly = Butterfly(input1, input2, name=f"bf: {stage*4+bf_index}")
+            output1, output2 = butterfly.outputs
+
+            ports[input1_index] = output1
+            ports[input2_index] = output2
+
+    return ports
+
+
+def _get_bit_reversed_number(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 _get_bit_reversed_ports(ports):
+    num_of_ports = len(ports)
+    bits = int(np.log2(num_of_ports))
+    return [ports[_get_bit_reversed_number(i, bits)] for i in range(num_of_ports)]
+
+
+def _generate_twiddles(points, number_of_stages):
+    twiddles = []
+    for stage in range(1, number_of_stages + 1):
+        stage_twiddles = []
+        for k in range(points // 2 ** (stage)):
+            twiddle = np.exp(-1j * 2 * np.pi * stage * k / points)
+            print("STAGE:", stage, "k:", k, "TW:", twiddle)
+            stage_twiddles.append(twiddle)
+        twiddles.append(stage_twiddles)
+    return twiddles