def test_integrate_stepwise(sample_dim_name, dt):
"""
We need an integrate_stepwise routine for scripts that evaluate the model
timeseries, since the saved model operates on a single timestep. This code
is hard to write for a variety of reasons.
The test here makes sure that the integration code performs an identical
integration to the one used to train the model.
train_keras_model does this by predicting a timeseries
directly, while integrate_stepwise integrates the stepwise model over
many timesteps.
"""
train_dataset = get_train_dataset(sample_dim_name, dt)
np.random.seed(0)
tf.random.set_seed(1)
model = _BPTTTrainer(
sample_dim_name,
["a", "b"],
n_units=32,
n_hidden_layers=4,
kernel_regularizer=None,
train_batch_size=48,
optimizer="adam",
)
model.fit_statistics(train_dataset)
model.fit(train_dataset)
air_temperature, specific_humidity = model.train_keras_model.predict(
x=model.get_keras_inputs(train_dataset)
)
out_ds = model.predictor_model.integrate_stepwise(train_dataset)
np.testing.assert_allclose(
out_ds["air_temperature"].values[:, 0, :], air_temperature[:, 0, :], atol=1e-6
)
np.testing.assert_allclose(
out_ds["specific_humidity"].values[:, 0, :],
specific_humidity[:, 0, :],
atol=1e-6,
)
np.testing.assert_allclose(
out_ds["air_temperature"].values[:, 1, :], air_temperature[:, 1, :], atol=1e-6
)
np.testing.assert_allclose(
out_ds["specific_humidity"].values[:, 1, :],
specific_humidity[:, 1, :],
atol=1e-6,
)
np.testing.assert_allclose(
out_ds["air_temperature"].values, air_temperature, atol=1e-5
)
np.testing.assert_allclose(
out_ds["specific_humidity"].values, specific_humidity, atol=1e-5
)
assert not np.all(air_temperature == specific_humidity)
def get_model(sample_dim_name, use_moisture_limiter):
return _BPTTTrainer(
sample_dim_name,
["a", "b"],
n_units=32,
n_hidden_layers=4,
kernel_regularizer=None,
train_batch_size=48,
optimizer="adam",
use_moisture_limiter=use_moisture_limiter,
)
def test_predict_model_gives_different_tendencies(sample_dim_name, dt):
train_dataset = get_train_dataset(sample_dim_name, dt)
model = _BPTTTrainer(
sample_dim_name,
["a", "b"],
n_units=32,
n_hidden_layers=4,
kernel_regularizer=None,
train_batch_size=48,
optimizer="adam",
)
model.fit_statistics(train_dataset)
model.fit(train_dataset)
tendency_ds = model.predictor_model.predict(train_dataset.isel(time=0))
assert not np.all(tendency_ds["dQ1"].values == tendency_ds["dQ2"].values)
def test_train_model_gives_different_outputs(sample_dim_name, dt):
train_dataset = get_train_dataset(sample_dim_name, dt)
model = _BPTTTrainer(
sample_dim_name,
["a", "b"],
n_units=32,
n_hidden_layers=4,
kernel_regularizer=None,
train_batch_size=48,
optimizer="adam",
)
model.fit_statistics(train_dataset)
model.fit(train_dataset)
air_temperature, specific_humidity = model.train_keras_model.predict(
x=model.get_keras_inputs(train_dataset)
)
assert not np.all(air_temperature == specific_humidity)
def test_reloaded_model_gives_same_outputs(sample_dim_name, dt):
train_dataset = get_train_dataset(sample_dim_name, dt)
model = _BPTTTrainer(
sample_dim_name,
["a", "b"],
n_units=32,
n_hidden_layers=4,
kernel_regularizer=None,
train_batch_size=48,
optimizer="adam",
)
model.fit_statistics(train_dataset)
model.fit(train_dataset)
with tempfile.TemporaryDirectory() as tmpdir:
fv3fit.dump(model.predictor_model, tmpdir)
loaded = fv3fit.load(tmpdir)
first_timestep = train_dataset.isel(time=0)
reference_output = model.predictor_model.predict(first_timestep)
# test that loaded model gives the same predictions
loaded_output = loaded.predict(first_timestep)
xr.testing.assert_equal(reference_output, loaded_output)
def test_train_model_uses_correct_given_tendency(sample_dim_name, dt):
train_dataset = get_train_dataset(sample_dim_name, dt)
np.random.seed(0)
tf.random.set_seed(1)
model = _BPTTTrainer(
sample_dim_name,
["a", "b"],
n_units=32,
n_hidden_layers=4,
kernel_regularizer=None,
train_batch_size=48,
optimizer="adam",
)
model.fit_statistics(train_dataset)
model.fit(train_dataset)
assert not np.all(
train_dataset["air_temperature_tendency_due_to_model"].values == 0.0
)
air_temperature, specific_humidity = model.train_keras_model.predict(
x=model.get_keras_inputs(train_dataset)
)
dQ1, dQ2 = model.train_tendency_model.predict(
x=model.get_keras_inputs(train_dataset)
)
# take difference of first output from initial input state
Q1_train_tendency = (
air_temperature[:, 0, :] - train_dataset["air_temperature"][:, 0, :]
) / (dt.total_seconds())
Q2_train_tendency = (
specific_humidity[:, 0, :] - train_dataset["specific_humidity"][:, 0, :]
) / (dt.total_seconds())
np.testing.assert_allclose(
Q1_train_tendency - dQ1[:, 0, :],
train_dataset["air_temperature_tendency_due_to_model"][:, 0, :].values,
atol=1e-6,
)
np.testing.assert_allclose(
Q2_train_tendency - dQ2[:, 0, :],
train_dataset["specific_humidity_tendency_due_to_model"][:, 0, :].values,
atol=1e-6,
)
# check for second timestep also
Q1_train_tendency = (air_temperature[:, 1, :] - air_temperature[:, 0, :]) / (
dt.total_seconds()
)
Q2_train_tendency = (specific_humidity[:, 1, :] - specific_humidity[:, 0, :]) / (
dt.total_seconds()
)
np.testing.assert_allclose(
Q1_train_tendency - dQ1[:, 1, :],
train_dataset["air_temperature_tendency_due_to_model"][:, 1, :].values,
atol=1e-6,
)
np.testing.assert_allclose(
Q2_train_tendency - dQ2[:, 1, :],
train_dataset["specific_humidity_tendency_due_to_model"][:, 1, :].values,
atol=1e-6,
)