Skip to content

Commit

Permalink
Chore: pre-commit autoupdate
Browse files Browse the repository at this point in the history
  • Loading branch information
pre-commit-ci[bot] committed Jun 3, 2024
1 parent 9c2203c commit f1b582c
Show file tree
Hide file tree
Showing 10 changed files with 190 additions and 85 deletions.
34 changes: 25 additions & 9 deletions tests/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@
import pandas as pd # type: ignore
import pytest
import xarray as xr
from pytest import approx

from hazard.protocols import OpenDataset, WriteDataset
from hazard.utilities import zarr_utilities
Expand Down Expand Up @@ -37,8 +36,12 @@ class TestSource(OpenDataset):
def __init__(self, datasets: Dict[Tuple[str, int], xr.Dataset]):
self.datasets = datasets

def open_dataset_year(self, gcm: str, scenario: str, quantity: str, year: int, chunks=None) -> xr.Dataset:
return self.datasets[(quantity, year)] # ignore scenario and gcm: we test just a single one
def open_dataset_year(
self, gcm: str, scenario: str, quantity: str, year: int, chunks=None
) -> xr.Dataset:
return self.datasets[
(quantity, year)
] # ignore scenario and gcm: we test just a single one


class TestTarget(WriteDataset):
Expand All @@ -47,7 +50,9 @@ class TestTarget(WriteDataset):
def __init__(self):
self.datasets = {}

def write(self, path: str, dataset: xr.Dataset, spatial_coords: Optional[bool] = False):
def write(
self, path: str, dataset: xr.Dataset, spatial_coords: Optional[bool] = False
):
self.datasets[path] = dataset

def read(self, path: str):
Expand All @@ -74,7 +79,9 @@ def _create_test_datasets_tas(

def _create_test_dataset_averaged() -> xr.Dataset:
"""An example 3x3 array that might result from some operation averaging over time."""
temperature = np.array([[293.0, 298.0, 310.0], [304.0, 302.0, 300.0], [308.0, 290.0, 294.0]])
temperature = np.array(
[[293.0, 298.0, 310.0], [304.0, 302.0, 300.0], [308.0, 290.0, 294.0]]
)
lat = np.arange(3.0, 0.0, -1.0)
lon = np.arange(0.0, 3.0, 1.0)
ds = xr.Dataset(
Expand All @@ -87,7 +94,9 @@ def _create_test_dataset_averaged() -> xr.Dataset:
return ds


def _create_test_dataset_tas(year: int, offset: float = 0, quantity: str = "tasmax") -> xr.Dataset:
def _create_test_dataset_tas(
year: int, offset: float = 0, quantity: str = "tasmax"
) -> xr.Dataset:
"""Create test xarray Dataset.
Convention is that data is arranged in image-like way:
- dimensions are ('latitude', 'longitude')
Expand All @@ -97,7 +106,10 @@ def _create_test_dataset_tas(year: int, offset: float = 0, quantity: str = "tasm
Returns:
xr.Dataset : test Dataset
"""
temperature_t1 = np.array([[293.0, 298.0, 310.0], [304.0, 302.0, 300.0], [308.0, 290.0, 294.0]]) + offset
temperature_t1 = (
np.array([[293.0, 298.0, 310.0], [304.0, 302.0, 300.0], [308.0, 290.0, 294.0]])
+ offset
)
temperature_t2 = temperature_t1 + 1.0 # temp at t1 + 1 degree
temperature_t3 = temperature_t2 + 2.0
# stack to give 3 time points
Expand All @@ -116,7 +128,9 @@ def _create_test_dataset_tas(year: int, offset: float = 0, quantity: str = "tasm
return ds


def _create_test_dataset_hurs(year: int, offset: float = 0, quantity="hurs") -> xr.Dataset:
def _create_test_dataset_hurs(
year: int, offset: float = 0, quantity="hurs"
) -> xr.Dataset:
"""Create test xarray Dataset.
Convention is that data is arranged in image-like way:
- dimensions are ('latitude', 'longitude')
Expand All @@ -126,7 +140,9 @@ def _create_test_dataset_hurs(year: int, offset: float = 0, quantity="hurs") ->
Returns:
xr.Dataset : test Dataset
"""
hurs_t1 = np.array([[70.0, 72.0, 69.0], [72.0, 71.0, 70.0], [80.0, 74.0, 68.0]]) + offset
hurs_t1 = (
np.array([[70.0, 72.0, 69.0], [72.0, 71.0, 70.0], [80.0, 74.0, 68.0]]) + offset
)
hurs_t2 = hurs_t1 + 5.0 # hurs at t1 + 5%
hurs_t3 = hurs_t2 + 10.0
# stack to give 3 time points
Expand Down
34 changes: 24 additions & 10 deletions tests/test_drought_indicators.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,8 +9,12 @@
import zarr # type: ignore

from hazard.docs_store import DocStore
from hazard.models.drought_index import DroughtIndicator, LocalZarrWorkingStore, ProgressStore, S3ZarrWorkingStore
from tests.conftest import TestTarget
from hazard.models.drought_index import (
DroughtIndicator,
LocalZarrWorkingStore,
ProgressStore,
S3ZarrWorkingStore,
)


@pytest.mark.skip(reason="incomplete")
Expand All @@ -21,7 +25,9 @@ def test_spei_indicator(test_output_dir, s3_credentials):
working_store = S3ZarrWorkingStore()
working_store = LocalZarrWorkingStore(test_output_dir)
model = DroughtIndicator(working_store)
model.calculate_spei(gcm, scenario, progress_store=ProgressStore(test_output_dir, "spei_prog_store"))
model.calculate_spei(
gcm, scenario, progress_store=ProgressStore(test_output_dir, "spei_prog_store")
)
# target = TestTarget()
data_chunks = model.get_datachunks()
test_chunk = data_chunks["Chunk_0255"]
Expand All @@ -33,9 +39,9 @@ def test_spei_indicator(test_output_dir, s3_credentials):
# model.calculate_spei("MIROC6", "ssp585")
# model.calculate_annual_average_spei("MIROC6", "ssp585", 2080, target)

ds_tas = model.read_quantity_from_s3_store(gcm, scenario, "tas", lat_min, lat_max, lon_min, lon_max).chunk(
{"time": 100000}
)
ds_tas = model.read_quantity_from_s3_store(
gcm, scenario, "tas", lat_min, lat_max, lon_min, lon_max
).chunk({"time": 100000})
ds_tas_local = ds_tas.compute()
series_tas = ds_tas_local["tas"][0, 0, :].values
# ds_pr = model.read_quantity_from_s3_store(gcm, scenario, "pr", lat_min, lat_max, lon_min, lon_max).chunk(
Expand All @@ -50,7 +56,9 @@ def test_spei_indicator(test_output_dir, s3_credentials):


def test_partial_write_zarr(test_output_dir):
zarr_store = zarr.DirectoryStore(os.path.join(test_output_dir, "drought", "hazard.zarr"))
zarr_store = zarr.DirectoryStore(
os.path.join(test_output_dir, "drought", "hazard.zarr")
)

lat = np.arange(-60 + 0.25 / 2, 90 + 0.25 / 2, 0.25)
lon = np.arange(0.25 / 2, 360 + 0.25 / 2, 0.25)
Expand All @@ -70,8 +78,12 @@ def test_partial_write_zarr(test_output_dir):
ds_spei.to_zarr(store=zarr_store, mode="w", compute=False)
# see https://docs.xarray.dev/en/stable/user-guide/io.html?appending-to-existing-zarr-stores=#appending-to-existing-zarr-stores # noqa: E501
sliced = ds_spei.sel(lat=slice(10, 20), lon=slice(30, 40))
lat_indexes = np.where(np.logical_and(ds_spei["lat"].values >= 10, ds_spei["lat"].values <= 20))[0]
lon_indexes = np.where(np.logical_and(ds_spei["lon"].values >= 30, ds_spei["lon"].values <= 40))[0]
lat_indexes = np.where(
np.logical_and(ds_spei["lat"].values >= 10, ds_spei["lat"].values <= 20)
)[0]
lon_indexes = np.where(
np.logical_and(ds_spei["lon"].values >= 30, ds_spei["lon"].values <= 40)
)[0]
ds_spei_slice = (
xr.DataArray(
coords={
Expand Down Expand Up @@ -109,6 +121,8 @@ def test_progress_store(test_output_dir):
def test_doc_store(test_output_dir, s3_credentials):
docs_store = DocStore()
docs_store.write_new_empty_inventory()
zarr_store = zarr.DirectoryStore(os.path.join(test_output_dir, "drought", "hazard.zarr"))
zarr_store = zarr.DirectoryStore(
os.path.join(test_output_dir, "drought", "hazard.zarr")
)
resource = DroughtIndicator(zarr_store).resource
docs_store.update_inventory([resource])
90 changes: 65 additions & 25 deletions tests/test_heat_indicators.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,8 @@
import json
import os
from typing import Dict, List
from typing import List

import fsspec.implementations.local as local # type: ignore
import numpy as np
import pandas as pd # type: ignore
import pytest
import s3fs # type: ignore
import xarray as xr
Expand All @@ -13,13 +11,17 @@
from pytest import approx

import hazard.utilities.zarr_utilities as zarr_utilities
from hazard.docs_store import DocStore, HazardResources # type: ignore
from hazard.docs_store import DocStore # type: ignore
from hazard.models.degree_days import BatchItem, DegreeDays, HeatingCoolingDegreeDays
from hazard.models.work_loss import WorkLossIndicator
from hazard.protocols import OpenDataset, WriteDataset
from hazard.sources.nex_gddp_cmip6 import NexGddpCmip6
from hazard.sources.osc_zarr import OscZarr
from tests.conftest import TestSource, TestTarget, _create_test_dataset_averaged, _create_test_datasets_tas
from tests.conftest import (
TestSource,
TestTarget,
_create_test_dataset_averaged,
_create_test_datasets_tas,
)

from .conftest import _create_test_datasets_hurs, test_output_dir # noqa: F401; pylint: disable=unused-variable

Expand All @@ -32,16 +34,24 @@ def test_degree_days_mocked():
source = TestSource(_create_test_datasets_tas())
target = TestTarget()
# cut down the transform
model = DegreeDays(window_years=2, gcms=[gcm], scenarios=[scenario], central_years=[year])
model = DegreeDays(
window_years=2, gcms=[gcm], scenarios=[scenario], central_years=[year]
)
model.run_all(source, target, debug_mode=True)
with source.open_dataset_year(gcm, scenario, "tasmax", 2029) as y0:
with source.open_dataset_year(gcm, scenario, "tasmax", 2030) as y1:
scale = 365.0 / len(y0.time)
deg0 = scale * xr.where(y0.tasmax > (32 + 273.15), y0.tasmax - (32 + 273.15), 0).sum(dim=["time"])
deg1 = scale * xr.where(y1.tasmax > (32 + 273.15), y1.tasmax - (32 + 273.15), 0).sum(dim=["time"])
deg0 = scale * xr.where(
y0.tasmax > (32 + 273.15), y0.tasmax - (32 + 273.15), 0
).sum(dim=["time"])
deg1 = scale * xr.where(
y1.tasmax > (32 + 273.15), y1.tasmax - (32 + 273.15), 0
).sum(dim=["time"])
expected = (deg0 + deg1) / 2
assert expected.values == approx(
target.datasets["chronic_heat/osc/v2/mean_degree_days_v2_above_32c_NorESM2-MM_ssp585_2030"].values
target.datasets[
"chronic_heat/osc/v2/mean_degree_days_v2_above_32c_NorESM2-MM_ssp585_2030"
].values
)


Expand All @@ -58,7 +68,9 @@ def test_work_loss_mocked():
source = TestSource(test_sets)
target = TestTarget()
# cut down the transform
model = WorkLossIndicator(window_years=2, gcms=[gcm], scenarios=[scenario], central_years=[year])
model = WorkLossIndicator(
window_years=2, gcms=[gcm], scenarios=[scenario], central_years=[year]
)
model.run_all(source, target, debug_mode=True)
expected: List[xr.DataArray] = []
with source.open_dataset_year(gcm, scenario, "tas", 2029).tas as t0:
Expand All @@ -67,20 +79,30 @@ def test_work_loss_mocked():
with source.open_dataset_year(gcm, scenario, "hurs", 2030).hurs as h1:
tas_c = t0 - 273.15 # convert from K to C
# vpp is water vapour partial pressure in kPa
vpp = (h0 / 100.0) * 6.105 * np.exp((17.27 * tas_c) / (237.7 + tas_c))
vpp = (
(h0 / 100.0) * 6.105 * np.exp((17.27 * tas_c) / (237.7 + tas_c))
)
wbgt0 = 0.567 * tas_c + 0.393 * vpp + 3.94
tas_c = t1 - 273.15
vpp = (h1 / 100.0) * 6.105 * np.exp((17.27 * tas_c) / (237.7 + tas_c))
vpp = (
(h1 / 100.0) * 6.105 * np.exp((17.27 * tas_c) / (237.7 + tas_c))
)
wbgt1 = 0.567 * tas_c + 0.393 * vpp + 3.94

for alpha1, alpha2 in [alpha_light, alpha_medium, alpha_heavy]:
wa0 = 0.1 + 0.9 / (1.0 + (wbgt0 / alpha1) ** alpha2) # work-ability
wa0 = 0.1 + 0.9 / (
1.0 + (wbgt0 / alpha1) ** alpha2
) # work-ability
wlm0 = 1.0 - wa0.mean(dim=["time"]) # work-loss
wa1 = 0.1 + 0.9 / (1.0 + (wbgt1 / alpha1) ** alpha2) # work-ability
wa1 = 0.1 + 0.9 / (
1.0 + (wbgt1 / alpha1) ** alpha2
) # work-ability
wlm1 = 1.0 - wa1.mean(dim=["time"]) # work-loss
expected.append((wlm0 + wlm1) / 2)
assert expected[0].values == approx(
target.datasets["chronic_heat/osc/v2/mean_work_loss_low_NorESM2-MM_ssp585_2030"].values
target.datasets[
"chronic_heat/osc/v2/mean_work_loss_low_NorESM2-MM_ssp585_2030"
].values
)


Expand All @@ -89,7 +111,9 @@ def test_zarr_read_write(test_output_dir): # noqa: F811
read from the zarr array alone using attributes and ignoring coordinates.
"""
ds = _create_test_dataset_averaged()
store = zarr.DirectoryStore(os.path.join(test_output_dir, "hazard_test", "hazard.zarr"))
store = zarr.DirectoryStore(
os.path.join(test_output_dir, "hazard_test", "hazard.zarr")
)
source = OscZarr(store=store)
source.write("test", ds.tasmax)
# ds.to_zarr(store, compute=True, group="test", mode="w", consolidated=False)
Expand All @@ -106,12 +130,16 @@ def test_degree_days(test_output_dir): # noqa: F811
download_test_datasets(test_output_dir, gcm, scenario, years)
# source: read downloaded datasets from local file system
fs = local.LocalFileSystem()
source = NexGddpCmip6(root=os.path.join(test_output_dir, NexGddpCmip6.bucket), fs=fs)
source = NexGddpCmip6(
root=os.path.join(test_output_dir, NexGddpCmip6.bucket), fs=fs
)
# target: write zarr to load fine system
store = zarr.DirectoryStore(os.path.join(test_output_dir, "hazard", "hazard.zarr"))
target = OscZarr(store=store)
# cut down the model and run
model = HeatingCoolingDegreeDays(window_years=1, gcms=[gcm], scenarios=[scenario], central_years=[years[0]])
model = HeatingCoolingDegreeDays(
window_years=1, gcms=[gcm], scenarios=[scenario], central_years=[years[0]]
)
# model = DegreeDays(window_years=1, gcms=[gcm], scenarios=[scenario], central_years=[years[0]])
model.run_all(source, target)
# check one point...
Expand All @@ -136,15 +164,21 @@ def test_work_loss(test_output_dir): # noqa: F811
gcm = "NorESM2-MM"
scenario = "ssp585"
years = [2028, 2029, 2030]
download_test_datasets(test_output_dir, gcm, scenario, years, indicators=["tas", "hurs"])
download_test_datasets(
test_output_dir, gcm, scenario, years, indicators=["tas", "hurs"]
)
# source: read downloaded datasets from local file system
fs = local.LocalFileSystem()
source = NexGddpCmip6(root=os.path.join(test_output_dir, NexGddpCmip6.bucket), fs=fs)
source = NexGddpCmip6(
root=os.path.join(test_output_dir, NexGddpCmip6.bucket), fs=fs
)
# target: write zarr to load fine system
store = zarr.DirectoryStore(os.path.join(test_output_dir, "hazard", "hazard.zarr"))
target = OscZarr(store=store)
# cut down the model and run
model = WorkLossIndicator(window_years=3, gcms=[gcm], scenarios=[scenario], central_years=[years[1]])
model = WorkLossIndicator(
window_years=3, gcms=[gcm], scenarios=[scenario], central_years=[years[1]]
)
# resources = list(model.inventory())
# models = HazardResources(resources=resources)
# json_str = json.dumps(models.dict(), indent=4) # pretty print
Expand All @@ -170,9 +204,13 @@ def test_example_run_degree_days():
scenario = "ssp585"
year = 2030
source = NexGddpCmip6()
target = OscZarr(prefix="hazard_test") # test prefix is "hazard_test"; main one "hazard"
target = OscZarr(
prefix="hazard_test"
) # test prefix is "hazard_test"; main one "hazard"
# cut down the transform
model = DegreeDays(window_years=1, gcms=[gcm], scenarios=[scenario], central_years=[year])
model = DegreeDays(
window_years=1, gcms=[gcm], scenarios=[scenario], central_years=[year]
)

docs_store.update_inventory(model.inventory())

Expand All @@ -181,7 +219,9 @@ def test_example_run_degree_days():
assert True


def download_test_datasets(test_output_dir, gcm, scenario, years, indicators=["tasmax"]): # noqa: F811
def download_test_datasets(
test_output_dir, gcm, scenario, years, indicators=["tasmax"]
): # noqa: F811
store = NexGddpCmip6()
s3 = s3fs.S3FileSystem(anon=True)
for year in years:
Expand Down
4 changes: 1 addition & 3 deletions tests/test_inventory.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
import fsspec.implementations.local as local # type: ignore
import pytest

from hazard.docs_store import DocStore, HazardResources
from hazard.docs_store import DocStore
from hazard.models.days_tas_above import DaysTasAboveIndicator
from hazard.models.degree_days import DegreeDays, HeatingCoolingDegreeDays
from hazard.models.drought_index import DroughtIndicator
Expand All @@ -17,8 +17,6 @@
from hazard.onboard.wri_aqueduct_water_risk import WRIAqueductWaterRisk
from hazard.utilities import zarr_utilities # type: ignore

from .conftest import test_output_dir


def test_create_inventory(test_output_dir): # noqa: F811
"""Create inventory for all indicators and write into this repo."""
Expand Down
Loading

0 comments on commit f1b582c

Please sign in to comment.