Fixes and things

polyharmonics/__main__.py

@@ -10,8 +10,8 @@
     benchmark,
     legendre_command,
     plot2d_command,
-    plot3d_command,
     plotbench_command,
+    plotsh_command,
     spherical_harmonic_command,
 )
 
@@ -48,7 +48,7 @@ def main(
 app.command(name="associated-legendre")(associated_legendre_command)
 app.command(name="spherical-harmonic")(spherical_harmonic_command)
 app.command(name="plot2d")(plot2d_command)
-app.command(name="plot3d")(plot3d_command)
+app.command(name="plot3d")(plotsh_command)
 app.command(name="plotbench")(plotbench_command)
 app.add_typer(benchmark, name="benchmark")
 

polyharmonics/cli/__init__.py

@@ -5,6 +5,6 @@
 from .benchmark_cmd import app as benchmark
 from .legendre_cmd import legendre_command
 from .plot2d_cmd import plot2d_command
-from .plot3d_cmd import plot3d_command
 from .plotbench_cmd import plotbench_command
+from .plotsh_cmd import plotsh_command
 from .spherical_harmonic_cmd import spherical_harmonic_command

polyharmonics/cli/benchmark/spherical_harmonic_bench.py

@@ -58,7 +58,7 @@ def spherical_harmonic_bench_command(
             n, m = value.split(":")
             if n is None or m is None or n == "" or m == "":
                 raise typer.BadParameter(
-                    "Between each ',' must be a pair of integers separated by ':'."  # noqa: E501
+                    "nm must either be a pair of integers separated by ':' or a list of such pairs separated by commas."  # noqa: E501
                 )
             else:
                 nm_values.append((int(n), int(m)))

polyharmonics/cli/legendre_cmd.py

@@ -1,3 +1,4 @@
+from math import cos as cosine
 from time import time
 from typing import List, Optional
 
@@ -48,6 +49,13 @@ def legendre_command(
         case_sensitive=False,
         help="Display the time taken to calculate the function(s).",
     ),
+    polar: bool = typer.Option(
+        False,
+        "-p",
+        "--polar",
+        case_sensitive=False,
+        help="Calculate the polynomial(s) with polar coordinates.",
+    ),
 ) -> None:
     """Calculate and print the Legendre polynomial(s)."""
 
@@ -91,10 +99,10 @@ def legendre_command(
     ):
         if x_values:
             result: List[List[float]] = [
-                [legendre(i, x) for x in x_values] for i in n_values
+                [legendre(i, eval=x, polar=polar) for x in x_values] for i in n_values
             ]
         else:
-            result: List[Expr] = [legendre(i) for i in n_values]
+            result: List[Expr] = [legendre(i, polar=polar) for i in n_values]
 
     if display_time:
         t_end = time()
@@ -104,14 +112,17 @@ def legendre_command(
 
     for n, pol in zip(n_values, result):
         if print_latex:
-            console.print(f"[bold {color}]P_{n}(x) = {latex(pol)}[/]\n")
+            cos_theta = "cos(\\theta)" if polar else "x"
+            console.print(f"[bold {color}]P_{n}({cos_theta}) = {latex(pol)}[/]\n")  # noqa: E501
         elif x_values:
             for i, x in enumerate(x_values):
                 console.print(
-                    f"[bold {color}]P{str(n).translate(SUB)}({x}) = {pol[i]}[/]\n"
+                    f"[bold {color}]P{str(n).translate(SUB)}({cosine(x) if polar else x}) = {pol[i]}[/]\n"  # noqa: E501
                 )
         else:
-            console.print(f"[bold {color}]P{str(n).translate(SUB)}(x) = [/]")
+            console.print(
+                f"[bold {color}]P{str(n).translate(SUB)}({'cos(θ)' if polar else 'x'}) = [/]"  # noqa: E501
+            )  # noqa: E501
             console.print(f"[bold {color}] {pretty(pol)}[/]\n")
 
     raise typer.Exit()

polyharmonics/cli/plot2d_cmd.py

@@ -1,14 +1,13 @@
+from math import pi
 from typing import List, Tuple
 
+import matplotlib.pyplot as plt
+import numpy as np
 import typer
 from rich.console import Console
-from sympy import Expr, latex, pretty
 
-from polyharmonics import legendre
+from polyharmonics import associated_legendre, legendre
 
-from .colors import Color
-
-SUB = str.maketrans("0123456789", "₀₁₂₃₄₅₆₇₈₉")
 console = Console()
 
 
@@ -31,9 +30,46 @@ def plot2d_command(
         Either a pair of integers separated by ':' or a comma-separated list of such pairs.""",  # noqa: E501
         metavar="SUB:SUP",
     ),
+    title: str = typer.Option(
+        None,
+        "--title",
+        case_sensitive=False,
+        help="Title of the plot.",
+    ),
+    x_lim: str = typer.Option(
+        None,
+        "--xlim",
+        case_sensitive=False,
+        help="Limits for the x-axis as two numbers separated by a comma.",
+    ),
+    y_lim: str = typer.Option(
+        None,
+        "--ylim",
+        case_sensitive=False,
+        help="Limits for the y-axis as two numbers separated by a comma.",
+    ),
+    equal_axes: bool = typer.Option(
+        False,
+        "--eq-axes",
+        case_sensitive=False,
+        help="Set equal scaling (proportion 1:1) on both axes.",
+    ),
+    center_axes: bool = typer.Option(
+        False,
+        "--center-axes",
+        case_sensitive=False,
+        help="Fix axes at the center of the plot, passing through (0, 0).",
+    ),
+    show_grid: bool = typer.Option(
+        False, "--grid", case_sensitive=False, help="Show grid lines."
+    ),
+    polar: bool = typer.Option(
+        False, "--polar", case_sensitive=False, help="Plot in polar coordinates."
+    ),
 ) -> None:
     """Plot the 2D functions given by this package."""
 
+    n_values: List[int] = []
     try:
         n_values = sorted([int(value) for value in plt_legendre.split(",")])
         if any(i < 0 for i in n_values):
@@ -44,6 +80,8 @@ def plot2d_command(
         raise typer.BadParameter(
             "legendre must be an integer or a comma-separated list of integers."  # noqa: E501
         )
+    except AttributeError:
+        pass
 
     nm_values: List[Tuple[int, int]] = []
     try:
@@ -65,8 +103,64 @@ def plot2d_command(
         raise typer.BadParameter(
             "associated-legendre must either be a pair of integers separated by ':' or a list of such pairs separated by commas."  # noqa: E501
         )
+    except AttributeError:
+        pass
 
     if not n_values and not nm_values:
         raise typer.BadParameter("No functions to plot.")
 
+    if x_lim is not None:
+        try:
+            x_lim = [float(value) for value in x_lim.split(",")]
+            if len(x_lim) != 2:
+                raise ValueError
+        except ValueError:
+            raise typer.BadParameter("x_lim must be two numbers separated by a comma.")
+    if y_lim is not None:
+        try:
+            y_lim = [float(value) for value in y_lim.split(",")]
+            if len(y_lim) != 2:
+                raise ValueError
+        except ValueError:
+            raise typer.BadParameter("y_lim must be two numbers separated by a comma.")
+
+    x = np.linspace(0 if polar else -1, pi if polar else 1, 500)
+
+    plt.figure(figsize=(10, 6))
+
+    for n in n_values:
+        y = [legendre(n, eval=xi, polar=polar) for xi in x]
+        plt.plot(x, y, label=f"$P_{n}(x)$")
+
+    for n, m in nm_values:
+        y = [associated_legendre(n, m, eval=xi, polar=polar) for xi in x]
+        plt.plot(x, y, label=f"$P_{n}^{m}(x)$")
+
+    if title:
+        plt.title(title)
+    if x_lim:
+        plt.xlim(x_lim)
+    if y_lim:
+        plt.ylim(y_lim)
+    if equal_axes:
+        plt.gca().set_aspect("equal", adjustable="box")
+    if center_axes:
+        ax = plt.gca()
+        ax.spines["left"].set_position("zero")
+        ax.spines["left"].set_color("black")
+        ax.spines["left"].set_linewidth(0.5)
+        ax.spines["bottom"].set_position("zero")
+        ax.spines["bottom"].set_color("black")
+        ax.spines["bottom"].set_linewidth(0.5)
+        ax.spines["right"].set_color("none")
+        ax.spines["top"].set_color("none")
+        ax.xaxis.set_ticks_position("bottom")
+        ax.yaxis.set_ticks_position("left")
+        ax.xaxis.set_tick_params(direction="in", which="both")
+        ax.yaxis.set_tick_params(direction="in", which="both")
+
+    plt.legend()
+    plt.grid(show_grid)
+    plt.show()
+
     raise typer.Exit()

polyharmonics/cli/plotsh_cmd.py

@@ -0,0 +1,42 @@
+from typing import List, Tuple
+
+import matplotlib.pyplot as plt
+import numpy as np
+import typer
+from rich.console import Console
+
+from polyharmonics import spherical_harmonic
+
+console = Console()
+
+
+def plotsh_command(
+    nm: str = typer.Argument(
+        ...,
+        help="""The corresponding subscript(s) and superscript(s) of the function(s).
+        Either a pair of integers separated by ':' or a comma-separated list of such pairs.""",  # noqa: E501
+        metavar="SUB:SUP",
+    ),
+) -> None:
+    """Plot the spherical harmonics given by this package."""
+
+    try:
+        n_values: List[int] = []
+        m_values: List[int] = []
+        for value in nm.split(","):
+            n, m = value.split(":")
+            if n is None or m is None or n == "" or m == "":
+                raise typer.BadParameter(
+                    "nm must either be a pair of integers separated by ':' or a list of such pairs separated by commas."  # noqa: E501
+                )
+            else:
+                n_values.append(int(n))
+                m_values.append(int(m))
+        if any(i < 0 for i in n_values):
+            raise typer.BadParameter("All subscripts must be greater or equal to 0.")
+    except ValueError:
+        raise typer.BadParameter(
+            "nm must either be a pair of integers separated by ':' or a list of such pairs separated by commas."  # noqa: E501
+        )
+
+    raise typer.Exit()

polyharmonics/cli/spherical_harmonic_cmd.py

@@ -65,7 +65,7 @@ def spherical_harmonic_command(
             n, m = value.split(":")
             if n is None or m is None or n == "" or m == "":
                 raise typer.BadParameter(
-                    "Between each ',' must be a pair of integers separated by ':'."
+                    "nm must either be a pair of integers separated by ':' or a list of such pairs separated by commas."  # noqa: E501
                 )
             else:
                 n_values.append(int(n))

polyharmonics/legendre_polynomials.py

@@ -1,8 +1,9 @@
+from math import cos as cosine
 from typing import List
 
-from sympy import Expr, Rational, diff, expand, factorial, floor, symbols
+from sympy import Expr, Rational, cos, diff, expand, factorial, floor, symbols
 
-x, t, k = symbols("x t k")
+x, t, k, th = symbols("x t k θ")
 gen_Legendre: Expr = (1 - 2 * x * t + t**2) ** Rational(-1, 2)
 
 
@@ -132,13 +133,14 @@ def legendre_exp(n: int, eval: float | None = None) -> Expr:
     return pol
 
 
-def legendre(n: int, eval: float | None = None) -> Expr:
+def legendre(n: int, polar: bool = False, eval: float | None = None) -> Expr:
     """
     Calculate the analytical expression of the Legendre polynomial.
 
     Args:
         n (int): The degree of the Legendre polynomial.
             Must be an integer greater than or equal to 0.
+        polar (bool): If True, the function is returned in polar coordinates.
         eval (float, None): The value at which the polynomial is evaluated.
             If None, the polynomial is returned as an expression.
             Default is None.
@@ -169,4 +171,10 @@ def legendre(n: int, eval: float | None = None) -> Expr:
             return 1
         elif eval == -1:
             return 1 if n % 2 == 0 else -1
-    return legendre_rec(n, eval=eval, store=False) if eval is not None else legendre_exp(n)
+    return (
+        legendre_rec(n, eval=cosine(eval) if polar else eval, store=False)
+        if eval is not None
+        else legendre_exp(n).subs(x, cos(th))
+        if polar
+        else legendre_exp(n)
+    )

pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "polyharmonics"
-version = "0.6.0"
+version = "0.7.0"
 description = "Ortogonal polynomials in the unit sphere"
 readme = "README.md"
 authors = [