Project: s41467-025-55873-3 | Function: prob_infect (r)

system: You are an expert R programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow R best practices and maintain consistent code style.
Return **only** the complete R function code.

user: Complete this R function:

**DESCRIPTION:** Core infection probability calculation implementing flexible functional forms (polynomial, logarithmic, reciprocal terms). Calculates probability of infection based on cell concentrations or CCR5 expression levels using exponential survival model.

**KEY TERMS:** {
  "theta": "Model parameters vector",
  "x": "CCR5 fraction or cell concentration",
  "exp_term": "Exponential hazard function",
  "R_0": "Basic reproduction number"
}

**CONTEXT CODE:** 
prob_infect_week <- function(week_of_infection, prob_inoc_infect, infect_cens){
  #function to calculate the probability of getting infected on a specific week
  #(i.e., after 'week_of_infection' inoculations)
  #prob_inoc_infect is the probability of a single inoculation causing an infection
  #infect_cens: 1 infection detected, 0 censoring occured at current week
  
  
  Prob_infect_week = (1-prob_inoc_infect)^(week_of_infection -infect_cens)*
    (prob_inoc_infect)^infect_cens
  
  
  return(Prob_infect_week)
  
}

prob_infect_core <- function(theta, x) {
  # theta: Model parameters vector (c(p1, p2, ..., pn, exp_term))
  # x: CCR5 fraction or cell concentration

  # Assuming polynomial functional form (default)
  core_prob <- 1 - exp(-theta[1] * x - theta[2] * x^2)

  # Check if logarithmic functional form
  if (length(theta) > 3 && theta[3] != 0) {
    core_prob <- 1 - exp(-theta[1] * log(x) - theta[2] * x^2)
  }

  # Check if reciprocal term functional form
  if (length(theta) > 4 && theta[4] != 0) {
    core_prob <- 1 - exp(-theta[1] * x - theta[2] * x^2 - theta[3] * 1 / x)
  }

  # Exponential hazard function term
  exp_term <- exp(theta[length(theta)]) # Assuming last element of theta is exp_term
  core_prob <- core_prob * exp_term

  # Basic reproduction number
  R_0 <- core_prob / (1 - core_prob)

  return(list(core_prob = core_prob, R_0 = R_0))
}

# Usage example:
# theta <- c(0.5, 0.1, 0.2, 0.3, 2) # Model parameters vector
# x <- 0.8 # CCR5 fraction or cell concentration
# result <- prob_infect_core(theta, x)
# print(paste("Core Infection Probability:", result$core_prob))
# print(paste("Basic Reproduction Number:", result$R_0))

prob_infect <- function(params, x, param.id = 1) {
  stop("Function implementation has been masked")
}


prob_infect_Arm <-function(params, Arm_ind){
  #params is the list of probability of getting infected after inoculation for
  # each Arm (100% KO, 75% KO, ...) and Arm is the index for which Arm the 
  # animal was part of (1=100% KO, 2 = 75% KO, ....)
  
  prob_infect_it = params[Arm_ind] #probability of getting infected for this
  #animal
  
  return(prob_infect_it)
}


Complete the function implementation replacing "stop("Not implemented")".