Rename inference I/O variables for clarity.

thebuckleylab · Nov 21, 2024 · e1b4d2f · e1b4d2f
1 parent b47ee87
commit e1b4d2f
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 12 deletions.
diff --git a/jpc/_core/_infer.py b/jpc/_core/_infer.py
@@ -20,8 +20,8 @@
 def solve_pc_inference(
         params: Tuple[PyTree[Callable], Optional[PyTree[Callable]]],
         activities: PyTree[ArrayLike],
-        y: ArrayLike,
-        x: Optional[ArrayLike] = None,
+        output: ArrayLike,
+        input: Optional[ArrayLike] = None,
         loss_id: str = "MSE",
         solver: AbstractSolver = Heun(),
         max_t1: int = 20,
@@ -49,8 +49,8 @@ def solve_pc_inference(
 
     - `params`: Tuple with callable model layers and optional skip connections.
     - `activities`: List of activities for each layer free to vary.
-    - `y`: Observation or target of the generative model.
-    - `x`: Optional prior of the generative model.
+    - `output`: Observation or target of the generative model.
+    - `input`: Optional prior of the generative model.
 
     **Other arguments:**
 
@@ -87,7 +87,7 @@ def solve_pc_inference(
         t1=max_t1,
         dt0=dt,
         y0=activities,
-        args=(params, y, x, loss_id, stepsize_controller),
+        args=(params, output, input, loss_id, stepsize_controller),
         stepsize_controller=stepsize_controller,
         event=Event(steady_state_event_with_timeout),
         saveat=saveat

diff --git a/jpc/_test.py b/jpc/_test.py
@@ -124,7 +124,7 @@ def test_generative_pc(
     input_preds = solve_pc_inference(
         params=params,
         activities=activities,
-        y=output,
+        output=output,
         solver=ode_solver,
         max_t1=max_t1,
         dt=dt,
@@ -207,7 +207,7 @@ def test_hpc(
     hpc_preds = solve_pc_inference(
         params=gen_params,
         activities=amort_activities,
-        y=output,
+        output=output,
         solver=ode_solver,
         max_t1=max_t1,
         dt=dt,
@@ -223,7 +223,7 @@ def test_hpc(
     gen_preds = solve_pc_inference(
         params=gen_params,
         activities=activities,
-        y=output,
+        output=output,
         solver=ode_solver,
         max_t1=max_t1,
         dt=dt,

diff --git a/jpc/_train.py b/jpc/_train.py
@@ -147,8 +147,8 @@ def make_pc_step(
     equilib_activities = solve_pc_inference(
         params=(model, skip_model),
         activities=activities,
-        y=output,
-        x=input,
+        output=output,
+        input=input,
         loss_id=loss_id,
         solver=ode_solver,
         max_t1=max_t1,
@@ -329,8 +329,8 @@ def make_hpc_step(
     equilib_activities = solve_pc_inference(
         params=gen_params,
         activities=amort_activities[1:] if input is not None else amort_activities,
-        y=output,
-        x=input,
+        output=output,
+        input=input,
         solver=ode_solver,
         max_t1=max_t1,
         dt=dt,