def completed_rundir(configuration, tmpdir_factory):
tendency_dataset_path = tmpdir_factory.mktemp("tendencies")
if configuration == ConfigEnum.predictor:
model = get_mock_predictor()
model_path = str(tmpdir_factory.mktemp("model"))
fv3fit.dump(model, str(model_path))
config = get_ml_config(model_path)
elif configuration == ConfigEnum.nudging:
tendencies = _tendency_dataset()
tendencies.to_zarr(
str(tendency_dataset_path.join("ds.zarr")), consolidated=True
)
config = get_nudging_config(str(tendency_dataset_path.join("ds.zarr")))
elif configuration == ConfigEnum.microphys_emulation:
model_path = str(tmpdir_factory.mktemp("model").join("model.tf"))
model = create_emulation_model()
model.save(model_path)
config = get_emulation_config(model_path)
else:
raise NotImplementedError()
rundir = tmpdir_factory.mktemp("rundir").join("subdir")
run_native(config, str(rundir))
return rundir
def test_ml_stepper_state_update(state):
ml_stepper = PureMLStepper(
get_mock_predictor("state_and_tendency"), timestep=900, hydrostatic=False
)
(tendencies, diagnostics, states) = ml_stepper(None, state)
assert set(states) == {"total_sky_downward_shortwave_flux_at_surface"}
assert set(tendencies) == {"dQ1"}
def test_adapter_regression(state, regtest, tmpdir_factory):
model_path = str(tmpdir_factory.mktemp("model"))
mock = get_mock_predictor(dQ1_tendency=1 / 86400)
fv3fit.dump(mock, model_path)
adapted_model = Adapter(
Config(model_path, {
"air_temperature": "dQ1",
"specific_humidity": "dQ2"
}), 900)
transform = StepTransformer(
adapted_model,
MockDerivedState(state),
"machine_learning",
diagnostic_variables={
"tendency_of_specific_humidity_due_to_machine_learning",
"tendency_of_air_temperature_due_to_machine_learning",
"tendency_of_internal_energy_due_to_machine_learning",
},
timestep=900,
)
def add_one_to_temperature():
state["air_temperature"] += 1
return {"some_diag": state["specific_humidity"]}
out = transform(add_one_to_temperature)()
# ensure tendency of internal energy is non-zero somewhere (GH#1433)
max = abs(out["tendency_of_internal_energy_due_to_machine_learning"]).max()
assert max.values.item() > 1e-6
# sort to make the check deterministic
regression_state(out, regtest)
def test_mock_predictor_checksum(output_type, state, regtest):
mock_model = get_mock_predictor(output_type)
predicted = mock_model.predict_columnwise(state, feature_dim="z")
test_state_regression(predicted, regtest)
def ml_stepper(model_output_type):
timestep = 900
mock_model = get_mock_predictor(model_output_type)
ml_stepper = PureMLStepper(mock_model, timestep, hydrostatic=False)
return ml_stepper