Add color option to plot_points (#357)

* Add color option to `plot_points`

* Make plot_points more accessible in api

* Fix formatting

Author: stefsmeets

src/gemdat/plots/__init__.py

@@ -21,6 +21,7 @@
     jumps_vs_time,
     msd_per_element,
     plot_3d,
+    plot_3d_points,
     radial_distribution,
     rectilinear,
     shape,
@@ -44,6 +45,7 @@
     'msd_per_element',
     'msd_per_element',
     'plot_3d',
+    'plot_3d_points',
     'polar',
     'radial_distribution',
     'rectilinear',

src/gemdat/plots/plotly/__init__.py

@@ -15,7 +15,7 @@
 from ._jumps_vs_distance import jumps_vs_distance
 from ._jumps_vs_time import jumps_vs_time
 from ._msd_per_element import msd_per_element
-from ._plot3d import plot_3d
+from ._plot3d import plot_3d, plot_3d_points
 from ._polar import polar
 from ._radial_distribution import radial_distribution
 from ._rectilinear import rectilinear
@@ -37,6 +37,7 @@
     'jumps_vs_time',
     'msd_per_element',
     'plot_3d',
+    'plot_3d_points',
     'polar',
     'radial_distribution',
     'rectilinear',

src/gemdat/plots/plotly/_plot3d.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from collections.abc import Collection
-from typing import TYPE_CHECKING, Sequence
+from typing import TYPE_CHECKING, Optional, Sequence
 
 import numpy as np
 import plotly.express as px
@@ -65,7 +65,14 @@ def plot_lattice_vectors(lattice: Lattice, *, fig: go.Figure):
         )
 
 
-def plot_points(points: np.ndarray, labels: Sequence, *, fig: go.Figure, point_size: int = 5):
+def plot_3d_points(
+    points: np.ndarray,
+    labels: Sequence,
+    *,
+    fig: go.Figure,
+    point_size: int = 5,
+    colors: Optional[dict[str, str]] = None,
+):
     """Plot points using plotly.
 
     Parameters
@@ -78,13 +85,18 @@ def plot_points(points: np.ndarray, labels: Sequence, *, fig: go.Figure, point_s
         Plotly figure to add traces to
     point_size : int, optional
         Size of the points
+    colors: dict, optional
+        Mapping of colors for the each label.
+        See the following link for a list of accepted colours:
+        https://developer.mozilla.org/en-US/docs/Web/CSS/named-color
     """
     assert len(points) == len(labels)
 
-    colors = {
-        label: px.colors.sample_colorscale('rainbow', [i / (len(labels) - 1)])
-        for i, label in enumerate(labels)
-    }
+    if not colors:
+        colors = {
+            label: px.colors.sample_colorscale('rainbow', [i / (len(labels) - 1)])
+            for i, label in enumerate(labels)
+        }
 
     for i, (x, y, z) in enumerate(points):
         label = labels[i]
@@ -120,7 +132,7 @@ def plot_structure(structure: Structure, *, lattice: Lattice | None = None, fig:
     else:
         cart_coords = structure.cart_coords
 
-    plot_points(cart_coords, labels=structure.labels, fig=fig)
+    plot_3d_points(cart_coords, labels=structure.labels, fig=fig)
     plot_lattice_vectors(structure.lattice, fig=fig)