Skip to content

fix type issues surfaced from pace integration #81

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 3 commits into from
Jan 15, 2026
Merged

fix type issues surfaced from pace integration #81

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
19a8405
fix type issues surfaced from pace integration
romanc Jan 13, 2026
5b43b0c
NDSL update: Quantity requires backend
romanc Jan 13, 2026
cf5390c
not a typo in pyrte, but with out code
romanc Jan 14, 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
3 changes: 0 additions & 3 deletions pyshield/_config.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -67,9 +67,6 @@ class PhysicsConfig:
nwat: int = DEFAULT_INT
schemes: List = None
ntracers: int = int(len(tracer_variables))
ntiw: int = DEFAULT_INT
ntcw: int = DEFAULT_INT
ntke: int = DEFAULT_INT
Comment on lines -70 to -72
Copy link
Copy Markdown
Collaborator Author

@romanc romanc Jan 13, 2026

Choose a reason for hiding this comment

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

Those are already defined below with a comment (same default values)

do_qa: bool = DEFAULT_BOOL
do_inline_mp: bool = False
"""Whether microphysics is inlined in the dycore"""
Expand Down
1 change: 1 addition & 0 deletions pyshield/physics_state.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -487,6 +487,7 @@ def init_from_storages(
_field.metadata["units"],
origin=sizer.get_origin(dims),
extent=sizer.get_extent(dims),
backend=quantity_factory.backend,
)
else:
quantity = quantity_factory.zeros(
Expand Down
7 changes: 3 additions & 4 deletions pyshield/radiation/rad_astro.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
import datetime
import re
from pathlib import Path
from typing import Union
Copy link
Copy Markdown
Collaborator Author

@romanc romanc Jan 13, 2026

Choose a reason for hiding this comment

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

modern typing uses str | float instead of Union[str, float].


import numpy as np

Expand Down Expand Up @@ -65,7 +64,7 @@ def date_to_julian(iyear: int, imonth: int, iday: int) -> int:

def read_NOAA_solar_file(
solar_fname: Path,
) -> dict[str : Union[float, int, dict[str : Union[int, float, dict[str:float]]]]]:
Copy link
Copy Markdown
Collaborator Author

@romanc romanc Jan 13, 2026

Choose a reason for hiding this comment

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

the format for typing dictionaries is dict[str, Any] not dict[str : Any]

) -> dict[str, float | int | dict[str, int | float | dict[str, float]]]:
"""
function to read in a file of solar constants structured as:
first_year last_year first_cycle last_cycle mean_value info_string
Expand Down Expand Up @@ -98,7 +97,7 @@ def read_NOAA_solar_file(
solar_constant_data["icy1"] = int(table_dat[2])
solar_constant_data["icy2"] = int(table_dat[3])
solar_constant_data["smean"] = float(table_dat[4])
elif re.fullmatch(r"^(\*)\1{1,}$", table_dat[0]):
Copy link
Copy Markdown
Collaborator Author

@romanc romanc Jan 13, 2026

Choose a reason for hiding this comment

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

always use is (not) None to compare against None to help the type system.

elif re.fullmatch(r"^(\*)\1{1,}$", table_dat[0]) is not None:
break # end at the asterisks
else:
year = int(table_dat[0])
Expand Down Expand Up @@ -350,7 +349,7 @@ def solar_update(
deltim: float,
lsol_chg: bool,
iyr_sav: int,
isolflg: bool,
isolflg: int,
solar_constant_data: dict = None,
) -> tuple[float, float, float, float, float, float, int, int]:
"""
Expand Down
10 changes: 5 additions & 5 deletions pyshield/radiation/rad_gases.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
import os
import re
from pathlib import Path
from typing import Union

import numpy as np

Expand Down Expand Up @@ -32,7 +31,7 @@
F113VMR_DEF = 8.2000e-11 # gfdl 1999 value


def read_global_annual_co2(co2gbl_file: Path) -> dict[str : Union[Float, Int]]:
def read_global_annual_co2(co2gbl_file: Path) -> dict[str, Float | Int]:
"""
Function to read a text file of CO2 global half-yearly means
and growth rates into a model.
Expand Down Expand Up @@ -62,7 +61,7 @@ def read_global_annual_co2(co2gbl_file: Path) -> dict[str : Union[Float, Int]]:

def read_monthly_resolved_co2(
co2dat_file: Path,
) -> tuple[int, dict[Union[str, Int] : Union[Float, list[Float]]]]:
) -> tuple[int, dict[str | Int, Float | list[Float]]]:
"""
Function to read a text file of 15-degree CO2 monthly means into a model.
Assumes a format of:
Expand Down Expand Up @@ -103,7 +102,7 @@ def read_monthly_resolved_co2(

def read_monthly_cycle_co2(
co2cyc_file: Path,
) -> dict[Union[str, Int] : Union[Int, Float, dict[str : Union[Float, list[Float]]]]]:
) -> dict[str | Int, Int | Float | dict[str, Float | list[Float]]]:
"""
Function to read a text file of 15-degree CO2 monthly deviations into a model.
Assumes a format of:
Expand Down Expand Up @@ -175,7 +174,6 @@ def read_co2_files(input_dir: Path, prefix="") -> tuple[dict, dict, dict]:
if co2_glob_file.is_file():
co2_glb_data = read_global_annual_co2(co2_glob_file)
if monthly_co2_files:
co2_mvr_data = {}
for monthly_file in monthly_co2_files:
year, data = read_monthly_resolved_co2(input_dir.joinpath(monthly_file))
co2_mvr_data[year] = data
Expand Down Expand Up @@ -242,6 +240,7 @@ def gas_init(
float,
float,
float,
float,
np.ndarray,
np.ndarray,
dict,
Expand Down Expand Up @@ -336,6 +335,7 @@ def gas_init(
cfc11,
cfc12,
cfc22,
cfc113,
Copy link
Copy Markdown
Collaborator Author

@romanc romanc Jan 13, 2026

Choose a reason for hiding this comment

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

there was an inconsistency with those 14 or 15 return values. From usage, I inferred that the 15 one is probably what we want ... please double-check

Copy link
Copy Markdown
Collaborator

@jjuyeonkim jjuyeonkim Jan 13, 2026

Choose a reason for hiding this comment

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

it looks like that's the missing one to me too based off of RTE_RRTMGPDriver in rte_rrtmgp.py.

Copy link
Copy Markdown
Collaborator

@oelbert oelbert Jan 15, 2026

Choose a reason for hiding this comment

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

lol that's correct, whoops! Now I'm worried how that isn't causing issues currently

Copy link
Copy Markdown
Collaborator Author

@romanc romanc Jan 19, 2026

Choose a reason for hiding this comment

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

@oelbert If I were you, I'd want #82 to merge as soon as possible. to avoid such issues. There are two science questions in that PR over there, which I think are for you to decide.

ccl4,
co2_glb,
co2_arr,
Expand Down
21 changes: 15 additions & 6 deletions pyshield/radiation/state.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -312,6 +312,7 @@ def init_from_storages(
_field.metadata["units"],
origin=sizer.get_origin(dims),
extent=sizer.get_extent(dims),
backend=quantity_factory.backend,
)
else:
quantity = quantity_factory.zeros(
Expand Down Expand Up @@ -358,30 +359,38 @@ def to_rterrtmgp_xr(self) -> xr.Dataset:
for name, field_info in self.__dataclass_fields__.items():
if name not in ["quantity_factory", "np_like"]:
if field_info.metadata["intent"] != "out":
if issubclass(field_info.type, Quantity):
if issubclass(field_info.type, Quantity): # type: ignore[arg-type]
dims = []
slice_list = []
ndims = len(field_info.metadata["dims"])
nz = self._nz
for dim_name in field_info.metadata["dims"]:
# dims.append(f"{dim_name}_{name}")
if dim_name == "z_interface":
slice_list.append(self._np.s_[:])
slice_list.append(
self._np.s_[:] # type: ignore[attr-defined]
)
nz = self._nz + 1
dims.append("level")
elif dim_name == "z":
slice_list.append(self._np.s_[:-1])
slice_list.append(
self._np.s_[:-1] # type: ignore[attr-defined]
)
dims.append("layer")
elif "INTERFACE" in dim_name:
slice_list.append(self._np.s_[3:-3])
slice_list.append(
self._np.s_[3:-3] # type: ignore[attr-defined]
)
else:
slice_list.append(self._np.s_[3:-4])
slice_list.append(
self._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:
newshape = (-1, nz)
dims.insert(0, "column")
elif ndims == 2: # x-y array:
newshape = (-1,) # noqa
newshape = (-1,) # type: ignore[assignment]
dims.insert(0, "column")
elif ndims == 1: # z-array
newshape == (nz,)
Expand Down
1 change: 1 addition & 0 deletions pyshield/stencils/gfdl_cld_microphysics/gfdl_cld_microphysics_state.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -519,6 +519,7 @@ def init_from_storages(
_field.metadata["units"],
origin=sizer.get_origin(_field.metadata["dims"]),
extent=sizer.get_extent(_field.metadata["dims"]),
backend=quantity_factory.backend,
)
return cls(**inputs)

Expand Down
1 change: 1 addition & 0 deletions pyshield/stencils/pbl/satmedmfvdiff_state.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -358,6 +358,7 @@ def init_from_storages(
_field.metadata["units"],
origin=sizer.get_origin(dims),
extent=sizer.get_extent(dims),
backend=quantity_factory.backend,
)
else:
quantity = quantity_factory.zeros(
Expand Down
1 change: 1 addition & 0 deletions pyshield/stencils/shallow_convection/samf_shalconv_state.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -222,6 +222,7 @@ def init_from_storages(
_field.metadata["units"],
origin=sizer.get_origin(dims),
extent=sizer.get_extent(dims),
backend=quantity_factory.backend,
)
else:
quantity = quantity_factory.zeros(dims, _field.metadata["units"])
Expand Down
1 change: 1 addition & 0 deletions pyshield/stencils/surface/sfc_state.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -481,6 +481,7 @@ def init_from_storages(
_field.metadata["units"],
origin=sizer.get_origin(dims),
extent=sizer.get_extent(dims),
backend=quantity_factory.backend,
)
else:
quantity = quantity_factory.zeros(
Expand Down
1 change: 1 addition & 0 deletions tests/savepoint/translate/translate_fv_update_phys.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -173,6 +173,7 @@ def compute_parallel(self, inputs, communicator):
units="test",
origin=(0, 0, 0),
extent=storage.shape,
backend=self.grid.quantity_factory.backend,
)
state = DycoreState(**inputs)
self._base.compute_func = ApplyPhysicsToDycore(
Expand Down
2 changes: 2 additions & 0 deletions tests/savepoint/translate/translate_sedimentation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -717,12 +717,14 @@ def compute(self, inputs):
inputs[var],
dims=[X_DIM, Y_DIM, Z_DIM],
units="unknown",
backend=self.quantity_factory.backend,
)
elif len(inputs[var].shape) == 2:
inputs[var] = Quantity(
inputs[var],
dims=[X_DIM, Y_DIM],
units="unknown",
backend=self.quantity_factory.backend,
)
else:
raise TypeError(
Expand Down