Skip to content

Commit

Permalink
atleast 50% performance increase by removing the range() (git add -u)…
Browse files Browse the repository at this point in the history 
… and replacing it with einsum
  • Loading branch information
@CamDavidsonPilon
CamDavidsonPilon committed Jan 22, 2019
1 parent f113c41 commit 7c9ebb0
Show file tree
Hide file tree
Showing 2 changed files with 59 additions and 72 deletions.
127 changes: 57 additions & 70 deletions lifelines/fitters/coxph_fitter.py
View file
Expand Up @@ -513,6 +513,7 @@ def _newton_rhaphson(
return beta

def _get_efron_values_single(self, X, T, E, weights, beta):
# TODO: push the beta.reshape up one level.
"""
Calculates the first and second order vector differentials, with respect to beta.
Note that X, T, E are assumed to be sorted on T!
Expand Down Expand Up @@ -554,6 +555,7 @@ def _get_efron_values_single(self, X, T, E, weights, beta):
hessian = np.zeros((d, d))
gradient = np.zeros((1, d))
log_lik = 0
beta = beta.reshape(d)

# Init risk and tie sums to zero
x_death_sum = np.zeros((1, d))
Expand All @@ -564,15 +566,15 @@ def _get_efron_values_single(self, X, T, E, weights, beta):
# Init number of ties and weights
weight_count = 0.0
tied_death_counts = 0
scores = weights[:, None] * exp(dot(X, beta))
scores = weights * exp(dot(X, beta))

# Iterate backwards to utilize recursive relationship
for i in range(n - 1, -1, -1):
# Doing it like this to preserve shape
ti = T[i]
ei = E[i]
xi = X[i : i + 1]
score = scores[i : i + 1]
score = scores[i]
w = weights[i]

# Calculate phi values
Expand Down Expand Up @@ -609,41 +611,36 @@ def _get_efron_values_single(self, X, T, E, weights, beta):
partial_ll = 0
weighted_average = weight_count / tied_death_counts

for l in range(tied_death_counts):
if tied_death_counts > 1:

if tied_death_counts > 1:
# A good explaination for how Efron handles ties. Consider three of five subjects who fail at the time.
# As it is not known a priori that who is the first to fail, so one-third of
# (φ1 + φ2 + φ3) is adjusted from sum_j^{5} φj after one fails. Similarly two-third
# of (φ1 + φ2 + φ3) is adjusted after first two individuals fail, etc.

# A good explaination for how Efron handles ties. Consider three of five subjects who fail at the time.
# As it is not known a priori that who is the first to fail, so one-third of
# (φ1 + φ2 + φ3) is adjusted from sum_j^{5} φj after one fails. Similarly two-third
# of (φ1 + φ2 + φ3) is adjusted after first two individuals fail, etc.
# alot of this is now in einstien notation for performance, but see original "expanded" code here
#

increasing_proportion = l / tied_death_counts
denom = risk_phi - increasing_proportion * tie_phi
numer = risk_phi_x - increasing_proportion * tie_phi_x
# Hessian
a1 = (risk_phi_x_x - increasing_proportion * tie_phi_x_x) / denom
else:
denom = risk_phi
numer = risk_phi_x
# Hessian
a1 = risk_phi_x_x / denom

t = numer / denom
# Gradient
partial_gradient = partial_gradient + t
# In case numer and denom both are really small numbers,
# make sure to do division before multiplications
# this is faster than an outerproduct
a2 = t.T.dot(t)

partial_hessian = partial_hessian - (a1 - a2)
partial_ll = partial_ll - np.log(denom[0][0])

# Values outside tie sum
gradient = gradient + (x_death_sum - weighted_average * partial_gradient)
log_lik = log_lik + (dot(x_death_sum, beta)[0][0] + weighted_average * partial_ll)
hessian = hessian + (weighted_average * partial_hessian)
increasing_proportion = np.arange(tied_death_counts) / tied_death_counts
denom = 1. / (risk_phi - increasing_proportion * tie_phi)
numer = risk_phi_x - np.outer(increasing_proportion, tie_phi_x)
# Hessian
# computes outer products and sums them together.
a1 = np.einsum('ab, i->ab', risk_phi_x_x, denom) - np.einsum('ab, i->ab', tie_phi_x_x, increasing_proportion * denom)
else:
# no tensors here, but do some casting to make it easier in the converging step next.
denom = 1. /np.array([risk_phi])
numer = risk_phi_x
# Hessian
a1 = risk_phi_x_x * denom

t = numer * denom[:, None]
a2 = np.einsum('Bi, Bj->ij', t, t) # batch outer product

gradient = gradient + x_death_sum - weighted_average * t.sum(0)

log_lik = log_lik + dot(x_death_sum, beta)[0] + weighted_average * np.log(denom).sum()
hessian = hessian + weighted_average * (a2 - a1)

# reset tie values
tied_death_counts = 0
Expand Down Expand Up @@ -757,48 +754,38 @@ def _get_efron_values_batch(self, X, T, E, weights, beta): # pylint: disable=to
tie_phi_x = matrix_axis_0_sum_to_array(phi_x_i[deaths])
tie_phi_x_x = dot(xi_deaths.T, phi_i[deaths] * xi_deaths)

partial_ll = 0
partial_gradient = np.zeros((1, d))
partial_hessian = np.zeros((d, d))

weight_count = array_sum_to_scalar(weights_deaths)
weighted_average = weight_count / tied_death_counts

for l in range(tied_death_counts):
if tied_death_counts > 1:

if tied_death_counts > 1:
# A good explaination for how Efron handles ties. Consider three of five subjects who fail at the time.
# As it is not known a priori that who is the first to fail, so one-third of
# (φ1 + φ2 + φ3) is adjusted from sum_j^{5} φj after one fails. Similarly two-third
# of (φ1 + φ2 + φ3) is adjusted after first two individuals fail, etc.

# A good explaination for how Efron handles ties. Consider three of five subjects who fail at the time.
# As it is not known a priori that who is the first to fail, so one-third of
# (φ1 + φ2 + φ3) is adjusted from sum_j^{5} φj after one fails. Similarly two-third
# of (φ1 + φ2 + φ3) is adjusted after first two individuals fail, etc.
# alot of this is now in einstien notation for performance, but see original "expanded" code here
#

increasing_proportion = l / tied_death_counts
denom = risk_phi - increasing_proportion * tie_phi
numer = risk_phi_x - increasing_proportion * tie_phi_x
# Hessian
a1 = (risk_phi_x_x - increasing_proportion * tie_phi_x_x) / denom
else:
denom = risk_phi
numer = risk_phi_x
# Hessian
a1 = risk_phi_x_x / denom

t = numer / denom
# Gradient
partial_gradient = partial_gradient + t
# In case numer and denom both are really small numbers,
# make sure to do division before multiplications
# this is faster than an outerproduct
a2 = t.T.dot(t)

partial_hessian = partial_hessian - (a1 - a2)
partial_ll = partial_ll - np.log(denom)

# Values outside tie sum
gradient = gradient + (x_death_sum - weighted_average * partial_gradient)
log_lik = log_lik + (dot(x_death_sum, beta)[0] + weighted_average * partial_ll)
hessian = hessian + (weighted_average * partial_hessian)
increasing_proportion = np.arange(tied_death_counts) / tied_death_counts
denom = 1. / (risk_phi - increasing_proportion * tie_phi)
numer = risk_phi_x - np.outer(increasing_proportion, tie_phi_x)
# Hessian
# computes outer products and sums them together.
a1 = np.einsum('ab,i->ab', risk_phi_x_x, denom) - np.einsum('ab,i->ab', tie_phi_x_x, increasing_proportion * denom)
else:
# no tensors here, but do some casting to make it easier in the converging step next.
denom = 1. /np.array([risk_phi])
numer = risk_phi_x
# Hessian
a1 = risk_phi_x_x * denom

t = numer * denom[:, None]
a2 = np.einsum('Bi, Bj->ij', t, t) # batch outer product

gradient = gradient + x_death_sum - weighted_average * t.sum(0)
log_lik = log_lik + dot(x_death_sum, beta)[0] + weighted_average * np.log(denom).sum()
hessian = hessian + weighted_average * (a2 - a1)

return hessian, gradient, log_lik

Expand Down
4 changes: 2 additions & 2 deletions tests/test_estimation.py
View file
Expand Up @@ -1123,7 +1123,7 @@ def test_print_summary_with_decimals(self, rossi, cph):
sys.stdout = out

cph = CoxPHFitter()
cph.fit(rossi, duration_col="week", event_col="arrest")
cph.fit(rossi, duration_col="week", event_col="arrest", batch_mode=True)
cph._time_fit_was_called = "2018-10-23 02:40:45 UTC"
cph.print_summary(decimals=1)
output_dec_1 = out.getvalue().strip().split()
Expand All @@ -1134,6 +1134,7 @@ def test_print_summary_with_decimals(self, rossi, cph):
assert output_dec_1 != output_dec_3
finally:
sys.stdout = saved_stdout
cph.fit(rossi, duration_col="week", event_col="arrest", batch_mode=False)

def test_print_summary(self, rossi, cph):

Expand Down Expand Up @@ -1645,7 +1646,6 @@ def test_cluster_option_with_strata(self, regression_dataset):
cph = CoxPHFitter()
cph.fit(df, "T", "E", cluster_col="id", strata=["strata"], show_progress=True)
expected = pd.Series({"var": 0.643})
cph.print_summary()
assert_series_equal(cph.summary["se(coef)"], expected, check_less_precise=2, check_names=False)

def test_robust_errors_with_less_trival_weights_is_the_same_as_R(self, regression_dataset):
Expand Down

0 comments on commit 7c9ebb0

Please sign in to comment.