Refactor observables initialization in BayesianModel

JR-1991 · JR-1991 · commit 77e9101dbc15 · 2025-09-30T00:36:02.000+02:00
Simplifies the construction of the observables array by using model.get_state_order with as_indices=True and modeled=True, replacing the previous list comprehension.
diff --git a/catalax/mcmc/mcmc.py b/catalax/mcmc/mcmc.py
@@ -329,11 +329,7 @@ def __init__(
         ]
 
         self.observables = jnp.array(
-            [
-                i
-                for i, state in enumerate(model.get_state_order())
-                if model.odes[state].observable
-            ]
+            model.get_state_order(as_indices=True, modeled=True)
         )
 
     def __call__(

Original file line number	Diff line number	Diff line change
`@@ -329,11 +329,7 @@ def __init__(`
`329`	`329`	`]`
`330`	`330`
`331`	`331`	`self.observables = jnp.array(`
`332`		`- [`
`333`		`- i`
`334`		`- for i, state in enumerate(model.get_state_order())`
`335`		`- if model.odes[state].observable`
`336`		`- ]`
	`332`	`+ model.get_state_order(as_indices=True, modeled=True)`
`337`	`333`	`)`
`338`	`334`
`339`	`335`	`def __call__(`