Models.py

# -*- coding: utf-8 -*-
"""
Created on Mon Oct  7 17:15:15 2024
"""

import os
dir_path = os.path.dirname(os.path.realpath(__file__))
os.chdir(dir_path)

# IMPORTING USEFUL LIBRARIES

import sys
import platform
import scipy
import datetime
import pandas as pd
import pymc as pm
import numpy as np

import arviz as az
from scipy.special import logit
import xarray as xr

from scipy.io import netcdf
from Utils import *
from Data import *

def Predation_Model(df_dig, df_pred,
                    Correction_Barber=None):
    """
    Parameters
    ----------
    df_dig: Dataframe
        Data from digestion experiment under controlled conditions.
    df_pred : Dataframe
        Prey detection data from carabid stomach contents in the field.
    Correction_Barber: STR,
        Correcting the time lag between the end of the digestion processes and predation. Here the 
        correction can be either "Average" or "Random" depending on the use of the average time lag
        or the explicit consideration of time lag as a random variable.

    Returns
    -------
    model: pyMC Model Object
        A pyMC model ready for training.
    """
    
    # PRE-PROCESSING OBSERVATIONS
    # -------------------------------------------------------------------------------------------------------------------------------------------------

    if len(pd.unique(df_dig["Couple"])) != len(pd.unique(df_dig["Predator"]))*len(pd.unique(df_dig["Prey_fed"])) :
        
        for prey in pd.unique(df_dig["Prey_fed"]) :
            for predator in pd.unique(df_dig["Predator"]) :
                
                if not ("_".join([predator, prey]) in pd.unique(df_dig["Couple"])) :
                    
                    couple_row = {'Couple': "_".join([predator, prey]), 
                                  'Predator': predator, 
                                  'digestion_time': 0,
                                  'Prey_fed': prey,
                                  'Detection': 1}
                    
                    df_row = pd.DataFrame(couple_row, index=[0])
                    df_dig = pd.concat([df_dig, df_row], ignore_index=True, sort=False)
                                    
    observations_predation_idx = np.asarray(df_pred.index)
    observations_digestion_idx = np.asarray(df_dig.index)

    time_idx, times = pd.factorize(df_dig["digestion_time"], sort=True)

    preys_digestion_idx, preys = pd.factorize(df_dig["Prey_fed"], sort=True)
    preys_predation_idx, _ = pd.factorize(df_pred["Prey"], sort=True)

    predators_digestion_idx, predators = pd.factorize(df_dig["Predator"], sort=True)
    predators_predation_idx, _ = pd.factorize(df_pred["Predator"], sort=True)

    couples_digestion_idx, couples = pd.factorize(df_dig["Couple"], sort=True)
    couples_predation_idx, _ = pd.factorize(df_pred["Couple"], sort=True)

    preys_de_couple_idx = np.asarray(
        [list(preys).index(couple.split("_")[1]) for couple in couples])
    predators_de_couple_idx = np.asarray(
        [list(predators).index(couple.split("_")[0]) for couple in couples])
    couples_de_couple_idx = np.asarray([list(couples).index(
        couple) for couple in couples])

    coords = {"observations_predation": observations_predation_idx,
              "observations_digestion": observations_digestion_idx,
              "preys": preys,
              "predators": predators,
              "couples": couples}

    # MODEL DEFINITION
    # -------------------------------------------------------------------------------------------------------------------------------------------------

    with pm.Model(coords=coords) as model:

        # Converting observed variables into np.array format
        test_predation_obs = pm.ConstantData("test_predation_obs", df_pred["Detection"],
                                             dims="observations_predation")
        test_digestion_obs = pm.ConstantData("test_obs", df_dig["Detection"],
                                             dims="observations_digestion")
        temps_digestion = pm.ConstantData("temps_digestion", df_dig["digestion_time"],
                                          dims="observations_digestion")

        # HYPERPRIORS

        # - On the intercept beta0 (intercept mu_beta0, prey effect alpha_beta0_p, 
        # predator effect delta_beta0_c, and effect of the gamma_beta0_pc couple)

        # prey effect
        alpha_beta0_p = pm.Normal("alpha_beta0_p", mu=0, tau=1, dims="preys")

        # predator effect
        mu_delta_beta0_c = pm.Normal("mu_delta_beta0_c", mu=0, tau=100)
        log_sigma_delta_beta0_c = pm.Uniform("log_sigma_delta_beta0_c", 0, 1)
        tau_delta_beta0_c = 1/np.exp(log_sigma_delta_beta0_c)

        # prey-predator interaction effect
        mu_gamma_beta0_pc = pm.Normal("mu_gamma_beta0_pc", mu=0, tau=100)
        log_sigma_gamma_beta0_pc = pm.Uniform("log_sigma_gamma_beta0_pc", 0, 1)
        tau_gamma_beta0_pc = 1/np.exp(log_sigma_gamma_beta0_pc)

        # On the slope of the beta1 digestion curve (mean value mu_beta1, and variability tau_beta1)
        
        mu_log_beta1 = pm.Normal("mu_log_beta1", mu=-2, tau=100)
        sigma_log_beta1 = pm.Uniform("sigma_log_beta1", 0, 1)
        tau_log_beta1 = 1/np.exp(sigma_log_beta1)

        # On the log of predation rate lambda (intecept mu_loglambda, prey effect 
        # alpha_loglambda_p, predator effect delta_loglambda_c,
        
        # prey effect
        alpha_loglambda_p = pm.Normal("alpha_loglambda_p", mu=0, tau=0.05, dims="preys")

        # predator effect
        mu_delta_loglambda_c = pm.Normal("mu_delta_loglambda_c", mu=0, tau=100)
        log_sigma_delta_loglambda_c = pm.Uniform("log_sigma_delta_loglambda_c", 0, 100)
        tau_delta_loglambda_c = 1/np.exp(log_sigma_delta_loglambda_c)

        # predator-prey interaction effect
        mu_gamma_loglambda_pc = pm.Normal("mu_gamma_loglambda_pc", mu=0, tau=100)
        log_sigma_gamma_loglambda_pc = pm.Uniform("log_sigma_gamma_loglambda_pc", 0, 100)
        tau_gamma_loglambda_pc = 1/np.exp(log_sigma_gamma_loglambda_pc)

        # PRIORS

        # On digestion curve parameter beta0
        delta_beta0_c = pm.Normal("delta_beta0_c", mu=mu_delta_beta0_c, 
                                  tau=tau_delta_beta0_c, dims="predators")
        
        gamma_beta0_pc = pm.Normal("gamma_beta0_pc", mu=mu_gamma_beta0_pc, 
                                   tau=tau_gamma_beta0_pc, dims="couples")

        # On predation rate
        delta_loglambda_c = pm.Normal("delta_loglambda_c", mu=mu_delta_loglambda_c,
                                      tau=tau_delta_loglambda_c, dims="predators")

        gamma_loglambda_pc = pm.Normal("gamma_loglambda_pc", mu=mu_gamma_loglambda_pc,
                                       tau=tau_gamma_loglambda_pc, dims="couples")

        # MODEL

        # On digestion
        
        # Beta1, ~ slope
        log_beta1_pc = pm.Normal("log_beta1_pc", mu=mu_log_beta1, 
                                  tau=tau_log_beta1, dims="couples")
            
        beta1_pc = pm.Deterministic("beta1_pc", (-1)*np.exp(log_beta1_pc), 
                                    dims="couples")
        
        # Beta0, ~ intercept
        beta0_pc = pm.Deterministic("beta0_pc", alpha_beta0_p[preys_de_couple_idx]
                                    + delta_beta0_c[predators_de_couple_idx]
                                    + gamma_beta0_pc[couples_de_couple_idx],
                                    dims="couples")
        
        # Resulting probability of detection as a function of time
        Pi_t = pm.Deterministic("Pi_t", beta0_pc[couples_digestion_idx]
                                + beta1_pc[couples_digestion_idx] *
                                temps_digestion,
                                dims="observations_digestion")

        # On predation

        lambda_pc = pm.Deterministic("lambda_pc", np.exp(alpha_loglambda_p[preys_de_couple_idx]
                                                         + delta_loglambda_c[predators_de_couple_idx]
                                                         + gamma_loglambda_pc[couples_de_couple_idx]),
                                     dims="couples")  # Intensité de predation (par heure).

        # On the combination of the two processes, with or without considering time lag in the Barber
        # trap
        
        Tmax = 2000
        if Correction_Barber == "Random" or "random":
            
            # In this version, the integral I_pc is calculated explicitly (no analytical formula).
            # In reality, we calculate the expectation of the integral I_pc.
            #
            # The calculation of the integral I_pc knowing T0 = t0 is as follows:
            # E(I_pc|T0 = t0) = scipy.integrate.quad(lambda t : inv_logit(beta0_pc + beta1_pc*t), t0, Tmax)
            # Then :
            # E(I_pc) = E(I_pc|T0 = t0)*P(T0 = t0) = scipy.integrate.quad(lambda t0 : (1/12) *
            # scipy.integrate.quad(lambda t : inv_logit(beta0_pc + beta1_pc*t), t0, Tmax), 0, 12)
            # Note in particular that E(I_pc) != scipy.integrate.quad(lambda t : inv_logit(beta0_pc + 
            # beta1_pc*t), E(TO), Tmax).
            
            # The computation of the integral being to slow, we used a Here's a discretized version 
            # with 13 possible values for T0.
            
            I_pc = pm.Deterministic("I_pc", pm.math.maximum(0, sum([(1/13*beta1_pc) *
                                                                (pm.math.log(1 + pm.math.exp(beta0_pc + beta1_pc*Tmax))
                                                                 - pm.math.log(1 + pm.math.exp(beta0_pc + beta1_pc*t0)))
                                                                for t0 in range(12)])), 
                                    dims="couples")
            
        elif Correction_Barber == "Average" or "average":
            
            T0 = 6
            I_pc = pm.Deterministic("I_pc", (1/beta1_pc)*(np.log(1 + np.exp(beta0_pc + beta1_pc*Tmax))
                                                      - np.log(1 + np.exp(beta0_pc + beta1_pc*T0))),
                                dims="couples")  # Digestion curve integral between t=0 et t=Tmax.
        
        else:
            
            T0 = 0
            I_pc = pm.Deterministic("I_pc", (1/beta1_pc)*(np.log(1 + np.exp(beta0_pc + beta1_pc*Tmax))
                                                      - np.log(1 + np.exp(beta0_pc + beta1_pc*T0))),
                                dims="couples")  # Digestion curve integral between t=0 et t=Tmax.
            
        pI = pm.Deterministic("pI", 1 - pm.math.maximum(0.001,
                                               np.exp(-lambda_pc[couples_predation_idx] *
                                                      I_pc[couples_predation_idx])),
                              dims="observations_predation")

        # OBSERVATION PROCESSES

        test_predation_pred = pm.Bernoulli("test_predation_pred", p=pI, observed=test_predation_obs,
                                           dims="observations_predation")
        test_digestion_pred = pm.Bernoulli("test_dig_pred", logit_p=Pi_t, observed=test_digestion_obs,
                                           dims="observations_digestion")

    return model

# -------------------------------------------------------------------------------------------------------------------------
# MODEL CALIBRATION

if __name__ == "__main__" :
    
    # Importing and pre-processing data

    dataset = Data(import_ = True,
                   preprocess_ = True)
    
    df_dig = dataset.get_data(data_type = "lab")
    df_pred = dataset.get_data(data_type = "field")
        
    # Model definition

    model = Predation_Model(df_dig, df_pred,
                            Correction_Barber='random')

    # Model calibration
       
    Sample = pm.sample(draws=50000,
                       chains=3,
                       tune=20000,
                       progressbar=True,
                       model=model,
                       compute_convergence_checks=True)
        
    # Saving the Sample
        
    Path_to_InferenceData = '/Samples/'
    InferenceData_tostock = 'Sample_{}.nc'.format(str(datetime.datetime.now().strftime("%Y%m%d%H%M")))

    DataSet_Sample = az.convert_to_dataset(Sample)
    DataSet_Sample.to_netcdf("".join([Path_to_InferenceData, InferenceData_tostock]))