Making a unit system based on prime factorisation (Prototype)

units_prime_proto.py

from math import gcd
from collections import namedtuple

UDIMLESS = 1

# must be prime
UTIME = 2
ULENGTH = 3
UMASS = 5
UCURRENT = 7
UTEMP = 11
USUBSTANCE = 13
ULUMINOSITY = 17

UnitValue = namedtuple('UnitValue', ['value', 'punit', 'nunit'])

def add(a, b):
    assert a.punit == b.punit
    assert a.nunit == b.nunit
    return UnitValue(a.value + b.value, a.punit, a.nunit)

def subtract(a, b):
    return add(a, UnitValue(-b.value, b.punit, b.nunit))

def multiply(a, b): 
    punit = a.punit*b.punit
    nunit = a.nunit*b.nunit
    # the clever bit - use Highest Common Factor (HCF or GCD) to simplify units
    # based on prime factorisation
    hcf = gcd(punit, nunit)
    return UnitValue(a.value*b.value, punit//hcf, nunit//hcf)

def divide(a, b):
    assert b.value != 0
    return multiply(a, UnitValue(1./b.value, b.nunit, b.punit))

# some simple tests
result = add(UnitValue(3.5, UTIME, UDIMLESS), UnitValue(1.5, UTIME, UDIMLESS))
assert result.value == 5.0
assert result.punit == UTIME
assert result.nunit == UDIMLESS

result = subtract(UnitValue(3.5, UCURRENT, UDIMLESS), UnitValue(1.5, UCURRENT, UDIMLESS))
assert result.value == 2.0
assert result.punit == UCURRENT
assert result.nunit == UDIMLESS

result = multiply(UnitValue(2.0, ULENGTH, UDIMLESS), UnitValue(1.5, UTIME, UDIMLESS))
assert result.value == 3.0
assert result.punit == UTIME*ULENGTH
assert result.nunit == UDIMLESS

result = multiply(UnitValue(2.0, ULENGTH, UDIMLESS), UnitValue(1.5, UDIMLESS, UTIME))
assert result.value == 3.0
assert result.punit == ULENGTH
assert result.nunit == UTIME

result = multiply(UnitValue(2.0, ULENGTH, UDIMLESS), UnitValue(1.5, UDIMLESS, ULENGTH))
assert result.value == 3.0
assert result.punit == UDIMLESS
assert result.nunit == UDIMLESS

result = multiply(UnitValue(2.0, ULENGTH*UMASS*UMASS, UTIME), UnitValue(1.5, UTIME*UTIME*UMASS, ULENGTH*ULENGTH*UMASS*UCURRENT))
assert result.value == 3.0
assert result.punit == UMASS*UMASS*UTIME
assert result.nunit == ULENGTH*UCURRENT

result = divide(UnitValue(10.6, ULENGTH*UMASS*UMASS, UTIME), UnitValue(2.0, UTIME*UTIME*UMASS, ULENGTH*ULENGTH*UMASS*UCURRENT))
assert result.value == 5.3
assert result.punit == UMASS*UMASS*ULENGTH*ULENGTH*ULENGTH*UCURRENT
assert result.nunit == UTIME*UTIME*UTIME

result = divide(UnitValue(4.5, ULENGTH, UDIMLESS), UnitValue(1.5, UDIMLESS, ULENGTH))
assert result.value == 3.0
assert result.punit == ULENGTH*ULENGTH
assert result.nunit == UDIMLESS

result = divide(UnitValue(4.5, ULENGTH, UDIMLESS), UnitValue(1.5, ULENGTH, UDIMLESS))
assert result.value == 3.0
assert result.punit == UDIMLESS
assert result.nunit == UDIMLESS

result = divide(UnitValue(4.5, ULENGTH*UTIME, USUBSTANCE), UnitValue(1.5, ULENGTH, USUBSTANCE))
assert result.value == 3.0
assert result.punit == UTIME
assert result.nunit == UDIMLESS

print("All pass!")