A script to visualize the Collatz Conjecture

collatz.py

#! /usr/bin/env python3
# Fun with the Collatz Conjecture

from argparse import ArgumentParser
import csv
import logging
import matplotlib.pyplot as plt
import sys


logger = logging.getLogger(__name__)
ENDINGS = (1,)


def fc(x):
    if x % 2 == 0:
        return x // 2
    else:
        return 3 * x + 1


def is_done(x):
    return x in ENDINGS


def get_sequence(start, log_mod=10, verbose=False):
    x = int(start)

    i = 0
    while not is_done(x):
        x_1 = fc(x)
        i += 1
        yield x
        x = x_1
        if verbose and i % log_mod == 0:
            logger.warning(
                f'Continuing sequence for {start=}. Total Points: {i}...'
            )
    yield x

def get_args():
    parser = ArgumentParser(
        'collatz.py',
        'A fun way to work with the numbers of the collatz conjecture.'
    )
    parser.add_argument(
        '-n', '--number',
        type=int,
        default=None,
        help='Calculate the sequence for the given single number.'
    )
    parser.add_argument(
        '-s', '--start',
        type=int,
        default=0,
        help='Define the start of a range to calculate the values for.'
    )
    parser.add_argument(
        '-e', '--end',
        type=int,
        default=0,
        help='Define the end of a range to calculate the values for.'
    )
    parser.add_argument(
        '-v', '--verbose',
        action='store_true',
        default=False,
        help='Log verbose output'
    )
    parser.add_argument(
        '-c', '--columns',
        action='store_true',
        default=False,
        help=(
            'Write the results as one column of newline separated values. '
            'The default is a one-line row of comma separated values.'
        )
    )
    parser.add_argument(
        '-p', '--plot',
        action='store_true',
        default=False,
        help=(
            'Plot the given values.'
        )
    )
    return parser.parse_args()


def main():
    args = get_args()

    if args.plot:
        values = [args.number] if args.number else range(args.start, args.end)

        for i in values:
            y = [*get_sequence(i, verbose=args.verbose)]
            x = [*range(len(y))]
            plt.plot(x, y)

        plt.show()

    else:
        writer = csv.writer(sys.stdout)
        values = [args.number] if args.number else range(args.start, args.end)
        write = writer.writerows if args.columns else writer.writerow
        for i in values:
            write([
                [n] if args.columns else n
                for n in get_sequence(i, verbose=args.verbose)
            ])



if __name__ == '__main__':
    main()

requirements.txt

contourpy==1.0.7
cycler==0.11.0
fonttools==4.39.2
kiwisolver==1.4.4
matplotlib==3.7.1
numpy==1.24.2
packaging==23.0
Pillow==9.4.0
pyparsing==3.0.9
python-dateutil==2.8.2
six==1.16.0