基本单位

  1. import math
  3. import numpy as np
  5. import matplotlib.units as units
  6. import matplotlib.ticker as ticker
  7. from matplotlib.cbook import iterable
  10. class ProxyDelegate(object):
  11.     def __init__(self, fn_name, proxy_type):
  12.         self.proxy_type = proxy_type
  13.         self.fn_name = fn_name
  15.     def __get__(self, obj, objtype=None):
  16.         return self.proxy_type(self.fn_name, obj)
  19. class TaggedValueMeta(type):
  20.     def __init__(self, name, bases, dict):
  21.         for fn_name in self._proxies:
  22.             try:
  23.                 dummy = getattr(self, fn_name)
  24.             except AttributeError:
  25.                 setattr(self, fn_name,
  26.                         ProxyDelegate(fn_name, self._proxies[fn_name]))
  29. class PassThroughProxy(object):
  30.     def __init__(self, fn_name, obj):
  31.         self.fn_name = fn_name
  32.         self.target = obj.proxy_target
  34.     def __call__(self, *args):
  35.         fn = getattr(self.target, self.fn_name)
  36.         ret = fn(*args)
  37.         return ret
  40. class ConvertArgsProxy(PassThroughProxy):
  41.     def __init__(self, fn_name, obj):
  42.         PassThroughProxy.__init__(self, fn_name, obj)
  43.         self.unit = obj.unit
  45.     def __call__(self, *args):
  46.         converted_args = []
  47.         for a in args:
  48.             try:
  49.                 converted_args.append(a.convert_to(self.unit))
  50.             except AttributeError:
  51.                 converted_args.append(TaggedValue(a, self.unit))
  52.         converted_args = tuple([c.get_value() for c in converted_args])
  53.         return PassThroughProxy.__call__(self, *converted_args)
  56. class ConvertReturnProxy(PassThroughProxy):
  57.     def __init__(self, fn_name, obj):
  58.         PassThroughProxy.__init__(self, fn_name, obj)
  59.         self.unit = obj.unit
  61.     def __call__(self, *args):
  62.         ret = PassThroughProxy.__call__(self, *args)
  63.         return (NotImplemented if ret is NotImplemented
  64.                 else TaggedValue(ret, self.unit))
  67. class ConvertAllProxy(PassThroughProxy):
  68.     def __init__(self, fn_name, obj):
  69.         PassThroughProxy.__init__(self, fn_name, obj)
  70.         self.unit = obj.unit
  72.     def __call__(self, *args):
  73.         converted_args = []
  74.         arg_units = [self.unit]
  75.         for a in args:
  76.             if hasattr(a, 'get_unit') and not hasattr(a, 'convert_to'):
  77.                 # if this arg has a unit type but no conversion ability,
  78.                 # this operation is prohibited
  79.                 return NotImplemented
  81.             if hasattr(a, 'convert_to'):
  82.                 try:
  83.                     a = a.convert_to(self.unit)
  84.                 except:
  85.                     pass
  86.                 arg_units.append(a.get_unit())
  87.                 converted_args.append(a.get_value())
  88.             else:
  89.                 converted_args.append(a)
  90.                 if hasattr(a, 'get_unit'):
  91.                     arg_units.append(a.get_unit())
  92.                 else:
  93.                     arg_units.append(None)
  94.         converted_args = tuple(converted_args)
  95.         ret = PassThroughProxy.__call__(self, *converted_args)
  96.         if ret is NotImplemented:
  97.             return NotImplemented
  98.         ret_unit = unit_resolver(self.fn_name, arg_units)
  99.         if ret_unit is NotImplemented:
  100.             return NotImplemented
  101.         return TaggedValue(ret, ret_unit)
  104. class TaggedValue(metaclass=TaggedValueMeta):
  106.     _proxies = {'__add__': ConvertAllProxy,
  107.                 '__sub__': ConvertAllProxy,
  108.                 '__mul__': ConvertAllProxy,
  109.                 '__rmul__': ConvertAllProxy,
  110.                 '__cmp__': ConvertAllProxy,
  111.                 '__lt__': ConvertAllProxy,
  112.                 '__gt__': ConvertAllProxy,
  113.                 '__len__': PassThroughProxy}
  115.     def __new__(cls, value, unit):
  116.         # generate a new subclass for value
  117.         value_class = type(value)
  118.         try:
  119.             subcls = type('TaggedValue_of_%s' % (value_class.__name__),
  120.                           tuple([cls, value_class]),
  121.                           {})
  122.             if subcls not in units.registry:
  123.                 units.registry[subcls] = basicConverter
  124.             return object.__new__(subcls)
  125.         except TypeError:
  126.             if cls not in units.registry:
  127.                 units.registry[cls] = basicConverter
  128.             return object.__new__(cls)
  130.     def __init__(self, value, unit):
  131.         self.value = value
  132.         self.unit = unit
  133.         self.proxy_target = self.value
  135.     def __getattribute__(self, name):
  136.         if name.startswith('__'):
  137.             return object.__getattribute__(self, name)
  138.         variable = object.__getattribute__(self, 'value')
  139.         if hasattr(variable, name) and name not in self.__class__.__dict__:
  140.             return getattr(variable, name)
  141.         return object.__getattribute__(self, name)
  143.     def __array__(self, dtype=object):
  144.         return np.asarray(self.value).astype(dtype)
  146.     def __array_wrap__(self, array, context):
  147.         return TaggedValue(array, self.unit)
  149.     def __repr__(self):
  150.         return 'TaggedValue({!r}, {!r})'.format(self.value, self.unit)
  152.     def __str__(self):
  153.         return str(self.value) + ' in ' + str(self.unit)
  155.     def __len__(self):
  156.         return len(self.value)
  158.     def __iter__(self):
  159.         # Return a generator expression rather than use `yield`, so that
  160.         # TypeError is raised by iter(self) if appropriate when checking for
  161.         # iterability.
  162.         return (TaggedValue(inner, self.unit) for inner in self.value)
  164.     def get_compressed_copy(self, mask):
  165.         new_value = np.ma.masked_array(self.value, mask=mask).compressed()
  166.         return TaggedValue(new_value, self.unit)
  168.     def convert_to(self, unit):
  169.         if unit == self.unit or not unit:
  170.             return self
  171.         new_value = self.unit.convert_value_to(self.value, unit)
  172.         return TaggedValue(new_value, unit)
  174.     def get_value(self):
  175.         return self.value
  177.     def get_unit(self):
  178.         return self.unit
  181. class BasicUnit(object):
  182.     def __init__(self, name, fullname=None):
  183.         self.name = name
  184.         if fullname is None:
  185.             fullname = name
  186.         self.fullname = fullname
  187.         self.conversions = dict()
  189.     def __repr__(self):
  190.         return 'BasicUnit(%s)' % self.name
  192.     def __str__(self):
  193.         return self.fullname
  195.     def __call__(self, value):
  196.         return TaggedValue(value, self)
  198.     def __mul__(self, rhs):
  199.         value = rhs
  200.         unit = self
  201.         if hasattr(rhs, 'get_unit'):
  202.             value = rhs.get_value()
  203.             unit = rhs.get_unit()
  204.             unit = unit_resolver('__mul__', (self, unit))
  205.         if unit is NotImplemented:
  206.             return NotImplemented
  207.         return TaggedValue(value, unit)
  209.     def __rmul__(self, lhs):
  210.         return self*lhs
  212.     def __array_wrap__(self, array, context):
  213.         return TaggedValue(array, self)
  215.     def __array__(self, t=None, context=None):
  216.         ret = np.array([1])
  217.         if t is not None:
  218.             return ret.astype(t)
  219.         else:
  220.             return ret
  222.     def add_conversion_factor(self, unit, factor):
  223.         def convert(x):
  224.             return x*factor
  225.         self.conversions[unit] = convert
  227.     def add_conversion_fn(self, unit, fn):
  228.         self.conversions[unit] = fn
  230.     def get_conversion_fn(self, unit):
  231.         return self.conversions[unit]
  233.     def convert_value_to(self, value, unit):
  234.         conversion_fn = self.conversions[unit]
  235.         ret = conversion_fn(value)
  236.         return ret
  238.     def get_unit(self):
  239.         return self
  242. class UnitResolver(object):
  243.     def addition_rule(self, units):
  244.         for unit_1, unit_2 in zip(units[:-1], units[1:]):
  245.             if unit_1 != unit_2:
  246.                 return NotImplemented
  247.         return units[0]
  249.     def multiplication_rule(self, units):
  250.         non_null = [u for u in units if u]
  251.         if len(non_null) > 1:
  252.             return NotImplemented
  253.         return non_null[0]
  255.     op_dict = {
  256.         '__mul__': multiplication_rule,
  257.         '__rmul__': multiplication_rule,
  258.         '__add__': addition_rule,
  259.         '__radd__': addition_rule,
  260.         '__sub__': addition_rule,
  261.         '__rsub__': addition_rule}
  263.     def __call__(self, operation, units):
  264.         if operation not in self.op_dict:
  265.             return NotImplemented
  267.         return self.op_dict[operation](self, units)
  270. unit_resolver = UnitResolver()
  272. cm = BasicUnit('cm', 'centimeters')
  273. inch = BasicUnit('inch', 'inches')
  274. inch.add_conversion_factor(cm, 2.54)
  275. cm.add_conversion_factor(inch, 1/2.54)
  277. radians = BasicUnit('rad', 'radians')
  278. degrees = BasicUnit('deg', 'degrees')
  279. radians.add_conversion_factor(degrees, 180.0/np.pi)
  280. degrees.add_conversion_factor(radians, np.pi/180.0)
  282. secs = BasicUnit('s', 'seconds')
  283. hertz = BasicUnit('Hz', 'Hertz')
  284. minutes = BasicUnit('min', 'minutes')
  286. secs.add_conversion_fn(hertz, lambda x: 1./x)
  287. secs.add_conversion_factor(minutes, 1/60.0)
  290. # radians formatting
  291. def rad_fn(x, pos=None):
  292.     if x >= 0:
  293.         n = int((x / np.pi) * 2.0 + 0.25)
  294.     else:
  295.         n = int((x / np.pi) * 2.0 - 0.25)
  297.     if n == 0:
  298.         return '0'
  299.     elif n == 1:
  300.         return r'$\pi/2$'
  301.     elif n == 2:
  302.         return r'$\pi$'
  303.     elif n == -1:
  304.         return r'$-\pi/2$'
  305.     elif n == -2:
  306.         return r'$-\pi$'
  307.     elif n % 2 == 0:
  308.         return r'$%s\pi$' % (n//2,)
  309.     else:
  310.         return r'$%s\pi/2$' % (n,)
  313. class BasicUnitConverter(units.ConversionInterface):
  314.     @staticmethod
  315.     def axisinfo(unit, axis):
  316.         'return AxisInfo instance for x and unit'
  318.         if unit == radians:
  319.             return units.AxisInfo(
  320.                 majloc=ticker.MultipleLocator(base=np.pi/2),
  321.                 majfmt=ticker.FuncFormatter(rad_fn),
  322.                 label=unit.fullname,
  323.             )
  324.         elif unit == degrees:
  325.             return units.AxisInfo(
  326.                 majloc=ticker.AutoLocator(),
  327.                 majfmt=ticker.FormatStrFormatter(r'$%i^\circ$'),
  328.                 label=unit.fullname,
  329.             )
  330.         elif unit is not None:
  331.             if hasattr(unit, 'fullname'):
  332.                 return units.AxisInfo(label=unit.fullname)
  333.             elif hasattr(unit, 'unit'):
  334.                 return units.AxisInfo(label=unit.unit.fullname)
  335.         return None
  337.     @staticmethod
  338.     def convert(val, unit, axis):
  339.         if units.ConversionInterface.is_numlike(val):
  340.             return val
  341.         if iterable(val):
  342.             return [thisval.convert_to(unit).get_value() for thisval in val]
  343.         else:
  344.             return val.convert_to(unit).get_value()
  346.     @staticmethod
  347.     def default_units(x, axis):
  348.         'return the default unit for x or None'
  349.         if iterable(x):
  350.             for thisx in x:
  351.                 return thisx.unit
  352.         return x.unit
  355. def cos(x):
  356.     if iterable(x):
  357.         return [math.cos(val.convert_to(radians).get_value()) for val in x]
  358.     else:
  359.         return math.cos(x.convert_to(radians).get_value())
  362. basicConverter = BasicUnitConverter()
  363. units.registry[BasicUnit] = basicConverter
  364. units.registry[TaggedValue] = basicConverter

下载这个示例