ENH: add freq rate type including rate het

README.md

@@ -19,15 +19,15 @@ This project is still in early development, if you encounter any problems or hav
 
 ```python
 from piqtree2 import build_tree
-from piqtree2.model import DnaModel
+from piqtree2.model import DnaModel, Model
 from cogent3 import load_aligned_seqs # Included with piqtree2!
 
 # Load Sequences
 aln = load_aligned_seqs("tests/data/example.fasta", moltype="dna")
 aln = aln.take_seqs(["Human", "Chimpanzee", "Rhesus", "Mouse"])
 
 # Reconstruct a phylogenetic tree with IQ-TREE!
-tree = build_tree(aln, DnaModel.JC, rand_seed=1) # Optionally specify a random seed.
+tree = build_tree(aln, Model(DnaModel.JC), rand_seed=1) # Optionally specify a random seed.
 
 print("Tree topology:", tree) # A cogent3 tree object
 print("Log-likelihood:", tree.params["lnL"])
@@ -41,7 +41,7 @@ print("Log-likelihood:", tree.params["lnL"])
 
 ```python
 from piqtree2 import fit_tree
-from piqtree2.model import DnaModel
+from piqtree2.model import DnaModel, Model
 from cogent3 import load_aligned_seqs, make_tree # Included with piqtree2!
 
 # Load Sequences
@@ -52,7 +52,7 @@ aln = aln.take_seqs(["Human", "Chimpanzee", "Rhesus", "Mouse"])
 tree = make_tree("(Human, Chimpanzee, (Rhesus, Mouse));")
 
 # Fit branch lengths with IQ-TREE!
-tree = fit_tree(aln, tree, DnaModel.JC, rand_seed=1) # Optionally specify a random seed.
+tree = fit_tree(aln, tree, Model(DnaModel.JC), rand_seed=1) # Optionally specify a random seed.
 
 print("Tree with branch lengths:", tree) # A cogent3 tree object
 print("Log-likelihood:", tree.params["lnL"])

pyproject.toml

@@ -93,7 +93,7 @@ target-version = "py39"
 # Unlike Flake8, Ruff doesn't enable pycodestyle warnings (`W`) or
 # McCabe complexity (`C901`) by default.
 select = ["ALL"]
-ignore = ["EXE002", "FA100", "E501", "D"]
+ignore = ["EXE002", "FA100", "E501", "D", "PLR2004"]
 
 # Allow fix for all enabled rules (when `--fix`) is provided.
 fixable = ["ALL"]
@@ -120,7 +120,7 @@ dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
     "N803",
     "D"
 ]
-"src/piqtree2/model/_model.py" = ["N815"] # use IQ-TREE naming scheme
+"src/piqtree2/model/_substitution_model.py" = ["N815"] # use IQ-TREE naming scheme
 
 [tool.ruff.format]
 # Like Black, use double quotes for strings.

src/piqtree2/__init__.py

@@ -7,12 +7,14 @@
     random_trees,
     robinson_foulds,
 )
-from piqtree2.model import available_models
+from piqtree2.model import available_freq_type, available_models, available_rate_type
 
 __version__ = "0.2.0"
 
 __all__ = [
+    "available_freq_type",
     "available_models",
+    "available_rate_type",
     "build_tree",
     "fit_tree",
     "random_trees",

src/piqtree2/iqtree/_tree.py

@@ -9,7 +9,7 @@
 
 from piqtree2.exceptions import ParseIqTreeError
 from piqtree2.iqtree._decorator import iqtree_func
-from piqtree2.model import SubstitutionModel
+from piqtree2.model import Model
 
 iq_build_tree = iqtree_func(iq_build_tree, hide_files=True)
 iq_fit_tree = iqtree_func(iq_fit_tree, hide_files=True)
@@ -42,7 +42,7 @@ def _process_tree_yaml(tree_yaml: dict, names: Sequence[str]) -> cogent3.PhyloNo
 
 def build_tree(
     aln: Union[cogent3.Alignment, cogent3.ArrayAlignment],
-    model: SubstitutionModel,
+    model: Model,
     rand_seed: Optional[int] = None,
 ) -> cogent3.PhyloNode:
     """Reconstruct a phylogenetic tree.
@@ -53,8 +53,8 @@ def build_tree(
     ----------
     aln : Union[cogent3.Alignment, cogent3.ArrayAlignment]
         The sequence alignment.
-    model : SubstitutionModel
-        The substitution model.
+    model : Model
+        The substitution model with base frequencies and rate heterogeneity.
     rand_seed : Optional[int], optional
         The random seed - 0 or None means no seed, by default None.
 
@@ -70,14 +70,14 @@ def build_tree(
     names = aln.names
     seqs = [str(seq) for seq in aln.iter_seqs(names)]
 
-    yaml_result = yaml.safe_load(iq_build_tree(names, seqs, model.value, rand_seed))
+    yaml_result = yaml.safe_load(iq_build_tree(names, seqs, str(model), rand_seed))
     return _process_tree_yaml(yaml_result, names)
 
 
 def fit_tree(
     aln: Union[cogent3.Alignment, cogent3.ArrayAlignment],
     tree: cogent3.PhyloNode,
-    model: SubstitutionModel,
+    model: Model,
     rand_seed: Optional[int] = None,
 ) -> cogent3.PhyloNode:
     """Fit branch lengths to a tree.
@@ -91,8 +91,8 @@ def fit_tree(
         The sequence alignment.
     tree : cogent3.PhyloNode
         The topology to fit branch lengths to.
-    model : SubstitutionModel
-        The substitution model.
+    model : Model
+        The substitution model with base frequencies and rate heterogeneity.
     rand_seed : Optional[int], optional
         The random seed - 0 or None means no seed, by default None.
 
@@ -110,6 +110,6 @@ def fit_tree(
     newick = str(tree)
 
     yaml_result = yaml.safe_load(
-        iq_fit_tree(names, seqs, model.value, newick, rand_seed),
+        iq_fit_tree(names, seqs, str(model), newick, rand_seed),
     )
     return _process_tree_yaml(yaml_result, names)

src/piqtree2/model/__init__.py

@@ -1,6 +1,22 @@
 """Models available in IQ-TREE."""
 
-from ._model import AaModel, DnaModel, SubstitutionModel
-from ._options import available_models
+from ._freq_type import FreqType
+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
 
-__all__ = ["available_models", "AaModel", "DnaModel", "SubstitutionModel"]
+__all__ = [
+    "available_freq_type",
+    "available_models",
+    "available_rate_type",
+    "AaModel",
+    "DiscreteGammaModel",
+    "DnaModel",
+    "FreeRateModel",
+    "FreqType",
+    "Model",
+    "RateModel",
+    "RateType",
+    "SubstitutionModel",
+]

src/piqtree2/model/_freq_type.py

@@ -0,0 +1,34 @@
+import functools
+from enum import Enum, unique
+
+from typing_extensions import Self
+
+
+@unique
+class FreqType(Enum):
+    """Types of base frequencies."""
+
+    F = "F"
+    FO = "FO"
+    FQ = "FQ"
+
+    @staticmethod
+    @functools.cache
+    def _descriptions() -> dict[Self, str]:
+        return {
+            FreqType.F: "Empirical state frequency observed from the data.",
+            FreqType.FO: "State frequency optimized by maximum-likelihood from the data. Note that this is with letter-O and not digit-0.",
+            FreqType.FQ: "Equal state frequency.",
+        }
+
+    @property
+    def description(self) -> str:
+        """The description of the FreqType.
+
+        Returns
+        -------
+        str
+            The description of the FreqType.
+
+        """
+        return self._descriptions()[self]