def spacetime2(reg, sReg, df, ko, omega_age_smooth, lambda_time_smooth,
lambda_time_smooth_nodata, zeta_space_smooth, zeta_space_smooth_nodata):
"""
Compute the spacetime weight matrix for a super region. Full data set tells
which values need weights, train data set are the residuals which need
weighting.
"""
full_sub = df[(df.region == reg)]
train_sub = df[(df.super_region == sReg) & (ko)]
year_start = np.min(df.year)
year_end = np.max(df.year)
Wat = timeW(full_sub, train_sub, omega_age_smooth, lambda_time_smooth,
lambda_time_smooth_nodata, year_start, year_end).astype("float32")
NR, SN, C, R, SR = matCRS(full_sub, train_sub)
xi_mat = calculate_xi_matrix(full_sub, train_sub, zeta_space_smooth,
zeta_space_smooth_nodata).astype("float32")
NR = weight_matrix(NR, xi_mat[:,0], Wat).astype("float32")
SN = weight_matrix(SN, xi_mat[:,1], Wat).astype("float32")
C = weight_matrix(C, xi_mat[:,2], Wat).astype("float32")
R = weight_matrix(R, xi_mat[:,3], Wat).astype("float32")
SR = weight_matrix(SR, xi_mat[:,4], Wat).astype("float32")
final = EV("NR + SN + C + R + SR").astype("float32")
del NR, SN, C, R, SR
account_missing = final.sum(0)
account_missing[account_missing == .0] = 1.
return EV("final / account_missing").astype("float32")
def weight_matrix(valid_positions, xi_vector, weight_matrix):
"""
(matrix, vector, matrix)
Given a matrix of valid positions for an analytic region (valid_positions),
a vector of appropriate xi weights to use for each column in that vector
(xi_vector), and an age year weighted matrix generated by timeW will return
a matrix re-weighted so that each column adds up to the corresponding xi
value in the xi_vector.
"""
weights = EV("valid_positions * weight_matrix")
sum_of_weights = weights.sum(0)
sum_of_weights[sum_of_weights == .0] = 1.
return EV("(weights / sum_of_weights) * xi_vector")