Hyperparameter Tuning using Optuna

.gitignore

@@ -31,3 +31,4 @@ MANIFEST
 *.vscode
 .flake8
 .coverage
+.serena

.serena/.gitignore

@@ -0,0 +1 @@
+/cache

.serena/project.yml

@@ -0,0 +1,67 @@
+# language of the project (csharp, python, rust, java, typescript, go, cpp, or ruby)
+#  * For C, use cpp
+#  * For JavaScript, use typescript
+# Special requirements:
+#  * csharp: Requires the presence of a .sln file in the project folder.
+language: python
+
+# whether to use the project's gitignore file to ignore files
+# Added on 2025-04-07
+ignore_all_files_in_gitignore: true
+# list of additional paths to ignore
+# same syntax as gitignore, so you can use * and **
+# Was previously called `ignored_dirs`, please update your config if you are using that.
+# Added (renamed) on 2025-04-07
+ignored_paths: []
+
+# whether the project is in read-only mode
+# If set to true, all editing tools will be disabled and attempts to use them will result in an error
+# Added on 2025-04-18
+read_only: false
+
+# list of tool names to exclude. We recommend not excluding any tools, see the readme for more details.
+# Below is the complete list of tools for convenience.
+# To make sure you have the latest list of tools, and to view their descriptions,
+# execute `uv run scripts/print_tool_overview.py`.
+#
+#  * `activate_project`: Activates a project by name.
+#  * `check_onboarding_performed`: Checks whether project onboarding was already performed.
+#  * `create_text_file`: Creates/overwrites a file in the project directory.
+#  * `delete_lines`: Deletes a range of lines within a file.
+#  * `delete_memory`: Deletes a memory from Serena's project-specific memory store.
+#  * `execute_shell_command`: Executes a shell command.
+#  * `find_referencing_code_snippets`: Finds code snippets in which the symbol at the given location is referenced.
+#  * `find_referencing_symbols`: Finds symbols that reference the symbol at the given location (optionally filtered by type).
+#  * `find_symbol`: Performs a global (or local) search for symbols with/containing a given name/substring (optionally filtered by type).
+#  * `get_current_config`: Prints the current configuration of the agent, including the active and available projects, tools, contexts, and modes.
+#  * `get_symbols_overview`: Gets an overview of the top-level symbols defined in a given file.
+#  * `initial_instructions`: Gets the initial instructions for the current project.
+#     Should only be used in settings where the system prompt cannot be set,
+#     e.g. in clients you have no control over, like Claude Desktop.
+#  * `insert_after_symbol`: Inserts content after the end of the definition of a given symbol.
+#  * `insert_at_line`: Inserts content at a given line in a file.
+#  * `insert_before_symbol`: Inserts content before the beginning of the definition of a given symbol.
+#  * `list_dir`: Lists files and directories in the given directory (optionally with recursion).
+#  * `list_memories`: Lists memories in Serena's project-specific memory store.
+#  * `onboarding`: Performs onboarding (identifying the project structure and essential tasks, e.g. for testing or building).
+#  * `prepare_for_new_conversation`: Provides instructions for preparing for a new conversation (in order to continue with the necessary context).
+#  * `read_file`: Reads a file within the project directory.
+#  * `read_memory`: Reads the memory with the given name from Serena's project-specific memory store.
+#  * `remove_project`: Removes a project from the Serena configuration.
+#  * `replace_lines`: Replaces a range of lines within a file with new content.
+#  * `replace_symbol_body`: Replaces the full definition of a symbol.
+#  * `restart_language_server`: Restarts the language server, may be necessary when edits not through Serena happen.
+#  * `search_for_pattern`: Performs a search for a pattern in the project.
+#  * `summarize_changes`: Provides instructions for summarizing the changes made to the codebase.
+#  * `switch_modes`: Activates modes by providing a list of their names
+#  * `think_about_collected_information`: Thinking tool for pondering the completeness of collected information.
+#  * `think_about_task_adherence`: Thinking tool for determining whether the agent is still on track with the current task.
+#  * `think_about_whether_you_are_done`: Thinking tool for determining whether the task is truly completed.
+#  * `write_memory`: Writes a named memory (for future reference) to Serena's project-specific memory store.
+excluded_tools: []
+
+# initial prompt for the project. It will always be given to the LLM upon activating the project
+# (contrary to the memories, which are loaded on demand).
+initial_prompt: ""
+
+project_name: "doubleml-for-py"

doubleml/did/did.py

@@ -370,7 +370,14 @@ def _sensitivity_element_est(self, preds):
         return element_dict
 
     def _nuisance_tuning(
-        self, smpls, param_grids, scoring_methods, n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search
+        self,
+        smpls,
+        param_grids,
+        scoring_methods,
+        n_folds_tune,
+        n_jobs_cv,
+        search_mode,
+        n_iter_randomized_search,
     ):
         x, y = check_X_y(self._dml_data.x, self._dml_data.y, force_all_finite=False)
         x, d = check_X_y(x, self._dml_data.d, force_all_finite=False)
@@ -394,6 +401,7 @@ def _nuisance_tuning(
             n_jobs_cv,
             search_mode,
             n_iter_randomized_search,
+            learner_name="ml_g",
         )
         g1_tune_res = _dml_tune(
             y,
@@ -406,6 +414,7 @@ def _nuisance_tuning(
             n_jobs_cv,
             search_mode,
             n_iter_randomized_search,
+            learner_name="ml_g",
         )
 
         g0_best_params = [xx.best_params_ for xx in g0_tune_res]
@@ -423,6 +432,7 @@ def _nuisance_tuning(
                 n_jobs_cv,
                 search_mode,
                 n_iter_randomized_search,
+                learner_name="ml_m",
             )
             m_best_params = [xx.best_params_ for xx in m_tune_res]
             params = {"ml_g0": g0_best_params, "ml_g1": g1_best_params, "ml_m": m_best_params}
@@ -435,6 +445,99 @@ def _nuisance_tuning(
 
         return res
 
+    def _nuisance_tuning_optuna(
+        self,
+        param_grids,
+        scoring_methods,
+        n_folds_tune,
+        n_jobs_cv,
+        optuna_settings,
+    ):
+        """
+        Optuna-based hyperparameter tuning for DID nuisance models.
+
+        Performs tuning once on the whole dataset using cross-validation,
+        returning the same optimal parameters for all folds.
+        """
+        from ..utils._tune_optuna import _dml_tune_optuna
+
+        x, y = check_X_y(self._dml_data.x, self._dml_data.y, force_all_finite=False)
+        x, d = check_X_y(x, self._dml_data.d, force_all_finite=False)
+
+        if scoring_methods is None:
+            if self.score == "observational":
+                scoring_methods = {"ml_g0": None, "ml_g1": None, "ml_m": None}
+            else:
+                scoring_methods = {"ml_g0": None, "ml_g1": None}
+
+        # Separate data by treatment status for conditional mean tuning
+        mask_d0 = d == 0
+        mask_d1 = d == 1
+
+        x_d0 = x[mask_d0, :]
+        y_d0 = y[mask_d0]
+        train_inds_d0 = [np.arange(x_d0.shape[0])]
+
+        g0_tune_res = _dml_tune_optuna(
+            y_d0,
+            x_d0,
+            train_inds_d0,
+            self._learner["ml_g"],
+            param_grids["ml_g0"],
+            scoring_methods["ml_g0"],
+            n_folds_tune,
+            n_jobs_cv,
+            optuna_settings,
+            learner_name="ml_g0",
+        )
+
+        x_d1 = x[mask_d1, :]
+        y_d1 = y[mask_d1]
+        train_inds_d1 = [np.arange(x_d1.shape[0])]
+
+        g1_tune_res = _dml_tune_optuna(
+            y_d1,
+            x_d1,
+            train_inds_d1,
+            self._learner["ml_g"],
+            param_grids["ml_g1"],
+            scoring_methods["ml_g1"],
+            n_folds_tune,
+            n_jobs_cv,
+            optuna_settings,
+            learner_name="ml_g1",
+        )
+
+        # Tune propensity score on full dataset for observational score
+        full_train_inds = [np.arange(x.shape[0])]
+        m_tune_res = None
+        if self.score == "observational":
+            m_tune_res = _dml_tune_optuna(
+                d,
+                x,
+                full_train_inds,
+                self._learner["ml_m"],
+                param_grids["ml_m"],
+                scoring_methods["ml_m"],
+                n_folds_tune,
+                n_jobs_cv,
+                optuna_settings,
+                learner_name="ml_m",
+            )
+
+        g0_best_params = [xx.best_params_ for xx in g0_tune_res]
+        g1_best_params = [xx.best_params_ for xx in g1_tune_res]
+
+        if self.score == "observational":
+            m_best_params = [xx.best_params_ for xx in m_tune_res]
+            params = {"ml_g0": g0_best_params, "ml_g1": g1_best_params, "ml_m": m_best_params}
+            tune_res = {"g0_tune": g0_tune_res, "g1_tune": g1_tune_res, "m_tune": m_tune_res}
+        else:
+            params = {"ml_g0": g0_best_params, "ml_g1": g1_best_params}
+            tune_res = {"g0_tune": g0_tune_res, "g1_tune": g1_tune_res}
+
+        return {"params": params, "tune_res": tune_res}
+
     def sensitivity_benchmark(self, benchmarking_set, fit_args=None):
         """
         Computes a benchmark for a given set of features.

doubleml/did/did_binary.py

@@ -568,7 +568,14 @@ def _score_elements(self, y, d, g_hat0, g_hat1, m_hat, p_hat):
         return psi_a, psi_b
 
     def _nuisance_tuning(
-        self, smpls, param_grids, scoring_methods, n_folds_tune, n_jobs_cv, search_mode, n_iter_randomized_search
+        self,
+        smpls,
+        param_grids,
+        scoring_methods,
+        n_folds_tune,
+        n_jobs_cv,
+        search_mode,
+        n_iter_randomized_search,
     ):
         x, y = check_X_y(self._x_data_subset, self._y_data_subset, force_all_finite=False)
         x, d = check_X_y(x, self._g_data_subset, force_all_finite=False)
@@ -634,6 +641,90 @@ def _nuisance_tuning(
 
         return res
 
+    def _nuisance_tuning_optuna(
+        self,
+        param_grids,
+        scoring_methods,
+        n_folds_tune,
+        n_jobs_cv,
+        optuna_settings,
+    ):
+        from ..utils._tune_optuna import _dml_tune_optuna
+
+        x, y = check_X_y(self._x_data_subset, self._y_data_subset, force_all_finite=False)
+        x, d = check_X_y(x, self._g_data_subset, force_all_finite=False)
+
+        if scoring_methods is None:
+            scoring_methods = {"ml_g": None, "ml_m": None}
+
+        mask_d0 = d == 0
+        mask_d1 = d == 1
+
+        x_d0 = x[mask_d0, :]
+        y_d0 = y[mask_d0]
+        train_inds_d0 = [np.arange(x_d0.shape[0])]
+        g0_param_grid = param_grids.get("ml_g0", param_grids["ml_g"])
+        g0_scoring = scoring_methods.get("ml_g0", scoring_methods["ml_g"])
+        g0_tune_res = _dml_tune_optuna(
+            y_d0,
+            x_d0,
+            train_inds_d0,
+            self._learner["ml_g"],
+            g0_param_grid,
+            g0_scoring,
+            n_folds_tune,
+            n_jobs_cv,
+            optuna_settings,
+            learner_name=("ml_g0", "ml_g"),
+        )
+
+        x_d1 = x[mask_d1, :]
+        y_d1 = y[mask_d1]
+        train_inds_d1 = [np.arange(x_d1.shape[0])]
+        g1_param_grid = param_grids.get("ml_g1", param_grids["ml_g"])
+        g1_scoring = scoring_methods.get("ml_g1", scoring_methods["ml_g"])
+        g1_tune_res = _dml_tune_optuna(
+            y_d1,
+            x_d1,
+            train_inds_d1,
+            self._learner["ml_g"],
+            g1_param_grid,
+            g1_scoring,
+            n_folds_tune,
+            n_jobs_cv,
+            optuna_settings,
+            learner_name=("ml_g1", "ml_g"),
+        )
+
+        full_train_inds = [np.arange(x.shape[0])]
+        m_tune_res = None
+        if self.score == "observational":
+            m_tune_res = _dml_tune_optuna(
+                d,
+                x,
+                full_train_inds,
+                self._learner["ml_m"],
+                param_grids["ml_m"],
+                scoring_methods["ml_m"],
+                n_folds_tune,
+                n_jobs_cv,
+                optuna_settings,
+                learner_name="ml_m",
+            )
+
+        g0_best_params = [xx.best_params_ for xx in g0_tune_res]
+        g1_best_params = [xx.best_params_ for xx in g1_tune_res]
+
+        if self.score == "observational":
+            m_best_params = [xx.best_params_ for xx in m_tune_res]
+            params = {"ml_g0": g0_best_params, "ml_g1": g1_best_params, "ml_m": m_best_params}
+            tune_res = {"g0_tune": g0_tune_res, "g1_tune": g1_tune_res, "m_tune": m_tune_res}
+        else:
+            params = {"ml_g0": g0_best_params, "ml_g1": g1_best_params}
+            tune_res = {"g0_tune": g0_tune_res, "g1_tune": g1_tune_res}
+
+        return {"params": params, "tune_res": tune_res}
+
     def _sensitivity_element_est(self, preds):
         y = self._y_data_subset
         d = self._g_data_subset