From efcb6e92a6e88086bc074b880bdbbc9a40c57e15 Mon Sep 17 00:00:00 2001 From: Charlie Moore Date: Fri, 19 Jun 2026 18:23:26 -0500 Subject: [PATCH 1/3] Unfinished WIP --- .../generator/SpecificationParameters.groovy | 2 + .../tag/generator/ai/OpenAiWrapper.groovy | 44 ++++++++++++------- .../ClassicReportWithKnownDiagnoses.groovy | 27 ++++++++++++ .../tag/generator/metadata/Patient.groovy | 1 + .../tag/generator/metadata/Protocol.groovy | 1 + .../metadata/cohorting/CodeStrategy.groovy | 26 +++++++++++ .../cohorting/SpecializedCohort.groovy | 10 +++++ .../metadata/cohorting/StudyRequest.groovy | 13 ++++++ .../metadata/protocols/SimpleCt.groovy | 44 +++++++++++++++++++ src/main/resources/procedure_codes.json | 9 ++++ 10 files changed, 162 insertions(+), 15 deletions(-) create mode 100644 src/main/groovy/edu/washu/tag/generator/ai/catalog/ClassicReportWithKnownDiagnoses.groovy create mode 100644 src/main/groovy/edu/washu/tag/generator/metadata/cohorting/CodeStrategy.groovy create mode 100644 src/main/groovy/edu/washu/tag/generator/metadata/cohorting/SpecializedCohort.groovy create mode 100644 src/main/groovy/edu/washu/tag/generator/metadata/cohorting/StudyRequest.groovy create mode 100644 src/main/groovy/edu/washu/tag/generator/metadata/protocols/SimpleCt.groovy diff --git a/src/main/groovy/edu/washu/tag/generator/SpecificationParameters.groovy b/src/main/groovy/edu/washu/tag/generator/SpecificationParameters.groovy index 1ba16bf..d1022f8 100644 --- a/src/main/groovy/edu/washu/tag/generator/SpecificationParameters.groovy +++ b/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,6 +45,7 @@ class SpecificationParameters { (new KoreanPatientRandomizer(randomizerWeight: 5)) ] List institutionOverrides = [] + List cohorts = [] void postprocess() { averageStudiesPerPatient = numStudies / numPatients diff --git a/src/main/groovy/edu/washu/tag/generator/ai/OpenAiWrapper.groovy b/src/main/groovy/edu/washu/tag/generator/ai/OpenAiWrapper.groovy index abc4373..b667086 100644 --- a/src/main/groovy/edu/washu/tag/generator/ai/OpenAiWrapper.groovy +++ b/src/main/groovy/edu/washu/tag/generator/ai/OpenAiWrapper.groovy @@ -39,19 +39,18 @@ class OpenAiWrapper { .build() } - PatientOutput generateReportsForPatient(Patient patient, Map> reports) { + PatientOutput generateReportsForPatientFromPlaceholders(Patient patient, Map reports) { final List uidMapping = [] final PatientRep patientRep = convertToPatientRep(patient, uidMapping) final PatientOutput patientOutput = new PatientOutput(patientId: patient.epicMrn) patientRep.studies.each { studyRep -> final String studyInstanceUid = uidMapping[Integer.parseInt(studyRep.uid)] - final Map.Entry> entry = reports.find { entry -> + final Map.Entry entry = reports.find { entry -> entry.key.studyInstanceUid == studyInstanceUid } final Study study = entry.key - final Class reportClass = entry.value - final GeneratedReport promptReport = reportClass.getDeclaredConstructor().newInstance() + final GeneratedReport promptReport = entry.value final String comparison = { if (promptReport instanceof WithComparison && studyRep.compareTo == null) { @@ -81,7 +80,7 @@ class OpenAiWrapper { BASE_RAD_CONTEXT_SINGULAR, mainPrompt, model, - reportClass + promptReport.class as Class ).withValidation({ GeneratedReport report -> promptReport.preserveState(report) report.validateReport() @@ -96,6 +95,15 @@ class OpenAiWrapper { patientOutput } + PatientOutput generateReportsForPatient(Patient patient, Map> reports) { + generateReportsForPatientFromPlaceholders( + patient, + reports.collectEntries { study, reportClass -> + [(study): reportClass.getDeclaredConstructor().newInstance()] + } + ) + } + List generateReportsForPatients(List patients) { final List uidMapping = [] @@ -136,17 +144,23 @@ class OpenAiWrapper { final Map comparisons = [:] final Map> studiesByStandardizedDescription = [:] - sortedStudies.each { study -> - uidMapping << study.studyInstanceUid - final String description = study.simpleDescription - final List previousStudies = studiesByStandardizedDescription.computeIfAbsent( - description, - { [] } - ) - if (!previousStudies.isEmpty()) { - comparisons.put(study.studyInstanceUid, previousStudies.last()) + if (patient.compareAdjacentStudies) { + sortedStudies.collate(2, 1, false).each { pair -> + comparisons.put(pair[1].studyInstanceUid, pair[0].studyInstanceUid) + } + } else { + sortedStudies.each { study -> + uidMapping << study.studyInstanceUid + final String description = study.simpleDescription + final List previousStudies = studiesByStandardizedDescription.computeIfAbsent( + description, + { [] } + ) + if (!previousStudies.isEmpty()) { + comparisons.put(study.studyInstanceUid, previousStudies.last()) + } + previousStudies << study.studyInstanceUid } - previousStudies << study.studyInstanceUid } patientRep.setStudies(sortedStudies.collect { sortedStudy -> diff --git a/src/main/groovy/edu/washu/tag/generator/ai/catalog/ClassicReportWithKnownDiagnoses.groovy b/src/main/groovy/edu/washu/tag/generator/ai/catalog/ClassicReportWithKnownDiagnoses.groovy new file mode 100644 index 0000000..5e002da --- /dev/null +++ b/src/main/groovy/edu/washu/tag/generator/ai/catalog/ClassicReportWithKnownDiagnoses.groovy @@ -0,0 +1,27 @@ +package edu.washu.tag.generator.ai.catalog + +import ca.uhn.hl7v2.model.v281.message.ORU_R01 +import edu.washu.tag.generator.ai.GeneratedReport +import edu.washu.tag.generator.ai.StudyRep +import edu.washu.tag.generator.ai.catalog.attribute.* +import edu.washu.tag.generator.ai.catalog.builder.HistoricalReportTextBuilder +import edu.washu.tag.generator.ai.catalog.builder.ModernReportTextBuilder +import edu.washu.tag.generator.hl7.v2.ReportVersion +import edu.washu.tag.generator.metadata.RadiologyReport +import edu.washu.tag.generator.metadata.Study + +class ClassicReportWithKnownDiagnoses extends ClassicReport { + + @Override + String getUserMessage(Study study, StudyRep studyRep) { + 'Be brief in the examination section. The findings and impressions should be significantly detailed. ' + diagnosisPrompt() + } + + @Override + Boolean validateReport() { + pruneCodes() + // TODO: ? overwrite diag prompt? + } + + +} diff --git a/src/main/groovy/edu/washu/tag/generator/metadata/Patient.groovy b/src/main/groovy/edu/washu/tag/generator/metadata/Patient.groovy index d602fba..ac95f90 100644 --- a/src/main/groovy/edu/washu/tag/generator/metadata/Patient.groovy +++ b/src/main/groovy/edu/washu/tag/generator/metadata/Patient.groovy @@ -37,6 +37,7 @@ class Patient implements DicomEncoder { @JsonIgnore double personalHeightMod @JsonIgnore double personalWeightMod @JsonIgnore boolean reportsMatched = false + @JsonIgnore boolean compareAdjacentStudies = false Patient randomize(GenerationContext generationContext, SequentialIdGenerator studyIdGenerator) { final LocalDate sixteenthBirthday = dateOfBirth.plusYears(16) diff --git a/src/main/groovy/edu/washu/tag/generator/metadata/Protocol.groovy b/src/main/groovy/edu/washu/tag/generator/metadata/Protocol.groovy index f1035a2..f573d2a 100644 --- a/src/main/groovy/edu/washu/tag/generator/metadata/Protocol.groovy +++ b/src/main/groovy/edu/washu/tag/generator/metadata/Protocol.groovy @@ -26,6 +26,7 @@ import java.util.concurrent.ThreadLocalRandom @JsonSubTypes.Type(value = MammogramSingleView, name = 'mg'), @JsonSubTypes.Type(value = MammogramFourView, name = 'mg4view'), @JsonSubTypes.Type(value = PetCt, name = 'petct'), + @JsonSubTypes.Type(value = SimpleCt, name = 'simplect'), @JsonSubTypes.Type(value = AnteroposteriorCr, name = 'ap_cr'), @JsonSubTypes.Type(value = SingleViewXRay, name = 'dx'), @JsonSubTypes.Type(value = MyocardialPerfusion, name = 'nm_mpi'), diff --git a/src/main/groovy/edu/washu/tag/generator/metadata/cohorting/CodeStrategy.groovy b/src/main/groovy/edu/washu/tag/generator/metadata/cohorting/CodeStrategy.groovy new file mode 100644 index 0000000..66fbce7 --- /dev/null +++ b/src/main/groovy/edu/washu/tag/generator/metadata/cohorting/CodeStrategy.groovy @@ -0,0 +1,26 @@ +package edu.washu.tag.generator.metadata.cohorting + +import edu.washu.tag.generator.util.RandomGenUtils +import org.apache.commons.lang3.RandomUtils + +import java.util.function.Function + +enum CodeStrategy { + + ONLY_ONE ({ List codes -> + RandomGenUtils.randomListEntry(codes) } + ), + STRICT_SUBSET ({ List codes -> + RandomGenUtils.randomSubset(codes, RandomUtils.insecure().randomInt(1, codes.size())) + }), + SUBSET ({ List codes -> + RandomGenUtils.randomSubset(codes, RandomUtils.insecure().randomInt(1, codes.size() + 1)) + }) + + private final Function, List> codeRandomizer + + CodeStrategy(Function, List> randomizer) { + codeRandomizer = randomizer + } + +} \ No newline at end of file diff --git a/src/main/groovy/edu/washu/tag/generator/metadata/cohorting/SpecializedCohort.groovy b/src/main/groovy/edu/washu/tag/generator/metadata/cohorting/SpecializedCohort.groovy new file mode 100644 index 0000000..6e4f9c5 --- /dev/null +++ b/src/main/groovy/edu/washu/tag/generator/metadata/cohorting/SpecializedCohort.groovy @@ -0,0 +1,10 @@ +package edu.washu.tag.generator.metadata.cohorting + +class SpecializedCohort { + + String name + int patientCount + List ageRange // hmm? + List trajectory + +} diff --git a/src/main/groovy/edu/washu/tag/generator/metadata/cohorting/StudyRequest.groovy b/src/main/groovy/edu/washu/tag/generator/metadata/cohorting/StudyRequest.groovy new file mode 100644 index 0000000..65850e4 --- /dev/null +++ b/src/main/groovy/edu/washu/tag/generator/metadata/cohorting/StudyRequest.groovy @@ -0,0 +1,13 @@ +package edu.washu.tag.generator.metadata.cohorting + +import edu.washu.tag.generator.metadata.Protocol + +class StudyRequest { + + String name + Protocol protocol + List diagnosisCodes + String additionalGenerationContext = '' + CodeStrategy codeStrategy = CodeStrategy.ONLY_ONE + +} diff --git a/src/main/groovy/edu/washu/tag/generator/metadata/protocols/SimpleCt.groovy b/src/main/groovy/edu/washu/tag/generator/metadata/protocols/SimpleCt.groovy new file mode 100644 index 0000000..b6fac70 --- /dev/null +++ b/src/main/groovy/edu/washu/tag/generator/metadata/protocols/SimpleCt.groovy @@ -0,0 +1,44 @@ +package edu.washu.tag.generator.metadata.protocols + +import edu.washu.tag.generator.metadata.* +import edu.washu.tag.generator.metadata.enums.BodyPart +import edu.washu.tag.generator.metadata.seriesTypes.ct.CtWithAttenuationCorrection +import edu.washu.tag.generator.metadata.seriesTypes.ct.PatientProtocol +import edu.washu.tag.generator.metadata.seriesTypes.ct.Topogram + +class SimpleCt extends PetCt { + + @Override + List getAllSeriesTypes() { + [ + new Topogram(), + new CtWithAttenuationCorrection(), + new PatientProtocol() + ] + } + + @Override + List getApplicableBodyParts() { + [ + BodyPart.BRAIN, + BodyPart.HEAD + ] + } + + @Override + ProcedureCode getProcedureCode(BodyPart bodyPart) { + ProcedureCode.lookup('ct head') + } + + @Override + GeneralizedProcedure getGeneralizedProcedure() { + GeneralizedProcedure.CT + } + + + @Override + String getSimpleDescription(BodyPart bodyPart) { + "${bodyPart.dicomRepresentation.toLowerCase()} CT" + } + +} diff --git a/src/main/resources/procedure_codes.json b/src/main/resources/procedure_codes.json index 054fb57..c97c83d 100644 --- a/src/main/resources/procedure_codes.json +++ b/src/main/resources/procedure_codes.json @@ -331,5 +331,14 @@ "codeMeaning": "IR INTRO CATHETER AORTA BILATERAL", "alternateText": "XRay Angiography" } + }, + "ct head": { + "impliedModality": "CT", + "codedTriplet": { + "codeValue": "CE704500", + "codingSchemeDesignator": "UNKDEV", + "codeMeaning": "CT HEAD WO", + "alternateText": "CT head w/o contrast" + } } } From bc8fa8bd71ff9150f0a8d698ba06b8a3a2d762a9 Mon Sep 17 00:00:00 2001 From: Charlie Moore Date: Tue, 23 Jun 2026 14:35:11 -0500 Subject: [PATCH 2/3] Additional WIP work including temporal refactor on supporting cohorts --- .../washu/tag/generator/BatchRequest.groovy | 91 +++++++++++++++++ .../tag/generator/BatchRequestContext.groovy | 23 +++++ .../edu/washu/tag/generator/Batcher.groovy | 85 ++++++++++++++++ .../tag/generator/PopulationGenerator.groovy | 98 ++----------------- .../tag/generator/ai/OpenAiWrapper.groovy | 2 +- .../tag/generator/metadata/Patient.groovy | 55 +++++++++-- .../washu/tag/generator/metadata/Study.groovy | 2 +- .../cohorting/SpecializedCohort.groovy | 4 +- .../metadata/cohorting/StudyRequest.groovy | 2 + .../module/study/GeneralStudyModule.groovy | 4 +- .../study/StudyDateDistribution.groovy | 7 +- .../study/UniformStudyDateDistribution.groovy | 5 +- .../EarlySetupHandlerActivityImpl.groovy | 3 +- .../GenerateDatasetWorkflowImpl.groovy | 27 ++--- .../generator/util/StringReplacements.groovy | 2 +- src/main/resources/rad_prompt.txt | 2 +- .../washu/tag/generator/BatchingTest.groovy | 78 ++++++++++++++- 17 files changed, 362 insertions(+), 128 deletions(-) create mode 100644 src/main/groovy/edu/washu/tag/generator/BatchRequestContext.groovy create mode 100644 src/main/groovy/edu/washu/tag/generator/Batcher.groovy diff --git a/src/main/groovy/edu/washu/tag/generator/BatchRequest.groovy b/src/main/groovy/edu/washu/tag/generator/BatchRequest.groovy index 0222105..20c50a6 100644 --- a/src/main/groovy/edu/washu/tag/generator/BatchRequest.groovy +++ b/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 patients = [] + + if (cohort == null) { + numPatients.times { generatedPatients -> + final Closure 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}") + } + } + } } diff --git a/src/main/groovy/edu/washu/tag/generator/BatchRequestContext.groovy b/src/main/groovy/edu/washu/tag/generator/BatchRequestContext.groovy new file mode 100644 index 0000000..2f576b9 --- /dev/null +++ b/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 patientRandomizers + Consumer handler + boolean generateReports = true + boolean temporalHeartbeat = false + +} diff --git a/src/main/groovy/edu/washu/tag/generator/Batcher.groovy b/src/main/groovy/edu/washu/tag/generator/Batcher.groovy new file mode 100644 index 0000000..81fb2e7 --- /dev/null +++ b/src/main/groovy/edu/washu/tag/generator/Batcher.groovy @@ -0,0 +1,85 @@ +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 resolveBatches() { + patientOffset = 0 + studyOffset = 0 + idOffset = 0 + + final List 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 -> + final int cohortPatients = cohort.numPatients + final int cohortStudies = cohortPatients * cohort.trajectory.size() + final List 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 computeBatchesFrom(int numPatients, int numStudies, int numSeries, Function summarizer) { + 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) + ) + } + } + +} diff --git a/src/main/groovy/edu/washu/tag/generator/PopulationGenerator.groovy b/src/main/groovy/edu/washu/tag/generator/PopulationGenerator.groovy index 9787ea5..c8f36f8 100644 --- a/src/main/groovy/edu/washu/tag/generator/PopulationGenerator.groovy +++ b/src/main/groovy/edu/washu/tag/generator/PopulationGenerator.groovy @@ -40,7 +40,7 @@ class PopulationGenerator { final IdOffsets idOffsets = new IdOffsets() CodeCache.initializeCache() - final List batchRequests = generator.resolveBatches() + final List 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 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 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 patients = [] - - batchRequest.numPatients.times { generatedPatients -> - final Closure 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 initPatientRandomizers() { diff --git a/src/main/groovy/edu/washu/tag/generator/ai/OpenAiWrapper.groovy b/src/main/groovy/edu/washu/tag/generator/ai/OpenAiWrapper.groovy index b667086..0771eb8 100644 --- a/src/main/groovy/edu/washu/tag/generator/ai/OpenAiWrapper.groovy +++ b/src/main/groovy/edu/washu/tag/generator/ai/OpenAiWrapper.groovy @@ -24,7 +24,7 @@ class OpenAiWrapper { private static final String BASE_RAD_CONTEXT = FileIOUtils.readResource('rad_context.txt') private static final String BASE_RAD_CONTEXT_SINGULAR = FileIOUtils.readResource('rad_context_singular.txt') private static final String BASE_RAD_PROMPT = FileIOUtils.readResource('rad_prompt.txt') - .replace(StringReplacements.DIAG_PROMPT, new ClassicReport().diagnosisPrompt()) + .replace(StringReplacements.PROMPT_ADDITIONS, new ClassicReport().diagnosisPrompt()) private static final String BASE_RAD_PROMPT_SINGULAR = FileIOUtils.readResource('rad_prompt_singular.txt') private static final ObjectMapper objectMapper = new ObjectMapper() diff --git a/src/main/groovy/edu/washu/tag/generator/metadata/Patient.groovy b/src/main/groovy/edu/washu/tag/generator/metadata/Patient.groovy index ac95f90..78323d6 100644 --- a/src/main/groovy/edu/washu/tag/generator/metadata/Patient.groovy +++ b/src/main/groovy/edu/washu/tag/generator/metadata/Patient.groovy @@ -2,13 +2,18 @@ package edu.washu.tag.generator.metadata import com.fasterxml.jackson.annotation.JsonIgnore import edu.washu.tag.generator.GenerationContext +import edu.washu.tag.generator.metadata.cohorting.SpecializedCohort +import edu.washu.tag.generator.metadata.cohorting.StudyRequest import edu.washu.tag.generator.metadata.enums.Nationality import edu.washu.tag.generator.metadata.enums.Race import edu.washu.tag.generator.metadata.enums.Sex import edu.washu.tag.generator.metadata.patient.EpicMrnGenerator import edu.washu.tag.generator.metadata.patient.MpiGenerator import edu.washu.tag.generator.metadata.patient.PatientIdGenerator +import edu.washu.tag.generator.metadata.study.StudyDateDistribution +import edu.washu.tag.generator.metadata.study.UniformStudyDateDistribution import edu.washu.tag.generator.util.SequentialIdGenerator +import org.apache.commons.lang3.RandomUtils import org.dcm4che3.data.Attributes import org.dcm4che3.data.Tag import org.dcm4che3.data.VR @@ -31,15 +36,15 @@ class Patient implements DicomEncoder { String mpi String epicMrn @JsonIgnore Map patientIds - @Deprecated @JsonIgnore Map idsForGenerators = [:] @JsonIgnore LocalDate earliestAvailableStudyDate @JsonIgnore Nationality nationality @JsonIgnore double personalHeightMod @JsonIgnore double personalWeightMod @JsonIgnore boolean reportsMatched = false @JsonIgnore boolean compareAdjacentStudies = false + @JsonIgnore SpecializedCohort parentCohort - Patient randomize(GenerationContext generationContext, SequentialIdGenerator studyIdGenerator) { + Patient randomize(GenerationContext generationContext, SequentialIdGenerator studyIdGenerator, SpecializedCohort parentCohort = null) { final LocalDate sixteenthBirthday = dateOfBirth.plusYears(16) patientInstanceUid = UIDUtils.createUID() patientIds = generationContext.patientIdGenerators.collectEntries { idGen -> @@ -52,14 +57,46 @@ class Patient implements DicomEncoder { [(idGen): id] } earliestAvailableStudyDate = sixteenthBirthday.isAfter(imagingDataEpoch) ? sixteenthBirthday : imagingDataEpoch - generationContext.calculateStudyCountForCurrentPatient().times { - final Protocol selectedProtocol = generationContext.chooseProtocol(this) - final Study study = new Study(protocol : selectedProtocol, patient : this) - .randomize(generationContext.specificationParameters, studyIdGenerator, selectedProtocol) - studies << study - generationContext.previouslyGeneratedStudies++ - generationContext.previouslyGeneratedSeries += study.series.size() + if (parentCohort == null) { + generationContext.calculateStudyCountForCurrentPatient().times { + final Study study = new Study(protocol : generationContext.chooseProtocol(this), patient : this) + .randomize(generationContext.specificationParameters, studyIdGenerator) + studies << study + generationContext.previouslyGeneratedStudies++ + generationContext.previouslyGeneratedSeries += study.series.size() + } + } else { + this.parentCohort = parentCohort + compareAdjacentStudies = true + int trajectoryDuration = 0 + final Map offsetMapForCohort = parentCohort.trajectory.collectEntries { studyReq -> + if (studyReq != parentCohort.trajectory[0]) { + final int studyOffset = RandomUtils.insecure().randomInt(studyReq.studyOffsetMin, studyReq.studyOffsetMax + 1) + trajectoryDuration += studyOffset + [(studyReq): trajectoryDuration] + } else { + [(studyReq): 0] + } + } + final StudyDateDistribution defaultDistribution = generationContext.specificationParameters.studyDateDistribution + final LocalDate latestPossibleStartDate = defaultDistribution.maxDate.minusDays(trajectoryDuration) + final LocalDate startDate = new UniformStudyDateDistribution( + minDate: defaultDistribution.minDate, + maxDate: latestPossibleStartDate + ).generateStudyDate(this) + + parentCohort.trajectory.each { studyRequest -> + final Study study = new Study( + protocol: studyRequest.protocol, + patient : this, + studyDate: startDate.plusDays(offsetMapForCohort.get(studyRequest)) + ).randomize(generationContext.specificationParameters, studyIdGenerator) + studies << study + generationContext.previouslyGeneratedStudies++ + generationContext.previouslyGeneratedSeries += study.series.size() + } } + this } diff --git a/src/main/groovy/edu/washu/tag/generator/metadata/Study.groovy b/src/main/groovy/edu/washu/tag/generator/metadata/Study.groovy index 1a08bcc..cee7e47 100644 --- a/src/main/groovy/edu/washu/tag/generator/metadata/Study.groovy +++ b/src/main/groovy/edu/washu/tag/generator/metadata/Study.groovy @@ -57,7 +57,7 @@ class Study implements DicomEncoder, PrivateElementContainer { @JsonIgnore List performingPhysiciansName // Series level field, but in most cases it's going to be fixed across a study private static final Logger logger = LoggerFactory.getLogger(Study) - Study randomize(SpecificationParameters specificationParameters, SequentialIdGenerator studyIdGenerator, Protocol protocol) { + Study randomize(SpecificationParameters specificationParameters, SequentialIdGenerator studyIdGenerator) { studyId = studyIdGenerator.get() accessionNumber = studyId protocol.setFieldsFor(specificationParameters, patient, this) diff --git a/src/main/groovy/edu/washu/tag/generator/metadata/cohorting/SpecializedCohort.groovy b/src/main/groovy/edu/washu/tag/generator/metadata/cohorting/SpecializedCohort.groovy index 6e4f9c5..42fc054 100644 --- a/src/main/groovy/edu/washu/tag/generator/metadata/cohorting/SpecializedCohort.groovy +++ b/src/main/groovy/edu/washu/tag/generator/metadata/cohorting/SpecializedCohort.groovy @@ -3,8 +3,8 @@ package edu.washu.tag.generator.metadata.cohorting class SpecializedCohort { String name - int patientCount - List ageRange // hmm? + int numPatients + List ageRange // todo hmm? List trajectory } diff --git a/src/main/groovy/edu/washu/tag/generator/metadata/cohorting/StudyRequest.groovy b/src/main/groovy/edu/washu/tag/generator/metadata/cohorting/StudyRequest.groovy index 65850e4..9c3dd60 100644 --- a/src/main/groovy/edu/washu/tag/generator/metadata/cohorting/StudyRequest.groovy +++ b/src/main/groovy/edu/washu/tag/generator/metadata/cohorting/StudyRequest.groovy @@ -9,5 +9,7 @@ class StudyRequest { List diagnosisCodes String additionalGenerationContext = '' CodeStrategy codeStrategy = CodeStrategy.ONLY_ONE + int studyOffsetMin + int studyOffsetMax } diff --git a/src/main/groovy/edu/washu/tag/generator/metadata/module/study/GeneralStudyModule.groovy b/src/main/groovy/edu/washu/tag/generator/metadata/module/study/GeneralStudyModule.groovy index 039bf90..51b1157 100644 --- a/src/main/groovy/edu/washu/tag/generator/metadata/module/study/GeneralStudyModule.groovy +++ b/src/main/groovy/edu/washu/tag/generator/metadata/module/study/GeneralStudyModule.groovy @@ -25,7 +25,9 @@ class GeneralStudyModule implements StudyLevelModule { final Protocol protocol = study.protocol final Equipment scanner = study.primaryEquipment // if more than one device, they should at least handle all study-level fields the same study.setStudyInstanceUid(UIDUtils.createUID()) - study.setStudyDate(specificationParameters.studyDateDistribution.generateStudyDate(patient)) + if (study.studyDate == null) { + study.setStudyDate(specificationParameters.studyDateDistribution.generateStudyDate(patient)) + } study.setStudyTime(RandomGenUtils.randomStudyTime()) if (protocol.includeMedicalStaff()) { diff --git a/src/main/groovy/edu/washu/tag/generator/metadata/study/StudyDateDistribution.groovy b/src/main/groovy/edu/washu/tag/generator/metadata/study/StudyDateDistribution.groovy index e0d1f60..999e137 100644 --- a/src/main/groovy/edu/washu/tag/generator/metadata/study/StudyDateDistribution.groovy +++ b/src/main/groovy/edu/washu/tag/generator/metadata/study/StudyDateDistribution.groovy @@ -14,8 +14,11 @@ import java.time.LocalDate @JsonSubTypes([ @JsonSubTypes.Type(value = UniformStudyDateDistribution, name = 'uniform') ]) -interface StudyDateDistribution { +abstract class StudyDateDistribution { - LocalDate generateStudyDate(Patient patient) + LocalDate minDate = Patient.imagingDataEpoch + LocalDate maxDate = LocalDate.now().minusMonths(6) + + abstract LocalDate generateStudyDate(Patient patient) } \ No newline at end of file diff --git a/src/main/groovy/edu/washu/tag/generator/metadata/study/UniformStudyDateDistribution.groovy b/src/main/groovy/edu/washu/tag/generator/metadata/study/UniformStudyDateDistribution.groovy index 9b764cf..ee14625 100644 --- a/src/main/groovy/edu/washu/tag/generator/metadata/study/UniformStudyDateDistribution.groovy +++ b/src/main/groovy/edu/washu/tag/generator/metadata/study/UniformStudyDateDistribution.groovy @@ -5,10 +5,7 @@ import edu.washu.tag.generator.util.RandomGenUtils import java.time.LocalDate -class UniformStudyDateDistribution implements StudyDateDistribution { - - LocalDate minDate = Patient.imagingDataEpoch - LocalDate maxDate = LocalDate.now().minusMonths(6) +class UniformStudyDateDistribution extends StudyDateDistribution { @Override LocalDate generateStudyDate(Patient patient) { diff --git a/src/main/groovy/edu/washu/tag/generator/temporal/activity/EarlySetupHandlerActivityImpl.groovy b/src/main/groovy/edu/washu/tag/generator/temporal/activity/EarlySetupHandlerActivityImpl.groovy index b099d2e..3e082d9 100644 --- a/src/main/groovy/edu/washu/tag/generator/temporal/activity/EarlySetupHandlerActivityImpl.groovy +++ b/src/main/groovy/edu/washu/tag/generator/temporal/activity/EarlySetupHandlerActivityImpl.groovy @@ -3,6 +3,7 @@ package edu.washu.tag.generator.temporal.activity import com.fasterxml.jackson.databind.ObjectMapper import edu.washu.tag.generator.BatchRequest import edu.washu.tag.generator.BatchProcessor +import edu.washu.tag.generator.Batcher import edu.washu.tag.generator.PopulationGenerator import edu.washu.tag.generator.metadata.GenerationCache import edu.washu.tag.generator.temporal.TemporalApplication @@ -25,7 +26,7 @@ class EarlySetupHandlerActivityImpl implements EarlySetupHandlerActivity { final PopulationGenerator generator = new PopulationGenerator() generator.readSpecificationParameters(specificationParamsPath) - final List batchRequests = generator.resolveBatches(patientsPerFullBatch) + final List batchRequests = new Batcher(generator.specificationParameters, patientsPerFullBatch).resolveBatches() logger.info("Request has been resolved into ${batchRequests.size()} standalone batches") batchRequests diff --git a/src/main/groovy/edu/washu/tag/generator/temporal/workflow/GenerateDatasetWorkflowImpl.groovy b/src/main/groovy/edu/washu/tag/generator/temporal/workflow/GenerateDatasetWorkflowImpl.groovy index 6f1d2d8..a775b3b 100644 --- a/src/main/groovy/edu/washu/tag/generator/temporal/workflow/GenerateDatasetWorkflowImpl.groovy +++ b/src/main/groovy/edu/washu/tag/generator/temporal/workflow/GenerateDatasetWorkflowImpl.groovy @@ -38,19 +38,6 @@ class GenerateDatasetWorkflowImpl implements GenerateDatasetWorkflow { ).build() ) - private final BatchHandlerActivity batchHandlerActivity = - Workflow.newActivityStub( - BatchHandlerActivity, - ActivityOptions.newBuilder() - .setStartToCloseTimeout(Duration.ofHours(24)) - .setHeartbeatTimeout(Duration.ofMinutes(15)) - .setRetryOptions(RetryOptions.newBuilder() - .setMaximumInterval(Duration.ofSeconds(1)) - .setMaximumAttempts(3) - .build()) - .build() - ) - @Override void generateDataset(GenerateDatasetInput input) { final WorkflowInfo workflowInfo = Workflow.getInfo() @@ -78,6 +65,20 @@ class GenerateDatasetWorkflowImpl implements GenerateDatasetWorkflow { // ~1-2k batches per run; beyond that, window the fan-out with continueAsNew (the window would be a history // checkpoint only, not a concurrency knob). Promise.allOf(batches.collect { batchRequest -> + final BatchHandlerActivity batchHandlerActivity = + Workflow.newActivityStub( + BatchHandlerActivity, + ActivityOptions.newBuilder() + .setStartToCloseTimeout(Duration.ofHours(24)) + .setHeartbeatTimeout(Duration.ofMinutes(15)) + .setSummary(batchRequest.temporalSummary) + .setRetryOptions(RetryOptions.newBuilder() + .setMaximumInterval(Duration.ofSeconds(1)) + .setMaximumAttempts(3) + .build()) + .build() + ) + Async.function( batchHandlerActivity::formAndWriteBatch, new BatchHandlerActivityInput(input, nameCache, idOffsets, batchRequest)) diff --git a/src/main/groovy/edu/washu/tag/generator/util/StringReplacements.groovy b/src/main/groovy/edu/washu/tag/generator/util/StringReplacements.groovy index b4726b4..6d0b59f 100644 --- a/src/main/groovy/edu/washu/tag/generator/util/StringReplacements.groovy +++ b/src/main/groovy/edu/washu/tag/generator/util/StringReplacements.groovy @@ -6,6 +6,6 @@ class StringReplacements { public static final String LATERALITY = '%LATERALITY%' public static final String VIEW = '%VIEW%' public static final String FPS = '%FPS%' - public static final String DIAG_PROMPT = '%DIAGNOSIS_PROMPT%' + public static final String PROMPT_ADDITIONS = '%PROMPT_ADDITIONS%' } diff --git a/src/main/resources/rad_prompt.txt b/src/main/resources/rad_prompt.txt index 2e07bdc..ec33ed5 100644 --- a/src/main/resources/rad_prompt.txt +++ b/src/main/resources/rad_prompt.txt @@ -1,3 +1,3 @@ Please provide reports of the images, not being afraid to use technological jargon. Include a good deal of variance in your prose. Be brief in the examination section. The findings and impressions should be significantly detailed. When comparing to a previous report, include the date explicitly. -Do not skip any patients or studies in your output. %DIAGNOSIS_PROMPT% The patients with imaging studies you are assessing are what follows. \ No newline at end of file +Do not skip any patients or studies in your output. %PROMPT_ADDITIONS% The patients with imaging studies you are assessing are what follows. \ No newline at end of file diff --git a/src/test/groovy/edu/washu/tag/generator/BatchingTest.groovy b/src/test/groovy/edu/washu/tag/generator/BatchingTest.groovy index 20a4291..6806ce4 100644 --- a/src/test/groovy/edu/washu/tag/generator/BatchingTest.groovy +++ b/src/test/groovy/edu/washu/tag/generator/BatchingTest.groovy @@ -1,5 +1,7 @@ package edu.washu.tag.generator +import edu.washu.tag.generator.metadata.cohorting.SpecializedCohort +import edu.washu.tag.generator.metadata.cohorting.StudyRequest import org.testng.annotations.Test import static org.testng.Assert.assertEquals @@ -8,8 +10,8 @@ class BatchingTest { @Test void testBatchingHappyPath() { - final List batchRequests = new PopulationGenerator( - specificationParameters: new SpecificationParameters( + final List batchRequests = new Batcher( + new SpecificationParameters( numPatients: 4000, numStudies: 10000, numSeries: 30000 @@ -29,8 +31,8 @@ class BatchingTest { @Test void testBatchingMessy() { - final List batchRequests = new PopulationGenerator( - specificationParameters: new SpecificationParameters( + final List batchRequests = new Batcher( + new SpecificationParameters( numPatients: 4100, numStudies: 10000, numSeries: 30001 @@ -59,4 +61,72 @@ class BatchingTest { assertEquals(batch2.studyOffset, 6667) } + @Test + void testBatchingWithCohorts() { + final List batchRequests = new Batcher( + new SpecificationParameters( + numPatients: 1000, + numStudies: 10000, + numSeries: 30000, + cohorts: [ + new SpecializedCohort( + name: 'cohort1', + numPatients: 4000, + trajectory: [ + new StudyRequest( + name: 'MR' + ), + new StudyRequest( + name: 'CT' + ) + ] + ), + new SpecializedCohort( + name: 'cohort2', + numPatients: 2000, + trajectory: [ + new StudyRequest( + name: 'MR' + ), + new StudyRequest( + name: 'CT' + ), + new StudyRequest( + name: 'CT2' + ) + ] + ) + ] + ) + ).resolveBatches() + assertEquals(batchRequests.size(), 4) + final BatchRequest batch0 = batchRequests[0] + assertEquals(batch0.numPatients, 1000) + assertEquals(batch0.numStudies, 10000) + assertEquals(batch0.numSeries, 30000) + assertEquals(batch0.patientOffset, 0) + assertEquals(batch0.studyOffset, 0) + + final BatchRequest batch1 = batchRequests[1] + assertEquals(batch1.numPatients, 2000) + assertEquals(batch1.numStudies, 4000) + assertEquals(batch1.patientOffset, 1000) + assertEquals(batch1.studyOffset, 10000) + assertEquals(batch1.cohort.name, 'cohort1') + + final BatchRequest batch2 = batchRequests[2] + assertEquals(batch2.numPatients, 2000) + assertEquals(batch2.numStudies, 4000) + assertEquals(batch2.patientOffset, 3000) + assertEquals(batch2.studyOffset, 14000) + assertEquals(batch2.cohort.name, 'cohort1') + + final BatchRequest batch3 = batchRequests[3] + assertEquals(batch3.numPatients, 2000) + assertEquals(batch3.numStudies, 6000) + assertEquals(batch3.patientOffset, 5000) + assertEquals(batch3.studyOffset, 18000) + assertEquals(batch3.cohort.name, 'cohort2') + } + } From c9fed119a65ba5089e2bb83c7a939690a5c2accb Mon Sep 17 00:00:00 2001 From: Charlie Moore Date: Thu, 25 Jun 2026 11:10:49 -0500 Subject: [PATCH 3/3] First draft of fully realized cohort feature --- .../edu/washu/tag/generator/Batcher.groovy | 6 ++ .../generator/SpecificationParameters.groovy | 3 + .../tag/generator/ai/OpenAiWrapper.groovy | 5 +- .../ClassicReportWithKnownDiagnoses.groovy | 27 --------- .../tag/generator/ai/catalog/CodeCache.groovy | 38 ++++++++++-- .../ai/catalog/TechnicianNoteReport.groovy | 2 +- .../hl7/v2/CyclicVariedGenerator.groovy | 1 + .../hl7/v2/CyclicVariedGptGenerator.groovy | 48 ++++++++++++--- .../hl7/v2/segment/Dg1Generator.groovy | 12 ++-- .../tag/generator/metadata/Diagnosis.groovy | 11 ++++ .../tag/generator/metadata/Patient.groovy | 2 + .../generator/metadata/RadiologyReport.groovy | 1 + .../washu/tag/generator/metadata/Study.groovy | 21 ++++++- .../metadata/cohorting/CodeStrategy.groovy | 58 ++++++++++++++++--- .../cohorting/SpecializedCohort.groovy | 1 - .../metadata/cohorting/StudyRequest.groovy | 17 +++++- .../reports/RadiologyReport2_7.groovy | 7 +-- .../tag/generator/util/RandomGenUtils.groovy | 11 ++++ 18 files changed, 203 insertions(+), 68 deletions(-) delete mode 100644 src/main/groovy/edu/washu/tag/generator/ai/catalog/ClassicReportWithKnownDiagnoses.groovy create mode 100644 src/main/groovy/edu/washu/tag/generator/metadata/Diagnosis.groovy diff --git a/src/main/groovy/edu/washu/tag/generator/Batcher.groovy b/src/main/groovy/edu/washu/tag/generator/Batcher.groovy index 81fb2e7..f421b47 100644 --- a/src/main/groovy/edu/washu/tag/generator/Batcher.groovy +++ b/src/main/groovy/edu/washu/tag/generator/Batcher.groovy @@ -39,6 +39,9 @@ class Batcher { 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 batchesForCohort = computeBatchesFrom( @@ -60,6 +63,9 @@ class Batcher { } private List computeBatchesFrom(int numPatients, int numStudies, int numSeries, Function summarizer) { + if (numPatients == 0) { + return [] + } final int totalNumBatches = Math.ceilDiv(numPatients, patientsPerFullBatch) final int patientsInSmallerBatches = numPatients / totalNumBatches diff --git a/src/main/groovy/edu/washu/tag/generator/SpecificationParameters.groovy b/src/main/groovy/edu/washu/tag/generator/SpecificationParameters.groovy index d1022f8..739aec2 100644 --- a/src/main/groovy/edu/washu/tag/generator/SpecificationParameters.groovy +++ b/src/main/groovy/edu/washu/tag/generator/SpecificationParameters.groovy @@ -48,6 +48,9 @@ class SpecificationParameters { List cohorts = [] void postprocess() { + if (numPatients == 0) { + return + } averageStudiesPerPatient = numStudies / numPatients averageSeriesPerStudy = (numSeries as double) / numStudies diff --git a/src/main/groovy/edu/washu/tag/generator/ai/OpenAiWrapper.groovy b/src/main/groovy/edu/washu/tag/generator/ai/OpenAiWrapper.groovy index 0771eb8..cb6c671 100644 --- a/src/main/groovy/edu/washu/tag/generator/ai/OpenAiWrapper.groovy +++ b/src/main/groovy/edu/washu/tag/generator/ai/OpenAiWrapper.groovy @@ -73,6 +73,7 @@ class OpenAiWrapper { " The patient is ${patient.sex.name().toLowerCase()} and was born on ${TimeUtils.UNAMBIGUOUS_DATE.format(patient.dateOfBirth)}." + " The study you are currently evaluating is a ${studyRep.description}. ${comparison}" + promptReport.getUserMessage(study, studyRep) + + (study.additionalGenerationContext ?: '') + fullComparisonSerialization final GeneratedReport generatedReport = new GenericLlmCall<>( @@ -144,13 +145,15 @@ class OpenAiWrapper { final Map comparisons = [:] final Map> studiesByStandardizedDescription = [:] + sortedStudies.each { study -> + uidMapping << study.studyInstanceUid + } if (patient.compareAdjacentStudies) { sortedStudies.collate(2, 1, false).each { pair -> comparisons.put(pair[1].studyInstanceUid, pair[0].studyInstanceUid) } } else { sortedStudies.each { study -> - uidMapping << study.studyInstanceUid final String description = study.simpleDescription final List previousStudies = studiesByStandardizedDescription.computeIfAbsent( description, diff --git a/src/main/groovy/edu/washu/tag/generator/ai/catalog/ClassicReportWithKnownDiagnoses.groovy b/src/main/groovy/edu/washu/tag/generator/ai/catalog/ClassicReportWithKnownDiagnoses.groovy deleted file mode 100644 index 5e002da..0000000 --- a/src/main/groovy/edu/washu/tag/generator/ai/catalog/ClassicReportWithKnownDiagnoses.groovy +++ /dev/null @@ -1,27 +0,0 @@ -package edu.washu.tag.generator.ai.catalog - -import ca.uhn.hl7v2.model.v281.message.ORU_R01 -import edu.washu.tag.generator.ai.GeneratedReport -import edu.washu.tag.generator.ai.StudyRep -import edu.washu.tag.generator.ai.catalog.attribute.* -import edu.washu.tag.generator.ai.catalog.builder.HistoricalReportTextBuilder -import edu.washu.tag.generator.ai.catalog.builder.ModernReportTextBuilder -import edu.washu.tag.generator.hl7.v2.ReportVersion -import edu.washu.tag.generator.metadata.RadiologyReport -import edu.washu.tag.generator.metadata.Study - -class ClassicReportWithKnownDiagnoses extends ClassicReport { - - @Override - String getUserMessage(Study study, StudyRep studyRep) { - 'Be brief in the examination section. The findings and impressions should be significantly detailed. ' + diagnosisPrompt() - } - - @Override - Boolean validateReport() { - pruneCodes() - // TODO: ? overwrite diag prompt? - } - - -} diff --git a/src/main/groovy/edu/washu/tag/generator/ai/catalog/CodeCache.groovy b/src/main/groovy/edu/washu/tag/generator/ai/catalog/CodeCache.groovy index ca974c4..5b9de4a 100644 --- a/src/main/groovy/edu/washu/tag/generator/ai/catalog/CodeCache.groovy +++ b/src/main/groovy/edu/washu/tag/generator/ai/catalog/CodeCache.groovy @@ -3,6 +3,7 @@ package edu.washu.tag.generator.ai.catalog import com.fasterxml.jackson.databind.ObjectMapper import com.google.common.util.concurrent.RateLimiter import edu.washu.tag.generator.ai.catalog.attribute.DiagnosisCodeDesignator +import edu.washu.tag.generator.metadata.Diagnosis import org.slf4j.Logger import org.slf4j.LoggerFactory @@ -18,7 +19,7 @@ class CodeCache { private static final int MAX_RETRIES = 5 private static final ObjectMapper objectMapper = new ObjectMapper() private static final Logger logger = LoggerFactory.getLogger(CodeCache) - public static final Map> codes = [:] + public static final Map> codes = [:] public static final Map> knownBadCodes = [:] static void initializeCache() { @@ -36,17 +37,38 @@ class CodeCache { cacheCodesFromSearch(designator, code) == QueryStatus.MATCH } + static Diagnosis lookupDiagnosis(DiagnosisCodeDesignator designator, String code) { + codes[designator].find { + it.code = code + } + } + static String lookupCode(DiagnosisCodeDesignator designator, String code) { - codes[designator][code] + lookupDiagnosis(designator, code).codeMeaning + } + + static String lookupCode(Diagnosis code) { + code.codeMeaning ?: lookupCode(code.designator, code.code) + } + + static List resolvePartialCode(DiagnosisCodeDesignator designator, String code) { + final QueryStatus queryStatus = cacheCodesFromSearch(designator, code) + if (queryStatus == QueryStatus.MATCH) { + [lookupDiagnosis(designator, code)] + } else { + codes[designator].findAll { + it.code.startsWith(code) + } + } } private static QueryStatus cacheCodesFromSearch(DiagnosisCodeDesignator designator, String search) { - final Map codesForDesignator = codes.get(designator, [:]) + final List codesForDesignator = codes.get(designator, []) final List knownBad = knownBadCodes.get(designator, []) if (knownBad.contains(search)) { return QueryStatus.NO_MATCH } - if (codesForDesignator.containsKey(search)) { + if (codesForDesignator.any { it.code == search }) { return QueryStatus.MATCH } final String url = designator.getUrlForSearch(search) @@ -61,9 +83,13 @@ class CodeCache { } else { final List parsed = objectMapper.readValue(response.body(), List) parsed[3].each { List code -> - codesForDesignator.put(code[0], code[1]) + codesForDesignator << new Diagnosis( + designator: designator, + code: code[0], + codeMeaning: code[1] + ) } - if (codesForDesignator.containsKey(search)) { + if (codesForDesignator.any { it.code == search }) { return QueryStatus.MATCH } else { knownBad << search diff --git a/src/main/groovy/edu/washu/tag/generator/ai/catalog/TechnicianNoteReport.groovy b/src/main/groovy/edu/washu/tag/generator/ai/catalog/TechnicianNoteReport.groovy index 157113e..fe0f74c 100644 --- a/src/main/groovy/edu/washu/tag/generator/ai/catalog/TechnicianNoteReport.groovy +++ b/src/main/groovy/edu/washu/tag/generator/ai/catalog/TechnicianNoteReport.groovy @@ -40,7 +40,7 @@ class TechnicianNoteReport extends GeneratedReport impleme 'Be brief in the examination section. The findings and impressions should be significantly detailed. ' + 'The technicianNote section should not be empty and could be anything the technician needs to note in the report. It should be less formal in tone than the other sections.' + 'It should focus on the examination itself, so it should not mention things such as image quality. Examples of possible values could be things from this non-exhaustive list:\n' + - examples.collect { '* ' + it }.join('\n') + examples.collect { '* ' + it }.join('\n') + '\n' } @Override diff --git a/src/main/groovy/edu/washu/tag/generator/hl7/v2/CyclicVariedGenerator.groovy b/src/main/groovy/edu/washu/tag/generator/hl7/v2/CyclicVariedGenerator.groovy index a967d90..05abbc1 100644 --- a/src/main/groovy/edu/washu/tag/generator/hl7/v2/CyclicVariedGenerator.groovy +++ b/src/main/groovy/edu/washu/tag/generator/hl7/v2/CyclicVariedGenerator.groovy @@ -9,6 +9,7 @@ import edu.washu.tag.generator.hl7.v2.model.TransportationMode import edu.washu.tag.generator.metadata.Patient import edu.washu.tag.generator.metadata.RadiologyReport import edu.washu.tag.generator.metadata.Study +import edu.washu.tag.generator.metadata.cohorting.SpecializedCohort import io.temporal.workflow.Workflow import org.slf4j.Logger diff --git a/src/main/groovy/edu/washu/tag/generator/hl7/v2/CyclicVariedGptGenerator.groovy b/src/main/groovy/edu/washu/tag/generator/hl7/v2/CyclicVariedGptGenerator.groovy index 0af3dc6..42a9d78 100644 --- a/src/main/groovy/edu/washu/tag/generator/hl7/v2/CyclicVariedGptGenerator.groovy +++ b/src/main/groovy/edu/washu/tag/generator/hl7/v2/CyclicVariedGptGenerator.groovy @@ -3,7 +3,10 @@ package edu.washu.tag.generator.hl7.v2 import edu.washu.tag.generator.ai.GeneratedReport import edu.washu.tag.generator.ai.OpenAiWrapper import edu.washu.tag.generator.ai.PatientOutput +import edu.washu.tag.generator.ai.StudyRep import edu.washu.tag.generator.ai.catalog.ReportRegistry +import edu.washu.tag.generator.ai.catalog.attribute.DiagnosisCodeDesignator +import edu.washu.tag.generator.ai.catalog.attribute.WithDiagnosisCodes import edu.washu.tag.generator.metadata.Patient import edu.washu.tag.generator.metadata.RadiologyReport import edu.washu.tag.generator.metadata.Study @@ -32,10 +35,13 @@ class CyclicVariedGptGenerator extends CyclicVariedGenerator { : null openAiWrapper = new OpenAiWrapper(endpoint, apiKeyEnvVar, model, queryParams, llmHeartbeat) } - List bulkPatients + final boolean isCohort = patients.any { it.parentCohort != null } + List bulkPatients = [] List customPatients = [] if (customReportFraction == null) { bulkPatients = patients + } else if (isCohort) { // for custom cohort, each patient is pre-assigned different diagnosis codes, don't want to overload model context with per-study complex specifications + customPatients = patients } else { (bulkPatients, customPatients) = patients.split { ThreadLocalRandom.current().nextDouble() > customReportFraction @@ -69,14 +75,38 @@ class CyclicVariedGptGenerator extends CyclicVariedGenerator { } } } - customPatients.eachWithIndex { patient, index -> - patientOutputs << openAiWrapper.generateReportsForPatient( - patient, - patient.studies.collectEntries { study -> - [(study): ReportRegistry.randomReportClass(study)] - } - ) - heartbeatAndLog("Generated customized reports for patient [${index + 1}/${customPatients.size()}]") + if (isCohort) { + customPatients.eachWithIndex { patient, index -> + patientOutputs << openAiWrapper.generateReportsForPatientFromPlaceholders( + patient, + patient.studies.collectEntries { study -> + final GeneratedReport placeholderReport = ReportRegistry.randomReportClass(study).getDeclaredConstructor().newInstance() + final String dxString = study.resolveDiagnoses().collect { it.code }.join(',') + if (placeholderReport instanceof WithDiagnosisCodes) { + placeholderReport.metaClass.diagnosisPrompt = { + "The diagnoses should be set to ${dxString} and the meaning of those codes used in the rest of the report. " + } + placeholderReport.designator = DiagnosisCodeDesignator.ICD_10 + } else { + study.setAdditionalGenerationContext( + "The diagnoses for this report are ${dxString}. The meaning of those codes should be used in the rest of the report. ${study.additionalGenerationContext ?: ''} " + ) + } + [(study): placeholderReport] + } + ) + heartbeatAndLog("Generated customized reports for cohort patient [${index + 1}/${customPatients.size()}]") + } + } else { + customPatients.eachWithIndex { patient, index -> + patientOutputs << openAiWrapper.generateReportsForPatient( + patient, + patient.studies.collectEntries { study -> + [(study): ReportRegistry.randomReportClass(study)] + } + ) + heartbeatAndLog("Generated customized reports for patient [${index + 1}/${customPatients.size()}]") + } } patientOutputs } diff --git a/src/main/groovy/edu/washu/tag/generator/hl7/v2/segment/Dg1Generator.groovy b/src/main/groovy/edu/washu/tag/generator/hl7/v2/segment/Dg1Generator.groovy index 83dc6fe..05fb5be 100644 --- a/src/main/groovy/edu/washu/tag/generator/hl7/v2/segment/Dg1Generator.groovy +++ b/src/main/groovy/edu/washu/tag/generator/hl7/v2/segment/Dg1Generator.groovy @@ -6,6 +6,7 @@ import ca.uhn.hl7v2.model.v281.message.ORU_R01 import ca.uhn.hl7v2.model.v281.segment.DG1 import edu.washu.tag.generator.ai.catalog.CodeCache import edu.washu.tag.generator.ai.catalog.attribute.DiagnosisCodeDesignator +import edu.washu.tag.generator.metadata.Diagnosis import edu.washu.tag.generator.metadata.RadiologyReport class Dg1Generator extends NonstandardSegmentGenerator { @@ -22,13 +23,11 @@ class Dg1Generator extends NonstandardSegmentGenerator { */ int index - String code - DiagnosisCodeDesignator designator + Diagnosis code - Dg1Generator(int index, String code, DiagnosisCodeDesignator designator) { + Dg1Generator(int index, Diagnosis code) { this.index = index this.code = code - this.designator = designator } @Override @@ -53,12 +52,13 @@ class Dg1Generator extends NonstandardSegmentGenerator { @Override void generateSegment(RadiologyReport radReport, DG1 baseSegment) { - final String codeMeaning = CodeCache.lookupCode(designator, code) + final String codeMeaning = code.codeMeaning ?: CodeCache.lookupCode(code) + final DiagnosisCodeDesignator designator = code.designator baseSegment.getDg11_SetIDDG1().setValue(String.valueOf(index + 1)) baseSegment.getDg12_DiagnosisCodingMethod().setValue(designator.hl7Representation) final CWE codeCwe = baseSegment.getDg13_DiagnosisCodeDG1() - codeCwe.getCwe1_Identifier().setValue(code) + codeCwe.getCwe1_Identifier().setValue(code.code) codeCwe.getCwe2_Text().setValue(codeMeaning) codeCwe.getCwe3_NameOfCodingSystem().setValue(designator.hl7Representation) baseSegment.getDg14_DiagnosisDescription().setValue(codeMeaning) diff --git a/src/main/groovy/edu/washu/tag/generator/metadata/Diagnosis.groovy b/src/main/groovy/edu/washu/tag/generator/metadata/Diagnosis.groovy new file mode 100644 index 0000000..8c1fe92 --- /dev/null +++ b/src/main/groovy/edu/washu/tag/generator/metadata/Diagnosis.groovy @@ -0,0 +1,11 @@ +package edu.washu.tag.generator.metadata + +import edu.washu.tag.generator.ai.catalog.attribute.DiagnosisCodeDesignator + +class Diagnosis { + + String code + String codeMeaning + DiagnosisCodeDesignator designator + +} diff --git a/src/main/groovy/edu/washu/tag/generator/metadata/Patient.groovy b/src/main/groovy/edu/washu/tag/generator/metadata/Patient.groovy index 78323d6..e433016 100644 --- a/src/main/groovy/edu/washu/tag/generator/metadata/Patient.groovy +++ b/src/main/groovy/edu/washu/tag/generator/metadata/Patient.groovy @@ -91,6 +91,8 @@ class Patient implements DicomEncoder { patient : this, studyDate: startDate.plusDays(offsetMapForCohort.get(studyRequest)) ).randomize(generationContext.specificationParameters, studyIdGenerator) + study.setDiagnoses(studyRequest.randomizeCodes(study)) + study.setAdditionalGenerationContext(studyRequest.additionalGenerationContext) studies << study generationContext.previouslyGeneratedStudies++ generationContext.previouslyGeneratedSeries += study.series.size() diff --git a/src/main/groovy/edu/washu/tag/generator/metadata/RadiologyReport.groovy b/src/main/groovy/edu/washu/tag/generator/metadata/RadiologyReport.groovy index 815b6d7..0369396 100644 --- a/src/main/groovy/edu/washu/tag/generator/metadata/RadiologyReport.groovy +++ b/src/main/groovy/edu/washu/tag/generator/metadata/RadiologyReport.groovy @@ -16,6 +16,7 @@ import ca.uhn.hl7v2.validation.impl.ValidationContextImpl import com.fasterxml.jackson.annotation.JsonIgnore import com.fasterxml.jackson.annotation.JsonTypeInfo import edu.washu.tag.generator.ai.GeneratedReport +import edu.washu.tag.generator.ai.catalog.attribute.WithDiagnosisCodes import edu.washu.tag.generator.hl7.v2.ReportVersion import edu.washu.tag.generator.hl7.v2.model.DoctorEncoder import edu.washu.tag.generator.hl7.v2.model.ReportStatus diff --git a/src/main/groovy/edu/washu/tag/generator/metadata/Study.groovy b/src/main/groovy/edu/washu/tag/generator/metadata/Study.groovy index cee7e47..bc3901c 100644 --- a/src/main/groovy/edu/washu/tag/generator/metadata/Study.groovy +++ b/src/main/groovy/edu/washu/tag/generator/metadata/Study.groovy @@ -1,8 +1,8 @@ package edu.washu.tag.generator.metadata import com.fasterxml.jackson.annotation.JsonIgnore -import edu.washu.tag.generator.metadata.patient.EpicMrnGenerator -import edu.washu.tag.generator.metadata.patient.MpiGenerator +import edu.washu.tag.generator.ai.GeneratedReport +import edu.washu.tag.generator.ai.catalog.attribute.WithDiagnosisCodes import edu.washu.tag.generator.metadata.patient.PatientId import edu.washu.tag.generator.metadata.patient.PatientIdGenerator import edu.washu.tag.generator.util.SequentialIdGenerator @@ -45,6 +45,7 @@ class Study implements DicomEncoder, PrivateElementContainer { String procedureCodeId GeneralizedProcedure generalizedProcedure BodyPart bodyPartExamined // TODO: this is a simplifying assumption. This is really a series-level field + List diagnoses @JsonIgnore Patient patient @JsonIgnore String ethnicGroup // Yes, this is a patient field, but we want the value to be inconsistently encoded *across studies*, but consistent within a study @JsonIgnore Protocol protocol @@ -55,6 +56,7 @@ class Study implements DicomEncoder, PrivateElementContainer { @JsonIgnore Map> operatorMap @JsonIgnore List primaryOperators @JsonIgnore List performingPhysiciansName // Series level field, but in most cases it's going to be fixed across a study + @JsonIgnore String additionalGenerationContext private static final Logger logger = LoggerFactory.getLogger(Study) Study randomize(SpecificationParameters specificationParameters, SequentialIdGenerator studyIdGenerator) { @@ -133,4 +135,19 @@ class Study implements DicomEncoder, PrivateElementContainer { patientIds } + @JsonIgnore + List resolveDiagnoses() { + final GeneratedReport genReport = radReport?.generatedReport + if (diagnoses != null && !diagnoses.isEmpty()) { + diagnoses + } else if (genReport instanceof WithDiagnosisCodes) { + final String designator = genReport.designator + genReport.parsedCodes.collect { + new Diagnosis(code: it, designator: designator) + } + } else { + [] + } + } + } diff --git a/src/main/groovy/edu/washu/tag/generator/metadata/cohorting/CodeStrategy.groovy b/src/main/groovy/edu/washu/tag/generator/metadata/cohorting/CodeStrategy.groovy index 66fbce7..2102b30 100644 --- a/src/main/groovy/edu/washu/tag/generator/metadata/cohorting/CodeStrategy.groovy +++ b/src/main/groovy/edu/washu/tag/generator/metadata/cohorting/CodeStrategy.groovy @@ -1,5 +1,9 @@ package edu.washu.tag.generator.metadata.cohorting +import edu.washu.tag.generator.ai.catalog.CodeCache +import edu.washu.tag.generator.ai.catalog.attribute.DiagnosisCodeDesignator +import edu.washu.tag.generator.metadata.Diagnosis +import edu.washu.tag.generator.metadata.Study import edu.washu.tag.generator.util.RandomGenUtils import org.apache.commons.lang3.RandomUtils @@ -7,20 +11,56 @@ import java.util.function.Function enum CodeStrategy { - ONLY_ONE ({ List codes -> - RandomGenUtils.randomListEntry(codes) } - ), - STRICT_SUBSET ({ List codes -> - RandomGenUtils.randomSubset(codes, RandomUtils.insecure().randomInt(1, codes.size())) + ONLY_ONE ({ Map codes -> + [RandomGenUtils.setupWeightedLottery(codes).sample()] }), - SUBSET ({ List codes -> - RandomGenUtils.randomSubset(codes, RandomUtils.insecure().randomInt(1, codes.size() + 1)) + ONLY_ONE_FIXED_FOR_PATIENT (ONLY_ONE.codeRandomizer) { + @Override + List randomizeAndActualizeCodes(Study study, Map weightedCodes) { + final String patientUid = study.patient.patientInstanceUid + if (!(fixedForPatientUidCodeLookup.containsKey(patientUid))) { + fixedForPatientUidCodeLookup.put( + patientUid, + actualizeCode(ONLY_ONE.codeRandomizer.apply(weightedCodes)[0]) + ) + } + [fixedForPatientUidCodeLookup[patientUid]] + } + }, + STRICT_SUBSET ({ Map codes -> + randomizeSubsetOfSize( + RandomUtils.insecure().randomInt(1, codes.size()), + + ) + }), + SUBSET ({ Map codes -> + randomizeSubsetOfSize( + RandomUtils.insecure().randomInt(1, codes.size() + 1) + ) }) - private final Function, List> codeRandomizer + private final Function, List> codeRandomizer + private static final Map fixedForPatientUidCodeLookup = [:] - CodeStrategy(Function, List> randomizer) { + CodeStrategy(Function, List> randomizer) { codeRandomizer = randomizer } + List randomizeAndActualizeCodes(Study study, Map weightedCodes) { + final List randomizedCodes = codeRandomizer.apply(weightedCodes) + randomizedCodes.collect { code -> + actualizeCode(code) + } + } + + private List randomizeSubsetOfSize(int subsetSize, Map weightedCodes) { + RandomGenUtils.sampleWithoutReplacement(subsetSize, weightedCodes) + } + + private Diagnosis actualizeCode(String code) { + RandomGenUtils.randomListEntry( + CodeCache.resolvePartialCode(DiagnosisCodeDesignator.ICD_10, code) + ) + } + } \ No newline at end of file diff --git a/src/main/groovy/edu/washu/tag/generator/metadata/cohorting/SpecializedCohort.groovy b/src/main/groovy/edu/washu/tag/generator/metadata/cohorting/SpecializedCohort.groovy index 42fc054..5db5f7d 100644 --- a/src/main/groovy/edu/washu/tag/generator/metadata/cohorting/SpecializedCohort.groovy +++ b/src/main/groovy/edu/washu/tag/generator/metadata/cohorting/SpecializedCohort.groovy @@ -4,7 +4,6 @@ class SpecializedCohort { String name int numPatients - List ageRange // todo hmm? List trajectory } diff --git a/src/main/groovy/edu/washu/tag/generator/metadata/cohorting/StudyRequest.groovy b/src/main/groovy/edu/washu/tag/generator/metadata/cohorting/StudyRequest.groovy index 9c3dd60..25e0806 100644 --- a/src/main/groovy/edu/washu/tag/generator/metadata/cohorting/StudyRequest.groovy +++ b/src/main/groovy/edu/washu/tag/generator/metadata/cohorting/StudyRequest.groovy @@ -1,15 +1,28 @@ package edu.washu.tag.generator.metadata.cohorting +import com.fasterxml.jackson.annotation.JsonIgnore +import edu.washu.tag.generator.metadata.Diagnosis import edu.washu.tag.generator.metadata.Protocol +import edu.washu.tag.generator.metadata.Study class StudyRequest { String name Protocol protocol - List diagnosisCodes + List diagnosisCodes = [] + Map diagnosisCodesWeighted = [:] String additionalGenerationContext = '' - CodeStrategy codeStrategy = CodeStrategy.ONLY_ONE + CodeStrategy codeStrategy = CodeStrategy.ONLY_ONE_FIXED_FOR_PATIENT int studyOffsetMin int studyOffsetMax + @JsonIgnore + List randomizeCodes(Study study) { + final Map combinedCodes = new HashMap<>(diagnosisCodesWeighted) + diagnosisCodes.each { code -> + combinedCodes.put(code, 10) + } + codeStrategy.randomizeAndActualizeCodes(study, combinedCodes) + } + } diff --git a/src/main/groovy/edu/washu/tag/generator/metadata/reports/RadiologyReport2_7.groovy b/src/main/groovy/edu/washu/tag/generator/metadata/reports/RadiologyReport2_7.groovy index f7a0830..11592c7 100644 --- a/src/main/groovy/edu/washu/tag/generator/metadata/reports/RadiologyReport2_7.groovy +++ b/src/main/groovy/edu/washu/tag/generator/metadata/reports/RadiologyReport2_7.groovy @@ -7,7 +7,6 @@ import ca.uhn.hl7v2.model.v281.group.ORU_R01_PATIENT import ca.uhn.hl7v2.model.v281.message.ORU_R01 import ca.uhn.hl7v2.model.v281.segment.MSH import ca.uhn.hl7v2.model.v281.segment.OBX -import edu.washu.tag.generator.ai.catalog.attribute.WithDiagnosisCodes import edu.washu.tag.generator.hl7.v2.ReportVersion import edu.washu.tag.generator.hl7.v2.model.DoctorEncoder import edu.washu.tag.generator.hl7.v2.model.DoctorEncoder2_7 @@ -39,9 +38,9 @@ class RadiologyReport2_7 extends RadiologyReport { new ZpfGenerator().generateNonstandardSegment(this, radReport) } - if (hl7Version == ReportVersion.V2_7 && generatedReport instanceof WithDiagnosisCodes) { - generatedReport.parsedCodes.eachWithIndex { diagnosis, index -> - new Dg1Generator(index, diagnosis, generatedReport.designator).generateNonstandardSegment(this, radReport) + if (hl7Version == ReportVersion.V2_7) { + study.resolveDiagnoses().eachWithIndex { diagnosis, index -> + new Dg1Generator(index, diagnosis).generateNonstandardSegment(this, radReport) } } diff --git a/src/main/groovy/edu/washu/tag/generator/util/RandomGenUtils.groovy b/src/main/groovy/edu/washu/tag/generator/util/RandomGenUtils.groovy index 7c1ab00..8b9805c 100644 --- a/src/main/groovy/edu/washu/tag/generator/util/RandomGenUtils.groovy +++ b/src/main/groovy/edu/washu/tag/generator/util/RandomGenUtils.groovy @@ -75,4 +75,15 @@ class RandomGenUtils { ) } + static List sampleWithoutReplacement(int subsetSize, Map weights) { + final List results = [] + Map filteredWeights = new HashMap<>(weights) + while (results.size() < subsetSize) { + final X result = setupWeightedLottery(filteredWeights).sample() + results << result + filteredWeights.remove(result) + } + results + } + }