Custom cohorts feature with temporal refactor to remove unnecessary batching

src/main/groovy/edu/washu/tag/generator/BatchRequest.groovy

@@ -1,5 +1,12 @@
 package edu.washu.tag.generator
 
+import edu.washu.tag.generator.metadata.Patient
+import edu.washu.tag.generator.metadata.cohorting.SpecializedCohort
+import edu.washu.tag.generator.util.SequentialIdGenerator
+import io.temporal.activity.Activity
+import org.slf4j.Logger
+import org.slf4j.LoggerFactory
+
 class BatchRequest {
 
     int id
@@ -8,5 +15,89 @@ class BatchRequest {
     int numSeries
     int patientOffset
     int studyOffset
+    String temporalSummary
+    SpecializedCohort cohort
+    private static final Logger logger = LoggerFactory.getLogger(BatchRequest)
+
+    void generateBatchWithHandler(BatchRequestContext batchRequestContext) {
+        final SpecificationParameters specificationParameters = batchRequestContext.specificationParameters
+        specificationParameters.postprocess()
+        batchRequestContext.generationCache.cache()
+
+        final GenerationContext generationContext = new GenerationContext(
+            specificationParameters: specificationParameters,
+            patientIdGenerators: batchRequestContext.idOffsets.getPatientIdGeneratorsFromOffset(patientOffset),
+            studyCountOverride: 0
+        )
+        int generatedStudies = 0
+        int generatedSeries = 0
+        final SequentialIdGenerator studyIdGenerator = batchRequestContext.idOffsets.getStudyIdEncoderFromOffset(studyOffset)
+        final List<Patient> patients = []
+
+        if (cohort == null) {
+            numPatients.times { generatedPatients ->
+                final Closure<Integer> remainingStudies = { numStudies - generatedStudies }
+
+                generationContext.setCurrentAverageStudiesPerPatient(generatedPatients == 0 ? 0.0 : generatedStudies / generatedPatients)
+                generationContext.setPreviouslyGeneratedSeries(generatedSeries)
+                generationContext.setPreviouslyGeneratedStudies(generatedStudies)
+                if (generatedPatients > numPatients - 4) {
+                    generationContext.setStudyCountMaximum(remainingStudies() - (numPatients - generatedPatients - 1))
+                } // 3rd to last must leave 2 reports for final 2 patients, 2nd to last must leave a report for final patient
+                if (generatedPatients == numPatients - 1) { // for last patient in batch, we need to ensure exact number of studies
+                    generationContext.setStudyCountOverride(remainingStudies())
+                }
+                final Patient patient = batchRequestContext.patientRandomizers.sample().createPatient(specificationParameters)
+                patient.randomize(generationContext, studyIdGenerator)
+                patient.studies.each { study ->
+                    generatedStudies++
+                    generatedSeries += study.series.size()
+                    if (!specificationParameters.generateRadiologyReports) {
+                        study.cachePatientIdsForStudy()
+                    }
+                }
+                if (batchRequestContext.temporalHeartbeat) {
+                    Activity.executionContext.heartbeat("Batch ${id}, patient ${generatedPatients + 1} post-DICOM")
+                }
+                patients << patient
+            }
+
+            logger.info("DICOM specs for batch ${id} with ${patients.size()} patients and ${patients*.studies*.size().sum()} studies")
+
+            if (specificationParameters.generateRadiologyReports && batchRequestContext.generateReports) {
+                specificationParameters.reportGeneratorImplementation.generateReportsForPatients(patients, batchRequestContext.temporalHeartbeat)
+            }
+        } else {
+            numPatients.times { generatedPatients ->
+                final Patient patient = batchRequestContext.patientRandomizers.sample().createPatient(specificationParameters)
+                patient.randomize(generationContext, studyIdGenerator, cohort)
+                patient.studies.each { study ->
+                    if (!specificationParameters.generateRadiologyReports) {
+                        study.cachePatientIdsForStudy()
+                    }
+                }
+                if (batchRequestContext.temporalHeartbeat) {
+                    Activity.executionContext.heartbeat("Cohort batch ${id}, patient ${generatedPatients + 1} post-DICOM")
+                }
+                patients << patient
+            }
+
+            logger.info("DICOM specs for cohort batch ${id} with ${patients.size()} patients and ${patients*.studies*.size().sum()} studies")
+
+            if (specificationParameters.generateRadiologyReports && batchRequestContext.generateReports) {
+                specificationParameters.reportGeneratorImplementation.generateReportsForPatients(patients, batchRequestContext.temporalHeartbeat)
+            }
+        }
+
+        patients.eachWithIndex { patient, generatedPatients ->
+            batchRequestContext.handler.accept(patient)
+            if ((generatedPatients + 1) % 100 == 0) {
+                logger.info("Generated DICOM and/or HL7 for patient ${patient.patientInstanceUid} in batch ${id} [${generatedPatients + 1}/${numPatients}]")
+            }
+            if (batchRequestContext.temporalHeartbeat) {
+                Activity.executionContext.heartbeat("Batch ${id}, patient ${generatedPatients + 1}")
+            }
+        }
+    }
 
 }

src/main/groovy/edu/washu/tag/generator/BatchRequestContext.groovy

@@ -0,0 +1,23 @@
+package edu.washu.tag.generator
+
+import edu.washu.tag.generator.metadata.GenerationCache
+import edu.washu.tag.generator.metadata.Patient
+import edu.washu.tag.generator.metadata.patient.PatientRandomizer
+import groovy.transform.builder.Builder
+import groovy.transform.builder.SimpleStrategy
+import org.apache.commons.math3.distribution.EnumeratedDistribution
+
+import java.util.function.Consumer
+
+@Builder(builderStrategy = SimpleStrategy, prefix = '')
+class BatchRequestContext {
+
+    SpecificationParameters specificationParameters
+    GenerationCache generationCache
+    IdOffsets idOffsets
+    EnumeratedDistribution<PatientRandomizer> patientRandomizers
+    Consumer<Patient> handler
+    boolean generateReports = true
+    boolean temporalHeartbeat = false
+
+}

src/main/groovy/edu/washu/tag/generator/Batcher.groovy

@@ -0,0 +1,91 @@
+package edu.washu.tag.generator
+
+import java.util.function.Function
+
+class Batcher {
+
+    SpecificationParameters specificationParameters
+    int patientsPerFullBatch
+    int patientOffset = 0
+    int studyOffset = 0
+    int idOffset = 0
+
+    Batcher(SpecificationParameters specificationParameters, int patientsPerFullBatch = BatchSpecification.MAX_PATIENTS) {
+        new File('batches').mkdir() // while we're still in a single process
+        this.specificationParameters = specificationParameters
+        this.patientsPerFullBatch = patientsPerFullBatch
+    }
+
+    /**
+     * Resolves the entire requested population into a flat list of standalone batches. Each {@link BatchRequest}
+     * carries its own patient/study offsets, so it can be fulfilled independently by any worker with no
+     * coordination between batches.
+     */
+    List<BatchRequest> resolveBatches() {
+        patientOffset = 0
+        studyOffset = 0
+        idOffset = 0
+
+        final List<BatchRequest> allBatches = computeBatchesFrom(
+            specificationParameters.numPatients,
+            specificationParameters.numStudies,
+            specificationParameters.numSeries,
+            { Integer id -> "standard batch ${id}" }
+        )
+
+        if (!specificationParameters.cohorts.isEmpty()) {
+            patientOffset += specificationParameters.numPatients
+            studyOffset += specificationParameters.numStudies
+            idOffset += allBatches.size()
+
+            specificationParameters.cohorts.each { cohort ->
+                cohort.trajectory.each { studyReq ->
+                    studyReq.protocol.postprocess(specificationParameters)
+                }
+                final int cohortPatients = cohort.numPatients
+                final int cohortStudies = cohortPatients * cohort.trajectory.size()
+                final List<BatchRequest> batchesForCohort = computeBatchesFrom(
+                    cohortPatients,
+                    cohortStudies,
+                    0, // obviously not true, but we don't need this
+                    { id -> "cohort ${cohort.name} batch ${id}" }
+                )
+                batchesForCohort.each { batch ->
+                    batch.setCohort(cohort)
+                    allBatches << batch
+                }
+                patientOffset += cohortPatients
+                studyOffset += cohortStudies
+                idOffset += batchesForCohort.size()
+            }
+        }
+        allBatches
+    }
+
+    private List<BatchRequest> computeBatchesFrom(int numPatients, int numStudies, int numSeries, Function<Integer, String> summarizer) {
+        if (numPatients == 0) {
+            return []
+        }
+        final int totalNumBatches = Math.ceilDiv(numPatients, patientsPerFullBatch)
+
+        final int patientsInSmallerBatches = numPatients / totalNumBatches
+        final int patientRemainder = numPatients % totalNumBatches
+        final int studiesInSmallerBatch = numStudies / totalNumBatches
+        final int studyRemainder = numStudies % totalNumBatches
+        final int seriesInSmallerBatch = numSeries / totalNumBatches
+        final int seriesRemainder = numSeries % totalNumBatches
+
+        (0 ..< totalNumBatches).collect { batchId ->
+            new BatchRequest(
+                id: idOffset + batchId,
+                numPatients: patientsInSmallerBatches + (batchId < patientRemainder ? 1 : 0),
+                numStudies: studiesInSmallerBatch + (batchId < studyRemainder ? 1 : 0),
+                numSeries: seriesInSmallerBatch + (batchId < seriesRemainder ? 1 : 0),
+                patientOffset: patientOffset + (patientsInSmallerBatches * batchId) + Math.min(batchId, patientRemainder),
+                studyOffset: studyOffset + (studiesInSmallerBatch * batchId) + Math.min(batchId, studyRemainder),
+                temporalSummary: summarizer.apply(batchId)
+            )
+        }
+    }
+
+}

src/main/groovy/edu/washu/tag/generator/PopulationGenerator.groovy

@@ -40,7 +40,7 @@ class PopulationGenerator {
         final IdOffsets idOffsets = new IdOffsets()
         CodeCache.initializeCache()
 
-        final List<BatchRequest> batchRequests = generator.resolveBatches()
+        final List<BatchRequest> batchRequests = new Batcher(generator.specificationParameters).resolveBatches()
         final String batchFulfillment = batchRequests.size() > 1 ? "split into ${batchRequests.size()} batches" : 'fulfilled in a single batch'
         println("STAGE 1: request will be ${batchFulfillment}")
 
@@ -62,39 +62,6 @@ class PopulationGenerator {
         }
     }
 
-    /**
-     * Resolves the entire requested population into a flat list of standalone batches. Each {@link BatchRequest}
-     * carries its own patient/study offsets, so it can be fulfilled independently by any worker with no
-     * coordination between batches.
-     */
-    List<BatchRequest> resolveBatches(int patientsPerFullBatch = BatchSpecification.MAX_PATIENTS) {
-        new File('batches').mkdir() // while we're still in a single process
-
-        final int numPatients = specificationParameters.numPatients
-        final int numStudies = specificationParameters.numStudies
-        final int numSeries = specificationParameters.numSeries
-
-        final int totalNumBatches = Math.ceilDiv(numPatients, patientsPerFullBatch)
-
-        final int patientsInSmallerBatches = numPatients / totalNumBatches
-        final int patientRemainder = numPatients % totalNumBatches
-        final int studiesInSmallerBatch = numStudies / totalNumBatches
-        final int studyRemainder = numStudies % totalNumBatches
-        final int seriesInSmallerBatch = numSeries / totalNumBatches
-        final int seriesRemainder = numSeries % totalNumBatches
-
-        (0 ..< totalNumBatches).collect { batchId ->
-            new BatchRequest(
-                id: batchId,
-                numPatients: patientsInSmallerBatches + (batchId < patientRemainder ? 1 : 0),
-                numStudies: studiesInSmallerBatch + (batchId < studyRemainder ? 1 : 0),
-                numSeries: seriesInSmallerBatch + (batchId < seriesRemainder ? 1 : 0),
-                patientOffset: patientsInSmallerBatches * batchId + Math.min(batchId, patientRemainder),
-                studyOffset: studiesInSmallerBatch * batchId + Math.min(batchId, studyRemainder)
-            )
-        }
-    }
-
     void readSpecificationParameters(String configName) {
         setSpecificationParameters(
             new YamlObjectMapper().readValue(new File(configName), SpecificationParameters)
@@ -146,61 +113,16 @@ class PopulationGenerator {
     }
 
     private void generateBatchWithHandler(GenerationCache generationCache, IdOffsets idOffsets, BatchRequest batchRequest, Consumer<Patient> handler, boolean generateReports, boolean temporalHeartbeat = false) {
-        specificationParameters.postprocess()
-        generationCache.cache()
-
-        final GenerationContext generationContext = new GenerationContext(
-            specificationParameters: specificationParameters,
-            patientIdGenerators: idOffsets.getPatientIdGeneratorsFromOffset(batchRequest.patientOffset),
-            studyCountOverride: 0
+        batchRequest.generateBatchWithHandler(
+            new BatchRequestContext()
+                .specificationParameters(specificationParameters)
+                .generationCache(generationCache)
+                .idOffsets(idOffsets)
+                .patientRandomizers(patientRandomizers)
+                .handler(handler)
+                .generateReports(generateReports)
+                .temporalHeartbeat(temporalHeartbeat)
         )
-        int generatedStudies = 0
-        int generatedSeries = 0
-        final SequentialIdGenerator studyIdGenerator = idOffsets.getStudyIdEncoderFromOffset(batchRequest.studyOffset)
-        final List<Patient> patients = []
-
-        batchRequest.numPatients.times { generatedPatients ->
-            final Closure<Integer> remainingStudies = { batchRequest.numStudies - generatedStudies }
-
-            generationContext.setCurrentAverageStudiesPerPatient(generatedPatients == 0 ? 0.0 : generatedStudies / generatedPatients)
-            generationContext.setPreviouslyGeneratedSeries(generatedSeries)
-            generationContext.setPreviouslyGeneratedStudies(generatedStudies)
-            if (generatedPatients > batchRequest.numPatients - 4) {
-                generationContext.setStudyCountMaximum(remainingStudies() - (batchRequest.numPatients - generatedPatients - 1))
-            } // 3rd to last must leave 2 reports for final 2 patients, 2nd to last must leave a report for final patient
-            if (generatedPatients == batchRequest.numPatients - 1) { // for last patient in batch, we need to ensure exact number of studies
-                generationContext.setStudyCountOverride(remainingStudies())
-            }
-            final Patient patient = patientRandomizers.sample().createPatient(specificationParameters)
-            patient.randomize(generationContext, studyIdGenerator)
-            patient.studies.each { study ->
-                generatedStudies++
-                generatedSeries += study.series.size()
-                if (!specificationParameters.generateRadiologyReports) {
-                    study.cachePatientIdsForStudy()
-                }
-            }
-            if (temporalHeartbeat) {
-                Activity.executionContext.heartbeat("Batch ${batchRequest.id}, patient ${generatedPatients + 1} post-DICOM")
-            }
-            patients << patient
-        }
-
-        logger.info("DICOM specs for batch ${batchRequest.id} with ${patients.size()} patients and ${patients*.studies*.size().sum()} studies")
-
-        if (specificationParameters.generateRadiologyReports && generateReports) {
-            specificationParameters.reportGeneratorImplementation.generateReportsForPatients(patients, temporalHeartbeat)
-        }
-
-        patients.eachWithIndex { patient, generatedPatients ->
-            handler.accept(patient)
-            if ((generatedPatients + 1) % 100 == 0) {
-                logger.info("Generated DICOM and/or HL7 for patient ${patient.patientInstanceUid} in batch ${batchRequest.id} [${generatedPatients + 1}/${batchRequest.numPatients}]")
-            }
-            if (temporalHeartbeat) {
-                Activity.executionContext.heartbeat("Batch ${batchRequest.id}, patient ${generatedPatients + 1}")
-            }
-        }
     }
 
     private EnumeratedDistribution<PatientRandomizer> initPatientRandomizers() {

src/main/groovy/edu/washu/tag/generator/SpecificationParameters.groovy

@@ -5,6 +5,7 @@ import edu.washu.tag.generator.hl7.v2.ReportGenerator
 import edu.washu.tag.generator.metadata.Institution
 import edu.washu.tag.generator.metadata.Patient
 import edu.washu.tag.generator.metadata.Protocol
+import edu.washu.tag.generator.metadata.cohorting.SpecializedCohort
 import edu.washu.tag.generator.metadata.patient.DefaultPatientRandomizer
 import edu.washu.tag.generator.metadata.patient.GreekPatientRandomizer
 import edu.washu.tag.generator.metadata.patient.JapanesePatientRandomizer
@@ -44,8 +45,12 @@ class SpecificationParameters {
         (new KoreanPatientRandomizer(randomizerWeight: 5))
     ]
     List<Institution> institutionOverrides = []
+    List<SpecializedCohort> cohorts = []
 
     void postprocess() {
+        if (numPatients == 0) {
+            return
+        }
         averageStudiesPerPatient = numStudies / numPatients
         averageSeriesPerStudy = (numSeries as double) / numStudies