Get survival function with Spark ML

Question

I am training an Accelerated failure time model with PySpark (from pyspark.ml.regression import AFTSurvivalRegression)

Now I want to apply the model to new data and get the probability that the event will happen before time t (survival function), which method should I use ? The documentation is not clear to me : https://spark.apache.org/docs/2.1.0/api/python/pyspark.ml.html#pyspark.ml.regression.AFTSurvivalRegression

For example, if I do the following:

from pyspark.ml.regression import AFTSurvivalRegression
from pyspark.ml.linalg import Vectors

training = spark.createDataFrame([
    (1.218, 1.0, Vectors.dense(1.560, -0.605)),
    (2.949, 0.0, Vectors.dense(0.346, 2.158)),
    (3.627, 0.0, Vectors.dense(1.380, 0.231)),
    (0.273, 1.0, Vectors.dense(0.520, 1.151)),
    (4.199, 0.0, Vectors.dense(0.795, -0.226))], ["label", "censor", "features"])
quantileProbabilities = [0.25, 0.75]
aft = AFTSurvivalRegression(quantileProbabilities=quantileProbabilities,
                        quantilesCol="quantiles")

model = aft.fit(training)
model.transform(training).show(truncate=False)

I get as an output :

Does it mean that for the first line, P(event happening between 0.832 and 9.48) = 50% ?

Thanks

Get survival function with Spark ML

0 Answers