Possible bug in schema.apply_time_conventions (?)

[timothyas]

First off, thanks to all the developers for WeatherBench2! It is making ML weather model comparisons very streamlined ...

I am curious if there is a bug in the following lines of schema.apply_time_conventions:

weatherbench2/weatherbench2/schema.py

Lines 32 to 43 in 16e0131

	if by_init:
	# Need to rename time dimension because different from time dimension in
	# truth dataset
	forecast = forecast.rename({'time': 'init_time'})
	valid_time = forecast.init_time + forecast.lead_time
	forecast.coords['valid_time'] = valid_time
	assert not hasattr(
	forecast, 'time'
	), f'Forecast should not have time dimension at this point: {forecast}'
	else:
	init_time = forecast.time - forecast.lead_time
	forecast.coords['init_time'] = init_time

Currently, we only get to this code block if prediction_timedelta is in the dataset, but not if that coordinate is already named lead_time. In that case, the coordinate time is never renamed to init_time, and the coordinate valid_time is never created, which messes up this line (at least).

A simple fix would be to untab this, and then change the "else" to an elif as I did here. I'd be happy to open a PR with that if you agree with this and I'm not missing something.