code: def effect_by_segment(
      df: pd.DataFrame,
      treatment: str,
      outcome: str,
      prediction: str,
      segments: int = 10,
      effect_fn: EffectFnType = linear_effect,
  ) -> pd.Series:
      """
      Segments the dataset by a prediction's quantile and estimates the treatment effect by segment.
  
      Parameters
      ----------
      df : Pandas' DataFrame
          A Pandas' DataFrame with target and prediction scores.
  
      treatment : Strings
          The name of the treatment column in `df`.
  
      outcome : Strings
          The name of the outcome column in `df`.
  
      prediction : Strings
          The name of the prediction column in `df`.
  
      segments : Integer
          The number of the segments to create. Uses Pandas' qcut under the hood.
  
      effect_fn : function (df: pandas.DataFrame, treatment: str, outcome: str) -> int or Array of int
          A function that computes the treatment effect given a dataframe, the name of the treatment column and the name
          of the outcome column.
  
      Returns
      ----------
      effect by band : Pandas' Series
          The effect stored in a Pandas' series were the indexes are the segments
      """
  
      effect_fn_partial = partial(
          effect_fn, treatment_column=treatment, outcome_column=outcome
      )
      return (
          df.assign(
              **{
                  f"{prediction}_band": pd.qcut(
                      df[prediction], q=segments
                  )
              }
          )
          .groupby(f"{prediction}_band")
          .apply(effect_fn_partial)
      )
  
file_path: fklearn/src/fklearn/causal/validation/curves.py
imports: from typing import List
  import numpy
  import pandas
  from toolz import curry, partial
  from fklearn.types import EffectFnType
  from fklearn.causal.effects import linear_effect
label_desc: Use image node_id 1 for calling a global function with example usage: effect_by_segment(df, treatment, outcome, prediction, segments, effect_fn) and returns: df