CategoricalVariable feature.

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 <barth.alexander@gmail.com> 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"

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
 

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

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"

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

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