Validate GroupSelector input before expanding to strata.

Author: szu-yun-ko

inst/include/population_dynamics/population/subpopulation.hpp

@@ -114,6 +114,26 @@ struct PartitionSpec {
       return strata;
     }
 
+    if (group.level.size() != axes.size()) {
+      throw std::invalid_argument(
+          "PartitionSpec::expand_group_to_strata: group.level size " +
+          std::to_string(group.level.size()) +
+          " does not match number of axes " + std::to_string(axes.size()));
+    }
+    for (size_t i = 0; i < axes.size(); ++i) {
+      if (group.level[i] == GroupSelector::kWildcard) {
+        continue;
+      }
+      if (group.level[i] < 0 ||
+          static_cast<size_t>(group.level[i]) >= axes[i].size()) {
+        throw std::invalid_argument(
+            "PartitionSpec::expand_group_to_strata: level index " +
+            std::to_string(group.level[i]) + " out of bounds for axis " +
+            std::to_string(i) + " (size " + std::to_string(axes[i].size()) +
+            ")");
+      }
+    }
+
     std::vector<size_t> current_levels(axes.size(), 0);
     expand_group_to_strata_recursive(group, 0, &current_levels, &strata);
     return strata;

tests/gtest/test_population_Subpopulation.cpp

@@ -57,6 +57,20 @@ TEST(PartitionSpec, ExpandGroupToStrataSelectsSingleStratum) {
             (std::vector<size_t>{0}));
 }
 
+TEST(PartitionSpec, ExpandGroupToStrataRejectsMismatchedGroupSize) {
+  fims_popdy::PartitionSpec spec = fims_popdy::MakeDefaultSexPartitionSpec();
+  fims_popdy::GroupSelector group;
+  group.level = {0, 1};
+  EXPECT_THROW(spec.expand_group_to_strata(group), std::invalid_argument);
+}
+
+TEST(PartitionSpec, ExpandGroupToStrataRejectsOutOfBoundsLevel) {
+  fims_popdy::PartitionSpec spec = fims_popdy::MakeDefaultSexPartitionSpec();
+  fims_popdy::GroupSelector group;
+  group.level = {2};
+  EXPECT_THROW(spec.expand_group_to_strata(group), std::invalid_argument);
+}
+
 TEST(IndexLayout, FoldedIndicesMatchYearAgeAndStratum) {
   fims_popdy::IndexLayout layout;
   layout.n_strata = 2;