add batch size as an argument, use nanmean for pretraining if allele gets pruned from the completed matrix