Skip to content

Commit

Permalink
Merge branch 'main' into dgsharp/gracefulFail
Browse files Browse the repository at this point in the history
  • Loading branch information
dannys4 authored Jul 26, 2023
2 parents 7c14dfd + d666310 commit 39b9aca
Show file tree
Hide file tree
Showing 11 changed files with 72 additions and 40 deletions.
6 changes: 3 additions & 3 deletions MParT/AffineFunction.h
Original file line number Diff line number Diff line change
Expand Up @@ -51,9 +51,9 @@ class AffineFunction : public ParameterizedFunctionBase<MemorySpace>


protected:

StridedMatrix<double,MemorySpace> A_;
StridedVector<double,MemorySpace> b_;
Kokkos::View<double**, Kokkos::LayoutLeft, MemorySpace> A_;
Kokkos::View<double*, Kokkos::LayoutLeft, MemorySpace> b_;

int ldA;

Expand Down
5 changes: 3 additions & 2 deletions MParT/AffineMap.h
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
namespace mpart{

/** @brief Defines transformations of the form \f$Ax+b\f$ for an invertible matrix \f$A\f$ and vector offset \f$b\f$.
* Makes a deep copy of any views passed to the constructor.
*/
template<typename MemorySpace>
class AffineMap : public ConditionalMapBase<MemorySpace>
Expand Down Expand Up @@ -64,8 +65,8 @@ class AffineMap : public ConditionalMapBase<MemorySpace>

protected:

StridedMatrix<double,MemorySpace> A_;
StridedVector<double,MemorySpace> b_;
Kokkos::View<double**, Kokkos::LayoutLeft, MemorySpace> A_;
Kokkos::View<double*, Kokkos::LayoutLeft, MemorySpace> b_;

mpart::PartialPivLU<MemorySpace> luSolver_;
double logDet_;
Expand Down
2 changes: 1 addition & 1 deletion bindings/julia/src/JlArrayConversions.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,7 @@ Eigen::Map<const Eigen::Matrix<int,Eigen::Dynamic, Eigen::Dynamic>,0,Eigen::Oute
int* mptr = mat.data();
unsigned int rows = size(mat,0);
unsigned int cols = size(mat,1);
return Eigen::Map<const Eigen::Matrix<int,Eigen::Dynamic, Eigen::Dynamic>,0,Eigen::OuterStride<>>(mptr, rows, cols, Eigen::OuterStride<>(std::max(rows,cols)));
return Eigen::Map<const Eigen::Matrix<int,Eigen::Dynamic, Eigen::Dynamic>,0,Eigen::OuterStride<>>(mptr, rows, cols, Eigen::OuterStride<>(rows));
}

/**
Expand Down
2 changes: 2 additions & 0 deletions bindings/julia/src/MultiIndex.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,8 @@ void mpart::binding::MultiIndexWrapper(jlcxx::Module &mod) {
.method("NumActiveForward", &MultiIndexSet::NumActiveForward)
.method("NumForward", &MultiIndexSet::NumForward)
.method("Size", &MultiIndexSet::Size)
.method("addto!", [](MultiIndexSet &mset, MultiIndex const& idx){ return mset += idx; })
.method("addto!", [](MultiIndexSet &mset, MultiIndexSet const& mset2){ return mset += mset2; })
;

mod.method("MultiIndexSet", [](jlcxx::ArrayRef<int,2> idxs) {
Expand Down
5 changes: 5 additions & 0 deletions bindings/matlab/mat/MultiIndexSet.m
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,11 @@ function delete(this)
multi = MultiIndex(multi_id,"id");
end

function mset = DeepCopy(this)
mset_id = MParT_('MultiIndexSet_DeepCopy', this.id_);
mset = MultiIndeSet(mset_id,"id");
end

function result = MultiToIndex(this,multi)
result= MParT_('MultiIndexSet_MultiToIndex',this.id_,multi.get_id());
result = result + 1;
Expand Down
10 changes: 10 additions & 0 deletions bindings/matlab/src/MultiIndexSet_mex.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -169,6 +169,16 @@ MEX_DEFINE(MultiIndexSet_ExpandAny) (int nlhs, mxArray* plhs[],
output.set(0, mset->Expand());
}

MEX_DEFINE(MultiIndexSet_DeepCopy) (int nlhs, mxArray* plhs[],
int nrhs, const mxArray* prhs[]) {
InputArguments input(nrhs, prhs, 1);
OutputArguments output(nlhs, plhs, 1);
const MultiIndexSet& mset = Session<MultiIndexSet>::getConst(input.get(0));
MultiIndexSet mset_copy = MultiIndexSet::CreateTotalOrder(mset.Length(), 0);
mset_copy += mset;
output.set(0, Session<MultiIndexSet>::create(new MultiIndexSet(mset_copy)));
}

MEX_DEFINE(MultiIndexSet_ForciblyExpand) (int nlhs, mxArray* plhs[],
int nrhs, const mxArray* prhs[]) {
InputArguments input(nrhs, prhs, 2);
Expand Down
12 changes: 6 additions & 6 deletions bindings/python/src/AffineMap.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -13,15 +13,15 @@ using namespace mpart::binding;
void mpart::binding::AffineMapWrapperHost(py::module &m)
{
py::class_<AffineMap<Kokkos::HostSpace>, ConditionalMapBase<Kokkos::HostSpace>, std::shared_ptr<AffineMap<Kokkos::HostSpace>>>(m, "AffineMap")
.def(py::init( [](Eigen::Ref<Eigen::VectorXd> const& b)
.def(py::init( [](py::EigenDRef<Eigen::VectorXd> const& b)
{
return new AffineMap<Kokkos::HostSpace>(VecToKokkos<double, Kokkos::HostSpace>(b));
}))
.def(py::init( [](Eigen::Ref<Eigen::MatrixXd> const& A, Eigen::Ref<Eigen::VectorXd> const& b)
.def(py::init( [](py::EigenDRef<Eigen::MatrixXd> const& A, py::EigenDRef<Eigen::VectorXd> const& b)
{
return new AffineMap<Kokkos::HostSpace>(MatToKokkos<double, Kokkos::HostSpace>(A), VecToKokkos<double, Kokkos::HostSpace>(b));
}))
.def(py::init( [](Eigen::Ref<Eigen::MatrixXd> const& A)
.def(py::init( [](py::EigenDRef<Eigen::MatrixXd> const& A)
{
return new AffineMap<Kokkos::HostSpace>(MatToKokkos<double, Kokkos::HostSpace>(A));
}));
Expand All @@ -31,15 +31,15 @@ void mpart::binding::AffineMapWrapperHost(py::module &m)
void mpart::binding::AffineFunctionWrapperHost(py::module &m)
{
py::class_<AffineFunction<Kokkos::HostSpace>, ParameterizedFunctionBase<Kokkos::HostSpace>, std::shared_ptr<AffineFunction<Kokkos::HostSpace>>>(m, "AffineFunction")
.def(py::init( [](Eigen::Ref<Eigen::VectorXd> const& b)
.def(py::init( [](py::EigenDRef<Eigen::VectorXd> const& b)
{
return new AffineFunction<Kokkos::HostSpace>(VecToKokkos<double, Kokkos::HostSpace>(b));
}))
.def(py::init( [](Eigen::Ref<Eigen::MatrixXd> const& A, Eigen::Ref<Eigen::VectorXd> const& b)
.def(py::init( [](py::EigenDRef<Eigen::MatrixXd> const& A, py::EigenDRef<Eigen::VectorXd> const& b)
{
return new AffineFunction<Kokkos::HostSpace>(MatToKokkos<double, Kokkos::HostSpace>(A), VecToKokkos<double, Kokkos::HostSpace>(b));
}))
.def(py::init( [](Eigen::Ref<Eigen::MatrixXd> const& A)
.def(py::init( [](py::EigenDRef<Eigen::MatrixXd> const& A)
{
return new AffineFunction<Kokkos::HostSpace>(MatToKokkos<double, Kokkos::HostSpace>(A));
}));
Expand Down
12 changes: 9 additions & 3 deletions bindings/python/src/MultiIndex.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -113,10 +113,10 @@ void mpart::binding::MultiIndexWrapper(py::module &m)
.def(py::init<const unsigned int>())
.def(py::init<Eigen::Ref<const Eigen::MatrixXi> const&>())
.def("fix", &MultiIndexSet::Fix)
.def("__len__", &MultiIndexSet::Length)
.def("__len__", &MultiIndexSet::Length, "Retrieves the length of _each_ multiindex within this set (i.e. the dimension of the input)")
.def("__getitem__", &MultiIndexSet::at)
.def("at", &MultiIndexSet::at)
.def("Size", &MultiIndexSet::Size)
.def("Size", &MultiIndexSet::Size, "Retrieves the number of elements in this MultiIndexSet")

.def_static("CreateTotalOrder", &MultiIndexSet::CreateTotalOrder, py::arg("length"), py::arg("maxOrder"), py::arg("limiter")=MultiIndexLimiter::None())
.def_static("CreateTensorProduct", &MultiIndexSet::CreateTensorProduct, py::arg("length"), py::arg("maxOrder"), py::arg("limiter")=MultiIndexLimiter::None())
Expand All @@ -127,9 +127,15 @@ void mpart::binding::MultiIndexWrapper(py::module &m)
.def("IndexToMulti",&MultiIndexSet::IndexToMulti)
.def("MultiToIndex", &MultiIndexSet::MultiToIndex)
.def("MaxOrders", &MultiIndexSet::MaxOrders)
.def("Expand", py::overload_cast<unsigned int>(&MultiIndexSet::Expand))
.def("Expand", py::overload_cast<unsigned int>(&MultiIndexSet::Expand), "Expand frontier w.r.t one multiindex")
.def("Expand", py::overload_cast<>(&MultiIndexSet::Expand), "Expand all frontiers of a MultiIndexSet")
.def("append", py::overload_cast<MultiIndex const&>(&MultiIndexSet::operator+=))
.def("__iadd__", py::overload_cast<MultiIndex const&>(&MultiIndexSet::operator+=))
.def("DeepCopy",[](const MultiIndexSet& mset)
{
MultiIndexSet mset_copy = mset;
return mset_copy;
})
.def("Activate", py::overload_cast<MultiIndex const&>(&MultiIndexSet::Activate))
.def("AddActive", &MultiIndexSet::AddActive)
.def("Frontier", &MultiIndexSet::Frontier)
Expand Down
23 changes: 15 additions & 8 deletions src/AffineFunction.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -11,15 +11,15 @@ using namespace mpart;

template<typename MemorySpace>
AffineFunction<MemorySpace>::AffineFunction(StridedVector<double,MemorySpace> b) : ParameterizedFunctionBase<MemorySpace>(b.size(),b.size(),0),
b_("b",b.layout())
b_("b",b.extent(0))
{

Kokkos::deep_copy(b_, b);
}

template<typename MemorySpace>
AffineFunction<MemorySpace>::AffineFunction(StridedMatrix<double,MemorySpace> A) : ParameterizedFunctionBase<MemorySpace>(A.extent(1),A.extent(0),0),
A_("A", A.layout())
A_("A", A.extent(0), A.extent(1))
{
Kokkos::deep_copy(A_, A);
assert(A_.extent(0)<=A_.extent(1));
Expand All @@ -29,8 +29,8 @@ AffineFunction<MemorySpace>::AffineFunction(StridedMatrix<double,MemorySpace> A)
template<typename MemorySpace>
AffineFunction<MemorySpace>::AffineFunction(StridedMatrix<double,MemorySpace> A,
StridedVector<double,MemorySpace> b) : ParameterizedFunctionBase<MemorySpace>(A.extent(1),A.extent(0),0),
A_("A", A.layout()),
b_("b",b.layout())
A_("A", A.extent(0), A.extent(1)),
b_("b", b.extent(0))
{
Kokkos::deep_copy(A_, A);
Kokkos::deep_copy(b_, b);
Expand All @@ -43,19 +43,26 @@ template<typename MemorySpace>
void AffineFunction<MemorySpace>::EvaluateImpl(StridedMatrix<const double, MemorySpace> const& pts,
StridedMatrix<double, MemorySpace> output)
{
unsigned int numPts = pts.extent(1);
Kokkos::MDRangePolicy<Kokkos::Rank<2>, typename MemoryToExecution<MemorySpace>::Space> policy({{0, 0}}, {{numPts, this->outputDim}});

// Linear part
if(A_.extent(0)>0){

// Initialize output to zeros
Kokkos::parallel_for(policy, KOKKOS_CLASS_LAMBDA(const int& j, const int& i) {
output(i,j) = 0.0;
});

dgemm<MemorySpace>(1.0, A_, pts, 0.0, output);

}else{
Kokkos::deep_copy(output, pts);
}

// Bias part
if(b_.size()>0){

unsigned int numPts = pts.extent(1);
Kokkos::MDRangePolicy<Kokkos::Rank<2>, typename MemoryToExecution<MemorySpace>::Space> policy({0, 0}, {numPts, this->outputDim});

Kokkos::parallel_for(policy, KOKKOS_CLASS_LAMBDA(const int& j, const int& i) {
output(i,j) += b_(i);
});
Expand Down
31 changes: 15 additions & 16 deletions src/AffineMap.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -9,15 +9,15 @@ using namespace mpart;

template<typename MemorySpace>
AffineMap<MemorySpace>::AffineMap(StridedVector<double,MemorySpace> b) : ConditionalMapBase<MemorySpace>(b.size(),b.size(),0),
b_("b",b.layout()),
b_("b", b.extent(0)),
logDet_(0.0)
{
Kokkos::deep_copy(b_, b);
}

template<typename MemorySpace>
AffineMap<MemorySpace>::AffineMap(StridedMatrix<double,MemorySpace> A) : ConditionalMapBase<MemorySpace>(A.extent(1),A.extent(0),0),
A_("A", A.layout())
A_("A", A.extent(0), A.extent(1))
{
Kokkos::deep_copy(A_, A);
assert(A_.extent(0)<=A_.extent(1));
Expand All @@ -28,8 +28,8 @@ AffineMap<MemorySpace>::AffineMap(StridedMatrix<double,MemorySpace> A) : Conditi
template<typename MemorySpace>
AffineMap<MemorySpace>::AffineMap(StridedMatrix<double,MemorySpace> A,
StridedVector<double,MemorySpace> b) : ConditionalMapBase<MemorySpace>(A.extent(1),A.extent(0),0),
A_("A", A.layout()),
b_("b",b.layout())
A_("A", A.extent(0), A.extent(1)),
b_("b", b.extent(0))
{
Kokkos::deep_copy(A_, A);
Kokkos::deep_copy(b_, b);
Expand All @@ -42,13 +42,6 @@ AffineMap<MemorySpace>::AffineMap(StridedMatrix<double,MemorySpace> A,
template<typename MemorySpace>
void AffineMap<MemorySpace>::Factorize(){

// If A_ is not column major, create a column major version
if(A_.stride_0()!=1){
Kokkos::View<double**, Kokkos::LayoutLeft, MemorySpace> anew("A_", A_.extent(0), A_.extent(1));
Kokkos::deep_copy(anew, A_);
A_ = anew;
}

if(A_.extent(0)!=A_.extent(1)){
StridedMatrix<const double, MemorySpace> Asub = Kokkos::subview(A_, Kokkos::ALL(), std::make_pair(A_.extent(1)-A_.extent(0),A_.extent(1)));
luSolver_.compute(Asub);
Expand Down Expand Up @@ -145,19 +138,25 @@ template<typename MemorySpace>
void AffineMap<MemorySpace>::EvaluateImpl(StridedMatrix<const double, MemorySpace> const& pts,
StridedMatrix<double, MemorySpace> output)
{
unsigned int numPts = pts.extent(1);
Kokkos::MDRangePolicy<Kokkos::Rank<2>, typename MemoryToExecution<MemorySpace>::Space> policy({{0, 0}}, {{numPts, this->outputDim}});

// Linear part
if(A_.extent(0)>0){

// Initialize output to zeros
Kokkos::parallel_for(policy, KOKKOS_CLASS_LAMBDA(const int& j, const int& i) {
output(i,j) = 0.0;
});

dgemm<MemorySpace>(1.0, A_, pts, 0.0, output);

}else{
Kokkos::deep_copy(output, pts);
}

// Bias part
if(b_.size()>0){

unsigned int numPts = pts.extent(1);
Kokkos::MDRangePolicy<Kokkos::Rank<2>, typename MemoryToExecution<MemorySpace>::Space> policy({{0, 0}}, {{numPts, this->outputDim}});

Kokkos::parallel_for(policy, KOKKOS_CLASS_LAMBDA(const int& j, const int& i) {
output(i,j) += b_(i);
});
Expand Down
4 changes: 3 additions & 1 deletion src/TrainMap.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,9 @@ nlopt::opt SetupOptimization(unsigned int dim, TrainOptions options) {
template<>
double mpart::TrainMap(std::shared_ptr<ConditionalMapBase<Kokkos::HostSpace>> map, std::shared_ptr<MapObjective<Kokkos::HostSpace>> objective, TrainOptions options) {
if(map->Coeffs().extent(0) == 0) {
std::cout << "TrainMap: Initializing map coeffs to 1." << std::endl;
if(options.verbose) {
std::cout << "TrainMap: Initializing map coeffs to 1." << std::endl;
}
Kokkos::View<double*, Kokkos::HostSpace> coeffs ("Default coeffs", map->numCoeffs);
Kokkos::parallel_for("Setting default coeff val", map->numCoeffs, KOKKOS_LAMBDA(const unsigned int i){
coeffs(i) = 1.;
Expand Down

0 comments on commit 39b9aca

Please sign in to comment.