lifespan_predictor_sdk
Lifespan Predictor SDK
1"""Lifespan Predictor SDK""" 2 3from .survival_model import ( 4 CategoricalRange, 5 ContinuousRange, 6 IntegerRange, 7 SurvivalModel, 8) 9from .hpo_io import TuningTrials, stratified_kfold 10from .training_log import EvaluationResult, TrainingLog 11from .data_loader import ( 12 download_nasa_dataset, 13 get_nasa_train_dataframe, 14 get_nasa_test_dataframe, 15) 16 17__all__ = [ 18 "SurvivalModel", 19 "ContinuousRange", 20 "IntegerRange", 21 "CategoricalRange", 22 "stratified_kfold", 23 "TrainingLog", 24 "EvaluationResult", 25 "TuningTrials", 26 "download_nasa_dataset", 27 "get_nasa_train_dataframe", 28 "get_nasa_test_dataframe", 29]
class
SurvivalModel:
370class SurvivalModel: 371 """ 372 A DeepSurv survival model: train it, then predict time to failure. 373 374 A model is built from a config dict and trained with `fit()`. Preprocessing 375 (scaling and the optional engineered features below) is part of the model, so 376 `predict_mean()` and `predict()` take raw feature rows in the same shape as 377 the training data. 378 379 ```python 380 model = SurvivalModel({ 381 "model_type": "deepsurv", 382 "algorithm_arn": os.environ["SAGEMAKER_ALGORITHM_ARN"], 383 "epochs": 30, 384 "num_nodes": [24], 385 "max_rul": 125, 386 }) 387 model.fit(train_df, val_split=0.2) 388 model.deploy() 389 mean_times = model.predict_mean(test_df) 390 ``` 391 392 ### Hyperparameters 393 394 Every key is optional. A key marked *Range* may also be given as a 395 `ContinuousRange`, `IntegerRange` or `CategoricalRange`, which makes `fit()` 396 tune it (see `fit()`). 397 398 | Key | Type | Default | Range | Meaning | 399 |-----|------|---------|-------|---------| 400 | `model_type` | str | `"deepsurv"` | no | The model. `"deepsurv"` is the one this release offers. | 401 | `epochs` | int | `100` | yes | Maximum training epochs. With validation data, training stops early when the validation loss stops improving. | 402 | `num_nodes` | list of int | `[32, 32]` | yes | Hidden layer sizes, one entry per layer. | 403 | `dropout` | float | `0.1` | yes | Dropout rate after each hidden layer. | 404 | `learning_rate` | float | `0.01` | yes | Initial learning rate. It is reduced when the validation loss plateaus. | 405 | `batch_size` | int | `256` | yes | Rows per training batch. | 406 | `max_rul` | float or None | `None` | yes | Caps the training durations at this value, and limits the C-index to units whose duration is at most this. Common for turbofan data, where early life carries no wear signal. | 407 | `use_delta` | bool | `False` | yes | Adds, for each feature, its change from the unit's previous row. | 408 | `use_per_unit_normalisation` | bool | `False` | yes | Adds, for each feature, its offset from the unit's first row. | 409 | `rolling_means` | list of int or None | `None` | yes | Window sizes; adds, for each feature and window, the unit's rolling mean. | 410 | `calibration` | str or None | `None` | yes | `"isotonic"` calibrates the predicted survival curves on the validation data, so `fit()` then needs `val_df` or `val_split`. | 411 | `n_trials` | int | `2` | no | Tuning only: how many configurations to try. | 412 | `hpo_metric` | str | `"c-index"` | no | Tuning only: the score trials are compared by: `"c-index"`, `"ibs"` (integrated Brier score) or `"antolini"` (time-dependent C-index). | 413 | `verbose` | bool | `False` | no | Tuning only: log each trial in detail. | 414 415 ### SageMaker settings 416 417 Where training runs. They are not hyperparameters and cannot be tuned. 418 419 | Key | Type | If left out | Meaning | 420 |-----|------|-------------|---------| 421 | `algorithm_arn` | str | `SAGEMAKER_ALGORITHM_ARN` is used; without it, `fit()` raises | The ARN of the algorithm you subscribed to. | 422 | `instance_type` | str | `"ml.m5.large"` | The instance type of the training job. | 423 | `bucket` | str | `SAGEMAKER_BUCKET` is used; without it, SageMaker's default bucket for your account and region | The S3 bucket for training data and model artifacts. | 424 425 The job's IAM role is read from the `SAGEMAKER_ROLE_ARN` environment variable, 426 which must be set. 427 """ 428 429 def __init__(self, config): 430 """ 431 Create an untrained model. 432 433 Args: 434 config: The model's configuration: hyperparameters and SageMaker 435 settings, as listed in the tables above. Missing keys take their 436 defaults. 437 """ 438 # Minimal shared defaults 439 defaults = { 440 "model_type": "deepsurv", 441 "epochs": 100, 442 "max_rul": None, 443 } 444 445 self.config = {**defaults, **config} 446 447 # Shared state 448 self._is_fitted = False 449 self._training_log = None 450 self._tuning_trials = None 451 self._evaluation_results = None 452 self.in_features = self.config.get("in_features") 453 454 # Stashed DataFrames from fit() for post-fit conversion methods 455 self._fit_train_df = None 456 self._fit_val_df = None 457 self._fit_test_df = None 458 self._fit_test_labels_df = None 459 460 # Create backend. The backend reports its own name back into the config, 461 # so `get_config()` still answers the question without this file having to 462 # spell out which backends exist — that is _make_backend's business. 463 self._backend = _make_backend(self.config.get("backend")) 464 self.config["backend"] = self._backend.name 465 466 # Let backend merge its defaults and initialize state 467 self._backend.init_state(self.config) 468 469 # ------------------------------------------------------------------ 470 # Property accessors — delegate to backend for test compatibility 471 # ------------------------------------------------------------------ 472 473 474 475 476 477 @property 478 def _calibrator(self): 479 return getattr(self._backend, "_calibrator", None) 480 481 @_calibrator.setter 482 def _calibrator(self, value): 483 self._backend._calibrator = value 484 485 @property 486 def _estimator(self): 487 return getattr(self._backend, "_estimator", None) 488 489 @property 490 def _model_data(self): 491 return getattr(self._backend, "_model_data", None) 492 493 @_model_data.setter 494 def _model_data(self, value): 495 self._backend._model_data = value 496 497 @property 498 def _training_job_name(self): 499 return getattr(self._backend, "_training_job_name", None) 500 501 @_training_job_name.setter 502 def _training_job_name(self, value): 503 self._backend._training_job_name = value 504 505 @property 506 def _sagemaker_model_path(self): 507 return getattr(self._backend, "_sagemaker_model_path", None) 508 509 @_sagemaker_model_path.setter 510 def _sagemaker_model_path(self, value): 511 self._backend._sagemaker_model_path = value 512 513 @property 514 def _sagemaker_model_name(self): 515 return getattr(self._backend, "_sagemaker_model_name", None) 516 517 @_sagemaker_model_name.setter 518 def _sagemaker_model_name(self, value): 519 self._backend._sagemaker_model_name = value 520 521 @property 522 def _sagemaker_model(self): 523 return getattr(self._backend, "_sagemaker_model", None) 524 525 @_sagemaker_model.setter 526 def _sagemaker_model(self, value): 527 self._backend._sagemaker_model = value 528 529 @property 530 def _fleet_sub_models(self): 531 return getattr(self._backend, "_fleet_sub_models", None) 532 533 @_fleet_sub_models.setter 534 def _fleet_sub_models(self, value): 535 self._backend._fleet_sub_models = value 536 537 @property 538 def _fleet_n_risks(self): 539 return getattr(self._backend, "_fleet_n_risks", None) 540 541 @_fleet_n_risks.setter 542 def _fleet_n_risks(self, value): 543 self._backend._fleet_n_risks = value 544 545 @property 546 def _cr_kmeans(self): 547 return getattr(self._backend, "_cr_kmeans", None) 548 549 @_cr_kmeans.setter 550 def _cr_kmeans(self, value): 551 self._backend._cr_kmeans = value 552 553 @property 554 def _cr_resid_mean(self): 555 return getattr(self._backend, "_cr_resid_mean", None) 556 557 @_cr_resid_mean.setter 558 def _cr_resid_mean(self, value): 559 self._backend._cr_resid_mean = value 560 561 @property 562 def _cr_resid_std(self): 563 return getattr(self._backend, "_cr_resid_std", None) 564 565 @_cr_resid_std.setter 566 def _cr_resid_std(self, value): 567 self._backend._cr_resid_std = value 568 569 @property 570 def _cr_feature_scaler(self): 571 return getattr(self._backend, "_cr_feature_scaler", None) 572 573 @_cr_feature_scaler.setter 574 def _cr_feature_scaler(self, value): 575 self._backend._cr_feature_scaler = value 576 577 @property 578 def _sagemaker_session(self): 579 return getattr(self._backend, "_sagemaker_session", None) 580 581 # ------------------------------------------------------------------ 582 # Static helpers 583 # ------------------------------------------------------------------ 584 585 @staticmethod 586 def _get_sdk_version(): 587 """This package's version, or ``None`` when it cannot be determined. 588 589 Read from the installed distribution's metadata, which is the one place 590 that has the answer at the time it is asked: the version is set at build 591 time from the release tag, so no file in the source tree carries it. 592 593 ``None`` rather than a placeholder. The container skips its compatibility 594 check when no version reaches it, and a parseable stand-in misses that path 595 and fails the comparison instead — telling the caller their version does 596 not match, on the strength of a number nobody chose. 597 """ 598 import importlib.metadata 599 600 try: 601 return importlib.metadata.version("lifespan-predictor") 602 except Exception: 603 return None 604 605 # ------------------------------------------------------------------ 606 # Core training 607 # ------------------------------------------------------------------ 608 609 def fit( 610 self, 611 train_df, 612 val_df=None, 613 val_split=None, 614 test_df=None, 615 test_labels_df=None, 616 ): 617 """ 618 Train the model, tuning its hyperparameters first if the config asks. 619 620 **Training once.** When every config value is a plain value, `fit()` 621 trains one model on `train_df` and returns its training log: 622 623 ```python 624 model = SurvivalModel({"model_type": "deepsurv", "epochs": 30}) 625 log = model.fit(train_df, val_split=0.2) 626 print(log.best_epochs) 627 ``` 628 629 To measure the model, give `fit()` a test set it does not train on, as 630 `test_df` and `test_labels_df`, and read `get_evaluation_result()`. 631 Validation data is no measure: it decides when training stops. 632 633 **Tuning.** When any config value is a `ContinuousRange`, `IntegerRange` 634 or `CategoricalRange`, `fit()` first searches for the best values: 635 636 1. It splits `train_df` into 5 folds with `stratified_kfold`. The number 637 of folds is fixed at 5. 638 2. It tries `n_trials` configurations, each drawn from the ranges and 639 scored by 5-fold cross-validation: trained on four folds, 640 early-stopped on the fifth, scored on that fifth fold by 641 `hpo_metric`, and averaged over the five. 642 3. It retrains on all the training data with the best configuration, 643 and returns that model's training log. 644 645 The trials and the best configuration are then available from 646 `get_tuning_trials()`. `val_df` and `val_split` do not change how trials 647 are scored: each trial is scored on its own folds, and the validation 648 data serves the final training only. With `val_split`, the held-out units 649 are set aside first, so the folds are made from the rest. A config 650 without ranges trains once, with no cross-validation at all. 651 652 Each fold that scores a trial also decided when that trial stopped 653 training, so trial scores lean optimistic. They rank configurations 654 against each other; they do not measure the final model, which a test 655 set does. 656 657 ```python 658 model = SurvivalModel({ 659 "model_type": "deepsurv", 660 "learning_rate": ContinuousRange(1e-4, 1e-2, log=True), 661 "num_nodes": CategoricalRange([[16], [32, 32]]), 662 "n_trials": 20, 663 }) 664 model.fit(train_df) 665 print(model.get_tuning_trials().best_config) 666 ``` 667 668 To search in a way of your own, such as a grid search, see 669 `stratified_kfold`. 670 671 **Data.** Each DataFrame has one row per unit and time step. `train_df` 672 and `val_df` hold `unit` (the unit's ID), `duration` (the time from that 673 row to the unit's failure or censoring), `event` (1 if the unit failed, 0 674 if it was censored) and the feature columns, and may hold `cycles` (the 675 time step). Every column other than `unit`, `cycles`, `duration`, 676 `event` and `group` is a feature. 677 678 Args: 679 train_df: The training data. 680 val_df: Validation data, with the same columns as `train_df`. It 681 drives early stopping and the learning-rate reductions, so it is 682 not a fair set to score the model on. Give either `val_df` or 683 `val_split`, not both. 684 val_split: The fraction of `train_df`'s units (0 to 1) to hold out 685 as validation data instead, split by unit and balanced on event 686 rate and duration. 687 test_df: Test samples: rows of feature columns with `unit`, without 688 labels. Needs `test_labels_df`. 689 test_labels_df: One row per test unit, with `unit`, `duration` (the 690 remaining life at the unit's last row in `test_df`) and `event`. 691 With `test_df`, the model is evaluated after training, and 692 `get_evaluation_result()` returns the scores. 693 694 Returns: 695 The `TrainingLog` of the model that was trained: of the single 696 training, or of the final training after tuning. 697 698 Raises: 699 ValueError: If both `val_df` and `val_split` are given, or if 700 `calibration` is `"isotonic"` and there is no validation data. 701 """ 702 if val_df is not None and val_split is not None: 703 raise ValueError( 704 "val_df and val_split are mutually exclusive. " 705 "Provide one or the other, not both." 706 ) 707 708 if val_df is None and val_split is not None: 709 from lifespan_predictor_sdk.hpo_io import _stratified_val_split 710 711 train_df, val_df = _stratified_val_split(train_df, val_fraction=val_split) 712 print( 713 f"val_split={val_split}: {train_df['unit'].nunique()} train units, " 714 f"{val_df['unit'].nunique()} val units " 715 f"({len(train_df)} train rows, {len(val_df)} val rows)" 716 ) 717 718 if self.config.get("calibration") == "isotonic" and val_df is None: 719 raise ValueError( 720 "calibration='isotonic' requires validation data. " 721 "Provide val_df or val_split to fit()." 722 ) 723 724 training_log = self._backend.fit( 725 self, 726 train_df, 727 val_df=val_df, 728 test_df=test_df, 729 test_labels_df=test_labels_df, 730 ) 731 732 self._training_log = training_log 733 734 # Whatever a backend does once training has finished is the backend's own 735 # business; the facade tells it that fitting is done and asks nothing about 736 # what it is. 737 self._backend._after_fit( 738 self, val_df=val_df, test_df=test_df, test_labels_df=test_labels_df 739 ) 740 741 return training_log 742 743 # ------------------------------------------------------------------ 744 # Prediction 745 # ------------------------------------------------------------------ 746 747 def predict_mean(self, X, **kwargs): 748 """ 749 Predict each row's mean time to failure. 750 751 Runs a batch transform job, or uses the real-time endpoint if `deploy()` 752 created one. 753 754 Args: 755 X: Feature rows, a DataFrame with `unit` and the training data's 756 feature columns, and no labels. 757 **kwargs: Batch transform options: `model_name` (default: the model 758 `deploy()` created), `instance_type` (default `"ml.m5.large"`), 759 `instance_count` (default 1), `wait` (default True) and 760 `max_payload` in MB (default 6). 761 762 Returns: 763 A numpy array of mean times to failure, one per row of `X`. 764 765 Raises: 766 ValueError: If the model has not been fitted. 767 """ 768 if not self._is_fitted: 769 raise ValueError( 770 "Model must be fitted before prediction. Call fit() first." 771 ) 772 return self._backend.predict_mean(self, X, **kwargs) 773 774 def predict(self, X, time_points=None): 775 """ 776 Predict each row's survival curve. 777 778 Runs the same job as `predict_mean()`. 779 780 Args: 781 X: Feature rows, as for `predict_mean()`. 782 time_points: Times to return the curves at. They select times from 783 the model's own time grid; values are not interpolated. Without 784 them, the whole grid is returned. 785 786 Returns: 787 A DataFrame of survival probabilities, indexed by time under the name 788 `duration`, with one column per row of `X`. 789 """ 790 curves = self._backend.predict(self, X, time_points=time_points) 791 792 # Named here rather than in either backend because this is the one 793 # function all of them pass through. Three paths build that frame — 794 # pycox for DeepSurv, the Weibull helper for WTTE-RNN, and the batch 795 # transform for SageMaker — and only the first named the index, having 796 # inherited the name from pycox rather than chosen it. The other two 797 # agreed with each other and disagreed with it, which is a difference a 798 # caller sees when comparing or concatenating results. 799 # 800 # Set only when absent, so a backend with a reason to name its own 801 # index keeps it. 802 if curves is not None and curves.index.name is None: 803 curves.index.name = SURVIVAL_TIME_INDEX_NAME 804 return curves 805 806 # ------------------------------------------------------------------ 807 # Evaluation (shared logic) 808 # ------------------------------------------------------------------ 809 810 def evaluate(self, X, durations, events): 811 """ 812 Score the model on labelled data. 813 814 Typically `X` is a test set's last row per unit, and `durations` and 815 `events` its labels. With `max_rul` set, the C-index counts only units 816 whose duration is at most `max_rul`. 817 818 Args: 819 X: Feature rows, as for `predict_mean()`. 820 durations: For each row of `X`, the time to failure or censoring. 821 events: For each row of `X`, 1 if the unit failed and 0 if it was 822 censored. 823 824 Returns: 825 An `EvaluationResult` with the C-index, RMSE, calibration and 826 integrated Brier score. 827 828 Raises: 829 ValueError: If the model has not been fitted. 830 """ 831 if not self._is_fitted: 832 raise ValueError( 833 "Model must be fitted before evaluation. Call fit() first." 834 ) 835 836 predictions = self.predict_mean(X) 837 838 max_rul = self.config.get("max_rul", None) 839 print(f"Evaluating metrics with max_rul: {max_rul}") 840 if max_rul is not None: 841 mask_cindex = durations <= max_rul 842 else: 843 mask_cindex = np.ones(len(durations), dtype=bool) 844 845 # C-index 846 c_index = c_index_score( 847 durations[mask_cindex], events[mask_cindex], predictions[mask_cindex] 848 ) 849 # The censoring-corrected companion. `predictions` are survival 850 # times, so they are negated into a risk for the IPCW estimator's 851 # higher-means-worse convention. 852 corrected_c_index = ipcw_c_index( 853 durations[mask_cindex], 854 events[mask_cindex], 855 -np.asarray(predictions)[mask_cindex], 856 ) 857 858 # RMSE (uncensored only) 859 mask_rmse = (np.asarray(events) > 0) & mask_cindex 860 rmse = rmse_score(durations[mask_rmse], predictions[mask_rmse]) 861 862 # Calibration metrics from survival curves 863 cal_bias = None 864 cal_mean_cdf = None 865 cal_classification = None 866 cal_sorted_cdf = None 867 cal_empirical_cdf = None 868 cal_sorted_cdf_events = None 869 cal_empirical_cdf_events = None 870 ibs_value = None 871 ibs_class = None 872 antolini_value = None 873 874 try: 875 import warnings as _warnings 876 877 from lifespan_predictor_sdk.calibration import ( 878 calibration_bias, 879 calibration_curve_data_from_survival, 880 integrated_brier_score, 881 interpret_calibration_bias, 882 interpret_integrated_brier_score, 883 ) 884 885 # Get survival curves from backend 886 survival_curves, time_grid, surv_df = self._backend.get_survival_curves( 887 self, X 888 ) 889 890 if survival_curves is not None: 891 sorted_cdf, empirical_cdf = calibration_curve_data_from_survival( 892 durations, survival_curves, time_grid 893 ) 894 event_mask = events.astype(bool) 895 if np.sum(event_mask) >= 2: 896 sorted_cdf_events, empirical_cdf_events = ( 897 calibration_curve_data_from_survival( 898 durations[event_mask], 899 survival_curves[event_mask], 900 time_grid, 901 ) 902 ) 903 else: 904 sorted_cdf_events, empirical_cdf_events = None, None 905 bias = calibration_bias(sorted_cdf, empirical_cdf) 906 mean_cdf = float(np.mean(sorted_cdf)) 907 bias_interp = interpret_calibration_bias(mean_cdf) 908 909 cal_bias = float(bias) 910 cal_mean_cdf = mean_cdf 911 cal_classification = bias_interp["classification"] 912 cal_sorted_cdf = sorted_cdf.tolist() 913 cal_empirical_cdf = empirical_cdf.tolist() 914 cal_sorted_cdf_events = ( 915 sorted_cdf_events.tolist() 916 if sorted_cdf_events is not None 917 else None 918 ) 919 cal_empirical_cdf_events = ( 920 empirical_cdf_events.tolist() 921 if empirical_cdf_events is not None 922 else None 923 ) 924 925 # IBS 926 try: 927 pos_mask = durations > 0 928 if np.sum(pos_mask) > 0 and np.sum(events[pos_mask]) > 0: 929 dur_pos = durations[pos_mask] 930 ev_pos = events[pos_mask] 931 sc_pos = survival_curves[pos_mask] 932 t_min, t_max = float(np.min(dur_pos)), float(np.max(dur_pos)) 933 tp_mask = (time_grid >= t_min) & (time_grid <= t_max) 934 if np.sum(tp_mask) >= 2: 935 with _warnings.catch_warnings(record=True) as caught: 936 _warnings.simplefilter("always") 937 ibs = integrated_brier_score( 938 dur_pos, 939 ev_pos, 940 sc_pos[:, tp_mask], 941 time_grid[tp_mask], 942 ) 943 for w in caught: 944 print(f" \u26a0\ufe0f {w.message}") 945 if not np.isnan(ibs): 946 ibs_interp = interpret_integrated_brier_score( 947 float(ibs) 948 ) 949 ibs_value = float(ibs) 950 ibs_class = ibs_interp["classification"] 951 except Exception as _e: 952 print(f" \u26a0\ufe0f Could not compute IBS: {_e}") 953 954 # Antolini time-dependent C-index (requires pycox \u2014 optional) 955 if surv_df is not None: 956 try: 957 from pycox.evaluation import EvalSurv 958 959 ev = EvalSurv(surv_df, durations, events, censor_surv="km") 960 score = ev.concordance_td("antolini") 961 if not np.isnan(score): 962 antolini_value = float(score) 963 except Exception as _e: 964 print( 965 f" \u26a0\ufe0f Could not compute Antolini C-index: {_e}" 966 ) 967 except Exception as e: 968 print(f" \u26a0\ufe0f Could not compute calibration metrics: {e}") 969 970 # Competing-risk metrics (backend-specific) 971 ( 972 cr_per_risk_c, 973 cr_per_risk_antolini, 974 cr_per_risk_predictive, 975 ) = self._backend.get_competing_risk_metrics( 976 self, X, durations, events, mask_cindex=mask_cindex 977 ) 978 979 return EvaluationResult( 980 c_index=c_index, 981 ipcw_c_index=corrected_c_index, 982 rmse=rmse, 983 predictions=predictions, 984 mask_cindex=mask_cindex, 985 mask_rmse=mask_rmse, 986 y_test=durations, 987 events=np.asarray(events), 988 calibration_bias=cal_bias, 989 calibration_mean_cdf=cal_mean_cdf, 990 calibration_classification=cal_classification, 991 calibration_sorted_cdf=cal_sorted_cdf, 992 calibration_empirical_cdf=cal_empirical_cdf, 993 calibration_sorted_cdf_events=cal_sorted_cdf_events, 994 calibration_empirical_cdf_events=cal_empirical_cdf_events, 995 ibs=ibs_value, 996 ibs_classification=ibs_class, 997 antolini_c_index=antolini_value, 998 competing_risk_per_risk_cause_specific_c_index=cr_per_risk_c, 999 competing_risk_per_risk_antolini=cr_per_risk_antolini, 1000 competing_risk_per_risk_predictive_c_index=cr_per_risk_predictive, 1001 ) 1002 1003 # ------------------------------------------------------------------ 1004 # Calibration 1005 # ------------------------------------------------------------------ 1006 1007 1008 # ------------------------------------------------------------------ 1009 # Delegated post-fit analysis methods 1010 # ------------------------------------------------------------------ 1011 1012 1013 1014 1015 def get_mean_time_to_failure(self, X): 1016 """ 1017 The same as `predict_mean()`. 1018 1019 Args: 1020 X: Feature rows, as for `predict_mean()`. 1021 1022 Returns: 1023 A numpy array of mean times to failure, one per row of `X`. 1024 """ 1025 return self.predict_mean(X) 1026 1027 def _predict_all_rows(self, df): 1028 """Return survival curves for every row in df (used by tuning).""" 1029 return self._backend._predict_all_rows(self, df) 1030 1031 def _compute_val_antolini(self, val_df): 1032 """Compute Antolini's time-dependent C-index on val_df.""" 1033 return self._backend._compute_val_antolini(self, val_df) 1034 1035 def _compute_val_ibs(self, val_df, train_df): 1036 """Compute IBS on val_df.""" 1037 return self._backend._compute_val_ibs(self, val_df, train_df) 1038 1039 def _fit_calibrator_on_df(self, df): 1040 """Fit isotonic calibrator on df.""" 1041 return self._backend._fit_calibrator_on_df(self, df) 1042 1043 def _apply_calibrator(self, surv_df): 1044 """Apply isotonic calibrator to survival curves.""" 1045 return self._backend._apply_calibrator(surv_df) 1046 1047 1048 1049 # ------------------------------------------------------------------ 1050 # Delegated SageMaker-only methods 1051 # ------------------------------------------------------------------ 1052 1053 def deploy(self, model_name=None, endpoint_instance_type=None): 1054 """ 1055 Make the trained model available for prediction on SageMaker. 1056 1057 By default this creates a SageMaker model that `predict_mean()` and 1058 `predict()` run batch transform jobs against. With 1059 `endpoint_instance_type`, it also hosts the model on a real-time endpoint, 1060 which predictions then use until `delete()`. An endpoint is billed for as 1061 long as it exists. 1062 1063 Args: 1064 model_name: The name of the SageMaker model. Without it, a name is 1065 derived from the algorithm's. 1066 endpoint_instance_type: The instance type of a real-time endpoint, 1067 such as `"ml.m5.large"`. Without it, no endpoint is created. 1068 1069 Returns: 1070 The name of the SageMaker model. 1071 """ 1072 return self._backend.deploy( 1073 self, model_name=model_name, endpoint_instance_type=endpoint_instance_type 1074 ) 1075 1076 def delete(self): 1077 """Delete what `deploy()` created: the SageMaker model, and its endpoint.""" 1078 self._backend.delete(self) 1079 1080 def print_training_job_summary(self): 1081 """Print a summary of the model's SageMaker training job.""" 1082 self._backend.print_training_job_summary(self) 1083 1084 # ------------------------------------------------------------------ 1085 # Shared accessors 1086 # ------------------------------------------------------------------ 1087 1088 def get_training_log(self): 1089 """ 1090 The training log of the last `fit()`. 1091 1092 Returns: 1093 The same `TrainingLog` `fit()` returned. 1094 1095 Raises: 1096 ValueError: If the model has not been fitted. 1097 """ 1098 if not self._is_fitted: 1099 raise ValueError( 1100 "Model must be fitted before accessing training log. Call fit() first." 1101 ) 1102 if self._training_log is None: 1103 raise ValueError( 1104 "No training log available. This should not happen if model was fitted properly." 1105 ) 1106 return self._training_log 1107 1108 def get_tuning_trials(self): 1109 """ 1110 The results of the tuning `fit()` ran. 1111 1112 Returns: 1113 A `TuningTrials`, or None if the config had no ranges and `fit()` 1114 did not tune. 1115 """ 1116 return getattr(self, "_tuning_trials", None) 1117 1118 def get_config(self): 1119 """ 1120 The model's config, with defaults filled in. 1121 1122 Returns: 1123 The config dict. 1124 """ 1125 return self.config 1126 1127 def print_config(self): 1128 """Print the model configuration.""" 1129 print("Model Configuration:") 1130 for key, value in self.config.items(): 1131 print(f" {key}: {value}") 1132 1133 def get_evaluation_result(self, pre_calibration=False): 1134 """ 1135 Get the evaluation result from training. 1136 1137 Args: 1138 pre_calibration: If True, return the pre-calibration result. 1139 1140 Returns: 1141 EvaluationResult if test data was provided during fit(), otherwise None. 1142 """ 1143 if pre_calibration: 1144 if self._evaluation_results is None: 1145 raise ValueError( 1146 "No pre-calibration result available: model has not been calibrated." 1147 ) 1148 return self._evaluation_results["before"] 1149 if self._evaluation_results is not None: 1150 return self._evaluation_results.get("after") 1151 return None
A DeepSurv survival model: train it, then predict time to failure.
A model is built from a config dict and trained with fit(). Preprocessing
(scaling and the optional engineered features below) is part of the model, so
predict_mean() and predict() take raw feature rows in the same shape as
the training data.
model = SurvivalModel({
"model_type": "deepsurv",
"algorithm_arn": os.environ["SAGEMAKER_ALGORITHM_ARN"],
"epochs": 30,
"num_nodes": [24],
"max_rul": 125,
})
model.fit(train_df, val_split=0.2)
model.deploy()
mean_times = model.predict_mean(test_df)
Hyperparameters
Every key is optional. A key marked Range may also be given as a
ContinuousRange, IntegerRange or CategoricalRange, which makes fit()
tune it (see fit()).
| Key | Type | Default | Range | Meaning |
|---|---|---|---|---|
model_type |
str | "deepsurv" |
no | The model. "deepsurv" is the one this release offers. |
epochs |
int | 100 |
yes | Maximum training epochs. With validation data, training stops early when the validation loss stops improving. |
num_nodes |
list of int | [32, 32] |
yes | Hidden layer sizes, one entry per layer. |
dropout |
float | 0.1 |
yes | Dropout rate after each hidden layer. |
learning_rate |
float | 0.01 |
yes | Initial learning rate. It is reduced when the validation loss plateaus. |
batch_size |
int | 256 |
yes | Rows per training batch. |
max_rul |
float or None | None |
yes | Caps the training durations at this value, and limits the C-index to units whose duration is at most this. Common for turbofan data, where early life carries no wear signal. |
use_delta |
bool | False |
yes | Adds, for each feature, its change from the unit's previous row. |
use_per_unit_normalisation |
bool | False |
yes | Adds, for each feature, its offset from the unit's first row. |
rolling_means |
list of int or None | None |
yes | Window sizes; adds, for each feature and window, the unit's rolling mean. |
calibration |
str or None | None |
yes | "isotonic" calibrates the predicted survival curves on the validation data, so fit() then needs val_df or val_split. |
n_trials |
int | 2 |
no | Tuning only: how many configurations to try. |
hpo_metric |
str | "c-index" |
no | Tuning only: the score trials are compared by: "c-index", "ibs" (integrated Brier score) or "antolini" (time-dependent C-index). |
verbose |
bool | False |
no | Tuning only: log each trial in detail. |
SageMaker settings
Where training runs. They are not hyperparameters and cannot be tuned.
| Key | Type | If left out | Meaning |
|---|---|---|---|
algorithm_arn |
str | SAGEMAKER_ALGORITHM_ARN is used; without it, fit() raises |
The ARN of the algorithm you subscribed to. |
instance_type |
str | "ml.m5.large" |
The instance type of the training job. |
bucket |
str | SAGEMAKER_BUCKET is used; without it, SageMaker's default bucket for your account and region |
The S3 bucket for training data and model artifacts. |
The job's IAM role is read from the SAGEMAKER_ROLE_ARN environment variable,
which must be set.
SurvivalModel(config)
429 def __init__(self, config): 430 """ 431 Create an untrained model. 432 433 Args: 434 config: The model's configuration: hyperparameters and SageMaker 435 settings, as listed in the tables above. Missing keys take their 436 defaults. 437 """ 438 # Minimal shared defaults 439 defaults = { 440 "model_type": "deepsurv", 441 "epochs": 100, 442 "max_rul": None, 443 } 444 445 self.config = {**defaults, **config} 446 447 # Shared state 448 self._is_fitted = False 449 self._training_log = None 450 self._tuning_trials = None 451 self._evaluation_results = None 452 self.in_features = self.config.get("in_features") 453 454 # Stashed DataFrames from fit() for post-fit conversion methods 455 self._fit_train_df = None 456 self._fit_val_df = None 457 self._fit_test_df = None 458 self._fit_test_labels_df = None 459 460 # Create backend. The backend reports its own name back into the config, 461 # so `get_config()` still answers the question without this file having to 462 # spell out which backends exist — that is _make_backend's business. 463 self._backend = _make_backend(self.config.get("backend")) 464 self.config["backend"] = self._backend.name 465 466 # Let backend merge its defaults and initialize state 467 self._backend.init_state(self.config)
Create an untrained model.
Arguments:
- config: The model's configuration: hyperparameters and SageMaker settings, as listed in the tables above. Missing keys take their defaults.
def
fit( self, train_df, val_df=None, val_split=None, test_df=None, test_labels_df=None):
609 def fit( 610 self, 611 train_df, 612 val_df=None, 613 val_split=None, 614 test_df=None, 615 test_labels_df=None, 616 ): 617 """ 618 Train the model, tuning its hyperparameters first if the config asks. 619 620 **Training once.** When every config value is a plain value, `fit()` 621 trains one model on `train_df` and returns its training log: 622 623 ```python 624 model = SurvivalModel({"model_type": "deepsurv", "epochs": 30}) 625 log = model.fit(train_df, val_split=0.2) 626 print(log.best_epochs) 627 ``` 628 629 To measure the model, give `fit()` a test set it does not train on, as 630 `test_df` and `test_labels_df`, and read `get_evaluation_result()`. 631 Validation data is no measure: it decides when training stops. 632 633 **Tuning.** When any config value is a `ContinuousRange`, `IntegerRange` 634 or `CategoricalRange`, `fit()` first searches for the best values: 635 636 1. It splits `train_df` into 5 folds with `stratified_kfold`. The number 637 of folds is fixed at 5. 638 2. It tries `n_trials` configurations, each drawn from the ranges and 639 scored by 5-fold cross-validation: trained on four folds, 640 early-stopped on the fifth, scored on that fifth fold by 641 `hpo_metric`, and averaged over the five. 642 3. It retrains on all the training data with the best configuration, 643 and returns that model's training log. 644 645 The trials and the best configuration are then available from 646 `get_tuning_trials()`. `val_df` and `val_split` do not change how trials 647 are scored: each trial is scored on its own folds, and the validation 648 data serves the final training only. With `val_split`, the held-out units 649 are set aside first, so the folds are made from the rest. A config 650 without ranges trains once, with no cross-validation at all. 651 652 Each fold that scores a trial also decided when that trial stopped 653 training, so trial scores lean optimistic. They rank configurations 654 against each other; they do not measure the final model, which a test 655 set does. 656 657 ```python 658 model = SurvivalModel({ 659 "model_type": "deepsurv", 660 "learning_rate": ContinuousRange(1e-4, 1e-2, log=True), 661 "num_nodes": CategoricalRange([[16], [32, 32]]), 662 "n_trials": 20, 663 }) 664 model.fit(train_df) 665 print(model.get_tuning_trials().best_config) 666 ``` 667 668 To search in a way of your own, such as a grid search, see 669 `stratified_kfold`. 670 671 **Data.** Each DataFrame has one row per unit and time step. `train_df` 672 and `val_df` hold `unit` (the unit's ID), `duration` (the time from that 673 row to the unit's failure or censoring), `event` (1 if the unit failed, 0 674 if it was censored) and the feature columns, and may hold `cycles` (the 675 time step). Every column other than `unit`, `cycles`, `duration`, 676 `event` and `group` is a feature. 677 678 Args: 679 train_df: The training data. 680 val_df: Validation data, with the same columns as `train_df`. It 681 drives early stopping and the learning-rate reductions, so it is 682 not a fair set to score the model on. Give either `val_df` or 683 `val_split`, not both. 684 val_split: The fraction of `train_df`'s units (0 to 1) to hold out 685 as validation data instead, split by unit and balanced on event 686 rate and duration. 687 test_df: Test samples: rows of feature columns with `unit`, without 688 labels. Needs `test_labels_df`. 689 test_labels_df: One row per test unit, with `unit`, `duration` (the 690 remaining life at the unit's last row in `test_df`) and `event`. 691 With `test_df`, the model is evaluated after training, and 692 `get_evaluation_result()` returns the scores. 693 694 Returns: 695 The `TrainingLog` of the model that was trained: of the single 696 training, or of the final training after tuning. 697 698 Raises: 699 ValueError: If both `val_df` and `val_split` are given, or if 700 `calibration` is `"isotonic"` and there is no validation data. 701 """ 702 if val_df is not None and val_split is not None: 703 raise ValueError( 704 "val_df and val_split are mutually exclusive. " 705 "Provide one or the other, not both." 706 ) 707 708 if val_df is None and val_split is not None: 709 from lifespan_predictor_sdk.hpo_io import _stratified_val_split 710 711 train_df, val_df = _stratified_val_split(train_df, val_fraction=val_split) 712 print( 713 f"val_split={val_split}: {train_df['unit'].nunique()} train units, " 714 f"{val_df['unit'].nunique()} val units " 715 f"({len(train_df)} train rows, {len(val_df)} val rows)" 716 ) 717 718 if self.config.get("calibration") == "isotonic" and val_df is None: 719 raise ValueError( 720 "calibration='isotonic' requires validation data. " 721 "Provide val_df or val_split to fit()." 722 ) 723 724 training_log = self._backend.fit( 725 self, 726 train_df, 727 val_df=val_df, 728 test_df=test_df, 729 test_labels_df=test_labels_df, 730 ) 731 732 self._training_log = training_log 733 734 # Whatever a backend does once training has finished is the backend's own 735 # business; the facade tells it that fitting is done and asks nothing about 736 # what it is. 737 self._backend._after_fit( 738 self, val_df=val_df, test_df=test_df, test_labels_df=test_labels_df 739 ) 740 741 return training_log
Train the model, tuning its hyperparameters first if the config asks.
Training once. When every config value is a plain value, fit()
trains one model on train_df and returns its training log:
model = SurvivalModel({"model_type": "deepsurv", "epochs": 30})
log = model.fit(train_df, val_split=0.2)
print(log.best_epochs)
To measure the model, give fit() a test set it does not train on, as
test_df and test_labels_df, and read get_evaluation_result().
Validation data is no measure: it decides when training stops.
Tuning. When any config value is a ContinuousRange, IntegerRange
or CategoricalRange, fit() first searches for the best values:
- It splits
train_dfinto 5 folds withstratified_kfold. The number of folds is fixed at 5. - It tries
n_trialsconfigurations, each drawn from the ranges and scored by 5-fold cross-validation: trained on four folds, early-stopped on the fifth, scored on that fifth fold byhpo_metric, and averaged over the five. - It retrains on all the training data with the best configuration, and returns that model's training log.
The trials and the best configuration are then available from
get_tuning_trials(). val_df and val_split do not change how trials
are scored: each trial is scored on its own folds, and the validation
data serves the final training only. With val_split, the held-out units
are set aside first, so the folds are made from the rest. A config
without ranges trains once, with no cross-validation at all.
Each fold that scores a trial also decided when that trial stopped training, so trial scores lean optimistic. They rank configurations against each other; they do not measure the final model, which a test set does.
model = SurvivalModel({
"model_type": "deepsurv",
"learning_rate": ContinuousRange(1e-4, 1e-2, log=True),
"num_nodes": CategoricalRange([[16], [32, 32]]),
"n_trials": 20,
})
model.fit(train_df)
print(model.get_tuning_trials().best_config)
To search in a way of your own, such as a grid search, see
stratified_kfold.
Data. Each DataFrame has one row per unit and time step. train_df
and val_df hold unit (the unit's ID), duration (the time from that
row to the unit's failure or censoring), event (1 if the unit failed, 0
if it was censored) and the feature columns, and may hold cycles (the
time step). Every column other than unit, cycles, duration,
event and group is a feature.
Arguments:
- train_df: The training data.
- val_df: Validation data, with the same columns as
train_df. It drives early stopping and the learning-rate reductions, so it is not a fair set to score the model on. Give eitherval_dforval_split, not both. - val_split: The fraction of
train_df's units (0 to 1) to hold out as validation data instead, split by unit and balanced on event rate and duration. - test_df: Test samples: rows of feature columns with
unit, without labels. Needstest_labels_df. - test_labels_df: One row per test unit, with
unit,duration(the remaining life at the unit's last row intest_df) andevent. Withtest_df, the model is evaluated after training, andget_evaluation_result()returns the scores.
Returns:
The
TrainingLogof the model that was trained: of the single training, or of the final training after tuning.
Raises:
- ValueError: If both
val_dfandval_splitare given, or ifcalibrationis"isotonic"and there is no validation data.
def
predict_mean(self, X, **kwargs):
747 def predict_mean(self, X, **kwargs): 748 """ 749 Predict each row's mean time to failure. 750 751 Runs a batch transform job, or uses the real-time endpoint if `deploy()` 752 created one. 753 754 Args: 755 X: Feature rows, a DataFrame with `unit` and the training data's 756 feature columns, and no labels. 757 **kwargs: Batch transform options: `model_name` (default: the model 758 `deploy()` created), `instance_type` (default `"ml.m5.large"`), 759 `instance_count` (default 1), `wait` (default True) and 760 `max_payload` in MB (default 6). 761 762 Returns: 763 A numpy array of mean times to failure, one per row of `X`. 764 765 Raises: 766 ValueError: If the model has not been fitted. 767 """ 768 if not self._is_fitted: 769 raise ValueError( 770 "Model must be fitted before prediction. Call fit() first." 771 ) 772 return self._backend.predict_mean(self, X, **kwargs)
Predict each row's mean time to failure.
Runs a batch transform job, or uses the real-time endpoint if deploy()
created one.
Arguments:
- X: Feature rows, a DataFrame with
unitand the training data's feature columns, and no labels. - **kwargs: Batch transform options:
model_name(default: the modeldeploy()created),instance_type(default"ml.m5.large"),instance_count(default 1),wait(default True) andmax_payloadin MB (default 6).
Returns:
A numpy array of mean times to failure, one per row of
X.
Raises:
- ValueError: If the model has not been fitted.
def
predict(self, X, time_points=None):
774 def predict(self, X, time_points=None): 775 """ 776 Predict each row's survival curve. 777 778 Runs the same job as `predict_mean()`. 779 780 Args: 781 X: Feature rows, as for `predict_mean()`. 782 time_points: Times to return the curves at. They select times from 783 the model's own time grid; values are not interpolated. Without 784 them, the whole grid is returned. 785 786 Returns: 787 A DataFrame of survival probabilities, indexed by time under the name 788 `duration`, with one column per row of `X`. 789 """ 790 curves = self._backend.predict(self, X, time_points=time_points) 791 792 # Named here rather than in either backend because this is the one 793 # function all of them pass through. Three paths build that frame — 794 # pycox for DeepSurv, the Weibull helper for WTTE-RNN, and the batch 795 # transform for SageMaker — and only the first named the index, having 796 # inherited the name from pycox rather than chosen it. The other two 797 # agreed with each other and disagreed with it, which is a difference a 798 # caller sees when comparing or concatenating results. 799 # 800 # Set only when absent, so a backend with a reason to name its own 801 # index keeps it. 802 if curves is not None and curves.index.name is None: 803 curves.index.name = SURVIVAL_TIME_INDEX_NAME 804 return curves
Predict each row's survival curve.
Runs the same job as predict_mean().
Arguments:
- X: Feature rows, as for
predict_mean(). - time_points: Times to return the curves at. They select times from the model's own time grid; values are not interpolated. Without them, the whole grid is returned.
Returns:
A DataFrame of survival probabilities, indexed by time under the name
duration, with one column per row ofX.
def
evaluate(self, X, durations, events):
810 def evaluate(self, X, durations, events): 811 """ 812 Score the model on labelled data. 813 814 Typically `X` is a test set's last row per unit, and `durations` and 815 `events` its labels. With `max_rul` set, the C-index counts only units 816 whose duration is at most `max_rul`. 817 818 Args: 819 X: Feature rows, as for `predict_mean()`. 820 durations: For each row of `X`, the time to failure or censoring. 821 events: For each row of `X`, 1 if the unit failed and 0 if it was 822 censored. 823 824 Returns: 825 An `EvaluationResult` with the C-index, RMSE, calibration and 826 integrated Brier score. 827 828 Raises: 829 ValueError: If the model has not been fitted. 830 """ 831 if not self._is_fitted: 832 raise ValueError( 833 "Model must be fitted before evaluation. Call fit() first." 834 ) 835 836 predictions = self.predict_mean(X) 837 838 max_rul = self.config.get("max_rul", None) 839 print(f"Evaluating metrics with max_rul: {max_rul}") 840 if max_rul is not None: 841 mask_cindex = durations <= max_rul 842 else: 843 mask_cindex = np.ones(len(durations), dtype=bool) 844 845 # C-index 846 c_index = c_index_score( 847 durations[mask_cindex], events[mask_cindex], predictions[mask_cindex] 848 ) 849 # The censoring-corrected companion. `predictions` are survival 850 # times, so they are negated into a risk for the IPCW estimator's 851 # higher-means-worse convention. 852 corrected_c_index = ipcw_c_index( 853 durations[mask_cindex], 854 events[mask_cindex], 855 -np.asarray(predictions)[mask_cindex], 856 ) 857 858 # RMSE (uncensored only) 859 mask_rmse = (np.asarray(events) > 0) & mask_cindex 860 rmse = rmse_score(durations[mask_rmse], predictions[mask_rmse]) 861 862 # Calibration metrics from survival curves 863 cal_bias = None 864 cal_mean_cdf = None 865 cal_classification = None 866 cal_sorted_cdf = None 867 cal_empirical_cdf = None 868 cal_sorted_cdf_events = None 869 cal_empirical_cdf_events = None 870 ibs_value = None 871 ibs_class = None 872 antolini_value = None 873 874 try: 875 import warnings as _warnings 876 877 from lifespan_predictor_sdk.calibration import ( 878 calibration_bias, 879 calibration_curve_data_from_survival, 880 integrated_brier_score, 881 interpret_calibration_bias, 882 interpret_integrated_brier_score, 883 ) 884 885 # Get survival curves from backend 886 survival_curves, time_grid, surv_df = self._backend.get_survival_curves( 887 self, X 888 ) 889 890 if survival_curves is not None: 891 sorted_cdf, empirical_cdf = calibration_curve_data_from_survival( 892 durations, survival_curves, time_grid 893 ) 894 event_mask = events.astype(bool) 895 if np.sum(event_mask) >= 2: 896 sorted_cdf_events, empirical_cdf_events = ( 897 calibration_curve_data_from_survival( 898 durations[event_mask], 899 survival_curves[event_mask], 900 time_grid, 901 ) 902 ) 903 else: 904 sorted_cdf_events, empirical_cdf_events = None, None 905 bias = calibration_bias(sorted_cdf, empirical_cdf) 906 mean_cdf = float(np.mean(sorted_cdf)) 907 bias_interp = interpret_calibration_bias(mean_cdf) 908 909 cal_bias = float(bias) 910 cal_mean_cdf = mean_cdf 911 cal_classification = bias_interp["classification"] 912 cal_sorted_cdf = sorted_cdf.tolist() 913 cal_empirical_cdf = empirical_cdf.tolist() 914 cal_sorted_cdf_events = ( 915 sorted_cdf_events.tolist() 916 if sorted_cdf_events is not None 917 else None 918 ) 919 cal_empirical_cdf_events = ( 920 empirical_cdf_events.tolist() 921 if empirical_cdf_events is not None 922 else None 923 ) 924 925 # IBS 926 try: 927 pos_mask = durations > 0 928 if np.sum(pos_mask) > 0 and np.sum(events[pos_mask]) > 0: 929 dur_pos = durations[pos_mask] 930 ev_pos = events[pos_mask] 931 sc_pos = survival_curves[pos_mask] 932 t_min, t_max = float(np.min(dur_pos)), float(np.max(dur_pos)) 933 tp_mask = (time_grid >= t_min) & (time_grid <= t_max) 934 if np.sum(tp_mask) >= 2: 935 with _warnings.catch_warnings(record=True) as caught: 936 _warnings.simplefilter("always") 937 ibs = integrated_brier_score( 938 dur_pos, 939 ev_pos, 940 sc_pos[:, tp_mask], 941 time_grid[tp_mask], 942 ) 943 for w in caught: 944 print(f" \u26a0\ufe0f {w.message}") 945 if not np.isnan(ibs): 946 ibs_interp = interpret_integrated_brier_score( 947 float(ibs) 948 ) 949 ibs_value = float(ibs) 950 ibs_class = ibs_interp["classification"] 951 except Exception as _e: 952 print(f" \u26a0\ufe0f Could not compute IBS: {_e}") 953 954 # Antolini time-dependent C-index (requires pycox \u2014 optional) 955 if surv_df is not None: 956 try: 957 from pycox.evaluation import EvalSurv 958 959 ev = EvalSurv(surv_df, durations, events, censor_surv="km") 960 score = ev.concordance_td("antolini") 961 if not np.isnan(score): 962 antolini_value = float(score) 963 except Exception as _e: 964 print( 965 f" \u26a0\ufe0f Could not compute Antolini C-index: {_e}" 966 ) 967 except Exception as e: 968 print(f" \u26a0\ufe0f Could not compute calibration metrics: {e}") 969 970 # Competing-risk metrics (backend-specific) 971 ( 972 cr_per_risk_c, 973 cr_per_risk_antolini, 974 cr_per_risk_predictive, 975 ) = self._backend.get_competing_risk_metrics( 976 self, X, durations, events, mask_cindex=mask_cindex 977 ) 978 979 return EvaluationResult( 980 c_index=c_index, 981 ipcw_c_index=corrected_c_index, 982 rmse=rmse, 983 predictions=predictions, 984 mask_cindex=mask_cindex, 985 mask_rmse=mask_rmse, 986 y_test=durations, 987 events=np.asarray(events), 988 calibration_bias=cal_bias, 989 calibration_mean_cdf=cal_mean_cdf, 990 calibration_classification=cal_classification, 991 calibration_sorted_cdf=cal_sorted_cdf, 992 calibration_empirical_cdf=cal_empirical_cdf, 993 calibration_sorted_cdf_events=cal_sorted_cdf_events, 994 calibration_empirical_cdf_events=cal_empirical_cdf_events, 995 ibs=ibs_value, 996 ibs_classification=ibs_class, 997 antolini_c_index=antolini_value, 998 competing_risk_per_risk_cause_specific_c_index=cr_per_risk_c, 999 competing_risk_per_risk_antolini=cr_per_risk_antolini, 1000 competing_risk_per_risk_predictive_c_index=cr_per_risk_predictive, 1001 )
Score the model on labelled data.
Typically X is a test set's last row per unit, and durations and
events its labels. With max_rul set, the C-index counts only units
whose duration is at most max_rul.
Arguments:
- X: Feature rows, as for
predict_mean(). - durations: For each row of
X, the time to failure or censoring. - events: For each row of
X, 1 if the unit failed and 0 if it was censored.
Returns:
An
EvaluationResultwith the C-index, RMSE, calibration and integrated Brier score.
Raises:
- ValueError: If the model has not been fitted.
def
get_mean_time_to_failure(self, X):
1015 def get_mean_time_to_failure(self, X): 1016 """ 1017 The same as `predict_mean()`. 1018 1019 Args: 1020 X: Feature rows, as for `predict_mean()`. 1021 1022 Returns: 1023 A numpy array of mean times to failure, one per row of `X`. 1024 """ 1025 return self.predict_mean(X)
The same as predict_mean().
Arguments:
- X: Feature rows, as for
predict_mean().
Returns:
A numpy array of mean times to failure, one per row of
X.
def
deploy(self, model_name=None, endpoint_instance_type=None):
1053 def deploy(self, model_name=None, endpoint_instance_type=None): 1054 """ 1055 Make the trained model available for prediction on SageMaker. 1056 1057 By default this creates a SageMaker model that `predict_mean()` and 1058 `predict()` run batch transform jobs against. With 1059 `endpoint_instance_type`, it also hosts the model on a real-time endpoint, 1060 which predictions then use until `delete()`. An endpoint is billed for as 1061 long as it exists. 1062 1063 Args: 1064 model_name: The name of the SageMaker model. Without it, a name is 1065 derived from the algorithm's. 1066 endpoint_instance_type: The instance type of a real-time endpoint, 1067 such as `"ml.m5.large"`. Without it, no endpoint is created. 1068 1069 Returns: 1070 The name of the SageMaker model. 1071 """ 1072 return self._backend.deploy( 1073 self, model_name=model_name, endpoint_instance_type=endpoint_instance_type 1074 )
Make the trained model available for prediction on SageMaker.
By default this creates a SageMaker model that predict_mean() and
predict() run batch transform jobs against. With
endpoint_instance_type, it also hosts the model on a real-time endpoint,
which predictions then use until delete(). An endpoint is billed for as
long as it exists.
Arguments:
- model_name: The name of the SageMaker model. Without it, a name is derived from the algorithm's.
- endpoint_instance_type: The instance type of a real-time endpoint,
such as
"ml.m5.large". Without it, no endpoint is created.
Returns:
The name of the SageMaker model.
def
delete(self):
1076 def delete(self): 1077 """Delete what `deploy()` created: the SageMaker model, and its endpoint.""" 1078 self._backend.delete(self)
Delete what deploy() created: the SageMaker model, and its endpoint.
def
print_training_job_summary(self):
1080 def print_training_job_summary(self): 1081 """Print a summary of the model's SageMaker training job.""" 1082 self._backend.print_training_job_summary(self)
Print a summary of the model's SageMaker training job.
def
get_training_log(self):
1088 def get_training_log(self): 1089 """ 1090 The training log of the last `fit()`. 1091 1092 Returns: 1093 The same `TrainingLog` `fit()` returned. 1094 1095 Raises: 1096 ValueError: If the model has not been fitted. 1097 """ 1098 if not self._is_fitted: 1099 raise ValueError( 1100 "Model must be fitted before accessing training log. Call fit() first." 1101 ) 1102 if self._training_log is None: 1103 raise ValueError( 1104 "No training log available. This should not happen if model was fitted properly." 1105 ) 1106 return self._training_log
The training log of the last fit().
Returns:
The same
TrainingLogfit()returned.
Raises:
- ValueError: If the model has not been fitted.
def
get_tuning_trials(self):
1108 def get_tuning_trials(self): 1109 """ 1110 The results of the tuning `fit()` ran. 1111 1112 Returns: 1113 A `TuningTrials`, or None if the config had no ranges and `fit()` 1114 did not tune. 1115 """ 1116 return getattr(self, "_tuning_trials", None)
The results of the tuning fit() ran.
Returns:
A
TuningTrials, or None if the config had no ranges andfit()did not tune.
def
get_config(self):
1118 def get_config(self): 1119 """ 1120 The model's config, with defaults filled in. 1121 1122 Returns: 1123 The config dict. 1124 """ 1125 return self.config
The model's config, with defaults filled in.
Returns:
The config dict.
def
print_config(self):
1127 def print_config(self): 1128 """Print the model configuration.""" 1129 print("Model Configuration:") 1130 for key, value in self.config.items(): 1131 print(f" {key}: {value}")
Print the model configuration.
def
get_evaluation_result(self, pre_calibration=False):
1133 def get_evaluation_result(self, pre_calibration=False): 1134 """ 1135 Get the evaluation result from training. 1136 1137 Args: 1138 pre_calibration: If True, return the pre-calibration result. 1139 1140 Returns: 1141 EvaluationResult if test data was provided during fit(), otherwise None. 1142 """ 1143 if pre_calibration: 1144 if self._evaluation_results is None: 1145 raise ValueError( 1146 "No pre-calibration result available: model has not been calibrated." 1147 ) 1148 return self._evaluation_results["before"] 1149 if self._evaluation_results is not None: 1150 return self._evaluation_results.get("after") 1151 return None
Get the evaluation result from training.
Arguments:
- pre_calibration: If True, return the pre-calibration result.
Returns:
EvaluationResult if test data was provided during fit(), otherwise None.
@dataclass
class
ContinuousRange:
341@dataclass 342class ContinuousRange: 343 """Continuous hyperparameter range for HPO.""" 344 345 low: float 346 high: float 347 log: bool = False
Continuous hyperparameter range for HPO.
@dataclass
class
IntegerRange:
350@dataclass 351class IntegerRange: 352 """Integer hyperparameter range for HPO.""" 353 354 low: int 355 high: int
Integer hyperparameter range for HPO.
@dataclass
class
CategoricalRange:
358@dataclass 359class CategoricalRange: 360 """Categorical hyperparameter range for HPO.""" 361 362 choices: list
Categorical hyperparameter range for HPO.
def
stratified_kfold(train_df, n_splits: int = 5, seed: int = 42):
236def stratified_kfold(train_df, n_splits: int = 5, seed: int = 42): 237 """ 238 Split training data into cross-validation folds, by unit. 239 240 These are the folds `SurvivalModel.fit()` tunes on. Use them to search 241 hyperparameters your own way, such as a grid search. 242 243 - **By unit.** All of a unit's rows go to the same side, so no unit is ever 244 in both halves of a fold, and each unit is in exactly one hold-out half. 245 - **Balanced.** Units are spread so that each fold has a similar event rate 246 and a similar spread of durations, which keeps fold scores comparable. 247 - **Deterministic.** The same data and `seed` give the same folds. 248 249 **Scoring a fold.** Score each model on its fold's hold-out half, which 250 must play no part in training. Give the model its own validation data for 251 early stopping, `val_split`, drawn from the training half. Do not pass the 252 hold-out half as `val_df`: validation data chooses the epoch the model 253 keeps and when its learning rate drops, so a score on it is optimistic. 254 255 The model predicts once per unit, at the unit's last row. A hold-out half 256 holds whole histories, so each unit's last row is at, or near, its end of 257 life, and scoring there says little. Cut each hold-out unit at a random 258 cycle first and label it with its remaining life at the cut, as a test set 259 is built. Then pass it to `fit()` as `test_df` and `test_labels_df`: the 260 model is scored on it after training, in the same training job, and never 261 trains on it. 262 263 ```python 264 from itertools import product 265 266 import numpy as np 267 import pandas as pd 268 269 from lifespan_predictor_sdk import SurvivalModel, stratified_kfold 270 271 272 def as_test_set(fold, seed=0): 273 # Cut each unit at a random cycle; label it with its remaining life there. 274 rng = np.random.default_rng(seed) 275 rows, labels = [], [] 276 for _, history in fold.groupby("unit"): 277 cut = history.iloc[: rng.integers(1, len(history) + 1)] 278 rows.append(cut.drop(columns=["duration", "event"])) 279 labels.append(cut.iloc[-1][["unit", "duration", "event"]]) 280 labels = pd.DataFrame(labels).astype({"unit": int, "event": int}) 281 return pd.concat(rows), labels 282 283 284 folds = [ 285 (fold_train, *as_test_set(hold_out)) 286 for fold_train, hold_out in stratified_kfold(train_df) 287 ] 288 grid = {"learning_rate": [1e-3, 1e-2], "num_nodes": [[16], [32, 32]]} 289 290 best_params, best_score = None, float("-inf") 291 for values in product(*grid.values()): 292 params = dict(zip(grid, values)) 293 scores = [] 294 for fold_train, test_df, test_labels_df in folds: 295 model = SurvivalModel({"model_type": "deepsurv", "epochs": 30, **params}) 296 model.fit( 297 fold_train, 298 val_split=0.2, 299 test_df=test_df, 300 test_labels_df=test_labels_df, 301 ) 302 scores.append(model.get_evaluation_result().c_index) 303 score = float(np.mean(scores)) 304 if score > best_score: 305 best_params, best_score = params, score 306 307 final = SurvivalModel({"model_type": "deepsurv", "epochs": 30, **best_params}) 308 final.fit(train_df, val_split=0.2) 309 ``` 310 311 On SageMaker every `fit()` is a separate training job, so this grid of 4 312 points over 5 folds runs 20 training jobs, plus the final one. Scoring 313 needs no `deploy()`: each job scores its own hold-out. 314 315 Args: 316 train_df: Training data with `unit`, `event` and `duration` columns, as 317 `SurvivalModel.fit()` takes it. 318 n_splits: The number of folds. 319 seed: Fixes how units are shuffled before they are assigned to folds. 320 321 Yields: 322 `n_splits` pairs `(fold_train_df, hold_out_df)` of DataFrames, each with 323 `train_df`'s columns. 324 325 Raises: 326 ValueError: If `train_df` lacks `unit`, `event` or `duration`. 327 """ 328 for col in ("unit", "event", "duration"): 329 if col not in train_df.columns: 330 raise ValueError( 331 f"stratified_kfold requires column '{col}' in train_df. " 332 f"Found columns: {list(train_df.columns)}" 333 ) 334 335 unit_stats = ( 336 train_df.groupby("unit") 337 .agg(event=("event", "last"), duration_mean=("duration", "mean")) 338 .reset_index() 339 ) 340 341 # Shuffle unit_stats with fixed seed so round-robin assignment is deterministic 342 rng = np.random.default_rng(seed=seed) 343 unit_stats = unit_stats.iloc[rng.permutation(len(unit_stats))].reset_index( 344 drop=True 345 ) 346 347 fold_assignments = _assign_units_to_folds(unit_stats, n_splits=n_splits) 348 349 for val_units in fold_assignments: 350 val_mask = train_df["unit"].isin(val_units) 351 fold_train = train_df[~val_mask].reset_index(drop=True) 352 fold_val = train_df[val_mask].reset_index(drop=True) 353 yield fold_train, fold_val
Split training data into cross-validation folds, by unit.
These are the folds SurvivalModel.fit() tunes on. Use them to search
hyperparameters your own way, such as a grid search.
- By unit. All of a unit's rows go to the same side, so no unit is ever in both halves of a fold, and each unit is in exactly one hold-out half.
- Balanced. Units are spread so that each fold has a similar event rate and a similar spread of durations, which keeps fold scores comparable.
- Deterministic. The same data and
seedgive the same folds.
Scoring a fold. Score each model on its fold's hold-out half, which
must play no part in training. Give the model its own validation data for
early stopping, val_split, drawn from the training half. Do not pass the
hold-out half as val_df: validation data chooses the epoch the model
keeps and when its learning rate drops, so a score on it is optimistic.
The model predicts once per unit, at the unit's last row. A hold-out half
holds whole histories, so each unit's last row is at, or near, its end of
life, and scoring there says little. Cut each hold-out unit at a random
cycle first and label it with its remaining life at the cut, as a test set
is built. Then pass it to fit() as test_df and test_labels_df: the
model is scored on it after training, in the same training job, and never
trains on it.
from itertools import product
import numpy as np
import pandas as pd
from lifespan_predictor_sdk import SurvivalModel, stratified_kfold
def as_test_set(fold, seed=0):
# Cut each unit at a random cycle; label it with its remaining life there.
rng = np.random.default_rng(seed)
rows, labels = [], []
for _, history in fold.groupby("unit"):
cut = history.iloc[: rng.integers(1, len(history) + 1)]
rows.append(cut.drop(columns=["duration", "event"]))
labels.append(cut.iloc[-1][["unit", "duration", "event"]])
labels = pd.DataFrame(labels).astype({"unit": int, "event": int})
return pd.concat(rows), labels
folds = [
(fold_train, *as_test_set(hold_out))
for fold_train, hold_out in stratified_kfold(train_df)
]
grid = {"learning_rate": [1e-3, 1e-2], "num_nodes": [[16], [32, 32]]}
best_params, best_score = None, float("-inf")
for values in product(*grid.values()):
params = dict(zip(grid, values))
scores = []
for fold_train, test_df, test_labels_df in folds:
model = SurvivalModel({"model_type": "deepsurv", "epochs": 30, **params})
model.fit(
fold_train,
val_split=0.2,
test_df=test_df,
test_labels_df=test_labels_df,
)
scores.append(model.get_evaluation_result().c_index)
score = float(np.mean(scores))
if score > best_score:
best_params, best_score = params, score
final = SurvivalModel({"model_type": "deepsurv", "epochs": 30, **best_params})
final.fit(train_df, val_split=0.2)
On SageMaker every fit() is a separate training job, so this grid of 4
points over 5 folds runs 20 training jobs, plus the final one. Scoring
needs no deploy(): each job scores its own hold-out.
Arguments:
- train_df: Training data with
unit,eventanddurationcolumns, asSurvivalModel.fit()takes it. - n_splits: The number of folds.
- seed: Fixes how units are shuffled before they are assigned to folds.
Yields:
n_splitspairs(fold_train_df, hold_out_df)of DataFrames, each withtrain_df's columns.
Raises:
- ValueError: If
train_dflacksunit,eventorduration.
class
TrainingLog:
443class TrainingLog: 444 """ 445 The history of one training run, returned by `SurvivalModel.fit()`. 446 447 Attributes: 448 train_loss: Training loss per epoch. 449 val_loss: Validation loss per epoch, or None without validation data. 450 learning_rate: Learning rate per epoch. 451 epochs: The configured maximum number of epochs. 452 num_epochs: The number of epochs trained, fewer than `epochs` if 453 training stopped early. 454 best_epochs: The epoch with the lowest validation loss, whose weights 455 the model keeps; None without validation data. 456 """ 457 458 def __init__(self, train_loss, learning_rate, epochs, val_loss=None): 459 self.train_loss = train_loss 460 self.learning_rate = learning_rate 461 self.epochs = epochs 462 self.val_loss = val_loss 463 464 # Automatically determine num_epochs. This is the number of epochs trained. 465 self.num_epochs = len(self.train_loss) 466 if self.val_loss: 467 self.best_epochs = int(self.val_loss.index(min(self.val_loss)) + 1) 468 else: 469 self.best_epochs = None 470 self._val_c_index = None 471 472 def to_pandas(self): 473 """ 474 The losses as a DataFrame. 475 476 Returns: 477 A DataFrame with `epoch`, `train_loss` and, with validation data, 478 `val_loss`. 479 """ 480 data = { 481 "epoch": list(range(1, len(self.train_loss) + 1)), 482 "train_loss": self.train_loss, 483 } 484 if self.val_loss: 485 data["val_loss"] = self.val_loss 486 return pd.DataFrame(data) 487 488 def plot(self, figsize=(12, 6), title="Training History"): 489 """ 490 Plot training and validation losses with learning rate changes and best epoch markers. 491 492 The plot automatically includes: 493 - Training and validation loss curves 494 - Learning rate reduction markers (vertical dashed lines) 495 - Best epoch marker (star on validation loss) 496 - Fit quality analysis in legend title (GOOD_FIT, OVER_FIT, etc.) 497 498 Args: 499 figsize: Tuple of (width, height) for the figure size 500 title: Title for the plot 501 502 Returns: 503 matplotlib Figure and Axes objects 504 505 Example: 506 >>> log = model.fit(X_train, durations, events, X_val, durations_val, events_val) 507 >>> fig, ax = log.plot() 508 >>> plt.savefig('training_history.png') 509 >>> # Legend will show fit quality automatically (e.g., "✓ Fit Quality: GOOD FIT") 510 """ 511 try: 512 import matplotlib.pyplot as plt 513 except ImportError: 514 raise ImportError( 515 "Plotting requires matplotlib. Install with: uv add matplotlib" 516 ) 517 518 fig, ax = plt.subplots(figsize=figsize) 519 520 epochs = list(range(1, len(self.train_loss) + 1)) 521 522 # Plot training loss 523 ax.plot( 524 epochs, 525 self.train_loss, 526 label="Training Loss", 527 color="#2E86AB", 528 linewidth=2, 529 marker="o", 530 markersize=4, 531 alpha=0.8, 532 ) 533 534 # Plot validation loss if available 535 if self.val_loss: 536 ax.plot( 537 epochs, 538 self.val_loss, 539 label="Validation Loss", 540 color="#A23B72", 541 linewidth=2, 542 marker="s", 543 markersize=4, 544 alpha=0.8, 545 ) 546 547 # Mark best epoch with a larger dot 548 if self.best_epochs and self.best_epochs <= len(self.val_loss): 549 best_val_loss = self.val_loss[self.best_epochs - 1] 550 ax.plot( 551 self.best_epochs, 552 best_val_loss, 553 marker="*", 554 markersize=20, 555 color="#F18F01", 556 label=f"Best Epoch ({self.best_epochs})", 557 zorder=5, 558 ) 559 560 # Detect and mark learning rate reductions 561 if self.learning_rate and len(self.learning_rate) > 1: 562 lr_changes = [] 563 for i in range(1, len(self.learning_rate)): 564 if self.learning_rate[i] < self.learning_rate[i - 1] * 0.95: 565 lr_changes.append(i + 1) # +1 for 1-based epoch numbering 566 567 # Draw vertical lines for LR changes 568 for epoch_idx in lr_changes: 569 ax.axvline( 570 x=epoch_idx, 571 color="#C73E1D", 572 linestyle="--", 573 linewidth=1.5, 574 alpha=0.7, 575 ) 576 577 # Add single legend entry for LR changes 578 if lr_changes: 579 ax.axvline( 580 x=-1, 581 color="#C73E1D", 582 linestyle="--", 583 linewidth=1.5, 584 alpha=0.7, 585 label="LR Reduction", 586 ) 587 588 # Analyze fit quality for legend 589 fit_quality = self.analyze_fit_quality() 590 591 # Create legend title with fit quality 592 legend_title = None 593 if fit_quality is not None: 594 quality_name = fit_quality.value.upper().replace("_", " ") 595 # Add symbol for visual indication (using text symbols for better compatibility) 596 quality_symbols = { 597 "GOOD_FIT": "✓", 598 "FULL_FIT": "★", 599 "OVER_FIT": "⚠", 600 "UNDER_FIT": "↓", 601 "UNKNOWN_FIT": "?", 602 } 603 symbol = quality_symbols.get(fit_quality.value.upper(), "") 604 legend_title = f"{symbol} Fit Quality: {quality_name}" 605 606 # Styling 607 ax.set_xlabel("Epoch", fontsize=12, fontweight="bold") 608 ax.set_ylabel("Loss", fontsize=12, fontweight="bold") 609 ax.set_title(title, fontsize=14, fontweight="bold", pad=20) 610 611 # Add legend with fit quality title 612 legend = ax.legend(loc="best", framealpha=0.9, fontsize=10) 613 if legend_title: 614 legend.set_title(legend_title, prop={"size": 10, "weight": "bold"}) 615 616 ax.grid(True, alpha=0.3, linestyle=":", linewidth=0.5) 617 ax.spines["top"].set_visible(False) 618 ax.spines["right"].set_visible(False) 619 620 # Set x-axis to show integer epochs 621 ax.set_xticks(range(1, len(epochs) + 1, max(1, len(epochs) // 10))) 622 623 plt.tight_layout() 624 return fig, ax 625 626 def analyze_fit_quality(self): 627 """ 628 Judge from the losses whether the model fits its data well. 629 630 Epochs after `best_epochs` are not considered, and a brief rise in 631 validation loss just after a learning-rate reduction is tolerated. 632 633 Returns: 634 None without validation data; otherwise a `FitQuality`, whose 635 `.value` is one of: 636 637 - `"good_fit"`: both losses fell and the gap between them did not 638 widen, and training ran to about its last epoch. 639 - `"full_fit"`: like a good fit, but training stopped early and the 640 last learning-rate reduction no longer helped. 641 - `"over_fit"`: validation loss rose while training loss fell for 642 more than three epochs in a row. 643 - `"under_fit"`: both losses still falling, but by less than 1% per 644 epoch over the last ten epochs. 645 - `"unknown_fit"`: none of the above. 646 """ 647 if self.val_loss is None or len(self.val_loss) == 0: 648 return None 649 # Detect learning rate reductions (95% threshold to account for floating point) 650 lr_reductions = [] 651 for i in range(1, len(self.learning_rate)): 652 if self.learning_rate[i] < self.learning_rate[i - 1] * 0.95: 653 lr_reductions.append(i) 654 655 # Truncate to best_epochs so post-best data doesn't skew fit analysis 656 if self.best_epochs is not None: 657 train_to_best = self.train_loss[: self.best_epochs] 658 val_to_best = self.val_loss[: self.best_epochs] 659 else: 660 train_to_best = self.train_loss 661 val_to_best = self.val_loss 662 663 # Check if GOOD_FIT conditions hold 664 is_good_fit = self._check_good_fit(train_to_best, val_to_best, lr_reductions) 665 666 if is_good_fit: 667 # Check if FULL_FIT: early stopping kicked in 668 if self.epochs > self.num_epochs: 669 return FitQuality.FULL_FIT 670 else: 671 return FitQuality.GOOD_FIT 672 673 # If not GOOD_FIT, check for OVER_FIT (already truncated to best_epochs above) 674 is_over_fit = self._check_over_fit(train_to_best, val_to_best, lr_reductions) 675 if is_over_fit: 676 return FitQuality.OVER_FIT 677 678 # If not OVER_FIT, check for UNDER_FIT 679 is_under_fit = self._check_under_fit(train_to_best, val_to_best) 680 if is_under_fit: 681 return FitQuality.UNDER_FIT 682 683 # Neither condition holds 684 return FitQuality.UNKNOWN_FIT 685 686 def _check_good_fit(self, train_loss, val_loss, lr_reductions): 687 """ 688 Check if training meets GOOD_FIT criteria: 689 1. Both train_loss and val_loss are decreasing overall 690 2. Gap between losses is not widening excessively (considering improvement rates) 691 3. Temporary val_loss increases after LR reductions are tolerated (1-2 epochs) 692 693 Returns: 694 bool: True if GOOD_FIT conditions are met 695 """ 696 if len(train_loss) < 5 or len(val_loss) < 5: 697 # Not enough data to determine 698 return False 699 700 # 1. Check if both losses are decreasing overall 701 # Compare first 5 epochs average with last 5 epochs average 702 train_early = sum(train_loss[:5]) / 5 703 train_late = sum(train_loss[-5:]) / 5 704 val_early = sum(val_loss[:5]) / 5 705 val_late = sum(val_loss[-5:]) / 5 706 707 if train_late >= train_early or val_late >= val_early: 708 # Losses not decreasing overall 709 return False 710 711 # Calculate improvement rates 712 train_improvement = (train_early - train_late) / train_early 713 val_improvement = (val_early - val_late) / val_early 714 715 # 2. Check if gap is not widening excessively 716 # Gap = val_loss - train_loss 717 gap_early = val_early - train_early 718 gap_late = val_late - train_late 719 720 # If both losses are improving significantly (>10%), allow more gap widening 721 # This handles the natural phenomenon where validation plateaus more than training 722 if train_improvement > 0.1 and val_improvement > 0.1: 723 # Both improving well, allow gap to widen up to 2.5x 724 gap_threshold = 2.5 725 else: 726 # Slow improvement, be stricter about gap widening 727 gap_threshold = 1.5 728 729 if gap_late > gap_early * gap_threshold: 730 # Gap widening too much even with lenient threshold 731 return False 732 733 # 3. Check for prolonged val_loss increases (ignoring LR reduction spikes) 734 # Mark epochs that are within 1-2 epochs after LR reduction as "tolerance zones" 735 tolerance_zones = set() 736 for lr_idx in lr_reductions: 737 tolerance_zones.add(lr_idx) # Epoch of LR reduction 738 if lr_idx + 1 < len(val_loss): 739 tolerance_zones.add(lr_idx + 1) # +1 epoch 740 if lr_idx + 2 < len(val_loss): 741 tolerance_zones.add(lr_idx + 2) # +2 epochs 742 743 # Count consecutive val_loss increases outside tolerance zones 744 consecutive_increases = 0 745 max_consecutive_increases = 0 746 747 for i in range(1, len(val_loss)): 748 if val_loss[i] > val_loss[i - 1]: 749 # Val loss increased 750 if i not in tolerance_zones: 751 consecutive_increases += 1 752 max_consecutive_increases = max( 753 max_consecutive_increases, consecutive_increases 754 ) 755 else: 756 consecutive_increases = 0 757 758 # If val_loss increased for >3 consecutive epochs (outside tolerance), not GOOD_FIT 759 if max_consecutive_increases > 3: 760 return False 761 762 # 4. Check improvement rate is reasonable (not too slow) 763 # GOOD_FIT should show meaningful progress, not just tiny decreases 764 # Analyze last 10 epochs (or fewer if not available) 765 if len(train_loss) >= 15: 766 window_size = min(10, len(train_loss) // 2) 767 train_recent = train_loss[-window_size:] 768 val_recent = val_loss[-window_size:] 769 770 train_start = train_recent[0] 771 train_end = train_recent[-1] 772 val_start = val_recent[0] 773 val_end = val_recent[-1] 774 775 if train_start > 0 and val_start > 0: 776 train_improvement_per_epoch = ( 777 (train_start - train_end) / train_start 778 ) / window_size 779 val_improvement_per_epoch = ( 780 (val_start - val_end) / val_start 781 ) / window_size 782 783 # Context matters: when losses are high (>1.0), we expect faster improvement 784 # When losses are low (<0.7), slower percentage improvement is acceptable 785 avg_loss = (train_end + val_end) / 2 786 787 if avg_loss > 1.0: 788 # High losses: expect at least 1% improvement per epoch 789 improvement_threshold = 0.01 790 else: 791 # Low losses: accept slower improvement (0.3% per epoch) 792 improvement_threshold = 0.003 793 794 # If BOTH improving less than threshold, it's too slow for GOOD_FIT 795 if ( 796 train_improvement_per_epoch < improvement_threshold 797 and val_improvement_per_epoch < improvement_threshold 798 ): 799 return False 800 801 # All conditions met 802 return True 803 804 def _check_over_fit(self, train_loss, val_loss, lr_reductions): 805 """ 806 Check if training shows OVER_FIT characteristics: 807 - Gap between train and val loss widening over time 808 - Val_loss increasing while train_loss decreasing for >3 consecutive epochs 809 - Tolerance for 1-2 epochs after LR reductions 810 811 Returns: 812 bool: True if OVER_FIT conditions are met 813 """ 814 if len(train_loss) < 5 or len(val_loss) < 5: 815 return False 816 817 # Mark tolerance zones around LR reductions 818 tolerance_zones = set() 819 for lr_idx in lr_reductions: 820 tolerance_zones.add(lr_idx) 821 if lr_idx + 1 < len(val_loss): 822 tolerance_zones.add(lr_idx + 1) 823 if lr_idx + 2 < len(val_loss): 824 tolerance_zones.add(lr_idx + 2) 825 826 # Count consecutive epochs where val_loss increases AND train_loss decreases 827 # (outside tolerance zones) 828 consecutive_divergence = 0 829 max_consecutive_divergence = 0 830 831 for i in range(1, min(len(train_loss), len(val_loss))): 832 val_increasing = val_loss[i] > val_loss[i - 1] 833 train_decreasing = train_loss[i] < train_loss[i - 1] 834 835 if val_increasing and train_decreasing: 836 # Diverging behavior 837 if i not in tolerance_zones: 838 consecutive_divergence += 1 839 max_consecutive_divergence = max( 840 max_consecutive_divergence, consecutive_divergence 841 ) 842 else: 843 consecutive_divergence = 0 844 845 # OVER_FIT if >3 consecutive epochs of divergence 846 if max_consecutive_divergence > 3: 847 return True 848 849 # Also check if gap is widening significantly 850 train_early = sum(train_loss[:5]) / 5 851 train_late = sum(train_loss[-5:]) / 5 852 val_early = sum(val_loss[:5]) / 5 853 val_late = sum(val_loss[-5:]) / 5 854 855 gap_early = val_early - train_early 856 gap_late = val_late - train_late 857 858 # OVER_FIT if gap more than tripled AND val_loss isn't improving much 859 if gap_late > gap_early * 3.0 and val_late > val_early * 0.95: 860 return True 861 862 return False 863 864 def _check_under_fit(self, train_loss, val_loss): 865 """ 866 Check if training shows UNDER_FIT characteristics: 867 - Both losses decreasing but average improvement per epoch <1% over last 10 epochs 868 869 Returns: 870 bool: True if UNDER_FIT conditions are met 871 """ 872 if len(train_loss) < 15 or len(val_loss) < 15: 873 # Need at least 15 epochs to analyze last 10 properly 874 return False 875 876 # Analyze last 10 epochs 877 window_size = min(10, len(train_loss) // 2) # Use at most half the data 878 train_recent = train_loss[-window_size:] 879 val_recent = val_loss[-window_size:] 880 881 # Check if both are still decreasing 882 train_start = train_recent[0] 883 train_end = train_recent[-1] 884 val_start = val_recent[0] 885 val_end = val_recent[-1] 886 887 if train_end >= train_start or val_end >= val_start: 888 # Not decreasing, can't be UNDER_FIT 889 return False 890 891 # Calculate average improvement per epoch (as percentage) 892 train_improvement_per_epoch = ( 893 (train_start - train_end) / train_start 894 ) / window_size 895 val_improvement_per_epoch = ((val_start - val_end) / val_start) / window_size 896 897 # UNDER_FIT if both improving less than 1% per epoch 898 if train_improvement_per_epoch < 0.01 and val_improvement_per_epoch < 0.01: 899 return True 900 901 return False
The history of one training run, returned by SurvivalModel.fit().
Attributes:
- train_loss: Training loss per epoch.
- val_loss: Validation loss per epoch, or None without validation data.
- learning_rate: Learning rate per epoch.
- epochs: The configured maximum number of epochs.
- num_epochs: The number of epochs trained, fewer than
epochsif training stopped early. - best_epochs: The epoch with the lowest validation loss, whose weights the model keeps; None without validation data.
TrainingLog(train_loss, learning_rate, epochs, val_loss=None)
458 def __init__(self, train_loss, learning_rate, epochs, val_loss=None): 459 self.train_loss = train_loss 460 self.learning_rate = learning_rate 461 self.epochs = epochs 462 self.val_loss = val_loss 463 464 # Automatically determine num_epochs. This is the number of epochs trained. 465 self.num_epochs = len(self.train_loss) 466 if self.val_loss: 467 self.best_epochs = int(self.val_loss.index(min(self.val_loss)) + 1) 468 else: 469 self.best_epochs = None 470 self._val_c_index = None
def
to_pandas(self):
472 def to_pandas(self): 473 """ 474 The losses as a DataFrame. 475 476 Returns: 477 A DataFrame with `epoch`, `train_loss` and, with validation data, 478 `val_loss`. 479 """ 480 data = { 481 "epoch": list(range(1, len(self.train_loss) + 1)), 482 "train_loss": self.train_loss, 483 } 484 if self.val_loss: 485 data["val_loss"] = self.val_loss 486 return pd.DataFrame(data)
The losses as a DataFrame.
Returns:
A DataFrame with
epoch,train_lossand, with validation data,val_loss.
def
plot(self, figsize=(12, 6), title='Training History'):
488 def plot(self, figsize=(12, 6), title="Training History"): 489 """ 490 Plot training and validation losses with learning rate changes and best epoch markers. 491 492 The plot automatically includes: 493 - Training and validation loss curves 494 - Learning rate reduction markers (vertical dashed lines) 495 - Best epoch marker (star on validation loss) 496 - Fit quality analysis in legend title (GOOD_FIT, OVER_FIT, etc.) 497 498 Args: 499 figsize: Tuple of (width, height) for the figure size 500 title: Title for the plot 501 502 Returns: 503 matplotlib Figure and Axes objects 504 505 Example: 506 >>> log = model.fit(X_train, durations, events, X_val, durations_val, events_val) 507 >>> fig, ax = log.plot() 508 >>> plt.savefig('training_history.png') 509 >>> # Legend will show fit quality automatically (e.g., "✓ Fit Quality: GOOD FIT") 510 """ 511 try: 512 import matplotlib.pyplot as plt 513 except ImportError: 514 raise ImportError( 515 "Plotting requires matplotlib. Install with: uv add matplotlib" 516 ) 517 518 fig, ax = plt.subplots(figsize=figsize) 519 520 epochs = list(range(1, len(self.train_loss) + 1)) 521 522 # Plot training loss 523 ax.plot( 524 epochs, 525 self.train_loss, 526 label="Training Loss", 527 color="#2E86AB", 528 linewidth=2, 529 marker="o", 530 markersize=4, 531 alpha=0.8, 532 ) 533 534 # Plot validation loss if available 535 if self.val_loss: 536 ax.plot( 537 epochs, 538 self.val_loss, 539 label="Validation Loss", 540 color="#A23B72", 541 linewidth=2, 542 marker="s", 543 markersize=4, 544 alpha=0.8, 545 ) 546 547 # Mark best epoch with a larger dot 548 if self.best_epochs and self.best_epochs <= len(self.val_loss): 549 best_val_loss = self.val_loss[self.best_epochs - 1] 550 ax.plot( 551 self.best_epochs, 552 best_val_loss, 553 marker="*", 554 markersize=20, 555 color="#F18F01", 556 label=f"Best Epoch ({self.best_epochs})", 557 zorder=5, 558 ) 559 560 # Detect and mark learning rate reductions 561 if self.learning_rate and len(self.learning_rate) > 1: 562 lr_changes = [] 563 for i in range(1, len(self.learning_rate)): 564 if self.learning_rate[i] < self.learning_rate[i - 1] * 0.95: 565 lr_changes.append(i + 1) # +1 for 1-based epoch numbering 566 567 # Draw vertical lines for LR changes 568 for epoch_idx in lr_changes: 569 ax.axvline( 570 x=epoch_idx, 571 color="#C73E1D", 572 linestyle="--", 573 linewidth=1.5, 574 alpha=0.7, 575 ) 576 577 # Add single legend entry for LR changes 578 if lr_changes: 579 ax.axvline( 580 x=-1, 581 color="#C73E1D", 582 linestyle="--", 583 linewidth=1.5, 584 alpha=0.7, 585 label="LR Reduction", 586 ) 587 588 # Analyze fit quality for legend 589 fit_quality = self.analyze_fit_quality() 590 591 # Create legend title with fit quality 592 legend_title = None 593 if fit_quality is not None: 594 quality_name = fit_quality.value.upper().replace("_", " ") 595 # Add symbol for visual indication (using text symbols for better compatibility) 596 quality_symbols = { 597 "GOOD_FIT": "✓", 598 "FULL_FIT": "★", 599 "OVER_FIT": "⚠", 600 "UNDER_FIT": "↓", 601 "UNKNOWN_FIT": "?", 602 } 603 symbol = quality_symbols.get(fit_quality.value.upper(), "") 604 legend_title = f"{symbol} Fit Quality: {quality_name}" 605 606 # Styling 607 ax.set_xlabel("Epoch", fontsize=12, fontweight="bold") 608 ax.set_ylabel("Loss", fontsize=12, fontweight="bold") 609 ax.set_title(title, fontsize=14, fontweight="bold", pad=20) 610 611 # Add legend with fit quality title 612 legend = ax.legend(loc="best", framealpha=0.9, fontsize=10) 613 if legend_title: 614 legend.set_title(legend_title, prop={"size": 10, "weight": "bold"}) 615 616 ax.grid(True, alpha=0.3, linestyle=":", linewidth=0.5) 617 ax.spines["top"].set_visible(False) 618 ax.spines["right"].set_visible(False) 619 620 # Set x-axis to show integer epochs 621 ax.set_xticks(range(1, len(epochs) + 1, max(1, len(epochs) // 10))) 622 623 plt.tight_layout() 624 return fig, ax
Plot training and validation losses with learning rate changes and best epoch markers.
The plot automatically includes:
- Training and validation loss curves
- Learning rate reduction markers (vertical dashed lines)
- Best epoch marker (star on validation loss)
- Fit quality analysis in legend title (GOOD_FIT, OVER_FIT, etc.)
Arguments:
- figsize: Tuple of (width, height) for the figure size
- title: Title for the plot
Returns:
matplotlib Figure and Axes objects
Example:
>>> log = model.fit(X_train, durations, events, X_val, durations_val, events_val) >>> fig, ax = log.plot() >>> plt.savefig('training_history.png') >>> # Legend will show fit quality automatically (e.g., "✓ Fit Quality: GOOD FIT")
def
analyze_fit_quality(self):
626 def analyze_fit_quality(self): 627 """ 628 Judge from the losses whether the model fits its data well. 629 630 Epochs after `best_epochs` are not considered, and a brief rise in 631 validation loss just after a learning-rate reduction is tolerated. 632 633 Returns: 634 None without validation data; otherwise a `FitQuality`, whose 635 `.value` is one of: 636 637 - `"good_fit"`: both losses fell and the gap between them did not 638 widen, and training ran to about its last epoch. 639 - `"full_fit"`: like a good fit, but training stopped early and the 640 last learning-rate reduction no longer helped. 641 - `"over_fit"`: validation loss rose while training loss fell for 642 more than three epochs in a row. 643 - `"under_fit"`: both losses still falling, but by less than 1% per 644 epoch over the last ten epochs. 645 - `"unknown_fit"`: none of the above. 646 """ 647 if self.val_loss is None or len(self.val_loss) == 0: 648 return None 649 # Detect learning rate reductions (95% threshold to account for floating point) 650 lr_reductions = [] 651 for i in range(1, len(self.learning_rate)): 652 if self.learning_rate[i] < self.learning_rate[i - 1] * 0.95: 653 lr_reductions.append(i) 654 655 # Truncate to best_epochs so post-best data doesn't skew fit analysis 656 if self.best_epochs is not None: 657 train_to_best = self.train_loss[: self.best_epochs] 658 val_to_best = self.val_loss[: self.best_epochs] 659 else: 660 train_to_best = self.train_loss 661 val_to_best = self.val_loss 662 663 # Check if GOOD_FIT conditions hold 664 is_good_fit = self._check_good_fit(train_to_best, val_to_best, lr_reductions) 665 666 if is_good_fit: 667 # Check if FULL_FIT: early stopping kicked in 668 if self.epochs > self.num_epochs: 669 return FitQuality.FULL_FIT 670 else: 671 return FitQuality.GOOD_FIT 672 673 # If not GOOD_FIT, check for OVER_FIT (already truncated to best_epochs above) 674 is_over_fit = self._check_over_fit(train_to_best, val_to_best, lr_reductions) 675 if is_over_fit: 676 return FitQuality.OVER_FIT 677 678 # If not OVER_FIT, check for UNDER_FIT 679 is_under_fit = self._check_under_fit(train_to_best, val_to_best) 680 if is_under_fit: 681 return FitQuality.UNDER_FIT 682 683 # Neither condition holds 684 return FitQuality.UNKNOWN_FIT
Judge from the losses whether the model fits its data well.
Epochs after best_epochs are not considered, and a brief rise in
validation loss just after a learning-rate reduction is tolerated.
Returns:
None without validation data; otherwise a
FitQuality, whose.valueis one of:
"good_fit": both losses fell and the gap between them did not widen, and training ran to about its last epoch."full_fit": like a good fit, but training stopped early and the last learning-rate reduction no longer helped."over_fit": validation loss rose while training loss fell for more than three epochs in a row."under_fit": both losses still falling, but by less than 1% per epoch over the last ten epochs."unknown_fit": none of the above.
class
EvaluationResult:
11class EvaluationResult: 12 """ 13 How well a model scored on labelled data. 14 15 Returned by `SurvivalModel.evaluate()`, and by 16 `SurvivalModel.get_evaluation_result()` when `fit()` was given test data. 17 A metric that could not be computed is None. 18 19 Attributes: 20 c_index: Harrell's concordance index: the fraction of comparable pairs 21 of units whose predicted times are in the right order. 0.5 is 22 chance, 1.0 is perfect. 23 ipcw_c_index: The concordance index corrected for censoring. 24 rmse: Root mean squared error of the predicted mean times, over the 25 units that failed. 26 ibs: Integrated Brier score of the survival curves; lower is better. 27 ibs_classification: A plain-words reading of `ibs`. 28 calibration_bias: How far the predicted probabilities are from the 29 observed ones, signed. 30 calibration_classification: A plain-words reading of 31 `calibration_bias`. 32 33 A result can also be read like a dict: `result["c_index"]`, 34 `result.get("ibs")`. 35 """ 36 37 def __init__( 38 self, 39 c_index, 40 rmse, 41 predictions=None, 42 mask_cindex=None, 43 mask_rmse=None, 44 y_test=None, 45 events=None, 46 calibration_bias=None, 47 calibration_mean_cdf=None, 48 calibration_classification=None, 49 calibration_sorted_cdf=None, 50 calibration_empirical_cdf=None, 51 calibration_sorted_cdf_events=None, 52 calibration_empirical_cdf_events=None, 53 ibs=None, 54 ibs_classification=None, 55 antolini_c_index=None, 56 competing_risk_per_risk_cause_specific_c_index=None, 57 competing_risk_per_risk_antolini=None, 58 competing_risk_per_risk_predictive_c_index=None, 59 ipcw_c_index=None, 60 ): 61 self.c_index = c_index 62 self.rmse = rmse 63 self.predictions = predictions 64 self.mask_cindex = mask_cindex 65 self.mask_rmse = mask_rmse 66 self.y_test = y_test 67 self.events = events 68 self.calibration_bias = calibration_bias 69 self.calibration_mean_cdf = calibration_mean_cdf 70 self.calibration_classification = calibration_classification 71 self.calibration_sorted_cdf = calibration_sorted_cdf 72 self.calibration_empirical_cdf = calibration_empirical_cdf 73 self.calibration_sorted_cdf_events = calibration_sorted_cdf_events 74 self.calibration_empirical_cdf_events = calibration_empirical_cdf_events 75 self.ibs = ibs 76 self.ibs_classification = ibs_classification 77 self.antolini_c_index = antolini_c_index 78 self.competing_risk_per_risk_cause_specific_c_index = ( 79 competing_risk_per_risk_cause_specific_c_index 80 ) 81 self.competing_risk_per_risk_antolini = competing_risk_per_risk_antolini 82 self.competing_risk_per_risk_predictive_c_index = ( 83 competing_risk_per_risk_predictive_c_index 84 ) 85 self.ipcw_c_index = ipcw_c_index 86 87 def __getitem__(self, key): 88 return getattr(self, key) 89 90 def __contains__(self, key): 91 return hasattr(self, key) 92 93 def get(self, key, default=None): 94 """ 95 A metric by name, like `dict.get`. 96 97 Args: 98 key: The attribute's name, such as `"c_index"`. 99 default: What to return if there is no such attribute. 100 101 Returns: 102 The metric's value, or `default`. 103 """ 104 return getattr(self, key, default) 105 106 def get_metrics(self): 107 """ 108 The scalar metrics, as a dict. 109 110 Returns: 111 A dict with `c_index` and `rmse`, and each other metric that was 112 computed. 113 """ 114 metrics = {"c_index": self.c_index, "rmse": self.rmse} 115 if self.ipcw_c_index is not None: 116 metrics["ipcw_c_index"] = self.ipcw_c_index 117 if self.antolini_c_index is not None: 118 metrics["antolini_c_index"] = self.antolini_c_index 119 if self.ibs is not None: 120 metrics["ibs"] = self.ibs 121 if self.calibration_bias is not None: 122 metrics["calibration_bias"] = self.calibration_bias 123 if self.competing_risk_per_risk_cause_specific_c_index is not None: 124 metrics["competing_risk_per_risk_cause_specific_c_index"] = list( 125 self.competing_risk_per_risk_cause_specific_c_index 126 ) 127 if self.competing_risk_per_risk_antolini is not None: 128 metrics["competing_risk_per_risk_antolini"] = list( 129 self.competing_risk_per_risk_antolini 130 ) 131 if self.competing_risk_per_risk_predictive_c_index is not None: 132 metrics["competing_risk_per_risk_predictive_c_index"] = list( 133 self.competing_risk_per_risk_predictive_c_index 134 ) 135 return metrics 136 137 def print_metrics(self): 138 """Print evaluation metrics in a formatted table.""" 139 # A competing-risks model's top block is scored over *all* causes, so it says 140 # so. Without the label it reads as though it described a single risk — which 141 # is exactly how it was misread while these fields held risk 1's numbers. 142 # 143 # The predictive figure is the one every competing-risks model reports, so it 144 # is what identifies one. The cause-specific figures are fleet-only diagnostics 145 # and are absent for DeepHit, which has no sub-models to diagnose. 146 is_competing_risk = ( 147 self.competing_risk_per_risk_predictive_c_index is not None 148 or self.competing_risk_per_risk_cause_specific_c_index is not None 149 or self.competing_risk_per_risk_antolini is not None 150 ) 151 header = ( 152 "All-Cause Evaluation Metrics" 153 if is_competing_risk 154 else "Evaluation Metrics" 155 ) 156 print(f"=== {header} ===") 157 # The IPCW-corrected figure leads when it is available. Harrell's C depends 158 # on the study's censoring distribution, so the same model on the same 159 # population scores differently under heavier drop-out; the corrected one 160 # does not. Harrell's stays on the result object and in get_metrics() for 161 # anyone tracking it, but it is not the headline. 162 if self.ipcw_c_index is not None: 163 print( 164 f" C-index (IPCW): {self.ipcw_c_index:.4f} (higher is better, " 165 "0.5 = random; corrected for censoring)" 166 ) 167 else: 168 print( 169 f" C-index (Harrell): {self.c_index:.4f} (higher is better, 0.5 = random, 1.0 = perfect)" 170 ) 171 if self.antolini_c_index is not None: 172 print( 173 f" C-index (Antolini): {self.antolini_c_index:.4f} (higher is better, time-dependent)" 174 ) 175 print(f" RMSE: {self.rmse:.4f} (lower is better)") 176 if self.ibs is not None: 177 label = f" ({self.ibs_classification})" if self.ibs_classification else "" 178 print( 179 f" IBS: {self.ibs:.4f} (lower is better, <0.1 = excellent){label}" 180 ) 181 if self.calibration_bias is not None: 182 label = ( 183 f" ({self.calibration_classification})" 184 if self.calibration_classification 185 else "" 186 ) 187 print( 188 f" Calibration bias: {self.calibration_bias:.4f} (closer to 0 is better){label}" 189 ) 190 if ( 191 self.competing_risk_per_risk_cause_specific_c_index is not None 192 or self.competing_risk_per_risk_antolini is not None 193 or self.competing_risk_per_risk_predictive_c_index is not None 194 ): 195 196 def _fmt(c): 197 return ( 198 "n/a" 199 if c is None or (isinstance(c, float) and c != c) 200 else f"{c:.4f}" 201 ) 202 203 print("=== Competing-Risk Metrics ===") 204 # The predictive figure leads, and is the only per-risk number that means 205 # the same thing for a DeepHit and for a CSH fleet — so it is the only one 206 # the two may be compared on. It keeps a competing-event unit comparable, 207 # because that unit is known never to fail from cause k. 208 # 209 # The figures below it are within-family diagnostics, printed under their 210 # own heading so they are not read as competing answers to the same 211 # question. A fleet's are its sub-models' own scores, in a counterfactual 212 # world where cause k is the only cause; DeepHit's Antolini scores the 213 # 1 - CIF_k curve in the real one. Useful for locating a weak sub-model, 214 # not for ranking one model family against the other. 215 if self.competing_risk_per_risk_predictive_c_index is not None: 216 print(" Per-risk C-index (predictive, competing-risks-aware):") 217 for k, c in enumerate( 218 self.competing_risk_per_risk_predictive_c_index, start=1 219 ): 220 print(f" Risk {k}: {_fmt(c)}") 221 if ( 222 self.competing_risk_per_risk_cause_specific_c_index is not None 223 or self.competing_risk_per_risk_antolini is not None 224 ): 225 print(" Diagnostics (within-family, not comparable across models):") 226 if self.competing_risk_per_risk_cause_specific_c_index is not None: 227 print(" Per-sub-model C-index (cause-specific, Harrell):") 228 for k, c in enumerate( 229 self.competing_risk_per_risk_cause_specific_c_index, start=1 230 ): 231 print(f" Risk {k}: {_fmt(c)}") 232 if self.competing_risk_per_risk_antolini is not None: 233 print(" Per-risk C-index (time-dependent, Antolini):") 234 for k, c in enumerate( 235 self.competing_risk_per_risk_antolini, start=1 236 ): 237 print(f" Risk {k}: {_fmt(c)}") 238 239 def plot_predictions(self, figsize=(8, 6)): 240 """Scatter plot of actual against predicted times, for the units that failed. 241 242 Args: 243 figsize: The figure's size in inches, as `(width, height)`. 244 245 Returns: 246 A matplotlib `(fig, ax)` tuple, or None if the per-unit predictions 247 were not kept, as for the result of a training job. 248 """ 249 import matplotlib.pyplot as plt 250 251 if self.predictions is None or self.mask_rmse is None or self.y_test is None: 252 print("Prediction arrays not available (e.g. SageMaker backend).") 253 return None 254 255 import numpy as np 256 257 y_test_masked = self.y_test[self.mask_rmse] 258 y_pred_masked = self.predictions[self.mask_rmse] 259 260 fig, ax = plt.subplots(figsize=figsize) 261 262 # Color by competing-risk class when events are multi-risk 263 events_masked = None 264 if self.events is not None: 265 events_masked = np.asarray(self.events)[self.mask_rmse].astype(int) 266 267 if events_masked is not None and events_masked.max(initial=0) > 1: 268 unique_risks = np.sort(np.unique(events_masked[events_masked > 0])) 269 cmap = plt.get_cmap("tab10") 270 for i, k in enumerate(unique_risks): 271 mask = events_masked == k 272 ax.scatter( 273 y_test_masked[mask], 274 y_pred_masked[mask], 275 color=cmap(i % 10), 276 label=f"Risk {int(k)}", 277 alpha=0.8, 278 ) 279 ax.legend(title="Competing risk", loc="best") 280 else: 281 ax.scatter(y_test_masked, y_pred_masked) 282 283 ax.set_xlabel("Actual Values (y_test)") 284 ax.set_ylabel("Predicted Values (y_pred)") 285 ax.set_title("Actual vs Predicted") 286 plt.show() 287 return fig, ax 288 289 def plot_calibration(self, figsize=(12, 5), show="uncensored"): 290 """Plot the calibration curve, and the integrated Brier score. 291 292 Args: 293 figsize: The figure's size in inches, as `(width, height)`. 294 show: `"uncensored"` draws the curve over the units that failed; 295 `"all"` draws it over all units, with the uncensored curve 296 beside it. 297 298 Returns: 299 A matplotlib `(fig, axes)` tuple, or None if there is no calibration 300 data. 301 """ 302 if self.calibration_sorted_cdf is None and self.calibration_bias is None: 303 print("No calibration data available.") 304 return None 305 306 import matplotlib.pyplot as plt 307 import numpy as np 308 309 fig, axes = plt.subplots(1, 2, figsize=figsize) 310 fig.suptitle("Model Calibration", fontsize=14) 311 312 ax1 = axes[0] 313 sorted_cdf = self.calibration_sorted_cdf 314 empirical_cdf = self.calibration_empirical_cdf 315 if sorted_cdf is not None and empirical_cdf is not None: 316 sorted_cdf = np.atleast_1d(np.asarray(sorted_cdf, dtype=float)) 317 empirical_cdf = np.atleast_1d(np.asarray(empirical_cdf, dtype=float)) 318 sorted_cdf_events = self.calibration_sorted_cdf_events 319 empirical_cdf_events = self.calibration_empirical_cdf_events 320 if sorted_cdf_events is not None: 321 sorted_cdf_events = np.atleast_1d( 322 np.asarray(sorted_cdf_events, dtype=float) 323 ) 324 if empirical_cdf_events is not None: 325 empirical_cdf_events = np.atleast_1d( 326 np.asarray(empirical_cdf_events, dtype=float) 327 ) 328 ax1.plot([0, 1], [0, 1], "k--", label="Perfect calibration") 329 if show in ("all", "both"): 330 ax1.plot( 331 sorted_cdf, empirical_cdf, "b-", lw=2, label="All observations" 332 ) 333 ax1.fill_between( 334 sorted_cdf, sorted_cdf, empirical_cdf, alpha=0.1, color="blue" 335 ) 336 if sorted_cdf_events is not None and empirical_cdf_events is not None: 337 ax1.plot( 338 sorted_cdf_events, 339 empirical_cdf_events, 340 "r-", 341 lw=2, 342 label="Uncensored only", 343 ) 344 elif show == "uncensored": 345 cdf = sorted_cdf_events if sorted_cdf_events is not None else sorted_cdf 346 ecdf = ( 347 empirical_cdf_events 348 if empirical_cdf_events is not None 349 else empirical_cdf 350 ) 351 label = ( 352 "Uncensored only" 353 if sorted_cdf_events is not None 354 else "All observations" 355 ) 356 ax1.plot(cdf, ecdf, "b-", lw=2, label=label) 357 ax1.fill_between(cdf, cdf, ecdf, alpha=0.15, color="blue") 358 else: 359 ax1.plot( 360 sorted_cdf, empirical_cdf, "b-", lw=2, label="All observations" 361 ) 362 ax1.fill_between( 363 sorted_cdf, sorted_cdf, empirical_cdf, alpha=0.15, color="blue" 364 ) 365 ax1.set_xlabel("Predicted CDF") 366 ax1.set_ylabel("Empirical CDF") 367 ax1.set_title("Calibration Curve") 368 ax1.legend() 369 if ( 370 self.calibration_classification is not None 371 or self.calibration_bias is not None 372 ): 373 parts = [] 374 if self.calibration_classification: 375 parts.append(self.calibration_classification) 376 if self.calibration_bias is not None: 377 parts.append(f"bias={self.calibration_bias:.3f}") 378 ax1.annotate( 379 "\n".join(parts), xy=(0.05, 0.85), xycoords="axes fraction" 380 ) 381 else: 382 ax1.text( 383 0.5, 384 0.5, 385 "No calibration data", 386 ha="center", 387 va="center", 388 transform=ax1.transAxes, 389 ) 390 ax1.set_title("Calibration Curve") 391 392 ax2 = axes[1] 393 if self.ibs is not None: 394 colors = { 395 "excellent": "green", 396 "good": "yellowgreen", 397 "acceptable": "orange", 398 "poor": "red", 399 } 400 color = colors.get(self.ibs_classification or "", "steelblue") 401 ax2.barh(["IBS"], [self.ibs], color=color) 402 ax2.axvline( 403 0.25, color="gray", linestyle="--", label="Random baseline (0.25)" 404 ) 405 ax2.set_xlim(0, max(0.4, self.ibs * 1.2)) 406 ax2.set_title("Integrated Brier Score") 407 ax2.legend() 408 ibs_label = f"{self.ibs:.3f}" 409 if self.ibs_classification: 410 ibs_label += f" — {self.ibs_classification}" 411 ax2.annotate(ibs_label, xy=(self.ibs + 0.01, 0), va="center") 412 else: 413 ax2.text( 414 0.5, 415 0.5, 416 "IBS not available", 417 ha="center", 418 va="center", 419 transform=ax2.transAxes, 420 ) 421 ax2.set_title("Integrated Brier Score") 422 423 plt.tight_layout() 424 return fig, axes 425 426 def report(self): 427 """Print the metrics and draw both plots.""" 428 self.print_metrics() 429 self.plot_predictions() 430 self.plot_calibration()
How well a model scored on labelled data.
Returned by SurvivalModel.evaluate(), and by
SurvivalModel.get_evaluation_result() when fit() was given test data.
A metric that could not be computed is None.
Attributes:
- c_index: Harrell's concordance index: the fraction of comparable pairs of units whose predicted times are in the right order. 0.5 is chance, 1.0 is perfect.
- ipcw_c_index: The concordance index corrected for censoring.
- rmse: Root mean squared error of the predicted mean times, over the units that failed.
- ibs: Integrated Brier score of the survival curves; lower is better.
- ibs_classification: A plain-words reading of
ibs. - calibration_bias: How far the predicted probabilities are from the observed ones, signed.
- calibration_classification: A plain-words reading of
calibration_bias.
A result can also be read like a dict: result["c_index"],
result.get("ibs").
EvaluationResult( c_index, rmse, predictions=None, mask_cindex=None, mask_rmse=None, y_test=None, events=None, calibration_bias=None, calibration_mean_cdf=None, calibration_classification=None, calibration_sorted_cdf=None, calibration_empirical_cdf=None, calibration_sorted_cdf_events=None, calibration_empirical_cdf_events=None, ibs=None, ibs_classification=None, antolini_c_index=None, competing_risk_per_risk_cause_specific_c_index=None, competing_risk_per_risk_antolini=None, competing_risk_per_risk_predictive_c_index=None, ipcw_c_index=None)
37 def __init__( 38 self, 39 c_index, 40 rmse, 41 predictions=None, 42 mask_cindex=None, 43 mask_rmse=None, 44 y_test=None, 45 events=None, 46 calibration_bias=None, 47 calibration_mean_cdf=None, 48 calibration_classification=None, 49 calibration_sorted_cdf=None, 50 calibration_empirical_cdf=None, 51 calibration_sorted_cdf_events=None, 52 calibration_empirical_cdf_events=None, 53 ibs=None, 54 ibs_classification=None, 55 antolini_c_index=None, 56 competing_risk_per_risk_cause_specific_c_index=None, 57 competing_risk_per_risk_antolini=None, 58 competing_risk_per_risk_predictive_c_index=None, 59 ipcw_c_index=None, 60 ): 61 self.c_index = c_index 62 self.rmse = rmse 63 self.predictions = predictions 64 self.mask_cindex = mask_cindex 65 self.mask_rmse = mask_rmse 66 self.y_test = y_test 67 self.events = events 68 self.calibration_bias = calibration_bias 69 self.calibration_mean_cdf = calibration_mean_cdf 70 self.calibration_classification = calibration_classification 71 self.calibration_sorted_cdf = calibration_sorted_cdf 72 self.calibration_empirical_cdf = calibration_empirical_cdf 73 self.calibration_sorted_cdf_events = calibration_sorted_cdf_events 74 self.calibration_empirical_cdf_events = calibration_empirical_cdf_events 75 self.ibs = ibs 76 self.ibs_classification = ibs_classification 77 self.antolini_c_index = antolini_c_index 78 self.competing_risk_per_risk_cause_specific_c_index = ( 79 competing_risk_per_risk_cause_specific_c_index 80 ) 81 self.competing_risk_per_risk_antolini = competing_risk_per_risk_antolini 82 self.competing_risk_per_risk_predictive_c_index = ( 83 competing_risk_per_risk_predictive_c_index 84 ) 85 self.ipcw_c_index = ipcw_c_index
def
get(self, key, default=None):
93 def get(self, key, default=None): 94 """ 95 A metric by name, like `dict.get`. 96 97 Args: 98 key: The attribute's name, such as `"c_index"`. 99 default: What to return if there is no such attribute. 100 101 Returns: 102 The metric's value, or `default`. 103 """ 104 return getattr(self, key, default)
A metric by name, like dict.get.
Arguments:
- key: The attribute's name, such as
"c_index". - default: What to return if there is no such attribute.
Returns:
The metric's value, or
default.
def
get_metrics(self):
106 def get_metrics(self): 107 """ 108 The scalar metrics, as a dict. 109 110 Returns: 111 A dict with `c_index` and `rmse`, and each other metric that was 112 computed. 113 """ 114 metrics = {"c_index": self.c_index, "rmse": self.rmse} 115 if self.ipcw_c_index is not None: 116 metrics["ipcw_c_index"] = self.ipcw_c_index 117 if self.antolini_c_index is not None: 118 metrics["antolini_c_index"] = self.antolini_c_index 119 if self.ibs is not None: 120 metrics["ibs"] = self.ibs 121 if self.calibration_bias is not None: 122 metrics["calibration_bias"] = self.calibration_bias 123 if self.competing_risk_per_risk_cause_specific_c_index is not None: 124 metrics["competing_risk_per_risk_cause_specific_c_index"] = list( 125 self.competing_risk_per_risk_cause_specific_c_index 126 ) 127 if self.competing_risk_per_risk_antolini is not None: 128 metrics["competing_risk_per_risk_antolini"] = list( 129 self.competing_risk_per_risk_antolini 130 ) 131 if self.competing_risk_per_risk_predictive_c_index is not None: 132 metrics["competing_risk_per_risk_predictive_c_index"] = list( 133 self.competing_risk_per_risk_predictive_c_index 134 ) 135 return metrics
def
print_metrics(self):
137 def print_metrics(self): 138 """Print evaluation metrics in a formatted table.""" 139 # A competing-risks model's top block is scored over *all* causes, so it says 140 # so. Without the label it reads as though it described a single risk — which 141 # is exactly how it was misread while these fields held risk 1's numbers. 142 # 143 # The predictive figure is the one every competing-risks model reports, so it 144 # is what identifies one. The cause-specific figures are fleet-only diagnostics 145 # and are absent for DeepHit, which has no sub-models to diagnose. 146 is_competing_risk = ( 147 self.competing_risk_per_risk_predictive_c_index is not None 148 or self.competing_risk_per_risk_cause_specific_c_index is not None 149 or self.competing_risk_per_risk_antolini is not None 150 ) 151 header = ( 152 "All-Cause Evaluation Metrics" 153 if is_competing_risk 154 else "Evaluation Metrics" 155 ) 156 print(f"=== {header} ===") 157 # The IPCW-corrected figure leads when it is available. Harrell's C depends 158 # on the study's censoring distribution, so the same model on the same 159 # population scores differently under heavier drop-out; the corrected one 160 # does not. Harrell's stays on the result object and in get_metrics() for 161 # anyone tracking it, but it is not the headline. 162 if self.ipcw_c_index is not None: 163 print( 164 f" C-index (IPCW): {self.ipcw_c_index:.4f} (higher is better, " 165 "0.5 = random; corrected for censoring)" 166 ) 167 else: 168 print( 169 f" C-index (Harrell): {self.c_index:.4f} (higher is better, 0.5 = random, 1.0 = perfect)" 170 ) 171 if self.antolini_c_index is not None: 172 print( 173 f" C-index (Antolini): {self.antolini_c_index:.4f} (higher is better, time-dependent)" 174 ) 175 print(f" RMSE: {self.rmse:.4f} (lower is better)") 176 if self.ibs is not None: 177 label = f" ({self.ibs_classification})" if self.ibs_classification else "" 178 print( 179 f" IBS: {self.ibs:.4f} (lower is better, <0.1 = excellent){label}" 180 ) 181 if self.calibration_bias is not None: 182 label = ( 183 f" ({self.calibration_classification})" 184 if self.calibration_classification 185 else "" 186 ) 187 print( 188 f" Calibration bias: {self.calibration_bias:.4f} (closer to 0 is better){label}" 189 ) 190 if ( 191 self.competing_risk_per_risk_cause_specific_c_index is not None 192 or self.competing_risk_per_risk_antolini is not None 193 or self.competing_risk_per_risk_predictive_c_index is not None 194 ): 195 196 def _fmt(c): 197 return ( 198 "n/a" 199 if c is None or (isinstance(c, float) and c != c) 200 else f"{c:.4f}" 201 ) 202 203 print("=== Competing-Risk Metrics ===") 204 # The predictive figure leads, and is the only per-risk number that means 205 # the same thing for a DeepHit and for a CSH fleet — so it is the only one 206 # the two may be compared on. It keeps a competing-event unit comparable, 207 # because that unit is known never to fail from cause k. 208 # 209 # The figures below it are within-family diagnostics, printed under their 210 # own heading so they are not read as competing answers to the same 211 # question. A fleet's are its sub-models' own scores, in a counterfactual 212 # world where cause k is the only cause; DeepHit's Antolini scores the 213 # 1 - CIF_k curve in the real one. Useful for locating a weak sub-model, 214 # not for ranking one model family against the other. 215 if self.competing_risk_per_risk_predictive_c_index is not None: 216 print(" Per-risk C-index (predictive, competing-risks-aware):") 217 for k, c in enumerate( 218 self.competing_risk_per_risk_predictive_c_index, start=1 219 ): 220 print(f" Risk {k}: {_fmt(c)}") 221 if ( 222 self.competing_risk_per_risk_cause_specific_c_index is not None 223 or self.competing_risk_per_risk_antolini is not None 224 ): 225 print(" Diagnostics (within-family, not comparable across models):") 226 if self.competing_risk_per_risk_cause_specific_c_index is not None: 227 print(" Per-sub-model C-index (cause-specific, Harrell):") 228 for k, c in enumerate( 229 self.competing_risk_per_risk_cause_specific_c_index, start=1 230 ): 231 print(f" Risk {k}: {_fmt(c)}") 232 if self.competing_risk_per_risk_antolini is not None: 233 print(" Per-risk C-index (time-dependent, Antolini):") 234 for k, c in enumerate( 235 self.competing_risk_per_risk_antolini, start=1 236 ): 237 print(f" Risk {k}: {_fmt(c)}")
Print evaluation metrics in a formatted table.
def
plot_predictions(self, figsize=(8, 6)):
239 def plot_predictions(self, figsize=(8, 6)): 240 """Scatter plot of actual against predicted times, for the units that failed. 241 242 Args: 243 figsize: The figure's size in inches, as `(width, height)`. 244 245 Returns: 246 A matplotlib `(fig, ax)` tuple, or None if the per-unit predictions 247 were not kept, as for the result of a training job. 248 """ 249 import matplotlib.pyplot as plt 250 251 if self.predictions is None or self.mask_rmse is None or self.y_test is None: 252 print("Prediction arrays not available (e.g. SageMaker backend).") 253 return None 254 255 import numpy as np 256 257 y_test_masked = self.y_test[self.mask_rmse] 258 y_pred_masked = self.predictions[self.mask_rmse] 259 260 fig, ax = plt.subplots(figsize=figsize) 261 262 # Color by competing-risk class when events are multi-risk 263 events_masked = None 264 if self.events is not None: 265 events_masked = np.asarray(self.events)[self.mask_rmse].astype(int) 266 267 if events_masked is not None and events_masked.max(initial=0) > 1: 268 unique_risks = np.sort(np.unique(events_masked[events_masked > 0])) 269 cmap = plt.get_cmap("tab10") 270 for i, k in enumerate(unique_risks): 271 mask = events_masked == k 272 ax.scatter( 273 y_test_masked[mask], 274 y_pred_masked[mask], 275 color=cmap(i % 10), 276 label=f"Risk {int(k)}", 277 alpha=0.8, 278 ) 279 ax.legend(title="Competing risk", loc="best") 280 else: 281 ax.scatter(y_test_masked, y_pred_masked) 282 283 ax.set_xlabel("Actual Values (y_test)") 284 ax.set_ylabel("Predicted Values (y_pred)") 285 ax.set_title("Actual vs Predicted") 286 plt.show() 287 return fig, ax
Scatter plot of actual against predicted times, for the units that failed.
Arguments:
- figsize: The figure's size in inches, as
(width, height).
Returns:
A matplotlib
(fig, ax)tuple, or None if the per-unit predictions were not kept, as for the result of a training job.
def
plot_calibration(self, figsize=(12, 5), show='uncensored'):
289 def plot_calibration(self, figsize=(12, 5), show="uncensored"): 290 """Plot the calibration curve, and the integrated Brier score. 291 292 Args: 293 figsize: The figure's size in inches, as `(width, height)`. 294 show: `"uncensored"` draws the curve over the units that failed; 295 `"all"` draws it over all units, with the uncensored curve 296 beside it. 297 298 Returns: 299 A matplotlib `(fig, axes)` tuple, or None if there is no calibration 300 data. 301 """ 302 if self.calibration_sorted_cdf is None and self.calibration_bias is None: 303 print("No calibration data available.") 304 return None 305 306 import matplotlib.pyplot as plt 307 import numpy as np 308 309 fig, axes = plt.subplots(1, 2, figsize=figsize) 310 fig.suptitle("Model Calibration", fontsize=14) 311 312 ax1 = axes[0] 313 sorted_cdf = self.calibration_sorted_cdf 314 empirical_cdf = self.calibration_empirical_cdf 315 if sorted_cdf is not None and empirical_cdf is not None: 316 sorted_cdf = np.atleast_1d(np.asarray(sorted_cdf, dtype=float)) 317 empirical_cdf = np.atleast_1d(np.asarray(empirical_cdf, dtype=float)) 318 sorted_cdf_events = self.calibration_sorted_cdf_events 319 empirical_cdf_events = self.calibration_empirical_cdf_events 320 if sorted_cdf_events is not None: 321 sorted_cdf_events = np.atleast_1d( 322 np.asarray(sorted_cdf_events, dtype=float) 323 ) 324 if empirical_cdf_events is not None: 325 empirical_cdf_events = np.atleast_1d( 326 np.asarray(empirical_cdf_events, dtype=float) 327 ) 328 ax1.plot([0, 1], [0, 1], "k--", label="Perfect calibration") 329 if show in ("all", "both"): 330 ax1.plot( 331 sorted_cdf, empirical_cdf, "b-", lw=2, label="All observations" 332 ) 333 ax1.fill_between( 334 sorted_cdf, sorted_cdf, empirical_cdf, alpha=0.1, color="blue" 335 ) 336 if sorted_cdf_events is not None and empirical_cdf_events is not None: 337 ax1.plot( 338 sorted_cdf_events, 339 empirical_cdf_events, 340 "r-", 341 lw=2, 342 label="Uncensored only", 343 ) 344 elif show == "uncensored": 345 cdf = sorted_cdf_events if sorted_cdf_events is not None else sorted_cdf 346 ecdf = ( 347 empirical_cdf_events 348 if empirical_cdf_events is not None 349 else empirical_cdf 350 ) 351 label = ( 352 "Uncensored only" 353 if sorted_cdf_events is not None 354 else "All observations" 355 ) 356 ax1.plot(cdf, ecdf, "b-", lw=2, label=label) 357 ax1.fill_between(cdf, cdf, ecdf, alpha=0.15, color="blue") 358 else: 359 ax1.plot( 360 sorted_cdf, empirical_cdf, "b-", lw=2, label="All observations" 361 ) 362 ax1.fill_between( 363 sorted_cdf, sorted_cdf, empirical_cdf, alpha=0.15, color="blue" 364 ) 365 ax1.set_xlabel("Predicted CDF") 366 ax1.set_ylabel("Empirical CDF") 367 ax1.set_title("Calibration Curve") 368 ax1.legend() 369 if ( 370 self.calibration_classification is not None 371 or self.calibration_bias is not None 372 ): 373 parts = [] 374 if self.calibration_classification: 375 parts.append(self.calibration_classification) 376 if self.calibration_bias is not None: 377 parts.append(f"bias={self.calibration_bias:.3f}") 378 ax1.annotate( 379 "\n".join(parts), xy=(0.05, 0.85), xycoords="axes fraction" 380 ) 381 else: 382 ax1.text( 383 0.5, 384 0.5, 385 "No calibration data", 386 ha="center", 387 va="center", 388 transform=ax1.transAxes, 389 ) 390 ax1.set_title("Calibration Curve") 391 392 ax2 = axes[1] 393 if self.ibs is not None: 394 colors = { 395 "excellent": "green", 396 "good": "yellowgreen", 397 "acceptable": "orange", 398 "poor": "red", 399 } 400 color = colors.get(self.ibs_classification or "", "steelblue") 401 ax2.barh(["IBS"], [self.ibs], color=color) 402 ax2.axvline( 403 0.25, color="gray", linestyle="--", label="Random baseline (0.25)" 404 ) 405 ax2.set_xlim(0, max(0.4, self.ibs * 1.2)) 406 ax2.set_title("Integrated Brier Score") 407 ax2.legend() 408 ibs_label = f"{self.ibs:.3f}" 409 if self.ibs_classification: 410 ibs_label += f" — {self.ibs_classification}" 411 ax2.annotate(ibs_label, xy=(self.ibs + 0.01, 0), va="center") 412 else: 413 ax2.text( 414 0.5, 415 0.5, 416 "IBS not available", 417 ha="center", 418 va="center", 419 transform=ax2.transAxes, 420 ) 421 ax2.set_title("Integrated Brier Score") 422 423 plt.tight_layout() 424 return fig, axes
Plot the calibration curve, and the integrated Brier score.
Arguments:
- figsize: The figure's size in inches, as
(width, height). - show:
"uncensored"draws the curve over the units that failed;"all"draws it over all units, with the uncensored curve beside it.
Returns:
A matplotlib
(fig, axes)tuple, or None if there is no calibration data.
class
TuningTrials:
137class TuningTrials: 138 """ 139 The trials of a tuning run, returned by `SurvivalModel.get_tuning_trials()`. 140 141 Iterate over it for the trials in the order they ran; `len()` counts them. 142 Each trial has: 143 144 - `number`: its position in the run, from 0. 145 - `params`: the values drawn from the config's ranges. 146 - `config`: the whole config it trained with. 147 - `value`: its mean score over the folds. Higher is better: for 148 `hpo_metric="ibs"` it is the integrated Brier score negated. 149 - `fold_scores()`: its score on each fold. 150 - `metric`: the `hpo_metric` it was scored by. 151 - `state`: `"COMPLETE"`, or `"FAIL"` if it could not be scored. 152 153 Attributes: 154 best_config: The config of the best trial, with every range replaced by 155 its best value. This is the config `fit()` retrained with. 156 """ 157 158 def __init__(self, trials: list): 159 self._trials = trials 160 self.best_config: dict | None = None 161 162 def __bool__(self): 163 return bool(self._trials) 164 165 def __len__(self): 166 return len(self._trials) 167 168 def __iter__(self): 169 return iter(self._trials) 170 171 def get_best_trial(self): 172 """ 173 The completed trial with the highest mean score. 174 175 Returns: 176 That trial, or None if no trial completed. 177 """ 178 completed = [ 179 t 180 for t in self._trials 181 if t.state == "COMPLETE" and t.value != float("-inf") 182 ] 183 if not completed: 184 return None 185 return max(completed, key=lambda t: t.value)
The trials of a tuning run, returned by SurvivalModel.get_tuning_trials().
Iterate over it for the trials in the order they ran; len() counts them.
Each trial has:
number: its position in the run, from 0.params: the values drawn from the config's ranges.config: the whole config it trained with.value: its mean score over the folds. Higher is better: forhpo_metric="ibs"it is the integrated Brier score negated.fold_scores(): its score on each fold.metric: thehpo_metricit was scored by.state:"COMPLETE", or"FAIL"if it could not be scored.
Attributes:
- best_config: The config of the best trial, with every range replaced by
its best value. This is the config
fit()retrained with.
def
get_best_trial(self):
171 def get_best_trial(self): 172 """ 173 The completed trial with the highest mean score. 174 175 Returns: 176 That trial, or None if no trial completed. 177 """ 178 completed = [ 179 t 180 for t in self._trials 181 if t.state == "COMPLETE" and t.value != float("-inf") 182 ] 183 if not completed: 184 return None 185 return max(completed, key=lambda t: t.value)
The completed trial with the highest mean score.
Returns:
That trial, or None if no trial completed.
def
download_nasa_dataset(save_dir, datasets=None, version='cmapss'):
396def download_nasa_dataset(save_dir, datasets=None, version="cmapss"): 397 """ 398 Download a NASA turbofan dataset, unless its files are already there. 399 400 Files already in `save_dir` are not downloaded again. `"ncmapss"` downloads 401 NASA's whole archive (about 16 GB) whatever `datasets` says, and needs 402 `bsdtar`; `"ncmapss-full"` downloads only the sub-datasets asked for. 403 404 Args: 405 save_dir: The directory to save the files in. 406 datasets: Which sub-datasets to make sure are present; None means all 407 of them. 408 version: Which dataset, as for `get_nasa_train_dataframe()`. 409 """ 410 if version in _NCMAPSS_HOSTED_VERSIONS: 411 _download_ncmapss_hosted( 412 save_dir, 413 datasets if datasets is not None else NCMAPSS_SUBSETS, 414 version, 415 ) 416 return 417 if version == "ncmapss": 418 download_ncmapss_dataset( 419 save_dir, 420 datasets=datasets if datasets is not None else NCMAPSS_SUBSETS, 421 ) 422 return 423 _download_cmapss(save_dir, datasets if datasets is not None else CMAPSS_SUBSETS)
Download a NASA turbofan dataset, unless its files are already there.
Files already in save_dir are not downloaded again. "ncmapss" downloads
NASA's whole archive (about 16 GB) whatever datasets says, and needs
bsdtar; "ncmapss-full" downloads only the sub-datasets asked for.
Arguments:
- save_dir: The directory to save the files in.
- datasets: Which sub-datasets to make sure are present; None means all of them.
- version: Which dataset, as for
get_nasa_train_dataframe().
def
get_nasa_train_dataframe(path, datasets=[1, 2, 3, 4], version='cmapss'):
114def get_nasa_train_dataframe(path, datasets=[1, 2, 3, 4], version="cmapss"): 115 """ 116 Load the training split of a NASA turbofan dataset, ready for `fit()`. 117 118 Download it first with `download_nasa_dataset()`. 119 120 Args: 121 path: The directory the dataset was downloaded to. 122 datasets: Which sub-datasets to include: 1 to 4 for C-MAPSS, 1 to 8 for 123 N-CMAPSS. 124 version: `"cmapss"` for the original C-MAPSS data, or an N-CMAPSS 125 variant: `"ncmapss"` (NASA's archive), `"ncmapss-full"` (the same 126 files, one download per sub-dataset) or `"ncmapss-mini"` (downsampled 127 to one row per flight cycle). In N-CMAPSS, `event` is the failure 128 mode: the sub-dataset's number. 129 130 Returns: 131 A DataFrame with one row per unit and cycle: `unit`, `cycles`, 132 `duration`, `event` and the sensor and setting columns. 133 """ 134 if version in _NCMAPSS_VERSIONS: 135 df = None 136 for i, filename in enumerate(_ncmapss_files_for(version)): 137 if (i + 1) not in datasets: 138 continue 139 h5_path = os.path.join(path, filename) 140 df2 = make_ncmapss_dataframe(h5_path, split="dev") 141 df2[group_col] = i 142 df2[event_col] = i + 1 # failure mode = dataset index (competing risks) 143 if df is not None: 144 df2[unit_col] = df2[unit_col] + df[unit_col].max() 145 df = pd.concat([df, df2], ignore_index=True) 146 gc.collect() 147 return df 148 149 # Original C-MAPSS (text format) 150 df = None 151 for i, file in enumerate( 152 [ 153 os.path.join(path, "train_FD001.txt"), 154 os.path.join(path, "train_FD002.txt"), 155 os.path.join(path, "train_FD003.txt"), 156 os.path.join(path, "train_FD004.txt"), 157 ] 158 ): 159 if (i + 1) not in datasets: 160 continue 161 df2 = make_turbofan_dataframe(file) 162 df2[group_col] = i 163 if df is not None: 164 df2[unit_col] = df2[unit_col] + df[unit_col].max() 165 df = pd.concat([df, df2]) 166 gc.collect() 167 168 train_max_cycles = df.groupby(unit_col)[step_col].max().reset_index() 169 train_max_cycles.columns = [unit_col, "max_cycles"] 170 df = df.merge(train_max_cycles, on=unit_col, how="left") 171 df[duration_col] = df["max_cycles"] - df[step_col] + 1 172 df[event_col] = 1 173 return df.drop(columns=["max_cycles"])
Load the training split of a NASA turbofan dataset, ready for fit().
Download it first with download_nasa_dataset().
Arguments:
- path: The directory the dataset was downloaded to.
- datasets: Which sub-datasets to include: 1 to 4 for C-MAPSS, 1 to 8 for N-CMAPSS.
- version:
"cmapss"for the original C-MAPSS data, or an N-CMAPSS variant:"ncmapss"(NASA's archive),"ncmapss-full"(the same files, one download per sub-dataset) or"ncmapss-mini"(downsampled to one row per flight cycle). In N-CMAPSS,eventis the failure mode: the sub-dataset's number.
Returns:
A DataFrame with one row per unit and cycle:
unit,cycles,duration,eventand the sensor and setting columns.
def
get_nasa_test_dataframe(path, datasets=[1, 2, 3, 4], version='cmapss'):
210def get_nasa_test_dataframe(path, datasets=[1, 2, 3, 4], version="cmapss"): 211 """ 212 Load the test split of a NASA turbofan dataset and its labels. 213 214 The pair is what `fit()` takes as `test_df` and `test_labels_df`. Each test 215 unit's rows stop before it failed, and its label is the remaining life at 216 its last row. 217 218 Args: 219 path: The directory the dataset was downloaded to. 220 datasets: Which sub-datasets to include, as for 221 `get_nasa_train_dataframe()`. 222 version: Which dataset, as for `get_nasa_train_dataframe()`. 223 224 Returns: 225 A tuple `(test_df, test_labels_df)`: the feature rows, one per unit and 226 cycle, and one label row per unit with `unit`, `duration` and `event`. 227 """ 228 if version in _NCMAPSS_VERSIONS: 229 test_df = None 230 label_rows = [] 231 232 for i, filename in enumerate(_ncmapss_files_for(version)): 233 if (i + 1) not in datasets: 234 continue 235 h5_path = os.path.join(path, filename) 236 df2 = make_ncmapss_dataframe(h5_path, split="test") 237 df2 = _truncate_ncmapss_test_split(df2, i) 238 df2[group_col] = i 239 240 if test_df is not None: 241 df2[unit_col] = df2[unit_col] + test_df[unit_col].max() 242 243 # RUL at the cut = label duration; event = failure mode 244 last = df2.groupby(unit_col)[duration_col].last().reset_index() 245 last[event_col] = i + 1 246 label_rows.append(last) 247 248 test_df = pd.concat([test_df, df2], ignore_index=True) 249 gc.collect() 250 251 test_labels_df = pd.concat(label_rows, ignore_index=True) 252 253 # The label belongs to test_labels_df, not to the feature frame. 254 # make_ncmapss_dataframe sets duration from Y, so it arrives here on 255 # both frames; SurvivalModel.predict rejects it as leakage, and the 256 # C-MAPSS branch below never produces it. Drop it so the two versions 257 # return the same shape. 258 test_df = test_df.drop(columns=[duration_col]) 259 260 return test_df, test_labels_df 261 262 test_df = None 263 all_durations = [] 264 265 for i, files in enumerate( 266 [ 267 (os.path.join(path, "test_FD001.txt"), os.path.join(path, "RUL_FD001.txt")), 268 (os.path.join(path, "test_FD002.txt"), os.path.join(path, "RUL_FD002.txt")), 269 (os.path.join(path, "test_FD003.txt"), os.path.join(path, "RUL_FD003.txt")), 270 (os.path.join(path, "test_FD004.txt"), os.path.join(path, "RUL_FD004.txt")), 271 ] 272 ): 273 if (i + 1) not in datasets: 274 continue 275 276 # Load test data for this dataset 277 df2 = make_turbofan_dataframe(files[0]) 278 279 # Load RUL values (one per unit in this dataset) 280 rul_values = np.loadtxt(files[1]) 281 282 # Store RUL values for later 283 all_durations.extend(rul_values) 284 285 # Adjust unit numbers to be globally unique 286 if test_df is not None: 287 df2[unit_col] = df2[unit_col] + test_df[unit_col].max() 288 289 # Concatenate 290 test_df = pd.concat([test_df, df2]) 291 292 # Force garbage collection after each file 293 gc.collect() 294 295 # Convert durations to array 296 durations = np.array(all_durations, dtype=np.float32) 297 298 # All test units eventually fail (event = 1) 299 events = np.ones(len(durations), dtype=np.int32) 300 301 # Create test_labels_df - one row per unit with duration and event 302 test_labels_df = test_df.groupby("unit").last().reset_index()[["unit"]] 303 test_labels_df[duration_col] = durations 304 test_labels_df[event_col] = events 305 306 return test_df, test_labels_df
Load the test split of a NASA turbofan dataset and its labels.
The pair is what fit() takes as test_df and test_labels_df. Each test
unit's rows stop before it failed, and its label is the remaining life at
its last row.
Arguments:
- path: The directory the dataset was downloaded to.
- datasets: Which sub-datasets to include, as for
get_nasa_train_dataframe(). - version: Which dataset, as for
get_nasa_train_dataframe().
Returns:
A tuple
(test_df, test_labels_df): the feature rows, one per unit and cycle, and one label row per unit withunit,durationandevent.