add atlas-densities combination manipulate (#59)

* add `atlas-densities combination manipulate`

CHANGELOG.rst

@@ -1,6 +1,9 @@
 Changelog
 =========
 
+Version 0.2.2
+* add ``atlas-densities combination manipulate``
+
 Version 0.2.1
 -------------
 * The following *cell-density* sub-commands can now optionally take a ``--group-ids-config-path``:

README.rst

@@ -305,7 +305,7 @@ and whole brain estimates from `Kim et al. (2017)`_ (located at
 BBP Cell Atlas version 2
 ~~~~~~~~~~~~~~~~~~~~~~~~
 
-Estimate excitatory, GAD67, Pvalb, SST, and VIP neuron densities from the literature and the
+Estimate GAD67, Pvalb, SST, and VIP neuron densities from the literature and the
 transfer functions computed previously (first density estimates).
 
 .. code-block:: bash
@@ -351,6 +351,19 @@ The molecular density of approx_lamp5 was calculated from the other molecular de
 
 which approximates the molecular density of lamp5.
 
+This can be calculated via command line via:
+
+.. code-block:: bash
+
+    atlas-densities combination manipulate             \
+       --clip                                          \
+       --base-nrrd data/molecular_densities/gad67.nrrd \
+       --subtract data/molecular_densities/vip.nrrd    \
+       --subtract data/molecular_densities/pv.nrrd     \
+       --subtract data/molecular_densities/sst.nrrd    \
+       --output-path approx_lamp5.nrrd
+
+
 The command outputs the density files in the output-dir and a metadata json file:
 
 .. code-block:: javascript

atlas_densities/app/combination.py

@@ -2,16 +2,20 @@
 
 Combination operates on two or more volumetric files with nrrd format.
 """
+from __future__ import annotations
+
 import json
 import logging
 from pathlib import Path
 
 import click
+import numpy as np
 import pandas as pd
 import voxcell  # type: ignore
 import yaml  # type: ignore
 from atlas_commons.app_utils import (
     EXISTING_FILE_PATH,
+    assert_properties,
     common_atlas_options,
     log_args,
     set_verbose,
@@ -226,3 +230,77 @@ def combine_markers(annotation_path, hierarchy_path, config):
     proportions = dict(glia_intensities.proportion.astype(str))
     with open(config["outputCellTypeProportionsPath"], "w", encoding="utf-8") as out:
         json.dump(proportions, out, indent=1, separators=(",", ": "))
+
+
+class OrderedParamsCommand(click.Command):
+    """Allow repeated params, but keeping their order
+
+    Inspired by:
+        https://stackoverflow.com/a/65744803
+    """
+
+    options: list[tuple[str, str]] = []
+
+    def parse_args(self, ctx, args):
+        parser = self.make_parser(ctx)
+        opts, _, param_order = parser.parse_args(args=list(args))
+        if opts.get("help", False):
+            click.utils.echo(ctx.get_help(), color=ctx.color)
+            ctx.exit()
+
+        for param in param_order:
+            if param.multiple:
+                type(self).options.append((param, opts[param.name].pop(0)))
+
+        return super().parse_args(ctx, args)
+
+
+@app.command(cls=OrderedParamsCommand)
+@click.option(
+    "--base-nrrd",
+    required=True,
+    type=EXISTING_FILE_PATH,
+    help="Path to nrrd file to which others are added/subtracted",
+)
+@click.option("--output-path", required=True, help="Path of the nrrd file to write")
+@click.option("--add", multiple=True, type=EXISTING_FILE_PATH, help="Add nrrd file to base-nrrd")
+@click.option(
+    "--subtract", multiple=True, type=EXISTING_FILE_PATH, help="Subtract nrrd file to base-nrrd"
+)
+@click.option(
+    "--clip", is_flag=True, default=False, help="Clip volume after each addition / subtraction"
+)
+def manipulate(base_nrrd, clip, add, subtract, output_path):  # pylint: disable=unused-argument
+    """Add and subtract NRRD files from the `base-nrrd`
+
+    Note: the `--add` and `--subtract` can be used multiple times.
+    """
+    volumes = []
+    operations = []
+    paths = []
+    for param, value in OrderedParamsCommand.options:
+        operations.append(param.name)
+        volumes.append(voxcell.VoxelData.load_nrrd(value))
+        paths.append(value)
+
+    L.debug("Loading base NRRD: %s", base_nrrd)
+    combined = voxcell.VoxelData.load_nrrd(base_nrrd)
+
+    assert_properties(
+        [
+            combined,
+        ]
+        + volumes
+    )
+
+    for operation, volume, path in zip(operations, volumes, paths):
+        L.debug("%s with %s", operation, path)
+        if operation == "add":
+            combined.raw += volume.raw
+        elif operation == "subtract":
+            combined.raw -= volume.raw
+
+        if clip:
+            combined.raw = np.clip(combined.raw, a_min=0, a_max=None)
+
+    combined.save_nrrd(output_path)

tests/app/test_combination.py

@@ -128,3 +128,95 @@ def test_combine_markers(tmp_path):
         for type_, arr in expected_glia_intensities.items():
             voxel_data = VoxelData.load_nrrd(f"{type_}.nrrd")
             npt.assert_array_almost_equal(voxel_data.raw, arr, decimal=5)
+
+
+def test_manipulate(tmp_path):
+    nrrd = tmp_path / "nrrd.nrrd"
+    array = np.array(
+        [
+            [[0.0, 0.0, 0.0]],
+            [[0.0, 1.0, 0.0]],
+            [[0.0, 1.0, 0.0]],
+        ]
+    )
+    VoxelData(array, voxel_dimensions=[25] * 3).save_nrrd(nrrd)
+
+    runner = CliRunner()
+    result = runner.invoke(
+        tested.app,
+        [
+            "manipulate",
+            "--clip",
+            "--base-nrrd",
+            nrrd,
+            "--add",
+            nrrd,
+            "--subtract",
+            nrrd,
+            "--add",
+            nrrd,
+            "--subtract",
+            nrrd,
+            "--subtract",
+            nrrd,
+            "--output-path",
+            tmp_path / "manipulate.nrrd",
+        ],
+        catch_exceptions=False,
+    )
+    assert result.exit_code == 0
+    output = VoxelData.load_nrrd(tmp_path / "manipulate.nrrd")
+    assert np.allclose(output.raw, 0.0)
+
+
+def test_manipulate_clip(tmp_path):
+    nrrd = tmp_path / "nrrd.nrrd"
+    array = np.array(
+        [
+            [[0.0, 0.0, 0.0]],
+            [[0.0, 1.0, 0.0]],
+            [[0.0, 1.0, 0.0]],
+        ]
+    )
+    VoxelData(array, voxel_dimensions=[25] * 3).save_nrrd(nrrd)
+
+    runner = CliRunner()
+    result = runner.invoke(
+        tested.app,
+        [
+            "manipulate",
+            "--clip",
+            "--base-nrrd",
+            nrrd,
+            "--subtract",
+            nrrd,
+            "--subtract",
+            nrrd,
+            "--output-path",
+            tmp_path / "manipulate.nrrd",
+        ],
+        catch_exceptions=False,
+    )
+    assert result.exit_code == 0
+    output = VoxelData.load_nrrd(tmp_path / "manipulate.nrrd")
+    assert np.allclose(output.raw, 0.0)
+
+    runner = CliRunner()
+    result = runner.invoke(
+        tested.app,
+        [
+            "manipulate",
+            "--base-nrrd",
+            nrrd,
+            "--subtract",
+            nrrd,
+            "--subtract",
+            nrrd,
+            "--output-path",
+            tmp_path / "manipulate.nrrd",
+        ],
+        catch_exceptions=False,
+    )
+    assert result.exit_code == 0
+    output = VoxelData.load_nrrd(tmp_path / "manipulate.nrrd")
+    assert np.count_nonzero(output.raw < 0) == 2