in causalml/inference/tree/uplift.pyx [0:0]
def arr_evaluate_Chi(np.ndarray[P_TYPE_t, ndim=1] node_summary_p,
np.ndarray[N_TYPE_t, ndim=1] node_summary_n):
'''
Calculate Chi-Square statistic as split evaluation criterion for a given node.
Args
----
node_summary_p : array of shape [n_class]
Has type numpy.double.
The positive probabilities of each of the control
and treament groups of the current node, i.e. [P(Y=1|T=i)...]
node_summary_n : array of shape [n_class]
Has type numpy.int32.
The counts of each of the control
and treament groups of the current node, i.e. [N(T=i)...]
Returns
-------
d_res : Chi-Square
'''
cdef int n_class = node_summary_p.shape[0]
cdef P_TYPE_t p_c = node_summary_p[0]
cdef P_TYPE_t d_res = 0.0
cdef int i = 0
cdef P_TYPE_t max_eps_pc = max(0.1 ** 6, p_c)
cdef P_TYPE_t max_eps_1_pc = max(0.1 ** 6, 1 - p_c)
cdef P_TYPE_t diff_sq = 0.0
for i in range(1, n_class):
diff_sq = (node_summary_p[i] - p_c) * (node_summary_p[i] - p_c)
d_res += (diff_sq / max_eps_pc + diff_sq / max_eps_1_pc)
return d_res