script

# Clean up the workspace
rm(list = ls())

# Plot layout
par(mfrow = c(1, 1))

suppressPackageStartupMessages({
  library(AlphaSimR)
  library(dplyr)
  library(ggplot2)
})

set.seed(20260709)

# This script tests four factors in an Atlantic salmon breeding scheme:
# founder Ne, number of families, progeny per family, and the proportion kept
# after the smolt stage.
#
# Selection is a simple two-step process:
# Stage A: select smolts for Trait 1 (growth) using phenotypic selection.
# Stage B: select the next parents for Trait 2 (disease challenge) using
#          genomic selection with RR-GBLUP.
#
# The number of replicates is kept small here because this is a test script.

#### User settings

nFounders <- 1000
nChr <- 29
segSites <- 600  # enough sites for 100 QTL and 500 SNP markers per chromosome
genLen <- 1

nQTLPerChr <- 100
h2Growth <- 0.35
h2SRS <- 0.43

nGenerations <- 20
generationIntervalYears <- 4
nReplicates <- 3

parentSelectionMode <- "rrgblup"
nSnpPerChr <- 500

outputDir <- file.path(getwd(), "outputs")
if (!dir.exists(outputDir)) {
  dir.create(outputDir, recursive = TRUE)
}

#### Scenario grid

NeLevels <- c(50, 100, 150)
nFamilyLevels <- c(50, 100, 300)
progenySizeLevels <- c(25, 30, 100, 150)
smoltPropLevels <- c(0.30, 0.40, 0.50)

scenarioGrid <- expand.grid(
  founderNe = NeLevels,
  nFamilies = nFamilyLevels,
  progenyPerCross = progenySizeLevels,
  smoltProp = smoltPropLevels,
  stringsAsFactors = FALSE
)

scenarioGrid$scenario <- paste0(
  "Ne", scenarioGrid$founderNe,
  "_NFam", scenarioGrid$nFamilies,
  "_Prog", scenarioGrid$progenyPerCross,
  "_Sm", as.integer(100 * scenarioGrid$smoltProp)
)

#### Function to run one replicate

runOneRep <- function(founderNe, nFamilies, progenyPerCross, smoltProp, repId, scenarioName) {
  set.seed(100000 + founderNe * 1000 + nFamilies * 10 + progenyPerCross + round(smoltProp * 100) + repId)

  founderGenome <- runMacs2(
    nInd = nFounders,
    nChr = nChr,
    segSites = segSites,
    Ne = founderNe,
    genLen = genLen
  )

  SP <- SimParam$new(founderGenome)
  SP$setSexes("yes_sys")
  SP$addTraitA(nQtlPerChr = nQTLPerChr, mean = 120, var = 30)
  SP$addTraitA(nQtlPerChr = nQTLPerChr, mean = 35, var = 30)
  SP$setVarE(h2 = c(h2Growth, h2SRS))
  SP$addSnpChip(nSnpPerChr = nSnpPerChr)

  founders <- newPop(founderGenome, simParam = SP)
  founders <- setPheno(founders, simParam = SP)

# I use dams as the family number. Each dam makes one family.
  nDams <- nFamilies
  nSires <- max(10, round(nDams / 3))

  currentDams <- selectInd(
    pop = founders,
    nInd = nDams,
    use = "pheno",
    sex = "F",
    trait = 2,
    simParam = SP
  )

  currentSires <- selectInd(
    pop = founders,
    nInd = nSires,
    use = "pheno",
    sex = "M",
    trait = 2,
    simParam = SP
  )

  generationData <- NULL

  for (generation in seq_len(nGenerations)) {
    yearStart <- (generation - 1) * generationIntervalYears

# Round-robin mating keeps the family set simple and controlled.
    crossPlan <- cbind(
      female = seq_len(currentDams@nInd),
      male = rep(seq_len(currentSires@nInd), length.out = currentDams@nInd)
    )

    offspring <- makeCross2(
      females = currentDams,
      males = currentSires,
      crossPlan = crossPlan,
      nProgeny = progenyPerCross,
      simParam = SP
    )
    offspring <- setPheno(offspring, simParam = SP)

# Stage A: smolt selection on growth.
    nSmoltSelected <- round(smoltProp * offspring@nInd)
    nSmoltSelected <- max(nSmoltSelected, nDams + nSires)
    nSmoltSelected <- min(nSmoltSelected, offspring@nInd)

    smoltSelected <- selectInd(
      pop = offspring,
      nInd = nSmoltSelected,
      use = "pheno",
      trait = 1,
      simParam = SP
    )

    nFemaleSmolts <- sum(smoltSelected@sex == "F")
    nMaleSmolts <- sum(smoltSelected@sex == "M")

    if (nFemaleSmolts < nDams || nMaleSmolts < nSires) {
      stop("Not enough male or female smolts to choose the next parents. Try a larger smoltProp or progenyPerCross.")
    }

# Stage B: parent selection on Trait 2 using RR-GBLUP.
    rrSol <- RRBLUP(
      pop = smoltSelected,
      traits = 2,
      use = "pheno",
      snpChip = 1,
      simParam = SP
    )

    candPop <- setEBV(smoltSelected, rrSol, simParam = SP)

    nextDams <- selectInd(
      pop = candPop,
      nInd = nDams,
      use = "ebv",
      sex = "F",
      simParam = SP
    )

    nextSires <- selectInd(
      pop = candPop,
      nInd = nSires,
      use = "ebv",
      sex = "M",
      simParam = SP
    )

    nextParents <- c(nextDams, nextSires)

# This is only possible in the simulation because true GV is known.
    accParentSelection <- cor(
      as.numeric(gv(candPop)[, 2]),
      as.numeric(ebv(candPop)),
      use = "complete.obs"
    )

    oneRow <- data.frame(
      scenario = scenarioName,
      parentSelectionMode = parentSelectionMode,
      founderNe = founderNe,
      nFamilies = nFamilies,
      nDams = nDams,
      nSires = nSires,
      progenyPerCross = progenyPerCross,
      smoltProp = smoltProp,
      rep = repId,
      generation = generation,
      yearStart = yearStart,
      yearEnd = yearStart + generationIntervalYears - 1,
      nOffspring = offspring@nInd,
      nSmoltSelected = smoltSelected@nInd,
      nFemaleSmolts = nFemaleSmolts,
      nMaleSmolts = nMaleSmolts,
      nNextParents = nextParents@nInd,
      meanGvTrait2_Offspring = mean(offspring@gv[, 2]),
      varGvTrait2_Offspring = var(offspring@gv[, 2]),
      meanGvTrait2_NextParents = mean(nextParents@gv[, 2]),
      accPhenoTrait2_Offspring = cor(offspring@gv[, 2], offspring@pheno[, 2], use = "complete.obs"),
      accParentSelection = accParentSelection,
      stringsAsFactors = FALSE
    )

    generationData <- rbind(generationData, oneRow)

    currentDams <- nextDams
    currentSires <- nextSires
  }

  generationData <- generationData %>%
    arrange(generation) %>%
    mutate(deltaG_Trait2 = meanGvTrait2_NextParents - lag(meanGvTrait2_NextParents))

  generationData
}

#### Run all scenarios

allGenerationData <- NULL

for (i in seq_len(nrow(scenarioGrid))) {
  scn <- scenarioGrid[i, ]

  for (repId in seq_len(nReplicates)) {
    cat("Running:", scn$scenario, "rep", repId, "\n")

    repData <- runOneRep(
      founderNe = scn$founderNe,
      nFamilies = scn$nFamilies,
      progenyPerCross = scn$progenyPerCross,
      smoltProp = scn$smoltProp,
      repId = repId,
      scenarioName = scn$scenario
    )

    allGenerationData <- rbind(allGenerationData, repData)
  }
}

#### Summarize results

scenarioTrend <- allGenerationData %>%
  group_by(scenario, parentSelectionMode, founderNe, nFamilies, progenyPerCross, smoltProp, generation) %>%
  summarise(
    meanGvTrait2_NextParents = mean(meanGvTrait2_NextParents),
    varGvTrait2_Offspring = mean(varGvTrait2_Offspring),
    deltaG_Trait2 = mean(deltaG_Trait2, na.rm = TRUE),
    accPhenoTrait2_Offspring = mean(accPhenoTrait2_Offspring),
    accParentSelection = mean(accParentSelection, na.rm = TRUE),
    .groups = "drop"
  )

endpointByRep <- allGenerationData %>%
  filter(generation == nGenerations) %>%
  select(
    scenario,
    parentSelectionMode,
    founderNe,
    nFamilies,
    progenyPerCross,
    smoltProp,
    rep,
    finalMeanGv = meanGvTrait2_NextParents,
    finalVarGv = varGvTrait2_Offspring,
    finalAcc = accPhenoTrait2_Offspring,
    finalAccParentSelection = accParentSelection
  )

endpointByScenario <- endpointByRep %>%
  group_by(scenario, parentSelectionMode, founderNe, nFamilies, progenyPerCross, smoltProp) %>%
  summarise(
    finalMeanGv = mean(finalMeanGv),
    finalVarGv = mean(finalVarGv),
    finalAcc = mean(finalAcc),
    finalAccParentSelection = mean(finalAccParentSelection, na.rm = TRUE),
    .groups = "drop"
  )

#### Simple models for the final generation

fitGain <- lm(finalMeanGv ~ founderNe + nFamilies + progenyPerCross + smoltProp, data = endpointByRep)
fitVar <- lm(finalVarGv ~ founderNe + nFamilies + progenyPerCross + smoltProp, data = endpointByRep)
fitAcc <- lm(finalAcc ~ founderNe + nFamilies + progenyPerCross + smoltProp, data = endpointByRep)

cat("\nModel summary: final mean GV for Trait 2\n")
print(summary(fitGain))

cat("\nModel summary: final GV variance for Trait 2\n")
print(summary(fitVar))

cat("\nModel summary: final phenotype accuracy proxy for Trait 2\n")
print(summary(fitAcc))

cat("\nEndpoint means by scenario:\n")
print(endpointByScenario)

#### Plots

pGain <- ggplot(endpointByScenario, aes(x = factor(nFamilies), y = finalMeanGv, fill = factor(founderNe))) +
  geom_col(position = "dodge") +
  facet_grid(smoltProp ~ progenyPerCross) +
  labs(
    x = "Number of families",
    y = "Final mean GV for Trait 2",
    fill = "Founder Ne",
    title = "Final genetic gain"
  ) +
  theme_minimal()
print(pGain)

pVar <- ggplot(endpointByScenario, aes(x = factor(nFamilies), y = finalVarGv, fill = factor(founderNe))) +
  geom_col(position = "dodge") +
  facet_grid(smoltProp ~ progenyPerCross) +
  labs(
    x = "Number of families",
    y = "Final GV variance for Trait 2",
    fill = "Founder Ne",
    title = "Final genetic variance"
  ) +
  theme_minimal()
print(pVar)

pAcc <- ggplot(endpointByScenario, aes(x = factor(nFamilies), y = finalAcc, fill = factor(founderNe))) +
  geom_col(position = "dodge") +
  facet_grid(smoltProp ~ progenyPerCross) +
  labs(
    x = "Number of families",
    y = "Final phenotype accuracy proxy",
    fill = "Founder Ne",
    title = "Final accuracy proxy"
  ) +
  theme_minimal()
print(pAcc)

#### Save outputs

write.csv(allGenerationData, file.path(outputDir, "salmon_hypothesis_allGenerationData.csv"), row.names = FALSE)
write.csv(scenarioTrend, file.path(outputDir, "salmon_hypothesis_scenarioTrend.csv"), row.names = FALSE)
write.csv(endpointByRep, file.path(outputDir, "salmon_hypothesis_endpointByRep.csv"), row.names = FALSE)
write.csv(endpointByScenario, file.path(outputDir, "salmon_hypothesis_endpointByScenario.csv"), row.names = FALSE)

ggsave(file.path(outputDir, "salmon_hypothesis_endpoint_gain.png"), pGain, width = 10, height = 6, dpi = 300)
ggsave(file.path(outputDir, "salmon_hypothesis_endpoint_variance.png"), pVar, width = 10, height = 6, dpi = 300)
ggsave(file.path(outputDir, "salmon_hypothesis_endpoint_accuracy.png"), pAcc, width = 10, height = 6, dpi = 300)