From bd9d002c4733d8557dc5cbf47f66b7210d9ff365 Mon Sep 17 00:00:00 2001 From: lupemba Date: Sun, 12 Jul 2026 13:52:56 +0200 Subject: [PATCH] CategoricalVariable feature. Read categorical variables using CategoricalArrays.jl --- Project.toml | 4 +- src/CommonDataModel.jl | 4 + src/categoricalvariable.jl | 148 +++++++++++++++++++++++++++++++++++++ test/Project.toml | 2 + test/runtests.jl | 4 + test/test_categorical.jl | 141 +++++++++++++++++++++++++++++++++++ 6 files changed, 302 insertions(+), 1 deletion(-) create mode 100644 src/categoricalvariable.jl create mode 100644 test/test_categorical.jl diff --git a/Project.toml b/Project.toml index 4dfaba4..ba1b48e 100644 --- a/Project.toml +++ b/Project.toml @@ -4,9 +4,10 @@ keywords = ["netcdf", "GRIB", "climate and forecast conventions", "oceanography" license = "MIT" desc = "CommonDataModel is a module that defines types common to NetCDF and GRIB data" authors = ["Alexander Barth and contributors (https://github.com/JuliaGeo/CommonDataModel.jl/graphs/contributors)"] -version = "0.4.4" +version = "0.4.5" [deps] +CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597" CFTime = "179af706-886a-5703-950a-314cd64e0468" DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8" Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" @@ -16,6 +17,7 @@ Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" [compat] +CategoricalArrays = "1" CFTime = "0.2.7" DataStructures = "0.17, 0.18, 0.19" Dates = "1" diff --git a/src/CommonDataModel.jl b/src/CommonDataModel.jl index 8c2fb86..05041b6 100644 --- a/src/CommonDataModel.jl +++ b/src/CommonDataModel.jl @@ -52,6 +52,9 @@ import Statistics: var +import CategoricalArrays: + CategoricalValue, + CategoricalArray include("CatArrays.jl") @@ -70,6 +73,7 @@ include("aggregation.jl") include("groupby.jl") include("rolling.jl") include("memory_dataset.jl") +include("categoricalvariable.jl") end # module CommonDataModel diff --git a/src/categoricalvariable.jl b/src/categoricalvariable.jl new file mode 100644 index 0000000..1e54534 --- /dev/null +++ b/src/categoricalvariable.jl @@ -0,0 +1,148 @@ + + +abstract type AbstractCategoricalVariable{V, N, R} <: AbstractVariable{CategoricalValue{V, UInt32}, N} end + +# special methods for AbstractCategoricalVariable +getvaluearray(a::AbstractCategoricalVariable)::AbstractVariable = throw(ArgumentError( + "getvaluearray is not implemented for $(typeof(a))." +)) +getmapping(a::AbstractCategoricalVariable)::AbstractDict = throw(ArgumentError( + "getmapping is not implemented for $(typeof(a))." +)) + +# forward CommonDataModel.API +name(v_category::AbstractCategoricalVariable) = name(getvaluearray(v_category)) +dimnames(v_category::AbstractCategoricalVariable) = dimnames(getvaluearray(v_category)) +dataset(v_category::AbstractCategoricalVariable) = dataset(getvaluearray(v_category)) +attribnames(v_category::AbstractCategoricalVariable) = attribnames(getvaluearray(v_category)) +attrib(v_category::AbstractCategoricalVariable, name::SymbolOrString) = attrib(getvaluearray(v_category),name) + +# forward other basic method +Base.size(a::AbstractCategoricalVariable) = size(getvaluearray(a)) +DiskArrays.haschunks(a::AbstractCategoricalVariable) = DiskArrays.haschunks(getvaluearray(a)) +DiskArrays.eachchunk(a::AbstractCategoricalVariable) = DiskArrays.eachchunk(getvaluearray(a)) +Base.getindex(a::AbstractCategoricalVariable, name::SymbolOrString) = getindex(getvaluearray(a),name) +Base.getindex(a::AbstractCategoricalVariable, name::CFStdName) = getindex(getvaluearray(a),name) + +# ---- internal helpers --------------------------------------------------------- + +function _sorted_labels(mapping::AbstractDict{R, V}) where {R, V} + sorted_codes = sort(collect(keys(mapping))) + return V[mapping[c] for c in sorted_codes] +end + +function _build_categorical_array( + raw::AbstractArray{R, N}, mapping::AbstractDict{R, V} +) where {V, N, R} + label_values = V[mapping[c] for c in raw] + return CategoricalArray{V, N, UInt32}(label_values; levels=_sorted_labels(mapping)) +end + +function _build_cat_value(code::R, mapping::AbstractDict{R, V}) where {R, V} + ca = CategoricalArray{V, 1, UInt32}(V[mapping[code]]; levels=_sorted_labels(mapping)) + return ca[1] +end + + +# ---- readblock! --------------------------------------------------------- + +function DiskArrays.readblock!( + a::AbstractCategoricalVariable{V, N, R}, aout, r::AbstractUnitRange... +) where {V, N, R} + raw = Array{R}(undef, length.(r)...) + DiskArrays.readblock!(getvaluearray(a), raw, r...) + ca = _build_categorical_array(raw, getmapping(a)) + for i in eachindex(aout, ca) + aout[i] = ca[i] + end + return ca +end + +function DiskArrays.writeblock!( + ::AbstractCategoricalVariable, ::Any, r::AbstractUnitRange... +) + throw(ArgumentError( + "Writing to a categorical CF variable is not supported yet." + )) +end + +# ---- getindex ----------------------------------------------------------------- +# +# Two methods via multiple dispatch — no runtime type check needed. + +# Scalar indices (all Integer or CartesianIndex) → CategoricalValue +function Base.getindex( + a::AbstractCategoricalVariable{V, N, R}, inds::Union{Integer, CartesianIndex}... +) where {V, N, R} + @boundscheck checkbounds(a, inds...) + DiskArrays.checkscalar(a, inds) + raw = getvaluearray(a)[inds...] + return _build_cat_value(raw, getmapping(a)) +end + +# Array indices (ranges, colons, vectors, …) → CategoricalArray +function Base.getindex(a::AbstractCategoricalVariable{V, N, R}, inds...) where {V, N, R} + @boundscheck checkbounds(a, inds...) + raw = getvaluearray(a)[inds...] + return _build_categorical_array(raw, getmapping(a)) +end + +# ---- CFVariable ----------------------------------------------------------------- +function _add_missing(data, mapping, f_m_vals) + pairs = [mapping[code] => missing for code in f_m_vals] + return isempty(pairs) ? data : replace(data, pairs...) +end + +# A custom `maskingvalue` (e.g. NaN, used for numeric CFVariables) does not +# make sense for categorical data, so masked entries always become `missing` +# here regardless of `maskingvalue(v)`. +function DiskArrays.readblock!( + v::CFVariable{T,N,TV}, aout, r::AbstractUnitRange... + ) where {T,N,TV<:AbstractCategoricalVariable} + + parent_var = parent(v) ## + data = similar(aout, eltype(parent_var)) + DiskArrays.readblock!(parent_var, data, r...) + + aout .= _add_missing(data, + getmapping(parent_var), + fill_and_missing_values(v)) + + return nothing +end + + +function Base.getindex(v::CFVariable{T,N,TV}, inds::Union{Integer, CartesianIndex}... + ) where {T,N,TV<:AbstractCategoricalVariable} + + parent_var = parent(v) + cat_val = Base.getindex(parent_var, inds...) + mapping = getmapping(parent_var) + is_missing = cat_val in (mapping[code] for code in fill_and_missing_values(v)) + return is_missing ? missing : cat_val +end + +function Base.getindex(v::CFVariable{T,N,TV}, inds... + ) where {T,N,TV<:AbstractCategoricalVariable} + + parent_var = parent(v) + data = Base.getindex(parent_var, inds...) + return _add_missing(data, + getmapping(parent_var), + fill_and_missing_values(v)) +end + + +function Base.getindex(v::CFVariable{T,N,TV}, name::SymbolOrString + ) where {T,N,TV<:AbstractCategoricalVariable} + + parent_var = parent(v) + return getindex(getvaluearray(parent_var),name) +end + +function Base.getindex(v::CFVariable{T,N,TV}, name::CFStdName + ) where {T,N,TV<:AbstractCategoricalVariable} + + parent_var = parent(v) + return getindex(getvaluearray(parent_var),name) +end diff --git a/test/Project.toml b/test/Project.toml index 542a8ae..98c40be 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,6 +1,7 @@ [deps] Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" CFTime = "179af706-886a-5703-950a-314cd64e0468" +CategoricalArrays = "324d7699-5711-5eae-9e2f-1d82baa6b597" DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8" Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" DiskArrays = "3c3547ce-8d99-4f5e-a174-61eb10b00ae3" @@ -11,3 +12,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [compat] Aqua = "0.8" +CategoricalArrays = "1" diff --git a/test/runtests.jl b/test/runtests.jl index 03cf8aa..41d1361 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -39,6 +39,10 @@ end include("test_rolling.jl") end +@testset "categorical variable" begin + include("test_categorical.jl") +end + @testset "aqua checks" begin include("test_aqua.jl") end diff --git a/test/test_categorical.jl b/test/test_categorical.jl new file mode 100644 index 0000000..25b60a3 --- /dev/null +++ b/test/test_categorical.jl @@ -0,0 +1,141 @@ +using Test +using CommonDataModel: + AbstractCategoricalVariable, + AbstractVariable, + MemoryDataset, + defVar, + cfvariable, + dataset + +import CommonDataModel as CDM +import DiskArrays +import CategoricalArrays: CategoricalValue, CategoricalArray, unwrap, levels + + +struct CategoricalVariable{V, N, R} <: AbstractCategoricalVariable{V, N, R} + data::AbstractVariable{R,N} + mapping::Dict{R,V} +end + +CDM.getvaluearray(a::CategoricalVariable) = a.data +CDM.getmapping(a::CategoricalVariable) = a.mapping + +const CLOUD_MAPPING = Dict{Int8, String}( + Int8(0) => "Not processed", + Int8(1) => "Cloud free", + Int8(2) => "Cloud contaminated", + Int8(3) => "Cloud filled", + Int8(4) => "Dust contaminated", +) + +# 3×4 grid of cloud codes, chunked 2×2 +const RAW_CODES = Int8[ + 0 1 2 3; + 1 2 3 4; + 0 0 1 1 +] + +function make_mock() + ds = MemoryDataset(tempname(), "c") + v = defVar(ds,"cloud",RAW_CODES,("lon","lat"), attrib = [ + "standard_name" => "cloud_mask", + "long_name" => "Cloud mask", + "_FillValue" => Int8(0) + ]) + + mock = CategoricalVariable(parent(v), CLOUD_MAPPING) + return mock +end + + +@testset "AbstractCategoricalVariable — eltype" begin + mock = make_mock() + @test eltype(mock) == CategoricalValue{String, UInt32} +end + +@testset "AbstractCategoricalVariable — collect" begin + mock = make_mock() + ca = collect(mock) + + @test ca isa CategoricalArray{String, 2, UInt32} + @test size(ca) == size(RAW_CODES) + + # Level ordering follows sorted raw codes (0,1,2,3,4) + expected_levels = [CLOUD_MAPPING[k] for k in sort(collect(keys(CLOUD_MAPPING)))] + @test levels(ca) == expected_levels + + # Values match the mapping + for i in eachindex(RAW_CODES) + @test unwrap(ca[i]) == CLOUD_MAPPING[RAW_CODES[i]] + end +end + + +@testset "AbstractCategoricalVariable — array getindex" begin + mock = make_mock() + + slice = mock[1:2, :] + @test slice isa CategoricalArray{String, 2, UInt32} + @test size(slice) == (2, 4) + @test slice == collect(mock)[1:2, :] + + row = mock[1, :] + @test row isa CategoricalArray{String, 1, UInt32} + @test size(row) == (4,) + @test row == collect(mock)[1, :] + + col = mock[:, 2] + @test col isa CategoricalArray{String, 1, UInt32} + @test unwrap.(col) == getindex.(Ref(CLOUD_MAPPING), RAW_CODES[:, 2]) +end + + +@testset "AbstractCategoricalVariable — scalar getindex" begin + mock = make_mock() + + val = mock[1, 1] + @test val isa CategoricalValue{String, UInt32} + @test unwrap(val) == CLOUD_MAPPING[RAW_CODES[1, 1]] + + val2 = mock[2, 4] + @test unwrap(val2) == CLOUD_MAPPING[RAW_CODES[2, 4]] +end + +@testset "AbstractCategoricalVariable — broadcasting" begin + mock = make_mock() + + # Broadcast a function element-wise: unwrap over all elements + broad_cast_var = mock .== "Cloud free" + @test broad_cast_var isa DiskArrays.BroadcastDiskArray + @test sum(broad_cast_var) == 4 +end + + +@testset "AbstractCategoricalVariable — in" begin + mock = make_mock() + + @test "Cloud free" in mock + @test "Not processed" in mock + @test !("Snowy" in mock) +end + + +@testset "categorical CFVariable" begin + mock = make_mock() + mock_cf = cfvariable(dataset(mock), "cloud_type";_v = mock) + + @test eltype(mock_cf) == Union{eltype(mock),Missing} + @test ismissing(mock_cf[1, 1]) + @test mock_cf[1, 2] == mock[1, 2] + @test mock_cf[2, 1] == mock[2, 1] + @test ismissing(mock_cf[3, 2]) + + date_read = mock_cf[:,:] + @test date_read isa CategoricalArray + @test count(ismissing, date_read) == 3 + + broad_cast_var = mock_cf .== "Cloud free" + @test broad_cast_var isa DiskArrays.BroadcastDiskArray + @test sum(skipmissing(broad_cast_var)) == 4 + +end