Skip to content

[Update] Numpy 2x #91

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
oelbert merged 9 commits into from
Apr 2, 2026
Merged

[Update] Numpy 2x #91

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
cd7fc58
Remove `NumpyModule`
FlorianDeconinck Mar 3, 2026
71c80b8
[TMP] Update CI
FlorianDeconinck Mar 3, 2026
1acff65
Update to numpy >-2
FlorianDeconinck Mar 3, 2026
f06a06d
Lint
FlorianDeconinck Mar 3, 2026
9801c12
Update pace test
FlorianDeconinck Mar 3, 2026
6db01bd
[TMP] Use `numpy_2x` branch of ndsl
FlorianDeconinck Mar 3, 2026
af36bcb
Merge remote-tracking branch 'NOAA/develop' into update/numpy_2x
FlorianDeconinck Mar 18, 2026
f9ffca8
Revert the `ndsl` configuration to `develop`
FlorianDeconinck Apr 1, 2026
02ae7b0
Restore `pyfv3` proper branch
FlorianDeconinck Apr 2, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .github/workflows/pace_tests.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ on:

jobs:
pace_unit_tests:
uses: NOAA-GFDL/pace/.github/workflows/main_unit_tests.yaml@develop
uses: floriandeconinck/pace/.github/workflows/main_unit_tests.yaml@update/numpy_2x
Copy link
Copy Markdown
Collaborator

@oelbert oelbert Apr 2, 2026

Choose a reason for hiding this comment

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

I assume we switch this back once we have a Pace PR in?

Copy link
Copy Markdown
Collaborator Author

@FlorianDeconinck FlorianDeconinck Apr 2, 2026

Choose a reason for hiding this comment

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

Yes, sorry I should have put that in the PR.

We are going up and down the repos to make sure everything keeps green. So once Pace is updated, we will go back down and re-hook develop

with:
component_trigger: true
component_name: pySHiELD
10 changes: 2 additions & 8 deletions .github/workflows/translate.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,8 +44,9 @@ jobs:
uses: actions/checkout@v6
with:
submodules: 'recursive'
repository: NOAA-GFDL/pySHiELD
repository: floriandeconinck/PySHiELD
Copy link
Copy Markdown
Collaborator

@oelbert oelbert Apr 2, 2026

Choose a reason for hiding this comment

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

Don't think this is intended?

Copy link
Copy Markdown
Collaborator Author

@FlorianDeconinck FlorianDeconinck Apr 2, 2026

Choose a reason for hiding this comment

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

It is - it's for the other repositories calling in.

path: pySHiELD
ref: update/numpy_2x

- name: 'Setup Python ${{ inputs.python_version && inputs.python_version || env.python_default }}'
uses: actions/setup-python@v6
Expand All @@ -55,13 +56,6 @@ jobs:
- name: Install mpi (MPICH flavor)
run: pip install mpich

- name: Install numpy==1.26.4
# Front-load installing numpy to force its usage in the radiation
# scheme of pySHiELD.
# TODO: This is an ugly hack and it should be cleaned up latest once
# we add support for numpy >= 2.0
run: pip install numpy==1.26.4

- name: External trigger remove existing component in pySHiELD/develop
if: ${{ inputs.component_name }}
run: rm -rf ${GITHUB_WORKSPACE}/pySHiELD/${{inputs.component_name}}
Expand Down
2 changes: 1 addition & 1 deletion examples/notebook/utilities.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,7 @@ def states_from_fortran_restarts(
tracer_data = xr.open_dataset(tracer_datafile)

state = PhysicsState.init_zeros(quantity_factory, schemes)
radstate = RTE_RRTMGPState.init_zeros(quantity_factory, np)
radstate = RTE_RRTMGPState.init_zeros(quantity_factory)
sstate = SurfaceState.init_zeros(quantity_factory)

buff_3d = np.zeros_like(state.prsi.field)
Expand Down
4 changes: 2 additions & 2 deletions pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ classifiers = [
]
dependencies = [
"f90nml>=1.1.0",
"numpy==1.26.4",
"numpy>=2",
"xarray",
"pyrte-rrtmgp @ git+https://github.com/NOAA-GFDL/pyRTE-RRTMGP.git@main"
]
Expand All @@ -40,7 +40,7 @@ extras = [
"pyshield[pyfv3]",
"pyshield[test]"
]
ndsl = ["ndsl @ git+https://github.com/NOAA-GFDL/NDSL.git@2026.02.00"]
ndsl = ["ndsl @ git+https://github.com/NOAA-GFDL/NDSL.git@develop"]
pyfv3 = ["pyfv3 @ git+https://github.com/NOAA-GFDL/pyFV3.git@develop"]
test = [
"coverage",
Expand Down
17 changes: 5 additions & 12 deletions pyshield/radiation/state.py
View file
Open in desktop
Copy link
Copy Markdown
Collaborator

@oelbert oelbert Apr 2, 2026

Choose a reason for hiding this comment

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

You don't like my pre-xumpy hack?

Copy link
Copy Markdown
Collaborator Author

@FlorianDeconinck FlorianDeconinck Apr 2, 2026

Choose a reason for hiding this comment

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

I loved it. It broke numpy 2 :)

Original file line number Diff line number Diff line change
@@ -1,13 +1,13 @@
from dataclasses import InitVar, dataclass, field, fields
from typing import Any, Dict, Mapping

import numpy as np
import xarray as xr

import ndsl.dsl.gt4py_utils as gt_utils
from ndsl import GridSizer, Quantity, QuantityFactory
from ndsl.constants import I_DIM, J_DIM, K_DIM, K_INTERFACE_DIM
from ndsl.dsl.typing import Float
from ndsl.types import NumpyModule


@dataclass()
Expand Down Expand Up @@ -261,14 +261,11 @@ class RTE_RRTMGPState:
}
)
quantity_factory: InitVar[QuantityFactory]
np_like: InitVar[NumpyModule]

def __post_init__(
self,
quantity_factory: QuantityFactory,
np_like: NumpyModule,
):
self._np = np_like
self._nz = quantity_factory.sizer.nz
self._nx = quantity_factory.sizer.nx
self._ny = quantity_factory.sizer.ny
Expand All @@ -277,7 +274,6 @@ def __post_init__(
def init_zeros(
cls,
quantity_factory,
np_like: NumpyModule,
) -> "RTE_RRTMGPState":
initial_arrays = {}
for _field in fields(cls):
Expand All @@ -290,7 +286,6 @@ def init_zeros(
return cls(
**initial_arrays,
quantity_factory=quantity_factory,
np_like=np_like,
)

@classmethod
Expand All @@ -299,7 +294,6 @@ def init_from_storages(
storages: Mapping[str, Any],
sizer: GridSizer,
quantity_factory: QuantityFactory,
np_like: NumpyModule,
) -> "RTE_RRTMGPState":
inputs: Dict[str, Quantity] = {}
for _field in fields(cls):
Expand All @@ -323,7 +317,6 @@ def init_from_storages(
return cls(
**inputs,
quantity_factory=quantity_factory,
np_like=np_like,
)

@property
Expand Down Expand Up @@ -368,22 +361,22 @@ def to_rterrtmgp_xr(self) -> xr.Dataset:
# dims.append(f"{dim_name}_{name}")
if dim_name == "z_interface":
slice_list.append(
self._np.s_[:] # type: ignore[attr-defined]
np.s_[:] # type: ignore[attr-defined]
)
nz = self._nz + 1
dims.append("level")
elif dim_name == "z":
slice_list.append(
self._np.s_[:-1] # type: ignore[attr-defined]
np.s_[:-1] # type: ignore[attr-defined]
)
dims.append("layer")
elif "INTERFACE" in dim_name:
slice_list.append(
self._np.s_[3:-3] # type: ignore[attr-defined]
np.s_[3:-3] # type: ignore[attr-defined]
)
else:
slice_list.append(
self._np.s_[3:-4] # type: ignore[attr-defined]
np.s_[3:-4] # type: ignore[attr-defined]
)
# We have to reshape to get the max 2D shape rterrtmgp expects:
if ndims == 3: # x-y-z array:
Expand Down
2 changes: 1 addition & 1 deletion tests/integration/test_all.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -82,7 +82,7 @@ def states_from_fortran_restarts(
sfc_data = xr.open_dataset(sfc_datafile)
state = PhysicsState.init_zeros(quantity_factory, schemes)
sstate = SurfaceState.init_zeros(qf_sfc)
radstate = RTE_RRTMGPState.init_zeros(quantity_factory, np)
radstate = RTE_RRTMGPState.init_zeros(quantity_factory)
buff_3d = np.zeros_like(state.prsi.field)
npz = buff_3d.shape[2]
for k in range(npz):
Expand Down
4 changes: 2 additions & 2 deletions tests/integration/test_radiation_driver.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -108,7 +108,7 @@ def states_from_fortran_restarts(
tracer_data = xr.open_dataset(tracer_datafile)

state = PhysicsState.init_zeros(quantity_factory, schemes)
radstate = RTE_RRTMGPState.init_zeros(quantity_factory, np)
radstate = RTE_RRTMGPState.init_zeros(quantity_factory)
sstate = SurfaceState.init_zeros(quantity_factory)

buff_3d = np.zeros_like(state.prsi.field)
Expand Down Expand Up @@ -179,7 +179,7 @@ def test_rte_rrtmgp(datapath: Path):
gridlon = grid_data.lon_agrid
gridlat = grid_data.lat_agrid

state = RTE_RRTMGPState.init_zeros(quantity_factory, np)
state = RTE_RRTMGPState.init_zeros(quantity_factory)
sstate = SurfaceState.init_zeros(quantity_factory)
fortran_restart_to_radstate(
dycore_datafile=dycore_data,
Expand Down