unitrefine 0.2.0__py3-none-any.whl

unitrefine/__init__.py

@@ -0,0 +1,75 @@
+from importlib.metadata import PackageNotFoundError, version
+
+try:
+    # the version in pyproject.toml, e.g. to say which version a problem was found with
+    __version__ = version("unitrefine")
+except PackageNotFoundError:
+    # run from the source folder without installing it
+    __version__ = "unknown"
+
+# The SpikeInterface-GUI scripts run in their own process. They print these markers so the main
+# window can show which step they are on while they open, and when their window is shown.
+SIGUI_SHOWN_MARKER = "UNITREFINE_SIGUI_WINDOW_SHOWN"
+SIGUI_STEP_MARKER = "UNITREFINE_SIGUI_STEP:"
+# printed when labels could not be saved where they belong and were saved to a backup file instead
+LABELS_BACKUP_MARKER = "UNITREFINE_LABELS_BACKUP:"
+
+
+# The labels are passed to the SpikeInterface-GUI in curation format 2, which spikeinterface-gui handles
+# correctly (also for units shown without a label) from this version on
+MINIMUM_SIGUI_VERSION = "0.14.0"
+
+
+def check_sigui_version():
+    """
+    Stops the SpikeInterface-GUI script with a clear message if spikeinterface-gui is too old: older versions
+    would not save the labels set with the n/g/m keys correctly.
+    """
+    import sys
+    from packaging.version import parse
+    import spikeinterface_gui
+
+    installed = getattr(spikeinterface_gui, "__version__", "0")
+    if parse(installed) < parse(MINIMUM_SIGUI_VERSION):
+        print(f"UnitRefine needs spikeinterface-gui {MINIMUM_SIGUI_VERSION} or newer, but {installed} is installed, "
+              f"which would not save your labels correctly. Please upgrade it, e.g. with "
+              f"'pip install --upgrade spikeinterface-gui' or 'uv sync'.", flush=True)
+        sys.exit(1)
+
+
+def unwritable_file_message(path, error):
+    """What to tell the user when a file can't be written, e.g. because it is read-only."""
+    advice = ("The file is read-only or open in another program. In File Explorer, right-click it, choose "
+              "Properties and untick 'Read-only' (or close the other program), then try again.")
+    return f"UnitRefine could not save {path}:\n{error}\n\n{advice}"
+
+
+def save_labels_or_backup(labels, path, write=None, **to_csv_kwargs):
+    """
+    Saves the labels set in the SpikeInterface-GUI to `path` (with `write(labels, path)` if given). If that
+    fails, e.g. because the file is read-only, the labels are saved to a backup file in the home folder, so
+    they are not lost, and the main window is told where (LABELS_BACKUP_MARKER). Returns where they were saved.
+    """
+    import time
+    from pathlib import Path
+
+    path = Path(path)
+    try:
+        if write is not None:
+            write(labels, path)
+        else:
+            labels.to_csv(path, **to_csv_kwargs)
+        return path
+    except OSError as error:
+        backup = Path.home() / f"UnitRefine_unsaved_{time.strftime('%Y%m%d_%H%M%S')}_{path.name}"
+        labels.to_csv(backup, **to_csv_kwargs)
+        print(unwritable_file_message(path, error), flush=True)
+        print(f"Your labels were saved to {backup} instead. Once {path.name} can be written, copy that file over it.",
+              flush=True)
+        print(f"{LABELS_BACKUP_MARKER}{backup}|{path}", flush=True)
+        return backup
+
+
+def report_sigui_step(step_name):
+    """Called by the SpikeInterface-GUI scripts when they start a step, e.g. 'loading'."""
+    print(f"{SIGUI_STEP_MARKER}{step_name}", flush=True)

unitrefine/apply_models.py

@@ -0,0 +1,111 @@
+from pathlib import Path
+
+import PyQt5.QtWidgets as QtWidgets
+from spikeinterface_gui.main import check_folder_is_analyzer
+
+from unitrefine.predictions import find_analyzers
+
+
+class ApplyModelsDialog(QtWidgets.QDialog):
+    """Pick the analyzers to apply the selected models to, and the folder to save the predictions in."""
+
+    def __init__(self, selection_text, default_output_folder, parent=None):
+
+        super().__init__(parent)
+
+        self.setWindowTitle("UNITREFINE: Apply models to analyzers")
+        self.setMinimumWidth(650)
+
+        layout = QtWidgets.QVBoxLayout(self)
+
+        selection_label = QtWidgets.QLabel(f"Models: {selection_text}")
+        selection_label.setWordWrap(True)
+        layout.addWidget(selection_label)
+
+        analyzers_label = QtWidgets.QLabel("Analyzers the models will be applied to:")
+        analyzers_label.setStyleSheet("font-weight: bold;")
+        layout.addWidget(analyzers_label)
+
+        self.analyzer_list = QtWidgets.QListWidget()
+        self.analyzer_list.setSelectionMode(QtWidgets.QAbstractItemView.ExtendedSelection)
+        layout.addWidget(self.analyzer_list)
+
+        buttons_layout = QtWidgets.QHBoxLayout()
+        add_button = QtWidgets.QPushButton("+ Add analyzer folder")
+        add_button.clicked.connect(self.add_analyzer_folder)
+        add_all_button = QtWidgets.QPushButton("+ Add all analyzers inside a folder")
+        add_all_button.setToolTip("Searches up to 3 folder levels down for Sorting Analyzer folders")
+        add_all_button.clicked.connect(self.add_analyzers_inside_folder)
+        remove_button = QtWidgets.QPushButton("Remove selected")
+        remove_button.clicked.connect(self.remove_selected)
+        buttons_layout.addWidget(add_button)
+        buttons_layout.addWidget(add_all_button)
+        buttons_layout.addWidget(remove_button)
+        layout.addLayout(buttons_layout)
+
+        output_label = QtWidgets.QLabel("Save the predictions (one csv per analyzer) in:")
+        output_label.setStyleSheet("font-weight: bold;")
+        layout.addWidget(output_label)
+
+        output_layout = QtWidgets.QHBoxLayout()
+        self.output_folder_form = QtWidgets.QLineEdit(str(default_output_folder))
+        browse_button = QtWidgets.QPushButton("Browse")
+        browse_button.clicked.connect(self.choose_output_folder)
+        output_layout.addWidget(self.output_folder_form)
+        output_layout.addWidget(browse_button)
+        layout.addLayout(output_layout)
+
+        button_box = QtWidgets.QDialogButtonBox()
+        self.apply_button = button_box.addButton("Apply", QtWidgets.QDialogButtonBox.AcceptRole)
+        button_box.addButton(QtWidgets.QDialogButtonBox.StandardButton.Cancel)
+        button_box.accepted.connect(self.accept_if_ready)
+        button_box.rejected.connect(self.reject)
+        layout.addWidget(button_box)
+
+    def add_paths(self, paths):
+        existing = set(self.analyzer_paths())
+        for path in paths:
+            if str(path) not in existing:
+                self.analyzer_list.addItem(str(path))
+                existing.add(str(path))
+
+    def add_analyzer_folder(self):
+        folder = QtWidgets.QFileDialog.getExistingDirectory(self, "Select a Sorting Analyzer folder")
+        if folder == "":
+            return
+        if check_folder_is_analyzer(folder):
+            self.add_paths([Path(folder)])
+        else:
+            print(f"Selected directory {folder} is not a SortingAnalyzer.")
+
+    def add_analyzers_inside_folder(self):
+        folder = QtWidgets.QFileDialog.getExistingDirectory(self, "Select a folder containing Sorting Analyzers")
+        if folder == "":
+            return
+        analyzers = find_analyzers(folder)
+        print(f"Found {len(analyzers)} Sorting Analyzers in {folder}.")
+        self.add_paths(analyzers)
+
+    def remove_selected(self):
+        for item in self.analyzer_list.selectedItems():
+            self.analyzer_list.takeItem(self.analyzer_list.row(item))
+
+    def choose_output_folder(self):
+        folder = QtWidgets.QFileDialog.getExistingDirectory(self, "Select the folder to save the predictions in")
+        if folder != "":
+            self.output_folder_form.setText(folder)
+
+    def analyzer_paths(self):
+        return [self.analyzer_list.item(row).text() for row in range(self.analyzer_list.count())]
+
+    def output_folder(self):
+        return self.output_folder_form.text().strip()
+
+    def accept_if_ready(self):
+        if len(self.analyzer_paths()) == 0:
+            QtWidgets.QMessageBox.warning(self, "No analyzers", "Add at least one analyzer to apply the models to.")
+            return
+        if self.output_folder() == "":
+            QtWidgets.QMessageBox.warning(self, "No output folder", "Choose a folder to save the predictions in.")
+            return
+        self.accept()

unitrefine/cli.py

@@ -0,0 +1,251 @@
+"""
+Applies models to analyzers without the GUI, e.g. on a computing cluster (HPC). Nothing here imports Qt, so it
+runs without a screen.
+
+    unitrefine-apply ANALYZERS --model MODEL [--model MODEL ...] [--cascade] [--output FOLDER] [--compute_missing_metrics]
+    unitrefine-download-models [REPO_ID ...] [--models_folder FOLDER]
+
+ANALYZERS can be analyzer folders, folders with analyzers inside, or text files with one analyzer path per line.
+A MODEL is a model folder (e.g. my_project/models/my_model) or a Hugging Face repo id, which is downloaded once
+into the models folder. The predictions are the ones the GUI makes, with the same labels and columns: one csv per
+analyzer, {output}/{analyzer name}__labels_from_{models}.csv. Running it again reuses the predictions which are up
+to date, so only new or changed analyzers are labelled. With --compute_missing_metrics, the metrics the models need
+which an analyzer doesn't have (e.g. the drift metrics) are computed first, and saved in the analyzer.
+
+With --task_index, only one of the analyzers is labelled, e.g. one per job of a SLURM job array
+(--task_index $SLURM_ARRAY_TASK_ID); --list shows the analyzers and their numbers. The csv files are named
+from the whole list, so the jobs' files don't clash.
+"""
+
+import argparse
+import os
+import shutil
+import sys
+from pathlib import Path
+
+from unitrefine import __version__
+from unitrefine.models_hub import downloaded_repo_id, download_model, hfh_model_groups, looks_like_repo_id
+from unitrefine.predictions import (
+    analyzer_output_names,
+    apply_to_analyzers,
+    find_analyzers,
+    get_model_classes,
+    is_layer1_model,
+    is_layer2_model,
+    is_remote_path,
+    selection_key,
+)
+
+# where models downloaded from the Hugging Face Hub are kept, unless --models_folder is given
+MODELS_FOLDER_VARIABLE = "UNITREFINE_MODELS_FOLDER"
+DEFAULT_MODELS_FOLDER = Path.home() / ".cache" / "unitrefine" / "models"
+
+
+class UsageError(ValueError):
+    """A problem with what was asked for, e.g. a model which doesn't exist: shown as a message, not a traceback."""
+
+
+def default_models_folder():
+    return Path(os.environ.get(MODELS_FOLDER_VARIABLE, DEFAULT_MODELS_FOLDER))
+
+
+def collect_analyzers(sources, max_depth=3):
+    """
+    The analyzers in `sources`: analyzer folders, folders with analyzers inside (up to `max_depth` levels down),
+    text files with one path per line (blank lines and lines starting with # are skipped), or remote (s3) paths.
+    """
+    analyzers = []
+    for source in sources:
+        source = str(source)
+        if is_remote_path(source):
+            found = [source]
+        elif Path(source).is_file():
+            lines = [line.strip() for line in Path(source).read_text().splitlines()]
+            found = collect_analyzers([line for line in lines if line != "" and not line.startswith("#")], max_depth)
+        elif Path(source).is_dir():
+            found = find_analyzers(source, max_depth)
+            if len(found) == 0:
+                raise UsageError(f"No analyzers found in {source} (searched {max_depth} folder levels down).")
+        else:
+            raise UsageError(f"{source} does not exist.")
+        for path in found:
+            if str(path) not in [str(analyzer) for analyzer in analyzers]:
+                analyzers.append(path)
+    return analyzers
+
+
+def resolve_model(model, models_folder=None, download=True):
+    """
+    The folder of `model`: a model folder, or a Hugging Face repo id, which is downloaded into `models_folder`
+    the first time. Several jobs can download the same model at once: each downloads it to its own folder, and
+    the first one to finish is kept.
+    """
+    if Path(model).is_dir():
+        model_folder = Path(model)
+        if not (model_folder / "best_model.skops").is_file():
+            raise UsageError(f"{model_folder} is not a UnitRefine model: it has no best_model.skops.")
+        return model_folder
+    if not looks_like_repo_id(model):
+        raise UsageError(f"The model {model} was not found: give a model folder, or a Hugging Face repo id "
+                         f"such as SpikeInterface/UnitRefine_sua_mua_classifier.")
+
+    models_folder = Path(models_folder) if models_folder is not None else default_models_folder()
+    model_folder = models_folder / model.split("/")[-1]
+    if model_folder.is_dir():
+        if downloaded_repo_id(model_folder) != model:
+            raise UsageError(f"{model_folder} exists, but was not downloaded from {model}. Give another --models_folder.")
+        return model_folder
+    if not download:
+        raise UsageError(f"{model} has not been downloaded into {models_folder}.")
+
+    models_folder.mkdir(parents=True, exist_ok=True)
+    temporary_folder = models_folder / f".{model_folder.name}.download-{os.getpid()}"
+    try:
+        download_model(model, temporary_folder)
+    except Exception as e:
+        raise UsageError(
+            f"Could not download {model} from the Hugging Face Hub: {e}\nIf this computer has no internet (e.g. a "
+            f"compute node), download the model first where there is, with: unitrefine-download-models {model} "
+            f"--models_folder {models_folder}"
+        ) from e
+    try:
+        temporary_folder.rename(model_folder)
+    except OSError:
+        if not model_folder.is_dir():
+            raise
+        # another job downloaded it at the same time
+        shutil.rmtree(temporary_folder, ignore_errors=True)
+    print(f"Saved {model} in {model_folder}.")
+    return model_folder
+
+
+def model_selection(model_folders, cascade):
+    """The mode ("binary" or "cascade") and the model folders in the order they are applied."""
+    if not cascade:
+        return "binary", list(model_folders)
+    if len(model_folders) != 2:
+        raise UsageError("A cascade needs two models: a noise vs neural model and a SUA vs MUA model.")
+    first, second = model_folders
+    if is_layer1_model(get_model_classes(first)) and is_layer2_model(get_model_classes(second)):
+        return "cascade", [first, second]
+    if is_layer1_model(get_model_classes(second)) and is_layer2_model(get_model_classes(first)):
+        return "cascade", [second, first]
+    raise UsageError(
+        f"A cascade needs a noise vs neural model and a SUA vs MUA model, but {Path(first).name} predicts "
+        f"{get_model_classes(first)} and {Path(second).name} predicts {get_model_classes(second)}."
+    )
+
+
+def apply_models(analyzers, models, cascade=False, output_folder="unitrefine_predictions", models_folder=None,
+                 task_index=None, max_depth=3, compute_missing_metrics=False, n_jobs=1):
+    """
+    Applies `models` (model folders or Hugging Face repo ids) to `analyzers` (see collect_analyzers), each on its
+    own or as a cascade, and saves one predictions csv per analyzer in `output_folder`. With `task_index`, only
+    that analyzer (counting from 0) is labelled. With `compute_missing_metrics`, the metrics the models need which
+    an analyzer doesn't have are computed first (with `n_jobs` processes) and saved in the analyzer. Returns a list of dicts with the analyzer's `path` and the saved
+    `file`, or the `error` for analyzers which could not be labelled, e.g. because a metric is missing.
+    """
+    analyzer_paths = collect_analyzers(analyzers, max_depth)
+    if len(analyzer_paths) == 0:
+        raise UsageError("No analyzers were given.")
+    # the names come from the whole list, so each job of an array gets a different file name
+    output_names = analyzer_output_names(analyzer_paths)
+    if task_index is not None:
+        if not 0 <= task_index < len(analyzer_paths):
+            raise UsageError(f"--task_index {task_index} is out of range: there are {len(analyzer_paths)} analyzers, "
+                             f"numbered 0 to {len(analyzer_paths) - 1}.")
+        analyzer_paths, output_names = [analyzer_paths[task_index]], [output_names[task_index]]
+
+    model_folders = [resolve_model(model, models_folder) for model in models]
+    mode, model_folders = model_selection(model_folders, cascade)
+    key = selection_key(mode, [model_folder.name for model_folder in model_folders])
+    return apply_to_analyzers(analyzer_paths, mode, model_folders, key, output_folder, output_names=output_names,
+                              compute_missing_metrics=compute_missing_metrics, job_kwargs=dict(n_jobs=n_jobs))
+
+
+def apply_main(argv=None):
+    parser = argparse.ArgumentParser(
+        prog="unitrefine-apply",
+        description="Applies UnitRefine models to Sorting Analyzers and saves their predictions, one csv per analyzer.",
+        epilog="Example: unitrefine-apply /data/sortings --model SpikeInterface/UnitRefine_noise_neural_classifier "
+               "--model SpikeInterface/UnitRefine_sua_mua_classifier --cascade --output predictions",
+    )
+    parser.add_argument("analyzers", nargs="+",
+                        help="analyzer folders, folders with analyzers inside, or text files with one analyzer per line")
+    parser.add_argument("--version", action="version", version=f"unitrefine {__version__}")
+    parser.add_argument("--model", action="append", default=[], dest="models",
+                        help="a model folder or a Hugging Face repo id; give --model again for more models")
+    parser.add_argument("--cascade", action="store_true",
+                        help="apply the two models as a cascade: noise vs neural, then SUA vs MUA on the neural units")
+    parser.add_argument("--output", default="unitrefine_predictions", help="folder to save the predictions in")
+    parser.add_argument("--models_folder", default=None,
+                        help=f"where models from the Hugging Face Hub are kept (default: ${MODELS_FOLDER_VARIABLE} "
+                             f"or {DEFAULT_MODELS_FOLDER})")
+    parser.add_argument("--task_index", type=int, default=None,
+                        help="only label this analyzer, counting from 0, e.g. $SLURM_ARRAY_TASK_ID")
+    parser.add_argument("--compute_missing_metrics", action="store_true",
+                        help="compute the metrics the models need which an analyzer doesn't have, and save them in "
+                             "the analyzer (needs its recording)")
+    parser.add_argument("--n_jobs", type=int, default=1,
+                        help="processes used to compute metrics, e.g. $SLURM_CPUS_PER_TASK")
+    parser.add_argument("--max_depth", type=int, default=3, help="how many folder levels down to look for analyzers")
+    parser.add_argument("--list", action="store_true", help="only list the analyzers and their numbers")
+    args = parser.parse_args(argv)
+
+    try:
+        if args.list:
+            analyzer_paths = collect_analyzers(args.analyzers, args.max_depth)
+            for index, path in enumerate(analyzer_paths):
+                print(f"{index}\t{path}")
+            print(f"{len(analyzer_paths)} analyzers. For a SLURM job array: --array=0-{len(analyzer_paths) - 1}")
+            return 0
+        if len(args.models) == 0:
+            parser.error("give at least one --model")
+        results = apply_models(args.analyzers, args.models, cascade=args.cascade, output_folder=args.output,
+                               models_folder=args.models_folder, task_index=args.task_index, max_depth=args.max_depth,
+                               compute_missing_metrics=args.compute_missing_metrics, n_jobs=args.n_jobs)
+    except UsageError as e:
+        print(f"unitrefine-apply: {e}", file=sys.stderr)
+        return 2
+
+    failed = [result for result in results if "error" in result]
+    print(f"\nLabelled {len(results) - len(failed)}/{len(results)} analyzers. Predictions saved in {Path(args.output).resolve()}.")
+    for result in failed:
+        print(f"Not labelled: {result['path']}: {result['error']}", file=sys.stderr)
+    if any("required metrics" in result["error"] for result in failed) and not args.compute_missing_metrics:
+        print("To compute the missing metrics and save them in the analyzers, add --compute_missing_metrics.", file=sys.stderr)
+    return 1 if len(failed) > 0 else 0
+
+
+def download_main(argv=None):
+    parser = argparse.ArgumentParser(
+        prog="unitrefine-download-models",
+        description="Downloads UnitRefine models from the Hugging Face Hub, e.g. before running unitrefine-apply on "
+                    "compute nodes without internet. Without repo ids, lists the published UnitRefine models.",
+    )
+    parser.add_argument("repo_ids", nargs="*", help="Hugging Face repo ids, e.g. SpikeInterface/UnitRefine_sua_mua_classifier")
+    parser.add_argument("--models_folder", default=None,
+                        help=f"where to save them (default: ${MODELS_FOLDER_VARIABLE} or {DEFAULT_MODELS_FOLDER})")
+    args = parser.parse_args(argv)
+
+    if len(args.repo_ids) == 0:
+        print("UnitRefine models on the Hugging Face Hub:")
+        for group in hfh_model_groups:
+            for repo_id, note in group["models"]:
+                print(f"  {repo_id}  {note}".rstrip())
+        print("\nDownload them with: unitrefine-download-models REPO_ID [REPO_ID ...]")
+        return 0
+
+    status = 0
+    for repo_id in args.repo_ids:
+        try:
+            model_folder = resolve_model(repo_id, args.models_folder)
+            print(f"{repo_id}: {model_folder}")
+        except UsageError as e:
+            print(f"unitrefine-download-models: {e}", file=sys.stderr)
+            status = 2
+    return status
+
+
+if __name__ == "__main__":
+    sys.exit(apply_main())