[RFC] 3D static structure factor

[ChiCheng45]

Description of work
An implementation of the 3D static structure factor. Additionally, it should reduce memory usage for multiprocess jobs if the output data is large, which can occur for this job on windows.

Tasks remaining

• Fix when a large trajectory is used. When a large number of atoms are used, memory usage is very high.
• Allow the options to change q-vector basis and etc.
• Change the current implementation so that it uses FFT to calculate the result? Perhaps use something like https://github.com/flatironinstitute/finufft?
• Check performance and memory usage. Fix if excessive.
• Update documentation.

[ChiCheng45]

Here is the result for SrTiO3. Scattering lengths are not implemented yet, and I will need to adjust and fix the grid spacing so it depends on the unit cell.

[ChiCheng45]

I've updated so that the grid depends on the reciprocal lattice. I've also added partial terms to the results. Here is the results for the SrTiO3.

Total

TiTi

Note the because of the choice we made in #963, the each point above averages the results over $-q$ and $q$.

[MBartkowiakSTFC]

I've updated so that the grid depends on the reciprocal lattice. I've also added partial terms to the results. Here is the results for the SrTiO3.

This is already looking realistic.

Probably the next thing that would be good to have is reciprocal lattice units instead of 1/nm in the plot axes. Since this is tetragonal I4/mcm (but with c dimension on the first axis), peaks like 010 or 001 should not be visible, but 011, 020 and 002 should be there. I think that's what we already see in the plot above, but it would be good for the axis labels to reflect this.

[ChiCheng45]

@MBartkowiakSTFC what do you reckon about adding finufft as a depedency? Seems to be a good use case for this problem, and it's massively faster than the current method. For a 100,000 atom system with 300,000 or more q-vectors its like going from hours to minutes.

direct (c1ef616)

    def run_step(self, index):
        frame_idx = index // self.q_vec_n_chunks
        chunk_idx = index % self.q_vec_n_chunks

        frame = self.configuration["frames"]["value"][frame_idx]
        coords = self.trajectory.configuration(frame)["coordinates"]
        start = chunk_idx * self.q_vec_chunk_size
        stop = min((chunk_idx + 1) * self.q_vec_chunk_size, self.n_q_vecs)

        q_vectors, q_idxs = q_vectors_in_cube(self.inverse, self.max_hkl, start, stop)

        rho = {}
        for element, idxs in self._indicesPerElement.items():
            selectedCoordinates = np.take(coords, idxs, axis=0)
            rho[element] = np.sum(
                np.exp(1j * np.dot(selectedCoordinates, q_vectors)), axis=0
            )

        return index, (q_idxs, rho)

    def combine(self, index, results):
        q_idxs, rho = results
        for pair_str, (label_i, label_j) in self.labels:
            np.add.at(
                self._outputData[f"ssf3d/{pair_str}"],
                q_idxs,
                (rho[label_i] * rho[label_j].conj()).real,
            )

non uniform fft

    def run_step(self, index):
        frame = self.configuration["frames"]["value"][index]
        coords = self.trajectory.configuration(frame)["coordinates"]

        rho = {}
        for element, idxs in self._indicesPerElement.items():
            pts = 2 * np.pi * coords[idxs] @ self.inverse
            rho[element] = finufft.nufft3d1(
                pts[:, 0],
                pts[:, 1],
                pts[:, 2],
                np.ones(len(pts), dtype=np.complex128),
                n_modes=(
                    2 * self.max_hkl[0] + 1,
                    2 * self.max_hkl[1] + 1,
                    2 * self.max_hkl[2] + 1
                ),
                eps=1e-3,
            )

        return index, rho

    def combine(self, index, rho):
        for pair_str, (label_i, label_j) in self.labels:
            self._outputData[f"ssf3d/{pair_str}"] += (
                rho[label_i] * rho[label_j].conj()
            ).real

[MBartkowiakSTFC]

@MBartkowiakSTFC what do you reckon about adding finufft as a depedency? Seems to be a good use case for this problem, and it's massively faster than the current method. For a 100,000 atom system with 300,000 or more q-vectors its like going from hours to minutes.

Since other packages for the same purpose (e.g. MP_tools mentioned in #1196) use finufft, we were going to look at it sooner or later.

It's a <3MB download, we might as well use it.

[oerc0122]

As we chatted about, if this becomes an optional dependency until we add more use-cases I'd be happier, but I'm also not opposed to just using it directly. It seems well resourced, well used and likely to stay alive in the long term, so yeah.