Consolidate docstrings and add type annotations

[dalonsoa]

Docstrings should follow the Google style... and have some content. While this might be difficult in some cases, requiring domain knowledge, a minimal description based on what the code is doing in the simple cases will be a time saver.

As a bonus, adding type annotations in those cases when they are clear will also be useful to clarify the meaning of the code.