think about adding support for imputation
think about adding support for imputation