Merge pull request #290 from MeasureTransport/282-construct-triangula…

…rmap-with-coefficients-from-its-given-components

Adding KeepCoeffs option to TriangularMap and ComposedMap

MParT/ComposedMap.h

@@ -34,10 +34,10 @@ class ComposedMap : public ConditionalMapBase<MemorySpace>
     /** @brief Construct a block triangular map from a collection of other ConditionalMapBase objects.
 
     @param maps A vector of ConditionalMapBase objects defining each \f$T_k\f$ in the composition. Note: each map must be square (i.e., have equal input and output dimension).
+    @param moveCoeffs Whether you want this ComposedMap object to be responsible for the coefficients already in maps. If true, the new object will take ownership of all coefficient vectors within maps. If false, the new object will not have any coefficients set.
     @param maxChecks The maximum number of checkpoints to use during gradient computations.  If maxChecks==1, then no checkpointing will be utilized and all forward states will be recomputed.  If maxChecks==components.size(), then all states will be stored and reused during the backward pass.  This is the most efficient option, but can require an intractable amount of memory for high-dimensional or deep parameterizations.   The default value is -1, which will set the maximum number of checkpoints to be equal to the number of layers (i.e., map.size()).
     */
-    ComposedMap(std::vector<std::shared_ptr<ConditionalMapBase<MemorySpace>>> const& maps,
-                int maxChecks=-1);
+    ComposedMap(std::vector<std::shared_ptr<ConditionalMapBase<MemorySpace>>> const& maps, bool moveCoeffs=false, int maxChecks=-1);
 
     virtual ~ComposedMap() = default;
 
@@ -51,8 +51,10 @@ class ComposedMap : public ConditionalMapBase<MemorySpace>
     */
     using ConditionalMapBase<MemorySpace>::SetCoeffs;
     void SetCoeffs(Kokkos::View<double*, Kokkos::HostSpace> coeffs) override;
+    void WrapCoeffs(Kokkos::View<double*, Kokkos::HostSpace> coeffs) override;
     #if defined(MPART_ENABLE_GPU)
     void SetCoeffs(Kokkos::View<double*, Kokkos::DefaultExecutionSpace::memory_space> coeffs) override;
+    void WrapCoeffs(Kokkos::View<double*, mpart::DeviceSpace> coeffs) override;
     #endif
 
     virtual std::shared_ptr<ConditionalMapBase<MemorySpace>> GetComponent(unsigned int i){ return maps_.at(i);}
@@ -144,4 +146,4 @@ class ComposedMap : public ConditionalMapBase<MemorySpace>
 
 }
 
-#endif
+#endif

MParT/ParameterizedFunctionBase.h

@@ -173,7 +173,10 @@ namespace mpart {
                                    StridedMatrix<const double, MemorySpace> const& sens,
                                    StridedMatrix<double, MemorySpace>              output) = 0;
 
+        /** Checks to see if the coefficients have been initialized yet, returns true if so, false if not */
+        bool CheckCoefficients() const;
 
+
         const unsigned int inputDim; /// The total dimension of the input N+M
         const unsigned int outputDim; /// The output dimension M
         const unsigned int numCoeffs; /// The number of coefficients used to parameterize this map.
@@ -194,6 +197,8 @@ namespace mpart {
          */
         void CheckDeviceMismatch(std::string functionName) const;
 
+
+
         /** Checks to see if the coefficients have been initialized yet. If not, an exception is thrown. */
         void CheckCoefficients(std::string const& functionName) const;
 

MParT/TriangularMap.h

@@ -36,8 +36,11 @@ class TriangularMap : public ConditionalMapBase<MemorySpace>{
 
     @param components A vector of ConditionalMapBase objects defining each \f$T_k\f$ in the block triangular map.
                       To maintain the correct block structure, the dimensions of the components must satisfy \f$N_k = N_{k-1}+M_{k}\f$.
+    @param moveCoeffs A boolean to choose whether to keep the coefficients from the maps in `components` or make new ones. If true,
+                      the new object takes ownership of all the coefficients vectors within all the maps in `components` (changing the coefficients
+                      in the new map will then change the coefficients in the original maps). If false, no coefficients are copied or created.
     */
-    TriangularMap(std::vector<std::shared_ptr<ConditionalMapBase<MemorySpace>>> const& components);
+    TriangularMap(std::vector<std::shared_ptr<ConditionalMapBase<MemorySpace>>> const& components, bool moveCoeffs=false);
 
     virtual ~TriangularMap() = default;
 
@@ -133,4 +136,4 @@ class TriangularMap : public ConditionalMapBase<MemorySpace>{
 
 }
 
-#endif
+#endif

bindings/python/src/ComposedMap.cpp

@@ -18,7 +18,7 @@ void mpart::binding::ComposedMapWrapper(py::module &m)
 
     // ComposedMap
     py::class_<ComposedMap<MemorySpace>, ConditionalMapBase<MemorySpace>, std::shared_ptr<ComposedMap<MemorySpace>>>(m, tName.c_str())
-        .def(py::init<std::vector<std::shared_ptr<ConditionalMapBase<MemorySpace>>>,int>())
+        .def(py::init<std::vector<std::shared_ptr<ConditionalMapBase<MemorySpace>>>,bool,int>(), py::arg("maps"), py::arg("moveCoeffs") = false, py::arg("maxChecks")=-1)
         ;
 
 }

bindings/python/src/TriangularMap.cpp

@@ -18,7 +18,7 @@ void mpart::binding::TriangularMapWrapper(py::module &m)
 
     // TriangularMap
     py::class_<TriangularMap<MemorySpace>, ConditionalMapBase<MemorySpace>, std::shared_ptr<TriangularMap<MemorySpace>>>(m, tName.c_str())
-        .def(py::init<std::vector<std::shared_ptr<ConditionalMapBase<MemorySpace>>>>())
+        .def(py::init<std::vector<std::shared_ptr<ConditionalMapBase<MemorySpace>>>, bool>(), py::arg("comps"), py::arg("moveCoeffs") = false)
         .def("GetComponent", &TriangularMap<MemorySpace>::GetComponent)
         ;
 

bindings/python/tests/test_ComposedMap.py

@@ -35,6 +35,11 @@ def test_CoeffsMap():
 
 def test_Evaluate():
     assert composed_map.Evaluate(x).shape == (dim, num_samples)
+    x_ = x.copy()
+    for map_ in maps:
+        x_ = map_.Evaluate(x_)
+    assert np.all(composed_map.Evaluate(x) == x_)
+
 
 
 def test_LogDeterminant():

bindings/python/tests/test_ComposedMap_moveCoeffs.py

@@ -0,0 +1,52 @@
+# Test methods of ComposedMap object
+import mpart
+import numpy as np
+print("made map")
+opts = mpart.MapOptions()
+
+num_maps = 4
+dim = 2
+maxOrder = 1
+
+maps = []
+coeffs_ = []
+numCoeffs = 0
+for i in range(num_maps):
+    map = mpart.CreateTriangular(dim,dim,maxOrder,opts)
+    coeffs = np.random.randn(map.numCoeffs)
+    map.SetCoeffs(coeffs)
+    maps.append(map)
+    coeffs_.append(coeffs)
+    numCoeffs += map.numCoeffs
+
+
+moveCoeffs = True
+composed_map = mpart.ComposedMap(maps, moveCoeffs, num_maps)
+all_coeffs = np.hstack(coeffs_).flatten()
+
+num_samples = 100
+x = np.random.randn(dim, num_samples)
+
+def test_numCoeffs():
+    assert composed_map.numCoeffs == numCoeffs
+
+def test_CoeffsMap():
+    assert np.all(composed_map.CoeffMap() == all_coeffs)
+
+
+def test_Evaluate():
+    assert composed_map.Evaluate(x).shape == (dim, num_samples)
+    x_ = x.copy()
+    for map_ in maps:
+        x_ = map_.Evaluate(x_)
+    assert np.all(composed_map.Evaluate(x) == x_)
+
+
+def test_LogDeterminant():
+    assert composed_map.LogDeterminant(x).shape == (num_samples,)
+
+def test_Inverse():
+
+    y = composed_map.Evaluate(x)
+    x_ = composed_map.Inverse(np.zeros((1,num_samples)),y)
+    assert np.allclose(x_, x, atol=1E-3)

bindings/python/tests/test_TriangularMap.py

@@ -16,15 +16,16 @@
 triangular = mpart.TriangularMap([map_1,map_2])
 num_samples = 100
 x = np.random.randn(2,num_samples)
+coeffs = np.random.randn(triangular.numCoeffs)
+triangular.SetCoeffs(coeffs)
 
 
 def test_numCoeffs():
     assert triangular.numCoeffs == 2 + 3
 
 
 def test_CoeffsMap():
-    coeffs = np.random.randn(triangular.numCoeffs)
-    triangular.SetCoeffs(coeffs)
+
     assert np.all(triangular.CoeffMap() == coeffs)
 
 
@@ -36,8 +37,7 @@ def test_LogDeterminant():
     assert triangular.LogDeterminant(x).shape == (num_samples,)
 
 def test_Inverse():
-    coeffs = np.random.randn(triangular.numCoeffs)
-    triangular.SetCoeffs(coeffs)
+
     y = triangular.Evaluate(x)
     x_ = triangular.Inverse(np.zeros((1,num_samples)),y)
     assert np.allclose(x_, x, atol=1E-3)

bindings/python/tests/test_TriangularMap_moveCoeffs.py

@@ -0,0 +1,42 @@
+# Test methods of TriangularMap object
+import mpart
+import numpy as np
+
+opts = mpart.MapOptions()
+
+multis_1 = np.array([[0],[1]])  # linear
+multis_2 = np.array([[0,0],[0,1],[2,0]])  # quadratic in x_1, linear in x_2, matches form of target 
+
+mset_1 = mpart.MultiIndexSet(multis_1).fix(True)
+mset_2 = mpart.MultiIndexSet(multis_2).fix(True)
+
+map_1 = mpart.CreateComponent(mset_1, opts)
+coeffs_1 = np.random.randn(map_1.numCoeffs)
+map_1.SetCoeffs(coeffs_1)
+
+map_2 = mpart.CreateComponent(mset_2, opts)
+coeffs_2 = np.random.randn(map_2.numCoeffs)
+map_2.SetCoeffs(coeffs_2)
+
+triangular = mpart.TriangularMap([map_1,map_2], moveCoeffs=True)
+num_samples = 100
+x = np.random.randn(2,num_samples)
+
+
+def test_numCoeffs():
+    assert triangular.numCoeffs == 2 + 3
+
+
+def test_Evaluate():
+    assert triangular.Evaluate(x).shape == (2,num_samples)
+
+
+def test_LogDeterminant():
+    assert triangular.LogDeterminant(x).shape == (num_samples,)
+
+def test_Inverse():
+    # coeffs = np.random.randn(triangular.numCoeffs)
+    # triangular.SetCoeffs(coeffs)
+    y = triangular.Evaluate(x)
+    x_ = triangular.Inverse(np.zeros((1,num_samples)),y)
+    assert np.allclose(x_, x, atol=1E-3)

src/ComposedMap.cpp

@@ -132,7 +132,7 @@ Kokkos::View<double**, Kokkos::LayoutLeft, MemorySpace> ComposedMap<MemorySpace>
 
 
 template<typename MemorySpace>
-ComposedMap<MemorySpace>::ComposedMap(std::vector<std::shared_ptr<ConditionalMapBase<MemorySpace>>> const& maps,
+ComposedMap<MemorySpace>::ComposedMap(std::vector<std::shared_ptr<ConditionalMapBase<MemorySpace>>> const& maps, bool moveCoeffs,
                                       int maxChecks) : ConditionalMapBase<MemorySpace>(maps.front()->inputDim, 
                                                                                        maps.front()->inputDim,
                                                                                        std::accumulate(maps.begin(), maps.end(), 0, [](size_t sum, std::shared_ptr<ConditionalMapBase<MemorySpace>> const& comp){ return sum + comp->numCoeffs; })),
@@ -153,6 +153,28 @@ ComposedMap<MemorySpace>::ComposedMap(std::vector<std::shared_ptr<ConditionalMap
         msg << "In ComposedMap constructor, each map in the composition must be square. Output dimension (" << maps_.at(maps_.size()-1)->outputDim << ") of component " << maps_.size()-1 << " is not equal to the input dimension (" << maps_.at(maps_.size()-1)->inputDim << ").";
         throw std::invalid_argument(msg.str());
     }
+
+    // if moveCoeffs is set to true, we check if each map's coeffs are set, and then copy them into the new composed map's coeffs
+    if(moveCoeffs){
+
+        Kokkos::View<double*,MemorySpace> coeffs("coeffs", this->numCoeffs);
+        unsigned int cumNumCoeffs = 0;
+
+        for(unsigned int i=0; i<maps_.size(); ++i){
+
+            if(!maps_.at(i)->CheckCoefficients()){
+                std::stringstream msg;
+                msg << "In ComposedMap constructor, moveCoeffs set to true, but map " << i <<" doesn't have coeffs set";
+                throw std::invalid_argument(msg.str());
+            }
+
+            Kokkos::View<double*,MemorySpace> subCoeffs = Kokkos::subview(coeffs, std::make_pair(cumNumCoeffs, cumNumCoeffs+maps_.at(i)->numCoeffs));
+            Kokkos::deep_copy(subCoeffs, maps_.at(i)->Coeffs());
+            cumNumCoeffs += maps_.at(i)->numCoeffs;
+        }
+
+        this->WrapCoeffs(coeffs);
+    }
 }
 
 template<typename MemorySpace>
@@ -172,6 +194,23 @@ void ComposedMap<MemorySpace>::SetCoeffs(Kokkos::View<double*, Kokkos::HostSpace
     }
 }
 
+template<typename MemorySpace>
+void ComposedMap<MemorySpace>::WrapCoeffs(Kokkos::View<double*, Kokkos::HostSpace> coeffs)
+{
+    // First, call the ConditionalMapBase version of this function to copy the view into the savedCoeffs member variable
+    ConditionalMapBase<MemorySpace>::WrapCoeffs(coeffs);
+
+    // Now create subviews for each of the components
+    unsigned int cumNumCoeffs = 0;
+    for(unsigned int i=0; i<maps_.size(); ++i){
+        assert(cumNumCoeffs+maps_.at(i)->numCoeffs <= this->savedCoeffs.size());
+
+        maps_.at(i)->WrapCoeffs( Kokkos::subview(this->savedCoeffs,
+            std::make_pair(cumNumCoeffs, cumNumCoeffs+maps_.at(i)->numCoeffs)));
+        cumNumCoeffs += maps_.at(i)->numCoeffs;
+    }
+}
+
 #if defined(MPART_ENABLE_GPU)
 template<typename MemorySpace>
 void ComposedMap<MemorySpace>::SetCoeffs(Kokkos::View<double*, Kokkos::DefaultExecutionSpace::memory_space> coeffs)

src/ParameterizedFunctionBase.cpp

@@ -340,12 +340,11 @@ Eigen::RowMatrixXd ParameterizedFunctionBase<mpart::DeviceSpace>::CoeffGrad(Eige
 
 
 #endif
-
 template<typename MemorySpace>
-void ParameterizedFunctionBase<MemorySpace>::CheckCoefficients(std::string const& functionName) const
+bool ParameterizedFunctionBase<MemorySpace>::CheckCoefficients() const
 {
     if(this->numCoeffs==0)
-        return;
+        return true;
 
     bool good = true;
 
@@ -355,6 +354,17 @@ void ParameterizedFunctionBase<MemorySpace>::CheckCoefficients(std::string const
         good = false;
     }
 
+    return good;
+
+}
+
+
+template<typename MemorySpace>
+void ParameterizedFunctionBase<MemorySpace>::CheckCoefficients(std::string const& functionName) const
+{
+
+    bool good = CheckCoefficients();
+
     if(!good){
         std::stringstream msg;
         msg << "Error in \"" << functionName << "\", the coefficients have not been set yet.  Make sure to call SetCoeffs() before calling this function.";

src/TriangularMap.cpp

@@ -7,19 +7,22 @@
 using namespace mpart;
 
 template<typename MemorySpace>
-TriangularMap<MemorySpace>::TriangularMap(std::vector<std::shared_ptr<ConditionalMapBase<MemorySpace>>> const& components) : ConditionalMapBase<MemorySpace>(components.back()->inputDim,
+TriangularMap<MemorySpace>::TriangularMap(std::vector<std::shared_ptr<ConditionalMapBase<MemorySpace>>> const& components, bool moveCoeffs) : ConditionalMapBase<MemorySpace>(components.back()->inputDim,
                         std::accumulate(components.begin(), components.end(), 0, [](size_t sum, std::shared_ptr<ConditionalMapBase<MemorySpace>> const& comp){ return sum + comp->outputDim; }),
                         std::accumulate(components.begin(), components.end(), 0, [](size_t sum, std::shared_ptr<ConditionalMapBase<MemorySpace>> const& comp){ return sum + comp->numCoeffs; })),
                         comps_(components)
 {
+
 
     // Check the sizes of all the inputs
+
     for(unsigned int i=0; i<comps_.size(); ++i){
         if(comps_.at(i)->outputDim > comps_.at(i)->inputDim){
             std::stringstream msg;
             msg << "In TriangularMap constructor, the output dimension (" << comps_.at(i)->outputDim << ") of component " << i << " is greater than the input dimension (" << comps_.at(i)->inputDim << ").";
             throw std::invalid_argument(msg.str());
         }
+
     }
 
     for(unsigned int i=1; i<comps_.size(); ++i){
@@ -32,8 +35,33 @@ TriangularMap<MemorySpace>::TriangularMap(std::vector<std::shared_ptr<Conditiona
             throw std::invalid_argument(msg.str());
         }
     }
+
+
+    // if moveCoeffs is set to true, we check if each component's coeffs are set, and then copy them into the new triangular map's coeffs
+    if(moveCoeffs){
+
+        Kokkos::View<double*,MemorySpace> coeffs("coeffs", this->numCoeffs);
+        unsigned int cumNumCoeffs = 0;
+
+        for(unsigned int i=0; i<comps_.size(); ++i){
+
+            if(!comps_.at(i)->CheckCoefficients()){
+                std::stringstream msg;
+                msg << "In TriangularMap constructor, moveCoeffs set to true, but component " << i <<" doesn't have coeffs set";
+                throw std::invalid_argument(msg.str());
+            }
+
+            Kokkos::View<double*,MemorySpace> subCoeffs = Kokkos::subview(coeffs, std::make_pair(cumNumCoeffs, cumNumCoeffs+comps_.at(i)->numCoeffs));
+            Kokkos::deep_copy(subCoeffs, comps_.at(i)->Coeffs());
+            cumNumCoeffs += comps_.at(i)->numCoeffs;
+        }
+
+        this->WrapCoeffs(coeffs);
+    }
 }
 
+
+
 template<typename MemorySpace>
 void TriangularMap<MemorySpace>::SetCoeffs(Kokkos::View<double*, Kokkos::HostSpace> coeffs)
 {