diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt
index 330208d0c..6c87b61de 100644
--- a/src/CMakeLists.txt
+++ b/src/CMakeLists.txt
@@ -41,7 +41,7 @@ set(_SOURCES
   "hubbard/hubbard_occupancies_derivatives.cpp"
   "hubbard/hubbard_potential_energy.cpp"
   "hubbard/hubbard_matrix.cpp"
-  "potential/generate_d_operator_matrix.cpp"
+  "potential/generate_d_mtrx.cpp"
   "potential/generate_pw_coeffs.cpp"
   "potential/paw_potential.cpp"
   "potential/poisson.cpp"
diff --git a/src/api/sirius_api.cpp b/src/api/sirius_api.cpp
index 00110685b..776696209 100644
--- a/src/api/sirius_api.cpp
+++ b/src/api/sirius_api.cpp
@@ -6421,7 +6421,7 @@ sirius_generate_d_operator_matrix(void* const* gs_handler__, int* error_code__)
     call_sirius(
             [&]() {
                 auto& gs = get_gs(gs_handler__);
-                gs.potential().generate_D_operator_matrix();
+                gs.potential().generate_d_mtrx();
             },
             error_code__);
 }
diff --git a/src/context/simulation_context.hpp b/src/context/simulation_context.hpp
index 6c5bcdb83..24b164d57 100644
--- a/src/context/simulation_context.hpp
+++ b/src/context/simulation_context.hpp
@@ -253,7 +253,8 @@ class Simulation_context : public Simulation_parameters
     memory_t host_memory_t_{memory_t::none};
 
     /// SPLA library context.
-    std::shared_ptr<::spla::Context> spla_ctx_{new ::spla::Context{SPLA_PU_HOST}};
+    /** Context is mutable to allow SPLA library to change it. */
+    mutable std::shared_ptr<::spla::Context> spla_ctx_{new ::spla::Context{SPLA_PU_HOST}};
 
     std::ostream* output_stream_{nullptr};
     std::ofstream output_file_stream_;
@@ -706,14 +707,8 @@ class Simulation_context : public Simulation_parameters
         return fft_coarse_grid_;
     }
 
-    auto const&
-    spla_context() const
-    {
-        return *spla_ctx_;
-    }
-
     auto&
-    spla_context()
+    spla_context() const
     {
         return *spla_ctx_;
     }
diff --git a/src/density/density.hpp b/src/density/density.hpp
index ca9d40014..838fa4438 100644
--- a/src/density/density.hpp
+++ b/src/density/density.hpp
@@ -70,7 +70,7 @@ inline auto
 get_rho_up_dn(int num_mag_dims__, double rho__, r3::vector<double> mag__)
 {
     if (rho__ < 0.0) {
-        return std::make_pair<double, double>(0.0, 0.0);
+        return std::make_pair(0.0, 0.0);
     }
 
     double mag{0};
@@ -85,7 +85,7 @@ get_rho_up_dn(int num_mag_dims__, double rho__, r3::vector<double> mag__)
         mag = std::min(mag__.length(), rho__);
     }
 
-    return std::make_pair<double, double>(0.5 * (rho__ + mag), 0.5 * (rho__ - mag));
+    return std::make_pair(0.5 * (rho__ + mag), 0.5 * (rho__ - mag));
 }
 
 /// PAW density storage.
diff --git a/src/geometry/force.cpp b/src/geometry/force.cpp
index c1df5f009..54f0cc4b5 100644
--- a/src/geometry/force.cpp
+++ b/src/geometry/force.cpp
@@ -95,7 +95,7 @@ Force::add_k_point_contribution(K_point<T>& kp__, mdarray<double, 2>& forces__)
     mdarray<real_type<F>, 2> f({3, ctx_.unit_cell().num_atoms()});
     f.zero();
 
-    add_k_point_contribution_nonlocal<T, F>(ctx_, bp_grad, kp__, f);
+    add_k_point_contribution_nonlocal<T, F>(potential_, bp_grad, kp__, f);
 
     for (int ia = 0; ia < ctx_.unit_cell().num_atoms(); ia++) {
         for (int x : {0, 1, 2}) {
@@ -275,7 +275,7 @@ Force::calc_forces_hubbard()
 
     if (ctx_.hubbard_correction()) {
         /* recompute the hubbard potential */
-        ::sirius::generate_potential(density_.occupation_matrix(), potential_.hubbard_potential());
+        //::sirius::generate_potential(density_.occupation_matrix(), potential_.hubbard_potential());
 
         Q_operator<double> q_op(ctx_);
 
@@ -385,9 +385,9 @@ Force::calc_forces_us()
     forces_us_ = mdarray<double, 2>({3, ctx_.unit_cell().num_atoms()});
     forces_us_.zero();
 
-    potential_.fft_transform(-1);
+    //potential_.fft_transform(-1);
 
-    Unit_cell& unit_cell = ctx_.unit_cell();
+    auto const& unit_cell = ctx_.unit_cell();
 
     double reduce_g_fact = ctx_.gvec().reduced() ? 2.0 : 1.0;
 
@@ -484,11 +484,11 @@ Force::calc_forces_scf_corr()
     forces_scf_corr_.zero();
 
     /* get main arrays */
-    auto& dveff = potential_.dveff();
+    auto const& dveff = potential_.dveff();
 
-    Unit_cell& unit_cell = ctx_.unit_cell();
+    auto const& unit_cell = ctx_.unit_cell();
 
-    fft::Gvec const& gvec = ctx_.gvec();
+    auto const& gvec = ctx_.gvec();
 
     int gvec_count  = gvec.count();
     int gvec_offset = gvec.offset();
@@ -499,7 +499,7 @@ Force::calc_forces_scf_corr()
 
     #pragma omp parallel for
     for (int ia = 0; ia < unit_cell.num_atoms(); ia++) {
-        Atom& atom = unit_cell.atom(ia);
+        auto const& atom = unit_cell.atom(ia);
 
         int iat = atom.type_id();
 
@@ -545,9 +545,9 @@ Force::calc_forces_core()
     /* transform from real space to reciprocal */
     xc_pot.rg().fft_transform(-1);
 
-    Unit_cell& unit_cell = ctx_.unit_cell();
+    auto const& unit_cell = ctx_.unit_cell();
 
-    fft::Gvec const& gvecs = ctx_.gvec();
+    auto const& gvecs = ctx_.gvec();
 
     int gvec_count  = gvecs.count();
     int gvec_offset = gvecs.offset();
@@ -557,7 +557,7 @@ Force::calc_forces_core()
     /* here the calculations are in lattice vectors space */
     #pragma omp parallel for
     for (int ia = 0; ia < unit_cell.num_atoms(); ia++) {
-        Atom& atom = unit_cell.atom(ia);
+        auto const& atom = unit_cell.atom(ia);
         if (atom.type().ps_core_charge_density().empty()) {
             continue;
         }
diff --git a/src/geometry/force.hpp b/src/geometry/force.hpp
index dbf1f6748..8faf5771c 100644
--- a/src/geometry/force.hpp
+++ b/src/geometry/force.hpp
@@ -35,7 +35,7 @@ class Force
   private:
     Simulation_context& ctx_;
 
-    const Density& density_;
+    Density const& density_;
 
     Potential& potential_;
 
diff --git a/src/geometry/non_local_functor.hpp b/src/geometry/non_local_functor.hpp
index 16840dc94..aa3bf892d 100644
--- a/src/geometry/non_local_functor.hpp
+++ b/src/geometry/non_local_functor.hpp
@@ -17,6 +17,7 @@
 #include "function3d/periodic_function.hpp"
 #include "k_point/k_point.hpp"
 #include "density/augmentation_operator.hpp"
+#include "potential/potential.hpp"
 #include "beta_projectors/beta_projectors_base.hpp"
 
 namespace sirius {
@@ -24,12 +25,14 @@ namespace sirius {
 /** \tparam T  Precision type of the wave-functions */
 template <typename T, typename F>
 void
-add_k_point_contribution_nonlocal(Simulation_context& ctx__, Beta_projectors_base<T>& bp_base__, K_point<T>& kp__,
+add_k_point_contribution_nonlocal(Potential& potential__, Beta_projectors_base<T>& bp_base__, K_point<T>& kp__,
                                   mdarray<real_type<F>, 2>& collect_res__)
 {
     PROFILE("sirius::add_k_point_contribution_nonlocal");
 
-    auto& uc = ctx__.unit_cell();
+    auto& ctx = potential__.ctx();
+
+    auto& uc = ctx.unit_cell();
 
     if (uc.max_mt_basis_size() == 0) {
         return;
@@ -39,7 +42,7 @@ add_k_point_contribution_nonlocal(Simulation_context& ctx__, Beta_projectors_bas
 
     double main_two_factor{-2};
 
-    auto mt          = ctx__.processing_unit_memory_t();
+    auto mt          = ctx.processing_unit_memory_t();
     auto bp_gen      = bp.make_generator(mt);
     auto beta_coeffs = bp_gen.prepare();
 
@@ -54,25 +57,25 @@ add_k_point_contribution_nonlocal(Simulation_context& ctx__, Beta_projectors_bas
         /* store <beta|psi> for spin up and down */
         matrix<F> beta_phi_chunks[2];
 
-        for (int ispn = 0; ispn < ctx__.num_spins(); ispn++) {
-            int nbnd              = kp__.num_occupied_bands(ispn);
-            beta_phi_chunks[ispn] = inner_prod_beta<F>(ctx__.spla_context(), mt, ctx__.host_memory_t(),
-                                                       is_device_memory(mt), beta_coeffs, kp__.spinor_wave_functions(),
-                                                       wf::spin_index(ispn), wf::band_range(0, nbnd));
+        for (int ispn = 0; ispn < ctx.num_spins(); ispn++) {
+            int nbnd = kp__.num_occupied_bands(ispn);
+            beta_phi_chunks[ispn] =
+                    inner_prod_beta<F>(ctx.spla_context(), mt, ctx.host_memory_t(), is_device_memory(mt), beta_coeffs,
+                                       kp__.spinor_wave_functions(), wf::spin_index(ispn), wf::band_range(0, nbnd));
         }
 
         for (int x = 0; x < bp_base__.num_comp(); x++) {
             /* generate chunk for inner product of beta gradient */
             bp_base_gen.generate(beta_coeffs_base, icnk, x);
 
-            for (int ispn = 0; ispn < ctx__.num_spins(); ispn++) {
+            for (int ispn = 0; ispn < ctx.num_spins(); ispn++) {
                 int spin_factor = (ispn == 0 ? 1 : -1);
 
                 int nbnd = kp__.num_occupied_bands(ispn);
 
                 /* inner product of beta gradient and WF */
                 auto bp_base_phi_chunk = inner_prod_beta<F>(
-                        ctx__.spla_context(), mt, ctx__.host_memory_t(), is_device_memory(mt), beta_coeffs_base,
+                        ctx.spla_context(), mt, ctx.host_memory_t(), is_device_memory(mt), beta_coeffs_base,
                         kp__.spinor_wave_functions(), wf::spin_index(ispn), wf::band_range(0, nbnd));
 
                 splindex_block<> spl_nbnd(nbnd, n_blocks(kp__.comm().size()), block_id(kp__.comm().rank()));
@@ -86,6 +89,8 @@ add_k_point_contribution_nonlocal(Simulation_context& ctx__, Beta_projectors_bas
                     int nbf  = beta_coeffs_base.beta_chunk_->desc_(beta_desc_idx::nbf, ia_chunk);
                     int iat  = uc.atom(ia).type_id();
 
+                    auto& d_mtrx = potential__.d_mtrx(ia);
+
                     if (uc.atom(ia).type().spin_orbit_coupling()) {
                         RTE_THROW("stress and forces with SO coupling are not upported");
                     }
@@ -120,14 +125,14 @@ add_k_point_contribution_nonlocal(Simulation_context& ctx__, Beta_projectors_bas
                             /* Qij exists only in the case of ultrasoft/PAW */
                             real_type<F> qij{0};
                             if (uc.atom(ia).type().augment()) {
-                                qij = ctx__.augmentation_op(iat).q_mtrx(ibf, jbf);
+                                qij = ctx.augmentation_op(iat).q_mtrx(ibf, jbf);
                             }
                             std::complex<real_type<F>> dij{0};
 
                             /* get non-magnetic or collinear spin parts of dij*/
-                            switch (ctx__.num_spins()) {
+                            switch (ctx.num_spins()) {
                                 case 1: {
-                                    dij = uc.atom(ia).d_mtrx(ibf, jbf, 0);
+                                    dij = d_mtrx(ibf, jbf, 0);
                                     if (lm1 == lm2) {
                                         dij += uc.atom(ia).type().d_mtrx_ion()(idxrf1, idxrf2);
                                     }
@@ -136,8 +141,7 @@ add_k_point_contribution_nonlocal(Simulation_context& ctx__, Beta_projectors_bas
 
                                 case 2: {
                                     /* Dij(00) = dij + dij_Z ;  Dij(11) = dij - dij_Z*/
-                                    dij = (uc.atom(ia).d_mtrx(ibf, jbf, 0) +
-                                           spin_factor * uc.atom(ia).d_mtrx(ibf, jbf, 1));
+                                    dij = (d_mtrx(ibf, jbf, 0) + spin_factor * d_mtrx(ibf, jbf, 1));
                                     if (lm1 == lm2) {
                                         dij += uc.atom(ia).type().d_mtrx_ion()(idxrf1, idxrf2);
                                     }
@@ -154,10 +158,10 @@ add_k_point_contribution_nonlocal(Simulation_context& ctx__, Beta_projectors_bas
                             for_bnd(ibf, jbf, dij, qij, beta_phi_chunks[ispn]);
 
                             /* for non-collinear case*/
-                            if (ctx__.num_mag_dims() == 3) {
+                            if (ctx.num_mag_dims() == 3) {
                                 /* Dij(10) = dij_X + i dij_Y ; Dij(01) = dij_X - i dij_Y */
-                                dij = std::complex<real_type<F>>(uc.atom(ia).d_mtrx(ibf, jbf, 2),
-                                                                 spin_factor * uc.atom(ia).d_mtrx(ibf, jbf, 3));
+                                dij = std::complex<real_type<F>>(d_mtrx(ibf, jbf, 2),
+                                                                 spin_factor * d_mtrx(ibf, jbf, 3));
                                 /* add non-diagonal spin components*/
                                 for_bnd(ibf, jbf, dij, 0.0, beta_phi_chunks[ispn + spin_factor]);
                             }
@@ -167,8 +171,6 @@ add_k_point_contribution_nonlocal(Simulation_context& ctx__, Beta_projectors_bas
             } // ispn
         } // x
     }
-
-    // bp_base__.dismiss();
 }
 
 } // namespace sirius
diff --git a/src/geometry/stress.cpp b/src/geometry/stress.cpp
index 140878764..a6fceb33c 100644
--- a/src/geometry/stress.cpp
+++ b/src/geometry/stress.cpp
@@ -45,7 +45,7 @@ Stress::calc_stress_nonloc_aux()
         auto mg  = kp->spinor_wave_functions().memory_guard(mem, wf::copy_to::device);
         Beta_projectors_strain_deriv<T> bp_strain_deriv(ctx_, kp->gkvec());
 
-        add_k_point_contribution_nonlocal<T, F>(ctx_, bp_strain_deriv, *kp, collect_result);
+        add_k_point_contribution_nonlocal<T, F>(potential_, bp_strain_deriv, *kp, collect_result);
     }
 
     #pragma omp parallel
diff --git a/src/hamiltonian/diagonalize_pp.hpp b/src/hamiltonian/diagonalize_pp.hpp
index 3186e8593..2da958b6b 100644
--- a/src/hamiltonian/diagonalize_pp.hpp
+++ b/src/hamiltonian/diagonalize_pp.hpp
@@ -146,8 +146,6 @@ diagonalize_pp_exact(int ispn__, Hamiltonian_k<T> const& Hk__, K_point<T>& kp__)
             }
         } // i (atoms in chunk)
     }
-    // kp.beta_projectors_row().dismiss();
-    // kp.beta_projectors_col().dismiss();
 
     if (ctx.cfg().control().verification() >= 1) {
         double max_diff = check_hermitian(ovlp, kp__.num_gkvec());
diff --git a/src/hamiltonian/hamiltonian.cpp b/src/hamiltonian/hamiltonian.cpp
index d7874784e..052830aad 100644
--- a/src/hamiltonian/hamiltonian.cpp
+++ b/src/hamiltonian/hamiltonian.cpp
@@ -32,7 +32,7 @@ Hamiltonian0<T>::Hamiltonian0(Potential& potential__, bool precompute_lapw__, bo
             new Local_operator<T>(ctx_, ctx_.spfft_coarse<T>(), ctx_.gvec_coarse_fft_sptr(), &potential__));
 
     if (!ctx_.full_potential()) {
-        d_op_ = std::unique_ptr<D_operator<T>>(new D_operator<T>(ctx_));
+        d_op_ = std::unique_ptr<D_operator<T>>(new D_operator<T>(potential__));
         q_op_ = std::unique_ptr<Q_operator<T>>(new Q_operator<T>(ctx_));
     }
     if (ctx_.full_potential()) {
diff --git a/src/hamiltonian/non_local_operator.cpp b/src/hamiltonian/non_local_operator.cpp
index 96ab0d47c..d790b3682 100644
--- a/src/hamiltonian/non_local_operator.cpp
+++ b/src/hamiltonian/non_local_operator.cpp
@@ -14,6 +14,7 @@
 #include "non_local_operator.hpp"
 #include "beta_projectors/beta_projectors_base.hpp"
 #include "hubbard/hubbard_matrix.hpp"
+#include "potential/potential.hpp"
 
 namespace sirius {
 
@@ -93,21 +94,21 @@ Non_local_operator<T>::get_matrix(int ispn, memory_t mem) const
 }
 
 template <typename T>
-D_operator<T>::D_operator(Simulation_context const& ctx_)
-    : Non_local_operator<T>(ctx_)
+D_operator<T>::D_operator(Potential& potential__)
+    : Non_local_operator<T>(potential__.ctx())
 {
-    if (ctx_.gamma_point()) {
-        this->op_ = mdarray<T, 3>({1, this->packed_mtrx_size_, ctx_.num_mag_dims() + 1});
+    if (potential__.ctx().gamma_point()) {
+        this->op_ = mdarray<T, 3>({1, this->packed_mtrx_size_, this->ctx_.num_mag_dims() + 1});
     } else {
-        this->op_ = mdarray<T, 3>({2, this->packed_mtrx_size_, ctx_.num_mag_dims() + 1});
+        this->op_ = mdarray<T, 3>({2, this->packed_mtrx_size_, this->ctx_.num_mag_dims() + 1});
     }
     this->op_.zero();
-    initialize();
+    initialize(potential__);
 }
 
 template <typename T>
 void
-D_operator<T>::initialize()
+D_operator<T>::initialize(Potential& potential__)
 {
     PROFILE("sirius::D_operator::initialize");
 
@@ -147,7 +148,7 @@ D_operator<T>::initialize()
                                             for (int alpha = 0; alpha < 4; alpha++) {
                                                 for (int sigma1 = 0; sigma1 < 2; sigma1++) {
                                                     for (int sigma2 = 0; sigma2 < 2; sigma2++) {
-                                                        result += atom.d_mtrx(xi1p, xi2p, alpha) *
+                                                        result += potential__.d_mtrx(ia)(xi1p, xi2p, alpha) *
                                                                   pauli_matrix[alpha][sigma1][sigma2] *
                                                                   atom.type().f_coefficients(xi1, xi1p, sigma, sigma1) *
                                                                   atom.type().f_coefficients(xi2p, xi2, sigma2, sigmap);
@@ -215,8 +216,8 @@ D_operator<T>::initialize()
                     int idx = xi2 * nbf + xi1;
                     switch (this->ctx_.num_mag_dims()) {
                         case 3: {
-                            T bx = uc.atom(ia).d_mtrx(xi1, xi2, 2);
-                            T by = uc.atom(ia).d_mtrx(xi1, xi2, 3);
+                            T bx = potential__.d_mtrx(ia)(xi1, xi2, 2);
+                            T by = potential__.d_mtrx(ia)(xi1, xi2, 3);
 
                             this->op_(0, this->packed_mtrx_offset_(ia) + idx, 2) = bx;
                             this->op_(1, this->packed_mtrx_offset_(ia) + idx, 2) = -by;
@@ -225,8 +226,8 @@ D_operator<T>::initialize()
                             this->op_(1, this->packed_mtrx_offset_(ia) + idx, 3) = by;
                         }
                         case 1: {
-                            T v  = uc.atom(ia).d_mtrx(xi1, xi2, 0);
-                            T bz = uc.atom(ia).d_mtrx(xi1, xi2, 1);
+                            T v  = potential__.d_mtrx(ia)(xi1, xi2, 0);
+                            T bz = potential__.d_mtrx(ia)(xi1, xi2, 1);
 
                             /* add ionic part */
                             if (lm1 == lm2) {
@@ -238,7 +239,7 @@ D_operator<T>::initialize()
                             break;
                         }
                         case 0: {
-                            this->op_(0, this->packed_mtrx_offset_(ia) + idx, 0) = uc.atom(ia).d_mtrx(xi1, xi2, 0);
+                            this->op_(0, this->packed_mtrx_offset_(ia) + idx, 0) = potential__.d_mtrx(ia)(xi1, xi2, 0);
                             /* add ionic part */
                             if (lm1 == lm2) {
                                 this->op_(0, this->packed_mtrx_offset_(ia) + idx, 0) += dion(idxrf1, idxrf2);
diff --git a/src/hamiltonian/non_local_operator.hpp b/src/hamiltonian/non_local_operator.hpp
index 46f235e5b..5c559c099 100644
--- a/src/hamiltonian/non_local_operator.hpp
+++ b/src/hamiltonian/non_local_operator.hpp
@@ -29,16 +29,17 @@ template <typename T>
 class Beta_projectors;
 template <typename T>
 class Beta_projectors_base;
+class Potential;
 
 template <typename T>
 class D_operator : public Non_local_operator<T>
 {
   private:
     void
-    initialize();
+    initialize(Potential& potential__);
 
   public:
-    D_operator(Simulation_context const& ctx_);
+    D_operator(Potential& potential__);
 };
 
 template <typename T>
@@ -57,7 +58,6 @@ class U_operator
 {
   private:
     Simulation_context const& ctx_;
-    // sddk::mdarray<std::complex<T>, 3> um_;
     std::array<la::dmatrix<std::complex<T>>, 4> um_;
     std::vector<int> offset_;
     std::vector<std::pair<int, int>> atomic_orbitals_;
diff --git a/src/hubbard/hubbard_matrix.hpp b/src/hubbard/hubbard_matrix.hpp
index 636f2a684..135efab9e 100644
--- a/src/hubbard/hubbard_matrix.hpp
+++ b/src/hubbard/hubbard_matrix.hpp
@@ -22,7 +22,7 @@ namespace sirius {
 class Hubbard_matrix
 {
   protected:
-    Simulation_context& ctx_;
+    Simulation_context const& ctx_;
     /// Local part of Hubbard matrix
     int num_steps_{0};
     double constraint_error_{1.0};
diff --git a/src/potential/generate_d_operator_matrix.cpp b/src/potential/generate_d_mtrx.cpp
similarity index 96%
rename from src/potential/generate_d_operator_matrix.cpp
rename to src/potential/generate_d_mtrx.cpp
index bfe8d8a52..2d53d39e1 100644
--- a/src/potential/generate_d_operator_matrix.cpp
+++ b/src/potential/generate_d_mtrx.cpp
@@ -23,9 +23,9 @@ mul_veff_with_phase_factors_gpu(int num_atoms__, int num_gvec_loc__, std::comple
 #endif
 
 void
-Potential::generate_D_operator_matrix()
+Potential::generate_d_mtrx()
 {
-    PROFILE("sirius::Potential::generate_D_operator_matrix");
+    PROFILE("sirius::Potential::generate_d_mtrx");
 
     /* local number of G-vectors */
     int gvec_count   = ctx_.gvec().count();
@@ -70,7 +70,7 @@ Potential::generate_D_operator_matrix()
 
                     for (int xi2 = 0; xi2 < nbf; xi2++) {
                         for (int xi1 = 0; xi1 < nbf; xi1++) {
-                            atom.d_mtrx(xi1, xi2, iv) = 0;
+                            d_mtrx_[ia](xi1, xi2, iv) = 0;
                         }
                     }
                 }
@@ -200,14 +200,13 @@ Potential::generate_D_operator_matrix()
 
             #pragma omp parallel for schedule(static)
             for (int i = 0; i < atom_type.num_atoms(); i++) {
-                int ia     = atom_type.atom_id(i);
-                auto& atom = unit_cell_.atom(ia);
+                int ia = atom_type.atom_id(i);
 
                 for (int xi2 = 0; xi2 < nbf; xi2++) {
                     for (int xi1 = 0; xi1 <= xi2; xi1++) {
                         int idx12 = xi2 * (xi2 + 1) / 2 + xi1;
                         /* D-matix is symmetric */
-                        atom.d_mtrx(xi1, xi2, iv) = atom.d_mtrx(xi2, xi1, iv) =
+                        d_mtrx_[ia](xi1, xi2, iv) = d_mtrx_[ia](xi2, xi1, iv) =
                                 d_tmp(idx12, i, iv) * unit_cell_.omega();
                     }
                 }
diff --git a/src/potential/paw_potential.cpp b/src/potential/paw_potential.cpp
index 247bc6713..1002d4ca1 100644
--- a/src/potential/paw_potential.cpp
+++ b/src/potential/paw_potential.cpp
@@ -35,11 +35,11 @@ Potential::init_PAW()
             [this](int ia) { return lmax_t(2 * this->unit_cell_.atom(ia).type().indexr().lmax()); });
 
     /* initialize dij matrix */
-    paw_dij_.resize(unit_cell_.num_paw_atoms());
+    d_mtrx_paw_.resize(unit_cell_.num_paw_atoms());
     for (int i = 0; i < unit_cell_.num_paw_atoms(); i++) {
-        int ia      = unit_cell_.paw_atom_index(paw_atom_index_t::global(i));
-        paw_dij_[i] = mdarray<double, 3>(
-                {unit_cell_.atom(ia).mt_basis_size(), unit_cell_.atom(ia).mt_basis_size(), ctx_.num_mag_dims() + 1});
+        int ia         = unit_cell_.paw_atom_index(paw_atom_index_t::global(i));
+        int nbf        = unit_cell_.atom(ia).mt_basis_size();
+        d_mtrx_paw_[i] = mdarray<double, 3>({nbf, nbf, ctx_.num_mag_dims() + 1});
     }
 }
 
@@ -90,11 +90,11 @@ Potential::generate_PAW_effective_potential(Density const& density)
     #pragma omp parallel for
     for (auto it : unit_cell_.spl_num_paw_atoms()) {
         auto ia = unit_cell_.paw_atom_index(it.i);
-        calc_PAW_local_Dij(ia, paw_dij_[it.i]);
+        calc_PAW_local_Dij(ia, d_mtrx_paw_[it.i]);
     }
     for (int i = 0; i < unit_cell_.num_paw_atoms(); i++) {
         auto location = unit_cell_.spl_num_paw_atoms().location(typename paw_atom_index_t::global(i));
-        comm_.bcast(paw_dij_[i].at(memory_t::host), paw_dij_[i].size(), location.ib);
+        comm_.bcast(d_mtrx_paw_[i].at(memory_t::host), d_mtrx_paw_[i].size(), location.ib);
     }
 
     /* add paw Dij to uspp Dij */
@@ -106,7 +106,7 @@ Potential::generate_PAW_effective_potential(Density const& density)
         for (int imagn = 0; imagn < ctx_.num_mag_dims() + 1; imagn++) {
             for (int ib2 = 0; ib2 < atom.mt_basis_size(); ib2++) {
                 for (int ib1 = 0; ib1 < atom.mt_basis_size(); ib1++) {
-                    atom.d_mtrx(ib1, ib2, imagn) += paw_dij_[i](ib1, ib2, imagn);
+                    d_mtrx_[ia](ib1, ib2, imagn) += d_mtrx_paw_[i](ib1, ib2, imagn);
                 }
             }
         }
@@ -318,4 +318,51 @@ Potential::calc_PAW_one_elec_energy(Atom const& atom__, mdarray<double, 2> const
     return energy;
 }
 
+double
+Potential::PAW_xc_total_energy(Density const& density__) const
+{
+    if (!unit_cell_.num_paw_atoms()) {
+        return 0;
+    }
+    /* compute contribution from the core */
+    double ecore{0};
+    #pragma omp parallel for reduction(+:ecore)
+    for (auto it : unit_cell_.spl_num_paw_atoms()) {
+        auto ia = unit_cell_.paw_atom_index(it.i);
+
+        auto& atom = unit_cell_.atom(ia);
+
+        auto& atom_type = atom.type();
+
+        auto& ps_core = atom_type.ps_core_charge_density();
+        auto& ae_core = atom_type.paw_ae_core_charge_density();
+
+        Spline<double> s(atom_type.radial_grid());
+        auto y00inv = 1.0 / y00;
+        for (int ir = 0; ir < atom_type.num_mt_points(); ir++) {
+            s(ir) = y00inv * ((*paw_ae_exc_)[ia](0, ir) * ae_core[ir] - (*paw_ps_exc_)[ia](0, ir) * ps_core[ir]) *
+                    std::pow(atom_type.radial_grid(ir), 2);
+        }
+        ecore += s.interpolate().integrate(0);
+    }
+    comm_.allreduce(&ecore, 1);
+
+    return inner(*paw_ae_exc_, density__.paw_density().ae_component(0)) -
+           inner(*paw_ps_exc_, density__.paw_density().ps_component(0)) + ecore;
+}
+
+double
+Potential::PAW_one_elec_energy(Density const& density__) const
+{
+    double e{0};
+    #pragma omp parallel for reduction(+:e)
+    for (auto it : unit_cell_.spl_num_paw_atoms()) {
+        auto ia = unit_cell_.paw_atom_index(it.i);
+        auto dm = density__.density_matrix_aux(atom_index_t::global(ia));
+        e += calc_PAW_one_elec_energy(unit_cell_.atom(ia), dm, d_mtrx_paw_[it.i]);
+    }
+    comm_.allreduce(&e, 1);
+    return e;
+}
+
 } // namespace sirius
diff --git a/src/potential/potential.cpp b/src/potential/potential.cpp
index 328af402b..82e926287 100644
--- a/src/potential/potential.cpp
+++ b/src/potential/potential.cpp
@@ -132,6 +132,15 @@ Potential::Potential(Simulation_context& ctx__)
         }
     }
 
+    if (!ctx_.full_potential()) {
+        d_mtrx_.resize(unit_cell_.num_atoms());
+        for (int ia = 0; ia < unit_cell_.num_atoms(); ia++) {
+            int nbf     = unit_cell_.atom(ia).mt_basis_size();
+            d_mtrx_[ia] = mdarray<double, 3>({nbf, nbf, ctx_.num_mag_dims() + 1}, mdarray_label("d_mtrx_"));
+            d_mtrx_[ia].zero();
+        }
+    }
+
     /* in case of PAW */
     init_PAW();
 
@@ -321,7 +330,7 @@ Potential::generate(Density const& density__, bool use_symmetry__, bool transfor
     }
 
     if (!ctx_.full_potential()) {
-        generate_D_operator_matrix();
+        generate_d_mtrx();
         generate_PAW_effective_potential(density__);
         if (ctx_.verbosity() >= 3) {
             rte::ostream out(ctx_.out(), "potential");
@@ -350,7 +359,7 @@ Potential::generate(Density const& density__, bool use_symmetry__, bool transfor
                     for (int ib2 = 0; ib2 < atom.mt_basis_size(); ib2++) {
                         out << "    ";
                         for (int ib1 = 0; ib1 < atom.mt_basis_size(); ib1++) {
-                            out << ffmt(8, 3) << atom.d_mtrx(ib1, ib2, imagn);
+                            out << ffmt(8, 3) << d_mtrx_[ia](ib1, ib2, imagn);
                         }
                         out << std::endl;
                     }
diff --git a/src/potential/potential.hpp b/src/potential/potential.hpp
index a3b83bd5d..79ee01a7d 100644
--- a/src/potential/potential.hpp
+++ b/src/potential/potential.hpp
@@ -103,6 +103,18 @@ class Potential : public Field4D
     /// Plane-wave coefficients of the squared inverse relativistic mass weighted by the unit step-function.
     mdarray<std::complex<double>, 1> rm2_inv_pw_;
 
+    /// Auxiliary form of the D_{ij} operator matrix of the pseudo-potential method.
+    /** The matrix is calculated for the scalar and vector effective fields (thus, it is real and symmetric).
+     *  \f[
+     *      D_{\xi \xi'}^{\alpha} = \int V({\bf r}) Q_{\xi \xi'}^{\alpha}({\bf r}) d{\bf r}
+     *  \f]
+     *
+     *  D-operator represents non-local part of pseudopotential.
+     *
+     *  The ionic part of the D-operator matrix is added in the D_operator class, when it is initialized.
+     */
+    std::vector<mdarray<double, 3>> d_mtrx_;
+
     /// Hartree contribution to total energy from PAW atoms.
     double paw_hartree_total_energy_{0.0};
 
@@ -116,7 +128,7 @@ class Potential : public Field4D
     std::unique_ptr<Spheric_function_set<double, paw_atom_index_t>> paw_ae_exc_;
 
     /// Contribution to D-operator matrix from the PAW atoms.
-    std::vector<mdarray<double, 3>> paw_dij_;
+    std::vector<mdarray<double, 3>> d_mtrx_paw_;
 
     mdarray<double, 2> aux_bf_;
 
@@ -147,14 +159,14 @@ class Potential : public Field4D
                              std::vector<Flm const*> ps_density);
 
     void
-    calc_PAW_local_Dij(typename atom_index_t::global ia__, mdarray<double, 3>& paw_dij__);
+    calc_PAW_local_Dij(typename atom_index_t::global ia__, mdarray<double, 3>& d_mtrx_paw__);
 
     double
     calc_PAW_hartree_potential(Atom& atom, Flm const& full_density, Flm& full_potential);
 
     double
     calc_PAW_one_elec_energy(Atom const& atom__, mdarray<double, 2> const& density_matrix__,
-                             mdarray<double, 3> const& paw_dij__) const;
+                             mdarray<double, 3> const& d_mtrx_paw__) const;
 
     /// Compute MT part of the potential and MT multipole moments
     inline auto
@@ -641,67 +653,26 @@ class Potential : public Field4D
      *  is calculated for all atoms of the same type.
      */
     void
-    generate_D_operator_matrix();
+    generate_d_mtrx();
+
+    void
+    check_potential_continuity_at_mt();
 
     void
     generate_PAW_effective_potential(Density const& density);
 
     double
-    PAW_xc_total_energy(Density const& density__) const
-    {
-        if (!unit_cell_.num_paw_atoms()) {
-            return 0;
-        }
-        /* compute contribution from the core */
-        double ecore{0};
-        #pragma omp parallel for reduction(+:ecore)
-        for (auto it : unit_cell_.spl_num_paw_atoms()) {
-            auto ia = unit_cell_.paw_atom_index(it.i);
-
-            auto& atom = unit_cell_.atom(ia);
-
-            auto& atom_type = atom.type();
-
-            auto& ps_core = atom_type.ps_core_charge_density();
-            auto& ae_core = atom_type.paw_ae_core_charge_density();
-
-            Spline<double> s(atom_type.radial_grid());
-            auto y00inv = 1.0 / y00;
-            for (int ir = 0; ir < atom_type.num_mt_points(); ir++) {
-                s(ir) = y00inv * ((*paw_ae_exc_)[ia](0, ir) * ae_core[ir] - (*paw_ps_exc_)[ia](0, ir) * ps_core[ir]) *
-                        std::pow(atom_type.radial_grid(ir), 2);
-            }
-            ecore += s.interpolate().integrate(0);
-        }
-        comm_.allreduce(&ecore, 1);
-
-        return inner(*paw_ae_exc_, density__.paw_density().ae_component(0)) -
-               inner(*paw_ps_exc_, density__.paw_density().ps_component(0)) + ecore;
-    }
+    PAW_xc_total_energy(Density const& density__) const;
 
     double
+    PAW_one_elec_energy(Density const& density__) const;
+
+    inline double
     PAW_total_energy(Density const& density__) const
     {
         return paw_hartree_total_energy_ + PAW_xc_total_energy(density__);
     }
 
-    double
-    PAW_one_elec_energy(Density const& density__) const
-    {
-        double e{0};
-        #pragma omp parallel for reduction(+:e)
-        for (auto it : unit_cell_.spl_num_paw_atoms()) {
-            auto ia = unit_cell_.paw_atom_index(it.i);
-            auto dm = density__.density_matrix_aux(atom_index_t::global(ia));
-            e += calc_PAW_one_elec_energy(unit_cell_.atom(ia), dm, paw_dij_[it.i]);
-        }
-        comm_.allreduce(&e, 1);
-        return e;
-    }
-
-    void
-    check_potential_continuity_at_mt();
-
     auto&
     effective_potential()
     {
@@ -732,6 +703,12 @@ class Potential : public Field4D
         return *dveff_;
     }
 
+    auto const&
+    dveff() const
+    {
+        return *dveff_;
+    }
+
     auto const&
     effective_potential_mt(atom_index_t::local ialoc) const
     {
@@ -906,6 +883,12 @@ class Potential : public Field4D
     {
         return ewald_energy_;
     }
+
+    inline auto const&
+    d_mtrx(int ia) const
+    {
+        return d_mtrx_[ia];
+    }
 };
 
 inline void
diff --git a/src/unit_cell/atom.hpp b/src/unit_cell/atom.hpp
index 91adce425..3c0451acb 100644
--- a/src/unit_cell/atom.hpp
+++ b/src/unit_cell/atom.hpp
@@ -64,16 +64,6 @@ class Atom
     /// Orbital quantum number for UJ correction.
     int uj_correction_l_{-1};
 
-    /// Auxiliary form of the D_{ij} operator matrix of the pseudo-potential method.
-    /** The matrix is calculated for the scalar and vector effective fields (thus, it is real and symmetric).
-     *  \f[
-     *      D_{\xi \xi'}^{\alpha} = \int V({\bf r}) Q_{\xi \xi'}^{\alpha}({\bf r}) d{\bf r}
-     *  \f]
-     *
-     *  The ionic part of the D-operator matrix is added in the D_operator class, when it is initialized.
-     */
-    mdarray<double, 3> d_mtrx_;
-
   public:
     /// Constructor.
     Atom(Atom_type const& type__, r3::vector<double> position__, r3::vector<double> vector_field__)
@@ -105,13 +95,6 @@ class Atom
 
             uj_correction_matrix_ = mdarray<std::complex<double>, 4>({16, 16, 2, 2});
         }
-
-        if (!type().parameters().full_potential()) {
-            int nbf = type().mt_basis_size();
-            d_mtrx_ = mdarray<double, 3>({nbf, nbf, type().parameters().num_mag_dims() + 1},
-                                         mdarray_label("Atom::d_mtrx_"));
-            d_mtrx_.zero();
-        }
     }
 
     /// Generate radial Hamiltonian and effective magnetic field integrals
@@ -611,30 +594,6 @@ class Atom
     {
         return uj_correction_matrix_(lm1, lm2, ispn1, ispn2);
     }
-
-    inline double&
-    d_mtrx(int xi1, int xi2, int iv)
-    {
-        return d_mtrx_(xi1, xi2, iv);
-    }
-
-    inline double const&
-    d_mtrx(int xi1, int xi2, int iv) const
-    {
-        return d_mtrx_(xi1, xi2, iv);
-    }
-
-    inline auto const&
-    d_mtrx() const
-    {
-        return d_mtrx_;
-    }
-
-    inline auto&
-    d_mtrx()
-    {
-        return d_mtrx_;
-    }
 };
 
 } // namespace sirius
diff --git a/src/unit_cell/atom_symmetry_class.cpp b/src/unit_cell/atom_symmetry_class.cpp
index a5cba7c74..77df675d3 100644
--- a/src/unit_cell/atom_symmetry_class.cpp
+++ b/src/unit_cell/atom_symmetry_class.cpp
@@ -856,33 +856,6 @@ Atom_symmetry_class::generate_core_charge_density(relativity_t core_rel__)
         }
     }
 
-    //#pragma omp parallel default(shared)
-    //{
-    //    std::vector<double> rho_t(rho.num_points());
-    //    std::memset(&rho_t[0], 0, rho.num_points() * sizeof(double));
-
-    //    #pragma omp for
-    //    for (int ist = 0; ist < atom_type_.num_atomic_levels(); ist++) {
-    //        if (atom_type_.atomic_level(ist).core) {
-    //            Bound_state bs(core_rel__, atom_type_.zn(), atom_type_.atomic_level(ist).n,
-    //            atom_type_.atomic_level(ist).l,
-    //                           atom_type_.atomic_level(ist).k, rgrid, veff, level_energy[ist]);
-
-    //            auto& rho = bs.rho();
-    //            for (int i = 0; i < rgrid.num_points(); i++) {
-    //                rho_t[i] += atom_type_.atomic_level(ist).occupancy * rho(i) / fourpi;
-    //            }
-
-    //            level_energy[ist] = bs.enu();
-    //        }
-    //    }
-
-    //    #pragma omp critical
-    //    for (int i = 0; i < rho.num_points(); i++) {
-    //        rho(i) += rho_t[i];
-    //    }
-    //}
-
     for (int ir = 0; ir < atom_type_.num_mt_points(); ir++) {
         ae_core_charge_density_[ir] = rho(ir);
     }
diff --git a/src/unit_cell/atom_type.hpp b/src/unit_cell/atom_type.hpp
index 121eedcc3..50ab7d837 100644
--- a/src/unit_cell/atom_type.hpp
+++ b/src/unit_cell/atom_type.hpp
@@ -224,7 +224,10 @@ class Atom_type
      *  oribtal occupancies. */
     std::vector<double> paw_wf_occ_;
 
-    /// Core electron contribution to all electron charge density in PAW method.
+    /// Core electron contribution to all-electron charge density in PAW method.
+    /** In PAW, core charge density is defined for each atom type, because it comes from the
+     *  pseudopotential file and is fixed during the calculation. This is different from LAPW
+     *  core charge density, which is recomputed for each atom symmetry class. */
     std::vector<double> paw_ae_core_charge_density_;
 
     /// True if the pseudo potential includes spin orbit coupling.