epgx 0.1.dev0__py3-none-any.whl

epgx/examples.py

@@ -0,0 +1,146 @@
+"""Utilities for the example scripts (examples/)
+
+    quick(): reduced sizes for fast runs (EPGX_QUICK=1 or `--quick`)
+    show(name): save the open figures if requested (EPGX_SAVEFIG=<folder> or
+        `--savefig [folder]`), then `plt.show()`
+    device_name(): readable name of the computing device (e.g. for timings)
+
+Example headers (module docstring):
+
+    Title
+
+    Description (free text)
+
+    Shows: what the example demonstrates
+    Reference: Author et al. Journal Year;vol:pages. doi:... (optional)
+    Reproduces: figure(s) of the reference (optional)
+    Differences: w.r.t. the reference (optional)
+
+Field values may continue on indented lines. The index of the examples
+(examples/README.md) is generated from the headers (tools/make_examples_index.py).
+"""
+
+import ast
+import os
+import pathlib
+import platform
+import re
+import subprocess
+import sys
+
+QUICK_VAR = "EPGX_QUICK"
+SAVEFIG_VAR = "EPGX_SAVEFIG"
+DEFAULT_FOLDER = "figures"
+HEADER_FIELDS = ("Shows", "Reference", "Reproduces", "Differences")
+
+
+def parse_header(docstring):
+    """parse an example header (cf. module docstring)
+
+    Returns:
+        dict with "title", "description" and the header fields found
+    """
+    lines = (docstring or "").strip().splitlines()
+    if not lines:
+        raise ValueError("Empty example header")
+    header = {"title": lines[0].strip()}
+    pattern = re.compile(rf"^({'|'.join(HEADER_FIELDS)}):\s*(.*)$")
+    description, field = [], None
+    for line in lines[1:]:
+        match = pattern.match(line)
+        if match:
+            field = match.group(1)
+            header[field] = match.group(2).strip()
+        elif field and line.startswith((" ", "\t")) and line.strip():
+            header[field] += " " + line.strip()
+        else:
+            field = None
+            description.append(line)
+    header["description"] = "\n".join(description).strip()
+    return header
+
+
+def read_header(path):
+    """parse the header of an example file"""
+    tree = ast.parse(pathlib.Path(path).read_text(encoding="utf-8"))
+    return parse_header(ast.get_docstring(tree))
+
+
+def quick():
+    """True in quick mode (EPGX_QUICK=1 or `--quick` on the command line)"""
+    value = os.environ.get(QUICK_VAR, "").strip().lower()
+    return value in ("1", "true", "yes") or "--quick" in sys.argv
+
+
+def savefig_folder():
+    """folder where to save figures (EPGX_SAVEFIG or `--savefig [folder]`), or None"""
+    if "--savefig" in sys.argv:
+        index = sys.argv.index("--savefig")
+        value = sys.argv[index + 1] if index + 1 < len(sys.argv) else ""
+        return pathlib.Path(
+            value if value and not value.startswith("-") else DEFAULT_FOLDER
+        )
+    value = os.environ.get(SAVEFIG_VAR)
+    return pathlib.Path(value) if value else None
+
+
+def show(name):
+    """save the open figures if requested, then show them
+
+    Args:
+        name: name of the example (e.g. `__file__`), used for the file names:
+            <folder>/<example>_<figure label or number>.png
+    """
+    from matplotlib import pyplot as plt
+
+    folder = savefig_folder()
+    if folder is not None:
+        folder.mkdir(parents=True, exist_ok=True)
+        stem = pathlib.Path(name).stem
+        for number in plt.get_fignums():
+            figure = plt.figure(number)
+            label = figure.get_label() or str(number)
+            figure.savefig(folder / f"{stem}_{label}.png", dpi=100)
+    plt.show()
+
+
+def device_name():
+    """readable name of the computing device and precision, e.g. for timings
+
+    GPU: model reported by JAX; CPU: processor name (registry on Windows,
+    /proc/cpuinfo on Linux, sysctl on macOS) and number of threads.
+    """
+    import jax
+
+    device = jax.devices()[0]
+    precision = "x64" if jax.config.jax_enable_x64 else "float32"
+    if device.platform != "cpu":
+        return f"{device.device_kind} ({device.platform}), {precision}"
+    return f"{_cpu_name()} ({os.cpu_count()} threads), {precision}"
+
+
+def _cpu_name():
+    """processor name, or the generic platform description"""
+    try:
+        system = platform.system()
+        if system == "Windows":
+            import winreg
+
+            key = r"HARDWARE\DESCRIPTION\System\CentralProcessor\0"
+            with winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, key) as handle:
+                name = winreg.QueryValueEx(handle, "ProcessorNameString")[0]
+        elif system == "Linux":
+            text = pathlib.Path("/proc/cpuinfo").read_text()
+            name = re.search(r"model name\s*:\s*(.*)", text).group(1)
+        elif system == "Darwin":
+            name = subprocess.run(
+                ["sysctl", "-n", "machdep.cpu.brand_string"],
+                capture_output=True,
+                text=True,
+                check=True,
+            ).stdout
+        else:
+            name = ""
+    except Exception:
+        name = ""
+    return " ".join(name.split()) or platform.processor() or "CPU"