Train CoxPH Model in Sparkling Water

Sparkling Water provides API for H2O CoxPH in Scala and Python. The following sections describe how to train the CoxPH model in Sparkling Water in both languages. See also Parameters of H2OCoxPH.

Scala

First, let’s start Sparkling Shell as

./bin/sparkling-shell

Start H2O cluster inside the Spark environment

import ai.h2o.sparkling._
import java.net.URI
val hc = H2OContext.getOrCreate()

Parse the data using H2O and convert them to Spark Frame

import org.apache.spark.SparkFiles
spark.sparkContext.addFile("https://raw.githubusercontent.com/h2oai/sparkling-water/master/examples/smalldata/coxph_test/heart.csv")
val heartDF = spark.read.option("header", "true").option("inferSchema", "true").csv(SparkFiles.get("heart.csv"))
val Array(trainingDF, testingDF) = heartDF.randomSplit(Array(0.8, 0.2), seed = 12345)

Train the model. You can configure all the available CoxPH parameters using provided setters.

import ai.h2o.sparkling.ml.algos.H2OCoxPH
val estimator = new H2OCoxPH().
    setStartCol("start").
    setStopCol("stop").
    setTies("breslow").
    setLabelCol("event")
val model = estimator.fit(trainingDF)

You can also get raw model details by calling the getModelDetails() method available on the model as:

model.getModelDetails()

Run Predictions

model.transform(testingDF).show(false)

Python

First, let’s start PySparkling Shell as

./bin/pysparkling

Start H2O cluster inside the Spark environment

from pysparkling import *
hc = H2OContext.getOrCreate()

Parse the data using H2O and convert them to Spark Frame

import h2o
heartFrame = h2o.import_file("https://raw.githubusercontent.com/h2oai/sparkling-water/master/examples/smalldata/coxph_test/heart.csv")
trainingFrame, testingFrame = heartFrame.split_frame(ratios = [.8], seed = 1234)
trainingDF = hc.asSparkFrame(trainingFrame)
testingDF = hc.asSparkFrame(testingFrame)

Train the model. You can configure all the available CoxPH arguments using provided setters or constructor parameters.

from pysparkling.ml import H2OCoxPH
estimator = H2OCoxPH()\
    .setStartCol('start')\
    .setStopCol('stop')\
    .setTies('breslow')\
    .setLabelCol('event')
model = estimator.fit(trainingDF)

You can also get raw model details by calling the getModelDetails() method available on the model as:

model.getModelDetails()

Run Predictions

model.transform(testingDF).show(truncate = False)