more improvements to arithmetic

Darren Dale · Darren Dale · commit 42166d48b8d0 · 2009-01-17T22:10:17.000-05:00
diff --git a/quantities/dimensionality.py b/quantities/dimensionality.py
@@ -19,7 +19,6 @@ class BaseDimensionality(object):
     def simplified(self):
         if len(self):
             rq = 1*unit_registry['dimensionless']
-            print type(rq), type(rq.dimensionality)
             for u, d in self.iteritems():
                 rq *= u.reference_quantity**d
             return rq.dimensionality
diff --git a/quantities/quantity.py b/quantities/quantity.py
@@ -10,6 +10,8 @@
 from quantities.registry import unit_registry
 
 def prepare_compatible_units(s, o):
+    if not isinstance(o, Quantity):
+        o = Quantity(o, copy=False)
     try:
         assert s.dimensionality.simplified == o.dimensionality.simplified
         return s.simplified, o.simplified
@@ -180,7 +182,6 @@ def __add__(self, other):
         dims = self.dimensionality + other.dimensionality
         ret = super(Quantity, self).__add__(other)
         ret._dimensionality = dims
-
         return ret
 
     # TODO: in-place arithmetic should check for .base, and raise if not None
@@ -189,50 +190,42 @@ def __iadd__(self, other):
         if not isinstance(other, Quantity):
             other = Quantity(other, copy=False)
 
-        sd = self.dimensionality
-        sd += other.dimensionality
-        m = self.magnitude
-        m += other.magnitude
-
-        return self
+        self._dimensionality += other.dimensionality
+        return super(Quantity, self).__iadd__(other)
 
     def __radd__(self, other):
         return self.__add__(other)
 
     def __sub__(self, other):
         if not isinstance(other, Quantity):
-            other = Quantity(other, copy=False)
+            other = numpy.asarray(other).view(Quantity)
 
         dims = self.dimensionality - other.dimensionality
-        magnitude = self.magnitude - other.magnitude
-
-        return Quantity(magnitude, dims, magnitude.dtype)
+        ret = super(Quantity, self).__sub__(other)
+        ret._dimensionality = dims
+        return ret
 
     def __isub__(self, other):
         if not isinstance(other, Quantity):
-            other = Quantity(other, copy=False)
+            other = numpy.asarray(other).view(Quantity)
 
-        sd = self.dimensionality
-        sd -= other.dimensionality
-        m = self.magnitude
-        m -= other.magnitude
-
-        return self
+        self._dimensionality -= other.dimensionality
+        return super(Quantity, self).__isub__(other)
 
     def __rsub__(self, other):
         if not isinstance(other, Quantity):
-            other = Quantity(other, copy=False)
+            other = numpy.asarray(other).view(Quantity)
 
         dims = other.dimensionality - self.dimensionality
-        magnitude = other.magnitude - self.magnitude
-
-        return Quantity(magnitude, dims, magnitude.dtype)
+        ret = super(Quantity, self).__rsub__(other)
+        ret._dimensionality = dims
+        return ret
 
     def __mul__(self, other):
         try:
             dims = self.dimensionality * other.dimensionality
         except AttributeError:
-            other = Quantity(other, copy=False)
+            other = numpy.asarray(other).view(Quantity)
             dims = Dimensionality(self.dimensionality)
 
         ret = super(Quantity, self).__mul__(other)
@@ -241,14 +234,11 @@ def __mul__(self, other):
 
     def __imul__(self, other):
         try:
-            sd = self.dimensionality
-            sd *= other.dimensionality
-            m = self.magnitude
-            m *= other.magnitude
+            self._dimensionality *= other.dimensionality
         except AttributeError:
-            m = self.magnitude
-            m *= other
-        return self
+            pass
+
+        return super(Quantity, self).__imul__(other)
 
     def __rmul__(self, other):
         return self.__mul__(other)
@@ -257,7 +247,7 @@ def __truediv__(self, other):
         try:
             dims = self.dimensionality / other.dimensionality
         except AttributeError:
-            other = Quantity(other, copy=False)
+            other = numpy.asarray(other).view(Quantity)
             dims = Dimensionality(self.dimensionality)
 
         ret = super(Quantity, self).__truediv__(other)
@@ -269,21 +259,24 @@ def __div__(self, other):
 
     def __itruediv__(self, other):
         try:
-            sd = self.dimensionality
-            sd /= other.dimensionality
-            m = self.magnitude
-            m /= other.magnitude
+            self._dimensionality /= other.dimensionality
         except AttributeError:
-            m = self.magnitude
-            m /= other
-        return self
+            pass
+        return super(Quantity, self).__itruediv__(other)
 
     def __idiv__(self, other):
         return self.__itruediv__(other)
 
     def __rtruediv__(self, other):
-        temp = Quantity(1/self.magnitude, self.dimensionality**-1, copy=False)
-        return other * temp
+        try:
+            dims = other.dimensionality / self.dimensionality
+        except AttributeError:
+            other = numpy.asarray(other).view(Quantity)
+            dims = Dimensionality(self.dimensionality**-1)
+
+        ret = super(Quantity, self).__rtruediv__(other)
+        ret._dimensionality = dims
+        return ret
 
     def __rdiv__(self, other):
         return self.__rtruediv__(other)
@@ -294,41 +287,39 @@ def __pow__(self, other):
                 raise ValueError("exponent must be dimensionless")
             other = other.simplified.magnitude
 
-        other = numpy.array(other)
+        other = numpy.asarray(other)
         try:
             assert other.min() == other.max()
             other = other.min()
         except AssertionError:
             raise ValueError('Quantities must be raised to a single power')
 
         dims = self.dimensionality**other
-        magnitude = self.magnitude**other
-        return Quantity(magnitude, dims, magnitude.dtype)
+        ret = super(Quantity, self).__pow__(other)
+        ret._dimensionality = dims
+        return ret
 
     def __ipow__(self, other):
         if isinstance(other, Quantity):
             if other.dimensionality.simplified:
                 raise ValueError("exponent must be dimensionless")
             other = other.simplified.magnitude
 
-        other = numpy.array(other)
+        other = numpy.asarray(other)
         try:
             assert other.min() == other.max()
             other = other.min()
         except AssertionError:
             raise ValueError('Quantities must be raised to a single power')
 
-        sd = self.dimensionality
-        sd **= other
-        m = self.magnitude
-        m **= other
-        return self
+        self._dimensionality **= other
+        return super(Quantity, self).__ipow__(other)
 
     def __rpow__(self, other):
         if self.dimensionality.simplified:
             raise ValueError("exponent must be dimensionless")
 
-        return other**self.simplified.magnitude
+        return super(Quantity, self.simplified).__rpow__(other)
 
     def __repr__(self):
         return '%s %s'%(numpy.ndarray.__str__(self), self.dimensionality)
diff --git a/quantities/tests/test_quantities.py b/quantities/tests/test_quantities.py
@@ -494,7 +494,7 @@ def test_multiplication(self):
             103.0 * q.kPa*q.inch
         )
 
-        self.assertAlmostEqual((5.2 * q.eV) * (300.2 * q.eV), 1561.04 * q.eV**2)
+        self.assertAlmostEqual((5.2 * q.J) * (300.2 * q.J), 1561.04 * q.J**2)
 
         # the formatting should be the same
         self.assertEqual(
diff --git a/quantities/uncertainquantity.py b/quantities/uncertainquantity.py
@@ -137,10 +137,10 @@ def __rtruediv__(self, other):
         return other * temp
 
     def __pow__(self, other):
-        res = Quantity.__pow__(self, other)
+        res = super(UncertainQuantity, self).__pow__(other)
         ru = other * self.relative_uncertainty
-        u = res * ru
-        return UncertainQuantity(res, uncertainty=u, copy=False)
+        res.uncertainty = res * ru
+#        return UncertainQuantity(res, uncertainty=u, copy=False)
 
     def __getitem__(self, key):
         return UncertainQuantity(
diff --git a/quantities/unitquantity.py b/quantities/unitquantity.py
@@ -14,6 +14,17 @@
 ]
 
 
+def quantity(f):
+
+    def wrapped(*args, **kwargs):
+        ret = f(*args, **kwargs)
+        if isinstance(ret, UnitQuantity):
+            return ret.view(Quantity).copy()
+        return ret
+
+    return wrapped
+
+
 class UnitQuantity(Quantity):
 
     _primary_order = 99
@@ -73,6 +84,39 @@ def __repr__(self):
             return s+'\nnote: %s'%self.note
         return s
 
+    @quantity
+    def __mul__(self, other):
+        return super(UnitQuantity, self).__mul__(other)
+
+    @quantity
+    def __rmul__(self, other):
+        return super(UnitQuantity, self).__rmul__(other)
+
+    @quantity
+    def __truediv__(self, other):
+        return super(UnitQuantity, self).__truediv__(other)
+
+    @quantity
+    def __rtruediv__(self, other):
+        return super(UnitQuantity, self).__rtruediv__(other)
+
+    @quantity
+    def __div__(self, other):
+        return super(UnitQuantity, self).__div__(other)
+    @quantity
+    def __rdiv__(self, other):
+        return super(UnitQuantity, self).__rdiv__(other)
+
+    def __imul__(self, other):
+        raise TypeError('can not modify protected units')
+
+    def __itruediv__(self, other):
+        raise TypeError('can not modify protected units')
+
+    def __idiv__(self, other):
+        raise TypeError('can not modify protected units')
+
+
     @property
     def format_order(self):
         return self._format_order

Original file line number	Diff line number	Diff line change
`@@ -494,7 +494,7 @@ def test_multiplication(self):`
`494`	`494`	`103.0 * q.kPa*q.inch`
`495`	`495`	`)`
`496`	`496`
`497`		`- self.assertAlmostEqual((5.2 * q.eV) * (300.2 * q.eV), 1561.04 * q.eV**2)`
	`497`	`+ self.assertAlmostEqual((5.2 * q.J) * (300.2 * q.J), 1561.04 * q.J**2)`
`498`	`498`
`499`	`499`	`# the formatting should be the same`
`500`	`500`	`self.assertEqual(`