bdext 0.1.65__py3-none-any.whl → 0.1.66__py3-none-any.whl

bdeissct_dl/estimator.py

@@ -1,100 +1,13 @@
-import numpy as np
 import pandas as pd
 
 from bdeissct_dl import MODEL_PATH
-from bdeissct_dl.bdeissct_model import MODEL2TARGET_COLUMNS, BD, MODELS, \
-    MODEL_FINDER, F_S, X_S, X_C, UPSILON, UPS_X_C, F_S_X_S, F_E
+from bdeissct_dl.bdeissct_model import MODEL2TARGET_COLUMNS, BD, MODELS
 from bdeissct_dl.model_serializer import load_model_keras, load_scaler_numpy
 from bdeissct_dl.training import get_test_data
 from bdeissct_dl.tree_encoder import forest2sumstat_df, scale_back
 from bdeissct_dl.tree_manager import read_forest
 
 
-def predict_parameters_mf(forest_sumstats, model_name=MODEL_FINDER, model_path=MODEL_PATH):
-    n_forests = len(forest_sumstats)
-    n_models = len(MODELS)
-
-    if MODEL_FINDER == model_name:
-        import bdeissct_dl.training_model_finder
-        X = bdeissct_dl.training_model_finder.get_test_data(df=forest_sumstats)
-        model_weights = load_model_keras(model_path, model_name).predict(X)
-    else:
-        model_weights = np.zeros((n_forests, n_models), dtype=float)
-        model_weights[:, MODELS.index(model_name)] = 1
-
-    scaler_x = load_scaler_numpy(model_path, suffix='x')
-    X, SF = get_test_data(dfs=[forest_sumstats], scaler_x=scaler_x)
-
-    results = []
-
-    # result = pd.DataFrame(index=np.arange(X.shape[0]))
-
-    model_ids = [i for i in range(n_models) if not np.all(model_weights[:, i] == 0)]
-    for model_id in model_ids:
-        model_name = MODELS[model_id]
-
-        X_cur, SF_cur = np.array(X), np.array(SF)
-
-        model = load_model_keras(model_path, model_name)
-        Y_pred = model.predict(X_cur)
-
-        target_columns = MODEL2TARGET_COLUMNS[model_name]
-        if F_S_X_S in Y_pred:
-            if F_S in target_columns:
-                Y_pred[F_S] = Y_pred[F_S_X_S][:, 0]
-            if X_S in target_columns:
-                Y_pred[X_S] = Y_pred[F_S_X_S][:, 1]
-            del Y_pred[F_S_X_S]
-        if UPS_X_C in Y_pred:
-            if UPSILON in target_columns:
-                Y_pred[UPSILON] = Y_pred[UPS_X_C][:, 0]
-            if X_C in target_columns:
-                Y_pred[X_C] = Y_pred[UPS_X_C][:, 1]
-            del Y_pred[UPS_X_C]
-
-        for col in target_columns:
-            if len(Y_pred[col].shape) == 2 and Y_pred[col].shape[1] == 1:
-                Y_pred[col] = Y_pred[col].squeeze(axis=1)
-
-        scale_back(Y_pred, SF_cur)
-
-        results.append(pd.DataFrame.from_dict(Y_pred, orient='columns'))
-
-    if len(model_ids) == 1:
-        result = results[0]
-    else:
-        bdei_ids = {_[0] for _ in enumerate(model_ids) if 'EI' in MODELS[_[1]]}
-        bdss_ids = {_[0] for _ in enumerate(model_ids) if 'SS' in MODELS[_[1]]}
-        ct_ids = {_[0] for _ in enumerate(model_ids) if 'CT' in MODELS[_[1]]}
-
-        if ct_ids and len(ct_ids) < len(model_ids):
-            for idx in range(len(model_ids)):
-                if idx not in ct_ids:
-                    results[idx].loc[:, UPSILON] = 0
-                    results[idx].loc[:, X_C] = 1
-
-        if bdei_ids and len(bdei_ids) < len(model_ids):
-            for idx in range(len(model_ids)):
-                if idx not in bdei_ids:
-                    results[idx].loc[:, F_E] = 0
-
-        if bdss_ids and len(bdss_ids) < len(model_ids):
-            for idx in range(len(model_ids)):
-                if not idx in bdss_ids:
-                    results[idx].loc[:, F_S] = 0
-                    results[idx].loc[:, X_S] = 1
-
-        columns = results[0].columns
-        result = pd.DataFrame(index=forest_sumstats.index)
-        for col in columns:
-            predictions = np.array([res[col].to_numpy(dtype=float, na_value=0) for res in results]).T
-            weights = model_weights[:, model_ids]
-            result[col] = np.average(predictions, weights=weights, axis=1)
-
-    return result
-
-
-
 def predict_parameters(forest_sumstats, model_name=BD, model_path=MODEL_PATH):
     scaler_x = load_scaler_numpy(model_path, suffix='x')
     X, SF = get_test_data(dfs=[forest_sumstats], scaler_x=scaler_x)
@@ -106,20 +19,9 @@ def predict_parameters(forest_sumstats, model_name=BD, model_path=MODEL_PATH):
         model = load_model_keras(model_path, f'{model_name}.{col}')
         Y_pred = model.predict(X)
 
-
-        # if F_S in target_columns:
-        #     Y_pred[F_S] = Y_pred[F_S_X_S][:, 0]
-        #     Y_pred[X_S] = Y_pred[F_S_X_S][:, 1]
-        #     del Y_pred[F_S_X_S]
-        # if UPSILON in target_columns:
-        #     Y_pred[UPSILON] = Y_pred[UPS_X_C][:, 0]
-        #     Y_pred[X_C] = Y_pred[UPS_X_C][:, 1]
-        #     del Y_pred[UPS_X_C]
-
         if len(Y_pred[col].shape) == 2 and Y_pred[col].shape[1] == 1:
             Y_pred[col] = Y_pred[col].squeeze(axis=1)
 
-        print(Y_pred)
         scale_back(Y_pred, SF)
         res_df = pd.DataFrame.from_dict(Y_pred, orient='columns')
         result = result.join(res_df, how='outer') if result is not None else res_df
@@ -136,19 +38,17 @@ def main():
 
     parser = \
         argparse.ArgumentParser(description="Estimate BD(EI)(SS)(CT) model parameters.")
-    parser.add_argument('--model_name', choices=MODELS + (MODEL_FINDER,), default=BD, type=str,
-                        help=f'BDEISSCT model flavour. If {MODEL_FINDER} is specified, '
-                             f'model finder will be used to pick the model.')
+    parser.add_argument('--model_name', choices=MODELS, default=BD, type=str,
+                        help=f'BDEISSCT model flavour')
     parser.add_argument('--model_path', default=MODEL_PATH,
                         help='By default our pretrained BD(EI)(SS)(CT) models are used, '
                              'but it is possible to specify a path to a custom folder here, '
                              'containing files "<model_name>.keras" (with the model), '
-                             'and scaler-related files to rescale the input data X, and the output Y: '
-                             'for X: "data_scalerx_mean.npy", "data_scalerx_scale.npy", "data_scalerx_var.npy" '
+                             'and scaler-related files to rescale the input data X: '
+                             '"data_scalerx_mean.npy", "data_scalerx_scale.npy", "data_scalerx_var.npy" '
                              '(unpickled numpy-saved arrays), '
                              'and "data_scalerx_n_samples_seen.txt" '
-                             'a text file containing the number of examples in the training set). '
-                             'For Y the file names are the same, just x replaced by y, e.g., "data_scalery_mean.npy".'
+                             'a text file containing the number of examples in the training set).'
                         )
     parser.add_argument('--p', default=0, type=float, help='sampling probability')
     parser.add_argument('--log', default=None, type=str, help="output log file")
@@ -159,10 +59,10 @@ def main():
 
     if not params.sumstats:
         if params.p <= 0 or params.p > 1:
-            raise ValueError('The sampling probability must be grater than 0 and not greater than 1.')
+            raise ValueError('The sampling probability must be between 0 (exclusive) and 1 (inclusive).')
 
         forest = read_forest(params.nwk)
-        print(f'Read a forest of {len(forest)} trees with {sum(len(_) for _ in forest)} tips in total')
+        print(f'Read a tree with {sum(len(_) for _ in forest)} tips.')
         forest_df = forest2sumstat_df(forest, rho=params.p)
     else:
         forest_df = pd.read_csv(params.sumstats)

bdeissct_dl/model_serializer.py

@@ -1,8 +1,8 @@
 
-import tensorflow as tf
 import os
-import joblib
+
 import numpy as np
+import tensorflow as tf
 from sklearn.preprocessing import StandardScaler
 
 from bdeissct_dl.dl_model import relu_plus_one, half_sigmoid, loss_ct, loss_ss, CTLayer, SSLayer, loss_prob
@@ -18,8 +18,7 @@ def save_model_keras(model, path, model_name):
 def load_model_keras(path, model_name):
     tf.keras.config.enable_unsafe_deserialization()
     return tf.keras.models.load_model(os.path.join(path, f'{model_name}.keras'),
-                                      custom_objects={"loss_ct": loss_ct, "loss_ss": loss_ss, "loss_prob": loss_prob, \
-                                                      "relu_plus_one": relu_plus_one, "half_sigmoid": half_sigmoid, "CTLayer": CTLayer, "SSLayer": SSLayer})
+                                      custom_objects={"relu_plus_one": relu_plus_one, "half_sigmoid": half_sigmoid})
 
 def save_model_h5(model, path, model_name):
     model.save(os.path.join(path, f'{model_name}.h5'), overwrite=True, zipped=True)
@@ -38,35 +37,6 @@ def load_model_json(path, model_name):
     model.load_weights(os.path.join(path, f'{model_name}.weights.h5'))
     return model
 
-def save_model_onnx(model, path, model_name):
-    import tf2onnx
-    import onnx
-
-    input_signature = [tf.TensorSpec(model.inputs[0].shape, model.inputs[0].dtype, name='x')]
-    model.output_names = ['output']
-    onnx_model, _ = tf2onnx.convert.from_keras(model, input_signature=input_signature)
-    onnx.save(onnx_model, os.path.join(path, f'{model_name}.onnx'))
-
-def load_model_onnx(path, model_name):
-    """
-    TODO: this does not work due to onnx vs keras naming issues
-        (keras does not accept slashes in names that onnx creates)
-
-    :param path:
-    :return:
-    """
-    import onnx
-    from onnx2keras import onnx_to_keras
-    onnx_model = onnx.load(os.path.join(path, f'{model_name}.onnx'))
-    return onnx_to_keras(onnx_model, ['x'])
-
-def save_scaler_joblib(scaler, prefix, suffix=''):
-    joblib.dump(scaler, os.path.join(prefix, f'data_scaler{suffix}.gz'))
-
-def load_scaler_joblib(prefix, suffix=''):
-    return joblib.load(os.path.join(prefix, f'data_scaler{suffix}.gz')) \
-        if os.path.exists(os.path.join(prefix, f'data_scaler{suffix}.gz')) else None
-
 def save_scaler_numpy(scaler, prefix, suffix=''):
     np.save(os.path.join(prefix, f'data_scaler{suffix}_mean.npy'), scaler.mean_, allow_pickle=False)
     np.save(os.path.join(prefix, f'data_scaler{suffix}_scale.npy'), scaler.scale_, allow_pickle=False)

bdeissct_dl/scaler_fitting.py

@@ -4,9 +4,8 @@ import pandas as pd
 from sklearn.preprocessing import StandardScaler
 
 from bdeissct_dl import MODEL_PATH
-from bdeissct_dl.bdeissct_model import BD, BDCT, BDEI, BDEICT, \
-    BDSS, BDSSCT, BDEISS, BDEISSCT, TARGET_COLUMNS_BDEISSCT
-from bdeissct_dl.model_serializer import save_scaler_joblib, save_scaler_numpy
+from bdeissct_dl.bdeissct_model import TARGET_COLUMNS_BDEISSCT
+from bdeissct_dl.model_serializer import save_scaler_numpy
 from bdeissct_dl.training import get_data_characteristics
 
 
@@ -28,8 +27,7 @@ def main():
     parser = \
         argparse.ArgumentParser(description="Fit a BD(EI)(SS)(CT) data scaler.")
     parser.add_argument('--train_data', type=str, nargs='+',
-                        default=[f'/home/azhukova/projects/bdeissct_dl/simulations_bdeissct/training/500_1000/{model}/{i}/trees.csv.xz' for i in range(120) for model in [BD, BDCT, BDEI, BDEICT, BDSS, BDSSCT, BDEISS, BDEISSCT]]
-                        ,
+                        # default=[f'/home/azhukova/projects/bdeissct_dl/simulations_bdeissct/training/500_1000/{model}/{i}/trees.csv.xz' for i in range(120) for model in [BD, BDCT, BDEI, BDEICT, BDSS, BDSSCT, BDEISS, BDEISSCT]],
                         help="path to the files where the encoded training data are stored")
     parser.add_argument('--model_path', default=MODEL_PATH, type=str,
                         help="path to the folder where the scaler should be stored.")
@@ -43,7 +41,6 @@ def main():
     fit_scalers(paths=params.train_data, x_indices=x_indices, scaler_x=scaler_x)
 
     if scaler_x is not None:
-        save_scaler_joblib(scaler_x, params.model_path, suffix='x')
         save_scaler_numpy(scaler_x, params.model_path, suffix='x')
 
 

bdeissct_dl/sumstat_checker.py

@@ -1,5 +1,5 @@
 from bdeissct_dl import MODEL_PATH
-from bdeissct_dl.training import get_test_data, get_X_columns
+from bdeissct_dl.training import get_test_data, FEATURE_COLUMNS
 from bdeissct_dl.tree_encoder import forest2sumstat_df
 from bdeissct_dl.tree_manager import read_forest
 from bdeissct_dl.model_serializer import load_scaler_numpy
@@ -9,7 +9,7 @@ def check_sumstats(forest_sumstats, model_path=MODEL_PATH):
     scaler_x = load_scaler_numpy(model_path, suffix='x')
     X, SF = get_test_data(dfs=[forest_sumstats], scaler_x=scaler_x)
 
-    feature_columns = get_X_columns(forest_sumstats.columns)
+    feature_columns = FEATURE_COLUMNS
 
     for i in range(len(feature_columns)):
         value = X[0, i]

bdeissct_dl/training.py

@@ -6,8 +6,8 @@ import pandas as pd
 import tensorflow as tf
 
 from bdeissct_dl import MODEL_PATH, BATCH_SIZE, EPOCHS
-from bdeissct_dl.bdeissct_model import MODEL2TARGET_COLUMNS, UPSILON, X_C, KAPPA, F_E, F_S, \
-    X_S, TARGET_COLUMNS_BDCT, UPS_X_C, F_S_X_S, REPRODUCTIVE_NUMBER, INFECTION_DURATION, BDEI, LA, PSI, RHO
+from bdeissct_dl.bdeissct_model import MODEL2TARGET_COLUMNS, UPSILON, X_C, KAPPA, INCUBATION_PERIOD, F_S, \
+    X_S, TARGET_COLUMNS_BDCT, REPRODUCTIVE_NUMBER, INFECTION_DURATION
 from bdeissct_dl.dl_model import build_model
 from bdeissct_dl.model_serializer import save_model_keras, load_scaler_numpy, \
     load_model_keras
@@ -16,7 +16,7 @@ from bdeissct_dl.tree_encoder import SCALING_FACTOR, STATS
 FEATURE_COLUMNS = [_ for _ in STATS if _ not in {'n_trees', 'n_tips', 'n_inodes', 'len_forest',
                                                  REPRODUCTIVE_NUMBER, INFECTION_DURATION,
                                                  UPSILON, X_C, KAPPA,
-                                                 F_E,
+                                                 INCUBATION_PERIOD,
                                                  F_S, X_S,
                                                  SCALING_FACTOR}]
 
@@ -29,14 +29,10 @@ def calc_validation_fraction(m):
     return 0.01
 
 
-def get_X_columns(columns):
-    return FEATURE_COLUMNS
-
-
 def get_test_data(dfs=None, paths=None, scaler_x=None):
     if not dfs:
         dfs = [pd.read_csv(path) for path in paths]
-    feature_columns = get_X_columns(dfs[0].columns)
+    feature_columns = FEATURE_COLUMNS
 
     Xs, SFs = [], []
     for df in dfs:
@@ -53,12 +49,11 @@ def get_test_data(dfs=None, paths=None, scaler_x=None):
     return X, SF
 
 
-def get_data_characteristics(paths, target_columns=TARGET_COLUMNS_BDCT, feature_columns=None):
+def get_data_characteristics(paths, target_columns=TARGET_COLUMNS_BDCT, feature_columns=FEATURE_COLUMNS):
     col2index_y = {}
     col2index_x = {}
 
     df = pd.read_csv(paths[0])
-    feature_columns = get_X_columns(df.columns) if feature_columns is None else feature_columns
     feature_column_set = set(feature_columns)
     target_columns = target_columns if target_columns is not None else []
     target_column_set = set(target_columns)
@@ -109,36 +104,21 @@ def get_train_data(target_columns, columns_x, columns_y, file_pattern=None, file
     if INFECTION_DURATION in target_columns:
         train_labels[INFECTION_DURATION] = Y[:, col_i]
         col_i += 1
-    # if UPSILON in target_columns:
-    #     train_labels[UPS_X_C] = Y[:, col_i: (col_i + 2)]
-    #     col_i += 2
     if UPSILON in target_columns:
         train_labels[UPSILON] = Y[:, col_i]
         col_i += 1
     if X_C in target_columns:
         train_labels[X_C] = Y[:, col_i]
         col_i += 1
-    if F_E in target_columns:
-        train_labels[F_E] = Y[:, col_i]
+    if INCUBATION_PERIOD in target_columns:
+        train_labels[INCUBATION_PERIOD] = Y[:, col_i]
         col_i += 1
-    # if F_S in target_columns:
-    #     train_labels[F_S_X_S] = Y[:, col_i: (col_i + 2)]
-    #     col_i += 2
     if F_S in target_columns:
         train_labels[F_S] = Y[:, col_i]
         col_i += 1
     if X_S in target_columns:
         train_labels[X_S] = Y[:, col_i]
         col_i += 1
-    if LA in target_columns:
-        train_labels[LA] = Y[:, col_i]
-        col_i += 1
-    if PSI in target_columns:
-        train_labels[PSI] = Y[:, col_i]
-        col_i += 1
-    if RHO in target_columns:
-        train_labels[RHO] = Y[:, col_i]
-        col_i += 1
 
     dataset = tf.data.Dataset.from_tensor_slices((X, train_labels))
 
@@ -209,9 +189,9 @@ def main():
         print(model.summary())
 
         ds_train = get_train_data([col], x_indices, [y_idx], file_pattern=None, filenames=params.train_data, \
-                                  scaler_x=scaler_x, batch_size=BATCH_SIZE * 8, shuffle=True)
+                                  scaler_x=scaler_x, batch_size=BATCH_SIZE, shuffle=True)
         ds_val = get_train_data([col], x_indices, [y_idx], file_pattern=None, filenames=params.val_data, \
-                                scaler_x=scaler_x, batch_size=BATCH_SIZE * 8, shuffle=True)
+                                scaler_x=scaler_x, batch_size=BATCH_SIZE, shuffle=True)
 
         #early stopping to avoid overfitting
         early_stop = tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=25)
@@ -220,7 +200,6 @@ def main():
         model.fit(ds_train, verbose=1, epochs=params.epochs, validation_data=ds_val, callbacks=[early_stop])
 
         print(f'Saving the trained model {params.model_name}.{col} to {params.model_path}...')
-
         save_model_keras(model, path=params.model_path, model_name=f'{params.model_name}.{col}')
 
 

bdeissct_dl/tree_encoder.py

@@ -2,7 +2,6 @@ import io
 import os
 from glob import iglob
 
-import numpy as np
 import pandas as pd
 from treesumstats import FeatureCalculator, FeatureRegistry, FeatureManager
 from treesumstats.balance_sumstats import BalanceFeatureCalculator
@@ -13,8 +12,8 @@ from treesumstats.ltt_sumstats import LTTFeatureCalculator
 from treesumstats.subtree_sumstats import SubtreeFeatureCalculator
 from treesumstats.transmission_chain_sumstats import TransmissionChainFeatureCalculator
 
-from bdeissct_dl.bdeissct_model import RHO, LA, PSI, F_E, UPSILON, X_C, KAPPA, F_S, X_S, RATE_PARAMETERS, \
-    TIME_PARAMETERS, PI_E, PI_I, PI_S, PI_EC, PI_IC, PI_SC, LA_AVG, INFECTION_DURATION, REPRODUCTIVE_NUMBER
+from bdeissct_dl.bdeissct_model import RHO, UPSILON, X_C, KAPPA, F_S, X_S, RATE_PARAMETERS, \
+    TIME_PARAMETERS, INFECTION_DURATION, REPRODUCTIVE_NUMBER, INCUBATION_PERIOD
 from bdeissct_dl.tree_manager import read_forest, rescale_forest_to_avg_brlen
 
 TARGET_AVG_BL = 1
@@ -72,14 +71,14 @@ def parse_parameters(log):
         R = df.loc[i, REPRODUCTIVE_NUMBER]
         d = df.loc[i, INFECTION_DURATION]
         rho = df.loc[i, RHO]
-        f_e = df.loc[i, F_E] if F_E in df.columns else 0
+        d_inc = df.loc[i, INCUBATION_PERIOD] if INCUBATION_PERIOD in df.columns else 0
         f_ss = df.loc[i, F_S] if F_S in df.columns else 0
         x_ss = df.loc[i, X_S] if X_S in df.columns else 1
         upsilon = df.loc[i, UPSILON] if UPSILON in df.columns else 0
         x_c = df.loc[i, X_C] if X_C in df.columns else 1
         kappa = df.loc[i, KAPPA] if KAPPA in df.columns else 0
 
-        yield R, d, rho, f_e, f_ss, x_ss, upsilon, x_c, kappa
+        yield R, d, rho, d_inc, f_ss, x_ss, upsilon, x_c, kappa
 
 
 class BDEISSCTFeatureCalculator(FeatureCalculator):
@@ -89,7 +88,7 @@ class BDEISSCTFeatureCalculator(FeatureCalculator):
         pass
 
     def feature_names(self):
-        return [REPRODUCTIVE_NUMBER, INFECTION_DURATION, RHO, F_E, F_S, X_S, UPSILON, X_C, KAPPA, \
+        return [REPRODUCTIVE_NUMBER, INFECTION_DURATION, RHO, INCUBATION_PERIOD, F_S, X_S, UPSILON, X_C, KAPPA, \
                 SCALING_FACTOR]
 
     def set_forest(self, forest, **kwargs):
@@ -99,12 +98,6 @@ class BDEISSCTFeatureCalculator(FeatureCalculator):
         return kwargs[feature_name] if feature_name in kwargs else None
 
     def help(self, feature_name, *args, **kwargs):
-        if LA == feature_name:
-            return 'transmission rate.'
-        if LA_AVG == feature_name:
-            return 'average transmission rate.'
-        if PSI == feature_name:
-            return 'removal rate.'
         if RHO == feature_name:
             return 'sampling probability.'
         if UPSILON == feature_name:
@@ -117,26 +110,14 @@ class BDEISSCTFeatureCalculator(FeatureCalculator):
             return 'super-spreading ratio.'
         if F_S == feature_name:
             return 'fraction of super-spreaders.'
-        if F_E == feature_name:
-            return 'fraction of incubation over total infected-to-removed time.'
-        if PI_E == feature_name:
-            return 'fraction of unnotified exposed individuals'
-        if PI_EC == feature_name:
-            return 'fraction of notified exposed individuals'
-        if PI_I == feature_name:
-            return 'fraction of unnotified infectious regular spreaders'
-        if PI_IC == feature_name:
-            return 'fraction of notified infectious regular spreaders'
-        if PI_S == feature_name:
-            return 'fraction of unnotified infectious superpreaders'
-        if PI_SC == feature_name:
-            return 'fraction of notified infectious superspreaders'
         if SCALING_FACTOR == feature_name:
             return 'tree scaling factor.'
         if REPRODUCTIVE_NUMBER == feature_name:
             return 'reproduction number.'
         if INFECTION_DURATION == feature_name:
             return 'infection duration.'
+        if INCUBATION_PERIOD == feature_name:
+            return 'incubation period.'
         return None
 
 
@@ -252,14 +233,14 @@ TIME_DIFF_STATS = ['time_diff_in_2_real_mean', 'time_diff_in_3L_real_mean', 'tim
 
 EPI_STATS = [REPRODUCTIVE_NUMBER, INFECTION_DURATION, RHO,
              UPSILON, X_C, KAPPA,
-             F_E,
+             INCUBATION_PERIOD,
              F_S, X_S]
 
 STATS = ['n_tips'] \
         + BRLEN_STATS + TIME_STATS + CHAIN_STATS + LTT_STATS + BALANCE_STATS + TOPOLOGY_STATS + TIME_DIFF_STATS \
         + EPI_STATS + [SCALING_FACTOR]
 
-def forest2sumstat_df(forest, rho, R=0, d=0, x_c=0, upsilon=0, kappa=1, f_e=0, f_ss=0, x_ss=1,
+def forest2sumstat_df(forest, rho, R=0, d=0, x_c=0, upsilon=0, kappa=1, d_inc=0, f_ss=0, x_ss=1,
                       target_avg_brlen=TARGET_AVG_BL):
     """
     Rescales the input forest to have mean branch lengths of 1, calculates its summary statistics,
@@ -269,7 +250,7 @@ def forest2sumstat_df(forest, rho, R=0, d=0, x_c=0, upsilon=0, kappa=1, f_e=0, f
     :param x_ss: presumed superspreading ratio (how many times superspreader's transmission rate is higher
         than that of a standard spreader, 1 by default)
     :param f_ss: presumed fraction of superspreaders in the infectious population (0 by default)
-    :param f_e: presumed fraction of incubation over total infected-to-removed time (0 by default)
+    :param d_inc: presumed incubation period length (0 by default)
     :param forest: list(ete3.Tree) forest to encode
     :param rho: presumed sampling probability
     :param upsilon: presumed notification probability
@@ -286,7 +267,7 @@ def forest2sumstat_df(forest, rho, R=0, d=0, x_c=0, upsilon=0, kappa=1, f_e=0, f
 
     kwargs = {SCALING_FACTOR: scaling_factor,
               REPRODUCTIVE_NUMBER: R, INFECTION_DURATION: d, RHO: rho,
-              F_E: f_e,
+              INCUBATION_PERIOD: d_inc,
               F_S: f_ss, X_S: x_ss,
               X_C: x_c, UPSILON: upsilon, KAPPA: kappa}
     scale(kwargs, scaling_factor)
@@ -337,11 +318,11 @@ def save_forests_as_sumstats(output, nwks=None, logs=None, patterns=None, target
             for ps, forest in zip(parameters, forests):
 
                 scaling_factor = rescale_forest_to_avg_brlen(forest, target_avg_length=target_avg_brlen)
-                R, d, rho, f_e, f_ss, x_ss, upsilon, x_c, kappa = ps
+                R, d, rho, d_inc, f_ss, x_ss, upsilon, x_c, kappa = ps
                 kwargs = {SCALING_FACTOR: scaling_factor}
                 kwargs[REPRODUCTIVE_NUMBER], kwargs[INFECTION_DURATION], kwargs[RHO] = R, d, rho
                 kwargs[UPSILON], kwargs[KAPPA], kwargs[X_C] = upsilon, kappa, x_c
-                kwargs[F_E] = f_e
+                kwargs[INCUBATION_PERIOD] = d_inc
                 kwargs[F_S], kwargs[X_S] = f_ss, x_ss
 
                 scale(kwargs, scaling_factor)