Add DimensionalData extension (#3413)

Co-authored-by: longemen3000 <[email protected]>
Co-authored-by: Andrés Riedemann <[email protected]>

Author: odow

Project.toml

@@ -12,15 +12,23 @@ Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 SnoopPrecompile = "66db9d55-30c0-4569-8b51-7e840670fc0c"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 
+[weakdeps]
+DimensionalData = "0703355e-b756-11e9-17c0-8b28908087d0"
+
+[extensions]
+JuMPDimensionalDataExt = "DimensionalData"
+
 [compat]
+DimensionalData = "0.24"
 MathOptInterface = "1.18"
 MutableArithmetics = "1"
 OrderedCollections = "1"
 SnoopPrecompile = "1"
 julia = "1.6"
 
 [extras]
+DimensionalData = "0703355e-b756-11e9-17c0-8b28908087d0"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test"]
+test = ["DimensionalData", "Test"]

docs/Project.toml

@@ -3,6 +3,7 @@ CDDLib = "3391f64e-dcde-5f30-b752-e11513730f60"
 CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b"
 Clarabel = "61c947e1-3e6d-4ee4-985a-eec8c727bd6e"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
+DimensionalData = "0703355e-b756-11e9-17c0-8b28908087d0"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6"
 Dualization = "191a621a-6537-11e9-281d-650236a99e60"
@@ -34,6 +35,7 @@ CDDLib = "=0.9.2"
 CSV = "0.10"
 Clarabel = "=0.5.1"
 DataFrames = "1"
+DimensionalData = "0.24"
 Documenter = "0.27.9, 0.28"
 Dualization = "0.5"
 GLPK = "=1.1.2"

docs/make.jl

@@ -187,11 +187,16 @@ for (solver, data) in TOML.parsefile(joinpath(@__DIR__, "packages.toml"))
     end
 end
 push!(_LIST_OF_SOLVERS, "JuliaOpt/NLopt.jl" => "packages/NLopt.md")
+push!(
+    _LIST_OF_EXTENSIONS,
+    "rafaqz/DimensionalData.jl" => "extensions/DimensionalData.md",
+)
+
 # Sort, with jump-dev repos at the start.
 sort!(_LIST_OF_SOLVERS; by = x -> (!startswith(x[1], "jump-dev/"), x[1]))
 sort!(_LIST_OF_EXTENSIONS; by = x -> (!startswith(x[1], "jump-dev/"), x[1]))
 pushfirst!(_LIST_OF_SOLVERS, "Introduction" => "packages/solvers.md")
-pushfirst!(_LIST_OF_EXTENSIONS, "Introduction" => "packages/extensions.md")
+pushfirst!(_LIST_OF_EXTENSIONS, "Introduction" => "extensions/introduction.md")
 
 # ==============================================================================
 #  JuMP API

docs/src/extensions/DimensionalData.md

@@ -0,0 +1,94 @@
+# DimensionalData.jl
+
+[DimensionalData.jl](https://github.com/rafaqz/DimensionalData.jl) provides
+tools and abstractions for working with rectangular arrays that have named
+dimensions.
+
+!!! compat
+    Using the DimensionalData extension with JuMP requires Julia v1.9 or later.
+
+The DimensionalData extension in JuMP lets you construct a `DimensionalData.DimArray`
+as an alternative to [`Containers.DenseAxisArray`](@ref) in the JuMP macros.
+
+## License
+
+DimensionalData.jl is licensed under the [MIT license](https://github.com/rafaqz/DimensionalData.jl/blob/main/LICENSE).
+
+## Installation
+
+Install DimensionalData using `Pkg.add`:
+
+```julia
+import Pkg
+Pkg.add("DimensionalData")
+```
+
+## Use with JuMP
+
+Activate the extension by loading both JuMP and DimensionalData:
+
+```jldoctest ext_dimensional_data
+julia> using JuMP, DimensionalData
+```
+
+Then, pass `container = DimensionalData.DimArray` in the [`@variable`](@ref),
+[`@constraint`](@ref), or [`@expression`](@ref) macros:
+```jldoctest ext_dimensional_data
+julia> model = Model();
+
+julia> @variable(
+           model,
+           x[i = 2:4, j = ["a", "b"]] >= i,
+           container = DimensionalData.DimArray,
+       )
+3×2 DimArray{VariableRef,2} with dimensions:
+  Dim{:i} Sampled{Int64} 2:4 ForwardOrdered Regular Points,
+  Dim{:j} Categorical{String} String["a", "b"] ForwardOrdered
+    "a"     "b"
+ 2  x[2,a]  x[2,b]
+ 3  x[3,a]  x[3,b]
+ 4  x[4,a]  x[4,b]
+```
+
+Here `x` is a `DimensionalData.Dim` array object, so indexing uses the
+DimensionalData syntax:
+```jldoctest ext_dimensional_data
+julia> x[At(2), At("a")]
+x[2,a]
+
+julia> x[2, 2]
+x[3,b]
+```
+
+You can use `container = DimensionalData.DimArray` in the [`@expression`](@ref)
+macro:
+```jldoctest ext_dimensional_data
+julia> @expression(
+           model,
+           expr[j = ["a", "b"]],
+           sum(x[At(i), At(j)] for i in 2:4),
+           container = DimensionalData.DimArray,
+       )
+2-element DimArray{AffExpr,1} with dimensions:
+  Dim{:j} Categorical{String} String["a", "b"] ForwardOrdered
+ "a"  x[2,a] + x[3,a] + x[4,a]
+ "b"  x[2,b] + x[3,b] + x[4,b]
+```
+and in [`@constraint`](@ref):
+```jldoctest ext_dimensional_data
+julia> @constraint(
+           model,
+           [j = ["a", "b"]],
+           expr[At(j)] <= 1,
+           container = DimensionalData.DimArray,
+       )
+2-element DimArray{ConstraintRef{Model, MathOptInterface.ConstraintIndex{MathOptInterface.ScalarAffineFunction{Float64}, MathOptInterface.LessThan{Float64}}, ScalarShape},1} with dimensions:
+  Dim{:j} Categorical{String} String["a", "b"] ForwardOrdered
+ "a"  x[2,a] + x[3,a] + x[4,a] ≤ 1
+ "b"  x[2,b] + x[3,b] + x[4,b] ≤ 1
+```
+
+## Documentation
+
+See the [DimensionalData.jl documentation](https://rafaqz.github.io/DimensionalData.jl/stable/)
+for more details on the syntax and features of `DimensionalData.DimArray`.

docs/src/packages/extensions.md renamed to docs/src/extensions/introduction.md

@@ -19,3 +19,13 @@ README contents in the JuMP documentation for the benefit of users.
 Written an extension? Add it to this section of the JuMP documentation by making
 a pull request to the [`docs/packages.toml`](https://github.com/jump-dev/JuMP.jl/blob/master/docs/packages.toml)
 file.
+
+## Weak dependencies
+
+Some extensions listed in this section are implemented using the [weak dependency](https://pkgdocs.julialang.org/v1/creating-packages/#Weak-dependencies)
+feature added to Julia in v1.9. These extensions are activated if and only if
+you have `JuMP` and the other package loaded into your current scope with
+`using` or `import`.
+
+!!! compat
+    Using a weak dependency requires Julia v1.9 or later.

ext/JuMPDimensionalDataExt.jl

@@ -0,0 +1,33 @@
+#  Copyright 2017, Iain Dunning, Joey Huchette, Miles Lubin, and contributors
+#  This Source Code Form is subject to the terms of the Mozilla Public
+#  License, v. 2.0. If a copy of the MPL was not distributed with this
+#  file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+module JuMPDimensionalDataExt
+
+import DimensionalData
+import JuMP
+
+function JuMP.Containers.container(
+    f::F,
+    indices::JuMP.Containers.VectorizedProductIterator,
+    c::Type{DimensionalData.DimArray},
+    names::AbstractVector,
+) where {F<:Function}
+    dims = NamedTuple(i => j for (i, j) in zip(names, indices.prod.iterators))
+    return DimensionalData.DimArray(map(i -> f(i...), indices), dims)
+end
+
+function JuMP.Containers.container(
+    ::Function,
+    ::JuMP.Containers.NestedIterator,
+    ::Type{DimensionalData.DimArray},
+    ::AbstractVector,
+)
+    return error(
+        "Unable to create a `DimArray` because the container does not form " *
+        "a dense rectangular array",
+    )
+end
+
+end #module

ext/test_DimensionalData.jl

@@ -0,0 +1,67 @@
+#  Copyright 2017, Iain Dunning, Joey Huchette, Miles Lubin, and contributors
+#  This Source Code Form is subject to the terms of the Mozilla Public
+#  License, v. 2.0. If a copy of the MPL was not distributed with this
+#  file, You can obtain one at https://mozilla.org/MPL/2.0/.
+
+module TestContainersDimensionalData
+
+using Test
+
+using DimensionalData
+using JuMP
+
+function test_dimension_data_variable()
+    model = Model()
+    @variable(model, x[i = 2:4, j = ["a", "b"]], container = DimArray)
+    @test x isa DimArray
+    @test x[At(2), At("a")] isa VariableRef
+    @test JuMP.name(x[At(4), At("b")]) == "x[4,b]"
+    @test @expression(model, sum(x[At(i), At("a")] for i in 2:4)) isa AffExpr
+    @constraint(model, c, sum(x[At(i), At("a")] for i in 2:4) <= 1)
+    @test c isa ConstraintRef
+    return
+end
+
+function test_dimension_data_expression()
+    model = Model()
+    B = ["a", "b"]
+    @variable(model, x[i = 2:4, j = B], container = DimArray)
+    @expression(
+        model,
+        expr[j = B],
+        sum(x[At(i), At(j)] for i in 2:4),
+        container = DimArray,
+    )
+    @test expr isa DimArray
+    @test expr[At("a")] isa AffExpr
+    return
+end
+
+function test_dimensional_data_missing_names()
+    model = Model()
+    @test @variable(model, [1:3, 1:2], container = DimArray) isa DimArray
+    @test @variable(model, [i = 1:3, 1:2], container = DimArray) isa DimArray
+    @test @variable(model, [1:3, j = 1:2], container = DimArray) isa DimArray
+    return
+end
+
+function test_dimensional_data_sparse()
+    model = Model()
+    @test_throws(
+        ErrorException(
+            "Unable to create a `DimArray` because the container does not form " *
+            "a dense rectangular array",
+        ),
+        @variable(model, [i = 1:3, i:2], container = DimArray),
+    )
+    @test_throws(
+        ErrorException(
+            "Unable to create a `DimArray` because the container does not form " *
+            "a dense rectangular array",
+        ),
+        @variable(model, [i = 1:3; isodd(i)], container = DimArray),
+    )
+    return
+end
+
+end

test/Kokako.jl

@@ -144,8 +144,12 @@ Calls `include_modules_to_test(dir)` where `dir` is the `/test` directory of the
 package `package`.
 """
 function include_modules_to_test(package::Module)
-    dir = joinpath(dirname(dirname(pathof(package))), "test")
-    return include_modules_to_test(dir)
+    root = dirname(dirname(pathof(package)))
+    modules = include_modules_to_test(joinpath(root, "test"))
+    if VERSION >= v"1.9"
+        append!(modules, include_modules_to_test(joinpath(root, "ext")))
+    end
+    return modules
 end
 
 end  # module