Thanks to visit codestin.com
Credit goes to github.com

Skip to content

🚀 LocalState class: a state-like class for organizing Locals #338

Merged
FlorianDeconinck merged 8 commits intoNOAA-GFDL:developfrom
FlorianDeconinck:feature/LocalState
Dec 9, 2025
Merged

🚀 LocalState class: a state-like class for organizing Locals #338
FlorianDeconinck merged 8 commits intoNOAA-GFDL:developfrom
FlorianDeconinck:feature/LocalState

Conversation

@FlorianDeconinck
Copy link
Collaborator

@FlorianDeconinck FlorianDeconinck commented Dec 8, 2025

Description

Introducing LocalState a State containing only Locals. This allows to organize the Locals in the model in a state-like nested dataclasses structure instead of allocating them directly with make_local.

Also:

  • Add checks for forbidden hybrid Quantity/Locals states in creating State and LocalState
  • Add many guardrails to force LocalState to fit Locals boudaries
  • Unit tests for forbidden behavior

How has this been tested?

New unit tests covering the basics cases.
Deployed on GEOS's microphysics code during optimization

Checklist

  • My code follows the style guidelines of this project
  • I have performed a self-review of my own code
  • I have commented my code, particularly in hard-to-understand areas
  • I have made corresponding changes to the documentation (e.g. add new modules to docs/docstrings/)
  • My changes generate no new warnings
  • Any dependent changes have been merged and published in downstream modules
  • New check tests, if applicable, are included

Add checks for forbidden hybrid `Quantity/Locals` states
Add many guardrails to force `LocalState` to fit `Local`s boudaries
@FlorianDeconinck
Copy link
Collaborator Author

Adding @fmalatino and @jjuyeonkim for backend visibility on complex NDSL concept that will be reused for optimization (Local allow us to identify which data can be manipulated because they won't be read outside of an NDSLRuntime object)

Adding @oelbert for frontend API visbility. Locals will be the backbone of the merge + cache optimize routine, especially for physics.

Copy link
Contributor

@fmalatino fmalatino left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for looping me in.

from ndsl.types import Number


from ndsl.quantity import Quantity, Local # isort:skip
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why do they need to be imported last?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Quantity, Local and State are all intertwined and need a severe clean up because Quantity is used everywhere else.

I am locking them at the bottom so all prior dependency run through before we ask for it - that way we stay clear of the recursion.

)

initial_quantities[_field.name] = quantity_factory_allocator(
allow_mismatch_float_precision = True
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

no need to assign, right?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am saving us from a bug where we change the first one but no the second one...

_field.metadata["units"],
dtype=_field.metadata["dtype"],
allow_mismatch_float_precision=True,
allow_mismatch_float_precision=allow_mismatch_float_precision,
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
allow_mismatch_float_precision=allow_mismatch_float_precision,
allow_mismatch_float_precision=True,

origin=quantity.origin,
extent=quantity.extent,
backend=quantity.backend,
allow_mismatch_float_precision=allow_mismatch_float_precision,
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
allow_mismatch_float_precision=allow_mismatch_float_precision,
allow_mismatch_float_precision=True,

allow_mismatch_float_precision=allow_mismatch_float_precision,
)
initial_quantities[_field.name] = local_
else:
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is the logic here slightly funky? We always create a quantity but only add it if the type is not local? Should we not only create it if it is non-local?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No we create a Quantity which is the only thing that truly hold memory then if we have been told to make Local, we wrap the quantity into a Local (Local derives from Quantity, but python sucks at OOP).

The allow_locals extra flags is only there because we are giving the capacity to allocate a LocalState as a pure Quantity state for testing and interfacing

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Local derives from Quantity, but python sucks at OOP

😱

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No lies here.

Copy link
Collaborator

@twicki twicki left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This all makes sense, thanks for talking me through it!

@FlorianDeconinck FlorianDeconinck added this pull request to the merge queue Dec 9, 2025
Merged via the queue into NOAA-GFDL:develop with commit 2724861 Dec 9, 2025
6 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants