Add assert_eltype helper

[charleskawczynski]

When working on #3817, I found some things that surprised me a bit. In particular, with the new helper (added in this PR), I believe that this will fail (let's see what CI says):

function surface_velocity(ᶠu₃, ᶠuₕ³)
    assert_eltype(ᶠu₃, Geometry.Covariant3Vector)
    assert_eltype(ᶠuₕ³, Geometry.Contravariant3Vector)
    sfc_u₃ = Fields.level(ᶠu₃.components.data.:1, half)
    sfc_uₕ³ = Fields.level(ᶠuₕ³.components.data.:1, half)
    sfc_g³³ = g³³_field(sfc_u₃)
    w₃ = @. lazy(-sfc_uₕ³ / sfc_g³³) # u³ = uₕ³ + w³ = uₕ³ + w₃ * g³³
    assert_eltype(w₃, Geometry.Covariant3Vector)
    return w₃
end

despite, where this function is called looks like sfc_u₃ .= surface_velocity(Y.f.u₃, ᶠuₕ³). The variable being assigned appears to be a Covariant3Vector, however, the RHS of that expression doesn't seem to be.

I think the reason is because we reach into the components to get floats, and then the type of the output is actually a Contravariant3Vector, as it adopts the type of g³³.

This doesn't error, however, because ClimaCore has conversion methods for this sort of thing. I'm not sure if this is intended behavior, or if this will yield the intended result or not.

Assuming ClimaCore does this automatically, I'm actually surprised that this method isn't simply written as

function surface_velocity(ᶠu₃, ᶠuₕ³)
    assert_eltype(ᶠu₃, Geometry.Covariant3Vector)
    assert_eltype(ᶠuₕ³, Geometry.Contravariant3Vector)
    sfc_u₃ = Fields.level(ᶠu₃, half)
    sfc_uₕ³ = Fields.level(ᶠuₕ³, half)
    w₃ = @. lazy(-sfc_uₕ³ / sfc_u₃)
    assert_eltype(w₃, Geometry.Covariant3Vector)
    return w₃
end

Also, the assert_eltype should work for DataLayouts, Fields, and Broadcasted objects, so we can safely remove some of the other type assertions, which currently prevents some of the lazy machinery from working.

[coderabbitai]

Walkthrough

This update adds runtime type assertions to velocity vector functions in src/cache/precomputed_quantities.jl, enforcing geometric vector type correctness. It also introduces a versatile assert_eltype function with multiple dispatch in src/utils/utilities.jl to check element types in various container-like objects. No public APIs or function signatures are changed; only internal checks and a new utility function are added, preserving existing logic and flow.

Note

⚡️ AI Code Reviews for VS Code, Cursor, Windsurf

CodeRabbit now has a plugin for VS Code, Cursor and Windsurf. This brings AI code reviews directly in the code editor. Each commit is reviewed immediately, finding bugs before the PR is raised. Seamless context handoff to your AI code agent ensures that you can easily incorporate review feedback.
Learn more here.

---

📜 Recent review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between f309f92 and 5e6f601.

📒 Files selected for processing (1)

• src/cache/precomputed_quantities.jl (3 hunks)

🚧 Files skipped from review as they are similar to previous changes (1)

• src/cache/precomputed_quantities.jl

⏰ Context from checks skipped due to timeout of 90000ms (10)

• GitHub Check: evaluate
• GitHub Check: docbuild
• GitHub Check: ci 1.11 - ubuntu-latest
• GitHub Check: ci 1.11 - windows-latest
• GitHub Check: ci 1.10 - ubuntu-latest
• GitHub Check: ci 1.10 - windows-latest
• GitHub Check: downstream ClimaCoupler.jl (1.10)
• GitHub Check: downstream ClimaCoupler.jl (1.11)
• GitHub Check: test (1.11)
• GitHub Check: test (1.10)

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

• Review comments: Directly reply to a review comment made by CodeRabbit. Example:
  • I pushed a fix in commit <commit_id>, please review it.
  • Explain this complex logic.
  • Open a follow-up GitHub issue for this discussion.
• Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query. Examples:
  • @coderabbitai explain this code block.
  • @coderabbitai modularize this function.
• PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
  • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
  • @coderabbitai read src/utils.ts and explain its main purpose.
  • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.
  • @coderabbitai help me debug CodeRabbit configuration file.

Support

Need help? Create a ticket on our support page for assistance with any issues or questions.

Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments.

CodeRabbit Commands (Invoked using PR comments)

• @coderabbitai pause to pause the reviews on a PR.
• @coderabbitai resume to resume the paused reviews.
• @coderabbitai review to trigger an incremental review. This is useful when automatic reviews are disabled for the repository.
• @coderabbitai full review to do a full review from scratch and review all the files again.
• @coderabbitai summary to regenerate the summary of the PR.
• @coderabbitai generate docstrings to generate docstrings for this PR.
• @coderabbitai generate sequence diagram to generate a sequence diagram of the changes in this PR.
• @coderabbitai resolve resolve all the CodeRabbit review comments.
• @coderabbitai configuration to show the current CodeRabbit configuration for the repository.
• @coderabbitai help to get help.

Other keywords and placeholders

• Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
• Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
• Add @coderabbitai anywhere in the PR title to generate the title automatically.

CodeRabbit Configuration File (.coderabbit.yaml)

• You can programmatically configure CodeRabbit by adding a .coderabbit.yaml file to the root of your repository.
• Please see the configuration documentation for more information.
• If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: # yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json

Documentation and Community

• Visit our Documentation for detailed information on how to use CodeRabbit.
• Join our Discord Community to get help, request features, and share feedback.
• Follow us on X/Twitter for updates and announcements.

[coderabbitai]

Actionable comments posted: 0

🧹 Nitpick comments (1)

src/cache/precomputed_quantities.jl (1)

290-297: Clean up commented code.

The commented alternative implementation should be removed if it's no longer needed, as it adds clutter to the codebase.

-    # sfc_u₃ = Fields.level(ᶠu₃, half)
-    # sfc_uₕ³ = Fields.level(ᶠuₕ³, half)
-    # w₃ = @. lazy(-sfc_uₕ³ / sfc_u₃) # are metric terms automatically applied here?
-

📜 Review details

Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between e743a79 and 6bbc0fc.

📒 Files selected for processing (2)

• src/cache/precomputed_quantities.jl (3 hunks)
• src/utils/utilities.jl (1 hunks)

⏰ Context from checks skipped due to timeout of 90000ms (12)

• GitHub Check: docbuild
• GitHub Check: test (1.10)
• GitHub Check: test (1.11)
• GitHub Check: ci 1.11 - windows-latest
• GitHub Check: ci 1.11 - macOS-latest
• GitHub Check: ci 1.10 - windows-latest
• GitHub Check: ci 1.10 - macOS-latest
• GitHub Check: evaluate
• GitHub Check: ci 1.10 - ubuntu-latest
• GitHub Check: ci 1.11 - ubuntu-latest
• GitHub Check: downstream ClimaCoupler.jl (1.10)
• GitHub Check: downstream ClimaCoupler.jl (1.11)

🔇 Additional comments (4)

src/utils/utilities.jl (1)

497-516: Well-designed utility function with proper multiple dispatch.

The assert_eltype implementation is clean and follows Julia best practices. The multiple dispatch design allows extensibility for different container types, and the error message is informative.

src/cache/precomputed_quantities.jl (3)

259-259: Good type safety addition.

Adding the assertion ensures ᶜuₕ conforms to the expected Covariant12Vector type, catching potential type mismatches early.

---

285-286: Proper input validation.

These assertions verify the geometric vector types match expectations, which aligns with the PR's goal of catching type inconsistencies.

---

326-328: Comprehensive type checking.

All three assertions properly validate the expected geometric vector types for the velocity computation inputs and outputs.

[dennisYatunin]

@trontrytel and I are confused about the purpose of this PR. Is there a bug that you have identified in surface_velocity? I am not sure what that bug could be, as the surface vertical velocity follows what is in the dycore paper. And this function needs to return a scalar Field, given how it is used in set_velocity_at_surface!.

Also, I am strongly opposed to adding type assertions throughout our code like this. They make it very challenging to properly use duck typing, e.g., to use Covariant1Vectors instead of Covariant12Vectors for horizontal velocity in 2D domains. Unnecessary type restrictions have been making it very difficult for me to implement automatic differentiation in ClimaAtmos, and I think we should avoid them wherever possible.

[charleskawczynski]

Is there a bug that you have identified in surface_velocity?

Potentially?

I am not sure what that bug could be, as the surface vertical velocity follows what is in the dycore paper. And this function needs to return a scalar Field, given how it is used in set_velocity_at_surface!.

Shouldn't it return a Covariant3Vector field? That's the type of the destination.

Also, I am strongly opposed to adding type assertions throughout our code like this. They make it very challenging to properly use duck typing, e.g., to use Covariant1Vectors instead of Covariant12Vectors for horizontal velocity in 2D domains. Unnecessary type restrictions have been making it very difficult for me to implement automatic differentiation in ClimaAtmos, and I think we should avoid them wherever possible.

I don't think that this type assertion will restrict autodiff stuff. Actually, I'd also like to remove some type assertions, like Field in compute_kinetic, as that doesn't allow it to work with broadcasted objects. This type assertion, however, will, as it also supports broadcasted objects, and there is no assertion on the underlying float type, just the axistensor type.

I see type assertions as a tool to help ensure correctness. I agree that we shouldn't use them if they restrict behavior, and as far as I can tell, this one should not.

[dennisYatunin]

The destination is actually a scalar field, not a Covariant3Vector field. See the following lines in set_velocity_at_surface!:

ClimaAtmos.jl/src/cache/precomputed_quantities.jl

Lines 273 to 274 in 5e6f601

	sfc_u₃ = Fields.level(Y.f.u₃.components.data.:1, half)
	sfc_u₃ .= surface_velocity(Y.f.u₃, ᶠuₕ³)

So returning a Covariant3Vector field will likely break set_velocity_at_surface!. If you're really keen on returning a vector field from this function, you'll need to change set_velocity_at_surface! appropriately. However, I would suggest sticking with what we have now, since it is exactly what we wrote in the dycore paper.

And yes, this particular type assertion does not restrict autodiff; it only restricts the horizontal velocity type. I was just pointing out that most type assertions constrain the amount of duck typing we can perform, with the types of real numbers and horizontal velocities being two examples. If we want to check for correct types, we should do so through specific unit tests, instead of scattering restrictive type assertions around our code.

[charleskawczynski]

I was under the assumption that the ₃subscript indicated Covariant vector. Assuming this is the case, I would suggest that we also change sfc_u₃ = Fields.level(Y.f.u₃.components.data.:1, half).

[dennisYatunin]

Yeah it would be more accurate to call it something like sfc_u₃_data (which would also match the naming convention for diagnostic EDMF).

[charleskawczynski]

I'll close, as it seems that people don't want this.