importitertoolsimportmathfrompackaging.versionimportparseasparse_versionimportnumpyasnpimportmatplotlib.tickerastickerimportmatplotlib.unitsasunitsclassProxyDelegate:def__init__(self,fn_name,proxy_type):self.proxy_type=proxy_typeself.fn_name=fn_namedef__get__(self,obj,objtype=None):returnself.proxy_type(self.fn_name,obj)classTaggedValueMeta(type):def__init__(self,name,bases,dict):forfn_nameinself._proxies:ifnothasattr(self,fn_name):setattr(self,fn_name,ProxyDelegate(fn_name,self._proxies[fn_name]))classPassThroughProxy:def__init__(self,fn_name,obj):self.fn_name=fn_nameself.target=obj.proxy_targetdef__call__(self,*args):fn=getattr(self.target,self.fn_name)ret=fn(*args)returnretclassConvertArgsProxy(PassThroughProxy):def__init__(self,fn_name,obj):super().__init__(fn_name,obj)self.unit=obj.unitdef__call__(self,*args):converted_args=[]forainargs:try:converted_args.append(a.convert_to(self.unit))exceptAttributeError:converted_args.append(TaggedValue(a,self.unit))converted_args=tuple([c.get_value()forcinconverted_args])returnsuper().__call__(*converted_args)classConvertReturnProxy(PassThroughProxy):def__init__(self,fn_name,obj):super().__init__(fn_name,obj)self.unit=obj.unitdef__call__(self,*args):ret=super().__call__(*args)return(NotImplementedifretisNotImplementedelseTaggedValue(ret,self.unit))classConvertAllProxy(PassThroughProxy):def__init__(self,fn_name,obj):super().__init__(fn_name,obj)self.unit=obj.unitdef__call__(self,*args):converted_args=[]arg_units=[self.unit]forainargs:ifhasattr(a,'get_unit')andnothasattr(a,'convert_to'):# If this argument has a unit type but no conversion ability,# this operation is prohibited.returnNotImplementedifhasattr(a,'convert_to'):try:a=a.convert_to(self.unit)exceptException:passarg_units.append(a.get_unit())converted_args.append(a.get_value())else:converted_args.append(a)ifhasattr(a,'get_unit'):arg_units.append(a.get_unit())else:arg_units.append(None)converted_args=tuple(converted_args)ret=super().__call__(*converted_args)ifretisNotImplemented:returnNotImplementedret_unit=unit_resolver(self.fn_name,arg_units)ifret_unitisNotImplemented:returnNotImplementedreturnTaggedValue(ret,ret_unit)classTaggedValue(metaclass=TaggedValueMeta):_proxies={'__add__':ConvertAllProxy,'__sub__':ConvertAllProxy,'__mul__':ConvertAllProxy,'__rmul__':ConvertAllProxy,'__cmp__':ConvertAllProxy,'__lt__':ConvertAllProxy,'__gt__':ConvertAllProxy,'__len__':PassThroughProxy}def__new__(cls,value,unit):# generate a new subclass for valuevalue_class=type(value)try:subcls=type(f'TaggedValue_of_{value_class.__name__}',(cls,value_class),{})returnobject.__new__(subcls)exceptTypeError:returnobject.__new__(cls)def__init__(self,value,unit):self.value=valueself.unit=unitself.proxy_target=self.valuedef__copy__(self):returnTaggedValue(self.value,self.unit)def__getattribute__(self,name):ifname.startswith('__'):returnobject.__getattribute__(self,name)variable=object.__getattribute__(self,'value')ifhasattr(variable,name)andnamenotinself.__class__.__dict__:returngetattr(variable,name)returnobject.__getattribute__(self,name)def__array__(self,dtype=object,copy=False):returnnp.asarray(self.value,dtype)def__array_wrap__(self,array,context=None,return_scalar=False):returnTaggedValue(array,self.unit)def__repr__(self):returnf'TaggedValue({self.value!r},{self.unit!r})'def__str__(self):returnf"{self.value} in{self.unit}"def__len__(self):returnlen(self.value)ifparse_version(np.__version__)>=parse_version('1.20'):def__getitem__(self,key):returnTaggedValue(self.value[key],self.unit)def__iter__(self):# Return a generator expression rather than use `yield`, so that# TypeError is raised by iter(self) if appropriate when checking for# iterability.return(TaggedValue(inner,self.unit)forinnerinself.value)defget_compressed_copy(self,mask):new_value=np.ma.masked_array(self.value,mask=mask).compressed()returnTaggedValue(new_value,self.unit)defconvert_to(self,unit):ifunit==self.unitornotunit:returnselftry:new_value=self.unit.convert_value_to(self.value,unit)exceptAttributeError:new_value=selfreturnTaggedValue(new_value,unit)defget_value(self):returnself.valuedefget_unit(self):returnself.unitclassBasicUnit:# numpy scalars convert eager and np.float64(2) * BasicUnit('cm')# would thus return a numpy scalar. To avoid this, we increase the# priority of the BasicUnit.__array_priority__=np.float64(0).__array_priority__+1def__init__(self,name,fullname=None):self.name=nameiffullnameisNone:fullname=nameself.fullname=fullnameself.conversions=dict()def__repr__(self):returnf'BasicUnit({self.name})'def__str__(self):returnself.fullnamedef__call__(self,value):returnTaggedValue(value,self)def__mul__(self,rhs):value=rhsunit=selfifhasattr(rhs,'get_unit'):value=rhs.get_value()unit=rhs.get_unit()unit=unit_resolver('__mul__',(self,unit))ifunitisNotImplemented:returnNotImplementedreturnTaggedValue(value,unit)def__rmul__(self,lhs):returnself*lhsdef__array_wrap__(self,array,context=None,return_scalar=False):returnTaggedValue(array,self)def__array__(self,t=None,context=None,copy=False):ret=np.array(1)iftisnotNone:returnret.astype(t)else:returnretdefadd_conversion_factor(self,unit,factor):defconvert(x):returnx*factorself.conversions[unit]=convertdefadd_conversion_fn(self,unit,fn):self.conversions[unit]=fndefget_conversion_fn(self,unit):returnself.conversions[unit]defconvert_value_to(self,value,unit):conversion_fn=self.conversions[unit]ret=conversion_fn(value)returnretdefget_unit(self):returnselfclassUnitResolver:defaddition_rule(self,units):forunit_1,unit_2initertools.pairwise(units):ifunit_1!=unit_2:returnNotImplementedreturnunits[0]defmultiplication_rule(self,units):non_null=[uforuinunitsifu]iflen(non_null)>1:returnNotImplementedreturnnon_null[0]op_dict={'__mul__':multiplication_rule,'__rmul__':multiplication_rule,'__add__':addition_rule,'__radd__':addition_rule,'__sub__':addition_rule,'__rsub__':addition_rule}def__call__(self,operation,units):ifoperationnotinself.op_dict:returnNotImplementedreturnself.op_dict[operation](self,units)unit_resolver=UnitResolver()cm=BasicUnit('cm','centimeters')inch=BasicUnit('inch','inches')inch.add_conversion_factor(cm,2.54)cm.add_conversion_factor(inch,1/2.54)radians=BasicUnit('rad','radians')degrees=BasicUnit('deg','degrees')radians.add_conversion_factor(degrees,180.0/np.pi)degrees.add_conversion_factor(radians,np.pi/180.0)secs=BasicUnit('s','seconds')hertz=BasicUnit('Hz','Hertz')minutes=BasicUnit('min','minutes')secs.add_conversion_fn(hertz,lambdax:1./x)secs.add_conversion_factor(minutes,1/60.0)# radians formattingdefrad_fn(x,pos=None):ifx>=0:n=int((x/np.pi)*2.0+0.25)else:n=int((x/np.pi)*2.0-0.25)ifn==0:return'0'elifn==1:returnr'$\pi/2$'elifn==2:returnr'$\pi$'elifn==-1:returnr'$-\pi/2$'elifn==-2:returnr'$-\pi$'elifn%2==0:returnfr'${n//2}\pi$'else:returnfr'${n}\pi/2$'classBasicUnitConverter(units.ConversionInterface):@staticmethoddefaxisinfo(unit,axis):"""Return AxisInfo instance for x and unit."""ifunit==radians:returnunits.AxisInfo(majloc=ticker.MultipleLocator(base=np.pi/2),majfmt=ticker.FuncFormatter(rad_fn),label=unit.fullname,)elifunit==degrees:returnunits.AxisInfo(majloc=ticker.AutoLocator(),majfmt=ticker.FormatStrFormatter(r'$%i^\circ$'),label=unit.fullname,)elifunitisnotNone:ifhasattr(unit,'fullname'):returnunits.AxisInfo(label=unit.fullname)elifhasattr(unit,'unit'):returnunits.AxisInfo(label=unit.unit.fullname)returnNone@staticmethoddefconvert(val,unit,axis):ifnp.iterable(val):ifisinstance(val,np.ma.MaskedArray):val=val.astype(float).filled(np.nan)out=np.empty(len(val))fori,thisvalinenumerate(val):ifnp.ma.is_masked(thisval):out[i]=np.nanelse:try:out[i]=thisval.convert_to(unit).get_value()exceptAttributeError:out[i]=thisvalreturnoutifnp.ma.is_masked(val):returnnp.nanelse:returnval.convert_to(unit).get_value()@staticmethoddefdefault_units(x,axis):"""Return the default unit for x or None."""ifnp.iterable(x):forthisxinx:returnthisx.unitreturnx.unitdefcos(x):ifnp.iterable(x):return[math.cos(val.convert_to(radians).get_value())forvalinx]else:returnmath.cos(x.convert_to(radians).get_value())units.registry[BasicUnit]=units.registry[TaggedValue]=BasicUnitConverter()

DownloadJupyternotebook:basic_units.ipynb

DownloadPythonsourcecode:basic_units.py

Downloadzipped:basic_units.zip