My Python math utilities

domath.py

# To use this, run `python3 -i domath.py`.

import math
from math import *
from matplotlib import pyplot
from numpy import inf
import numpy as np
from numpy.linalg import det, inv, matrix_power, eigvals, eig, solve, matrix_rank
from sympy import Matrix, flatten

import os

def sci(num):
    return "{:e}".format(num)

def plot_it(x_begin, x_end, *funcs, step=0.0001, upper_bound=inf, lower_bound=-inf, equalize_axes=False):
    for f in funcs:
        x_values = []
        y_values = []
        i = x_begin
        while i <= x_end:
            i += step
            x_values.append(i)
            if f(i) > upper_bound:
                y_values.append(upper_bound)
                continue
            if f(i) < lower_bound:
                y_values.append(lower_bound)
                continue
            y_values.append(f(i))
        pyplot.plot(x_values, y_values)
    pyplot.grid(True)
    if equalize_axes: pyplot.gca().set_aspect('equal', adjustable='box')
    pyplot.show()

# Helper methods for using degrees instead of radians
mode = os.getenv('TRIGMODE')
if mode and mode.startswith("deg"):
    def sin(x):
        return math.sin(radians(x))
    def cos(x):
        return math.cos(radians(x))
    def tan(x):
        return math.tan(radians(x))

    def asin(x):
        return degrees(math.asin(x))
    def acos(x):
        return degrees(math.acos(x))
    def atan(x):
        return degrees(math.atan(x))



def sec(x):
    return 1.0 / cos(x)
def csc(x):
    return 1.0 / sin(x)
def cot(x):
    return 1.0 / tan(x)


# TODO: optimize i guess
def factorize(x):
    if x != floor(x): raise RuntimeError("factorize accepts only whole numbers")

    if x == 0: return [0]
    if x == 1: return [1]

    factors = []
    i = 1
    orig_x = x

    if x < 0:
        factors.append(-1)
        x *= -1

    while prod(factors) != orig_x:
        # if len(factors) > 0 and i > orig_x / factors[len(factors)-1]:
        #     break
        while x % i == 0 and i != 1:
            factors.append(i)
            x /= i
        i += 1

    return factors


from time import time_ns
def factorize_benchmark():
    times = []
    for i in range(10000):
        start_time = time_ns()
        for i in range(10): factorize(i)
        end_time = time_ns()
        times.append(end_time - start_time)
    def f(x):
        time_s = times[x]/1e9
        # if x > 0 and times[x-1]/times[x] > 3:
        #     return times[x-1]/1e9
        return time_s
    plot_it(0, 998, f, step=1)



# Chemistry

# Avogadro's number (1 mole)
avogadro = 6.02214076e23

# specific heat capacity of water (J/g/K)
c_water = 4.184

# atomic masses
H = 1.00794
Li = 6.941
Be = 9.0122
B = 10.811
C = 12.011
N = 14.0067
O = 15.9994
F = 18.9984
Na = 22.9898
Mg = 24.305
Al = 26.9815
P = 30.9738
S = 32.066
Cl = 35.453
K = 39.0983
Ca = 40.078
V = 50.9415
Cr = 51.9961
Fe = 55.845
Ni = 58.6934
Cu = 63.546
Zn = 65.38
Br = 79.904
Ag = 107.8682
I = 126.90447
Ba = 137.327
W = 183.84
Pb = 207.2

# molecular masses
H2 = H*2
N2 = N*2
O2 = O*2
H2O = H*2 + O
OH = O + H
NaOH = Na + O + H
H2O2 = H*2 + O*2
NH3 = N + H*3
NH2 = N + H*2
CO2 = C + O*2
CO = C + O
H2SO4 = H*2 + S + O*4
H3PO4 = H*3 + P + O*4
KOH = K + O + H
C2H4O2 = C*2 + H*4 + O*2
C11H22O2 = C*11 + H*22 + O*2
C6H12O6 = C*6 + H*12 + O*6
C3H8O3 = C*3 + H*8 + O*3
NaHCO3 = Na + H + C + O*3
NiCl2 = Ni + Cl*2
AgCl = Ag + Cl
NaCl = Na + Cl
KCl = K + Cl
HBr = H + Br


# Physics

# electrostatics

# Coulomb's constant
k = 8.9875517923e9

# elementary charge
q_0 = 1.602176634e-19

# permittivity of free space
epsilon_0 = 8.8541878128e-12

# permeability of free space
mu_0 = 1.25663706212e-6

def coulomb_force(q1, q2, r):
    return k * q1 * q2 / r**2


# Linear Algebra

def _3x3(*nums):
    return np.array(nums).reshape(3, 3)

def _2x2(*nums):
    return np.array(nums).reshape(2, 2)

def _(dims, *nums):
    return np.array(nums).reshape(*dims)

def rref(matrix):
    shape = matrix.shape
    return np.array(
        flatten( Matrix(matrix).rref()[0] )  # ignore pivot columns for now
    ).astype(np.float64).reshape(shape)