Ragged Vectors and Orix

[CSSFrancis]

@hakonanes @argerlt @tinabe I'm curious on your thoughts on some (larger) changes to the underlying strucuture of data stored in orix.

The (big) problem is that a lot of the operations and data in orix would be widely useful in different packages like pyxem, diffsims, py4dstem etc but all of those packages require a ragged data strucuture or more truely n dimensional objects.

For example consider in diffsims where we have a ReciprocalLatticeVector. For a simulation of different orientations we return a list of ReciprocalLatticeVector objects but we could just as easily return a ragged array with a shape (num_simulations, ragged (num vectors per sim), 3).

For something like pyxem, when we find diffraction peaks these diffraction peaks are essentially two dimensional slices in the full 3D reciporical space and it would be nice to be able to efficiently rotate them using orix. Same with orientation mapping, it would be nice to be able to have the pyxem results for orientation mapping to be a crystal map where you could do something like restrict the results to only return orientations close to some zone axis etc.

I was thinking about implementing something like https://docs.pytorch.org/docs/stable/nested.html where data is stored like:

This would be a lightweight package that Hyperspy, orix, pyxem could depend on and could easily be replaced with the pytorch version if people want to do more intensive ML.

[argerlt]

I'm brand new to this concept, so I've got a base level question: what is the advantage to doing this over, say ,[x.some_function(y) for x,y in zip(vector1,vector2)]? Or, if you really are thinking ahead and want some speed,

threads = [threading.Thread(target=x.somefunction(y)) for x,y in zip(vector1,vector2) ]
[x.start() for x in threads]
[x.join() for x in threads]

Or the equivalent idea but with Pools.

Is it just the mental clarity of not having to think through list comprehension every time, or is there some other added advantages I'm not seeing? Would a jagged array be considerably faster for doing the calculations you want? I could understand it working better with dask arrays or numba functions than the hacked-together threading I'm suggesting...

Also, have you used this before? pytorch's website does a good job making me believe this feature would have no impact for people who don't use it, but they also use the phrase "prototype", which is always....concerning.

[hakonanes]

Is this implementation enough for you, @CSSFrancis? A new minimal ragged vector class stores flattened vectors and the consecutive sizes of each set of vectors. We provide access via a getter, accepting an indexing into the sets, and iteration by iterating over the sets.

import numpy as np

import orix.vector as ove


class RaggedVector3d:
    def __init__(self, v: ove.Vector3d, sizes: np.ndarray) -> None:
        v = v.flatten()
        sizes = sizes.ravel()

        sizes_sum = np.sum(sizes)
        if sizes_sum != v.size:
            raise ValueError(
                f"Sum of sizes {sizes_sum} must equal number of vectors {v.size}"
            )

        self._sizes = sizes
        self._v = v

    @property
    def v(self) -> ove.Vector3d:
        return self._v

    @property
    def sizes(self) -> np.ndarray:
        return self._sizes

    @property
    def offsets(self) -> np.ndarray:
        return np.concatenate(([0], np.cumsum(self._sizes)))

    @property
    def size(self) -> int:
        return self._sizes.size

    def __repr__(self) -> str:
        return (
            f"{self.__class__.__name__}(size={self.size}, sizes={self.sizes.tolist()})"
        )

    def __getitem__(self, idx: int) -> ove.Vector3d:
        if idx < 0:
            start, end = idx - 1, idx
        else:
            start, end = idx, idx + 1
        offsets = self.offsets
        offset_start, offset_end = offsets[start], offsets[end]
        v = self._v[offset_start:offset_end]
        return v

    def __iter__(self) -> ove.Vector3d:
        for i in range(self.size):
            yield self[i]

In use:

import matplotlib.pyplot as plt
import numpy as np

import orix.plot as opl
import orix.vector as ove


vz = ove.Vector3d([0, 0, 1])
v = ove.Vector3d.get_circle(vz, steps=100)
sizes = np.array([10, 20, 30, 40])

v_ragged = RaggedVector3d(v, sizes)
print(v_ragged)
# RaggedVector3d(size=4, sizes=[10, 20, 30, 40])

# Plot sets of vectors
fig = plt.figure()
ax = fig.add_subplot(111, projection="stereographic")
for i, v_i in enumerate(v_ragged):
    ax.scatter(v_i, c=f"C{i}", label=i)
_ = fig.legend()

I'm OK with orix providing such a minimal class. However, I'm more reluctant to the class providing more methods like OffsetVector3d.mean() etc. I think downstream packages should iterate and do the calculation themselves.

What do you think?