Skip to content

Commit

Permalink
Update test_eva assertions and update fit_gev return type
Browse files Browse the repository at this point in the history
  • Loading branch information
@stellema
stellema committed Nov 16, 2023
1 parent dfc770a commit c249008
Show file tree
Hide file tree
Showing 2 changed files with 55 additions and 56 deletions.
28 changes: 23 additions & 5 deletions unseen/general_utils.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -182,7 +182,9 @@ def _goodness_of_fit(data, theta):
scale = scale + scale1 * t

res = goodness_of_fit(
genextreme, test_data, known_params=dict(c=shape, loc=loc, scale=scale)
genextreme,
test_data,
known_params=dict(c=shape, loc=loc, scale=scale),
)
return res.pvalue

Expand Down Expand Up @@ -394,14 +396,27 @@ def _fit(
theta_i = -shape, loc, loc1, scale, scale1

# Optimisation bounds (scale parameter must be non-negative).
bounds = [(None, None), (None, None), (None, None), (0, None), (None, None)]
bounds = [
(None, None),
(None, None),
(None, None),
(0, None),
(None, None),
]

# Minimise the negative log-likelihood function to get optimal theta.
res = minimize(nllf, theta_i, args=(data, x), method=method, bounds=bounds)
res = minimize(
nllf,
theta_i,
args=(data, x),
method=method,
bounds=bounds,
)
theta = res.x

# Restore 'shape' sign for consistency with scipy.stats results.
theta[0] *= -1
theta = np.array([i for i in theta], dtype="float64")
return theta

def fit(data, **kwargs):
Expand Down Expand Up @@ -450,7 +465,7 @@ def fit(data, **kwargs):

# Return a tuple of scalars instead of a 1D data array
if len(data.shape) == 1:
theta = tuple([i.item() for i in theta])
theta = np.array([i for i in theta], dtype="float64")
return theta


Expand Down Expand Up @@ -597,7 +612,10 @@ def plot_gev_return_curve(
) = event_data

ax.plot(
curve_return_periods, curve_values, color="tab:blue", label="GEV fit to data"
curve_return_periods,
curve_values,
color="tab:blue",
label="GEV fit to data",
)
ax.fill_between(
curve_return_periods,
Expand Down
83 changes: 32 additions & 51 deletions unseen/tests/test_eva.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,12 +8,12 @@

rtol = 0.3 # relative tolerance
alpha = 0.05
np.random.seed(0)
np.random.seed(1)


def example_da_gev_1d():
"""An example 1D GEV DataArray and its distribution parameters."""
time = np.arange("2000-01-01", "2002-01-01", dtype=np.datetime64)
time = np.arange("2000-01-01", "2005-01-01", dtype=np.datetime64)

# Generate shape, location and scale parameters.
shape = np.random.uniform()
Expand All @@ -35,7 +35,7 @@ def example_da_gev_1d_dask():

def example_da_gev_3d():
"""An example 3D GEV DataArray and its distribution parameters."""
time = np.arange("2000-01-01", "2002-01-01", dtype=np.datetime64)
time = np.arange("2000-01-01", "2005-01-01", dtype=np.datetime64)
lats = np.arange(0, 2)
lons = np.arange(0, 2)
shape = (len(lats), len(lons))
Expand All @@ -46,7 +46,11 @@ def example_da_gev_3d():
theta = np.stack([c, loc, scale], axis=-1)

rvs = genextreme.rvs(
c, loc=loc, scale=scale, size=(time.size, *shape), random_state=0
c,
loc=loc,
scale=scale,
size=(time.size, *shape),
random_state=0,
)
data = xr.DataArray(rvs, coords=[time, lats, lons], dims=["time", "lat", "lon"])
return data, theta
Expand All @@ -65,6 +69,10 @@ def add_example_gev_trend(data):
return data + trend


data, theta_i = example_da_gev_3d()
theta = fit_gev(data, stationary=True, check_fit=False)


def test_fit_gev_1d():
"""Run stationary fit using 1D array & check results."""
data, theta_i = example_da_gev_1d()
Expand All @@ -89,92 +97,65 @@ def test_fit_gev_3d():
npt.assert_allclose(theta, theta_i, rtol=rtol)


# def test_fit_gev_3d_dask():
# """Run stationary fit using 3D dask array & check results."""
# data, theta_i = example_da_gev_3d_dask()
# theta = fit_gev(data, stationary=True, check_fit=False)
# # Check fitted params match params used to create data.
# npt.assert_allclose(theta, theta_i, rtol=rtol)
def test_fit_gev_3d_dask():
"""Run stationary fit using 3D dask array & check results."""
data, theta_i = example_da_gev_3d_dask()
theta = fit_gev(data, stationary=True, check_fit=False)
# Check fitted params match params used to create data.
npt.assert_allclose(theta, theta_i, rtol=rtol)


def test_fit_ns_gev_1d():
"""Run stationary fit using 1D array & check results."""
data, theta_i = example_da_gev_1d()
data, _ = example_da_gev_1d()
data = add_example_gev_trend(data)

# Check function runs.
theta = fit_gev(data, stationary=False, check_fit=False)

check_gev_fit(data, theta, time_dim="time")

# Check fitted params match params used to create data.
shape, loc, loc1, scale, scale1 = theta
npt.assert_allclose((shape, loc, scale), theta_i, rtol=rtol)

# Check it detected a positive treand
assert loc1 > 0
assert scale1 > 0
pvalue = check_gev_fit(data, theta, time_dim="time")
assert np.all(pvalue) > alpha


def test_fit_ns_gev_1d_dask():
"""Run stationary fit using 1D dask array & check results."""
data, theta_i = example_da_gev_1d_dask()
data, _ = example_da_gev_1d_dask()

# Add a positive linear trend.
data = add_example_gev_trend(data)

# Check function runs.
theta = fit_gev(data, stationary=False, check_fit=False)

# Check fitted params match params used to create data.
shape, loc, loc1, scale, scale1 = theta
npt.assert_allclose((shape, loc, scale), theta_i, rtol=rtol)

# Check it fitted a positive trend.
assert np.all(loc1) > 0
assert np.all(scale1) > 0

pvalue = check_gev_fit(data, theta, time_dim="time")
assert np.all(pvalue) > alpha


def test_fit_ns_gev_3d():
"""Run stationary fit using 3D array & check results."""
data, theta_i = example_da_gev_3d()
data, _ = example_da_gev_3d()

# Add a positive linear trend.
data = add_example_gev_trend(data)

# Check function runs.
theta = fit_gev(data, stationary=False, check_fit=False)

# Check fitted params match params used to create data.
npt.assert_allclose(theta.isel(theta=[0, 1, 3]), theta_i, rtol=rtol)

# Check it fitted a positive trend.
assert np.all(theta.isel(theta=2)) > 0
assert np.all(theta.isel(theta=4)) > 0

pvalue = check_gev_fit(data, theta, time_dim="time")
assert np.all(pvalue) > alpha


# def test_fit_gev_3d_dask():
# """Run stationary fit using 3D dask array & check results."""
# data, theta_i = example_da_gev_3d_dask()

# # Add a positive linear trend.
# data = add_example_gev_trend(data)

# # Check function runs.
# theta = fit_gev(data, stationary=False, check_fit=False)
def test_fit_ns_gev_3d_dask():
"""Run stationary fit using 3D dask array & check results."""
data, _ = example_da_gev_3d_dask()

# # Check fitted params match params used to create data.
# npt.assert_allclose(theta.isel(theta=[0, 1, 3]), theta_i, rtol=rtol)
# Add a positive linear trend.
data = add_example_gev_trend(data)

# # Check it fitted a positive trend.
# assert np.all(theta.isel(theta=2)) > 0
# assert np.all(theta.isel(theta=4)) > 0
# Check function runs.
theta = fit_gev(data, stationary=False, check_fit=False)

# pvalue = check_gev_fit(data, theta, time_dim="time")
# assert np.all(pvalue) > alpha
pvalue = check_gev_fit(data, theta, time_dim="time")
assert np.all(pvalue) > alpha

0 comments on commit c249008

Please sign in to comment.