[codex] Harden federated longitudinal SuStaIn

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

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)]])

pySuStaIn/__init__.py

@@ -7,3 +7,4 @@
 from .OrdinalSustain import *
 from .ZScoreSustainMissingData import *
 from .MixedTypeSustain import *
+from .LongitudinalZscoreSustain import *

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)

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)

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.

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}

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

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")
+print("Finished pySuStaIn setup.py")