pydrex.cli

  1"""> PyDRex: Entry points and argument handling for command line tools.
  2
  3All CLI handlers should be registered in the `CLI_HANDLERS` namedtuple,
  4which ensures that they will be installed as executable scripts alongside the package.
  5
  6"""
  7
  8import argparse
  9import os
 10from collections import namedtuple
 11from zipfile import ZipFile
 12
 13from pydrex import exceptions as _err
 14from pydrex import io as _io
 15from pydrex import logger as _log
 16from pydrex import minerals as _minerals
 17from pydrex import stats as _stats
 18from pydrex import visualisation as _vis
 19
 20
 21class CliTool:
 22    """Base class for CLI tools defining the required interface."""
 23
 24    def __call__(self):
 25        return NotImplementedError
 26
 27    def _get_args(self) -> argparse.Namespace | type[NotImplementedError]:
 28        return NotImplementedError
 29
 30
 31class MeshGenerator(CliTool):
 32    """PyDRex script to generate various simple meshes.
 33
 34    Only rectangular (2D) meshes are currently supported. The RESOLUTION must be a comma
 35    delimited set of directives of the form `<LOC>:<RES>` where `<LOC>` is a location
 36    specifier, i.e. either "G" (global) or a compas direction like "N", "S", "NE", etc.,
 37    and `<RES>` is a floating point value to be set as the resolution at that location.
 38
 39    """
 40
 41    def __call__(self):
 42        try:  # This one is dangerous, especially in CI.
 43            from pydrex import mesh as _mesh
 44        except ImportError:
 45            raise _err.MissingDependencyError(
 46                "missing optional meshing dependencies."
 47                + " Have you installed the package with 'pip install pydrex[mesh]'?"
 48            )
 49
 50        args = self._get_args()
 51
 52        if args.kind == "rectangle":
 53            if args.center is None:
 54                center = (0, 0)
 55            else:
 56                center = [float(s) for s in args.center.split(",")]
 57                assert len(center) == 2
 58
 59            width, height = map(float, args.size.split(","))
 60            _loc_map = {
 61                "G": "global",
 62                "N": "north",
 63                "S": "south",
 64                "E": "east",
 65                "W": "west",
 66                "NE": "north-east",
 67                "NW": "north-west",
 68                "SE": "south-east",
 69                "SW": "south-west",
 70            }
 71            try:
 72                resolution = {
 73                    _loc_map[k]: float(v)
 74                    for k, v in map(lambda s: s.split(":"), args.resolution.split(","))
 75                }
 76            except KeyError:
 77                raise KeyError(
 78                    "invalid or unsupported location specified in resolution directive"
 79                ) from None
 80            except ValueError:
 81                raise ValueError(
 82                    "invalid resolution value. The format should be '<LOC1>:<RES1>,<LOC2>:<RES2>,...'"
 83                ) from None
 84            _mesh.rectangle(
 85                args.output[:-4],
 86                (args.ref_axes[0], args.ref_axes[1]),
 87                center,
 88                width,
 89                height,
 90                resolution,
 91            )
 92
 93    def _get_args(self) -> argparse.Namespace:
 94        assert self.__doc__ is not None, f"missing docstring for {self}"
 95        description, epilog = self.__doc__.split(os.linesep + os.linesep, 1)
 96        parser = argparse.ArgumentParser(description=description, epilog=epilog)
 97        parser.add_argument("size", help="width,height[,depth] of the mesh")
 98        parser.add_argument(
 99            "-r",
100            "--resolution",
101            help="resolution for the mesh (edge length hint(s) for gmsh)",
102            required=True,
103        )
104        parser.add_argument("output", help="output file (.msh)")
105        parser.add_argument(
106            "-c",
107            "--center",
108            help="center of the mesh as 2 or 3 comma-separated coordinates. default: (0, 0[, 0])",
109            default=None,
110        )
111        parser.add_argument(
112            "-a",
113            "--ref-axes",
114            help=(
115                "two letters from {'x', 'y', 'z'} that specify"
116                + " the horizontal and vertical axes of the mesh"
117            ),
118            default="xz",
119        )
120        parser.add_argument(
121            "-k", "--kind", help="kind of mesh, e.g. 'rectangle'", default="rectangle"
122        )
123        return parser.parse_args()
124
125
126class H5partExtractor(CliTool):
127    """PyDRex script to extract raw CPO data from Fluidity .h5part files.
128
129    Fluidity saves data stored on model `particles` to an `.h5part` file.
130    This script converts that file to canonical serialisation formats:
131    - a `.npz` file containing the raw CPO orientations and (surrogate) grain sizes
132    - an `.scsv` file containing the pathline positions and accumulated strain
133
134    It is assumed that CPO data is stored in keys called 'CPO_<N>' in the .h5part
135    data, where `<N>` is an integer in the range 1—`n_grains`. The accumulated strain is
136    read from the attribute `CPO_<S>` where S=`ngrains`+1. Particle positions are read
137    from the attributes `x`, `y`, and `z`.
138
139    At the moment, dynamic changes in fabric or phase are not supported.
140
141    """
142
143    def __call__(self):
144        args = self._get_args()
145        _io.extract_h5part(
146            args.input, args.phase, args.fabric, args.ngrains, args.output
147        )
148
149    def _get_args(self) -> argparse.Namespace:
150        assert self.__doc__ is not None, f"missing docstring for {self}"
151        description, epilog = self.__doc__.split(os.linesep + os.linesep, 1)
152        parser = argparse.ArgumentParser(description=description, epilog=epilog)
153        parser.add_argument("input", help="input file (.h5part)")
154        parser.add_argument(
155            "-p",
156            "--phase",
157            help="type of `pydrex.MineralPhase` (as an ordinal number); 0 by default",
158            default=0,
159        )
160        parser.add_argument(
161            "-f",
162            "--fabric",
163            type=int,
164            help="type of `pydrex.MineralFabric` (as an ordinal number); 0 by default",
165            default=0,
166        )
167        parser.add_argument(
168            "-n",
169            "--ngrains",
170            help="number of grains used in the Fluidity simulation",
171            type=int,
172            required=True,
173        )
174        parser.add_argument(
175            "-o",
176            "--output",
177            help="filename for the output NPZ file (stem also used for the .scsv)",
178            required=True,
179        )
180        return parser.parse_args()
181
182
183class NPZFileInspector(CliTool):
184    """PyDRex script to show information about serialized CPO data.
185
186    Lists the keys that should be used for the `postfix` in `pydrex.Mineral.load` and
187    `pydrex.Mineral.from_file`.
188
189    """
190
191    def __call__(self):
192        args = self._get_args()
193        with ZipFile(args.input) as npz:
194            names = npz.namelist()
195            print("NPZ file with keys:")
196            for name in names:
197                if not (
198                    name.startswith("meta")
199                    or name.startswith("fractions")
200                    or name.startswith("orientations")
201                ):
202                    _log.warning(f"found unknown NPZ key '{name}' in '{args.input}'")
203                print(f" - {name}")
204
205    def _get_args(self) -> argparse.Namespace:
206        assert self.__doc__ is not None, f"missing docstring for {self}"
207        description, epilog = self.__doc__.split(os.linesep + os.linesep, 1)
208        parser = argparse.ArgumentParser(description=description, epilog=epilog)
209        parser.add_argument("input", help="input file (.npz)")
210        return parser.parse_args()
211
212
213class PoleFigureVisualiser(CliTool):
214    """PyDRex script to plot pole figures of serialized CPO data.
215
216    Produces [100], [010] and [001] pole figures for serialized `pydrex.Mineral`s.
217    If the range of indices is not specified,
218    a maximum of 25 of each pole figure will be produced by default.
219
220    """
221
222    def __call__(self):
223        try:
224            args = self._get_args()
225            if args.range is None:
226                i_range = None
227            else:
228                start, stop_ex, step = (int(s) for s in args.range.split(":"))
229                # Make command line start:stop:step stop-inclusive, it's more intuitive.
230                i_range = range(start, stop_ex + step, step)
231
232            density_kwargs = {"kernel": args.kernel}
233            if args.smoothing is not None:
234                density_kwargs["σ"] = args.smoothing
235
236            mineral = _minerals.Mineral.from_file(args.input, postfix=args.postfix)
237            if i_range is None:
238                i_range = range(0, len(mineral.orientations))
239                if len(i_range) > 25:
240                    _log.warning(
241                        "truncating to 25 timesteps (out of %s total)", len(i_range)
242                    )
243                    i_range = range(0, 25)
244
245            orientations_resampled, _ = _stats.resample_orientations(
246                mineral.orientations[i_range.start : i_range.stop : i_range.step],
247                mineral.fractions[i_range.start : i_range.stop : i_range.step],
248            )
249            if args.scsv is None:
250                strains = None
251            else:
252                strains = _io.read_scsv(args.scsv).strain[
253                    i_range.start : i_range.stop : i_range.step
254                ]
255            _vis.polefigures(
256                orientations_resampled,
257                ref_axes=args.ref_axes,
258                i_range=i_range,
259                density=args.density,
260                savefile=args.out,
261                strains=strains,
262                **density_kwargs,
263            )
264        except (argparse.ArgumentError, ValueError, _err.Error) as e:
265            _log.error(str(e))
266
267    def _get_args(self) -> argparse.Namespace:
268        assert self.__doc__ is not None, f"missing docstring for {self}"
269        description, epilog = self.__doc__.split(os.linesep + os.linesep, 1)
270        parser = argparse.ArgumentParser(description=description, epilog=epilog)
271        parser.add_argument("input", help="input file (.npz)")
272        parser.add_argument(
273            "-r",
274            "--range",
275            help="range of strain indices to be plotted, in the format start:stop:step",
276            default=None,
277        )
278        parser.add_argument(
279            "-f",
280            "--scsv",
281            help=(
282                "path to SCSV file with a column named 'strain'"
283                + " that lists shear strain percentages for each strain index"
284            ),
285            default=None,
286        )
287        parser.add_argument(
288            "-p",
289            "--postfix",
290            help=(
291                "postfix of the mineral to load,"
292                + " required if the input file contains data for multiple minerals"
293            ),
294            default=None,
295        )
296        parser.add_argument(
297            "-d",
298            "--density",
299            help="toggle contouring of pole figures using point density estimation",
300            default=False,
301            action="store_true",
302        )
303        parser.add_argument(
304            "-k",
305            "--kernel",
306            help=(
307                "kernel function for point density estimation, one of:"
308                + f" {list(_stats.SPHERICAL_COUNTING_KERNELS.keys())}"
309            ),
310            default="linear_inverse_kamb",
311        )
312        parser.add_argument(
313            "-s",
314            "--smoothing",
315            help="smoothing parameter for Kamb type density estimation kernels",
316            default=None,
317            type=float,
318            metavar="σ",
319        )
320        parser.add_argument(
321            "-a",
322            "--ref-axes",
323            help=(
324                "two letters from {'x', 'y', 'z'} that specify"
325                + " the horizontal and vertical axes of the pole figures"
326            ),
327            default="xz",
328        )
329        parser.add_argument(
330            "-o",
331            "--out",
332            help="name of the output file, with either .png or .pdf extension",
333            default="polefigures.png",
334        )
335        return parser.parse_args()
336
337
338# These are not the final names of the executables (those are set in pyproject.toml).
339_CLI_HANDLERS = namedtuple(
340    "_CLI_HANDLERS",
341    (
342        "pole_figure_visualiser",
343        "npz_file_inspector",
344        "mesh_generator",
345        "h5part_extractor",
346    ),
347)
348CLI_HANDLERS = _CLI_HANDLERS(
349    pole_figure_visualiser=PoleFigureVisualiser(),
350    npz_file_inspector=NPZFileInspector(),
351    mesh_generator=MeshGenerator(),
352    h5part_extractor=H5partExtractor(),
353)