Source code for mpu.math

#!/usr/bin/env python

"""
Mathematical functions which are not adequately covered by standard libraries.

Standard libraries are:

* `math <https://docs.python.org/3/library/math.html>`_
* `scipy <https://docs.scipy.org/doc/scipy/reference/>`_
* `sympy <http://docs.sympy.org/latest/index.html>`_

"""


# Core Library
import math as math_stl
import operator
from functools import reduce


[docs]def generate_primes(): """ Generate an infinite sequence of prime numbers. The algorithm was originally written by David Eppstein, UC Irvine. See: http://code.activestate.com/recipes/117119/ Examples -------- >>> g = generate_primes() >>> next(g) 2 >>> next(g) 3 >>> next(g) 5 """ divisors = {} # map number to at least one divisor candidate = 2 # next potential prime while True: if candidate in divisors: # candidate is composite. divisors[candidate] is the list of primes # that divide it. Since we've reached candidate, we no longer need # it in the map, but we'll mark the next multiples of its witnesses # to prepare for larger numbers for p in divisors[candidate]: divisors.setdefault(p + candidate, []).append(p) del divisors[candidate] else: # candidate is a new prime yield candidate # mark its first multiple that isn't # already marked in previous iterations divisors[candidate * candidate] = [candidate] candidate += 1
[docs]def factorize(number): """ Get the prime factors of an integer except for 1. Parameters ---------- number : int Returns ------- primes : iterable Examples -------- >>> factorize(-17) [-1, 17] >>> factorize(8) [2, 2, 2] >>> factorize(3**25) [3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3] >>> factorize(1) [1] """ if not isinstance(number, int): raise ValueError("integer expected, but type(number)={}".format(type(number))) if number < 0: return [-1] + factorize(number * (-1)) elif number == 0: raise ValueError("All primes are prime factors of 0.") else: for i in range(2, int(math_stl.ceil(number ** 0.5)) + 1): if number % i == 0: if i == number: return [i] else: return [i] + factorize(int(number / i)) return [number]
[docs]def is_prime(number): """ Check if a number is prime. Parameters ---------- number : int Returns ------- is_prime_number : bool Examples -------- >>> is_prime(-17) False >>> is_prime(17) True >>> is_prime(47055833459) True """ return len(factorize(number)) == 1
[docs]def product(iterable, start=1): """ Calculate the product of the iterables. Parameters ---------- iterable : iterable List, tuple or similar which contains numbers start : number, optional (default: 1) Returns ------- product : number Examples -------- >>> product([1, 2, 3, 4, 5]) 120 >>> product([]) 1 """ return reduce(operator.mul, iterable, start)
[docs]def argmax(iterable): """ Find the first index of the biggest value in the iterable. Parameters ---------- iterable : iterable Returns ------- argmax : int Examples -------- >>> argmax([0, 0, 0]) 0 >>> argmax([1, 0, 0]) 0 >>> argmax([0, 1, 0]) 1 >>> argmax([]) """ max_value = None max_index = None for index, value in enumerate(iterable): if (max_value is None) or max_value < value: max_value = value max_index = index return max_index
[docs]def round_up(x, decimal_places): """ Round a float up to decimal_places. Parameters ---------- x : float decimal_places : int Returns ------- rounded_float : float Examples -------- >>> round_up(1.2344, 3) 1.235 >>> round_up(1.234, 3) 1.234 >>> round_up(1.23456, 3) 1.235 >>> round_up(1.23456, 2) 1.24 """ return round(x + 5 * 10 ** (-1 * (decimal_places + 1)), decimal_places)
[docs]def round_down(x, decimal_places): """ Round a float down to decimal_places. Parameters ---------- x : float decimal_places : int Returns ------- rounded_float : float Examples -------- >>> round_down(1.23456, 3) 1.234 >>> round_down(1.23456, 2) 1.23 """ from math import floor d = int("1" + ("0" * decimal_places)) return floor(x * d) / d
[docs]def gcd(a: int, b: int) -> int: """ Calculate the greatest common divisor. Currently, this uses the Euclidean algorithm. Parameters ---------- a : int Non-zero b : int Returns ------- greatest_common_divisor : int Examples -------- >>> gcd(1, 7) 1 >>> gcd(-1, -1) 1 >>> gcd(1337, 42) 7 >>> gcd(-1337, -42) 7 >>> gcd(120, 364) 4 >>> gcd(273, 1870) 1 """ if a == 0 or b == 0: raise ValueError("gcd(a={a}, b={b}) is undefined") while b != 0: a, b = b, a % b return abs(a)