Implement generalized skyline and arbitrary clock models

Author: 4ment

.github/workflows/python-package.yml

@@ -16,7 +16,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: [3.7, 3.8, 3.9, '3.10', '3.11']
+        python-version: [3.8, 3.9, '3.10', '3.11']
     env:
       # https://github.com/pytest-dev/pytest/issues/2042
       PY_IGNORE_IMPORTMISMATCH: 1

environment-dev.yml

@@ -4,13 +4,13 @@ channels:
   - bioconda
   - conda-forge
 dependencies:
-  - python>=3.7
+  - python>=3.10
   - numpy>=1.7
   - dendropy
   - pytorch>=1.9
   - pip
 # Dev dependencies
-  - black
+  - black=23.3.0
   - docformatter
   - flake8
   - isort

environment.yml

@@ -3,7 +3,7 @@ channels:
   - pytorch
   - bioconda
 dependencies:
-  - python=3.7
+  - python=3.10
   - numpy>=1.7
   - dendropy
   - pytorch>=1.9

setup.cfg

@@ -13,14 +13,14 @@ classifiers =
     Intended Audience :: Science/Research
     License :: OSI Approved :: GNU General Public License v3 (GPLv3)
     Operating System :: OS Independent
-    Programming Language :: Python :: 3.7
     Programming Language :: Python :: 3.8
     Programming Language :: Python :: 3.9
     Programming Language :: Python :: 3.10
+    Programming Language :: Python :: 3.11
     Topic :: Scientific/Engineering :: Bio-Informatics
 
 [options]
-python_requires = >=3.7
+python_requires = >=3.8
 packages = find:
 package_dir =
     =.

torchtree/__init__.py

@@ -1,4 +1,5 @@
 """This is the root package of the torchtree framework."""
+
 from ._version import __version__
 from .core.parameter import CatParameter, Parameter, TransformedParameter, ViewParameter
 

torchtree/cli/__init__.py

@@ -1,5 +1,6 @@
 """The cli package contains modules for creating JSON configuration files
 through a command-line interface."""
+
 from torchtree.cli.plugin_manager import PluginManager
 
 PLUGIN_MANAGER = PluginManager()

torchtree/cli/advi.py

@@ -840,14 +840,19 @@ def create_logger(id_, parameters, arg):
     if arg.coalescent:
         models.append('coalescent')
         if arg.coalescent in COALESCENT_PIECEWISE:
-            models.append('gmrf')
             models.append(
                 {
                     'id': arg.coalescent,
                     'type': 'JointDistributionModel',
-                    'distributions': ['coalescent', 'gmrf'],
+                    'distributions': ['coalescent'],
                 }
             )
+            if arg.theta_prior == 'eml':
+                models[-1]['distributions'].append('coalescent.theta.eml')
+                models.append('coalescent.theta.eml')
+            else:
+                models[-1]['distributions'].append('gmrf')
+                models.append('gmrf')
 
     return {
         "id": id_,
@@ -877,14 +882,19 @@ def create_sampler(id_, var_id, parameters, arg):
     if arg.coalescent:
         models.append('coalescent')
         if arg.coalescent in COALESCENT_PIECEWISE:
-            models.append('gmrf')
             models.append(
                 {
                     'id': arg.coalescent,
                     'type': 'JointDistributionModel',
-                    'distributions': ['coalescent', 'gmrf'],
+                    'distributions': ['coalescent'],
                 }
             )
+            if arg.theta_prior == 'eml':
+                models[-1]['distributions'].append('coalescent.theta.eml')
+                models.append('coalescent.theta.eml')
+            else:
+                models[-1]['distributions'].append('gmrf')
+                models.append('gmrf')
 
     return {
         "id": id_,
@@ -937,7 +947,7 @@ def build_advi(arg):
     if arg.clock is not None and arg.heights == 'ratio':
         jacobians_list.append('tree')
 
-    if arg.coalescent in COALESCENT_PIECEWISE:
+    if arg.coalescent in COALESCENT_PIECEWISE and arg.theta_prior != 'eml':
         jacobians_list.remove("coalescent.theta")
 
     joint_jacobian = {
@@ -984,13 +994,20 @@ def build_advi(arg):
             parameters.extend(
                 (
                     f'{branch_model_id}.rates.prior.mean',
-                    f'{branch_model_id}.rates.prior.scale',
+                    f'{branch_model_id}.rates.prior.stdev',
+                )
+            )
+        elif arg.clock == 'ncln':
+            parameters.extend(
+                (
+                    f'{branch_model_id}.location',
+                    f'{branch_model_id}.scale',
                 )
             )
         else:
             parameters.append(f"{branch_model_id}.rate")
 
-        if arg.clock == 'horseshoe' or arg.clock == 'ucln':
+        if arg.clock == 'horseshoe' or arg.clock in ('ucln', 'ncln'):
             parameters.append(f'{branch_model_id}.rates')
     else:
         parameters = ['tree.blens']
@@ -999,7 +1016,11 @@ def build_advi(arg):
         if arg.coalescent_integrated is None:
             parameters.append("coalescent.theta")
 
-        if arg.coalescent in COALESCENT_PIECEWISE and not arg.gmrf_integrated:
+        if (
+            arg.coalescent in COALESCENT_PIECEWISE
+            and not arg.gmrf_integrated
+            and arg.theta_prior is None
+        ):
             parameters.append('gmrf.precision')
         elif arg.coalescent == 'exponential':
             parameters.append('coalescent.growth')

torchtree/cli/argparse_utils.py

@@ -31,6 +31,19 @@ def str_or_float(arg, choices):
             )
 
 
+def str_or_int(arg):
+    """Used by argparse when the argument can be either an integer or a string."""
+    try:
+        return int(arg)
+    except ValueError:
+        if isinstance(arg, str):
+            return arg
+        else:
+            raise argparse.ArgumentTypeError(
+                'invalid choice (choose from an integer or a string)'
+            )
+
+
 def list_of_float(arg, length):
     """Used by argparse when the argument should be a list of floats."""
     values = arg.split(",")