update __init__s to use __all__ and make all module functionality ava…

…ilable

sparsecoding/__init__.py

@@ -1,6 +1,42 @@
-"Modules for sparse coding."
+from .models import SparseCoding
+from .inference import LCA, IHT, ISTA, LSM, MP, OMP, Vanilla, PyTorchOptimizer
+from .visualization import plot_dictionary, plot_patches
+from .priors import SpikeSlabPrior, L0Prior
+from .datasets import BarsDataset, FieldDataset
+from .dictionaries import (
+    load_dictionary_from_pickle,
+    load_bars_dictionary,
+    load_olshausen_dictionary,
+)
 
-from . import models
-from . import inference
-from . import visualization
-from . import priors
+__all__ = [
+    # Models
+    "SparseCoding",
+
+    # Inference
+    "LCA",
+    "IHT",
+    "ISTA",
+    "LSM",
+    "MP",
+    "OMP",
+    "Vanilla",
+    "PyTorchOptimizer",
+
+    # Visualization
+    "plot_dictionary",
+    "plot_patches",
+
+    # Priors
+    "SpikeSlabPrior",
+    "L0Prior",
+
+    # Dictionaries
+    "load_dictionary_from_pickle",
+    "load_bars_dictionary",
+    "load_olshausen_dictionary",
+
+    # Datasets
+    "BarsDataset",
+    "FieldDataset",
+]

sparsecoding/transforms/__init__.py

@@ -1,2 +1,12 @@
-from . import patch
-from . import whiten
+from .patch import sample_random_patches, patchify, quilt
+from .whiten import Whitener
+
+__all__ = [
+    # From patch module
+    "sample_random_patches",
+    "patchify",
+    "quilt",
+
+    # From whiten module
+    "Whitener",
+]