| 4581 | @pytest.mark.parametrize("use_datetime", [True, False]) |
| 4582 | @pytest.mark.filterwarnings("ignore:overflow encountered in multiply") |
| 4583 | def test_polyfit(self, use_dask, use_datetime) -> None: |
| 4584 | if use_dask and not has_dask: |
| 4585 | pytest.skip("requires dask") |
| 4586 | xcoord = xr.DataArray( |
| 4587 | pd.date_range("1970-01-01", freq="D", periods=10), dims=("x",), name="x" |
| 4588 | ) |
| 4589 | x = xr.core.missing.get_clean_interp_index(xcoord, "x") |
| 4590 | if not use_datetime: |
| 4591 | xcoord = x |
| 4592 | |
| 4593 | da_raw = DataArray( |
| 4594 | np.stack((10 + 1e-15 * x + 2e-28 * x**2, 30 + 2e-14 * x + 1e-29 * x**2)), |
| 4595 | dims=("d", "x"), |
| 4596 | coords={"x": xcoord, "d": [0, 1]}, |
| 4597 | ) |
| 4598 | |
| 4599 | if use_dask: |
| 4600 | da = da_raw.chunk({"d": 1}) |
| 4601 | else: |
| 4602 | da = da_raw |
| 4603 | |
| 4604 | out = da.polyfit("x", 2) |
| 4605 | expected = DataArray( |
| 4606 | [[2e-28, 1e-15, 10], [1e-29, 2e-14, 30]], |
| 4607 | dims=("d", "degree"), |
| 4608 | coords={"degree": [2, 1, 0], "d": [0, 1]}, |
| 4609 | ).T |
| 4610 | assert_allclose(out.polyfit_coefficients, expected, rtol=1e-3) |
| 4611 | |
| 4612 | # Full output and deficient rank |
| 4613 | with warnings.catch_warnings(): |
| 4614 | warnings.simplefilter("ignore", RankWarning) |
| 4615 | out = da.polyfit("x", 12, full=True) |
| 4616 | assert out.polyfit_residuals.isnull().all() |
| 4617 | |
| 4618 | # With NaN |
| 4619 | da_raw[0, 1:3] = np.nan |
| 4620 | if use_dask: |
| 4621 | da = da_raw.chunk({"d": 1}) |
| 4622 | else: |
| 4623 | da = da_raw |
| 4624 | out = da.polyfit("x", 2, skipna=True, cov=True) |
| 4625 | assert_allclose(out.polyfit_coefficients, expected, rtol=1e-3) |
| 4626 | assert "polyfit_covariance" in out |
| 4627 | |
| 4628 | # Skipna + Full output |
| 4629 | out = da.polyfit("x", 2, skipna=True, full=True) |
| 4630 | assert_allclose(out.polyfit_coefficients, expected, rtol=1e-3) |
| 4631 | assert out.x_matrix_rank == 3 |
| 4632 | np.testing.assert_almost_equal(out.polyfit_residuals, [0, 0]) |
| 4633 | |
| 4634 | with warnings.catch_warnings(): |
| 4635 | warnings.simplefilter("ignore", RankWarning) |
| 4636 | out = da.polyfit("x", 8, full=True) |
| 4637 | np.testing.assert_array_equal(out.polyfit_residuals.isnull(), [True, False]) |
| 4638 | |
| 4639 | @requires_dask |
| 4640 | def test_polyfit_nd_dask(self) -> None: |