Merge pull request #136 from r-xue/skycoord-scalar

properly save and retrieve a SkyCoord object when it's a scalar or N-D array
telegraphic · Jul 28, 2020 · 2be935f · 2be935f
2 parents 8063836 + ef02def
commit 2be935f
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 7 deletions.
diff --git a/hickle/loaders/load_astropy.py b/hickle/loaders/load_astropy.py
@@ -61,7 +61,7 @@ def create_astropy_skycoord(py_obj, h_group, name, **kwargs):
 
     lat = py_obj.data.lat.value
     lon = py_obj.data.lon.value
-    dd = np.column_stack((lon, lat))
+    dd = np.stack((lon, lat), axis=-1)
 
     d = h_group.create_dataset(name, data=dd, dtype='float64', **kwargs)
     lon_unit = str(py_obj.data.lon.unit).encode('ascii')
@@ -172,7 +172,7 @@ def load_astropy_skycoord_dataset(h_node):
     py_type, _, data = get_type_and_data(h_node)
     lon_unit = h_node.attrs["lon_unit"]
     lat_unit = h_node.attrs["lat_unit"]
-    q = py_type(data[:, 0], data[:, 1], unit=(lon_unit, lat_unit))
+    q = py_type(data[..., 0], data[..., 1], unit=(lon_unit, lat_unit))
     return q
 
 

diff --git a/hickle/loaders/load_numpy.py b/hickle/loaders/load_numpy.py
@@ -42,7 +42,7 @@ def create_np_scalar_dataset(py_obj, h_group, name, **kwargs):
             iterable.
     """
 
-    d = h_group.create_dataset(name, data=py_obj, **kwargs)
+    d = h_group.create_dataset(name, data=py_obj)
 
     d.attrs["np_dtype"] = bytes(str(d.dtype), 'ascii')
     return(d)

diff --git a/hickle/tests/test_astropy.py b/hickle/tests/test_astropy.py
@@ -120,21 +120,43 @@ def test_astropy_angle_array():
 
 
 def test_astropy_skycoord():
-    ra = Angle(['1d20m', '1d21m'], unit='degree')
-    dec = Angle(['33d0m0s', '33d01m'], unit='degree')
+    ra = Angle('1d20m', unit='degree')
+    dec = Angle('33d0m0s', unit='degree')
+    radec = SkyCoord(ra, dec)
+    hkl.dump(radec, "test_ap.h5")
+    radec2 = hkl.load("test_ap.h5")
+    assert radec.ra == radec2.ra
+    assert radec.dec == radec2.dec
+
+    ra = Angle('1d20m', unit='hourangle')
+    dec = Angle('33d0m0s', unit='degree')
+    radec = SkyCoord(ra, dec)
+    hkl.dump(radec, "test_ap.h5")
+    radec2 = hkl.load("test_ap.h5")
+    assert radec.ra == radec2.ra
+    assert radec.dec == radec2.dec
+
+
+def test_astropy_skycoord_array():
+    ra = Angle(['1d20m', '0d21m'], unit='degree')
+    dec = Angle(['33d0m0s', '-33d01m'], unit='degree')
     radec = SkyCoord(ra, dec)
     hkl.dump(radec, "test_ap.h5")
     radec2 = hkl.load("test_ap.h5")
     assert np.allclose(radec.ra.value, radec2.ra.value)
     assert np.allclose(radec.dec.value, radec2.dec.value)
+    assert radec.ra.shape == radec2.ra.shape
+    assert radec.dec.shape == radec2.dec.shape
 
-    ra = Angle(['1d20m', '1d21m'], unit='hourangle')
-    dec = Angle(['33d0m0s', '33d01m'], unit='degree')
+    ra = Angle([['1d20m', '0d21m'], ['1d20m', '0d21m']], unit='hourangle')
+    dec = Angle([['33d0m0s', '33d01m'], ['33d0m0s', '33d01m']], unit='degree')
     radec = SkyCoord(ra, dec)
     hkl.dump(radec, "test_ap.h5")
     radec2 = hkl.load("test_ap.h5")
     assert np.allclose(radec.ra.value, radec2.ra.value)
     assert np.allclose(radec.dec.value, radec2.dec.value)
+    assert radec.ra.shape == radec2.ra.shape
+    assert radec.dec.shape == radec2.dec.shape
 
 
 # %% MAIN SCRIPT
@@ -147,3 +169,4 @@ def test_astropy_skycoord():
     test_astropy_angle()
     test_astropy_angle_array()
     test_astropy_skycoord()
+    test_astropy_skycoord_array()