imas-python 0.1.dev1618__py3-none-any.whl

imas/backends/netcdf/nc_validate.py

@@ -0,0 +1,53 @@
+from imas.backends.netcdf.db_entry_nc import NCDBEntryImpl
+from imas.backends.netcdf.nc2ids import NC2IDS
+from imas.db_entry import DBEntry
+from imas.exception import InvalidNetCDFEntry
+
+
+def validate_netcdf_file(filename: str) -> None:
+    """Validate if the provided netCDF file adheres to the IMAS conventions."""
+    if not filename.endswith(".nc"):
+        raise InvalidNetCDFEntry(
+            f"Invalid filename `{filename}` provided: "
+            "an IMAS netCDF file should end with `.nc`"
+        )
+
+    with DBEntry(filename, "r") as entry:
+        entry_impl: NCDBEntryImpl = entry._dbe_impl
+        dataset = entry_impl._dataset
+        factory = entry_impl._ds_factory
+
+        ids_names = factory.ids_names()
+
+        # Check that groups in the dataset correspond to an IDS/occurrence and no
+        # additional variables are smuggled inside:
+        groups = [dataset] + [dataset[group] for group in dataset.groups]
+        for group in groups:
+            if group.variables or group.dimensions:
+                raise InvalidNetCDFEntry(
+                    "NetCDF file should not have variables or dimensions in the "
+                    f"{group.name} group."
+                )
+            if group is dataset:
+                continue
+            if group.name not in ids_names:
+                raise InvalidNetCDFEntry(
+                    f"Invalid group name {group.name}: there is no IDS with this name."
+                )
+            for subgroup in group.groups:
+                try:
+                    int(subgroup)
+                except ValueError:
+                    raise InvalidNetCDFEntry(
+                        f"Invalid group name {group.name}/{subgroup}: "
+                        f"{subgroup} is not a valid occurrence number."
+                    )
+
+        for ids_name in ids_names:
+            for occurrence in entry.list_all_occurrences(ids_name):
+                group = dataset[f"{ids_name}/{occurrence}"]
+                try:
+                    NC2IDS(group, factory.new(ids_name)).validate_variables()
+                except InvalidNetCDFEntry as exc:
+                    occ = f":{occurrence}" if occurrence else ""
+                    raise InvalidNetCDFEntry(f"Invalid IDS {ids_name}{occ}: {exc}")

imas/command/cli.py

@@ -0,0 +1,239 @@
+# This file is part of imas-python.
+# You should have received the imas-python LICENSE file with this project.
+""" Main CLI entry point """
+
+import logging
+import sys
+from contextlib import ExitStack
+from pathlib import Path
+
+import click
+from packaging.version import Version
+from rich import box, console, traceback
+from rich.progress import (
+    BarColumn,
+    Progress,
+    SpinnerColumn,
+    TaskProgressColumn,
+    TextColumn,
+    TimeElapsedColumn,
+)
+from rich.table import Table
+
+import imas
+import imas.backends.imas_core.imas_interface
+from imas import DBEntry, dd_zip
+from imas.backends.imas_core.imas_interface import ll_interface
+from imas.command.db_analysis import analyze_db, process_db_analysis
+from imas.command.helpers import min_version_guard, setup_rich_log_handler
+from imas.command.timer import Timer
+from imas.exception import UnknownDDVersion
+
+logger = logging.getLogger(__name__)
+
+
+def _excepthook(type_, value, tb):
+    logger.debug("Suppressed traceback:", exc_info=(type_, value, tb))
+    # Only display the last traceback frame:
+    if tb is not None:
+        while tb.tb_next:
+            tb = tb.tb_next
+    rich_tb = traceback.Traceback.from_exception(type_, value, tb, extra_lines=0)
+    console.Console(stderr=True).print(rich_tb)
+
+
+@click.group("imas", invoke_without_command=True, no_args_is_help=True)
+def cli():
+    """imas-python command line interface.
+
+    Please use one of the available commands listed below. You can get help for each
+    command by executing:
+
+        imas <command> --help
+    """
+    # Limit the traceback to 1 item: avoid scaring CLI users with long traceback prints
+    # and let them focus on the actual error message
+    sys.excepthook = _excepthook
+
+
+cli.add_command(analyze_db)
+cli.add_command(process_db_analysis)
+
+
+@cli.command("version")
+def print_version():
+    """Print version information of imas-python."""
+    cons = console.Console()
+    grid = Table(
+        title="imas-python version info", show_header=False, title_style="bold"
+    )
+    grid.box = box.HORIZONTALS
+    if cons.size.width > 120:
+        grid.width = 120
+    grid.add_row("imas-python version:", imas.__version__)
+    grid.add_section()
+    grid.add_row("Default data dictionary version:", imas.IDSFactory().dd_version)
+    dd_versions = ", ".join(imas.dd_zip.dd_xml_versions())
+    grid.add_row("Available data dictionary versions:", dd_versions)
+    grid.add_section()
+    grid.add_row("Access Layer core version:", ll_interface.get_al_version() or "N/A")
+    console.Console().print(grid)
+
+
+@cli.command("print", no_args_is_help=True)
+@click.argument("uri")
+@click.argument("ids")
+@click.argument("occurrence", default=0)
+@click.option(
+    "--all",
+    "-a",
+    "print_all",
+    is_flag=True,
+    help="Also show nodes with empty/default values",
+)
+def print_ids(uri, ids, occurrence, print_all):
+    """Pretty print the contents of an IDS.
+
+    \b
+    uri         URI of the Data Entry (e.g. "imas:mdsplus?path=testdb").
+    ids         Name of the IDS to print (e.g. "core_profiles").
+    occurrence  Which occurrence to print (defaults to 0).
+    """
+    min_version_guard(Version("5.0"))
+    setup_rich_log_handler(False)
+
+    with DBEntry(uri, "r") as dbentry:
+        ids_obj = dbentry.get(ids, occurrence, autoconvert=False)
+        imas.util.print_tree(ids_obj, not print_all)
+
+
+@cli.command("convert", no_args_is_help=True)
+@click.argument("uri_in")
+@click.argument("dd_version")
+@click.argument("uri_out")
+@click.option(
+    "--ids",
+    default="*",
+    help="Specify which IDS to convert. \
+If not provided, all IDSs in the data entry are converted.",
+)
+@click.option("--occurrence", default=-1, help="Specify which occurrence to convert.")
+@click.option("--quiet", "-q", is_flag=True, help="Suppress progress output.")
+@click.option("--timeit", is_flag=True, help="Show timing information.")
+@click.option(
+    "--no-provenance",
+    is_flag=True,
+    help="Don't add provenance metadata to the converted IDS.",
+)
+def convert_ids(
+    uri_in, dd_version, uri_out, ids, occurrence, quiet, timeit, no_provenance
+):
+    """Convert a Data Entry (or a single IDS) to the target DD version.
+
+    Provide a different backend to URI_OUT than URI_IN to convert between backends.
+    For example:
+
+        imas convert imas:mdsplus?path=db-in 3.41.0 imas:hdf5?path=db-out
+
+    \b
+    uri_in      URI of the input Data Entry.
+    dd_version  Data dictionary version to convert to. Can also be the path to an
+                IDSDef.xml to convert to a custom/unreleased DD version.
+    uri_out     URI of the output Data Entry.
+    """
+    min_version_guard(Version("5.1"))
+    setup_rich_log_handler(quiet)
+
+    # Check if we can load the requested version
+    if dd_version in dd_zip.dd_xml_versions():
+        version_params = dict(dd_version=dd_version)
+    elif Path(dd_version).exists():
+        version_params = dict(xml_path=dd_version)
+    else:
+        raise UnknownDDVersion(dd_version, dd_zip.dd_xml_versions())
+
+    provenance_origin_uri = ""
+    if not no_provenance:
+        provenance_origin_uri = uri_in
+
+    # Use an ExitStack to avoid three nested with-statements
+    with ExitStack() as stack:
+        entry_in = stack.enter_context(DBEntry(uri_in, "r"))
+        entry_out = stack.enter_context(DBEntry(uri_out, "x", **version_params))
+
+        # First build IDS/occurrence list so we can show a decent progress bar
+        ids_list = [ids] if ids != "*" else entry_out.factory.ids_names()
+        idss_with_occurrences = []
+        for ids_name in ids_list:
+            if occurrence == -1:
+                idss_with_occurrences.extend(
+                    (ids_name, occ) for occ in entry_in.list_all_occurrences(ids_name)
+                )
+            else:
+                idss_with_occurrences.append((ids_name, occurrence))
+
+        # Create progress bar and task
+        columns = (
+            TimeElapsedColumn(),
+            BarColumn(),
+            TaskProgressColumn(),
+            TextColumn("[progress.description]{task.description}"),
+            SpinnerColumn("simpleDots", style="[white]"),
+        )
+        progress = stack.enter_context(Progress(*columns, disable=quiet))
+        task = progress.add_task("Converting", total=len(idss_with_occurrences) * 3)
+        # Create timer for timing get/convert/put
+        timer = Timer("Operation", "IDS/occurrence")
+
+        # Convert all IDSs
+        for ids_name, occurrence in idss_with_occurrences:
+            name = f"[bold green]{ids_name}[/][green]/{occurrence}[/]"
+
+            progress.update(task, description=f"Reading {name}")
+            with timer("Get", name):
+                ids = entry_in.get(ids_name, occurrence, autoconvert=False)
+
+            progress.update(task, description=f"Converting {name}", advance=1)
+            # Explicitly convert instead of auto-converting during put. This is a bit
+            # slower, but gives better diagnostics:
+            if ids._dd_version == entry_out.dd_version:
+                ids2 = ids
+            else:
+                with timer("Convert", name):
+                    ids2 = imas.convert_ids(
+                        ids,
+                        None,
+                        factory=entry_out.factory,
+                        provenance_origin_uri=provenance_origin_uri,
+                    )
+
+            # Store in output entry:
+            progress.update(task, description=f"Storing {name}", advance=1)
+            with timer("Put", name):
+                entry_out.put(ids2, occurrence)
+
+            # Update progress bar
+            progress.update(task, advance=1)
+
+    # Display timing information
+    if timeit:
+        console.Console().print(timer.get_table("Time required per IDS"))
+
+
+@cli.command("validate_nc", no_args_is_help=True)
+@click.argument("filename", type=click.Path(exists=True, dir_okay=False))
+def validate_nc(filename):
+    """Validate if the provided netCDF file adheres to the IMAS conventions."""
+    from imas.backends.netcdf.nc_validate import validate_netcdf_file
+
+    try:
+        validate_netcdf_file(filename)
+    except Exception as exc:
+        click.echo(f"File `{filename}` does not adhere to the IMAS conventions:")
+        click.echo(exc)
+        sys.exit(1)
+    click.echo(f"File `{filename}` is a valid IMAS netCDF file.")
+
+
+if __name__ == "__main__":
+    cli()

imas/command/db_analysis.py

@@ -0,0 +1,290 @@
+"""imas-python-based command line tool for analysing fields in a database."""
+
+import gzip
+import json
+import logging
+import platform
+import re
+import readline
+import sys
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Dict, Iterable, List, Optional
+
+import click
+import rich
+import rich.panel
+import rich.progress
+import rich.table
+import rich.text
+import rich.tree
+
+import imas
+from imas.command.helpers import setup_rich_log_handler
+from imas.ids_metadata import IDSMetadata
+
+directory_path = click.Path(exists=True, file_okay=False, path_type=Path)
+outfile_path = click.Path(dir_okay=False, writable=True, path_type=Path)
+infile_path = click.Path(exists=True, dir_okay=False, readable=True, path_type=Path)
+
+logger = logging.getLogger(__name__)
+
+
+@click.command("analyze-db")
+@click.argument("dbentry", nargs=-1, type=directory_path)
+@click.option(
+    "--output",
+    "-o",
+    type=outfile_path,
+    default="imas-db-analysis.json.gz",
+    help="Output file",
+)
+def analyze_db(dbentry: Iterable[Path], output: Path) -> None:
+    """Analyze one or more Data Entries stored in the HDF5 backend format.
+
+    This analysis will collect the following data:
+
+    \b
+    - The (host)name of the machine
+    - The path to the Data Entry
+    - Which IDSs are stored in the Data Entry
+    - Which fields in the IDSs are filled
+
+    \b
+    Arguments:
+    DBENTRY     Folder(s) containing IMAS data stored in the HDF5 backend format.
+
+    \b
+    Notes:
+    1. This tool does not accept IMAS URIs, only folder names. To avoid loading all
+       data, the IDSs are inspected by looking at the HDF5 files directly.
+    2. This tool uses the optional `h5py` dependency. An error is raised when this
+       package is not available.
+    3. If your data is stored in another format than HDF5, you may use `imas convert`
+       to convert the data into the HDF5 backend format first.
+    """
+    # Test if h5py is available
+    try:
+        import h5py  # noqa
+    except ModuleNotFoundError:
+        rich.print(
+            "[red]Module [bold]h5py[/bold] is not available.[/]\n"
+            "Please load the [bold]h5py[/] HPC module, or install "
+            "[bold]h5py[/] with `[bold]pip install h5py[/bold]`.",
+            file=sys.stderr,
+        )
+        sys.exit(1)
+    setup_rich_log_handler(False)
+
+    # Load existing data?
+    if output.exists():
+        logger.info("Existing data found. New data will be appended to it.")
+        try:
+            with gzip.open(output, "rt", encoding="utf-8") as file:
+                data = json.load(file)
+        except Exception:
+            logger.error("Could not read data from '%s'.", output, exc_info=True)
+            logger.error("Please remove the file '%s' and try again.", output)
+            sys.exit(1)
+    else:
+        data = []
+
+    # Analyze provided paths
+    for entry in rich.progress.track(dbentry, transient=True):
+        # Check if a master.h5 is present
+        if not (entry / "master.h5").exists:
+            logger.error("Could not find '%s/master.h5', skipping this folder.", entry)
+            continue
+        result = analyze_folder(entry)
+        data.append(result)
+
+    if data:
+        with gzip.open(output, "wt", encoding="utf-8") as file:
+            json.dump(data, file)
+        logger.info(f"Output data is stored in {output}")
+
+
+def analyze_folder(entry: Path):
+    """Analyze a folder with HDF5 files, returning a structure containing IDS usage."""
+    return {
+        "machine_name": platform.node(),
+        "dbentry_path": str(entry),
+        "ids_info": [
+            ids_info(idsfile)
+            for idsfile in entry.glob("*.h5")
+            if idsfile.stem != "master"
+        ],
+    }
+
+
+def ids_info(idsfile: Path):
+    """Open HDF5 ids file and return structure containing the filled data paths."""
+    import h5py
+
+    name_and_occurrence = idsfile.stem
+    name, occurrence = re.split("_(?=[0-9]+)|$", name_and_occurrence, 1)
+    filled_data = []
+    with h5py.File(idsfile, "r") as f:
+        group = f[name_and_occurrence]
+        for dataset in group.keys():
+            if dataset.endswith("SHAPE"):
+                continue
+            path = dataset.replace("[", "").replace("]", "").replace("&", "/")
+            filled_data.append(path)
+
+    return {
+        "name": name,
+        "occurrence": occurrence,
+        "filled_data": filled_data,
+    }
+
+
+@click.command("process-db-analysis")
+@click.argument(
+    "infiles", metavar="INPUT_FILES...", nargs=-1, type=infile_path, required=True
+)
+@click.option("--show-empty-ids", is_flag=True, help="Show empty IDSs in the overview.")
+def process_db_analysis(infiles, show_empty_ids):
+    """Process supplied Data Entry analyses, and display statistics.
+
+    \b
+    Arguments:
+    INPUT_FILES     File(s) produced by `imas analyze-db` to process.
+    """
+    setup_rich_log_handler(False)
+
+    factory = imas.IDSFactory()
+    filled_per_ids = {ids_name: set() for ids_name in factory.ids_names()}
+    logger.info("Using Data Dictionary version %s.", factory.dd_version)
+    logger.info("Reading %d input files...", len(infiles))
+
+    # Read input data and collate usage info per IDS
+    for fname in infiles:
+        with gzip.open(fname, "rt", encoding="utf-8") as file:
+            data = json.load(file)
+
+        for entry in data:
+            for ids_info in entry["ids_info"]:
+                fill_info = filled_per_ids[ids_info["name"]]
+                fill_info.update(ids_info["filled_data"])
+
+    logger.info("Done reading input files.")
+    logger.info("Analyzing filled data...")
+
+    # Construct AnalysisNodes per IDS
+    analysis_nodes: Dict[str, _AnalysisNode] = {}
+    for ids_name, filled in filled_per_ids.items():
+        metadata = factory.new(ids_name).metadata
+        ids_analysis_node = _AnalysisNode("")
+
+        def walk_metadata_tree(metadata: IDSMetadata, node: _AnalysisNode):
+            """Recursively walk the IDSMetadata tree to check which nodes are filled."""
+            for meta_child in metadata._children.values():
+                path = meta_child.path_string
+                childnode = _AnalysisNode(path, path in filled)
+                node.children.append(childnode)
+
+                if meta_child._children:  # Childnode is a structure/AoS itself
+                    # Recurse
+                    walk_metadata_tree(meta_child, childnode)
+                    # Update node statistics
+                    node.num_desc_nodes += childnode.num_desc_nodes
+                    node.num_desc_nodes_filled += childnode.num_desc_nodes_filled
+                else:  # Childnode is a data node
+                    # Update node statistics
+                    node.num_desc_nodes += 1
+                    if childnode.used:
+                        node.num_desc_nodes_filled += 1
+            node.used = node.num_desc_nodes_filled > 0
+
+        walk_metadata_tree(metadata, ids_analysis_node)
+        analysis_nodes[ids_name] = ids_analysis_node
+
+    # Display summary results per IDS
+    logger.info("Analysis done. Printing results:")
+    sorted_ids_names = sorted(
+        analysis_nodes,
+        key=lambda ids_name: -analysis_nodes[ids_name].fill_fraction,
+    )
+    table = rich.table.Table("IDS", "Filled nodes", caption="Usage summary")
+    for ids_name in sorted_ids_names:
+        node = analysis_nodes[ids_name]
+        if node.num_desc_nodes_filled == 0 and not show_empty_ids:
+            continue  # hide IDSs without data
+        dim = "" if node.num_desc_nodes_filled > 0 else "[dim]"
+        table.add_row(
+            f"{dim}[bold]{ids_name}",
+            f"{dim}{node.num_desc_nodes_filled: >4} / {node.num_desc_nodes: >4} "
+            f"({node.fill_fraction: >6.2%})",
+        )
+    rich.print(table)
+
+    # Use readline for autocompleting IDS names when pressing <tab>
+    def rlcompleter(text, state) -> Optional[str]:
+        matching_idss = [name for name in sorted_ids_names if name.startswith(text)]
+        if state < len(matching_idss):
+            return matching_idss[state]
+        return None
+
+    readline.set_completer(rlcompleter)
+    readline.parse_and_bind("tab: complete")
+
+    # Display input prompt for detailed results per IDS
+    while True:
+        # click.prompt doesn't work nicely with readline, so use builtin input():
+        ids_name = input("Enter IDS name to show detailed usage [exit]: ").strip()
+
+        if ids_name == "exit" or not ids_name:
+            break
+        elif ids_name not in analysis_nodes:
+            rich.print(f"[red]Unknown IDS name: [bold]{ids_name}")
+            continue
+
+        node = analysis_nodes[ids_name]
+        tree = rich.tree.Tree(
+            f"[bold red]{ids_name}[/]: [bold cyan]{node.num_desc_nodes_filled}[/]/"
+            f"[bold cyan]{node.num_desc_nodes}[/] data nodes filled"
+        )
+        for childnode in node.children:
+            childnode.fill_tree(tree)
+        rich.print(rich.panel.Panel(tree, expand=False))
+
+
+@dataclass
+class _AnalysisNode:
+    """Data class to store analysis results for a node in an IDS."""
+
+    path: str
+    """Path of this node, e.g. "ids_properties/comment"."""
+    used: bool = False
+    """True when any analysed Data Entry had this node populated."""
+    children: List["_AnalysisNode"] = field(repr=False, default_factory=list)
+    """List of child nodes (when this is a structure or AoS node)."""
+    num_desc_nodes = 0
+    """Total descendent data (i.e. no structure/AoS) nodes."""
+    num_desc_nodes_filled = 0
+    """Total number of descendent data nodes that are used."""
+
+    @property
+    def fill_fraction(self):
+        """The fill fraction of all descendent nodes."""
+        return self.num_desc_nodes_filled / self.num_desc_nodes
+
+    def fill_tree(self, tree: rich.tree.Tree):
+        """Fill provided rich.tree.Tree with analysis data."""
+        if not self.children:  # leaf node
+            if not self.used:
+                return  # skip empty leaf nodes
+            label = f"[green]{self.path}[/] is used"
+        else:
+            label = (
+                f"[blue]{self.path}[/]: [bold cyan]{self.num_desc_nodes_filled}[/]/"
+                f"[bold cyan]{self.num_desc_nodes}[/] descendent nodes filled"
+            )
+            if self.num_desc_nodes_filled == 0:
+                label = rich.text.Text.from_markup(label, style="dim")
+
+        subtree = tree.add(label)
+        if self.used:
+            for child in self.children:
+                child.fill_tree(subtree)

imas/command/helpers.py

@@ -0,0 +1,45 @@
+import logging
+import sys
+
+import click
+from packaging.version import Version
+from rich.logging import RichHandler
+
+from imas.backends.imas_core.imas_interface import ll_interface
+
+
+def setup_rich_log_handler(quiet: bool):
+    """Setup rich.logging.RichHandler on the root logger.
+
+    Args:
+        quiet: When True: set log level of the `imas` logger to WARNING or higher.
+    """
+    # Disable default imas log handler
+    imas_logger = logging.getLogger("imas")
+    for handler in imas_logger.handlers:
+        imas_logger.removeHandler(handler)
+    # Disable any root log handlers
+    root_logger = logging.getLogger()
+    for handler in root_logger.handlers:
+        root_logger.removeHandler(handler)
+    # Install rich handler on the root logger:
+    root_logger.addHandler(RichHandler())
+    if quiet:  # Silence imas-python INFO messages
+        # If loglevel is less than WARNING, set it to WARNING:
+        imas_logger.setLevel(max(logging.WARNING, imas_logger.getEffectiveLevel()))
+
+
+def min_version_guard(al_version: Version):
+    """Print an error message if the loaded AL version is too old.
+
+    Args:
+        al_version: Minimum imas_core version required for this command.
+    """
+    used_version = ll_interface._al_version
+    if used_version >= al_version:
+        return
+    click.echo(
+        f"This command requires at least version {al_version} of the Access Layer."
+    )
+    click.echo(f"The current loaded version is {used_version}, which is too old.")
+    sys.exit(1)