From de17b878e35ca7fb05c5af2a262f1c9ef799c51e Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Mon, 6 Jan 2025 12:31:29 +0530
Subject: [PATCH] feat: validate causality of analysis points

---
 src/systems/analysis_points.jl | 54 ++++++++++++++++++++++++++++------
 test/analysis_points.jl        | 10 +++++++
 2 files changed, 55 insertions(+), 9 deletions(-)

diff --git a/src/systems/analysis_points.jl b/src/systems/analysis_points.jl
index 9ff06109e5..a2b7fb7ea8 100644
--- a/src/systems/analysis_points.jl
+++ b/src/systems/analysis_points.jl
@@ -72,6 +72,36 @@ struct AnalysisPoint
     these are all `RealInput` connectors.
     """
     outputs::Union{Nothing, Vector{Any}}
+
+    function AnalysisPoint(input, name::Symbol, outputs; verbose = true)
+        # input to analysis point should be an output variable
+        if verbose && input !== nothing
+            var = ap_var(input)
+            isoutput(var) || ap_warning(1, name, true)
+        end
+        # outputs of analysis points should be input variables
+        if verbose && outputs !== nothing
+            for (i, output) in enumerate(outputs)
+                var = ap_var(output)
+                isinput(var) || ap_warning(2 + i, name, false)
+            end
+        end
+
+        return new(input, name, outputs)
+    end
+end
+
+function ap_warning(arg::Int, name::Symbol, should_be_output)
+    causality = should_be_output ? "output" : "input"
+    @warn """
+    The $(arg)-th argument to analysis point $(name) was not a $causality. This is supported in \
+    order to handle inverse models, but may not be what you intended.
+
+    If you are building a forward mode (causal), you may want to swap this argument with \
+    one on the opposite side of the name of the analysis point provided to `connect`. \
+    Learn more about the causality of analysis points in the docstring for `AnalysisPoint`. \
+    Silence this message using `connect(out, :name, in...; warn = false)`.
+    """
 end
 
 AnalysisPoint() = AnalysisPoint(nothing, Symbol(), nothing)
@@ -133,8 +163,8 @@ function renamespace(sys, ap::AnalysisPoint)
 end
 
 # create analysis points via `connect`
-function Symbolics.connect(in, ap::AnalysisPoint, outs...)
-    return AnalysisPoint() ~ AnalysisPoint(in, ap.name, collect(outs))
+function Symbolics.connect(in, ap::AnalysisPoint, outs...; verbose = true)
+    return AnalysisPoint() ~ AnalysisPoint(in, ap.name, collect(outs); verbose)
 end
 
 """
@@ -161,15 +191,21 @@ connect(P.input, :plant_input, C.output)
 
 typically is not (unless the model is an inverse model).
 
-# Arguments:
+# Arguments
+
+- `output_connector`: An output connector
+- `input_connector`: An input connector
+- `ap`: An explicitly created [`AnalysisPoint`](@ref)
+- `ap_name`: If a name is given, an [`AnalysisPoint`](@ref) with the given name will be
+  created automatically.
+
+# Keyword arguments
 
-  - `output_connector`: An output connector
-  - `input_connector`: An input connector
-  - `ap`: An explicitly created [`AnalysisPoint`](@ref)
-  - `ap_name`: If a name is given, an [`AnalysisPoint`](@ref) with the given name will be created automatically.
+- `verbose`: Warn if an input is connected to an output (reverse causality). Silence this
+  warning if you are analyzing an inverse model.
 """
-function Symbolics.connect(in::AbstractSystem, name::Symbol, out, outs...)
-    return AnalysisPoint() ~ AnalysisPoint(in, name, [out; collect(outs)])
+function Symbolics.connect(in::AbstractSystem, name::Symbol, out, outs...; verbose = true)
+    return AnalysisPoint() ~ AnalysisPoint(in, name, [out; collect(outs)]; verbose)
 end
 
 """
diff --git a/test/analysis_points.jl b/test/analysis_points.jl
index 9361b5a686..8c93e919b9 100644
--- a/test/analysis_points.jl
+++ b/test/analysis_points.jl
@@ -24,6 +24,16 @@ using Symbolics: NAMESPACE_SEPARATOR
     @test isequal(sys_ap2, sys_normal2)
 end
 
+@testset "Inverse causality throws a warning" begin
+    @named P = FirstOrder(k = 1, T = 1)
+    @named C = Gain(; k = -1)
+
+    ap = AnalysisPoint(:plant_input)
+    @test_warn ["1-th argument", "plant_input", "not a output"] connect(
+        P.input, ap, C.output)
+    @test_nowarn connect(P.input, ap, C.output; verbose = false)
+end
+
 # also tests `connect(input, name::Symbol, outputs...)` syntax
 @testset "AnalysisPoint is accessible via `getproperty`" begin
     @named P = FirstOrder(k = 1, T = 1)