|
| 1 | +""" |
| 2 | +``` |
| 3 | +struct VarName{sym} |
| 4 | + indexing :: String |
| 5 | +end |
| 6 | +``` |
| 7 | +
|
| 8 | +A variable identifier. Every variable has a symbol `sym` and `indices `indexing`. |
| 9 | +The Julia variable in the model corresponding to `sym` can refer to a single value or |
| 10 | +to a hierarchical array structure of univariate, multivariate or matrix variables. `indexing` stores the indices that can access the random variable from the Julia |
| 11 | +variable. |
| 12 | +
|
| 13 | +Examples: |
| 14 | +
|
| 15 | +- `x[1] ~ Normal()` will generate a `VarName` with `sym == :x` and `indexing == "[1]"`. |
| 16 | +- `x[:,1] ~ MvNormal(zeros(2))` will generate a `VarName` with `sym == :x` and |
| 17 | + `indexing == "[Colon(),1]"`. |
| 18 | +- `x[:,1][2] ~ Normal()` will generate a `VarName` with `sym == :x` and |
| 19 | + `indexing == "[Colon(),1][2]"`. |
| 20 | +""" |
| 21 | +struct VarName{sym} |
| 22 | + indexing::String |
| 23 | +end |
| 24 | + |
| 25 | + |
| 26 | + |
| 27 | +""" |
| 28 | + @varname(var) |
| 29 | +
|
| 30 | +A macro that returns an instance of `VarName` given the symbol or expression of a Julia variable, e.g. `@varname x[1,2][1+5][45][3]` returns `VarName{:x}("[1,2][6][45][3]")`. |
| 31 | +""" |
| 32 | +macro varname(expr::Union{Expr, Symbol}) |
| 33 | + expr |> varname |> esc |
| 34 | +end |
| 35 | +function varname(expr) |
| 36 | + ex = deepcopy(expr) |
| 37 | + (ex isa Symbol) && return quote |
| 38 | + DynamicPPL.VarName{$(QuoteNode(ex))}("") |
| 39 | + end |
| 40 | + (ex.head == :ref) || throw("VarName: Mis-formed variable name $(expr)!") |
| 41 | + inds = :(()) |
| 42 | + while ex.head == :ref |
| 43 | + if length(ex.args) >= 2 |
| 44 | + strs = map(x -> :($x === (:) ? "Colon()" : string($x)), ex.args[2:end]) |
| 45 | + pushfirst!(inds.args, :("[" * join($(Expr(:vect, strs...)), ",") * "]")) |
| 46 | + end |
| 47 | + ex = ex.args[1] |
| 48 | + isa(ex, Symbol) && return quote |
| 49 | + DynamicPPL.VarName{$(QuoteNode(ex))}(foldl(*, $inds, init = "")) |
| 50 | + end |
| 51 | + end |
| 52 | + throw("VarName: Mis-formed variable name $(expr)!") |
| 53 | +end |
| 54 | + |
| 55 | +macro vsym(expr::Union{Expr, Symbol}) |
| 56 | + expr |> vsym |
| 57 | +end |
| 58 | + |
| 59 | +""" |
| 60 | + vsym(expr::Union{Expr, Symbol}) |
| 61 | +
|
| 62 | +Returns the variable symbol given the input variable expression `expr`. For example, if the input `expr = :(x[1])`, the output is `:x`. |
| 63 | +""" |
| 64 | +function vsym(expr::Union{Expr, Symbol}) |
| 65 | + ex = deepcopy(expr) |
| 66 | + (ex isa Symbol) && return QuoteNode(ex) |
| 67 | + (ex.head == :ref) || throw("VarName: Mis-formed variable name $(expr)!") |
| 68 | + while ex.head == :ref |
| 69 | + ex = ex.args[1] |
| 70 | + isa(ex, Symbol) && return QuoteNode(ex) |
| 71 | + end |
| 72 | + throw("VarName: Mis-formed variable name $(expr)!") |
| 73 | +end |
| 74 | + |
| 75 | +""" |
| 76 | + @vinds(expr) |
| 77 | +
|
| 78 | +Returns a tuple of tuples of the indices in `expr`. For example, `@vinds x[1,:][2]` returns |
| 79 | +`((1, Colon()), (2,))`. |
| 80 | +""" |
| 81 | +macro vinds(expr::Union{Expr, Symbol}) |
| 82 | + expr |> vinds |> esc |
| 83 | +end |
| 84 | +function vinds(expr::Union{Expr, Symbol}) |
| 85 | + ex = deepcopy(expr) |
| 86 | + inds = Expr(:tuple) |
| 87 | + (ex isa Symbol) && return inds |
| 88 | + (ex.head == :ref) || throw("VarName: Mis-formed variable name $(expr)!") |
| 89 | + while ex.head == :ref |
| 90 | + pushfirst!(inds.args, Expr(:tuple, ex.args[2:end]...)) |
| 91 | + ex = ex.args[1] |
| 92 | + isa(ex, Symbol) && return inds |
| 93 | + end |
| 94 | + throw("VarName: Mis-formed variable name $(expr)!") |
| 95 | +end |
| 96 | + |
| 97 | +""" |
| 98 | + split_var_str(var_str, inds_as = Vector) |
| 99 | +
|
| 100 | +This function splits a variable string, e.g. `"x[1:3,1:2][3,2]"` to the variable's symbol `"x"` and the indexing `"[1:3,1:2][3,2]"`. If `inds_as = String`, the indices are returned as a string, e.g. `"[1:3,1:2][3,2]"`. If `inds_as = Vector`, the indices are returned as a vector of vectors of strings, e.g. `[["1:3", "1:2"], ["3", "2"]]`. |
| 101 | +""" |
| 102 | +function split_var_str(var_str, inds_as = Vector) |
| 103 | + ind = findfirst(c -> c == '[', var_str) |
| 104 | + if inds_as === String |
| 105 | + if ind === nothing |
| 106 | + return var_str, "" |
| 107 | + else |
| 108 | + return var_str[1:ind-1], var_str[ind:end] |
| 109 | + end |
| 110 | + end |
| 111 | + @assert inds_as === Vector |
| 112 | + inds = Vector{String}[] |
| 113 | + if ind === nothing |
| 114 | + return var_str, inds |
| 115 | + end |
| 116 | + sym = var_str[1:ind-1] |
| 117 | + ind = length(sym) |
| 118 | + while ind < length(var_str) |
| 119 | + ind += 1 |
| 120 | + @assert var_str[ind] == '[' |
| 121 | + push!(inds, String[]) |
| 122 | + while var_str[ind] != ']' |
| 123 | + ind += 1 |
| 124 | + if var_str[ind] == '[' |
| 125 | + ind2 = findnext(c -> c == ']', var_str, ind) |
| 126 | + push!(inds[end], strip(var_str[ind:ind2])) |
| 127 | + ind = ind2+1 |
| 128 | + else |
| 129 | + ind2 = findnext(c -> c == ',' || c == ']', var_str, ind) |
| 130 | + push!(inds[end], strip(var_str[ind:ind2-1])) |
| 131 | + ind = ind2 |
| 132 | + end |
| 133 | + end |
| 134 | + end |
| 135 | + return sym, inds |
| 136 | +end |
0 commit comments