diff --git a/CHANGELOG.md b/CHANGELOG.md
index 5081b5fc..73e8f319 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,20 @@
 
 ## [0.2.2] - 2024-XX-XX
 
+**Documentation**
+
+- Various fixes in documentation, including documenting returned fits.
+
+**Breaking changes**
+
+- The `return_posteriors` argument has been removed and replaced with `return_fit`.
+  An instance of one of two previously internal classes, `ExpectationPropagation`
+  and `BeliefPropagation`, are now returned when `return_fit=True`, and posteriors can
+  be obtained using `fit.node_posteriors()`.
+
+- Topology-only dating (setting `mutation_rate=None`) has been removed for tree sequences
+  of more than one tree, as tests have found that span-weighting the conditional coalescent
+  causes substantial bias.
 
 ## [0.2.1] - 2024-07-31
 
diff --git a/docs/Makefile b/docs/Makefile
index 61dc187d..1b909c7d 100644
--- a/docs/Makefile
+++ b/docs/Makefile
@@ -1,7 +1,7 @@
 # Need to set PYTHONPATH so that we pick up the local tsdate
 PYPATH=$(shell pwd)/../
 TSDATE_VERSION:=$(shell PYTHONPATH=${PYPATH} \
-   python3 -c 'import tsdate; print(tsdate.__version__.split("+")[0])')
+   python -c 'import tsdate; print(tsdate.__version__.split("+")[0])')
 
 BUILDDIR      = _build
 
@@ -16,4 +16,4 @@ dist:
 	PYTHONPATH=${PYPATH} ./build.sh
 
 clean:
-	rm -fR $(BUILDDIR)
\ No newline at end of file
+	rm -fR $(BUILDDIR)
diff --git a/docs/_config.yml b/docs/_config.yml
index 747e3b2a..4b39fcda 100644
--- a/docs/_config.yml
+++ b/docs/_config.yml
@@ -35,6 +35,7 @@ html:
 sphinx:
     extra_extensions:
     - sphinx.ext.autodoc
+    - sphinx.ext.mathjax
     - sphinx.ext.autosummary
     - sphinx.ext.todo
     - sphinx.ext.viewcode
diff --git a/docs/python-api.md b/docs/python-api.md
index 43b51c6b..847a95ec 100644
--- a/docs/python-api.md
+++ b/docs/python-api.md
@@ -32,13 +32,28 @@ This page provides formal documentation for the _tsdate_ Python API.
 .. autofunction:: tsdate.maximization
 ```
 
+## Underlying fit objects
+
+Instances of the classes below are returned by setting `return_fit=True`
+when dating. The fits can be inspected to obtain more detailed results than
+may be present in the returned tree sequence and its metadata. The classes
+are not intended to be instantiated directly.
+
+```{eval-rst}
+.. autoclass:: tsdate.discrete.BeliefPropagation()
+   :members:
+
+.. autoclass:: tsdate.variational.ExpectationPropagation()
+   :members:
+```
+
 ## Prior and Time Discretisation Options
 
 ```{eval-rst}
 .. autofunction:: tsdate.build_prior_grid
 .. autofunction:: tsdate.build_parameter_grid
-.. autoclass:: tsdate.base.NodeGridValues
-.. autodata:: tsdate.base.DEFAULT_APPROX_PRIOR_SIZE
+.. autoclass:: tsdate.node_time_class.NodeTimeValues
+.. autodata:: tsdate.prior.DEFAULT_APPROX_PRIOR_SIZE
 ```
 
 ## Variable population sizes
diff --git a/docs/usage.md b/docs/usage.md
index d89f2385..d51a14b1 100644
--- a/docs/usage.md
+++ b/docs/usage.md
@@ -175,23 +175,20 @@ running _tsdate_:
     The metadata values (currently saved as `mn` and `vr`) need not be constrained by
     the topology of the tree sequence, and should be used in preference
     e.g. to `nodes_time` and `mutations_time` when evaluating the accuracy of _tsdate_. 
-2. The `return_posteriors` parameter can be used when calling {func}`tsdate.date`, which
-    then returns both the dated tree sequence and a dictionary specifying the posterior
-    distributions.
-
-<!--
-The returned posterior is a dictionary keyed by integer node ID, with values representing the
-probability distribution of times. This can be read in to a [pandas](https://pandas.pydata.org)
-dataframe:
+2. The `return_fit` parameter can be used when calling {func}`tsdate.date`, which
+    then returns both the dated tree sequence and a fit object. This object can then be
+    queried for the unconstrained posterior distributions using e.g. `.node_posteriors()`
+    which can be read in to a [pandas](https://pandas.pydata.org) dataframe, as below:
 
 ```{code-cell} ipython3
 import pandas as pd
-redated_ts, posteriors = tsdate.date(
-    sim_ts, mutation_rate=1e-6, method="inside_outside", return_posteriors=True)
-posteriors_df = pd.DataFrame(posteriors)
-posteriors_df.head()  # Show the dataframe
+redated_ts, fit = tsdate.date(
+    sim_ts, mutation_rate=1e-6, return_fit=True)
+posteriors_df = pd.DataFrame(fit.node_posteriors())  # mutation_posteriors() also available
+posteriors_df.tail()  # Show the dataframe
 ```
 
+<!--
 Since we are using a {ref}`sec_methods_discrete_time` method, each node
 (numbered column of the dataframe) is associated with a vector of probabilities
 that sum to one: each cell gives the probability that the time of the node
diff --git a/tests/test_accuracy.py b/tests/test_accuracy.py
index 8e06ad3d..82324751 100644
--- a/tests/test_accuracy.py
+++ b/tests/test_accuracy.py
@@ -108,9 +108,7 @@ def test_basic(
         assert sim_mutations_parameters["command"] == "sim_mutations"
         mu = sim_mutations_parameters["rate"]
 
-        dts, posteriors = tsdate.inside_outside(
-            ts, population_size=Ne, mutation_rate=mu, return_posteriors=True
-        )
+        dts = tsdate.inside_outside(ts, population_size=Ne, mutation_rate=mu)
         # make sure we can read node metadata - old tsdate versions didn't set a schema
         if dts.table_metadata_schemas.node.schema is None:
             tables = dts.dump_tables()
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 2d1d5d5e..05183f8c 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -232,10 +232,9 @@ def test_no_output_variational_gamma(self, tmp_path, capfd):
 
     @pytest.mark.parametrize(("flag", "log_status"), logging_flags.items())
     def test_verbosity(self, tmp_path, caplog, flag, log_status):
-        popsize = 10000
         ts = msprime.simulate(
             10,
-            Ne=popsize,
+            Ne=10000,
             mutation_rate=1e-8,
             recombination_rate=1e-8,
             length=2e4,
@@ -246,7 +245,9 @@ def test_verbosity(self, tmp_path, caplog, flag, log_status):
         caplog.set_level(getattr(logging, log_status))
         # either tsdate preprocess or tsdate date (in_out method has debug asserts)
         self.run_tsdate_cli(tmp_path, ts, flag, cmd="preprocess")
-        self.run_tsdate_cli(tmp_path, ts, f"-n 10 --method inside_outside {flag}")
+        self.run_tsdate_cli(
+            tmp_path, ts, f"--mutation-rate 1e-8 --rescaling-intervals 0 {flag}"
+        )
         assert log_status in caplog.text
 
     @pytest.mark.parametrize(
diff --git a/tests/test_functions.py b/tests/test_functions.py
index e68bddca..ece17320 100644
--- a/tests/test_functions.py
+++ b/tests/test_functions.py
@@ -38,15 +38,10 @@
 from utility_functions import constrain_ages_topo
 
 import tsdate
-from tsdate import base
-from tsdate.core import (
-    DiscreteTimeMethod,
-    InOutAlgorithms,
-    InsideOutsideMethod,
-    Likelihoods,
-    LogLikelihoods,
-)
+from tsdate.core import DiscreteTimeMethod, InsideOutsideMethod
 from tsdate.demography import PopulationSizeHistory
+from tsdate.discrete import BeliefPropagation, Likelihoods, LogLikelihoods
+from tsdate.node_time_class import LIN_GRID, LOG_GRID, NodeTimeValues
 from tsdate.prior import (
     ConditionalCoalescentTimes,
     MixturePrior,
@@ -744,19 +739,19 @@ def test_log_tri_functions(self):
                     )
 
 
-class TestNodeGridValuesClass:
+class TestNodeTimeValuesClass:
     def test_init(self):
         num_nodes = 5
         ids = np.array([3, 4])
         timepoints = np.array(range(10))
-        store = base.NodeGridValues(num_nodes, ids, timepoints, fill_value=6)
+        store = NodeTimeValues(num_nodes, ids, timepoints, fill_value=6)
         assert store.grid_data.shape == (len(ids), len(timepoints))
         assert len(store.fixed_data) == (num_nodes - len(ids))
         assert np.all(store.grid_data == 6)
         assert np.all(store.fixed_data == 6)
 
         ids = np.array([3, 4], dtype=np.int32)
-        store = base.NodeGridValues(num_nodes, ids, timepoints, fill_value=5)
+        store = NodeTimeValues(num_nodes, ids, timepoints, fill_value=5)
         assert store.grid_data.shape == (len(ids), len(timepoints))
         assert len(store.fixed_data) == num_nodes - len(ids)
         assert np.all(store.fixed_data == 5)
@@ -767,7 +762,7 @@ def test_set_and_get(self):
         fill = {}
         rng = np.random.default_rng(1)
         for ids in ([3, 4], []):
-            store = base.NodeGridValues(
+            store = NodeTimeValues(
                 num_nodes, np.array(ids, dtype=np.int32), np.array(range(grid_size))
             )
             for i in range(num_nodes):
@@ -781,52 +776,50 @@ def test_set_and_get(self):
     def test_bad_init(self):
         ids = [3, 4]
         with pytest.raises(ValueError):
-            base.NodeGridValues(3, np.array(ids), np.array([0, 1.2, 2]))
+            NodeTimeValues(3, np.array(ids), np.array([0, 1.2, 2]))
         with pytest.raises(AttributeError):
-            base.NodeGridValues(5, np.array(ids), -1)
+            NodeTimeValues(5, np.array(ids), -1)
         with pytest.raises(ValueError):
-            base.NodeGridValues(5, np.array([-1]), np.array([0, 1.2, 2]))
+            NodeTimeValues(5, np.array([-1]), np.array([0, 1.2, 2]))
         with pytest.raises(ValueError, match="must be a 1D numpy array"):
-            base.NodeGridValues(5, np.array([[1, 2], [3, 4]]), np.array([0, 1.2, 2]))
+            NodeTimeValues(5, np.array([[1, 2], [3, 4]]), np.array([0, 1.2, 2]))
 
     def test_clone(self):
         num_nodes = 10
         grid_size = 2
         ids = [3, 4]
-        orig = base.NodeGridValues(num_nodes, np.array(ids), np.array(range(grid_size)))
+        orig = NodeTimeValues(num_nodes, np.array(ids), np.array(range(grid_size)))
         orig[3] = np.array([1, 2])
         orig[4] = np.array([4, 3])
         orig[0] = 1.5
         orig[9] = 2.5
         # test with np.zeros
-        clone = base.NodeGridValues.clone_with_new_data(orig, 0)
+        clone = NodeTimeValues.clone_with_new_data(orig, 0)
         assert clone.grid_data.shape == orig.grid_data.shape
         assert clone.fixed_data.shape == orig.fixed_data.shape
         assert np.all(clone.grid_data == 0)
         assert np.all(clone.fixed_data == 0)
         # test with something else
-        clone = base.NodeGridValues.clone_with_new_data(orig, 5)
+        clone = NodeTimeValues.clone_with_new_data(orig, 5)
         assert clone.grid_data.shape == orig.grid_data.shape
         assert clone.fixed_data.shape == orig.fixed_data.shape
         assert np.all(clone.grid_data == 5)
         assert np.all(clone.fixed_data == 5)
         # test with different
         scalars = np.arange(num_nodes - len(ids))
-        clone = base.NodeGridValues.clone_with_new_data(orig, 0, scalars)
+        clone = NodeTimeValues.clone_with_new_data(orig, 0, scalars)
         assert clone.grid_data.shape == orig.grid_data.shape
         assert clone.fixed_data.shape == orig.fixed_data.shape
         assert np.all(clone.grid_data == 0)
         assert np.all(clone.fixed_data == scalars)
 
-        clone = base.NodeGridValues.clone_with_new_data(orig, np.array([[1, 2], [4, 3]]))
+        clone = NodeTimeValues.clone_with_new_data(orig, np.array([[1, 2], [4, 3]]))
         for i in range(num_nodes):
             if i in ids:
                 assert np.all(clone[i] == orig[i])
             else:
                 assert np.isnan(clone[i])
-        clone = base.NodeGridValues.clone_with_new_data(
-            orig, np.array([[1, 2], [4, 3]]), 0
-        )
+        clone = NodeTimeValues.clone_with_new_data(orig, np.array([[1, 2], [4, 3]]), 0)
         for i in range(num_nodes):
             if i in ids:
                 assert np.all(clone[i] == orig[i])
@@ -836,14 +829,14 @@ def test_clone(self):
     def test_bad_clone(self):
         num_nodes = 10
         ids = [3, 4]
-        orig = base.NodeGridValues(num_nodes, np.array(ids), np.array([0, 1.2]))
+        orig = NodeTimeValues(num_nodes, np.array(ids), np.array([0, 1.2]))
         with pytest.raises(ValueError):
-            base.NodeGridValues.clone_with_new_data(
+            NodeTimeValues.clone_with_new_data(
                 orig,
                 np.array([[1, 2, 3], [4, 5, 6]]),
             )
         with pytest.raises(ValueError):
-            base.NodeGridValues.clone_with_new_data(
+            NodeTimeValues.clone_with_new_data(
                 orig,
                 0,
                 np.array([[1, 2], [4, 5]]),
@@ -853,10 +846,10 @@ def test_convert_to_probs(self):
         num_nodes = 10
         ids = [3, 4]
         make_nan_row = 4
-        orig = base.NodeGridValues(num_nodes, np.array(ids), np.array([0, 1.2]), 1)
+        orig = NodeTimeValues(num_nodes, np.array(ids), np.array([0, 1.2]), 1)
         orig[make_nan_row][0] = np.nan
         assert np.all(np.isnan(orig[make_nan_row]) == [True, False])
-        orig.force_probability_space(base.LIN)
+        orig.force_probability_space(LIN_GRID)
         orig.to_probabilities()
         for n in orig.nonfixed_nodes:
             if n == make_nan_row:
@@ -871,41 +864,39 @@ def test_convert_back(self):
         ids = [3, 4]
         timepoints = np.array([0, 1.2])
         fill_val = 0.1
-        orig = base.NodeGridValues(num_nodes, np.array(ids), timepoints, fill_val)
+        orig = NodeTimeValues(num_nodes, np.array(ids), timepoints, fill_val)
         assert np.allclose(orig.grid_data, fill_val)
-        assert orig.probability_space == base.LIN
-        orig.force_probability_space(base.LOG)
-        assert orig.probability_space == base.LOG
-        orig.force_probability_space(base.LOG)  # repeat, to check nothing changes
-        assert orig.probability_space == base.LOG
-        orig.force_probability_space(base.LIN)
-        assert orig.probability_space == base.LIN
-        orig.force_probability_space(base.LIN)  # repeat, to check nothing changes
-        assert orig.probability_space == base.LIN
+        assert orig.probability_space == LIN_GRID
+        orig.force_probability_space(LOG_GRID)
+        assert orig.probability_space == LOG_GRID
+        orig.force_probability_space(LOG_GRID)  # repeat, to check nothing changes
+        assert orig.probability_space == LOG_GRID
+        orig.force_probability_space(LIN_GRID)
+        assert orig.probability_space == LIN_GRID
+        orig.force_probability_space(LIN_GRID)  # repeat, to check nothing changes
+        assert orig.probability_space == LIN_GRID
         assert np.allclose(orig.grid_data, fill_val)
 
     def test_cannot_convert_to_probs(self):
         # No class implemention of logsumexp to convert to probabilities in log space
         num_nodes = 10
         ids = [3, 4]
-        orig = base.NodeGridValues(num_nodes, np.array(ids), np.array([0, 1.2]))
-        orig.force_probability_space(base.LOG)
+        orig = NodeTimeValues(num_nodes, np.array(ids), np.array([0, 1.2]))
+        orig.force_probability_space(LOG_GRID)
         with pytest.raises(NotImplementedError, match="linear space"):
             orig.to_probabilities()
 
     def test_cannot_force_probability_space(self):
         num_nodes = 10
         ids = [3, 4]
-        orig = base.NodeGridValues(num_nodes, np.array(ids), np.array([0, 1.2]))
-        assert orig.probability_space == base.LIN
-        with pytest.raises(TypeError, match="Cannot force linear"):
-            orig.force_probability_space(base.GAMMA_PAR)
+        orig = NodeTimeValues(num_nodes, np.array(ids), np.array([0, 1.2]))
+        assert orig.probability_space == LIN_GRID
 
     def test_bad_probability_space(self):
         num_nodes = 10
         ids = [3, 4]
-        orig = base.NodeGridValues(num_nodes, np.array(ids), np.array([0, 1.2]))
-        assert orig.probability_space == base.LIN
+        orig = NodeTimeValues(num_nodes, np.array(ids), np.array([0, 1.2]))
+        assert orig.probability_space == LIN_GRID
         with pytest.raises(ValueError, match="Bad probability space"):
             orig.force_probability_space("bad_space")
 
@@ -919,7 +910,7 @@ def test_nonmatching_prior_vs_lik_timepoints(self):
         priors = tsdate.build_prior_grid(ts, Ne, timepoints1)
         lls = Likelihoods(ts, timepoints2)
         with pytest.raises(ValueError, match="timepoints"):
-            InOutAlgorithms(priors, lls)
+            BeliefPropagation(priors, lls)
 
     def test_nonmatching_prior_vs_lik_fixednodes(self):
         ts1 = utility_functions.single_tree_ts_n3()
@@ -929,7 +920,7 @@ def test_nonmatching_prior_vs_lik_fixednodes(self):
         priors = tsdate.build_prior_grid(ts1, Ne, timepoints)
         lls = Likelihoods(ts2, priors.timepoints)
         with pytest.raises(ValueError, match="fixed"):
-            InOutAlgorithms(priors, lls)
+            BeliefPropagation(priors, lls)
 
 
 class TestInsideAlgorithm:
@@ -952,7 +943,7 @@ def run_inside_algorithm(
         else:
             lls = Likelihoods(ts, priors.timepoints, mut_rate, eps=eps)
         lls.precalculate_mutation_likelihoods()
-        algo = InOutAlgorithms(priors, lls)
+        algo = BeliefPropagation(priors, lls)
         marg_lik = algo.inside_pass(standardize=standardize)
         return algo, priors, marg_lik
 
@@ -1077,17 +1068,15 @@ def test_two_tree_mutation_ts(self):
 
     def test_dangling_fails(self):
         ts = utility_functions.single_tree_ts_n2_dangling()
-        # print(ts.draw_text())
-        # print("Samples:", ts.samples())
         Ne = 0.5
         with pytest.raises(ValueError, match="simplified"):
             tsdate.build_prior_grid(ts, Ne, timepoints=np.array([0, 1.2, 2]))
         # mut_rate = 1
         # eps = 1e-6
         # lls = Likelihoods(ts, priors.timepoints, mut_rate, eps)
-        # algo = InOutAlgorithms(priors, lls)
+        # fit = BeliefPropagation(priors, lls)
         # with pytest.raises(ValueError, match="dangling"):
-        #     algo.inside_pass()
+        #     fit.inside_pass()
 
     def test_standardize_marginal_likelihood(self):
         ts = utility_functions.two_tree_mutation_ts()
@@ -1117,7 +1106,7 @@ def run_outside_algorithm(
         eps = 1e-6
         lls = Likelihoods(ts, grid, mut_rate, eps=eps)
         lls.precalculate_mutation_likelihoods()
-        algo = InOutAlgorithms(prior_vals, lls)
+        algo = BeliefPropagation(prior_vals, lls)
         algo.inside_pass()
         algo.outside_pass(standardize=standardize, ignore_oldest_root=ignore_oldest_root)
         return algo
@@ -1159,7 +1148,7 @@ def test_outside_before_inside_fails(self):
         mut_rate = 1
         lls = Likelihoods(ts, priors.timepoints, mut_rate)
         lls.precalculate_mutation_likelihoods()
-        algo = InOutAlgorithms(priors, lls)
+        algo = BeliefPropagation(priors, lls)
         with pytest.raises(RuntimeError):
             algo.outside_pass()
 
@@ -1217,9 +1206,9 @@ def find_posterior(self, ts, prior_distr):
         eps = 1e-6
         lls = Likelihoods(ts, grid, mut_rate, eps=eps)
         lls.precalculate_mutation_likelihoods()
-        algo = InOutAlgorithms(prior_vals, lls)
+        algo = BeliefPropagation(prior_vals, lls)
         algo.inside_pass()
-        posterior = algo.outside_pass(standardize=False)
+        algo.outside_pass(standardize=False)
         assert np.array_equal(
             np.sum(algo.inside.grid_data * algo.outside.grid_data, axis=1),
             np.sum(algo.inside.grid_data * algo.outside.grid_data, axis=1),
@@ -1228,38 +1217,38 @@ def find_posterior(self, ts, prior_distr):
             np.sum(algo.inside.grid_data * algo.outside.grid_data, axis=1),
             np.sum(algo.inside.grid_data[-1]),
         )
-        return posterior, algo
+        return algo
 
     def test_one_tree_n2(self):
         ts = utility_functions.single_tree_ts_n2()
         for distr in ("gamma", "lognorm"):
-            posterior, algo = self.find_posterior(ts, distr)
+            self.find_posterior(ts, distr)
 
     def test_one_tree_n3(self):
         ts = utility_functions.single_tree_ts_n3()
         for distr in ("gamma", "lognorm"):
-            posterior, algo = self.find_posterior(ts, distr)
+            self.find_posterior(ts, distr)
 
     def test_one_tree_n4(self):
         ts = utility_functions.single_tree_ts_n4()
         for distr in ("gamma", "lognorm"):
-            posterior, algo = self.find_posterior(ts, distr)
+            self.find_posterior(ts, distr)
 
     def test_one_tree_n3_mutation(self):
         ts = utility_functions.single_tree_ts_mutation_n3()
         for distr in ("gamma", "lognorm"):
-            posterior, algo = self.find_posterior(ts, distr)
+            self.find_posterior(ts, distr)
 
     def test_polytomy_tree(self):
         ts = utility_functions.polytomy_tree_ts()
         for distr in ("gamma", "lognorm"):
-            posterior, algo = self.find_posterior(ts, distr)
+            self.find_posterior(ts, distr)
 
     def test_tree_with_unary_nodes(self):
         ts = utility_functions.single_tree_ts_with_unary()
         for distr in ("gamma", "lognorm"):
             with pytest.raises(ValueError, match="unary"):
-                posterior, algo = self.find_posterior(ts, distr)
+                self.find_posterior(ts, distr)
 
 
 class TestGilTree:
@@ -1284,7 +1273,7 @@ def test_gil_tree(self):
             eps = 0.01
             lls = Likelihoods(ts, grid, mut_rate, eps=eps, standardize=False)
             lls.precalculate_mutation_likelihoods()
-            algo = InOutAlgorithms(prior_vals, lls)
+            algo = BeliefPropagation(prior_vals, lls)
             algo.inside_pass(standardize=False, cache_inside=cache_inside)
             algo.outside_pass(standardize=False)
             assert np.allclose(
@@ -1316,7 +1305,7 @@ def edges_ordering(self, ts, fn):
             fixed_node_set=fixed_nodes,
             progress=False,
         )
-        dynamic_prog = InOutAlgorithms(priors, liklhd, progress=False)
+        dynamic_prog = BeliefPropagation(priors, liklhd, progress=False)
         if fn == "outside_pass":
             edges_by_child = dynamic_prog.edges_by_child_desc()
             seen_children = list()
@@ -1388,16 +1377,17 @@ def run_outside_maximization(self, ts, prior_distr="lognorm"):
         eps = 1e-6
         lls = Likelihoods(ts, priors.timepoints, mut_rate, eps=eps)
         lls.precalculate_mutation_likelihoods()
-        algo = InOutAlgorithms(priors, lls)
+        algo = BeliefPropagation(priors, lls)
         algo.inside_pass()
-        return lls, algo, algo.outside_maximization(eps=eps)
+        algo.outside_maximization(eps=eps)
+        return lls, algo, algo.posterior_mean
 
     def test_outside_maximization_error(self):
         ts = utility_functions.single_tree_ts_n2()
         priors = tsdate.build_prior_grid(ts, 0.5)
         lls = Likelihoods(ts, priors.timepoints, 1)
         lls.precalculate_mutation_likelihoods()
-        algo = InOutAlgorithms(priors, lls)
+        algo = BeliefPropagation(priors, lls)
         with pytest.raises(RuntimeError, match="not yet run"):
             algo.outside_maximization(eps=1e-6)
 
@@ -1541,29 +1531,26 @@ class TestDiscretisedMeanVar:
     def test_discretised_mean_var(self):
         ts = utility_functions.single_tree_ts_n2()
         for distr in ("gamma", "lognorm"):
-            posterior, algo = TestTotalFunctionalValueTree().find_posterior(ts, distr)
-            mn_post, vr_post = DiscreteTimeMethod.mean_var(ts, posterior)
+            algo = TestTotalFunctionalValueTree().find_posterior(ts, distr)
+            p_grid = algo.posterior_grid
+            mn_post, vr_post = DiscreteTimeMethod.mean_var(ts, p_grid)
             assert np.array_equal(
                 mn_post,
-                [
-                    0,
-                    0,
-                    np.sum(posterior.timepoints * posterior[2]) / np.sum(posterior[2]),
-                ],
+                [0, 0, np.sum(p_grid.timepoints * p_grid[2]) / np.sum(p_grid[2])],
             )
 
     def test_node_metadata_inside_outside(self):
         ts = msprime.simulate(
             10, mutation_rate=1, recombination_rate=1, length=20, random_seed=12
         )
-        algorithm = InsideOutsideMethod(ts, mutation_rate=None, population_size=10000)
+        algorithm = InsideOutsideMethod(ts, mutation_rate=1, population_size=10000)
         mn_post, *_ = algorithm.run(
             eps=1e-6,
             outside_standardize=True,
             ignore_oldest_root=False,
-            probability_space=tsdate.base.LOG,
+            probability_space=tsdate.node_time_class.LOG_GRID,
         )
-        dts = tsdate.inside_outside(ts, population_size=10000, mutation_rate=None)
+        dts = tsdate.inside_outside(ts, population_size=10000, mutation_rate=1)
         unconstr_mn = [nd.metadata["mn"] for nd in dts.nodes() if "mn" in nd.metadata]
         assert np.allclose(unconstr_mn, mn_post)
         assert np.all(dts.tables.nodes.time >= mn_post)
@@ -1572,7 +1559,7 @@ def test_node_metadata_variational_gamma(self):
         ts = msprime.simulate(
             10, mutation_rate=1, recombination_rate=1, length=20, random_seed=12
         )
-        dts, _ = tsdate.variational_gamma(ts, mutation_rate=1, return_posteriors=True)
+        dts = tsdate.variational_gamma(ts, mutation_rate=1)
         unconstr_mn = [nd.metadata["mn"] for nd in dts.nodes() if "mn" in nd.metadata]
         # TODO: compare against the returned posterior means too
         # see
@@ -1830,13 +1817,13 @@ def test_node_selection_param(self):
 
     def test_sites_time_insideoutside(self):
         ts = utility_functions.two_tree_mutation_ts()
-        dated = tsdate.inside_outside(ts, mutation_rate=None, population_size=1)
-        algorithm = InsideOutsideMethod(ts, mutation_rate=None, population_size=1)
+        dated = tsdate.inside_outside(ts, mutation_rate=1, population_size=1)
+        algorithm = InsideOutsideMethod(ts, mutation_rate=1, population_size=1)
         mn_post, *_ = algorithm.run(
             eps=1e-6,
             outside_standardize=True,
             ignore_oldest_root=False,
-            probability_space=tsdate.base.LOG,
+            probability_space=LOG_GRID,
         )
         assert np.array_equal(
             mn_post[ts.tables.mutations.node],
@@ -1936,14 +1923,14 @@ def test_sites_time_simulated_inside_outside(self):
         ts = msprime.simulate(
             10, mutation_rate=1, recombination_rate=1, length=20, random_seed=12
         )
-        algorithm = InsideOutsideMethod(ts, mutation_rate=None, population_size=10000)
+        algorithm = InsideOutsideMethod(ts, mutation_rate=1, population_size=10000)
         mn_post, *_ = algorithm.run(
             eps=1e-6,
             outside_standardize=True,
             ignore_oldest_root=False,
-            probability_space=tsdate.base.LOG,
+            probability_space=tsdate.node_time_class.LOG_GRID,
         )
-        dts = tsdate.inside_outside(ts, mutation_rate=None, population_size=10000)
+        dts = tsdate.inside_outside(ts, mutation_rate=1, population_size=10000)
         assert np.allclose(
             mn_post[ts.tables.mutations.node],
             tsdate.sites_time_from_ts(dts, unconstrained=True, min_time=0),
@@ -1958,9 +1945,10 @@ def test_sites_time_simulated_variational_gamma(self):
         ts = msprime.simulate(
             10, mutation_rate=1, recombination_rate=1, length=20, random_seed=12
         )
-        dts, pst = tsdate.variational_gamma(ts, mutation_rate=1, return_posteriors=True)
+        dts, fit = tsdate.variational_gamma(ts, mutation_rate=1, return_fit=True)
+        posteriors = fit.node_posteriors()
         assert np.allclose(
-            pst[ts.tables.mutations.node],
+            posteriors[ts.tables.mutations.node],
             tsdate.sites_time_from_ts(dts, unconstrained=True, min_time=0),
         )
         assert np.allclose(
diff --git a/tests/test_inference.py b/tests/test_inference.py
index cd53105e..ae6aa262 100644
--- a/tests/test_inference.py
+++ b/tests/test_inference.py
@@ -34,9 +34,9 @@
 import utility_functions
 
 import tsdate
-from tsdate.base import LIN, LOG
 from tsdate.demography import PopulationSizeHistory
 from tsdate.evaluation import remove_edges, unsupported_edges
+from tsdate.node_time_class import LIN_GRID, LOG_GRID
 
 
 class TestConstants:
@@ -65,11 +65,25 @@ def test_no_population_size(self):
             tsdate.inside_outside(ts, mutation_rate=None)
 
     def test_no_mutation(self):
-        ts = utility_functions.two_tree_mutation_ts()
-        with pytest.raises(ValueError, match="method requires mutation rate"):
-            tsdate.date(ts, method="maximization", population_size=1, mutation_rate=None)
-        with pytest.raises(ValueError, match="method requires mutation rate"):
-            tsdate.date(ts, method="variational_gamma", mutation_rate=None)
+        for ts in (
+            utility_functions.two_tree_mutation_ts(),
+            utility_functions.single_tree_ts_mutation_n3(),
+        ):
+            with pytest.raises(ValueError, match="method requires mutation rate"):
+                tsdate.date(
+                    ts, method="maximization", population_size=1, mutation_rate=None
+                )
+            with pytest.raises(ValueError, match="method requires mutation rate"):
+                tsdate.date(ts, method="variational_gamma", mutation_rate=None)
+            if ts.num_trees > 1:
+                with pytest.raises(NotImplementedError, match="more than one tree"):
+                    tsdate.date(
+                        ts, method="inside_outside", population_size=1, mutation_rate=None
+                    )
+            else:
+                tsdate.date(
+                    ts, method="inside_outside", population_size=1, mutation_rate=None
+                )
 
     def test_not_needed_population_size(self):
         ts = utility_functions.two_tree_mutation_ts()
@@ -87,9 +101,9 @@ def test_both_ne_and_population_size_specified(self):
         ts = utility_functions.two_tree_mutation_ts()
         with pytest.raises(ValueError, match="Only provide one of Ne"):
             tsdate.inside_outside(
-                ts, mutation_rate=None, population_size=PopulationSizeHistory(1), Ne=1
+                ts, mutation_rate=1, population_size=PopulationSizeHistory(1), Ne=1
             )
-        tsdate.inside_outside(ts, mutation_rate=None, Ne=PopulationSizeHistory(1))
+        tsdate.inside_outside(ts, mutation_rate=1, Ne=PopulationSizeHistory(1))
 
     def test_inside_outside_dangling_failure(self):
         ts = utility_functions.single_tree_ts_n2_dangling()
@@ -115,7 +129,7 @@ def test_variational_gamma_unary_failure(self):
         with pytest.raises(ValueError, match="unary"):
             tsdate.variational_gamma(ts, mutation_rate=1)
 
-    @pytest.mark.parametrize("probability_space", [LOG, LIN])
+    @pytest.mark.parametrize("probability_space", [LOG_GRID, LIN_GRID])
     @pytest.mark.parametrize("mu", [None, 1])
     def test_fails_with_recombination(self, probability_space, mu):
         ts = utility_functions.two_tree_mutation_ts()
@@ -138,55 +152,103 @@ def test_default_alternative_time_units(self):
         ts = tsdate.date(ts, mutation_rate=1, time_units="years")
         assert ts.time_units == "years"
 
-    def test_no_posteriors(self):
+    def test_deprecated_return_posteriors(self):
         ts = utility_functions.two_tree_mutation_ts()
-        with pytest.raises(ValueError, match="Cannot return posterior"):
-            tsdate.date(
-                ts,
-                population_size=1,
-                return_posteriors=True,
-                method="maximization",
-                mutation_rate=1,
-            )
+        with pytest.raises(ValueError, match="deprecated"):
+            tsdate.date(ts, return_posteriors=True, mutation_rate=1)
+
+    def test_return_fit(self):
+        ts = utility_functions.two_tree_mutation_ts()
+        _, fit = tsdate.date(ts, return_fit=True, mutation_rate=1)
+        assert hasattr(fit, "node_posteriors")
+
+    def test_no_maximization_posteriors(self):
+        ts = utility_functions.two_tree_mutation_ts()
+        _, fit = tsdate.date(
+            ts,
+            population_size=1,
+            return_fit=True,
+            method="maximization",
+            mutation_rate=1,
+        )
+        with pytest.raises(ValueError, match="outside"):
+            fit.node_posteriors()
 
     def test_discretised_posteriors(self):
         ts = utility_functions.two_tree_mutation_ts()
-        ts, posteriors = tsdate.inside_outside(
-            ts, mutation_rate=None, population_size=1, return_posteriors=True
+        ts, fit = tsdate.inside_outside(
+            ts, mutation_rate=1, population_size=1, return_fit=True
         )
-        assert len(posteriors) == ts.num_nodes - ts.num_samples + 1
-        assert len(posteriors["time"]) > 0
+        posteriors = fit.node_posteriors()
+        assert len(posteriors) == ts.num_nodes
+        assert len(posteriors[0]) > 0
         for node in ts.nodes():
-            if not node.is_sample():
-                assert node.id in posteriors
-                assert len(posteriors[node.id]) == len(posteriors["time"])
-                assert np.isclose(np.sum(posteriors[node.id]), 1)
-
-    def test_variational_posteriors(self):
-        """
-        There are no time-gridded posteriors returned by variational gamma,
-        Output is currently None, but see https://github.com/tskit-dev/tsdate/issues/388
-        """
+            nd_vals = np.array(list(posteriors[node.id]))
+            if node.is_sample():
+                assert np.all(np.isnan(nd_vals))
+            else:
+                assert np.isclose(np.sum(nd_vals), 1)
+
+    def test_variational_node_posteriors(self):
         ts = utility_functions.two_tree_mutation_ts()
-        ts, posteriors = tsdate.date(
+        ts, fit = tsdate.date(
             ts,
             mutation_rate=1e-2,
             method="variational_gamma",
-            return_posteriors=True,
+            return_fit=True,
         )
-        assert posteriors is None
+        posteriors = fit.node_posteriors()
+        for nd in ts.nodes():
+            mn, vr = posteriors[nd.id]
+            assert posteriors["mean"][nd.id] == mn
+            assert posteriors["variance"][nd.id] == vr
+            assert np.isclose(nd.metadata["mn"], mn)
+            assert np.isclose(nd.metadata["vr"], vr)
+
+    def test_variational_mutation_posteriors(self):
+        ts = utility_functions.two_tree_mutation_ts()
+        ts, fit = tsdate.date(
+            ts,
+            mutation_rate=1e-2,
+            method="variational_gamma",
+            return_fit=True,
+        )
+        posteriors = fit.mutation_posteriors()
+        for mut in ts.mutations():
+            mn, vr = posteriors[mut.id]
+            assert posteriors["mean"][mut.id] == mn
+            assert posteriors["variance"][mut.id] == vr
+            assert np.isclose(mut.metadata["mn"], mn)
+            assert np.isclose(mut.metadata["vr"], vr)
+
+    def test_variational_mean_edge_logconst(self):
+        # This should give a guide to EP convergence
+        ts = utility_functions.two_tree_mutation_ts()
+        ts, fit = tsdate.date(
+            ts,
+            mutation_rate=1e-2,
+            method="variational_gamma",
+            return_fit=True,
+        )
+        # Should give a guide to EP convergence. Testing on this example gives
+        # [16.082488, 15.586616, 15.589837, 15.589840, 15.589840, ...] (6 d.p.)
+        # check to that DP by multiplying by 1e6 and rounding
+        test_vals = np.array([16082488, 15586616, 15589837, 15589840, 15589840])
+        obs = np.round(np.array(fit.mean_edge_logconst) * 1e6)
+        assert np.all(obs[0:5] == test_vals)
+        assert np.all(obs[5:] == test_vals[-1])
 
     def test_marginal_likelihood(self):
         ts = utility_functions.two_tree_mutation_ts()
         _, _, marg_lik = tsdate.inside_outside(
             ts,
-            mutation_rate=None,
+            mutation_rate=1,
             population_size=1,
-            return_posteriors=True,
+            return_fit=True,
             return_likelihood=True,
         )
         _, marg_lik_again = tsdate.inside_outside(
-            ts, mutation_rate=None, population_size=1, return_likelihood=True
+            ts, mutation_rate=1, population_size=1, return_likelihood=True
         )
         assert marg_lik == marg_lik_again
 
@@ -194,20 +256,12 @@ def test_intervals(self):
         ts = utility_functions.two_tree_ts()
         long_ts = utility_functions.two_tree_ts_extra_length()
         keep_ts = long_ts.keep_intervals([[0.0, 1.0]])
-        delete_ts = long_ts.delete_intervals([[1.0, 1.5]])
-        dated_ts = tsdate.inside_outside(ts, mutation_rate=None, population_size=1)
-        dated_keep_ts = tsdate.inside_outside(
-            keep_ts, mutation_rate=None, population_size=1
-        )
-        dated_deleted_ts = tsdate.inside_outside(
-            delete_ts, mutation_rate=None, population_size=1
-        )
-        assert np.allclose(
-            dated_ts.tables.nodes.time[:], dated_keep_ts.tables.nodes.time[:]
-        )
-        assert np.allclose(
-            dated_ts.tables.nodes.time[:], dated_deleted_ts.tables.nodes.time[:]
-        )
+        del_ts = long_ts.delete_intervals([[1.0, 1.5]])
+        dat_ts = tsdate.inside_outside(ts, mutation_rate=1, population_size=1)
+        dat_keep_ts = tsdate.inside_outside(keep_ts, mutation_rate=1, population_size=1)
+        dat_del_ts = tsdate.inside_outside(del_ts, mutation_rate=1, population_size=1)
+        assert np.allclose(dat_ts.tables.nodes.time[:], dat_keep_ts.tables.nodes.time[:])
+        assert np.allclose(dat_ts.tables.nodes.time[:], dat_del_ts.tables.nodes.time[:])
 
 
 class TestSimulated:
@@ -284,13 +338,13 @@ def test_linear_space(self):
             ts, population_size=10000, timepoints=10, approximate_priors=None
         )
         dated_ts = tsdate.inside_outside(
-            ts, mutation_rate=1e-8, priors=priors, probability_space=LIN
+            ts, mutation_rate=1e-8, priors=priors, probability_space=LIN_GRID
         )
         maximized_ts = tsdate.maximization(
             ts,
             mutation_rate=1e-8,
             priors=priors,
-            probability_space=LIN,
+            probability_space=LIN_GRID,
         )
         self.ts_equal_except_times(ts, dated_ts)
         self.ts_equal_except_times(ts, maximized_ts)
diff --git a/tsdate/cli.py b/tsdate/cli.py
index 541ffd13..6cf3b792 100644
--- a/tsdate/cli.py
+++ b/tsdate/cli.py
@@ -155,7 +155,7 @@ def tsdate_cli_parser():
     )
     parser.add_argument(
         "--rescaling-intervals",
-        type=float,
+        type=int,
         help=(
             "The number of time intervals within which to estimate a time scaling "
             f"parameter. Default: None treated as {core.DEFAULT_RESCALING_INTERVALS}"
diff --git a/tsdate/core.py b/tsdate/core.py
index bd5e3605..0cd6ea11 100644
--- a/tsdate/core.py
+++ b/tsdate/core.py
@@ -21,24 +21,18 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 """
-Infer the age of nodes conditional on a tree sequence topology.
+Infer the age of nodes from mutational data, conditional on a tree sequence topology.
 """
 
-import functools
-import itertools
 import logging
-import multiprocessing
-import operator
 import time  # DEBUG
-from collections import defaultdict, namedtuple
+from collections import namedtuple
 
-import numba
 import numpy as np
-import scipy.stats
 import tskit
-from tqdm.auto import tqdm
 
-from . import base, demography, prior, provenance, schemas, util, variational
+from . import demography, discrete, prior, provenance, schemas, util, variational
+from .node_time_class import LIN_GRID, LOG_GRID
 
 logger = logging.getLogger(__name__)
 
@@ -49,830 +43,18 @@
 DEFAULT_EPSILON = 1e-6
 
 
-class Likelihoods:
-    """
-    A class to store and process likelihoods. Likelihoods for edges are stored as a
-    flattened lower triangular matrix of all the possible delta t's. This class also
-    provides methods for accessing this lower triangular matrix, multiplying it, etc.
-
-    If ``standardize`` is true, routines will operate to standardize the likelihoods
-    such that their maximum is one (in linear space) or zero (in log space)
-    """
-
-    probability_space = base.LIN
-    identity_constant = 1.0
-    null_constant = 0.0
-
-    def __init__(
-        self,
-        ts,
-        timepoints,
-        mutation_rate=None,
-        recombination_rate=None,
-        *,
-        eps=0,
-        fixed_node_set=None,
-        standardize=False,
-        progress=False,
-    ):
-        self.ts = ts
-        self.timepoints = timepoints
-        self.fixednodes = set(ts.samples()) if fixed_node_set is None else fixed_node_set
-        self.mut_rate = mutation_rate
-        self.rec_rate = recombination_rate
-        self.standardize = standardize
-        self.grid_size = len(timepoints)
-        self.tri_size = self.grid_size * (self.grid_size + 1) / 2
-        self.ll_mut = {}
-        self.mut_edges = self.get_mut_edges(ts)
-        self.progress = progress
-        # Need to set eps properly in the 2 lines below, to account for values in the
-        # same timeslice
-        self.timediff_lower_tri = np.concatenate(
-            [
-                self.timepoints[time_index] - self.timepoints[0 : time_index + 1] + eps
-                for time_index in np.arange(len(self.timepoints))
-            ]
-        )
-        self.timediff = self.timepoints - self.timepoints[0] + eps
-
-        # The mut_ll contains unpacked (1D) lower triangular matrices. We need to
-        # index this by row and by column index.
-        self.row_indices = []
-        for t in range(self.grid_size):
-            n = np.arange(self.grid_size)
-            self.row_indices.append((((n * (n + 1)) // 2) + t)[t:])
-        self.col_indices = []
-        running_sum = 0  # use this to find the index of the last element of
-        # each column in order to appropriately sum the vv by columns.
-        for i in np.arange(self.grid_size):
-            arr = np.arange(running_sum, running_sum + self.grid_size - i)
-            index = arr[-1]
-            running_sum = index + 1
-            val = arr[0]
-            self.col_indices.append(val)
-
-        # These are used for transforming an array of grid_size into one of tri_size
-        # By repeating elements over rows to form an upper or a lower triangular matrix
-        self.to_lower_tri = np.concatenate(
-            [np.arange(time_idx + 1) for time_idx in np.arange(self.grid_size)]
-        )
-        self.to_upper_tri = np.concatenate(
-            [
-                np.arange(time_idx, self.grid_size)
-                for time_idx in np.arange(self.grid_size + 1)
-            ]
-        )
-
-    @staticmethod
-    def get_mut_edges(ts):
-        """
-        Get the number of mutations on each edge in the tree sequence.
-        """
-        mut_edges = np.zeros(ts.num_edges, dtype=np.int64)
-        for m in ts.mutations():
-            if m.edge != tskit.NULL:
-                mut_edges[m.edge] += 1
-        return mut_edges
-
-    @staticmethod
-    def _lik(muts, span, dt, mutation_rate, standardize=True):
-        """
-        The likelihood of an edge given a number of mutations, as set of time deltas (dt)
-        and a span. This is a static function to allow parallelization
-        """
-        ll = scipy.stats.poisson.pmf(muts, dt * mutation_rate * span)
-        if standardize:
-            return ll / np.max(ll)
-        else:
-            return ll
-
-    @staticmethod
-    def _lik_wrapper(muts_span, dt, mutation_rate, standardize=True):
-        """
-        A wrapper to allow this _lik to be called by pool.imap_unordered, returning the
-        mutation and span values
-        """
-        return muts_span, Likelihoods._lik(
-            muts_span[0], muts_span[1], dt, mutation_rate, standardize=standardize
-        )
-
-    def precalculate_mutation_likelihoods(self, num_threads=None, unique_method=0):
-        """
-        We precalculate these because the pmf function is slow, but can be trivially
-        parallelised. We store the likelihoods in a cache because they only depend on
-        the number of mutations and the span, so can potentially be reused.
-
-        However, we don't bother storing the likelihood for edges above a *fixed* node,
-        because (a) these are only used once per node and (b) sample edges are often
-        long, and hence their span will be unique. This also allows us to deal easily
-        with fixed nodes at explicit times (rather than in time slices)
-        """
-
-        if self.mut_rate is None:
-            raise RuntimeError(
-                "Cannot calculate mutation likelihoods with no mutation_rate set"
-            )
-        if unique_method == 0:
-            self.unfixed_likelihood_cache = {
-                (muts, e.span): None
-                for muts, e in zip(self.mut_edges, self.ts.edges())
-                if e.child not in self.fixednodes
-            }
-        else:
-            fixed_nodes = np.array(list(self.fixednodes))
-            keys = np.unique(
-                np.core.records.fromarrays(
-                    (self.mut_edges, self.ts.edges_right - self.ts.edges_left),
-                    names="muts,span",
-                )[np.logical_not(np.isin(self.ts.edges_child, fixed_nodes))]
-            )
-            if unique_method == 1:
-                self.unfixed_likelihood_cache = dict.fromkeys({tuple(t) for t in keys})
-            else:
-                self.unfixed_likelihood_cache = {tuple(t): None for t in keys}
-
-        if num_threads:
-            f = functools.partial(  # Set constant values for params for static _lik
-                self._lik_wrapper,
-                dt=self.timediff_lower_tri,
-                mutation_rate=self.mut_rate,
-                standardize=self.standardize,
-            )
-            if num_threads == 1:
-                # Useful for testing
-                for key in tqdm(
-                    self.unfixed_likelihood_cache.keys(),
-                    disable=not self.progress,
-                    desc="Precalculating Likelihoods",
-                ):
-                    returned_key, likelihoods = f(key)
-                    self.unfixed_likelihood_cache[returned_key] = likelihoods
-            else:
-                with tqdm(
-                    total=len(self.unfixed_likelihood_cache.keys()),
-                    disable=not self.progress,
-                    desc="Precalculating Likelihoods",
-                ) as prog_bar:
-                    with multiprocessing.Pool(processes=num_threads) as pool:
-                        for key, pmf in pool.imap_unordered(
-                            f, self.unfixed_likelihood_cache.keys()
-                        ):
-                            self.unfixed_likelihood_cache[key] = pmf
-                            prog_bar.update()
-        else:
-            for muts, span in tqdm(
-                self.unfixed_likelihood_cache.keys(),
-                disable=not self.progress,
-                desc="Precalculating Likelihoods",
-            ):
-                self.unfixed_likelihood_cache[muts, span] = self._lik(
-                    muts,
-                    span,
-                    dt=self.timediff_lower_tri,
-                    mutation_rate=self.mut_rate,
-                    standardize=self.standardize,
-                )
-
-    def get_mut_lik_fixed_node(self, edge):
-        """
-        Get the mutation likelihoods for an edge whose child is at a
-        fixed time, but whose parent may take any of the time slices in the timepoints
-        that are equal to or older than the child age. This is not cached, as it is
-        likely to be unique for each edge
-        """
-        assert (
-            edge.child in self.fixednodes
-        ), "Wrongly called fixed node function on non-fixed node"
-        assert (
-            self.mut_rate is not None
-        ), "Cannot calculate mutation likelihoods with no mutation_rate set"
-
-        mutations_on_edge = self.mut_edges[edge.id]
-        child_time = self.ts.node(edge.child).time
-        assert child_time == 0
-        # Temporary hack - we should really take a more precise likelihood
-        return self._lik(
-            mutations_on_edge,
-            edge.span,
-            self.timediff,
-            self.mut_rate,
-            standardize=self.standardize,
-        )
-
-    def get_mut_lik_lower_tri(self, edge):
-        """
-        Get the cached mutation likelihoods for an edge with non-fixed parent and child
-        nodes, returning values for all the possible time differences between timepoints
-        These values are returned as a flattened lower triangular matrix, the
-        form required in the inside algorithm.
-
-        """
-        # Debugging asserts - should probably remove eventually
-        assert (
-            edge.child not in self.fixednodes
-        ), "Wrongly called lower_tri function on fixed node"
-        assert hasattr(
-            self, "unfixed_likelihood_cache"
-        ), "Must call `precalculate_mutation_likelihoods()` before getting likelihoods"
-
-        mutations_on_edge = self.mut_edges[edge.id]
-        return self.unfixed_likelihood_cache[mutations_on_edge, edge.span]
-
-    def get_mut_lik_upper_tri(self, edge):
-        """
-        Same as :meth:`get_mut_lik_lower_tri`, but the returned array is ordered as
-        flattened upper triangular matrix (suitable for the outside algorithm), rather
-        than a lower triangular one
-        """
-        return self.get_mut_lik_lower_tri(edge)[np.concatenate(self.row_indices)]
-
-    # The following functions don't access the likelihoods directly, but allow
-    # other input arrays of length grid_size to be repeated in such a way that they can
-    # be directly multiplied by the unpacked lower triangular matrix, or arrays of length
-    # of the number of cells in the lower triangular matrix to be summed (e.g. by row)
-    # to give a shorter array of length grid_size
-
-    def make_lower_tri(self, input_array):
-        """
-        Repeat the input array row-wise to make a flattened lower triangular matrix
-        """
-        assert len(input_array) == self.grid_size
-        return input_array[self.to_lower_tri]
-
-    def rowsum_lower_tri(self, input_array):
-        """
-        Describe the reduceat trickery here. Presumably the opposite of make_lower_tri
-        """
-        assert len(input_array) == self.tri_size
-        return np.add.reduceat(input_array, self.row_indices[0])
-
-    def make_upper_tri(self, input_array):
-        """
-        Repeat the input array row-wise to make a flattened upper triangular matrix
-        """
-        assert len(input_array) == self.grid_size
-        return input_array[self.to_upper_tri]
-
-    def rowsum_upper_tri(self, input_array):
-        """
-        Describe the reduceat trickery here. Presumably the opposite of make_upper_tri
-        """
-        assert len(input_array) == self.tri_size
-        return np.add.reduceat(input_array, self.col_indices)
-
-    # Mutation & recombination algorithms on a tree sequence
-
-    def n_breaks(self, edge):
-        """
-        Number of known breakpoints, only used in recombination likelihood calc
-        """
-        return (edge.left != 0) + (edge.right != self.ts.get_sequence_length())
-
-    def combine(self, lik_1, lik_2):
-        return lik_1 * lik_2
-
-    def ratio(self, lik_1, lik_2, div_0_null=False):
-        """
-        Return the ratio of lik_1 to lik_2. In linear space, this divides lik_1 by lik_2
-        If div_0_null==True, then 0/0 is set to the null_constant
-        """
-        with np.errstate(divide="ignore", invalid="ignore"):
-            ret = lik_1 / lik_2
-        if div_0_null:
-            ret[np.isnan(ret)] = self.null_constant
-        return ret
-
-    def marginalize(self, lik):
-        """
-        Return the sum of likelihoods
-        """
-        return np.sum(lik)
-
-    def _recombination_lik(self, edge, fixed=True):
-        # Needs to return a lower tri *or* flattened array depending on `fixed`
-        raise NotImplementedError(
-            "Using the recombination clock is not currently supported"
-            ". See https://github.com/awohns/tsdate/issues/5 for details"
-        )
-        # return (
-        #     np.power(prev_state, self.n_breaks(edge)) *
-        #     np.exp(-(prev_state * self.rec_rate * edge.span * 2)))
-
-    def get_inside(self, arr, edge):
-        liks = self.identity_constant
-        if self.rec_rate is not None:
-            liks = self._recombination_lik(edge)
-        if self.mut_rate is not None:
-            liks *= self.get_mut_lik_lower_tri(edge)
-        return self.rowsum_lower_tri(arr * liks)
-
-    def get_outside(self, arr, edge):
-        liks = self.identity_constant
-        if self.rec_rate is not None:
-            liks = self._recombination_lik(edge)
-        if self.mut_rate is not None:
-            liks *= self.get_mut_lik_upper_tri(edge)
-        return self.rowsum_upper_tri(arr * liks)
-
-    def get_fixed(self, arr, edge):
-        liks = self.identity_constant
-        if self.rec_rate is not None:
-            liks = self._recombination_lik(edge, fixed=True)
-        if self.mut_rate is not None:
-            liks *= self.get_mut_lik_fixed_node(edge)
-        return arr * liks
-
-    def scale_geometric(self, fraction, value):
-        return value**fraction
-
-
-class LogLikelihoods(Likelihoods):
-    """
-    Identical to the Likelihoods class but stores and returns log likelihoods
-    """
-
-    probability_space = base.LOG
-    identity_constant = 0.0
-    null_constant = -np.inf
-
-    """
-    Uses an alternative to logsumexp, useful for large grid sizes, see
-    http://www.nowozin.net/sebastian/blog/streaming-log-sum-exp-computation.html
-    """
-
-    @staticmethod
-    @numba.jit(nopython=True)
-    def logsumexp(X):
-        alpha = -np.inf
-        r = 0.0
-        for x in X:
-            if x != -np.inf:
-                if x <= alpha:
-                    r += np.exp(x - alpha)
-                else:
-                    r *= np.exp(alpha - x)
-                    r += 1.0
-                    alpha = x
-        return -np.inf if r == 0 else np.log(r) + alpha
-
-    @staticmethod
-    def _lik(muts, span, dt, mutation_rate, standardize=True):
-        """
-        The likelihood of an edge given a number of mutations, as set of time deltas (dt)
-        and a span. This is a static function to allow parallelization
-        """
-        ll = scipy.stats.poisson.logpmf(muts, dt * mutation_rate * span)
-        if standardize:
-            return ll - np.max(ll)
-        else:
-            return ll
-
-    @staticmethod
-    def _lik_wrapper(muts_span, dt, mutation_rate, standardize=True):
-        """
-        Needs redefining to refer to the LogLikelihoods class
-        """
-        return muts_span, LogLikelihoods._lik(
-            muts_span[0], muts_span[1], dt, mutation_rate, standardize=standardize
-        )
-
-    def rowsum_lower_tri(self, input_array):
-        """
-        The function below is equivalent to (but numba makes it faster than)
-        np.logaddexp.reduceat(input_array, self.row_indices[0])
-        """
-        assert len(input_array) == self.tri_size
-        res = list()
-        i_start = self.row_indices[0][0]
-        for i in self.row_indices[0][1:]:
-            res.append(self.logsumexp(input_array[i_start:i]))
-            i_start = i
-        res.append(self.logsumexp(input_array[i:]))
-        return np.array(res)
-
-    def rowsum_upper_tri(self, input_array):
-        """
-        The function below is equivalent to (but numba makes it faster than)
-        np.logaddexp.reduceat(input_array, self.col_indices)
-        """
-        assert len(input_array) == self.tri_size
-        res = list()
-        i_start = self.col_indices[0]
-        for i in self.col_indices[1:]:
-            res.append(self.logsumexp(input_array[i_start:i]))
-            i_start = i
-        res.append(self.logsumexp(input_array[i:]))
-        return np.array(res)
-
-    def _recombination_loglik(self, edge, fixed=True):
-        # Needs to return a lower tri *or* flattened array depending on `fixed`
-        raise NotImplementedError(
-            "Using the recombination clock is not currently supported"
-            ". See https://github.com/awohns/tsdate/issues/5 for details"
-        )
-        # return (
-        #     np.power(prev_state, self.n_breaks(edge)) *
-        #     np.exp(-(prev_state * self.rec_rate * edge.span * 2)))
-
-    def combine(self, loglik_1, loglik_2):
-        return loglik_1 + loglik_2
-
-    def ratio(self, loglik_1, loglik_2, div_0_null=False):
-        """
-        In log space, likelihood ratio is loglik_1 - loglik_2
-        If div_0_null==True, then if either is -inf it returns -inf (the null_constant)
-        """
-        with np.errstate(divide="ignore", invalid="ignore"):
-            ret = loglik_1 - loglik_2
-        if div_0_null:
-            ret[np.isnan(ret)] = self.null_constant
-        return ret
-
-    def marginalize(self, loglik):
-        """
-        Return the logged sum of likelihoods
-        """
-        return self.logsumexp(loglik)
-
-    def get_inside(self, arr, edge):
-        log_liks = self.identity_constant
-        if self.rec_rate is not None:
-            log_liks = self._recombination_loglik(edge)
-        if self.mut_rate is not None:
-            log_liks += self.get_mut_lik_lower_tri(edge)
-        return self.rowsum_lower_tri(arr + log_liks)
-
-    def get_outside(self, arr, edge):
-        log_liks = self.identity_constant
-        if self.rec_rate is not None:
-            log_liks = self._recombination_loglik(edge)
-        if self.mut_rate is not None:
-            log_liks += self.get_mut_lik_upper_tri(edge)
-        return self.rowsum_upper_tri(arr + log_liks)
-
-    def get_fixed(self, arr, edge):
-        log_liks = self.identity_constant
-        if self.rec_rate is not None:
-            log_liks = self._recombination_loglik(edge, fixed=True)
-        if self.mut_rate is not None:
-            log_liks += self.get_mut_lik_fixed_node(edge)
-        return arr + log_liks
-
-    def scale_geometric(self, fraction, value):
-        return fraction * value
-
-
-class InOutAlgorithms:
-    """
-    Contains the inside and outside algorithms
-    """
-
-    def __init__(self, priors, lik, *, progress=False):
-        if (
-            lik.fixednodes.intersection(priors.nonfixed_nodes)
-            or len(lik.fixednodes) + len(priors.nonfixed_nodes) != lik.ts.num_nodes
-        ):
-            raise ValueError(
-                "The prior and likelihood objects disagree on which nodes are fixed"
-            )
-        if not np.allclose(lik.timepoints, priors.timepoints):
-            raise ValueError(
-                "The prior and likelihood objects disagree on the timepoints used"
-            )
-
-        self.priors = priors
-        self.nonfixed_nodes = priors.nonfixed_nodes
-        self.lik = lik
-        self.ts = lik.ts
-
-        self.fixednodes = lik.fixednodes
-        self.progress = progress
-        # If necessary, convert priors to log space
-        self.priors.force_probability_space(lik.probability_space)
-
-        self.spans = np.bincount(
-            self.ts.edges_child,
-            weights=self.ts.edges_right - self.ts.edges_left,
-        )
-        self.spans = np.pad(self.spans, (0, self.ts.num_nodes - len(self.spans)))
-
-        self.root_spans = defaultdict(float)
-        for tree in self.ts.trees(root_threshold=2):
-            if tree.has_single_root:
-                self.root_spans[tree.root] += tree.span
-        # Add on the spans when this is a root
-        for root, span_when_root in self.root_spans.items():
-            self.spans[root] += span_when_root
-
-    # === Grouped edge iterators ===
-
-    def edges_by_parent_asc(self, grouped=True):
-        """
-        Return an itertools.groupby object of edges grouped by parent in ascending order
-        of the time of the parent. Since tree sequence properties guarantee that edges
-        are listed in nondecreasing order of parent time
-        (https://tskit.readthedocs.io/en/latest/data-model.html#edge-requirements)
-        we can simply use the standard edge order
-        """
-        if grouped:
-            return itertools.groupby(self.ts.edges(), operator.attrgetter("parent"))
-        else:
-            return self.ts.edges()
-
-    def edges_by_child_desc(self, grouped=True):
-        """
-        Return an itertools.groupby object of edges grouped by child in descending order
-        of the time of the child.
-        """
-        it = (
-            self.ts.edge(u)
-            for u in np.lexsort(
-                (self.ts.edges_child, -self.ts.nodes_time[self.ts.edges_child])
-            )
-        )
-        if grouped:
-            return itertools.groupby(it, operator.attrgetter("child"))
-        else:
-            return it
-
-    def edges_by_child_then_parent_desc(self, grouped=True):
-        """
-        Return an itertools.groupby object of edges grouped by child in descending order
-        of the time of the child, then by descending order of age of child
-        """
-        wtype = np.dtype(
-            [
-                ("child_age", self.ts.nodes_time.dtype),
-                ("child_node", self.ts.edges_child.dtype),
-                ("parent_age", self.ts.nodes_time.dtype),
-            ]
-        )
-        w = np.empty(self.ts.num_edges, dtype=wtype)
-        w["child_age"] = self.ts.nodes_time[self.ts.edges_child]
-        w["child_node"] = self.ts.edges_child
-        w["parent_age"] = -self.ts.nodes_time[self.ts.edges_parent]
-        sorted_child_parent = (
-            self.ts.edge(i)
-            for i in reversed(
-                np.argsort(w, order=("child_age", "child_node", "parent_age"))
-            )
-        )
-        if grouped:
-            return itertools.groupby(sorted_child_parent, operator.attrgetter("child"))
-        else:
-            return sorted_child_parent
-
-    # === MAIN ALGORITHMS ===
-
-    def inside_pass(self, *, standardize=True, cache_inside=False, progress=None):
-        """
-        Use dynamic programming to find approximate posterior to sample from
-        """
-        if progress is None:
-            progress = self.progress
-        inside = self.priors.clone_with_new_data(  # store inside matrix values
-            grid_data=np.nan, fixed_data=self.lik.identity_constant
-        )
-        if cache_inside:
-            g_i = np.full(
-                (self.ts.num_edges, self.lik.grid_size), self.lik.identity_constant
-            )
-        denominator = np.full(self.ts.num_nodes, np.nan)
-        assert (
-            self.lik.standardize is False
-        ), "Marginal likelihood requires unstandardized mutation likelihoods"
-        marginal_lik = self.lik.identity_constant
-        # Iterate through the nodes via groupby on parent node
-        for parent, edges in tqdm(
-            self.edges_by_parent_asc(),
-            desc="Inside",
-            total=inside.num_nonfixed,
-            disable=not progress,
-        ):
-            """
-            for each node, find the conditional prob of age at every time
-            in time grid
-            """
-            if parent in self.fixednodes:
-                continue  # there is no hidden state for this parent - it's fixed
-            val = self.priors[parent].copy()
-            for edge in edges:
-                spanfrac = edge.span / self.spans[edge.child]
-                # Calculate vals for each edge
-                if edge.child in self.fixednodes:
-                    # NB: geometric scaling works exactly when all nodes fixed in graph
-                    # but is an approximation when times are unknown.
-                    daughter_val = self.lik.scale_geometric(spanfrac, inside[edge.child])
-                    edge_lik = self.lik.get_fixed(daughter_val, edge)
-                else:
-                    inside_values = inside[edge.child]
-                    if np.ndim(inside_values) == 0 or np.all(np.isnan(inside_values)):
-                        # Child appears fixed, or we have not visited it. Either our
-                        # edge order is wrong (bug) or we have hit a dangling node
-                        raise ValueError(
-                            "The input tree sequence includes "
-                            "dangling nodes: please simplify it"
-                        )
-                    daughter_val = self.lik.scale_geometric(
-                        spanfrac, self.lik.make_lower_tri(inside[edge.child])
-                    )
-                    edge_lik = self.lik.get_inside(daughter_val, edge)
-                val = self.lik.combine(val, edge_lik)
-                if cache_inside:
-                    g_i[edge.id] = edge_lik
-            denominator[parent] = (
-                np.max(val) if standardize else self.lik.identity_constant
-            )
-            inside[parent] = self.lik.ratio(val, denominator[parent])
-            if standardize:
-                marginal_lik = self.lik.combine(marginal_lik, denominator[parent])
-        if cache_inside:
-            self.g_i = self.lik.ratio(g_i, denominator[self.ts.edges_child, None])
-        # Keep the results in this object
-        self.inside = inside
-        self.denominator = denominator
-        # Calculate marginal likelihood
-        for root, span_when_root in self.root_spans.items():
-            spanfrac = span_when_root / self.spans[root]
-            root_val = self.lik.scale_geometric(spanfrac, inside[root])
-            marginal_lik = self.lik.combine(marginal_lik, self.lik.marginalize(root_val))
-        return marginal_lik
-
-    def outside_pass(
-        self,
-        *,
-        standardize=False,
-        ignore_oldest_root=False,
-        progress=None,
-    ):
-        """
-        Computes the full posterior distribution on nodes, returning the
-        posterior values. These are *not* probabilities, as they do not sum to one:
-        to convert to probabilities, call posterior.to_probabilities()
-
-        Standardizing *during* the outside process may be necessary if there is
-        overflow, but means that we cannot check the total functional value at each node
-
-        Ignoring the oldest root may also be necessary when the oldest root node
-        causes numerical stability issues.
-
-        The rows in the posterior returned correspond to node IDs as given by
-        self.nodes
-        """
-        if progress is None:
-            progress = self.progress
-        if not hasattr(self, "inside"):
-            raise RuntimeError("You have not yet run the inside algorithm")
-
-        outside = self.inside.clone_with_new_data(grid_data=0, probability_space=base.LIN)
-        for root, span_when_root in self.root_spans.items():
-            outside[root] = span_when_root / self.spans[root]
-        outside.force_probability_space(self.inside.probability_space)
-
-        for child, edges in tqdm(
-            self.edges_by_child_desc(),
-            desc="Outside",
-            total=len(np.unique(self.ts.edges_child)),
-            disable=not progress,
-        ):
-            if child in self.fixednodes:
-                continue
-            val = np.full(self.lik.grid_size, self.lik.identity_constant)
-            for edge in edges:
-                if ignore_oldest_root:
-                    if edge.parent == self.ts.num_nodes - 1:
-                        continue
-                if edge.parent in self.fixednodes:
-                    raise RuntimeError(
-                        "Fixed nodes cannot currently be parents in the TS"
-                    )
-                # Geometric scaling works exactly for all nodes fixed in graph
-                # but is an approximation when times are unknown.
-                spanfrac = edge.span / self.spans[child]
-                try:
-                    inside_div_gi = self.lik.ratio(
-                        self.inside[edge.parent], self.g_i[edge.id], div_0_null=True
-                    )
-                except AttributeError:  # we haven't cached g_i so we recalculate
-                    daughter_val = self.lik.scale_geometric(
-                        spanfrac, self.lik.make_lower_tri(self.inside[edge.child])
-                    )
-                    edge_lik = self.lik.get_inside(daughter_val, edge)
-                    cur_g_i = self.lik.ratio(edge_lik, self.denominator[child])
-                    inside_div_gi = self.lik.ratio(
-                        self.inside[edge.parent], cur_g_i, div_0_null=True
-                    )
-                parent_val = self.lik.scale_geometric(
-                    spanfrac,
-                    self.lik.make_upper_tri(
-                        self.lik.combine(outside[edge.parent], inside_div_gi)
-                    ),
-                )
-                if standardize:
-                    parent_val = self.lik.ratio(parent_val, np.max(parent_val))
-                edge_lik = self.lik.get_outside(parent_val, edge)
-                val = self.lik.combine(val, edge_lik)
-
-            # vv[0] = 0  # Seems a hack: internal nodes should be allowed at time 0
-            assert self.denominator[edge.child] > self.lik.null_constant
-            outside[child] = self.lik.ratio(val, self.denominator[child])
-            if standardize:
-                outside[child] = self.lik.ratio(val, np.max(val))
-        self.outside = outside
-        posterior = outside.clone_with_new_data(
-            grid_data=self.lik.combine(self.inside.grid_data, outside.grid_data),
-            fixed_data=np.nan,
-        )  # We should never use the posterior for a fixed node
-        return posterior
-
-    def outside_maximization(self, *, eps, progress=None):
-        if progress is None:
-            progress = self.progress
-        if not hasattr(self, "inside"):
-            raise RuntimeError("You have not yet run the inside algorithm")
-
-        maximized_node_times = np.zeros(self.ts.num_nodes, dtype="int")
-
-        if self.lik.probability_space == base.LOG:
-            poisson = scipy.stats.poisson.logpmf
-        elif self.lik.probability_space == base.LIN:
-            poisson = scipy.stats.poisson.pmf
-
-        mut_edges = self.lik.mut_edges
-        mrcas = np.where(
-            np.isin(np.arange(self.ts.num_nodes), self.ts.edges_child, invert=True)
-        )[0]
-        for i in mrcas:
-            if i not in self.fixednodes:
-                maximized_node_times[i] = np.argmax(self.inside[i])
-
-        for child, edges in tqdm(
-            self.edges_by_child_then_parent_desc(),
-            desc="Maximization",
-            total=len(np.unique(self.ts.edges_child)),
-            disable=not progress,
-        ):
-            if child in self.fixednodes:
-                continue
-            for edge_index, edge in enumerate(edges):
-                if edge_index == 0:
-                    youngest_par_index = maximized_node_times[edge.parent]
-                    parent_time = self.lik.timepoints[maximized_node_times[edge.parent]]
-                    ll_mut = poisson(
-                        mut_edges[edge.id],
-                        (
-                            parent_time
-                            - self.lik.timepoints[: youngest_par_index + 1]
-                            + eps
-                        )
-                        * self.lik.mut_rate
-                        * edge.span,
-                    )
-                    result = self.lik.ratio(ll_mut, np.max(ll_mut))
-                else:
-                    cur_parent_index = maximized_node_times[edge.parent]
-                    if cur_parent_index < youngest_par_index:
-                        youngest_par_index = cur_parent_index
-                    parent_time = self.lik.timepoints[maximized_node_times[edge.parent]]
-                    ll_mut = poisson(
-                        mut_edges[edge.id],
-                        (
-                            parent_time
-                            - self.lik.timepoints[: youngest_par_index + 1]
-                            + eps
-                        )
-                        * self.lik.mut_rate
-                        * edge.span,
-                    )
-                    result[: youngest_par_index + 1] = self.lik.combine(
-                        self.lik.ratio(
-                            ll_mut[: youngest_par_index + 1],
-                            np.max(ll_mut[: youngest_par_index + 1]),
-                        ),
-                        result[: youngest_par_index + 1],
-                    )
-            inside_val = self.inside[child][: (youngest_par_index + 1)]
-
-            maximized_node_times[child] = np.argmax(
-                self.lik.combine(result[: youngest_par_index + 1], inside_val)
-            )
-
-        return self.lik.timepoints[np.array(maximized_node_times).astype("int")]
-
-
 # Classes for each method
 Results = namedtuple(
     "Results",
     [
         "posterior_mean",
         "posterior_var",
-        "posterior_obj",
         "mutation_mean",
         "mutation_var",
         "mutation_lik",
         "mutation_edge",
         "mutation_node",
+        "fit_object",
     ],
 )
 
@@ -900,18 +82,27 @@ def __init__(
         recombination_rate=None,
         time_units=None,
         priors=None,
-        return_posteriors=None,
         return_likelihood=None,
+        return_fit=None,
         record_provenance=None,
         constr_iterations=None,
         progress=None,
+        # Deprecated params
+        return_posteriors=None,
     ):
         # Set up all the generic params describe in the tsdate.date function, and define
         # priors if not passed-in already
+        if return_posteriors is not None:
+            raise ValueError(
+                'The "return_posteriors" parameter has been deprecated. Either use the '
+                "posterior values encoded in node metadata or set ``return_fit=True`` "
+                "then access `fit.node_posteriors()` to obtain a transposed version "
+                "of the matrix previously returned when ``return_posteriors=True.``"
+            )
         self.ts = ts
         self.mutation_rate = mutation_rate
         self.recombination_rate = recombination_rate
-        self.return_posteriors = return_posteriors
+        self.return_fit = return_fit
         self.return_likelihood = return_likelihood
         self.pbar = progress
         self.time_units = "generations" if time_units is None else time_units
@@ -964,9 +155,7 @@ def __init__(
                 mk_prior = getattr(prior, self.prior_grid_func_name)
                 # Default to not creating approximate priors unless ts has
                 # greater than DEFAULT_APPROX_PRIOR_SIZE samples
-                approx = (
-                    True if ts.num_samples > base.DEFAULT_APPROX_PRIOR_SIZE else False
-                )
+                approx = ts.num_samples > prior.DEFAULT_APPROX_PRIOR_SIZE
                 self.priors = mk_prior(
                     ts, Ne, approximate_priors=approx, progress=progress
                 )
@@ -1061,16 +250,12 @@ def set_time_metadata(self, table, mean, var, default_schema, overwrite=False):
             except tskit.MetadataValidationError as e:
                 logger.warning(f"Could not set time metadata in {table_name}: {e}")
 
-    def parse_result(self, result, epsilon, extra_posterior_cols=None):
+    def parse_result(self, result, epsilon):
         # Construct the tree sequence to return and add other stuff we might want to
         # return. pst_cols is a dict to be appended to the output posterior dict
         ret = [self.get_modified_ts(result, epsilon)]
-        if self.return_posteriors:
-            pst_dict = None
-            if result.posterior_obj is not None:
-                pst_dict = result.posterior_obj.nonfixed_dict()
-                pst_dict.update(extra_posterior_cols or {})
-            ret.append(pst_dict)
+        if self.return_fit:
+            ret.append(result.fit_object)
         if self.return_likelihood:
             ret.append(result.mutation_lik)
         return tuple(ret) if len(ret) > 1 else ret.pop()
@@ -1112,8 +297,8 @@ def mean_var(ts, posterior):
 
     def main_algorithm(self, probability_space, epsilon, num_threads):
         # Algorithm class is shared by inside-outside & outside-maximization methods
-        if probability_space != base.LOG:
-            liklhd = Likelihoods(
+        if probability_space == LIN_GRID:
+            liklhd = discrete.Likelihoods(
                 self.ts,
                 self.priors.timepoints,
                 self.mutation_rate,
@@ -1122,8 +307,8 @@ def main_algorithm(self, probability_space, epsilon, num_threads):
                 fixed_node_set=self.get_fixed_nodes_set(),
                 progress=self.pbar,
             )
-        else:
-            liklhd = LogLikelihoods(
+        elif probability_space == LOG_GRID:
+            liklhd = discrete.LogLikelihoods(
                 self.ts,
                 self.priors.timepoints,
                 self.mutation_rate,
@@ -1132,10 +317,15 @@ def main_algorithm(self, probability_space, epsilon, num_threads):
                 fixed_node_set=self.get_fixed_nodes_set(),
                 progress=self.pbar,
             )
+        else:
+            raise ValueError(
+                f"Invalid discrete probability space: {probability_space}. Must be "
+                f"one of {LIN_GRID} or {LOG_GRID}"
+            )
         if self.mutation_rate is not None:
             liklhd.precalculate_mutation_likelihoods(num_threads=num_threads)
 
-        return InOutAlgorithms(self.priors, liklhd, progress=self.pbar)
+        return discrete.BeliefPropagation(self.priors, liklhd, progress=self.pbar)
 
 
 class InsideOutsideMethod(DiscreteTimeMethod):
@@ -1150,32 +340,41 @@ def run(
         num_threads=None,
         cache_inside=None,
     ):
+        if self.mutation_rate is None and self.recombination_rate is None:
+            if self.ts.num_trees > 1:
+                raise NotImplementedError(
+                    "Specifying no mutation or recombination rate implies dating using "
+                    "the topology-only clock. This produces biased results under "
+                    "recombination (https://github.com/tskit-dev/tsdate/issues/292). "
+                    "The topology-only clock has therefore been deprecated for tree "
+                    "sequences representing more than one tree."
+                )
         if self.provenance_params is not None:
             self.provenance_params.update(
                 {k: v for k, v in locals().items() if k != "self"}
             )
-        dynamic_prog = self.main_algorithm(probability_space, eps, num_threads)
-        marginal_likl = dynamic_prog.inside_pass(cache_inside=cache_inside)
-        posterior_obj = dynamic_prog.outside_pass(
+        fit_obj = self.main_algorithm(probability_space, eps, num_threads)
+        marginal_likl = fit_obj.inside_pass(cache_inside=cache_inside)
+        fit_obj.outside_pass(
             standardize=outside_standardize, ignore_oldest_root=ignore_oldest_root
         )
         # Turn the posterior into probabilities
-        posterior_obj.standardize()  # Just to ensure there are no floating point issues
-        posterior_obj.force_probability_space(base.LIN)
-        posterior_obj.to_probabilities()
+        fit_obj.posterior_grid.standardize()  # Just to ensure no floating point issues
+        fit_obj.posterior_grid.force_probability_space(LIN_GRID)
+        fit_obj.posterior_grid.to_probabilities()
 
-        posterior_mean, posterior_var = self.mean_var(self.ts, posterior_obj)
+        posterior_mean, posterior_var = self.mean_var(self.ts, fit_obj.posterior_grid)
         mut_edge = np.full(self.ts.num_mutations, tskit.NULL)
         mut_node = self.ts.mutations_node
         return Results(
             posterior_mean,
             posterior_var,
-            posterior_obj,
             None,
             None,
             marginal_likl,
             mut_edge,
             mut_node,
+            fit_obj,
         )
 
 
@@ -1184,8 +383,6 @@ class MaximizationMethod(DiscreteTimeMethod):
 
     def __init__(self, ts, **kwargs):
         super().__init__(ts, **kwargs)
-        if self.return_posteriors:
-            raise ValueError("Cannot return posterior with maximization method")
 
     def run(
         self,
@@ -1194,19 +391,26 @@ def run(
         num_threads=None,
         cache_inside=None,
     ):
-        if self.mutation_rate is None:
+        if self.mutation_rate is None and self.recombination_rate is None:
             raise ValueError("Outside maximization method requires mutation rate")
         if self.provenance_params is not None:
             self.provenance_params.update(
                 {k: v for k, v in locals().items() if k != "self"}
             )
-        dynamic_prog = self.main_algorithm(probability_space, eps, num_threads)
-        marginal_likl = dynamic_prog.inside_pass(cache_inside=cache_inside)
-        posterior_mean = dynamic_prog.outside_maximization(eps=eps)
+        fit_obj = self.main_algorithm(probability_space, eps, num_threads)
+        marginal_likl = fit_obj.inside_pass(cache_inside=cache_inside)
+        fit_obj.outside_maximization(eps=eps)
         mut_edge = np.full(self.ts.num_mutations, tskit.NULL)
         mut_node = self.ts.mutations_node
         return Results(
-            posterior_mean, None, None, None, None, marginal_likl, mut_edge, mut_node
+            fit_obj.posterior_mean,
+            None,
+            None,
+            None,
+            marginal_likl,
+            mut_edge,
+            mut_node,
+            fit_obj,
         )
 
 
@@ -1237,12 +441,12 @@ def run(
         if self.mutation_rate is None:
             raise ValueError("Variational gamma method requires mutation rate")
 
-        posterior = variational.ExpectationPropagation(
+        fit_obj = variational.ExpectationPropagation(
             self.ts,
             mutation_rate=self.mutation_rate,
             singletons_phased=singletons_phased,
         )
-        posterior.run(
+        fit_obj.infer(
             ep_iterations=max_iterations,
             max_shape=max_shape,
             rescale_intervals=rescaling_intervals,
@@ -1251,20 +455,20 @@ def run(
             rescale_segsites=match_segregating_sites,
             progress=self.pbar,
         )
-
-        node_mn, node_va = posterior.node_moments()
-        mutation_mn, mutation_va = posterior.mutation_moments()
-        mutation_edge, mutation_node = posterior.mutation_mapping()
+        marginal_likl = fit_obj.marginal_likelihood()
+        node_mn, node_va = fit_obj.node_moments()
+        mutation_mn, mutation_va = fit_obj.mutation_moments()
+        mutation_edge, mutation_node = fit_obj.mutation_mapping()
 
         return Results(
             node_mn,
             node_va,
-            None,
             mutation_mn,
             mutation_va,
-            None,
+            marginal_likl,
             mutation_edge,
             mutation_node,
+            fit_obj,
         )
 
 
@@ -1334,10 +538,10 @@ def maximization(
         object. The ``population_size`` parameter is only used when ``priors`` is
         ``None``. Conversely, if ``priors`` is not ``None``, no ``population_size``
         value should be specified.
-    :param tsdate.base.NodeGridValues priors: NodeGridValues object containing the prior
-        parameters for each node-to-be-dated. Note that different estimation methods may
-        require different types of prior, as described in the documentation for each
-        estimation method.
+    :param tsdate.node_time_class.NodeTimeValues priors: NodeTimeValues object containing
+        the prior parameters for each node-to-be-dated. Note that different estimation
+        methods may require different types of prior, as described in the documentation
+        for each estimation method.
     :param float eps: The error factor in time difference calculations, and the
         minimum distance separating parent and child ages in the returned tree sequence.
         Default: None, treated as 1e-6.
@@ -1349,9 +553,8 @@ def maximization(
     :param \\**kwargs: Other keyword arguments as described in the :func:`date` wrapper
         function, notably ``mutation_rate``, and ``population_size`` or ``priors``.
         Further arguments include ``time_units``, ``progress``, and
-        ``record_provenance``. The additional ``return_likelihood`` argument can be used
-        to return additional information (see below). Posteriors cannot be returned using
-        this estimation method.
+        ``record_provenance``.  The additional arguments ``return_fit`` and
+        ``return_likelihood`` can be used to return additional information (see below).
     :return:
         - **ts** (:class:`~tskit.TreeSequence`) -- a copy of the input tree sequence with
           updated node times based on the posterior mean, corrected where necessary to
@@ -1369,7 +572,7 @@ def maximization(
     if eps is None:
         eps = DEFAULT_EPSILON
     if probability_space is None:
-        probability_space = base.LOG
+        probability_space = LOG_GRID
 
     dating_method = MaximizationMethod(
         tree_sequence,
@@ -1457,10 +660,10 @@ def inside_outside(
         object. The ``population_size`` parameter is only used when ``priors`` is
         ``None``. Conversely, if ``priors`` is not ``None``, no ``population_size``
         value should be specified.
-    :param tsdate.base.NodeGridValues priors: NodeGridValues object containing the prior
-        parameters for each node-to-be-dated. Note that different estimation methods may
-        require different types of prior, as described in the documentation for each
-        estimation method.
+    :param tsdate.node_time_class.NodeTimeValues priors: NodeTimeValues object containing
+        the prior parameters for each node-to-be-dated. Note that different estimation
+        methods may require different types of prior, as described in the documentation
+        for each estimation method.
     :param float eps: The error factor in time difference calculations, and the
         minimum distance separating parent and child ages in the returned tree sequence.
         Default: None, treated as 1e-6.
@@ -1483,22 +686,17 @@ def inside_outside(
     :param \\**kwargs: Other keyword arguments as described in the :func:`date` wrapper
         function, notably ``mutation_rate``, and ``population_size`` or ``priors``.
         Further arguments include ``time_units``, ``progress``, and
-        ``record_provenance``. The additional arguments ``return_posteriors`` and
+        ``record_provenance``. The additional arguments ``return_fit`` and
         ``return_likelihood`` can be used to return additional information (see below).
     :return:
         - **ts** (:class:`~tskit.TreeSequence`) -- a copy of the input tree sequence with
           updated node times based on the posterior mean, corrected where necessary to
           ensure that parents are strictly older than all their children by an amount
           given by the ``eps`` parameter.
-        - **posteriors** (:py:class:`dict`) -- (Only returned if ``return_posteriors``
-          is ``True``) A dictionary of posterior probabilities.
-          Each node whose time was inferred corresponds to an item in this dictionary
-          whose key is the node ID and value is an array of probabilities of the node
-          being at a list of timepoints. Timepoint values are provided in the
-          returned dictionary under the key named "time". When read
-          as a pandas ``DataFrame`` object using ``pd.DataFrame(posteriors)``,
-          the rows correspond to labelled timepoints and columns are
-          headed by their respective node ID.
+        - **fit** (:class:`~discrete.BeliefPropagation`) -- (Only returned if
+          ``return_fit`` is ``True``) The underlying object used to run the dating
+          inference. This can then be queried e.g. using
+          :meth:`~discrete.BeliefPropagation.node_posteriors()`
         - **marginal_likelihood** (:py:class:`float`) -- (Only returned if
           ``return_likelihood`` is ``True``) The marginal likelihood of
           the mutation data given the inferred node times.
@@ -1511,7 +709,7 @@ def inside_outside(
     if eps is None:
         eps = DEFAULT_EPSILON
     if probability_space is None:
-        probability_space = base.LOG
+        probability_space = LOG_GRID
     if outside_standardize is None:
         outside_standardize = True
     if ignore_oldest_root is None:
@@ -1531,9 +729,7 @@ def inside_outside(
         cache_inside=cache_inside,
         probability_space=probability_space,
     )
-    return dating_method.parse_result(
-        result, eps, {"time": result.posterior_obj.timepoints}
-    )
+    return dating_method.parse_result(result, eps)
 
 
 def variational_gamma(
@@ -1589,21 +785,17 @@ def variational_gamma(
         are polytomies. Default ``False``.
     :param \\**kwargs: Other keyword arguments as described in the :func:`date` wrapper
         function, including ``time_units``, ``progress``, and ``record_provenance``.
-        The arguments ``return_posteriors`` and ``return_likelihood`` can be
+        The arguments ``return_fit`` and ``return_likelihood`` can be
         used to return additional information (see below).
     :return:
         - **ts** (:class:`~tskit.TreeSequence`) -- a copy of the input tree sequence with
           updated node times based on the posterior mean, corrected where necessary to
           ensure that parents are strictly older than all their children by an amount
           given by the ``eps`` parameter.
-        - **posteriors** (:py:class:`dict`) -- (Only returned if ``return_posteriors``
-          is ``True``) A dictionary of posterior probabilities.
-          Each node whose time was inferred corresponds to an item in this dictionary
-          whose key is the node ID and value is an array of the ``[shape, rate]``
-          parameters of the posterior gamma distribution for that node. When read
-          as a pandas ``DataFrame`` object using ``pd.DataFrame(posteriors)``,
-          the first row of the data frame is the shape and the second the rate
-          parameter, each column being headed by the respective node ID.
+        - **fit** (:class:`~variational.ExpectationPropagation`) -- (Only returned
+          if ``return_fit`` is ``True``). The underlying object used to run the dating
+          inference. This can then be queried e.g. using
+          :meth:`~variational.ExpectationPropagation.node_posteriors()`
         - **marginal_likelihood** (:py:class:`float`) -- (Only returned if
           ``return_likelihood`` is ``True``) The marginal likelihood of
           the mutation data given the inferred node times. Not currently
@@ -1644,7 +836,7 @@ def variational_gamma(
         regularise_roots=regularise_roots,
         singletons_phased=singletons_phased,
     )
-    return dating_method.parse_result(result, eps, {"parameter": ["shape", "rate"]})
+    return dating_method.parse_result(result, eps)
 
 
 estimation_methods = {
@@ -1672,7 +864,7 @@ def date(
     time_units=None,
     method=None,
     constr_iterations=None,
-    return_posteriors=None,
+    return_fit=None,
     return_likelihood=None,
     progress=None,
     record_provenance=True,
@@ -1683,17 +875,15 @@ def date(
     Infer dates for nodes in a genealogical graph (or :ref:`ARG<tutorials:sec_args>`)
     stored in the :ref:`succinct tree sequence<tskit:sec_introduction>` format.
     New times are assigned to nodes using the estimation algorithm specified by
-    ``method`` (see note below). If a ``mutation_rate`` is given,
-    the mutation clock is used. The recombination clock is unsupported at this
-    time.  If neither a ``mutation_rate`` nor a ``recombination_rate`` is given, a
-    topology-only clock is used. Times associated with mutations and times associated
+    ``method`` (see note below). A ``mutation_rate`` must be given (the recombination_rate
+    parameter, implementing a recombination clock, is unsupported at this
+    time). Times associated with mutations and times associated
     with non-fixed (non-sample) nodes are overwritten. For example:
 
     .. code-block:: python
 
       mu = 1e-8
-      Ne = ts.diversity()/4/mu  #  In the absence of external info, use ts for prior Ne
-      new_ts = tsdate.date(ts, mutation_rate=mu, population_size=Ne)
+      new_ts = tsdate.date(ts, mutation_rate=mu)
 
     .. note::
         This is a wrapper for the named functions that are listed in
@@ -1719,15 +909,15 @@ def date(
     :param string method: What estimation method to use. See
         :data:`~tsdate.core.estimation_methods` for possible values.
         If ``None`` (default) the "variational_gamma" method is currently chosen.
-    :param bool return_posteriors: If ``True``, instead of returning just a dated tree
-        sequence, return a tuple of ``(dated_ts, posteriors)``.
+    :param bool return_fit: If ``True``, instead of returning just a dated tree
+        sequence, return a tuple of ``(dated_ts, fit)``.
         Default: None, treated as False.
     :param int constr_iterations: The maximum number of constrained least
         squares iterations to use prior to forcing positive branch lengths.
         Default: None, treated as 0.
     :param bool return_likelihood: If ``True``, return the log marginal likelihood
         from the inside algorithm in addition to the dated tree sequence. If
-        ``return_posteriors`` is also ``True``, then the marginal likelihood
+        ``return_fit`` is also ``True``, then the marginal likelihood
         will be the last element of the tuple. Default: None, treated as False.
     :param bool progress: Show a progress bar. Default: None, treated as False.
     :param bool record_provenance: Should the tsdate command be appended to the
@@ -1738,8 +928,8 @@ def date(
         documented in those specific functions.
     :return:
         A copy of the input tree sequence but with updated node times, or (if
-        ``return_posteriors`` or ``return_likelihood`` is True) a tuple of that
-        tree sequence plus a dictionary of posterior probabilities and/or the
+        ``return_fit`` or ``return_likelihood`` is True) a tuple of that
+        tree sequence plus a fit object and/or the
         marginal likelihood given the mutations on the tree sequence.
     """
     # Only the .date() wrapper needs to consider the deprecated "Ne" param
@@ -1755,7 +945,7 @@ def date(
         time_units=time_units,
         progress=progress,
         constr_iterations=constr_iterations,
-        return_posteriors=return_posteriors,
+        return_fit=return_fit,
         return_likelihood=return_likelihood,
         record_provenance=record_provenance,
         **kwargs,
diff --git a/tsdate/discrete.py b/tsdate/discrete.py
new file mode 100644
index 00000000..3350490a
--- /dev/null
+++ b/tsdate/discrete.py
@@ -0,0 +1,838 @@
+# Classes used for discete time algorithms (inside-outside and outsize_maximization)
+# Note that these methods are no longer the default method used by tsdate
+import functools
+import itertools
+import multiprocessing
+import operator
+from collections import defaultdict
+
+import numba
+import numpy as np
+import scipy.stats
+import tskit
+from tqdm.auto import tqdm
+
+from .node_time_class import LIN_GRID, LOG_GRID
+
+
+class Likelihoods:
+    """
+    A class to store and process likelihoods. Likelihoods for edges are stored as a
+    flattened lower triangular matrix of all the possible delta t's. This class also
+    provides methods for accessing this lower triangular matrix, multiplying it, etc.
+
+    If ``standardize`` is true, routines will operate to standardize the likelihoods
+    such that their maximum is one (in linear space) or zero (in log space)
+    """
+
+    probability_space = LIN_GRID
+    identity_constant = 1.0
+    null_constant = 0.0
+
+    def __init__(
+        self,
+        ts,
+        timepoints,
+        mutation_rate=None,
+        recombination_rate=None,
+        *,
+        eps=0,
+        fixed_node_set=None,
+        standardize=False,
+        progress=False,
+    ):
+        self.ts = ts
+        self.timepoints = timepoints
+        self.fixednodes = set(ts.samples()) if fixed_node_set is None else fixed_node_set
+        self.mut_rate = mutation_rate
+        self.rec_rate = recombination_rate
+        self.standardize = standardize
+        self.grid_size = len(timepoints)
+        self.tri_size = self.grid_size * (self.grid_size + 1) / 2
+        self.ll_mut = {}
+        self.mut_edges = self.get_mut_edges(ts)
+        self.progress = progress
+        # Need to set eps properly in the 2 lines below, to account for values in the
+        # same timeslice
+        self.timediff_lower_tri = np.concatenate(
+            [
+                self.timepoints[time_index] - self.timepoints[0 : time_index + 1] + eps
+                for time_index in np.arange(len(self.timepoints))
+            ]
+        )
+        self.timediff = self.timepoints - self.timepoints[0] + eps
+
+        # The mut_ll contains unpacked (1D) lower triangular matrices. We need to
+        # index this by row and by column index.
+        self.row_indices = []
+        for t in range(self.grid_size):
+            n = np.arange(self.grid_size)
+            self.row_indices.append((((n * (n + 1)) // 2) + t)[t:])
+        self.col_indices = []
+        running_sum = 0  # use this to find the index of the last element of
+        # each column in order to appropriately sum the vv by columns.
+        for i in np.arange(self.grid_size):
+            arr = np.arange(running_sum, running_sum + self.grid_size - i)
+            index = arr[-1]
+            running_sum = index + 1
+            val = arr[0]
+            self.col_indices.append(val)
+
+        # These are used for transforming an array of grid_size into one of tri_size
+        # By repeating elements over rows to form an upper or a lower triangular matrix
+        self.to_lower_tri = np.concatenate(
+            [np.arange(time_idx + 1) for time_idx in np.arange(self.grid_size)]
+        )
+        self.to_upper_tri = np.concatenate(
+            [
+                np.arange(time_idx, self.grid_size)
+                for time_idx in np.arange(self.grid_size + 1)
+            ]
+        )
+
+    @staticmethod
+    def get_mut_edges(ts):
+        """
+        Get the number of mutations on each edge in the tree sequence.
+        """
+        mut_edges = np.zeros(ts.num_edges, dtype=np.int64)
+        for m in ts.mutations():
+            if m.edge != tskit.NULL:
+                mut_edges[m.edge] += 1
+        return mut_edges
+
+    @staticmethod
+    def _lik(muts, span, dt, mutation_rate, standardize=True):
+        """
+        The likelihood of an edge given a number of mutations, as set of time deltas (dt)
+        and a span. This is a static function to allow parallelization
+        """
+        ll = scipy.stats.poisson.pmf(muts, dt * mutation_rate * span)
+        if standardize:
+            return ll / np.max(ll)
+        else:
+            return ll
+
+    @staticmethod
+    def _lik_wrapper(muts_span, dt, mutation_rate, standardize=True):
+        """
+        A wrapper to allow this _lik to be called by pool.imap_unordered, returning the
+        mutation and span values
+        """
+        return muts_span, Likelihoods._lik(
+            muts_span[0], muts_span[1], dt, mutation_rate, standardize=standardize
+        )
+
+    def precalculate_mutation_likelihoods(self, num_threads=None, unique_method=0):
+        """
+        We precalculate these because the pmf function is slow, but can be trivially
+        parallelised. We store the likelihoods in a cache because they only depend on
+        the number of mutations and the span, so can potentially be reused.
+
+        However, we don't bother storing the likelihood for edges above a *fixed* node,
+        because (a) these are only used once per node and (b) sample edges are often
+        long, and hence their span will be unique. This also allows us to deal easily
+        with fixed nodes at explicit times (rather than in time slices)
+        """
+
+        if self.mut_rate is None:
+            raise RuntimeError(
+                "Cannot calculate mutation likelihoods with no mutation_rate set"
+            )
+        if unique_method == 0:
+            self.unfixed_likelihood_cache = {
+                (muts, e.span): None
+                for muts, e in zip(self.mut_edges, self.ts.edges())
+                if e.child not in self.fixednodes
+            }
+        else:
+            fixed_nodes = np.array(list(self.fixednodes))
+            keys = np.unique(
+                np.core.records.fromarrays(
+                    (self.mut_edges, self.ts.edges_right - self.ts.edges_left),
+                    names="muts,span",
+                )[np.logical_not(np.isin(self.ts.edges_child, fixed_nodes))]
+            )
+            if unique_method == 1:
+                self.unfixed_likelihood_cache = dict.fromkeys({tuple(t) for t in keys})
+            else:
+                self.unfixed_likelihood_cache = {tuple(t): None for t in keys}
+
+        if num_threads:
+            f = functools.partial(  # Set constant values for params for static _lik
+                self._lik_wrapper,
+                dt=self.timediff_lower_tri,
+                mutation_rate=self.mut_rate,
+                standardize=self.standardize,
+            )
+            if num_threads == 1:
+                # Useful for testing
+                for key in tqdm(
+                    self.unfixed_likelihood_cache.keys(),
+                    disable=not self.progress,
+                    desc="Precalculating Likelihoods",
+                ):
+                    returned_key, likelihoods = f(key)
+                    self.unfixed_likelihood_cache[returned_key] = likelihoods
+            else:
+                with tqdm(
+                    total=len(self.unfixed_likelihood_cache.keys()),
+                    disable=not self.progress,
+                    desc="Precalculating Likelihoods",
+                ) as prog_bar:
+                    with multiprocessing.Pool(processes=num_threads) as pool:
+                        for key, pmf in pool.imap_unordered(
+                            f, self.unfixed_likelihood_cache.keys()
+                        ):
+                            self.unfixed_likelihood_cache[key] = pmf
+                            prog_bar.update()
+        else:
+            for muts, span in tqdm(
+                self.unfixed_likelihood_cache.keys(),
+                disable=not self.progress,
+                desc="Precalculating Likelihoods",
+            ):
+                self.unfixed_likelihood_cache[muts, span] = self._lik(
+                    muts,
+                    span,
+                    dt=self.timediff_lower_tri,
+                    mutation_rate=self.mut_rate,
+                    standardize=self.standardize,
+                )
+
+    def get_mut_lik_fixed_node(self, edge):
+        """
+        Get the mutation likelihoods for an edge whose child is at a
+        fixed time, but whose parent may take any of the time slices in the timepoints
+        that are equal to or older than the child age. This is not cached, as it is
+        likely to be unique for each edge
+        """
+        assert (
+            edge.child in self.fixednodes
+        ), "Wrongly called fixed node function on non-fixed node"
+        assert (
+            self.mut_rate is not None
+        ), "Cannot calculate mutation likelihoods with no mutation_rate set"
+
+        mutations_on_edge = self.mut_edges[edge.id]
+        child_time = self.ts.node(edge.child).time
+        assert child_time == 0
+        # Temporary hack - we should really take a more precise likelihood
+        return self._lik(
+            mutations_on_edge,
+            edge.span,
+            self.timediff,
+            self.mut_rate,
+            standardize=self.standardize,
+        )
+
+    def get_mut_lik_lower_tri(self, edge):
+        """
+        Get the cached mutation likelihoods for an edge with non-fixed parent and child
+        nodes, returning values for all the possible time differences between timepoints
+        These values are returned as a flattened lower triangular matrix, the
+        form required in the inside algorithm.
+
+        """
+        # Debugging asserts - should probably remove eventually
+        assert (
+            edge.child not in self.fixednodes
+        ), "Wrongly called lower_tri function on fixed node"
+        assert hasattr(
+            self, "unfixed_likelihood_cache"
+        ), "Must call `precalculate_mutation_likelihoods()` before getting likelihoods"
+
+        mutations_on_edge = self.mut_edges[edge.id]
+        return self.unfixed_likelihood_cache[mutations_on_edge, edge.span]
+
+    def get_mut_lik_upper_tri(self, edge):
+        """
+        Same as :meth:`get_mut_lik_lower_tri`, but the returned array is ordered as
+        flattened upper triangular matrix (suitable for the outside algorithm), rather
+        than a lower triangular one
+        """
+        return self.get_mut_lik_lower_tri(edge)[np.concatenate(self.row_indices)]
+
+    # The following functions don't access the likelihoods directly, but allow
+    # other input arrays of length grid_size to be repeated in such a way that they can
+    # be directly multiplied by the unpacked lower triangular matrix, or arrays of length
+    # of the number of cells in the lower triangular matrix to be summed (e.g. by row)
+    # to give a shorter array of length grid_size
+
+    def make_lower_tri(self, input_array):
+        """
+        Repeat the input array row-wise to make a flattened lower triangular matrix
+        """
+        assert len(input_array) == self.grid_size
+        return input_array[self.to_lower_tri]
+
+    def rowsum_lower_tri(self, input_array):
+        """
+        Describe the reduceat trickery here. Presumably the opposite of make_lower_tri
+        """
+        assert len(input_array) == self.tri_size
+        return np.add.reduceat(input_array, self.row_indices[0])
+
+    def make_upper_tri(self, input_array):
+        """
+        Repeat the input array row-wise to make a flattened upper triangular matrix
+        """
+        assert len(input_array) == self.grid_size
+        return input_array[self.to_upper_tri]
+
+    def rowsum_upper_tri(self, input_array):
+        """
+        Describe the reduceat trickery here. Presumably the opposite of make_upper_tri
+        """
+        assert len(input_array) == self.tri_size
+        return np.add.reduceat(input_array, self.col_indices)
+
+    # Mutation & recombination algorithms on a tree sequence
+
+    def n_breaks(self, edge):
+        """
+        Number of known breakpoints, only used in recombination likelihood calc
+        """
+        return (edge.left != 0) + (edge.right != self.ts.get_sequence_length())
+
+    def combine(self, lik_1, lik_2):
+        return lik_1 * lik_2
+
+    def ratio(self, lik_1, lik_2, div_0_null=False):
+        """
+        Return the ratio of lik_1 to lik_2. In linear space, this divides lik_1 by lik_2
+        If div_0_null==True, then 0/0 is set to the null_constant
+        """
+        with np.errstate(divide="ignore", invalid="ignore"):
+            ret = lik_1 / lik_2
+        if div_0_null:
+            ret[np.isnan(ret)] = self.null_constant
+        return ret
+
+    def marginalize(self, lik):
+        """
+        Return the sum of likelihoods
+        """
+        return np.sum(lik)
+
+    def _recombination_lik(self, edge, fixed=True):
+        # Needs to return a lower tri *or* flattened array depending on `fixed`
+        raise NotImplementedError(
+            "Using the recombination clock is not currently supported"
+            ". See https://github.com/awohns/tsdate/issues/5 for details"
+        )
+        # return (
+        #     np.power(prev_state, self.n_breaks(edge)) *
+        #     np.exp(-(prev_state * self.rec_rate * edge.span * 2)))
+
+    def get_inside(self, arr, edge):
+        liks = self.identity_constant
+        if self.rec_rate is not None:
+            liks *= self._recombination_lik(edge)
+        if self.mut_rate is not None:
+            liks *= self.get_mut_lik_lower_tri(edge)
+        return self.rowsum_lower_tri(arr * liks)
+
+    def get_outside(self, arr, edge):
+        liks = self.identity_constant
+        if self.rec_rate is not None:
+            liks *= self._recombination_lik(edge)
+        if self.mut_rate is not None:
+            liks *= self.get_mut_lik_upper_tri(edge)
+        return self.rowsum_upper_tri(arr * liks)
+
+    def get_fixed(self, arr, edge):
+        liks = self.identity_constant
+        if self.rec_rate is not None:
+            liks *= self._recombination_lik(edge, fixed=True)
+        if self.mut_rate is not None:
+            liks *= self.get_mut_lik_fixed_node(edge)
+        return arr * liks
+
+    def scale_geometric(self, fraction, value):
+        return value**fraction
+
+
+class LogLikelihoods(Likelihoods):
+    """
+    Identical to the Likelihoods class but stores and returns log likelihoods
+    """
+
+    probability_space = LOG_GRID
+    identity_constant = 0.0
+    null_constant = -np.inf
+
+    """
+    Uses an alternative to logsumexp, useful for large grid sizes, see
+    http://www.nowozin.net/sebastian/blog/streaming-log-sum-exp-computation.html
+    """
+
+    @staticmethod
+    @numba.jit(nopython=True)
+    def logsumexp(X):
+        alpha = -np.inf
+        r = 0.0
+        for x in X:
+            if x != -np.inf:
+                if x <= alpha:
+                    r += np.exp(x - alpha)
+                else:
+                    r *= np.exp(alpha - x)
+                    r += 1.0
+                    alpha = x
+        return -np.inf if r == 0 else np.log(r) + alpha
+
+    @staticmethod
+    def _lik(muts, span, dt, mutation_rate, standardize=True):
+        """
+        The likelihood of an edge given a number of mutations, as set of time deltas (dt)
+        and a span. This is a static function to allow parallelization
+        """
+        ll = scipy.stats.poisson.logpmf(muts, dt * mutation_rate * span)
+        if standardize:
+            return ll - np.max(ll)
+        else:
+            return ll
+
+    @staticmethod
+    def _lik_wrapper(muts_span, dt, mutation_rate, standardize=True):
+        """
+        Needs redefining to refer to the LogLikelihoods class
+        """
+        return muts_span, LogLikelihoods._lik(
+            muts_span[0], muts_span[1], dt, mutation_rate, standardize=standardize
+        )
+
+    def rowsum_lower_tri(self, input_array):
+        """
+        The function below is equivalent to (but numba makes it faster than)
+        np.logaddexp.reduceat(input_array, self.row_indices[0])
+        """
+        assert len(input_array) == self.tri_size
+        res = list()
+        i_start = self.row_indices[0][0]
+        for i in self.row_indices[0][1:]:
+            res.append(self.logsumexp(input_array[i_start:i]))
+            i_start = i
+        res.append(self.logsumexp(input_array[i:]))
+        return np.array(res)
+
+    def rowsum_upper_tri(self, input_array):
+        """
+        The function below is equivalent to (but numba makes it faster than)
+        np.logaddexp.reduceat(input_array, self.col_indices)
+        """
+        assert len(input_array) == self.tri_size
+        res = list()
+        i_start = self.col_indices[0]
+        for i in self.col_indices[1:]:
+            res.append(self.logsumexp(input_array[i_start:i]))
+            i_start = i
+        res.append(self.logsumexp(input_array[i:]))
+        return np.array(res)
+
+    def _recombination_loglik(self, edge, fixed=True):
+        # Needs to return a lower tri *or* flattened array depending on `fixed`
+        raise NotImplementedError(
+            "Using the recombination clock is not currently supported"
+            ". See https://github.com/awohns/tsdate/issues/5 for details"
+        )
+        # return (
+        #     np.power(prev_state, self.n_breaks(edge)) *
+        #     np.exp(-(prev_state * self.rec_rate * edge.span * 2)))
+
+    def combine(self, loglik_1, loglik_2):
+        return loglik_1 + loglik_2
+
+    def ratio(self, loglik_1, loglik_2, div_0_null=False):
+        """
+        In log space, likelihood ratio is loglik_1 - loglik_2
+        If div_0_null==True, then if either is -inf it returns -inf (the null_constant)
+        """
+        with np.errstate(divide="ignore", invalid="ignore"):
+            ret = loglik_1 - loglik_2
+        if div_0_null:
+            ret[np.isnan(ret)] = self.null_constant
+        return ret
+
+    def marginalize(self, loglik):
+        """
+        Return the logged sum of likelihoods
+        """
+        return self.logsumexp(loglik)
+
+    def get_inside(self, arr, edge):
+        log_liks = self.identity_constant
+        if self.rec_rate is not None:
+            log_liks += self._recombination_loglik(edge)
+        if self.mut_rate is not None:
+            log_liks += self.get_mut_lik_lower_tri(edge)
+        return self.rowsum_lower_tri(arr + log_liks)
+
+    def get_outside(self, arr, edge):
+        log_liks = self.identity_constant
+        if self.rec_rate is not None:
+            log_liks += self._recombination_loglik(edge)
+        if self.mut_rate is not None:
+            log_liks += self.get_mut_lik_upper_tri(edge)
+        return self.rowsum_upper_tri(arr + log_liks)
+
+    def get_fixed(self, arr, edge):
+        log_liks = self.identity_constant
+        if self.rec_rate is not None:
+            log_liks += self._recombination_loglik(edge, fixed=True)
+        if self.mut_rate is not None:
+            log_liks += self.get_mut_lik_fixed_node(edge)
+        return arr + log_liks
+
+    def scale_geometric(self, fraction, value):
+        return fraction * value
+
+
+class BeliefPropagation:
+    """
+    The class that encapsulates running exact belief propagation models,
+    in particular the discrete-time inside and outside algorithms.
+    """
+
+    def __init__(self, priors, lik, *, progress=False):
+        if (
+            lik.fixednodes.intersection(priors.nonfixed_nodes)
+            or len(lik.fixednodes) + len(priors.nonfixed_nodes) != lik.ts.num_nodes
+        ):
+            raise ValueError(
+                "The prior and likelihood objects disagree on which nodes are fixed"
+            )
+        if not np.allclose(lik.timepoints, priors.timepoints):
+            raise ValueError(
+                "The prior and likelihood objects disagree on the timepoints used"
+            )
+
+        self.priors = priors
+        self.nonfixed_nodes = priors.nonfixed_nodes
+        self.lik = lik
+        self.ts = lik.ts
+
+        self.fixednodes = lik.fixednodes
+        self.progress = progress
+        # If necessary, convert priors to log space
+        self.priors.force_probability_space(lik.probability_space)
+
+        self.spans = np.bincount(
+            self.ts.edges_child,
+            weights=self.ts.edges_right - self.ts.edges_left,
+        )
+        self.spans = np.pad(self.spans, (0, self.ts.num_nodes - len(self.spans)))
+
+        self.root_spans = defaultdict(float)
+        for tree in self.ts.trees(root_threshold=2):
+            if tree.has_single_root:
+                self.root_spans[tree.root] += tree.span
+        # Add on the spans when this is a root
+        for root, span_when_root in self.root_spans.items():
+            self.spans[root] += span_when_root
+
+    # === Grouped edge iterators ===
+
+    def edges_by_parent_asc(self, grouped=True):
+        # Return an itertools.groupby object of edges grouped by parent in ascending
+        # order of the time of the parent. As tree sequence properties guarantee that
+        # edges are listed in nondecreasing order of parent time
+        # (https://tskit.readthedocs.io/en/latest/data-model.html#edge-requirements)
+        # we can simply use the standard edge order
+        if grouped:
+            return itertools.groupby(self.ts.edges(), operator.attrgetter("parent"))
+        else:
+            return self.ts.edges()
+
+    def edges_by_child_desc(self, grouped=True):
+        # Return an itertools.groupby object of edges grouped by child in descending
+        # order of the time of the child.
+        it = (
+            self.ts.edge(u)
+            for u in np.lexsort(
+                (self.ts.edges_child, -self.ts.nodes_time[self.ts.edges_child])
+            )
+        )
+        if grouped:
+            return itertools.groupby(it, operator.attrgetter("child"))
+        else:
+            return it
+
+    def edges_by_child_then_parent_desc(self, grouped=True):
+        # Return an itertools.groupby object of edges grouped by child in descending
+        # order of the time of the child, then by descending order of age of child
+        wtype = np.dtype(
+            [
+                ("child_age", self.ts.nodes_time.dtype),
+                ("child_node", self.ts.edges_child.dtype),
+                ("parent_age", self.ts.nodes_time.dtype),
+            ]
+        )
+        w = np.empty(self.ts.num_edges, dtype=wtype)
+        w["child_age"] = self.ts.nodes_time[self.ts.edges_child]
+        w["child_node"] = self.ts.edges_child
+        w["parent_age"] = -self.ts.nodes_time[self.ts.edges_parent]
+        sorted_child_parent = (
+            self.ts.edge(i)
+            for i in reversed(
+                np.argsort(w, order=("child_age", "child_node", "parent_age"))
+            )
+        )
+        if grouped:
+            return itertools.groupby(sorted_child_parent, operator.attrgetter("child"))
+        else:
+            return sorted_child_parent
+
+    # === MAIN ALGORITHMS ===
+
+    def inside_pass(self, *, standardize=True, cache_inside=False, progress=None):
+        # Use dynamic programming to find approximate posterior to sample from
+        if progress is None:
+            progress = self.progress
+        inside = self.priors.clone_with_new_data(  # store inside matrix values
+            grid_data=np.nan, fixed_data=self.lik.identity_constant
+        )
+        if cache_inside:
+            g_i = np.full(
+                (self.ts.num_edges, self.lik.grid_size), self.lik.identity_constant
+            )
+        denominator = np.full(self.ts.num_nodes, np.nan)
+        assert (
+            self.lik.standardize is False
+        ), "Marginal likelihood requires unstandardized mutation likelihoods"
+        marginal_lik = self.lik.identity_constant
+        # Iterate through the nodes via groupby on parent node
+        for parent, edges in tqdm(
+            self.edges_by_parent_asc(),
+            desc="Inside",
+            total=inside.num_nonfixed,
+            disable=not progress,
+        ):
+            """
+            for each node, find the conditional prob of age at every time
+            in time grid
+            """
+            if parent in self.fixednodes:
+                continue  # there is no hidden state for this parent - it's fixed
+            val = self.priors[parent].copy()
+            for edge in edges:
+                spanfrac = edge.span / self.spans[edge.child]
+                # Calculate vals for each edge
+                if edge.child in self.fixednodes:
+                    # NB: geometric scaling works exactly when all nodes fixed in graph
+                    # but is an approximation when times are unknown.
+                    daughter_val = self.lik.scale_geometric(spanfrac, inside[edge.child])
+                    edge_lik = self.lik.get_fixed(daughter_val, edge)
+                else:
+                    inside_values = inside[edge.child]
+                    if np.ndim(inside_values) == 0 or np.all(np.isnan(inside_values)):
+                        # Child appears fixed, or we have not visited it. Either our
+                        # edge order is wrong (bug) or we have hit a dangling node
+                        raise ValueError(
+                            "The input tree sequence includes "
+                            "dangling nodes: please simplify it"
+                        )
+                    daughter_val = self.lik.scale_geometric(
+                        spanfrac, self.lik.make_lower_tri(inside[edge.child])
+                    )
+                    edge_lik = self.lik.get_inside(daughter_val, edge)
+                val = self.lik.combine(val, edge_lik)
+                if cache_inside:
+                    g_i[edge.id] = edge_lik
+            denominator[parent] = (
+                np.max(val) if standardize else self.lik.identity_constant
+            )
+            inside[parent] = self.lik.ratio(val, denominator[parent])
+            if standardize:
+                marginal_lik = self.lik.combine(marginal_lik, denominator[parent])
+        if cache_inside:
+            self.g_i = self.lik.ratio(g_i, denominator[self.ts.edges_child, None])
+        # Keep the results in this object
+        self.inside = inside
+        self.denominator = denominator
+        # Calculate marginal likelihood
+        for root, span_when_root in self.root_spans.items():
+            spanfrac = span_when_root / self.spans[root]
+            root_val = self.lik.scale_geometric(spanfrac, inside[root])
+            marginal_lik = self.lik.combine(marginal_lik, self.lik.marginalize(root_val))
+        return marginal_lik
+
+    def outside_pass(
+        self,
+        *,
+        standardize=False,
+        ignore_oldest_root=False,
+        progress=None,
+    ):
+        # Computes the full posterior distribution on nodes, returning the
+        # posterior values. These are *not* probabilities, as they do not sum to one:
+        # to convert to probabilities, call posterior.to_probabilities()
+        #
+        # Standardizing *during* the outside process may be necessary if there is
+        # overflow, but means that we cannot check the total functional value at each node
+        #
+        # Ignoring the oldest root may also be necessary when the oldest root node
+        # causes numerical stability issues.
+        if progress is None:
+            progress = self.progress
+        if not hasattr(self, "inside"):
+            raise RuntimeError("You have not yet run the inside algorithm")
+
+        outside = self.inside.clone_with_new_data(grid_data=0, probability_space=LIN_GRID)
+        for root, span_when_root in self.root_spans.items():
+            outside[root] = span_when_root / self.spans[root]
+        outside.force_probability_space(self.inside.probability_space)
+
+        for child, edges in tqdm(
+            self.edges_by_child_desc(),
+            desc="Outside",
+            total=len(np.unique(self.ts.edges_child)),
+            disable=not progress,
+        ):
+            if child in self.fixednodes:
+                continue
+            val = np.full(self.lik.grid_size, self.lik.identity_constant)
+            for edge in edges:
+                if ignore_oldest_root:
+                    if edge.parent == self.ts.num_nodes - 1:
+                        continue
+                if edge.parent in self.fixednodes:
+                    raise RuntimeError(
+                        "Fixed nodes cannot currently be parents in the TS"
+                    )
+                # Geometric scaling works exactly for all nodes fixed in graph
+                # but is an approximation when times are unknown.
+                spanfrac = edge.span / self.spans[child]
+                try:
+                    inside_div_gi = self.lik.ratio(
+                        self.inside[edge.parent], self.g_i[edge.id], div_0_null=True
+                    )
+                except AttributeError:  # we haven't cached g_i so we recalculate
+                    daughter_val = self.lik.scale_geometric(
+                        spanfrac, self.lik.make_lower_tri(self.inside[edge.child])
+                    )
+                    edge_lik = self.lik.get_inside(daughter_val, edge)
+                    cur_g_i = self.lik.ratio(edge_lik, self.denominator[child])
+                    inside_div_gi = self.lik.ratio(
+                        self.inside[edge.parent], cur_g_i, div_0_null=True
+                    )
+                parent_val = self.lik.scale_geometric(
+                    spanfrac,
+                    self.lik.make_upper_tri(
+                        self.lik.combine(outside[edge.parent], inside_div_gi)
+                    ),
+                )
+                if standardize:
+                    parent_val = self.lik.ratio(parent_val, np.max(parent_val))
+                edge_lik = self.lik.get_outside(parent_val, edge)
+                val = self.lik.combine(val, edge_lik)
+
+            # vv[0] = 0  # Seems a hack: internal nodes should be allowed at time 0
+            assert self.denominator[edge.child] > self.lik.null_constant
+            outside[child] = self.lik.ratio(val, self.denominator[child])
+            if standardize:
+                outside[child] = self.lik.ratio(val, np.max(val))
+        self.outside = outside  # useful to access for testing purposes
+        self.posterior_grid = outside.clone_with_new_data(
+            grid_data=self.lik.combine(self.inside.grid_data, outside.grid_data),
+            fixed_data=np.nan,
+        )  # NB: we should never use the posterior for a fixed node
+
+    def node_posteriors(self):
+        """
+        Return the distribution of posterior node times as a structured array. The
+        returned value can be e.g. read into ``pandas.DataFrame`` for further analysis.
+
+        .. note::
+            The ``outside_maximization`` method does not provide node time posteriors.
+
+        :return: The distribution of posterior node times as a structured array with
+            columns as timepoints. Row ``i`` corresponds to the probabilities of
+            node ``i`` lying at each timepoint. Nodes with fixed times are set to
+            ``np.nan`` for the entire row.
+        :rtype: numpy.ndarray
+        """
+        try:
+            return self.posterior_grid.node_probability_array()
+        except AttributeError:
+            raise ValueError(
+                "Cannot get posteriors without running the outside algorithm"
+            ) from None
+
+    def outside_maximization(self, *, eps, progress=None):
+        if progress is None:
+            progress = self.progress
+        if not hasattr(self, "inside"):
+            raise RuntimeError("You have not yet run the inside algorithm")
+
+        maximized_node_times = np.zeros(self.ts.num_nodes, dtype="int")
+
+        if self.lik.probability_space == LOG_GRID:
+            poisson = scipy.stats.poisson.logpmf
+        elif self.lik.probability_space == LIN_GRID:
+            poisson = scipy.stats.poisson.pmf
+
+        mut_edges = self.lik.mut_edges
+        mrcas = np.where(
+            np.isin(np.arange(self.ts.num_nodes), self.ts.edges_child, invert=True)
+        )[0]
+        for i in mrcas:
+            if i not in self.fixednodes:
+                maximized_node_times[i] = np.argmax(self.inside[i])
+
+        for child, edges in tqdm(
+            self.edges_by_child_then_parent_desc(),
+            desc="Maximization",
+            total=len(np.unique(self.ts.edges_child)),
+            disable=not progress,
+        ):
+            if child in self.fixednodes:
+                continue
+            for edge_index, edge in enumerate(edges):
+                if edge_index == 0:
+                    youngest_par_index = maximized_node_times[edge.parent]
+                    parent_time = self.lik.timepoints[maximized_node_times[edge.parent]]
+                    ll_mut = poisson(
+                        mut_edges[edge.id],
+                        (
+                            parent_time
+                            - self.lik.timepoints[: youngest_par_index + 1]
+                            + eps
+                        )
+                        * self.lik.mut_rate
+                        * edge.span,
+                    )
+                    result = self.lik.ratio(ll_mut, np.max(ll_mut))
+                else:
+                    cur_parent_index = maximized_node_times[edge.parent]
+                    if cur_parent_index < youngest_par_index:
+                        youngest_par_index = cur_parent_index
+                    parent_time = self.lik.timepoints[maximized_node_times[edge.parent]]
+                    ll_mut = poisson(
+                        mut_edges[edge.id],
+                        (
+                            parent_time
+                            - self.lik.timepoints[: youngest_par_index + 1]
+                            + eps
+                        )
+                        * self.lik.mut_rate
+                        * edge.span,
+                    )
+                    result[: youngest_par_index + 1] = self.lik.combine(
+                        self.lik.ratio(
+                            ll_mut[: youngest_par_index + 1],
+                            np.max(ll_mut[: youngest_par_index + 1]),
+                        ),
+                        result[: youngest_par_index + 1],
+                    )
+            inside_val = self.inside[child][: (youngest_par_index + 1)]
+
+            maximized_node_times[child] = np.argmax(
+                self.lik.combine(result[: youngest_par_index + 1], inside_val)
+            )
+        # The outside_maximization method does not provide a full posterior but
+        # simply the means of the posterior distributions
+        self.posterior_mean = self.lik.timepoints[
+            np.array(maximized_node_times).astype("int")
+        ]
diff --git a/tsdate/base.py b/tsdate/node_time_class.py
similarity index 81%
rename from tsdate/base.py
rename to tsdate/node_time_class.py
index 3129edba..e5b4a209 100644
--- a/tsdate/base.py
+++ b/tsdate/node_time_class.py
@@ -27,26 +27,23 @@
 import numpy as np
 
 FLOAT_DTYPE = np.float64
-LIN = "linear"
-LOG = "logarithmic"
+LIN_GRID = "linear"
+LOG_GRID = "logarithmic"
 GAMMA_PAR = "gamma_parameter"
 
-#: The default value for `approx_prior_size` (see :func:`~tsdate.build_prior_grid` and
-#: :func:`~tsdate.build_parameter_grid`)
-DEFAULT_APPROX_PRIOR_SIZE = 10000
 
-
-class NodeGridValues:
+class NodeTimeValues:
     """
     A class to store times or discretised distributions of times for node ids. For nodes
     with fixed times, only a single time value needs to be stored. For non-fixed nodes,
-    an array of len(timepoints) probabilies is required.
+    an array of either len(timepoints) probabilties or a set of (gamma) distribution
+    parameters is required.
 
     .. note::
 
         This class is not intended to be used directly by users and may be subject
         to change of name or internal structure in future versions. For details on
-        how to create a ``NodeGridValues`` object to be used as a prior, see
+        how to create a ``NodeTimeValues`` object to be used as a prior, see
         :ref:`sec_priors`.
 
     :ivar num_nodes: The number of nodes that will be stored in this object
@@ -79,7 +76,12 @@ def __init__(
             raise ValueError(
                 "All non fixed node ids must be between zero and the total node number"
             )
-        grid_size = len(timepoints) if type(timepoints) is np.ndarray else timepoints
+        try:
+            grid_size = len(timepoints)
+            self.probability_space = LIN_GRID
+        except TypeError:
+            grid_size = timepoints
+            self.probability_space = GAMMA_PAR
         self.timepoints = timepoints
         # Make timepoints immutable so no risk of overwritting them with copy
         self.timepoints.setflags(write=False)
@@ -95,7 +97,6 @@ def __init__(
         self.row_lookup[np.logical_not(np.isin(np.arange(num_nodes), nonfixed_nodes))] = (
             -np.arange(num_nodes - self.num_nonfixed) - 1
         )
-        self.probability_space = LIN
 
     def force_probability_space(self, probability_space):
         """
@@ -108,23 +109,23 @@ def force_probability_space(self, probability_space):
             probability_space,
             "space",
         )
-        if probability_space == LIN:
-            if self.probability_space == LIN:
+        if probability_space == LIN_GRID:
+            if self.probability_space == LIN_GRID:
                 pass
-            elif self.probability_space == LOG:
+            elif self.probability_space == LOG_GRID:
                 self.grid_data = np.exp(self.grid_data)
                 self.fixed_data = np.exp(self.fixed_data)
-                self.probability_space = LIN
+                self.probability_space = LIN_GRID
             else:
                 raise TypeError("Cannot force " + " ".join(descr))
-        elif probability_space == LOG:
-            if self.probability_space == LOG:
+        elif probability_space == LOG_GRID:
+            if self.probability_space == LOG_GRID:
                 pass
-            elif self.probability_space == LIN:
+            elif self.probability_space == LIN_GRID:
                 with np.errstate(divide="ignore", invalid="ignore"):
                     self.grid_data = np.log(self.grid_data)
                     self.fixed_data = np.log(self.fixed_data)
-                self.probability_space = LOG
+                self.probability_space = LOG_GRID
             else:
                 raise TypeError("Cannot force " + " ".join(descr))
         elif probability_space == GAMMA_PAR:
@@ -140,22 +141,22 @@ def standardize(self):
         Standardize grid data so the max for each row is one (in linear space) or zero
         (in logarithmic space)
 
-        TODO - is it clear why we omit the first element of the
+        TODO - is it clear why we omit the first element of the grid?
         """
         rowmax = self.grid_data[:, 1:].max(axis=1)
-        if self.probability_space == LIN:
+        if self.probability_space == LIN_GRID:
             self.grid_data = self.grid_data / rowmax[:, np.newaxis]
-        elif self.probability_space == LOG:
+        elif self.probability_space == LOG_GRID:
             self.grid_data = self.grid_data - rowmax[:, np.newaxis]
         else:
-            raise RuntimeError("Probability space is not", LIN, "or", LOG)
+            raise RuntimeError("Probability space is not", LIN_GRID, "or", LOG_GRID)
 
     def to_probabilities(self):
         """
         Change grid data into probabilities (i.e. each row sums to one in linear or zero
         in logarithmic space)
         """
-        if self.probability_space != LIN:
+        if self.probability_space != LIN_GRID:
             raise NotImplementedError("Can only convert to probabilities in linear space")
         assert not np.any(self.grid_data < 0)
         self.grid_data = self.grid_data / self.grid_data.sum(axis=1)[:, np.newaxis]
@@ -178,7 +179,7 @@ def clone_with_new_data(
         self, grid_data=np.nan, fixed_data=None, probability_space=None
     ):
         """
-        Take the row indices etc from an existing NodeGridValues object and make a new
+        Take the row indices etc from an existing NodeTimeValues object and make a new
         similar one but with different data. If grid_data is a single number, fill the
         entire data array with that, otherwise assume the data is a numpy array of the
         correct size to fill the gridded data. If grid_data is None, fill with NaN
@@ -200,7 +201,7 @@ def fill_fixed(orig, fixed_data):
                     orig.fixed_data.shape, fixed_data, dtype=orig.fixed_data.dtype
                 )
 
-        new_obj = NodeGridValues.__new__(NodeGridValues)
+        new_obj = NodeTimeValues.__new__(NodeTimeValues)
         new_obj.num_nodes = self.num_nodes
         new_obj.nonfixed_nodes = self.nonfixed_nodes
         new_obj.num_nonfixed = self.num_nonfixed
@@ -233,8 +234,15 @@ def fill_fixed(orig, fixed_data):
             new_obj.probability_space = probability_space
         return new_obj
 
-    def nonfixed_dict(self):
+    def node_probability_array(self):
         """
-        Return a dictionary mapping integer node ids to their data.
+        Return a structured array with columns for each timepoint and rows for each node.
+        Fixed nodes have np.nan in all columns, non-fixed nodes have the grid data.
         """
-        return {n: self[n] for n in self.nonfixed_nodes}
+        if self.probability_space != LIN_GRID:
+            raise ValueError("Only valid for linear-space probabilities")
+        dtype = self.grid_data.dtype
+        struct_dtype = [(str(t), self.grid_data.dtype) for t in self.timepoints]
+        result = np.full((self.num_nodes, len(self.timepoints)), np.nan, dtype=dtype)
+        result[self.nonfixed_nodes, :] = self.grid_data[:, :]
+        return result.ravel().view(dtype=struct_dtype)
diff --git a/tsdate/prior.py b/tsdate/prior.py
index 9a24207c..45211de1 100644
--- a/tsdate/prior.py
+++ b/tsdate/prior.py
@@ -36,7 +36,11 @@
 import tskit
 from tqdm.auto import tqdm
 
-from . import base, cache, demography, provenance, util
+from . import cache, demography, node_time_class, provenance, util
+
+#: The default value for `approx_prior_size` (see :func:`~tsdate.build_prior_grid` and
+#: :func:`~tsdate.build_parameter_grid`)
+DEFAULT_APPROX_PRIOR_SIZE = 10000
 
 
 class PriorParams(namedtuple("PriorParamsBase", "alpha, beta, mean, var")):
@@ -191,7 +195,7 @@ def add(self, total_tips, approximate=None):
         if approximate is not None:
             self.approximate = approximate
         else:
-            if total_tips >= base.DEFAULT_APPROX_PRIOR_SIZE:
+            if total_tips >= DEFAULT_APPROX_PRIOR_SIZE:
                 self.approximate = True
             else:
                 self.approximate = False
@@ -202,10 +206,10 @@ def add(self, total_tips, approximate=None):
                 " the ConditionalCoalescentTimes object with a non-zero number"
             )
 
-        if not self.approximate and total_tips >= base.DEFAULT_APPROX_PRIOR_SIZE:
+        if not self.approximate and total_tips >= DEFAULT_APPROX_PRIOR_SIZE:
             logging.warning(
                 "Calculating exact priors for more than "
-                f"{base.DEFAULT_APPROX_PRIOR_SIZE} tips. Consider "
+                f"{DEFAULT_APPROX_PRIOR_SIZE} tips. Consider "
                 "setting `approximate=True` for a faster calculation."
             )
 
@@ -217,7 +221,9 @@ def add(self, total_tips, approximate=None):
         # they are stored separately to obviate need to move between them
         # We should only use prior[2] upwards
         priors = np.full(
-            (total_tips + 1, len(PriorParams._fields)), np.nan, dtype=base.FLOAT_DTYPE
+            (total_tips + 1, len(PriorParams._fields)),
+            np.nan,
+            dtype=node_time_class.FLOAT_DTYPE,
         )
 
         if self.approximate:
@@ -374,7 +380,9 @@ def get_mixture_prior_params(self, spans_by_samples):
         seen_mixtures = {}
         # allocate space for params for all nodes, even though we only use nodes_to_date
         num_nodes, num_params = spans_by_samples.ts.num_nodes, len(param_cols)
-        priors = np.full((num_nodes + 1, num_params), np.nan, dtype=base.FLOAT_DTYPE)
+        priors = np.full(
+            (num_nodes + 1, num_params), np.nan, dtype=node_time_class.FLOAT_DTYPE
+        )
         for node in tqdm(
             spans_by_samples.nodes_to_date,
             total=len(spans_by_samples.nodes_to_date),
@@ -467,7 +475,7 @@ def __init__(self, tree_sequence, *, progress=False, allow_unary=False):
 
         # We will store the spans in here, and normalize them at the end
         self._spans = defaultdict(
-            lambda: defaultdict(lambda: defaultdict(base.FLOAT_DTYPE))
+            lambda: defaultdict(lambda: defaultdict(node_time_class.FLOAT_DTYPE))
         )
 
         if not allow_unary:
@@ -853,7 +861,7 @@ def finalize(self):
         spans_dtype = np.dtype(
             {
                 "names": ("descendant_tips", "span"),
-                "formats": (np.uint64, base.FLOAT_DTYPE),
+                "formats": (np.uint64, node_time_class.FLOAT_DTYPE),
             }
         )
 
@@ -995,7 +1003,7 @@ def fill_priors(
     """
     Take the alpha and beta values from the node_parameters array, which contains
     one row for each node in the TS (including fixed nodes)
-    and fill out a NodeGridValues object with the prior values from the
+    and fill out a NodeTimeValues object with the prior values from the
     gamma or lognormal distribution with those parameters.
 
     The `population_size` can be a scalar, or an object with a `.to_natural_timescale`
@@ -1021,7 +1029,7 @@ def fill_priors(
     datable_nodes = np.where(datable_nodes)[0]
 
     # convert timepoints to generational timescale
-    prior_times = base.NodeGridValues(
+    prior_times = node_time_class.NodeTimeValues(
         ts.num_nodes,
         datable_nodes[np.argsort(ts.nodes_time[datable_nodes])].astype(np.int32),
         population_size.to_natural_timescale(timepoints),
@@ -1060,7 +1068,7 @@ def __init__(
     ):
         if approximate_priors:
             if not approx_prior_size:
-                approx_prior_size = base.DEFAULT_APPROX_PRIOR_SIZE
+                approx_prior_size = DEFAULT_APPROX_PRIOR_SIZE
         else:
             if approx_prior_size is not None:
                 raise ValueError(
@@ -1107,7 +1115,9 @@ def make_discretised_prior(self, population_size, timepoints=20, progress=False)
             timepoints = create_timepoints(self.base_priors, timepoints + 1)
         elif isinstance(timepoints, np.ndarray):
             try:
-                timepoints = np.sort(timepoints.astype(base.FLOAT_DTYPE, casting="safe"))
+                timepoints = np.sort(
+                    timepoints.astype(node_time_class.FLOAT_DTYPE, casting="safe")
+                )
             except TypeError:
                 raise TypeError(
                     "Timepoints array cannot be converted to float dtype"
@@ -1152,12 +1162,12 @@ def make_parameter_grid(self, population_size, progress=False):
         datable_nodes[ts.samples()] = False
         datable_nodes = np.where(datable_nodes)[0]
 
-        prior_pars = base.NodeGridValues(
+        prior_pars = node_time_class.NodeTimeValues(
             self.tree_sequence.num_nodes,
             datable_nodes[np.argsort(ts.nodes_time[datable_nodes])].astype(np.int32),
             np.array([0, np.inf]),
         )
-        prior_pars.probability_space = base.GAMMA_PAR
+        prior_pars.probability_space = node_time_class.GAMMA_PAR
 
         shape = self.prior_params[:, PriorParams.field_index("alpha")]
         rate = self.prior_params[:, PriorParams.field_index("beta")]
@@ -1204,7 +1214,8 @@ def prior_grid(
         the result. Default: False
     :param int approx_prior_size: Number of samples above which a precalculated prior is
         used. Only valid if ``approximate_priors`` is True. Default: ``None``, treated as
-        :data:`~tsdate.base.DEFAULT_APPROX_PRIOR_SIZE` if ``approximate_priors`` is True.
+        :data:`~tsdate.prior.DEFAULT_APPROX_PRIOR_SIZE` if
+        ``approximate_priors`` is True.
     :param string prior_distr: What distribution to use to approximate the conditional
         coalescent prior. Can be "lognorm" for the lognormal distribution (generally a
         better fit, but slightly slower to calculate) or "gamma" for the gamma
@@ -1212,7 +1223,7 @@ def prior_grid(
         "lognorm"
     :return: A prior object to pass to :func:`date` and similar functions containing
         prior values for inference and a discretised time grid
-    :rtype:  base.NodeGridValues
+    :rtype:  node_time_class.NodeTimeValues
     """
 
     mixture_prior = MixturePrior(
@@ -1254,8 +1265,8 @@ def parameter_grid(
         the result. Default: False
     :param int approx_prior_size: Number of samples above which a precalculated prior is
         used. Only valid if ``approximate_priors`` is True. Default: ``None``, treated as
-        :data:`~tsdate.base.DEFAULT_APPROX_PRIOR_SIZE` if ``approximate_priors`` is True.
-    :rtype:  base.NodeGridValues
+        :data:`~tsdate.prior.DEFAULT_APPROX_PRIOR_SIZE` if ``approximate_priors`` is True.
+    :rtype:  node_time_class.NodeTimeValues
     """
 
     mixture_prior = MixturePrior(
diff --git a/tsdate/variational.py b/tsdate/variational.py
index 0aa51910..a80cd3e4 100644
--- a/tsdate/variational.py
+++ b/tsdate/variational.py
@@ -109,14 +109,15 @@ def _rescale(x, s):
 
 class ExpectationPropagation:
     r"""
-    Expectation propagation (EP) algorithm to infer approximate marginal
-    distributions for node ages.
+    The class that encapsulates running the variational gamma approach to
+    tsdate fitting. This contains the Expectation propagation (EP) algorithm
+    to infer approximate marginal distributions for node ages.
 
     The probability model has the form,
 
     .. math::
 
-        \prod_{i \in \mathcal{N} f(t_i | \theta_i)
+        \prod_{i \in \mathcal{N}} f(t_i | \theta_i)
         \prod_{(i,j) \in \mathcal{E}} g(y_ij | t_i - t_j)
 
     where :math:`f(.)` is a prior distribution on node ages with parameters
@@ -125,7 +126,7 @@ class ExpectationPropagation:
 
     .. math::
 
-        \prod_{i \in \mathcal{N} q(t_i | \eta_i)
+        \prod_{i \in \mathcal{N}} q(t_i | \eta_i)
         \prod_{(i,j) \in \mathcal{E}} q(t_i | \gamma_{ij}) q(t_j | \kappa_{ij})
 
     where :math:`q(.)` are pseudo-gamma distributions (termed 'factors'), and
@@ -143,10 +144,8 @@ class ExpectationPropagation:
     @staticmethod
     @numba.njit(_void(_f2r, _i1r, _i1r))
     def _check_valid_constraints(constraints, edges_parent, edges_child):
-        """
-        Check that upper-bound on node age is greater than maximum lower-bound
-        for ages of descendants
-        """
+        # Check that upper-bound on node age is greater than maximum lower-bound
+        # for ages of descendants
         lower_max = constraints[:, LOWER].copy()
         for p, c in zip(edges_parent, edges_child):
             lower_max[p] = max(lower_max[c], lower_max[p])
@@ -173,7 +172,7 @@ def _check_valid_state(
         edge_factors,
         block_factors,
     ):
-        """Check that the messages sum to the posterior (debugging only)"""
+        # Check that the messages sum to the posterior (debugging only)
         block_one, block_two = block_nodes
         posterior_check = np.zeros(posterior.shape)
         for i, (p, c) in enumerate(zip(edges_parent, edges_child)):
@@ -223,7 +222,7 @@ def __init__(self, ts, *, mutation_rate, singletons_phased=True):
         self.sizebiased_likelihoods, _ = count_mutations(ts, size_biased=True)
         self.sizebiased_likelihoods[:, 1] *= mutation_rate
         count_timing -= time.time()
-        logger.info(f"Extracted mutations in {abs(count_timing):.2f} seconds")
+        logger.debug(f"Extracted mutations in {abs(count_timing):.2f} seconds")
 
         # count mutations in singleton blocks
         phase_timing = time.time()
@@ -239,7 +238,7 @@ def __init__(self, ts, *, mutation_rate, singletons_phased=True):
         phase_timing -= time.time()
         logger.info(f"Found {num_unphased} unphased singleton mutations")
         logger.info(f"Split unphased singleton edges into {num_blocks} blocks")
-        logger.info(f"Phased singletons in {abs(phase_timing):.2f} seconds")
+        logger.debug(f"Phased singletons in {abs(phase_timing):.2f} seconds")
 
         # mutable
         self.node_factors = np.zeros((ts.num_nodes, 2, 2))
@@ -288,30 +287,28 @@ def propagate_likelihood(
         min_step,
         unphased,
     ):
-        """
-        Update approximating factors for Poisson mutation likelihoods on edges.
-
-        :param ndarray edges_parent: integer array of parent ids per edge
-        :param ndarray edges_child: integer array of child ids per edge
-        :param ndarray likelihoods: array of dimension `[num_edges, 2]`
-            containing mutation count and mutational target size per edge.
-        :param ndarray constraints: array of dimension `[num_nodes, 2]`
-            containing lower and upper bounds for each node.
-        :param ndarray posterior: array of dimension `[num_nodes, 2]`
-            containing natural parameters for each node, updated in-place.
-        :param ndarray factors: array of dimension `[num_edges, 2, 2]`
-            containing parent and child factors (natural parameters) for each
-            edge, updated in-place.
-        :param ndarray lognorm: array of dimension `[num_edges]`
-            containing the approximate normalizing constants per edge,
-            updated in-place.
-        :param ndarray scale: array of dimension `[num_nodes]` containing a
-            scaling factor for the posteriors, updated in-place.
-        :param float max_shape: the maximum allowed shape for node posteriors.
-        :param float min_step: the minimum allowed step size in (0, 1).
-        :param bool unphased: if True, edges are treated as blocks of unphased
-            singletons in contemporary individuals
-        """
+        # Update approximating factors for Poisson mutation likelihoods on edges.
+        #
+        # :param numpy.ndarray edges_parent: integer array of parent ids per edge
+        # :param numpy.ndarray edges_child: integer array of child ids per edge
+        # :param numpy.ndarray likelihoods: array of dimension `[num_edges, 2]`
+        #     containing mutation count and mutational target size per edge.
+        # :param numpy.ndarray constraints: array of dimension `[num_nodes, 2]`
+        #     containing lower and upper bounds for each node.
+        # :param numpy.ndarray posterior: array of dimension `[num_nodes, 2]`
+        #     containing natural parameters for each node, updated in-place.
+        # :param numpy.ndarray factors: array of dimension `[num_edges, 2, 2]`
+        #     containing parent and child factors (natural parameters) for each
+        #     edge, updated in-place.
+        # :param numpy.ndarray lognorm: array of dimension `[num_edges]`
+        #     containing the approximate normalizing constants per edge,
+        #     updated in-place.
+        # :param numpy.ndarray scale: array of dimension `[num_nodes]` containing a
+        #     scaling factor for the posteriors, updated in-place.
+        # :param float max_shape: the maximum allowed shape for node posteriors.
+        # :param float min_step: the minimum allowed step size in (0, 1).
+        # :param bool unphased: if True, edges are treated as blocks of unphased
+        #     singletons in contemporary individuals
 
         assert constraints.shape == posterior.shape
         assert edges_child.size == edges_parent.size
@@ -435,25 +432,21 @@ def twin_projection(x, y):
     @staticmethod
     @numba.njit(_void(_b1r, _f2w, _f3w, _f1w, _f, _i, _f))
     def propagate_prior(free, posterior, factors, scale, max_shape, em_maxitt, em_reltol):
-        """
-        Update approximating factors for global prior.
-
-        :param ndarray free: boolean array indicating if prior should be
-            applied to node
-        :param ndarray penalty: initial value for regularisation penalty
-        :param ndarray posterior: rows are nodes, columns are first and
-            second natural parameters of gamma posteriors. Updated in
-            place.
-        :param ndarray factors: rows are nodes, columns index different
-            types of updates. Updated in place.
-        :param ndarray scale: array of dimension `[num_nodes]` containing a
-            scaling factor for the posteriors, updated in-place.
-        :param float max_shape: the maximum allowed shape for node posteriors.
-        :param int em_maxitt: the maximum number of EM iterations to use when
-            fitting the regularisation.
-        :param int em_reltol: the termination criterion for relative change in
-            log-likelihood.
-        """
+        # Update approximating factors for global prior.
+        #
+        # :param ndarray free: boolean array for if prior should be applied to node
+        # :param ndarray penalty: initial value for regularisation penalty
+        # :param ndarray posterior: rows are nodes, columns are first and
+        #     second natural parameters of gamma posteriors. Updated in place.
+        # :param ndarray factors: rows are nodes, columns index different
+        #     types of updates. Updated in place.
+        # :param ndarray scale: array of dimension `[num_nodes]` containing a
+        #     scaling factor for the posteriors, updated in-place.
+        # :param float max_shape: the maximum allowed shape for node posteriors.
+        # :param int em_maxitt: the maximum number of EM iterations to use when
+        #     fitting the regularisation.
+        # :param int em_reltol: the termination criterion for relative change in
+        #     log-likelihood.
 
         assert free.size == posterior.shape[0]
         assert factors.shape == (free.size, 2, 2)
@@ -503,33 +496,29 @@ def propagate_mutations(
         scale,
         unphased,
     ):
-        """
-        Calculate posteriors for mutations.
-
-        :param ndarray mutations_order: integer array giving order in
-            which to traverse mutations
-        :param ndarray mutations_posterior: array of dimension `(num_mutations, 2)`
-            containing natural parameters for each mutation, modified in place
-        :param ndarray mutations_phase: array of dimension `(num_mutations, )`
-            containing mutation phase, modified in place
-        :param ndarray mutations_edge: integer array giving edge for each
-            mutation
-        :param ndarray edges_parent: integer array of parent ids per edge
-        :param ndarray edges_child: integer array of child ids per edge
-        :param ndarray likelihoods: array of dimension `[num_edges, 2]`
-            containing mutation count and mutational target size per edge.
-        :param ndarray constraints: array of dimension `[num_nodes, 2]`
-            containing lower and upper bounds for each node.
-        :param ndarray posterior: array of dimension `[num_nodes, 2]`
-            containing natural parameters for each node
-        :param ndarray factors: array of dimension `[num_edges, 2, 2]`
-            containing parent and child factors (natural parameters) for each
-            edge
-        :param ndarray scale: array of dimension `(num_nodes, )` containing a
-            scaling factor for the posteriors
-        :param bool unphased: if True, edges are treated as blocks of unphased
-            singletons in contemporary individuals
-        """
+        # Calculate posteriors for mutations. (publicly undocumented)
+        #
+        # :param ndarray mutations_order: integer array giving order in
+        #     which to traverse mutations
+        # :param ndarray mutations_posterior: array of dimension `(num_mutations, 2)`
+        #     containing natural parameters for each mutation, modified in place
+        # :param ndarray mutations_phase: array of dimension `(num_mutations, )`
+        #     containing mutation phase, modified in place
+        # :param ndarray mutations_edge: integer array giving edge for each mutation
+        # :param ndarray edges_parent: integer array of parent ids per edge
+        # :param ndarray edges_child: integer array of child ids per edge
+        # :param ndarray likelihoods: array of dimension `[num_edges, 2]`
+        #     containing mutation count and mutational target size per edge.
+        # :param ndarray constraints: array of dimension `[num_nodes, 2]`
+        #     containing lower and upper bounds for each node.
+        # :param ndarray posterior: array of dimension `[num_nodes, 2]`
+        #     containing natural parameters for each node
+        # :param ndarray factors: array of dimension `[num_edges, 2, 2]`
+        #     containing parent and child factors (natural parameters) for each edge
+        # :param ndarray scale: array of dimension `(num_nodes, )` containing a
+        #     scaling factor for the posteriors
+        # :param bool unphased: if True, edges are treated as blocks of unphased
+        #     singletons in contemporary individuals
 
         # TODO: scale should be 1.0, can we delete
         # TODO: we don't seem to need to damp?
@@ -633,7 +622,7 @@ def rescale_factors(
         block_factors,
         scale,
     ):
-        """Incorporate scaling term into factors and reset"""
+        # Incorporate scaling term into factors and reset
         p, c = edges_parent, edges_child
         j, k = block_nodes
         edge_factors[:, ROOTWARD] *= scale[p, np.newaxis]
@@ -654,7 +643,7 @@ def iterate(
         regularise=True,
         check_valid=False,  # for debugging
     ):
-        # pass through singleton blocks
+        logger.debug("Passing through singleton blocks")
         self.propagate_likelihood(
             self.block_order,
             self.block_nodes[ROOTWARD],
@@ -670,7 +659,7 @@ def iterate(
             USE_BLOCK_LIKELIHOOD,
         )
 
-        # rootward + leafward pass through edges
+        logger.debug("Rootward + leafward pass through edges")
         self.propagate_likelihood(
             self.edge_order,
             self.edge_parents,
@@ -686,8 +675,8 @@ def iterate(
             USE_EDGE_LIKELIHOOD,
         )
 
-        # exponential regularization on roots
         if regularise:
+            logger.debug("Exponential regularization on roots")
             self.propagate_prior(
                 self.roots,
                 self.node_posterior,
@@ -698,7 +687,7 @@ def iterate(
                 em_reltol,
             )
 
-        # absorb the scaling term into the factors
+        logger.debug("Absorbing scaling term into the factors")
         self.rescale_factors(
             self.edge_parents,
             self.edge_children,
@@ -729,7 +718,7 @@ def rescale(
         quantile_width=0.5,
         progress=False,
     ):
-        """Normalise posteriors so that empirical mutation rate is constant"""
+        # Normalise posteriors so that empirical mutation rate is constant
         likelihoods = self.edge_likelihoods if rescale_segsites \
             else self.sizebiased_likelihoods  # fmt: skip
         reallocate_unphased(  # correct mutation counts for unphased singletons
@@ -769,7 +758,7 @@ def rescale(
             quantile_width,
         )
 
-    def run(
+    def infer(
         self,
         *,
         ep_iterations=10,
@@ -781,6 +770,8 @@ def run(
         regularise=True,
         progress=None,
     ):
+        # Run multiple rounds of expectation propagation, and return stats
+        self.mean_edge_logconst = []  # Undocumented: can be used to assess convergence
         nodes_timing = time.time()
         for _ in tqdm(
             np.arange(ep_iterations),
@@ -792,6 +783,8 @@ def run(
                 min_step=min_step,
                 regularise=regularise,
             )
+            self.mean_edge_logconst.append(np.mean(self.edge_logconst))
+
         nodes_timing -= time.time()
         skipped_edges = np.sum(np.isnan(self.edge_logconst))
         logger.info(f"Skipped {skipped_edges} edges with invalid factors")
@@ -799,6 +792,7 @@ def run(
 
         muts_timing = time.time()
         mutations_phased = self.mutation_blocks == tskit.NULL
+        logger.debug("Passing through unphased singletons")
         self.propagate_mutations(  # unphased singletons
             self.mutation_order[~mutations_phased],
             self.mutation_posterior,
@@ -813,6 +807,7 @@ def run(
             self.node_scale,
             USE_BLOCK_LIKELIHOOD,
         )
+        logger.debug("Passing through phased mutations")
         self.propagate_mutations(  # phased mutations
             self.mutation_order[mutations_phased],
             self.mutation_posterior,
@@ -857,9 +852,7 @@ def run(
             logger.info(f"Timescale rescaled in {abs(rescale_timing):.2f} seconds")
 
     def node_moments(self):
-        """
-        Posterior mean and variance of node ages
-        """
+        # Posterior mean and variance of node ages (equivalent to node_posteriors)
         alpha, beta = self.node_posterior.T
         nodes_mn = np.ascontiguousarray(self.node_constraints[:, 0])
         nodes_va = np.zeros(nodes_mn.size)
@@ -869,9 +862,7 @@ def node_moments(self):
         return nodes_mn, nodes_va
 
     def mutation_moments(self):
-        """
-        Posterior mean and variance of mutation ages
-        """
+        # Posterior mean and variance of mutation ages
         alpha, beta = self.mutation_posterior.T
         muts_mn = np.full(alpha.size, np.nan)
         muts_va = np.full(alpha.size, np.nan)
@@ -881,13 +872,59 @@ def mutation_moments(self):
         return muts_mn, muts_va
 
     def mutation_mapping(self):
-        """
-        Map from mutations to edges and subtended node, using estimated singleton
-        phase (if singletons were unphased)
-        """
+        # Map from mutations to edges and subtended node, using estimated singleton
+        # phase (if singletons were unphased)
+
         # TODO: should these be copies? Should members be readonly?
         return self.mutation_edges, self.mutation_nodes
 
+    def marginal_likelihood(self):
+        # Return the marginal likelihood of the data given the model
+
+        # TODO: implement
+        return None
+
+    def node_posteriors(self):
+        """
+        Return parameters specifying the inferred posterior distribution of node
+        times which can be e.g. read into a ``pandas.DataFrame`` for further analysis.
+        The mean times are not strictly constrained by topology, so unlike the
+        ``nodes_time`` attribute of a tree sequence, the mean time of a parent node
+        may occasionally be less than that of one of its children.
+
+        :return: The distribution of posterior node times as a structured array of
+            mean and variance. Row ``i`` gives the mean and variance of inferred
+            node times for node ``i``.
+        :rtype: numpy.ndarray
+        """
+        node_mn, node_va = self.node_moments()
+        dtype = [("mean", node_mn.dtype), ("variance", node_va.dtype)]
+        data = np.empty(node_mn.size, dtype=dtype)
+        data["mean"] = node_mn
+        data["variance"] = node_va
+        return data
+
+    def mutation_posteriors(self):
+        """
+        Return parameters specifying the inferred posterior distribution of mutation
+        times which can be e.g. read into a ``pandas.DataFrame`` for further analysis.
+        These are calculated as the midpoint distribution of the posterior node time
+        distributions of the node above and below the mutation. Note that this means
+        it is possible for a mean mutation time not to lie between the mean values of
+        its parent and child nodes.
+
+        :return: The distribution of posterior mutation times as a structured array
+            of mean and variance. Row ``i`` gives the mean and variance of inferred
+            mutations times for mutation ``i``.
+        :rtype: numpy.ndarray
+        """
+        mut_mn, mut_va = self.mutation_moments()
+        dtype = [("mean", mut_mn.dtype), ("variance", mut_va.dtype)]
+        data = np.empty(mut_mn.size, dtype=dtype)
+        data["mean"] = mut_mn
+        data["variance"] = mut_va
+        return data
+
 
 # NB: used for debugging
 # def date(
@@ -913,7 +950,7 @@ def mutation_mapping(self):
 #         mutation_rate=mutation_rate,
 #         singletons_phased=singletons_phased,
 #     )
-#     posterior.run(
+#     posterior.infer(
 #         ep_maxitt=max_iterations,
 #         max_shape=max_shape,
 #         rescale_intervals=rescaling_intervals,