fixes and code cleanup in stochastic module