ecmwf / earthkit-meteo

Apache License 2.0
4 stars 0 forks source link

WIP: Support multiple array formats #13

Open sandorkertesz opened 6 months ago

sandorkertesz commented 6 months ago

Fixes #9

As an experiment , I converted thermo.specific_humidity_from_vapour_pressure to be array format agnostic but still allowing numbers as input. The code goes like this:

from earthkit.meteo.utils.array import array_namespace
...
def specific_humidity_from_vapour_pressure(e, p, eps=1e-4):
   ....

    ns = array_namespace(e, p)
    v = ns.asarray(p + (constants.epsilon - 1) * e)
    v[ns.asarray(p - e) < eps] = np.nan
    return constants.epsilon * e / v

I also modified the test for this method to make it work for multiple array formats and numbers too:

from earthkit.meteo.utils.testing import ARRAY_BACKENDS
...

@pytest.mark.parametrize(
    "vp, p, v_ref",
    [
        ([895.992614, 2862.662152, 10000], [700, 1000, 50], [0.008, 0.018, np.nan]),
        ([895.992614, 2862.662152, 100000], 700, [0.008, 0.0258354146, np.nan]),
        (895.992614, 700, 0.008),
        (100000, 700, np.nan),
    ],
)
@pytest.mark.parametrize("array_backend", ARRAY_BACKENDS)
def test_specific_humidity_from_vapour_pressure(vp, p, v_ref, array_backend):
    vp, p, v_ref = array_backend.asarray(vp, p, v_ref)
    p = p * 100
    q = thermo.array.specific_humidity_from_vapour_pressure(vp, p)

    assert array_backend.allclose(q, v_ref, equal_nan=True)

If we can agree on this approach all the other methods can be converted similarly.

codecov-commenter commented 6 months ago

Welcome to Codecov :tada:

Once you merge this PR into your default branch, you're all set! Codecov will compare coverage reports and display results in all future pull requests.

Thanks for integrating Codecov - We've got you covered :open_umbrella:

sandorkertesz commented 2 months ago

Hi @corentincarton, @oiffrig, any thought on this?