Merge pull request #727 from European-XFEL/display-inst-data-counts

lsxfel / .info() option to show instrument data counts

Author: takluyver

docs/cli.rst

@@ -22,19 +22,39 @@ file:
 
 .. program:: lsxfel
 
+The options below are only relevant when inspecting a single run or file.
+
+.. option:: --source <source-pattern>, -s <source-pattern>
+
+   Select which sources to show, using a pattern like ``*/BAM*:output`` or
+   a substring like ``BAM``. Can be used several times to select different
+   patterns.
+
 .. option:: --detail <source-pattern>
 
-   Show more detail on the keys and data of the sources selected by a pattern
-   like ``*/XGM/*``. Only applies when inspecting a single run or file.
-   Can be used several times to select different patterns.
+   Show more detail on the keys and data of the sources selected by a pattern,
+   in the same formats as ``--source``. Can be used several times.
 
    This option can make ``lsxfel`` considerably slower.
 
+.. option:: --counts
+
+   Show data counts for instrument sources.
+
+.. option:: --no-group
+
+   Don't try to group similar source names; show each individually.
+
 .. option:: --aggregators, -g
 
    Show the data aggregator in which each source is saved. This option
    can make ``lsxfel`` considerably slower.
 
+.. option:: --auxiliary, -x
+
+   Include auxiliary sources alongside regular sources. This can make ``lsxfel``
+   slower.
+
 .. _cmd-validate:
 
 ``extra-data-validate``

extra_data/cli/lsxfel.py

@@ -8,17 +8,7 @@
 import sys
 
 from ..read_machinery import FilenameInfo
-from ..reader import H5File, RunDirectory
-
-
-def describe_file(path, details_for_sources=(), with_aggregators=False,
-                  with_auxiliary=False, group_sources=True):
-    """Describe a single HDF5 data file"""
-    basename = os.path.basename(path)
-    print(basename, ": Data file")
-
-    h5file = H5File(path)
-    h5file.info(details_for_sources, with_aggregators, with_auxiliary, group_sources)
+from ..reader import DataCollection, H5File, RunDirectory
 
 
 def summarise_file(path):
@@ -29,16 +19,6 @@ def summarise_file(path):
     print(f"  {len(f.train_ids)} trains, {len(f.all_sources)} sources")
 
 
-def describe_run(path, details_for_sources=(), with_aggregators=False,
-                 with_auxiliary=False, group_sources=True):
-    basename = os.path.basename(path)
-    print(basename, ": Run directory")
-    print()
-
-    run = RunDirectory(path)
-    run.info(details_for_sources, with_aggregators, with_auxiliary, group_sources)
-
-
 def summarise_run(path, indent=''):
     basename = os.path.basename(path)
 
@@ -74,6 +54,10 @@ def main(argv=None):
         prog='lsxfel', description="Summarise XFEL data in files or folders"
     )
     ap.add_argument('paths', nargs='*', help="Files/folders to look at")
+    ap.add_argument('--source', '-s', action='append', default=[],
+        help="Filter the list of sources, with wildcard patterns like "
+             "'*XGM/*:output' or partial matches like 'XGM'. Can be used more "
+             "than once.")
     ap.add_argument('--detail', '-d', action='append', default=[],
         help="Show details on keys & data for specified sources. "
              "This can slow down lsxfel considerably. "
@@ -82,6 +66,10 @@ def main(argv=None):
              "Can be used more than once to include several patterns. "
              "Only used when inspecting a single run or file."
     )
+    ap.add_argument('--counts', '-c', action='store_true',
+        help="Show data counts for instrument sources.")
+    ap.add_argument('--no-group', action='store_true',
+        help="Don't try to group similar source names; show each individually.")
     ap.add_argument('--aggregators', '-g', action='store_true',
         help="Include the aggregator each source is saved in. "
              "Only used when inspecting a single run or file and can slow "
@@ -90,16 +78,32 @@ def main(argv=None):
         help="Include auxiliary sources alongside regular sources. "
              "Only used when inspecting a single run or file and can slow "
              "down lsxfel.")
-    ap.add_argument('--no-group', action='store_true',
-        help="Don't try to group similar source names; show each individually.")
     args = ap.parse_args(argv)
     paths = args.paths or [os.path.abspath(os.getcwd())]
 
     # If --detail doesn't contain glob wildcards, treat it as a substring
     detail_patterns = [p if re.search(r'[*?[]', p) else f'*{p}*'
                        for p in args.detail]
-
-    group_sources = not args.no_group
+    source_patterns = [p if re.search(r'[*?[]', p) else f'*{p}*'
+                       for p in args.source]
+
+    def dc_info(dc: DataCollection):
+        if source_patterns:
+            try:
+                sel = dc.select(source_patterns)
+            except ValueError as e:  # No matching sources
+                sys.exit(str(e))
+            print(f"Showing {len(sel.all_sources)} / {len(dc.all_sources)} sources")
+        else:
+            sel = dc
+        print()
+        sel.info(
+            detail_patterns,
+            counts=args.counts,
+            group_sources=(not args.no_group),
+            with_aggregators=args.aggregators,
+            with_auxiliary=args.auxiliary,
+        )
 
     if len(paths) == 1:
         path = paths[0]
@@ -109,8 +113,8 @@ def main(argv=None):
             contents = sorted(os.listdir(path))
             if any(f.endswith('.h5') for f in contents):
                 # Run directory
-                describe_run(path, detail_patterns, args.aggregators,
-                             args.auxiliary, group_sources=group_sources)
+                print(basename, ": Run directory")
+                dc_info(RunDirectory(path))
             elif any(re.match(r'r\d+', f) for f in contents):
                 # Proposal directory, containing runs
                 print(basename, ": Proposal data directory")
@@ -131,8 +135,8 @@ def main(argv=None):
                 print(basename, ": Unrecognised directory")
         elif os.path.isfile(path):
             if path.endswith('.h5'):
-                describe_file(path, detail_patterns, args.aggregators,
-                              args.auxiliary, group_sources=group_sources)
+                print(basename, ": Data file")
+                dc_info(H5File(path))
             else:
                 print(basename, ": Unrecognised file")
                 return 2

extra_data/display.py

@@ -13,14 +13,15 @@
 class InfoPrinter:
     def __init__(
             self, dc: DataCollection, details_for_sources=(), with_aggregators=False,
-            group_sources=True,
+            data_counts=False, group_sources=True,
     ):
         self.dc = dc
         self.details_for_sources = details_for_sources
         self.details_sources_re = [re.compile(fnmatch.translate(p))
                                    for p in details_for_sources]
         self.with_aggregators = with_aggregators
-        self.group_sources = group_sources and not with_aggregators
+        self.data_counts = data_counts
+        self.group_sources = group_sources and not (with_aggregators or data_counts)
 
         # Invert aliases for faster lookup.
         self.src_aliases = defaultdict(set)
@@ -92,7 +93,7 @@ def source_line(self, s):
         agg_str = f" [{self.dc[s].aggregator}]" if self.with_aggregators else ""
         print(f"  -{agg_str} {s} {self.src_alias_list(s)}")
 
-    def list_sources(self, srcs: list, detail: Callable):
+    def list_sources(self, srcs: list, detail: Callable, data_counts=False):
         current_group = SourceGroup()
         def flush_group():
             nonlocal current_group
@@ -108,6 +109,8 @@ def flush_group():
                 self.source_line(s)
                 if show_detail:
                     detail(s)
+                elif data_counts:
+                    self.inst_data(s)
             else:
                 if not current_group.add(s):
                     flush_group()
@@ -116,7 +119,7 @@ def flush_group():
 
     def inst_sources(self):
         srcs = self.dc.instrument_sources
-        if self.details_sources_re:
+        if (self.details_sources_re or self.data_counts):
             # All instrument sources with details enabled.
             displayed_inst_srcs = srcs - self.dc.legacy_sources.keys()
             print(len(displayed_inst_srcs), "instrument sources:")
@@ -127,17 +130,28 @@ def inst_sources(self):
             )
             print(len(displayed_inst_srcs), "instrument sources (excluding XTDF detectors):")
 
-        self.list_sources(displayed_inst_srcs, self.inst_detail)
+        self.list_sources(displayed_inst_srcs, self.inst_detail, self.data_counts)
         print()
 
+    def inst_data(self, s):
+        # Show data volume, without list of keys
+        sd = self.dc[s]
+        sif = SourceInfoFormatter(self.dc[s])
+        index_groups = sorted(sd.index_groups)
+        if len(index_groups) == 1:
+            print(f"    {sif.src_data_detail(index_groups[0])}")
+        else:
+            for group in index_groups:
+                print(f"    {group}: {sif.src_data_detail(group)}")
+
     def inst_detail(self, s):
         # Detail for instrument sources:
         sd = self.dc[s]
-        sif = SourceInfoFormatter(self.dc[s], self.srckey_aliases[s])
+        sif = SourceInfoFormatter(sd, self.srckey_aliases[s])
         for group, keys in groupby(sorted(sd.keys()), key=lambda k: k.split(".")[0]):
             print(f"    - {group}:")
             keys = list(keys)
-            print("      " + sif.src_data_detail(keys))
+            print("      " + sif.src_data_detail(group))
             for l in sif.keys_detail(keys):
                 print("      " + l)
 
@@ -182,7 +196,7 @@ def legacy_sources(self):
 
     def show(self, with_auxiliary=False):
         self.trains()
-        if not self.details_sources_re:
+        if not (self.details_sources_re or self.data_counts):
             self.xtdf()
         self.inst_sources()
         self.ctrl_sources()
@@ -308,16 +322,17 @@ def __init__(self, src: SourceData, key_aliases=None):
         self.src = src
         self.key_aliases = key_aliases or {}
 
-    def src_data_detail(self, keys):
+    def src_data_detail(self, index_group):
         """Detail for how much data is present for an instrument group"""
-        if not keys:
-            return
-        counts = self.src[list(keys)[0]].data_counts()
-        ntrains_data = (counts > 0).sum()
+        counts = self.src.data_counts(index_group=index_group)
+        counts = counts[counts > 0]
+        ntrains_data = len(counts)
+        cmin, cmax = (counts.min(), counts.max()) if ntrains_data else (0, 0)
+        count_range = f"{cmin}" if (cmin == cmax) else f"{cmin}–{cmax}"
         return (
             f'data for {ntrains_data} trains '
             f'({ntrains_data / len(self.src.train_ids):.2%}), '
-            f'up to {counts.max()} entries per train'
+            f'{count_range} entries per train'
         )
 
     def keys_detail(self, keys=None):

extra_data/reader.py

@@ -1304,14 +1304,30 @@ def __repr__(self):
         return f"<extra_data.DataCollection for {len(self.all_sources)} " \
                f"sources and {len(self.train_ids)} trains>"
 
-    def info(self, details_for_sources=(), with_aggregators=False,
-             with_auxiliary=False, group_sources=True):
-        """Show information about the selected data."""
+    def info(self, details_for_sources=(), *, counts=False, group_sources=True,
+             with_aggregators=False, with_auxiliary=False):
+        """Show information about the selected data.
+
+        Parameters
+        ----------
+
+        details_for_sources: list of str
+            Glob patterns selecting sources to show additional information.
+        counts: bool
+            Show data counts for all instrument sources.
+        group_sources: bool
+            Group similar source names together if possible (on by default).
+        with_aggregators: bool
+            Show which data aggregator each source was saved by.
+        with_auxiliary: bool
+            Show auxiliary sources (REDUCTION & ERRATA)
+        """
         from .display import InfoPrinter
         InfoPrinter(
             self,
             details_for_sources=details_for_sources,
             with_aggregators=with_aggregators,
+            data_counts=counts,
             group_sources=group_sources,
         ).show(with_auxiliary)
 

extra_data/tests/test_display.py

@@ -109,3 +109,11 @@ def test_show_ungrouped(mock_fxe_jungfrau_run, capsys):
     out, err = capsys.readouterr()
     assert "FXE_XAD_JF1M/DET/JNGFR{01-02}" not in out
     assert "FXE_XAD_JF1M/DET/JNGFR01" in out
+
+
+def test_show_counts(mock_fxe_raw_run, capsys):
+    run = RunDirectory(mock_fxe_raw_run)
+    run.info(counts=True)
+    out, err = capsys.readouterr()
+    assert "data for 0 trains" in out  # FXE_XAD_GEC/CAM/CAMERA_NODATA
+    assert "data for 480 trains" in out

extra_data/tests/test_lsxfel.py

@@ -1,3 +1,5 @@
+import re
+
 from extra_data.cli import lsxfel
 
 
@@ -13,3 +15,11 @@ def test_lsxfel_run(mock_fxe_raw_run, capsys):
 
     assert "480 trains" in out
     assert "16 detector files" in out
+
+
+def test_lsxfel_main(mock_fxe_raw_run, capsys):
+    lsxfel.main([mock_fxe_raw_run, "--source", "XGM"])
+    out, err = capsys.readouterr()
+    assert re.search(r"trains:\s*480", out)
+    assert "SA1_XTD2_XGM/DOOCS/MAIN:output" in out
+    assert "FXE_XAD_GEC/CAM/CAMERA" not in out  # Selected out by --source