code: def check_cox(
      shelter, x, expected_formula, stratify_by=None, weight=None
  ):
      shelter = shelter[x + ["event"] + ([weight] if weight else [])]
  
      cph_py = (
          CoxPHFitter(strata=stratify_by)
          if stratify_by
          else CoxPHFitter()
      )
  
      for col in stratify_by:
          shelter[col] = shelter[col].astype("category")
  
      cph_py.fit(
          shelter,
          duration_col="surv_hours",
          event_col="event",
          weights_col=weight,
      )
      cph_py.print_summary()
  
      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 != ""
      assert cph_h2o.formula() == expected_formula, (
          "Expected formula to be '"
          + expected_formula
          + "' but it was "
          + cph_h2o.formula()
      )
  
      pred_h2o = cph_h2o.predict(test_data=shelter_h2o)
      assert len(pred_h2o) == len(shelter)
      metrics_h2o = cph_h2o.model_performance(shelter_h2o)
      concordance_py = concordance_for_lifelines(cph_py)
      assert abs(concordance_py - metrics_h2o.concordance()) < 0.001
      hazard_h2o_as_pandas = (
          cph_h2o.baseline_hazard_frame.as_data_frame(use_pandas=True)
      )
  
      hazard_py = cph_py.baseline_hazard_
  
      for col_name in hazard_py.columns:
          if isinstance(col_name, int):
              new_name = "({0})".format(col_name)
          else:
              new_name = str(col_name)
          hazard_py.rename(columns={col_name: new_name}, inplace=True)
  
      hazard_py_reordered_columns = hazard_py.reset_index(
          drop=True
      ).sort_index(axis=1)
      hazard_h2o_reordered_columns = (
          hazard_h2o_as_pandas.drop("t", axis="columns")
          .reset_index(drop=True)
          .sort_index(axis=1)
      )
  
      assert_frame_equal(
          hazard_py_reordered_columns,
          hazard_h2o_reordered_columns,
          check_dtype=False,
          check_index_type=False,
          check_column_type=False,
          check_less_precise=True,
      )
  
      survival_h2o_as_pandas = (
          cph_h2o.baseline_survival_frame.as_data_frame(use_pandas=True)
      )
  
      survival_py = cph_py.baseline_survival_
  
      for col_name in survival_py.columns:
          if isinstance(col_name, int):
              new_name = "({0})".format(col_name)
          else:
              new_name = str(col_name)
          survival_py.rename(columns={col_name: new_name}, inplace=True)
  
      survival_py_reordered_columns = survival_py.reset_index(
          drop=True
      ).sort_index(axis=1)
      survival_h2o_reordered_columns = (
          survival_h2o_as_pandas.drop("t", axis="columns")
          .reset_index(drop=True)
          .sort_index(axis=1)
      )
  
      assert_frame_equal(
          survival_py_reordered_columns,
          survival_h2o_reordered_columns,
          check_dtype=False,
          check_index_type=False,
          check_column_type=False,
          check_less_precise=True,
      )
  
file_path: h2o-3/h2o-py/tests/testdir_algos/coxph/pyunit_coxph_concordance_survival_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 6 for calling a global function with example usage: check_cox(shelter, x, expected_formula, stratify_by, weight) without return types