# Single Trait Animal Model - Date : 2020-11-11

# A model for Animal Evaluation

# Linear Models for the Prediction of Animal Breeding Values, Raphael Mrode, 3rd Edition

# Julius Mugambe(PhD)


# Clear workspace
rm(list = ls())
# Set working directory
# Setting working directory
setwd("G:/R-MME") 
getwd()
## [1] "G:/R-MME"
# Functions

# install.packages('MASS', dependencies = TRUE)
library(MASS)

# make design matrix 
desgn <- function(v) {
  if (is.numeric(v)) {
    va = v
    mrow = length(va)
    mcol = max(va)
  }
  if (is.character(v)) {
    vf = factor(v)
    # Save the index value of each level
    va = as.numeric(vf)
    mrow = length(va)
    mcol = length(levels(vf))
  }
  
  # Initializing X-matrix
  X = matrix(data = c(0), nrow = mrow, ncol = mcol)
  
  for (i in 1:mrow) {
    ic = va[i]
    X[i, ic] = 1
  }
  return(X)
}

# function to make inverse of numerator relationship matrix
ainv = function(ped) {
  n = nrow(ped)
  Ainv = matrix(c(0), nrow = n, ncol = n)
  
  for (i in 1:n) {
    animal = ped[i, 1]
    sire = ped[i, 2]
    dam = ped[i, 3]
    
    if (sire == 0 & dam == 0) {
      # both parents unknown
      alpha = 1
      Ainv[animal, animal] = alpha + Ainv[animal, animal]
    } else if (sire != 0 & dam == 0) {
      # sire known
      alpha = 4/3
      Ainv[animal, animal] = alpha + Ainv[animal, animal]
      Ainv[animal, sire] = -alpha/2 + Ainv[animal, sire]
      Ainv[sire, animal] = -alpha/2 + Ainv[sire, animal]
      Ainv[sire, sire] = alpha/4 + Ainv[sire, sire]
    } else if (sire == 0 & dam != 0) {
      # dam known
      alpha = 4/3
      Ainv[animal, animal] = alpha + Ainv[animal, animal]
      Ainv[animal, dam] = -alpha/2 + Ainv[animal, dam]
      Ainv[dam, animal] = -alpha/2 + Ainv[dam, animal]
      Ainv[dam, dam] = alpha/4 + Ainv[dam, dam]
    } else {
      # both parents known
      alpha = 2
      Ainv[animal, animal] = alpha + Ainv[animal, animal]
      Ainv[animal, sire] = -alpha/2 + Ainv[animal, sire]
      Ainv[sire, animal] = -alpha/2 + Ainv[sire, animal]
      Ainv[animal, dam] = -alpha/2 + Ainv[animal, dam]
      Ainv[dam, animal] = -alpha/2 + Ainv[dam, animal]
      Ainv[sire, sire] = alpha/4 + Ainv[sire, sire]
      Ainv[sire, dam] = alpha/4 + Ainv[sire, dam]
      Ainv[dam, sire] = alpha/4 + Ainv[dam, sire]
      Ainv[dam, dam] = alpha/4 + Ainv[dam, dam]
    }
  }
  return(Ainv)
}

# Univariate Model with one Random Effect.
#inputting variance components
sigma_2_a = 20
sigma_2_e = 40
alpha = sigma_2_e/sigma_2_a

# load pedigree data
pedigree <- read.csv("pedigree.txt", sep="")
# Calculate the AINV from the pedigree using the ainv function
A <- ainv(pedigree)

# Load data
data  <- read.csv("data.txt", sep="")

# Making design matrices 
X = desgn(data$sex)
Z = desgn(data$id)
y = data$wwg


#Mixed model equations
XtX <- crossprod(X)
XtZ <- crossprod(X, Z)
ZtX <- crossprod(Z, X)
ZtZ <- crossprod(Z)

# The addition of Ainv_alpha to ZtZ in the LSE yeilds the MME
ZtZA <- crossprod(Z) + (A) * alpha
Xty <- crossprod(X, y)
Zty <- crossprod(Z, y)
LHS1 <- cbind(XtX, XtZ)
LHS2 <- cbind(ZtX, ZtZA)
LHS <- rbind(LHS1, LHS2)
RHS <- rbind(Xty, Zty)

#solutions
sol <- solve(LHS) %*% RHS
round(sol, digits=3)
##         [,1]
##  [1,]  3.404
##  [2,]  4.359
##  [3,]  0.098
##  [4,] -0.019
##  [5,] -0.041
##  [6,] -0.009
##  [7,] -0.186
##  [8,]  0.177
##  [9,] -0.249
## [10,]  0.183