From b2ad806be27b8dd73b7eeea829ee3b547cd70a39 Mon Sep 17 00:00:00 2001
From: Marta Alcalde-Herraiz <91142894+martaalcalde@users.noreply.github.com>
Date: Fri, 13 Dec 2024 19:01:57 +0000
Subject: [PATCH] Group age distribution by cohort

---
 R/cohortDiagnostics.R                   | 3 +++
 tests/testthat/test-cohortDiagnostics.R | 9 +++++++++
 2 files changed, 12 insertions(+)

diff --git a/R/cohortDiagnostics.R b/R/cohortDiagnostics.R
index 59c51ea..80bb272 100644
--- a/R/cohortDiagnostics.R
+++ b/R/cohortDiagnostics.R
@@ -70,9 +70,12 @@ cohortDiagnostics <- function(cohort){
 
   cli::cli_bullets(c("*" = "Getting age density"))
   results[["cohort_density"]] <- cdm[[tempCohortName]] |>
+    PatientProfiles::addCohortName() |>
     PatientProfiles::summariseResult(
       strata    = "sex",
       includeOverallStrata = FALSE,
+      group     = "cohort_name",
+      includeOverallGroup  = FALSE,
       variables = "age",
       estimates = "density") |>
     suppressMessages()
diff --git a/tests/testthat/test-cohortDiagnostics.R b/tests/testthat/test-cohortDiagnostics.R
index dcccbec..6bc1f66 100644
--- a/tests/testthat/test-cohortDiagnostics.R
+++ b/tests/testthat/test-cohortDiagnostics.R
@@ -58,6 +58,15 @@ test_that("run with multiple cohorts", {
       dplyr::pull("result_type"),
     "table")))
 
+  # Check density is calculated by cohort
+  expect_identical(result |>
+      dplyr::filter(variable_name == "age") |>
+      dplyr::select("group_level") |>
+      dplyr::distinct() |>
+      dplyr::pull() |>
+      dplyr::arrange(),
+    c("cohort_1", "cohort_2"))
+
   # cohort and timing and overlap should have been estimated now we have more than one cohort
   expect_true(any(stringr::str_detect(
     omopgenerics::settings(result) |>