Skip to content

Commit 0fa5d78

Browse files
ddaleddale
committed
Merge pull request python-quantities#87 from burnpanck/fix-min-max-axis
Fixes axis and out parameters on Quantity methods
2 parents e86419a + 4e1a3ee commit 0fa5d78

2 files changed

Lines changed: 235 additions & 78 deletions

File tree

‎quantities/quantity.py‎

Lines changed: 109 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -422,11 +422,14 @@ def _tolist(self, work_list):
422422

423423
@with_doc(np.ndarray.sum)
424424
def sum(self, axis=None, dtype=None, out=None):
425-
return Quantity(
426-
self.magnitude.sum(axis, dtype, out),
427-
self.dimensionality,
428-
copy=False
429-
)
425+
ret = self.magnitude.sum(axis, dtype, None if out is None else out.magnitude)
426+
dim = self.dimensionality
427+
if out is None:
428+
return Quantity(ret, dim, copy=False)
429+
if not isinstance(out, Quantity):
430+
raise TypeError("out parameter must be a Quantity")
431+
out._dimensionality = dim
432+
return out
430433

431434
@with_doc(np.nansum)
432435
def nansum(self, axis=None, dtype=None, out=None):
@@ -480,11 +483,18 @@ def nonzero(self):
480483

481484
@with_doc(np.ndarray.max)
482485
def max(self, axis=None, out=None):
483-
return Quantity(
484-
self.magnitude.max(),
485-
self.dimensionality,
486-
copy=False
487-
)
486+
ret = self.magnitude.max(axis, None if out is None else out.magnitude)
487+
dim = self.dimensionality
488+
if out is None:
489+
return Quantity(ret, dim, copy=False)
490+
if not isinstance(out, Quantity):
491+
raise TypeError("out parameter must be a Quantity")
492+
out._dimensionality = dim
493+
return out
494+
495+
@with_doc(np.ndarray.argmin)
496+
def argmax(self, axis=None, out=None):
497+
return self.magnitude.argmax(axis, out)
488498

489499
@with_doc(np.nanmax)
490500
def nanmax(self, axis=None, out=None):
@@ -496,11 +506,14 @@ def nanmax(self, axis=None, out=None):
496506

497507
@with_doc(np.ndarray.min)
498508
def min(self, axis=None, out=None):
499-
return Quantity(
500-
self.magnitude.min(),
501-
self.dimensionality,
502-
copy=False
503-
)
509+
ret = self.magnitude.min(axis, None if out is None else out.magnitude)
510+
dim = self.dimensionality
511+
if out is None:
512+
return Quantity(ret, dim, copy=False)
513+
if not isinstance(out, Quantity):
514+
raise TypeError("out parameter must be a Quantity")
515+
out._dimensionality = dim
516+
return out
504517

505518
@with_doc(np.nanmin)
506519
def nanmin(self, axis=None, out=None):
@@ -511,8 +524,8 @@ def nanmin(self, axis=None, out=None):
511524
)
512525

513526
@with_doc(np.ndarray.argmin)
514-
def argmin(self,axis=None, out=None):
515-
return self.magnitude.argmin()
527+
def argmin(self, axis=None, out=None):
528+
return self.magnitude.argmin(axis, out)
516529

517530
@with_doc(np.ndarray.argmax)
518531
def argmax(self,axis=None, out=None):
@@ -528,11 +541,14 @@ def nanargmax(self,axis=None, out=None):
528541

529542
@with_doc(np.ndarray.ptp)
530543
def ptp(self, axis=None, out=None):
531-
return Quantity(
532-
self.magnitude.ptp(),
533-
self.dimensionality,
534-
copy=False
535-
)
544+
ret = self.magnitude.ptp(axis, None if out is None else out.magnitude)
545+
dim = self.dimensionality
546+
if out is None:
547+
return Quantity(ret, dim, copy=False)
548+
if not isinstance(out, Quantity):
549+
raise TypeError("out parameter must be a Quantity")
550+
out._dimensionality = dim
551+
return out
536552

537553
@with_doc(np.ndarray.clip)
538554
def clip(self, min=None, max=None, out=None):
@@ -553,23 +569,35 @@ def clip(self, min=None, max=None, out=None):
553569
max.rescale(self._dimensionality).magnitude,
554570
out
555571
)
556-
return Quantity(clipped, self.dimensionality, copy=False)
572+
dim = self.dimensionality
573+
if out is None:
574+
return Quantity(clipped, dim, copy=False)
575+
if not isinstance(out, Quantity):
576+
raise TypeError("out parameter must be a Quantity")
577+
out._dimensionality = dim
578+
return out
557579

558580
@with_doc(np.ndarray.round)
559581
def round(self, decimals=0, out=None):
560-
return Quantity(
561-
self.magnitude.round(decimals, out),
562-
self.dimensionality,
563-
copy=False
564-
)
582+
ret = self.magnitude.round(decimals, None if out is None else out.magnitude)
583+
dim = self.dimensionality
584+
if out is None:
585+
return Quantity(ret, dim, copy=False)
586+
if not isinstance(out, Quantity):
587+
raise TypeError("out parameter must be a Quantity")
588+
out._dimensionality = dim
589+
return out
565590

566591
@with_doc(np.ndarray.trace)
567592
def trace(self, offset=0, axis1=0, axis2=1, dtype=None, out=None):
568-
return Quantity(
569-
self.magnitude.trace(offset, axis1, axis2, dtype, out),
570-
self.dimensionality,
571-
copy=False
572-
)
593+
ret = self.magnitude.trace(offset, axis1, axis2, dtype, None if out is None else out.magnitude)
594+
dim = self.dimensionality
595+
if out is None:
596+
return Quantity(ret, dim, copy=False)
597+
if not isinstance(out, Quantity):
598+
raise TypeError("out parameter must be a Quantity")
599+
out._dimensionality = dim
600+
return out
573601

574602
@with_doc(np.ndarray.squeeze)
575603
def squeeze(self, axis=None):
@@ -581,10 +609,14 @@ def squeeze(self, axis=None):
581609

582610
@with_doc(np.ndarray.mean)
583611
def mean(self, axis=None, dtype=None, out=None):
584-
return Quantity(
585-
self.magnitude.mean(axis, dtype, out),
586-
self.dimensionality,
587-
copy=False)
612+
ret = self.magnitude.mean(axis, dtype, None if out is None else out.magnitude)
613+
dim = self.dimensionality
614+
if out is None:
615+
return Quantity(ret, dim, copy=False)
616+
if not isinstance(out, Quantity):
617+
raise TypeError("out parameter must be a Quantity")
618+
out._dimensionality = dim
619+
return out
588620

589621
@with_doc(np.nanmean)
590622
def nanmean(self, axis=None, dtype=None, out=None):
@@ -596,19 +628,25 @@ def nanmean(self, axis=None, dtype=None, out=None):
596628

597629
@with_doc(np.ndarray.var)
598630
def var(self, axis=None, dtype=None, out=None, ddof=0):
599-
return Quantity(
600-
self.magnitude.var(axis, dtype, out, ddof),
601-
self._dimensionality**2,
602-
copy=False
603-
)
631+
ret = self.magnitude.var(axis, dtype, out, ddof)
632+
dim = self._dimensionality**2
633+
if out is None:
634+
return Quantity(ret, dim, copy=False)
635+
if not isinstance(out, Quantity):
636+
raise TypeError("out parameter must be a Quantity")
637+
out._dimensionality = dim
638+
return out
604639

605640
@with_doc(np.ndarray.std)
606641
def std(self, axis=None, dtype=None, out=None, ddof=0):
607-
return Quantity(
608-
self.magnitude.std(axis, dtype, out, ddof),
609-
self._dimensionality,
610-
copy=False
611-
)
642+
ret = self.magnitude.std(axis, dtype, out, ddof)
643+
dim = self.dimensionality
644+
if out is None:
645+
return Quantity(ret, dim, copy=False)
646+
if not isinstance(out, Quantity):
647+
raise TypeError("out parameter must be a Quantity")
648+
out._dimensionality = dim
649+
return out
612650

613651
@with_doc(np.nanstd)
614652
def nanstd(self, axis=None, dtype=None, out=None, ddof=0):
@@ -625,15 +663,25 @@ def prod(self, axis=None, dtype=None, out=None):
625663
else:
626664
power = self.shape[axis]
627665

628-
return Quantity(
629-
self.magnitude.prod(axis, dtype, out),
630-
self._dimensionality**power,
631-
copy=False
632-
)
666+
ret = self.magnitude.prod(axis, dtype, None if out is None else out.magnitude)
667+
dim = self._dimensionality**power
668+
if out is None:
669+
return Quantity(ret, dim, copy=False)
670+
if not isinstance(out, Quantity):
671+
raise TypeError("out parameter must be a Quantity")
672+
out._dimensionality = dim
673+
return out
633674

634675
@with_doc(np.ndarray.cumsum)
635676
def cumsum(self, axis=None, dtype=None, out=None):
636-
return super(Quantity, self).cumsum(axis, dtype, out)*self.units
677+
ret = self.magnitude.cumsum(axis, dtype, None if out is None else out.magnitude)
678+
dim = self.dimensionality
679+
if out is None:
680+
return Quantity(ret, dim, copy=False)
681+
if not isinstance(out, Quantity):
682+
raise TypeError("out parameter must be a Quantity")
683+
out._dimensionality = dim
684+
return out
637685

638686
@with_doc(np.ndarray.cumprod)
639687
def cumprod(self, axis=None, dtype=None, out=None):
@@ -642,8 +690,14 @@ def cumprod(self, axis=None, dtype=None, out=None):
642690
raise ValueError(
643691
"Quantity must be dimensionless, try using simplified"
644692
)
645-
else:
646-
return super(Quantity, self).cumprod(axis, dtype, out)
693+
694+
ret = self.magnitude.cumprod(axis, dtype, out)
695+
dim = self.dimensionality
696+
if out is None:
697+
return Quantity(ret, dim, copy=False)
698+
if isinstance(out, Quantity):
699+
out._dimensionality = dim
700+
return out
647701

648702
# list of unsupported functions: [choose]
649703

0 commit comments

Comments
 (0)