model_utils.py

''' County and State Data Processing network'''

import numpy as np
import torch
import torch.nn as nn
from data_utils import get_state_train_data_flu, create_window_seqs, get_county_train_data, counties
from sklearn.preprocessing import StandardScaler
import torch.nn as nn
import numpy as np
import torch
import pdb 
import math
import pandas as pd
import os
import yaml
cuda = torch.device('cuda')
dtype = torch.float
SMOOTH_WINDOW = 7

class TransformerAttn(nn.Module):
    """
    Module that calculates self-attention weights using transformer like attention
    """

    def __init__(self, dim_in=40, value_dim=40, key_dim=40) -> None:
        """
        param dim_in: Dimensionality of input sequence
        param value_dim: Dimension of value transform
        param key_dim: Dimension of key transform
        """
        super(TransformerAttn, self).__init__()
        self.value_layer = nn.Linear(dim_in, value_dim)
        self.query_layer = nn.Linear(dim_in, value_dim)
        self.key_layer = nn.Linear(dim_in, key_dim)

    def forward(self, seq):
        """
        param seq: Sequence in dimension [Seq len, Batch, Hidden size]
        """
        seq_in = seq.transpose(0, 1)
        value = self.value_layer(seq_in)
        query = self.query_layer(seq_in)
        keys = self.key_layer(seq_in)
        weights = (value @ query.transpose(1, 2)) / math.sqrt(seq.shape[-1])
        weights = torch.softmax(weights, -1)
        return (weights @ keys).transpose(1, 0)

    def forward_mask(self, seq, mask):
        """
        param seq: Sequence in dimension [Seq len, Batch, Hidden size]
        """
        seq_in = seq.transpose(0, 1)
        value = self.value_layer(seq_in)
        query = self.query_layer(seq_in)
        keys = self.key_layer(seq_in)
        weights = (value @ query.transpose(1, 2)) / math.sqrt(seq.shape[-1])
        weights = torch.exp(weights)
        weights = (weights.transpose(1, 2) * mask.transpose(1, 0)).transpose(1, 2)
        weights = weights / (weights.sum(-1, keepdim=True))
        return (weights @ keys).transpose(1, 0) * mask

class EmbedAttenSeq(nn.Module):
    """
    Module to embed a sequence. Adds Attention modul
    """

    def __init__(
        self,
        dim_seq_in: int = 5,
        dim_metadata: int = 3,
        rnn_out: int = 40,
        dim_out: int = 50,
        n_layers: int = 1,
        bidirectional: bool = False,
        attn=TransformerAttn,
        dropout=0.0,
    ) -> None:
        """
        param dim_seq_in: Dimensionality of input vector (no. of age groups)
        param dim_out: Dimensionality of output vector
        param dim_metadata: Dimensions of metadata for all sequences
        param rnn_out: output dimension for rnn
        """
        super(EmbedAttenSeq, self).__init__()

        self.dim_seq_in = dim_seq_in
        self.dim_metadata = dim_metadata
        self.rnn_out = rnn_out
        self.dim_out = dim_out
        self.bidirectional = bidirectional

        self.rnn = nn.GRU(
            input_size=self.dim_seq_in,
            hidden_size=self.rnn_out // 2 if self.bidirectional else self.rnn_out,
            bidirectional=bidirectional,
            num_layers=n_layers,
            dropout=dropout,
        )
        self.attn_layer = attn(self.rnn_out, self.rnn_out, self.rnn_out)
        self.out_layer = [
            nn.Linear(
                in_features=self.rnn_out + self.dim_metadata, out_features=self.dim_out
            ),
            nn.Tanh(),
            nn.Dropout(dropout),
        ]
        self.out_layer = nn.Sequential(*self.out_layer)

        def init_weights(m):
            if isinstance(m, nn.Linear):
                torch.nn.init.xavier_uniform_(m.weight)
                m.bias.data.fill_(0.01)
        self.out_layer.apply(init_weights)

    def forward_mask(self, seqs, metadata, mask):
        # Take last output from GRU
        latent_seqs = self.rnn(seqs)[0]
        latent_seqs = latent_seqs
        latent_seqs = self.attn_layer.forward_mask(latent_seqs, mask)
        latent_seqs = latent_seqs.sum(0)
        out = self.out_layer(torch.cat([latent_seqs, metadata], dim=1))
        return out

    def forward(self, seqs, metadata):
        # Take last output from GRU
        latent_seqs, encoder_hidden = self.rnn(seqs)
        latent_seqs = self.attn_layer(latent_seqs).sum(0)
        out = self.out_layer(torch.cat([latent_seqs, metadata], dim=1))
        return out, encoder_hidden


class DecodeSeq(nn.Module):
    """
    Module to embed a sequence. Adds Attention modul
    """

    def __init__(
        self,
        dim_seq_in: int = 5,
        dim_metadata: int = 3,
        rnn_out: int = 40,
        dim_out: int = 5,
        n_layers: int = 1,
        bidirectional: bool = False,
        dropout=0.0,
    ) -> None:
        """
        param dim_seq_in: Dimensionality of input vector (no. of age groups)
        param dim_out: Dimensionality of output vector
        param dim_metadata: Dimensions of metadata for all sequences
        param rnn_out: output dimension for rnn
        """
        super(DecodeSeq, self).__init__()

        self.dim_seq_in = dim_seq_in
        self.dim_metadata = dim_metadata
        self.rnn_out = rnn_out
        self.dim_out = dim_out
        self.bidirectional = bidirectional

        self.act_fcn = nn.Tanh()

        # to embed input
        self.embed_input = nn.Linear(self.dim_seq_in, self.rnn_out) 

        # to combine input and context
        self.attn_combine = nn.Linear(2*self.rnn_out, self.rnn_out)

        self.rnn = nn.GRU(
            input_size=self.rnn_out,
            hidden_size=self.rnn_out // 2 if self.bidirectional else self.rnn_out,
            bidirectional=bidirectional,
            num_layers=n_layers,
            dropout=dropout,
        )
        self.out_layer = [
            nn.Linear(
                in_features=self.rnn_out, out_features=self.dim_out
            ),
            nn.Tanh(),
            nn.Dropout(dropout),
        ]
        self.out_layer = nn.Sequential(*self.out_layer)

        # initialize
        def init_weights(m):
            if isinstance(m, nn.Linear):
                torch.nn.init.xavier_uniform_(m.weight)
                m.bias.data.fill_(0.01)
        self.out_layer.apply(init_weights)
        self.embed_input.apply(init_weights)
        self.attn_combine.apply(init_weights)

    def forward(self, Hi_data, encoder_hidden, context):
        # Hi_data is scaled time
        inputs = Hi_data.transpose(1,0)
        if self.bidirectional:
            h0 = encoder_hidden[2:]
        else:
            h0 = encoder_hidden[2:].sum(0).unsqueeze(0)
        # combine input and context
        inputs = self.embed_input(inputs)
        # repeat context for each item in sequence  
        context = context.repeat(inputs.shape[0],1,1)
        inputs = torch.cat((inputs, context), 2) 
        inputs = self.attn_combine(inputs)
        # Take last output from GRU
        latent_seqs = self.rnn(inputs, h0)[0]
        latent_seqs = latent_seqs.transpose(1,0)
        latent_seqs = self.out_layer(latent_seqs)
        return latent_seqs

''' smooth data with moving average (common with fitting mechanistic models) '''
def moving_average(x, w):
    return pd.Series(x).rolling(w, min_periods=1).mean().values

''' Specify which state '''
def fetch_county_data_covid(state='MA', county_id='25005', pred_week='202021', batch_size=32, noise_level=0):
    ''' Import COVID data for counties '''
    np.random.seed(17)

    if county_id == 'all':
        all_counties = counties[state]
    else:
        all_counties = [county_id]

    c_seqs = []  # county sequences of features
    c_ys = []  # county targets
    for county in all_counties:
        X_county, y = get_county_train_data(county,pred_week,noise_level=noise_level)
        y = moving_average(y[:,1].ravel(),SMOOTH_WINDOW).reshape(-1,1)
        c_seqs.append(X_county.to_numpy())
        c_ys.append(y)
    c_seqs = np.array(c_seqs)  # Shape: [regions, time, features]
    c_ys = np.array(c_ys)  # Shape: [regions, time, 1]

    # Normalize
    # One scaler per county
    scalers = [StandardScaler() for _ in range(len(all_counties))]
    c_seqs_norm = []
    for i, scaler in enumerate(scalers):
        c_seqs_norm.append(scaler.fit_transform(c_seqs[i]))
    c_seqs_norm = np.array(c_seqs_norm)

    ''' Create static metadata data for each county '''

    county_idx = {r: i for i, r in enumerate(all_counties)}
    def one_hot(idx, dim=len(county_idx)):
        ans = np.zeros(dim, dtype="float32")
        ans[idx] = 1.0
        return ans
    metadata = np.array([one_hot(county_idx[r]) for r in all_counties])

    ''' Prepare train and validation dataset '''

    min_sequence_length = 20
    metas, seqs, y, y_mask = [], [], [], []
    for meta, seq, ys in zip(metadata, c_seqs_norm, c_ys):
        seq, ys, ys_mask = create_window_seqs(seq,ys,min_sequence_length)
        metas.append(meta)
        seqs.append(seq[[-1]])
        y.append(ys[[-1]])
        y_mask.append(ys_mask[[-1]])

    all_metas = np.array(metas, dtype="float32")
    all_county_seqs = torch.cat(seqs,axis=0)
    all_county_ys = torch.cat(y,axis=0)
    all_county_y_mask = torch.cat(y_mask,axis=0)
    
    counties_train, metas_train, X_train, y_train, y_mask_train = \
        all_counties, all_metas, all_county_seqs, all_county_ys, all_county_y_mask
    
    train_dataset = SeqData(counties_train, metas_train, X_train, y_train, y_mask_train)
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

    assert all_county_seqs.shape[1] == all_county_ys.shape[1]
    seqlen = all_county_seqs.shape[1]
    return train_loader, metas_train.shape[1], X_train.shape[2], seqlen
    

def fetch_county_data_flu(state='MA', county_id='25005', pred_week='202021', batch_size=32, noise_level=0):
    ''' in flu, our features are state-level and target ILI is only available at state'''
    np.random.seed(17)

    ''' Import data for all counties '''

    if county_id == 'all':
        all_counties = counties[state]
    else:
        all_counties = [county_id]

    X_state, y = get_state_train_data_flu(state, pred_week, noise_level=noise_level)
    y = moving_average(y.ravel(),SMOOTH_WINDOW).reshape(-1,1)
    c_seqs = []  # county sequences of features
    c_ys = []  # county targets
    for _ in all_counties:
        c_seqs.append(X_state.to_numpy())
        c_ys.append(y)
    c_seqs = np.array(c_seqs)  # Shape: [regions, time, features]
    c_ys = np.array(c_ys)  # Shape: [regions, time, 1]

    # Normalize
    # One scaler per county
    scalers = [StandardScaler() for _ in range(len(all_counties))]
    c_seqs_norm = []
    for i, scaler in enumerate(scalers):
        c_seqs_norm.append(scaler.fit_transform(c_seqs[i]))
    c_seqs_norm = np.array(c_seqs_norm)

    ''' Create static metadata data for each county '''

    county_idx = {r: i for i, r in enumerate(all_counties)}
    def one_hot(idx, dim=len(county_idx)):
        ans = np.zeros(dim, dtype="float32")
        ans[idx] = 1.0
        return ans
    metadata = np.array([one_hot(county_idx[r]) for r in all_counties])

    ''' Prepare train and validation dataset '''
    min_sequence_length = 5
    metas, seqs, y, y_mask = [], [], [], []
    for meta, seq, ys in zip(metadata, c_seqs_norm, c_ys):
        seq, ys, ys_mask = create_window_seqs(seq,ys,min_sequence_length)
        metas.append(meta)
        seqs.append(seq[[-1]])
        y.append(ys[[-1]])
        y_mask.append(ys_mask[[-1]])

    all_metas = np.array(metas, dtype="float32")
    all_county_seqs = torch.cat(seqs,axis=0)
    all_county_ys = torch.cat(y,axis=0)
    all_county_y_mask = torch.cat(y_mask,axis=0)
    
    counties_train, metas_train, X_train, y_train, y_mask_train = \
        all_counties, all_metas, all_county_seqs, all_county_ys, all_county_y_mask
    train_dataset = SeqData(counties_train, metas_train, X_train, y_train, y_mask_train)
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    
    assert all_county_seqs.shape[1] == all_county_ys.shape[1]
    seqlen = all_county_seqs.shape[1]

    return train_loader, metas_train.shape[1], X_train.shape[2], seqlen

# dataset class
class SeqData(torch.utils.data.Dataset):
    def __init__(self, region, meta, X, y, mask_y):
        self.region = region
        self.meta = meta
        self.X = X
        self.y = y
        # self.mask_y = mask_y

    def __len__(self):
        return self.X.shape[0]

    def __getitem__(self, idx):
        return (
            self.region[idx],
            self.meta[idx],
            self.X[idx, :, :],
            self.y[idx]
        )

class ODE(nn.Module):
    def __init__(self, params, device):
        super(ODE, self).__init__()
        county_id = params['county_id']
        abm_params = f'Data/{county_id}_generated_params.yaml'
        #Reading params
        with open(abm_params, 'r') as stream:
            try:
                abm_params = yaml.safe_load(stream)
            except yaml.YAMLError as exc:
                print('Error in reading parameters file')
                print(exc)
        params.update(abm_params)
        self.params = params
        self.device = device
        self.num_agents = self.params['num_agents'] # Population


class SEIRM(ODE):
    def __init__(self, params, device):
        super().__init__(params,device)
    
    def init_compartments(self,learnable_params):
        ''' let's get initial conditions '''
        initial_infections_percentage = learnable_params['initial_infections_percentage']
        initial_conditions = torch.empty((5)).to(self.device)
        no_infected = (initial_infections_percentage / 100) * self.num_agents # 1.0 is ILI
        initial_conditions[2] = no_infected
        initial_conditions[0] = self.num_agents - no_infected
        print('initial infected',no_infected)
        self.state = initial_conditions

    def step(self, t, values):
        """
        Computes ODE states via equations       
            state is the array of state value (S,E,I,R,M)
        """
        params = {
            'beta':values[0],
            'alpha':values[1],
            'gamma':values[2],
            'mu':values[3],
            'initial_infections_percentage': values[4], 
        }
        if t==0:
            self.init_compartments(params)
        # to make the NN predict lower numbers, we can make its prediction to be N-Susceptible
        dSE = params['beta'] * self.state[0] * self.state[2] / self.num_agents
        dEI = params['alpha'] * self.state[1]
        dIR = params['gamma'] * self.state[2]
        dIM = params['mu'] * self.state[2]

        dS  = -1.0 * dSE
        dE  = dSE - dEI
        dI = dEI - dIR - dIM
        dR  = dIR
        dM  = dIM

        # concat and reshape to make it rows as obs, cols as states
        self.dstate = torch.stack([dS, dE, dI, dR, dM], 0)
        NEW_INFECTIONS_TODAY = dEI
        NEW_DEATHS_TODAY = dIM
        # update state
        self.state = self.state + self.dstate
        
        return NEW_INFECTIONS_TODAY, NEW_DEATHS_TODAY

class SIRS(ODE):
    def __init__(self, params, device):
        super().__init__(params,device)

    def init_compartments(self,learnable_params):
        ''' let's get initial conditions '''
        initial_infections_percentage = learnable_params['initial_infections_percentage']
        initial_conditions = torch.empty((2)).to(self.device)
        no_infected = (initial_infections_percentage / 100) * self.num_agents # 1.0 is ILI
        initial_conditions[1] = no_infected
        initial_conditions[0] = self.num_agents - no_infected
        print('initial infected',no_infected)

        self.state = initial_conditions

    def step(self, t, values):
        """
        Computes ODE states via equations       
            state is the array of state value (S,I)
        """
        params = {
            'beta':values[0],  # contact rate, range: 0-1
            'initial_infections_percentage': values[1], 
        }
        # set from expertise
        params['D'] = 3.5
        params['L'] = 2000
        if t==0:
            self.init_compartments(params)
        dS =  (self.num_agents - self.state[0] - self.state[1]) / params['L'] -  params['beta'] * self.state[0] * self.state[1] / self.num_agents
        dSI = params['beta'] * self.state[0] * self.state[1] / self.num_agents
        dI = dSI - self.state[1] / params['D']

        # concat and reshape to make it rows as obs, cols as states
        self.dstate = torch.stack([dS, dI], 0)

        NEW_INFECTIONS_TODAY = dSI
        # ILI is percentage of outpatients with influenza-like illness
        # ILI = params['lambda'] * dSI / self.num_agents
        # this is what Shaman and Pei do https://github.com/SenPei-CU/Multi-Pathogen_ILI_Forecast/blob/master/code/SIRS_AH.m
        ILI =  dSI / self.num_agents * 100 # multiply 100 because it is percentage

        # update state
        self.state = self.state + self.dstate
        return NEW_INFECTIONS_TODAY, ILI

if __name__ == '__main__':
    print("THIS SHOULD NOT EXECUTE!")
    """ Create model """