ENH: modifications in prep for the documentation

src/piqtree2/_app/__init__.py

@@ -2,35 +2,70 @@
 
 import cogent3
 import cogent3.app.typing as c3_types
-from cogent3.app.composable import define_app
+from cogent3.app import composable
 from cogent3.util.misc import extend_docstring_from
 
 from piqtree2 import TreeGenMode, build_tree, fit_tree, nj_tree, random_trees
 from piqtree2.model import Model
 
 
-@define_app
-@extend_docstring_from(build_tree)
-def piqtree_phylo(
-    aln: cogent3.Alignment | cogent3.ArrayAlignment,
-    model: Model,
-    rand_seed: int | None = None,
-) -> cogent3.PhyloNode | cogent3.app.typing.SerialisableType:
-    return build_tree(aln, model, rand_seed)
+@composable.define_app
+class piqtree_phylo:
+    @extend_docstring_from(build_tree)
+    def __init__(
+        self,
+        substitution_model: str,
+        freq_type: str | None = None,
+        rate_model: str | None = None,
+        *,
+        invariant_sites: bool = False,
+        rand_seed: int | None = None,
+    ) -> None:
+        self._model = Model(
+            substitution_model=substitution_model,
+            invariant_sites=invariant_sites,
+            rate_model=rate_model,
+            freq_type=freq_type,
+        )
+        self._rand_seed = rand_seed
 
+    def main(
+        self,
+        aln: cogent3.Alignment | cogent3.ArrayAlignment,
+    ) -> cogent3.PhyloNode | cogent3.app.typing.SerialisableType:
+        return build_tree(aln, self._model, self._rand_seed)
 
-@define_app
-@extend_docstring_from(fit_tree)
-def piqtree_fit(
-    aln: cogent3.Alignment | cogent3.ArrayAlignment,
-    tree: cogent3.PhyloNode,
-    model: Model,
-    rand_seed: int | None = None,
-) -> cogent3.PhyloNode | cogent3.app.typing.SerialisableType:
-    return fit_tree(aln, tree, model, rand_seed)
+
+@composable.define_app
+class piqtree_fit:
+    @extend_docstring_from(build_tree)
+    def __init__(
+        self,
+        tree: cogent3.PhyloNode,
+        substitution_model: str,
+        freq_type: str | None = None,
+        rate_model: str | None = None,
+        *,
+        rand_seed: int | None = None,
+        invariant_sites: bool = False,
+    ) -> None:
+        self._tree = tree
+        self._model = Model(
+            substitution_model=substitution_model,
+            invariant_sites=invariant_sites,
+            rate_model=rate_model,
+            freq_type=freq_type,
+        )
+        self._rand_seed = rand_seed
+
+    def main(
+        self,
+        aln: cogent3.Alignment | cogent3.ArrayAlignment,
+    ) -> cogent3.PhyloNode | cogent3.app.typing.SerialisableType:
+        return fit_tree(aln, self._tree, self._model, self._rand_seed)
 
 
-@define_app
+@composable.define_app
 @extend_docstring_from(random_trees)
 def piqtree_random_trees(
     num_taxa: int,
@@ -41,10 +76,15 @@ def piqtree_random_trees(
     return random_trees(num_taxa, tree_mode, num_trees, rand_seed)
 
 
-@define_app
+@composable.define_app
 @extend_docstring_from(nj_tree)
 def piqtree_nj(dists: c3_types.PairwiseDistanceType) -> cogent3.PhyloNode:
     return nj_tree(dists)
 
 
-_ALL_APP_NAMES = ["piqtree_phylo", "piqtree_fit", "piqtree_random_trees", "piqtree_nj"]
+_ALL_APP_NAMES = [
+    "piqtree_phylo",
+    "piqtree_fit",
+    "piqtree_random_trees",
+    "piqtree_nj",
+]

src/piqtree2/model/__init__.py

@@ -1,10 +1,21 @@
 """Models available in IQ-TREE."""
 
-from ._freq_type import FreqType
+from ._freq_type import FreqType, get_freq_type
 from ._model import Model
 from ._options import available_freq_type, available_models, available_rate_type
-from ._rate_type import DiscreteGammaModel, FreeRateModel, RateModel, RateType
-from ._substitution_model import AaModel, DnaModel, SubstitutionModel
+from ._rate_type import (
+    DiscreteGammaModel,
+    FreeRateModel,
+    RateModel,
+    RateType,
+    get_rate_type,
+)
+from ._substitution_model import (
+    AaModel,
+    DnaModel,
+    SubstitutionModel,
+    get_substitution_model,
+)
 
 __all__ = [
     "available_freq_type",
@@ -15,6 +26,9 @@
     "DnaModel",
     "FreeRateModel",
     "FreqType",
+    "get_freq_type",
+    "get_rate_type",
+    "get_substitution_model",
     "Model",
     "RateModel",
     "RateType",

src/piqtree2/model/_freq_type.py

@@ -1,3 +1,4 @@
+import contextlib
 import functools
 from enum import Enum, unique
 
@@ -32,3 +33,20 @@ def description(self) -> str:
 
         """
         return self._descriptions()[self]
+
+    def iqtree_str(self) -> str:
+        return self.value
+
+
+def get_freq_type(name: str | FreqType) -> FreqType:
+    """returns the substitution model enum for name."""
+    if isinstance(name, FreqType):
+        return name
+
+    name = name.lstrip("+")
+
+    with contextlib.suppress(KeyError):
+        return FreqType[name]
+
+    msg = f"Unknown state frequency type: {name!r}"
+    raise ValueError(msg)

src/piqtree2/model/_model.py

@@ -1,6 +1,6 @@
-from piqtree2.model._freq_type import FreqType
-from piqtree2.model._rate_type import RateType
-from piqtree2.model._substitution_model import SubstitutionModel
+from piqtree2.model._freq_type import FreqType, get_freq_type
+from piqtree2.model._rate_type import RateModel, get_rate_type
+from piqtree2.model._substitution_model import SubstitutionModel, get_substitution_model
 
 
 class Model:
@@ -11,26 +11,34 @@ class Model:
 
     def __init__(
         self,
-        substitution_model: SubstitutionModel,
-        freq_type: FreqType | None = None,
-        rate_type: RateType | None = None,
+        substitution_model: str | SubstitutionModel,
+        freq_type: str | FreqType | None = None,
+        rate_model: str | RateModel | None = None,
+        *,
+        invariant_sites: bool = False,
     ) -> None:
         """Constructor for the model.
 
         Parameters
         ----------
-        substitution_model : SubstitutionModel
+        substitution_model : str |SubstitutionModel
             The substitution model to use
-        freq_type : Optional[FreqType], optional
-            Base frequency specification, by default None. (defaults
+        freq_type : str | FreqType | None, optional
+            State frequency specification, by default None. (defaults
             to empirical base frequencies if not specified by model).
-        rate_type : Optional[FreqType], optional
+        rate_model : str | RateModel | None, optional
             Rate heterogeneity across sites model, by default
-            no invariable sites, no Gamma, and no FreeRate.
+            no Gamma, and no FreeRate.
+        rate_type : bool, optional
+            Invariable sites.
         """
-        self.substitution_model = substitution_model
-        self.freq_type = freq_type
-        self.rate_type = rate_type
+        self.substitution_model = get_substitution_model(substitution_model)
+        self.freq_type = get_freq_type(freq_type) if freq_type else None
+        self.rate_type = (
+            get_rate_type(rate_model, invariant_sites=invariant_sites)
+            if rate_model is not None or invariant_sites
+            else None
+        )
 
     def __str__(self) -> str:
         """Convert the model into the IQ-TREE representation.
@@ -40,6 +48,10 @@ def __str__(self) -> str:
         str
             The IQ-TREE representation of the mode.
         """
-        freq_str = "" if self.freq_type is None else "+" + self.freq_type.value
-        rate_str = "" if self.rate_type is None else self.rate_type.iqtree_str()
-        return self.substitution_model.value + freq_str + rate_str
+        iqtree_extra_args = filter(
+            lambda x: x is not None,
+            (self.freq_type, self.rate_type),
+        )
+        return "+".join(
+            x.iqtree_str() for x in [self.substitution_model, *iqtree_extra_args]
+        )

src/piqtree2/model/_options.py

@@ -25,7 +25,7 @@ def _make_models(model_type: type[SubstitutionModel]) -> dict[str, list[str]]:
     for model_class in model_classes:
         for model in model_class:
             data["Model Type"].append(model.model_type())
-            data["Abbreviation"].append(model.name)
+            data["Abbreviation"].append(model.value)
             data["Description"].append(model.description)
 
     return data
@@ -40,12 +40,19 @@ def available_models(model_type: str | None = None) -> _Table:
         either "nucleotide", "protein" or None. If None, all models are returned.
 
     """
+    template = "Available {}substitution models"
     if model_type == "dna":
-        table = make_table(data=_make_models(DnaModel))
+        table = make_table(
+            data=_make_models(DnaModel), title=template.format("nucleotide ")
+        )
     elif model_type == "protein":
-        table = make_table(data=_make_models(AaModel))
+        table = make_table(
+            data=_make_models(AaModel), title=template.format("protein ")
+        )
     else:
-        table = make_table(data=_make_models(SubstitutionModel))
+        table = make_table(
+            data=_make_models(SubstitutionModel), title=template.format("")
+        )
 
     return table
 
@@ -55,10 +62,10 @@ def available_freq_type() -> _Table:
     data = {"Freq Type": [], "Description": []}
 
     for freq_type in FreqType:
-        data["Freq Type"].append(freq_type.name)
+        data["Freq Type"].append(freq_type.value)
         data["Description"].append(freq_type.description)
 
-    return make_table(data=data)
+    return make_table(data=data, title="Available frequency types")
 
 
 def available_rate_type() -> _Table:
@@ -69,4 +76,4 @@ def available_rate_type() -> _Table:
         data["Rate Type"].append(rate_type.iqtree_str())
         data["Description"].append(get_description(rate_type))
 
-    return make_table(data=data)
+    return make_table(data=data, title="Available rate heterogeneity types")