Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add option to use eval to compute derived values for Catalog arrays. #173

Merged
merged 6 commits into from
Mar 6, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions CHANGELOG.rst
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ Dependency Change
-----------------

- Switched from cffi to pybind11 for the C++ bindings. (#155)
- If using fitsio, it now must be version > 1.0.6. (#173)


API Changes
Expand Down Expand Up @@ -59,6 +60,8 @@ New features
------------

- Give a better error message when patch is given as an integer, but npatch is not provided. (#150)
- Added ``x_eval``, ``y_eval``, etc. which let you calculate a derived quantity from an input
catalog using Python eval on columns in the file. (#151, #173)
- Allow vark, varg, varv for a Catalog be specifiable on input, rather than calculated directly
from the corresponding values. (#154)
- Allow numpy.random.Generator for rng arguments (in addition to legacy RandomState). (#157)
Expand Down
2 changes: 1 addition & 1 deletion test_requirements.txt
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
fitsio>=0.9
fitsio>=1.2
pandas>=0.20
scipy>=1.2
mockmpi>=0.8
101 changes: 53 additions & 48 deletions tests/test_catalog.py
Original file line number Diff line number Diff line change
Expand Up @@ -318,6 +318,23 @@ def test_ascii():
np.testing.assert_almost_equal(cat6b.ra, ra * (pi/12.))
np.testing.assert_almost_equal(cat6b.dec, dec * (pi/180.))

# Check using eval feature to apply the units rather than ra_units/dec_units
config_names['ra_units'] = 'rad'
config_names['dec_units'] = 'rad'
config_names['ra_eval'] = 'ra * np.pi/12.'
config_names['dec_eval'] = 'dec * math.pi/180.'
cat6c = treecorr.Catalog(file_name, config_names)
np.testing.assert_almost_equal(cat6c.ra, ra * (pi/12.))
np.testing.assert_almost_equal(cat6c.dec, dec * (pi/180.))

# Can also skip ra_col, dec_col and specify them in extra_cols.
del config_names['ra_col']
del config_names['dec_col']
config_names['extra_cols'] = ['ra', 'dec']
cat6d = treecorr.Catalog(file_name, config_names)
np.testing.assert_almost_equal(cat6c.ra, ra * (pi/12.))
np.testing.assert_almost_equal(cat6c.dec, dec * (pi/180.))

assert_raises(ValueError, treecorr.Catalog, file_name, config, ra_col=-1)
assert_raises(ValueError, treecorr.Catalog, file_name, config, ra_col=100)
assert_raises(ValueError, treecorr.Catalog, file_name, config, ra_col='invalid')
Expand Down Expand Up @@ -765,6 +782,42 @@ def _test_aardvark(filename, file_type, ext):
np.testing.assert_almost_equal(cat2.w[46392], 0.) # index = 1200379
np.testing.assert_almost_equal(cat2.w[46393], 0.9995946) # index = 1200386

# Test some eval calculations
center_ra = np.mean(cat1.ra) * coord.radians
center_dec = np.mean(cat1.dec) * coord.radians
center = coord.CelestialCoord(center_ra, center_dec)
center_repr = repr(center)
ntot = cat1.ntot
cat3 = treecorr.Catalog(
file_name,
# tangent plane projection for ra,dec -> x,y
x_eval=f"{center_repr}.project_rad(RA*math.pi/180, DEC*math.pi/180)[0]",
y_eval=f"{center_repr}.project_rad(RA*math.pi/180, DEC*math.pi/180)[1]",
# distortion rather than shear
g1_eval="GAMMA1 * 2 / (1 + GAMMA1**2 + GAMMA2**2)",
g2_eval="GAMMA2 * 2 / (1 + GAMMA1**2 + GAMMA2**2)",
# random spin-3 directions with mu for the amplitude. (This one is pretty contrived.)
t1_eval=f"KAPPA * np.random.default_rng(1234).normal(0,0.3,{ntot})",
t2_eval=f"KAPPA * np.random.default_rng(1235).normal(0,0.3,{ntot})",
extra_cols=['RA', 'DEC', 'GAMMA1', 'GAMMA2', 'KAPPA'])
print('made cat3')
print('cat3 = ',cat3)
print('x = ',cat3.x)
print('y = ',cat3.y)
print('g1 = ',cat3.g1)
print('g2 = ',cat3.g2)
print('t1 = ',cat3.t1)
print('t2 = ',cat3.t2)
x1, y1 = center.project_rad(cat1.ra, cat1.dec)
np.testing.assert_allclose(cat3.x, x1, atol=1.e-6)
np.testing.assert_allclose(cat3.y, y1, atol=1.e-6)
gsq = cat1.g1**2 + cat1.g2**2
np.testing.assert_allclose(cat3.g1, cat1.g1 * 2 / (1+gsq), rtol=1.e-6)
np.testing.assert_allclose(cat3.g2, cat1.g2 * 2 / (1+gsq), rtol=1.e-6)
np.testing.assert_allclose(np.mean(cat3.t1), 0., atol=1.e-3)
np.testing.assert_allclose(np.mean(cat3.t2), 0., atol=1.e-3)
np.testing.assert_allclose(np.std(cat3.t1), np.std(cat3.t2), rtol=0.1)

assert_raises(ValueError, treecorr.Catalog, file_name, config, x_col='invalid')
assert_raises(ValueError, treecorr.Catalog, file_name, config, y_col='invalid')
assert_raises(ValueError, treecorr.Catalog, file_name, config, z_col='invalid')
Expand Down Expand Up @@ -1193,54 +1246,6 @@ def test_ext():
t1_col='t1', t2_col='t2', q1_col='q1', q2_col='q2',
ext=1)

# test that the case where we can't slice works
# by pretending that we are using an old fitsio version,
# temporarily.
with mock.patch('fitsio.__version__', '1.0.6'):
cat11 = treecorr.Catalog(fname, allow_xyz=True,
x_col='x', y_col='y', z_col='z',
ra_col='ra', dec_col='dec', r_col='r',
ra_units='rad', dec_units='rad',
w_col='w', wpos_col='wpos', flag_col='flag',
k_col='k', g1_col='g1', g2_col='g2', v1_col='v1', v2_col='v2',
t1_col='t1', t2_col='t2', q1_col='q1', q2_col='q2',
x_ext=3, y_ext=3, z_ext=3,
ra_ext=4, dec_ext=4, r_ext=4,
w_ext=5, wpos_ext=5, flag_ext=5,
k_ext=6, g1_ext=6, g2_ext=6, v1_ext=6, v2_ext=6,
t1_ext=6, t2_ext=6, q1_ext=6, q2_ext=6)
assert cat11 == cat1
# and equiv for RA
cat12 = treecorr.Catalog(fname,
ra_col='ra', dec_col='dec', r_col='r',
ra_units='rad', dec_units='rad',
ext=4, ra_ext=4, last_row=120)
cat13 = treecorr.Catalog(fname,
ra_col='ra', dec_col='dec', r_col='r',
ra_units='rad', dec_units='rad',
ext=4, ra_ext=4, last_row=100)

np.testing.assert_allclose(cat12.ra[:99], cat13.ra[:99])
# and equiv for RA
cat14 = treecorr.Catalog(fname,
x_col='ra', y_col='dec',
x_units='rad', y_units='rad',
ext=4, ra_ext=4, last_row=120)
cat15 = treecorr.Catalog(fname,
x_col='ra', y_col='dec',
x_units='rad', y_units='rad',
ext=4, x_ext=4, last_row=100)

np.testing.assert_allclose(cat14.x[:99], cat15.x[:99])

cat16 = treecorr.Catalog(fname,
ra_col='ra', dec_col='dec',
ra_units='rad', dec_units='rad',
ext=4, last_row=100,
patch_col='patch', patch_ext='patch',
patch=0)

np.testing.assert_allclose(cat16.ra, cat4.ra[:100][::5])

@timer
def test_direct():
Expand Down
19 changes: 0 additions & 19 deletions tests/test_reader.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,9 +60,6 @@ def test_fits_reader():
# Default ext is "in" reader
assert 1 in r

# Probably can slice, but depends on installed fitsio version
assert r.can_slice == (fitsio.__version__ > '1.0.6')

s = slice(0, 10, 2)
for ext in [1, 'AARDWOLF']:
data = r.read(['RA'], s, ext=ext)
Expand Down Expand Up @@ -160,13 +157,6 @@ def test_fits_reader():
with assert_raises(RuntimeError):
w.write(['KAPPA', 'MU'], [kappa, mu], params={'test': True}, ext='KM')

# Regardless of the system's fitsio version, check the two cases in code.
with mock.patch('fitsio.__version__', '1.0.6'):
with FitsReader(os.path.join('data','Aardvark.fit')) as r:
assert not r.can_slice
with mock.patch('fitsio.__version__', '1.1.0'):
with FitsReader(os.path.join('data','Aardvark.fit')) as r:
assert r.can_slice
with mock.patch.dict(sys.modules, {'fitsio':None}):
with CaptureLog() as cl:
with assert_raises(ImportError):
Expand Down Expand Up @@ -217,9 +207,6 @@ def test_hdf_reader():
# Default ext is "in" reader
assert '/' in r

# Can always slice
assert r.can_slice

s = slice(0, 10, 2)
data = r.read(['RA'], s)
dec = r.read('DEC', s)
Expand Down Expand Up @@ -329,9 +316,6 @@ def test_parquet_reader():
# Default ext is "in" reader
assert None in r

# Can always slice
assert r.can_slice

s = slice(0, 10, 2)
data = r.read(['RA'], s)
dec = r.read('DEC', s)
Expand Down Expand Up @@ -393,9 +377,6 @@ def _test_ascii_reader(r, has_names=True):
# Default ext is "in" reader
assert None in r

# Can always slice
assert r.can_slice

# cols are: ra, dec, x, y, k, g1, g2, w, z, r, wpos, flag
s = slice(0, 10, 2)
data = r.read([1,3,9], s)
Expand Down
Loading