[petsc-users] Petsc4py Simulation: Mat.axpy() slow

Donald Rex Planalp Donald.Planalp at colorado.edu
Fri Nov 10 20:53:41 CST 2023


Hello,

I am trying to use petsc4py to conduct a quantum mechanics simulation. I've
been able to construct all of the relevant matrices, however I am reaching
a gigantic bottleneck.

For the simplest problem I am running I have a few matrices each about
5000x5000. In order to begin time propagation I need to add these matrices
together. However, on 6 cores of my local machine it is taking
approximately 1-2 seconds per addition. Since I need to do this for each
time step in my simulation it is prohibitively slow since there could be
upwards of 10K time steps.

Below is the relevant code:

structure = structure=PETSc.Mat.Structure.DIFFERENT_NONZERO_PATTERN
    if test2:
        def makeLeft(S,MIX,ANG,ATOM,i):


            S.axpy(-Field.pulse[i],MIX,structure)
            S.axpy(-Field.pulse[i],ANG,structure)
            S.axpy(-1,ATOM,structure)
            return S
        def makeRight(S,MIX,ANG,ATOM,i):



            S.axpy(Field.pulse[i],MIX,structure)
            S.axpy(Field.pulse[i],ANG,structure)
            S.axpy(1,ATOM,structure)

            return S

        H_mix = Int.H_mix
        H_mix.scale(1j * dt /2)

        H_ang = Int.H_ang
        H_ang.scale(1j * dt /2)

        H_atom = Int.H_atom
        H_atom.scale(1j * dt /2)

        S = Int.S_total

        psi_initial = psi.psi_initial.copy()
        ksp = PETSc.KSP().create(PETSc.COMM_WORLD)


        for i,t in enumerate(box.t):
            print(i,L)



            O_L = makeLeft(S,H_mix,H_ang,H_atom,i)
            O_R = makeRight(S,H_mix,H_ang,H_atom,i)



            if i == 0:
                known = O_R.getVecRight()
                sol = O_L.getVecRight()

            O_R.mult(psi_initial,known)

            ksp.setOperators(O_L)


            ksp.solve(known,sol)




            psi_initial.copy(sol)


I need to clean it up a bit, but the main point is that I need to add the
matrices many times for a single time step. I can't preallocate memory very
well since some of the matrices aren't the most sparse either. It seems if
I cannot speed up the addition it will be difficult to continue so I was
wondering if you had any insights.

Thank you for your time
-------------- next part --------------
An HTML attachment was scrubbed...
URL: <http://lists.mcs.anl.gov/pipermail/petsc-users/attachments/20231110/810ca56f/attachment.html>


More information about the petsc-users mailing list