diff --git a/pySuStaIn/LongitudinalZscoreSustain.py b/pySuStaIn/LongitudinalZscoreSustain.py index 87b0db4..fd481f0 100644 --- a/pySuStaIn/LongitudinalZscoreSustain.py +++ b/pySuStaIn/LongitudinalZscoreSustain.py @@ -24,15 +24,22 @@ ``subject_ids`` (n_visits,) groups visits by subject, and rows for a subject are assumed to be in chronological order. """ +from collections import defaultdict + import numpy as np +from pySuStaIn.AbstractSustain import AbstractSustainData from pySuStaIn.ZscoreSustain import ZscoreSustain, ZScoreSustainData -class LongitudinalZScoreSustainData: +class LongitudinalZScoreSustainData(AbstractSustainData): def __init__(self, visit_data, subject_ids, numStages): self.visit_data = np.asarray(visit_data, dtype=float) self.subject_ids = np.asarray(subject_ids) + if self.visit_data.ndim != 2: + raise ValueError("visit_data must be a 2D array of shape (visits, biomarkers)") + if self.subject_ids.shape[0] != self.visit_data.shape[0]: + raise ValueError("subject_ids must have one entry per visit_data row") self.__numStages = numStages # group visit rows by subject, preserving first-appearance order and the # within-subject row order (assumed chronological) @@ -98,12 +105,18 @@ def _calculate_likelihood_stage(self, sustainData, S): M = len(groups) Np1 = E.shape[1] J = np.zeros((M, Np1)) + # Run the monotone-path DP vectorised over subjects, grouped by their + # number of visits V (same operations as the per-subject recursion, + # batched). g_v(k) = sum_{k'>=k} e_v(k') g_{v+1}(k') is a reverse cumsum. + by_v = defaultdict(list) for m, rows in enumerate(groups): - e = E[rows] # (V, N+1), chronological - V = e.shape[0] - g = np.ones(Np1) # g_{V+1} - for v in range(V - 1, 0, -1): # visits V..2 (0-indexed V-1..1) - cont = e[v] * g - g = np.cumsum(cont[::-1])[::-1] # g_v(k) = sum_{k'>=k} e_v(k') g_{v+1}(k') - J[m] = e[0] * g # joint with baseline stage k_1 + by_v[len(rows)].append(m) + for V, subj_idx in by_v.items(): + subj_idx = np.asarray(subj_idx) + A = np.stack([E[groups[m]] for m in subj_idx], axis=0) # (n_g, V, N+1) + g = np.ones((len(subj_idx), Np1)) # g_{V+1} + for v in range(V - 1, 0, -1): # visits V..2 + cont = A[:, v, :] * g + g = np.cumsum(cont[:, ::-1], axis=1)[:, ::-1] + J[subj_idx] = A[:, 0, :] * g # joint with baseline stage return J diff --git a/pySuStaIn/ZscoreSustain.py b/pySuStaIn/ZscoreSustain.py index 9f570fb..9c15fe9 100644 --- a/pySuStaIn/ZscoreSustain.py +++ b/pySuStaIn/ZscoreSustain.py @@ -286,13 +286,13 @@ def _optimise_parameters(self, sustainData, S_init, f_init, rng): if np.any(min_filter): min_zscore_bound = max(possible_zscores_biomarker[min_filter]) min_zscore_bound_event = events[((self.stage_zscore[0] == min_zscore_bound).astype(int) + (self.stage_biomarker_index[0] == selected_biomarker).astype(int)) == 2] - move_event_to_lower_bound = current_location[min_zscore_bound_event] + 1 + move_event_to_lower_bound = int(current_location[min_zscore_bound_event[0]] + 1) else: move_event_to_lower_bound = 0 if np.any(max_filter): max_zscore_bound = min(possible_zscores_biomarker[max_filter]) max_zscore_bound_event = events[((self.stage_zscore[0] == max_zscore_bound).astype(int) + (self.stage_biomarker_index[0] == selected_biomarker).astype(int)) == 2] - move_event_to_upper_bound = current_location[max_zscore_bound_event] + move_event_to_upper_bound = int(current_location[max_zscore_bound_event[0]]) else: move_event_to_upper_bound = N # FIXME: hack because python won't produce an array in range (N,N), while matlab will produce an array (N)... urgh @@ -389,14 +389,14 @@ def _perform_mcmc(self, sustainData, seq_init, f_init, n_iterations, seq_sigma, if np.any(min_filter): min_zscore_bound = max(possible_zscores_biomarker[min_filter]) min_zscore_bound_event = events[((self.stage_zscore[0] == min_zscore_bound).astype(int) + (self.stage_biomarker_index[0] == selected_biomarker).astype(int)) == 2] - move_event_to_lower_bound = current_location[min_zscore_bound_event] + 1 + move_event_to_lower_bound = int(current_location[min_zscore_bound_event[0]] + 1) else: move_event_to_lower_bound = 0 if np.any(max_filter): max_zscore_bound = min(possible_zscores_biomarker[max_filter]) max_zscore_bound_event = events[((self.stage_zscore[0] == max_zscore_bound).astype(int) + (self.stage_biomarker_index[0] == selected_biomarker).astype(int)) == 2] - move_event_to_upper_bound = current_location[max_zscore_bound_event] + move_event_to_upper_bound = int(current_location[max_zscore_bound_event[0]]) else: move_event_to_upper_bound = N @@ -418,7 +418,7 @@ def _perform_mcmc(self, sustainData, seq_init, f_init, n_iterations, seq_sigma, weight /= np.sum(weight) index = self.global_rng.choice(range(len(possible_positions)), 1, replace=True, p=weight) # FIXME: difficult to check this because random.choice is different to Matlab randsample - move_event_to = possible_positions[index] + move_event_to = int(possible_positions[index[0]]) current_sequence = np.delete(current_sequence, move_event_from, 0) new_sequence = np.concatenate([current_sequence[np.arange(move_event_to)], [selected_event], current_sequence[np.arange(move_event_to, N - 1)]]) diff --git a/pySuStaIn/__init__.py b/pySuStaIn/__init__.py index c44a4a0..6ff7532 100644 --- a/pySuStaIn/__init__.py +++ b/pySuStaIn/__init__.py @@ -7,3 +7,4 @@ from .OrdinalSustain import * from .ZScoreSustainMissingData import * from .MixedTypeSustain import * +from .LongitudinalZscoreSustain import * diff --git a/pySuStaIn/federated/README.md b/pySuStaIn/federated/README.md index 422bf32..42dd01a 100644 --- a/pySuStaIn/federated/README.md +++ b/pySuStaIn/federated/README.md @@ -51,8 +51,8 @@ S, f, loglike = fed.fit_em(S_init, f_init, rng) # multi-start ML fit for a fixed number of subtypes: S, f, loglike = fed.fit(N_S=2, n_startpoints=25, seed=0) -# per-centre ML subtype + stage (assignments computed locally): -assignments = fed.subtype_and_stage(S, f) +# aggregate per-centre subtype/stage counts; row-level assignments stay local: +summaries = fed.subtype_and_stage(S, f) ``` ## Longitudinal data (interdependent visits per subject) diff --git a/pySuStaIn/federated/client.py b/pySuStaIn/federated/client.py index 37b6501..755e0b9 100644 --- a/pySuStaIn/federated/client.py +++ b/pySuStaIn/federated/client.py @@ -15,6 +15,7 @@ """ import os import tempfile +from collections import OrderedDict import numpy as np @@ -24,24 +25,43 @@ class FederatedClient: """Generic centre wrapping a local model + its data object.""" - def __init__(self, local_model, sustain_data, name="centre"): + def __init__(self, local_model, sustain_data, name="centre", cache_size=512): self._model = local_model self._data = sustain_data self.name = name self._N = local_model.stage_zscore.shape[1] # number of events self.M = int(sustain_data.getNumSamples()) # number of subjects + # Memoise the per-sequence stage-likelihood. It is a pure function of the + # event ordering (the model params and data are fixed), so caching is + # numerically exact. Bounded LRU keyed by the integer ordering; keep the + # default conservative because each cached value is subjects x stages. + self._cache = OrderedDict() + self._cache_max = int(cache_size) @property def num_samples(self): return self.M + def _stage(self, seq): + """Memoised ``_calculate_likelihood_stage`` for one subtype ordering.""" + key = np.asarray(seq).astype(int).tobytes() + cached = self._cache.get(key) + if cached is not None: + self._cache.move_to_end(key) + return cached + val = self._model._calculate_likelihood_stage(self._data, np.asarray(seq)) + self._cache[key] = val + if len(self._cache) > self._cache_max: + self._cache.popitem(last=False) + return val + # --- per-subject stage likelihoods for all subtypes (local only) ---------- def _pperm_all(self, S): S = np.asarray(S) N_S = S.shape[0] out = np.zeros((self.M, self._N + 1, N_S)) for s in range(N_S): - out[:, :, s] = self._model._calculate_likelihood_stage(self._data, S[s]) + out[:, :, s] = self._stage(S[s]) return out # --- aggregate statistics returned to the server ------------------------- @@ -74,12 +94,10 @@ def score_candidates(self, S_current, f, s, candidate_seqs): for sp in range(N_S): if sp == s: continue - p_sp = self._model._calculate_likelihood_stage(self._data, S_current[sp]) - wsum_others += f[sp] * np.sum(p_sp, axis=1) + wsum_others += f[sp] * np.sum(self._stage(S_current[sp]), axis=1) scores = np.zeros(len(candidate_seqs)) for idx, seq in enumerate(candidate_seqs): - p_s = self._model._calculate_likelihood_stage(self._data, np.asarray(seq)) - tps = wsum_others + f[s] * np.sum(p_s, axis=1) + tps = wsum_others + f[s] * np.sum(self._stage(seq), axis=1) scores[idx] = np.sum(np.log(tps + 1e-250)) return scores @@ -95,6 +113,19 @@ def subtype_and_stage(self, S, f): ml_stage = np.array([int(np.argmax(w[m, :, ml_subtype[m]])) for m in range(self.M)]) return ml_subtype, ml_stage, prob_cluster + def subtype_stage_summary(self, S, f): + """Aggregate local assignment counts, without returning row-level labels.""" + ml_subtype, ml_stage, _ = self.subtype_and_stage(S, f) + N_S = np.asarray(S).shape[0] + counts = np.zeros((N_S, self._N + 1), dtype=int) + np.add.at(counts, (ml_subtype, ml_stage), 1) + return { + "num_samples": self.M, + "subtype_stage_counts": counts, + "subtype_counts": counts.sum(axis=1), + "stage_counts": counts.sum(axis=0), + } + def _mk_output(folder, prefix): if folder is None: @@ -107,7 +138,7 @@ class ZscoreFederatedClient(FederatedClient): """Cross-sectional Z-score centre (one row per subject).""" def __init__(self, data, Z_vals, Z_max, biomarker_labels, name="centre", - seed=0, output_folder=None): + seed=0, output_folder=None, cache_size=512): data = np.asarray(data, dtype=float) model = ZscoreSustain( data, Z_vals, Z_max, biomarker_labels, @@ -116,14 +147,14 @@ def __init__(self, data, Z_vals, Z_max, biomarker_labels, name="centre", dataset_name=name, use_parallel_startpoints=False, seed=seed, ) sustain_data = ZScoreSustainData(data, model.stage_zscore.shape[1]) - super().__init__(model, sustain_data, name=name) + super().__init__(model, sustain_data, name=name, cache_size=cache_size) class LongitudinalFederatedClient(FederatedClient): """Longitudinal centre (multiple interdependent visits per subject).""" def __init__(self, visit_data, subject_ids, Z_vals, Z_max, biomarker_labels, - name="centre", seed=0, output_folder=None): + name="centre", seed=0, output_folder=None, cache_size=512): # imported here to avoid a hard dependency when only cross-sectional is used from pySuStaIn.LongitudinalZscoreSustain import LongitudinalZscoreSustain model = LongitudinalZscoreSustain( @@ -132,4 +163,4 @@ def __init__(self, visit_data, subject_ids, Z_vals, Z_max, biomarker_labels, output_folder=_mk_output(output_folder, "fed_client_long_"), dataset_name=name, use_parallel_startpoints=False, seed=seed, ) - super().__init__(model, model._AbstractSustain__sustainData, name=name) + super().__init__(model, model._AbstractSustain__sustainData, name=name, cache_size=cache_size) diff --git a/pySuStaIn/federated/experiments/RESULTS.md b/pySuStaIn/federated/experiments/RESULTS.md index 8f92a5c..0e0d187 100644 --- a/pySuStaIn/federated/experiments/RESULTS.md +++ b/pySuStaIn/federated/experiments/RESULTS.md @@ -21,6 +21,36 @@ Reproduce with `run_validation.py` (see README). Kendall tau federated vs pooled : [1.000 1.000] (mean 1.000) ``` +## Longitudinal (interdependent visits per subject) + +Subjects have multiple visits sharing one subtype and a monotonically +non-decreasing stage (`LongitudinalZscoreSustain`). Reproduce with +`run_validation_longitudinal.py`. + +Full-scale equivalence — 600 subjects × 3 visits, 10 biomarkers, **3 subtypes**, +10 centres, 15 starts: + +``` +=== EQUIVALENCE: federated-long EM vs pooled-long EM (same init + RNG) === + pooled loglike = -29448.392220 + fed loglike = -29448.392220 + |loglike diff| = 7.276e-12 |f diff| = 1.665e-16 + sequences identical = True +``` + +Recovery — 150 subjects × 3 visits, 5 biomarkers, 2 subtypes, 5 centres: + +``` + fed loglike = -3872.983 fractions = [0.557 0.443] + Kendall tau federated vs ground truth : [0.924 0.905] (mean 0.914) + Kendall tau federated vs pooled : [1.000 1.000] (mean 1.000) +``` + +Sanity: with **one visit per subject** the longitudinal model is *identical* to +cross-sectional (`_calculate_likelihood_stage` diff `0`, pinned by +`tests/test_regression_golden.py` and `tests/test_longitudinal_zscore.py`). The +opt-in `SUSTAIN_FULL=1` test runs the full 10-biomarker / 3-subtype recovery. + ## Takeaways - **Federation is exact.** From an identical init + RNG, federated EM reproduces @@ -35,10 +65,11 @@ Reproduce with `run_validation.py` (see README). ## Notes / limitations -- This validates the **ML fit** (federated EM). Federated MCMC uncertainty, - federated number-of-subtypes selection (per-centre `f` + federated CVIC) and - longitudinal handling are roadmap items (see README). -- The federated fit is currently **correct but slow** (Python per-centre loops + - recomputation in the sequence search); inner-loop caching is the obvious next - optimisation before large real-data runs. +- This validates the **ML fit** (federated EM), cross-sectional **and + longitudinal**. Federated MCMC uncertainty and federated number-of-subtypes + selection (per-centre `f` + federated CVIC) remain roadmap items (see README). +- The hot paths are now optimised (bounded per-sequence stage-likelihood cache; + vectorised longitudinal monotone-path DP), verified numerically exact by the + golden regression tests. Further speed-ups (e.g. batching the sequence search) + are possible before very large real-data runs. - Harmonisation across centres is assumed upstream and out of scope. diff --git a/pySuStaIn/federated/server.py b/pySuStaIn/federated/server.py index 14fe516..7ffe7bd 100644 --- a/pySuStaIn/federated/server.py +++ b/pySuStaIn/federated/server.py @@ -88,14 +88,14 @@ def _candidate_sequences(self, S_opt, s, selected_event): min_zscore_bound = max(possible_zscores_biomarker[min_filter]) min_zscore_bound_event = events[((self.stage_zscore[0] == min_zscore_bound).astype(int) + (self.stage_biomarker_index[0] == selected_biomarker).astype(int)) == 2] - move_event_to_lower_bound = current_location[min_zscore_bound_event] + 1 + move_event_to_lower_bound = int(current_location[min_zscore_bound_event[0]] + 1) else: move_event_to_lower_bound = 0 if np.any(max_filter): max_zscore_bound = min(possible_zscores_biomarker[max_filter]) max_zscore_bound_event = events[((self.stage_zscore[0] == max_zscore_bound).astype(int) + (self.stage_biomarker_index[0] == selected_biomarker).astype(int)) == 2] - move_event_to_upper_bound = current_location[max_zscore_bound_event] + move_event_to_upper_bound = int(current_location[max_zscore_bound_event[0]]) else: move_event_to_upper_bound = N if move_event_to_lower_bound == move_event_to_upper_bound: @@ -164,9 +164,22 @@ def fit(self, N_S, n_startpoints=25, seed=0): best = (ml_seq, ml_f, ml_like) return best - def subtype_and_stage(self, S, f): - """Per-centre ML subtype+stage (assignments computed locally).""" - out = {} - for c in self.clients: - out[c.name] = c.subtype_and_stage(S, np.asarray(f).reshape(-1)) - return out + def run_sustain_algorithm(self, *args, **kwargs): + raise NotImplementedError( + "FederatedZscoreSustain currently supports ML fitting via fit() and " + "fit_em() only. The inherited run_sustain_algorithm() also runs " + "MCMC uncertainty and per-subject central staging, which are not " + "implemented for the aggregate-only federated path." + ) + + def subtype_and_stage(self, S, f, *, return_individual=False): + """Per-centre subtype/stage summaries by default. + + Individual assignments are row-level outputs. Keep them at the centre in + a real federation; ``return_individual=True`` is only for local + in-process simulations or debugging where row-level return is allowed. + """ + f = np.asarray(f).reshape(-1) + if return_individual: + return {c.name: c.subtype_and_stage(S, f) for c in self.clients} + return {c.name: c.subtype_stage_summary(S, f) for c in self.clients} diff --git a/pySuStaIn/federated/simulate.py b/pySuStaIn/federated/simulate.py index 6e80d6a..514d9be 100644 --- a/pySuStaIn/federated/simulate.py +++ b/pySuStaIn/federated/simulate.py @@ -29,8 +29,8 @@ def simulate_zscore(n_biomarkers=10, n_samples=1000, n_subtypes=2, gt_subtypes = rng.choice(n_subtypes, size=n_samples, p=subtype_fractions) n_controls = int(round(n_samples * frac_controls)) - gt_stages = np.zeros((n_samples, 1), dtype=int) - gt_stages[n_controls:, 0] = rng.integers(1, N_stages + 1, size=n_samples - n_controls) + gt_stages = np.zeros(n_samples, dtype=int) + gt_stages[n_controls:] = rng.integers(1, N_stages + 1, size=n_samples - n_controls) data, data_denoised, stage_value = ZscoreSustain.generate_data( gt_subtypes, gt_stages, gt_sequences, Z_vals, Z_max @@ -92,7 +92,7 @@ def simulate_longitudinal(n_biomarkers=10, n_subjects=600, n_subtypes=3, n_visit # expand to visit level subtypes_visit = np.repeat(gt_subtypes, n_visits) subject_ids = np.repeat(np.arange(n_subjects), n_visits) - stages_visit = subj_stages.reshape(-1, 1) + stages_visit = subj_stages.reshape(-1) visit_data, _, _ = ZscoreSustain.generate_data( subtypes_visit, stages_visit, gt_sequences, Z_vals, Z_max diff --git a/setup.py b/setup.py index 08caac2..d62ed10 100644 --- a/setup.py +++ b/setup.py @@ -2,7 +2,7 @@ # License: TBC __version__ = '0.1' -from setuptools import setup +from setuptools import find_packages, setup #parse the requirement.txt file, ignoring commented lines, placing results in install_reqs with open('requirements.txt', 'r') as f: @@ -26,10 +26,10 @@ maintainer= 'Leon Aksman', maintainer_email= 'l.aksman@ucl.ac.uk', license= 'TBC', - packages= ['pySuStaIn', 'sim'], + packages= find_packages(), python_requires= '>=3.7', install_requires = install_reqs, #the parsed requirements from requirements.txt entry_points= {}, zip_safe= False) -print("Finished pySuStaIn setup.py") \ No newline at end of file +print("Finished pySuStaIn setup.py") diff --git a/tests/test_federated_zscore.py b/tests/test_federated_zscore.py index 237047a..4646e28 100644 --- a/tests/test_federated_zscore.py +++ b/tests/test_federated_zscore.py @@ -7,6 +7,7 @@ import tempfile import numpy as np +from setuptools import find_packages from pySuStaIn.ZscoreSustain import ZscoreSustain from pySuStaIn.federated.client import ZscoreFederatedClient @@ -57,3 +58,35 @@ def test_single_subtype_equivalence(): fs, ff, fl = fed.fit_em(S0.copy(), f0.copy(), np.random.default_rng(5)) assert abs(pl - fl) < 1e-6 assert np.array_equal(ps.astype(int), fs.astype(int)) + + +def test_package_discovery_includes_federated_subpackages(): + packages = set(find_packages()) + assert "pySuStaIn.federated" in packages + assert "pySuStaIn.federated.experiments" in packages + + +def test_server_assignment_default_returns_aggregate_counts(): + sim = simulate_zscore(n_biomarkers=4, n_samples=80, n_subtypes=2, seed=13) + data, Zv, Zm, labels = sim["data"], sim["Z_vals"], sim["Z_max"], sim["labels"] + shards = split_into_centres(data.shape[0], n_centres=4, seed=13) + pooled = _pooled(data, Zv, Zm, labels) + clients = [ZscoreFederatedClient(data[ix], Zv, Zm, labels, name=f"c{i}") + for i, ix in enumerate(shards)] + fed = FederatedZscoreSustain(clients, Zv, Zm, labels, N_S_max=2) + + sd = pooled._AbstractSustain__sustainData + S0 = np.array([pooled._initialise_sequence(sd, np.random.default_rng(13))[0] + for _ in range(2)]) + f0 = np.ones(2) / 2 + + summaries = fed.subtype_and_stage(S0, f0) + for client, summary in zip(clients, summaries.values()): + counts = summary["subtype_stage_counts"] + assert counts.shape == (2, fed._N + 1) + assert counts.sum() == client.num_samples + assert "subtype_counts" in summary + assert "stage_counts" in summary + + individual = fed.subtype_and_stage(S0, f0, return_individual=True) + assert len(individual["c0"][0]) == clients[0].num_samples diff --git a/tests/test_regression_golden.py b/tests/test_regression_golden.py new file mode 100644 index 0000000..38f3693 --- /dev/null +++ b/tests/test_regression_golden.py @@ -0,0 +1,47 @@ +"""Golden regression: pin the exact current fit outputs so optimisations cannot +silently change the numerics. + +These deterministic fits (fixed init + RNG) must keep producing the same +sequences and log-likelihood. They guard the longitudinal monotone-path DP and +the cross-sectional EM independently of the fed==pooled checks (a DP change that +affects pooled and federated equally would pass fed==pooled but fail here). + +If you intentionally change the model maths, re-snapshot these values. +""" +import tempfile + +import numpy as np + +from pySuStaIn.ZscoreSustain import ZscoreSustain +from pySuStaIn.LongitudinalZscoreSustain import LongitudinalZscoreSustain +from pySuStaIn.federated.simulate import simulate_zscore, simulate_longitudinal + +GOLDEN_XS_LOGLIKE = -1723.134839679955 +GOLDEN_XS_SEQ = [[4, 2, 0, 7, 3, 5, 10, 8, 1, 12, 13, 6, 11, 9, 14], + [2, 4, 1, 6, 11, 9, 0, 3, 8, 7, 5, 10, 12, 14, 13]] + +GOLDEN_LONG_LOGLIKE = -3182.007273642864 +GOLDEN_LONG_SEQ = [[0, 2, 3, 8, 1, 6, 13, 11, 7, 4, 5, 9, 10, 12, 14], + [4, 2, 9, 7, 12, 0, 1, 14, 3, 6, 11, 5, 10, 8, 13]] + + +def test_cross_sectional_golden(): + s = simulate_zscore(n_biomarkers=5, n_samples=200, n_subtypes=2, seed=3) + xs = ZscoreSustain(s["data"], s["Z_vals"], s["Z_max"], s["labels"], 1, 2, 1, + tempfile.mkdtemp(), "x", False, 0) + sd = xs._AbstractSustain__sustainData + S0 = np.array([xs._initialise_sequence(sd, np.random.default_rng(3))[0] for _ in range(2)]) + seq, f, ll, *_ = xs._perform_em(sd, S0.copy(), np.ones(2) / 2, np.random.default_rng(7)) + assert abs(float(ll) - GOLDEN_XS_LOGLIKE) < 1e-6 + assert seq.astype(int).tolist() == GOLDEN_XS_SEQ + + +def test_longitudinal_golden(): + L = simulate_longitudinal(n_biomarkers=5, n_subjects=120, n_subtypes=2, n_visits=3, seed=4) + lg = LongitudinalZscoreSustain(L["visit_data"], L["subject_ids"], L["Z_vals"], L["Z_max"], + L["labels"], 1, 2, 1, tempfile.mkdtemp(), "l", False, 0) + lsd = lg._AbstractSustain__sustainData + S0 = np.array([lg._initialise_sequence(lsd, np.random.default_rng(4))[0] for _ in range(2)]) + seq, f, ll, *_ = lg._perform_em(lsd, S0.copy(), np.ones(2) / 2, np.random.default_rng(9)) + assert abs(float(ll) - GOLDEN_LONG_LOGLIKE) < 1e-6 + assert seq.astype(int).tolist() == GOLDEN_LONG_SEQ