Refactor parts of demos/demo_training.ipynb into the library, and make good tests

[phinate]

The working example in demos/ doesn't use half of the modules, and function within the demo are good generalisations of the attempted functionality some of the modules (e.g. hists) try to provide.

• Refactor demos/demo_training.ipynb into library functions, keeping in mind the balance of user friendliness and flexibility
• Trim down library to get rid of unused functions
• Write good tests (notebook cells in nbdev)