Add demes to gadma and option to change dynamics domain and choose en…

…gine to draw

gadma/__init__.py

@@ -29,6 +29,15 @@
     import dadi
 except ImportError:
     dadi = None
+try:
+    import demes
+except ImportError:
+    demes = None
+try:
+    import demesdraw
+except ImportError:
+    demesdraw = None
+
 try:
     import GPy
 except ImportError:
@@ -38,9 +47,9 @@
 except ImportError:
     GPyOpt = None
 try:
-    import smac
-    import ConfigSpace
-    import bayesmark
+    import smac  # NOQA
+    import ConfigSpace  # NOQA
+    import bayesmark  # NOQA
     smac_available = True
 except ImportError:
     smac_available = False
@@ -51,11 +60,16 @@
 matplotlib_available = matplotlib is not None
 moments_available = moments is not None
 dadi_available = dadi is not None
+demes_available = demes is not None
+demesdraw_available = demesdraw is not None
+
 GPy_available = GPy is not None
 GPyOpt_available = GPyOpt is not None
 
 from .data import DataHolder, SFSDataHolder, VCFDataHolder  # NOQA
 from .engines import get_engine, all_engines  # NOQA
+from .engines import all_available_engines, all_simulation_engines  # NOQA
+from .engines import all_drawing_engines  # NOQA
 
 from .models import DemographicModel, EpochDemographicModel  # NOQA
 from .models import CustomDemographicModel, StructureDemographicModel  # NOQA

gadma/cli/settings.py

@@ -25,6 +25,8 @@
     engine = 'moments'
 elif dadi_available:
     engine = 'dadi'
+model_plot_engine = "moments"
+sfs_plot_engine = None  # None means the same as engine
 relative_parameters = False
 no_migrations = False
 symmetric_migrations = False
@@ -101,6 +103,7 @@
 max_t = TimeVariable.default_domain[1]
 min_m = MigrationVariable.default_domain[0]
 max_m = MigrationVariable.default_domain[1]
+dynamics = list(DynamicVariable.default_domain)
 
 # Parameters for local search alg
 # ls_verbose = None

gadma/cli/settings_storage.py

@@ -4,13 +4,14 @@
 from . import settings
 from ..data import SFSDataHolder
 from ..engines import get_engine, MomentsEngine
+from ..engines import all_engines, all_drawing_engines
 from ..models import StructureDemographicModel, CustomDemographicModel
 from ..optimizers import get_local_optimizer, get_global_optimizer
 from ..optimizers import LinearConstrain
 from ..utils import check_dir_existence, check_file_existence, abspath,\
     module_name_from_path, custom_generator
 from ..utils import PopulationSizeVariable, TimeVariable, MigrationVariable,\
-    ContinuousVariable
+    ContinuousVariable, DynamicVariable
 import warnings
 import importlib.util
 import sys
@@ -125,13 +126,15 @@ def __setattr__(self, name, value):
         empty_dir_attrs = ['output_directory']
         data_holder_attrs = ['projections', 'outgroup',
                              'population_labels', 'sequence_length']
-        bounds_attrs = ['min_n', 'max_n', 'min_t', 'max_t', 'min_m', 'max_m']
+        bounds_attrs = ['min_n', 'max_n', 'min_t', 'max_t', 'min_m', 'max_m',
+                        'dynamics']
         bounds_lists = ['lower_bound', 'upper_bound', 'parameter_identifiers']
         missed_attrs = ['engine', 'global_optimizer', 'local_optimizer',
                         '_inner_data', '_bootstrap_data', 'X_init', 'Y_init',
                         'model_func', 'get_engine_args', 'data_holder',
                         'units_of_time_in_drawing', 'resume_from_settings',
-                        'dadi_available', 'moments_available']
+                        'dadi_available', 'moments_available',
+                        'model_plot_engine', 'sfs_plot_engine']
 
         super_hasattr = True
         setattr_at_the_end = True
@@ -344,7 +347,22 @@ def __setattr__(self, name, value):
                 setattr(self.data_holder, name, value)
         # 3.3 For engine we need check it exists
         elif name == 'engine':
-            get_engine(value)
+            engine = get_engine(value)
+            if not engine.can_evaluate:
+                raise ValueError(f"Engine {value} cannot evaluate "
+                                 f"log-likelihood. Available engines are: "
+                                 f"{[engine.id in all_engines()]}")
+        elif name == 'model_plot_engine':
+            engine = get_engine(value)
+            if not engine.can_draw:
+                raise ValueError(f"Engine {value} cannot draw model plots. "
+                                 f"Available engines are: "
+                                 f"{[engine.id in all_drawing_engines()]}")
+        elif name == 'sfs_plot_engine' and value is not None:
+            engine = get_engine(value)
+            if not hasattr(engine, "draw_sfs_plots"):
+                raise ValueError(f"Engine {value} cannot draw sfs plots. "
+                                 f"Available engines are: dadi, moments")
         # 3.4 For local and global optimizer we need check existence
         elif name == 'global_optimizer':
             get_global_optimizer(value)
@@ -444,6 +462,13 @@ def __setattr__(self, name, value):
                 cls = TimeVariable
             elif name.endswith('m'):
                 cls = MigrationVariable
+            elif name == "dynamics":
+                cls = DynamicVariable
+                if isinstance(value, str):
+                    value = [x.strip() for x in value.split(",")]
+                    for i in range(len(value)):
+                        if value[i].isdigit():
+                            value[i] = int(value[i])
             else:
                 raise AttributeError("Check for supported variables")
 
@@ -452,8 +477,17 @@ def __setattr__(self, name, value):
                 cls.default_domain = [value, cls.default_domain[1]]
             elif name.startswith('max'):
                 cls.default_domain = [cls.default_domain[0], value]
+            else:
+                assert name == "dynamics"
+                for dyn in value:
+                    if dyn not in cls._help_dict:
+                        raise ValueError(f"Unknown dynamic {value}. Available "
+                                         "dynamics are: "
+                                         f"{list(cls._help_dict.keys())}")
+                cls.default_domain = value
 
-            domain_changed = np.any(old_domain != np.array(cls.default_domain))
+            domain_changed = len(old_domain) != len(cls.default_domain) or\
+                np.any(old_domain != np.array(cls.default_domain))
             if domain_changed:
                 if name.endswith('n'):
                     warnings.warn(f"Domain of PopulationSizeVariable changed "
@@ -464,6 +498,9 @@ def __setattr__(self, name, value):
                 if name.endswith('m'):
                     warnings.warn(f"Domain of MigrationVariable changed to "
                                   f"{cls.default_domain}")
+                if name == "dynamics":
+                    warnings.warn(f"Domain of DynamicVariable changed to "
+                                  f"{cls.default_domain}")
 
         # 3.10 If we set custom filename with model we should check it is
         # valid python code
@@ -966,8 +1003,10 @@ def get_engine_args(self, engine_id=None):
             engine_id = self.engine
         if engine_id == 'dadi':
             args = (self.pts,)
-        else:
+        elif engine_id == "moments":
             args = (MomentsEngine.default_dt_fac,)
+        else:
+            args = ()
         return args
 
     def get_linear_constrain_for_model(self, model):

gadma/core/core_run.py

@@ -1,4 +1,4 @@
-from ..engines import get_engine, all_engines
+from ..engines import get_engine, all_available_engines
 from ..utils import sort_by_other_list, ensure_dir_existence,\
                     ensure_file_existence, check_file_existence,\
                     check_dir_existence
@@ -7,6 +7,7 @@
 from ..utils import get_aic_score, get_claic_score, ident_transform, bcolors
 from ..models import EpochDemographicModel, StructureDemographicModel
 from .draw_and_generate_code import draw_plots_to_file, generate_code_to_file
+from .draw_and_generate_code import get_Nanc_gen_time_and_units
 from ..cli import SettingsStorage
 import os
 import numpy as np
@@ -103,7 +104,15 @@ def __init__(self, index, shared_dict, settings):
                 # If our model is built via gadma then we can write code for
                 # all engines.
                 if isinstance(self.model, EpochDemographicModel):
-                    for engine in all_engines():
+                    demes_will_not_work = False
+                    mu_and_L = self.engine.model.mu is not None and \
+                        self.settings.sequence_length is not None
+                    if not (self.engine.model.has_anc_size or
+                            self.engine.model.theta0 is not None or mu_and_L):
+                        demes_will_not_work = True
+                    for engine in all_available_engines():
+                        if engine.id == "demes" and demes_will_not_work:
+                            continue
                         engine_dir = os.path.join(self.code_dir, engine.id)
                         ensure_dir_existence(engine_dir)
         # Set counters to zero for callbacks to count number of their calls
@@ -182,7 +191,7 @@ def base_callback(self, x, y):
             prefix = (self.settings.LOCAL_OUTPUT_DIR_PREFIX +
                       self.settings.LONG_NAME_2_SHORT.get(best_by, best_by))
             save_code_file = os.path.join(self.output_dir,
-                                          prefix + "_model.py")
+                                          prefix + "_model")
             try:
                 generate_code_to_file(x, self.engine,
                                       self.settings, save_code_file)
@@ -222,21 +231,34 @@ def code_iter_callback(self, x, y):
         """
         n_iter = self.code_iter_callback_counter
         verbose = self.settings.print_models_code_every_n_iteration
-        filename = f"iteration_{n_iter}.py"
+        filename = f"iteration_{n_iter}"
         if verbose != 0 and n_iter % verbose == 0:
+            Nanc, gen_time, gen_time_units = get_Nanc_gen_time_and_units(
+                x=x,
+                engine=self.engine,
+                settings=self.settings,
+            )
             if isinstance(self.model, EpochDemographicModel):
-                for engine in all_engines():
+                for engine_id in os.listdir(self.code_dir):
+                    engine = get_engine(engine_id)
                     save_file = os.path.join(self.code_dir, engine.id,
                                              filename)
                     args = self.settings.get_engine_args(engine.id)
                     engine.set_data(self.engine.data)
                     engine.data_holder = self.engine.data_holder
                     engine.set_model(self.engine.model)
-                    engine.generate_code(x, save_file, *args)
+                    try:
+                        engine.generate_code(x, save_file, *args, Nanc,
+                                             gen_time=gen_time,
+                                             gen_time_units=gen_time_units)
+                    except Exception as e:
+                        pass
             else:
                 save_file = os.path.join(self.code_dir, filename)
                 args = self.settings.get_engine_args()
-                self.engine.generate_code(x, save_file, *args)
+                self.engine.generate_code(x, save_file, *args,  Nanc,
+                                          gen_time=gen_time,
+                                          gen_time_units=gen_time_units)
         self.code_iter_callback_counter += 1
 
     def callback(self, x, y):
@@ -257,10 +279,7 @@ def callback(self, x, y):
             self.draw_iter_callback(x, y)
         except Exception:
             pass
-        try:
-            self.code_iter_callback(x, y)
-        except Exception:
-            pass
+        self.code_iter_callback(x, y)
 
     def intermediate_callback(self, x, y):
         """
@@ -312,19 +331,20 @@ def intermediate_callback(self, x, y):
                   self.settings.LONG_NAME_2_SHORT.get(best_by.lower(),
                                                       best_by.lower()))
         save_plot_file = os.path.join(self.output_dir, prefix + "_model.png")
-        save_code_file = os.path.join(self.output_dir, prefix + "_model.py")
+        save_code_file = os.path.join(self.output_dir, prefix + "_model")
         try:
             draw_plots_to_file(x, self.engine, self.settings,
                                save_plot_file, fig_title)
         except Exception as e:
-            print(f"{bcolors.FAIL}Run {self.index}: failed to draw model due "
-                  f"to the following exception: {e}{bcolors.ENDC}")
+            raise e
+            print(f"{bcolors.WARNING}Run {self.index}: failed to draw model "
+                  f"due to the following exception: {e}{bcolors.ENDC}")
         try:
             generate_code_to_file(x, self.engine,
                                   self.settings, save_code_file)
         except Exception as e:
-            print(f"{bcolors.FAIL}Run {self.index}: failed to generate code "
-                  f"due to the following exception: {e}{bcolors.ENDC}")
+            print(f"{bcolors.WARNING}Run {self.index}: failed to generate code"
+                  f" due to the following exception: {e}{bcolors.ENDC}")
 
     def draw_model_in_output_dir(self, x, y,
                                  best_by='log-likelihood', final=True):
@@ -346,7 +366,7 @@ def draw_model_in_output_dir(self, x, y,
             prefix = self.settings.LOCAL_OUTPUT_DIR_PREFIX
         prefix += self.settings.LONG_NAME_2_SHORT.get(best_by, best_by)
         save_plot_file = os.path.join(self.output_dir, prefix + "_model.png")
-        save_code_file = os.path.join(self.output_dir, prefix + "_model.py")
+        save_code_file = os.path.join(self.output_dir, prefix + "_model")
         try:
             generate_code_to_file(x, self.engine,
                                   self.settings, save_code_file)