From 2e14121c84ebc0cb2c93415d034bcbb7aa88e831 Mon Sep 17 00:00:00 2001
From: fverdugo <fverdugo@cimne.upc.edu>
Date: Mon, 22 Jul 2019 15:06:03 +0200
Subject: [PATCH] Added QGradMonomialBasis

---
 Manifest.toml                   |   4 +-
 src/QGradMonomialBases.jl       | 168 ++++++++++++++++++++++++++++++++
 src/TensorPolynomialBases.jl    |  10 ++
 test/QGradMonomialBasesTests.jl |  52 ++++++++++
 test/runtests.jl                |   6 +-
 5 files changed, 237 insertions(+), 3 deletions(-)
 create mode 100644 src/QGradMonomialBases.jl
 create mode 100644 test/QGradMonomialBasesTests.jl

diff --git a/Manifest.toml b/Manifest.toml
index f65b653..1f32272 100644
--- a/Manifest.toml
+++ b/Manifest.toml
@@ -51,9 +51,9 @@ uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [[TensorValues]]
 deps = ["LinearAlgebra", "StaticArrays"]
-git-tree-sha1 = "057e34fa9adeaa725391c311ac0ec22c1e429a9d"
+git-tree-sha1 = "9d192922681bdc4f3294b5d565d2d4b37d21fa18"
 uuid = "31c64edf-cdeb-50e4-845d-ed2334360c20"
-version = "0.3.1"
+version = "0.3.3"
 
 [[Test]]
 deps = ["Distributed", "InteractiveUtils", "Logging", "Random"]
diff --git a/src/QGradMonomialBases.jl b/src/QGradMonomialBases.jl
new file mode 100644
index 0000000..1a7d5ee
--- /dev/null
+++ b/src/QGradMonomialBases.jl
@@ -0,0 +1,168 @@
+
+# Types and constructors
+
+struct QGradMonomialBasis{P,G,D} <: TensorPolynomialBasis{P,P,G}
+  order::Int
+  terms::CartesianIndices{D}
+  perms::Matrix{Int}
+end
+
+function (::Type{QGradMonomialBasis{P}})(order::Int) where P
+  D = _length(P)
+  G = _gradient_type(P,P)
+  n1d = order+1
+  t = fill(n1d,D)
+  t[1] = order
+  terms = CartesianIndices(tuple(t...))
+  perms = _prepare_perms(D)
+  QGradMonomialBasis{P,G,D}(order,terms,perms)
+end
+
+# Implementation of the interface
+
+length(b::QGradMonomialBasis{P,G,D}) where {P,G,D} = D*b.order*(b.order+1)^(D-1)
+
+ndims(b::QGradMonomialBasis{P,G,D}) where {P,G,D} = D
+
+function ScratchData(b::QGradMonomialBasis{P}) where P
+  V = P
+  T = eltype(V)
+  dim = _length(P)
+  n1d = b.order+1
+  c = zeros(dim,n1d)
+  g = zeros(dim,n1d)
+  MonomialBasisCache{T}(c,g)
+end
+
+function evaluate!(
+  v::AbstractVector{P},
+  b::QGradMonomialBasis{P,G,D},
+  x::P,
+  cache::MonomialBasisCache) where {P,G,D}
+  _evaluate_nd_qgrad!(v,x,b.order,b.terms,b.perms,cache.c)
+end
+
+function gradient!(
+  v::AbstractVector{G},
+  b::QGradMonomialBasis{P,G,D},
+  x::P,
+  cache::MonomialBasisCache) where {P,G,D}
+  V = P
+  _gradient_nd_qgrad!(v,x,b.order,b.terms,b.perms,cache.c,cache.g,V)
+end
+
+# Helpers
+
+function _prepare_perms(D)
+  perms = zeros(Int,D,D)
+  for j in 1:D
+    for d in j:D
+      perms[d,j] =  d-j+1
+    end
+    for d in 1:(j-1)
+      perms[d,j] =  d+(D-j)+1
+    end
+  end
+  perms
+end
+
+function _evaluate_nd_qgrad!(
+  v::AbstractVector{V},
+  x,
+  order,
+  terms::CartesianIndices{D},
+  perms::Matrix{Int},
+  c::AbstractMatrix{T}) where {V,T,D}
+
+  dim = D
+  for d in 1:dim
+    _evaluate_1d!(c,x,order,d)
+  end
+
+  o = one(T)
+  k = 1
+  m = zero(_mutable(V))
+  js = eachindex(m)
+  z = zero(T)
+
+  for ci in terms
+
+    for j in js
+
+      @inbounds for i in js
+        m[i] = z
+      end
+
+      s = o
+      @inbounds for d in 1:dim
+        s *= c[d,ci[perms[d,j]]]
+      end
+
+      m[j] = s
+      v[k] = m
+      k += 1
+
+    end
+
+  end
+
+end
+
+function _gradient_nd_qgrad!(
+  v::AbstractVector{G},
+  x,
+  order,
+  terms::CartesianIndices{D},
+  perms::Matrix{Int},
+  c::AbstractMatrix{T},
+  g::AbstractMatrix{T},
+  ::Type{V}) where {G,T,D,V}
+
+  dim = D
+  for d in 1:dim
+    _evaluate_1d!(c,x,order,d)
+    _gradient_1d!(g,x,order,d)
+  end
+
+  z = zero(_mutable(V))
+  m = zero(_mutable(G))
+  js = eachindex(z)
+  mjs = eachindex(m)
+  o = one(T)
+  zi = zero(T)
+  k = 1
+
+  for ci in terms
+
+    for j in js
+
+      s = z
+      for i in js
+        s[i] = o
+      end
+
+      for q in 1:dim
+        for d in 1:dim
+          if d != q
+            s[q] *= c[d,ci[perms[d,j]]]
+          else
+            s[q] *= g[d,ci[perms[d,j]]]
+          end
+        end
+      end
+
+      @inbounds for i in mjs
+        m[i] = zi
+      end
+
+      for i in js
+        m[i,j] = s[i]
+      end
+      v[k] = m
+      k += 1
+
+    end
+
+  end
+
+end
diff --git a/src/TensorPolynomialBases.jl b/src/TensorPolynomialBases.jl
index 1c9f0ef..33d6cff 100644
--- a/src/TensorPolynomialBases.jl
+++ b/src/TensorPolynomialBases.jl
@@ -4,9 +4,17 @@ using StaticArrays
 using TensorValues
 using Test
 
+# Move to VectorValues (begin)
+import Base: convert
+function convert(::Type{<:MultiValue{S,T,N,L}},a::NTuple{L,T}) where {S,T,N,L}
+  MultiValue(SArray{S,T}(a))
+end
+# Move to VectorValues (end)
+
 export ScratchData
 export TensorPolynomialBasis
 export MonomialBasis
+export QGradMonomialBasis
 export test_polynomial_basis_without_gradient
 export test_polynomial_basis
 export gradient_type, value_type, point_type
@@ -20,4 +28,6 @@ include("Interfaces.jl")
 
 include("MonomialBases.jl")
 
+include("QGradMonomialBases.jl")
+
 end # module
diff --git a/test/QGradMonomialBasesTests.jl b/test/QGradMonomialBasesTests.jl
new file mode 100644
index 0000000..2c6b4ce
--- /dev/null
+++ b/test/QGradMonomialBasesTests.jl
@@ -0,0 +1,52 @@
+module QGradMonomialBasesTests
+
+using Test
+using TensorPolynomialBases
+using TensorValues
+
+P = VectorValue{3,Float64}
+order = 1
+basis = QGradMonomialBasis{P}(order)
+G = gradient_type(basis)
+
+x = VectorValue(2.0,3.0,5.0)
+v = P[
+  (1.0, 0.0, 0.0), (0.0, 1.0, 0.0), (0.0, 0.0, 1.0),
+  (3.0, 0.0, 0.0), (0.0, 5.0, 0.0), (0.0, 0.0, 2.0),
+  (5.0, 0.0, 0.0), (0.0, 2.0, 0.0), (0.0, 0.0, 3.0),
+  (15.0, 0.0, 0.0), (0.0, 10.0, 0.0), (0.0, 0.0, 6.0)]
+g = G[
+  (0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0),
+  (0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0),
+  (0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0),
+  (0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0),
+  (0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0),
+  (0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0),
+  (0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0),
+  (0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0),
+  (0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0),
+  (0.0, 5.0, 3.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0),
+  (0.0, 0.0, 0.0, 5.0, 0.0, 2.0, 0.0, 0.0, 0.0),
+  (0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 3.0, 2.0, 0.0)]
+test_polynomial_basis(basis,x,v,g)
+
+P = VectorValue{2,Float64}
+order = 2
+basis = QGradMonomialBasis{P}(order)
+G = gradient_type(basis)
+
+x = VectorValue(2.0,3.0)
+
+v = P[
+  (1.0, 0.0), (0.0, 1.0), (2.0, 0.0), (0.0, 3.0),
+  (3.0, 0.0), (0.0, 2.0), (6.0, 0.0), (0.0, 6.0),
+  (9.0, 0.0), (0.0, 4.0), (18.0, 0.0), (0.0, 12.0)]
+
+g = G[
+  (0.0, 0.0, 0.0, 0.0), (0.0, 0.0, 0.0, 0.0), (1.0, 0.0, 0.0, 0.0),
+  (0.0, 0.0, 0.0, 1.0), (0.0, 1.0, 0.0, 0.0), (0.0, 0.0, 1.0, 0.0),
+  (3.0, 2.0, 0.0, 0.0), (0.0, 0.0, 3.0, 2.0), (0.0, 6.0, 0.0, 0.0),
+  (0.0, 0.0, 4.0, 0.0), (9.0, 12.0, 0.0, 0.0), (0.0, 0.0, 12.0, 4.0)]
+test_polynomial_basis(basis,x,v,g)
+
+end # module
diff --git a/test/runtests.jl b/test/runtests.jl
index 73b03cc..37edea2 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -15,4 +15,8 @@ end
   include("MonomialBasesTests.jl")
 end
 
-end # module TensorPolynomialBasesTests
+@testset "QGradMonomialBasesTests" begin
+  include("QGradMonomialBasesTests.jl")
+end
+
+end # module