From f46198d41d9081dae9e4e3c1c9bd8d6165d65faa Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sebastian=20Miclu=C8=9Ba-C=C3=A2mpeanu?=
 <sebastian.mc95@proton.me>
Date: Fri, 2 Aug 2024 03:07:55 +0300
Subject: [PATCH] feat: add `ParametrizedInterpolation`

Co-authored-by: Fredrik Bagge Carlson <baggepinnen@gmail.com>
---
 Project.toml                                  |  8 +++++-
 ...kitStandardLibraryDataInterpolationsExt.jl | 10 ++++++++
 src/Blocks/Blocks.jl                          |  2 +-
 src/Blocks/sources.jl                         | 25 +++++++++++++++++++
 4 files changed, 43 insertions(+), 2 deletions(-)
 create mode 100644 ext/ModelingToolkitStandardLibraryDataInterpolationsExt.jl

diff --git a/Project.toml b/Project.toml
index e28a6868c..181d519da 100644
--- a/Project.toml
+++ b/Project.toml
@@ -11,11 +11,17 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 
+[weakdeps]
+DataInterpolations = "82cc6244-b520-54b8-b5a6-8a565e85f1d0"
+
+[extensions]
+ModelingToolkitStandardLibraryDataInterpolationsExt = "DataInterpolations"
+
 [compat]
 Aqua = "0.8"
 ChainRulesCore = "1.18"
 ControlSystemsBase = "1.4"
-DataInterpolations = "4.6"
+DataInterpolations = "4.6, 5, 6"
 DiffEqBase = "6.147"
 IfElse = "0.1"
 LinearAlgebra = "1.10"
diff --git a/ext/ModelingToolkitStandardLibraryDataInterpolationsExt.jl b/ext/ModelingToolkitStandardLibraryDataInterpolationsExt.jl
new file mode 100644
index 000000000..ff6431cd7
--- /dev/null
+++ b/ext/ModelingToolkitStandardLibraryDataInterpolationsExt.jl
@@ -0,0 +1,10 @@
+module ModelingToolkitStandardLibraryDataInterpolationsExt
+
+using ModelingToolkitStandardLibrary
+using ModelingToolkitStandardLibrary.Blocks.Symbolics
+using DataInterpolations
+
+@register_symbolic ModelingToolkitStandardLibrary.Blocks.apply_interpolation(
+    t, i::DataInterpolations.AbstractInterpolation)
+
+end
diff --git a/src/Blocks/Blocks.jl b/src/Blocks/Blocks.jl
index 3eceab637..0bf733c92 100644
--- a/src/Blocks/Blocks.jl
+++ b/src/Blocks/Blocks.jl
@@ -16,7 +16,7 @@ export Log, Log10
 include("math.jl")
 
 export Constant, TimeVaryingFunction, Sine, Cosine, ContinuousClock, Ramp, Step, ExpSine,
-       Square, Triangular, Parameter, SampledData
+       Square, Triangular, Parameter, SampledData, ParametrizedInterpolation
 include("sources.jl")
 
 export Limiter, DeadZone, SlewRateLimiter
diff --git a/src/Blocks/sources.jl b/src/Blocks/sources.jl
index 82cde2bb1..877ad01f5 100644
--- a/src/Blocks/sources.jl
+++ b/src/Blocks/sources.jl
@@ -730,3 +730,28 @@ end
 function SampledData(; name, buffer, sample_time, circular_buffer)
     SampledData(SampledDataType.vector_based; name, buffer, sample_time, circular_buffer)
 end
+
+# This needs to be extend for interpolation types
+apply_interpolation(t, interp) = interp(t)
+
+@register_symbolic build_interpolation(
+    interpolation_type::UnionAll, u::AbstractArray, x::AbstractArray)
+build_interpolation(interpolation_type, u, x) = interpolation_type(u, x)
+
+function ParametrizedInterpolation(interp_type::T, u, x, t = t; name) where {T}
+    @parameters data[1:length(x)] = u
+    @parameters ts[1:length(x)] = x
+    @parameters interpolation_type::T=interp_type [tunable = false]
+    @parameters interpolator::interp_type
+
+    @named output = RealOutput()
+
+    eqs = [output.u ~ apply_interpolation(t, interpolator)]
+
+    ODESystem(eqs, t, [], [u, x, interpolation_type, interpolator];
+        parameter_dependencies = [
+            interpolator => build_interpolation(interpolation_type, u, x)
+        ],
+        systems = [output],
+        name)
+end