nanovna-saver/src/NanoVNASaver/Calibration.py

#  NanoVNASaver
#
#  A python program to view and export Touchstone data from a NanoVNA
#  Copyright (C) 2019, 2020  Rune B. Broberg
#  Copyright (C) 2020,2021 NanoVNA-Saver Authors
#
#  This program is free software: you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program.  If not, see <https://www.gnu.org/licenses/>.
import logging
import cmath
import math
import os
import re
from collections import defaultdict, UserDict
from dataclasses import dataclass
from typing import List

from scipy.interpolate import interp1d

from NanoVNASaver.RFTools import Datapoint


IDEAL_SHORT = complex(-1, 0)
IDEAL_OPEN = complex(1, 0)
IDEAL_LOAD = complex(0, 0)
IDEAL_THROUGH = complex(1, 0)

RXP_CAL_HEADER = re.compile(r"""
    ^ \# \s+ Hz \s+
    ShortR \s+ ShortI \s+ OpenR \s+ OpenI \s+
    LoadR \s+ LoadI
    (?P<through> \s+ ThroughR \s+ ThroughI)?
    (?P<thrurefl> \s+ ThrureflR \s+ ThrureflI)?
    (?P<isolation> \s+ IsolationR \s+ IsolationI)?
    \s* $
""", re.VERBOSE | re.IGNORECASE)

RXP_CAL_LINE = re.compile(r"""
    ^ \s*
    (?P<freq>\d+) \s+
    (?P<shortr>[-0-9Ee.]+) \s+ (?P<shorti>[-0-9Ee.]+) \s+
    (?P<openr>[-0-9Ee.]+) \s+ (?P<openi>[-0-9Ee.]+) \s+
    (?P<loadr>[-0-9Ee.]+) \s+ (?P<loadi>[-0-9Ee.]+)
    ( \s+ (?P<throughr>[-0-9Ee.]+) \s+ (?P<throughi>[-0-9Ee.]+))?
    ( \s+ (?P<thrureflr>[-0-9Ee.]+) \s+ (?P<thrurefli>[-0-9Ee.]+))?
    ( \s+ (?P<isolationr>[-0-9Ee.]+) \s+ (?P<isolationi>[-0-9Ee.]+))?
    \s* $
""", re.VERBOSE)

logger = logging.getLogger(__name__)


def correct_delay(d: Datapoint, delay: float, reflect: bool = False):
    mult = 2 if reflect else 1
    corr_data = d.z * cmath.exp(
        complex(0, 1) * 2 * math.pi * d.freq * delay * -1 * mult)
    return Datapoint(d.freq, corr_data.real, corr_data.imag)


@dataclass
class CalData:
    # pylint: disable=too-many-instance-attributes
    short: complex = complex(0.0, 0.0)
    open: complex = complex(0.0, 0.0)
    load: complex = complex(0.0, 0.0)
    through: complex = complex(0.0, 0.0)
    thrurefl: complex = complex(0.0, 0.0)
    isolation: complex = complex(0.0, 0.0)
    freq: int = 0
    e00: float = 0.0  # Directivity
    e11: float = 0.0  # Port1 match
    delta_e: float = 0.0  # Tracking
    e10e01: float = 0.0  # Forward Reflection Tracking
    # 2 port
    e30: float = 0.0  # Forward isolation
    e22: float = 0.0  # Port2 match
    e10e32: float = 0.0  # Forward transmission

    def __str__(self):
        return (
            f'{self.freq}'
            f' {self.short.real} {self.short.imag}'
            f' {self.open.real} {self.open.imag}'
            f' {self.load.real} {self.load.imag}' + (
                f' {self.through.real} {self.through.imag}'
                f' {self.thrurefl.real} {self.thrurefl.imag}'
                f' {self.isolation.real} {self.isolation.imag}'
                if self.through else ''
            )
        )


@dataclass
class CalElement:
    # pylint: disable=too-many-instance-attributes
    short_is_ideal: bool = True
    short_l0: float = 5.7e-12
    short_l1: float = -8.96e-20
    short_l2: float = -1.1e-29
    short_l3: float = -4.12e-37
    short_length: float = -34.2  # ps

    open_is_ideal: bool = True
    open_c0: float = 2.1e-14
    open_c1: float = 5.67e-23
    open_c2: float = -2.39e-31
    open_c3: float = 2.0e-40
    open_length: float = 0.0

    load_is_ideal: bool = True
    load_r: float = 50.0
    load_l: float = 0.0
    load_c: float = 0.0
    load_length: float = 0.0

    through_is_ideal: bool = True
    through_length: float = 0.0


class CalDataSet(UserDict):
    def __init__(self):
        super().__init__()
        self.notes = ""
        self.data: defaultdict[int, CalData] = defaultdict(CalData)

    def __str__(self):
        return (
            (
                "# Calibration data for NanoVNA-Saver\n"
                + "\n".join([f"! {note}" for note in self.notes.splitlines()])
                + "\n" + "# Hz ShortR ShortI OpenR OpenI LoadR LoadI"
                + (" ThroughR ThroughI ThrureflR"
                   " ThrureflI IsolationR IsolationI\n"
                   if self.complete2port() else "\n")
                + "\n".join([
                    f"{self.data.get(freq)}" for freq in self.frequencies()
                ]) + "\n"
            )
            if self.complete1port() else ""
        )

    def _append_match(self, m: re.Match, header: str,
                      line_nr: int, line: str) -> None:
        cal = m.groupdict()
        columns = {
            col[:-1] for col in cal.keys() if cal[col] and col != "freq"
        }
        if "through" in columns and header == "sol":
            logger.warning("Through data with sol header. %i: %s",
                           line_nr, line)
        # fix short data (without thrurefl)
        if "thrurefl" in columns and "isolation" not in columns:
            cal["isolationr"] = cal["thrureflr"]
            cal["isolationi"] = cal["thrurefli"]
            cal["thrureflr"], cal["thrurefli"] = None, None
        for name in columns:
            self.insert(
                name,
                Datapoint(int(cal["freq"]),
                          float(cal[f"{name}r"]),
                          float(cal[f"{name}i"])))

    def from_str(self, text: str) -> 'CalDataSet':
        # reset data
        self.notes = ""
        self.data = defaultdict(CalData)
        header = ""
        # parse text
        for i, line in enumerate(text.splitlines(), 1):
            line = line.strip()

            if line.startswith("!"):
                self.notes += f"{line[2:]}\n"
                continue
            if m := RXP_CAL_HEADER.search(line):
                if header:
                    logger.warning(
                        "Duplicate header in cal data. %i: %s", i, line)
                header = "through" if m.group("through") else "sol"
                continue
            if not line or line.startswith("#"):
                continue

            m = RXP_CAL_LINE.search(line)
            if not m:
                logger.warning("Illegal caldata. Line %i: %s", i, line)
                continue
            if not header:
                logger.warning(
                    "Caldata without having read header: %i: %s", i, line)
            self._append_match(m, header, line, i)
        return self

    def insert(self, name: str, dp: Datapoint):
        if name not in {'short', 'open', 'load',
                        'through', 'thrurefl', 'isolation'}:
            raise KeyError(name)
        freq = dp.freq
        setattr(self.data[freq], name, (dp.z))
        self.data[freq].freq = freq

    def frequencies(self) -> List[int]:
        return sorted(self.data.keys())

    def get(self, key: int, default: CalData = None) -> CalData:
        return self.data.get(key, default)

    def items(self):
        yield from self.data.items()

    def values(self):
        for freq in self.frequencies():
            yield self.get(freq)

    def size_of(self, name: str) -> int:
        return len(
            [True for val in self.data.values() if getattr(val, name)]
        )

    def complete1port(self) -> bool:
        for val in self.data.values():
            if not all((val.short, val.open, val.load)):
                return False
        return any(self.data)

    def complete2port(self) -> bool:
        if not self.complete1port():
            return False
        for val in self.data.values():
            if not all((val.through, val.thrurefl, val.isolation)):
                return False
        return any(self.data)


class Calibration:
    def __init__(self):

        self.notes = []
        self.dataset = CalDataSet()
        self.cal_element = CalElement()
        self.interp = {}
        self.isCalculated = False

        self.source = "Manual"

    def insert(self, name: str, data: List[Datapoint]):
        for dp in data:
            self.dataset.insert(name, dp)

    def size(self) -> int:
        return len(self.dataset.frequencies())

    def data_size(self, name) -> int:
        return self.dataset.size_of(name)

    def isValid1Port(self) -> bool:
        return self.dataset.complete1port()

    def isValid2Port(self) -> bool:
        return self.dataset.complete2port()

    def _calc_port_1(self, freq: int, cal: CalData):
        g1 = self.gamma_short(freq)
        g2 = self.gamma_open(freq)
        g3 = self.gamma_load(freq)

        gm1 = cal.short
        gm2 = cal.open
        gm3 = cal.load

        denominator = (g1 * (g2 - g3) * gm1 +
                       g2 * g3 * gm2 - g2 * g3 * gm3 -
                       (g2 * gm2 - g3 * gm3) * g1)
        cal.e00 = - ((g2 * gm3 - g3 * gm3) * g1 * gm2 -
                     (g2 * g3 * gm2 - g2 * g3 * gm3 -
                     (g3 * gm2 - g2 * gm3) * g1) * gm1
                     ) / denominator
        cal.e11 = ((g2 - g3) * gm1 - g1 * (gm2 - gm3) +
                   g3 * gm2 - g2 * gm3) / denominator
        cal.delta_e = - ((g1 * (gm2 - gm3) - g2 * gm2 + g3 *
                          gm3) * gm1 + (g2 * gm3 - g3 * gm3) *
                         gm2) / denominator

    def _calc_port_2(self, freq: int, cal: CalData):
        gt = self.gamma_through(freq)

        gm4 = cal.through
        gm5 = cal.thrurefl
        gm6 = cal.isolation
        gm7 = gm5 - cal.e00

        cal.e30 = cal.isolation
        cal.e10e01 = cal.e00 * cal.e11 - cal.delta_e
        cal.e22 = gm7 / (
            gm7 * cal.e11 * gt ** 2 + cal.e10e01 * gt ** 2)
        cal.e10e32 = (gm4 - gm6) * (
            1 - cal.e11 * cal.e22 * gt ** 2) / gt

    def calc_corrections(self):
        if not self.isValid1Port():
            logger.warning(
                "Tried to calibrate from insufficient data.")
            raise ValueError(
                "All of short, open and load calibration steps"
                "must be completed for calibration to be applied.")
        logger.debug("Calculating calibration for %d points.", self.size())

        for freq, caldata in self.dataset.items():
            try:
                self._calc_port_1(freq, caldata)
                if self.isValid2Port():
                    self._calc_port_2(freq, caldata)
            except ZeroDivisionError as exc:
                self.isCalculated = False
                logger.error(
                    "Division error - did you use the same measurement"
                    " for two of short, open and load?")
                raise ValueError(
                    f"Two of short, open and load returned the same"
                    f" values at frequency {freq}Hz.") from exc

        self.gen_interpolation()
        self.isCalculated = True
        logger.debug("Calibration correctly calculated.")

    def gamma_short(self, freq: int) -> complex:
        if self.cal_element.short_is_ideal:
            return IDEAL_SHORT
        logger.debug("Using short calibration set values.")
        cal_element = self.cal_element
        Zsp = complex(0.0, 2.0 * math.pi * freq * (
            cal_element.short_l0 + cal_element.short_l1 * freq +
            cal_element.short_l2 * freq**2 + cal_element.short_l3 * freq**3))
        # Referencing https://arxiv.org/pdf/1606.02446.pdf (18) - (21)
        return (Zsp / 50.0 - 1.0) / (Zsp / 50.0 + 1.0) * cmath.exp(
            complex(0.0,
                    -4.0 * math.pi * freq * cal_element.short_length))

    def gamma_open(self, freq: int) -> complex:
        if self.cal_element.open_is_ideal:
            return IDEAL_OPEN
        logger.debug("Using open calibration set values.")
        cal_element = self.cal_element
        Zop = complex(0.0, 2.0 * math.pi * freq * (
            cal_element.open_c0 + cal_element.open_c1 * freq +
            cal_element.open_c2 * freq**2 + cal_element.open_c3 * freq**3))
        return ((1.0 - 50.0 * Zop) / (1.0 + 50.0 * Zop)) * cmath.exp(
            complex(0.0,
                    -4.0 * math.pi * freq * cal_element.open_length))

    def gamma_load(self, freq: int) -> complex:
        if self.cal_element.load_is_ideal:
            return IDEAL_LOAD
        logger.debug("Using load calibration set values.")
        cal_element = self.cal_element
        Zl = complex(cal_element.load_r, 0.0)
        if cal_element.load_c > 0.0:
            Zl = cal_element.load_r / complex(
                1.0,
                2.0 * cal_element.load_r * math.pi * freq * cal_element.load_c)
        if cal_element.load_l > 0.0:
            Zl = Zl + complex(0.0, 2 * math.pi * freq * cal_element.load_l)
        return (Zl / 50.0 - 1.0) / (Zl / 50.0 + 1.0) * cmath.exp(
            complex(0.0, -4 * math.pi * freq * cal_element.load_length))

    def gamma_through(self, freq: int) -> complex:
        if self.cal_element.through_is_ideal:
            return IDEAL_THROUGH
        logger.debug("Using through calibration set values.")
        cal_element = self.cal_element
        return cmath.exp(
            complex(0.0, -2.0 * math.pi * cal_element.through_length * freq))

    def gen_interpolation(self):
        (freq, e00, e11, delta_e, e10e01, e30, e22, e10e32) = zip(*[
            (c.freq, c.e00, c.e11, c.delta_e, c.e10e01, c.e30, c.e22, c.e10e32)
            for c in self.dataset.values()])

        self.interp = {
            "e00": interp1d(freq, e00,
                            kind="slinear", bounds_error=False,
                            fill_value=(e00[0], e00[-1])),
            "e11": interp1d(freq, e11,
                            kind="slinear", bounds_error=False,
                            fill_value=(e11[0], e11[-1])),
            "delta_e": interp1d(freq, delta_e,
                                kind="slinear", bounds_error=False,
                                fill_value=(delta_e[0], delta_e[-1])),
            "e10e01": interp1d(freq, e10e01,
                               kind="slinear", bounds_error=False,
                               fill_value=(e10e01[0], e10e01[-1])),
            "e30": interp1d(freq, e30,
                            kind="slinear", bounds_error=False,
                            fill_value=(e30[0], e30[-1])),
            "e22": interp1d(freq, e22,
                            kind="slinear", bounds_error=False,
                            fill_value=(e22[0], e22[-1])),
            "e10e32": interp1d(freq, e10e32,
                               kind="slinear", bounds_error=False,
                               fill_value=(e10e32[0], e10e32[-1])),
        }

    def correct11(self, dp: Datapoint):
        i = self.interp
        s11 = (dp.z - i["e00"](dp.freq)) / (
            (dp.z * i["e11"](dp.freq)) - i["delta_e"](dp.freq))
        return Datapoint(dp.freq, s11.real, s11.imag)

    def correct21(self, dp: Datapoint, dp11: Datapoint):
        i = self.interp
        s21 = (dp.z - i["e30"](dp.freq)) / i["e10e32"](dp.freq)
        s21 = s21 * (i["e10e01"](dp.freq) / (i["e11"](dp.freq)
                                             * dp11.z - i["delta_e"](dp.freq)))
        return Datapoint(dp.freq, s21.real, s21.imag)

    def save(self, filename: str):
        self.dataset.notes = "\n".join(self.notes)
        if not self.isValid1Port():
            raise ValueError("Not a valid calibration")
        with open(filename, mode="w", encoding='utf-8') as calfile:
            calfile.write(str(self.dataset))

    def load(self, filename):
        self.source = os.path.basename(filename)
        with open(filename, encoding='utf-8') as calfile:
            self.dataset = CalDataSet().from_str(calfile.read())
            self.notes = self.dataset.notes.splitlines()