added other preprocessing functions

kosuri-indu · kosuri-indu · commit fae302699d3b · 2025-06-27T18:41:16.000+05:30
diff --git a/Project.toml b/Project.toml
@@ -4,6 +4,8 @@ authors = ["Dilum Aluthge", "contributors"]
 version = "2.0.0"
 
 [deps]
+InlineStrings = "842dd82b-1e85-43dc-bf29-5d0ee9dffc48"
+OMOPCommonDataModel = "ba65db9e-6590-4054-ab8a-101ed9124986"
 Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
 
 [weakdeps]
@@ -12,15 +14,15 @@ Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 DrWatson = "634d3b9d-ee7a-5ddf-bec9-22491ea816e1"
 Serialization = "9e88b42a-f829-5b0c-bbe9-9e923198166b"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
-InlineStrings = "842dd82b-1e85-43dc-bf29-5d0ee9dffc48"
-OMOPCommonDataModel = "ba65db9e-6590-4054-ab8a-101ed9124986"
 
 [extensions]
 HealthBaseDrWatsonExt = "DrWatson"
 HealthBaseOMOPCDMExt = ["DataFrames", "OMOPCommonDataModel", "InlineStrings", "Serialization", "Statistics", "Dates"]
 
 [compat]
 Dates = "1.11.0"
+InlineStrings = "1.4.4"
+OMOPCommonDataModel = "0.1.4"
 Tables = "1.12.1"
 julia = "1.10"
 
diff --git a/ext/HealthBaseOMOPCDMExt.jl b/ext/HealthBaseOMOPCDMExt.jl
@@ -87,7 +87,11 @@ coltypes = eltype.(eachcol(ht.source))
 4. Quick-fail/warning for bad data:
 # TODO: Will finalize the implementation soon and then add an example
 """
-function HealthBase.HealthTable(df::DataFrame; omop_cdm_version::String="v5.4.0", disable_type_enforcement=false, collect_errors=true)
+function HealthBase.HealthTable(
+    df::DataFrame; omop_cdm_version::String="v5.4.0", 
+    disable_type_enforcement=false, 
+    collect_errors=true
+)
     if !haskey(OMOPCDM_VERSIONS, omop_cdm_version)
         throw(ArgumentError("OMOP CDM version '$(omop_cdm_version)' is not supported. Available versions: $(keys(OMOPCDM_VERSIONS))"))
     end
@@ -142,7 +146,11 @@ function HealthBase.HealthTable(df::DataFrame; omop_cdm_version::String="v5.4.0"
 end
 
 # TODO: Add Documentation
-function HealthBase.one_hot_encode(ht::HealthTable; cols::Vector{Symbol}, drop_original::Bool=true)
+function HealthBase.one_hot_encode(
+    ht::HealthTable; 
+    cols::Vector{Symbol}, 
+    drop_original::Bool=true
+)
     df = copy(ht.source)
     for col in cols
         unique_vals = unique(skipmissing(df[!, col]))
@@ -158,21 +166,87 @@ function HealthBase.one_hot_encode(ht::HealthTable; cols::Vector{Symbol}, drop_o
 end
 
 # TODO: Add Documentation
-function HealthBase.impute_missing(ht::HealthTable; cols::Vector{Symbol}, strategy::Symbol=:mean)
+function HealthBase.impute_missing(
+    ht::HealthTable;
+    cols::Union{Vector{Symbol}, Vector{Pair{Symbol,Symbol}}},
+    strategy::Symbol=:mean,
+)
     df = copy(ht.source)
-    for col in cols
-        if strategy == :mean
-            non_missing_vals = skipmissing(df[!, col])
-            if isempty(non_missing_vals)
-                throw(ArgumentError("Column '$col' has only missing values."))
+
+    strat_pairs = cols isa Vector{Symbol} ? [c => strategy for c in cols] : cols
+
+    for (col, strat) in strat_pairs
+        @assert col in propertynames(df) "Column '$(col)' not found in table."
+        vals = df[!, col]
+        nonmiss = collect(skipmissing(vals))
+        if isempty(nonmiss)
+            throw(ArgumentError("Column '$(col)' has only missing values – cannot impute."))
+        end
+
+        replacement = begin
+            if strat == :mean
+                mean(nonmiss)
+            elseif strat == :median
+                median(nonmiss)
+            elseif strat == :mode
+                mode_val = nothing
+                counts = Dict{Any,Int}()
+                for v in nonmiss
+                    counts[v] = get(counts,v,0)+1
+                    if mode_val === nothing || counts[v] > counts[mode_val]
+                        mode_val = v
+                    end
+                end
+                mode_val
+            else
+                throw(ArgumentError("Unsupported imputation strategy '$(strat)'. Supported: :mean, :median, :mode."))
             end
-            mean_val = mean(non_missing_vals)
-            df[!, col] = coalesce.(df[!, col], mean_val)
-        else
-            throw(ArgumentError("Unsupported imputation strategy: $strategy"))
         end
+        df[!, col] = coalesce.(vals, replacement)
     end
-    return HealthBase.HealthTable(df; omop_cdm_version=ht.omop_cdm_version)
+
+    return HealthBase.HealthTable(source=df, omop_cdm_version=ht.omop_cdm_version)
+end
+
+# TODO: Add Documentation
+function HealthBase.apply_vocabulary_compression(
+    ht::HealthTable; 
+    cols::Vector{Symbol}, 
+    min_freq::Integer=10, 
+    other_label::AbstractString="Other"
+)
+    df = copy(ht.source)
+    for col in cols
+        @assert col in propertynames(df) "Column '$(col)' not found in table."
+        counts = combine(groupby(df, col), nrow => :freq)
+        to_compress = counts[counts.freq .< min_freq, col]
+        if !isempty(to_compress)
+            mask = in(to_compress).(df[!, col])
+            df[mask, col] .= other_label
+        end
+    end
+    return HealthBase.HealthTable(source=df, omop_cdm_version=ht.omop_cdm_version)
+end
+
+# TODO: Add Documentation
+function HealthBase.map_concepts(
+    ht::HealthTable;
+    col::Symbol,
+    mapping::AbstractDict,
+    new_col::Union{Symbol,Nothing}=nothing,
+    drop_original::Bool=false,
+)
+    @assert col in propertynames(ht.source) "Column '$(col)' not found in table."
+    df = copy(ht.source)
+
+    target_col = isnothing(new_col) ? col : new_col
+    df[!, target_col] = get.(Ref(mapping), df[!, col], df[!, col])
+
+    if drop_original && !isnothing(new_col)
+        select!(df, Not(col))
+    end
+
+    return HealthBase.HealthTable(source=df, omop_cdm_version=ht.omop_cdm_version)
 end
 
 # TODO: Add Documentation
@@ -194,6 +268,4 @@ function HealthBase.normalize_column(ht::HealthTable; cols::Vector{Symbol}, meth
     return HealthBase.HealthTable(df; omop_cdm_version=ht.omop_cdm_version)
 end
 
-# TODO: Add Other Preprocessing Utilities
-
 end
diff --git a/src/omopcdm_stub.jl b/src/omopcdm_stub.jl
@@ -1,7 +1,11 @@
 function one_hot_encode end
 function normalize_column end
 function impute_missing end
+function apply_vocabulary_compression end
+function map_concepts end
 
 export one_hot_encode
 export normalize_column
 export impute_missing
+export apply_vocabulary_compression
+export map_concepts
diff --git a/test/omopcdmext.jl b/test/omopcdmext.jl
@@ -1,4 +1,5 @@
 using Statistics
+using DataFrames
 
 @testset "HealthBaseOMOPCDMExt" begin
     # This DataFrame is compliant with the OMOP CDM v5.4.0 PERSON table schema.
@@ -57,48 +58,98 @@ using Statistics
         @test df2 == person_df_good
     end
 
-    @testset "Preprocessing Utilities" begin
+    @testset "Preprocessing Utilities (single example DataFrame)" begin
         df = DataFrame(
-            id = 1:4,
-            cat = ["a", "b", "a", "c"],
-            num1 = [1.0, 2.5, missing, 4.0],
-            num2 = [10.0, missing, 30.0, 40.0]
+            person_id = 1:10,
+            condition_source_value = [
+                "Hypertension", "Diabetes", "Asthma", "Asthma", "Hypertension",
+                "Fibromyalgia", "Hyperlipidemia", "RareDisease1", "RareDisease2", "RareDisease3"
+            ],
+            condition_concept_id = [
+                316866,   # Hypertension
+                201826,   # Diabetes
+                317009,   # Asthma (mild)
+                317010,   # Asthma (severe)
+                316866,   # Hypertension
+                317707,   # Fibromyalgia
+                4329058,  # Hyperlipidemia
+                1234567, 2345678, 3456789  # Rare
+            ],
+            systolic_bp = [140.0, 130.0, 110.0, missing, 150.0, 120.0, missing, 135.0, 128.0, 145.0],
+            diastolic_bp = [90.0, 85.0, 70.0, missing, 95.0, 80.0, missing, 88.0, 82.0, 92.0]
         )
-        base_ht = HealthTable(df)
+        base_ht = HealthTable(source=df; omop_cdm_version="v5.4.1")
     
         @testset "one_hot_encode - drop_original=true" begin
-            ht_oh = one_hot_encode(base_ht; cols=[:cat], drop_original=true)
-            @test "cat" ∉ names(ht_oh.source)
-            expected_cols = Set([:cat_a, :cat_b, :cat_c])
+            ht_oh = one_hot_encode(base_ht; cols=[:condition_source_value], drop_original=true)
+            @test :condition_source_value ∉ names(ht_oh.source)
+            expected_cols = Set(Symbol.("condition_source_value_" .* ["Hypertension", "Diabetes", "Asthma", "Fibromyalgia", "Hyperlipidemia"]))
             @test expected_cols ⊆ Set(Symbol.(names(ht_oh.source)))
-            @test ht_oh.source.cat_a == [true, false, true, false]
-            @test ht_oh.source.cat_b == [false, true, false, false]
-            @test ht_oh.source.cat_c == [false, false, false, true]
+            @test ht_oh.source.condition_source_value_Hypertension == (df.condition_source_value .== "Hypertension")
+            @test ht_oh.source.condition_source_value_Diabetes == (df.condition_source_value .== "Diabetes")
         end
         
         @testset "one_hot_encode - drop_original=false" begin
-            ht_oh = one_hot_encode(base_ht; cols=[:cat], drop_original=false)
-            @test "cat" in names(ht_oh.source)
-            expected_cols = Set([:cat_a, :cat_b, :cat_c])
+            ht_oh = one_hot_encode(base_ht; cols=[:condition_source_value], drop_original=false)
+            @test :condition_source_value in Symbol.(names(ht_oh.source))
+            expected_cols = Set(Symbol.("condition_source_value_" .* ["Hypertension", "Diabetes", "Asthma", "Fibromyalgia", "Hyperlipidemia"]))
             @test expected_cols ⊆ Set(Symbol.(names(ht_oh.source)))
         end               
     
         @testset "impute_missing - mean" begin
-            ht_imp = impute_missing(base_ht; cols=[:num1, :num2], strategy=:mean)
-            @test all(!ismissing, ht_imp.source.num1)
-            @test all(!ismissing, ht_imp.source.num2)
-            @test ht_imp.source.num1[3] ≈ mean(skipmissing(base_ht.source.num1))
-            @test ht_imp.source.num2[2] ≈ mean(skipmissing(base_ht.source.num2))
+            ht_imp = impute_missing(base_ht; cols=[:systolic_bp, :diastolic_bp], strategy=:mean)
+            @test all(!ismissing, ht_imp.source.systolic_bp)
+            @test all(!ismissing, ht_imp.source.diastolic_bp)
+            @test ht_imp.source.systolic_bp[4] ≈ mean(skipmissing(base_ht.source.systolic_bp))
+            @test ht_imp.source.diastolic_bp[4] ≈ mean(skipmissing(base_ht.source.diastolic_bp))
         end
     
         @testset "normalize_column - standard" begin
-            ht_imp = impute_missing(base_ht; cols=[:num1, :num2])
-            ht_norm = normalize_column(ht_imp; cols=[:num1, :num2])
-            @test isapprox(mean(ht_norm.source.num1), 0.0; atol=1e-8)
-            @test isapprox(std(ht_norm.source.num1), 1.0; atol=1e-8)
-            @test isapprox(mean(ht_norm.source.num2), 0.0; atol=1e-8)
-            @test isapprox(std(ht_norm.source.num2), 1.0; atol=1e-8)
-            @test all(x -> x isa Float64, vec(Matrix(ht_norm.source[:, [:num1, :num2]])))
+            ht_imp = impute_missing(base_ht; cols=[:systolic_bp, :diastolic_bp])
+            ht_norm = normalize_column(ht_imp; cols=[:systolic_bp, :diastolic_bp])
+            @test isapprox(mean(ht_norm.source.systolic_bp), 0.0; atol=1e-8)
+            @test isapprox(std(ht_norm.source.systolic_bp), 1.0; atol=1e-8)
+            @test isapprox(mean(ht_norm.source.diastolic_bp), 0.0; atol=1e-8)
+            @test isapprox(std(ht_norm.source.diastolic_bp), 1.0; atol=1e-8)
+            @test all(x -> x isa Float64, vec(Matrix(ht_norm.source[:, [:systolic_bp, :diastolic_bp]])))
+        end
+
+        @testset "impute_missing - median" begin
+            df_mid = DataFrame(num = [1.0, missing, 3.0, missing])
+            ht_mid = HealthTable(df_mid)
+            ht_imp = impute_missing(ht_mid; cols=[:num], strategy=:median)
+            @test all(!ismissing, ht_imp.source.num)
+            @test ht_imp.source.num[2] == median([1.0, 3.0])
+        end
+
+        @testset "impute_missing - mixed strategies" begin
+            ht_mix = impute_missing(base_ht; cols=[:systolic_bp => :mean, :diastolic_bp => :median])
+            @test ht_mix.source.systolic_bp[4] ≈ mean(skipmissing(base_ht.source.systolic_bp))
+            @test ht_mix.source.diastolic_bp[4] == median(skipmissing(base_ht.source.diastolic_bp))
+        end
+
+        @testset "apply_vocabulary_compression" begin
+            ht_comp = apply_vocabulary_compression(base_ht; cols=[:condition_source_value], min_freq=2, other_label="Other")
+            @test ht_comp.source.condition_source_value == [
+                "Hypertension", "Other", "Asthma", "Asthma", "Hypertension", "Other", "Other", "Other", "Other", "Other"
+            ]
+        end
+
+        @testset "map_concepts" begin
+            mapping = Dict(
+                316866 => "Hypertension",
+                201826 => "Diabetes",
+                317009 => "Asthma",
+                317010 => "Asthma"
+            )
+            ht_m1 = map_concepts(base_ht; col=:condition_concept_id, mapping=mapping, new_col=:condition_group)
+            @test ht_m1.source.condition_group == [
+                "Hypertension", "Diabetes", "Asthma", "Asthma", "Hypertension", 317707, 4329058, 1234567, 2345678, 3456789
+            ]
+            @test :condition_concept_id in Symbol.(names(ht_m1.source))
+            ht_m2 = map_concepts(base_ht; col=:condition_concept_id, mapping=mapping, new_col=:condition_group, drop_original=true)
+            @test :condition_concept_id ∉ names(ht_m2.source)
+            @test ht_m2.source.condition_group[1] == "Hypertension"
         end
     end
 
