code: def check_cox(
      shelter,
      x,
      expected_hazard_cols,
      expected_survival_cols,
      stratify_by=None,
      weight=None,
  ):
      shelter = shelter[x + ["event"] + ([weight] if weight else [])]
  
      for col in stratify_by:
          shelter[col] = shelter[col].astype("category")
  
      shelter_h2o = h2o.H2OFrame(shelter)
  
      for col in stratify_by:
          shelter_h2o[col] = shelter_h2o[col].asfactor()
  
      cph_h2o = H2OCoxProportionalHazardsEstimator(
          stop_column="surv_hours", stratify_by=stratify_by
      )
      cph_h2o.train(
          x=x,
          y="event",
          weights_column=weight,
          training_frame=shelter_h2o,
      )
  
      assert cph_h2o.model_id != ""
  
      pred_h2o = cph_h2o.predict(test_data=shelter_h2o)
      assert len(pred_h2o) == len(shelter)
      hazard_h2o_as_pandas = (
          cph_h2o.baseline_hazard_frame.as_data_frame(use_pandas=True)
      )
  
      hazard_h2o_reordered_columns = (
          hazard_h2o_as_pandas.drop("t", axis="columns")
          .reset_index(drop=True)
          .sort_index(axis=1)
      )
  
      assert (
          hazard_h2o_reordered_columns.columns.values.tolist()
          == expected_hazard_cols
      )
  
      survival_h2o_as_pandas = (
          cph_h2o.baseline_survival_frame.as_data_frame(use_pandas=True)
      )
  
      survival_h2o_reordered_columns = (
          survival_h2o_as_pandas.drop("t", axis="columns")
          .reset_index(drop=True)
          .sort_index(axis=1)
      )
  
      assert (
          survival_h2o_reordered_columns.columns.values.tolist()
          == expected_survival_cols
      )
  
file_path: h2o-3/h2o-py/tests/testdir_algos/coxph/pyunit_coxph_baseline_names_shelter.py
imports: import sys
  import h2o
  from lifelines import CoxPHFitter
  from pandas.testing import assert_frame_equal
  from pandas import read_csv
  import numpy
  from tests import pyunit_utils
  from h2o.estimators.coxph import H2OCoxProportionalHazardsEstimator
label_desc: Use image node_id 5 for calling a global function with example usage: check_cox(shelter, x, expected_hazard_cols, expected_survival_cols, stratify_by, weight) without return types