Parallelization issues with OrbitalMoment() in QuantumATK

[corsicATK]

Hello everyone,

I need to calculate atom-resolved orbital moments using OrbitalMoment() in QuantumATK and obtain the values per atom and per k-point. My first question is: is it possible to directly get this information from the output?

If not, my fallback solution is to compute the function for a single k-point at a time and loop over all k-points. This works fine sequentially.

However, when I try to parallelize my code with mpi4py by distributing k-points over several ranks, I run into problems:

Deadlocks can occur during the computation.

At the end of the loop, all ranks return the same orbital moment values, even though each rank is supposed to work on a different k-point.

It seems that OrbitalMoment() does not execute independently on each rank and requires some form of internal synchronization or communication.

I would appreciate any guidance on:

Whether it is possible to obtain atom-resolved orbital moments for each k-point directly.

If independent parallel computation per k-point is possible, and how to do it safely with QuantumATK.

Thanks a lot for your help!

[filipr]

Most QuantumATK operations do collective calls and cannot be manually mpi parallelized by e.g. just splitting a list of k-points among mpi processes. What you can do instead is that you can first do the DFT ground state in a separate calculation and store the configuration to an HDF5 file. Then you can make a separate post-processing script that calculate OrbitalMoment() for a single k-point in a loop over a subset of all the k-points you are interested in, where the subset is determined from an input argument or environment variable (e.g. job array index). Here's an unfinished example to give you the idea:

Code

import sys
my_process_index = int(sys.argv[1])
total_num_procs = ??

all_kpoints = numpy.array([[...], [...], ..., [...]])

my_kpoints = # logic to get local sequence of k-points here

configuration = nlread("dft_ground_state_calc.hdf5", BulkConfiguration)[-1]

orbital_moments = []
for kpoint in my_kpoints:
    orb_moment_analysis = OrbitalMoment(configuration, kpoints=MonkhorstPackGrid(1, 1, 1, k_point_shift=kpoint))
    orbital_moments.append(orb_moment_analysis.atomResolvedOrbitalMoment())

# Save 'orbital_moments' to file somehow, e.g. pickle

Then you can run this script either in a bash loop in a single script or submit it as multiple jobs or as a job array. Then you can collect the result files from each job and gather into a single array.

[corsicATK]

Thank you very much for your reply! That really confirms what I was thinking.
I actually did something similar (running several independent jobs), and it works. Thanks for your help.