diff --git a/include/GooseFEM/MatrixDiagonal.h b/include/GooseFEM/MatrixDiagonal.h
index 033c208..ae1e41f 100644
--- a/include/GooseFEM/MatrixDiagonal.h
+++ b/include/GooseFEM/MatrixDiagonal.h
@@ -1,178 +1,179 @@
 /**
 Diagonal matrix.
 
 \file MatrixDiagonal.h
 \copyright Copyright 2017. Tom de Geus. All rights reserved.
 \license This project is released under the GNU Public License (GPLv3).
 */
 
 #ifndef GOOSEFEM_MATRIXDIAGONAL_H
 #define GOOSEFEM_MATRIXDIAGONAL_H
 
 #include "config.h"
+#include "Element.h"
 
 namespace GooseFEM {
 
 /**
 Diagonal matrix.
 
 See Vector() for bookkeeping definitions.
 */
 class MatrixDiagonal {
 public:
 
     MatrixDiagonal() = default;
 
     virtual ~MatrixDiagonal() = default;
 
     /**
     Constructor.
 
     \param conn connectivity [#nelem, #nne].
     \param dofs DOFs per node [#nnode, #ndim].
     */
     MatrixDiagonal(const xt::xtensor<size_t, 2>& conn, const xt::xtensor<size_t, 2>& dofs);
 
     /**
     \return Number of elements.
     */
     size_t nelem() const;
 
     /**
     \return Number of nodes per element.
     */
     size_t nne() const;
 
     /**
     \return Number of nodes.
     */
     size_t nnode() const;
 
     /**
     \return Number of dimensions.
     */
     size_t ndim() const;
 
     /**
     \return Number of DOFs.
     */
     size_t ndof() const;
 
     /**
     \return DOFs per node [#nnode, #ndim]
     */
     xt::xtensor<size_t, 2> dofs() const;
 
     /**
     Assemble from matrices stored per element.
     \warning Ignores any off-diagonal terms.
 
     \param elemmat [#nelem, #nne * #ndim, #nne * #ndim].
     */
     virtual void assemble(const xt::xtensor<double, 3>& elemmat);
 
     /**
     Set all (diagonal) matrix components.
 
     \param A The matrix [#ndof].
     */
-    void set(const xt::xtensor<double, 1>& A);
+    virtual void set(const xt::xtensor<double, 1>& A);
 
     /**
     Return matrix as diagonal matrix.
 
     \param [#ndof].
     */
     virtual xt::xtensor<double, 1> Todiagonal() const;
 
     /**
     Dot-product \f$ b_i = A_{ij} x_j \f$.
 
     \param x nodevec [#nelem, #ndim].
     \return b nodevec overwritten [#nelem, #ndim].
     */
     xt::xtensor<double, 2> Dot(const xt::xtensor<double, 2>& x) const;
 
     /**
     Same as Dot(const xt::xtensor<double, 2>&, xt::xtensor<double, 2>& b)
     but writing to preallocated data.
 
     \param x nodevec [#nelem, #ndim].
     \param b nodevec overwritten [#nelem, #ndim].
     */
     virtual void dot(const xt::xtensor<double, 2>& x, xt::xtensor<double, 2>& b) const;
 
     /**
     Dot-product \f$ b_i = A_{ij} x_j \f$.
 
     \param x dofval [#ndof].
     \return b dofval overwritten [#ndof].
     */
     xt::xtensor<double, 1> Dot(const xt::xtensor<double, 1>& x) const;
 
     /**
     Same as Dot(const xt::xtensor<double, 1>&, xt::xtensor<double, 1>& b)
     but writing to preallocated data.
 
     \param x dofval [#ndof].
     \param b dofval overwritten [#ndof].
     */
     virtual void dot(const xt::xtensor<double, 1>& x, xt::xtensor<double, 1>& b) const;
 
     /**
     Solve \f$ x = A^{-1} b \f$.
 
     \param b nodevec [nelem, ndim].
     \return x nodevec [nelem, ndim].
     */
     xt::xtensor<double, 2> Solve(const xt::xtensor<double, 2>& b);
 
     /**
     Same as Solve(const xt::xtensor<double, 2>&)
     but writing to preallocated data.
 
     \param b nodevec [nelem, ndim].
     \param x nodevec overwritten [nelem, ndim].
     */
     virtual void solve(const xt::xtensor<double, 2>& b, xt::xtensor<double, 2>& x);
 
     /**
     Same as Solve(const xt::xtensor<double, 2>&)
     but for "dofval" input and output.
 
     \param b dofval [ndof].
     \return x dofval [ndof].
     */
     xt::xtensor<double, 1> Solve(const xt::xtensor<double, 1>& b);
 
     /**
     Same as Solve(const xt::xtensor<double, 1>&)
     but writing to preallocated data.
 
     \param b dofval [ndof].
     \param x dofval overwritten [ndof].
     */
     virtual void solve(const xt::xtensor<double, 1>& b, xt::xtensor<double, 1>& x);
 
 protected:
 
     xt::xtensor<size_t, 2> m_conn; ///< Connectivity [#nelem, #nne].
     xt::xtensor<size_t, 2> m_dofs; ///< DOF-numbers per node [#nnode, #ndim].
     size_t m_nelem; ///< See nelem().
     size_t m_nne; ///< See nne().
     size_t m_nnode; ///< See nnode().
     size_t m_ndim; ///< See ndim().
     size_t m_ndof; ///< See ndof().
     bool m_changed = true; ///< Signal changes to data.
 
 private:
     xt::xtensor<double, 1> m_A; ///< The matrix.
     xt::xtensor<double, 1> m_inv; /// Inverse of the matrix.
     void factorize(); ///< Compute inverse (automatically evaluated by "solve").
 };
 
 } // namespace GooseFEM
 
 #include "MatrixDiagonal.hpp"
 
 #endif
diff --git a/include/GooseFEM/MatrixDiagonalPartitioned.h b/include/GooseFEM/MatrixDiagonalPartitioned.h
index 9d64799..da3d2f7 100644
--- a/include/GooseFEM/MatrixDiagonalPartitioned.h
+++ b/include/GooseFEM/MatrixDiagonalPartitioned.h
@@ -1,226 +1,241 @@
 /**
 Diagonal matrix that is partitioned in:
 -   unknown DOFs
 -   prescribed DOFs
 
 \file MatrixDiagonalPartitioned.h
 \copyright Copyright 2017. Tom de Geus. All rights reserved.
 \license This project is released under the GNU Public License (GPLv3).
 */
 
 #ifndef GOOSEFEM_MATRIXDIAGONALPARTITIONED_H
 #define GOOSEFEM_MATRIXDIAGONALPARTITIONED_H
 
 #include "config.h"
 #include "MatrixDiagonal.h"
 
 namespace GooseFEM {
 
 /**
 Diagonal and partitioned matrix.
 
 See Vector() for bookkeeping definitions.
 */
 class MatrixDiagonalPartitioned : public MatrixDiagonal {
 public:
 
     MatrixDiagonalPartitioned() = default;
 
     /**
     Constructor.
 
     \param conn connectivity [#nelem, #nne].
     \param dofs DOFs per node [#nnode, #ndim].
     \param iip prescribed DOFs [#nnp].
     */
     MatrixDiagonalPartitioned(
         const xt::xtensor<size_t, 2>& conn,
         const xt::xtensor<size_t, 2>& dofs,
         const xt::xtensor<size_t, 1>& iip);
 
     /**
     Number of unknown DOFs.
     */
     size_t nnu() const;
 
     /**
     Number of prescribed DOFs.
     */
     size_t nnp() const;
 
     /**
     Unknown DOFs [#nnu].
     */
     xt::xtensor<size_t, 1> iiu() const;
 
     /**
     Prescribed DOFs [#nnp].
     */
     xt::xtensor<size_t, 1> iip() const;
 
     void assemble(const xt::xtensor<double, 3>& elemmat) override;
 
+    void set(const xt::xtensor<double, 1>& A) override;
+
     xt::xtensor<double, 1> Todiagonal() const override;
 
     void dot(const xt::xtensor<double, 2>& x, xt::xtensor<double, 2>& b) const override;
     void dot(const xt::xtensor<double, 1>& x, xt::xtensor<double, 1>& b) const override;
 
     /**
     \param x_u dofval [#nnu].
     \param x_p dofval [#nnp].
     \return b_u dofval [#nnu].
     */
     xt::xtensor<double, 1> Dot_u(
         const xt::xtensor<double, 1>& x_u,
         const xt::xtensor<double, 1>& x_p) const;
 
     /**
     \param x_u dofval [#nnu].
     \param x_p dofval [#nnp].
     \param b_u (overwritten) dofval [#nnu].
     */
     void dot_u(
         const xt::xtensor<double, 1>& x_u,
         const xt::xtensor<double, 1>& x_p,
               xt::xtensor<double, 1>& b_u) const;
 
     /**
     \param x_u dofval [#nnu].
     \param x_p dofval [#nnp].
     \return b_p dofval [#nnp].
     */
     xt::xtensor<double, 1> Dot_p(
         const xt::xtensor<double, 1>& x_u,
         const xt::xtensor<double, 1>& x_p) const;
 
     /**
     \param x_u dofval [#nnu].
     \param x_p dofval [#nnp].
     \param b_p (overwritten) dofval [#nnp].
     */
     void dot_p(
         const xt::xtensor<double, 1>& x_u,
         const xt::xtensor<double, 1>& x_p,
               xt::xtensor<double, 1>& b_p) const;
 
-    void solve(const xt::xtensor<double, 2>& b, xt::xtensor<double, 2>& x) override; // modified with "x_u"
-    void solve(const xt::xtensor<double, 1>& b, xt::xtensor<double, 1>& x) override; // modified with "x_u"
+    /**
+    Solve \f$ x_u = A_{uu}^{-1} (b_u - A_{up} * x_p) \equiv A_{uu}^{-1} b_u \f$.
+
+    \param b nodevec [#nelem, #ndim].
+    \param x nodevec, modified with `x_u` [#nelem, #ndim].
+    */
+    void solve(const xt::xtensor<double, 2>& b, xt::xtensor<double, 2>& x) override;
+
+    /**
+    Solve \f$ x_u = A_{uu}^{-1} (b_u - A_{up} * x_p) \equiv A_{uu}^{-1} b_u \f$.
+
+    \param b dofval [#ndof].
+    \param x dofval, modified with `x_u` [#ndof].
+    */
+    void solve(const xt::xtensor<double, 1>& b, xt::xtensor<double, 1>& x) override;
 
     /**
     \param b_u dofval [#nnu].
     \param x_p dofval [#nnp].
     \return x_u dofval [#nnu].
     */
     xt::xtensor<double, 1> Solve_u(
         const xt::xtensor<double, 1>& b_u,
         const xt::xtensor<double, 1>& x_p);
 
     /**
     \param b_u dofval [#nnu].
     \param x_p dofval [#nnp].
     \param x_u (overwritten) dofval [#nnu].
     */
     void solve_u(
         const xt::xtensor<double, 1>& b_u,
         const xt::xtensor<double, 1>& x_p,
               xt::xtensor<double, 1>& x_u);
 
     /**
     Get right-hand-size for corresponding to the prescribed DOFs.
 
     \f$ b_p = A_{pu} * x_u + A_{pp} * x_p = A_{pp} * x_p \equiv A_{pp} * x_p \f$
 
     and assemble them to the appropriate places in "nodevec".
 
     \param x "nodevec" [#nnode, #ndim].
     \param b "nodevec" [#nnode, #ndim].
     \return Copy of `b` with \f$ b_p \f$ overwritten.
     */
     xt::xtensor<double, 2> Reaction(
         const xt::xtensor<double, 2>& x,
         const xt::xtensor<double, 2>& b) const;
 
     /**
     Same as Reaction(const xt::xtensor<double, 2>&, const xt::xtensor<double, 2>&),
     but inserting in-place.
 
     \param x "nodevec" [#nnode, #ndim].
     \param b "nodevec" [#nnode, #ndim], \f$ b_p \f$ overwritten.
     */
     void reaction(
         const xt::xtensor<double, 2>& x,
               xt::xtensor<double, 2>& b) const;
 
     /**
     Same as Reaction(const xt::xtensor<double, 2>&, const xt::xtensor<double, 2>&),
     but of "dofval" input and output.
 
     \param x "dofval" [#ndof].
     \param b "dofval" [#ndof].
     \return Copy of `b` with \f$ b_p \f$ overwritten.
     */
     xt::xtensor<double, 1> Reaction(
         const xt::xtensor<double, 1>& x,
         const xt::xtensor<double, 1>& b) const;
 
     /**
     Same as Reaction(const xt::xtensor<double, 1>&, const xt::xtensor<double, 1>&),
     but inserting in-place.
 
     \param x "dofval" [#ndof].
     \param b "dofval" [#ndof], \f$ b_p \f$ overwritten.
     */
     void reaction(
         const xt::xtensor<double, 1>& x,
               xt::xtensor<double, 1>& b) const; // modified with "b_p"
 
     /**
     Same as Reaction(const xt::xtensor<double, 1>&, const xt::xtensor<double, 1>&),
     but with partitioned input and output.
 
     \param x_u unknown "dofval" [#nnu].
     \param x_p prescribed "dofval" [#nnp].
     \return b_p prescribed "dofval" [#nnp].
     */
     xt::xtensor<double, 1> Reaction_p(
         const xt::xtensor<double, 1>& x_u,
         const xt::xtensor<double, 1>& x_p) const;
 
     /**
     Same as Reaction_p(const xt::xtensor<double, 1>&, const xt::xtensor<double, 1>&),
     but writing to preallocated output.
 
     \param x_u unknown "dofval" [#nnu].
     \param x_p prescribed "dofval" [#nnp].
     \param b_p (overwritten) prescribed "dofval" [#nnp].
     */
     void reaction_p(
         const xt::xtensor<double, 1>& x_u,
         const xt::xtensor<double, 1>& x_p,
               xt::xtensor<double, 1>& b_p) const;
 
 private:
     // The diagonal matrix, and its inverse (re-used to solve different RHS)
     xt::xtensor<double, 1> m_Auu;
     xt::xtensor<double, 1> m_App;
     xt::xtensor<double, 1> m_inv_uu;
 
 
     // Bookkeeping
     xt::xtensor<size_t, 2> m_part; // DOF-numbers per node, renumbered  [nnode, ndim]
     xt::xtensor<size_t, 1> m_iiu;  // DOF-numbers that are unknown      [nnu]
     xt::xtensor<size_t, 1> m_iip;  // DOF-numbers that are prescribed   [nnp]
 
     // Dimensions
     size_t m_nnu;   // number of unknown DOFs
     size_t m_nnp;   // number of prescribed DOFs
 
     // Compute inverse (automatically evaluated by "solve")
     void factorize();
 };
 
 } // namespace GooseFEM
 
 #include "MatrixDiagonalPartitioned.hpp"
 
 #endif
diff --git a/include/GooseFEM/MatrixDiagonalPartitioned.hpp b/include/GooseFEM/MatrixDiagonalPartitioned.hpp
index 786829e..5d8e7e6 100644
--- a/include/GooseFEM/MatrixDiagonalPartitioned.hpp
+++ b/include/GooseFEM/MatrixDiagonalPartitioned.hpp
@@ -1,332 +1,349 @@
 /**
 Implementation of MatrixDiagonalPartitioned.h
 
 \file MatrixDiagonalPartitioned.hpp
 \copyright Copyright 2017. Tom de Geus. All rights reserved.
 \license This project is released under the GNU Public License (GPLv3).
 */
 
 #ifndef GOOSEFEM_MATRIXDIAGONALPARTITIONED_HPP
 #define GOOSEFEM_MATRIXDIAGONALPARTITIONED_HPP
 
 #include "MatrixDiagonalPartitioned.h"
 #include "Mesh.h"
 
 namespace GooseFEM {
 
 inline MatrixDiagonalPartitioned::MatrixDiagonalPartitioned(
     const xt::xtensor<size_t, 2>& conn,
     const xt::xtensor<size_t, 2>& dofs,
     const xt::xtensor<size_t, 1>& iip)
     : MatrixDiagonal(conn, dofs)
 {
     m_iip = iip;
     m_iiu = xt::setdiff1d(dofs, iip);
     m_nnp = m_iip.size();
     m_nnu = m_iiu.size();
     m_part = Mesh::Reorder({m_iiu, m_iip}).apply(m_dofs);
     m_Auu = xt::empty<double>({m_nnu});
     m_App = xt::empty<double>({m_nnp});
     m_inv_uu = xt::empty<double>({m_nnu});
 
     GOOSEFEM_ASSERT(xt::amax(m_iip)() <= xt::amax(m_dofs)());
 }
 
 inline size_t MatrixDiagonalPartitioned::nnu() const
 {
     return m_nnu;
 }
 
 inline size_t MatrixDiagonalPartitioned::nnp() const
 {
     return m_nnp;
 }
 
 inline xt::xtensor<size_t, 1> MatrixDiagonalPartitioned::iiu() const
 {
     return m_iiu;
 }
 
 inline xt::xtensor<size_t, 1> MatrixDiagonalPartitioned::iip() const
 {
     return m_iip;
 }
 
 inline void MatrixDiagonalPartitioned::factorize()
 {
     if (!m_changed) {
         return;
     }
 
     #pragma omp parallel for
     for (size_t d = 0; d < m_nnu; ++d) {
         m_inv_uu(d) = 1.0 / m_Auu(d);
     }
 
     m_changed = false;
 }
 
 inline void MatrixDiagonalPartitioned::assemble(const xt::xtensor<double, 3>& elemmat)
 {
     GOOSEFEM_ASSERT(xt::has_shape(elemmat, {m_nelem, m_nne * m_ndim, m_nne * m_ndim}));
     GOOSEFEM_ASSERT(Element::isDiagonal(elemmat));
 
     m_Auu.fill(0.0);
     m_App.fill(0.0);
 
     for (size_t e = 0; e < m_nelem; ++e) {
         for (size_t m = 0; m < m_nne; ++m) {
             for (size_t i = 0; i < m_ndim; ++i) {
 
                 size_t d = m_part(m_conn(e, m), i);
 
                 if (d < m_nnu) {
                     m_Auu(d) += elemmat(e, m * m_ndim + i, m * m_ndim + i);
                 }
                 else {
                     m_App(d - m_nnu) += elemmat(e, m * m_ndim + i, m * m_ndim + i);
                 }
             }
         }
     }
 
     m_changed = true;
 }
 
+inline void MatrixDiagonalPartitioned::set(const xt::xtensor<double, 1>& A)
+{
+    GOOSEFEM_ASSERT(A.size() == m_ndof);
+
+    #pragma omp parallel for
+    for (size_t d = 0; d < m_nnu; ++d) {
+        m_Auu(d) = A(m_iiu(d));
+    }
+
+    #pragma omp parallel for
+    for (size_t d = 0; d < m_nnp; ++d) {
+        m_App(d) = A(m_iip(d));
+    }
+
+    m_changed = true;
+}
+
 inline void
 MatrixDiagonalPartitioned::dot(const xt::xtensor<double, 2>& x, xt::xtensor<double, 2>& b) const
 {
     GOOSEFEM_ASSERT(xt::has_shape(x, {m_nnode, m_ndim}));
     GOOSEFEM_ASSERT(xt::has_shape(b, {m_nnode, m_ndim}));
 
     #pragma omp parallel for
     for (size_t m = 0; m < m_nnode; ++m) {
         for (size_t i = 0; i < m_ndim; ++i) {
 
             size_t d = m_part(m, i);
 
             if (d < m_nnu) {
                 b(m, i) = m_Auu(d) * x(m, i);
             }
             else {
                 b(m, i) = m_App(d - m_nnu) * x(m, i);
             }
         }
     }
 }
 
 inline void
 MatrixDiagonalPartitioned::dot(const xt::xtensor<double, 1>& x, xt::xtensor<double, 1>& b) const
 {
     GOOSEFEM_ASSERT(x.size() == m_ndof);
     GOOSEFEM_ASSERT(b.size() == m_ndof);
 
     #pragma omp parallel for
     for (size_t d = 0; d < m_nnu; ++d) {
         b(m_iiu(d)) = m_Auu(d) * x(m_iiu(d));
     }
 
     #pragma omp parallel for
     for (size_t d = 0; d < m_nnp; ++d) {
         b(m_iip(d)) = m_App(d) * x(m_iip(d));
     }
 }
 
 inline void MatrixDiagonalPartitioned::dot_u(
     const xt::xtensor<double, 1>& x_u,
     const xt::xtensor<double, 1>& x_p,
     xt::xtensor<double, 1>& b_u) const
 {
     UNUSED(x_p);
 
     GOOSEFEM_ASSERT(x_u.size() == m_nnu);
     GOOSEFEM_ASSERT(x_p.size() == m_nnp);
     GOOSEFEM_ASSERT(b_u.size() == m_nnu);
 
     #pragma omp parallel for
     for (size_t d = 0; d < m_nnu; ++d) {
         b_u(d) = m_Auu(d) * x_u(d);
     }
 }
 
 inline void MatrixDiagonalPartitioned::dot_p(
     const xt::xtensor<double, 1>& x_u,
     const xt::xtensor<double, 1>& x_p,
     xt::xtensor<double, 1>& b_p) const
 {
     UNUSED(x_u);
 
     GOOSEFEM_ASSERT(x_u.size() == m_nnu);
     GOOSEFEM_ASSERT(x_p.size() == m_nnp);
     GOOSEFEM_ASSERT(b_p.size() == m_nnp);
 
     #pragma omp parallel for
     for (size_t d = 0; d < m_nnp; ++d) {
         b_p(d) = m_App(d) * x_p(d);
     }
 }
 
 inline void
 MatrixDiagonalPartitioned::solve(const xt::xtensor<double, 2>& b, xt::xtensor<double, 2>& x)
 {
     GOOSEFEM_ASSERT(xt::has_shape(b, {m_nnode, m_ndim}));
     GOOSEFEM_ASSERT(xt::has_shape(x, {m_nnode, m_ndim}));
 
     this->factorize();
 
     #pragma omp parallel for
     for (size_t m = 0; m < m_nnode; ++m) {
         for (size_t i = 0; i < m_ndim; ++i) {
             if (m_part(m, i) < m_nnu) {
                 x(m, i) = m_inv_uu(m_part(m, i)) * b(m, i);
             }
         }
     }
 }
 
 inline void
 MatrixDiagonalPartitioned::solve(const xt::xtensor<double, 1>& b, xt::xtensor<double, 1>& x)
 {
     GOOSEFEM_ASSERT(b.size() == m_ndof);
     GOOSEFEM_ASSERT(x.size() == m_ndof);
 
     this->factorize();
 
     #pragma omp parallel for
     for (size_t d = 0; d < m_nnu; ++d) {
         x(m_iiu(d)) = m_inv_uu(d) * b(m_iiu(d));
     }
 }
 
 inline void MatrixDiagonalPartitioned::solve_u(
     const xt::xtensor<double, 1>& b_u,
     const xt::xtensor<double, 1>& x_p,
     xt::xtensor<double, 1>& x_u)
 {
     UNUSED(x_p);
 
     GOOSEFEM_ASSERT(b_u.size() == m_nnu);
     GOOSEFEM_ASSERT(x_p.size() == m_nnp);
     GOOSEFEM_ASSERT(x_u.size() == m_nnu);
 
     this->factorize();
 
     #pragma omp parallel for
     for (size_t d = 0; d < m_nnu; ++d) {
         x_u(d) = m_inv_uu(d) * b_u(d);
     }
 }
 
 inline void MatrixDiagonalPartitioned::reaction(
     const xt::xtensor<double, 2>& x, xt::xtensor<double, 2>& b) const
 {
     GOOSEFEM_ASSERT(xt::has_shape(x, {m_nnode, m_ndim}));
     GOOSEFEM_ASSERT(xt::has_shape(b, {m_nnode, m_ndim}));
 
     #pragma omp parallel for
     for (size_t m = 0; m < m_nnode; ++m) {
         for (size_t i = 0; i < m_ndim; ++i) {
             if (m_part(m, i) >= m_nnu) {
                 b(m, i) = m_App(m_part(m, i) - m_nnu) * x(m, i);
             }
         }
     }
 }
 
 inline void MatrixDiagonalPartitioned::reaction(
     const xt::xtensor<double, 1>& x, xt::xtensor<double, 1>& b) const
 {
     GOOSEFEM_ASSERT(x.size() == m_ndof);
     GOOSEFEM_ASSERT(b.size() == m_ndof);
 
     #pragma omp parallel for
     for (size_t d = 0; d < m_nnp; ++d) {
         b(m_iip(d)) = m_App(d) * x(m_iip(d));
     }
 }
 
 inline void MatrixDiagonalPartitioned::reaction_p(
     const xt::xtensor<double, 1>& x_u,
     const xt::xtensor<double, 1>& x_p,
     xt::xtensor<double, 1>& b_p) const
 {
     UNUSED(x_u);
 
     GOOSEFEM_ASSERT(x_u.size() == m_nnu);
     GOOSEFEM_ASSERT(x_p.size() == m_nnp);
     GOOSEFEM_ASSERT(b_p.size() == m_nnp);
 
     #pragma omp parallel for
     for (size_t d = 0; d < m_nnp; ++d) {
         b_p(d) = m_App(d) * x_p(d);
     }
 }
 
 inline xt::xtensor<double, 1> MatrixDiagonalPartitioned::Todiagonal() const
 {
     xt::xtensor<double, 1> ret = xt::zeros<double>({m_ndof});
 
     #pragma omp parallel for
     for (size_t d = 0; d < m_nnu; ++d) {
         ret(m_iiu(d)) = m_Auu(d);
     }
 
     #pragma omp parallel for
     for (size_t d = 0; d < m_nnp; ++d) {
         ret(m_iip(d)) = m_App(d);
     }
 
     return ret;
 }
 
 inline xt::xtensor<double, 1> MatrixDiagonalPartitioned::Dot_u(
     const xt::xtensor<double, 1>& x_u, const xt::xtensor<double, 1>& x_p) const
 {
     xt::xtensor<double, 1> b_u = xt::empty<double>({m_nnu});
     this->dot_u(x_u, x_p, b_u);
     return b_u;
 }
 
 inline xt::xtensor<double, 1> MatrixDiagonalPartitioned::Dot_p(
     const xt::xtensor<double, 1>& x_u, const xt::xtensor<double, 1>& x_p) const
 {
     xt::xtensor<double, 1> b_p = xt::empty<double>({m_nnp});
     this->dot_p(x_u, x_p, b_p);
     return b_p;
 }
 
 inline xt::xtensor<double, 1> MatrixDiagonalPartitioned::Solve_u(
     const xt::xtensor<double, 1>& b_u, const xt::xtensor<double, 1>& x_p)
 {
     xt::xtensor<double, 1> x_u = xt::empty<double>({m_nnu});
     this->solve_u(b_u, x_p, x_u);
     return x_u;
 }
 
 inline xt::xtensor<double, 2> MatrixDiagonalPartitioned::Reaction(
     const xt::xtensor<double, 2>& x, const xt::xtensor<double, 2>& b) const
 {
     xt::xtensor<double, 2> ret = b;
     this->reaction(x, ret);
     return ret;
 }
 
 inline xt::xtensor<double, 1> MatrixDiagonalPartitioned::Reaction(
     const xt::xtensor<double, 1>& x, const xt::xtensor<double, 1>& b) const
 {
     xt::xtensor<double, 1> ret = b;
     this->reaction(x, ret);
     return ret;
 }
 
 inline xt::xtensor<double, 1> MatrixDiagonalPartitioned::Reaction_p(
     const xt::xtensor<double, 1>& x_u, const xt::xtensor<double, 1>& x_p) const
 {
     xt::xtensor<double, 1> b_p = xt::empty<double>({m_nnp});
     this->reaction_p(x_u, x_p, b_p);
     return b_p;
 }
 
 } // namespace GooseFEM
 
 #endif
diff --git a/include/GooseFEM/MatrixPartitioned.h b/include/GooseFEM/MatrixPartitioned.h
index 18a2fc7..3845a52 100644
--- a/include/GooseFEM/MatrixPartitioned.h
+++ b/include/GooseFEM/MatrixPartitioned.h
@@ -1,295 +1,295 @@
 /**
 Sparse matrix that is partitioned in:
 -   unknown DOFs
 -   prescribed DOFs
 
 \file MatrixPartitioned.h
 \copyright Copyright 2017. Tom de Geus. All rights reserved.
 \license This project is released under the GNU Public License (GPLv3).
 */
 
 #ifndef GOOSEFEM_MATRIXPARTITIONED_H
 #define GOOSEFEM_MATRIXPARTITIONED_H
 
 #include "config.h"
 #include "Matrix.h"
 
 #include <Eigen/Eigen>
 #include <Eigen/Sparse>
 #include <Eigen/SparseCholesky>
 
 namespace GooseFEM {
 
 // forward declaration
 template <class> class MatrixPartitionedSolver;
 
 /**
 Sparse matrix partitioned in an unknown and a prescribed part. In particular:
 
 \f$ \begin{bmatrix} A_{uu} & A_{up} \\ A_{pu} & A_{pp} \end{bmatrix} \f$
 
 See VectorPartitioned() for bookkeeping definitions.
 */
 class MatrixPartitioned : public Matrix {
 public:
 
     MatrixPartitioned() = default;
 
     /**
     Constructor.
 
     \param conn connectivity [#nelem, #nne].
     \param dofs DOFs per node [#nnode, #ndim].
     \param iip prescribed DOFs [#nnp].
     */
     MatrixPartitioned(
         const xt::xtensor<size_t, 2>& conn,
         const xt::xtensor<size_t, 2>& dofs,
         const xt::xtensor<size_t, 1>& iip);
 
     /**
     Number of unknown DOFs.
     */
     size_t nnu() const;
 
     /**
     Number of prescribed DOFs.
     */
     size_t nnp() const;
 
     /**
     Unknown DOFs [#nnu].
     */
     xt::xtensor<size_t, 1> iiu() const;
 
     /**
     Prescribed DOFs [#nnp].
     */
     xt::xtensor<size_t, 1> iip() const;
 
     void assemble(const xt::xtensor<double, 3>& elemmat) override;
 
     void set(
         const xt::xtensor<size_t, 1>& rows,
         const xt::xtensor<size_t, 1>& cols,
         const xt::xtensor<double, 2>& matrix) override;
 
     void add(
         const xt::xtensor<size_t, 1>& rows,
         const xt::xtensor<size_t, 1>& cols,
         const xt::xtensor<double, 2>& matrix) override;
 
     void todense(xt::xtensor<double, 2>& ret) const override;
     void dot(const xt::xtensor<double, 2>& x, xt::xtensor<double, 2>& b) const override;
     void dot(const xt::xtensor<double, 1>& x, xt::xtensor<double, 1>& b) const override;
 
     /**
     Get right-hand-size for corresponding to the prescribed DOFs.
 
-    \f$ b_p = A_{pu} * x_u + A_{pp} * x_p = A_{pp} * x_p \f$
+    \f$ b_p = A_{pu} * x_u + A_{pp} * x_p \f$
 
     and assemble them to the appropriate places in "nodevec".
 
     \param x "nodevec" [#nnode, #ndim].
     \param b "nodevec" [#nnode, #ndim].
     \return Copy of `b` with \f$ b_p \f$ overwritten.
     */
     xt::xtensor<double, 2> Reaction(
         const xt::xtensor<double, 2>& x,
         const xt::xtensor<double, 2>& b) const;
 
     /**
     Same as Reaction(const xt::xtensor<double, 2>&, const xt::xtensor<double, 2>&),
     but inserting in-place.
 
     \param x "nodevec" [#nnode, #ndim].
     \param b "nodevec" [#nnode, #ndim], \f$ b_p \f$ overwritten.
     */
     void reaction(
         const xt::xtensor<double, 2>& x,
               xt::xtensor<double, 2>& b) const;
 
     /**
     Same as Reaction(const xt::xtensor<double, 2>&, const xt::xtensor<double, 2>&),
     but of "dofval" input and output.
 
     \param x "dofval" [#ndof].
     \param b "dofval" [#ndof].
     \return Copy of `b` with \f$ b_p \f$ overwritten.
     */
     xt::xtensor<double, 1> Reaction(
         const xt::xtensor<double, 1>& x,
         const xt::xtensor<double, 1>& b) const;
 
     /**
     Same as Reaction(const xt::xtensor<double, 1>&, const xt::xtensor<double, 1>&),
     but inserting in-place.
 
     \param x "dofval" [#ndof].
     \param b "dofval" [#ndof], \f$ b_p \f$ overwritten.
     */
     void reaction(
         const xt::xtensor<double, 1>& x,
               xt::xtensor<double, 1>& b) const;
 
     /**
     Same as Reaction(const xt::xtensor<double, 1>&, const xt::xtensor<double, 1>&),
     but with partitioned input and output.
 
     \param x_u unknown "dofval" [#nnu].
     \param x_p prescribed "dofval" [#nnp].
     \return b_p prescribed "dofval" [#nnp].
     */
     xt::xtensor<double, 1> Reaction_p(
         const xt::xtensor<double, 1>& x_u,
         const xt::xtensor<double, 1>& x_p) const;
 
     /**
     Same as Reaction_p(const xt::xtensor<double, 1>&, const xt::xtensor<double, 1>&),
     but writing to preallocated output.
 
     \param x_u unknown "dofval" [#nnu].
     \param x_p prescribed "dofval" [#nnp].
     \param b_p (overwritten) prescribed "dofval" [#nnp].
     */
     void reaction_p(
         const xt::xtensor<double, 1>& x_u,
         const xt::xtensor<double, 1>& x_p,
               xt::xtensor<double, 1>& b_p) const;
 
 private:
 
     Eigen::SparseMatrix<double> m_Auu; ///< The matrix.
     Eigen::SparseMatrix<double> m_Aup; ///< The matrix.
     Eigen::SparseMatrix<double> m_Apu; ///< The matrix.
     Eigen::SparseMatrix<double> m_App; ///< The matrix.
     std::vector<Eigen::Triplet<double>> m_Tuu; ///< Matrix entries.
     std::vector<Eigen::Triplet<double>> m_Tup; ///< Matrix entries.
     std::vector<Eigen::Triplet<double>> m_Tpu; ///< Matrix entries.
     std::vector<Eigen::Triplet<double>> m_Tpp; ///< Matrix entries.
     xt::xtensor<size_t, 1> m_iiu; ///< See iiu()
     xt::xtensor<size_t, 1> m_iip; ///< See iip()
     size_t m_nnu; ///< See #nnu
     size_t m_nnp; ///< See #nnp
 
     /**
     Renumbered DOFs per node, such that
 
         iiu = arange(nnu)
         iip = nnu + arange(nnp)
 
     making is much simpler to slice.
     */
     xt::xtensor<size_t, 2> m_part;
 
     // grant access to solver class
     template <class> friend class MatrixPartitionedSolver;
 
 private:
 
     // Convert arrays (Eigen version of VectorPartitioned, which contains public functions)
     Eigen::VectorXd AsDofs_u(const xt::xtensor<double, 1>& dofval) const;
     Eigen::VectorXd AsDofs_u(const xt::xtensor<double, 2>& nodevec) const;
     Eigen::VectorXd AsDofs_p(const xt::xtensor<double, 1>& dofval) const;
     Eigen::VectorXd AsDofs_p(const xt::xtensor<double, 2>& nodevec) const;
 };
 
 /**
 Solver for MatrixPartitioned().
 The idea is that this solver class can be used to solve for multiple right-hand-sides
 using one factorisation.
 */
 template <class Solver = Eigen::SimplicialLDLT<Eigen::SparseMatrix<double>>>
 class MatrixPartitionedSolver {
 public:
 
     MatrixPartitionedSolver() = default;
 
     /**
     Solve \f$ x_u = A_{uu}^{-1} (b_u - A_{up} * x_p) \f$.
 
     \param A sparse matrix, see MatrixPartitioned().
     \param b nodevec [nelem, ndim].
     \param x nodevec [nelem, ndim], used to read \f$ x_p \f$.
     \return x nodevec [nelem, ndim], \f$ x_u \f$ filled, \f$ x_p \f$ copied.
     */
     xt::xtensor<double, 2> Solve(
         MatrixPartitioned& A,
         const xt::xtensor<double, 2>& b,
         const xt::xtensor<double, 2>& x);
 
     /**
     Same as Solve(MatrixPartitioned&, const xt::xtensor<double, 2>&, const xt::xtensor<double, 2>&),
     but filling \f$ x_u \f$ in place.
 
     \param A sparse matrix, see MatrixPartitioned().
     \param b nodevec [nelem, ndim].
     \param x nodevec [nelem, ndim], \f$ x_p \f$ read, \f$ x_u \f$ filled.
     */
     void solve(MatrixPartitioned& A, const xt::xtensor<double, 2>& b, xt::xtensor<double, 2>& x);
 
     /**
     Same as Solve(MatrixPartitioned&, const xt::xtensor<double, 2>&, const xt::xtensor<double, 2>&),
     but for "dofval" input and output.
 
     \param A sparse matrix, see MatrixPartitioned().
     \param b dofval [ndof].
     \param x dofval [ndof], used to read \f$ x_p \f$.
     \return x dofval [ndof], \f$ x_u \f$ filled, \f$ x_p \f$ copied.
     */
     xt::xtensor<double, 1> Solve(
         MatrixPartitioned& A,
         const xt::xtensor<double, 1>& b,
         const xt::xtensor<double, 1>& x);
 
     /**
     Same as Solve(MatrixPartitioned&, const xt::xtensor<double, 1>&, const xt::xtensor<double, 1>&),
     but filling \f$ x_u \f$ in place.
 
     \param A sparse matrix, see MatrixPartitioned().
     \param b dofval [ndof].
     \param x dofval [ndof], \f$ x_p \f$ read, \f$ x_u \f$ filled.
     */
     void solve(MatrixPartitioned& A, const xt::xtensor<double, 1>& b, xt::xtensor<double, 1>& x);
 
     /**
     Same as Solve(MatrixPartitioned&, const xt::xtensor<double, 2>&, const xt::xtensor<double, 2>&),
     but with partitioned input and output.
 
     \param A sparse matrix, see MatrixPartitioned().
     \param b_u unknown dofval [nnu].
     \param x_p prescribed dofval [nnp]
     \return x_u unknown dofval [nnu].
     */
     xt::xtensor<double, 1> Solve_u(
         MatrixPartitioned& A,
         const xt::xtensor<double, 1>& b_u,
         const xt::xtensor<double, 1>& x_p);
 
     /**
     Same as
     Solve_u(MatrixPartitioned&, const xt::xtensor<double, 1>&, const xt::xtensor<double, 1>&),
     but writing to pre-allocated output.
 
     \param A sparse matrix, see MatrixPartitioned().
     \param b_u unknown dofval [nnu].
     \param x_p prescribed dofval [nnp]
     \param x_u (overwritten) unknown dofval [nnu].
     */
     void solve_u(
         MatrixPartitioned& A,
         const xt::xtensor<double, 1>& b_u,
         const xt::xtensor<double, 1>& x_p,
         xt::xtensor<double, 1>& x_u);
 
 private:
     Solver m_solver; ///< solver
     bool m_factor = true; ///< signal to force factorization
     void factorize(MatrixPartitioned& matrix); ///< compute inverse (evaluated by "solve")
 };
 
 } // namespace GooseFEM
 
 #include "MatrixPartitioned.hpp"
 
 #endif
diff --git a/python/Matrix.hpp b/python/Matrix.hpp
index a009545..7a0360c 100644
--- a/python/Matrix.hpp
+++ b/python/Matrix.hpp
@@ -1,97 +1,104 @@
 /* =================================================================================================
 
 (c - GPLv3) T.W.J. de Geus (Tom) | tom@geus.me | www.geus.me | github.com/tdegeus/GooseFEM
 
 ================================================================================================= */
 
 #include <Eigen/Eigen>
-#include <GooseFEM/GooseFEM.h>
+#include <GooseFEM/Matrix.h>
 #include <pybind11/eigen.h>
 #include <pybind11/pybind11.h>
 #include <pyxtensor/pyxtensor.hpp>
 
 namespace py = pybind11;
 
 void init_Matrix(py::module& m)
 {
 
     py::class_<GooseFEM::Matrix>(m, "Matrix")
 
         .def(
             py::init<const xt::xtensor<size_t, 2>&, const xt::xtensor<size_t, 2>&>(),
             "Sparse matrix",
             py::arg("conn"),
             py::arg("dofs"))
 
-        .def("nelem", &GooseFEM::Matrix::nelem, "Number of element")
-
-        .def("nne", &GooseFEM::Matrix::nne, "Number of nodes per element")
-
-        .def("nnode", &GooseFEM::Matrix::nnode, "Number of nodes")
-
-        .def("ndim", &GooseFEM::Matrix::ndim, "Number of dimensions")
-
-        .def("ndof", &GooseFEM::Matrix::ndof, "Number of degrees-of-freedom")
-
-        .def("dofs", &GooseFEM::Matrix::dofs, "Return degrees-of-freedom")
-
-        .def(
-            "assemble",
-            &GooseFEM::Matrix::assemble,
-            "Assemble matrix from 'elemmat",
-            py::arg("elemmat"))
-
-        .def(
-            "set",
-            &GooseFEM::Matrix::set,
-            "Overwrite with a dense (sub-) matrix",
-            py::arg("rows"),
-            py::arg("cols"),
-            py::arg("matrix"))
-
-        .def(
-            "add",
-            &GooseFEM::Matrix::add,
-            "Add a dense (sub-) matrix to the current matrix",
-            py::arg("rows"),
-            py::arg("cols"),
-            py::arg("matrix"))
-
-        .def("Todense", &GooseFEM::Matrix::Todense, "Return as dense matrix")
-
-        .def(
-            "Dot",
-            py::overload_cast<const xt::xtensor<double, 1>&>(&GooseFEM::Matrix::Dot, py::const_),
-            "Dot",
-            py::arg("x"))
-
-        .def(
-            "Dot",
-            py::overload_cast<const xt::xtensor<double, 2>&>(&GooseFEM::Matrix::Dot, py::const_),
-            "Dot",
-            py::arg("x"))
+        .def("nelem",
+             &GooseFEM::Matrix::nelem,
+             "See :cpp:func:`GooseFEM::Matrix::nelem`.")
+
+        .def("nne",
+             &GooseFEM::Matrix::nne,
+             "See :cpp:func:`GooseFEM::Matrix::nne`.")
+
+        .def("nnode",
+             &GooseFEM::Matrix::nnode,
+             "See :cpp:func:`GooseFEM::Matrix::nnode`.")
+
+        .def("ndim",
+             &GooseFEM::Matrix::ndim,
+             "See :cpp:func:`GooseFEM::Matrix::ndim`.")
+
+        .def("ndof",
+             &GooseFEM::Matrix::ndof,
+             "See :cpp:func:`GooseFEM::Matrix::ndof`.")
+
+        .def("dofs",
+             &GooseFEM::Matrix::dofs,
+             "See :cpp:func:`GooseFEM::Matrix::dofs`.")
+
+        .def("assemble",
+             &GooseFEM::Matrix::assemble,
+             "See :cpp:func:`GooseFEM::Matrix::assemble`.",
+             py::arg("elemmat"))
+
+        .def("set",
+             &GooseFEM::Matrix::set,
+             "See :cpp:func:`GooseFEM::Matrix::set`.",
+             py::arg("rows"),
+             py::arg("cols"),
+             py::arg("matrix"))
+
+        .def("add",
+             &GooseFEM::Matrix::add,
+             "See :cpp:func:`GooseFEM::Matrix::add`.",
+             py::arg("rows"),
+             py::arg("cols"),
+             py::arg("matrix"))
+
+        .def("Todense",
+             &GooseFEM::Matrix::Todense,
+             "See :cpp:func:`GooseFEM::Matrix::Todense`.")
+
+        .def("Dot",
+             py::overload_cast<const xt::xtensor<double, 1>&>(&GooseFEM::Matrix::Dot, py::const_),
+             "See :cpp:func:`GooseFEM::Matrix::Dot`.",
+             py::arg("x"))
+
+        .def("Dot",
+             py::overload_cast<const xt::xtensor<double, 2>&>(&GooseFEM::Matrix::Dot, py::const_),
+             "See :cpp:func:`GooseFEM::Matrix::Dot`.",
+             py::arg("x"))
 
         .def("__repr__", [](const GooseFEM::Matrix&) { return "<GooseFEM.Matrix>"; });
 
     py::class_<GooseFEM::MatrixSolver<>>(m, "MatrixSolver")
 
         .def(py::init<>(), "Sparse matrix solver")
 
-        .def(
-            "Solve",
-            py::overload_cast<GooseFEM::Matrix&, const xt::xtensor<double, 1>&>(
+        .def("Solve",
+             py::overload_cast<GooseFEM::Matrix&, const xt::xtensor<double, 1>&>(
                 &GooseFEM::MatrixSolver<>::Solve),
-            "Solve",
-            py::arg("matrix"),
-            py::arg("b"))
+             "See :cpp:func:`GooseFEM::MatrixSolver::Solve`.",
+             py::arg("A"),
+             py::arg("b"))
 
-        .def(
-            "Solve",
-            py::overload_cast<GooseFEM::Matrix&, const xt::xtensor<double, 2>&>(
+        .def("Solve",
+             py::overload_cast<GooseFEM::Matrix&, const xt::xtensor<double, 2>&>(
                 &GooseFEM::MatrixSolver<>::Solve),
-            "Solve",
-            py::arg("matrix"),
-            py::arg("b"))
+             "See :cpp:func:`GooseFEM::MatrixSolver::Solve`.",
+             py::arg("A"),
+             py::arg("b"))
 
         .def("__repr__", [](const GooseFEM::MatrixSolver<>&) { return "<GooseFEM.MatrixSolver>"; });
 }
diff --git a/python/MatrixDiagonalPartitioned.hpp b/python/MatrixDiagonalPartitioned.hpp
index 21871dd..90b43be 100644
--- a/python/MatrixDiagonalPartitioned.hpp
+++ b/python/MatrixDiagonalPartitioned.hpp
@@ -1,95 +1,88 @@
 /* =================================================================================================
 
 (c - GPLv3) T.W.J. de Geus (Tom) | tom@geus.me | www.geus.me | github.com/tdegeus/GooseFEM
 
 ================================================================================================= */
 
 #include <GooseFEM/MatrixDiagonalPartitioned.h>
 #include <pybind11/pybind11.h>
 #include <pyxtensor/pyxtensor.hpp>
 
 namespace py = pybind11;
 
 void init_MatrixDiagonalPartitioned(py::module& m)
 {
 
     py::class_<GooseFEM::MatrixDiagonalPartitioned, GooseFEM::MatrixDiagonal>(m, "MatrixDiagonalPartitioned")
 
         .def(py::init<
                 const xt::xtensor<size_t, 2>&,
                 const xt::xtensor<size_t, 2>&,
                 const xt::xtensor<size_t, 1>&>(),
              "Diagonal, partitioned, matrix",
              py::arg("conn"),
              py::arg("dofs"),
              py::arg("iip"))
 
-        .def(
-            "nnu",
-            &GooseFEM::MatrixDiagonalPartitioned::nnu,
-            "Number of unknown degrees-of-freedom")
-
-        .def(
-            "nnp",
-            &GooseFEM::MatrixDiagonalPartitioned::nnp,
-            "Number of prescribed degrees-of-freedom")
-
-        .def("iiu", &GooseFEM::MatrixDiagonalPartitioned::iiu, "Return unknown degrees-of-freedom")
-
-        .def(
-            "iip",
-            &GooseFEM::MatrixDiagonalPartitioned::iip,
-            "Return prescribed degrees-of-freedom")
-
-        .def(
-            "Dot_u",
-            py::overload_cast<const xt::xtensor<double, 1>&, const xt::xtensor<double, 1>&>(
+        .def("nnu",
+             &GooseFEM::MatrixDiagonalPartitioned::nnu,
+             "See :cpp:func:`GooseFEM::MatrixDiagonalPartitioned::nnu`.")
+
+        .def("nnp",
+             &GooseFEM::MatrixDiagonalPartitioned::nnp,
+             "See :cpp:func:`GooseFEM::MatrixDiagonalPartitioned::nnp`.")
+
+        .def("iiu",
+             &GooseFEM::MatrixDiagonalPartitioned::iiu,
+             "See :cpp:func:`GooseFEM::MatrixDiagonalPartitioned::iiu`.")
+
+        .def("iip",
+             &GooseFEM::MatrixDiagonalPartitioned::iip,
+             "See :cpp:func:`GooseFEM::MatrixDiagonalPartitioned::iip`.")
+
+        .def("Dot_u",
+             py::overload_cast<const xt::xtensor<double, 1>&, const xt::xtensor<double, 1>&>(
                 &GooseFEM::MatrixDiagonalPartitioned::Dot_u, py::const_),
-            "Dot product 'b_i = A_ij * x_j (b_u = A_uu * x_u + A_up * x_p == A_uu * x_u)",
-            py::arg("x_u"),
-            py::arg("x_p"))
+             "See :cpp:func:`GooseFEM::MatrixDiagonalPartitioned::Dot_u`.",
+             py::arg("x_u"),
+             py::arg("x_p"))
 
-        .def(
-            "Dot_p",
-            py::overload_cast<const xt::xtensor<double, 1>&, const xt::xtensor<double, 1>&>(
+        .def("Dot_p",
+             py::overload_cast<const xt::xtensor<double, 1>&, const xt::xtensor<double, 1>&>(
                 &GooseFEM::MatrixDiagonalPartitioned::Dot_p, py::const_),
-            "Dot product 'b_i = A_ij * x_j (b_p = A_pu * x_u + A_pp * x_p == A_pp * x_p)",
-            py::arg("x_u"),
-            py::arg("x_p"))
+             "See :cpp:func:`GooseFEM::MatrixDiagonalPartitioned::Dot_p`.",
+             py::arg("x_u"),
+             py::arg("x_p"))
 
-        .def(
-            "Solve_u",
-            py::overload_cast<const xt::xtensor<double, 1>&, const xt::xtensor<double, 1>&>(
+        .def("Solve_u",
+             py::overload_cast<const xt::xtensor<double, 1>&, const xt::xtensor<double, 1>&>(
                 &GooseFEM::MatrixDiagonalPartitioned::Solve_u),
-            "Solve_u",
-            py::arg("b_u"),
-            py::arg("x_p"))
+             "See :cpp:func:`GooseFEM::MatrixDiagonalPartitioned::Solve_u`.",
+             py::arg("b_u"),
+             py::arg("x_p"))
 
-        .def(
-            "Reaction",
-            py::overload_cast<const xt::xtensor<double, 1>&, const xt::xtensor<double, 1>&>(
+        .def("Reaction",
+             py::overload_cast<const xt::xtensor<double, 1>&, const xt::xtensor<double, 1>&>(
                 &GooseFEM::MatrixDiagonalPartitioned::Reaction, py::const_),
-            "Reaction",
-            py::arg("x"),
-            py::arg("b"))
+             "See :cpp:func:`GooseFEM::MatrixDiagonalPartitioned::Reaction`.",
+             py::arg("x"),
+             py::arg("b"))
 
-        .def(
-            "Reaction",
-            py::overload_cast<const xt::xtensor<double, 2>&, const xt::xtensor<double, 2>&>(
+        .def("Reaction",
+             py::overload_cast<const xt::xtensor<double, 2>&, const xt::xtensor<double, 2>&>(
                 &GooseFEM::MatrixDiagonalPartitioned::Reaction, py::const_),
-            "Reaction",
-            py::arg("x"),
-            py::arg("b"))
+             "See :cpp:func:`GooseFEM::MatrixDiagonalPartitioned::Reaction`.",
+             py::arg("x"),
+             py::arg("b"))
 
-        .def(
-            "Reaction_p",
-            py::overload_cast<const xt::xtensor<double, 1>&, const xt::xtensor<double, 1>&>(
+        .def("Reaction_p",
+             py::overload_cast<const xt::xtensor<double, 1>&, const xt::xtensor<double, 1>&>(
                 &GooseFEM::MatrixDiagonalPartitioned::Reaction_p, py::const_),
-            "Reaction_p",
-            py::arg("x_u"),
-            py::arg("x_p"))
+             "See :cpp:func:`GooseFEM::MatrixDiagonalPartitioned::Reaction_p`.",
+             py::arg("x_u"),
+             py::arg("x_p"))
 
         .def("__repr__", [](const GooseFEM::MatrixDiagonalPartitioned&) {
             return "<GooseFEM.MatrixDiagonalPartitioned>";
         });
 }
diff --git a/python/MatrixPartitioned.hpp b/python/MatrixPartitioned.hpp
index 4b6c789..6c847ec 100644
--- a/python/MatrixPartitioned.hpp
+++ b/python/MatrixPartitioned.hpp
@@ -1,154 +1,107 @@
 /* =================================================================================================
 
 (c - GPLv3) T.W.J. de Geus (Tom) | tom@geus.me | www.geus.me | github.com/tdegeus/GooseFEM
 
 ================================================================================================= */
 
 #include <Eigen/Eigen>
-#include <GooseFEM/GooseFEM.h>
+#include <GooseFEM/MatrixPartitioned.h>
 #include <pybind11/eigen.h>
 #include <pybind11/pybind11.h>
 #include <pyxtensor/pyxtensor.hpp>
 
 namespace py = pybind11;
 
 void init_MatrixPartitioned(py::module& m)
 {
 
-    py::class_<GooseFEM::MatrixPartitioned>(m, "MatrixPartitioned")
+    py::class_<GooseFEM::MatrixPartitioned, GooseFEM::Matrix>(m, "MatrixPartitioned")
 
-        .def(
-            py::init<
+        .def(py::init<
                 const xt::xtensor<size_t, 2>&,
                 const xt::xtensor<size_t, 2>&,
                 const xt::xtensor<size_t, 1>&>(),
-            "Sparse, partitioned, matrix",
-            py::arg("conn"),
-            py::arg("dofs"),
-            py::arg("iip"))
+             "Sparse, partitioned, matrix",
+             py::arg("conn"),
+             py::arg("dofs"),
+             py::arg("iip"))
 
-        .def("nelem", &GooseFEM::MatrixPartitioned::nelem, "Number of element")
+        .def("nnu",
+             &GooseFEM::MatrixPartitioned::nnu,
+             "See :cpp:func:`GooseFEM::MatrixPartitioned::nnu`.")
 
-        .def("nne", &GooseFEM::MatrixPartitioned::nne, "Number of nodes per element")
+        .def("nnp",
+             &GooseFEM::MatrixPartitioned::nnp,
+             "See :cpp:func:`GooseFEM::MatrixPartitioned::nnp`.")
 
-        .def("nnode", &GooseFEM::MatrixPartitioned::nnode, "Number of nodes")
+        .def("iiu",
+             &GooseFEM::MatrixPartitioned::iiu,
+             "See :cpp:func:`GooseFEM::MatrixPartitioned::iiu`.")
 
-        .def("ndim", &GooseFEM::MatrixPartitioned::ndim, "Number of dimensions")
+        .def("iip",
+             &GooseFEM::MatrixPartitioned::iip,
+             "See :cpp:func:`GooseFEM::MatrixPartitioned::iip`.")
 
-        .def("ndof", &GooseFEM::MatrixPartitioned::ndof, "Number of degrees-of-freedom")
-
-        .def("nnu", &GooseFEM::MatrixPartitioned::nnu, "Number of unknown degrees-of-freedom")
-
-        .def("nnp", &GooseFEM::MatrixPartitioned::nnp, "Number of prescribed degrees-of-freedom")
-
-        .def("dofs", &GooseFEM::MatrixPartitioned::dofs, "Return degrees-of-freedom")
-
-        .def("iiu", &GooseFEM::MatrixPartitioned::iiu, "Return unknown degrees-of-freedom")
-
-        .def("iip", &GooseFEM::MatrixPartitioned::iip, "Return prescribed degrees-of-freedom")
-
-        .def(
-            "assemble",
-            &GooseFEM::MatrixPartitioned::assemble,
-            "Assemble matrix from 'elemmat",
-            py::arg("elemmat"))
-
-        .def(
-            "set",
-            &GooseFEM::MatrixPartitioned::set,
-            "Overwrite with a dense (sub-) matrix",
-            py::arg("rows"),
-            py::arg("cols"),
-            py::arg("matrix"))
-
-        .def(
-            "add",
-            &GooseFEM::MatrixPartitioned::add,
-            "Add a dense (sub-) matrix to the current matrix",
-            py::arg("rows"),
-            py::arg("cols"),
-            py::arg("matrix"))
-
-        .def("Todense", &GooseFEM::MatrixPartitioned::Todense, "Return as dense matrix")
-
-        .def(
-            "Dot",
-            py::overload_cast<const xt::xtensor<double, 1>&>(&GooseFEM::MatrixPartitioned::Dot, py::const_),
-            "Dot",
-            py::arg("x"))
-
-        .def(
-            "Dot",
-            py::overload_cast<const xt::xtensor<double, 2>&>(&GooseFEM::MatrixPartitioned::Dot, py::const_),
-            "Dot",
-            py::arg("x"))
-
-        .def(
-            "Reaction",
-            py::overload_cast<const xt::xtensor<double, 1>&, const xt::xtensor<double, 1>&>(
+        .def("Reaction",
+             py::overload_cast<const xt::xtensor<double, 1>&, const xt::xtensor<double, 1>&>(
                 &GooseFEM::MatrixPartitioned::Reaction, py::const_),
-            "Reaction",
-            py::arg("x"),
-            py::arg("b"))
+             "See :cpp:func:`GooseFEM::MatrixPartitioned::Reaction`.",
+             py::arg("x"),
+             py::arg("b"))
 
-        .def(
-            "Reaction",
-            py::overload_cast<const xt::xtensor<double, 2>&, const xt::xtensor<double, 2>&>(
+        .def("Reaction",
+             py::overload_cast<const xt::xtensor<double, 2>&, const xt::xtensor<double, 2>&>(
                 &GooseFEM::MatrixPartitioned::Reaction, py::const_),
-            "Reaction",
-            py::arg("x"),
-            py::arg("b"))
+             "See :cpp:func:`GooseFEM::MatrixPartitioned::Reaction`.",
+             py::arg("x"),
+             py::arg("b"))
 
-        .def(
-            "Reaction_p",
-            py::overload_cast<const xt::xtensor<double, 1>&, const xt::xtensor<double, 1>&>(
+        .def("Reaction_p",
+             py::overload_cast<const xt::xtensor<double, 1>&, const xt::xtensor<double, 1>&>(
                 &GooseFEM::MatrixPartitioned::Reaction_p, py::const_),
-            "Reaction_p",
-            py::arg("x_u"),
-            py::arg("x_p"))
+             "See :cpp:func:`GooseFEM::MatrixPartitioned::Reaction_p`.",
+             py::arg("x_u"),
+             py::arg("x_p"))
 
         .def("__repr__", [](const GooseFEM::MatrixPartitioned&) {
             return "<GooseFEM.MatrixPartitioned>";
         });
 
     py::class_<GooseFEM::MatrixPartitionedSolver<>>(m, "MatrixPartitionedSolver")
 
         .def(py::init<>(), "Sparse, partitioned, matrix solver")
 
-        .def(
-            "Solve",
-            py::overload_cast<
+        .def("Solve",
+             py::overload_cast<
                 GooseFEM::MatrixPartitioned&,
                 const xt::xtensor<double, 1>&,
                 const xt::xtensor<double, 1>&>(&GooseFEM::MatrixPartitionedSolver<>::Solve),
-            "Solve",
-            py::arg("matrix"),
-            py::arg("b"),
-            py::arg("x"))
-
-        .def(
-            "Solve",
-            py::overload_cast<
+             "See :cpp:func:`GooseFEM::MatrixPartitionedSolver::Solve`.",
+             py::arg("matrix"),
+             py::arg("b"),
+             py::arg("x"))
+
+        .def("Solve",
+             py::overload_cast<
                 GooseFEM::MatrixPartitioned&,
                 const xt::xtensor<double, 2>&,
                 const xt::xtensor<double, 2>&>(&GooseFEM::MatrixPartitionedSolver<>::Solve),
-            "Solve",
-            py::arg("matrix"),
-            py::arg("b"),
-            py::arg("x"))
-
-        .def(
-            "Solve_u",
-            py::overload_cast<
+             "See :cpp:func:`GooseFEM::MatrixPartitionedSolver::Solve`.",
+             py::arg("matrix"),
+             py::arg("b"),
+             py::arg("x"))
+
+        .def("Solve_u",
+             py::overload_cast<
                 GooseFEM::MatrixPartitioned&,
                 const xt::xtensor<double, 1>&,
                 const xt::xtensor<double, 1>&>(&GooseFEM::MatrixPartitionedSolver<>::Solve_u),
-            "Solve_u",
-            py::arg("matrix"),
-            py::arg("b_u"),
-            py::arg("x_p"))
+             "See :cpp:func:`GooseFEM::MatrixPartitionedSolver::Solve_u`.",
+             py::arg("matrix"),
+             py::arg("b_u"),
+             py::arg("x_p"))
 
         .def("__repr__", [](const GooseFEM::MatrixPartitionedSolver<>&) {
             return "<GooseFEM.MatrixPartitionedSolver>";
         });
 }
diff --git a/test/basic/CMakeLists.txt b/test/basic/CMakeLists.txt
index 725bba6..b160aff 100644
--- a/test/basic/CMakeLists.txt
+++ b/test/basic/CMakeLists.txt
@@ -1,54 +1,55 @@
 
 cmake_minimum_required(VERSION 3.0)
 
 if(CMAKE_CURRENT_SOURCE_DIR STREQUAL CMAKE_SOURCE_DIR)
     project(GooseFEM-test-basic)
     find_package(GooseFEM REQUIRED CONFIG)
 endif()
 
 set(ASSERT ON)
 set(DEBUG ON)
 option(XSIMD "Use xsimd optimisations" ON)
 
 set(test_name "test-basic")
 
 find_package(Catch2 REQUIRED)
 find_package(xtensor REQUIRED)
 
 add_executable(${test_name}
     main.cpp
     version.cpp
     Allocate.cpp
     Element.cpp
     ElementHex8.cpp
     ElementQuad4.cpp
     ElementQuad4Planar.cpp
     Iterate.cpp
     Matrix.cpp
     MatrixDiagonal.cpp
+    MatrixDiagonalPartitioned.cpp
     Mesh.cpp
     MeshQuad4.cpp
     MeshHex8.cpp
     Vector.cpp
     VectorPartitioned.cpp)
 
 target_link_libraries(${test_name} PRIVATE
     Catch2::Catch2
     GooseFEM
     GooseFEM::compiler_warnings
     xtensor::optimize)
 
 if(ASSERT)
     target_link_libraries(${test_name} PRIVATE GooseFEM::assert)
 endif()
 
 if(DEBUG)
     target_link_libraries(${test_name} PRIVATE GooseFEM::debug)
 endif()
 
 if(SIMD)
     find_package(xsimd REQUIRED)
     target_link_libraries(${test_name} PRIVATE xtensor::use_xsimd)
 endif()
 
 add_test(NAME ${test_name} COMMAND ${test_name})
diff --git a/test/basic/Matrix.cpp b/test/basic/Matrix.cpp
index 70efe7b..d22f5ac 100644
--- a/test/basic/Matrix.cpp
+++ b/test/basic/Matrix.cpp
@@ -1,97 +1,95 @@
-
 #include <catch2/catch.hpp>
 #include <xtensor/xrandom.hpp>
-#include <xtensor/xmath.hpp>
-#include <Eigen/Eigen>
-#include <GooseFEM/GooseFEM.h>
+#include <GooseFEM/Matrix.h>
+#include <GooseFEM/MeshQuad4.h>
 
 #define ISCLOSE(a,b) REQUIRE_THAT((a), Catch::WithinAbs((b), 1.e-12));
 
 TEST_CASE("GooseFEM::Matrix", "Matrix.h")
 {
     SECTION("solve")
     {
         GooseFEM::Mesh::Quad4::Regular mesh(2, 2);
 
         size_t nne = mesh.nne();
         size_t ndim = mesh.ndim();
         size_t nelem = mesh.nelem();
         size_t nnode = mesh.nnode();
 
         xt::xtensor<double, 3> a = xt::empty<double>({nelem, nne * ndim, nne * ndim});
         xt::xtensor<double, 1> b = xt::random::rand<double>({nnode * ndim});
 
         for (size_t e = 0; e < nelem; ++e) {
             xt::xtensor<double, 2> ae = xt::random::rand<double>({nne * ndim, nne * ndim});
             ae = (ae + xt::transpose(ae)) / 2.0;
             xt::view(a, e, xt::all(), xt::all()) = ae;
         }
 
         GooseFEM::Matrix A(mesh.conn(), mesh.dofs());
         GooseFEM::MatrixSolver<> Solver;
         A.assemble(a);
         xt::xtensor<double, 1> C = A.Dot(b);
         xt::xtensor<double, 1> B = Solver.Solve(A, C);
 
         REQUIRE(B.size() == b.size());
         REQUIRE(xt::allclose(B, b));
     }
 
     SECTION("set/add/dot/solve - dofval")
     {
         xt::xtensor<double, 2> a = xt::random::rand<double>({10, 10});
         xt::xtensor<double, 1> x = xt::random::rand<double>({10});
         xt::xtensor<double, 1> b = xt::zeros<double>({10});
 
         xt::xtensor<double, 2> A = a + xt::transpose(a);
 
         for (size_t i = 0; i < A.shape(0); ++i) {
             for (size_t j = 0; j < A.shape(1); ++j) {
                 b(i) += A(i, j) * x(j);
             }
         }
 
         xt::xtensor<size_t, 2> conn = xt::zeros<size_t>({1, 5});
         xt::xtensor<size_t, 2> dofs = xt::arange<size_t>(10).reshape({5, 2});
 
         GooseFEM::Matrix K(conn, dofs);
         GooseFEM::MatrixSolver<> Solver;
         K.set(xt::arange<size_t>(10), xt::arange<size_t>(10), a);
         K.add(xt::arange<size_t>(10), xt::arange<size_t>(10), xt::transpose(a));
 
         REQUIRE(xt::allclose(A, K.Todense()));
         REQUIRE(xt::allclose(b, K.Dot(x)));
         REQUIRE(xt::allclose(x, Solver.Solve(K, b)));
     }
 
     SECTION("set/add/dot/solve - nodevec")
     {
         xt::xtensor<double, 2> a = xt::random::rand<double>({10, 10});
         xt::xtensor<double, 2> x = xt::random::rand<double>({5, 2});
         xt::xtensor<double, 2> b = xt::zeros<double>({5, 2});
 
         xt::xtensor<double, 2> A = a + xt::transpose(a);
 
         for (size_t m = 0; m < x.shape(0); ++m) {
             for (size_t n = 0; n < x.shape(0); ++n) {
                 for (size_t i = 0; i < x.shape(1); ++i) {
                     for (size_t j = 0; j < x.shape(1); ++j) {
                         b(m, i) += A(m * x.shape(1) + i, n * x.shape(1) + j) * x(n, j);
                     }
                 }
             }
         }
 
         xt::xtensor<size_t, 2> conn = xt::zeros<size_t>({1, 5});
         xt::xtensor<size_t, 2> dofs = xt::arange<size_t>(10).reshape({5, 2});
 
         GooseFEM::Matrix K(conn, dofs);
         GooseFEM::MatrixSolver<> Solver;
         K.set(xt::arange<size_t>(10), xt::arange<size_t>(10), a);
         K.add(xt::arange<size_t>(10), xt::arange<size_t>(10), xt::transpose(a));
 
         REQUIRE(xt::allclose(A, K.Todense()));
         REQUIRE(xt::allclose(b, K.Dot(x)));
         REQUIRE(xt::allclose(x, Solver.Solve(K, b)));
     }
 }
diff --git a/test/basic/MatrixDiagonal.cpp b/test/basic/MatrixDiagonal.cpp
index df81af1..939551e 100644
--- a/test/basic/MatrixDiagonal.cpp
+++ b/test/basic/MatrixDiagonal.cpp
@@ -1,44 +1,78 @@
-
 #include <catch2/catch.hpp>
 #include <xtensor/xrandom.hpp>
-#include <xtensor/xmath.hpp>
-#include <GooseFEM/GooseFEM.h>
+#include <GooseFEM/MatrixDiagonal.h>
+#include <GooseFEM/MeshQuad4.h>
+#include <GooseFEM/ElementQuad4.h>
+#include <GooseFEM/Vector.h>
 
 #define ISCLOSE(a,b) REQUIRE_THAT((a), Catch::WithinAbs((b), 1.e-12));
 
 TEST_CASE("GooseFEM::MatrixDiagonal", "MatrixDiagonal.h")
 {
+    SECTION("assemble")
+    {
+        GooseFEM::Mesh::Quad4::Regular mesh(2, 2);
+        GooseFEM::Vector vector(mesh.conn(), mesh.dofs());
+        GooseFEM::MatrixDiagonal A(mesh.conn(), mesh.dofs());
+
+        GooseFEM::Element::Quad4::Quadrature quad(
+            vector.AsElement(mesh.coor()),
+            GooseFEM::Element::Quad4::Nodal::xi(),
+            GooseFEM::Element::Quad4::Nodal::w());
+
+        xt::xtensor<double, 2> val_quad = xt::empty<double>({mesh.nelem(), quad.nip()});
+        for (size_t q = 0; q < quad.nip(); ++q) {
+            xt::view(val_quad, xt::all(), q) = 1.0;
+        }
+
+        A.assemble(quad.Int_N_scalar_NT_dV(val_quad));
+
+        xt::xtensor<double, 1> a = {0.25, 0.25,
+                                    0.5 , 0.5 ,
+                                    0.25, 0.25,
+                                    0.5 , 0.5 ,
+                                    1.0 , 1.0 ,
+                                    0.5 , 0.5 ,
+                                    0.25, 0.25,
+                                    0.5 , 0.5 ,
+                                    0.25, 0.25};
+
+        REQUIRE(xt::allclose(A.Todiagonal(), a));
+    }
 
     SECTION("dot")
     {
         GooseFEM::Mesh::Quad4::Regular mesh(2, 2);
 
         xt::xtensor<double, 1> a = xt::random::rand<double>({mesh.nnode() * mesh.ndim()});
-        xt::xtensor<double, 1> b = xt::random::rand<double>({mesh.nnode() * mesh.ndim()});
-        xt::xtensor<double, 1> c = a * b;
+        xt::xtensor<double, 1> x = xt::random::rand<double>({mesh.nnode() * mesh.ndim()});
+        xt::xtensor<double, 1> b = a * x;
 
         GooseFEM::MatrixDiagonal A(mesh.conn(), mesh.dofs());
         A.set(a);
-        xt::xtensor<double, 1> C = A.Dot(b);
+        xt::xtensor<double, 1> B = A.Dot(x);
 
-        REQUIRE(C.size() == c.size());
-        REQUIRE(xt::allclose(C, c));
+        REQUIRE(B.size() == b.size());
+        REQUIRE(xt::allclose(B, b));
+        REQUIRE(xt::allclose(A.Todiagonal(), a));
     }
 
     SECTION("solve")
     {
         GooseFEM::Mesh::Quad4::Regular mesh(2, 2);
 
         xt::xtensor<double, 1> a = xt::random::rand<double>({mesh.nnode() * mesh.ndim()});
-        xt::xtensor<double, 1> b = xt::random::rand<double>({mesh.nnode() * mesh.ndim()});
-        xt::xtensor<double, 1> c = a * b;
+        xt::xtensor<double, 1> x = xt::random::rand<double>({mesh.nnode() * mesh.ndim()});
+        xt::xtensor<double, 1> b = a * x;
 
         GooseFEM::MatrixDiagonal A(mesh.conn(), mesh.dofs());
         A.set(a);
-        xt::xtensor<double, 1> C = A.Dot(b);
-        xt::xtensor<double, 1> B = A.Solve(C);
+        xt::xtensor<double, 1> B = A.Dot(x);
+        xt::xtensor<double, 1> X = A.Solve(B);
 
         REQUIRE(B.size() == b.size());
+        REQUIRE(X.size() == x.size());
         REQUIRE(xt::allclose(B, b));
+        REQUIRE(xt::allclose(X, x));
     }
 }
diff --git a/test/basic/MatrixDiagonalPartitioned.cpp b/test/basic/MatrixDiagonalPartitioned.cpp
new file mode 100644
index 0000000..d7affd2
--- /dev/null
+++ b/test/basic/MatrixDiagonalPartitioned.cpp
@@ -0,0 +1,82 @@
+#include <catch2/catch.hpp>
+#include <xtensor/xrandom.hpp>
+#include <GooseFEM/MatrixDiagonalPartitioned.h>
+#include <GooseFEM/MeshQuad4.h>
+#include <GooseFEM/ElementQuad4.h>
+#include <GooseFEM/VectorPartitioned.h>
+
+#define ISCLOSE(a,b) REQUIRE_THAT((a), Catch::WithinAbs((b), 1.e-12));
+
+TEST_CASE("GooseFEM::MatrixDiagonalPartitioned", "MatrixDiagonalPartitioned.h")
+{
+    SECTION("assemble")
+    {
+        GooseFEM::Mesh::Quad4::Regular mesh(2, 2);
+        xt::xtensor<size_t, 1> iip = {0, 2};
+        GooseFEM::VectorPartitioned vector(mesh.conn(), mesh.dofs(), iip);
+        GooseFEM::MatrixDiagonalPartitioned A(mesh.conn(), mesh.dofs(), iip);
+
+        GooseFEM::Element::Quad4::Quadrature quad(
+            vector.AsElement(mesh.coor()),
+            GooseFEM::Element::Quad4::Nodal::xi(),
+            GooseFEM::Element::Quad4::Nodal::w());
+
+        xt::xtensor<double, 2> val_quad = xt::empty<double>({mesh.nelem(), quad.nip()});
+        for (size_t q = 0; q < quad.nip(); ++q) {
+            xt::view(val_quad, xt::all(), q) = 1.0;
+        }
+
+        A.assemble(quad.Int_N_scalar_NT_dV(val_quad));
+
+        xt::xtensor<double, 1> a = {0.25, 0.25,
+                                    0.5 , 0.5 ,
+                                    0.25, 0.25,
+                                    0.5 , 0.5 ,
+                                    1.0 , 1.0 ,
+                                    0.5 , 0.5 ,
+                                    0.25, 0.25,
+                                    0.5 , 0.5 ,
+                                    0.25, 0.25};
+
+        REQUIRE(xt::allclose(A.Todiagonal(), a));
+    }
+
+    SECTION("dot")
+    {
+        GooseFEM::Mesh::Quad4::Regular mesh(2, 2);
+        xt::xtensor<size_t, 1> iip = {0, 2};
+
+        xt::xtensor<double, 1> a = xt::random::rand<double>({mesh.nnode() * mesh.ndim()});
+        xt::xtensor<double, 1> x = xt::random::rand<double>({mesh.nnode() * mesh.ndim()});
+        xt::xtensor<double, 1> b = a * x;
+
+        GooseFEM::MatrixDiagonalPartitioned A(mesh.conn(), mesh.dofs(), iip);
+        A.set(a);
+        xt::xtensor<double, 1> B = A.Dot(x);
+
+        REQUIRE(B.size() == b.size());
+        REQUIRE(xt::allclose(B, b));
+        REQUIRE(xt::allclose(A.Todiagonal(), a));
+    }
+
+    SECTION("solve")
+    {
+        GooseFEM::Mesh::Quad4::Regular mesh(2, 2);
+        xt::xtensor<size_t, 1> iip = {0, 2};
+
+        xt::xtensor<double, 1> a = xt::random::rand<double>({mesh.nnode() * mesh.ndim()});
+        xt::xtensor<double, 1> x = xt::random::rand<double>({mesh.nnode() * mesh.ndim()});
+        xt::xtensor<double, 1> b = a * x;
+
+        GooseFEM::MatrixDiagonalPartitioned A(mesh.conn(), mesh.dofs(), iip);
+        A.set(a);
+        xt::xtensor<double, 1> B = A.Dot(x);
+        xt::xtensor<double, 1> X = A.Solve(B);
+        xt::view(X, xt::keep(iip)) = xt::view(x, xt::keep(iip));
+
+        REQUIRE(B.size() == b.size());
+        REQUIRE(X.size() == x.size());
+        REQUIRE(xt::allclose(B, b));
+        REQUIRE(xt::allclose(X, x));
+    }
+}
diff --git a/test/basic/MatrixParitioned.cpp b/test/basic/MatrixParitioned.cpp
index 7a5736a..eae90f2 100644
--- a/test/basic/MatrixParitioned.cpp
+++ b/test/basic/MatrixParitioned.cpp
@@ -1,110 +1,110 @@
 
 #include <catch2/catch.hpp>
 #include <xtensor/xrandom.hpp>
 #include <xtensor/xmath.hpp>
 #include <Eigen/Eigen>
 #include <GooseFEM/GooseFEM.h>
 
 #define ISCLOSE(a,b) REQUIRE_THAT((a), Catch::WithinAbs((b), 1.e-12));
 
 TEST_CASE("GooseFEM::MatrixPartitioned", "MatrixPartitioned.h")
 {
     SECTION("solve")
     {
         GooseFEM::Mesh::Quad4::Regular mesh(2, 2);
 
         size_t nne = mesh.nne();
         size_t ndim = mesh.ndim();
         size_t nelem = mesh.nelem();
         size_t nnode = mesh.nnode();
         auto dofs = mesh.dofs();
-        size_t npp = xt::amax(dofs)();
-        npp = (npp - npp % 2) / 2;
-        xt::xtensor<size_t, 1> iip = xt::arange<size_t<(npp);
+        size_t nnp = xt::amax(dofs)();
+        nnp = (nnp - nnp % 2) / 2;
+        xt::xtensor<size_t, 1> iip = xt::arange<size_t<(nnp);
 
         xt::xtensor<double, 3> a = xt::empty<double>({nelem, nne * ndim, nne * ndim});
         xt::xtensor<double, 1> b = xt::random::rand<double>({nnode * ndim});
 
         for (size_t e = 0; e < nelem; ++e) {
             xt::xtensor<double, 2> ae = xt::random::rand<double>({nne * ndim, nne * ndim});
             ae = (ae + xt::transpose(ae)) / 2.0;
             xt::view(a, e, xt::all(), xt::all()) = ae;
         }
 
         GooseFEM::MatrixPartitioned A(mesh.conn(), dofs, iip);
         GooseFEM::MatrixPartitionedSolver<> Solver;
         A.assemble(a);
         xt::xtensor<double, 1> C = A.Dot(b);
         xt::xtensor<double, 1> B = Solver.Solve(A, C);
 
         REQUIRE(B.size() == b.size());
         REQUIRE(xt::allclose(B, b));
 
         // check that allocating a different Solver instance still works
         GooseFEM::MatrixPartitionedSolver<> NewSolver;
         xt::xtensor<double, 1> NB = NewSolver.Solve(A, C);
 
         REQUIRE(NB.size() == b.size());
         REQUIRE(xt::allclose(NB, b));
     }
 
     SECTION("set/add/dot/solve - dofval")
     {
         xt::xtensor<double, 2> a = xt::random::rand<double>({10, 10});
         xt::xtensor<double, 1> x = xt::random::rand<double>({10});
         xt::xtensor<double, 1> b = xt::zeros<double>({10});
 
         xt::xtensor<double, 2> A = a + xt::transpose(a);
 
         for (size_t i = 0; i < A.shape(0); ++i) {
             for (size_t j = 0; j < A.shape(1); ++j) {
                 b(i) += A(i, j) * x(j);
             }
         }
 
         xt::xtensor<size_t, 2> conn = xt::zeros<size_t>({1, 5});
         xt::xtensor<size_t, 2> dofs = xt::arange<size_t>(10).reshape({5, 2});
-        xt::xtensor<size_t, 1> iip = xt::arange<size_t<(5);
+        xt::xtensor<size_t, 1> iip = {0, 2, 4};
 
         GooseFEM::MatrixPartitioned K(conn, dofs, iip);
         GooseFEM::MatrixPartitionedSolver<> Solver;
         K.set(xt::arange<size_t>(10), xt::arange<size_t>(10), a);
         K.add(xt::arange<size_t>(10), xt::arange<size_t>(10), xt::transpose(a));
 
         REQUIRE(xt::allclose(A, K.Todense()));
         REQUIRE(xt::allclose(b, K.Dot(x)));
         REQUIRE(xt::allclose(x, Solver.Solve(K, b)));
     }
 
     SECTION("set/add/dot/solve - nodevec")
     {
         xt::xtensor<double, 2> a = xt::random::rand<double>({10, 10});
         xt::xtensor<double, 2> x = xt::random::rand<double>({5, 2});
         xt::xtensor<double, 2> b = xt::zeros<double>({5, 2});
 
         xt::xtensor<double, 2> A = a + xt::transpose(a);
 
         for (size_t m = 0; m < x.shape(0); ++m) {
             for (size_t n = 0; n < x.shape(0); ++n) {
                 for (size_t i = 0; i < x.shape(1); ++i) {
                     for (size_t j = 0; j < x.shape(1); ++j) {
                         b(m, i) += A(m * x.shape(1) + i, n * x.shape(1) + j) * x(n, j);
                     }
                 }
             }
         }
 
         xt::xtensor<size_t, 2> conn = xt::zeros<size_t>({1, 5});
         xt::xtensor<size_t, 2> dofs = xt::arange<size_t>(10).reshape({5, 2});
-        xt::xtensor<size_t, 1> iip = xt::arange<size_t<(5);
+        xt::xtensor<size_t, 1> iip = {0, 2, 4};
 
         GooseFEM::MatrixPartitioned K(conn, dofs, iip);
         GooseFEM::MatrixPartitionedSolver<> Solver;
         K.set(xt::arange<size_t>(10), xt::arange<size_t>(10), a);
         K.add(xt::arange<size_t>(10), xt::arange<size_t>(10), xt::transpose(a));
 
         REQUIRE(xt::allclose(A, K.Todense()));
         REQUIRE(xt::allclose(b, K.Dot(x)));
         REQUIRE(xt::allclose(x, Solver.Solve(K, b)));
     }
 }