Speed up simulations

[viljarjf]

Description of the change

Speed up diffraction simulation.

With the new RotatedPhase class and pyxem/orix#533, I get around 6x speedup with a simple cubic crystal.

I want to look into SimulationGenerator.get_intersecting_reflections too, maybe we can rotate Ewald's sphere instead of the lattice to reduce computation.
Additionally, iterating over the rotations should be possible to parallelize.
@CSSFrancis maybe you have looked at these things already?

Progress of the PR

• Avoid deepcopy when rotating basis in DiffractingVector.rotate_with_basis
• Avoid expensive deepcopy when rotating Miller orix#533
• Look into rotating Ewald's sphere instead of the lattice
• Parallelization
• Docstrings for all functions
• Unit tests with pytest for all lines - we could add some test where we store the result from the main branch and compare with this and future outputs
• Clean code style by running black

For reviewers

• The PR title is short, concise, and will make sense 1 year later.
• New functions are imported in corresponding __init__.py.
• New features, API changes, and deprecations are mentioned in the
  unreleased section in CHANGELOG.rst.
• Contributor(s) are listed correctly in credits in diffsims/release_info.py and
  in .zenodo.json.

[viljarjf]

Performance testing

Using the following script to test before/after, testing different crystals (hexagonal, cubic and triclinic) and both small and large unit cells (to decrease/increase the number of reflections - which should be where the greatest performance gain is realized):

from diffsims.generators.simulation_generator import SimulationGenerator
from orix.crystal_map import Phase
from orix.quaternion import Rotation
from diffpy.structure import Lattice, Atom, Structure

import numpy as np
import timeit

gold = [Atom("Au", [0, 0, 0])]
hexagonal = Phase("test", 161, structure=Structure(gold,
        Lattice(4, 4, 5, 90, 90, 120)
))
large_hexagonal = Phase("test", 161, structure=Structure(gold,
        Lattice(20, 20, 25, 90, 90, 120)
))
cubic = Phase("test", 221, structure=Structure(gold,
        Lattice(4, 4, 4, 90, 90, 90)
))
large_cubic = Phase("test", 221, structure=Structure(gold,
        Lattice(20, 20, 20, 90, 90, 90)
))
triclinic = Phase("test", 1, structure=Structure(gold,
        Lattice(4, 5, 6, 80, 90, 130)
))
large_triclinic = Phase("test", 1, structure=Structure(gold,
        Lattice(20, 25, 30, 80, 90, 130)
))

gen = SimulationGenerator()
from numpy import random
random.seed(0)
rot = Rotation.random(1000)
kwargs = {
    "rotation": rot,
    "with_direct_beam": False,
    "reciprocal_radius": 5,
}

res = timeit.repeat(
    "gen.calculate_diffraction2d(hexagonal, **kwargs)",
    globals=globals(),
    number=2,
)
print(f"{'hexagonal' :<18}: {np.mean(res) :.2f} ± {np.std(res) :.2f} s")

res = timeit.repeat(
    "gen.calculate_diffraction2d(cubic, **kwargs)",
    globals=globals(),
    number=2,
)
print(f"{'cubic' :<18}: {np.mean(res) :.2f} ± {np.std(res) :.2f} s")

res = timeit.repeat(
    "gen.calculate_diffraction2d(triclinic, **kwargs)",
    globals=globals(),
    number=2,
)
print(f"{'triclinic' :<18}: {np.mean(res) :.2f} ± {np.std(res) :.2f} s")

kwargs["reciprocal_radius"] = 2 # we still get a huge number of reflections
res = timeit.repeat(
    "gen.calculate_diffraction2d(large_hexagonal, **kwargs)",
    globals=globals(),
    number=2,
)
print(f"{'large_hexagonal' :<18}: {np.mean(res) :.2f} ± {np.std(res) :.2f} s")

res = timeit.repeat(
    "gen.calculate_diffraction2d(large_cubic, **kwargs)",
    globals=globals(),
    number=2,
)
print(f"{'large_cubic' :<18}: {np.mean(res) :.2f} ± {np.std(res) :.2f} s")

res = timeit.repeat(
    "gen.calculate_diffraction2d(large_triclinic, **kwargs)",
    globals=globals(),
    number=2,
)
print(f"{'large_triclinic' :<18}: {np.mean(res) :.2f} ± {np.std(res) :.2f} s")

Running on current main branch of diffsims and orix:

hexagonal         : 7.16 ± 0.36 s
cubic             : 7.29 ± 0.48 s
triclinic         : 9.94 ± 0.12 s
large_hexagonal   : 80.11 ± 2.59 s
large_cubic       : 79.15 ± 5.94 s
large_triclinic   : 94.56 ± 3.78 s

And with the full changes in this branch + Orix:

hexagonal         : 3.32 ± 0.15 s
cubic             : 3.16 ± 0.11 s
triclinic         : 3.53 ± 0.13 s
large_hexagonal   : 29.00 ± 0.07 s
large_cubic       : 28.82 ± 0.23 s
large_triclinic   : 35.59 ± 0.12 s

So roughly 2-3x speedup. Not sure what I did to get the 5-6x speedup I saw before, probably some mistake.

When adding precession, the speedup is negligible. It might even be slower. Below is the script run with 1 degree precession, run on a different computer than above, so with/without precession times are not comparable.
main branch:

hexagonal         : 4.71 ± 0.12 s
cubic             : 4.84 ± 0.08 s
triclinic         : 5.08 ± 0.03 s
large_hexagonal   : 70.00 ± 5.09 s
large_cubic       : 67.83 ± 3.98 s
large_triclinic   : 82.24 ± 7.92 s

This branch:

hexagonal         : 4.13 ± 0.06 s
cubic             : 3.97 ± 0.02 s
triclinic         : 4.93 ± 0.20 s
large_hexagonal   : 85.45 ± 3.43 s
large_cubic       : 82.18 ± 3.69 s
large_triclinic   : 86.97 ± 13.73 s