Default NaN-Checker for Particle Tracking models (#3376)

* no NaNs * default * comment * punctuation * better comment * comment * comment * comment * grammar * update the OnlyParticleTrackingModel * Update src/Models/Models.jl Co-authored-by: Gregory L. Wagner <[email protected]> * Update src/Models/Models.jl Co-authored-by: Gregory L. Wagner <[email protected]> --------- Co-authored-by: Gregory L. Wagner <[email protected]>
CliMA · Nov 7, 2023 · 615b2a8 · 615b2a8
1 parent 92c2cb9
commit 615b2a8
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 2 deletions.
diff --git a/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl b/src/Models/HydrostaticFreeSurfaceModels/prescribed_hydrostatic_velocity_fields.jl
@@ -117,9 +117,12 @@ Adapt.adapt_structure(to, velocities::PrescribedVelocityFields) =
 const OnlyParticleTrackingModel = HydrostaticFreeSurfaceModel{TS, E, A, S, G, T, V, B, R, F, P, U, C} where
                  {TS, E, A, S, G, T, V, B, R, F, P<:AbstractLagrangianParticles, U<:PrescribedVelocityFields, C<:NamedTuple{(), Tuple{}}}                 
 
-function time_step!(model::OnlyParticleTrackingModel, Δt; euler=false) 
+function time_step!(model::OnlyParticleTrackingModel, Δt; callbacks = [], kwargs...) 
     model.timestepper.previous_Δt = Δt
     tick!(model.clock, Δt)
     step_lagrangian_particles!(model, Δt)
+    update_state!(model, callbacks)
 end
 
+update_state!(model::OnlyParticleTrackingModel, callbacks) = 
+    [callback(model) for callback in callbacks if callback.callsite isa UpdateStateCallsite]
diff --git a/src/Models/Models.jl b/src/Models/Models.jl
@@ -162,13 +162,25 @@ end
 # Check for NaNs in the first prognostic field (generalizes to prescribed velocities).
 function default_nan_checker(model::OceananigansModels)
     model_fields = prognostic_fields(model)
+
+    if isempty(model_fields) 
+        return nothing
+    end
+
     first_name = first(keys(model_fields))
     field_to_check_nans = NamedTuple{tuple(first_name)}(model_fields)
     nan_checker = NaNChecker(field_to_check_nans)
     return nan_checker
 end
 
-# This is here so that `NonhydrostaticModel` and  `HydrsostaticFreeSurfaceModel`
+using Oceananigans.Models.HydrostaticFreeSurfaceModels: OnlyParticleTrackingModel
+
+# Particle tracking models with prescribed velocities (and no tracers) 
+# have no prognostic fields and no chance to producing a NaN.
+default_nan_checker(::OnlyParticleTrackingModel) = nothing
+
+# Implementation of a `seawater_density` `KernelFunctionOperation
+# applicable to both `NonhydrostaticModel` and  `HydrostaticFreeSurfaceModel`
 include("seawater_density.jl")
 
 end # module