Skip to content

Adding New Functionality (Example)

Jump to bottom
william-dawson edited this page Jun 16, 2017 · 8 revisions

Overview

On this Wiki page, we will sketch out how one might go about adding new functionality to NTPoly. The example we give will be based on trying to add a McWeeny purification module.

Theory

McWeeny purification can be used to compute the density matrix from a Hamiltonian using a very tight recursive scheme:

P_{k+1} = 3P_k^2 - 2P_k^3

P_0 = \frac{\lambda}{2}(\mu I - H) + \frac{1}{2}I

Here is a potential python implementation:

p_0 = (lambda_val/2.0) * (chemical_potential*identity_mat - hamiltonian) + 0.5*identity_mat
for i in range(0,100):
  p2 = numpy.dot(p_0,p_0)
  p3 = numpy.dot(p_0,p2)
  p0 = 3.0*p2 - 2.0*p3

Getting Started

First, it is important to build an outline of the routine so that it will be easy to test. Let's begin by building a new Fortran module in Source/Fortran. We'll call it McWeenyModule.f90.

Fortran Code Outline

module McWeenyModule
  use IterativeSolversModule
  use DistributedBlockedSparseMatrixModule
  use DistributedMatrixMemoryPoolModule
  use DataTypesModule
  use mpi
implicit none
  public :: Solve
contains
  subroutine SolveMcWeeny(Hamiltonian, InverseSquareRoot, Density, &
    & chemical_potential, solver_parameters_in)
    !! Parameters
    type(DistributedSparseMatrix), intent(in)  :: Hamiltonian, InverseSquareRoot
    type(DistributedSparseMatrix), intent(inout) :: Density
    real(NTREAL), intent(in) :: chemical_potential
    type(IterativeSolverParameters), intent(in), optional :: solver_parameters_in
    !! Handling Optional Parameters
    type(IterativeSolverParameters) :: solver_parameters

    !! Optional Parameters
    if (present(solver_parameters_in)) then
      solver_parameters = solver_parameters_in
    else
      solver_parameters = IterativeSolverParameters()
    end if

  end subroutine SolveMcWeeny
end module McWeenyModule

This new module will take in the necessary matrices along with the solver parameters and the chemical potential. Before coding the actual solver, we need to wrap up this module so it an be called from C++. This time, we go to the Source/Wrappers/ directory, and add a McWeenySolverModule_wrp.f90 file.

Wrapper Code Outline

module McWeenySolversModule_wrp
  use IterativeSolversModule_wrp, only : IterativeSolverParameters_wrp
  use DensityMatrixSolversModule, only : SolveMcweeny
  use DistributedBlockedSparseMatrixModule_wrp, only : &
    & DistributedSparseMatrix_wrp
  use DataTypesModule, only : NTREAL
  use WrapperModule, only : SIZE_wrp
  use iso_c_binding, only : c_int
  implicit none
  private
  public :: SolveMcWeeny_wrp
contains
  subroutine SolveMcWeeny_wrp(ih_Hamiltonian, ih_InverseSquareRoot, &
    & ih_Density, chemical_potential, ih_solver_parameters) &
    & bind(c,name="SolveMcWeeny")
    integer(kind=c_int), intent(in) :: ih_Hamiltonian(SIZE_wrp)
    integer(kind=c_int), intent(in) :: ih_InverseSquareRoot(SIZE_wrp)
    integer(kind=c_int), intent(inout) :: ih_Density(SIZE_wrp)
    real(NTREAL), intent(in) :: chemical_potential
    integer(kind=c_int), intent(in) :: ih_solver_parameters(SIZE_wrp)
    type(DistributedSparseMatrix_wrp) :: h_Hamiltonian
    type(DistributedSparseMatrix_wrp) :: h_InverseSquareRoot
    type(DistributedSparseMatrix_wrp) :: h_Density
    type(IterativeSolverParameters_wrp) :: h_solver_parameters

    h_Hamiltonian = transfer(ih_Hamiltonian,h_Hamiltonian)
    h_InverseSquareRoot = transfer(ih_InverseSquareRoot,h_InverseSquareRoot)
    h_Density = transfer(ih_Density,h_Density)
    h_solver_parameters = transfer(ih_solver_parameters, h_solver_parameters)

    call SolveMcWeeny(h_Hamiltonian%data, h_InverseSquareRoot%data, &
         & h_Density%data, chemical_potential, h_solver_parameters%data)
  end subroutine SolveMcWeeny_wrp
end module McWeenySolversModule_wrp

This wrapper is necessary to hide the distributed datatypes from C. Writing a module like this is mostly just copy and pasting. For all the derived types, you need to pass an integer handle to it, and then use the transfer function to convert handle to the real object. Next we add an object oriented C++ layer in Source/CPlusPlus.

C++ Header File

#include "SolverBase.h"
namespace NTPoly
{
  class IterativeSolverParameters;
  class DistributedSparseMatrix;
  class McWeenySolvers : public SolverBase
  {
  public:
    static void Solve(const DistributedSparseMatrix& Hamiltonian, \
      const DistributedSparseMatrix& InverseSquareRoot, \
      DistributedSparseMatrix& Density, double chemical_potential, \
      const IterativeSolverParameters& solver_parameters);
  };
}

C++ Source File

#include "McWeenySolvers.h"
#include "IterativeSolversParameters.h"
#include "DistributedBlockedSparseMatrix.h"

extern "C"
{
  void McWeenySolvers_wrp(const int* ih_Hamiltonian, const int* ih_InverseSquareRoot, \
    int* ih_Density, const double* chemical_potential_out,
    const int* ih_solver_parameters);
}

namespace NTPoly
{
  void DensityMatrixSolvers::Solvers(const DistributedSparseMatrix& Hamiltonian, \
    const DistributedSparseMatrix& Overlap, \
    DistributedSparseMatrix& Density, double chemical_potential, \
    const IterativeSolverParameters& solver_parameters)
  {
    McWeenySolvers_wrp(GetIH(Hamiltonian), GetIH(Overlap), GetIH(Density), \
      &chemical_potential, GetIH(solver_parameters));
  }
}

This is also pretty much just copy and pasting.

Compilation and Testing

We've now added four new files. Each one of those files need to be added to the CMakeLists.txt files in their individual folders. Next, in the Swig directory, there is a file NTPolySwig.i . Here you need to add include statements for the new C++ header files.

#include "McWeenySolvers.h"

%include "McWeenySolvers.h"

Now let's add a new test that calls and verifies our McWeeny routine. In the UnitTests/Tests/ folder, there is testChemistry.py.in file. Here if we add a new test, Python will automatically add it to our unit testing setup. The Chemistry solvers have been specially set up to compare to some precomputed ground truth data.

def test_mcweeny(self):
  fock_matrix = nt.DistributedSparseMatrix(self.hamiltonian)
  overlap_matrix = nt.DistributedSparseMatrix(self.overlap)
  inverse_sqrt_matrix = nt.DistributedSparseMatrix(
    fock_matrix.GetActualDimension())
  density_matrix = nt.DistributedSparseMatrix(
    fock_matrix.GetActualDimension())

  permutation = nt.Permutation(fock_matrix.GetLogicalDimension())
  permutation.SetRandomPermutation()
  self.solver_parameters.SetLoadBalance(permutation)

  nt.SquareRootSolvers.InverseSquareRoot(overlap_matrix,inverse_sqrt_matrix,
    self.solver_parameters)
  chemical_potential = 0.0
  nt.McWeenySolvers.Solve(fock_matrix, inverse_sqrt_matrix,
    density_matrix, chemical_potential, self.solver_parameters)

  density_matrix.WriteToMatrixMarket(self.result_file)
  comm.barrier()

  self.check_full()

Now when one runs make test, this test will be automatically executed. Note that this is a .py.in file. Any changes here require you to recompile the code, even though it is python.

Core Iteration

Finally we can return to the first Fortran file and implement the actual McWeeny purification routine. First let's set up the local variables we'll need.

Clone this wiki locally