test: fix redshift bounds in tests

src/py21cmfast_tools/lc2ps.py

@@ -120,8 +120,6 @@ def calculate_ps( # noqa: C901
  # Infer HII_DIM from lc side shape
  if box_side_shape is None:
  box_side_shape = lc.shape[0]
- if get_variance and postprocess:
- raise NotImplementedError("Cannot get variance and rebin it properly.")
  if get_variance and interp is not None:
  raise NotImplementedError("Cannot get variance while interpolating.")
  if zs is None:
@@ -473,6 +471,12 @@ def postprocess_ps(
  log_kpar[crop[2] : crop[3]],
  nmodes,
  )
+ else:
+ return (
+ rebinned_ps[crop[0] : crop[1]][:, crop[2] : crop[3]],
+ kperp[crop[0] : crop[1]],
+ log_kpar[crop[2] : crop[3]],
+ )
  else:
  return (
  rebinned_ps[crop[0] : crop[1]][:, crop[2] : crop[3]],

tests/test_ps.py

@@ -8,7 +8,7 @@ def test_calculate_ps():
  rng = np.random.default_rng()
  test_lc = rng.random((100, 100, 1000))
  test_redshifts = np.logspace(np.log10(5), np.log10(30), 1000)
- zs = [5.0, 6.0, 10.0, 27.0]
+ zs = [6.0, 10.0, 27.0]
 
  calculate_ps(
  test_lc,
@@ -86,17 +86,18 @@ def test_calculate_ps_w_var():
  )
  out["full_var_2D"]
  out["var_1D"]
- with np.testing.assert_raises(NotImplementedError):
- calculate_ps(
- test_lc,
- test_redshifts,
- box_length=200,
- zs=zs,
- calc_1d=True,
- calc_global=True,
- get_variance=True,
- postprocess=True,
- )
+ out = calculate_ps(
+ test_lc,
+ test_redshifts,
+ box_length=200,
+ zs=zs,
+ calc_1d=True,
+ calc_2d=True,
+ calc_global=True,
+ get_variance=True,
+ postprocess=True,
+ )
+ out["final_var_2D"]
  with np.testing.assert_raises(NotImplementedError):
  calculate_ps(
  test_lc,