add type annotation

mllam · Nov 29, 2024 · cf8e3e4 · cf8e3e4
1 parent c1f706c
commit cf8e3e4
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/neural_lam/vis.py b/neural_lam/vis.py
@@ -2,6 +2,7 @@
 import matplotlib
 import matplotlib.pyplot as plt
 import numpy as np
+import xarray as xr
 
 # Local
 from . import utils
@@ -66,8 +67,8 @@ def plot_error_map(errors, datastore: BaseRegularGridDatastore, title=None):
 @matplotlib.rc_context(utils.fractional_plot_bundle(1))
 def plot_prediction(
     datastore: BaseRegularGridDatastore,
-    da_prediction=None,
-    da_target=None,
+    da_prediction: xr.DataArray = None,
+    da_target: xr.DataArray = None,
     title=None,
     vrange=None,
 ):