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 and Details of H2OCoxPHMOJOModel.

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)

Run Predictions

model.transform(testingDF).show(false)

You can also get model details via calling methods listed in Details of H2OCoxPHMOJOModel.

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)

Run Predictions

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

You can also get model details via calling methods listed in Details of H2OCoxPHMOJOModel.