Merge pull request #112 from adigitoleo/tools

feat: Add command line polefigure plotting tool.

pyproject.toml

@@ -51,7 +51,7 @@ repository = "https://github.com/Patol75/PyDRex/"
 documentation = "https://patol75.github.io/PyDRex/"
 
 [project.scripts]
-pydrex = "pydrex.run:main"
+pydrex-polefigures = "pydrex.cli:CLI_HANDLERS.pole_figure_visualiser"
 
 # Make setuptools include datafiles in wheels/packages.
 # Data files must be inside the package directory.

src/pydrex/cli.py

@@ -0,0 +1,133 @@
+"""PyDRex: Entry points for command line tools."""
+import argparse
+import os
+from collections import namedtuple
+from dataclasses import dataclass
+
+import numpy as np
+
+from pydrex import exceptions as _err
+from pydrex import logger as _log
+from pydrex import minerals as _minerals
+from pydrex import stats as _stats
+from pydrex import visualisation as _vis
+
+# NOTE: Register all cli handlers in the namedtuple at the end of the file.
+
+
+@dataclass
+class PoleFigureVisualiser:
+    """PyDRex script to plot pole figures of serialized CPO data.
+
+    Produces [100], [010] and [001] pole figures for serialized `pydrex.Mineral`s.
+    If the range of indices is not specified,
+    a maximum of 25 of each pole figure will be produced.
+
+    """
+
+    def __call__(self):
+        try:
+            args = self._get_args()
+            if args.range is None:
+                i_range = None
+            else:
+                i_range = range(*(int(s) for s in args.range.split(":")))
+
+            density_kwargs = {"kernel": args.kernel}
+            if args.smoothing is not None:
+                density_kwargs["σ"] = args.smoothing
+
+            mineral = _minerals.Mineral.from_file(args.input, postfix=args.postfix)
+            if i_range is None:
+                i_range = range(0, len(mineral.orientations))
+                if len(i_range) > 25:
+                    _log.warning(
+                        "truncating to 25 timesteps (out of %s total)", len(i_range)
+                    )
+                    i_range = range(0, 25)
+
+            orientations_resampled = [
+                _stats.resample_orientations(
+                    mineral.orientations[i], mineral.fractions[i]
+                )[0]
+                for i in np.arange(i_range.start, i_range.stop, i_range.step, dtype=int)
+            ]
+            _vis.polefigures(
+                orientations_resampled,
+                ref_axes=args.ref_axes,
+                i_range=i_range,
+                density=args.density,
+                savefile=args.out,
+                **density_kwargs,
+            )
+        except (argparse.ArgumentError, ValueError, _err.Error) as e:
+            _log.error(str(e))
+
+    def _get_args(self) -> argparse.Namespace:
+        description, epilog = self.__doc__.split(os.linesep + os.linesep, 1)
+        parser = argparse.ArgumentParser(description=description, epilog=epilog)
+        parser.add_argument("input", help="input file (.npz)")
+        parser.add_argument(
+            "-r",
+            "--range",
+            help="range of strain indices to be plotted, in the format start:stop:step",
+            default=None,
+        )
+        parser.add_argument(
+            "-p",
+            "--postfix",
+            help=(
+                "postfix of the mineral to load,"
+                + " required if the input file contains data for multiple minerals"
+            ),
+            default=None,
+        )
+        parser.add_argument(
+            "-d",
+            "--density",
+            help="toggle contouring of pole figures using point density estimation",
+            default=False,
+            action="store_true",
+        )
+        parser.add_argument(
+            "-k",
+            "--kernel",
+            help=(
+                "kernel function for point density estimation, one of:"
+                + f" {list(_stats.SPHERICAL_COUNTING_KERNELS.keys())}"
+            ),
+            default="linear_inverse_kamb",
+        )
+        parser.add_argument(
+            "-s",
+            "--smoothing",
+            help="smoothing parameter for Kamb type density estimation kernels",
+            default=None,
+            type=float,
+            metavar="σ",
+        )
+        parser.add_argument(
+            "-a",
+            "--ref-axes",
+            help=(
+                "two letters from {'x', 'y', 'z'} that specify"
+                + " the horizontal and vertical axes of the pole figures"
+            ),
+            default="xz",
+        )
+        parser.add_argument(
+            "-o",
+            "--out",
+            help="name of the output file, with either .png or .pdf extension",
+            default="polefigures.png",
+        )
+        return parser.parse_args()
+
+
+_CLI_HANDLERS = namedtuple(
+    "CLI_HANDLERS",
+    {
+        "pole_figure_visualiser",
+    },
+)
+CLI_HANDLERS = _CLI_HANDLERS(pole_figure_visualiser=PoleFigureVisualiser())

src/pydrex/visualisation.py

@@ -9,8 +9,6 @@
 from pydrex import axes as _axes
 from pydrex import io as _io
 from pydrex import logger as _log
-from pydrex import minerals as _minerals
-from pydrex import stats as _stats
 
 # Always show XY grid by default.
 plt.rcParams["axes.grid"] = True
@@ -25,43 +23,16 @@
 
 
 def polefigures(
-    datafile,
-    i_range=None,
-    postfix=None,
-    density=False,
-    ref_axes="xz",
-    savefile="polefigures.png",
-    **kwargs,
+    orientations, ref_axes, i_range, density=False, savefile="polefigures.png", **kwargs
 ):
-    """Plot [100], [010] and [001] pole figures for CPO data.
+    """Plot pole figures of a series of (Nx3x3) orientation matrix stacks.
 
-    The data is read from fields ending with the optional `postfix` in the NPZ file
-    `datafile`. Use `i_range` to specify the indices of the timesteps to be plotted,
-    which can be any valid Python range object, e.g. `range(0, 12, 2)` for a step of 2.
-    By default (`i_range=None`), a maximum of 25 timesteps are plotted.
-    If the number would exceed this, a warning is printed,
-    which signals the complete number of timesteps found in the file.
-
-    Use `density=True` to plot contoured pole figures instead of raw points.
-    In this case, any additional keyword arguments are passed to
-    `pydrex.stats.point_density`.
-
-    See also: `pydrex.minerals.Mineral.save`, `pydrex.axes.PoleFigureAxes.polefigure`.
+    Produces [100], [010] and [001] pole figures for (resampled) orientations.
+    For the argument specification, check the output of `pydrex-polefigures --help`
+    on the command line.
 
     """
-    mineral = _minerals.Mineral.from_file(datafile, postfix=postfix)
-    if i_range is None:
-        i_range = range(0, len(mineral.orientations))
-        if len(i_range) > 25:
-            _log.warning("truncating to 25 timesteps (out of %s total)", len(i_range))
-            i_range = range(0, 25)
-
-    orientations_resampled = [
-        _stats.resample_orientations(mineral.orientations[i], mineral.fractions[i])[0]
-        for i in np.arange(i_range.start, i_range.stop, i_range.step, dtype=int)
-    ]
-    n_orientations = len(orientations_resampled)
-
+    n_orientations = len(orientations)
     fig = plt.figure(figsize=(n_orientations, 4), dpi=600)
 
     if len(i_range) == 1:
@@ -96,7 +67,7 @@ def polefigures(
         grid[first_row + 2, :], edgecolor=plt.rcParams["grid.color"], linewidth=1
     )
     fig001.suptitle("[001]", fontsize="small")
-    for n, orientations in enumerate(orientations_resampled):
+    for n, orientations in enumerate(orientations):
         ax100 = fig100.add_subplot(
             1, n_orientations, n + 1, projection="pydrex.polefigure"
         )
@@ -132,7 +103,6 @@ def polefigures(
                 cbar = fig.colorbar(
                     pf,
                     ax=ax,
-                    ref_axes=ref_axes,
                     fraction=0.05,
                     location="bottom",
                     orientation="horizontal",