code: """
  Estimates covariance matrix for KW94 Dataset 1 with
  Simulated Max. Likelihood.
  
  """
  import os
  
  import numpy as np
  import pandas as pd
  import respy as rp
  from estimagic.differentiation.differentiation import jacobian
  from estimagic.inference.likelihood_covs import cov_jacobian
  from estimagic.optimization.optimize import maximize
  
  
  def jac_estimation_chol(save=False):
      """
      Estimates covariance matrix for KW94 Dataset 1 with Simulated Max. Likelihood.
      The Jacobian matrix is used instead of Hessian because it yields no inversion
      error.
      The parameters contain Cholesky factors instead of SD-Corr-Matrix because
      these factors are unconstrained. Therefore, their distribution can be estimated
      by an unconstrained normal distribution.
  
      Parameters
      ----------
      save : bool
          Indicates wether to save data.
  
      Returns
      -------
      par_estimates_chol_df : DataFrame
          Contains the estimates parameters and the not estimates fixed parameters
          in respy format.
      rand_par_chol_df : DataFrame
          Df containing variable parameters, SDs and lower and upper bound in
          estimagic format. It can be post-processed with surface/topography plot.
      cov_chol_df : DataFrame
          Df containing the covariance matrix.
      corr_chol_df : DataFrame
          DF containing the correlation matrix.
  
      """
      # Df is sample of 1000 agents in 40 periods.
      sim_params_sdcorr, options, df = rp.get_example_model("kw_94_one")
  
      # Write params in terms of Cholesky factors instead of SD-Corr-matrix.
      # This transformation holds only true for the parametrization in KW94 Dataset 1.
      # Simply change SD-Corr indices to cholesky indices.
      sim_params_chol = chol_reindex_params(sim_params_sdcorr)
  
      # Estimate parameters.
      # log_like = log_like_obs.mean(). Used for consistency with optimizers.
      # Gives log-likelihood function for mean agent.
      crit_func = rp.get_crit_func(sim_params_chol, options, df, "log_like")
  
      # Get constraint for parameter estimation
      constr = rp.get_parameter_constraints("kw_94_one")
      # Kick out constraints for SD-Corr-Matrix. Cholesky factors are unconstrained.
      constr_chol = constr[1:4]
  
      _, par_estimates_chol_df = maximize(
          crit_func,
          sim_params_chol,
          "scipy_L-BFGS-B",
          db_options={"rollover": 200},
          algo_options={"maxfun": 1},
          constraints=constr_chol,
          dashboard=False,
      )
  
      # df  will take lower and upper bounds after standard error esitmation
      # so that cols fit topography plot requirements.
      rand_par_chol_df = pd.DataFrame(
          data=par_estimates_chol_df["value"].values[:27],
          index=par_estimates_chol_df[:27].index,
          columns=["value"],
      )
  
      # The rest of this function estimates the variation of the estimates.
      # Log-likelihood function for sample of agents.
      log_like_obs_func = rp.get_crit_func(
          par_estimates_chol_df, options, df, version="log_like_obs"
      )
  
      # Jacobian matrix.
      jacobian_matrix = jacobian(
          log_like_obs_func, par_estimates_chol_df, extrapolation=False
      )
  
      # Drop zero lines to avoid multicollinearity for matrix inversion.
      jacobian_matrix = jacobian_matrix.loc[:, (jacobian_matrix != 0).any(axis=0)]
  
      jacobian_cov_matrix = cov_jacobian(jacobian_matrix.to_numpy())
  
      jacobian_cov_matrix = cov_jacobian(jacobian_matrix.to_numpy())
  
      cov_chol_df = pd.DataFrame(
          data=jacobian_cov_matrix,
          index=par_estimates_chol_df[:27].index,
          columns=par_estimates_chol_df[:27].index,
      )
  
      corr_chol_df = cov_chol_df.copy(deep=True)
      for i in range(0, len(cov_chol_df)):
          for j in range(0, len(cov_chol_df)):
              corr_chol_df.iloc[i, j] = cov_chol_df.iloc[i, j] / (
                  np.sqrt(cov_chol_df.iloc[i, i] * cov_chol_df.iloc[j, j])
              )
  
      assert -1 <= corr_chol_df.values.any() <= 1, "Corrs must be inside [-1,1]"
  
      # Estimate parameters.
      # log_like = log_like_obs.mean(). Used for consistency with optimizers.
      # Gives log-likelihood function for mean agent.
      crit_func = rp.get_crit_func(par_estimates_chol_df, options, df, "log_like")
  
      constr = rp.get_parameter_constraints("kw_94_one")
      # Kick out constraints for SD-Corr-Matrix. Cholesky factors are unconstrained.
      constr_chol = constr[1:4]
  
      # Include upper and lower bounds to par_df for surface/topography plot.
      rand_par_chol_df["sd"] = np.sqrt(np.diag(jacobian_cov_matrix))
      rand_par_chol_df["lower"] = rand_par_chol_df["value"] - 2 * rand_par_chol_df["sd"]
      rand_par_chol_df["upper"] = rand_par_chol_df["value"] + 2 * rand_par_chol_df["sd"]
  
      # Define the script path relative to the jupyter notebook that calls the script.
      abs_dir = os.path.dirname(__file__)
      if save is True:
          # Contains 3 fixed respy parameters.
          par_estimates_chol_df.to_pickle(
              os.path.join(abs_dir, "input/est_rp_params_chol.pkl")
          )
          # Contains only flexible parametes. Can be used for surface/topography plot.
          rand_par_chol_df.to_pickle(
              os.path.join(abs_dir, "input/est_rand_params_chol.pkl")
          )
          cov_chol_df.to_pickle(os.path.join(abs_dir, "input/est_cov_chol.pkl"))
          corr_chol_df.to_pickle(os.path.join(abs_dir, "input/est_corr_chol.pkl"))
      else:
          pass
  
      return par_estimates_chol_df, rand_par_chol_df, cov_chol_df, corr_chol_df
  
  
  def chol_reindex_params(params_sdcorr):
      """
      Creates the params Df with Cholesky factors and the right indices for
      respy. This transformation holds only true for the parametrization
      in KW94 Dataset 1.
      Thus, this function simply changes SD-Corr indices to cholesky indices.
      Without the slicing and merging, index ('maximum_exp', 'edu') yields
      an uniqueness error for the second index when (..., 'sd_edu') is set to
      (..., 'edu'). Yet, because we have double_indices the indices ARE unique.
  
      Parameters
      ----------
      params_sdcorr : DataFrame
          Parameters DataFrame in respy format with SD-Corr matrix elements
  
      Returns
      -------
      params_chol : DataFrame
          Parameters DataFrame in respy format with matrix elements from Choleksy
          decomposition of covariance matrix that underlies the SD-Corr matrix.
  
      """
      p_chol_slice = params_sdcorr.iloc[17:27, :]
      # Remove unused inherited index levels.
      p_chol_slice.index = p_chol_slice.index.remove_unused_levels()
      # Use the SPECIFIC property of Dataset 1 in KW94 where SD-Corr-Matrix
      # equals Cholesky maxtrix.
      # This mean we just need to, firstly, rename the first index.
      p_chol_slice.index = p_chol_slice.index.set_levels(
          p_chol_slice.index.levels[0].str.replace("shocks_sdcorr", "shocks_chol"),
          level=0,
      )
  
      # And secondly we need to convert the second index to respy cholesky format.
      dic = {"sd": "chol", "corr": "chol"}
      for i, j in dic.items():
          p_chol_slice.index = p_chol_slice.index.set_levels(
              p_chol_slice.index.levels[1].str.replace(i, j), level=1
          )
  
      # Insert params_chol with index in params_sdcorr by merging slices.
      part_1 = params_sdcorr.iloc[0:17, :]
      part_1.index = part_1.index.remove_unused_levels()
      part_3 = params_sdcorr.iloc[27:31, :]
      part_3.index = part_3.index.remove_unused_levels()
  
      parts = [part_1, p_chol_slice, part_3]
      params_chol = pd.concat(parts)
  
      return params_chol
  
  
  if __name__ == "__main__":
      jac_estimation_chol(save=True)
  
apis: numpy.sqrt pandas.DataFrame pandas.concat numpy.diag