95 lines
2.1 KiB
Python
95 lines
2.1 KiB
Python
# %% imports
|
|
import logging
|
|
from abc import ABC, abstractmethod
|
|
from pathlib import Path
|
|
from typing import Tuple, Union
|
|
|
|
import numpy as np
|
|
from matplotlib import pyplot as plt
|
|
from matplotlib.animation import FuncAnimation
|
|
from matplotlib.axes import Axes
|
|
from matplotlib.figure import Figure
|
|
|
|
# %% logging
|
|
log = logging.Logger(__name__)
|
|
log.setLevel(logging.INFO)
|
|
|
|
# %% paths
|
|
dir_ = Path(__file__).parent.resolve()
|
|
dir_root = dir_ / ".."
|
|
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]]
|
|
|
|
def __init__(self, frames: int = 100):
|
|
self.frames = int(frames)
|
|
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__}...")
|
|
an = FuncAnimation(
|
|
self.fig,
|
|
self.update,
|
|
frames=np.linspace(0, 1, self.frames, endpoint=False),
|
|
init_func=self.init,
|
|
blit=False,
|
|
)
|
|
an.save(str(filename), writer="pillow", fps=framerate)
|
|
log.info(f"Generating animation: {self.__class__}...Done")
|
|
|
|
def init(self) -> None:
|
|
pass
|
|
|
|
@abstractmethod
|
|
def update(self, t: float) -> None:
|
|
"""
|
|
Parameters
|
|
----------
|
|
t
|
|
[0, 1]
|
|
"""
|
|
pass
|
|
|
|
|
|
def db20(v):
|
|
return 20 * np.log10(np.abs(v))
|
|
|
|
|
|
def db10(v):
|
|
return 10 * np.log10(np.abs(v))
|
|
|
|
|
|
def db2v(db):
|
|
return 10 ** (db / 20)
|
|
|
|
|
|
def db2w(db):
|
|
return 10 ** (db / 10)
|
|
|
|
|
|
def wrap_phase(phase, deg: bool):
|
|
if deg:
|
|
maxval = 180
|
|
else:
|
|
maxval = np.pi
|
|
|
|
return (phase + maxval) % (2 * maxval) - maxval
|