quantity.py 31.4 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
#!/bin/env python
"""
Module simtk.unit.quantity

Physical quantities with units, intended to produce similar functionality
to Boost.Units package in C++ (but with a runtime cost).
Uses similar API as Scientific.Physics.PhysicalQuantities
but different internals to satisfy our local requirements.
In particular, there is no underlying set of 'canonical' base
units, whereas in Scientific.Physics.PhysicalQuantities all
units are secretly in terms of SI units.  Also, it is easier
to add new fundamental dimensions to simtk.dimensions.  You
Justin MacCallum's avatar
Justin MacCallum committed
13
might want to make new dimensions for, say, "currency" or
14
15
16
17
18
19
20
21
22
23
24
25
"information".

Some features of this implementation:
  * Quantities are a combination of a value and a unit.  The value
    part can be any python type, including numbers, lists, numpy
    arrays, and anything else.  The unit part must be a simtk.unit.Unit.
  * Operations like adding incompatible units raises an error.
  * Multiplying or dividing units/quantities creates new units.
  * Users can create new Units and Dimensions, but most of the useful
    ones are predefined.
  * Conversion factors between units are applied transitively, so all
    possible conversions are available.
Justin MacCallum's avatar
Justin MacCallum committed
26
27
28
  * I want dimensioned Quantities that are compatible with numpy arrays,
    but do not necessarily require the python numpy package. In other
    words, Quantities can be based on either numpy arrays or on built in
29
    python types.
Justin MacCallum's avatar
Justin MacCallum committed
30
31
32
33
34
  * Units are NOT necessarily stored in terms of SI units internally.
    This is very important for me, because one important application
    area for us is at the molecular scale. Using SI units internally
    can lead to exponent overflow in commonly used molecular force
    calculations. Internally, all unit systems are equally fundamental
35
36
37
38
39
    in SimTK.

Two possible enhancements that have not been implemented are
  1) Include uncertainties with propagation of errors
  2) Incorporate offsets for celsius <-> kelvin conversion
40
41
42
43
44
45
46
47
48
49
50
51



This is part of the OpenMM molecular simulation toolkit originating from
Simbios, the NIH National Center for Physics-Based Simulation of
Biological Structures at Stanford, funded under the NIH Roadmap for
Medical Research, grant U54 GM072970. See https://simtk.org.

Portions copyright (c) 2012 Stanford University and the Authors.
Authors: Christopher M. Bruns
Contributors: Peter Eastman

Justin MacCallum's avatar
Justin MacCallum committed
52
Permission is hereby granted, free of charge, to any person obtaining a
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
copy of this software and associated documentation files (the "Software"),
to deal in the Software without restriction, including without limitation
the rights to use, copy, modify, merge, publish, distribute, sublicense,
and/or sell copies of the Software, and to permit persons to whom the
Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
THE AUTHORS, CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE
USE OR OTHER DEALINGS IN THE SOFTWARE.
69
"""
70
from __future__ import division, print_function, absolute_import
Peter Eastman's avatar
Peter Eastman committed
71

72
73
74
75
76
77
__author__ = "Christopher M. Bruns"
__version__ = "0.5"


import math
import copy
78
79
from .standard_dimensions import *
from .unit import Unit, is_unit, dimensionless
80
81
82

class Quantity(object):
    """Physical quantity, such as 1.3 meters per second.
Justin MacCallum's avatar
Justin MacCallum committed
83

84
85
86
87
88
89
90
91
92
93
94
    Quantities contain both a value, such as 1.3; and a unit,
    such as 'meters per second'.

    Supported value types include:
      1 - numbers (float, int, long)
      2 - lists of numbers, e.g. [1,2,3]
      3 - tuples of numbers, e.g. (1,2,3)
            Note - unit conversions will cause tuples to be converted to lists
      4 - lists of tuples of numbers, lists of lists of ... etc. of numbers
      5 - numpy.arrays
    """
95
    __array_priority__ = 99
Justin MacCallum's avatar
Justin MacCallum committed
96

97
98
99
    def __init__(self, value=None, unit=None):
        """
        Create a new Quantity from a value and a unit.
Justin MacCallum's avatar
Justin MacCallum committed
100

101
102
103
104
105
        Parameters
         - value: (any type, usually a number) Measure of this quantity
         - unit: (Unit) the physical unit, e.g. simtk.unit.meters.
        """
        # When no unit is specified, bend over backwards to handle all one-argument possibilities
106
        if unit is None: # one argument version, copied from UList
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
            if is_unit(value):
                # Unit argument creates an empty list with that unit attached
                unit = value
                value = []
            elif is_quantity(value):
                # Ulist of a Quantity is just the Quantity itself
                unit = value.unit
                value = value._value
            elif _is_string(value):
                unit = dimensionless
            else:
                # Is value a container?
                is_container = True
                try:
                    i = iter(value)
                except TypeError:
                    is_container = False
                if is_container:
                    if len(value) < 1:
                        unit = dimensionless
                    else:
128
                        first_item = next(iter(value))
129
130
                        # Avoid infinite recursion for string, because a one-character
                        # string is its own first element
131
132
133
134
135
136
137
138
139
                        try:
                            isstr = bool(value == first_item)
                        except ValueError:
                            # For numpy, value == first_item returns a numpy
                            # array of booleans, which cannot be evaluated for
                            # truthiness (a ValueError is raised). So in this
                            # case, we don't have a string
                            isstr = False
                        if isstr:
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
                            unit = dimensionless
                        else:
                            unit = Quantity(first_item).unit
                     # Notice that tuples, lists, and numpy.arrays can all be initialized with a list
                    new_container = Quantity([], unit)
                    for item in value:
                        new_container.append(Quantity(item)) # Strips off units into list new_container._value
                    # __class__ trick does not work for numpy.arrays
                    try:
                        import numpy
                        if isinstance(value, numpy.ndarray):
                            value = numpy.array(new_container._value)
                        else:
                            # delegate contruction to container class from list
                            value = value.__class__(new_container._value)
                    except ImportError:
                        # delegate contruction to container class from list
                        value = value.__class__(new_container._value)
                else:
                    # Non-Quantity, non container
                    # Wrap in a dimensionless Quantity
Justin MacCallum's avatar
Justin MacCallum committed
161
                    unit = dimensionless
162
163
164
165
166
        # Accept simple scalar quantities as units
        if is_quantity(unit):
            value = value * unit._value
            unit = unit.unit
        # Use empty list for unspecified values
167
        if value is None:
Justin MacCallum's avatar
Justin MacCallum committed
168
169
            value = []

170
171
        self._value = value
        self.unit = unit
Justin MacCallum's avatar
Justin MacCallum committed
172

173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
    def __getstate__(self):
        state = dict()
        state['_value'] = self._value
        state['unit'] = self.unit
        return state

    def __setstate__(self, state):
        self._value = state['_value']
        self.unit = state['unit']
        return

    def __copy__(self):
        """
        Shallow copy produces a new Quantity with the shallow copy of value and the same unit.
        Because we want copy operations to work just the same way they would on the underlying value.
        """
        return Quantity(copy.copy(self._value), self.unit)

    def __deepcopy__(self, memo):
        """
        Deep copy produces a new Quantity with a deep copy of the value, and the same unit.
        Because we want copy operations to work just the same way they would on the underlying value.
        """
        return Quantity(copy.deepcopy(self._value, memo), self.unit)

    def __getattr__(self, attribute):
        """
        Delegate unrecognized attribute calls to the underlying value type.
        """
        ret_val = getattr(self._value, attribute)
        return ret_val
Justin MacCallum's avatar
Justin MacCallum committed
204

205
206
    def __str__(self):
        """Printable string version of this Quantity.
Justin MacCallum's avatar
Justin MacCallum committed
207

208
209
210
        Returns a string consisting of quantity number followed by unit abbreviation.
        """
        return str(self._value) + ' ' + str(self.unit.get_symbol())
Justin MacCallum's avatar
Justin MacCallum committed
211

212
213
214
215
216
    def __repr__(self):
        """
        """
        return (Quantity.__name__ + '(value=' + repr(self._value) + ', unit=' +
                str(self.unit) + ')')
Justin MacCallum's avatar
Justin MacCallum committed
217

218
219
220
221
222
    def format(self, format_spec):
        return format_spec % self._value + ' ' + str(self.unit.get_symbol())

    def __add__(self, other):
        """Add two Quantities.
Justin MacCallum's avatar
Justin MacCallum committed
223

224
225
        Only Quantities with the same dimensions (e.g. length)
        can be added.  Raises TypeError otherwise.
Justin MacCallum's avatar
Justin MacCallum committed
226

227
228
229
        Parameters
         - self: left hand member of sum
         - other: right hand member of sum
Justin MacCallum's avatar
Justin MacCallum committed
230

231
232
233
234
235
236
237
238
239
240
241
        Returns a new Quantity that is the sum of the two arguments.
        """
        # can only add using like units
        if not self.unit.is_compatible(other.unit):
            raise TypeError('Cannot add two quantities with incompatible units "%s" and "%s".' % (self.unit, other.unit))
        value = self._value + other.value_in_unit(self.unit)
        unit = self.unit
        return Quantity(value, unit)

    def __sub__(self, other):
        """Subtract two Quantities.
Justin MacCallum's avatar
Justin MacCallum committed
242

243
244
        Only Quantities with the same dimensions (e.g. length)
        can be subtracted.  Raises TypeError otherwise.
Justin MacCallum's avatar
Justin MacCallum committed
245

246
247
248
        Parameters
         - self: left hand member (a) of a - b.
         - other: right hand member (b) of a - b.
Justin MacCallum's avatar
Justin MacCallum committed
249

250
251
252
253
254
255
256
        Returns a new Quantity that is the difference of the two arguments.
        """
        if not self.unit.is_compatible(other.unit):
            raise TypeError('Cannot subtract two quantities with incompatible units "%s" and "%s".' % (self.unit, other.unit))
        value = self._value - other.value_in_unit(self.unit)
        unit = self.unit
        return Quantity(value, unit)
Justin MacCallum's avatar
Justin MacCallum committed
257

258
259
260
261
262
    def __eq__(self, other):
        """
        """
        if not is_quantity(other):
            return False
Peter Eastman's avatar
Peter Eastman committed
263
264
265
        if not self.unit.is_compatible(other.unit):
            return False
        return self.value_in_unit(other.unit) == other._value
Justin MacCallum's avatar
Justin MacCallum committed
266

267
268
269
    def __ne__(self, other):
        """
        """
Peter Eastman's avatar
Peter Eastman committed
270
        return not self.__eq__(other)
271

Peter Eastman's avatar
Peter Eastman committed
272
    def __lt__(self, other):
273
        """Compares two quantities.
Justin MacCallum's avatar
Justin MacCallum committed
274

275
        Raises TypeError if the Quantities are of different dimension (e.g. length vs. mass)
Justin MacCallum's avatar
Justin MacCallum committed
276

Peter Eastman's avatar
Peter Eastman committed
277
        Returns True if self < other, False otherwise.
278
        """
Peter Eastman's avatar
Peter Eastman committed
279
        return self._value < other.value_in_unit(self.unit)
280
281

    def __ge__(self, other):
Justin MacCallum's avatar
Justin MacCallum committed
282
        return self._value >= (other.value_in_unit(self.unit))
283
    def __gt__(self, other):
Justin MacCallum's avatar
Justin MacCallum committed
284
        return self._value > (other.value_in_unit(self.unit))
285
    def __le__(self, other):
Justin MacCallum's avatar
Justin MacCallum committed
286
        return self._value <= (other.value_in_unit(self.unit))
287
    def __lt__(self, other):
Justin MacCallum's avatar
Justin MacCallum committed
288
        return self._value < (other.value_in_unit(self.unit))
289

290
291
    _reduce_cache = {}

292
293
294
295
    def reduce_unit(self, guide_unit=None):
        """
        Combine similar component units and scale, to form an
        equal Quantity in simpler units.
Justin MacCallum's avatar
Justin MacCallum committed
296

297
298
        Returns underlying value type if unit is dimensionless.
        """
299
300
301
302
303
304
305
        key = (self.unit, guide_unit)
        if key in Quantity._reduce_cache:
            (unit, value_factor) = Quantity._reduce_cache[key]
        else:
            value_factor = 1.0
            canonical_units = {} # dict of dimensionTuple: (Base/ScaledUnit, exponent)
            # Bias result toward guide units
306
            if guide_unit is not None:
307
308
309
310
311
                for u, exponent in guide_unit.iter_base_or_scaled_units():
                    d = u.get_dimension_tuple()
                    if d not in canonical_units:
                        canonical_units[d] = [u, 0]
            for u, exponent in self.unit.iter_base_or_scaled_units():
312
                d = u.get_dimension_tuple()
313
                # Take first unit found in a dimension as canonical
314
                if d not in canonical_units:
315
316
317
318
319
320
321
322
323
324
325
326
327
                    canonical_units[d] = [u, exponent]
                else:
                    value_factor *= (u.conversion_factor_to(canonical_units[d][0])**exponent)
                    canonical_units[d][1] += exponent
            new_base_units = {}
            for d in canonical_units:
                u, exponent = canonical_units[d]
                if exponent != 0:
                    assert u not in new_base_units
                    new_base_units[u] = exponent
            # Create new unit
            if len(new_base_units) == 0:
                unit = dimensionless
328
            else:
329
330
331
332
333
334
335
336
337
338
                unit = Unit(new_base_units)
            # There might be a factor due to unit conversion, even though unit is dimensionless
            # e.g. suppose unit is meter/centimeter
            if unit.is_dimensionless():
                unit_factor = unit.conversion_factor_to(dimensionless)
                if unit_factor != 1.0:
                    value_factor *= unit_factor
                    # print "value_factor = %s" % value_factor
                unit = dimensionless
            Quantity._reduce_cache[key] = (unit, value_factor)
339
340
341
342
343
344
345
346
347
        # Create Quantity, then scale (in case value is a container)
        # That's why we don't just scale the value.
        result = Quantity(self._value, unit)
        if value_factor != 1.0:
            # __mul__ strips off dimensionless, if appropriate
            result = result * value_factor
        if unit.is_dimensionless():
            assert unit is dimensionless # should have been set earlier in this method
            if is_quantity(result):
348
                result = copy.deepcopy(result._value)
349
350
351
352
        return result

    def __mul__(self, other):
        """Multiply a quantity by another object
Justin MacCallum's avatar
Justin MacCallum committed
353

354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
        Returns a new Quantity that is the product of the self * other,
        unless the resulting unit is dimensionless, in which case the
        underlying value type is returned, instead of a Quantity.
        """
        if is_unit(other):
            # print "quantity * unit"
            # Many other mul/div operations delegate to here because I was debugging
            # a dimensionless unit conversion problem, which I ended up fixing within
            # the reduce_unit() method.
            unit = self.unit * other
            return Quantity(self._value, unit).reduce_unit(self.unit)
        elif is_quantity(other):
            # print "quantity * quantity"
            # Situations where the units cancel can result in scale factors from the unit cancellation.
            # To simplify things, delegate Quantity * Quantity to (Quantity * scalar) * unit
            return (self * other._value) * other.unit
        else:
            # print "quantity * scalar"
            return self._change_units_with_factor(self.unit, other, post_multiply=False)
Justin MacCallum's avatar
Justin MacCallum committed
373

374
375
376
    # value type might not be commutative for multiplication
    def __rmul__(self, other):
        """Multiply a scalar by a Quantity
Justin MacCallum's avatar
Justin MacCallum committed
377
378

        Returns a new Quantity with the same units as self, but with the value
379
380
381
382
383
384
385
386
387
388
389
390
        multiplied by other.
        """
        if is_unit(other):
            raise NotImplementedError('programmer is surprised __rmul__ was called instead of __mul__')
            # print "R unit * quantity"
        elif is_quantity(other):
            # print "R quantity * quantity"
            raise NotImplementedError('programmer is surprised __rmul__ was called instead of __mul__')
        else:
            # print "scalar * quantity"
            return self._change_units_with_factor(self.unit, other, post_multiply=True)
            # return Quantity(other * self._value, self.unit)
Justin MacCallum's avatar
Justin MacCallum committed
391

Peter Eastman's avatar
Peter Eastman committed
392
    def __truediv__(self, other):
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
        """Divide a Quantity by another object

        Returns a new Quantity, unless the resulting unit type is dimensionless,
        in which case the underlying value type is returned.
        """
        if is_unit(other):
            # print "quantity / unit"
            return self * pow(other, -1.0)
            # unit = self.unit / other
            # return Quantity(self._value, unit).reduce_unit(self.unit)
        elif is_quantity(other):
            # print "quantity / quantity"
            # Delegate quantity/quantity to (quantity/scalar)/unit
            return (self/other._value) / other.unit
        else:
            # print "quantity / scalar"
            return self * pow(other, -1.0)
            # return Quantity(self._value / other, self.unit)

412
413
    __div__ = __truediv__

Peter Eastman's avatar
Peter Eastman committed
414
    def __rtruediv__(self, other):
415
        """Divide a scalar by a quantity.
Justin MacCallum's avatar
Justin MacCallum committed
416

417
418
419
420
        Returns a new Quantity.  The resulting units are the inverse of the self argument units.
        """
        if is_unit(other):
            # print "R unit / quantity"
Peter Eastman's avatar
Peter Eastman committed
421
            raise NotImplementedError('programmer is surprised __rtruediv__ was called instead of __truediv__')
422
        elif is_quantity(other):
Peter Eastman's avatar
Peter Eastman committed
423
            raise NotImplementedError('programmer is surprised __rtruediv__ was called instead of __truediv__')
424
425
426
427
428
        else:
            # print "R scalar / quantity"
            return other * pow(self, -1.0)
            # return Quantity(other / self._value, pow(self.unit, -1.0))

429
430
    __rdiv__ = __rtruediv__

431
432
    def __pow__(self, exponent):
        """Raise a Quantity to a power.
Justin MacCallum's avatar
Justin MacCallum committed
433

434
435
436
437
438
        Generally both the value and the unit of the Quantity are affected by this operation.

        Returns a new Quantity equal to self**exponent.
        """
        return Quantity(pow(self._value, exponent), pow(self.unit, exponent))
Justin MacCallum's avatar
Justin MacCallum committed
439

440
441
442
    def sqrt(self):
        """
        Returns square root of a Quantity.
Justin MacCallum's avatar
Justin MacCallum committed
443

444
445
446
447
448
449
450
451
452
453
454
        Raises ArithmeticError if component exponents are not even.
        This behavior can be changed if you present a reasonable real life case to me.
        """
        # There might be a conversion factor from taking the square root of the unit
        new_value = math.sqrt(self._value)
        new_unit = self.unit.sqrt()
        unit_factor = self.unit.conversion_factor_to(new_unit*new_unit)
        if unit_factor != 1.0:
            new_value *= math.sqrt(unit_factor)
        return Quantity(value=new_value, unit=new_unit)

455
    def sum(self, *args, **kwargs):
456
457
458
459
460
461
        """
        Computes the sum of a sequence, with the result having the same unit as
        the current sequence.

        If the value is not iterable, it raises a TypeError (same behavior as if
        you tried to iterate over, for instance, an integer).
462
463
464
465

        This function can take as arguments any arguments recognized by
        `numpy.sum`. If arguments are passed to a non-numpy array, a TypeError
        is raised
466
467
468
        """
        try:
            # This will be much faster for numpy arrays
469
            mysum = self._value.sum(*args, **kwargs)
470
        except AttributeError:
471
472
            if args or kwargs:
                raise TypeError('Unsupported arguments for Quantity.sum')
Peter Eastman's avatar
Peter Eastman committed
473
474
475
476
477
478
            if len(self._value) == 0:
                mysum = 0
            else:
                mysum = self._value[0]
                for i in range(1, len(self._value)):
                    mysum += self._value[i]
479
480
        return Quantity(mysum, self.unit)

481
    def mean(self, *args, **kwargs):
482
483
484
485
486
        """
        Computes the mean of a sequence, with the result having the same unit as
        the current sequence.

        If the value is not iterable, it raises a TypeError
487
488
489
490

        This function can take as arguments any arguments recognized by
        `numpy.mean`. If arguments are passed to a non-numpy array, a TypeError
        is raised
491
492
493
        """
        try:
            # Faster for numpy arrays
494
            mean = self._value.mean(*args, **kwargs)
495
        except AttributeError:
496
497
            if args or kwargs:
                raise TypeError('Unsupported arguments for Quantity.mean')
498
            mean = (self.sum() / len(self._value))._value
499
500
        return Quantity(mean, self.unit)

501
    def std(self, *args, **kwargs):
502
503
504
505
506
        """
        Computes the square root of the variance of a sequence, with the result
        having the same unit as the current sequence.

        If the value is not iterable, it raises a TypeError
507
508
509
510

        This function can take as arguments any arguments recognized by
        `numpy.std`. If arguments are passed to a non-numpy array, a TypeError
        is raised
511
512
513
        """
        try:
            # Faster for numpy arrays
514
            std = self._value.std(*args, **kwargs)
515
        except AttributeError:
516
517
            if args or kwargs:
                raise TypeError('Unsupported arguments for Quantity.std')
518
519
            mean = self.mean()._value
            var = 0
520
            for val in self._value:
Jason Swails's avatar
Jason Swails committed
521
522
523
524
                res = mean - val
                var += res * res
            var /= len(self._value)
            std = math.sqrt(var)
525
526
        return Quantity(std, self.unit)

527
    def max(self, *args, **kwargs):
528
529
530
531
532
        """
        Computes the maximum value of the sequence, with the result having the
        same unit as the current sequence.

        If the value is not iterable, it raises a TypeError
533
534
535
536

        This function can take as arguments any arguments recognized by
        `numpy.max`. If arguments are passed to a non-numpy array, a TypeError
        is raised
537
538
539
        """
        try:
            # Faster for numpy arrays
540
            mymax = self._value.max(*args, **kwargs)
541
        except AttributeError:
542
543
            if args or kwargs:
                raise TypeError('Unsupported arguments for Quantity.max')
544
545
546
            mymax = max(self._value)
        return Quantity(mymax, self.unit)

547
    def min(self, *args, **kwargs):
548
549
550
551
552
        """
        Computes the minimum value of the sequence, with the result having the
        same unit as the current sequence.

        If the value is not iterable, it raises a TypeError
553
554
555
556

        This function can take as arguments any arguments recognized by
        `numpy.min`. If arguments are passed to a non-numpy array, a TypeError
        is raised
557
558
559
        """
        try:
            # Faster for numpy arrays
560
            mymin = self._value.min(*args, **kwargs)
561
        except AttributeError:
562
563
            if args or kwargs:
                raise TypeError('Unsupported arguments for Quantity.min')
564
565
566
            mymin = min(self._value)
        return Quantity(mymin, self.unit)

567
568
569
570
571
572
    def reshape(self, shape, order='C'):
        """
        Same as numpy.ndarray.reshape, except the result is a Quantity with the
        same units as the current object rather than a plain numpy.ndarray
        """
        try:
573
            return Quantity(self._value.reshape(shape, order=order), self.unit)
574
575
576
577
        except AttributeError:
            raise AttributeError('Only numpy array Quantity objects can be '
                                 'reshaped')

578
579
580
    def __abs__(self):
        """
        Return absolute value of a Quantity.
Justin MacCallum's avatar
Justin MacCallum committed
581

582
583
584
585
        The unit is unchanged.  A negative value of self will result in a positive value
        in the result.
        """
        return Quantity(abs(self._value), self.unit)
Justin MacCallum's avatar
Justin MacCallum committed
586

587
588
589
590
591
    def __pos__(self):
        """
        Returns a reference to self.
        """
        return Quantity(+(self._value), self.unit)
Justin MacCallum's avatar
Justin MacCallum committed
592

593
594
    def __neg__(self):
        """Negate a Quantity.
Justin MacCallum's avatar
Justin MacCallum committed
595

596
597
598
        Returns a new Quantity with a different sign on the value.
        """
        return Quantity(-(self._value), self.unit)
Justin MacCallum's avatar
Justin MacCallum committed
599

600
601
602
603
604
    def __nonzero__(self):
        """Returns True if value underlying Quantity is zero, False otherwise.
        """
        return bool(self._value)

605
    def __bool__(self):
Jason Swails's avatar
Jason Swails committed
606
        return bool(self._value)
607

608
609
610
611
612
613
614
615
    def __complex__(self):
        return Quantity(complex(self._value), self.unit)
    def __float__(self):
        return Quantity(float(self._value), self.unit)
    def __int__(self):
        return Quantity(int(self._value), self.unit)
    def __long__(self):
        return Quantity(int(self._value), self.unit)
Justin MacCallum's avatar
Justin MacCallum committed
616

617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
    def value_in_unit(self, unit):
        """
        Returns underlying value, in the specified units.
        """
        val = self.in_units_of(unit)
        if is_quantity(val):
            return val._value
        else: # naked dimensionless
            return val

    def value_in_unit_system(self, system):
        """
        Returns the underlying value type, after conversion to a particular unit system.
        """
        result = self.in_unit_system(system)
        if is_quantity(result):
            return result._value
        else:
            return result # dimensionless
Justin MacCallum's avatar
Justin MacCallum committed
636

637
638
639
640
    def in_unit_system(self, system):
        """
        Returns a new Quantity equal to this one, expressed in a particular unit system.
        """
641
642
643
        new_units = system.express_unit(self.unit)
        f = self.unit.conversion_factor_to(new_units)
        return self._change_units_with_factor(new_units, f)
644
645
646
647
648
649
650

    def in_units_of(self, other_unit):
        """
        Returns an equal Quantity expressed in different units.

        If the units are the same as those in self, a reference to self is returned.
        Raises a TypeError if the new unit is not compatible with the original unit.
Justin MacCallum's avatar
Justin MacCallum committed
651

652
653
654
655
656
657
658
659
        The post_multiply argument is used in case the multiplication operation is not commutative.
          i.e. result = factor * value when post_multiply is False
          and  result = value * factor when post_multiply is True
        """
        if not self.unit.is_compatible(other_unit):
            raise TypeError('Unit "%s" is not compatible with Unit "%s".' % (self.unit, other_unit))
        f = self.unit.conversion_factor_to(other_unit)
        return self._change_units_with_factor(other_unit, f)
Justin MacCallum's avatar
Justin MacCallum committed
660

661
662
663
664
665
666
667
668
669
670
    def _change_units_with_factor(self, new_unit, factor, post_multiply=True):
        # numpy arrays cannot be compared with 1.0, so just "try"
        factor_is_identity = False
        try:
            if (factor == 1.0):
                factor_is_identity = True
        except ValueError:
            pass
        if factor_is_identity:
            # No multiplication required
671
            result = Quantity(copy.deepcopy(self._value), new_unit)
672
673
674
675
676
677
        else:
            try:
                # multiply operator, if it exists, is preferred
                if post_multiply:
                    value = self._value * factor # works for number, numpy.array, or vec3, e.g.
                else:
Justin MacCallum's avatar
Justin MacCallum committed
678
                    value = factor * self._value # works for number, numpy.array, or vec3, e.g.
679
680
                result = Quantity(value, new_unit)
            except TypeError:
681
                value = copy.deepcopy(self._value)
682
                result = Quantity(self._scale_sequence(value, factor, post_multiply), new_unit)
683
684
685
686
687
        if (new_unit.is_dimensionless()):
            return result._value
        else:
            return result

688
689
690
    def _scale_sequence(self, value, factor, post_multiply):
        try:
            if post_multiply:
691
                value = value*factor
692
            else:
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
                value = factor*value
        except TypeError:
            try:
                if post_multiply:
                    if isinstance(value, tuple):
                        value = tuple([x*factor for x in value])
                    else:
                        for i in range(len(value)):
                            value[i] = value[i]*factor
                else:
                    if isinstance(value, tuple):
                        value = tuple([factor*x for x in value])
                    else:
                        for i in range(len(value)):
                            value[i] = factor*value[i]
708
            except TypeError:
709
710
                if isinstance(value, tuple):
                    value = tuple([self._scale_sequence(x, factor, post_multiply) for x in value])
711
712
                else:
                    for i in range(len(value)):
713
                        value[i] = self._scale_sequence(value[i], factor, post_multiply)
714
715
716
        return value


717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734

    ####################################
    ### Sequence methods of Quantity ###
    ###  in case value is a sequence ###
    ####################################

    def __len__(self):
        """
        Return size of internal value type.
        """
        return len(self._value)

    def __getitem__(self, key):
        """
        Keep the same units on contained elements.
        """
        assert not is_quantity(self._value[key])
        return Quantity(self._value[key], self.unit)
Justin MacCallum's avatar
Justin MacCallum committed
735

736
737
738
739
740
741
742
743
744
745
746
747
748
749
    def __setitem__(self, key, value):
        # Delegate slices to one-at-a time ___setitem___
        if isinstance(key, slice): # slice
            indices = key.indices(len(self))
            for i in range(*indices):
                self[i] = value[i]
        else: # single index
            # Check unit compatibility
            if self.unit.is_dimensionless() and is_dimensionless(value):
                pass # OK
            elif not self.unit.is_compatible(value.unit):
                raise TypeError('Unit "%s" is not compatible with Unit "%s".' % (self.unit, value.unit))
            self._value[key] = value / self.unit
            assert not is_quantity(self._value[key])
Justin MacCallum's avatar
Justin MacCallum committed
750

751
752
753
754
755
    def __delitem__(self, key):
        del(self._value[key])

    def __contains__(self, item):
        return self._value.__contains__(item.value_in_unit(self.unit))
Justin MacCallum's avatar
Justin MacCallum committed
756

757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
    def __iter__(self):
        for item in self._value:
            yield Quantity(item, self.unit)

    def count(self, item):
        return self._value.count(item.value_in_unit(self.unit))
    def index(self, item):
        return self._value.index(item.value_in_unit(self.unit))
    def append(self, item):
        if is_quantity(item):
            return self._value.append(item.value_in_unit(self.unit))
        elif is_dimensionless(self.unit):
            return self._value.append(item)
        else:
            raise TypeError("Cannot append item without units into list with units")
    def extend(self, rhs):
        self._value.extend(rhs.value_in_unit(self.unit))
    def insert(self, index, item):
        self._value.insert(index, item.value_in_unit(self.unit))
    def remove(self, item):
        self._value.remove(item)
    def pop(self, *args):
        return self._value.pop(*args) * self.unit
    # list.reverse will automatically delegate correctly
    # list.sort with no arguments will delegate correctly
    # list.sort with a comparison function cannot be done correctly


def is_quantity(x):
    """
    Returns True if x is a Quantity, False otherwise.
    """
    return isinstance(x, Quantity)

def is_dimensionless(x):
    """
    """
    if is_unit(x):
        return x.is_dimensionless()
    elif is_quantity(x):
        return x.unit.is_dimensionless()
    else:
        # everything else in the universe is dimensionless
        return True

# Strings can cause trouble
# as can any container that has infinite levels of containment
def _is_string(x):
     # step 1) String is always a container
     # and its contents are themselves containers.
     if isinstance(x, str):
         return True
     try:
         first_item = iter(x).next()
         inner_item = iter(first_item).next()
         if first_item is inner_item:
             return True
         else:
             return False
     except TypeError:
         return False
     except StopIteration:
         return False

# run module directly for testing
if __name__=='__main__':
    # Test the examples in the docstrings
    import doctest, sys
    doctest.testmod(sys.modules[__name__])