Merge branch 'main' of https://github.com/ks905383/xagg

tests/test_export.py

@@ -11,6 +11,55 @@
 from xagg.core import (process_weights,create_raster_polygons,get_pixel_overlaps,aggregate,read_wm)
 from xagg.wrappers import (pixel_overlaps)
 
+
+##### to_dataset() tests #####
+# build raster polygons from a simple 2x2 grid of lat/lon pixels
+ds = xr.Dataset({'test':(['lon','lat','run'],np.array([[[0,1],[2,3]],[[0,1],[2,3]]])),
+				 'lat_bnds':(['lat','bnds'],np.array([[-0.5,0.5],[0.5,1.5]])),
+				 'lon_bnds':(['lon','bnds'],np.array([[-0.5,0.5],[0.5,1.5]]))},
+				coords={'lat':(['lat'],np.array([0,1])),
+						'lon':(['lon'],np.array([0,1])),
+						'run':(['run'],np.array([0,1])),
+						'bnds':(['bnds'],np.array([0,1]))})
+
+# Create polygon covering multiple pixels
+gdf = {'name':['test'],
+			'geometry':[Polygon([(0,0),(0,1),(1,1),(1,0),(0,0)])]}
+gdf = gpd.GeoDataFrame(gdf,crs="EPSG:4326")
+
+# Get pixel overlaps
+wm = pixel_overlaps(ds,gdf)
+
+# Get aggregate
+agg = aggregate(ds,wm)
+
+def test_to_dataset(agg=agg):
+	# Change to dataset
+	ds_out = agg.to_dataset()
+
+	# Build reference output dataset
+	ds_ref = xr.Dataset({'name':(['poly_idx'],np.array(['test']).astype(object)),
+						 'test':(['poly_idx','run'],np.array([[1.0,2.0]]))},
+					coords={'poly_idx':(['poly_idx'],np.array([0])),
+							'run':(['run'],np.array([0,1]))})
+
+	# Assert equal within tolerance, again likely due to very slight 
+	# variation off from actual 1.0, 2.0 due to crs
+	xr.testing.assert_allclose(ds_out,ds_ref,atol=0.0001)
+
+def test_to_dataframe(agg=agg):
+	# Change to dataframe
+	df_out = agg.to_dataframe()
+
+	# Build reference output dataframe
+	df_ref = pd.DataFrame({'poly_idx':[0,0],'run':[0,1],'name':['test','test'],'test':[0.9999,1.9999]})
+	df_ref = df_ref.set_index(['poly_idx','run'])
+
+	# Assert equal within tolerance, again likely due to very slight 
+	# variation off from actual 1.0, 2.0 due to crs
+	pd.testing.assert_frame_equal(df_out,df_ref,atol=0.0001)
+
+
 ##### pixel_overlaps() export tests #####
 
 def test_pixel_overlaps_export_and_import():

xagg/__init__.py

@@ -4,4 +4,4 @@
 # two functions)
 from .wrappers import pixel_overlaps
 from .aux import (normalize,fix_ds,get_bnds,subset_find)
-from .core import aggregate
+from .core import (aggregate,read_wm)

xagg/export.py

@@ -106,7 +106,7 @@ def prep_for_nc(agg_obj,loc_dim='poly_idx'):
 
             # Now insert aggregated values 
             for poly_idx in agg_obj.agg.poly_idx:
-                ds_out[var].loc[{'poly_idx':poly_idx}] = agg_obj.agg.loc[poly_idx,var][0]
+                ds_out[var].loc[{'poly_idx':poly_idx}] = np.squeeze(agg_obj.agg.loc[poly_idx,var])
 
     # Add non-geographic coordinates for the variables to be aggregated
     for crd in [k for k in agg_obj.ds_in.sizes.keys() if (k not in ['lat','lon','loc','bnds'])]: