some changes

[aljabadi]

Hi @EvaYiwenWang,

Hope you find these suggestions useful. Also, you can start adding unit tests using testthat package.