Remove use of matching subfields

[charleskawczynski]

No description provided.

[charleskawczynski]

Hi @dennisYatunin, any idea why these unrolled_foreach calls don't seem to be constant propagating the symbols? This PR is passing CI but failing allocation / JET tests.

[tapios]

Please encapsulate the logic for looping over property names in some ways that we do not need to repeat it several times.

[dennisYatunin]

Hi @dennisYatunin, any idea why these unrolled_foreach calls don't seem to be constant propagating the symbols? This PR is passing CI but failing allocation / JET tests.

Yeah it's unfortunate that the symbols aren't being propagated. Maybe we're too deep in the compilation stacktrace for constant propagation to be triggered? I think the best workaround here is to use singleton types instead of symbols, since unrolling over singleton types is guaranteed to generate inferable code. In this case, we should use MatrixFields.top_level_names(Yₜ.c) instead of propertynames(Yₜ.c), and then the corresponding field views can be extracted as MatrixFields.get_field(Y.c, ρχ_name) instead of getproperty(Y.c, ρχ_name). This is the pattern used throughout the implicit solver.

[charleskawczynski]

@tapios / @dennisYatunin

    foreach_tracer(Yₜ, Y) do ᶜρχₜ, ᶜρχ, ρχ_name, is_ρq_tot
        if !is_ρq_tot
            ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ))
            vtt = vertical_transport(ᶜρ, ᶠu³, ᶜχ, float(dt), tracer_upwinding)
            @. ᶜρχₜ += vtt
        end
    end

does not seem to be equivalent to the original form:

    for (ᶜρχₜ, ᶜχ, χ_name) in matching_subfields(Yₜ.c, ᶜspecific)
        χ_name in (:e_tot, :q_tot) && continue
        vtt = vertical_transport(ᶜρ, ᶠu³, ᶜχ, float(dt), tracer_upwinding)
        @. ᶜρχₜ += vtt
    end

The error is:

ERROR: LoadError: MethodError: no method matching /(::@NamedTuple{ρatke::Float32}, ::Float32)
--
  | The function `/` exists, but no method is defined for this combination of argument types.
  |  
  | Closest candidates are:
  | /(::TaylorSeries.TaylorN{T}, ::S) where {T<:Union{Real, Complex}, S<:Union{Real, Complex}}
  | @ TaylorSeries /central/scratch/esm/slurm-buildkite/climaatmos-ci/depot/default/packages/TaylorSeries/azAPg/src/arithmetic.jl:801
  | /(::TaylorSeries.HomogeneousPolynomial{T}, ::S) where {T<:Union{Real, Complex}, S<:Union{Real, Complex}}
  | @ TaylorSeries /central/scratch/esm/slurm-buildkite/climaatmos-ci/depot/default/packages/TaylorSeries/azAPg/src/arithmetic.jl:801
  | /(::TaylorSeries.Taylor1{Rational{T}}, ::S) where {T<:Integer, S<:Union{Real, Complex}}
  | @ TaylorSeries /central/scratch/esm/slurm-buildkite/climaatmos-ci/depot/default/packages/TaylorSeries/azAPg/src/arithmetic.jl:793
  | ...
  |  
  | Stacktrace:
  | [1] specific
...

Should I change the condition to ρχ_name in (@name(e_tot), @name(q_tot)) && continue? (I'm wondering about the utility of is_ρq_tot if it's not covering all of the cases.

[tapios]

@tapios / @dennisYatunin

    foreach_tracer(Yₜ, Y) do ᶜρχₜ, ᶜρχ, ρχ_name, is_ρq_tot
        if !is_ρq_tot
            ᶜχ = @. lazy(specific(ᶜρχ, Y.c.ρ))
            vtt = vertical_transport(ᶜρ, ᶠu³, ᶜχ, float(dt), tracer_upwinding)
            @. ᶜρχₜ += vtt
        end
    end

does not seem to be equivalent to the original form:

    for (ᶜρχₜ, ᶜχ, χ_name) in matching_subfields(Yₜ.c, ᶜspecific)
        χ_name in (:e_tot, :q_tot) && continue
        vtt = vertical_transport(ᶜρ, ᶠu³, ᶜχ, float(dt), tracer_upwinding)
        @. ᶜρχₜ += vtt
    end

The error is:

ERROR: LoadError: MethodError: no method matching /(::@NamedTuple{ρatke::Float32}, ::Float32)
--
  | The function `/` exists, but no method is defined for this combination of argument types.
  |  
  | Closest candidates are:
  | /(::TaylorSeries.TaylorN{T}, ::S) where {T<:Union{Real, Complex}, S<:Union{Real, Complex}}
  | @ TaylorSeries /central/scratch/esm/slurm-buildkite/climaatmos-ci/depot/default/packages/TaylorSeries/azAPg/src/arithmetic.jl:801
  | /(::TaylorSeries.HomogeneousPolynomial{T}, ::S) where {T<:Union{Real, Complex}, S<:Union{Real, Complex}}
  | @ TaylorSeries /central/scratch/esm/slurm-buildkite/climaatmos-ci/depot/default/packages/TaylorSeries/azAPg/src/arithmetic.jl:801
  | /(::TaylorSeries.Taylor1{Rational{T}}, ::S) where {T<:Integer, S<:Union{Real, Complex}}
  | @ TaylorSeries /central/scratch/esm/slurm-buildkite/climaatmos-ci/depot/default/packages/TaylorSeries/azAPg/src/arithmetic.jl:793
  | ...
  |  
  | Stacktrace:
  | [1] specific
...

Should I change the condition to ρχ_name in (@name(e_tot), @name(q_tot)) && continue? (I'm wondering about the utility of is_ρq_tot if it's not covering all of the cases.

The problem seems to be that the function treats ρatke like other GS variables and then tries to extract atke, which doesn't make sense. Only GS variables should be included in the foreach_tracer loop. I think you need to replace Y in foreach_tracer by Y.c and Yₜ by Yₜ.c, i.e., foreach_tracer(Yₜ.c, Y.c) do.... (or whatever other substitution is required to exclude the SGS variables).

Also, given that ρχ_name is now emitted by foreach_tracer, we should remove the now-redundant is_ρq_tot output argument from foreach_tracer and use if ρχ_name == :q_tot instead of if is_ρq_tot

[charleskawczynski]

The problem seems to be that the function treats ρatke like other GS variables and then tries to extract atke, which doesn't make sense. Only GS variables should be included in the foreach_tracer loop. I think you need to replace Y in foreach_tracer by Y.c and Yₜ by Yₜ.c, i.e., foreach_tracer(Yₜ.c, Y.c) do.... (or whatever other substitution is required to exclude the SGS variables).

I don't think that will work because foreach_tracer already indexes into Y.c: ᶜρχ = MatrixFields.get_field(Y.c, scalar_name). I'll try with if ρχ_name in (@name(ρe_tot), @name(ρq_tot)) ... and see if that works.

I don't think this addresses the core issue of what will go wrong here: You need to restrict the loop to grid-scale variables, not those starting with sgs. This should be an easy fix either through change of input arguments or of foreach_tracer. (In other words, even if you exclude :e_tot, :q_tot correctly, you will still get wrong variables extracted here.)

[charleskawczynski]

I'm still getting the same error, so the issue seems to be with foreach_tracer / tracer_names. I think we need to update tracer_names to match the previous symbol-based approach:

is_energy_var(name) = name in (@name(ρe_tot), @name(ρae_tot))
is_momentum_var(name) = name in (@name(uₕ), @name(ρuₕ), @name(u₃), @name(ρw))
is_turbconv_var(name) = name in (@name(turbconv), @name(sgsʲs), @name(sgs⁰))
is_tracer_var(name) = !(
    name == @name(ρ) ||
    name == @name(ρa) ||
    is_energy_var(name) ||
    is_momentum_var(name) ||
    is_turbconv_var(name)
)

is_tracer_var(symbol::Symbol) = is_tracer_var(@name(symbol))

tracer_names(field) =
    unrolled_filter(is_tracer_var, MatrixFields.top_level_names(field))

Going to try this now, does this look okay?

[tapios]

I'm still getting the same error, so the issue seems to be with foreach_tracer / tracer_names. I think we need to update tracer_names to match the previous symbol-based approach:

is_energy_var(name) = name in (@name(ρe_tot), @name(ρae_tot))
is_momentum_var(name) = name in (@name(uₕ), @name(ρuₕ), @name(u₃), @name(ρw))
is_turbconv_var(name) = name in (@name(turbconv), @name(sgsʲs), @name(sgs⁰))
is_tracer_var(name) = !(
    name == @name(ρ) ||
    name == @name(ρa) ||
    is_energy_var(name) ||
    is_momentum_var(name) ||
    is_turbconv_var(name)
)

is_tracer_var(symbol::Symbol) = is_tracer_var(@name(symbol))

tracer_names(field) =
    unrolled_filter(is_tracer_var, MatrixFields.top_level_names(field))

Going to try this now, does this look okay?

You don't need ρae_tot in the first line (it's captured by the sgs* and by ρa exclusions). Given you have ρa in the exclusion list, you can also omit is_turbconv_var. Otherwise ok.

[tapios]

I'm still getting the same error, so the issue seems to be with foreach_tracer / tracer_names. I think we need to update tracer_names to match the previous symbol-based approach:

is_energy_var(name) = name in (@name(ρe_tot), @name(ρae_tot))
is_momentum_var(name) = name in (@name(uₕ), @name(ρuₕ), @name(u₃), @name(ρw))
is_turbconv_var(name) = name in (@name(turbconv), @name(sgsʲs), @name(sgs⁰))
is_tracer_var(name) = !(
    name == @name(ρ) ||
    name == @name(ρa) ||
    is_energy_var(name) ||
    is_momentum_var(name) ||
    is_turbconv_var(name)
)

is_tracer_var(symbol::Symbol) = is_tracer_var(@name(symbol))

tracer_names(field) =
    unrolled_filter(is_tracer_var, MatrixFields.top_level_names(field))

Going to try this now, does this look okay?

You don't need ρae_tot in the first line (it's captured by the sgs* and by ρa exclusions). Given you have ρa in the exclusion list, you can also omit is_turbconv_var. Otherwise ok.

What would be cleaner is positively identifying the GS variables, i.e., those of the form Y.c.ρ*, except ρ and ρe_tot. Could you do that? Then this function could remain unchanged if one day we want to add other categories of tracers.

[charleskawczynski]

What would be cleaner is positively identifying the GS variables, i.e., those of the form Y.c.ρ*, except ρ and ρe_tot. Could you do that? Then this function could remain unchanged if one day we want to add other categories of tracers.

A regex-style scan of the FieldName's name_chain would require converting the symbols to strings inside a generated function, which I wouldn't recommend. Also, it'd be difficult printing useful error messages with that code as that could cause undefined behavior (which we've run into in the past).

I kind of get the feeling that we're trying to make too many changes at once. fe9e66d worked, aside from a few jet failures, I feel like it'd be easier to bump the alloc / jet limits, merge that commit, and make these cleanup changes / reduce some duplication later so that we can make progress on the immediate goal of eliminating the dependency on ᶜspecific.

[tapios]

What would be cleaner is positively identifying the GS variables, i.e., those of the form Y.c.ρ*, except ρ and ρe_tot. Could you do that? Then this function could remain unchanged if one day we want to add other categories of tracers.

A regex-style scan of the FieldName's name_chain would require converting the symbols to strings inside a generated function, which I wouldn't recommend. Also, it'd be difficult printing useful error messages with that code as that could cause undefined behavior (which we've run into in the past).

I kind of get the feeling that we're trying to make too many changes at once. fe9e66d worked, aside from a few jet failures, I feel like it'd be easier to bump the alloc / jet limits, merge that commit, and make these cleanup changes / reduce some duplication later so that we can make progress on the immediate goal of eliminating the dependency on ᶜspecific.

As I said, I do not like code duplication; it just creates problems down the line. What we are doing here is not complicated.

I thought there were problems in the previous commit with getproperty(Y.c, ρχ_name)? If not, you can simply use

if is_tracer_var(ρχ_name) && !(ρχ_name in (:ρe_tot, :ρq_tot))Add commentMore actions
            ᶜρχ = getproperty(Y.c, ρχ_name)
            ᶜρχₜ = getproperty(Yₜ.c, ρχ_name)

inside a redefined foreach_tracer.

In any case, one way or another we need to replicate what specific_gs(Y.c) was doing before, which is extracting grid-scale tracers only. One wouldn't think this is a hard problem.

[charleskawczynski]

As I said, I do not like code duplication; it just creates problems down the line. What we are doing here is not complicated.

I thought there were problems in the previous commit with getproperty(Y.c, ρχ_name)? If not, you can simply use

if is_tracer_var(ρχ_name) && !(ρχ_name in (:ρe_tot, :ρq_tot))Add commentMore actions
            ᶜρχ = getproperty(Y.c, ρχ_name)
            ᶜρχₜ = getproperty(Yₜ.c, ρχ_name)

inside a redefined foreach_tracer.

In any case, one way or another we need to replicate what specific_gs(Y.c) was doing before, which is extracting grid-scale tracers only. One wouldn't think this is a hard problem.

In principle, the commit before using/modifying is_tracer_var seemed like it should work. Looking at the error from that:

MethodError: no method matching /(::@NamedTuple{ρatke::Float32}, ::Float32)

It seems like tracer_names is somehow including @name(sgs⁰.ρatke) in the iterator, which (I think) it shouldn't. I'll have a closer look at MatrixFields to see how MatrixFields.top_level_names works.

[tapios]

As I said, I do not like code duplication; it just creates problems down the line. What we are doing here is not complicated.
I thought there were problems in the previous commit with getproperty(Y.c, ρχ_name)? If not, you can simply use

if is_tracer_var(ρχ_name) && !(ρχ_name in (:ρe_tot, :ρq_tot))Add commentMore actions
            ᶜρχ = getproperty(Y.c, ρχ_name)
            ᶜρχₜ = getproperty(Yₜ.c, ρχ_name)

inside a redefined foreach_tracer.
In any case, one way or another we need to replicate what specific_gs(Y.c) was doing before, which is extracting grid-scale tracers only. One wouldn't think this is a hard problem.

In principle, the commit before using/modifying is_tracer_var seemed like it should work. Looking at the error from that:

MethodError: no method matching /(::@NamedTuple{ρatke::Float32}, ::Float32)

It seems like tracer_names is somehow including @name(sgs⁰.ρatke) in the iterator, which (I think) it shouldn't. I'll have a closer look at MatrixFields to see how MatrixFields.top_level_names works.

Ok. Same problem as for foreach_tracer, then.

[charleskawczynski]

Ok, I was able to reproduce it locally and I see:

julia> CA.tracer_names(Y.c)
(@name(ρq_tot), @name(sgs⁰))

So, I think we want tracer_names to not use top_level_names and instead use something like the leaves?

[dennisYatunin]

You could also keep using top_level_names and just filter out all names that start with sgs. There isn't a predefined function that makes an iterator of leaf names in MatrixFields (but it should be easy to write if you want it).

[charleskawczynski]

You could also keep using top_level_names and just filter out all names that start with sgs. There isn't a predefined function that makes an iterator of leaf names in MatrixFields (but it should be easy to write if you want it).

As in

tracer_names(field) =
    unrolled_filter(MatrixFields.top_level_names(field)) do name
        !(name in (@name(ρ), @name(ρe_tot), @name(uₕ), @name(sgs⁰), @name(sgsʲs)))
    end

?

[tapios]

LGTM. Thanks!

One minor detail: you may wish to consider renaming foreach_tracer to foreach_gs_tracer for added clarity, but it's not crucial.

[tapios]

LGTM. Thanks, @charleskawczynski!