matplotlib-multicolored-line 0.2.0__tar.gz → 0.3.0__tar.gz

{matplotlib_multicolored_line-0.2.0 → matplotlib_multicolored_line-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: matplotlib-multicolored-line
-Version: 0.2.0
+Version: 0.3.0
 Summary: Plot multicolored lines in Matplotlib
 Author-email: 34j <34j.95a2p@simplelogin.com>
 License-Expression: CC0-1.0

{matplotlib_multicolored_line-0.2.0 → matplotlib_multicolored_line-0.3.0}/pyproject.toml

@@ -4,7 +4,7 @@ requires = [ "setuptools" ]
 
 [project]
 name = "matplotlib-multicolored-line"
-version = "0.2.0"
+version = "0.3.0"
 description = "Plot multicolored lines in Matplotlib"
 readme = "README.md"
 license = "CC0-1.0"

{matplotlib_multicolored_line-0.2.0 → matplotlib_multicolored_line-0.3.0}/src/matplotlib_multicolored_line/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.2.0"
+__version__ = "0.3.0"
 from ._main import colored_line
 
 __all__ = ["colored_line"]

{matplotlib_multicolored_line-0.2.0 → matplotlib_multicolored_line-0.3.0}/src/matplotlib_multicolored_line/_main.py

@@ -31,8 +31,13 @@ def colored_line(
     *args : array-like
         The horizontal and vertical coordinates of the data points
         of shape (N,) or (N, m).
-    c : array-like
-        The color values, which should be the same size as x and y.
+    c : array-like or array-like of color, optional
+        The line values. Possible values:
+        - array-like of numbers of shape (N,) or (N, m) to be mapped
+            to colors using cmap and norm.
+        - array-like of numbers of shape (N, m, 3) or (N, m, 4),
+            where the last dimension is RGB or RGBA.
+        - array-like of colors of shape (N,) or (N, m).
     ax : matplotlib.axes.Axes, optional
         The axes to plot on. If not provided, the current axes will be used.
     start : array-like, optional
@@ -52,12 +57,13 @@ def colored_line(
         The generated line collection representing the colored line.
 
     """
-    if "array" in kwargs:
-        warnings.warn(
-            'The provided "array" keyword argument will be overridden',
-            UserWarning,
-            stacklevel=2,
-        )
+    for key in ["array", "colors"]:
+        if key in kwargs:
+            warnings.warn(
+                f'The provided "{key}" keyword argument will be overridden',
+                UserWarning,
+                stacklevel=2,
+            )
 
     # parse args
     if len(args) == 1:
@@ -72,10 +78,33 @@ def colored_line(
     x, y, c = map(np.asarray, (x, y, c))
     if x.ndim == 1:
         x = x[:, None]
+    elif x.ndim == 2:
+        pass
+    else:
+        raise ValueError(f"{x.ndim=} should be 1 or 2")
     if y.ndim == 1:
         y = y[:, None]
+    elif y.ndim == 2:
+        pass
+    else:
+        raise ValueError(f"{y.ndim=} should be 1 or 2")
     x, y = np.broadcast_arrays(x, y)
 
+    if c.ndim == 1:
+        c = c[:, None]
+    if c.ndim == 2:
+        c = np.broadcast_to(c, x.shape).reshape((-1,))
+        if np.issubdtype(c.dtype, np.number):
+            kwargs["array"] = c
+        else:
+            kwargs["colors"] = c
+    elif c.ndim == 3:
+        c = np.broadcast_to(c, (*x.shape, c.shape[-1])).reshape((-1, c.shape[-1]))
+        c = [tuple(c[i]) for i in range(c.shape[0])]
+        kwargs["colors"] = c
+    else:
+        raise ValueError(f"{c.ndim=} should be 2 or 3")
+
     # (N, m, 2)
     xy = np.stack((x, y), axis=-1)
     # (N, m, 2)
@@ -88,7 +117,6 @@ def colored_line(
     segments = np.stack((xy_start, xy, xy_end), axis=-2)
     segments = np.reshape(segments, (-1, segments.shape[-2], segments.shape[-1]))
 
-    kwargs["array"] = c.reshape(-1)
     lc = LineCollection(segments, **kwargs)
 
     # Plot the line collection to the axes

{matplotlib_multicolored_line-0.2.0 → matplotlib_multicolored_line-0.3.0}/src/matplotlib_multicolored_line.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: matplotlib-multicolored-line
-Version: 0.2.0
+Version: 0.3.0
 Summary: Plot multicolored lines in Matplotlib
 Author-email: 34j <34j.95a2p@simplelogin.com>
 License-Expression: CC0-1.0

matplotlib_multicolored_line-0.3.0/tests/test_main.py

@@ -0,0 +1,114 @@
+from collections.abc import Sequence
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pytest
+from numpy.typing import ArrayLike
+
+from matplotlib_multicolored_line import colored_line
+
+CACHE_PATH = Path(__file__).parent / ".cache"
+
+
+@pytest.fixture(autouse=True)
+def setup_cache() -> None:
+    CACHE_PATH.mkdir(parents=True, exist_ok=True)
+
+
+t = np.linspace(-7.4, -0.5, 200)
+
+
+@pytest.mark.parametrize(
+    "name, args, c",
+    [
+        ("notebook1", (0.9 * np.sin(t), 0.5 * np.sin(t)), t),
+        ("notebook2", ([0, 1, 2, 3, 4], [0, 1, 2, 1, 1]), [1, 2, 3, 4, 5]),
+        ("singley", (np.random.normal(size=(5,)),), np.random.normal(size=(5,))),
+        ("singley.rgb", (np.random.normal(size=(5,)),), np.random.normal(size=(5, 1, 3))),
+        ("singley.rgba", (np.random.normal(size=(5,)),), np.random.normal(size=(5, 1, 4))),
+        (
+            "singley.colors",
+            (np.random.normal(size=(5,)),),
+            np.random.choice(["black", "red"], size=(5,)),
+        ),
+        (
+            "singlexy",
+            (
+                np.random.normal(size=(5,)),
+                np.random.normal(size=(5,)),
+            ),
+            np.random.normal(size=(5,)),
+        ),
+        (
+            "singlexy.rgb",
+            (
+                np.random.normal(size=(5,)),
+                np.random.normal(size=(5,)),
+            ),
+            np.random.normal(size=(5, 1, 3)),
+        ),
+        (
+            "singlexy.rgba",
+            (
+                np.random.normal(size=(5,)),
+                np.random.normal(size=(5,)),
+            ),
+            np.random.normal(size=(5, 1, 4)),
+        ),
+        (
+            "singlexy.colors",
+            (
+                np.random.normal(size=(5,)),
+                np.random.normal(size=(5,)),
+            ),
+            np.random.choice(["black", "red"], size=(5,)),
+        ),
+        ("multiy", (np.random.normal(size=(5, 2)),), np.random.normal(size=(5, 2))),
+        ("multiy.rgb", (np.random.normal(size=(5, 2)),), np.random.uniform(size=(5, 2, 3))),
+        ("multiy.rgba", (np.random.normal(size=(5, 2)),), np.random.uniform(size=(5, 2, 4))),
+        (
+            "multiy.colors",
+            (np.random.normal(size=(5, 2)),),
+            np.random.choice(["black", "red"], size=(5, 2)),
+        ),
+        (
+            "multixy",
+            (
+                np.random.normal(size=(5, 2)),
+                np.random.normal(size=(5, 2)),
+            ),
+            np.random.normal(size=(5, 2)),
+        ),
+        (
+            "multixy.rgb",
+            (
+                np.random.normal(size=(5, 2)),
+                np.random.normal(size=(5, 2)),
+            ),
+            np.random.uniform(size=(5, 2, 3)),
+        ),
+        (
+            "multixy.rgba",
+            (
+                np.random.normal(size=(5, 2)),
+                np.random.normal(size=(5, 2)),
+            ),
+            np.random.uniform(size=(5, 2, 4)),
+        ),
+        (
+            "multixy.colors",
+            (
+                np.random.normal(size=(5, 2)),
+                np.random.normal(size=(5, 2)),
+            ),
+            np.random.choice(["black", "red"], size=(5, 2)),
+        ),
+    ],
+)
+def test_multicolor(name: str, args: Sequence[ArrayLike], c: ArrayLike) -> None:
+    fig, ax = plt.subplots()
+    lc = colored_line(*args, c=c, ax=ax, linewidth=10)
+    if "rgb" not in name:
+        fig.colorbar(lc)
+    fig.savefig(CACHE_PATH / f"example.{name}.jpg")

matplotlib_multicolored_line-0.2.0/tests/test_main.py

@@ -1,47 +0,0 @@
-from pathlib import Path
-
-import matplotlib.pyplot as plt
-import numpy as np
-import pytest
-
-from matplotlib_multicolored_line import colored_line
-
-CACHE_PATH = Path(__file__).parent / ".cache"
-
-
-@pytest.fixture(autouse=True)
-def setup_cache() -> None:
-    CACHE_PATH.mkdir(parents=True, exist_ok=True)
-
-
-def test_main() -> None:
-    t = np.linspace(-7.4, -0.5, 200)
-    x = 0.9 * np.sin(t)
-    y = 0.9 * np.cos(1.6 * t)
-
-    fig, ax = plt.subplots()
-    lc = colored_line(x, y, c=t, ax=ax, linewidth=10)
-    fig.colorbar(lc)
-    fig.savefig(CACHE_PATH / "example.jpg")
-
-
-def test_main_small() -> None:
-    x = [0, 1, 2, 3, 4]
-    y = [0, 1, 2, 1, 1]
-    c = [1, 2, 3, 4, 5]
-
-    fig, ax = plt.subplots()
-    lc = colored_line(x, y, c=c, ax=ax, linewidth=10)
-    fig.colorbar(lc)
-    fig.savefig(CACHE_PATH / "example_small.jpg")
-
-
-def test_multiple() -> None:
-    np.random.seed(0)
-    y = np.random.normal(size=(5, 2))
-    c = np.random.normal(size=(5, 2))
-
-    fig, ax = plt.subplots()
-    lc = colored_line(y, c=c, ax=ax, linewidth=10)
-    fig.colorbar(lc)
-    fig.savefig(CACHE_PATH / "example_multiple.jpg")