Implement linear extrapolation

[timmens]

In this PR, I add capabilities to perform linear extrapolation.

In particular, I

• Implement a custom map_coordinates function that allows for linear extrapolation. (A discussion on the JAX repo implied that they would not integrate this feature)
• Remove interpolation options. Since we only want to perform linear inter- and extrapolation, there is no need for added complexity.
• Update the regression tests since the results are different. Notably, the results now look better (more correct) than before. This also shows that the linear extrapolation helps with the solution and simulation.

Closes #83

How it works

To extrapolate linearly, one has to adjust the function that computes weights and indices of the coordinate points. For the linear interpolation, the JAX code can be found here.

This implementation implicitly assumes that coordinate is in the domain (interpolation). For values outside of the domain (extrapolation), one has to adjust the first line to

lower = jnp.clip(jnp.floor(coordinate), 0, input_size - 2)

where input_size corresponds to a dimension of the input shape. This leads to the same behavior for interpolation but assigns the correct indices and weights needed for linear extrapolation for coordinates outside the domain.

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[hmgaudecker]

Very cool. Had a brief look and leaving the note so I can see down the road what has changed from this point onwards.

[hmgaudecker]

Excellent, thank you!!!

Just a tiny comment.