Merge pull request #335 from ImperialCollegeLondon/334-standardise-nd…

…array-typing

Updating and systematising NDArray typing

docs/source/conf.py

@@ -99,6 +99,7 @@ class MyReferenceStyle(AuthorYearReferenceStyle):
     ("py:class", "numpy._typing._array_like._ScalarType_co"),
     ("py:class", "numpy._typing._generic_alias.ScalarType"),
     ("py:class", "numpy.float32"),
+    ("py:class", "numpy.float64"),
     ("py:class", "numpy.int64"),
     ("py:class", "numpy.timedelta64"),
     ("py:class", "numpy.bool_"),

pyrealm/constants/core_const.py

@@ -114,7 +114,7 @@ class CoreConst(ConstantsClass):
     :cite:t:`berger:1978a`."""
 
     # Hygro constants
-    magnus_coef: NDArray[np.float32] = field(
+    magnus_coef: NDArray[np.float64] = field(
         default_factory=lambda: np.array((611.2, 17.62, 243.12))
     )
     """Three coefficients of the Magnus equation for saturated vapour pressure,
@@ -132,23 +132,23 @@ class CoreConst(ConstantsClass):
     """Set the method used for calculating water density ('fisher' or 'chen')."""
 
     # Fisher Dial
-    fisher_dial_lambda: NDArray[np.float32] = field(
+    fisher_dial_lambda: NDArray[np.float64] = field(
         default_factory=lambda: np.array(
             [1788.316, 21.55053, -0.4695911, 0.003096363, -7.341182e-06]
         )
     )
     r"""Coefficients of the temperature dependent polynomial for :math:`\lambda`
      in the Tumlirz equation."""
 
-    fisher_dial_Po: NDArray[np.float32] = field(
+    fisher_dial_Po: NDArray[np.float64] = field(
         default_factory=lambda: np.array(
             [5918.499, 58.05267, -1.1253317, 0.0066123869, -1.4661625e-05]
         )
     )
     """Coefficients of the temperature dependent polynomial for :math:`P_0` in the
     Tumlirz equation."""
 
-    fisher_dial_Vinf: NDArray[np.float32] = field(
+    fisher_dial_Vinf: NDArray[np.float64] = field(
         default_factory=lambda: np.array(
             [
                 0.6980547,
@@ -168,7 +168,7 @@ class CoreConst(ConstantsClass):
     in the Tumlirz equation."""
 
     # Chen water density
-    chen_po: NDArray[np.float32] = field(
+    chen_po: NDArray[np.float64] = field(
         default_factory=lambda: np.array(
             [
                 0.99983952,
@@ -186,23 +186,23 @@ class CoreConst(ConstantsClass):
     r"""Coefficients of the polynomial relationship of water density with temperature at
     1 atm (:math:`P^0`, kg/m^3) from :cite:t:`chen:2008a`."""
 
-    chen_ko: NDArray[np.float32] = field(
+    chen_ko: NDArray[np.float64] = field(
         default_factory=lambda: np.array(
             [19652.17, 148.1830, -2.29995, 0.01281, -4.91564e-5, 1.035530e-7]
         )
     )
     r"""Polynomial relationship of bulk modulus of water with temperature at 1 atm
      (:math:`K^0`, kg/m^3) from :cite:t:`chen:2008a`."""
 
-    chen_ca: NDArray[np.float32] = field(
+    chen_ca: NDArray[np.float64] = field(
         default_factory=lambda: np.array(
             [3.26138, 5.223e-4, 1.324e-4, -7.655e-7, 8.584e-10]
         )
     )
     r"""Coefficients of the polynomial temperature dependent coefficient :math:`A` from
      :cite:t:`chen:2008a`."""
 
-    chen_cb: NDArray[np.float32] = field(
+    chen_cb: NDArray[np.float64] = field(
         default_factory=lambda: np.array(
             [7.2061e-5, -5.8948e-6, 8.69900e-8, -1.0100e-9, 4.3220e-12]
         )
@@ -220,11 +220,11 @@ class CoreConst(ConstantsClass):
     huber_mu_ast: float = 1e-06
     r"""Huber reference pressure (:math:`\mu_{ast}` 1.0e-6, Pa s)"""
 
-    huber_H_i: NDArray[np.float32] = field(
+    huber_H_i: NDArray[np.float64] = field(
         default_factory=lambda: np.array([1.67752, 2.20462, 0.6366564, -0.241605])
     )
     """Temperature dependent parameterisation of Hi in Huber."""
-    huber_H_ij: NDArray[np.float32] = field(
+    huber_H_ij: NDArray[np.float64] = field(
         default_factory=lambda: np.array(
             [
                 [0.520094, 0.0850895, -1.08374, -0.289555, 0.0, 0.0],

pyrealm/constants/pmodel_const.py

@@ -139,11 +139,11 @@ class PModelConst(ConstantsClass):
     # - note that kphio_C4 has been updated to account for an unintended double
     #   8 fold downscaling to account for the fraction of light reaching PS2.
     #   from original values of [-0.008, 0.00375, -0.58e-4]
-    kphio_C4: NDArray[np.float32] = field(
+    kphio_C4: NDArray[np.float64] = field(
         default_factory=lambda: np.array((-0.064, 0.03, -0.000464))
     )
     """Quadratic scaling of Kphio with temperature for C4 plants"""
-    kphio_C3: NDArray[np.float32] = field(
+    kphio_C3: NDArray[np.float64] = field(
         default_factory=lambda: np.array((0.352, 0.022, -0.00034))
     )
     """Quadratic scaling of Kphio with temperature for C3 plants"""
@@ -193,11 +193,11 @@ class PModelConst(ConstantsClass):
     """Exponent of the threshold function for Mengoli soil moisture"""
 
     # Unit cost ratio (beta) values for different CalcOptimalChi methods
-    beta_cost_ratio_prentice14: NDArray[np.float32] = field(
+    beta_cost_ratio_prentice14: NDArray[np.float64] = field(
         default_factory=lambda: np.array([146.0])
     )
     r"""Unit cost ratio for C3 plants (:math:`\beta`, 146.0)."""
-    beta_cost_ratio_c4: NDArray[np.float32] = field(
+    beta_cost_ratio_c4: NDArray[np.float64] = field(
         default_factory=lambda: np.array([146.0 / 9])
     )
     r"""Unit cost ratio for C4 plants (:math:`\beta`, 16.222)."""

pyrealm/core/hygro.py

@@ -12,7 +12,9 @@
 from pyrealm.core.utilities import bounds_checker, evaluate_horner_polynomial
 
 
-def calc_vp_sat(ta: NDArray, core_const: CoreConst = CoreConst()) -> NDArray:
+def calc_vp_sat(
+    ta: NDArray[np.float64], core_const: CoreConst = CoreConst()
+) -> NDArray[np.float64]:
     r"""Calculate vapour pressure of saturated air.
 
     This function calculates the vapour pressure of saturated air in kPa at a given
@@ -56,8 +58,10 @@ def calc_vp_sat(ta: NDArray, core_const: CoreConst = CoreConst()) -> NDArray:
 
 
 def convert_vp_to_vpd(
-    vp: NDArray, ta: NDArray, core_const: CoreConst = CoreConst()
-) -> NDArray:
+    vp: NDArray[np.float64],
+    ta: NDArray[np.float64],
+    core_const: CoreConst = CoreConst(),
+) -> NDArray[np.float64]:
     """Convert vapour pressure to vapour pressure deficit.
 
     Args:
@@ -86,8 +90,10 @@ def convert_vp_to_vpd(
 
 
 def convert_rh_to_vpd(
-    rh: NDArray, ta: NDArray, core_const: CoreConst = CoreConst()
-) -> NDArray:
+    rh: NDArray[np.float64],
+    ta: NDArray[np.float64],
+    core_const: CoreConst = CoreConst(),
+) -> NDArray[np.float64]:
     """Convert relative humidity to vapour pressure deficit.
 
     Args:
@@ -124,8 +130,10 @@ def convert_rh_to_vpd(
 
 
 def convert_sh_to_vp(
-    sh: NDArray, patm: NDArray, core_const: CoreConst = CoreConst()
-) -> NDArray:
+    sh: NDArray[np.float64],
+    patm: NDArray[np.float64],
+    core_const: CoreConst = CoreConst(),
+) -> NDArray[np.float64]:
     """Convert specific humidity to vapour pressure.
 
     Args:
@@ -149,8 +157,11 @@ def convert_sh_to_vp(
 
 
 def convert_sh_to_vpd(
-    sh: NDArray, ta: NDArray, patm: NDArray, core_const: CoreConst = CoreConst()
-) -> NDArray:
+    sh: NDArray[np.float64],
+    ta: NDArray[np.float64],
+    patm: NDArray[np.float64],
+    core_const: CoreConst = CoreConst(),
+) -> NDArray[np.float64]:
     """Convert specific humidity to vapour pressure deficit.
 
     Args:
@@ -185,7 +196,9 @@ def convert_sh_to_vpd(
 # The following functions are integrated from the evap.py implementation of SPLASH v1.
 
 
-def calc_saturation_vapour_pressure_slope(tc: NDArray) -> NDArray:
+def calc_saturation_vapour_pressure_slope(
+    tc: NDArray[np.float64],
+) -> NDArray[np.float64]:
     """Calculate the slope of the saturation vapour pressure curve.
 
     Calculates the slope of the saturation pressure temperature curve, following
@@ -207,7 +220,7 @@ def calc_saturation_vapour_pressure_slope(tc: NDArray) -> NDArray:
     )
 
 
-def calc_enthalpy_vaporisation(tc: NDArray) -> NDArray:
+def calc_enthalpy_vaporisation(tc: NDArray[np.float64]) -> NDArray[np.float64]:
     """Calculate the enthalpy of vaporization.
 
     Calculates the latent heat of vaporization of water as a function of
@@ -224,7 +237,7 @@ def calc_enthalpy_vaporisation(tc: NDArray) -> NDArray:
     return 1.91846e6 * ((tc + 273.15) / (tc + 273.15 - 33.91)) ** 2
 
 
-def calc_specific_heat(tc: NDArray) -> NDArray:
+def calc_specific_heat(tc: NDArray[np.float64]) -> NDArray[np.float64]:
     """Calculate the specific heat of air.
 
     Calculates the specific heat of air at a constant pressure (:math:`c_{pm}`, J/kg/K)
@@ -257,8 +270,8 @@ def calc_specific_heat(tc: NDArray) -> NDArray:
 
 
 def calc_psychrometric_constant(
-    tc: NDArray, p: NDArray, core_const: CoreConst = CoreConst()
-) -> NDArray:
+    tc: NDArray[np.float64], p: NDArray[np.float64], core_const: CoreConst = CoreConst()
+) -> NDArray[np.float64]:
     r"""Calculate the psychrometric constant.
 
     Calculates the psychrometric constant (:math:`\lambda`, Pa/K) given the temperature

pyrealm/core/pressure.py

@@ -2,12 +2,15 @@
 atmospheric pressure.
 """  # noqa D210, D415
 
+import numpy as np
 from numpy.typing import NDArray
 
 from pyrealm.constants import CoreConst
 
 
-def calc_patm(elv: NDArray, core_const: CoreConst = CoreConst()) -> NDArray:
+def calc_patm(
+    elv: NDArray[np.float64], core_const: CoreConst = CoreConst()
+) -> NDArray[np.float64]:
     r"""Calculate atmospheric pressure from elevation.
 
     Calculates atmospheric pressure as a function of elevation with reference to the

pyrealm/core/solar.py

@@ -24,7 +24,9 @@
 
 
 def calc_heliocentric_longitudes(
-    julian_day: NDArray, n_days: NDArray, core_const: CoreConst = CoreConst()
+    julian_day: NDArray[np.float64],
+    n_days: NDArray[np.float64],
+    core_const: CoreConst = CoreConst(),
 ) -> tuple[NDArray, NDArray]:
     """Calculate heliocentric longitude and anomaly.
 

pyrealm/core/utilities.py

@@ -211,13 +211,13 @@ def _get_interval_functions(interval_type: str = "[]") -> tuple[np.ufunc, np.ufu
 
 
 def bounds_checker(
-    values: NDArray,
+    values: NDArray[np.float64],
     lower: float = -np.inf,
     upper: float = np.inf,
     interval_type: str = "[]",
     label: str = "",
     unit: str = "",
-) -> NDArray:
+) -> NDArray[np.float64]:
     r"""Check inputs fall within bounds.
 
     This is a simple pass through function that tests whether the values fall within
@@ -255,7 +255,7 @@ def bounds_checker(
 
 
 def bounds_mask(
-    inputs: NDArray,
+    inputs: NDArray[np.float64],
     lower: float = -np.inf,
     upper: float = np.inf,
     interval_type: str = "[]",
@@ -312,7 +312,7 @@ def bounds_mask(
     # modifying the original input. Using type
     if not np.issubdtype(inputs.dtype, np.floating):
         # Copies implicitly
-        outputs = inputs.astype(np.float32)
+        outputs = inputs.astype(np.float64)
     else:
         outputs = inputs.copy()
 
@@ -336,7 +336,9 @@ def bounds_mask(
     return outputs
 
 
-def evaluate_horner_polynomial(x: NDArray, cf: list | NDArray) -> NDArray:
+def evaluate_horner_polynomial(
+    x: NDArray[np.float64], cf: list | NDArray
+) -> NDArray[np.float64]:
     r"""Evaluates a polynomial with coefficients `cf` at `x` using Horner's method.
 
     Horner's method is a fast way to evaluate polynomials, especially for large degrees,

pyrealm/core/water.py

@@ -10,10 +10,10 @@
 
 
 def calc_density_h2o_chen(
-    tc: NDArray,
-    p: NDArray,
+    tc: NDArray[np.float64],
+    p: NDArray[np.float64],
     core_const: CoreConst = CoreConst(),
-) -> NDArray:
+) -> NDArray[np.float64]:
     """Calculate the density of water using Chen et al 2008.
 
     This function calculates the density of water at a given temperature and pressure
@@ -65,10 +65,10 @@ def calc_density_h2o_chen(
 
 
 def calc_density_h2o_fisher(
-    tc: NDArray,
-    patm: NDArray,
+    tc: NDArray[np.float64],
+    patm: NDArray[np.float64],
     core_const: CoreConst = CoreConst(),
-) -> NDArray:
+) -> NDArray[np.float64]:
     """Calculate water density.
 
     Calculates the density of water as a function of temperature and atmospheric
@@ -124,11 +124,11 @@ def calc_density_h2o_fisher(
 
 
 def calc_density_h2o(
-    tc: NDArray,
-    patm: NDArray,
+    tc: NDArray[np.float64],
+    patm: NDArray[np.float64],
     core_const: CoreConst = CoreConst(),
     safe: bool = True,
-) -> NDArray:
+) -> NDArray[np.float64]:
     """Calculate water density.
 
     Calculates the density of water as a function of temperature and atmospheric
@@ -179,11 +179,11 @@ def calc_density_h2o(
 
 
 def calc_viscosity_h2o(
-    tc: NDArray,
-    patm: NDArray,
+    tc: NDArray[np.float64],
+    patm: NDArray[np.float64],
     core_const: CoreConst = CoreConst(),
     simple: bool = False,
-) -> NDArray:
+) -> NDArray[np.float64]:
     r"""Calculate the viscosity of water.
 
     Calculates the viscosity of water (:math:`\eta`) as a function of temperature and
@@ -247,11 +247,11 @@ def calc_viscosity_h2o(
 
 
 def calc_viscosity_h2o_matrix(
-    tc: NDArray,
-    patm: NDArray,
+    tc: NDArray[np.float64],
+    patm: NDArray[np.float64],
     core_const: CoreConst = CoreConst(),
     simple: bool = False,
-) -> NDArray:
+) -> NDArray[np.float64]:
     r"""Calculate the viscosity of water.
 
     Calculates the viscosity of water (:math:`\eta`) as a function of temperature and