improve saved models

Author: ceesem

models/model_ds15_us0_bd0.json

@@ -0,0 +1 @@
+{"downstream_hops": 15, "upstream_hops": 0, "bidirectional_hops": 0, "feature_columns": ["area_um2", "vol_um3", "max_dt_um", "vol_to_area", "syn_in", "syn_out", "down_area_um2", "down_vol_um3", "down_max_dt_um", "down_vol_to_area", "down_syn_in", "down_syn_out"], "spread_alpha": 0.5, "model_file": "xgb/model_ds15_us0_bd0.ubj"}

models/xgb/model_ds15_us0_bd0.ubj

@@ -9,7 +9,6 @@ description = "Using level-2 features for skeleton classification"
 readme = "README.md"
 requires-python = ">=3.9"
 dependencies = [
-    "joblib>=1.4.2",
     "meshparty>=1.18.2",
     "numpy>=2.0.2",
     "pandas>=2.2.3",
@@ -103,3 +102,12 @@ help = "Profile cpu and memory of task with scalene"
 [tool.poe.tasks.profile]
 cmd = "uv run pyinstrument -r html"
 help = "Profile cpu of task with pyinstrument"
+
+[tool.hatch.build.targets.sdist]
+only-include = ["src/l2label", "models"]
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/l2label"]
+
+[tool.hatch.build.targets.wheel.force-include]
+"models" = "l2label/models"

pyproject.toml

@@ -2,13 +2,14 @@
 import pathlib
 from typing import Literal, Optional, Union
 
-import joblib
 import numpy as np
 import pandas as pd
 import xgboost
 from meshparty import meshwork
 from scipy import sparse
 
+from . import models
+
 __all__ = [
     "make_skel_prop_df",
     "process_neuron",
@@ -312,6 +313,8 @@ def save_model(
         modelpath
         / f"{model_name}_ds{downstream_hops}_us{upstream_hops}_bd{bidirectional_hops}.joblib"
     )
+    model.save_model(model_filename)
+
     outdata = {
         "downstream_hops": downstream_hops,
         "upstream_hops": upstream_hops,
@@ -320,7 +323,6 @@ def save_model(
         "spread_alpha": spread_alpha,
         "model_file": str(model_filename.absolute()),
     }
-    joblib.dump(model, outdata["model_file"])
     description_file = (
         filepath
         / f"{model_name}_ds{downstream_hops}_us{upstream_hops}_bd{bidirectional_hops}.json"
@@ -343,13 +345,28 @@ def __init__(
     ):
         if config is not None:
             if not isinstance(config, dict):
-                with open(config, "r") as f:
-                    config = json.load(f)
+                config = models.load_model_config(config, dir="")
+            self._downstream_hops = config.get("downstream_hops")
+            self._upstream_hops = config.get("upstream_hops")
+            self._bidirectional_hops = config.get("bidirectional_hops")
+            self._spread_alpha = config.get("spread_alpha")
+            self._load_model(config.get("model_file"))
+            self._feature_columns = config.get("feature_columns")
+
+        elif (
+            config is None
+            and downstream_hops is None
+            and upstream_hops is None
+            and bidirectional_hops is None
+            and model is None
+            and spread_alpha is None
+        ):
+            config = models.load_model_config()
             self._downstream_hops = config.get("downstream_hops")
             self._upstream_hops = config.get("upstream_hops")
             self._bidirectional_hops = config.get("bidirectional_hops")
             self._spread_alpha = config.get("spread_alpha")
-            self._model = joblib.load(config.get("model_file"))
+            self._load_model(config.get("model_file"))
             self._feature_columns = config.get("feature_columns")
         else:
             if (
@@ -603,6 +620,7 @@ def predict_axon_mask(
         root_is_soma: bool = False,
         is_axon_seg: Optional[np.ndarray] = None,
         evaluate_isolated_dendrites: bool = False,
+        to_mesh_index: bool = False,
     ):
         """Predict a dendrite mask based on the smoothed label spreading and segment-level vote.
 
@@ -647,4 +665,13 @@ def predict_axon_mask(
             in_root_comp = comp_labels == comp_labels[nrn.skeleton.root]
         else:
             in_root_comp = np.full(len(nrn.skeleton.vertices), True)
-        return ~np.logical_and(~is_axon_base, in_root_comp)
+        is_axon = ~np.logical_and(~is_axon_base, in_root_comp)
+        if to_mesh_index:
+            is_axon_sk = nrn.SkeletonIndex(np.flatnonzero(is_axon))
+            is_axon = is_axon_sk.to_mesh_mask
+        return is_axon
+
+    def _load_model(self, filepath: str):
+        """Load a model from a file."""
+        self._model = xgboost.XGBClassifier()
+        self._model.load_model(filepath)

src/l2label/compartments.py

@@ -0,0 +1,52 @@
+import glob
+import json
+import pathlib
+from typing import Optional
+
+import xgboost
+
+model_dir = pathlib.Path(__file__).parents[2] / "models"
+
+current_model = "model_ds15_us0_bd0.json"
+
+
+def get_models(
+    dir: Optional[str] = None,
+):
+    """
+    Get a list of all models in the model directory.
+    """
+    if dir is None:
+        dir = model_dir
+
+    return [f.name for f in glob.glob(f"{dir}/*.json")]
+
+
+def load_model_config(
+    model_name: Optional[str] = None,
+    dir: Optional[str] = None,
+):
+    """
+    Load a model configuration file.
+
+    Parameters
+    ----------
+    model_name : str
+        The filename name of the model to load.
+    dir : str, optional
+        The directory to load the model from. If None, the default model directory is used.
+    """
+    if dir is None:
+        dir = model_dir
+        if model_name is None:
+            model_name = current_model
+
+    if model_name is None:
+        raise ValueError("model_name must be specified if dir is not provided")
+    if not model_name.endswith(".json"):
+        model_name += ".json"
+    with open(pathlib.Path(dir) / f"{model_name}", "r") as f:
+        model_config = json.load(f)
+
+    model_config["model_file"] = pathlib.Path(dir) / f"{model_config['model_file']}"
+    return model_config