Save predictions as an nc file at test time

[patel-zeel]

Is your feature request related to a problem? Please describe.
The feature request is motivated by the following problems:

• I was testing climate-learn downscaling against my models and wanted to see if there is a pattern in errors, e.g., during certain hours of the day, predictions are more erroneous. trainer.test(model, dm) provides a nice summary of metrics but does not save the predictions along with lat, lon, time.
• If someone wants to visualize the predictions side-by-side for a particular date from different models, it seems difficult to do it currently.

Describe the solution you'd like
A possible way could be to use a flow similar to cl.utils.visualize_at_index function.

climate-learn/src/climate_learn/utils/visualize.py

Lines 10 to 12 in 1a46b08

	def visualize_at_index(mm, dm, in_transform, out_transform, variable, src, index=0):
	lat, lon = dm.get_lat_lon()
	extent = [lon.min(), lon.max(), lat.min(), lat.max()]

There could be a function cl.utils.save_nc which may look like:

def save_nc(mm, dm, in_transform, out_transform, variable, src, save_dir):
    ...

This function can save the predictions in exactly the same format as the data nc files with lat, lon and time co-ordinates. It should be able to retrieve lat, lon, time from data module dm.

Willingness to work on a PR
I'll be happy to work on a PR to make this happen!

Additional context
This feature may also be useful to climate researchers who want to produce a time-lapse video of predictions with other libraries with additional geolayers similar to these examples in geemap library.

[jasonjewik]

Hi @patel-zeel, thank you for the suggestion. Because of travel, I will not be able to work on this for the next month. However, a feature like what you have suggested about saving latitude, longitude, and time is on the horizon for ClimateLearn. You can read more about my plans for the future of the library on this Google Doc. I plan to make a more formal post about these proposals when I return from my trip. In the meantime, I would very much appreciate you opening a PR to resolve this issue.