Vector Field Point Cloud Export

src/vectorose/plotting.py

@@ -2924,3 +2924,120 @@ def produce_3d_confidence_cone_plot(
     ax.set_aspect("equal")
 
     return ax
+
+
+def get_orientation_rgb(vectors: np.ndarray) -> np.ndarray:
+    """Get RGB colours corresponding to orientations.
+
+    Compute the RGB colour values reflecting the ``X,Y,Z`` components of
+    vector orientation, respectively. See **Notes** for additional details.
+
+    Parameters
+    ----------
+    vectors
+        Array of shape ``(n, 3)`` containing ``n`` **unit vectors** in 3D.
+
+    Returns
+    -------
+    numpy.ndarray
+        Array of shape ``(n, 3)`` containing the RGB colour values
+        corresponding to the orientations provided in `vectors`. These
+        values are **floating-point** representations, with each entry in
+        the range ``[0, 1]``.
+
+    Warnings
+    --------
+    The `vectors` must be normalised to unit length.
+
+    Colour is not unique! Vectors with very different directions, and even
+    different orientations may be assigned the same colour.
+
+    Notes
+    -----
+    The colour reflects the orientation of the provided vectors. The red,
+    green and blue channels reflect the strength of the vector in the
+    ``X``, ``Y`` and ``Z`` directions, respectively. The colour is **not**
+    a direct mapping of the vector components to these colour channels. Nor
+    is it a direct mapping of the absolute values of the components.
+    Otherwise, we would get unexpected results. For example, if a vector is
+    oriented in the ``XY`` plane at a 45\u00b0 angle, we would expect the
+    vector to appear as yellow. However, it would actually appear as a
+    darker mustard. To achieve the desired vibrant colours, we normalise
+    the absolute values of the components with respect to the largest
+    component.
+
+    This orientation representation is inspired by the visualisations in
+    the work by Reznikov et al. [#reznikov2022]_
+
+    References
+    ----------
+    .. [#reznikov2022] Reznikov, N., Liang, H., McKee, M. D., & Piché, N.
+        (2022). Technical note: Mapping of trabecular bone anisotropy and
+        volume fraction in 3D using μCT images of the human calcaneus.
+        *American Journal of Biological Anthropology*, *177* (3), 566–580.
+        https://doi.org/10.1002/ajpa.24474
+
+    """
+
+    abs_vectors = np.abs(vectors)
+    vector_max = abs_vectors.max(axis=-1)
+
+    colours = abs_vectors / vector_max[:, None]
+
+    return colours
+
+
+def create_pointcloud(
+        vectors: np.ndarray,
+        magnitude_key: str = "magnitude",
+        direction_key: str = "orientation",
+        colour_key: str = "colour",
+) -> pv.PolyData:
+    """Convert a vector field to a point cloud.
+
+    Using the provided vectors, produce a point cloud that can be
+    visualised in third-party software.
+
+    Parameters
+    ----------
+    vectors
+        Array of shape ``(n, 6)`` containing the vectors from which to
+        construct a point cloud. The first three columns are considered the
+        spatial locations, ordered ``XYZ``, and the last three columns are
+        considered the vector components, ordered ``XYZ``.
+    magnitude_key
+        Name to give the scalar slot containing the vector magnitudes.
+    direction_key
+        Name to give the vector slot containing the vector directions.
+    colour_key
+        Name to give the vector slot containing the RGB representation of the
+        direction.
+
+    Returns
+    -------
+    pyvista.PolyData
+        Point cloud containing ``n`` points, with point data containing the
+        vector magnitudes and unit direction vectors in the respective keys
+        passed to `magnitude_key` and `direction_key`. The RGB values
+        reflecting the direction are stored in point data with `colour_key`.
+
+    Warnings
+    --------
+    Zero-vectors should be excluded **before** calling this function.
+
+    """
+
+    locations = vectors[:, :3]
+    components = vectors[:, 3:]
+
+    pc = pv.PolyData(locations)
+
+    directions, magnitudes = util.normalise_vectors(components)
+
+    colours = get_orientation_rgb(directions)
+
+    pc.point_data[magnitude_key] = magnitudes
+    pc.point_data[direction_key] = directions
+    pc.point_data[colour_key] = colours
+
+    return pc

tests/test_plotting.py

@@ -1568,3 +1568,111 @@ def test_produce_3d_confidence_cone_plot(tmp_path):
 
     # Check that everything is plotted properly (n patches + 1 for sphere)
     assert len(ax.collections) == len(confidence_cone_patches) + 1
+
+
+def test_orientation_rgb():
+    """Test for getting the RGB values from vector orientation."""
+
+    vectors = np.array(
+        [
+            [1, 0, 0],
+            [0, 1, 0],
+            [0, 0, 1],
+            [1 / np.sqrt(2), 1 / np.sqrt(2), 0],
+            [1 / np.sqrt(2), -1 / np.sqrt(2), 0],
+            [-1, 0, 0],
+            [0, -1, 0],
+            [1 / np.sqrt(3), -1 / np.sqrt(3), -1 / np.sqrt(3)],
+        ]
+    )
+
+    expected_colours = np.array(
+        [
+            [1, 0, 0],
+            [0, 1, 0],
+            [0, 0, 1],
+            [1, 1, 0],
+            [1, 1, 0],
+            [1, 0, 0],
+            [0, 1, 0],
+            [1, 1, 1],
+        ]
+    )
+
+    colours = vr.plotting.get_orientation_rgb(vectors)
+
+    # Check the array shape
+    assert colours.shape == vectors.shape
+
+    # Check that the maximum in each row is 1
+    max_row_entries = colours.max(axis=-1)
+
+    assert np.all(np.isclose(max_row_entries, 1))
+
+    # Check the contents
+    assert np.all(np.isclose(colours, expected_colours))
+
+
+def test_create_pointcloud_no_keys(vectors):
+    """Test for creating point clouds."""
+
+    locations = np.random.default_rng(RANDOM_SEED).uniform(
+        low=0,
+        high=300,
+        size=vectors.shape,
+    )
+
+    vector_array = np.concatenate([locations, vectors], axis=-1)
+
+    expected_points = len(vector_array)
+
+    directions, magnitudes = vr.util.normalise_vectors(vector_array)
+
+    pc = vr.plotting.create_pointcloud(vector_array)
+
+    assert pc.n_points == expected_points
+
+    assert "magnitude" in pc.point_data.keys()
+    assert "orientation" in pc.point_data.keys()
+    assert "colour" in pc.point_data.keys()
+
+    # Check the magnitudes
+    assert np.all(np.isclose(pc.point_data["magnitude"], magnitudes))
+    assert np.all(np.isclose(pc.point_data["orientation"], directions[:, 3:]))
+
+
+def test_create_pointcloud_custom_keys(vectors):
+    """Test for creating point clouds with custom keys."""
+
+    locations = np.random.default_rng(RANDOM_SEED).uniform(
+        low=0,
+        high=300,
+        size=vectors.shape,
+    )
+
+    vector_array = np.concatenate([locations, vectors], axis=-1)
+
+    expected_points = len(vector_array)
+
+    magnitude_key = "length"
+    orientation_key = "dir"
+    colour_key = "col"
+
+    directions, magnitudes = vr.util.normalise_vectors(vector_array)
+
+    pc = vr.plotting.create_pointcloud(
+        vector_array,
+        magnitude_key=magnitude_key,
+        direction_key=orientation_key,
+        colour_key=colour_key,
+    )
+
+    assert pc.n_points == expected_points
+
+    assert magnitude_key in pc.point_data.keys()
+    assert orientation_key in pc.point_data.keys()
+    assert colour_key in pc.point_data.keys()
+
+    # Check the magnitudes
+    assert np.all(np.isclose(pc.point_data[magnitude_key], magnitudes))
+    assert np.all(np.isclose(pc.point_data[orientation_key], directions[:, 3:]))