cgse-coordinates 0.17.3__py3-none-any.whl → 0.18.0__py3-none-any.whl

egse/coordinates/pyplot.py

@@ -1,190 +1,202 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
-"""
-This module defines some plotting functions for Reference Frames.
-
-@author: Pierre Royer
-"""
+""" " This module contains plotting utility functions for reference frames."""
 
+import matplotlib
 import matplotlib.pyplot as plt
-from mpl_toolkits.mplot3d import Axes3D  # needed for 3d projection
+from egse.coordinates.reference_frame import ReferenceFrame
 import numpy as np
 import egse
-from egse.coordinates.referenceFrame import ReferenceFrame
 from egse.coordinates.point import Point, Points
 
 
-def plot_reference_frame(frame, master=None, figname=None, **kwargs):
-    """Plot a Reference Frame.
+def plot_reference_frame(
+    frame: ReferenceFrame, master: ReferenceFrame | None = None, fig_name: str | None = None, **kwargs
+) -> matplotlib.axes._subplots.Axes3DSubplot:
+    """Plots a reference frame.
+
+    Use ax.set_xlim3d(min, max) to properly set the ranges of the display.
 
     Args:
-        frame  : egse.coordinates.referenceFrame.ReferenceFrame
-        master : master ReferenceFrame (optional)
-        figname: string. Name of matplotlib figure (can be pre-existing)
-        kwargs : passed to matplotlib.axes._subplots.Axes3DSubplot.quiver
+        frame (ReferenceFrame): Reference frame
+        master (ReferenceFrame | None): Optional master reference frame.
+        fig_name (str): Name of the plot figure.
+        kwargs : Keyword arguments passed to matplotlib.axes._subplots.Axes3DSubplot.quiver.
 
     Returns:
-        matplotlib.axes._subplots.Axes3DSubplot displaying the reference frame.
-
-        The three unit vectors are shown with the following colors ('RGB'):
+        matplotlib.axes._subplots.Axes3DSubplot displaying the reference frame.  The three unit vectors are shown in
+        the following colours (RGB):
             x: Red
             y: Green
             z: Blue
-
-    .. note::
-        Use ax.set_xlim3d(min,max) to properly set the ranges of the display
-
     """
+
     if master is None:
-        tmpmaster = ReferenceFrame.createMaster()
+        temp_master = ReferenceFrame.create_master()
     else:
-        tmpmaster = master.__copy__()
-
-    f0 = frame.getOrigin()
-    fx = frame.getAxis("x", name="fx")
-    fy = frame.getAxis("y", name="fy")
-    fz = frame.getAxis("z", name="fz")
-    f0m = f0.expressIn(tmpmaster)[:3]
-    fxm = fx.expressIn(tmpmaster)[:3]
-    fym = fy.expressIn(tmpmaster)[:3]
-    fzm = fz.expressIn(tmpmaster)[:3]
-    del tmpmaster
-
-    # Origin of the X,Y and Z vectors (x = the 'x' coordinates of the origin of all 3 vectors)
-    # Every vector independently (--> plot in diff. colors)
-    x, y, z = np.array([f0m[0]]), np.array([f0m[1]]), np.array([f0m[2]])
-
-    # Orientation of the X,Y and Z vectors
-    vecxx, vecyx, veczx = np.array([fxm[0] - f0m[0]]), np.array([fym[0] - f0m[0]]), np.array([fzm[0] - f0m[0]])
-    vecxy, vecyy, veczy = np.array([fxm[1] - f0m[1]]), np.array([fym[1] - f0m[1]]), np.array([fzm[1] - f0m[1]])
-    vecxz, vecyz, veczz = np.array([fxm[2] - f0m[2]]), np.array([fym[2] - f0m[2]]), np.array([fzm[2] - f0m[2]])
+        temp_master = master.__copy__()
+
+    # Origin and unit axes of the reference frame
+
+    origin = frame.get_origin()
+    unit_axis_x = frame.get_axis("x", name="fx")
+    unit_axis_y = frame.get_axis("y", name="fy")
+    unit_axis_z = frame.get_axis("z", name="fz")
+
+    # Now expressed in the master reference frame
+
+    origin_master = origin.express_in(temp_master)[:3]
+    unit_axis_x_master = unit_axis_x.express_in(temp_master)[:3]
+    unit_axis_y_master = unit_axis_y.express_in(temp_master)[:3]
+    z_axis_master = unit_axis_z.express_in(temp_master)[:3]
+
+    del temp_master
+
+    # Coordinates of the origin, expressed in the master reference frame
+    origin_x_master, origin_y_master, origin_z_master = (
+        np.array([origin_master[0]]),
+        np.array([origin_master[1]]),
+        np.array([origin_master[2]]),
+    )
+
+    # Orientation of the unit vectors of the reference frame, expressed in the master reference frame
+
+    vecxx, vecyx, veczx = (
+        np.array([unit_axis_x_master[0] - origin_master[0]]),
+        np.array([unit_axis_y_master[0] - origin_master[0]]),
+        np.array([z_axis_master[0] - origin_master[0]]),
+    )
+    vecxy, vecyy, veczy = (
+        np.array([unit_axis_x_master[1] - origin_master[1]]),
+        np.array([unit_axis_y_master[1] - origin_master[1]]),
+        np.array([z_axis_master[1] - origin_master[1]]),
+    )
+    vecxz, vecyz, veczz = (
+        np.array([unit_axis_x_master[2] - origin_master[2]]),
+        np.array([unit_axis_y_master[2] - origin_master[2]]),
+        np.array([z_axis_master[2] - origin_master[2]]),
+    )
 
     kwargs.setdefault("length", 1)
     kwargs.setdefault("normalize", True)
     # kwargs.setdefault('figsize', (10,10))
 
-    fig = plt.figure(figname, figsize=plt.figaspect(1.0))
+    fig = plt.figure(fig_name, figsize=plt.figaspect(1.0))
     ax = fig.add_subplot(projection="3d")
-    ax.quiver(x, y, z, vecxx, vecxy, vecxz, color="r", **kwargs)
-    ax.quiver(x, y, z, vecyx, vecyy, vecyz, color="g", **kwargs)
-    ax.quiver(x, y, z, veczx, veczy, veczz, color="b", **kwargs)
+    ax.quiver(origin_x_master, origin_y_master, origin_z_master, vecxx, vecxy, vecxz, color="r", **kwargs)
+    ax.quiver(origin_x_master, origin_y_master, origin_z_master, vecyx, vecyy, vecyz, color="g", **kwargs)
+    ax.quiver(origin_x_master, origin_y_master, origin_z_master, veczx, veczy, veczz, color="b", **kwargs)
     # ax.axis('equal')
 
     return ax
 
 
-def plot_points(points, master=None, figname=None, **kwargs):
-    """Plot the Points object.
+def plot_points(
+    points: Points | list[Point], master=None, fig_name: str | None = None, **kwargs
+) -> matplotlib.axes._subplots.Axes3DSubplot:
+    """Plots the given collection of points.
+
+    Use ax.set_xlim3d(min, max) to properly set the ranges of the display.
 
     Args:
-        points : either a (egse.coordinate.point.)Points object or a list of Point objects
-        master : master ReferenceFrame (optional)
-        figname: string. Name of matplotlib figure (can be pre-existing)
-        kwargs : passed to matplotlib.axes._subplots.Axes3DSubplot.scatter
+        points (Points | list[Point]): Collection of points to plot.
+        master (ReferenceFrame | None): Optional master reference frame.
+        fig_name (str): Name of the plot figure.
+        kwargs : Keyword arguments passed to matplotlib.axes._subplots.Axes3DSubplot.scatter.
 
     Returns:
-        matplotlib.axes._subplots.Axes3DSubplot displaying the reference frame.
-
-    .. note::
-        Use ax.set_xlim3d(min,max) to properly set the ranges of the display.
-
+        matplotlib.axes._subplots.Axes3DSubplot displaying the given collection of points
     """
+
     if master is None:
-        tmpmaster = ReferenceFrame.createMaster()
+        temp_master = ReferenceFrame.create_master()
     else:
-        tmpmaster = master.__copy__()
-    #
+        temp_master = master.__copy__()
+
     if isinstance(points, list):
-        allpoints = Points(points, ref=tmpmaster)
+        all_points = Points(points, reference_frame=temp_master)
     elif isinstance(points, Points) or isinstance(points, egse.coordinates.point.Points):
-        allpoints = points
+        all_points = points
     else:
         raise ValueError("If the input is a list, all items in it must be Point objects")
-    #
-    del tmpmaster
-    #
-    coordinates = allpoints.coordinates
+
+    del temp_master
+
+    coordinates = all_points.coordinates
     xs = coordinates[0, :]
     ys = coordinates[1, :]
     zs = coordinates[2, :]
-    #
+
     kwargs.setdefault("s", 50)
     kwargs.setdefault("marker", "o")
     kwargs.setdefault("color", "k")
-    #
-    fig = plt.figure(figname)
+
+    fig = plt.figure(fig_name)
     ax = fig.add_subplot(projection="3d")
     ax.scatter(xs, ys, zs, **kwargs)
 
     return ax
 
 
-def plot_vectors(points, master=None, figname=None, fromorigin=True, **kwargs):
-    """Plot the Points object.
+def plot_vectors(
+    points: Points | list[Point], master=None, fig_name: str | None = None, from_origin: bool = True, **kwargs
+) -> matplotlib.axes._subplots.Axes3DSubplot:
+    """Plots the given collection of vectors.
+
+    Use ax.set_xlim3d(min, max) to properly set the ranges of the display.
 
     Args:
-        points : either a (egse.coordinate.point.)Points object or a list of Point objects
-        master : master ReferenceFrame (optional)
-        figname: string. Name of matplotlib figure (can be pre-existing)
-        fromorigin: bool
-                    if True, all vectors are displayed starting from the origin
-                    if False, all vectors go towards the origin
-        kwargs : passed to matplotlib.axes._subplots.Axes3DSubplot.scatter
+        points (Points | list[Point]): Collection of vectors to plot.
+        master (ReferenceFrame | None): Optional master reference frame.
+        fig_name (str): Name of the plot figure.
+        kwargs : Keyword arguments passed to matplotlib.axes._subplots.Axes3DSubplot.scatter.
 
     Returns:
-        matplotlib.axes._subplots.Axes3DSubplot displaying the reference frame.
-
-    .. note::
-        Use ax.set_xlim3d(min,max) to properly set the ranges of the display.
-
+        matplotlib.axes._subplots.Axes3DSubplot displaying the given collection of points
     """
 
     if master is None:
-        tmpmaster = ReferenceFrame.createMaster()
+        temp_master = ReferenceFrame.create_master()
     else:
-        tmpmaster = master.__copy__()
-    #
+        temp_master = master.__copy__()
+
     if isinstance(points, list):
-        allpoints = Points(points, ref=tmpmaster)
+        all_points = Points(points, reference_frame=temp_master)
     elif isinstance(points, Points) or isinstance(points, egse.coordinates.point.Points):
-        allpoints = points
+        all_points = points
     else:
         raise ValueError("If the input is a list, all items in it must be Point objects")
-    #
-    del tmpmaster
-    #
 
-    # SET DEFAULTS
+    del temp_master
+
     kwargs.setdefault("color", "k")
-    #
 
-    # PREPARE VECTOR COORDINATES
-    coordinates = allpoints.coordinates
+    # Prepare vector coordinates
+
+    coordinates = all_points.coordinates
     xs = coordinates[0, :]
     ys = coordinates[1, :]
     zs = coordinates[2, :]
 
-    # Origin of the X,Y and Z vectors
-    # ==> x = the 'x' coordinates of the origin of all vectors)
-    # ==> [x,y,z] = the origin of points.ref
-    x, y, z = points.ref.getOrigin().coordinates[:3]
+    # Origin of the x, y, and z vectors
+    # -> x = The x coordinates of the origin of all vectors
+    # -> [x,y,z] = Origin of points.reference_frame
+    x, y, z = points.reference_frame.get_origin().coordinates[:3]
     x = np.ones_like(xs) * x
     y = np.ones_like(xs) * y
     z = np.ones_like(xs) * z
 
-    # PLOT
+    # Plot
 
-    fig = plt.figure(figname)
+    fig = plt.figure(fig_name)
     ax = fig.gca(projection="3d")
 
-    if fromorigin:
+    if from_origin:
         ax.quiver(x, y, z, xs - x, ys - y, zs - z, **kwargs)
 
-    elif not fromorigin:
+    elif not from_origin:
         ax.quiver(xs, ys, zs, x - xs, y - ys, z - zs, **kwargs)
 
     else:
-        print("Parameter 'fromorigin' must be True or False")
+        print("Parameter 'from_origin' must be True or False")
         print("Setting it to True by default")
         ax.quiver(x, y, z, xs - x, ys - y, zs - z, **kwargs)