How to condense a mat using np

[pkozlows]

I am having some success at getting correct 1e HF energy, but in my implementation, I am using a lot of for loops, which I know isn't very efficient. Ideally, I would be able to use np.einsum() to get the trace of h1e.npy, so I need to figure out a way to condense this array down into the relevant spatial orbital indices, if that makes sense. For example, I would use:

for spacial_orb in spacial_or_petals: one_elec_mel += integrals[0][spacial_orb][spacial_orb]

is there an existing way to do this task in np or what I just want to slice the array using the relevent indices?

[Walter-Feng]

You can use np.ix_ to get rid of for loops.