"""Factory functions for vector and vector field objects."""
from typing import Any, Optional, Tuple, Union
from typing import Dict as TypingDict
from typing import List as TypingList
import numpy as np
from ..objects import VectorField, Vectors
from ..transform import process_transform_arguments
from .common import _default_color
# Type aliases for better readability
ArrayLike = Union[TypingList, np.ndarray, Tuple]
[docs]
def vector_field(
vectors: ArrayLike,
colors: TypingList[int] = None,
origin_color: Optional[int] = None,
head_color: Optional[int] = None,
color: int = _default_color,
use_head: bool = True,
head_size: float = 1.0,
scale: float = 1.0,
line_width: float = 0.01,
name: Optional[str] = None,
group: Optional[str] = None,
custom_data: Optional[TypingDict[str, Any]] = None,
compression_level: int = 0,
**kwargs: Any,
) -> VectorField:
if colors is None:
colors = []
return process_transform_arguments(
VectorField(
vectors=vectors,
colors=colors,
use_head=use_head,
head_size=head_size,
line_width=line_width,
head_color=head_color if head_color is not None else color,
origin_color=origin_color if origin_color is not None else color,
scale=scale,
name=name,
group=group,
custom_data=custom_data,
compression_level=compression_level,
),
**kwargs,
)
[docs]
def vectors(
origins: ArrayLike,
vectors: Optional[ArrayLike] = None,
colors: TypingList[int] = None,
origin_color: Optional[int] = None,
head_color: Optional[int] = None,
color: int = _default_color,
use_head: bool = True,
head_size: float = 1.0,
labels: TypingList[str] = None,
label_size: float = 1.0,
line_width: float = 0.01,
name: Optional[str] = None,
group: Optional[str] = None,
custom_data: Optional[TypingDict[str, Any]] = None,
compression_level: int = 0,
**kwargs: Any,
) -> Vectors:
if colors is None:
colors = []
if labels is None:
labels = []
return process_transform_arguments(
Vectors(
vectors=vectors if vectors is not None else origins,
origins=origins
if vectors is not None
else np.zeros_like(origins, dtype=np.float32),
colors=colors,
origin_color=origin_color if origin_color is not None else color,
head_color=head_color if head_color is not None else color,
use_head=use_head,
head_size=head_size,
labels=labels,
label_size=label_size,
line_width=line_width,
name=name,
group=group,
custom_data=custom_data,
compression_level=compression_level,
),
**kwargs,
)