bella-companion 0.0.47__tar.gz → 0.1.1__tar.gz

{bella_companion-0.0.47 → bella_companion-0.1.1}/PKG-INFO

@@ -1,12 +1,15 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.3
 Name: bella-companion
-Version: 0.0.47
+Version: 0.1.1
 Summary: Bayesian Evolutionary Layered Learning Architectures (BELLA) companion
-Requires-Python: >=3.10
-Description-Content-Type: text/markdown
+Author: gabriele-marino
+Author-email: gabriele-marino <gabmarino.8601@gmail.com>
 Requires-Dist: arviz>=0.22.0
-Requires-Dist: bella-lumiere>=0.0.14
+Requires-Dist: bella-lumiere>=0.1.0
 Requires-Dist: dotenv>=0.9.9
 Requires-Dist: phylogenie>=3.1.4
 Requires-Dist: poetry>=2.2.1
 Requires-Dist: seaborn>=0.13.2
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+

{bella_companion-0.0.47 → bella_companion-0.1.1}/pyproject.toml

@@ -1,18 +1,28 @@
 [project]
 name = "bella-companion"
-version = "0.0.47"
+version = "0.1.1"
 description = "Bayesian Evolutionary Layered Learning Architectures (BELLA) companion"
 readme = "README.md"
+authors = [
+    { name = "gabriele-marino", email = "gabmarino.8601@gmail.com" }
+]
 requires-python = ">=3.10"
 dependencies = [
     "arviz>=0.22.0",
-    "bella-lumiere>=0.0.14",
+    "bella-lumiere>=0.1.0",
     "dotenv>=0.9.9",
     "phylogenie>=3.1.4",
     "poetry>=2.2.1",
     "seaborn>=0.13.2",
 ]
 
+[project.scripts]
+bella = "bella_companion.cli:main"
+
+[build-system]
+requires = ["uv_build>=0.9.5,<0.10.0"]
+build-backend = "uv_build"
+
 [dependency-groups]
 dev = [
     "joblib-stubs>=1.5.2.0.20250831",
@@ -22,15 +32,3 @@ dev = [
 [tool.pyright]
 typeCheckingMode = "strict"
 
-[project.scripts]
-bella = "bella_companion.cli:main"
-
-[build-system]
-requires = ["setuptools>=42"]
-build-backend = "setuptools.build_meta"
-
-[tool.setuptools]
-packages = ["bella_companion"]
-
-[tool.setuptools.package-data]
-bella_companion = ["**/*"]

{bella_companion-0.0.47 → bella_companion-0.1.1/src}/bella_companion/eucovid/beast_configs/GLM.xml

@@ -48,7 +48,7 @@
                 <populationModel spec="ConstantPopulation" popSize="0.1"/>
             </stateNode>
             <stateNode id="Re" spec="RealParameter" value="1.0 1.1 1.2 1.3 1.4 1.01 1.1 1.2 1.3 1.4"/>
-            <stateNode id="samplingProportion" spec="RealParameter" value="1.11E-5 1.2E-5 1.3E-5 1.4E-5 1.5E-5 1.1E-5 1.2E-5 1.3E-5 1.4E-5 1.5E-5 0.0 0.0 0.0 0.0 0.0"/>
+            <stateNode id="samplingProportion" spec="RealParameter" value="1.1E-4 1.2E-4 1.3E-4 1.4E-4 1.5E-4 1.11E-4 1.2E-4 1.3E-4 1.4E-4 1.5E-4 0.0 0.0 0.0 0.0 0.0"/>
             <stateNode id="migrationRateW" spec="RealParameter" value="0"/>
             <stateNode id="migrationRateScaler" spec="RealParameter" value="0.1"/>
             <stateNode id="processLength" spec="RealParameter" value="10.0"/>
@@ -76,9 +76,7 @@
                         <becomeUninfectiousRate spec="SkylineVectorParameter" skylineValues="36.5" typeSet="@typeSet"/>
                         <migrationRate id="migrationRateSP" spec="SkylineMatrixParameter" timesAreAges="true" processLength="@processLength" typeSet="@typeSet">
                             <skylineValues spec="glmprior.util.GLMLogLinear" coefficients="@migrationRateW" scaleFactor="@migrationRateScaler" transform="false">
-                                <plate var="predictorFile" range="$(predictorFiles)">
-                                    <predictor spec="RealParameterFromXSV" fileName="$(predictorFile)"/>
-                                </plate>
+                                <predictor spec="RealParameterFromXSV" fileName="$(flightsPredictorFile)"/>
                                 <indicators spec="BooleanParameter" value="1"/>
                             </skylineValues>
                             <changeTimes spec="RealParameterFromXSV" fileName="$(changeTimesFile)"/>
@@ -94,26 +92,37 @@
                 <distribution spec="SmartZeroExcludingPrior" x="@Re">
                     <LogNormal name="distr" M="0.8" S="0.5"/>
                 </distribution>
-                <distribution spec="SmartZeroExcludingPrior" x="@samplingProportion" classesToExclude="1.2E-5 1.3E-5 1.4E-5 1.5E-5">
-                    <Uniform name="distr" lower="0" upper="0.15"/>
+
+                <!-- Sampling proportion in China -->
+                <distribution spec="SmartZeroExcludingPrior" x="@samplingProportionValues" classesToExclude="1.1E-4 1.2E-4 1.3E-4 1.4E-4 1.5E-4">
+                    <Uniform name="distr" upper="0.024714828897338403"/> 
+                </distribution>
+                <!-- Sampling proportion in China -->
+                <distribution spec="SmartZeroExcludingPrior" x="@samplingProportionValues" classesToExclude="1.11E-4 1.2E-4 1.3E-4 1.4E-4 1.5E-4">
+                    <Uniform name="distr" upper="$1.8682742128337982e-4"/>
                 </distribution>
-                <distribution spec="SmartZeroExcludingPrior" x="@samplingProportion" classesToExclude="1.1E-5 1.3E-5 1.4E-5 1.5E-5">
-                    <Uniform name="distr" lower="0" upper="0.093"/>
+                 <!-- Sampling proportion in France -->
+                <distribution spec="SmartZeroExcludingPrior" x="@samplingProportionValues" classesToExclude="1.1E-4 1.11E-4 1.3E-4 1.4E-4 1.5E-4">
+                    <Uniform name="distr" upper="0.09217877094972067"/>
                 </distribution>
-                <distribution spec="SmartZeroExcludingPrior" x="@samplingProportion" classesToExclude="1.1E-5 1.2E-5 1.4E-5 1.5E-5">
-                    <Uniform name="distr" lower="0" upper="0.10"/>
+                <!-- Sampling proportion in Germany -->
+                <distribution spec="SmartZeroExcludingPrior" x="@samplingProportionValues" classesToExclude="1.1E-4 1.11E-4 1.2E-4 1.4E-4 1.5E-4">
+                    <Uniform name="distr" upper="0.01770956316410862"/>
                 </distribution>
-                <distribution spec="SmartZeroExcludingPrior" x="@samplingProportion" classesToExclude="1.1E-5 1.2E-5 1.3E-5 1.5E-5">
-                    <Uniform name="distr" lower="0" upper="0.005"/>
+                <!-- Sampling proportion in Germany -->
+                <distribution spec="SmartZeroExcludingPrior" x="@samplingProportionValues" classesToExclude="1.1E-4 1.11E-4 1.2E-4 1.3E-4 1.5E-4">
+                    <Uniform name="distr" upper="0.002209756926738059"/>
                 </distribution>
-                <distribution spec="SmartZeroExcludingPrior" x="@samplingProportion" classesToExclude="1.1E-5 1.2E-5 1.3E-5 1.4E-5">
-                    <Uniform name="distr" lower="0" upper="0.057"/>
+                <!-- Sampling proportion in EU -->
+                <distribution spec="SmartZeroExcludingPrior" x="@samplingProportionValues" classesToExclude="1.1E-4 1.11E-4 1.2E-4 1.3E-4 1.4E-4">
+                    <Uniform name="distr" lower="0" upper="0.012770137524557957"/>
                 </distribution>
+
                 <prior name="distribution" x="@migrationRateW">
                     <Normal name="distr" mean="0.0" sigma="1.0"/>
                 </prior>
                 <prior name="distribution" x="@migrationRateScaler">
-                    <LogNormal name="distr" M="0.0" S="0.8"/>
+                    <LogNormal name="distr" M="0.0" S="1"/>
                 </prior>
                 <prior name="distribution" x="@gammaShape">
                     <Exponential name="distr" mean="0.5"/>
@@ -136,14 +145,15 @@
         <operator spec="BactrianScaleOperator" parameter="@processLength" scaleFactor="0.25" weight="3.0"/>
         <operator spec="SmartScaleOperator" parameter="@Re" scaleFactor="0.25" weight="10.0"/>
         <operator spec="SmartScaleOperator" parameter="@samplingProportion" weight="3.0"/>
-        <operator spec="BactrianRandomWalkOperator" parameter="@migrationRateW" weight="5.0" scaleFactor="0.1"/>
-        <operator spec="AdaptableVarianceMultivariateNormalOperator" weight="5.0" coefficient="1.0" scaleFactor="0.5" beta="0.05" initial="800" burnin="400" every="1"> 
+        <operator spec="BactrianScaleOperator" parameter="@migrationRateScaler" scaleFactor="0.25" weight="3.0"/>
+        <operator spec="BactrianRandomWalkOperator" parameter="@migrationRateW" scaleFactor="0.1" weight="5.0"/>
+        <operator spec="AdaptableVarianceMultivariateNormalOperator" weight="5.0" scaleFactor="0.5" coefficient="1.0" beta="0.05" initial="800" burnin="400" every="1"> 
             <transformations spec="Transform$NoTransform" f="@migrationRateW"/>
         </operator>
         <operator spec="BactrianScaleOperator" parameter="@gammaShape" weight="1.0"/>
         <operator spec="BactrianScaleOperator" parameter="@kappa" weight="1.0"/>
 
-        <logger spec="Logger" fileName="MCMC.log" logEvery="1000" model="@posterior" sanitiseHeaders="true" sort="smart">
+        <logger spec="Logger" fileName="MCMC.log" logEvery="1000" model="@posterior">
             <log idref="posterior"/>
             <log idref="likelihood"/>
             <log idref="prior"/>

{bella_companion-0.0.47 → bella_companion-0.1.1/src}/bella_companion/eucovid/beast_configs/MLP.xml

@@ -77,12 +77,11 @@
                         <becomeUninfectiousRate spec="SkylineVectorParameter" skylineValues="36.5" typeSet="@typeSet"/>
                         <migrationRate id="migrationRateSP" spec="SkylineMatrixParameter" timesAreAges="true" processLength="@processLength" typeSet="@typeSet">
                             <skylineValues id="migrationRate" spec="bella.BayesMLP" nodes="$(nodes)">
-                                <plate var="predictorFile" range="$(predictorFiles)">
-                                    <predictor spec="RealParameterFromXSV" fileName="$(predictorFile)"/>
-                                </plate>
+                                <predictor spec="RealParameterFromXSV" fileName="$(flightsPredictorFile)"/>
                                 <plate var="n" range="$(layersRange)">
                                     <weights idref="migrationRateW$(n)"/>
                                 </plate>
+                                <activationFunctionsOutput spec="bella.SoftPlus"/>
                             </skylineValues>
                             <changeTimes spec="RealParameterFromXSV" fileName="$(changeTimesFile)"/>
                         </migrationRate> 
@@ -144,7 +143,7 @@
         <operator spec="BactrianScaleOperator" parameter="@gammaShape" weight="1.0"/>
         <operator spec="BactrianScaleOperator" parameter="@kappa" weight="1.0"/>
 
-        <logger spec="Logger" fileName="MCMC.log" logEvery="1000" model="@posterior" sanitiseHeaders="true" sort="smart">
+        <logger spec="Logger" fileName="MCMC.log" logEvery="1000" model="@posterior">
             <log idref="posterior"/>
             <log idref="likelihood"/>
             <log idref="prior"/>

bella_companion-0.1.1/src/bella_companion/eucovid/run_beast.py

@@ -0,0 +1,36 @@
+import os
+from itertools import product
+from pathlib import Path
+
+from bella_companion.utils import submit_job
+
+
+def run_beast():
+    base_output_dir = Path(os.environ["BELLA_BEAST_OUTPUT_DIR"])
+    output_dir = base_output_dir / "eucovid"
+    os.makedirs(output_dir, exist_ok=True)
+
+    base_log_dir = Path(os.environ["BELLA_SBATCH_LOG_DIR"]) / "eucovid"
+    data_dir = Path(__file__).parent / "data"
+    beast_configs_dir = Path(__file__).parent / "beast_configs"
+
+    for seed, model in product(range(3), ["GLM", "MLP"]):
+        extra_data = '-D layersRange="0,1,2",nodes="16 8"' if model == "MLP" else ""
+        os.makedirs(output_dir / model, exist_ok=True)
+        submit_job(
+            " ".join(
+                [
+                    *os.environ["BELLA_RUN_BEAST_CMD"],
+                    f"-seed {seed}",
+                    f"-D msa_file={data_dir / 'msa.fasta'}",
+                    f"-D changeTimesFile={data_dir / 'changetimes.csv'}",
+                    f"-D flightsPredictorFile={data_dir / 'flights.csv'}",
+                    extra_data,
+                    f"-prefix {output_dir / model}{os.sep}{seed}-",
+                    str(beast_configs_dir / f"{model}.xml"),
+                ]
+            ),
+            base_log_dir / model / str(seed),
+            cpus=128,
+            mem_per_cpu=12000,
+        )

bella_companion-0.1.1/src/bella_companion/features.py

@@ -0,0 +1,13 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class Feature:
+    is_binary: bool
+    is_relevant: bool | None = None
+
+    @property
+    def color(self) -> str:
+        if self.is_relevant is None:
+            return "purple"
+        return "#E74C3C" if self.is_relevant else "gray"

{bella_companion-0.0.47 → bella_companion-0.1.1/src}/bella_companion/simulations/figures/epi_multitype.py

@@ -6,18 +6,18 @@ import joblib
 import matplotlib.pyplot as plt
 import numpy as np
 import polars as pl
-from lumiere.backend import relu, sigmoid
+from bella_lumiere.backend import relu, sigmoid
 
-from bella_companion.simulations.figures.explain import (
-    plot_partial_dependencies,
-    plot_shap_features_importance,
-)
 from bella_companion.simulations.scenarios.epi_multitype import (
     MIGRATION_PREDICTOR,
     MIGRATION_RATE_UPPER,
     MIGRATION_RATES,
     SCENARIO,
 )
+from bella_companion.utils import (
+    plot_partial_dependencies,
+    plot_shap_features_importance,
+)
 
 
 def _plot_predictions(log_summary: pl.DataFrame, output_dir: Path):

{bella_companion-0.0.47 → bella_companion-0.1.1/src}/bella_companion/simulations/figures/fbd_2traits.py

@@ -6,12 +6,8 @@ import joblib
 import matplotlib.pyplot as plt
 import numpy as np
 import polars as pl
-from lumiere.backend import relu, sigmoid
+from bella_lumiere.backend import relu, sigmoid
 
-from bella_companion.simulations.figures.explain import (
-    plot_partial_dependencies,
-    plot_shap_features_importance,
-)
 from bella_companion.simulations.figures.utils import step
 from bella_companion.simulations.scenarios.fbd_2traits import (
     FBD_RATE_UPPER,
@@ -20,6 +16,10 @@ from bella_companion.simulations.scenarios.fbd_2traits import (
     SCENARIO,
     STATES,
 )
+from bella_companion.utils import (
+    plot_partial_dependencies,
+    plot_shap_features_importance,
+)
 
 
 def _plot_predictions(log_summary: pl.DataFrame, output_dir: Path):

{bella_companion-0.0.47 → bella_companion-0.1.1/src}/bella_companion/simulations/scenarios/epi_multitype.py

@@ -2,7 +2,7 @@ import numpy as np
 from numpy.random import Generator
 from phylogenie import get_epidemiological_events
 
-from bella_companion.simulations.features import Feature
+from bella_companion.features import Feature
 from bella_companion.simulations.scenarios.common import (
     BECOME_UNINFECTIOUS_RATE,
     EPI_MAX_TIME,

{bella_companion-0.0.47 → bella_companion-0.1.1/src}/bella_companion/simulations/scenarios/epi_skyline.py

@@ -3,7 +3,7 @@ from functools import partial
 import numpy as np
 from phylogenie import SkylineParameter, get_epidemiological_events
 
-from bella_companion.simulations.features import Feature
+from bella_companion.features import Feature
 from bella_companion.simulations.scenarios.common import (
     BECOME_UNINFECTIOUS_RATE,
     EPI_MAX_TIME,

{bella_companion-0.0.47 → bella_companion-0.1.1/src}/bella_companion/simulations/scenarios/fbd_2traits.py

@@ -2,7 +2,7 @@ import numpy as np
 from numpy.random import Generator
 from phylogenie import SkylineVector, get_canonical_events
 
-from bella_companion.simulations.features import Feature
+from bella_companion.features import Feature
 from bella_companion.simulations.scenarios.common import (
     FBD_MAX_TIME,
     FBD_RATE_UPPER,

{bella_companion-0.0.47 → bella_companion-0.1.1/src}/bella_companion/simulations/scenarios/fbd_no_traits.py

@@ -3,7 +3,7 @@ from functools import partial
 import numpy as np
 from phylogenie import SkylineParameter, get_canonical_events
 
-from bella_companion.simulations.features import Feature
+from bella_companion.features import Feature
 from bella_companion.simulations.scenarios.common import (
     FBD_MAX_TIME,
     FBD_RATE_UPPER,

{bella_companion-0.0.47 → bella_companion-0.1.1/src}/bella_companion/simulations/scenarios/scenario.py

@@ -6,7 +6,7 @@ from typing import Any
 from numpy.random import Generator
 from phylogenie.treesimulator import Event
 
-from bella_companion.simulations.features import Feature
+from bella_companion.features import Feature
 
 
 class ScenarioType(Enum):

{bella_companion-0.0.47 → bella_companion-0.1.1/src}/bella_companion/utils/__init__.py

@@ -4,8 +4,8 @@ from bella_companion.utils.beast import (
     summarize_logs_dir,
 )
 from bella_companion.utils.explain import (
-    get_median_partial_dependence_values,
-    get_median_shap_features_importance,
+    plot_partial_dependencies,
+    plot_shap_features_importance,
 )
 from bella_companion.utils.slurm import submit_job
 
@@ -13,7 +13,7 @@ __all__ = [
     "summarize_log",
     "summarize_logs_dir",
     "read_weights_dir",
-    "get_median_partial_dependence_values",
-    "get_median_shap_features_importance",
+    "plot_partial_dependencies",
+    "plot_shap_features_importance",
     "submit_job",
 ]

{bella_companion-0.0.47 → bella_companion-0.1.1/src}/bella_companion/utils/beast.py

@@ -7,9 +7,9 @@ from typing import Any
 import arviz as az
 import numpy as np
 import polars as pl
+from bella_lumiere import read_log_file, read_weights
+from bella_lumiere.typing import Weights
 from joblib import Parallel, delayed
-from lumiere import read_log_file, read_weights
-from lumiere.typing import Weights
 from tqdm import tqdm
 
 from bella_companion.utils.slurm import get_job_metadata

bella_companion-0.1.1/src/bella_companion/utils/explain/__init__.py

@@ -0,0 +1,7 @@
+from bella_companion.utils.explain.pdp import plot_partial_dependencies
+from bella_companion.utils.explain.shap import plot_shap_features_importance
+
+__all__ = [
+    "plot_partial_dependencies",
+    "plot_shap_features_importance",
+]

{bella_companion-0.0.47/bella_companion/simulations/figures → bella_companion-0.1.1/src/bella_companion/utils}/explain/pdp.py

@@ -1,25 +1,53 @@
 import os
 from functools import partial
+from itertools import product
 from pathlib import Path
 
 import matplotlib.pyplot as plt
 import numpy as np
 import seaborn as sns
+from bella_lumiere import get_partial_dependence_values
+from bella_lumiere.typing import ActivationFunction, Weights
 from joblib import Parallel, delayed
-from lumiere.typing import ActivationFunction, Weights
+from numpy.typing import ArrayLike
 from tqdm import tqdm
 
-from bella_companion.simulations.features import Feature
-from bella_companion.utils import get_median_partial_dependence_values
+import bella_companion.utils.explain.typeguards as tg
+from bella_companion.features import Feature
+
+
+def _get_median_partial_dependence_values(
+    weights: list[Weights],  # shape: (n_weight_samples, ...)
+    inputs: ArrayLike,
+    features_grid: list[list[float]],
+    hidden_activation: ActivationFunction,
+    output_activation: ActivationFunction,
+) -> list[list[float]]:  # shape: (n_features, n_grid_points)
+    pdvalues = [
+        get_partial_dependence_values(
+            weights=w,
+            inputs=inputs,
+            features_grid=features_grid,
+            hidden_activation=hidden_activation,
+            output_activation=output_activation,
+        )
+        for w in weights
+    ]
+    return [
+        np.median([pd[feature_idx] for pd in pdvalues], axis=0).tolist()
+        for feature_idx in range(len(features_grid))
+    ]
 
 
 def plot_partial_dependencies(
-    weights: list[list[Weights]],  # shape: (n_runs, n_weights_samples, ...)
+    weights: list[Weights] | list[list[Weights]],
     features: dict[str, Feature],
-    output_dir: Path,
+    output_dir: str | Path,
     hidden_activation: ActivationFunction,
     output_activation: ActivationFunction,
+    inputs: ArrayLike | None = None,
 ):
+    output_dir = Path(output_dir)
     os.makedirs(output_dir, exist_ok=True)
 
     continuous_grid: list[float] = np.linspace(0, 1, 10).tolist()
@@ -27,10 +55,23 @@ def plot_partial_dependencies(
         [0, 1] if feature.is_binary else continuous_grid
         for feature in features.values()
     ]
+    if inputs is None:
+        inputs = list(product(*features_grid))
+
+    if tg.is_weights_list(weights):
+        func = get_partial_dependence_values
+    elif tg.is_nested_weights_list(weights):
+        func = _get_median_partial_dependence_values
+    else:
+        raise ValueError(
+            "weights must be a list of Weights or a list of list of Weights"
+        )
+
     jobs = Parallel(n_jobs=-1, return_as="generator_unordered")(
         delayed(
             partial(
-                get_median_partial_dependence_values,
+                func,
+                inputs=inputs,
                 features_grid=features_grid,
                 hidden_activation=hidden_activation,
                 output_activation=output_activation,
@@ -48,25 +89,28 @@ def plot_partial_dependencies(
     if any(not f.is_binary for f in features.values()):
         for (label, feature), feature_pdvalues in zip(features.items(), pdvalues):
             if not feature.is_binary:
-                color = "#E74C3C" if feature.is_relevant else "gray"
                 median = np.median(feature_pdvalues, axis=1)
                 lower = np.percentile(feature_pdvalues, 2.5, axis=1)
                 high = np.percentile(feature_pdvalues, 100 - 2.5, axis=1)
                 plt.fill_between(  # pyright: ignore
-                    continuous_grid, lower, high, alpha=0.25, color=color
+                    continuous_grid, lower, high, alpha=0.25, color=feature.color
                 )
                 for mcmc_pds in feature_pdvalues.T:
                     plt.plot(  # pyright: ignore
-                        continuous_grid, mcmc_pds, color=color, alpha=0.2, linewidth=1
+                        continuous_grid,
+                        mcmc_pds,
+                        color=feature.color,
+                        alpha=0.2,
+                        linewidth=1,
                     )
                 plt.plot(  # pyright: ignore
-                    continuous_grid, median, color=color, label=label
+                    continuous_grid, median, color=feature.color, label=label
                 )
-        plt.xlabel("Feature value")  # pyright: ignore
-        plt.ylabel("MLP Output")  # pyright: ignore
-        plt.legend()  # pyright: ignore
-        plt.savefig(output_dir / "PDPs-continuous.svg")  # pyright: ignore
-        plt.close()
+                plt.xlabel("Feature value")  # pyright: ignore
+                plt.ylabel("MLP Output")  # pyright: ignore
+                plt.legend()  # pyright: ignore
+                plt.savefig(output_dir / f"PDPs-continuous-{label}.svg")  # pyright: ignore
+                plt.close()
 
     if any(f.is_binary for f in features.values()):
         data: list[float] = []
@@ -84,10 +128,9 @@ def plot_partial_dependencies(
         )
         ax.get_legend().remove()  # pyright: ignore
 
-        for i, f in enumerate([f for f in features.values() if f.is_binary]):
-            color = "#E74C3C" if f.is_relevant else "gray"
+        for i, feature in enumerate([f for f in features.values() if f.is_binary]):
             for coll in ax.collections[i * 2 : i * 2 + 2]:
-                coll.set_facecolor(color)
+                coll.set_facecolor(feature.color)
 
         plt.xlabel("Feature")  # pyright: ignore
         plt.ylabel("MLP Output")  # pyright: ignore

bella_companion-0.1.1/src/bella_companion/utils/explain/shap.py

@@ -0,0 +1,104 @@
+from functools import partial
+from itertools import product
+from pathlib import Path
+from typing import Callable
+
+import matplotlib.pyplot as plt
+import numpy as np
+import seaborn as sns
+from bella_lumiere import get_shap_features_importance
+from bella_lumiere.typing import ActivationFunction, Array, Weights
+from joblib import Parallel, delayed
+from numpy.typing import ArrayLike
+from tqdm import tqdm
+
+import bella_companion.utils.explain.typeguards as tg
+from bella_companion.features import Feature
+
+
+def _get_median_shap_features_importance(
+    weights: list[Weights],
+    inputs: ArrayLike,
+    hidden_activation: ActivationFunction,
+    output_activation: ActivationFunction,
+) -> list[float]:  # length: n_features
+    features_importance = np.array(
+        [
+            get_shap_features_importance(
+                weights=w,
+                inputs=inputs,
+                hidden_activation=hidden_activation,
+                output_activation=output_activation,
+            )
+            for w in weights
+        ]
+    )
+    return np.median(features_importance, axis=0).tolist()
+
+
+def _plot_shap_violins(
+    features: dict[str, Feature],
+    features_importance: Array,  # shape: (n_samples, n_features)
+    output_file: str | Path,
+):
+    features_importance = np.asarray(features_importance)
+    features_importance /= features_importance.sum(axis=1, keepdims=True)
+
+    def _plot_violins(group_check: Callable[[Feature], bool]):
+        for i, (feature_name, feature) in enumerate(features.items()):
+            if group_check(feature):
+                sns.violinplot(
+                    y=features_importance[:, i],
+                    x=[feature_name] * len(features_importance),
+                    cut=0,
+                    color=feature.color,
+                )
+
+    _plot_violins(lambda f: not f.is_binary)
+    _plot_violins(lambda f: f.is_binary)
+    plt.xlabel("Feature")  # pyright: ignore
+    plt.ylabel("Importance")  # pyright: ignore
+    plt.savefig(output_file)  # pyright: ignore
+    plt.close()
+
+
+def plot_shap_features_importance(
+    weights: list[Weights] | list[list[Weights]],
+    features: dict[str, Feature],
+    output_file: str | Path,
+    hidden_activation: ActivationFunction,
+    output_activation: ActivationFunction,
+    inputs: ArrayLike | None = None,
+):
+    if inputs is None:
+        continuous_grid: list[float] = np.linspace(0, 1, 10).tolist()
+        features_grid: list[list[float]] = [
+            [0, 1] if feature.is_binary else continuous_grid
+            for feature in features.values()
+        ]
+        inputs = list(product(*features_grid))
+
+    if tg.is_weights_list(weights):
+        func = get_shap_features_importance
+    elif tg.is_nested_weights_list(weights):
+        func = _get_median_shap_features_importance
+    else:
+        raise ValueError(
+            "weights must be a list of Weights or a list of list of Weights"
+        )
+
+    jobs = Parallel(n_jobs=-1, return_as="generator_unordered")(
+        delayed(
+            partial(
+                func,
+                inputs=inputs,
+                hidden_activation=hidden_activation,
+                output_activation=output_activation,
+            )
+        )(w)
+        for w in weights
+    )
+    features_importance = np.array(
+        [job for job in tqdm(jobs, total=len(weights), desc="Evaluating SHAPs")]
+    )  # shape: (n_runs, n_features)
+    _plot_shap_violins(features, features_importance, output_file)

bella_companion-0.1.1/src/bella_companion/utils/explain/typeguards.py

@@ -0,0 +1,20 @@
+from typing import Any, TypeGuard
+
+import numpy as np
+from bella_lumiere.typing import Weights
+
+
+def is_list(obj: Any) -> TypeGuard[list[Any]]:
+    return isinstance(obj, list)
+
+
+def is_weights(obj: Any) -> TypeGuard[Weights]:
+    return is_list(obj) and all(isinstance(layer, np.ndarray) for layer in obj)
+
+
+def is_weights_list(weights: Any) -> TypeGuard[list[Weights]]:
+    return is_list(weights) and all(is_weights(w) for w in weights)
+
+
+def is_nested_weights_list(weights: Any) -> TypeGuard[list[list[Weights]]]:
+    return is_list(weights) and all(is_weights_list(w) for w in weights)