diff --git a/src/ProjectedMatrices.cc b/src/ProjectedMatrices.cc
index e50e98bc..74ff01ed 100644
--- a/src/ProjectedMatrices.cc
+++ b/src/ProjectedMatrices.cc
@@ -745,6 +745,23 @@ int ProjectedMatrices<MatrixType>::writeDM(HDFrestart& h5f_file)
     return dm_->write(h5f_file, name);
 }
 
+template <class MatrixType>
+int ProjectedMatrices<MatrixType>::writeSavedDM(HDFrestart& h5f_file)
+{
+    std::string name("/Density_Matrix_WF");
+
+    ReplicatedWorkSpace<double>& wspace(
+        ReplicatedWorkSpace<double>::instance());
+
+    const MatrixType* matrix = mat_X_old_.get();
+    wspace.initSquareMatrix(*matrix);
+
+    DISTMATDTYPE* work_matrix = wspace.square_matrix();
+
+    hid_t file_id = h5f_file.file_id();
+    return mgmol_tools::write_matrix(file_id, name, work_matrix, dim_);
+}
+
 template <class MatrixType>
 int ProjectedMatrices<MatrixType>::readDM(HDFrestart& h5f_file)
 {
@@ -752,6 +769,16 @@ int ProjectedMatrices<MatrixType>::readDM(HDFrestart& h5f_file)
     return dm_->read(h5f_file, name);
 }
 
+template <class MatrixType>
+int ProjectedMatrices<MatrixType>::readWFDM(HDFrestart& h5f_file)
+{
+    MGmol_MPI& mmpi = *(MGmol_MPI::instance());
+    mmpi.barrier();
+    if (mmpi.PE0()) std::cout << "ProjectedMatrices::readWFDM..." << std::endl;
+    std::string name("/Density_Matrix_WF");
+    return dm_->read(h5f_file, name);
+}
+
 template <class MatrixType>
 void ProjectedMatrices<MatrixType>::printEigenvalues(std::ostream& os) const
 {
diff --git a/src/ProjectedMatrices.h b/src/ProjectedMatrices.h
index 9bc5e3c2..f2881763 100644
--- a/src/ProjectedMatrices.h
+++ b/src/ProjectedMatrices.h
@@ -319,7 +319,9 @@ class ProjectedMatrices : public ProjectedMatricesInterface
     double computeEntropyWithCheb(const double kbt);
     double checkCond(const double tol, const bool flag = true) override;
     int writeDM(HDFrestart& h5f_file) override;
+    int writeSavedDM(HDFrestart& h5f_file);
     int readDM(HDFrestart& h5f_file) override;
+    int readWFDM(HDFrestart& h5f_file);
     void printEigenvalues(std::ostream& os) const;
     void updateDM(const int iterative_index) override;
     void updateDMwithEigenstates(const int iterative_index);
diff --git a/src/ProjectedMatricesInterface.h b/src/ProjectedMatricesInterface.h
index e02e0b45..c52595e1 100644
--- a/src/ProjectedMatricesInterface.h
+++ b/src/ProjectedMatricesInterface.h
@@ -273,6 +273,14 @@ class ProjectedMatricesInterface : public ChebyshevApproximationFunction
 
         return 0;
     }
+    virtual int readWFDM(HDFrestart& h5f_file)
+    {
+        (void)h5f_file;
+
+        exitWithErrorMessage("readWFDM");
+
+        return 0;
+    }
     virtual int writeDM(HDFrestart& h5f_file)
     {
         (void)h5f_file;
@@ -281,6 +289,14 @@ class ProjectedMatricesInterface : public ChebyshevApproximationFunction
 
         return 0;
     }
+    virtual int writeSavedDM(HDFrestart& h5f_file)
+    {
+        (void)h5f_file;
+
+        exitWithErrorMessage("writeSavedDM");
+
+        return 0;
+    }
     virtual void updateDMwithChebApproximation(const int iterative_index)
     {
         (void)iterative_index;
diff --git a/src/md.cc b/src/md.cc
index b773c27c..b50618c8 100644
--- a/src/md.cc
+++ b/src/md.cc
@@ -275,6 +275,10 @@ int MGmol<OrbitalsType>::dumpMDrestartFile(OrbitalsType& orbitals, Ions& ions,
                                  << std::endl;
             return ierr;
         }
+
+        // write DM associated with non-extrapolated wavefunctions
+        // (last computed solution of KS equations)
+        proj_matrices_->writeSavedDM(h5file);
     }
 
     ierr = h5file.close();
@@ -579,6 +583,13 @@ void MGmol<OrbitalsType>::md(OrbitalsType** orbitals, Ions& ions)
             lrs_->clearOldCenters();
         }
 
+        // save DM for possible restart write
+        // note: extrapolation is going to modify it!
+        if ((ct.out_restart_info > 2)
+            && (((md_iteration_ % ct.checkpoint) == 0)
+                   || (mdstep == ct.num_MD_steps)))
+            proj_matrices_->saveDM();
+
         preWFextrapolation();
 
         if (ct.dt > 0.
@@ -656,6 +667,7 @@ void MGmol<OrbitalsType>::md(OrbitalsType** orbitals, Ions& ions)
 template <class OrbitalsType>
 void MGmol<OrbitalsType>::loadRestartFile(const std::string filename)
 {
+    if (onpe0) std::cout << "loadRestartFile..." << std::endl;
     MGmol_MPI& mmpi(*(MGmol_MPI::instance()));
     Control& ct              = *(Control::instance());
     Mesh* mymesh             = Mesh::instance();
@@ -678,6 +690,12 @@ void MGmol<OrbitalsType>::loadRestartFile(const std::string filename)
 
         global_exit(0);
     }
+    if (!ct.fullyOccupied())
+    {
+        // overwrite DM with restart data in dataset Density_Matrix_WF
+        if (h5file.checkDataExists("Density_Matrix_WF"))
+            ierr = proj_matrices_->readWFDM(h5file);
+    }
 
     ierr = h5file.close();
     mmpi.allreduce(&ierr, 1, MPI_MIN);
diff --git a/src/restart.cc b/src/restart.cc
index c12bb258..2fa76ecb 100644
--- a/src/restart.cc
+++ b/src/restart.cc
@@ -168,7 +168,7 @@ int MGmol<OrbitalsType>::write_hdf5(HDFrestart& h5f_file,
 
         if (!ct.fullyOccupied())
         {
-            int ierr = proj_matrices_->writeDM(h5f_file);
+            ierr = proj_matrices_->writeDM(h5f_file);
             if (ierr < 0) return ierr;
         }
         if (ct.isLocMode()