diff --git a/quantities/quantity.py b/quantities/quantity.py index 135c18d4..ec783855 100644 --- a/quantities/quantity.py +++ b/quantities/quantity.py @@ -13,6 +13,10 @@ from .registry import unit_registry from .decorators import with_doc +PREFERRED = [] # List of preferred quanitities for each symbol, + # e.g. PREFERRED = [pq.mV, pq.pA, pq.UnitQuantity('femtocoulomb', 1e-15*pq.C, 'fC')] + # Intended to be overwritten in down-stream packages + if sys.version.startswith('3'): unicode = str @@ -195,10 +199,17 @@ def units(self, units): mag *= cf self._dimensionality = to_u.dimensionality - def rescale(self, units): + def rescale(self, units=None): """ - Return a copy of the quantity converted to the specified units + Return a copy of the quantity converted to the specified units. + If `units` is `None`, an attempt will be made to rescale the quantity + to preferred units (see `rescale_preferred`). """ + if units is None: + try: + return self.rescale_preferred() + except Exception as e: + raise Exception('No argument passed to `.rescale` and %s' % e) to_dims = validate_dimensionality(units) if self.dimensionality == to_dims: return self.astype(self.dtype) @@ -212,7 +223,27 @@ def rescale(self, units): %(from_u._dimensionality, to_u._dimensionality) ) return Quantity(cf*self.magnitude, to_u) - + + def rescale_preferred(self): + """ + Return a copy of the quantity converted to the preferred units and scale. + These will be identified from among the compatible units specified in the + list PREFERRED in this module. For example, a voltage quantity might be + converted to `mV`: + ``` + import quantities as pq + pq.quantity.PREFERRED = [pq.mV, pq.pA] + old = 3.1415 * pq.V + new = old.rescale_preferred() # `new` will be 3141.5 mV. + ``` + """ + units_str = str(self.simplified.dimensionality) + for preferred in PREFERRED: + if units_str == str(preferred.simplified.dimensionality): + return self.rescale(preferred) + raise Exception(("Preferred units for '%s' (or equivalent) not specified in " + "quantites.quantity.PREFERRED." % self.dimensionality)) + @with_doc(np.ndarray.astype) def astype(self, dtype=None, **kwargs): '''Scalars are returned as scalar Quantity arrays.''' diff --git a/quantities/tests/test_conversion.py b/quantities/tests/test_conversion.py index 7fb8502e..fb042a34 100644 --- a/quantities/tests/test_conversion.py +++ b/quantities/tests/test_conversion.py @@ -1,6 +1,8 @@ # -*- coding: utf-8 -*- +import unittest from .. import units as pq +from .. import quantity from .common import TestCase @@ -15,6 +17,44 @@ def test_inplace_conversion(self): def test_rescale(self): for u in ('ft', 'feet', pq.ft): self.assertQuantityEqual((10*pq.m).rescale(u), 32.80839895 * pq.ft) + + def test_rescale_preferred(self): + quantity.PREFERRED = [pq.mV, pq.pA] + q = 10*pq.V + self.assertQuantityEqual(q.rescale_preferred(), q.rescale(pq.mV)) + q = 5*pq.A + self.assertQuantityEqual(q.rescale_preferred(), q.rescale(pq.pA)) + quantity.PREFERRED = [] + + def test_rescale_preferred_failure(self): + quantity.PREFERRED = [pq.pA] + q = 10*pq.V + try: + self.assertQuantityEqual(q.rescale_preferred(), q.rescale(pq.mV)) + except: + self.assertTrue(True) + else: + self.assertTrue(False) + quantity.PREFERRED = [] + + def test_rescale_noargs(self): + quantity.PREFERRED = [pq.mV, pq.pA] + q = 10*pq.V + self.assertQuantityEqual(q.rescale(), q.rescale(pq.mV)) + q = 5*pq.A + self.assertQuantityEqual(q.rescale(), q.rescale(pq.pA)) + quantity.PREFERRED = [] + + def test_rescale_noargs_failure(self): + quantity.PREFERRED = [pq.pA] + q = 10*pq.V + try: + self.assertQuantityEqual(q.rescale_preferred(), q.rescale(pq.mV)) + except: + self.assertTrue(True) + else: + self.assertTrue(False) + quantity.PREFERRED = [] def test_compound_reduction(self): pc_per_cc = pq.CompoundUnit("pc/cm**3")