pintcompact.py

# -*- coding: utf-8 -*-

from __future__ import (division, unicode_literals,
                        print_function, absolute_import)

import unittest

import math
import copy
import re
import pint

base_powers = {-24: 'yocto',
               -21: 'zepto',
               -18: 'atto',
               -15: 'femto',
               -12: 'pico',
               -9:  'nano',
               -6:  'micro',
               -3:  'milli',
               0:   '',
               3:   'kilo',
               6:   'mega',
               9:   'giga',
               12:  'tera',
               15:  'peta',
               18:  'exa',
               21:  'zetta',
               24:  'yotta'}

prefixes = {'atto',
            'exa',
            'femto',
            'giga',
            'kilo',
            'mega',
            'micro',
            'milli',
            'nano',
            'peta',
            'pico',
            'tera',
            'yocto',
            'yotta',
            'zepto',
            'zetta',
            'centi',
            'deci',
            'deca',
            'hecto'}

prefixes_regex = '^('+'|'.join(prefixes)+')'
find_prefixes = re.compile(prefixes_regex)


def infer_base_unit(q):
    # Possibly change to use ureg.parse_unit_name
    ureg = q._REGISTRY
    base_units = [(find_prefixes.sub('', unit), power)
                  for unit, power in q.units.items()]

    combined_base_unit = ureg.dimensionless
    for unit, power in base_units:
        combined_base_unit *= getattr(ureg, unit) ** power

    return combined_base_unit


def compact(q, unit=None):
    """Convert a quantity to readable base units.

    For dimensionally simple inputs, no unit needs to be provided.
    To get output in terms of a derived unit (such as newtons), use
    the unit parameter.

    >>> import pint
    >>> ureg = pint.UnitRegistry()
    >>> compact(200e-9*ureg.s)
    <Quantity(200.0, 'nanosecond')>
    >>> compact(1e-2*ureg.N)
    <Quantity(10.0, 'millinewton')>
    """
    if q.unitless:
        return q

    if unit is None:
        unit = infer_base_unit(q)

    q_base = q.to(unit)

    magnitude = q_base.magnitude
    unit_items = q_base.units.items()
    unit_str = unit_items[0][0]
    unit_power = unit_items[0][1]

    if unit_power > 0:
        power = math.floor((math.log(magnitude, 1000) / unit_power)) * 3
    else:
        power = math.ceil((math.log(magnitude, 1000) / unit_power)) * 3

    if power <= -24:
        prefix = 'yocto'
    elif power >= 24:
        prefix = 'yotta'
    else:
        prefix = base_powers[power]

    new_unit = "{prefix}{unit_str}^{unit_power}".format(prefix=prefix,
                                                        unit_str=unit_str,
                                                        unit_power=unit_power)
    if len(unit_items) == 1:
        return q.to(new_unit)
    else:
        # If there are more units in the registry, combine them now;
        # This unit might not make a lot of sense, but it will at least be valid
        combined_new_unit = q._REGISTRY(new_unit)
        for unit, power in unit_items[1:]:
            combined_new_unit *= getattr(q._REGISTRY, unit) ** power

        return q.to(combined_new_unit)


class Test_compact(unittest.TestCase):

    def setUp(self):
        self.ureg = pint.UnitRegistry()
        self.Q_ = self.ureg.Quantity

    def assertQuantityAlmostIdentical(self, q1, q2):
        self.assertAlmostEqual(q1.magnitude, q2.magnitude)
        self.assertEqual(q1.units, q2.units)

    def test_dimensionally_simple_units(self):
        ureg = self.ureg
        inital_converted_quantities = [
            (1*ureg.m, 1*ureg.m),
            (1e-9*ureg.m, 1*ureg.nm)]

        for q, expected_compact_q in inital_converted_quantities:
            compact_q = compact(q)
            self.assertQuantityAlmostIdentical(compact_q, expected_compact_q)

    def test_power_units(self):
        ureg = self.ureg
        inital_converted_quantities = [
            (900*ureg.m**2, 900*ureg.m**2),
            (1e7*ureg.m**2, 10*ureg.km**2)]

        for q, expected_compact_q in inital_converted_quantities:
            compact_q = compact(q)
            self.assertQuantityAlmostIdentical(compact_q, expected_compact_q)

    def test_inverse_units(self):
        ureg = self.ureg
        inital_converted_quantities = [
            (1/ureg.m, 1/ureg.m),
            (100e9/ureg.m, 100/ureg.nm)]

        for q, expected_compact_q in inital_converted_quantities:
            compact_q = compact(q)
            self.assertQuantityAlmostIdentical(compact_q, expected_compact_q)

    def test_inverse_square_units(self):
        ureg = self.ureg
        inital_converted_quantities = [
            (1/ureg.m**2, 1/ureg.m**2),
            (1e11/ureg.m**2, 1e5/ureg.mm**2)]

        for q, expected_compact_q in inital_converted_quantities:
            compact_q = compact(q)
            self.assertQuantityAlmostIdentical(compact_q, expected_compact_q)

    def test_fractional_exponent_units(self):
        ureg = self.ureg
        inital_converted_quantities = [
            (1*ureg.m**0.5, 1*ureg.m**0.5),
            (1e-2*ureg.m**0.5, 10*ureg.um**0.5)]

        for q, expected_compact_q in inital_converted_quantities:
            compact_q = compact(q)
            self.assertQuantityAlmostIdentical(compact_q, expected_compact_q)

    def test_derived_units(self):
        ureg = self.ureg
        inital_converted_quantities = [
            (0.5*ureg.megabyte, 500*ureg.kilobyte),
            (1e-11*ureg.N, 10*ureg.pN)]

        for q, expected_compact_q in inital_converted_quantities:
            compact_q = compact(q)
            self.assertQuantityAlmostIdentical(compact_q, expected_compact_q)

    def test_unit_parameter(self):
        ureg = self.ureg
        inital_converted_quantities = [
            (self.Q_(100e-9, 'kg m / s^2'), ureg.N, 100*ureg.nN),
            (self.Q_(101.3e3, 'kg/m/s^2'), ureg.Pa, 101.3*ureg.kPa)]

        for q, unit, expected_compact_q in inital_converted_quantities:
            compact_q = compact(q, unit=unit)
            self.assertQuantityAlmostIdentical(compact_q, expected_compact_q)


def main():
    """Run the testsuite as command line application.
    """
    try:
        unittest.main()
    except Exception as e:
        print('Error: %s' % e)


if __name__ == '__main__':
    main()