# Retrieve default pattern fitters
fitters <- defaultPatternFitters()
# Add a custom diagonal pattern fitter
diagonalFitter <- function(mat, trainMask) {
# Extract diagonal values from training data
minDim <- min(nrow(mat), ncol(mat))
diagIndices <- cbind(1:minDim, 1:minDim)
# Only use diagonal elements that are in the training mask
validDiag <- trainMask[diagIndices]
if (any(validDiag)) {
diagVal <- mean(mat[diagIndices][validDiag])
} else {
diagVal <- mean(mat[trainMask]) # fallback to overall mean
}
matrix(diagVal, nrow = nrow(mat), ncol = ncol(mat))
}
# Extend the list with your own pattern
fitters$diagonal <- diagonalFitter
Run the code above in your browser using DataLab