From 3ce56c75c99fc3871145c2404dfbbb1bd69df02a Mon Sep 17 00:00:00 2001
From: Delattre Emilie <emilie.delattre@epfl.ch>
Date: Wed, 5 Apr 2023 10:53:54 +0200
Subject: [PATCH 01/12] Start adding a new command linde

---
 .../console/cmd_extract_features_and_train.py | 134 ++++++++++++++++++
 src/morphoclass/console/cmd_train.py          |   9 +-
 src/morphoclass/console/main.py               |   2 +
 3 files changed, 143 insertions(+), 2 deletions(-)
 create mode 100644 src/morphoclass/console/cmd_extract_features_and_train.py

diff --git a/src/morphoclass/console/cmd_extract_features_and_train.py b/src/morphoclass/console/cmd_extract_features_and_train.py
new file mode 100644
index 0000000..ccdd5f3
--- /dev/null
+++ b/src/morphoclass/console/cmd_extract_features_and_train.py
@@ -0,0 +1,134 @@
+# Copyright © 2022-2022 Blue Brain Project/EPFL
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Implementation of the `morphoclass train` CLI command."""
+from __future__ import annotations
+
+import logging
+import pathlib
+
+import click
+import yaml
+
+from morphoclass.types import StrPath
+
+logger = logging.getLogger(__name__)
+
+
+@click.command(name="train", help="Train a morphology classification model.")
+@click.argument("csv_path", type=click.Path(dir_okay=False))
+@click.argument("neurite_type", type=click.Choice(["apical", "axon", "basal", "all"]))
+@click.argument(
+    "feature",
+    type=click.Choice(
+        [
+            "graph-rd",
+            "graph-proj",
+            "diagram-tmd-rd",
+            "diagram-tmd-proj",
+            "diagram-deepwalk",
+            "image-tmd-rd",
+            "image-tmd-proj",
+            "image-deepwalk",
+        ]
+    ),
+)
+@click.option(
+    "--orient",
+    is_flag=True,
+    help="Orient the neurons so that the apicals are aligned with the positive y-axis.",
+)
+@click.option(
+    "--no-simplify-graph",
+    is_flag=True,
+    help="""
+    By default the neurite graph is reduced to branching nodes only. With this
+    flag the full neurite graph will be preserved.
+    """,
+)
+@click.option(
+    "--keep-diagram",
+    is_flag=True,
+    help="After converting the diagram to persistence image don't discard the diagram.",
+)
+@click.option(
+    "--model-config",
+    type=click.Path(exists=True, dir_okay=False),
+    required=True,
+    help="The model configuration file.",
+)
+@click.option(
+    "--splitter-config",
+    type=click.Path(exists=True, dir_okay=False),
+    required=True,
+    help="The splitter configuration file.",
+)
+@click.option(
+    "--output-dir",
+    type=click.Path(file_okay=False),
+    required=True,
+    help="The output directory.",
+)
+@click.option(
+    "-f",
+    "--force",
+    type=click.BOOL,
+    default=False,
+    is_flag=True,
+    help="Don't ask for overwriting existing output files.",
+)
+def cli(
+    csv_path: StrPath,
+    neurite_type: Literal["apical", "axon", "basal", "all"],
+    feature: Literal[
+        "graph-rd",
+        "graph-proj",
+        "diagram-tmd-rd",
+        "diagram-tmd-proj",
+        "diagram-deepwalk",
+        "image-tmd-rd",
+        "image-tmd-proj",
+        "image-deepwalk",
+    ],
+    orient: bool,
+    no_simplify_graph: bool,
+    keep_diagram: bool,
+    model_config: StrPath,
+    splitter_config: StrPath,
+    output_dir: StrPath,
+    force: bool,
+) -> None:
+    """Extract features and train the model."""
+    from morphoclass.console.cmd_extract_features import extract_features
+    from morphoclass.console.cmd_train import train
+
+    input_csv = pathlib.Path(csv_path).resolve()
+    output_dir = pathlib.Path(output_dir).resolve()
+
+    extract_features(
+        input_csv,
+        neurite_type[0],
+        feature[0],
+        output_dir / "features",
+        orient,
+        no_simplify_graph,
+        keep_diagram,
+        force,
+    )
+
+    train(
+        output_dir / "features",
+        model_config,
+        splitter_config,
+        output_dir / "checkpoints",
+    )
diff --git a/src/morphoclass/console/cmd_train.py b/src/morphoclass/console/cmd_train.py
index 72d9ee0..535e2b6 100644
--- a/src/morphoclass/console/cmd_train.py
+++ b/src/morphoclass/console/cmd_train.py
@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Implementation of the `morphoclass train` CLI command."""
+"""Implementation of the `morphoclass train-after-extraction` CLI command."""
 from __future__ import annotations
 
 import logging
@@ -25,7 +25,12 @@
 logger = logging.getLogger(__name__)
 
 
-@click.command(name="train", help="Train a morphology classification model.")
+@click.command(
+    name="train-after-extraction",
+    help=
+    """Train a morphology classification model.
+    Features need to be first extracted."""
+)
 @click.option(
     "--features-dir",
     type=click.Path(exists=True, file_okay=False),
diff --git a/src/morphoclass/console/main.py b/src/morphoclass/console/main.py
index 886d2f8..3566c7a 100644
--- a/src/morphoclass/console/main.py
+++ b/src/morphoclass/console/main.py
@@ -22,6 +22,7 @@
 import morphoclass
 from morphoclass.console import cmd_evaluate
 from morphoclass.console import cmd_extract_features
+from morphoclass.console import cmd_extract_features_and_train
 from morphoclass.console import cmd_morphometrics
 from morphoclass.console import cmd_organise_dataset
 from morphoclass.console import cmd_performance_table
@@ -143,3 +144,4 @@ def cli(verbose: int, log_file_path: pathlib.Path | None) -> None:
 cli.add_command(cmd_performance_table.cli)
 cli.add_command(cmd_extract_features.cli)
 cli.add_command(cmd_morphometrics.cli)
+cli.add_command(cmd_extract_features_and_train.cli)

From c26a95f0189df5f199c4377307180e66a200cc63 Mon Sep 17 00:00:00 2001
From: Delattre Emilie <emilie.delattre@epfl.ch>
Date: Wed, 5 Apr 2023 11:04:07 +0200
Subject: [PATCH 02/12] Make linting happy

---
 src/morphoclass/console/cmd_extract_features_and_train.py | 2 +-
 src/morphoclass/console/cmd_train.py                      | 4 ++--
 src/morphoclass/data/tns_dataset.py                       | 4 +---
 src/morphoclass/vis.py                                    | 2 +-
 tests/unit/test_metrics.py                                | 2 +-
 5 files changed, 6 insertions(+), 8 deletions(-)

diff --git a/src/morphoclass/console/cmd_extract_features_and_train.py b/src/morphoclass/console/cmd_extract_features_and_train.py
index ccdd5f3..c7197b4 100644
--- a/src/morphoclass/console/cmd_extract_features_and_train.py
+++ b/src/morphoclass/console/cmd_extract_features_and_train.py
@@ -16,9 +16,9 @@
 
 import logging
 import pathlib
+from typing import Literal
 
 import click
-import yaml
 
 from morphoclass.types import StrPath
 
diff --git a/src/morphoclass/console/cmd_train.py b/src/morphoclass/console/cmd_train.py
index 535e2b6..3f1c031 100644
--- a/src/morphoclass/console/cmd_train.py
+++ b/src/morphoclass/console/cmd_train.py
@@ -27,8 +27,8 @@
 
 @click.command(
     name="train-after-extraction",
-    help=
-    """Train a morphology classification model.
+    help="""
+    Train a morphology classification model.
     Features need to be first extracted."""
 )
 @click.option(
diff --git a/src/morphoclass/data/tns_dataset.py b/src/morphoclass/data/tns_dataset.py
index 6be3bb2..c5a190c 100644
--- a/src/morphoclass/data/tns_dataset.py
+++ b/src/morphoclass/data/tns_dataset.py
@@ -135,9 +135,7 @@ def __init__(
                 f"No data corresponding to layer {layer} found in data_path"
             )
         else:
-            self.class_dict = {
-                n: m_type for n, m_type in enumerate(sorted(self.m_types))
-            }
+            self.class_dict = dict(enumerate(sorted(self.m_types)))
             self.class_dict_inv = {v: k for k, v in self.class_dict.items()}
 
         self.distributions = input_distributions
diff --git a/src/morphoclass/vis.py b/src/morphoclass/vis.py
index 008afef..2b7e45c 100644
--- a/src/morphoclass/vis.py
+++ b/src/morphoclass/vis.py
@@ -820,7 +820,7 @@ def plot_neurite(
     nx.draw(
         g,
         ax=ax,
-        pos={n: xy for n, xy in enumerate(zip(px, py))},
+        pos=dict(enumerate(zip(px, py))),
         nodelist=[0],
         node_color="red",
         node_size=soma_size,
diff --git a/tests/unit/test_metrics.py b/tests/unit/test_metrics.py
index 06043b4..dfb987d 100644
--- a/tests/unit/test_metrics.py
+++ b/tests/unit/test_metrics.py
@@ -158,7 +158,7 @@ def test_inter_rater_score(targets, predictions, kind, score):
 
 
 def test_inter_rater_score_fail(targets, predictions):
-    with pytest.raises(Exception):
+    with pytest.raises(ValueError):
         morphoclass.metrics.inter_rater_score(targets, predictions, kind="invalid")
 
 

From e8aeacc71a67480ca71ef8678c867d58dc979244 Mon Sep 17 00:00:00 2001
From: Delattre Emilie <emilie.delattre@epfl.ch>
Date: Wed, 5 Apr 2023 11:06:29 +0200
Subject: [PATCH 03/12] Make black happy

---
 src/morphoclass/console/cmd_train.py                            | 2 +-
 src/morphoclass/models/concatenet.py                            | 1 -
 .../augmentors/add_random_points_to_reduction_mask.py           | 1 -
 3 files changed, 1 insertion(+), 3 deletions(-)

diff --git a/src/morphoclass/console/cmd_train.py b/src/morphoclass/console/cmd_train.py
index 3f1c031..a20d451 100644
--- a/src/morphoclass/console/cmd_train.py
+++ b/src/morphoclass/console/cmd_train.py
@@ -29,7 +29,7 @@
     name="train-after-extraction",
     help="""
     Train a morphology classification model.
-    Features need to be first extracted."""
+    Features need to be first extracted.""",
 )
 @click.option(
     "--features-dir",
diff --git a/src/morphoclass/models/concatenet.py b/src/morphoclass/models/concatenet.py
index 6ab610b..9953f16 100644
--- a/src/morphoclass/models/concatenet.py
+++ b/src/morphoclass/models/concatenet.py
@@ -45,7 +45,6 @@ class ConcateNet(nn.Module):
     """
 
     def __init__(self, n_node_features, n_classes, n_features_perslay, bn=False):
-
         super().__init__()
         self.n_node_features = n_node_features
         self.n_classes = n_classes
diff --git a/src/morphoclass/transforms/augmentors/add_random_points_to_reduction_mask.py b/src/morphoclass/transforms/augmentors/add_random_points_to_reduction_mask.py
index 54178cd..8279c1e 100644
--- a/src/morphoclass/transforms/augmentors/add_random_points_to_reduction_mask.py
+++ b/src/morphoclass/transforms/augmentors/add_random_points_to_reduction_mask.py
@@ -32,7 +32,6 @@ class AddRandomPointsToReductionMask:
     """
 
     def __init__(self, n_points):
-
         self.n_points = n_points
 
     @require_field("tmd_neurites_masks")

From 216d89e26b57ec337c110ee3851fc86cf1820b7a Mon Sep 17 00:00:00 2001
From: Delattre Emilie <emilie.delattre@epfl.ch>
Date: Wed, 5 Apr 2023 11:07:01 +0200
Subject: [PATCH 04/12] Make apidoc happy

---
 .../morphoclass.console.cmd_extract_features_and_train.rst | 7 +++++++
 docs/source/api/morphoclass.console.rst                    | 1 +
 2 files changed, 8 insertions(+)
 create mode 100644 docs/source/api/morphoclass.console.cmd_extract_features_and_train.rst

diff --git a/docs/source/api/morphoclass.console.cmd_extract_features_and_train.rst b/docs/source/api/morphoclass.console.cmd_extract_features_and_train.rst
new file mode 100644
index 0000000..f41e304
--- /dev/null
+++ b/docs/source/api/morphoclass.console.cmd_extract_features_and_train.rst
@@ -0,0 +1,7 @@
+morphoclass.console.cmd\_extract\_features\_and\_train module
+=============================================================
+
+.. automodule:: morphoclass.console.cmd_extract_features_and_train
+   :members:
+   :undoc-members:
+   :show-inheritance:
diff --git a/docs/source/api/morphoclass.console.rst b/docs/source/api/morphoclass.console.rst
index 67415de..ab3970c 100644
--- a/docs/source/api/morphoclass.console.rst
+++ b/docs/source/api/morphoclass.console.rst
@@ -9,6 +9,7 @@ Submodules
 
    morphoclass.console.cmd_evaluate
    morphoclass.console.cmd_extract_features
+   morphoclass.console.cmd_extract_features_and_train
    morphoclass.console.cmd_morphometrics
    morphoclass.console.cmd_organise_dataset
    morphoclass.console.cmd_performance_table

From d4d851aa0e44bf511ba2cf353e4edfd6fe9fb94f Mon Sep 17 00:00:00 2001
From: Delattre Emilie <emilie.delattre@epfl.ch>
Date: Wed, 5 Apr 2023 11:11:56 +0200
Subject: [PATCH 05/12] Try to make mypy happy

---
 .../console/cmd_extract_features.py           | 22 +++++++++++++++++++
 src/morphoclass/console/cmd_train.py          | 16 ++++++++++++++
 2 files changed, 38 insertions(+)

diff --git a/src/morphoclass/console/cmd_extract_features.py b/src/morphoclass/console/cmd_extract_features.py
index c952be5..7a76584 100644
--- a/src/morphoclass/console/cmd_extract_features.py
+++ b/src/morphoclass/console/cmd_extract_features.py
@@ -114,6 +114,28 @@ def cli(
     keep_diagram: bool,
     force: bool,
 ) -> None:
+    """Extract morphology features."""
+    return extract_features(
+        csv_path,
+        neurite_type,
+        feature,
+        output_dir,
+        orient,
+        no_simplify_graph,
+        keep_diagram,
+        force,
+    )
+
+def extract_features(
+    csv_path: StrPath,
+    neurite_type: str,
+    feature: str,
+    output_dir: StrPath,
+    orient: bool,
+    no_simplify_graph: bool,
+    keep_diagram: bool,
+    force: bool,
+):
     """Extract morphology features."""
     output_dir = pathlib.Path(output_dir)
     if output_dir.exists() and not force:
diff --git a/src/morphoclass/console/cmd_train.py b/src/morphoclass/console/cmd_train.py
index a20d451..35e0d9b 100644
--- a/src/morphoclass/console/cmd_train.py
+++ b/src/morphoclass/console/cmd_train.py
@@ -69,6 +69,22 @@ def cli(
     splitter_config: StrPath,
     checkpoint_dir: StrPath,
     force: bool,
+) -> None:
+    """Training and evaluation of the model."""
+    return train(
+        features_dir,
+        model_config,
+        splitter_config,
+        checkpoint_dir,
+        force,
+    )
+
+def train(
+    features_dir: StrPath,
+    model_config: StrPath,
+    splitter_config: StrPath,
+    checkpoint_dir: StrPath,
+    force: bool,
 ) -> None:
     """Training and evaluation of the model.
 

From cfb176064f94016be0d90fe0fb7610b1d9749c0c Mon Sep 17 00:00:00 2001
From: Delattre Emilie <emilie.delattre@epfl.ch>
Date: Wed, 5 Apr 2023 13:42:05 +0200
Subject: [PATCH 06/12] Try to make lint happt

---
 src/morphoclass/console/cmd_extract_features.py | 1 +
 src/morphoclass/console/cmd_train.py            | 1 +
 2 files changed, 2 insertions(+)

diff --git a/src/morphoclass/console/cmd_extract_features.py b/src/morphoclass/console/cmd_extract_features.py
index 7a76584..6b4ec3d 100644
--- a/src/morphoclass/console/cmd_extract_features.py
+++ b/src/morphoclass/console/cmd_extract_features.py
@@ -126,6 +126,7 @@ def cli(
         force,
     )
 
+
 def extract_features(
     csv_path: StrPath,
     neurite_type: str,
diff --git a/src/morphoclass/console/cmd_train.py b/src/morphoclass/console/cmd_train.py
index 35e0d9b..56edb3d 100644
--- a/src/morphoclass/console/cmd_train.py
+++ b/src/morphoclass/console/cmd_train.py
@@ -79,6 +79,7 @@ def cli(
         force,
     )
 
+
 def train(
     features_dir: StrPath,
     model_config: StrPath,

From 54ac4f3bfb5e23830e6039bf22b1b8a23653a6ad Mon Sep 17 00:00:00 2001
From: Delattre Emilie <emilie.delattre@epfl.ch>
Date: Wed, 5 Apr 2023 13:43:02 +0200
Subject: [PATCH 07/12] Try to make mypy happy

---
 src/morphoclass/console/cmd_extract_features.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/morphoclass/console/cmd_extract_features.py b/src/morphoclass/console/cmd_extract_features.py
index 6b4ec3d..93eda5e 100644
--- a/src/morphoclass/console/cmd_extract_features.py
+++ b/src/morphoclass/console/cmd_extract_features.py
@@ -136,7 +136,7 @@ def extract_features(
     no_simplify_graph: bool,
     keep_diagram: bool,
     force: bool,
-):
+) -> None:
     """Extract morphology features."""
     output_dir = pathlib.Path(output_dir)
     if output_dir.exists() and not force:

From 03da543589466266c4153c344dc38a555eed512e Mon Sep 17 00:00:00 2001
From: Delattre Emilie <emilie.delattre@epfl.ch>
Date: Wed, 5 Apr 2023 13:51:08 +0200
Subject: [PATCH 08/12] Try to make mypy happy (2)

---
 src/morphoclass/console/cmd_extract_features_and_train.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/morphoclass/console/cmd_extract_features_and_train.py b/src/morphoclass/console/cmd_extract_features_and_train.py
index c7197b4..2ab39c1 100644
--- a/src/morphoclass/console/cmd_extract_features_and_train.py
+++ b/src/morphoclass/console/cmd_extract_features_and_train.py
@@ -131,4 +131,5 @@ def cli(
         model_config,
         splitter_config,
         output_dir / "checkpoints",
+        force,
     )

From 9c87dde7d9d39927072ba10a1ba8f7d855486e63 Mon Sep 17 00:00:00 2001
From: Delattre Emilie <emilie.delattre@epfl.ch>
Date: Wed, 26 Apr 2023 12:48:57 +0200
Subject: [PATCH 09/12] Fix features and neurites

---
 src/morphoclass/console/cmd_extract_features_and_train.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/morphoclass/console/cmd_extract_features_and_train.py b/src/morphoclass/console/cmd_extract_features_and_train.py
index 2ab39c1..ad7c7db 100644
--- a/src/morphoclass/console/cmd_extract_features_and_train.py
+++ b/src/morphoclass/console/cmd_extract_features_and_train.py
@@ -117,8 +117,8 @@ def cli(
 
     extract_features(
         input_csv,
-        neurite_type[0],
-        feature[0],
+        neurite_type,
+        feature,
         output_dir / "features",
         orient,
         no_simplify_graph,

From 76171f46fc953fc6bd6d0318764bdbfe3ded63e0 Mon Sep 17 00:00:00 2001
From: Delattre Emilie <emilie.delattre@epfl.ch>
Date: Wed, 26 Apr 2023 12:51:00 +0200
Subject: [PATCH 10/12] Add info about model-config and splitter-config

---
 .../console/cmd_extract_features_and_train.py        | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/src/morphoclass/console/cmd_extract_features_and_train.py b/src/morphoclass/console/cmd_extract_features_and_train.py
index ad7c7db..8ca87b8 100644
--- a/src/morphoclass/console/cmd_extract_features_and_train.py
+++ b/src/morphoclass/console/cmd_extract_features_and_train.py
@@ -65,13 +65,21 @@
     "--model-config",
     type=click.Path(exists=True, dir_okay=False),
     required=True,
-    help="The model configuration file.",
+    help="""
+    The model configuration file.
+    For inspiration, model configuration files can be found under 
+    dvc/training/configs/
+    """,
 )
 @click.option(
     "--splitter-config",
     type=click.Path(exists=True, dir_okay=False),
     required=True,
-    help="The splitter configuration file.",
+    help="""
+    The splitter configuration file.
+    For inspiration, splitter configuration files can be found under 
+    dvc/training/configs/
+    """,
 )
 @click.option(
     "--output-dir",

From 6d1398cf8c2ae8a48d2c3a16a828eea5606e72a5 Mon Sep 17 00:00:00 2001
From: Delattre Emilie <emilie.delattre@epfl.ch>
Date: Wed, 26 Apr 2023 13:35:46 +0200
Subject: [PATCH 11/12] Fix linting

---
 src/morphoclass/console/cmd_extract_features_and_train.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/morphoclass/console/cmd_extract_features_and_train.py b/src/morphoclass/console/cmd_extract_features_and_train.py
index 8ca87b8..dfa4338 100644
--- a/src/morphoclass/console/cmd_extract_features_and_train.py
+++ b/src/morphoclass/console/cmd_extract_features_and_train.py
@@ -67,7 +67,7 @@
     required=True,
     help="""
     The model configuration file.
-    For inspiration, model configuration files can be found under 
+    For inspiration, model configuration files can be found under
     dvc/training/configs/
     """,
 )
@@ -77,7 +77,7 @@
     required=True,
     help="""
     The splitter configuration file.
-    For inspiration, splitter configuration files can be found under 
+    For inspiration, splitter configuration files can be found under
     dvc/training/configs/
     """,
 )

From 61e74c8956940a348ae7736e33179860280d7cc7 Mon Sep 17 00:00:00 2001
From: Delattre Emilie <emilie.delattre@epfl.ch>
Date: Wed, 26 Apr 2023 13:43:45 +0200
Subject: [PATCH 12/12] Add cleanlab in setup.cfg

---
 setup.cfg | 1 +
 1 file changed, 1 insertion(+)

diff --git a/setup.cfg b/setup.cfg
index 5bb4436..20f60a4 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -26,6 +26,7 @@ python_requires = >=3.8,<3.9
 install_requires =
     PyYAML
     captum
+    cleanlab==1.0.1
     click
     dash
     dash-bootstrap-components