Skip to content

Commit

Permalink
Merge pull request #226 from MeasureTransport/full-jacobian
Browse files Browse the repository at this point in the history 
Functionality for computing off-diagonal derivatives
  • Loading branch information
@michael-c-brennan
michael-c-brennan authored Sep 7, 2022
2 parents ceea9d4 + e53288f commit ffcf2b4
Show file tree
Hide file tree
Showing 21 changed files with 901 additions and 203 deletions.
4 changes: 3 additions & 1 deletion MParT/DerivativeFlags.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,9 @@ namespace DerivativeFlags{
Parameters, //<- Deriv wrt coeffs
Diagonal, //<- first derivative wrt diagonal
Diagonal2, //<- second derivative wrt diagonal
Mixed //<- gradient wrt coeffs of first derivative wrt x_d
Mixed, //<- gradient wrt coeffs of first derivative wrt x_d
Input, //<- gradient wrt map input
MixedInput //<- gradient of diagonal wrt map input
};

}
Expand Down
205 changes: 153 additions & 52 deletions MParT/MonotoneComponent.h
View file

Large diffs are not rendered by default.

135 changes: 74 additions & 61 deletions MParT/MonotoneIntegrand.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ class MonotoneIntegrand{

/**
@param cache A pointer to memory storing evaluations of \phi_{i,p}(x_i) for each i. These terms need
to be evaluated outside this class (e.g., using `_expansion.FillCache1` for \f$i\in\{0,\ldots,D-1\}\f$.
to be evaluated outside this class (e.g., using `expansion_.FillCache1` for \f$i\in\{0,\ldots,D-1\}\f$.
@param expansion
@param xd
@param coeffs
Expand All @@ -57,44 +57,45 @@ class MonotoneIntegrand{
}

KOKKOS_INLINE_FUNCTION MonotoneIntegrand(double* cache,
ExpansionType const& expansion,
PointType const& pt,
double xd,
CoeffsType const& coeffs,
DerivativeFlags::DerivativeType derivType) : _dim(pt.extent(0)),
_cache(cache),
_expansion(expansion),
_pt(pt),
_xd(xd),
_coeffs(coeffs),
_derivType(derivType)
ExpansionType const& expansion,
PointType const& pt,
double xd,
CoeffsType const& coeffs,
DerivativeFlags::DerivativeType derivType) : dim_(pt.extent(0)),
cache_(cache),
expansion_(expansion),
pt_(pt),
xd_(xd),
coeffs_(coeffs),
derivType_(derivType)
{
assert(derivType!=DerivativeFlags::Mixed);
assert(derivType!=DerivativeFlags::MixedInput);
}

KOKKOS_INLINE_FUNCTION MonotoneIntegrand(double* cache,
ExpansionType const& expansion,
PointType const& pt,
CoeffsType const& coeffs,
DerivativeFlags::DerivativeType derivType,
Kokkos::View<double*, MemorySpace> workspace) : MonotoneIntegrand(cache, expansion, pt, pt(pt.extent(0)-1), coeffs, derivType, workspace)
ExpansionType const& expansion,
PointType const& pt,
CoeffsType const& coeffs,
DerivativeFlags::DerivativeType derivType,
Kokkos::View<double*, MemorySpace> workspace) : MonotoneIntegrand(cache, expansion, pt, pt(pt.extent(0)-1), coeffs, derivType, workspace)
{
}

KOKKOS_INLINE_FUNCTION MonotoneIntegrand(double* cache,
ExpansionType const& expansion,
PointType const& pt,
double xd,
CoeffsType const& coeffs,
DerivativeFlags::DerivativeType derivType,
Kokkos::View<double*, MemorySpace> workspace) : _dim(pt.extent(0)),
_cache(cache),
_expansion(expansion),
_pt(pt),
_xd(xd),
_coeffs(coeffs),
_derivType(derivType),
_workspace(workspace)
ExpansionType const& expansion,
PointType const& pt,
double xd,
CoeffsType const& coeffs,
DerivativeFlags::DerivativeType derivType,
Kokkos::View<double*, MemorySpace> workspace) : dim_(pt.extent(0)),
cache_(cache),
expansion_(expansion),
pt_(pt),
xd_(xd),
coeffs_(coeffs),
derivType_(derivType),
_workspace(workspace)
{
if(derivType==DerivativeFlags::Mixed)
assert(workspace.extent(0)>=coeffs.extent(0));
Expand All @@ -109,58 +110,69 @@ class MonotoneIntegrand{
*/
KOKKOS_INLINE_FUNCTION void operator()(double t, double* output) const
{
const unsigned int numTerms = _expansion.NumCoeffs();
const unsigned int numTerms = expansion_.NumCoeffs();
const unsigned int dim = pt_.size();

unsigned int numOutputs = 1;
if(_derivType==DerivativeFlags::Diagonal)
if(derivType_==DerivativeFlags::Diagonal)
numOutputs++;
if((_derivType==DerivativeFlags::Parameters) || (_derivType==DerivativeFlags::Mixed))
if((derivType_==DerivativeFlags::Parameters) || (derivType_==DerivativeFlags::Mixed))
numOutputs += numTerms;
if((derivType_==DerivativeFlags::Input) || (derivType_==DerivativeFlags::MixedInput))
numOutputs += dim;

// Finish filling in the cache at the quadrature point (FillCache1 is called outside this class)
if((_derivType==DerivativeFlags::Diagonal)||(_derivType==DerivativeFlags::Mixed)){
_expansion.FillCache2(_cache, _pt, t*_xd, DerivativeFlags::Diagonal2);
if((derivType_==DerivativeFlags::Diagonal)||(derivType_==DerivativeFlags::Mixed)||(derivType_==DerivativeFlags::Input)){
expansion_.FillCache2(cache_, pt_, t*xd_, DerivativeFlags::Diagonal2);
}else{
_expansion.FillCache2(_cache, _pt, t*_xd, DerivativeFlags::Diagonal);
expansion_.FillCache2(cache_, pt_, t*xd_, DerivativeFlags::Diagonal);
}

// Use the cache to evaluate \partial_d f and, optionally, the gradient of \partial_d f wrt the coefficients.
// Use the cache to evaluate \partial_d f and, optionally, the gradient of \partial_d f wrt the coefficients or input.
double df = 0;
if(_derivType==DerivativeFlags::Parameters){
if(derivType_==DerivativeFlags::Parameters){
Kokkos::View<double*,MemorySpace,Kokkos::MemoryTraits<Kokkos::Unmanaged>> gradSeg(&output[1], numTerms);
df = _expansion.MixedDerivative(_cache, _coeffs, 1, gradSeg);
df = expansion_.MixedCoeffDerivative(cache_, coeffs_, 1, gradSeg);

double scale = _xd*PosFuncType::Derivative(df);
double scale = xd_*PosFuncType::Derivative(df);
for(unsigned int i=0; i<numTerms;++i)
gradSeg(i) *= scale;

}else if(_derivType==DerivativeFlags::Mixed){
}else if(derivType_==DerivativeFlags::Mixed){

df = _expansion.DiagonalDerivative(_cache, _coeffs, 1);
df = expansion_.DiagonalDerivative(cache_, coeffs_, 1);

double dgdf = PosFuncType::Derivative(df);
double df2 = _expansion.MixedDerivative(_cache, _coeffs, 2, _workspace);
double df2 = expansion_.MixedCoeffDerivative(cache_, coeffs_, 2, _workspace);

double scale = _xd* t * dgdf;
double scale = xd_* t * dgdf;
for(unsigned int i=0; i<numTerms; ++i)
_workspace(i) *= scale;

Kokkos::View<double*,MemorySpace,Kokkos::MemoryTraits<Kokkos::Unmanaged>> gradSeg(&output[1], numTerms);
df = _expansion.MixedDerivative(_cache, _coeffs, 1, gradSeg);
df = expansion_.MixedCoeffDerivative(cache_, coeffs_, 1, gradSeg);

scale = _xd*t*df2*PosFuncType::SecondDerivative(df) + dgdf;
scale = xd_*t*df2*PosFuncType::SecondDerivative(df) + dgdf;
for(unsigned int i=0; i<numTerms; ++i)
gradSeg(i) = scale*gradSeg(i) + _workspace(i);

}else if(derivType_==DerivativeFlags::Input){

Kokkos::View<double*,MemorySpace,Kokkos::MemoryTraits<Kokkos::Unmanaged>> gradSeg(&output[1], dim);
df = expansion_.MixedInputDerivative(cache_, coeffs_, gradSeg);

double scale = xd_*PosFuncType::Derivative(df);
for(unsigned int i=0; i<dim-1;++i)
gradSeg(i) *= scale;

}else{
df = _expansion.DiagonalDerivative(_cache, _coeffs, 1);
df = expansion_.DiagonalDerivative(cache_, coeffs_, 1);
}

// First output is always the integrand itself
double gf = PosFuncType::Evaluate(df);
output[0] = _xd*gf;

output[0] = xd_*gf;
// Check for infs or nans
if(std::isinf(gf)){
printf("\nERROR: In MonotoneIntegrand, value of g(df(...)) is inf. The value of df(...) is %0.4f, and the value of f(df(...)) is %0.4f.\n\n", df, gf);
Expand All @@ -169,26 +181,27 @@ class MonotoneIntegrand{
}

// Compute the derivative with respect to x_d
if(_derivType==DerivativeFlags::Diagonal){

if((derivType_==DerivativeFlags::Diagonal) || (derivType_==DerivativeFlags::Input)){

unsigned int ind = (derivType_==DerivativeFlags::Diagonal) ? 1 : dim;
// Compute \partial^2_d f
output[1] = _expansion.DiagonalDerivative(_cache, _coeffs, 2);
output[ind] = expansion_.DiagonalDerivative(cache_, coeffs_, 2);

// Use the chain rule to get \partial_d g(f)
output[1] *= _xd*t*PosFuncType::Derivative(df);
output[1] += gf;
output[ind] *= xd_*t*PosFuncType::Derivative(df);
output[ind] += gf;
}
}

private:

const unsigned int _dim;
double* _cache;
ExpansionType const& _expansion;
PointType const& _pt;
double _xd;
CoeffsType const& _coeffs;
DerivativeFlags::DerivativeType _derivType;
const unsigned int dim_;
double* cache_;
ExpansionType const& expansion_;
PointType const& pt_;
double xd_;
CoeffsType const& coeffs_;
DerivativeFlags::DerivativeType derivType_;
Kokkos::View<double*,MemorySpace> _workspace;

}; // class MonotoneIntegrand
Expand Down
62 changes: 61 additions & 1 deletion MParT/MultivariateExpansion.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -93,7 +93,67 @@ namespace mpart{
Kokkos::fence();
}

void CoeffGradImpl(StridedMatrix<const double, MemorySpace> const& pts,
virtual void GradientImpl(StridedMatrix<const double, MemorySpace> const& pts,
StridedMatrix<const double, MemorySpace> const& sens,
StridedMatrix<double, MemorySpace> output) override
{
using ExecutionSpace = typename MemoryToExecution<MemorySpace>::Space;

const unsigned int numPts = pts.extent(1);
const unsigned int inDim = pts.extent(0);

// Figure out how much memory we'll need in the cache
unsigned int cacheSize = worker.CacheSize();

auto functor = KOKKOS_CLASS_LAMBDA (typename Kokkos::TeamPolicy<ExecutionSpace>::member_type team_member) {

unsigned int ptInd = team_member.league_rank () * team_member.team_size () + team_member.team_rank ();

if(ptInd<numPts){

// Create a subview containing only the current point
auto pt = Kokkos::subview(pts, Kokkos::ALL(), ptInd);

// Get a pointer to the shared memory that Kokkos set up for this team
Kokkos::View<double*,MemorySpace> cache(team_member.thread_scratch(1), cacheSize);
Kokkos::View<double*,MemorySpace> grad(team_member.thread_scratch(1), inDim);

// Fill in entries in the cache that are independent of x_d. By passing DerivativeFlags::None, we are telling the expansion that no derivatives with wrt x_1,...x_{d-1} will be needed.
worker.FillCache1(cache.data(), pt, DerivativeFlags::Input);
worker.FillCache2(cache.data(), pt, pt(pt.size()-1), DerivativeFlags::Input);

unsigned int coeffStartInd = 0;
for(unsigned int i=0; i<inDim; ++i)
output(i, ptInd) = 0.0;

for(unsigned int d=0; d<this->outputDim; ++d){

// Extract the coefficients for this output dimension
auto coeffs = Kokkos::subview(this->savedCoeffs, std::make_pair(coeffStartInd, coeffStartInd+worker.NumCoeffs()));

// Evaluate the expansion
worker.InputDerivative(cache.data(), coeffs, grad);

for(unsigned int i=0; i<inDim; ++i)
output(i, ptInd) += sens(d,ptInd) * grad(i);

coeffStartInd += worker.NumCoeffs();
}
}
};


auto cacheBytes = Kokkos::View<double*,MemorySpace>::shmem_size(cacheSize + inDim);

// Paralel loop over each point computing T(x_1,...,x_D) for that point
auto policy = GetCachedRangePolicy<ExecutionSpace>(numPts, cacheBytes, functor);
Kokkos::parallel_for(policy, functor);

Kokkos::fence();
}


virtual void CoeffGradImpl(StridedMatrix<const double, MemorySpace> const& pts,
StridedMatrix<const double, MemorySpace> const& sens,
StridedMatrix<double, MemorySpace> output) override
{
Expand Down
Loading

0 comments on commit ffcf2b4

Please sign in to comment.