photonforge 1.3.0__cp310-cp310-win_amd64.whl → 1.3.2__cp310-cp310-win_amd64.whl

photonforge/time_steppers/filter.py

@@ -0,0 +1,400 @@
+from collections.abc import Sequence
+from typing import Literal
+
+import numpy
+from scipy import signal
+
+from ..extension import (
+    Component,
+    Interpolator,
+    PoleResidueMatrix,
+    TimeDomainModel,
+    TimeStepper,
+    register_time_stepper_class,
+)
+from ..typing import (
+    Frequency,
+    Impedance,
+    Loss,
+    PositiveFloat,
+    PositiveInt,
+    TimeDelay,
+    annotate,
+)
+from .modulator import _get_impedance_runner
+
+
+class FilterTimeStepper(TimeStepper):
+    """Filter time stepper for electrical signals
+
+    Args:
+        family: Filter family.
+        shape: Filter shape.
+        f_cutoff: Cutoff frequency for low- and high-pass filter shapes. For
+          ``family=="tunable_lp_rc"``, this can be an :class:`Interpolator`
+          to provide the dependency with the input voltage. For band-pass
+          and band-stop shapes, this is a 2-value sequence with low and high
+          cutoff frequencies.
+        order: Filter order.
+        ripple: Maximum ripple for Chebyshev filters.
+        window: Window specification for rectangular filters. Please consult
+          the help for ``scipy.signal.firwin`` for valid options.
+        a: Recursive (denominator) coefficients for a digital IIR filter.
+        b: Direct (numerator) coefficients for a digital FIR or IIR filter.
+        taps: Length of rectangular or Gaussian filters.
+        insertion_loss: Insertion loss added to the filter response.
+        ports: Input and output port names. If not set, the *sorted* list of
+          port names from the component is used.
+        z0: Characteristic impedance of the electrical port used to convert
+          the input control field amplitude to voltage for the
+          ``"tunable_lp_rc"`` family. If ``None``, derived from port
+          impedance, calculated by mode-solving, or set to 50 Ω.
+        mesh_refinement: Minimal number of mesh elements per wavelength used
+          for mode solving.
+        verbose: Flag setting the verbosity of mode solver runs.
+
+    Note:
+        Filtering is applied to input signals from the first port (sorted
+        alphabetically) and output on the second port. For the
+        ``"tunable_lp_rc"``, the third port contains the control signal.
+    """
+
+    def __init__(
+        self,
+        *,
+        family: Literal[
+            "tunable_lp_rc",
+            "digital",
+            "rc",
+            "butterworth",
+            "bessel",
+            "cheby1",
+            "rectangular",
+            "gaussian",
+        ],
+        shape: Literal["lp", "hp", "bp", "bs"] = "lp",
+        f_cutoff: Frequency
+        | annotate(Sequence[Frequency], minItems=2, maxItems=2)
+        | Interpolator = 0,
+        order: PositiveInt = 1,
+        ripple: Loss = 0,
+        window: str
+        | tuple[str, float]
+        | tuple[str, float, float]
+        | tuple[str, Sequence[float]] = "hann",
+        a: Sequence[complex] = (1.0,),
+        b: Sequence[complex] = (),
+        taps: PositiveInt = 101,
+        insertion_loss: Loss = 0,
+        ports: annotate(Sequence[str], minItems=2, maxItems=2) | None = None,
+        z0: Impedance | Interpolator | None = None,
+        mesh_refinement: PositiveFloat | None = None,
+        verbose: bool = True,
+    ):
+        super().__init__(
+            family=family,
+            shape=shape,
+            f_cutoff=f_cutoff,
+            order=order,
+            ripple=ripple,
+            window=window,
+            a=a,
+            b=b,
+            taps=taps,
+            insertion_loss=insertion_loss,
+            ports=ports,
+            z0=z0,
+            mesh_refinement=mesh_refinement,
+            verbose=verbose,
+        )
+
+    def setup_state(
+        self,
+        *,
+        component: Component,
+        time_step: TimeDelay,
+        verbose: bool | None = None,
+        **kwargs,
+    ):
+        """Initialize internal state.
+
+        Args:
+            component: Component representing the filter.
+            time_step: The interval between time steps (in seconds).
+            kwargs: Unused.
+        """
+        p = self.parametric_kwargs
+        self._family = p["family"]
+
+        required_ports = 3 if self._family == "tunable_lp_rc" else 2
+
+        ports = p["ports"]
+        component_ports = sorted(component.select_ports("electrical"))
+        if ports is None:
+            ports = component_ports
+            if len(ports) != required_ports:
+                raise RuntimeError(
+                    f"FilterTimeStepper of family {self._family!r} can only be used in components "
+                    f"with {required_ports} electrical ports."
+                )
+        elif len(ports) != 2:
+            raise ValueError("'Argument 'ports' must be a sequence of 2 port names.")
+        elif not all(name in component_ports for name in ports):
+            raise RuntimeError(
+                f"Not all port names defined in FilterTimeStepper match the electrical port "
+                f"names in component '{component.name}'."
+            )
+        elif required_ports == 3:
+            component_ports.remove(ports[0])
+            component_ports.remove(ports[1])
+            ports = (*ports, component_ports[0])
+
+        self.keys = tuple(p + "@0" for p in ports)
+
+        self._gain = 10.0 ** (p["insertion_loss"] / -20.0)
+
+        f_nyquist = 0.5 / time_step
+
+        if self._family == "tunable_lp_rc":
+            self._time_step = time_step
+            self._f_cutoff = p["f_cutoff"]
+            if not isinstance(self._f_cutoff, Interpolator):
+                self._f_cutoff = Interpolator([0], [self._f_cutoff])
+
+            if _get_impedance_runner(self, component, ports[2], time_step, verbose) is not None:
+                return self
+
+        elif self._family == "digital":
+            a = numpy.asarray(p["a"], dtype=float).ravel()
+            b = numpy.asarray(p["b"], dtype=float).ravel()
+            if b.size < 1 or a.size < 1:
+                raise RuntimeError(
+                    "Digital filter requires 'b' and 'a' coefficients of size greater than 0."
+                )
+            self._b = b / a[0]
+            self._a = a[1:] / a[0]
+            self._x = numpy.empty(b.size - 1, dtype=float)
+            self._y = numpy.empty(self._a.size, dtype=float)
+
+        elif self._family in ("rectangular", "gaussian"):
+            if self._family == "rectangular":
+                shape = p["shape"]
+                if shape == "lp":
+                    b = signal.firwin(
+                        p["taps"], p["f_cutoff"] / f_nyquist, pass_zero=True, window=p["window"]
+                    )
+                elif shape == "hp":
+                    b = signal.firwin(
+                        p["taps"], p["f_cutoff"] / f_nyquist, pass_zero=False, window=p["window"]
+                    )
+                elif shape == "bp":
+                    f_low, f_high = p["f_cutoff"]
+                    b = signal.firwin(
+                        p["taps"],
+                        [f_low / f_nyquist, f_high / f_nyquist],
+                        pass_zero=False,
+                        window=p["window"],
+                    )
+                elif shape == "bs":
+                    f_low, f_high = p["f_cutoff"]
+                    b = signal.firwin(
+                        p["taps"],
+                        [f_low / f_nyquist, f_high / f_nyquist],
+                        pass_zero=True,
+                        window=p["window"],
+                    )
+                else:
+                    raise ValueError(f"Unrecognized filter shape {shape!r}.")
+
+            else:
+                shape = p["shape"]
+                taps = p["taps"]
+                k = numpy.arange(taps)
+                fs = 1.0 / time_step
+                fpos = k / taps * fs
+                f_signed = numpy.where(k <= taps // 2, fpos, fpos - fs)
+
+                def lp(freqs, f_cutoff, m):
+                    return numpy.exp(-0.5 * numpy.log(2.0) * numpy.abs(freqs / f_cutoff) ** m)
+
+                def hp(freqs, f_cutoff, m):
+                    non_zero = freqs != 0
+                    transformed = numpy.zeros(freqs.size)
+                    transformed[non_zero] = f_cutoff**2 / numpy.abs(freqs[non_zero])
+                    h = lp(transformed, f_cutoff, m)
+                    h[numpy.logical_not(non_zero)] = 0.0
+                    return h
+
+                m = 2 * p["order"]
+                if shape == "lp":
+                    h = lp(f_signed, p["f_cutoff"], m)
+                elif shape == "hp":
+                    h = hp(f_signed, p["f_cutoff"], m)
+                else:
+                    f_low, f_high = p["f_cutoff"]
+                    f0 = 0.5 * (f_high + f_low)
+                    df = 0.5 * (f_high - f_low)
+                    if shape == "bp":
+                        h = lp(f_signed - f0, df, m) + lp(f_signed + f0, df, m)
+                    elif shape == "bs":
+                        h = hp(f_signed - f0, df, m) * hp(f_signed + f0, df, m)
+                    else:
+                        raise ValueError(f"Unrecognized filter shape {shape!r}.")
+
+                b = numpy.roll(numpy.fft.ifft(h).real, (taps - 1) // 2)
+
+            if b.size < 1:
+                raise RuntimeError("FIR kernel design failed.")
+
+            self._b = b.real
+            self._a = numpy.empty(0)
+            self._x = numpy.empty(b.size - 1, dtype=float)
+            self._y = numpy.empty(0, dtype=float)
+
+            self._family = "digital"
+
+        elif self._family in ("rc", "butterworth", "bessel", "cheby1"):
+            shape = p["shape"]
+            if shape == "lp":
+                f_cutoff = p["f_cutoff"]
+                if not (0 < f_cutoff < f_nyquist):
+                    raise ValueError(
+                        f"'f_cutoff' ({f_cutoff}) must be positive and less than "
+                        f"'1 / (2 * time_step)' ({f_nyquist})"
+                    )
+                w0 = 0.0
+                wc = 2.0 * numpy.pi * f_cutoff
+            elif shape == "hp":
+                f_cutoff = p["f_cutoff"]
+                if not (0 < f_cutoff < f_nyquist):
+                    raise ValueError(
+                        f"'f_cutoff' ({f_cutoff}) must be positive and less than "
+                        f"'1 / (2 * time_step)' ({f_nyquist})"
+                    )
+                w0 = 0.0
+                wc = 2.0 * numpy.pi * f_cutoff
+            elif shape == "bp":
+                f_low, f_high = p["f_cutoff"]
+                if not (0 < f_low < f_high < f_nyquist):
+                    raise ValueError(
+                        f"'f_low' ({f_low}) and 'f_high' ({f_high}) must be positive and "
+                        f"satisfy 'f_low < f_high < 1 / (2 * time_step)' ({f_nyquist})"
+                    )
+                w0 = numpy.pi * (f_low + f_high)
+                wc = numpy.pi * (f_high - f_low)
+            elif shape == "bs":
+                f_low, f_high = p["f_cutoff"]
+                if not (0 < f_low < f_high < f_nyquist):
+                    raise ValueError(
+                        f"'f_low' ({f_low}) and 'f_high' ({f_high}) must be positive and "
+                        f"satisfy 'f_low < f_high < 1 / (2 * time_step)' ({f_nyquist})"
+                    )
+                w0 = numpy.pi * (f_low + f_high)
+                wc = numpy.pi * (f_high - f_low)
+            else:
+                raise ValueError(f"Unrecognized filter shape {shape!r}.")
+
+            if self._family == "rc":
+                if shape == "lp":
+                    poles = [-wc]
+                    residues = [wc, 0.0]
+                elif shape == "hp":
+                    poles = [-wc]
+                    residues = [-wc, 1.0]
+                elif shape == "bp":
+                    poles = [-wc - 1j * w0]
+                    residues = [wc, 0.0]
+                elif shape == "bs":
+                    poles = [-wc - 1j * w0]
+                    residues = [-wc, 1.0]
+
+            else:
+                btype = "lowpass" if shape in ("lp", "bp") else "highpass"
+                if self._family == "butterworth":
+                    b, a = signal.butter(p["order"], wc, btype=btype, analog=True)
+                elif self._family == "bessel":
+                    b, a = signal.bessel(p["order"], wc, btype=btype, analog=True, norm="mag")
+                else:
+                    b, a = signal.cheby1(p["order"], p["ripple"], wc, btype=btype, analog=True)
+                residues, poles, k = signal.residue(b, a)
+                if btype == "lowpass":
+                    k = numpy.zeros(1)
+                residues = numpy.concatenate((residues, k))
+                poles = poles - 1j * w0
+
+            self._time_domain_model = TimeDomainModel(
+                PoleResidueMatrix(poles, {tuple(self.keys[:2]): residues}),
+                time_step,
+            )
+
+            self._family = "time_domain_model"
+
+        else:
+            raise ValueError(f"Unrecognized filter family {self._family!r}.")
+
+        self.reset()
+
+    @property
+    def status(self):
+        if not self._setup_thread.is_alive() and self._status["message"] == "running":
+            self._status["message"] = "error"
+        with self._lock:
+            return self._status
+
+    def reset(self):
+        """Reset internal state."""
+        if self._family == "tunable_lp_rc":
+            self._sqrt_r_in = numpy.real(self._z0) ** 0.5
+            self._y = 0.0
+        elif self._family == "digital":
+            self._x.fill(0.0)
+            self._y.fill(0.0)
+        elif self._family == "time_domain_model":
+            self._time_domain_model.reset()
+
+    def step_single(
+        self,
+        inputs: numpy.ndarray,
+        outputs: numpy.ndarray,
+        time_index: int,
+        update_state: bool,
+        shutdown: bool,
+    ) -> None:
+        """Take a single time step on the given inputs.
+
+        Args:
+            inputs: Input values at the current time step. Must be a 1D array of
+              complex values ordered
+              according to :attr:`keys`.
+            outputs: Pre-allocated output array where results will be stored.
+              Same size and type as ``inputs``.
+            time_index: Time series index for the current input.
+            update_state: Whether to update the internal stepper state.
+            shutdown: Whether this is the last call to the single stepping
+              function for the provided :class:`TimeSeries`.
+        """
+        if self._family == "tunable_lp_rc":
+            v_ctrl = self._sqrt_r_in * inputs[2].real
+            f_cutoff = self._f_cutoff(v_ctrl)
+            alpha = min(1.0, max(0.0, numpy.exp(-2.0 * numpy.pi * f_cutoff * self._time_step)))
+            y = self._y * alpha + inputs[0].real * (1.0 - alpha)
+            outputs[1] = self._gain * y
+            if update_state:
+                self._y = y
+        elif self._family == "digital":
+            x = inputs[0].real
+            y = self._b[0] * x + numpy.dot(self._b[1:], self._x) - numpy.dot(self._a, self._y)
+            outputs[1] = self._gain * y
+            if update_state:
+                if self._x.size > 0:
+                    self._x[1:] = self._x[:-1]
+                    self._x[0] = x
+                if self._y.size > 0:
+                    self._y[1:] = self._y[:-1]
+                    self._y[0] = y
+        elif self._family == "time_domain_model":
+            self._time_domain_model.step_single(inputs, outputs, time_index, update_state, shutdown)
+            outputs[1] *= self._gain
+
+
+register_time_stepper_class(FilterTimeStepper)