Optimizing Kwant computations

[akhmerov]

Hi all! A couple of quick suggestions for speeding up the Kwant simulations:

Potential computation

Kwant evaluates potentials one site at a time, which is inefficient because it maximizes Python overhead, itself quite appreciable in this function because it uses a bunch of Python calls.

https://github.com/condensedAI/KwantRL/blob/42be1b6eb1d90f51a65c98cae22fabd2ef000536/simulations/pixel_array_sim_2.py#L25-L30

Additionally site.pos, while it looks very similar to site.tag takes a fair bit more time because it multiplies the site lattice coordinates by the lattice vectors.

In order to avoid this bottleneck, I propose to:

1. Precompute each gate potential at a fixed voltage V=1 as a numpy array, and using numpy functions.
2. (optional) multiply each potential with the corresponding voltage and sum them together into the final potential.
3. Read off the value from the array

The main point here is to use large numpy arrays to compute the potentials instead of one site at a time operations.

Lead modes

Computation of lead modes may take a fair fraction of the total computation time for the scattering matrix. Now you are computing the lead modes each time you are computing the scattering matrix. You can "freeze" the lead Hamiltonian by calling syst_precalculated = syst.precalculate(energy). This precomputes the lead modes at a specific energy. You can either use it to cache all energies of interest, or even fix a single energy for the leads and use that one across all simulations.

[everthemore]

Fantastic, thank you for taking the time in going through the code. Can I adding you to this, @Torbjorn-Rasmussen ?

[akhmerov]

Ah, another thing that I didn't mention, of course, is parallelization, which you can also do e.g. using adaptive.

[Torbjorn-Rasmussen]

Hi i have been writing this up today and tested for a bit, thanks for taking a look at the code. The potential calculation gave a very nice speed up with a factor of a bit more than 2.

Precalculating the lead modes didn't give a noticable difference, but it might not be implemented correctly. I just added it to the finalized system so whenever the finalized system was called previously it now calls syst_precalculated, is this the correct implementation?

self.fqpc = self.qpc.finalized() self.fqpc = self.fqpc.precalculate(self.energy)

Regarding adaptive i think we did try to implement long ago but never finished, i will try to dig out the old files and move ahead with that.

[akhmerov]

Well, I didn't run a profile so I'm not sure whether precalculating modes would give a speedup. Can you share what you have e.g. via a branch, so that I can check the implementation?

[Torbjorn-Rasmussen]

I have pushed to main now, in the interest of cleaning up the repository down the line, i have however made a seperate script for the updated simulation, also to easily compare the two. The simulation with precalculated potentials is in https://github.com/condensedAI/KwantRL/blob/main/simulations/pixel_array.py. I can add some comments where i changed things to make it easier to identify changes now that i made a seperate file.

edit: i can see that linking does not work, the correct file is located in simulations/pixel_array.py https://github.com/condensedAI/KwantRL/blob/5324af4d5e0beab7b836471d722a3efc3853d6dd/simulations/pixel_array.py

[Ulfgard]

you can just link to a commit hash and we can look up exactly which lines got changed via the generated diff :-)

---

From: Torbjorn-Rasmussen @.***> Sent: Friday, December 3, 2021 4:45:14 PM To: condensedAI/KwantRL Cc: Subscribed Subject: Re: [condensedAI/KwantRL] Optimizing Kwant computations (Issue #4)

I have pushed to main now, in the interest of cleaning up the repository down the line, i have however made a seperate script for the updated simulation, also to easily compare the two. The simulation with precalculated potentials is in https://github.com/condensedAI/KwantRL/blob/main/simulations/pixel_array.py. I can add some comments where i changed things to make it easier to identify changes now that i made a seperate file.

— You are receiving this because you are subscribed to this thread. Reply to this email directly, view it on GitHubhttps://eur02.safelinks.protection.outlook.com/?url=https%3A%2F%2Fgithub.com%2FcondensedAI%2FKwantRL%2Fissues%2F4%23issuecomment-985624099&data=04%7C01%7Coswin.krause%40di.ku.dk%7Cce6c73623c834295c6e008d9b673e69d%7Ca3927f91cda14696af898c9f1ceffa91%7C0%7C0%7C637741431166825678%7CUnknown%7CTWFpbGZsb3d8eyJWIjoiMC4wLjAwMDAiLCJQIjoiV2luMzIiLCJBTiI6Ik1haWwiLCJXVCI6Mn0%3D%7C3000&sdata=v6m6BIH1j9PAHf3V5qmYsRDlEWhVEQqWuwZxFDpiGfM%3D&reserved=0, or unsubscribehttps://eur02.safelinks.protection.outlook.com/?url=https%3A%2F%2Fgithub.com%2Fnotifications%2Funsubscribe-auth%2FADSZGBRCDTVX3DVL3J3KF5TUPDQYVANCNFSM5JE7HXLQ&data=04%7C01%7Coswin.krause%40di.ku.dk%7Cce6c73623c834295c6e008d9b673e69d%7Ca3927f91cda14696af898c9f1ceffa91%7C0%7C0%7C637741431166825678%7CUnknown%7CTWFpbGZsb3d8eyJWIjoiMC4wLjAwMDAiLCJQIjoiV2luMzIiLCJBTiI6Ik1haWwiLCJXVCI6Mn0%3D%7C3000&sdata=2NC52lVD6IGeIn4cTDB2QU3El9Jq5kL%2F%2FNkNGrXHxlk%3D&reserved=0. Triage notifications on the go with GitHub Mobile for iOShttps://eur02.safelinks.protection.outlook.com/?url=https%3A%2F%2Fapps.apple.com%2Fapp%2Fapple-store%2Fid1477376905%3Fct%3Dnotification-email%26mt%3D8%26pt%3D524675&data=04%7C01%7Coswin.krause%40di.ku.dk%7Cce6c73623c834295c6e008d9b673e69d%7Ca3927f91cda14696af898c9f1ceffa91%7C0%7C0%7C637741431166835672%7CUnknown%7CTWFpbGZsb3d8eyJWIjoiMC4wLjAwMDAiLCJQIjoiV2luMzIiLCJBTiI6Ik1haWwiLCJXVCI6Mn0%3D%7C3000&sdata=jrUlj6BopsoHF2vjOLaSe8ByZ0UlgJI%2FtMrBInCJ2B4%3D&reserved=0 or Androidhttps://eur02.safelinks.protection.outlook.com/?url=https%3A%2F%2Fplay.google.com%2Fstore%2Fapps%2Fdetails%3Fid%3Dcom.github.android%26referrer%3Dutm_campaign%253Dnotification-email%2526utm_medium%253Demail%2526utm_source%253Dgithub&data=04%7C01%7Coswin.krause%40di.ku.dk%7Cce6c73623c834295c6e008d9b673e69d%7Ca3927f91cda14696af898c9f1ceffa91%7C0%7C0%7C637741431166835672%7CUnknown%7CTWFpbGZsb3d8eyJWIjoiMC4wLjAwMDAiLCJQIjoiV2luMzIiLCJBTiI6Ik1haWwiLCJXVCI6Mn0%3D%7C3000&sdata=isTfKCHfYaWiiuMYUr8WCcXJoKGHMsUyV2DHe%2F4vvL8%3D&reserved=0.

[Torbjorn-Rasmussen]

Yes, the commit hash is: 5324af4d5e0beab7b836471d722a3efc3853d6dd, unfortunately i made it in a seperate file so it's not really useful for seeing changes, the script i made is pixel_array.py but that is all additions.

[Torbjorn-Rasmussen]

Ah, another thing that I didn't mention, of course, is parallelization, which you can also do e.g. using adaptive.

I am slightly familiar with your library adaptive, we used a bit when Amber was a Masters student in our group, i had not noticed the implementation of using multiple cores it comes with. Are you suggesting just using the multiprocessing implementation or also trying to learn the loss function using the full library?

[Torbjorn-Rasmussen]

The class is now able to run in parallel using the multiprocessing package, it has been implemented as a function of the class in this commit: 066bbbe426d96e2c552dda180b9b1bbf5e6aeeb1

For my personal computer with 4 cores, this, in combination with the precalculation of some of the potentials, brought down the computational time by almost a factor 10.