Skip to content

Commit

Permalink
maybe fixed scaling
Browse files Browse the repository at this point in the history
  • Loading branch information
jagoosw committed Sep 14, 2023
1 parent 67f9365 commit ff0e343
Showing 1 changed file with 24 additions and 15 deletions.
39 changes: 24 additions & 15 deletions src/Utils/negative_tracers.jl
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ struct ScaleNegativeTracers{FA, SA, W}
warn :: W

ScaleNegativeTracers(tracers::FA, scalefactors::SA, warn::W) where {FA, SA, W} =
new{FA, SA, W}(tracers, scalefactors, warn)
warn ? error("Warning not currently implemented") : new{FA, SA, W}(tracers, scalefactors, warn)
end

adapt_structure(to, snt::ScaleNegativeTracers) = ScaleNegativeTracers(adapt(to, snt.tracers),
Expand All @@ -62,7 +62,7 @@ Future plans include implement a positivity-preserving timestepping scheme as th
If `warn` is true then scaling will raise a warning.
"""
function ScaleNegativeTracers(tracers; scalefactors = NamedTuple{tracers}(ones(length(tracers))), warn = false)
function ScaleNegativeTracers(tracers; scalefactors = ones(length(tracers)), warn = false)

Check warning on line 65 in src/Utils/negative_tracers.jl

View check run for this annotation

Codecov / codecov/patch

src/Utils/negative_tracers.jl#L65

Added line #L65 was not covered by tests
if length(scalefactors) != length(tracers)
error("Incorrect number of scale factors provided")
end
Expand All @@ -79,40 +79,49 @@ If `warn` is true then scaling will raise a warning.
"""
function ScaleNegativeTracers(model::AbstractBiogeochemistry; warn = false)
tracers = conserved_tracers(model)
scalefactors = NamedTuple{tracers}(ones(length(tracers)))
scalefactors = ones(length(tracers))

return ScaleNegativeTracers(tracers, scalefactors, warn)
end

function update_biogeochemical_state!(model, scale::ScaleNegativeTracers)
workgroup, worksize = work_layout(model.grid, :xyz)

dev = device(model.grid.architecture)

scale_for_negs_kernel! = scale_for_negs!(dev, workgroup, worksize)
scale_for_negs_kernel!(model.tracers, scale.tracers, scale.scalefactors)

tracers_to_scale = Tuple(model.tracers[tracer_name] for tracer_name in keys(scale.tracers))

scale_for_negs_kernel!(tracers_to_scale, scale.scalefactors)
end

@kernel function scale_for_negs!(fields, tracers, scalefactors)
@kernel function scale_for_negs!(tracers, scalefactors)

Check warning on line 99 in src/Utils/negative_tracers.jl

View check run for this annotation

Codecov / codecov/patch

src/Utils/negative_tracers.jl#L99

Added line #L99 was not covered by tests
i, j, k = @index(Global, NTuple)

t, p = 0.0, 0.0

@unroll for (idx, tracer) in enumerate(tracers)
field = @inbounds fields[tracer][i, j, k]
scalefactor = @inbounds scalefactors[tracer]
value = @inbounds tracer[i, j, k]
scalefactor = @inbounds scalefactors[idx]

t += field * scalefactor
if field > 0
p += field * scalefactor
t += value * scalefactor
if value > 0
p += value * scalefactor
end
end

t < 0 && error("Cell total < 0, can not scale negative tracers.")

@unroll for tracer in tracers
field = @inbounds fields[tracer][i, j, k]
value = @inbounds tracer[i, j, k]

if field > 0
field *= t / p
if value > 0
value *= t / p
else
field = 0
value = 0
end

@inbounds fields[tracer][i, j, k] = field
@inbounds tracer[i, j, k] = value
end
end

0 comments on commit ff0e343

Please sign in to comment.