Make ᶜspecific lazy and remove matching_subfields

[charleskawczynski]

This PR extends foreach_gs_tracer to handle specific tracer variables (like the ones in ᶜ∇²specific_tracers), and uses it to replace the for-loop over matching_subfields in the hyperdiffusion tendency. It also replaces specific_gs.(Y.c) with the lazy function ᶜspecific_gs_tracers(Y), which we can use to remove ᶜspecific from the cache in a future PR. Since matching_subfields, specific_gs, and some related functions are no longer needed, their definitions are removed.

This PR also fixes a minor bug and updates an inaccurate code comment in the hyperdiffusion tendency for precipitating species. Details can be found in the comment below.

[charleskawczynski]

@tapios, I tried lumping everything into the one function, but I'm getting an error:

ERROR: LoadError: The function body AST defined by this @generated function is not pure. This likely means it contains a closure, a comprehension or a generator.
  | Stacktrace:
  | [1] macro expansion
  | @ /central/scratch/esm/slurm-buildkite/climaatmos-ci/24693/climaatmos-ci/src/prognostic_equations/hyperdiffusion.jl:245 [inlined]
  |  

Is it alright if I break this back down into smaller function calls (which worked) and only do the generated function where it's needed? Or maybe @dennisYatunin has a better idea?

[tapios]

@tapios, I tried lumping everything into the one function, but I'm getting an error:

ERROR: LoadError: The function body AST defined by this @generated function is not pure. This likely means it contains a closure, a comprehension or a generator.
  | Stacktrace:
  | [1] macro expansion
  | @ /central/scratch/esm/slurm-buildkite/climaatmos-ci/24693/climaatmos-ci/src/prognostic_equations/hyperdiffusion.jl:245 [inlined]
  |  

Is it alright if I break this back down into smaller function calls (which worked) and only do the generated function where it's needed? Or maybe @dennisYatunin has a better idea?

I'd like to understand what fails here. The problem is in hyperdiffusion, l. 245. Nothing here needs to change if what we use is the old specific_gs. It looks like the problem is with the changes from specific_gs to all_specific_gs (which weren't needed here from what I know).

I'm fine with breaking this down into smaller functions, but I'm afraid the problem is something else. @dennisYatunin can weigh in.

[tapios]

@tapios, I tried lumping everything into the one function, but I'm getting an error:

ERROR: LoadError: The function body AST defined by this @generated function is not pure. This likely means it contains a closure, a comprehension or a generator.
  | Stacktrace:
  | [1] macro expansion
  | @ /central/scratch/esm/slurm-buildkite/climaatmos-ci/24693/climaatmos-ci/src/prognostic_equations/hyperdiffusion.jl:245 [inlined]
  |  

Is it alright if I break this back down into smaller function calls (which worked) and only do the generated function where it's needed? Or maybe @dennisYatunin has a better idea?

Isn't the right thing to do here is using the same logic of looping over scalar names that you used several times in #3858?

[charleskawczynski]

Is it alright if I break this back down into smaller function calls (which worked) and only do the generated function where it's needed? Or maybe @dennisYatunin has a better idea?

@dennisYatunin ?

[tapios]

Is it alright if I break this back down into smaller function calls (which worked) and only do the generated function where it's needed? Or maybe @dennisYatunin has a better idea?

@dennisYatunin ?

As I mentioned, you'd want to use a (corrected form) of foreach_tracer here, restricted to grid-scale variables. Then this issue will disappear.

[charleskawczynski]

As I mentioned, you'd want to use a (corrected form) of foreach_tracer here, restricted to grid-scale variables. Then this issue will disappear.

It's not clear to me how to use foreach_tracer in this function. Are you suggesting to scrap the function altogether and just use foreach_tracer? If so, we'll need foreach_tracer to emit both ρχ_name and χ_symbol, as both are needed in the hyperdiffusion. Specifically, how do we use foreach_tracer for:

    for χ_name in propertynames(ᶜ∇²specific_tracers)
        @. ᶜ∇²specific_tracers.:($$χ_name) = wdivₕ(gradₕ(ᶜspecific.:($$χ_name)))
    end

?

[dennisYatunin]

@tapios, I've modified this PR to include the change you requested above, replacing the last use of matching_subfields in with an extended version of foreach_gs_tracer that can handle specific tracer variables. I also defined a lazy alternative to specific_gs.(Y.c), which we can use to remove ᶜspecific from the cache in a future PR. The definitions of matching_subfields, specific_gs, and related functions have been removed since they are no longer necessary.

In the process of updating the hyperdiffusion tendency, I found a bug on lines 283--298 in the current version of hyperdiffusion.jl:

    ν₄_scalar = ν₄_scalar_coeff * h_length_scale^3
    (; ᶜ∇²specific_tracers) = p.hyperdiff
    for (ᶜρχₜ, ᶜ∇²χ, χ_name) in matching_subfields(Yₜ.c, ᶜ∇²specific_tracers)
        ν₄_scalar = ifelse(
            χ_name in (:q_rai, :q_sno, :n_rai),
            α_hyperdiff_tracer * ν₄_scalar,
            ν₄_scalar,
        )
        ...
    end

This loop modifies the value of ν₄_scalar whenever χ_name in (:q_rai, :q_sno, :n_rai), changing the diffusion coefficient for all subsequent tracer variables. I've fixed this by assigning a new name ν₄_scalar_for_χ to the variable in the loop. @trontrytel, @akshaysridhar, this bugfix may be relevant for you.

I also changed the docstring on lines 301--302 to more accurately describe what the hyperdiffusion tendency is doing. @tapios, please take a look and let me know if the new version of this comment on lines 300--302 is incorrect.

[tapios]

LGTM. Good catch on the hyperdiffusivity bug; that's important.

We should isolate the hyperviscosity/-diffusivity calculation to a function, and determine scalar diffusivity from hyperdiffusivity = hyperviscosity/Prandtl_number_hyperdiff, with the Prandtl number in the parameter file. You may consider introducing a function now, to avoid possible errors from the code duplication and inconsistent edits. We can also do it later.

[dennisYatunin]

Ok, I've moved the computation of ν₄_scalar and ν₄_vorticity into a new function ν₄(hyperdiff, Y), and I added your reference to its docstring. We can change this function to use the Prandtl number in a future PR.

[tapios]

Ok, I've moved the computation of ν₄_scalar and ν₄_vorticity into a new function ν₄(hyperdiff, Y), and I added your reference to its docstring. We can change this function to use the Prandtl number in a future PR.

Sounds good. Thanks.