Merge pull request #18 from bowenszhu/PolyChaos-create-var

Change to use `Symbolics.variables` to create an array of symbolic vars
FHoltorf · Oct 28, 2022 · f6931d4 · f6931d4
2 parents 96b9082 + bf1addb
commit f6931d4
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 14 deletions.
diff --git a/src/PCE/PCE.jl b/src/PCE/PCE.jl
@@ -60,9 +60,13 @@ function PCE(states, uni_basis::AbstractVector{<:Pair})
     moments = Vector{Num}[]
     for (i, state) in enumerate(collect(states))
         ind_vars = get_independent_vars(state)
-        create_var = isempty(ind_vars) ? name -> (@variables $(name))[1] :
-                     name -> (@variables $(name)(ind_vars...))[1]
-        push!(moments, [create_var(moment_name(i, j)) for j in 1:n_basis])
+        new_vars = if isempty(ind_vars)
+            Symbolics.variables(:z, i:i, 1:n_basis)
+        else
+            map(z -> z(ind_vars...),
+                Symbolics.variables(:z, i:i, 1:n_basis; T = Symbolics.FnType))
+        end
+        push!(moments, vec(new_vars))
     end
     ansatz = [states[i] => sum(moments[i][j] * sym_basis[j] for j in 1:n_basis)
               for i in 1:n_states]

diff --git a/src/PCE/PCE_utils.jl b/src/PCE/PCE_utils.jl
@@ -1,11 +1,5 @@
 import PolyChaos: computeSP2, computeSP, dim, deg
 
-# moment names
-function moment_name(i, j)
-    return Symbol("z" * Symbolics.map_subscripts(i) *
-                  "₋" * Symbolics.map_subscripts(j))
-end
-
 # getting independent variables
 function get_independent_vars(var)
     return []
@@ -109,7 +103,7 @@ end
 
 """
 `TensorProductOrthoPoly` objects represent bases formed as the tensor product of univariate `PolyChaos.AbstractOrthoPoly` bases.
-By default the basis elements of the tensor product are restricted to polynomials with total degree up to the maximum degree among the 
+By default the basis elements of the tensor product are restricted to polynomials with total degree up to the maximum degree among the
 univariate bases. This maximum degree can be manually specified, however.
 """
 struct TensorProductOrthoPoly{M, U}
@@ -149,8 +143,8 @@ end
 """
 $(TYPEDSIGNATURES)
 
-computes inner product between basis functions of a `TensorProductOrthoPoly` via 
-`PolyChaos`'s infrastructure (exploiting the tensor product form). 
+computes inner product between basis functions of a `TensorProductOrthoPoly` via
+`PolyChaos`'s infrastructure (exploiting the tensor product form).
 """
 function computeSP(basis_fxns, tpop::TensorProductOrthoPoly,
                    integrators = tpop.uni)
@@ -195,8 +189,8 @@ end
 """
 $(TYPEDSIGNATURES)
 
-Compute the an ascending list of `n`-dimensional multi-indices with fixed `grade` (= sum of entries) 
-in graded reverse lexicographic order. Constraints on the degrees considered can be incorporated. 
+Compute the an ascending list of `n`-dimensional multi-indices with fixed `grade` (= sum of entries)
+in graded reverse lexicographic order. Constraints on the degrees considered can be incorporated.
 """
 function grevlex(n::Int, grade::Int)
     if n == 1