"""
------------------------------------------
File name: stat_utils.py
Author: Clare Huang
"""
import itertools
import numpy as np
[docs]def calculate_covariance(var_a, var_b):
"""
Calculate covariance of two variables in time.
Args:
var_a: a numpy array or handle that can access elements via [time, lat, lon]
var_b: a numpy array or handle that can access elements via [time, lat, lon]
Returns:
cov_map in dimension of (lat, lon)
"""
lat_dim = var_a.shape[1]
lon_dim = var_a.shape[2]
cov_map = np.zeros((lat_dim, lon_dim))
for j in range(lat_dim): # has to loop through a dimension to conserve memory
cov_matrix = np.cov(m=var_a[:, j, :], y=var_b[:, j, :], rowvar=False)
row_cov = np.diagonal(cov_matrix, offset=lon_dim)
cov_map[j, :] = row_cov
return cov_map