def earnn(
experiment="one_month_forecast",
include_pred_month=True,
surrounding_pixels=None,
pretrained=True,
ignore_vars=None,
include_static=True,
):
# if the working directory is alread ml_drought don't need ../data
if Path(".").absolute().as_posix().split("/")[-1] == "ml_drought":
data_path = Path("data")
else:
data_path = Path("../data")
if not pretrained:
predictor = EARecurrentNetwork(
hidden_size=128,
data_folder=data_path,
experiment=experiment,
include_pred_month=include_pred_month,
surrounding_pixels=surrounding_pixels,
ignore_vars=ignore_vars,
include_static=include_static,
)
predictor.train(num_epochs=50, early_stopping=10)
predictor.evaluate(save_preds=True)
predictor.save_model()
else:
predictor = load_model(data_path /
f"models/{experiment}/ealstm/model.pt")
def earnn(
experiment="one_month_forecast",
include_pred_month=True,
surrounding_pixels=None,
pretrained=True,
ignore_vars=None,
):
data_path = get_data_path()
if not pretrained:
predictor = EARecurrentNetwork(
hidden_size=128,
data_folder=data_path,
experiment=experiment,
include_pred_month=include_pred_month,
surrounding_pixels=surrounding_pixels,
ignore_vars=ignore_vars,
)
predictor.train(num_epochs=50, early_stopping=5)
predictor.evaluate(save_preds=True)
predictor.save_model()
else:
predictor = load_model(data_path /
f"models/{experiment}/ealstm/model.pt")
test_file = data_path / f"features/{experiment}/test/2018_3"
assert test_file.exists()
all_explanations_for_file(test_file, predictor, batch_size=100)
def earnn(experiment='one_month_forecast',
include_pred_month=True,
surrounding_pixels=None,
pretrained=True):
# if the working directory is alread ml_drought don't need ../data
if Path('.').absolute().as_posix().split('/')[-1] == 'ml_drought':
data_path = Path('data')
else:
data_path = Path('../data')
if not pretrained:
predictor = EARecurrentNetwork(hidden_size=128,
data_folder=data_path,
experiment=experiment,
include_pred_month=include_pred_month,
surrounding_pixels=surrounding_pixels)
predictor.train(num_epochs=50, early_stopping=5)
predictor.evaluate(save_preds=True)
predictor.save_model()
else:
predictor = load_model(data_path /
f'models/{experiment}/ealstm/model.pt')
test_file = data_path / f'features/{experiment}/test/2018_3'
assert test_file.exists()
all_shap_for_file(test_file, predictor, batch_size=100)
def earnn(
experiment="one_month_forecast",
include_pred_month=True,
surrounding_pixels=None,
pretrained=False,
explain=False,
static="features",
ignore_vars=None,
num_epochs=50,
early_stopping=5,
static_embedding_size=10,
hidden_size=128,
predict_delta=False,
spatial_mask=None,
include_latlons=False,
normalize_y=True,
include_prev_y=True,
include_yearly_aggs=True, # new
clear_nans=True,
weight_observations=False,
pred_month_static=False,
):
data_path = get_data_path()
if not pretrained:
predictor = EARecurrentNetwork(
hidden_size=hidden_size,
data_folder=data_path,
experiment=experiment,
include_pred_month=include_pred_month,
surrounding_pixels=surrounding_pixels,
static=static,
static_embedding_size=static_embedding_size,
ignore_vars=ignore_vars,
predict_delta=predict_delta,
spatial_mask=spatial_mask,
include_latlons=include_latlons,
normalize_y=normalize_y,
include_prev_y=include_prev_y,
include_yearly_aggs=include_yearly_aggs,
clear_nans=clear_nans,
weight_observations=weight_observations,
pred_month_static=pred_month_static,
)
predictor.train(num_epochs=num_epochs, early_stopping=early_stopping)
predictor.evaluate(save_preds=True)
predictor.save_model()
else:
predictor = load_model(data_path /
f"models/{experiment}/ealstm/model.pt")
if explain:
test_file = data_path / f"features/{experiment}/test/2018_3"
assert test_file.exists()
all_explanations_for_file(test_file, predictor, batch_size=100)
def test_save(self, tmp_path, monkeypatch):
features_per_month = 5
dense_features = [10]
input_dense_features = copy(dense_features)
hidden_size = 128
rnn_dropout = 0.25
include_latlons = True
include_pred_month = True
include_yearly_aggs = True
yearly_agg_size = 3
def mocktrain(self):
self.model = EALSTM(features_per_month,
dense_features,
hidden_size,
rnn_dropout,
include_latlons,
include_pred_month,
experiment='one_month_forecast',
yearly_agg_size=yearly_agg_size)
self.features_per_month = features_per_month
self.yearly_agg_size = yearly_agg_size
monkeypatch.setattr(EARecurrentNetwork, 'train', mocktrain)
model = EARecurrentNetwork(hidden_size=hidden_size,
dense_features=dense_features,
include_pred_month=include_pred_month,
include_latlons=include_latlons,
rnn_dropout=rnn_dropout,
data_folder=tmp_path,
include_yearly_aggs=include_yearly_aggs)
model.train()
model.save_model()
assert (tmp_path / 'models/one_month_forecast/ealstm/model.pt').exists(), \
f'Model not saved!'
model_dict = torch.load(model.model_dir / 'model.pt',
map_location='cpu')
for key, val in model_dict['model']['state_dict'].items():
assert (model.model.state_dict()[key] == val).all()
assert model_dict['model']['features_per_month'] == features_per_month
assert model_dict['model']['yearly_agg_size'] == yearly_agg_size
assert model_dict['hidden_size'] == hidden_size
assert model_dict['rnn_dropout'] == rnn_dropout
assert model_dict['dense_features'] == input_dense_features
assert model_dict['include_pred_month'] == include_pred_month
assert model_dict['include_latlons'] == include_latlons
assert model_dict['include_yearly_aggs'] == include_yearly_aggs
assert model_dict['experiment'] == 'one_month_forecast'
def test_save(self, tmp_path, monkeypatch):
features_per_month = 5
dense_features = [10]
input_dense_features = copy(dense_features)
hidden_size = 128
rnn_dropout = 0.25
include_latlons = True
include_pred_month = True
include_yearly_aggs = True
yearly_agg_size = 3
include_prev_y = True
normalize_y = False
def mocktrain(self):
self.model = EALSTM(
features_per_month,
dense_features,
hidden_size,
rnn_dropout,
include_latlons,
include_pred_month,
experiment="one_month_forecast",
yearly_agg_size=yearly_agg_size,
include_prev_y=include_prev_y,
)
self.features_per_month = features_per_month
self.yearly_agg_size = yearly_agg_size
monkeypatch.setattr(EARecurrentNetwork, "train", mocktrain)
model = EARecurrentNetwork(
hidden_size=hidden_size,
dense_features=dense_features,
include_pred_month=include_pred_month,
include_latlons=include_latlons,
rnn_dropout=rnn_dropout,
data_folder=tmp_path,
include_yearly_aggs=include_yearly_aggs,
normalize_y=normalize_y,
)
model.train()
model.save_model()
assert (tmp_path / "models/one_month_forecast/ealstm/model.pt"
).exists(), f"Model not saved!"
model_dict = torch.load(model.model_dir / "model.pt",
map_location="cpu")
for key, val in model_dict["model"]["state_dict"].items():
assert (model.model.state_dict()[key] == val).all()
assert model_dict["model"]["features_per_month"] == features_per_month
assert model_dict["model"]["yearly_agg_size"] == yearly_agg_size
assert model_dict["hidden_size"] == hidden_size
assert model_dict["rnn_dropout"] == rnn_dropout
assert model_dict["dense_features"] == input_dense_features
assert model_dict["include_pred_month"] == include_pred_month
assert model_dict["include_latlons"] == include_latlons
assert model_dict["include_yearly_aggs"] == include_yearly_aggs
assert model_dict["experiment"] == "one_month_forecast"
assert model_dict["include_prev_y"] == include_prev_y
assert model_dict["normalize_y"] == normalize_y
)
# test the prediction functionality
# test_arrays_dict, preds_dict = ealstm.predict()
ealstm.evaluate(
spatial_unit_name='station_id',
save_preds=True
)
results_dict = json.load(open('data/models/one_timestep_forecast/ealstm/results.json', 'rb'))
print("Overall RMSE: ", results_dict['total'])
# 1 epoch = 1.04 / 2.81
# 100 epochs = 0.72 / 2.04
# save the model
ealstm.save_model()
# load the model
from src.models import load_model
ealstm2 = load_model(Path('data/models/one_timestep_forecast/ealstm/model.pt'))
# ------------------------------------------------------------------------
## ANALYSIS
# ------------------------------------------------------------------------
import matplotlib.pyplot as plt
# checking the performance of the models
EXPERIMENT = 'one_timestep_forecast'
from src.analysis import read_pred_data
from src.analysis.evaluation import join_true_pred_da