code: def predict_ci(new_df, ci_model, x_mat=None):
      """Predicts the quantiles of the ``offset_col`` (defined in ``ci_model``) in ``new_df``.
  
      Parameters
      ----------
      new_df : `pd.Dataframe`
          Prediction dataframe which minimally includes ``offset_col``
          and ``conditional_cols`` defined in the ``ci_model``.
      ci_model : `dict`
          Returned CI model from ``conf_interval``.
      x_mat : `np.ndarray` or None, default None
          The design matrix from the model fitted on ``new_df``.
          ``x_mat.shape[0]`` must match that of ``new_df``, and
          ``x_mat.shape[1]`` must match the number of features in ``ml_model`` where ``ci_model`` is fitted.
  
      Returns
      -------
      pred_df : `pd.Dataframe`
          A dataframe which includes ``new_df`` and new columns containing
          the quantiles.
      """
  
      ecdf_df = ci_model["ecdf_df"]
      ecdf_df_fallback = ci_model["ecdf_df_fallback"]
      offset_col = ci_model["offset_col"]
      min_admissible_value = ci_model["min_admissible_value"]
      max_admissible_value = ci_model["max_admissible_value"]
      conditional_cols = ci_model["conditional_cols"]
      std_col = ci_model["std_col"]
      quantile_summary_col = ci_model["quantile_summary_col"]
      fall_back_for_all = ci_model["fall_back_for_all"]
      lu_d_sqrt = ci_model["lu_d_sqrt"]
      n_train = ci_model["n_train"]
      x_mean = ci_model["x_train_mean"]
  
      # Computes the scaling factor for prediction interval standard errors.
      # We will use `pi_se_scaler` to scale `quantile_summary_col` and `std_col` columns.
      pi_se_scaler = 1
      # If `x_mat` is provided in case of linear or ridge regression,
      # we know the closed form expression of the prediction intervals.
      # Therefore, we need to scale the prediction intervals, regardless of train / test.
      if x_mat is not None and lu_d_sqrt is not None:
          assert (
              new_df.shape[0] == x_mat.shape[0]
          ), "In `predict_ci`, `new_df` need to be the same length as `x_mat`."
          # `h_mat @ h_mat.T = lu_d_sqrt @ lu_d_sqrt.T`, where `lu_d_sqrt` is p by p and is pre-calculated.
          assert (
              lu_d_sqrt.shape[0] == x_mat.shape[1]
          ), "Feature dimension p must match `lu_d_sqrt` and `x_mat`."
          n_pred, p = x_mat.shape
          X_pred = np.array(x_mat).reshape((n_pred, p))
          # Variance from coefficients contains 3 terms.
          # The first term is true for both linear and ridge regression, while the other two terms are for ridge only.
          # Let's use `L @ L.T` to replace `H @ H.T`,
          # and denote `x_mean_mat = np.repeat(x_mean, repeats=n_pred, axis=0)` (`n_pred` by `p`),
          # then we have the following formulas.
          # Term 1: `X_pred @ L @ L.T @ X_pred.T`.
          # Term 2: `x_mean_mat @ L @ L.T @ X_pred.T` plus its transpose.
          # Term 3: `(1 / n_train + x_mean @ L @ L.T @ x_mean.T) * np.ones((n_pred, n_pred))`.
          # (1) In ridge case, three terms can be simplified to a quadratic form plus a constant matrix:
          #     `(X_pred - x_mean_mat) @ L @ L.T @ (X_pred - x_mean_mat).T + (1 / n_train) * np.ones((n_pred, n_pred))`.
          # (2) In linear case, we only keep term 1: `X_pred @ L @ L.T @ X_pred.T`.
          # Since we only need the diagonal elements, they can be obtained in the following ways.
          # Define `A = L.T @ X_pred.T` (`p` by `n_pred`), then the diagonal of `A.T @ A` is equivalent
          # to `(A ** 2).sum(axis=0)`, length `n_pred`, which is much more efficient to compute.
          if x_mean is not None:  # Ridge.
              # In case of ridge regression, we use `x_mean` to center `X_pred` first.
              X_pred = X_pred - x_mean
          A = lu_d_sqrt.T @ X_pred.T  # `p` by `n_pred`.
          # We are ready to compute the variance scaler,
          # The last term of 1 below is from `np.ones((n_pred,))`, assuming i.i.d. errors with constant variance.
          if x_mean is not None:  # Ridge.
              pi_variance = (A**2).sum(axis=0) + (1 / n_train + 1)
          else:  # Linear.
              pi_variance = (A**2).sum(axis=0) + 1
          pi_se_scaler = np.sqrt(pi_variance)
  
      # Copies ``pred_df`` so that input df to predict is not altered
      pred_df = new_df.reset_index(drop=True)
      pred_df["temporary_overall_dummy"] = 0
      ecdf_df_fallback_dummy = ecdf_df_fallback.copy()
      ecdf_df_fallback_dummy["temporary_overall_dummy"] = 0
      pred_df_fallback = pd.merge(
          pred_df,
          ecdf_df_fallback_dummy,
          on=["temporary_overall_dummy"],
          how="left",
      )
      del pred_df["temporary_overall_dummy"]
      del pred_df_fallback["temporary_overall_dummy"]
  
      if (
          (conditional_cols is None)
          or (conditional_cols == [])
          or fall_back_for_all
      ):
          pred_df_conditional = pred_df_fallback.copy()
      else:
          pred_df_conditional = pd.merge(
              pred_df, ecdf_df, on=conditional_cols, how="left"
          )
  
      # When we have missing in the grouped case (which can happen if a level
      # in ``conditional_cols`` didn't appear in train dataset),
      # we fall back to the overall case.
      for col in [quantile_summary_col, std_col]:
          na_index = pred_df_conditional[col].isnull()
          pred_df_conditional.loc[na_index, col] = pred_df_fallback.loc[
              na_index, col
          ]
  
      # Before offsetting, applies `pi_se_scaler` to all rows in `quantile_summary_col` and `std_col` columns.
      # This scaling only happens in uncertainty model's predict phase (`predict_ci`).
      pred_df_conditional[std_col] *= pi_se_scaler
      # Since values in `quantile_summary_col` are tuples, we need to convert it to `np.array` first.
      quantile_summary = pred_df_conditional[
          quantile_summary_col
      ].apply(lambda value_tuple: np.array(value_tuple))
      quantile_summary *= pi_se_scaler
      pred_df_conditional[
          quantile_summary_col
      ] = quantile_summary.apply(lambda array: tuple(array))
      ci_model["pi_se_scaler"] = pi_se_scaler
  
      # Offsets the values in `distribution_col` by `offset_col`.
      if offset_col is None:
          pred_df_conditional[
              QUANTILE_SUMMARY_COL
          ] = pred_df_conditional[quantile_summary_col]
      else:
          pred_df_conditional[QUANTILE_SUMMARY_COL] = offset_tuple_col(
              df=pred_df_conditional,
              offset_col=offset_col,
              tuple_col=quantile_summary_col,
          )
  
      pred_df_conditional = limit_tuple_col(
          df=pred_df_conditional,
          tuple_col=QUANTILE_SUMMARY_COL,
          lower=min_admissible_value,
          upper=max_admissible_value,
      )
  
      # Only returns needed cols.
      returned_cols = [QUANTILE_SUMMARY_COL, std_col]
      if conditional_cols is not None:
          returned_cols = conditional_cols + returned_cols
  
      pred_df[returned_cols] = pred_df_conditional[returned_cols]
      # Standardizes `std_col` column name.
      pred_df.rename(columns={std_col: ERR_STD_COL}, inplace=True)
  
      return pred_df
  
file_path: greykite/greykite/algo/uncertainty/conditional/conf_interval.py
imports: import warnings
  import numpy
  import pandas
  import scipy
  from greykite.algo.uncertainty.conditional.dataframe_utils import limit_tuple_col
  from greykite.algo.uncertainty.conditional.dataframe_utils import offset_tuple_col
  from greykite.algo.uncertainty.conditional.estimate_distribution import estimate_empirical_distribution
  from greykite.algo.uncertainty.conditional.normal_quantiles import normal_quantiles_df
  from greykite.common.constants import ERR_STD_COL
  from greykite.common.constants import QUANTILE_SUMMARY_COL
label_desc: Use image node_id 2 for calling a global function with example usage: predict_ci(new_df, ci_model, x_mat) and returns: pred_df