text: # ./experiments/strength_of_effect_experiments.py
  """ Experiments with polynomial synthetic data generation.
  
  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. 
  SPDX-License-Identifier: Apache-2.0
  """
  from functools import partial
  import numpy as np
  from typing import Callable, Dict, Tuple
  from tqdm import trange
  
  from experiments.synthetic_data_generation import *
  from src.cmaxent.maxent import *
  from src.utils.plot_utils import *
  from src.utils.utils import *
  
  # Globals
  np.set_printoptions(precision=2)
  
  
  def strength_experiment(data_generating_function: Callable, plot_dir: str) -> Dict:
      """ Run CMAXENT ROC curve.
      """
      threshold = 0.15
      results = {}
  
      # Function used in the estimation of relative difference:
      vectorized_relative_difference = np.vectorize(relative_difference)
  
      # Arrays to save results
      horizontal = []
      known = []
      unknown = []
      
      for i in trange(200):
          data = data_generating_function(seed=i)
          
          # Example using simulated data and joint maxent:
          joint = JointMAXENT(x=data["x"],
                              moment_dict=data["moment_constraints_joint"],
                              seed=i)
          joint.fit()
          
          if not joint.solution.success:
              print("Optimization for Joint MaxEnt was not successful")
              continue
  
          # Conditional MAXENT with learned distribution over the causes
          conditional = ConditionalMAXENT(x = data["y"] + data["x"],
                                              moment_dict=data["moment_constraints_conditional"],
                                              probability_x=joint.compute_exponential_family(joint.multipliers),
                                              seed=i)
  
          conditional.fit()
          if not conditional.solution.success:
              print("Optimization for Conditional MaxEnt (with learned p(cause)) was not successful")
              continue
  
          # Conditional MAXENT with true probability over the causes.
          conditional_known_px = ConditionalMAXENT(x = data["y"] + data["x"],
                                                  moment_dict=data["moment_constraints_conditional"],
                                                  probability_x=data["px"],
                                                  seed=i)
  
          conditional_known_px.fit()
          if not conditional_known_px.solution.success:
              print("Optimization for Conditional MaxEnt (with true p(cause)) was not successful")
              continue
          
  
          # ACE value for this dataset, this are the horizontal locations on the plot
          horizontal.append(float(ace(data, 1, 1, [1,0])))
  
          # For unknown distribution of causes
          multipliers = conditional.multipliers
          estimators = vectorized_relative_difference(multipliers[1::2], multipliers[::2])
          unknown.append(estimators[0]>threshold)
  
          # For known distribution of causes
          multipliers = conditional_known_px.multipliers
          estimators = vectorized_relative_difference(multipliers[1::2], multipliers[::2])
          known.append(estimators[0]>threshold)
  
      nan_loc_horizontal = np.where(np.isnan(horizontal))[0]
      results["horizontal"] = np.delete(horizontal, nan_loc_horizontal)
      results["unknown"] = np.delete(unknown, nan_loc_horizontal)
      results["known"] = np.delete(known, nan_loc_horizontal)
  
      np.save(f"./experiments/results/rate_vs_ace_{data_generating_function.func.__name__}", results)
  
      plot_strength_scm(results, plot_dir)
      
      return results
  
  
  def main():
      """ Main function for the polynomial synthetic data experiment.
      """
      ### Experiments with polynomials:
      partial_first = partial(scm_polynomial_thresholding_first, always_include=[0])
      _ = strength_experiment(partial_first, "./experiments/results/rate_vs_strength_first.png")
      partial_second = partial(scm_polynomial_thresholding_second, always_include=[0])
      _ = strength_experiment(partial_second, "./experiments/results/rate_vs_strength_second.png")
      partial_third = partial(scm_polynomial_thresholding_third, always_include=[0])
      _ = strength_experiment(partial_third, "./experiments/results/rate_vs_strength_third.png")
      partial_fourth = partial(scm_polynomial_thresholding_fourth, always_include=[0])
      _ = strength_experiment(partial_fourth, "./experiments/results/rate_vs_strength_fourth.png")
  
  
  if __name__ == "__main__":
      main()
  