T operator()(RawArray<const T,2> x) const {
// Temporary arrays and views
GEODE_ASSERT(x.sizes()==vec(n+3,d));
// Collect quadrature points
Array<T,2> tq(n,quads,uninit);
Array<T,3> xq(n,quads,d,uninit);
Array<T,3> vq(n,quads,d,uninit);
for (int i=0;i<n;i++) {
T_INFO(i)
for (int q=0;q<quads;q++) {
const T s = samples[q];
tq(i,q) = t1+dt*s;
SPLINE_INFO(s)
for (int a=0;a<d;a++) {
X_INFO(i,a)
xq(i,q,a) = a0*x0+a1*x1+a2*x2+a3*x3;
vq(i,q,a) = b0*x0+b1*x1+b2*x2+b3*x3;
}
}
}
// Compute energy
const auto Uq_ = U(tq.reshape_own(n*quads),NdArray<const T>(qshape,xq.flat),NdArray<const T>(qshape,vq.flat));
GEODE_ASSERT(Uq_.size()==n*quads);
const auto Uq = Uq_.reshape(n,quads);
// Accumulate
T sum = 0;
for (int i=0;i<n;i++) {
const T dt = t[i+2]-t[i+1];
for (int q=0;q<quads;q++)
sum += weights[q]*dt*Uq(i,q);
}
return sum;
}
void hessian(RawArray<const T,2> x, RawArray<T,4> hess) const {
// Temporary arrays and views
GEODE_ASSERT(x.sizes()==vec(n+3,d) && hess.sizes()==vec(n+3,4,d,d));
const auto sx = smallx.flat.raw(),
sv = smallv.flat.raw();
// Collect quadrature points
const int e = 1+8*d+8*d*(d-1);
Array<T,3> tq( n,quads,e,uninit);
Array<T,4> xq(vec(n,quads,e,d),uninit);
Array<T,4> vq(vec(n,quads,e,d),uninit);
for (int i=0;i<n;i++) {
T_INFO(i)
for (int q=0;q<quads;q++) {
const T s = samples[q],
t = t1+dt*s;
for (int j=0;j<e;j++)
tq(i,q,j) = t;
SPLINE_INFO(s)
for (int a=0;a<d;a++) {
X_INFO(i,a)
const T x = a0*x0+a1*x1+a2*x2+a3*x3,
v = b0*x0+b1*x1+b2*x2+b3*x3;
for (int j=0;j<e;j++) {
xq(i,q,j,a) = x;
vq(i,q,j,a) = v;
}
int j = 1;
for (int b=0;b<d;b++) {
const T xb = sx[b],
vb = sv[b];
xq(i,q,j++,a) -= xb;
xq(i,q,j++,a) += xb;
vq(i,q,j++,a) -= vb;
vq(i,q,j++,a) += vb;
xq(i,q,j ,a) -= xb;
vq(i,q,j++,a) -= vb;
xq(i,q,j ,a) -= xb;
vq(i,q,j++,a) += vb;
xq(i,q,j ,a) += xb;
vq(i,q,j++,a) -= vb;
xq(i,q,j ,a) += xb;
vq(i,q,j++,a) += vb;
for (int c=b+1;c<d;c++) {
const T xc = sx[c],
vc = sv[c];
xq(i,q,j++,a) -= xb+xc;
xq(i,q,j++,a) -= xb-xc;
xq(i,q,j++,a) += xb-xc;
xq(i,q,j++,a) += xb+xc;
vq(i,q,j++,a) -= vb+vc;
vq(i,q,j++,a) -= vb-vc;
vq(i,q,j++,a) += vb-vc;
vq(i,q,j++,a) += vb+vc;
vq(i,q,j++,a) -= sv[b];
xq(i,q,j ,a) -= sx[b];
vq(i,q,j++,a) += sv[b];
xq(i,q,j ,a) += sx[b];
vq(i,q,j++,a) -= sv[b];
xq(i,q,j ,a) += sx[b];
vq(i,q,j++,a) += sv[b];
}
}
}
}
}
// Compute energies
const auto Uq_ = U(tq.reshape_own(n*quads*d4),NdArray<const T>(q2shape,xq.flat),NdArray<const T>(q2shape,vq.flat));
GEODE_ASSERT(Uq_.size()==n*quads*d4);
const auto Uq = Uq_.reshape(n,quads,d4);
// Accumulate
grad.fill(0);
const auto inv_2s = GEODE_RAW_ALLOCA(d,Vector<T,2>);
for (int a=0;a<d;a++)
inv_2s[a] = vec(.5/sx[a],.5/sv[a]);
for (int i=0;i<n;i++) {
T_INFO(i)
for (int q=0;q<quads;q++) {
const T s = samples[q],
w = dt*weights[q];
SPLINE_INFO(s)
for (int b=0;b<d;b++) {
const T wx = w*inv_2s[b].x*(Uq(i,q,4*b+1)-Uq(i,q,4*b )),
wv = w*inv_2s[b].y*(Uq(i,q,4*b+3)-Uq(i,q,4*b+2));
grad(i ,b) += a0*wx+b0*wv;
grad(i+1,b) += a1*wx+b1*wv;
grad(i+2,b) += a2*wx+b2*wv;
grad(i+3,b) += a3*wx+b3*wv;
}
}
}
}
void gradient(RawArray<const T,2> x, RawArray<T,2> grad) const {
// Temporary arrays and views
GEODE_ASSERT(x.sizes()==vec(n+3,d) && grad.sizes()==x.sizes());
const auto sx = smallx.flat.raw(),
sv = smallv.flat.raw();
// Collect quadrature points
const int e = 4*d;
Array<T,3> tq( n,quads,e,uninit);
Array<T,4> xq(vec(n,quads,e,d),uninit);
Array<T,4> vq(vec(n,quads,e,d),uninit);
for (int i=0;i<n;i++) {
T_INFO(i)
for (int q=0;q<quads;q++) {
const T s = samples[q],
t = t1+dt*s;
for (int j=0;j<e;j++)
tq(i,q,j) = t;
SPLINE_INFO(s)
for (int a=0;a<d;a++) {
X_INFO(i,a)
const T x = a0*x0+a1*x1+a2*x2+a3*x3,
v = b0*x0+b1*x1+b2*x2+b3*x3;
for (int j=0;j<e;j++) {
xq(i,q,j,a) = x;
vq(i,q,j,a) = v;
}
}
for (int a=0;a<d;a++) {
xq(i,q,4*a ,a) -= sx[a];
xq(i,q,4*a+1,a) += sx[a];
vq(i,q,4*a+2,a) -= sv[a];
vq(i,q,4*a+3,a) += sv[a];
}
}
}
// Compute energies
const auto Uq_ = U(tq.reshape_own(n*quads*e),NdArray<const T>(q2shape,xq.flat),NdArray<const T>(q2shape,vq.flat));
GEODE_ASSERT(Uq_.size()==n*quads*e);
const auto Uq = Uq_.reshape(n,quads,e);
// Accumulate
grad.fill(0);
const auto inv_2s = GEODE_RAW_ALLOCA(d,Vector<T,2>);
for (int a=0;a<d;a++)
inv_2s[a] = vec(.5/sx[a],.5/sv[a]);
for (int i=0;i<n;i++) {
T_INFO(i)
for (int q=0;q<quads;q++) {
const T s = samples[q],
w = dt*weights[q];
SPLINE_INFO(s)
for (int a=0;a<d;a++) {
const T wx = w*inv_2s[a].x*(Uq(i,q,4*a+1)-Uq(i,q,4*a )),
wv = w*inv_2s[a].y*(Uq(i,q,4*a+3)-Uq(i,q,4*a+2));
grad(i ,a) += a0*wx+b0*wv;
grad(i+1,a) += a1*wx+b1*wv;
grad(i+2,a) += a2*wx+b2*wv;
grad(i+3,a) += a3*wx+b3*wv;
}
}
}
}
