Support multiple additive terms

docs/_templates/autosummary/class.rst

@@ -39,7 +39,7 @@ Attributes
 
 {% for item in attributes %}
 
-.. autoattribute:: {{ [objname, item] | join(".") }}
+.. autoattribute:: {{ [fullname, item] | join(".") }}
 {%- endfor %}
 
 {% endif %}
@@ -52,7 +52,7 @@ Methods
 
 {% for item in methods %}
 
-.. automethod:: {{ [objname, item] | join(".") }}
+.. automethod:: {{ [fullname, item] | join(".") }}
 {%- endfor %}
 
 {% endif %}

docs/api.md

@@ -10,15 +10,20 @@
 .. autosummary::
     :toctree: generated
 
-    DataOptions
-    ModelOptions
-    TrainingOptions
-    SmoothOptions
     MOFAFLEX
     FeatureSet
     FeatureSets
 ```
 
+### Terms
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :recursive:
+
+    terms
+```
+
 ### Priors
 ```{eval-rst}
 .. autosummary::
@@ -28,6 +33,15 @@
     priors
 ```
 
+### Likelihoods
+```{eval-rst}
+.. autosummary::
+    :toctree: generated
+    :recursive:
+
+    likelihoods
+```
+
 ### Settings
 
 An instance of the [](#_core.settings.Settings) class is available as `mofaflex.settings` and allows configuring MOFA-FLEX.

docs/conf.py

@@ -6,10 +6,13 @@
 
 # -- Path setup --------------------------------------------------------------
 import sys
+import os
 from datetime import datetime
 from importlib.metadata import metadata
 from pathlib import Path
 
+os.environ["MOFAFLEX_DOCS"] = "1"
+
 HERE = Path(__file__).parent
 sys.path.insert(0, str(HERE / "extensions"))
 
@@ -90,7 +93,7 @@
     "numpy": ("https://numpy.org/doc/stable/", None),
     "pandas": ("https://pandas.pydata.org/docs/", None),
     "plotnine": ("https://plotnine.org/", None),
-    "pytorch": ("https://pytorch.org/docs/stable/", None),
+    "pytorch": ("https://docs.pytorch.org/docs/stable/", None),
     "muon-tutorials": ("https://muon-tutorials.readthedocs.io/en/latest", None),
 }
 

pyproject.toml

@@ -57,7 +57,7 @@ optional-dependencies.doc = [
   "myst-nb>=1.1",
   "pandas",
   "setuptools",                      # Until pybtex >0.24.0 releases: https://bitbucket.org/pybtex-devs/pybtex/issues/169/
-  "sphinx>=8.1",
+  "sphinx>=8.1,<9",                  # https://github.com/executablebooks/sphinx-tabs/issues/209
   "sphinx-autodoc-typehints",
   "sphinx-automodapi",
   "sphinx-book-theme>=1",

src/mofaflex/__init__.py

@@ -1,8 +1,8 @@
 import logging
 
 from . import pl, tl
-from ._core import MOFAFLEX, DataOptions, FeatureSet, FeatureSets, ModelOptions, TrainingOptions, settings
-from ._core.api import priors
+from ._core import MOFAFLEX, FeatureSet, FeatureSets, settings
+from ._core.api import likelihoods, priors, terms
 from ._version import __version__, __version_tuple__
 
 _logger = logging.getLogger(__name__)

src/mofaflex/_core/__init__.py

@@ -1,5 +1,5 @@
 from .datasets import MofaFlexDataset
 from .feature_sets import FeatureSet, FeatureSets
-from .mofaflex import MOFAFLEX, DataOptions, ModelOptions, TrainingOptions
+from .mofaflex import MOFAFLEX
 from .pcgse import pcgse_test
 from .settings import settings

src/mofaflex/_core/api/__init__.py

@@ -1 +1 @@
-from . import priors
+from . import likelihoods, priors, terms

src/mofaflex/_core/api/_generate.py

@@ -0,0 +1,77 @@
+import sys
+from abc import ABC, abstractmethod
+from collections.abc import Mapping
+from inspect import Parameter, signature
+
+
+class APIWrapper(ABC):
+    @abstractmethod
+    def __call__(self, axis, names):
+        pass
+
+    def __init__(self, *args, **kwargs):
+        self._args = args
+        self._kwargs = kwargs
+
+    def __eq__(self, other):
+        if not isinstance(other, __class__):
+            return NotImplemented
+        elif other.__class__ != self.__class__:
+            return False
+        else:
+            return self._args == other._args and self._kwargs == other._kwargs
+
+    def __hash__(self):
+        return hash((self.__class__, self._args, tuple(sorted(self._kwargs.items()))))
+
+
+def init_api(module: str, basecls: type, subclss: Mapping[str, type]):
+    mod = sys.modules[module]
+    coreinit = basecls.__dict__.get("__init__", None)
+    if coreinit is not None:
+        coresig = signature(coreinit)
+
+    all_ = []
+
+    basewrapper = type(basecls.__name__, (APIWrapper,), {"__module__": module})
+
+    for subname, subcls in subclss.items():
+        sig = signature(subcls.__init__)
+        annots = subcls.__init__.__annotations__
+        if coreinit is not None:
+            params = [
+                param
+                for i, param in enumerate(sig.parameters.values())
+                if i == 0 or param.name not in coresig.parameters
+            ]
+            sig = sig.replace(parameters=params)
+            annots = {param.name: param.annotation for param in params if param.annotation is not Parameter.empty}
+
+        def init(self, *args, **kwargs):
+            self.__init__.__signature__.bind(self, *args, **kwargs)  # check for argument compatibility
+            super(self.__class__, self).__init__(*args, **kwargs)
+
+        def call(self, *args, **kwargs):
+            return self._cls(*args, *self._args, **kwargs, **self._kwargs)
+
+        if coreinit is not None:
+            call.__signature__ = coresig
+            call.__annotations__ = coreinit.__annotations__
+
+        init.__signature__ = sig
+        init.__annotations__ = annots
+        init.__name__ = "__init__"
+        init.__qualname__ = f"{subname}.__init__"
+        call.__name__ = "__call__"
+        call.__qualname__ = f"{subname}.__call__"
+        apicls = type(
+            subname, (basewrapper,), {"_cls": subcls, "__init__": init, "__call__": call, "__module__": module}
+        )
+        apicls.__doc__ = subcls.__doc__
+
+        setattr(mod, subname, apicls)
+        all_.append(subname)
+
+    setattr(mod, basewrapper.__name__, basewrapper)
+    mod.__all__ = all_
+    mod.__dir__ = lambda: all_

src/mofaflex/_core/api/likelihoods.py

@@ -0,0 +1,4 @@
+from ..likelihoods import Likelihood
+from ._generate import init_api
+
+init_api(__name__, Likelihood, Likelihood.known_likelihoods())

src/mofaflex/_core/api/priors.py

@@ -1,76 +1,4 @@
-from abc import ABC, abstractmethod
-from collections.abc import Sequence
-from inspect import Parameter, signature
-from types import MappingProxyType
-from typing import Literal
+from ..priors import Prior
+from ._generate import init_api
 
-from ..priors import Prior as PriorCore
-
-
-class Prior(ABC):
-    @abstractmethod
-    def __call__(self, axis, names):
-        pass
-
-    def __init__(self, *args, **kwargs):
-        self._args = args
-        self._kwargs = MappingProxyType(kwargs)
-
-    def __eq__(self, other):
-        if not isinstance(other, __class__):
-            return NotImplemented
-        elif other.__class__ != self.__class__:
-            return False
-        else:
-            return self._args == other._args and self._kwargs == other._kwargs
-
-    def __hash__(self):
-        return hash((self.__class__, self._args, tuple(sorted(self._kwargs.items()))))
-
-
-__all__ = []
-
-
-def _init_priors():
-    for priorname in PriorCore.known_priors():
-        priorcls = PriorCore.class_(priorname)
-        sig = signature(priorcls.__init__)
-        params = [param for param in sig.parameters.values() if param.name not in ("axis", "names")]
-        sig = sig.replace(parameters=params)
-
-        def init(self, *args, **kwargs):
-            self.__init__.__signature__.bind(self, *args, **kwargs)  # check for argument compatibility
-            super(self.__class__, self).__init__(*args, **kwargs)
-
-        if priorcls is not PriorCore:
-
-            def call(self, axis: Literal[0, 1, "samples", "features"], names: str | Sequence[str]):
-                return self._cls(axis, names, *self._args, **self._kwargs)
-        else:
-
-            def call(self, axis: Literal[0, 1, "samples", "features"], names: str | Sequence[str]):
-                return PriorCore(self.__class__.__name__, axis, names, *self._args, **self._kwargs)
-
-        init.__signature__ = sig
-        init.__annotations__ = {
-            param.name: param.annotation for param in params if param.annotation is not Parameter.empty
-        }
-        init.__name__ = "__init__"
-        init.__qualname__ = f"{priorname}.__init__"
-        call.__name__ = "__call__"
-        call.__qualname__ = f"{priorname}.__call__"
-        apicls = type(
-            priorname, (Prior,), {"_cls": priorcls, "__init__": init, "__call__": call, "__module__": __name__}
-        )
-        if priorcls is not PriorCore:
-            apicls.__doc__ = priorcls.__doc__
-
-        globals()[priorname] = apicls
-        __all__.append(priorname)
-
-
-_init_priors()
-
-
-def __dir__():
-    return __all__
+init_api(__name__, Prior, Prior.known_priors())

src/mofaflex/_core/api/terms.py

@@ -0,0 +1,47 @@
+from inspect import Parameter, Signature, signature
+from typing import TYPE_CHECKING
+
+from ..terms import Term
+from ..utils import building_docs
+
+__all__ = []
+
+if TYPE_CHECKING:
+    pass
+
+
+def _init_api():
+    def make_wrapper(term: Term):  # required due to Python's late-binding closures
+        def wrapper(name="_", /, **kwargs):
+            from ..mofaflex import MOFAFLEX
+
+            return MOFAFLEX(**{name: term(**kwargs)})
+
+        return wrapper
+
+    for termname, term in Term.known_terms().items():
+        if not building_docs():
+            wrapper = make_wrapper(term)
+            sig = signature(term.__init__)
+            params = [signature(wrapper).parameters["name"]] + [
+                Parameter(param.name, Parameter.KEYWORD_ONLY, default=param.default, annotation=param.annotation)
+                for param in sig.parameters.values()
+            ]
+            wrapper.__signature__ = Signature(params)
+            wrapper.__annotations__ = term.__init__.__annotations__ | {"name": str, "return": "MOFAFLEX"}
+            wrapper.__doc__ = term.__doc__
+        else:
+            wrapper = type(termname, (), {"__module__": __name__, "__doc__": term.__doc__})
+            wrapper.__init__ = term.__init__
+            for api in term.api():
+                setattr(wrapper, api, getattr(term, api))
+
+        globals()[termname] = wrapper
+        __all__.append(termname)
+
+
+def __dir__():
+    return __all__
+
+
+_init_api()

src/mofaflex/_core/api/types.py

@@ -0,0 +1,10 @@
+from ..terms import TermWrapper
+from ..utils import building_docs
+from . import terms
+
+if building_docs():
+    for term in dir(terms):
+        globals()[term] = getattr(terms, term)
+else:
+    for term in dir(terms):
+        globals()[term] = TermWrapper

src/mofaflex/_core/datasets/__init__.py

@@ -1,4 +1,4 @@
 from .anndatadictdataset import AnnDataDictDataset
 from .base import MofaFlexDataset, Preprocessor
-from .misc import CovariatesDataset, GuidingVarsDataset, MofaFlexBatchSampler, StackDataset
+from .misc import CovariatesDataset, MofaFlexBatchSampler, StackDataset, merge_covariates
 from .mudatadataset import MuDataDataset