Make inferencer less dependent on sbi internals

brian2modelfitting/inferencer.py

@@ -23,6 +23,7 @@
                                      likelihood_nn,
                                      classifier_nn)
 from sbi.utils.torchutils import BoxUniform
+from sbi.inference.posteriors.direct_posterior import DirectPosterior
 import sbi.analysis
 import sbi.inference
 import torch
@@ -907,10 +908,10 @@ def infer(self, n_samples=None, theta=None, x=None, n_rounds=1,
         else:  # `.infer_step` has been called manually
             x_o = torch.tensor(self.x_o, dtype=torch.float32)
             prior = self.posterior.set_default_x(x_o)
-            if self.posterior._method_family == 'snpe':
-                args = [prior, ]
+            if isinstance(self.posterior, DirectPosterior):
+                args = (prior, )
             else:
-                args = [None, ]
+                args = (None, )
             # generate data if the posterior already exist given proposal
             self.theta = self.generate_training_data(self.n_samples, prior)
             self.x = self.extract_summary_statistics(self.theta, level=1)
@@ -939,10 +940,10 @@ def infer(self, n_samples=None, theta=None, x=None, n_rounds=1,
             if n_rounds > 1 and round < n_rounds - 1:
                 x_o = torch.tensor(self.x_o, dtype=torch.float32)
                 proposal = posterior.set_default_x(x_o)
-                if posterior._method_family == 'snpe':
-                    args = [proposal, ]
+                if isinstance(posterior, DirectPosterior):
+                    args = (proposal, )
                 else:
-                    args = [None, ]
+                    args = (None, )
                 self.theta = self.generate_training_data(self.n_samples,
                                                          proposal)
                 self.x = self.extract_summary_statistics(self.theta, level=1)

brian2modelfitting/tests/test_inferencer.py

@@ -3,6 +3,8 @@
 from brian2modelfitting.inferencer import (Inferencer,
                                            get_param_dict,
                                            calc_prior)
+from sbi.inference.posteriors.direct_posterior import DirectPosterior
+
 import numpy as np
 from numpy.testing import assert_equal, assert_almost_equal
 import pytest
@@ -316,7 +318,7 @@ def test_infer_step(setup_full):
     posterior = inferencer.infer_step(proposal=prior,
                                       n_samples=10,
                                       inference=inference)
-    assert posterior._method_family == 'snpe'
+    assert isinstance(posterior, DirectPosterior)
     assert_equal(np.array(posterior._x_shape), np.array([1, 5]))
 
 
@@ -348,4 +350,4 @@ def test_load_posterior(setup):
     inferencer.infer(n_samples=10, g=[1*nS, 100*nS])
     inferencer.save_posterior('posterior.pth')
     posterior_load = inferencer.load_posterior('posterior.pth')
-    assert inferencer.posterior.net == posterior_load.net
+    assert inferencer.posterior.__dict__ == posterior_load.__dict__