Source code for mpu.math

"""
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
from typing import Dict, Iterable, Iterator, List, Optional


[docs]def generate_primes() -> Iterator[int]: """ 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: Dict[int, List[int]] = {} # 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: int) -> List[int]: """ Get the prime factors of an integer except for 1. Parameters ---------- number : int Returns ------- primes : List[int] 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(f"integer expected, but type(number)={type(number)}") if number < 0: return [-1] + factorize(number * (-1)) elif number == 0: raise ValueError("All primes are prime factors of 0.") else: factors = [] factor = 2 while number % factor == 0: factors.append(factor) number = number // factor if number == 1: if len(factors) > 0: return factors else: return [1] for factor in range(3, int(math_stl.ceil(number ** 0.5)) + 1, 2): if number % factor == 0: return factors + [factor] + factorize(number // factor) return factors + [number]
[docs]def is_prime(number: int) -> bool: """ 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: Iterable, start: int = 1) -> int: """ 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: Iterable) -> Optional[int]: """ Find the first index of the biggest value in the iterable. Parameters ---------- iterable : Iterable Returns ------- argmax : Optional[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: float, decimal_places: int) -> float: """ 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: float, decimal_places: int) -> float: """ 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 """ d = int("1" + ("0" * decimal_places)) return math_stl.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 Non-zero 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(f"gcd(a={a}, b={b}) is undefined") while b != 0: a, b = b, a % b return abs(a)