start adding beamforming stuff

plots/09_beamforming.py

@@ -0,0 +1,123 @@
+# %% imports
+import numpy as np
+from utils import AnimatedPlot, dir_assets
+
+
+# %%
+class PcolorPlot(AnimatedPlot):
+    def __init__(self, elements: int, angle_deg: float = 0, spacing_lambda: float = 0.5, **kwargs):
+        super().__init__(**kwargs)
+
+        self.ax.axes.xaxis.set_visible(False)
+        self.ax.axes.yaxis.set_visible(False)
+        self.ax.set_aspect("equal")
+
+        self.spacing_lambda = spacing_lambda
+        x_range = spacing_lambda * (elements - 1)
+        self.element_x = np.linspace(-x_range / 2, x_range / 2, elements)
+        self.element_y = np.zeros(elements)
+        self.angle = np.deg2rad(angle_deg)
+        self.element_phase = np.zeros(elements)
+        self.element_phase = (
+            np.dot(np.array([self.element_x, self.element_y]).T, [np.sin(self.angle), np.cos(self.angle)]) * 2 * np.pi
+        )
+
+        x_max = np.max([x_range * 1.5, 10])
+        x_pixels = 200
+        self.x = np.linspace(-1, 1, x_pixels) * x_max
+        self.y = np.linspace(-0.2, 1, x_pixels) * x_max
+
+    def update(self, t: float):
+        self.ax.clear()
+
+        x = self.x
+        y = self.y
+
+        xx, yy = np.meshgrid(x, y)
+
+        cdata = []
+        for x, y, phase in zip(self.element_x, self.element_y, self.element_phase):
+            distance = np.sqrt((xx - x) ** 2 + (yy - y) ** 2) + float(np.finfo(float).tiny)
+            voltage = np.exp(1j * (phase + 2 * np.pi * distance - 2 * np.pi * t))  # * (1 / distance**2)
+            cdata.append(voltage)
+
+        cc = np.array(cdata).sum(axis=0)
+
+        self.ax.scatter(self.element_x, self.element_y, marker="x", color="k")
+        x_max = np.max(self.x)
+        self.ax.arrow(
+            x=np.sin(self.angle) * 0.05 * x_max,
+            y=np.cos(self.angle) * 0.05 * x_max,
+            dx=np.sin(self.angle) * 0.9 * x_max,
+            dy=np.cos(self.angle) * 0.9 * x_max,
+            head_width=0.05 * x_max,
+            head_length=0.05 * x_max,
+            width=0.01 * x_max,
+            fc="k",
+            ec="k",
+        )
+
+        return xx, yy, cc
+
+
+class PcolorPhasePlot(PcolorPlot):
+    def update(self, t: float):
+        xx, yy, cc = super().update(t)
+
+        self.ax.set_title(
+            "Nearfield Phase\n"
+            f"{len(self.element_x)} Element{'s' if len(self.element_x) > 1 else ''}, "
+            f"{self.spacing_lambda}$\\lambda$ Spacing, "
+            f"{np.rad2deg(self.angle):.0f}° Steer"
+        )
+        self.ax.pcolormesh(
+            xx,
+            yy,
+            np.angle(cc, deg=False),
+            clim=(-np.pi, np.pi),
+            cmap="twilight",
+            zorder=0,
+        )
+
+
+# TODO: don't animate this, it is constant
+class PcolorMagPlot(PcolorPlot):
+    def update(self, t: float):
+        xx, yy, cc = super().update(t)
+
+        self.ax.set_title(
+            # Note that this ignores 1/r**2 losses
+            "Nearfield Power Density\n"
+            f"{len(self.element_x)} Element{'s' if len(self.element_x) > 1 else ''}, "
+            f"{self.spacing_lambda}$\\lambda$ Spacing, "
+            f"{np.rad2deg(self.angle):.0f}° Steer"
+        )
+        self.ax.pcolormesh(
+            xx,
+            yy,
+            np.abs(cc) ** 2,
+            # 20 * np.log10(np.abs(cc / len(self.element_x))),
+            clim=(0, len(self.element_x) ** 2),
+            # clim=(-30, 0),
+            cmap="viridis",
+            zorder=0,
+        )
+
+
+# %%
+def generate():
+    for elements in [
+        1,
+        # 2,
+        4,
+        # 10,
+    ]:
+        # PcolorPhasePlot(elements=elements, angle_deg=45).save(dir_assets / "beamforming" / f"phase_xz_{elements}el.gif")
+        PcolorMagPlot(elements=elements, angle_deg=45).save(
+            dir_assets / "beamforming" / f"magnitude_xz_{elements}el.gif"
+        )
+
+
+# %%
+if __name__ == "__main__":
+    generate()

plots/utils.py

@@ -41,7 +41,7 @@ class AnimatedPlot(ABC):
 
     def __init__(self, frames: int = 100):
         self.frames = int(frames)
-        self.fig, self.ax = plt.subplots(1, 1)
+        self.fig, self.ax = plt.subplots(1, 1, tight_layout=True)
 
     def save(self, filename: Union[Path, str], framerate: int = 30):
         log.info(f"Generating animation: {self.__class__}...")
@@ -60,4 +60,10 @@ class AnimatedPlot(ABC):
 
     @abstractmethod
     def update(self, t: float) -> None:
+        """
+        Parameters
+        ----------
+        t
+            [0, 1]
+        """
         pass