From a6037a9d5b07f5dd3634ae816d0c8f8334696474 Mon Sep 17 00:00:00 2001
From: Christopher Tessum <ctessum@gmail.com>
Date: Sun, 27 Oct 2024 16:45:42 -0500
Subject: [PATCH] Fix advection coupling issue

---
 src/advection.jl       |  3 +--
 src/coupled_system.jl  | 12 ++++++------
 test/advection_test.jl |  6 +++---
 3 files changed, 10 insertions(+), 11 deletions(-)

diff --git a/src/advection.jl b/src/advection.jl
index ab8c993e..6bdf8259 100644
--- a/src/advection.jl
+++ b/src/advection.jl
@@ -66,7 +66,7 @@ function advection(vars, di::DomainInfo)
     eqs
 end
 
-function couple(c::CoupledSystem, _::Advection)::CoupledSystem
+function couple2(c::CoupledSystem, _::Advection)::CoupledSystem
     dom = domain(c)
 
     # Add in a model component to allow the specification of the wind velocity.
@@ -79,7 +79,6 @@ function couple(c::CoupledSystem, _::Advection)::CoupledSystem
     push!(c.pdefunctions, f)
     c
 end
-couple(a::Advection, c::CoupledSystem)::CoupledSystem = couple(c, a)
 
 struct ConstantWindCoupler
     sys
diff --git a/src/coupled_system.jl b/src/coupled_system.jl
index 5747354f..04be79f4 100644
--- a/src/coupled_system.jl
+++ b/src/coupled_system.jl
@@ -15,7 +15,7 @@ Things that can be added to a `CoupledSystem`:
     * [Callbacks](https://docs.sciml.ai/DiffEqDocs/stable/features/callback_functions/)
     * Types `X` that implement a `EarthSciMLBase.init_callback(::X, sys::CoupledSystem, sys_mtk, obs_eqs, domain::DomainInfo)::DECallback` method
     * Other `CoupledSystem`s
-    * Types `X` that implement a `EarthSciMLBase.couple(::X, ::CoupledSystem)` or `EarthSciMLBase.couple(::CoupledSystem, ::X)` method.
+    * Types `X` that implement a `EarthSciMLBase.couple2(::X, ::CoupledSystem)` or `EarthSciMLBase.couple2(::CoupledSystem, ::X)` method.
     * `Tuple`s or `AbstractVector`s of any of the things above.
 """
 mutable struct CoupledSystem
@@ -83,11 +83,11 @@ function couple(systems...)::CoupledSystem
             push!(o.callbacks, sys)
         elseif (sys isa Tuple) || (sys isa AbstractVector)
             o = couple(o, sys...)
-        elseif hasmethod(init_callback, (typeof(sys), Any, Any, Any, Any))
+        elseif hasmethod(init_callback, (typeof(sys), CoupledSystem, Any, Any, DomainInfo))
             push!(o.init_callbacks, sys)
-        elseif applicable(couple, o, sys)
-            o = couple(o, sys)
-        elseif applicable(couple, sys, o)
+        elseif hasmethod(couple2, (CoupledSystem, typeof(sys)))
+            o = couple2(o, sys)
+        elseif hasmethod(couple2, (typeof(sys), CoupledSystem))
             o = couple(sys, o)
         else
             error("Cannot couple a $(typeof(sys)).")
@@ -219,7 +219,7 @@ end
 """
 Types that implement an:
 
-`init_callback(x, sys, obs_eqs, domain, u0, p)::DECallback`
+`init_callback(x, sys::CoupledSystem, obs_eqs, domain::DomainInfo)::DECallback`
 
 method can also be coupled into a `CoupledSystem`.
 The `init_callback` function will be run before the simulator is run
diff --git a/test/advection_test.jl b/test/advection_test.jl
index c76b9a98..5b6f6810 100644
--- a/test/advection_test.jl
+++ b/test/advection_test.jl
@@ -1,4 +1,4 @@
-using Main.EarthSciMLBase
+using EarthSciMLBase
 using Test
 using DomainSets, MethodOfLines, ModelingToolkit, DifferentialEquations
 using ModelingToolkit: t, D
@@ -46,7 +46,7 @@ using Dates, DomainSets
     eq = equations(combined_mtk)
     eqstr = string(eq)
 
-    @test occursin("- MeanWind₊v_x(t, x)*Differential(x)(examplesys₊y(t, x))", eqstr) || 
+    @test occursin("- MeanWind₊v_x(t, x)*Differential(x)(examplesys₊y(t, x))", eqstr) ||
         occursin("- Differential(x)(examplesys₊y(t, x))*MeanWind₊v_x(t, x)", eqstr)
 
     @test_broken begin # Test fails because PDEs don't currently work with units.
@@ -94,4 +94,4 @@ end
     for term ∈ want_terms
         @test any(occursin.((term,), have_eqs))
     end
-end
\ No newline at end of file
+end