code: def marginal_expectation(
      prediction_method: Callable[[np.ndarray], np.ndarray],
      feature_samples: np.ndarray,
      baseline_samples: np.ndarray,
      baseline_feature_indices: List[int],
      return_averaged_results: bool = True,
      feature_perturbation: str = "randomize_columns_jointly",
      max_batch_size: int = -1,
  ) -> np.ndarray:
      """Estimates the marginal expectation for samples in baseline_noise_samples when randomizing features that are not
      part of baseline_feature_indices. This is, this function estimates
          y^i = E[Y | do(x^i_s)] := \\int_x_s' E[Y | x^i_s, x_s'] p(x_s') d x_s',
      where x^i_s is the i-th sample from baseline_noise_samples, s denotes the baseline_feature_indices and
      x_s' ~ X_s' denotes the randomized features that are not in s. For an approximation of the integral, the given
      prediction_method is evaluated multiple times for the same x^i_s, but different x_s' ~ X_s'.
  
      :param prediction_method: Prediction method of interest. This should expect a numpy array as input for making
      predictions.
      :param feature_samples: Samples from the joint distribution. These are used for randomizing the features that are not in
                              baseline_feature_indices.
      :param baseline_samples: Samples for which the marginal expectation should be estimated.
      :param baseline_feature_indices: Column indices of the features in s. These values for these features are remain constant
                                       when estimating the expectation.
      :param return_averaged_results: If set to True, the expectation over all evaluated samples for the i-th
      baseline_noise_samples is returned. If set to False, all corresponding results for the i-th sample are returned.
      :param feature_perturbation: Type of feature permutation:
          'randomize_columns_independently': Each feature not in s is randomly permuted separately.
          'randomize_columns_jointly': All features not in s are jointly permuted. Note that this still represents an
          interventional distribution.
      :param max_batch_size: Maximum batch size for a estimating the predictions. This has a significant influence on the
      overall memory usage. If set to -1, all samples are used in one batch.
      :return: If return_averaged_results is False, a numpy array where the i-th entry belongs to the marginal expectation
      of x^i_s when randomizing the remaining features.
      If return_averaged_results is True, a two dimensional numpy array where the i-th entry contains all
      predictions for x^i_s when randomizing the remaining features.
      """
      feature_samples, baseline_samples = shape_into_2d(
          feature_samples, baseline_samples
      )
  
      batch_size = (
          baseline_samples.shape[0]
          if max_batch_size == -1
          else max_batch_size
      )
      result = [np.nan] * baseline_samples.shape[0]
  
      # Make copy to avoid manipulating the original matrix.
      feature_samples = np.array(feature_samples)
  
      features_to_randomize = np.delete(
          np.arange(0, feature_samples.shape[1]),
          baseline_feature_indices,
      )
  
      if feature_perturbation == "randomize_columns_independently":
          feature_samples = permute_features(
              feature_samples, features_to_randomize, False
          )
      elif feature_perturbation == "randomize_columns_jointly":
          feature_samples = permute_features(
              feature_samples, features_to_randomize, True
          )
      else:
          raise ValueError(
              "Unknown argument %s as feature_perturbation type!"
              % feature_perturbation
          )
  
      # The given prediction method has to be evaluated multiple times on a large amount of different inputs. Typically,
      # the batch evaluation of a prediction model on multiple inputs at the same time is significantly faster
      # than evaluating it on single simples in a for-loop. To make use of this, we try to evaluate as many samples as
      # possible in one batch call of the prediction method. However, this also requires a lot of memory for many samples.
      # To overcome potential memory issues, multiple batch calls are performed, each with at most batch_size many
      # samples. The number of samples that are evaluated is normally
      # baseline_noise_samples.shape[0] * feature_samples.shape[0]. Here, we reduce it to
      # batch_size * feature_samples.shape[0]. If the batch_size would be set 1, then each baseline_noise_samples is
      # evaluated one by one in a for-loop.
      inputs = repmat(feature_samples, batch_size, 1)
      for offset in range(0, baseline_samples.shape[0], batch_size):
          # Each batch consist of at most batch_size * feature_samples.shape[0] many samples. If there are multiple
          # batches, the offset indicates the index of the current baseline_noise_samples that has not been evaluated yet.
          if offset + batch_size > baseline_samples.shape[0]:
              # If the batch size would be larger than the remaining amount of samples, it is reduced to only include the
              # remaining baseline_noise_samples.
              adjusted_batch_size = baseline_samples.shape[0] - offset
              inputs = inputs[
                  : adjusted_batch_size * feature_samples.shape[0]
              ]
          else:
              adjusted_batch_size = batch_size
  
          for index in range(adjusted_batch_size):
              # The inputs consist of batch_size many copies of feature_samples. Here, we set the columns of the features
              # in baseline_feature_indices to their respective values in baseline_noise_samples.
              inputs[
                  index
                  * feature_samples.shape[0] : (index + 1)
                  * feature_samples.shape[0],
                  baseline_feature_indices,
              ] = baseline_samples[
                  offset + index, baseline_feature_indices
              ]
  
          # After creating the (potentially large) input data matrix, we can evaluate the prediction method.
          predictions = np.array(prediction_method(inputs))
  
          for index in range(adjusted_batch_size):
              # Here, offset + index now indicates the sample index in baseline_noise_samples.
              if return_averaged_results:
                  # This would average all prediction results obtained for the 'offset + index'-th sample in
                  # baseline_noise_samples. This is, y^(offset + index) = E[Y | do(x^(offset + index)_s)].
                  result[offset + index] = np.mean(
                      predictions[
                          index
                          * feature_samples.shape[0] : (index + 1)
                          * feature_samples.shape[0]
                      ],
                      axis=0,
                  )
              else:
                  # This would return all prediction results obtained for the 'offset + index'-th sample in
                  # baseline_noise_samples, i.e. the results are not averaged.
                  result[offset + index] = predictions[
                      index
                      * feature_samples.shape[0] : (index + 1)
                      * feature_samples.shape[0]
                  ]
  
      return np.array(result)
  
file_path: dowhy/dowhy/gcm/stats.py
imports: from typing import Callable, List, Optional, Union
  import numpy
  from numpy.matlib import repmat
  from scipy import stats
  from sklearn.linear_model import LinearRegression
  from dowhy.gcm.constant import EPS
  from dowhy.gcm.util.general import shape_into_2d
label_desc: Use image node_id 2 for calling a global function with example usage: marginal_expectation(prediction_method, feature_samples, baseline_samples, baseline_feature_indices, return_averaged_results, feature_perturbation, max_batch_size) and returns: np