Implement suggestions

Author: Martin Ingram

pymc_extras/inference/__init__.py

@@ -12,7 +12,7 @@
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
 
-from pymc_extras.inference.deterministic_advi.dadvi import fit_deterministic_advi
+from pymc_extras.inference.dadvi.dadvi import fit_dadvi
 from pymc_extras.inference.fit import fit
 from pymc_extras.inference.laplace_approx.find_map import find_MAP
 from pymc_extras.inference.laplace_approx.laplace import fit_laplace
@@ -23,5 +23,5 @@
     "fit",
     "fit_laplace",
     "fit_pathfinder",
-    "fit_deterministic_advi",
+    "fit_dadvi",
 ]

pymc_extras/inference/deterministic_advi/__init__.py renamed to pymc_extras/inference/dadvi/__init__.py

@@ -5,6 +5,7 @@
 import pytensor.tensor as pt
 import xarray
 
+from better_optimize.constants import minimize_method
 from pymc import DictToArrayBijection, Model, join_nonshared_inputs
 from pymc.backends.arviz import (
     PointFunc,
@@ -21,16 +22,18 @@
 )
 
 
-def fit_deterministic_advi(
+def fit_dadvi(
     model: Model | None = None,
     n_fixed_draws: int = 30,
     random_seed: RandomSeed = None,
     n_draws: int = 1000,
     keep_untransformed: bool = False,
+    method: minimize_method = "trust-ncg",
+    **minimize_kwargs,
 ) -> az.InferenceData:
     """
     Does inference using deterministic ADVI (automatic differentiation
-    variational inference).
+    variational inference), DADVI for short.
 
     For full details see the paper cited in the references:
     https://www.jmlr.org/papers/v25/23-1015.html
@@ -57,6 +60,19 @@ def fit_deterministic_advi(
         Whether or not to keep the unconstrained variables (such as
         logs of positive-constrained parameters) in the output.
 
+    method: str
+        Which optimization method to use. The function calls
+        ``scipy.optimize.minimize``, so any of the methods there can
+        be used. The default is trust-ncg, which uses second-order
+        information and is generally very reliable. Other methods such
+        as L-BFGS-B might be faster but potentially more brittle and
+        may not converge exactly to the optimum.
+
+    minimize_kwargs:
+        Additional keyword arguments to pass to the
+        ``scipy.optimize.minimize`` function. See the documentation of
+        that function for details.
+
     Returns
     -------
     :class:`~arviz.InferenceData`
@@ -90,7 +106,14 @@ def fit_deterministic_advi(
         compute_hess=False,
     )
 
-    result = minimize(f_fused, np.zeros(2 * n_params), method="trust-ncg", jac=True, hessp=f_hessp)
+    result = minimize(
+        f_fused,
+        np.zeros(2 * n_params),
+        method=method,
+        jac=True,
+        hessp=f_hessp,
+        **minimize_kwargs,
+    )
 
     opt_var_params = result.x
     opt_means, opt_log_sds = np.split(opt_var_params, 2)
@@ -151,8 +174,7 @@ def create_dadvi_graph(
     )
 
     var_params = pt.vector(name="eta", shape=(2 * n_params,))
-
-    means , log_sds= pt.split(var_params, 2)
+    means, log_sds = var_params[:n_params], var_params[n_params:]
 
     draw_matrix = pt.constant(draws)
     samples = means + pt.exp(log_sds) * draw_matrix

pymc_extras/inference/deterministic_advi/dadvi.py renamed to pymc_extras/inference/dadvi/dadvi.py

@@ -42,6 +42,6 @@ def fit(method: str, **kwargs) -> az.InferenceData:
         return fit_laplace(**kwargs)
 
     if method == "deterministic_advi":
-        from pymc_extras.inference import fit_deterministic_advi
+        from pymc_extras.inference import fit_dadvi
 
-        return fit_deterministic_advi(**kwargs)
+        return fit_dadvi(**kwargs)