Refactor contact_matrix() to use assign_age_groups()

NAMESPACE

@@ -2,6 +2,7 @@
 
 S3method(check,contact_survey)
 S3method(clean,contact_survey)
+export(agegroups_to_limits)
 export(as_contact_survey)
 export(assign_age_groups)
 export(check)

NEWS.md

@@ -1,5 +1,11 @@
 # socialmixr (development version)
 
+* `contact_matrix()` now uses `assign_age_groups()` internally, reducing code
+  duplication and demonstrating the modular workflow (#227).
+
+* New `agegroups_to_limits()` function converts age group labels back to lower
+  age limits, the inverse of `limits_to_agegroups()`.
+
 * `get_survey()`, `download_survey()`, `list_surveys()`, `get_citation()`, and
   `survey_countries()` now warn unconditionally when called. These functions
   were soft-deprecated in 0.5.0 and users should switch to the

R/contact-matrix-utils.R

@@ -486,10 +486,7 @@ survey_pop_from_countries <- function(
   }
 
   if (survey_representative) {
-    survey_pop <- participants[,
-      lower.age.limit := reduce_agegroups(part_age, age_limits)
-    ]
-    survey_pop <- survey_pop[, list(population = .N), by = lower.age.limit]
+    survey_pop <- participants[, list(population = .N), by = lower.age.limit]
     survey_pop <- survey_pop[!is.na(lower.age.limit)]
     if ("year" %in% colnames(participants)) {
       survey_year <- participants[, median(year, na.rm = TRUE)]

R/contact_matrix.R

@@ -258,60 +258,14 @@ contact_matrix <- function(
   # If a survey contains data from multiple countries or if countries specified
   survey$participants <- filter_countries(survey$participants, countries)
 
-  ## Process participant ages: deal with ranges and missing data ---------------
-  survey$participants <- add_part_age(survey$participants)
-
-  ## sample estimated participant ages
-  survey$participants <- impute_participant_ages(
-    participants = survey$participants,
-    estimate = estimated_participant_age
-  )
-
-  # define age limits if not given
-  age_limits <- age_limits %||% get_age_limits(survey)
-
-  survey$participants <- drop_invalid_ages(
-    participants = survey$participants,
-    missing_action = missing_participant_age,
-    age_limits = age_limits
-  )
-
-  ## Process contact ages: deal with ranges and missing data -------------------
-  ## set contact age if it's not in the data
-  survey$contacts <- add_contact_age(survey$contacts)
-
-  ## convert factors to integers, preserving numeric values
-  survey$contacts <- convert_factor_to_integer(
-    data = survey$contacts,
-    cols = c(
-      "cnt_age",
-      "cnt_age_est_min",
-      "cnt_age_est_max",
-      "cnt_age_exact"
-    )
-  )
-
-  ## sample estimated contact ages
-  survey$contacts <- impute_contact_ages(
-    contacts = survey$contacts,
-    estimate = estimated_contact_age
-  )
-
-  # remove contact ages below the age limit, before dealing with missing contact ages
-  survey$contacts <- drop_ages_below_age_limit(
-    data = survey$contacts,
-    age_limits = age_limits
-  )
-
-  survey$participants <- drop_invalid_contact_ages(
-    contacts = survey$contacts,
-    participants = survey$participants,
-    missing_action = missing_contact_age
-  )
-
-  survey$contacts <- drop_missing_contact_ages(
-    contacts = survey$contacts,
-    missing_action = missing_contact_age
+  ## Process ages: impute from ranges, handle missing, assign age groups -------
+  survey <- assign_age_groups(
+    survey,
+    age_limits = age_limits,
+    estimated_participant_age = estimated_participant_age,
+    estimated_contact_age = estimated_contact_age,
+    missing_participant_age = missing_participant_age,
+    missing_contact_age = missing_contact_age
   )
 
   ## check if any filters have been requested ----------------------------------
@@ -321,11 +275,8 @@ contact_matrix <- function(
     filter = filter
   )
 
-  ## adjust age.group.breaks to the lower and upper ages in the survey ---------
-  survey$participants <- adjust_ppt_age_group_breaks(
-    participants = survey$participants,
-    age_limits = age_limits
-  )
+  ## recover resolved age_limits from the assigned age groups ------------------
+  age_limits <- agegroups_to_limits(survey$participants$age.group)
 
   ## ---------------------------------------------------------------------------
   ## if split, symmetric, or age weights are requested, get demographic data

R/globals.R

@@ -26,7 +26,6 @@ utils::globalVariables(c(
   "country", # <get_survey_countries>
   "population", # <survey_pop_from_countries>
   "lower.age.limit", # <survey_pop_from_countries>
-  "part_age", # <survey_pop_from_countries>
   "lower.age.limit", # <add_survey_upper_age_limit>
   "population", # <add_survey_upper_age_limit>
   "upper.age.limit", # <adjust_survey_age_groups>

R/limits_to_agegroups.R

@@ -50,5 +50,20 @@ limits_to_agegroups <- function(
   agegroups <- c(agegroups, paste0(limits[length(limits)], "+"))
   agegroups <- factor(agegroups, levels = agegroups, ordered = TRUE)
   names(agegroups) <- limits
-  return(unname(agegroups[as.character(x)]))
+  unname(agegroups[as.character(x)])
+}
+
+#' Convert age groups to lower age limits
+#'
+#' Inverse of [limits_to_agegroups()]. Extracts lower age limits from age group
+#'   labels.
+#' @param x age groups (a factor, as produced by [limits_to_agegroups()] or
+#'   [assign_age_groups()])
+#' @return a numeric vector of lower age limits
+#' @examples
+#' agegroups_to_limits(limits_to_agegroups(c(0, 5, 10), notation = "brackets"))
+#' @export
+agegroups_to_limits <- function(x) {
+  lvls <- if (is.factor(x)) levels(x) else unique(as.character(x))
+  as.numeric(sub("^\\[?(\\d+).*", "\\1", lvls))
 }

tests/testthat/test-agegroups.r

@@ -108,3 +108,29 @@ test_that("pop_age throws warnings or errors", {
     wpp_age("Germany", 2011)
   )
 })
+
+test_that("agegroups_to_limits round-trips with limits_to_agegroups (brackets)", {
+  limits <- c(0, 5, 10)
+  groups <- limits_to_agegroups(limits, notation = "brackets")
+  result <- agegroups_to_limits(groups)
+  expect_identical(result, limits)
+})
+
+test_that("agegroups_to_limits round-trips with limits_to_agegroups (dashes)", {
+  limits <- c(0, 5, 10)
+  groups <- limits_to_agegroups(limits, notation = "dashes")
+  result <- agegroups_to_limits(groups)
+  expect_identical(result, limits)
+})
+
+test_that("agegroups_to_limits works with character input", {
+  groups <- c("[0,5)", "[5,10)", "10+")
+  result <- agegroups_to_limits(groups)
+  expect_identical(result, c(0, 5, 10))
+})
+
+test_that("agegroups_to_limits works with single age group", {
+  groups <- factor("0+", levels = "0+", ordered = TRUE)
+  result <- agegroups_to_limits(groups)
+  expect_identical(result, 0)
+})