separate setitem_masked functions for scalars and vectors

Author: arahlin

maps/src/FlatSkyMap.cxx

@@ -1089,28 +1089,31 @@ flatskymap_getitem_masked(const FlatSkyMap &skymap, const G3SkyMapMask &m)
 	return out;
 }
 
+static void
+flatskymap_setitem_masked_scalar(FlatSkyMap &skymap, const G3SkyMapMask &m,
+    double dval)
+{
+	g3_assert(m.IsCompatible(skymap));
+
+	for (auto i : skymap) {
+		if (m.at(i.first))
+			skymap[i.first] = dval;
+	}
+}
+
 static void
 flatskymap_setitem_masked(FlatSkyMap &skymap, const G3SkyMapMask &m,
-    py::object val)
+    const std::vector<double> &val)
 {
 	g3_assert(m.IsCompatible(skymap));
 
-	if (py::extract<double>(val).check()) {
-		double dval = py::extract<double>(val)();
-		for (auto i : skymap) {
-			if (m.at(i.first))
-				skymap[i.first] = dval;
-		}
-	} else {
-		// XXX: the iterable case probably be optimized for numpy arrays
-		// XXX: check for size congruence first?
-		size_t j = 0;
-		for (auto i : skymap) {
-			if (m.at(i.first)) {
-				skymap[i.first] = py::extract<double>(val[j])();
-				j++;
-			}
-		}
+	if (val.size() != m.sum())
+		throw py::value_error("Item dimensions do not match masked area");
+
+	size_t j = 0;
+	for (auto i : skymap) {
+		if (m.at(i.first))
+			skymap[i.first] = val[j++];
 	}
 }
 
@@ -1373,6 +1376,7 @@ PYBINDINGS("maps", scope)
 	    .def("__setitem__", flatskymap_setslice_1d)
 	    .def("__getitem__", flatskymap_getitem_masked)
 	    .def("__setitem__", flatskymap_setitem_masked)
+	    .def("__setitem__", flatskymap_setitem_masked_scalar)
 	;
 
 	// Add buffer protocol interface

maps/src/HealpixSkyMap.cxx

@@ -1101,28 +1101,31 @@ HealpixSkyMap_getitem_masked(const HealpixSkyMap &skymap, const G3SkyMapMask &m)
 	return out;
 }
 
+static void
+HealpixSkyMap_setitem_masked_scalar(HealpixSkyMap &skymap, const G3SkyMapMask &m,
+    double dval)
+{
+	g3_assert(m.IsCompatible(skymap));
+
+	for (auto i : skymap) {
+		if (m.at(i.first))
+			skymap[i.first] = dval;
+	}
+}
+
 static void
 HealpixSkyMap_setitem_masked(HealpixSkyMap &skymap, const G3SkyMapMask &m,
-    py::object val)
+    const std::vector<double> &val)
 {
 	g3_assert(m.IsCompatible(skymap));
 
-	if (py::extract<double>(val).check()) {
-		double dval = py::extract<double>(val)();
-		for (auto i : skymap) {
-			if (m.at(i.first))
-				skymap[i.first] = dval;
-		}
-	} else {
-		// XXX: the iterable case probably be optimized for numpy arrays
-		// XXX: check for size congruence first?
-		size_t j = 0;
-		for (auto i : skymap) {
-			if (m.at(i.first)) {
-				skymap[i.first] = py::extract<double>(val[j])();
-				j++;
-			}
-		}
+	if (val.size() != m.sum())
+		throw py::value_error("Item dimensions do not match masked area");
+
+	size_t j = 0;
+	for (auto i : skymap) {
+		if (m.at(i.first))
+			skymap[i.first] = val[j++];
 	}
 }
 
@@ -1226,9 +1229,9 @@ PYBINDINGS("maps", scope)
 	    .def(py::init<const std::vector<uint64_t> &, const std::vector<double> &,
 	        size_t, bool, bool, MapCoordReference, G3Timestream::TimestreamUnits,
 		G3SkyMap::MapPolType, G3SkyMap::MapPolConv>(),
-	        py::arg("nside"),
 	        py::arg("index"),
 	        py::arg("data"),
+	        py::arg("nside"),
 	        py::arg("weighted") = true,
 	        py::arg("nested") = false,
 	        py::arg("coord_ref") = MapCoordReference::Equatorial,
@@ -1280,6 +1283,7 @@ PYBINDINGS("maps", scope)
 	    .def("__setitem__", &skymap_setitem)
 	    .def("__getitem__", HealpixSkyMap_getitem_masked)
 	    .def("__setitem__", HealpixSkyMap_setitem_masked)
+	    .def("__setitem__", HealpixSkyMap_setitem_masked_scalar)
 
 	    .def("nonzero_pixels", &HealpixSkyMap_nonzeropixels,
 		"Returns a list of the indices of the non-zero pixels in the "