orbitals.cc

Go to the documentation of this file.

/*
 *            Copyright 2009-2023 The VOTCA Development Team
 *                       (http://www.votca.org)
 *
 *      Licensed under the Apache License, Version 2.0 (the "License")
 *
 * You may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *              http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */
 
// Standard includes
#include <cstdio>
#include <fstream>
#include <iomanip>
#include <numeric>
 
// Local VOTCA includes
#include "votca/tools/version.h"
#include "votca/xtp/aomatrix.h"
#include "votca/xtp/orbitals.h"
#include "votca/xtp/orbreorder.h"
#include "votca/xtp/qmstate.h"
#include "votca/xtp/vc2index.h"
 
namespace votca {
namespace xtp {

 
Orbitals::Orbitals() : atoms_("", 0) { ; }

std::vector<Index> Orbitals::CheckDegeneracy(Index level,
                                             double energy_difference) const {
  if (this->hasUnrestrictedOrbitals()) {
    throw std::runtime_error(
        "Checking degeneracy not implemented for open-shell systems.");
  }
  std::vector<Index> result;
  if (level > mos_.eigenvalues().size()) {
    throw std::runtime_error(
        "Level for degeneracy is higher than maximum level");
  }
  double MOEnergyLevel = mos_.eigenvalues()(level);
 
  for (Index i = 0; i < mos_.eigenvalues().size(); ++i) {
    if (std::abs(mos_.eigenvalues()(i) - MOEnergyLevel) < energy_difference) {
      result.push_back(i);
    }
  }
 
  if (result.empty()) {
    result.push_back(level);
  }
  return result;
}
 
std::vector<Index> Orbitals::SortEnergies() {
  if (this->hasUnrestrictedOrbitals()) {
    throw std::runtime_error(
        "MO sorting not implemented for open-shell systems.");
  }
  std::vector<Index> index = std::vector<Index>(mos_.eigenvalues().size());
  std::iota(index.begin(), index.end(), 0);
  std::stable_sort(index.begin(), index.end(), [this](Index i1, Index i2) {
    return this->MOs().eigenvalues()[i1] < this->MOs().eigenvalues()[i2];
  });
  return index;
}

void Orbitals::SetupDftBasis(std::string basis_name) {
  if (this->QMAtoms().size() == 0) {
    throw std::runtime_error("Can't setup AOBasis without atoms");
  }
  BasisSet bs;
  bs.Load(basis_name);
  dftbasis_.Fill(bs, this->QMAtoms());
}
 
void Orbitals::SetupAuxBasis(std::string aux_basis_name) {
  if (this->QMAtoms().size() == 0) {
    throw std::runtime_error("Can't setup Aux AOBasis without atoms");
  }
  BasisSet bs;
  bs.Load(aux_basis_name);
  auxbasis_.Fill(bs, this->QMAtoms());
}
 
// Return the density change relative to the ground state. For excited states
// this corresponds to the response density Delta P, whereas for charged states
// it returns the projector that must be added to or removed from the ground-
// state density.
Eigen::MatrixXd Orbitals::DensityMatrixWithoutGS(const QMState& state) const {
  if (this->hasUnrestrictedOrbitals() &&
      state.Type() != QMStateType::ExcitonUKS) {
    throw std::runtime_error(
        "DensityMatrixWithoutGS not implemented for open-shell systems "
        "except ExcitonUKS.");
  }
  if (state.Type().isExciton()) {
    std::array<Eigen::MatrixXd, 2> DMAT = DensityMatrixExcitedState(state);
    return DMAT[1] - DMAT[0];
  } else if (state.Type().isKSState() || state.Type().isPQPState()) {
    return DensityMatrixKSstate(state);
  } else if (state.Type() == QMStateType::DQPstate) {
    Eigen::MatrixXd DMATQP = DensityMatrixQuasiParticle(state);
    if (state.StateIdx() > getHomo()) {
      return DMATQP;
    } else {
      return -DMATQP;
    }
  } else {
    throw std::runtime_error(
        "DensityMatrixWithoutGS does not yet implement QMStateType:" +
        state.Type().ToLongString());
  }
}
 
// Return the full AO density associated with the requested state. Depending
// on the state type this is the ground-state density with the corresponding
// particle, hole, or excitonic correction added on top.
Eigen::MatrixXd Orbitals::DensityMatrixFull(const QMState& state) const {
  if (state.isTransition()) {
    return this->TransitionDensityMatrix(state);
  }
  Eigen::MatrixXd result = this->DensityMatrixGroundState();
 
  if (getCalculationType() != "NoEmbedding") {
    result += getInactiveDensity();
  }
  if (state.Type().isExciton()) {
    std::array<Eigen::MatrixXd, 2> DMAT = DensityMatrixExcitedState(state);
    result = result - DMAT[0] + DMAT[1];  // Ground state + hole_contribution +
                                          // electron contribution
  } else if (state.Type() == QMStateType::DQPstate) {
    Eigen::MatrixXd DMATQP = DensityMatrixQuasiParticle(state);
    if (state.StateIdx() > getHomo()) {
      result += DMATQP;
    } else {
      result -= DMATQP;
    }
  } else if (state.Type() == QMStateType::KSstate ||
             state.Type() == QMStateType::PQPstate) {
    Eigen::MatrixXd DMATKS = DensityMatrixKSstate(state);
    if (state.StateIdx() <= getHomo()) {
      result -= DMATKS;
    } else {
      result += DMATKS;
    }
  } else if (state.Type() == QMStateType::Hole ||
             state.Type() == QMStateType::Electron) {
    if (!this->hasUnrestrictedOrbitals()) {
      throw std::runtime_error("QMStateType:" + state.Type().ToLongString() +
                               " requires an openshell calculation");
    }
  } else if (state.Type() != QMStateType::Gstate) {
    throw std::runtime_error(
        "DensityMatrixFull does not yet implement QMStateType:" +
        state.Type().ToLongString());
  }
  return result;
}
 
// Determine ground state density matrix
/*Eigen::MatrixXd Orbitals::DensityMatrixGroundState() const {
  if (!hasMOs()) {
    throw std::runtime_error("Orbitals file does not contain MO coefficients");
  }
  Eigen::MatrixXd occstates = mos_.eigenvectors().leftCols(occupied_levels_);
  Eigen::MatrixXd dmatGS = occstates * occstates.transpose();
  if (this->hasUnrestrictedOrbitals()) {
    Eigen::MatrixXd occstates_beta =
        mos_beta_.eigenvectors().leftCols(occupied_levels_beta_);
    dmatGS += occstates_beta * occstates_beta.transpose();
  } else {
    dmatGS *= 2.0;
  }
  return dmatGS;
}*/
 
// Ground-state one-particle density in the AO basis.
//
// For unrestricted calculations the total density is assembled as
//
//   P = P^alpha + P^beta,
//
// with P^sigma = C_occ^sigma (C_occ^sigma)^T.
//
// For restricted cases the spin-resolved helper reconstructs the same object
// from the shared spatial orbitals together with the stored alpha/beta
// occupation pattern.
Eigen::MatrixXd Orbitals::DensityMatrixGroundState() const {
  auto spin = DensityMatrixGroundStateSpinResolved();
  return spin[0] + spin[1];
}
 
// Density matrix for a single KS orbital
// One-orbital projector in the AO basis,
//
//   P^(k)_(mu,nu) = C_(mu,k) C_(nu,k),
//
// used for adding or removing a single KS occupation relative to the ground
// state.
Eigen::MatrixXd Orbitals::DensityMatrixKSstate(const QMState& state) const {
  if (this->hasUnrestrictedOrbitals()) {
    throw std::runtime_error(
        "DensityMatrixKSstate not implemented for open-shell systems.");
  }
  if (!hasMOs()) {
    throw std::runtime_error("Orbitals file does not contain MO coefficients");
  }
  if (state.Type() != QMStateType::KSstate &&
      state.Type() != QMStateType::PQPstate) {
    throw std::runtime_error("State:" + state.ToString() +
                             " is not a Kohn Sham state");
  }
  Eigen::VectorXd KSstate = mos_.eigenvectors().col(state.StateIdx());
  Eigen::MatrixXd dmatKS = KSstate * KSstate.transpose();
  return dmatKS;
}
 
Eigen::MatrixXd Orbitals::CalculateQParticleAORepresentation() const {
  if (!hasQPdiag()) {
    throw std::runtime_error("Orbitals file does not contain QP coefficients");
  }
  return mos_.eigenvectors().middleCols(qpmin_, qpmax_ - qpmin_ + 1) *
         QPdiag_.eigenvectors();
}
 
// Determine QuasiParticle Density Matrix
// Quasiparticle projector in the AO basis,
//
//   P^QP_(mu,nu) = Lambda_(mu,n) Lambda_(nu,n),
//
// where Lambda is the AO representation of the quasiparticle eigenvectors in
// the selected GW subspace.
Eigen::MatrixXd Orbitals::DensityMatrixQuasiParticle(
    const QMState& state) const {
  if (this->hasUnrestrictedOrbitals()) {
    throw std::runtime_error(
        "DensityMatrixQuasiParticle not implemented for open-shell systems.");
  }
  if (state.Type() != QMStateType::DQPstate) {
    throw std::runtime_error("State:" + state.ToString() +
                             " is not a quasiparticle state");
  }
  Eigen::MatrixXd lambda = CalculateQParticleAORepresentation();
  Eigen::MatrixXd dmatQP = lambda.col(state.StateIdx() - qpmin_) *
                           lambda.col(state.StateIdx() - qpmin_).transpose();
  return dmatQP;
}
 
// Electric dipole moment relative to the molecular origin,
//
//   mu = sum_A Z_A (R_A - R_0) - Tr[P M],
//
// where M contains the AO dipole integrals for x, y, and z and P is the full
// state density matrix. Transition densities omit the nuclear term and return
// only the electronic transition dipole.
Eigen::Vector3d Orbitals::CalcElDipole(const QMState& state) const {
  Eigen::Vector3d nuclei_dip = Eigen::Vector3d::Zero();
  if (!state.isTransition()) {
    for (const QMAtom& atom : atoms_) {
      nuclei_dip += (atom.getPos() - atoms_.getPos()) * atom.getNuccharge();
    }
  }
  AOBasis basis = getDftBasis();
  AODipole dipole;
  dipole.setCenter(atoms_.getPos());
  dipole.Fill(basis);
 
  Eigen::MatrixXd dmat = this->DensityMatrixFull(state);
  Eigen::Vector3d electronic_dip;
  for (Index i = 0; i < 3; ++i) {
    electronic_dip(i) = dmat.cwiseProduct(dipole.Matrix()[i]).sum();
  }
  return nuclei_dip - electronic_dip;
}
 
// AO transition density for a singlet BSE excitation.
//
// In TDA form this evaluates
//
//   gamma_(mu,nu) = sqrt(2) sum_(v,c) A_(v,c) C_(mu,v) C_(nu,c),
//
// and for full BSE the resonant and anti-resonant amplitudes are added before
// the AO transformation.
Eigen::MatrixXd Orbitals::TransitionDensityMatrix(const QMState& state) const {
 
  if (state.Type() == QMStateType::ExcitonUKS) {
    if (!this->hasUnrestrictedOrbitals()) {
      throw std::runtime_error(
          "ExcitonUKS transition density requested, but no unrestricted "
          "orbitals are present.");
    }
 
    const Eigen::MatrixXd& BSECoefs = BSE_uks_.eigenvectors();
    if (BSECoefs.cols() < state.StateIdx() + 1 || BSECoefs.rows() < 2) {
      throw std::runtime_error(
          "Orbitals object has no information about state:" + state.ToString());
    }
 
    Eigen::VectorXd coeffs = BSECoefs.col(state.StateIdx());
    if (!useTDA_) {
      coeffs += BSE_uks_.eigenvectors2().col(state.StateIdx());
    }
 
    const Index alpha_vtotal = getHomoAlpha() - bse_vmin_ + 1;
    const Index alpha_ctotal = bse_cmax_ - getLumoAlpha() + 1;
    const Index beta_vtotal = getHomoBeta() - bse_vmin_ + 1;
    const Index beta_ctotal = bse_cmax_ - getLumoBeta() + 1;
 
    const Index alpha_size = alpha_vtotal * alpha_ctotal;
    const Index beta_size = beta_vtotal * beta_ctotal;
 
    if (coeffs.size() != alpha_size + beta_size) {
      throw std::runtime_error(
          "ExcitonUKS coefficient vector size does not match expected "
          "alpha+beta BSE space dimensions.");
    }
 
    Eigen::MatrixXd tdm = Eigen::MatrixXd::Zero(mos_.eigenvectors().rows(),
                                                mos_.eigenvectors().rows());
 
    if (alpha_size > 0) {
      const Eigen::VectorXd coeffs_alpha = coeffs.head(alpha_size);
      const auto occ_alpha =
          mos_.eigenvectors().middleCols(bse_vmin_, alpha_vtotal);
      const auto virt_alpha =
          mos_.eigenvectors().middleCols(getLumoAlpha(), alpha_ctotal);
      const Eigen::Map<const Eigen::MatrixXd> mat_alpha(
          coeffs_alpha.data(), alpha_ctotal, alpha_vtotal);
 
      tdm += occ_alpha * mat_alpha.transpose() * virt_alpha.transpose();
    }
 
    if (beta_size > 0) {
      const Eigen::VectorXd coeffs_beta = coeffs.tail(beta_size);
      const auto occ_beta =
          mos_beta_.eigenvectors().middleCols(bse_vmin_, beta_vtotal);
      const auto virt_beta =
          mos_beta_.eigenvectors().middleCols(getLumoBeta(), beta_ctotal);
      const Eigen::Map<const Eigen::MatrixXd> mat_beta(
          coeffs_beta.data(), beta_ctotal, beta_vtotal);
 
      tdm += occ_beta * mat_beta.transpose() * virt_beta.transpose();
    }
 
    return tdm;
  }
 
  if (this->hasUnrestrictedOrbitals()) {
    throw std::runtime_error(
        "TransitionDensityMatrix not implemented for open-shell systems "
        "except ExcitonUKS.");
  }
  if (state.Type() != QMStateType::Singlet) {
    throw std::runtime_error(
        "Spin type not known for transition density matrix. Available only for "
        "singlet");
  }
  const Eigen::MatrixXd& BSECoefs = BSE_singlet_.eigenvectors();
  if (BSECoefs.cols() < state.StateIdx() + 1 || BSECoefs.rows() < 2) {
    throw std::runtime_error("Orbitals object has no information about state:" +
                             state.ToString());
  }
 
  Eigen::VectorXd coeffs = BSECoefs.col(state.StateIdx());
 
  if (!useTDA_) {
    coeffs += BSE_singlet_.eigenvectors2().col(state.StateIdx());
  }
  coeffs *= std::sqrt(2.0);
  auto occlevels = mos_.eigenvectors().middleCols(bse_vmin_, bse_vtotal_);
  auto virtlevels = mos_.eigenvectors().middleCols(bse_cmin_, bse_ctotal_);
  Eigen::Map<const Eigen::MatrixXd> mat(coeffs.data(), bse_ctotal_,
                                        bse_vtotal_);
 
  return occlevels * mat.transpose() * virtlevels.transpose();
}
 
// Spin-summed excited-state density correction split into hole and electron
// blocks. For full BSE the anti-resonant contribution is subtracted from the
// resonant part, yielding the usual RPA/BSE excited-state density.
std::array<Eigen::MatrixXd, 2> Orbitals::DensityMatrixExcitedState(
    const QMState& state) const {
 
  if (state.Type() == QMStateType::ExcitonUKS) {
    if (!this->hasUnrestrictedOrbitals()) {
      throw std::runtime_error(
          "ExcitonUKS excited-state density requested, but no unrestricted "
          "orbitals are present.");
    }
 
    const Eigen::MatrixXd& X = BSE_uks_.eigenvectors();
    if (X.cols() < state.StateIdx() + 1 || X.rows() < 2) {
      throw std::runtime_error(
          "Orbitals object has no information about state:" + state.ToString());
    }
 
    const Index alpha_vtotal = getHomoAlpha() - bse_vmin_ + 1;
    const Index alpha_ctotal = bse_cmax_ - getLumoAlpha() + 1;
    const Index beta_vtotal = getHomoBeta() - bse_vmin_ + 1;
    const Index beta_ctotal = bse_cmax_ - getLumoBeta() + 1;
 
    const Index alpha_size = alpha_vtotal * alpha_ctotal;
    const Index beta_size = beta_vtotal * beta_ctotal;
 
    if (X.rows() != alpha_size + beta_size) {
      throw std::runtime_error(
          "ExcitonUKS coefficient vector size does not match expected "
          "alpha+beta BSE space dimensions.");
    }
 
    const Index nao = mos_.eigenvectors().rows();
    std::array<Eigen::MatrixXd, 2> dmat;
    dmat[0] = Eigen::MatrixXd::Zero(nao, nao);  // hole
    dmat[1] = Eigen::MatrixXd::Zero(nao, nao);  // electron
 
    auto add_resonant = [&](const Eigen::VectorXd& coeffs, bool alpha_spin) {
      if (coeffs.size() == 0) {
        return;
      }
 
      const Index vtotal = alpha_spin ? alpha_vtotal : beta_vtotal;
      const Index ctotal = alpha_spin ? alpha_ctotal : beta_ctotal;
 
      const auto& mos_spin =
          alpha_spin ? mos_.eigenvectors() : mos_beta_.eigenvectors();
      const Index lumo = alpha_spin ? getLumoAlpha() : getLumoBeta();
 
      const auto occlevels = mos_spin.middleCols(bse_vmin_, vtotal);
      const auto virtlevels = mos_spin.middleCols(lumo, ctotal);
 
      const Eigen::Map<const Eigen::MatrixXd> mat(coeffs.data(), ctotal,
                                                  vtotal);
 
      dmat[0] += occlevels * (mat.transpose() * mat) * occlevels.transpose();
      dmat[1] += virtlevels * (mat * mat.transpose()) * virtlevels.transpose();
    };
 
    auto subtract_antiresonant = [&](const Eigen::VectorXd& coeffs,
                                     bool alpha_spin) {
      if (coeffs.size() == 0) {
        return;
      }
 
      const Index vtotal = alpha_spin ? alpha_vtotal : beta_vtotal;
      const Index ctotal = alpha_spin ? alpha_ctotal : beta_ctotal;
 
      const auto& mos_spin =
          alpha_spin ? mos_.eigenvectors() : mos_beta_.eigenvectors();
      const Index lumo = alpha_spin ? getLumoAlpha() : getLumoBeta();
 
      const auto occlevels = mos_spin.middleCols(bse_vmin_, vtotal);
      const auto virtlevels = mos_spin.middleCols(lumo, ctotal);
 
      const Eigen::Map<const Eigen::MatrixXd> mat(coeffs.data(), ctotal,
                                                  vtotal);
 
      // Full-BSE anti-resonant correction:
      // hole  -= virt * (B B^T) * virt^T
      // electron -= occ * (B^T B) * occ^T
      dmat[0] -= virtlevels * (mat * mat.transpose()) * virtlevels.transpose();
      dmat[1] -= occlevels * (mat.transpose() * mat) * occlevels.transpose();
    };
 
    const Eigen::VectorXd Xstate = X.col(state.StateIdx());
    add_resonant(Xstate.head(alpha_size), true);
    add_resonant(Xstate.tail(beta_size), false);
 
    if (!useTDA_) {
      const Eigen::MatrixXd& Y = BSE_uks_.eigenvectors2();
      if (Y.cols() < state.StateIdx() + 1 || Y.rows() != X.rows()) {
        throw std::runtime_error(
            "Orbitals object has inconsistent anti-resonant ExcitonUKS "
            "coefficients for state:" +
            state.ToString());
      }
 
      const Eigen::VectorXd Ystate = Y.col(state.StateIdx());
      subtract_antiresonant(Ystate.head(alpha_size), true);
      subtract_antiresonant(Ystate.tail(beta_size), false);
    }
 
    return dmat;
  }
 
  if (this->hasUnrestrictedOrbitals()) {
    throw std::runtime_error(
        "DensityMatrixExcitedState not implemented for open-shell systems "
        "except ExcitonUKS.");
  }
 
  std::array<Eigen::MatrixXd, 2> dmat = DensityMatrixExcitedState_R(state);
  if (!useTDA_) {
    std::array<Eigen::MatrixXd, 2> dmat_AR =
        DensityMatrixExcitedState_AR(state);
    dmat[0] -= dmat_AR[0];
    dmat[1] -= dmat_AR[1];
  }
  return dmat;
}
 
// Excited state density matrix
 
std::array<Eigen::MatrixXd, 2> Orbitals::DensityMatrixExcitedState_R(
    const QMState& state) const {
  if (!state.Type().isExciton()) {
    throw std::runtime_error(
        "Spin type not known for density matrix. Available are singlet and "
        "triplet");
  }
 
  const Eigen::MatrixXd& BSECoefs = (state.Type() == QMStateType::Singlet)
                                        ? BSE_singlet_.eigenvectors()
                                        : BSE_triplet_.eigenvectors();
  if (BSECoefs.cols() < state.StateIdx() + 1 || BSECoefs.rows() < 2) {
    throw std::runtime_error("Orbitals object has no information about state:" +
                             state.ToString());
  }
  /******
   *
   *    Density matrix for GW-BSE based excitations
   *
   *    - electron contribution
   *      D_ab = \sum{vc} \sum{c'} A_{vc}A_{vc'} mo_a(c)mo_b(c')
   *
   *    - hole contribution
   *      D_ab = \sum{vc} \sum{v'} A_{vc}A_{v'c} mo_a(v)mo_b(v')
   *
   */
 
  Eigen::VectorXd coeffs = BSECoefs.col(state.StateIdx());
 
  std::array<Eigen::MatrixXd, 2> dmatEX;
  // hole part as matrix products
  Eigen::MatrixXd occlevels =
      mos_.eigenvectors().middleCols(bse_vmin_, bse_vtotal_);
  dmatEX[0] = occlevels * CalcAuxMat_vv(coeffs) * occlevels.transpose();
 
  // electron part as matrix products
  Eigen::MatrixXd virtlevels =
      mos_.eigenvectors().middleCols(bse_cmin_, bse_ctotal_);
  dmatEX[1] = virtlevels * CalcAuxMat_cc(coeffs) * virtlevels.transpose();
 
  return dmatEX;
}
 
Eigen::MatrixXd Orbitals::CalcAuxMat_vv(const Eigen::VectorXd& coeffs) const {
  const Eigen::Map<const Eigen::MatrixXd> mat(coeffs.data(), bse_ctotal_,
                                              bse_vtotal_);
  return mat.transpose() * mat;
}
 
Eigen::MatrixXd Orbitals::CalcAuxMat_cc(const Eigen::VectorXd& coeffs) const {
  const Eigen::Map<const Eigen::MatrixXd> mat(coeffs.data(), bse_ctotal_,
                                              bse_vtotal_);
  return mat * mat.transpose();
}
 
std::array<Eigen::MatrixXd, 2> Orbitals::DensityMatrixExcitedState_AR(
    const QMState& state) const {
 
  if (!state.Type().isExciton()) {
    throw std::runtime_error(
        "Spin type not known for density matrix. Available are singlet and "
        "triplet");
  }
 
  const Eigen::MatrixXd& BSECoefs_AR = (state.Type() == QMStateType::Singlet)
                                           ? BSE_singlet_.eigenvectors2()
                                           : BSE_triplet_.eigenvectors2();
  if (BSECoefs_AR.cols() < state.StateIdx() + 1 || BSECoefs_AR.rows() < 2) {
    throw std::runtime_error("Orbitals object has no information about state:" +
                             state.ToString());
  }
  /******
   *
   *    Density matrix for GW-BSE based excitations
   *
   *    - electron contribution
   *      D_ab = \sum{vc} \sum{v'} B_{vc}B_{v'c} mo_a(v)mo_b(v')
   *
   *    - hole contribution
   *      D_ab = \sum{vc} \sum{c'} B_{vc}B_{vc'} mo_a(c)mo_b(c')
   *
   *
   *   more efficient:
   *
   *   - electron contribution
   *      D_ab = \sum{v} \sum{v'} mo_a(v)mo_b(v') [ \sum{c} B_{vc}B_{v'c} ]
   *           = \sum{v} \sum{v'} mo_a(v)mo_b(v') B_{vv'}
   *
   *   - hole contribution
   *      D_ab = \sum{c} \sum{c'} mo_a(c)mo_b(c') [ \sum{v} B_{vc}B_{vc'} ]
   *           = \sum{c} \sum{c'} mo_a(c)mo_b(c') B_{cc'}
   *
   */
 
  Eigen::VectorXd coeffs = BSECoefs_AR.col(state.StateIdx());
 
  std::array<Eigen::MatrixXd, 2> dmatAR;
  Eigen::MatrixXd virtlevels =
      mos_.eigenvectors().middleCols(bse_cmin_, bse_ctotal_);
  dmatAR[0] = virtlevels * CalcAuxMat_cc(coeffs) * virtlevels.transpose();
  // electron part as matrix products
  Eigen::MatrixXd occlevels =
      mos_.eigenvectors().middleCols(bse_vmin_, bse_vtotal_);
  dmatAR[1] = occlevels * CalcAuxMat_vv(coeffs) * occlevels.transpose();
 
  return dmatAR;
}
 
Eigen::VectorXd Orbitals::Oscillatorstrengths() const {
  return Oscillatorstrengths(QMStateType(QMStateType::Singlet));
}
 
Eigen::VectorXd Orbitals::Oscillatorstrengths(const QMStateType& type) const {
 
  const tools::EigenSystem* es = nullptr;
 
  if (type == QMStateType::Singlet) {
    es = &BSE_singlet_;
  } else if (type == QMStateType::ExcitonUKS) {
    es = &BSE_uks_;
  } else {
    throw std::runtime_error(
        "Orbitals::Oscillatorstrengths only implemented for "
        "restricted singlets and combined UKS excitons.");
  }
 
  Index size = Index(transition_dipoles_.size());
  if (size > es->eigenvalues().size()) {
    size = es->eigenvalues().size();
  }
 
  Eigen::VectorXd oscs = Eigen::VectorXd::Zero(size);
  for (Index i = 0; i < size; ++i) {
    // f = (2/3) * Omega * |d|^2
    // Omega must be in Hartree, d in e*bohr
    oscs(i) =
        transition_dipoles_[i].squaredNorm() * 2.0 / 3.0 * es->eigenvalues()(i);
  }
  return oscs;
}
 
double Orbitals::getTotalStateEnergy(const QMState& state) const {
  double total_energy = getDFTTotalEnergy();
  if (state.Type() == QMStateType::Gstate) {
    return total_energy;
  }
  total_energy += getExcitedStateEnergy(state);
  return total_energy;
}
 
double Orbitals::getExcitedStateEnergy(const QMState& state) const {
 
  double omega = 0.0;
  if (state.isTransition()) {
    throw std::runtime_error(
        "Total Energy does not exist for transition state");
  }
 
  if (state.Type() == QMStateType::Singlet) {
    if (BSE_singlet_.eigenvalues().size() < state.StateIdx() + 1) {
      throw std::runtime_error("Orbitals::getTotalEnergy You want " +
                               state.ToString() +
                               " which has not been calculated");
    }
    omega = BSE_singlet_.eigenvalues()[state.StateIdx()];
  } else if (state.Type() == QMStateType::Triplet) {
    if (BSE_triplet_.eigenvalues().size() < state.StateIdx() + 1) {
      throw std::runtime_error("Orbitals::getTotalEnergy You want " +
                               state.ToString() +
                               " which has not been calculated");
    }
    omega = BSE_triplet_.eigenvalues()[state.StateIdx()];
  } else if (state.Type() == QMStateType::DQPstate) {
    if (QPdiag_.eigenvalues().size() < state.StateIdx() + 1 - getGWAmin()) {
      throw std::runtime_error("Orbitals::getTotalEnergy You want " +
                               state.ToString() +
                               " which has not been calculated");
    }
    return QPdiag_.eigenvalues()[state.StateIdx() - getGWAmin()];
  } else if (state.Type() == QMStateType::KSstate) {
    if (mos_.eigenvalues().size() < state.StateIdx() + 1) {
      throw std::runtime_error("Orbitals::getTotalEnergy You want " +
                               state.ToString() +
                               " which has not been calculated");
    }
    return mos_.eigenvalues()[state.StateIdx()];
  } else if (state.Type() == QMStateType::PQPstate) {
    if (this->QPpert_energies_.rows() < state.StateIdx() + 1 - getGWAmin()) {
      throw std::runtime_error("Orbitals::getTotalEnergy You want " +
                               state.ToString() +
                               " which has not been calculated");
    }
    return QPpert_energies_(state.StateIdx() - getGWAmin(), 3);
  } else if (state.Type() == QMStateType::LMOstate) {
    if (lmos_energies_.size() < state.StateIdx() + 1) {
      throw std::runtime_error(
          "Orbitals::getTotalEnergy You want " + state.ToString() +
          " which is a LMO for virtual orbitals. Not implemented.");
    }
    return lmos_energies_(state.StateIdx());
  } else {
    throw std::runtime_error(
        "GetTotalEnergy only knows states:singlet,triplet,KS,DQP,PQP,LMOs");
  }
  return omega;  //  e.g. hartree
}
 
std::array<Eigen::MatrixXd, 3> Orbitals::CalcFreeTransition_Dipoles() const {
  const Eigen::MatrixXd& dft_orbitals = mos_.eigenvectors();
  AOBasis basis = getDftBasis();
  // Testing electric dipole AOMatrix
  AODipole dft_dipole;
  dft_dipole.Fill(basis);
 
  // now transition dipole elements for free interlevel transitions
  std::array<Eigen::MatrixXd, 3> interlevel_dipoles;
 
  Eigen::MatrixXd empty = dft_orbitals.middleCols(bse_cmin_, bse_ctotal_);
  Eigen::MatrixXd occ = dft_orbitals.middleCols(bse_vmin_, bse_vtotal_);
  for (Index i = 0; i < 3; i++) {
    interlevel_dipoles[i] = empty.transpose() * dft_dipole.Matrix()[i] * occ;
  }
  return interlevel_dipoles;
}
 
void Orbitals::CalcCoupledTransition_Dipoles() {
  CalcCoupledTransition_Dipoles(QMStateType(QMStateType::Singlet));
}
 
void Orbitals::CalcCoupledTransition_Dipoles(const QMStateType& type) {
 
  transition_dipoles_.clear();
 
  if (type == QMStateType::Singlet) {
    std::array<Eigen::MatrixXd, 3> interlevel_dipoles =
        CalcFreeTransition_Dipoles();
 
    Index numofstates = BSE_singlet_.eigenvalues().size();
    transition_dipoles_.reserve(numofstates);
 
    const double sqrt2 = std::sqrt(2.0);
 
    for (Index i_exc = 0; i_exc < numofstates; ++i_exc) {
      Eigen::VectorXd coeffs = BSE_singlet_.eigenvectors().col(i_exc);
      if (!useTDA_) {
        coeffs += BSE_singlet_.eigenvectors2().col(i_exc);
      }
 
      Eigen::Map<Eigen::MatrixXd> mat(coeffs.data(), bse_ctotal_, bse_vtotal_);
 
      Eigen::Vector3d tdipole = Eigen::Vector3d::Zero();
      for (Index i = 0; i < 3; ++i) {
        tdipole[i] = mat.cwiseProduct(interlevel_dipoles[i]).sum();
      }
 
      // Restricted singlet:
      // the excitation is the spin-adapted combination of alpha and beta
      // determinants, so the total transition dipole picks up sqrt(2).
      transition_dipoles_.push_back(-sqrt2 * tdipole);
    }
    return;
  }
 
  if (type == QMStateType::ExcitonUKS) {
    if (!hasUnrestrictedOrbitals()) {
      throw std::runtime_error(
          "Orbitals::CalcCoupledTransition_Dipoles(ExcitonUKS) requires "
          "unrestricted alpha/beta orbitals.");
    }
 
    const Index alpha_homo = getHomoAlpha();
    const Index beta_homo = getHomoBeta();
 
    const Index alpha_vtotal = alpha_homo - bse_vmin_ + 1;
    const Index alpha_ctotal = bse_cmax_ - (alpha_homo + 1) + 1;
    const Index alpha_size = alpha_vtotal * alpha_ctotal;
 
    const Index beta_vtotal = beta_homo - bse_vmin_ + 1;
    const Index beta_ctotal = bse_cmax_ - (beta_homo + 1) + 1;
    const Index beta_size = beta_vtotal * beta_ctotal;
 
    if (alpha_vtotal < 0 || alpha_ctotal < 0 || beta_vtotal < 0 ||
        beta_ctotal < 0) {
      throw std::runtime_error(
          "Orbitals::CalcCoupledTransition_Dipoles(ExcitonUKS): "
          "invalid UKS BSE window.");
    }
 
    AOBasis basis = getDftBasis();
    AODipole dft_dipole;
    dft_dipole.Fill(basis);
 
    std::array<Eigen::MatrixXd, 3> interlevel_dipoles_alpha;
    std::array<Eigen::MatrixXd, 3> interlevel_dipoles_beta;
 
    const Eigen::MatrixXd virt_alpha =
        mos_.eigenvectors().middleCols(alpha_homo + 1, alpha_ctotal);
    const Eigen::MatrixXd occ_alpha =
        mos_.eigenvectors().middleCols(bse_vmin_, alpha_vtotal);
 
    const Eigen::MatrixXd virt_beta =
        mos_beta_.eigenvectors().middleCols(beta_homo + 1, beta_ctotal);
    const Eigen::MatrixXd occ_beta =
        mos_beta_.eigenvectors().middleCols(bse_vmin_, beta_vtotal);
 
    for (Index i = 0; i < 3; ++i) {
      interlevel_dipoles_alpha[i] =
          virt_alpha.transpose() * dft_dipole.Matrix()[i] * occ_alpha;
      interlevel_dipoles_beta[i] =
          virt_beta.transpose() * dft_dipole.Matrix()[i] * occ_beta;
    }
 
    const Index numofstates = BSE_uks_.eigenvalues().size();
    transition_dipoles_.reserve(numofstates);
 
    for (Index i_exc = 0; i_exc < numofstates; ++i_exc) {
      Eigen::VectorXd coeffs_x = BSE_uks_.eigenvectors().col(i_exc);
      Eigen::VectorXd coeffs = coeffs_x;
      if (!useTDA_) {
        coeffs += BSE_uks_.eigenvectors2().col(i_exc);
      }
 
      if (coeffs.size() != alpha_size + beta_size) {
        throw std::runtime_error(
            "Orbitals::CalcCoupledTransition_Dipoles(ExcitonUKS): "
            "exciton vector size does not match alpha+beta sector sizes.");
      }
 
      Eigen::Map<const Eigen::MatrixXd> mat_alpha(coeffs.data(), alpha_ctotal,
                                                  alpha_vtotal);
      Eigen::Map<const Eigen::MatrixXd> mat_beta(coeffs.data() + alpha_size,
                                                 beta_ctotal, beta_vtotal);
 
      Eigen::Vector3d tdipole = Eigen::Vector3d::Zero();
      for (Index i = 0; i < 3; ++i) {
        tdipole[i] = mat_alpha.cwiseProduct(interlevel_dipoles_alpha[i]).sum() +
                     mat_beta.cwiseProduct(interlevel_dipoles_beta[i]).sum();
      }
 
      // Combined UKS exciton:
      // alpha and beta sectors are already explicit components of the
      // eigenvector, so there is NO extra sqrt(2) factor here.
      transition_dipoles_.push_back(-tdipole);
    }
    return;
  }
 
  throw std::runtime_error(
      "Orbitals::CalcCoupledTransition_Dipoles only implemented for "
      "restricted singlets and combined UKS excitons.");
}
 
void Orbitals::OrderMOsbyEnergy() {
  std::vector<Index> sort_index = SortEnergies();
  tools::EigenSystem MO_copy = mos_;
  Index size = mos_.eigenvalues().size();
  for (Index i = 0; i < size; ++i) {
    mos_.eigenvalues()(i) = MO_copy.eigenvalues()(sort_index[i]);
  }
  for (Index i = 0; i < size; ++i) {
    mos_.eigenvectors().col(i) = MO_copy.eigenvectors().col(sort_index[i]);
  }
}

void Orbitals::PrepareDimerGuess(const Orbitals& orbitalsA,
                                 const Orbitals& orbitalsB) {
 
  // This routine constructs a block-diagonal guess only from the restricted
  // MO coefficient matrices mos_. It therefore only makes sense for
  // restricted closed-shell inputs.
  if (orbitalsA.hasUnrestrictedOrbitals() ||
      orbitalsB.hasUnrestrictedOrbitals()) {
    throw std::runtime_error(
        "PrepareDimerGuess currently supports only restricted orbitals.");
  }
 
  if (orbitalsA.getSpin() != 1 || orbitalsB.getSpin() != 1) {
    throw std::runtime_error(
        "PrepareDimerGuess currently supports only closed-shell singlets.");
  }
 
  // constructing the direct product orbA x orbB
  Index basisA = orbitalsA.getBasisSetSize();
  Index basisB = orbitalsB.getBasisSetSize();
 
  // In the restricted closed-shell case, getLumo() is the number of occupied
  // spatial orbitals, which is exactly what this block-merging logic assumes.
  Index occA = orbitalsA.getLumo();
  Index occB = orbitalsB.getLumo();
 
  mos_.eigenvectors() = Eigen::MatrixXd::Zero(basisA + basisB, basisA + basisB);
 
  // AxB = | A 0 |
  //       | 0 B |
  if (orbitalsA.getDFTbasisName() != orbitalsB.getDFTbasisName()) {
    throw std::runtime_error("Basissets of Orbitals A and B differ " +
                             orbitalsA.getDFTbasisName() + ":" +
                             orbitalsB.getDFTbasisName());
  }
  this->SetupDftBasis(orbitalsA.getDFTbasisName());
 
  if (orbitalsA.getECPName() != orbitalsB.getECPName()) {
    throw std::runtime_error("ECPs of Orbitals A and B differ " +
                             orbitalsA.getECPName() + ":" +
                             orbitalsB.getECPName());
  }
  this->setECPName(orbitalsA.getECPName());
 
  // Keep occupation metadata fully consistent for a restricted singlet.
  this->setNumberOfOccupiedLevels(occA + occB);
  this->setNumberOfOccupiedLevelsBeta(occA + occB);
  this->setNumberOfAlphaElectrons(occA + occB);
  this->setNumberOfBetaElectrons(occA + occB);
  this->setChargeAndSpin(orbitalsA.getCharge() + orbitalsB.getCharge(), 1);
 
  this->MOs().eigenvectors().block(0, 0, basisA, basisA) =
      orbitalsA.MOs().eigenvectors();
  this->MOs().eigenvectors().block(basisA, basisA, basisB, basisB) =
      orbitalsB.MOs().eigenvectors();
 
  this->MOs().eigenvalues() = Eigen::VectorXd::Zero(basisA + basisB);
  this->MOs().eigenvalues().segment(0, basisA) = orbitalsA.MOs().eigenvalues();
  this->MOs().eigenvalues().segment(basisA, basisB) =
      orbitalsB.MOs().eigenvalues();
 
  this->OrderMOsbyEnergy();
}
 
void Orbitals::WriteToCpt(const std::string& filename) const {
  std::lock_guard<std::recursive_mutex> lock(
      votca::xtp::checkpoint_utils::Hdf5Mutex());
  CheckpointFile cpf(filename, CheckpointAccessLevel::CREATE);
  WriteToCpt(cpf);
}
 
void Orbitals::WriteToCpt(CheckpointFile& f) const {
  CheckpointWriter writer = f.getWriter("/QMdata");
  WriteToCpt(writer);
  WriteBasisSetsToCpt(writer);
}
 
void Orbitals::WriteBasisSetsToCpt(CheckpointWriter& w) const {
  CheckpointWriter dftWriter = w.openChild("dft");
  dftbasis_.WriteToCpt(dftWriter);
  CheckpointWriter auxWriter = w.openChild("aux");
  auxbasis_.WriteToCpt(auxWriter);
}
 
void Orbitals::WriteToCpt(CheckpointWriter& w) const {
  w(votca::tools::ToolsVersionStr(), "XTPVersion");
  w(orbitals_version(), "version");
  w(occupied_levels_, "occupied_levels");
  w(occupied_levels_beta_, "occupied_levels_beta");
  w(number_alpha_electrons_, "number_alpha_electrons");
  w(number_beta_electrons_, "number_beta_electrons");
  w(total_charge_, "charge");
  w(total_spin_, "spin");
 
  w(mos_, "mos");
  w(mos_beta_, "mos_beta");
  w(occupations_, "occupations");
  w(active_electrons_, "active_electrons");
  w(mos_embedding_, "mos_embedding");
  w(lmos_, "LMOs");
  w(lmos_energies_, "LMOs_energies");
  w(inactivedensity_, "inactivedensity");
  w(expandedMOs_, "TruncMOsFullBasis");
 
  CheckpointWriter molgroup = w.openChild("qmmolecule");
  atoms_.WriteToCpt(molgroup);
 
  w(qm_energy_, "qm_energy");
  w(qm_package_, "qm_package");
 
  w(rpamin_, "rpamin");
  w(rpamax_, "rpamax");
  w(qpmin_, "qpmin");
  w(qpmax_, "qpmax");
  w(bse_vmin_, "bse_vmin");
  w(bse_cmax_, "bse_cmax");
  w(functionalname_, "XCFunctional");
  w(grid_quality_, "XC_grid_quality");
  w(ScaHFX_, "ScaHFX");
 
  w(useTDA_, "useTDA");
  w(ECP_, "ECP");
 
  w(rpa_inputenergies_, "RPA_inputenergies");
  w(QPpert_energies_, "QPpert_energies");
 
  w(QPdiag_, "QPdiag");
 
  w(BSE_singlet_, "BSE_singlet");
 
  w(transition_dipoles_, "transition_dipoles");
 
  w(BSE_triplet_, "BSE_triplet");
 
  std::uint8_t tmp = use_Hqp_offdiag_ ? 1u : 0u;
  w(tmp, "use_Hqp_offdiag");
 
  w(BSE_singlet_energies_dynamic_, "BSE_singlet_dynamic");
 
  w(BSE_triplet_energies_dynamic_, "BSE_triplet_dynamic");
 
  w(CalcType_, "CalcType");
 
  // spin GW
  w(rpa_inputenergies_alpha_, "RPA_inputenergies_alpha");
  w(rpa_inputenergies_beta_, "RPA_inputenergies_beta");
 
  w(QPpert_energies_alpha_, "QPpert_energies_alpha");
  w(QPpert_energies_beta_, "QPpert_energies_beta");
 
  w(QPdiag_alpha_, "QPdiag_alpha");
  w(QPdiag_beta_, "QPdiag_beta");
 
  w(BSE_uks_, "BSE_uks");
  w(BSE_uks_energies_dynamic_, "BSE_uks_dynamic");
}
 
void Orbitals::ReadFromCpt(const std::string& filename) {
  std::lock_guard<std::recursive_mutex> lock(
      votca::xtp::checkpoint_utils::Hdf5Mutex());
  CheckpointFile cpf(filename, CheckpointAccessLevel::READ);
  ReadFromCpt(cpf);
}
 
void Orbitals::ReadFromCpt(CheckpointFile& f) {
  CheckpointReader reader = f.getReader("/QMdata");
  ReadFromCpt(reader);
  ReadBasisSetsFromCpt(reader);
}
 
void Orbitals::ReadBasisSetsFromCpt(CheckpointReader& r) {
  CheckpointReader dftReader = r.openChild("dft");
  dftbasis_.ReadFromCpt(dftReader);
  CheckpointReader auxReader = r.openChild("aux");
  auxbasis_.ReadFromCpt(auxReader);
}
 
void Orbitals::ReadFromCpt(CheckpointReader& r) {
  r(occupied_levels_, "occupied_levels");
  r(number_alpha_electrons_, "number_alpha_electrons");
 
  int version;
  r(version, "version");
 
  auto r_optional = [&r](auto& obj, const std::string& key) {
    try {
      r(obj, key);
    } catch (std::runtime_error&) {
      // Older checkpoints may not contain this field.
    }
  };
 
  // Read qmatoms
  CheckpointReader molgroup = r.openChild("qmmolecule");
  atoms_.ReadFromCpt(molgroup);
 
  r(qm_energy_, "qm_energy");
  r(qm_package_, "qm_package");
  try {
    r(lmos_, "LMOs");
    r(lmos_energies_, "LMOs_energies");
  } catch (std::runtime_error& e) {
    ;
  }
 
  r(version, "version");
  r(mos_, "mos");
  r(mos_embedding_, "mos_embedding");
  r(active_electrons_, "active_electrons");
  r(inactivedensity_, "inactivedensity");
  r(CalcType_, "CalcType");
  r(expandedMOs_, "TruncMOsFullBasis");
 
  // spin info available from version 6 or higher
  if (version > 5) {
    r(occupied_levels_beta_, "occupied_levels_beta");
    r(number_beta_electrons_, "number_beta_electrons");
    r(total_charge_, "charge");
    r(total_spin_, "spin");
    r(mos_beta_, "mos_beta");
    r(occupations_, "occupations");
  }
 
  if (version < 3) {
    // clang-format off
    std::array<Index, 49> votcaOrder_old = {
        0,                             // s
        0, -1, 1,                      // p
        0, -1, 1, -2, 2,               // d
        0, -1, 1, -2, 2, -3, 3,        // f
        0, -1, 1, -2, 2, -3, 3, -4, 4,  // g
        0, -1, 1, -2, 2, -3, 3, -4, 4,-5,5,  // h
        0, -1, 1, -2, 2, -3, 3, -4, 4,-5,5,-6,6  // i
    };
    // clang-format on
 
    std::array<Index, 49> multiplier;
    multiplier.fill(1);
    OrbReorder ord(votcaOrder_old, multiplier);
    ord.reorderOrbitals(mos_.eigenvectors(), this->getDftBasis());
  }
 
  if (version < 5) {  // we need to construct the basissets, NB. can only be
                      // done after reading the atoms.
    std::string dft_basis_name;
    std::string aux_basis_name;
    r(dft_basis_name, "dftbasis");
    r(aux_basis_name, "auxbasis");
    this->SetupDftBasis(dft_basis_name);
    this->SetupAuxBasis(aux_basis_name);
  }
 
  r(rpamin_, "rpamin");
  r(rpamax_, "rpamax");
  r(qpmin_, "qpmin");
  r(qpmax_, "qpmax");
  r(bse_vmin_, "bse_vmin");
  r(bse_cmax_, "bse_cmax");
  setBSEindices(bse_vmin_, bse_cmax_);
  try {
    r(functionalname_, "XCFunctional");
    r(grid_quality_, "XC_grid_quality");
  } catch (std::runtime_error& e) {
    grid_quality_ = "medium";
  }
  r(ScaHFX_, "ScaHFX");
  r(useTDA_, "useTDA");
  r(ECP_, "ECP");
 
  r(rpa_inputenergies_, "RPA_inputenergies");
 
  r(QPpert_energies_, "QPpert_energies");
  r(QPdiag_, "QPdiag");
 
  // Optional unrestricted GW data
  r_optional(rpa_inputenergies_alpha_, "RPA_inputenergies_alpha");
  r_optional(rpa_inputenergies_beta_, "RPA_inputenergies_beta");
 
  r_optional(QPpert_energies_alpha_, "QPpert_energies_alpha");
  r_optional(QPpert_energies_beta_, "QPpert_energies_beta");
 
  r_optional(QPdiag_alpha_, "QPdiag_alpha");
  r_optional(QPdiag_beta_, "QPdiag_beta");
  r_optional(BSE_uks_, "BSE_uks");
  r_optional(BSE_uks_energies_dynamic_, "BSE_uks_dynamic");
  r(BSE_singlet_, "BSE_singlet");
 
  r(transition_dipoles_, "transition_dipoles");
 
  r(BSE_triplet_, "BSE_triplet");
 
  std::uint8_t tmp = 0;
  r(tmp, "use_Hqp_offdiag");
  use_Hqp_offdiag_ = (tmp != 0);
 
  if (version > 1) {
    r(BSE_singlet_energies_dynamic_, "BSE_singlet_dynamic");
 
    r(BSE_triplet_energies_dynamic_, "BSE_triplet_dynamic");
  }
  // Backward compatibility for older restricted checkpoints that may not
  // contain explicit beta-electron information.
  if (!hasUnrestrictedOrbitals() && total_spin_ == 1 && occupied_levels_ > 0) {
    if (number_alpha_electrons_ == 0) {
      number_alpha_electrons_ = occupied_levels_;
    }
    if (number_beta_electrons_ == 0) {
      number_beta_electrons_ = occupied_levels_;
    }
  }
}
 
/*********************************************
 * Extensions Spin-DFT
 **********************************************/
 
// Spin-resolved ground-state density matrices.
//
// The returned pair is (P^alpha, P^beta) with
//
//   P^sigma = C_occ^sigma (C_occ^sigma)^T
//
// for unrestricted calculations. In the restricted branch the same spatial MO
// coefficients are reused and distributed over alpha/beta occupations according
// to the stored electron counts.
std::array<Eigen::MatrixXd, 2> Orbitals::DensityMatrixGroundStateSpinResolved()
    const {
  if (!hasMOs()) {
    throw std::runtime_error("Orbitals file does not contain MO coefficients");
  }
 
  std::array<Eigen::MatrixXd, 2> result;
  result[0] = Eigen::MatrixXd::Zero(mos_.eigenvectors().rows(),
                                    mos_.eigenvectors().rows());
  result[1] = Eigen::MatrixXd::Zero(mos_.eigenvectors().rows(),
                                    mos_.eigenvectors().rows());
 
  if (hasUnrestrictedOrbitals()) {
    if (occupied_levels_ > 0) {
      const Eigen::MatrixXd occ_a =
          mos_.eigenvectors().leftCols(occupied_levels_);
      result[0] = occ_a * occ_a.transpose();
    }
    if (occupied_levels_beta_ > 0) {
      const Eigen::MatrixXd occ_b =
          mos_beta_.eigenvectors().leftCols(occupied_levels_beta_);
      result[1] = occ_b * occ_b.transpose();
    }
    return result;
  }
 
  // Restricted branch:
  // Prefer explicit alpha/beta electron counts, but fall back to the legacy
  // occupied_levels_ convention when those counts were never set.
  Index n_alpha = number_alpha_electrons_;
  Index n_beta = number_beta_electrons_;
 
  if (n_alpha == 0 && n_beta == 0 && occupied_levels_ > 0) {
    n_alpha = occupied_levels_;
    n_beta = occupied_levels_;
  }
 
  const Index n_docc = std::min(n_alpha, n_beta);
  const Index n_socc_alpha = n_alpha - n_docc;
  const Index n_socc_beta = n_beta - n_docc;
 
  if (n_docc > 0) {
    const Eigen::MatrixXd docc = mos_.eigenvectors().leftCols(n_docc);
    const Eigen::MatrixXd d_docc = docc * docc.transpose();
    result[0] += d_docc;
    result[1] += d_docc;
  }
 
  if (n_socc_alpha > 0) {
    const Eigen::MatrixXd socc_a =
        mos_.eigenvectors().middleCols(n_docc, n_socc_alpha);
    result[0] += socc_a * socc_a.transpose();
  }
 
  if (n_socc_beta > 0) {
    const Eigen::MatrixXd socc_b =
        mos_.eigenvectors().middleCols(n_docc + n_socc_alpha, n_socc_beta);
    result[1] += socc_b * socc_b.transpose();
  }
 
  return result;
}
 
}  // namespace xtp
}  // namespace votca