✨ MOVE update

.gitignore

@@ -5,8 +5,9 @@
 __pycache__/
 *.py[cod]
 
-# NumPy binary files
-data*/*.npy
+# NumPy/PyTorch binary files
+*.npy
+*.pt
 
 # Distribution and packaging files
 build/
@@ -31,15 +32,29 @@ outputs/
 *.log
 
 # Tutorial files
-**/interim_data/
-**/processed_data/
-**/results/
-tutorial/maize/data
+tutorial/*
+!tutorial/config/*maize*.yaml
+!tutorial/config/*random_small*.yaml
+!tutorial/data
+!tutorial/maize/maize_dataset.py
+!tutorial/notebooks/*.ipynb
+!tutorial/README.md
 
 # Virtual environment
 venv/
 virtualvenv/
 
 # docs files
 docs/build/
-docs/source/_templates/
+docs/source/_templates/
+
+# Root folder
+/*.*
+!/.gitignore
+!/.readthedocs.yaml
+!/LICENSE
+!/MANIFEST.in
+!/README.md
+!/pyproject.toml
+!/requirements.txt
+!/setup.cfg

src/move/__init__.py

@@ -1,11 +1,10 @@
 from __future__ import annotations
 
 __license__ = "MIT"
-__version__ = (1, 4, 9)
-__all__ = ["conf", "data", "models", "training_loop", "VAE"]
+__version__ = (2, 0, 0)
+__all__ = ["conf", "data", "models", "tasks", "viz"]
 
 HYDRA_VERSION_BASE = "1.2"
 
-from move import conf, data, models
-from move.models.vae import VAE
-from move.training.training_loop import training_loop
+import move.visualization as viz
+from move import conf, data, models, tasks

src/move/__main__.py

@@ -3,16 +3,11 @@
 import hydra
 from omegaconf import OmegaConf
 
-import move.tasks
 from move import HYDRA_VERSION_BASE
-from move.conf.schema import (
-    AnalyzeLatentConfig,
-    EncodeDataConfig,
-    IdentifyAssociationsConfig,
-    MOVEConfig,
-    TuneModelConfig,
-)
+from move.conf.schema import SUPPORTED_TASKS, MOVEConfig
 from move.core.logging import get_logger
+from move.core.seed import set_global_seed
+from move.tasks.base import Task
 
 
 @hydra.main(
@@ -32,14 +27,11 @@ def main(config: MOVEConfig) -> None:
     if task_type is None:
         logger = get_logger("move")
         logger.info("No task specified.")
-    elif task_type is EncodeDataConfig:
-        move.tasks.encode_data(config.data)
-    elif issubclass(task_type, TuneModelConfig):
-        move.tasks.tune_model(config)
-    elif task_type is AnalyzeLatentConfig:
-        move.tasks.analyze_latent(config)
-    elif issubclass(task_type, IdentifyAssociationsConfig):
-        move.tasks.identify_associations(config)
+    elif issubclass(task_type, SUPPORTED_TASKS):
+        if config.seed is not None:
+            set_global_seed(config.seed)
+        task: Task = hydra.utils.instantiate(config.task, _recursive_=False)
+        task.run()
     else:
         raise ValueError("Unsupported type of task.")
 

src/move/analysis/fdr.py

@@ -0,0 +1,18 @@
+import math
+from typing import cast
+
+import numpy as np
+from numpy.typing import NDArray
+
+
+def argnearest(array: NDArray, target: float) -> int:
+    """Find value in array closest to target. Assumes array is sorted in
+    ascending order."""
+    idx = np.searchsorted(array, target, side="left")
+    if idx > 0 and (
+        idx == len(array)
+        or math.fabs(target - array[idx - 1]) < math.fabs(target - array[idx])
+    ):
+        return cast(int, idx - 1)
+    else:
+        return cast(int, idx)

src/move/analysis/feature_importance.py

@@ -0,0 +1,99 @@
+__all__ = ["FeatureImportance"]
+
+from typing import TYPE_CHECKING
+
+import pandas as pd
+import torch
+
+import move.visualization as viz
+from move.core.exceptions import UnsetProperty
+from move.data.io import sanitize_filename
+from move.tasks.base import CsvWriterMixin, ParentTask, SubTask
+
+if TYPE_CHECKING:
+    from move.data.dataloader import MoveDataLoader
+    from move.models.base import BaseVae
+
+
+class FeatureImportance(CsvWriterMixin, SubTask):
+    """Compute feature importance in latent space.
+
+    Feature importance is computed per feature per dataset. For each dataset,
+    a file will be created.
+
+    Feature importance is computed as the sum of differences in latent
+    variables generated when a feature is present/removed."""
+
+    data_filename_fmt: str = "feature_importance_{}.csv"
+    plot_filename_fmt: str = "feature_importance_{}.png"
+
+    def __init__(
+        self, parent: ParentTask, model: "BaseVae", dataloader: "MoveDataLoader"
+    ) -> None:
+        self.parent = parent
+        self.model = model
+        self.dataloader = dataloader
+
+    def plot(self) -> None:
+        if self.parent is None:
+            return
+        for dataset in self.dataloader.datasets:
+            csv_filename = self.data_filename_fmt.format(dataset.name)
+            csv_filepath = self.parent.output_dir / sanitize_filename(csv_filename)
+            fig_filename = self.plot_filename_fmt.format(dataset.name)
+            fig_filepath = self.parent.output_dir / sanitize_filename(fig_filename)
+
+            diffs = pd.read_csv(csv_filepath)
+
+            if dataset.data_type == "continuous":
+                fig = viz.plot_continuous_feature_importance(
+                    diffs.values, dataset.tensor.numpy(), dataset.feature_names
+                )
+            else:
+                # Categorical dataset is re-shaped to 3D shape
+                dataset_shape = getattr(dataset, "original_shape")
+                fig = viz.plot_categorical_feature_importance(
+                    diffs.values,
+                    dataset.tensor.reshape(-1, *dataset_shape).numpy(),
+                    dataset.feature_names,
+                    getattr(dataset, "mapping"),
+                )
+
+            fig.savefig(fig_filepath, bbox_inches="tight")
+
+    @torch.no_grad()
+    def run(self) -> None:
+        for dataset in self.dataloader.datasets:
+            self.log(f"Computing feature importance: '{dataset}'")
+            # Create a file for each dataset
+            # File is transposed; each column is a sample, each row a feature
+            if self.parent:
+                csv_filename = sanitize_filename(self.data_filename_fmt.format(dataset))
+                csv_filepath = self.parent.output_dir / csv_filename
+                colnames = ["feature_name"] + [""] * len(self.dataloader.dataset)
+                self.init_csv_writer(
+                    csv_filepath, fieldnames=colnames, extrasaction="ignore"
+                )
+            else:
+                raise UnsetProperty("Parent task")
+
+            # Make a perturbation for each feature
+            for feature_name in dataset.feature_names:
+                value = None if dataset.data_type == "discrete" else 0.0
+                self.dataloader.dataset.perturb(dataset.name, feature_name, value)
+                row = [feature_name]
+                for tup in self.dataloader:
+                    batch, pert_batch, _ = tup
+                    z = self.model.project(batch)
+                    z_pert = self.model.project(pert_batch)
+                    diff = torch.sum(z_pert - z, dim=-1)
+                    row.extend(diff.tolist())
+                self.write_row(row)
+
+            self.close_csv_writer(clear=True)
+
+            # Transpose CSV file, so each row is a sample, each column a feature
+            pd.read_csv(csv_filepath).T.to_csv(csv_filepath, index=False, header=False)
+
+        # Clear perturbation
+        self.dataloader.dataset.perturbation = None

src/move/analysis/hdi.py

@@ -0,0 +1,36 @@
+import math
+
+import torch
+
+
+def hdi_bounds(
+    x: torch.Tensor, hdi_prob: float = 0.95
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Return highest density interval (HDI) of a samples-features matrix.
+    The HDI represents the range within which most of the samples are located.
+
+    Args:
+        x: Matrix (`num_samples` x `num_features`)
+        hdi_prob: Percentage of samples inside the HDI
+
+    Returns:
+        Lower and upper bounds of HDI
+    """
+    # adapated from arviz
+
+    if x.dim() != 2:
+        raise ValueError("Can only calculate for matrices with two dimensions")
+
+    n = x.size(0)
+    x, _ = torch.sort(x, dim=0)
+
+    interval_idx_inc = math.floor(hdi_prob * n)
+    num_intervals = n - interval_idx_inc
+
+    interval_width = x[interval_idx_inc:] - x[:num_intervals]
+    min_idx = torch.argmin(interval_width, dim=0)
+
+    hdi_min = torch.diag(x[min_idx])
+    hdi_max = torch.diag(x[min_idx + interval_idx_inc])
+
+    return hdi_min, hdi_max

src/move/analysis/metrics.py

@@ -1,8 +1,19 @@
 __all__ = ["calculate_accuracy", "calculate_cosine_similarity"]
 
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, cast
+
 import numpy as np
+import pandas as pd
+import torch
 
+import move.visualization as viz
 from move.core.typing import FloatArray
+from move.tasks.base import CsvWriterMixin, ParentTask, SubTask
+
+if TYPE_CHECKING:
+    from move.data.dataloader import MoveDataLoader
+    from move.models.base import BaseVae
 
 
 def calculate_accuracy(
@@ -33,7 +44,7 @@ def calculate_accuracy(
     y_pred = np.ma.masked_array(reconstruction, mask=is_nan)
 
     num_features = np.ma.count(y_true, axis=1)
-    scores = np.ma.filled(np.sum(y_true == y_pred, axis=1) / num_features, 0)
+    scores = np.ma.filled(np.sum(y_true == y_pred, axis=1) / num_features, np.nan)
 
     return scores
 
@@ -64,7 +75,7 @@ def calculate_cosine_similarity(
 
     # Equivalent to `np.diag(sklearn.metrics.pairwise.cosine_similarity(x, y))`
     # But can handle masked arrays
-    scores = np.ma.compressed(np.sum(x * y, axis=1)) / (norm(x) * norm(y))
+    scores = np.ma.filled(np.sum(x * y, axis=1), np.nan) / (norm(x) * norm(y))
 
     return scores
 
@@ -80,4 +91,64 @@ def norm(x: np.ma.MaskedArray, axis: int = 1) -> FloatArray:
     Returns:
         1D array with the specified axis removed.
     """
-    return np.ma.compressed(np.sqrt(np.sum(x**2, axis=axis)))
+    return np.ma.filled(np.sqrt(np.sum(x**2, axis=axis)), np.nan)
+
+
+class ComputeAccuracyMetrics(CsvWriterMixin, SubTask):
+    """Compute accuracy metrics between original input and reconstruction (use
+    cosine similarity for continuous dataset reconstructions)."""
+
+    data_filename: str = "reconstruction_metrics.csv"
+    plot_filename: str = "reconstruction_metrics.png"
+
+    def __init__(
+        self, parent: ParentTask, model: "BaseVae", dataloader: "MoveDataLoader"
+    ) -> None:
+        self.parent = parent
+        self.model = model
+        self.dataloader = dataloader
+
+    def plot(self) -> None:
+        if self.parent and self.csv_filepath:
+            scores = pd.read_csv(self.csv_filepath, index_col=None)
+            fig = viz.plot_metrics_boxplot(scores, labels=None)
+            fig_path = self.parent.output_dir / self.plot_filename
+            fig.savefig(fig_path, bbox_inches="tight")
+
+    @torch.no_grad()
+    def run(self) -> None:
+        if self.parent:
+            csv_filepath = self.parent.output_dir / self.data_filename
+            colnames = self.dataloader.dataset.dataset_names
+            self.init_csv_writer(
+                csv_filepath, fieldnames=colnames, extrasaction="ignore"
+            )
+        else:
+            self.log("No parent task, metrics will not be saved.", "WARNING")
+
+        self.log("Computing accuracy metrics")
+
+        datasets = self.dataloader.datasets
+        for batch in self.dataloader:
+            batch_disc, batch_cont = self.model.split_input(batch[0])
+            recon = self.model.reconstruct(batch[0], as_one=True)
+            recon_disc, recon_cont = self.model.split_input(recon)
+
+            scores_per_dataset = {}
+            for i, dataset in enumerate(datasets[: len(batch_disc)]):
+                target = batch_disc[i].numpy()
+                preds = torch.argmax(
+                    (torch.log_softmax(recon_disc[i], dim=-1)), dim=-1
+                ).numpy()
+                scores = calculate_accuracy(target, preds)
+                scores_per_dataset[dataset.name] = scores
+
+            for i, dataset in enumerate(datasets[len(batch_disc) :]):
+                target = batch_cont[i].numpy()
+                preds = recon_cont[i].numpy()
+                scores = calculate_cosine_similarity(target, preds)
+                scores_per_dataset[dataset.name] = scores
+
+            self.write_cols(scores_per_dataset)
+
+        self.close_csv_writer()

src/move/conf/__init__.py

@@ -1,3 +1,15 @@
-__all__ = ["MOVEConfig"]
+__all__ = [
+    "AdamConfig",
+    "AdamWConfig",
+    "ProdigyConfig",
+    "SgdConfig",
+    "TrainingDataLoaderConfig",
+    "TrainingLoopConfig",
+    "VaeConfig",
+    "VaeNormalConfig",
+    "VaeTConfig",
+]
 
-from move.conf.schema import MOVEConfig
+from move.conf.models import VaeConfig, VaeNormalConfig, VaeTConfig
+from move.conf.optim import AdamConfig, AdamWConfig, ProdigyConfig, SgdConfig
+from move.conf.training import TrainingDataLoaderConfig, TrainingLoopConfig

src/move/conf/config_store.py

@@ -0,0 +1,6 @@
+__all__ = ["config_store"]
+
+from hydra.core.config_store import ConfigStore
+
+config_store = ConfigStore.instance()
+"""Hydra's config store singleton"""