Fix functions to work with re-organized data