Adding uncertainty to repr for Spectrum1D

astropy · Jun 28, 2019 · e8f00b1 · e8f00b1
1 parent 5d3e844
commit e8f00b1
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 24 deletions.
diff --git a/docs/arithmetic.rst b/docs/arithmetic.rst
@@ -28,22 +28,19 @@ Arithmetic support includes addition, subtract, multiplication, and division.
     >>> spec2 = Spectrum1D(spectral_axis=np.arange(1, 50) * u.nm, flux=np.random.sample(49)*u.Jy)
     >>> spec3 = spec1 + spec2
     >>> spec3 #doctest:+SKIP
-    Spectrum1D([1.42980955, 0.76450583, 0.53973912, 1.12714653, 1.46747729,
-            0.98485104, 1.13618017, 1.02447445, 0.68610084, 0.85083215,
-            0.57521794, 1.16854341, 1.05139223, 1.44550638, 1.67533841,
-            1.41277807, 0.46008942, 1.1399328 , 0.41708163, 0.61282024,
-            0.40034388, 0.95204057, 0.98280167, 1.30647318, 1.43317265,
-            0.71426198, 0.58622459, 1.17063336, 1.37244261, 1.06886942,
-            1.54349149, 1.15019089, 0.51719866, 1.23114699, 1.16464384,
-            0.90833751, 1.04018595, 1.34931354, 1.01936352, 0.39543304,
-            1.22407522, 0.34658842, 1.18760707, 1.38161461, 1.05829078,
-            1.57852604, 1.13365571, 0.59304282, 1.3913748 ])
-    >>> spec3.wavelength
-    <Quantity [ 10.,  20.,  30.,  40.,  50.,  60.,  70.,  80.,  90., 100.,
-           110., 120., 130., 140., 150., 160., 170., 180., 190., 200.,
-           210., 220., 230., 240., 250., 260., 270., 280., 290., 300.,
-           310., 320., 330., 340., 350., 360., 370., 380., 390., 400.,
-           410., 420., 430., 440., 450., 460., 470., 480., 490.] Angstrom>
+    <Spectrum1D(flux=<Quantity [1.85874504, 0.61621214, 0.83535041, 0.79151086, 0.45719958,
+               1.31989271, 0.27674835, 0.70582565, 1.32792166, 1.2417933 ,
+               0.80759589, 0.30630131, 0.67328771, 1.00480227, 0.34217072,
+               0.35003   , 1.2829507 , 1.47479322, 0.69394109, 1.85822987,
+               0.98505397, 1.77968666, 0.86440386, 0.8547567 , 1.49315833,
+               0.66677101, 1.26206613, 1.27884653, 1.42401291, 1.256702  ,
+               0.97971943, 0.82923743, 1.34281845, 0.70388224, 1.19961318,
+               0.98068789, 1.02384543, 0.31453759, 1.8461661 , 0.79451612,
+               1.24044751, 0.48436633, 0.31850599, 0.82791702, 1.23084396,
+               0.40616336, 0.60637889, 1.88098282, 0.76899619] Jy>, spectral_axis=<Quantity [ 1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9., 10., 11., 12., 13.,
+               14., 15., 16., 17., 18., 19., 20., 21., 22., 23., 24., 25., 26.,
+               27., 28., 29., 30., 31., 32., 33., 34., 35., 36., 37., 38., 39.,
+               40., 41., 42., 43., 44., 45., 46., 47., 48., 49.] nm>)>
 
 
 Propagation of Uncertainties

diff --git a/specutils/spectra/spectrum1d.py b/specutils/spectra/spectrum1d.py
@@ -262,11 +262,14 @@ def __str__(self):
         return result
 
     def __repr__(self):
-        if self.wcs:
-            result = "<Spectrum1D(flux={}, spectral_axis={})>".format(
-                repr(self.flux), repr(self.spectral_axis))
-        else:
-            result = "<Spectrum1D(flux={})>".format(repr(self.flux))
+        inner_str = "flux={}, spectral_axis={}".format(repr(self.flux),
+                                                       repr(self.spectral_axis))
+
+        if self.uncertainty is not None:
+            inner_str += ", uncertainty={}".format(repr(self.uncertainty))
+
+        result = "<Spectrum1D({})>".format(inner_str)
+
         return result
 
 

diff --git a/specutils/tests/test_spectrum1d.py b/specutils/tests/test_spectrum1d.py
@@ -234,9 +234,14 @@ def test_repr():
     assert result.startswith('<Spectrum1D(flux=<Quantity [')
     assert 'spectral_axis=<Quantity [' in result
 
-    spec_without_wcs = Spectrum1D(np.random.random(10) * u.Jy)
-    result = repr(spec_without_wcs)
-    assert result == '<Spectrum1D(flux={})>'.format(repr(spec_without_wcs.flux))
+    spec_with_unc = Spectrum1D(spectral_axis=np.linspace(100, 1000, 10) * u.nm,
+                               flux=np.random.random(10) * u.Jy,
+                               uncertainty=StdDevUncertainty(
+                                   np.random.sample(50), unit='Jy'))
+    result = repr(spec_with_unc)
+    assert result.startswith('<Spectrum1D(flux=<Quantity [')
+    assert 'spectral_axis=<Quantity [' in result
+    assert 'uncertainty=StdDevUncertainty(' in result
 
 
 def test_str():