Merge branch 'main' into jsw/fix-hard-coded-number-types

OceanBioME · Dec 2, 2024 · 0e12486 · 0e12486
2 parents dbce30e + dae84ae
commit 0e12486
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 6 deletions.
diff --git a/src/Particles/Particles.jl b/src/Particles/Particles.jl
@@ -135,8 +135,11 @@ const MODEL_WITH_BGC_PARTICLES =
 @inline function update_lagrangian_particle_properties!(particles::BiogeochemicalParticles, model, bgc, Δt)
     advect_particles!(particles.advection, particles, model, Δt)
     time_step_particle_fields!(particles.timestepper, particles, model, Δt)
+    update_particle_state!(particles, model, Δt)
 end
 
+@inline update_particle_state!(particles, model, Δt) = nothing
+
 size(particles::BiogeochemicalParticles) = (length(particles), )
 length(::BiogeochemicalParticles{N}) where N = N
 

diff --git a/src/Particles/tendencies.jl b/src/Particles/tendencies.jl
@@ -29,7 +29,7 @@ end
 
     field_values = ntuple(nf->particles.fields[nf][n], Val(Nf))
 
-    tracer_values = extract_tracer_values(particles.field_interpolation, particles, grid, fields, n)
+    tracer_values = extract_tracer_values(val_field_name, particles.field_interpolation, particles, grid, fields, n)
 
     tendencies[field_name][n] = particles.biogeochemistry(val_field_name, t, field_values..., tracer_values...)
 end
diff --git a/src/Particles/tracer_interpolation.jl b/src/Particles/tracer_interpolation.jl
@@ -2,6 +2,7 @@ using Oceananigans.Operators: volume
 using Oceananigans.Fields: fractional_indices, _interpolate, interpolator, Center
 using Oceananigans.Grids: AbstractGrid, Flat, Bounded, Periodic
 
+@inline get_node(::Flat, i, N) = one(i)
 @inline get_node(::Bounded,  i, N) = min(max(i, 1), N)
 @inline get_node(::Periodic, i, N) = ifelse(i < 1, N, ifelse(i > N, 1, i))
 
@@ -12,7 +13,7 @@ Specifies that tracer values should be taken from the nearst center point.
 """
 struct NearestPoint end
 
-@inline function extract_tracer_values(::NearestPoint, particles, grid, fields, n)
+@inline function extract_tracer_values(val_field_name, ::NearestPoint, particles, grid, fields, n)
     x = @inbounds particles.x[n]
     y = @inbounds particles.y[n]
     z = @inbounds particles.z[n]
@@ -29,11 +30,21 @@ struct NearestPoint end
     return field_values
 end
 
+# feels like fractional_indices could work more intuativly
+@inline collapse_position(x, y, z, ℓx, ℓy, ℓz) = (x, y, z)
+@inline collapse_position(x, y, z, ::Nothing, ℓy, ℓz) = (y, z)
+@inline collapse_position(x, y, z, ℓx, ::Nothing, ℓz) = (x, z)
+@inline collapse_position(x, y, z, ℓx, ℓy, ::Nothing) = (x, y)
+@inline collapse_position(x, y, z, ℓx, ::Nothing, ::Nothing) = (x, )
+@inline collapse_position(x, y, z, ::Nothing, ℓy, ::Nothing) = (y, )
+@inline collapse_position(x, y, z, ::Nothing, ::Nothing, ℓz) = (z, )
+
 @inline function nearest_node(x, y, z, grid::AbstractGrid{FT, TX, TY, TZ}) where {FT, TX, TY, TZ}
     # messy
-    ii, jj, kk = fractional_indices((x, y, z), grid, ifelse(isa(TX(), Flat), nothing, Center()),
-                                                     ifelse(isa(TY(), Flat), nothing, Center()),
-                                                     ifelse(isa(TZ(), Flat), nothing, Center()))
+    ℓx = ifelse(isa(TX(), Flat), nothing, Center())
+    ℓy = ifelse(isa(TY(), Flat), nothing, Center())
+    ℓz = ifelse(isa(TZ(), Flat), nothing, Center())
+    ii, jj, kk = fractional_indices(collapse_position(x, y, z, ℓx, ℓy, ℓz), grid, ℓx, ℓy, ℓz)
 
     ix = interpolator(ii)
     iy = interpolator(jj)

diff --git a/src/Particles/update_tracer_tendencies.jl b/src/Particles/update_tracer_tendencies.jl
@@ -36,7 +36,7 @@ possible_tuple_value(::Tuple, field_name) = field_name
 
     field_values = ntuple(nf->particles.fields[nf][n], Val(Nf))
 
-    tracer_values = extract_tracer_values(particles.field_interpolation, particles, grid, fields, n)
+    tracer_values = extract_tracer_values(val_field_name, particles.field_interpolation, particles, grid, fields, n)
 
     particle_tendency = particles.biogeochemistry(val_field_name, t, field_values..., tracer_values...)