Merge pull request #302 from epatters/cset-indexing-interface

Indexing notation for C-sets

src/categorical_algebra/CSetDataStructures.jl

@@ -3,7 +3,7 @@
 module CSetDataStructures
 export AbstractACSet, ACSet, AbstractCSet, CSet, Schema, FreeSchema,
   AbstractACSetType, ACSetType, ACSetTableType, AbstractCSetType, CSetType,
-  tables, nparts, has_part, subpart, has_subpart, incident,
+  tables, parts, nparts, has_part, subpart, has_subpart, incident,
   add_part!, add_parts!, set_subpart!, set_subparts!, rem_part!, rem_parts!,
   copy_parts!, copy_parts_only!, disjoint_union
 
@@ -255,7 +255,7 @@ function Base.show(io::IO, ::MIME"text/plain", acs::T) where {T<:AbstractACSet}
       # TODO: Set option `row_number_column_title=name` when next version of
       # PrettyTables is released, instead of making new table.
       cols = map(col -> replace_unassigned(col, "#undef"), fieldarrays(table))
-      table = StructArray((; ob => 1:nparts(acs,ob), cols...))
+      table = StructArray((; ob => parts(acs, ob), cols...))
       pretty_table(io, table, nosubheader=true)
     end
   end
@@ -273,7 +273,7 @@ function Base.show(io::IO, ::MIME"text/html", acs::T) where {T<:AbstractACSet}
     if !(eltype(table) <: EmptyTuple || isempty(table))
       # TODO: Set option `row_number_column_title`. See above.
       cols = map(col -> replace_unassigned(col, "#undef"), fieldarrays(table))
-      table = StructArray((; ob => 1:nparts(acs,ob), cols...))
+      table = StructArray((; ob => parts(acs, ob), cols...))
       pretty_table(io, table, backend=:html, standalone=false, nosubheader=true)
     end
   end
@@ -298,6 +298,10 @@ directly mutate these tables, especially when indexing is enabled!
 """
 tables(acs::ACSet) = acs.tables
 
+""" Parts of given type in a C-set.
+"""
+parts(acs::ACSet, type) = 1:nparts(acs, type)
+
 """ Number of parts of given type in a C-set.
 """
 nparts(acs::ACSet, type) = length(acs.tables[type])
@@ -325,13 +329,24 @@ end
 """ Get subpart of part in C-set.
 
 Both single and vectorized access are supported. Chaining, or composition, of
-parts is also supported. For example, given a vertex-attributed graph `g`,
+subparts is also supported. For example, given a vertex-attributed graph `g`,
 
 ```
 subpart(g, e, [:src, :vattr])
 ```
 
-returns the vertex attribute of the source vertex of the edge `e`.
+returns the vertex attribute of the source vertex of the edge `e`. As a
+shorthand, subparts can also be accessed by indexing:
+
+```
+g[e, :src] == subpart(g, e, :src)
+```
+
+Be warned that indexing with lists of subparts works as above:
+`g[e,[:src,:vattr]]` is equivalent to `subpart(g, e, [:src,:vattr])`. This
+differs from DataFrames but note that the alternative interpretation of
+`[:src,:vattr]` as two independent columns does not even make sense, since they
+have different domains (belong to different tables).
 """
 subpart(acs::ACSet, part, name::Symbol) = subpart(acs,name)[part]
 subpart(acs::ACSet, name::Symbol) = _subpart(acs,Val(name))
@@ -354,10 +369,12 @@ subpart_name(expr::GATExpr{:compose}) = mapreduce(subpart_name, vcat, args(expr)
   elseif name ∈ AD.attr
     :(acs.tables.$(dom(AD,name)).$name)
   else
-    throw(KeyError(name))
+    throw(ArgumentError("$(repr(name)) not in $(CD.hom) or $(AD.attr)"))
   end
 end
 
+Base.getindex(acs::ACSet, args...) = subpart(acs, args...)
+
 """ Get superparts incident to part in C-set.
 
 If the subpart is indexed, this takes constant time; otherwise, it takes linear
@@ -410,7 +427,7 @@ incident(acs::ACSet, part, expr::GATExpr; kw...) =
       :(broadcast_findall(part, acs.tables.$(dom(AD,name)).$name))
     end
   else
-    throw(KeyError(name))
+    throw(ArgumentError("$(repr(name)) not in $(CD.hom) or $(AD.attr)"))
   end
 end
 
@@ -521,7 +538,7 @@ set_subpart!(acs::ACSet, name::Symbol, new_subpart) =
       :(acs.tables.$ob.$name[part] = subpart)
     end
   else
-    throw(KeyError(name))
+    throw(ArgumentError("$(repr(name)) not in $(CD.hom) or $(AD.attr)"))
   end
 end
 

src/categorical_algebra/CSets.jl

@@ -182,7 +182,7 @@ function colimit(diagram::AbstractFreeDiagram{ACS}) where
     T = Ts.parameters[d]
     data = Vector{Union{Some{T},Nothing}}(nothing, nparts(Y, c))
     for (ι, X) in zip(ιs, Xs)
-      for i in 1:nparts(X, c)
+      for i in parts(X, c)
         j = ι[c](i)
         if isnothing(data[j])
           data[j] = Some(subpart(X, i, attr))

src/graphs/BasicGraphs.jl

@@ -51,8 +51,8 @@ src(g::AbstractACSet, args...) = subpart(g, args..., :src)
 tgt(g::AbstractACSet, args...) = subpart(g, args..., :tgt)
 dst(g::AbstractACSet, args...) = tgt(g, args...) # LightGraphs compatibility
 
-vertices(g::AbstractACSet) = 1:nv(g)
-edges(g::AbstractACSet) = 1:ne(g)
+vertices(g::AbstractACSet) = parts(g, :V)
+edges(g::AbstractACSet) = parts(g, :E)
 edges(g::AbstractACSet, src::Int, tgt::Int) =
   (e for e in incident(g, src, :src) if subpart(g, e, :tgt) == tgt)
 
@@ -279,7 +279,7 @@ end
 
 vertex(g::AbstractACSet, args...) = subpart(g, args..., :vertex)
 
-half_edges(g::AbstractACSet) = 1:nparts(g, :H)
+half_edges(g::AbstractACSet) = parts(g, :H)
 half_edges(g::AbstractACSet, v) = incident(g, v, :vertex)
 
 function half_edge_pairs(g::AbstractACSet, src::Int, tgt::Int)

src/wiring_diagrams/Undirected.jl

@@ -81,13 +81,12 @@ end
 
 nboxes(d::AbstractUWD) = nparts(d, :Box)
 njunctions(d::AbstractUWD) = nparts(d, :Junction)
-boxes(d::AbstractUWD) = 1:nboxes(d)
-junctions(d::AbstractUWD) = 1:njunctions(d)
+boxes(d::AbstractUWD) = parts(d, :Box)
+junctions(d::AbstractUWD) = parts(d, :Junction)
 
-ports(d::AbstractUWD; outer::Bool=false) =
-  1:nparts(d, outer ? :OuterPort : :Port)
+ports(d::AbstractUWD; outer::Bool=false) = parts(d, outer ? :OuterPort : :Port)
 ports(d::AbstractUWD, box) =
-  box == outer_box(d) ? (1:nparts(d, :OuterPort)) : incident(d, box, :box)
+  box == outer_box(d) ? parts(d, :OuterPort) : incident(d, box, :box)
 ports_with_junction(d::AbstractUWD, junction; outer::Bool=false) =
   incident(d, junction, outer ? :outer_junction : :junction)
 

test/categorical_algebra/CSetDataStructures.jl

@@ -44,8 +44,9 @@ set_subpart!(dds, 1, :Φ, 1)
 
 @test has_subpart(dds, :Φ)
 @test !has_subpart(dds, :nonsubpart)
-@test_throws KeyError subpart(dds, 1, :nonsubpart)
-@test_throws KeyError set_subpart!(dds, 1, :nonsubpart, 1)
+@test_throws ArgumentError subpart(dds, 1, :nonsubpart)
+@test_throws ArgumentError incident(dds, 1, :nonsuppart)
+@test_throws ArgumentError set_subpart!(dds, 1, :nonsubpart, 1)
 
 # Deletion.
 dds = DDS()
@@ -158,6 +159,11 @@ X, parent, height = TheoryDendrogram[[:X, :parent, :height]]
 @test subpart(d, 3, id(X)) == 3
 @test incident(d, 10, compose(parent, height)) == [1,2,3]
 
+# Indexing syntax.
+@test d[3, :parent] == 4
+@test d[3, [:parent, :height]] == 10
+@test d[:, :parent] == [4,4,4,5,5]
+
 # Copying parts.
 d2 = Dendrogram{Int}()
 copy_parts!(d2, d, X=[4,5])