<commit_before>'''
ODEs representing the HIV model.
'''

import warnings

import numpy
from scipy import integrate
import pandas

from . import control_rates


variables = (
    'susceptible',        # S
    'vaccinated',         # Q
    'acute',              # A
    'undiagnosed',        # U
    'diagnosed',          # D
    'treated',            # T
    'viral_suppression',  # V
    'AIDS',               # W
    'dead',               # Z
    'new_infections',     # R
)

alive = ('susceptible', 'vaccinated', 'acute', 'undiagnosed',
         'diagnosed', 'treated', 'viral_suppression', 'AIDS')

infected = ('acute', 'undiagnosed', 'diagnosed', 'treated',
            'viral_suppression', 'AIDS')


def get_variable(state, name):
    i = variables.index(name)
    return state[..., i]


def get_alive(state):
    return sum(get_variable(state, k) for k in alive)


def get_infected(state):
    return sum(get_variable(state, k) for k in infected)


# Variables to log transform: S, U, D, T, V, W
vars_log = [0, 3, 4, 5, 6, 7]
# Variables to not log transform: Q, A, Z, R
vars_nonlog = [1, 2, 8, 9]


def transform(state, _log0 = -20):
    '''
    Log transform some of the state variables.
    '''
    state_trans = numpy.empty(numpy.shape(state))
    state_trans[..., vars_log] = numpy.ma.log(state[..., vars_log]).filled(
        _log0)
    state_trans[..., vars_nonlog] = state[..., vars_nonlog]
    return state_trans


def transform_inv(state_trans):
    '''
    Inverse log transform some of the state variables.
    '''
    state = numpy.empty(numpy.shape(state_trans))
    state[..., vars_log] = numpy.exp(state_trans[..., vars_log])
    state[..., vars_nonlog] = state_trans[..., vars_nonlog]
    return state


def split_state(state):
    if isinstance(state, (pandas.Series, pandas.DataFrame)):
        state = state.values
    return map(numpy.squeeze, numpy.split(state, state.shape[-1], -1))


def rhs(t, state, target, parameters):
    # Force the state variables to be non-negative.
    # The last two state variables, dead from AIDS and new infections,
    # are cumulative numbers that are set to 0 at t = 0: these
    # can be negative if time goes backwards.
    state[ : -2] = state[ : - 2].clip(0, numpy.inf)

    S, Q, A, U, D, T, V, W, Z, R = state

    # Total sexually active population.
    N = S + Q + A + U + D + T + V

    control_rates_ =  control_rates.get(t, state, target, parameters)

    force_of_infection = (
        parameters.transmission_rate_acute * A
        + parameters.transmission_rate_unsuppressed * (U + D + T)
        + parameters.transmission_rate_suppressed * V) / N

    dS = (parameters.birth_rate * N
          - control_rates_.vaccination * S
          - force_of_infection * S
          - parameters.death_rate * S)

    dQ = (control_rates_.vaccination * S
          - (1 - target.vaccine_efficacy) * force_of_infection * Q
          - parameters.death_rate * Q)
    
    dA = (force_of_infection * S
          + (1 - target.vaccine_efficacy) * force_of_infection * Q
          - parameters.progression_rate_acute * A
          - parameters.death_rate * A)

    dU = (parameters.progression_rate_acute * A
          - control_rates_.diagnosis * U
          - parameters.death_rate * U
          - parameters.progression_rate_unsuppressed * U)

    dD = (control_rates_.diagnosis * U
          + control_rates_.nonadherence * (T + V)
          - control_rates_.treatment * D
          - parameters.death_rate * D
          - parameters.progression_rate_unsuppressed * D)
<commit_msg>Whitespace<commit_after>'''
ODEs representing the HIV model.
'''

import warnings

import numpy
from scipy import integrate
import pandas

from . import control_rates


variables = (
    'susceptible',        # S
    'vaccinated',         # Q
    'acute',              # A
    'undiagnosed',        # U
    'diagnosed',          # D
    'treated',            # T
    'viral_suppression',  # V
    'AIDS',               # W
    'dead',               # Z
    'new_infections',     # R
)

alive = ('susceptible', 'vaccinated', 'acute', 'undiagnosed',
         'diagnosed', 'treated', 'viral_suppression', 'AIDS')

infected = ('acute', 'undiagnosed', 'diagnosed', 'treated',
            'viral_suppression', 'AIDS')


def get_variable(state, name):
    i = variables.index(name)
    return state[..., i]


def get_alive(state):
    return sum(get_variable(state, k) for k in alive)


def get_infected(state):
    return sum(get_variable(state, k) for k in infected)


# Variables to log transform: S, U, D, T, V, W
vars_log = [0, 3, 4, 5, 6, 7]
# Variables to not log transform: Q, A, Z, R
vars_nonlog = [1, 2, 8, 9]


def transform(state, _log0 = -20):
    '''
    Log transform some of the state variables.
    '''
    state_trans = numpy.empty(numpy.shape(state))
    state_trans[..., vars_log] = numpy.ma.log(state[..., vars_log]).filled(
        _log0)
    state_trans[..., vars_nonlog] = state[..., vars_nonlog]
    return state_trans


def transform_inv(state_trans):
    '''
    Inverse log transform some of the state variables.
    '''
    state = numpy.empty(numpy.shape(state_trans))
    state[..., vars_log] = numpy.exp(state_trans[..., vars_log])
    state[..., vars_nonlog] = state_trans[..., vars_nonlog]
    return state


def split_state(state):
    if isinstance(state, (pandas.Series, pandas.DataFrame)):
        state = state.values
    return map(numpy.squeeze, numpy.split(state, state.shape[-1], -1))


def rhs(t, state, target, parameters):
    # Force the state variables to be non-negative.
    # The last two state variables, dead from AIDS and new infections,
    # are cumulative numbers that are set to 0 at t = 0: these
    # can be negative if time goes backwards.
    state[ : -2] = state[ : - 2].clip(0, numpy.inf)

    S, Q, A, U, D, T, V, W, Z, R = state

    # Total sexually active population.
    N = S + Q + A + U + D + T + V

    control_rates_ =  control_rates.get(t, state, target, parameters)

    force_of_infection = (
        parameters.transmission_rate_acute * A
        + parameters.transmission_rate_unsuppressed * (U + D + T)
        + parameters.transmission_rate_suppressed * V) / N

    dS = (parameters.birth_rate * N
          - control_rates_.vaccination * S
          - force_of_infection * S
          - parameters.death_rate * S)

    dQ = (control_rates_.vaccination * S
          - (1 - target.vaccine_efficacy) * force_of_infection * Q
          - parameters.death_rate * Q)

    dA = (force_of_infection * S
          + (1 - target.vaccine_efficacy) * force_of_infection * Q
          - parameters.progression_rate_acute * A
          - parameters.death_rate * A)

    dU = (parameters.progression_rate_acute * A
          - control_rates_.diagnosis * U
          - parameters.death_rate * U
          - parameters.progression_rate_unsuppressed * U)

    dD = (control_rates_.diagnosis * U
          + control_rates_.nonadherence * (T + V)
          - control_rates_.treatment * D
          - parameters.death_rate * D
          - parameters.progression_rate_unsuppressed * D)

{'fraction_non_alphanumeric': 0.12279463655610445, 'fraction_numerical': 0.006492589978828511, 'mean_word_length': 2.458272327964861, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1299341', 'n_tokens_mistral': 2444, 'n_tokens_neox': 2250, 'n_words': 614}