*** empty log message ***

mat.cc

@@ -1109,9 +1109,7 @@ const complex<double> NRMat< complex<double> >::trace() const
 template<>
 const double * NRMat<double>::diagonalof(NRVec<double> &r, const bool divide, bool cache) const
 {
-#ifdef DEBUG
 	if (r.size() != nn) laerror("diagonalof() incompatible vector");
-#endif
 
 double a;
 
@@ -1145,6 +1143,23 @@ return divide?NULL:&r[0];
 }
 
 
+//set diagonal
+template<>
+void NRMat<double>::diagonalset(const NRVec<double> &r)
+{
+	if (r.size() != nn) laerror("diagonalset() incompatible vector");
+	if(nn!=mm) laerror("diagonalset only for square matrix");
+
+copyonwrite();
+
+#ifdef MATPTR
+for (int i=0; i< nn; i++) v[i][i] = r[i];
+#else
+{int i,j; for (i=j=0; j< nn; ++j, i+=nn+1) v[i] = r[j];}
+#endif
+}
+
+
 
 
 

mat.h

@@ -92,6 +92,7 @@ public:
 	const NRVec<T> row(const int i, int l= -1) const; //row of, efficient
 	const NRVec<T> column(const int j, int l= -1) const {if(l<0) l=nn; NRVec<T> r(l); for(int i=0; i<l; ++i) r[i]= (*this)(i,j); return r;}; //column of, general but not very efficient
 	const T* diagonalof(NRVec<T> &, const bool divide=0, bool cache=false) const; //get diagonal 
+	void diagonalset(const NRVec<T> &); //set diagonal elements
 	void gemv(const T beta, NRVec<T> &r, const char trans, const T alpha, const NRVec<T> &x) const {r.gemv(beta,*this,trans,alpha,x);};
 	inline T* operator[](const int i);  //subscripting: pointer to row i
 	inline const T* operator[](const int i) const;