Plot the transfer function for simple RL and RC circuits

low_high_pass_filters.py

#! /usr/bin/env python3

import argparse
import math
from sys import argv

import matplotlib.pyplot as plt
import numpy as np

MIN_FREQ = 1
MAX_FREQ = 1e6

ARGS_RESISTOR = ('-R', '--res')
ARGS_CAPACITOR = ('-C', '--cap')
ARGS_INDUCTOR = ('-L', '--ind')

DESCRIPTION = '''
Plot the transfer function for simple RL and RC circuits.

Ordering of the elements (R, C, L) matter. For example, passing -R then -C
values will create a low pass circuit, while -C then -R values will create
a high pass circuit.

Example usage:
    ./low_high_pass_filters.py -R 1000 -C 2.2e-6
    ./low_high_pass_filters.py -L 10e-2 -R 500
'''


PREFIXES = {
    'T': 1e12,
    'G': 1e9,
    'M': 1e6,
    'k': 1e3,
    'm': 1e-3,
    'u': 1e-6,
    'n': 1e-9,
    'p': 1e-12,
}


def low_pass_transfer_function(freq, tau):
    omega_tau = 2 * math.pi * freq * tau
    return 1 / math.sqrt(1 + math.pow(omega_tau, 2))


def high_pass_transfer_function(freq, tau):
    omega_tau = 2 * math.pi * freq * tau
    return omega_tau / math.sqrt(1 + math.pow(omega_tau, 2))


def tau_inductor(resistor, inductor):
    return inductor / resistor


def tau_capacitor(resistor, capacitor):
    return capacitor * resistor


def plot(_type, transfer_function, max_freq, tau):
    frequencies = np.logspace(np.log10(MIN_FREQ), np.log10(max_freq), 50)

    values = [transfer_function(f, tau) for f in frequencies]

    # plot x values as log scale
    plt.semilogx(frequencies, values)

    plt.title('%s-Pass Filter' % _type)
    plt.ylabel(r'$\frac{V_{out}}{V_{in}}$', rotation=0, fontsize=15, labelpad=15)
    plt.grid(axis='x', which='both', linestyle='--')
    plt.xlim(frequencies[0], frequencies[-1])

    plt.hlines(0.707, MIN_FREQ, max_freq)
    plt.text(MIN_FREQ, 0.707, '0.707 ', verticalalignment='center',
             horizontalalignment='right')

    plt.show()


def _determine_filter_type(args):
    def is_capacitor(arg):
         return arg in ARGS_CAPACITOR

    # RC or RL circuit
    if args[1] in ARGS_RESISTOR:
        return 'low' if is_capacitor(args[3]) else 'high'
    # CR or LR circut
    else:
        return 'high' if is_capacitor(args[1]) else 'low'


class convert_to_number(argparse.Action):
    def __call__(self, parser, namespace, number_str, option_string=None):
        last_digit = number_str[-1]
        if last_digit in PREFIXES:
            multiplier = PREFIXES[last_digit]
            final = multiplier * float(number_str[:-1])
        else:
            final = float(number_str)

        setattr(namespace, self.dest, final)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description=DESCRIPTION,
        formatter_class=argparse.RawTextHelpFormatter,
    )

    parser.add_argument(
        '-m',
        '--max',
        help='Max frequency',
        default=MAX_FREQ,
        type=float,
        metavar='',
        dest='max_freq',
    )

    parser.add_argument(
        *ARGS_RESISTOR,
        help='Resistor value (in Ohms)',
        action=convert_to_number,
        metavar='',
        required=True,
        dest='resistor',
    )

    cap_or_ind = parser.add_mutually_exclusive_group(required=True)
    cap_or_ind.add_argument(
        *ARGS_CAPACITOR,
        help='Capacitor value (in Farads)',
        action=convert_to_number,
        metavar='',
        dest='capacitor',
    )

    cap_or_ind.add_argument(
        *ARGS_INDUCTOR,
        help='Inductor value (in Henries)',
        action=convert_to_number,
        metavar='',
        dest='inductor',
    )

    args = parser.parse_args()

    if args.capacitor:
        tau = tau_capacitor(args.resistor, args.capacitor)
    else:
        tau = tau_inductor(args.resistor, args.inductor)

    filter_type = _determine_filter_type(argv)

    if filter_type == 'low':
        plot('Low', low_pass_transfer_function, args.max_freq, tau)
    else:
        plot('High', high_pass_transfer_function, args.max_freq, tau)