reduce code duplication

This commit is contained in:
Brendan Haines 2025-03-22 21:59:22 -06:00
parent 7350978c9b
commit 8a9f8b8dde
2 changed files with 17 additions and 29 deletions

View File

@ -3,7 +3,7 @@ import logging
import numpy as np
from matplotlib import pyplot as plt
from utils import AnimatedPlot, dir_assets
from utils import AnimatedPlot, dir_assets, pi_ticks
# %% logging
log = logging.Logger(__name__)
@ -68,19 +68,6 @@ class SuperpositionPlot(AnimatedPlot):
self.ax.set_ylabel("V")
self.ax.set_title(f"$t={t:0.2f}$")
def pi_ticks(value, tick_number):
# find number of multiples of pi/2
N = int(np.round(2 * value / np.pi))
if N == 0:
return "0"
elif np.abs(N) == 2:
# +/- 1 * pi
return "$" + ("-" if N < 0 else "") + r"\pi$"
elif N % 2 == 0:
return "$" + f"{N // 2}" + r"\pi$"
else:
return "$" + ("-" if N < 0 else "") + r"\frac{" + f"{np.abs(N)}" + r"}{2}\pi$"
self.ax.xaxis.set_major_locator(plt.MultipleLocator(np.pi / 2))
self.ax.xaxis.set_major_formatter(plt.FuncFormatter(pi_ticks))
self.ax.grid(True)
@ -91,11 +78,11 @@ class SuperpositionPlot(AnimatedPlot):
class ReflectionPlot(AnimatedPlot):
x = np.linspace(-4 * np.pi, 0, 500)
def __init__(self, zl: complex, **kwargs):
def __init__(self, zl: complex, z0: complex = 1 + 0j, **kwargs):
super().__init__(**kwargs)
self.zl = zl
self.z0 = 1
self.z0 = z0
def update(self, t: float):
v = 2 * np.pi
@ -119,19 +106,6 @@ class ReflectionPlot(AnimatedPlot):
self.ax.set_ylabel("V")
self.ax.set_title(f"$Z_L/Z_0={self.zl}$, $t={t:0.2f}$")
def pi_ticks(value, tick_number):
# find number of multiples of pi/2
N = int(np.round(2 * value / np.pi))
if N == 0:
return "0"
elif np.abs(N) == 2:
# +/- 1 * pi
return "$" + ("-" if N < 0 else "") + r"\pi$"
elif N % 2 == 0:
return "$" + f"{N // 2}" + r"\pi$"
else:
return "$" + ("-" if N < 0 else "") + r"\frac{" + f"{np.abs(N)}" + r"}{2}\pi$"
self.ax.xaxis.set_major_locator(plt.MultipleLocator(np.pi / 2))
self.ax.xaxis.set_major_formatter(plt.FuncFormatter(pi_ticks))
self.ax.grid(True)

View File

@ -21,6 +21,20 @@ dir_assets = dir_root / "assets"
# %%
def pi_ticks(value, tick_number):
# find number of multiples of pi/2
N = int(np.round(2 * value / np.pi))
if N == 0:
return "0"
elif np.abs(N) == 2:
# +/- 1 * pi
return "$" + ("-" if N < 0 else "") + r"\pi$"
elif N % 2 == 0:
return "$" + f"{N // 2}" + r"\pi$"
else:
return "$" + ("-" if N < 0 else "") + r"\frac{" + f"{np.abs(N)}" + r"}{2}\pi$"
class AnimatedPlot(ABC):
fig: Figure
ax: Union[Axes, Tuple[Axes]]