<jupyter_start><jupyter_text>SINGLE SUBJECT<jupyter_code>def make_parameter_recov_data_single_subj(n_param_sets = 1000,
                                          n_samples_per_subject = 100,
                                          model = 'ddm_par2_no_bias',
                                          save = True,
                                          save_folder = 'data_parameter_recov',
                                          print_freq = 10):
    
    data_dict = {}
    for i in range(n_param_sets):
        if i % print_freq == 0:
            print(i)
        data_dict[i] = {}
        data_tmp, full_parameter_dict_tmp = hddm.simulators.hddm_dataset_generators.simulator_h_c(n_subjects = 1,
                                                                                                  n_trials_per_subject = n_samples_per_subject,
                                                                                                  model = model,
                                                                                                  p_outlier = 0.00,
                                                                                                  conditions = None,
                                                                                                  depends_on = None,
                                                                                                  regression_models = None,
                                                                                                  regression_covariates = None,
                                                                                                  group_only_regressors = False,
                                                                                                  group_only = None,
                                                                                                  fixed_at_default = None) #['z'])

        data_dict[i]['data'] = data_tmp
        data_dict[i]['gt_param_dict'] = full_parameter_dict_tmp

    if not os.path.isdir(save_folder):
        os.mkdir(save_folder)
    
    pickle.dump(data_dict, open(save_folder + '/' + model + '_npsets_' + str(n_param_sets) + '_ns_' + str(n_samples_per_subject) + '.pickle', 'wb'))
    
    return 'Finished generating data'


for nps in [100, 200, 400, 800, 1600]:
    for model in ['ddm_par2_no_bias', 'ddm_seq2_no_bias', 'ddm_mic2_adj_no_bias']:
        make_parameter_recov_data_single_subj(n_param_sets = 1000,
                                              n_samples_per_subject = nps,
                                              model = model,
                                              save = True,
                                              save_folder = 'data_parameter_recov/gt_data',
                                              print_freq = 100)<jupyter_output>0
new round of data simulation because parameter bounds where violated
new round of data simulation because parameter bounds where violated
new round of data simulation because parameter bounds where violated
new round of data simulation because parameter bounds where violated
new round of data simulation because parameter bounds where violated
new round of data simulation because parameter bounds where violated
100
new round of data simulation because parameter bounds where violated
200
new round of data simulation because parameter bounds where violated
new round of data simulation because parameter bounds where violated
new round of data simulation because parameter bounds where violated
new round of data simulation because parameter bounds where violated
300
new round of data simulation because parameter bounds where violated
new round of data simulation because parameter bounds where violated
new round of data simulation because parameter bounds where violated
new round of data si[...]<jupyter_text>CHONG DATA  MODEL FORMULATION:<jupyter_code>os.path.isdir('data_parameter_recov/gt_data')
# Simulate DATA
model = 'ddm_par2_no_bias'
# Test Regressor
n_samples_per_subject = 1000
n_subjects = None

data, full_parameter_dict = hddm.simulators.hddm_dataset_generators.simulator_h_c(n_subjects = n_subjects,
                                                                                  n_trials_per_subject = n_samples_per_subject,
                                                                                  model = model,
                                                                                  p_outlier = 0.00,
                                                                                  conditions = None,
                                                                                  depends_on = None,
                                                                                  regression_models = None,
                                                                                  regression_covariates = None,
                                                                                  group_only_regressors = False,
                                                                                  group_only = None,
                                                                                  fixed_at_default = None) #['z'])
full_parameter_dict<jupyter_output><empty_output><jupyter_text>GET CHONG DATA METADATA<jupyter_code>chong_data = pd.read_csv('data/chong_data_hddm_ready.csv', index_col = 0)
chong_data = chong_data.drop(['rxtime','Unnamed: 19', 'subj'], axis = 1)
np.allclose(chong_data['rxtime'].values, chong_data['rt'].values)
chong_data[['subj', 'subj_idx']]
chong_data.keys()
covariate_names = dmatrix('C(cond) + cond2', chong_data)
covariate_names
                separator = reg_model.find("~")
                outcome = reg_model[:separator].strip(" ")
                reg_model_stripped = reg_model[(separator + 1) :]

                if 

                # Run through patsy dmatrix to get the covariate names 
                # that patsy assigns !
                covariate_names = dmatrix(
                    reg_model_stripped, cov_df
                ).design_info.column_names

chong_data = pd.read_csv
def real_dat_to_sim(data = None,
                    model = None,
                    depends_on = {}):
    for param in hddm.model_config.model_config[model]['params']:
        data[param] = 0

    full_param_set = set(hddm.model_config.model_config[model]['params'])
    remainder = full_param_set - list(depends_on.keys())

    for remainder_tmp in remainder:
        
    

    
def make_conditions_df(conditions=None):
    """
    Makes a dataframe out of the supplied condition dictionary, that stores each combination as a row.
    """
    arg_tuple = tuple([conditions[key] for key in conditions.keys()])
    condition_rows = np.meshgrid(*arg_tuple)
    return pd.DataFrame(
                        np.column_stack([x_tmp.flatten() for x_tmp in condition_rows]),
                        columns=[key for key in conditions.keys()],
                        )
df_test = make_conditions_df(conditions={
        "c_one": ["high", "low"],
        "c_two": ["high", "low"],
        "c_three": ["high", "medium", "low"],
    })
df_test[['c_one', 'c_two']].drop_duplicates()
tmp_cond = chong_data[['isHighCorrect', 'isLowCorrect']].drop_duplicates()
tmp_cond
for cond_id in range(tmp_cond.shape[0]):
    bool_  = np.ones(chong_data.shape[0]).astype(np.bool)
    for key_ in tmp_cond.keys():
        bool_ = bool_ & (chong_data[key_] == tmp_cond.iloc[cond_id][key_])
    print(np.sum(bool_))
np.sum(bool_)



chong_data.loc[chong_data['subj_idx'] == 0, ['isHighCorrect']]
def simulator_h_c(data = None,
    n_subjects=10,
    n_trials_per_subject=100,
    model="ddm_vanilla",
    conditions={
        "c_one": ["high", "low"],
        "c_two": ["high", "low"],
        "c_three": ["high", "medium", "low"],
    },
    depends_on={"v": ["c_one", "c_two"]},
    regression_models=["z ~ covariate_name"],
    regression_covariates={
        "covariate_name": {"type": "categorical", "range": (0, 4)}
    },  # need this to make initial covariate matrix from which to use dmatrix (patsy)
    group_only_regressors=True,
    group_only=["z"],
    fixed_at_default=["t"],
    p_outlier=0.0,
    outlier_max_t=10.0,
    **kwargs,
):

    """Flexible simulator that allows specification of models very similar to the hddm model classes.

    :Arguments:
        data: pd.DataFrame <default=None>
            Actual covariate dataset. If data is supplied its covariates are used instead of generated.
        n_subjects: int <default=5>
            Number of subjects in the datasets
        n_trials_per_subject: int <default=500>
            Number of trials for each subject
        model: str <default = 'ddm_vanilla'>
            Model to sample from. For traditional hddm supported models, append '_vanilla' to the model. Omitting 'vanilla'
            imposes constraints on the parameter sets to not violate the trained parameter space of our LANs.
        conditions: dict <default={'c_one': ['high', 'low'], 'c_two': ['high', 'low'], 'c_three': ['high', 'medium', 'low']}>
            Keys represent condition relevant columns, and values are lists of unique items for each condition relevant column.
        depends_on: dict <default={'v': ['c_one', 'c_two']}>
            Keys specify model parameters that depend on the values --> lists of condition relevant columns.
        regression_models: list or strings <default = ['z ~ covariate_name']>
            Specify regression model formulas for one or more dependent parameters in a list.
        regression_covariates: dict <default={'covariate_name': {'type': 'categorical', 'range': (0, 4)}}>
            Dictionary in dictionary. Specify the name of the covariate column as keys, and for each key supply the 'type' (categorical, continuous) and
            'range' ((lower bound, upper bound)) of the covariate.
        group_only_regressors: bin <default=True>
            Should regressors only be specified at the group level? If true then only intercepts are specified subject wise.
            Other covariates act globally.
        group_only: list <default = ['z']>
            List of parameters that are specified only at the group level.
        fixed_at_default: list <default=['t']>
            List of parameters for which defaults are to be used. These defaults are specified in the model_config dictionary, which you can access via: hddm.simulators.model_config.
        p_outlier: float <default = 0.0>
            Specifies the proportion of outliers in the data.
        outlier_max_t: float <default = 10.0>
            Outliers are generated from np.random.uniform(low = 0, high = outlier_max_t) with random choices.
    Returns:
        (pandas.DataFrame, dict): The Dataframe holds the generated dataset, ready for constuction of an hddm model. The dictionary holds the groundtruth parameter (values) and parameter names (keys). Keys match
                                  the names of traces when fitting the equivalent hddm model. The parameter dictionary is useful for some graphs, otherwise not neccessary.
    """

    meta_params = {'group_param_dist': 'normal',
                   'gen_norm_std': 1 / 3,
                   'uniform_buffer': 1 / 5,
                   'gen_std_std': 1 / 8,
                   'covariate_range': 1 / 4
                  } 

    for key_ in kwargs.keys():
        meta_params[key] = kwargs[key]


    def check_params(data=None, model=None, is_nn=True):
        """
        Function checks if parameters are within legal bounds
        """
        for key in data.keys():
            if key in model_config[model]["params"]:
                if (
                    np.sum(
                        data[key]
                        < model_config[model]["param_bounds"][0][
                            model_config[model]["params"].index(key)
                        ]
                    )
                    > 0
                ):
                    return 0
                elif (
                    np.sum(
                        data[key]
                        > model_config[model]["param_bounds"][1][
                            model_config[model]["params"].index(key)
                        ]
                    )
                    > 0
                ):
                    return 0
        return 1

    def get_parameter_remainder(
        regression_models=None, group_only=None, depends_on=None, fixed_at_default=None
    ):

        """
        The arguments supplied to the simulator implicitly specify how we should handle a bunch of model parameters.
        If there remain model parameters that did not receive implicit instructions, we call these 'remainder' parameters
        and sample them randomly for our simulations.
        """

        # Add subject parameters to full_parameter_dict
        total_param_list = model_config[model]["params"]
        params_utilized = []

        # Regression Part
        # reg_df = make_covariate_df(regression_covariates, n_trials_per_subject)
        if regression_models is not None:
            for regression_model in regression_models:
                separator = regression_model.find("~")
                assert separator != -1, "No outcome variable specified."
                params_utilized += regression_model[:separator].strip(" ")

        # Group only Part
        if group_only is not None:
            params_utilized += group_only

        # Fixed Part
        if fixed_at_default is not None:
            params_utilized += fixed_at_default

        # Depends on Part
        if depends_on is not None:
            for depends_on_key in depends_on.keys():
                params_utilized += [depends_on_key]

        params_utilized = list(set(params_utilized))

        # Rest of Params
        remainder = set(total_param_list) - set(params_utilized)

        return remainder

    def make_covariate_df(regression_covariates, n_trials_per_subject):
        """
        Goes through the supplied covariate data, and turns it into a dataframe, with randomly generated covariate values.
        Each column refers to one covariate.
        """

        cov_df = pd.DataFrame(
            np.zeros((n_trials_per_subject, len(list(regression_covariates.keys())))),
            columns=[key for key in regression_covariates.keys()],
        )

        for covariate in regression_covariates.keys():
            tmp = regression_covariates[covariate]
            if tmp["type"] == "categorical":
                cov_df[covariate] = (
                    np.random.choice(
                        np.arange(tmp["range"][0], tmp["range"][1] + 1, 1),
                        replace = True,
                        size = n_trials_per_subject,
                    )
                    / (tmp["range"][1])
                )
            else:
                cov_df[covariate] = np.random.uniform(
                    low=tmp["range"][0], high=tmp["range"][1], size=n_trials_per_subject
                ) / (tmp["range"][1] - tmp["range"][0])

        return cov_df

    def make_conditions_df(conditions=None):
        """
        Makes a dataframe out of the supplied condition dictionary, that stores each combination as a row.
        """
        arg_tuple = tuple([conditions[key] for key in conditions.keys()])
        condition_rows = np.meshgrid(*arg_tuple)
        return pd.DataFrame(
                            np.column_stack([x_tmp.flatten() for x_tmp in condition_rows]),
                            columns=[key for key in conditions.keys()],
                            )

    def make_single_sub_cond_df(
        conditions_df,
        depends_on,
        regression_models,
        regression_covariates,
        group_only_regressors,
        group_only,
        fixed_at_default,
        remainder,
        model,
        group_level_parameter_dict,
        n_subjects,
        n_trials_per_subject,
        ):

        # Construct subject data
        full_parameter_dict = group_level_parameter_dict.copy()

        # Subject part -----------------------
        full_data = []
        # Condition --------------------------
        if conditions_df is None:
            n_conditions = 1
        else:
            n_conditions = conditions_df.shape[0]

        for condition_id in range(n_conditions):
            remainder_set = 0
            regressor_set = 0

            for subj_idx in range(n_subjects):
                # Parameter vector
                subj_data = pd.DataFrame(index = np.arange(0, n_trials_per_subject, 1))
                subj_data["subj_idx"] = str(subj_idx)

                # Fixed part
                if fixed_at_default is not None:
                    for fixed_tmp in fixed_at_default:
                        subj_data[fixed_tmp] = group_level_parameter_dict[fixed_tmp]

                # Group only part
                if group_only is not None:
                    for group_only_tmp in group_only:
                        if group_only_tmp in list(depends_on.keys()):
                            pass
                        else:
                            subj_data[group_only_tmp] = group_level_parameter_dict[
                                group_only_tmp
                            ]

                # Remainder part
                if remainder is not None:
                    for remainder_tmp in remainder:
                        if not remainder_set:
                            #print('group_level_parameter_dict')
                            #print(group_level_parameter_dict)
                            tmp_mean = group_level_parameter_dict[remainder_tmp]
                            tmp_std = group_level_parameter_dict[remainder_tmp + "_std"]
                            full_parameter_dict[
                                remainder_tmp + "_subj." + str(subj_idx)
                            ] = np.random.normal(loc=tmp_mean, scale=tmp_std)
                            subj_data[remainder_tmp] = full_parameter_dict[
                                remainder_tmp + "_subj." + str(subj_idx)
                            ]

                        if remainder_set:
                            subj_data[remainder_tmp] = full_parameter_dict[
                                remainder_tmp + "_subj." + str(subj_idx)
                            ]

                # Depends on part
                if depends_on is not None:
                    # conditions_tmp = conditions_df.iloc[condition_id]
                    for depends_tmp in depends_on.keys():
                        conditions_df_tmp = conditions_df[depends_on[depends_tmp]].iloc[
                            condition_id
                        ]
                        condition_elem = ".".join(conditions_df_tmp)

                        if depends_tmp not in group_only:
                            tmp_mean = group_level_parameter_dict[
                                depends_tmp + "(" + condition_elem + ")"
                            ]
                            tmp_std = group_level_parameter_dict[depends_tmp + "_std"]

                            full_parameter_dict[
                                depends_tmp
                                + "_subj("
                                + condition_elem
                                + ")."
                                + str(subj_idx)
                            ] = np.random.normal(loc=tmp_mean, scale=tmp_std)
                            
                            subj_data[depends_tmp] = full_parameter_dict[
                                depends_tmp
                                + "_subj("
                                + condition_elem
                                + ")."
                                + str(subj_idx)
                            ]
                        else:
                            subj_data[depends_tmp] = full_parameter_dict[
                                depends_tmp + "(" + condition_elem + ")"
                            ]

                        for condition_key_tmp in conditions_df_tmp.keys():
                            subj_data[condition_key_tmp] = conditions_df_tmp[
                                condition_key_tmp
                            ]  ##############################################################

                # Regressor part
                if regression_covariates is not None:
                    cov_df = make_covariate_df(
                        regression_covariates, n_trials_per_subject
                    )

                    # Add cov_df to subject data
                    for key_tmp in cov_df.keys():
                        subj_data[key_tmp] = cov_df[key_tmp].copy()

                if regression_models is not None:
                    for reg_model in regression_models:
                        
                        # Make Design Matrix
                        separator = reg_model.find("~")
                        outcome = reg_model[:separator].strip(" ")
                        reg_model_stripped = reg_model[(separator + 1) :]
                        design_matrix = dmatrix(reg_model_stripped, cov_df)

                        reg_params_tmp = []
                        reg_param_names_tmp = []
                        for reg_param_key in group_level_parameter_dict[
                            outcome + "_reg"
                        ].keys():
                            if (
                                group_only_regressors and "Intercept" in reg_param_key
                            ) or (not group_only_regressors):
                                reg_params_tmp.append(
                                    np.random.normal(
                                        loc=group_level_parameter_dict[
                                            outcome + "_reg"
                                        ][reg_param_key],
                                        scale=group_level_parameter_dict[
                                            outcome + "_reg_std"
                                        ][reg_param_key + "_std"],
                                    )
                                )

                                reg_param_names_tmp.append(
                                    reg_param_key + "_subj." + str(subj_idx)
                                )
                            else:
                                reg_params_tmp.append(
                                    group_level_parameter_dict[outcome + "_reg"][
                                        reg_param_key
                                    ]
                                )
                                reg_param_names_tmp.append(reg_param_key)

                        reg_params_tmp = np.array(reg_params_tmp)

                        for key in group_level_parameter_dict[outcome + "_reg"].keys():
                            full_parameter_dict[key] = group_level_parameter_dict[
                                outcome + "_reg"
                            ][key]
                        for key in group_level_parameter_dict[
                            outcome + "_reg_std"
                        ].keys():
                            full_parameter_dict[key] = group_level_parameter_dict[
                                outcome + "_reg_std"
                            ][key]

                        if not regressor_set:
                            for k in range(len(reg_param_names_tmp)):
                                full_parameter_dict[
                                    reg_param_names_tmp[k]
                                ] = reg_params_tmp[k]

                        subj_data[outcome] = (design_matrix * reg_params_tmp).sum(
                            axis=1
                        ) # AF-TD: This should probably include a noise term here (parameter really defined as coming from a linear model + noise)

                # Append full data:
                full_data.append(subj_data.copy())

            remainder_set = 1
            regressor_set = 1

        full_data = pd.concat(full_data)
        parameters = full_data[model_config[model]["params"]]

        # Run the actual simulations
        # print(parameters)

        sim_data = simulator(
            theta=parameters.values,
            model=model,
            n_samples=1,
            delta_t=0.001,
            max_t=20,
            no_noise=False,
            bin_dim=None,
            bin_pointwise=False,
        )

        # Post-processing
        full_data["rt"] = sim_data[0].astype(np.float64)
        full_data["response"] = sim_data[1].astype(np.float64)
        full_data.loc[full_data["response"] < 0, ["response"]] = 0.0

        # Add in outliers
        if p_outlier > 0:
            # print('passing through outlier creation')
            outlier_idx = np.random.choice(
                list(data.index),
                replace=False,
                size=int(p_outlier * len(list(data.index))),
            )
            outlier_data = np.zeros((outlier_idx.shape[0], 2))

            # Outlier rts
            outlier_data[:, 0] = np.random.uniform(
                low=0.0, high=outlier_max_t, size=outlier_data.shape[0]
            )

            # Outlier choices
            outlier_data[:, 1] = np.random.choice(
                sim_data[2]["possible_choices"], size=outlier_data.shape[0]
            )

            # Exchange data for outliers
            full_data.iloc[
                outlier_idx,
                [
                    list(full_data.keys()).index("rt"),
                    list(full_data.keys()).index("response"),
                ],
            ] = outlier_data

            # Identify outliers in dataframe
            full_data["outlier"] = 0
            full_data[outlier_idx, [list(full_data.keys()).index("outlier")]] = 1

        full_data_cols = ["rt", "response", "subj_idx"]

        if regression_covariates is not None:
            full_data_cols += [key for key in regression_covariates.keys()]
        if conditions is not None:
            full_data_cols += [key for key in conditions.keys()]

        full_data_cols += model_config[model]["params"]
        full_data = full_data[full_data_cols]
        full_data.reset_index(drop=True, inplace=True)

        # AF-Comment: Does this cover all corner cases?
        # If n_subjects is 1 --> we overwrite the group parameters with the subj.0 parameters
        if n_subjects == 1:
            new_param_dict = {}
            for key, value in full_parameter_dict.items():
                if "subj" in key:
                    new_key = key
                    new_key = new_key.replace("_subj", "")
                    new_key = new_key[: new_key.find(".")]
                    new_param_dict[new_key] = value
                elif "_std" in key:
                    pass
                else:
                    new_param_dict[key] = value
            full_parameter_dict = new_param_dict

        return full_data, full_parameter_dict

    def make_single_sub_cond_df_from_gt(data,
                                        conditions_df,
                                        depends_on,
                                        regression_models,
                                        regression_covariates,
                                        group_only_regressors,
                                        group_only,
                                        fixed_at_default,
                                        remainder,
                                        model,
                                        group_level_parameter_dict,
                                        ):

        # Construct subject data
        full_parameter_dict = group_level_parameter_dict.copy()

        # Subject part -----------------------
        full_data = []
        # Condition --------------------------

        # Initialize parameter columns in data
        for param in model_config[model]['params']:
            data[param] = 0

        # if conditions_df is None:
        #     n_conditions = 1
        # else:
        #     n_conditions = conditions_df.shape[0]

        for subj_idx in data['subj_idx'].unique(): #range(n_subjects):
            # Parameter vector

            #subj_data = pd.DataFrame(index = np.arange(0, n_trials_per_subject, 1))
            #subj_data["subj_idx"] = str(subj_idx)

            # Fixed part
            if fixed_at_default is not None:
                for fixed_tmp in fixed_at_default:
                    data.loc[data['subj_idx'] == subj_idx, ['fixed_tmp']] = group_level_parameter_dict[fixed_tmp]
                    #subj_data.loc[fixed_tmp] = group_level_parameter_dict[fixed_tmp]

            # Group only part
            if group_only is not None:
                for group_only_tmp in group_only:
                    if group_only_tmp in list(depends_on.keys()):
                        pass
                    else:
                        data.loc[data['subj_idx'] == subj_idx, [group_only_tmp]] = group_level_parameter_dict[
                            group_only_tmp
                        ]

            # Remainder part
            if remainder is not None:
                for remainder_tmp in remainder:
                    #print('group_level_parameter_dict')
                    #print(group_level_parameter_dict)
                    tmp_mean = group_level_parameter_dict[remainder_tmp]
                    tmp_std = group_level_parameter_dict[remainder_tmp + "_std"]
                    full_parameter_dict[
                        remainder_tmp + "_subj." + str(subj_idx)
                    ] = np.random.normal(loc=tmp_mean, scale=tmp_std)
                        
                    data.loc[data['subj_idx'] == subj_idx, [remainder_tmp]] = full_parameter_dict[
                            remainder_tmp + "_subj." + str(subj_idx)
                        ]

            # AF CONTINUE HERE !
            # Depends on part

            if depends_on is not None:
                # conditions_tmp = conditions_df.iloc[condition_id]
                # Go through depends_on variables:
                for depends_tmp in depends_on.keys():
                    conditions_df_tmp = conditions_df[depends_on[depends_tmp]].drop_duplicates()
                    for condition_id in range(conditions_df_tmp.shape[0]):
                    # conditions_df_tmp = conditions_df[depends_on[depends_tmp]].iloc[
                    #     condition_id
                    # ]

                        condition_elem = ".".join(conditions_df_tmp.iloc[condition_id])
                        bool_ = data['subj_idx'] = subj_idx

                        for key_ in conditions_df_tmp.keys():
                            bool_ = bool_ & (data[key_] == conditions_df_tmp.iloc[conditions_id][key_])

                    # Check if there is data which adheres to the condition currently active
                    # Otherwise there is nothing to update
                    # AF COMMENT: This check should already be applied at the point of generating the condition_df dataframe 
                    if np.sum(bool_) > 0:
                        if depends_tmp not in group_only:
                            tmp_mean = group_level_parameter_dict[
                                depends_tmp + "(" + condition_elem + ")"
                            ]
                            tmp_std = group_level_parameter_dict[depends_tmp + "_std"]

                            full_parameter_dict[
                                depends_tmp
                                + "_subj("
                                + condition_elem
                                + ")."
                                + str(subj_idx)
                            ] = np.random.normal(loc=tmp_mean, scale=tmp_std)
                            
                            data.loc[bool_, depends_tmp] = full_parameter_dict[
                                depends_tmp
                                + "_subj("
                                + condition_elem
                                + ")."
                                + str(subj_idx)
                            ]

                            # subj_data[depends_tmp] = full_parameter_dict[
                            #     depends_tmp
                            #     + "_subj("
                            #     + condition_elem
                            #     + ")."
                            #     + str(subj_idx)
                            # ]
                        else:
                            data.loc[bool_, depends_tmp] = full_parameter_dict[
                                depends_tmp + "(" + condition_elem + ")"
                            ]
                            # subj_data[depends_tmp] = full_parameter_dict[
                            #     depends_tmp + "(" + condition_elem + ")"
                            # ]
                        
                        # AF COMMENT: This assigns the actual condition columns to the data
                        # Not necessary when data is supplied instead
                        # 
                        # for condition_key_tmp in conditions_df_tmp.keys():
                        #     subj_data[condition_key_tmp] = conditions_df_tmp[
                        #         condition_key_tmp
                        #     ]

            # Regressor part
            if regression_covariates is not None:
                cov_df = make_covariate_df(
                    regression_covariates, n_trials_per_subject
                )

                # Add cov_df to subject data
                # AF COMMENT: Not necessary if 
                # for key_tmp in cov_df.keys():
                #     subj_data[key_tmp] = cov_df[key_tmp].copy()

            if regression_models is not None:
                for reg_model in regression_models:
                    
                    # Make Design Matrix
                    separator = reg_model.find("~")
                    outcome = reg_model[:separator].strip(" ")
                    reg_model_stripped = reg_model[(separator + 1) :]
                    design_matrix = dmatrix(reg_model_stripped, data.loc[data['subj_idx'] == subj_idx, :])

                    reg_params_tmp = []
                    reg_param_names_tmp = []
                    for reg_param_key in group_level_parameter_dict[
                        outcome + "_reg"
                    ].keys():
                        if (
                            group_only_regressors and "Intercept" in reg_param_key
                        ) or (not group_only_regressors):
                            reg_params_tmp.append(
                                np.random.normal(
                                    loc=group_level_parameter_dict[
                                        outcome + "_reg"
                                    ][reg_param_key],
                                    scale=group_level_parameter_dict[
                                        outcome + "_reg_std"
                                    ][reg_param_key + "_std"],
                                )
                            )

                            reg_param_names_tmp.append(
                                reg_param_key + "_subj." + str(subj_idx)
                            )
                        else:
                            reg_params_tmp.append(
                                group_level_parameter_dict[outcome + "_reg"][
                                    reg_param_key
                                ]
                            )
                            reg_param_names_tmp.append(reg_param_key)

                    reg_params_tmp = np.array(reg_params_tmp)

                    for key in group_level_parameter_dict[outcome + "_reg"].keys():
                        full_parameter_dict[key] = group_level_parameter_dict[
                            outcome + "_reg"
                        ][key]
                    
                    for key in group_level_parameter_dict[
                        outcome + "_reg_std"
                    ].keys():
                        full_parameter_dict[key] = group_level_parameter_dict[
                            outcome + "_reg_std"
                        ][key]

                    if not regressor_set:
                        for k in range(len(reg_param_names_tmp)):
                            full_parameter_dict[
                                reg_param_names_tmp[k]
                            ] = reg_params_tmp[k]

                    data.loc[data['subj_idx'] == subj_idx, [outcome]] = (design_matrix * reg_params_tmp).sum(axis=1) # AF-TD: This should probably include a noise term here (parameter really defined as coming from a linear model + noise)
                    
                    # subj_data[outcome] = (design_matrix * reg_params_tmp).sum(
                    #     axis=1
                    # ) # AF-TD: This should probably include a noise term here (parameter really defined as coming from a linear model + noise)

            # Append full data:
            # full_data.append(subj_data.copy())

            remainder_set = 1
            regressor_set = 1

        # full_data = pd.concat(full_data)
        parameters = data[model_config[model]["params"]]

        # Run the actual simulations
        # print(parameters)

        sim_data = simulator(
            theta=parameters.values,
            model=model,
            n_samples=1,
            delta_t=0.001,
            max_t=20,
            no_noise=False,
            bin_dim=None,
            bin_pointwise=False,
        )

        # Post-processing
        data["rt"] = sim_data[0].astype(np.float64)
        data["response"] = sim_data[1].astype(np.float64)
        data.loc[data["response"] < 0, ["response"]] = 0.0

        # Add in outliers
        if p_outlier > 0:
            outlier_idx = np.random.choice(
                list(data.index),
                replace=False,
                size=int(p_outlier * len(list(data.index))),
            )
            outlier_data = np.zeros((outlier_idx.shape[0], 2))

            # Outlier rts
            outlier_data[:, 0] = np.random.uniform(
                low=0.0, high=outlier_max_t, size=outlier_data.shape[0]
            )

            # Outlier choices
            outlier_data[:, 1] = np.random.choice(
                sim_data[2]["possible_choices"], size=outlier_data.shape[0]
            )

            # Exchange data for outliers
            data.loc[outlier_idx, ["rt","response"]] = outlier_data

            # Identify outliers in dataframe
            data["outlier"] = 0
            data.loc[outlier_idx, [list(full_data.keys()).index("outlier")]] = 1

        # AF COMMENT: Prob not useful
        # data_cols = ["rt", "response", "subj_idx"]

        # if regression_covariates is not None:
        #     data_cols += [key for key in regression_covariates.keys()]
        # if conditions is not None:
        #     data_cols += [key for key in conditions.keys()]

        # data_cols += model_config[model]["params"]
        # data = data[data_cols]
        # data.reset_index(drop=True, inplace=True)

        # AF-Comment: Does this cover all corner cases?
        # If n_subjects is 1 --> we overwrite the group parameters with the subj.0 parameters
        if len(data['subj_idx'].unique()) == 1:
            new_param_dict = {}
            for key, value in full_parameter_dict.items():
                if "subj" in key:
                    new_key = key
                    new_key = new_key.replace("_subj", "")
                    new_key = new_key[: new_key.find(".")]
                    new_param_dict[new_key] = value
                elif "_std" in key:
                    pass
                else:
                    new_param_dict[key] = value
            full_parameter_dict = new_param_dict

        return data, full_parameter_dict

    def make_group_level_params(data = None,
        conditions_df,
        group_only,
        depends_on,
        model,
        fixed_at_default,
        remainder,
        group_only_regressors,
        regression_models,
        regression_covariates,
        group_param_dist = 'normal',
        gen_norm_std = 1/4,
        uniform_buffer = 1/5,
        gen_std_std = 1/8,
        covariate_range = 1/4, # multiplied by range of parameter bounds to give size of covariate
    ):
        """
        Make group level parameters from the information supplied.
        """

        group_level_parameter_dict = {}

        # COLLECT PARAMETER WISE DATA AND ON CONSTRAINTS AND RV-GENERATORS ------
        param_gen_info = {}
        for param_name in model_config[model]["params"]:
            idx = model_config[model]["params"].index(param_name)
            
            param_gen_info[param_name] = {}
            param_gen_info[param_name]['range'] = model_config[model]["param_bounds"][1][idx] - \
                model_config[model]["param_bounds"][0][idx]
            
            param_gen_info[param_name]['mid'] = model_config[model]["param_bounds"][0][idx] + (param_gen_info[param_name]['range'] / 2)
            param_gen_info[param_name]['gen_norm_std'] = gen_norm_std * (param_gen_info[param_name]['range'] / 2)
            param_gen_info[param_name]['uniform_buffer'] = uniform_buffer * (param_gen_info[param_name]['range'] / 2)
            param_gen_info[param_name]['std_gen_std'] = gen_std_std * param_gen_info[param_name]['range']
            param_gen_info[param_name]['covariate_range'] = covariate_range * param_gen_info[param_name]['range']

            if  group_param_dist == 'normal':
                param_gen_info[param_name]['rv'] = partial(np.random.normal, 
                                                              loc = param_gen_info[param_name]['mid'],
                                                              scale = param_gen_info[param_name]['gen_norm_std'])
            elif group_param_dist == 'uniform':
                param_gen_info[param_name]['rv'] = partial(np.random.uniform,
                                                           low = model_config[model]["param_bounds"][0][param_name] + \
                                                               param_gen_info[param_name]['uniform_buffer'],
                                                           high = model_config[model]["param_bounds"][1][param_name] - \
                                                               param_gen_info[param_name]['uniform_buffer'])

            param_gen_info[param_name]['std_rv'] = partial(np.random.uniform,
                                                           low = 0,
                                                           high = param_gen_info[param_name]['std_gen_std'])

            param_gen_info[param_name]['covariate_rv'] = partial(np.random.uniform,
                                                                 low = - param_gen_info[param_name]['covariate_range'],
                                                                 high = param_gen_info[param_name]['covariate_range'])
        # -----------------------------------------------
        
        
        # Fixed part --------------------------------------------------------
        if fixed_at_default is not None:
            for fixed_tmp in fixed_at_default:
                group_level_parameter_dict[fixed_tmp] = model_config[model]["default_params"][model_config[model]["params"].index(fixed_tmp)]
                
        # Group only part (excluding depends on) ----------------------------
        if len(group_only) > 0:
            for group_only_tmp in group_only:
                if group_only_tmp in list(depends_on.keys()):
                    pass
                else:
                    group_level_parameter_dict[group_only_tmp] = param_gen_info[group_only_tmp]['rv']() 

        # Remainder part -----------------------------------------------------
        if remainder is not None:
            for remainder_tmp in remainder:
                group_level_parameter_dict[remainder_tmp] = param_gen_info[remainder_tmp]['rv']() 
                group_level_parameter_dict[remainder_tmp + "_std"] = param_gen_info[remainder_tmp]['std_rv']()
                
        # Depends on part ----------------------------------------------------
        if depends_on is not None:
            for depends_tmp in depends_on.keys():
                conditions_df_tmp = conditions_df[depends_on[depends_tmp]]

                # Get unique elements:
                unique_elems = []
                for i in range(conditions_df_tmp.shape[0]):
                    unique_elems.append(".".join(conditions_df_tmp.iloc[i]))
                unique_elems = np.unique(np.array(unique_elems))

                for unique_elem in unique_elems:
                    group_level_parameter_dict[
                        depends_tmp + "(" + unique_elem + ")"
                    ] = param_gen_info[depends_tmp]['rv']()
                    
                if depends_tmp not in group_only:
                    group_level_parameter_dict[
                        depends_tmp + "_std"
                    ] = param_gen_info[remainder_tmp]['std_rv']()
                    
                    #np.random.uniform(low=0, high=(1 / 10) * bound_to_bound_tmp)

        # Regressor part ------------------------------------------------------
        if regression_covariates is not None:
            # AF ADDED:
            # IF covariates supplied: skip generation
            if data is None:
                cov_df = make_covariate_df(regression_covariates, n_trials_per_subject)
            else:
                cov_df = data

        if regression_models is not None:
            for reg_model in regression_models:
                separator = reg_model.find("~")
                outcome = reg_model[:separator].strip(" ")
                reg_model_stripped = reg_model[(separator + 1) :]

                # Run through patsy dmatrix to get the covariate names 
                # that patsy assigns !
                covariate_names = dmatrix(
                    reg_model_stripped, cov_df
                ).design_info.column_names

                reg_trace_dict = OrderedDict()
                reg_std_trace_dict = OrderedDict()

                for covariate in covariate_names:
                    if ("Intercept" in covariate) or (covariate == '1'):

                        # AF-COMMENT: Here instead of covariate_rv --> just use
                        reg_trace_dict[outcome + "_" + covariate] = param_gen_info[outcome]['rv']()
                        
                        # print(reg_trace_dict[outcome + "_" + covariate])

                        # Intercept is always fit subject wise
                        reg_std_trace_dict[
                            outcome + "_" + covariate + "_" + "std"
                        ] = param_gen_info[outcome]['std_rv']()
                        
                    else:
                        reg_trace_dict[outcome + "_" + covariate] = param_gen_info[outcome]['covariate_rv']()
                        
                        if not group_only_regressors:
                            reg_std_trace_dict[
                                outcome + "_" + covariate + "_" + "std"
                            ] = param_gen_info[outcome]['std_rv']()
                            
                group_level_parameter_dict[outcome + "_reg"] = reg_trace_dict.copy()

                # AF-COMMENT: Is this necessary ?
                # if not group_only_regressors:
                group_level_parameter_dict[
                    outcome + "_reg" + "_std"
                ] = reg_std_trace_dict.copy()

        return group_level_parameter_dict

    # MAIN PART OF THE FUNCTION -----------------------------------------------------------------

    # Some checks
    if group_only is None:
        group_only = []

    # Specify 'remainder' parameters --> will be sampled randomly from the allowed range
    remainder = get_parameter_remainder(
        regression_models=regression_models,
        group_only=group_only,
        depends_on=depends_on,
        fixed_at_default=fixed_at_default,
    )
    #print(fixed_at_default)
    #print(remainder)

    # Make conditions df
    if depends_on is not None:
        conditions_df = make_conditions_df(conditions=conditions)
        #print("Conditions created...")
        #print(conditions_df)
    else:
        conditions_df = None

    params_ok_all = 0
    cnt = 0
    while params_ok_all == 0:
        if cnt > 0:
            print(
                "new round of data simulation because parameter bounds where violated"
            )

        group_level_param_dict = make_group_level_params(data = data,
            conditions_df=conditions_df,
            group_only=group_only,
            depends_on=depends_on,
            model=model,
            fixed_at_default=fixed_at_default,
            remainder=remainder,
            group_only_regressors=group_only_regressors,
            regression_models=regression_models,
            regression_covariates=regression_covariates,
            group_param_dist = meta_params['group_param_dist'],
            gen_norm_std = meta_params['gen_norm_std'],
            uniform_buffer = meta_params['uniform_buffer'],
            gen_std_std = meta_params['gen_std_std'],
            covariate_range = meta_params['covariate_range'],
            )

        if data is None:
            data_, full_parameter_dict=make_single_sub_cond_df(
                                                              conditions_df=conditions_df,
                                                              group_only=group_only,
                                                              depends_on=depends_on,
                                                              model=model,
                                                              fixed_at_default=fixed_at_default,
                                                              remainder=remainder,
                                                              regression_models=regression_models,
                                                              regression_covariates=regression_covariates,
                                                              group_only_regressors=group_only_regressors,
                                                              group_level_parameter_dict=group_level_param_dict,
                                                              n_trials_per_subject=n_trials_per_subject,
                                                              n_subjects=n_subjects,
                                                              )
        else:
            data_, full_parameter_dict=make_single_sub_cond_df_from_gt(data = data,
                                                                      conditions_df=conditions_df,
                                                                      group_only=group_only,
                                                                      depends_on=depends_on,
                                                                      model=model,
                                                                      fixed_at_default=fixed_at_default,
                                                                      remainder=remainder,
                                                                      regression_models=regression_models,
                                                                      regression_covariates=regression_covariates,
                                                                      group_only_regressors=group_only_regressors,
                                                                      group_level_parameter_dict=group_level_param_dict,
                                                                      n_trials_per_subject=n_trials_per_subject,
                                                                      n_subjects=n_subjects,
                                                                      )

        params_ok_all = check_params(data=data_, 
                                     model=model)
        cnt += 1

    return data_, full_parameter_dict<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.08354638073797077, 'fraction_numerical': 0.0044262983172434845, 'mean_word_length': 1.2264463512021069, 'pattern_counts': {'":': 11, '<': 28, '<?xml version=': 0, '>': 38, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 234, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '603140', 'license': 'MIT', 'n_tokens_mistral': 13214, 'n_tokens_neox': 12421, 'n_words': 2675, 'path': 'make_parameter_recov_data.ipynb', 'repo_name': 'AlexanderFengler/hierarchical_decision_making', 'size': '52414'}