def test_pairwise(): # Test the pairwise option. ngal = 1000 s = 10. np.random.seed(8675309) x1 = np.random.normal(0,s, (ngal,) ) y1 = np.random.normal(0,s, (ngal,) ) w1 = np.random.random(ngal) x2 = np.random.normal(0,s, (ngal,) ) y2 = np.random.normal(0,s, (ngal,) ) w2 = np.random.random(ngal) k2 = np.random.normal(0,3, (ngal,) ) w1 = np.ones_like(w1) w2 = np.ones_like(w2) cat1 = treecorr.Catalog(x=x1, y=y1, w=w1) cat2 = treecorr.Catalog(x=x2, y=y2, w=w2, k=k2) min_sep = 5. max_sep = 50. nbins = 10 bin_size = np.log(max_sep/min_sep) / nbins nk = treecorr.NKCorrelation(min_sep=min_sep, max_sep=max_sep, nbins=nbins) nk.process_pairwise(cat1, cat2) nk.finalize(cat2.vark) true_npairs = np.zeros(nbins, dtype=int) true_weight = np.zeros(nbins, dtype=float) true_xi = np.zeros(nbins, dtype=float) rsq = (x1-x2)**2 + (y1-y2)**2 r = np.sqrt(rsq) logr = np.log(r) ww = w1 * w2 xi = ww * k2 index = np.floor(np.log(r/min_sep) / bin_size).astype(int) mask = (index >= 0) & (index < nbins) np.add.at(true_npairs, index[mask], 1) np.add.at(true_weight, index[mask], ww[mask]) np.add.at(true_xi, index[mask], xi[mask]) true_xi /= true_weight np.testing.assert_array_equal(nk.npairs, true_npairs) np.testing.assert_allclose(nk.weight, true_weight, rtol=1.e-5, atol=1.e-8) np.testing.assert_allclose(nk.xi, true_xi, rtol=1.e-4, atol=1.e-8) # If cats have names, then the logger will mention them. # Also, test running with optional args. cat1.name = "first" cat2.name = "second" with CaptureLog() as cl: nk.logger = cl.logger nk.process_pairwise(cat1, cat2, metric='Euclidean', num_threads=2) assert "for cats first, second" in cl.output
def test_single(): # Use kappa(r) = kappa0 exp(-r^2/2r0^2) (1-r^2/2r0^2) around a single lens nsource = 100000 kappa0 = 0.05 r0 = 10. L = 5. * r0 np.random.seed(8675309) x = (np.random.random_sample(nsource)-0.5) * L y = (np.random.random_sample(nsource)-0.5) * L r2 = (x**2 + y**2) k = kappa0 * np.exp(-0.5*r2/r0**2) * (1.-0.5*r2/r0**2) lens_cat = treecorr.Catalog(x=[0], y=[0], x_units='arcmin', y_units='arcmin') source_cat = treecorr.Catalog(x=x, y=y, k=k, x_units='arcmin', y_units='arcmin') nk = treecorr.NKCorrelation(bin_size=0.1, min_sep=1., max_sep=25., sep_units='arcmin', verbose=1) nk.process(lens_cat, source_cat) r = nk.meanr true_k = kappa0 * np.exp(-0.5*r**2/r0**2) * (1.-0.5*r**2/r0**2) print('nk.xi = ',nk.xi) print('true_kappa = ',true_k) print('ratio = ',nk.xi / true_k) print('diff = ',nk.xi - true_k) print('max diff = ',max(abs(nk.xi - true_k))) # Note: there is a zero crossing, so need to include atol as well as rtol np.testing.assert_allclose(nk.xi, true_k, rtol=1.e-2, atol=1.e-4) # Check that we get the same result using the corr2 function lens_cat.write(os.path.join('data','nk_single_lens.dat')) source_cat.write(os.path.join('data','nk_single_source.dat')) config = treecorr.read_config('configs/nk_single.yaml') config['verbose'] = 0 treecorr.corr2(config) corr2_output = np.genfromtxt(os.path.join('output','nk_single.out'), names=True, skip_header=1) print('nk.xi = ',nk.xi) print('from corr2 output = ',corr2_output['kappa']) print('ratio = ',corr2_output['kappa']/nk.xi) print('diff = ',corr2_output['kappa']-nk.xi) np.testing.assert_allclose(corr2_output['kappa'], nk.xi, rtol=1.e-3) # There is special handling for single-row catalogs when using np.genfromtxt rather # than pandas. So mock it up to make sure we test it. if sys.version_info < (3,): return # mock only available on python 3 from unittest import mock with mock.patch.dict(sys.modules, {'pandas':None}): with CaptureLog() as cl: treecorr.corr2(config, logger=cl.logger) assert "Unable to import pandas" in cl.output corr2_output = np.genfromtxt(os.path.join('output','nk_single.out'), names=True, skip_header=1) np.testing.assert_allclose(corr2_output['kappa'], nk.xi, rtol=1.e-3)
def test_omp(): """Test setting the number of omp threads. """ import multiprocessing # If num_threads <= 0 or None, get num from cpu_count cpus = multiprocessing.cpu_count() assert treecorr.set_omp_threads(0) > 0 assert treecorr.set_omp_threads(0) <= cpus assert treecorr.set_omp_threads(None) > 0 assert treecorr.set_omp_threads(None) <= cpus # If num_threads == 1, it should always set to 1 assert treecorr.set_omp_threads(1) == 1 # If num_threads > 1, it could be 1 or up to the input num_threads assert treecorr.set_omp_threads(2) >= 1 assert treecorr.set_omp_threads(2) <= 2 assert treecorr.set_omp_threads(20) >= 1 assert treecorr.set_omp_threads(20) <= 20 # Repeat and check that appropriate messages are emitted with CaptureLog() as cl: num_threads = treecorr.set_omp_threads(0, logger=cl.logger) assert "multiprocessing.cpu_count() = " in cl.output assert "Telling OpenMP to use %s threads" % cpus in cl.output with CaptureLog() as cl: treecorr.set_omp_threads(None, logger=cl.logger) assert "multiprocessing.cpu_count() = " in cl.output assert "Telling OpenMP to use %s threads" % cpus in cl.output with CaptureLog() as cl: treecorr.set_omp_threads(1, logger=cl.logger) assert "multiprocessing.cpu_count() = " not in cl.output assert "Telling OpenMP to use 1 threads" in cl.output assert "Using %s threads" % num_threads not in cl.output assert "Unable to use multiple threads" not in cl.output with CaptureLog() as cl: treecorr.set_omp_threads(2, logger=cl.logger) assert "multiprocessing.cpu_count() = " not in cl.output assert "Telling OpenMP to use 2 threads" in cl.output # It's hard to tell what happens in the next step, since we can't control what # treecorr._lib.SetOMPThreads does. It depends on whether OpenMP is enabled and # how many cores are available. So let's mock it up. if sys.version_info < (3, ): return # mock only available on python 3 from unittest import mock with mock.patch('treecorr._lib') as _lib: # First mock with OpenMP enables and able to use lots of threads _lib.SetOMPThreads = lambda x: x assert treecorr.set_omp_threads(20) == 20 with CaptureLog() as cl: treecorr.set_omp_threads(20, logger=cl.logger) assert "OpenMP reports that it will use 20 threads" in cl.output assert "Using 20 threads" in cl.output # Next only 4 threads available _lib.SetOMPThreads = lambda x: 4 if x > 4 else x assert treecorr.set_omp_threads(20) == 4 with CaptureLog() as cl: treecorr.set_omp_threads(20, logger=cl.logger) assert "OpenMP reports that it will use 4 threads" in cl.output assert "Using 4 threads" in cl.output assert treecorr.set_omp_threads(2) == 2 with CaptureLog() as cl: treecorr.set_omp_threads(2, logger=cl.logger) assert "OpenMP reports that it will use 2 threads" in cl.output # Finally, no OpenMP _lib.SetOMPThreads = lambda x: 1 assert treecorr.set_omp_threads(20) == 1 with CaptureLog() as cl: treecorr.set_omp_threads(20, logger=cl.logger) assert "OpenMP reports that it will use 1 threads" in cl.output assert "Unable to use multiple threads" in cl.output
def test_check(): """Test checking the validity of config values. """ # First a simple case with no conflicts config1 = treecorr.read_config('configs/kg.yaml') valid_params = treecorr.corr2_valid_params config2 = treecorr.config.check_config(config1.copy(), valid_params) # Just check a few values assert config2['x_col'] == '1' assert config2['k_col'] == ['3', '0'] assert config2['verbose'] == 1 assert config2['kg_file_name'] == 'output/kg.out' # Will also have other parameters filled from the valid_params dict for key in config2: assert key in valid_params if key in config1: if isinstance(config1[key], list): assert [str(v) for v in config2[key] ] == [str(v) for v in config1[key]] else: assert config2[key] == config1[key] or str( config2[key]) == str(config1[key]) else: assert config2[key] == valid_params[key][2] # Check list of bool config1['flip_g1'] = [True, 0] config2 = treecorr.config.check_config(config1.copy(), valid_params) assert config2['flip_g1'] == [True, False] # Longer names are allowed config1['x_units'] = 'arcminutes' config1['y_units'] = 'arcminute' config2 = treecorr.config.check_config(config1.copy(), valid_params) assert config2['x_units'] == 'arcmin' assert config2['y_units'] == 'arcmin' # Also other aliases, but you need to list them explicitly. config1['reverse_g1'] = True with assert_raises(TypeError): treecorr.config.check_config(config1.copy(), valid_params) config2 = treecorr.config.check_config(config1.copy(), valid_params, aliases={'reverse_g1': 'flip_g1'}) assert config2['flip_g1'] == True assert 'reverse_g1' not in config2 del config1['reverse_g1'] # Invalid values raise errors config1['verbose'] = -1 with assert_raises(ValueError): treecorr.config.check_config(config1.copy(), valid_params) config1['verbose'] = 1 config1['metric'] = 'hyperbolic' with assert_raises(ValueError): treecorr.config.check_config(config1.copy(), valid_params) del config1['metric'] # With a logger, aliases emit a warning. config1['n2_file_name'] = 'output/n2.out' with CaptureLog() as cl: config2 = treecorr.config.check_config( config1.copy(), valid_params, logger=cl.logger, aliases={'n2_file_name': 'nn_file_name'}) assert "The parameter n2_file_name is deprecated." in cl.output assert "You should use nn_file_name instead." in cl.output # corr2 has a list of standard aliases # It is currently empty, but let's mock it up to test the functionality. if sys.version_info < (3, ): return # mock only available on python 3 from unittest import mock with mock.patch('treecorr.corr2_aliases', {'n2_file_name': 'nn_file_name'}): config2 = treecorr.config.check_config(config1.copy(), valid_params, aliases=treecorr.corr2_aliases) assert 'n2_file_name' not in config2 assert config2['nn_file_name'] == 'output/n2.out' del config1['n2_file_name']
def test_pairwise(): # Test the pairwise option. ngal = 1000 s = 10. rng = np.random.RandomState(8675309) x1 = rng.normal(0,s, (ngal,) ) y1 = rng.normal(0,s, (ngal,) ) w1 = rng.random_sample(ngal) k1 = rng.normal(5,1, (ngal,) ) x2 = rng.normal(0,s, (ngal,) ) y2 = rng.normal(0,s, (ngal,) ) w2 = rng.random_sample(ngal) g12 = rng.normal(0,0.2, (ngal,) ) g22 = rng.normal(0,0.2, (ngal,) ) w1 = np.ones_like(w1) w2 = np.ones_like(w2) cat1 = treecorr.Catalog(x=x1, y=y1, w=w1, k=k1) cat2 = treecorr.Catalog(x=x2, y=y2, w=w2, g1=g12, g2=g22) min_sep = 5. max_sep = 50. nbins = 10 bin_size = np.log(max_sep/min_sep) / nbins kg = treecorr.KGCorrelation(min_sep=min_sep, max_sep=max_sep, nbins=nbins) kg.process_pairwise(cat1, cat2) kg.finalize(cat1.vark, cat2.varg) true_npairs = np.zeros(nbins, dtype=int) true_weight = np.zeros(nbins, dtype=float) true_xi = np.zeros(nbins, dtype=complex) rsq = (x1-x2)**2 + (y1-y2)**2 r = np.sqrt(rsq) logr = np.log(r) expmialpha = ((x1-x2) - 1j*(y1-y2)) / r ww = w1 * w2 xi = -ww * k1 * (g12 + 1j*g22) * expmialpha**2 index = np.floor(np.log(r/min_sep) / bin_size).astype(int) mask = (index >= 0) & (index < nbins) np.add.at(true_npairs, index[mask], 1) np.add.at(true_weight, index[mask], ww[mask]) np.add.at(true_xi, index[mask], xi[mask]) true_xi /= true_weight np.testing.assert_array_equal(kg.npairs, true_npairs) np.testing.assert_allclose(kg.weight, true_weight, rtol=1.e-5, atol=1.e-8) np.testing.assert_allclose(kg.xi, true_xi.real, rtol=1.e-4, atol=1.e-8) np.testing.assert_allclose(kg.xi_im, true_xi.imag, rtol=1.e-4, atol=1.e-8) # If cats have names, then the logger will mention them. # Also, test running with optional args. cat1.name = "first" cat2.name = "second" with CaptureLog() as cl: kg.logger = cl.logger kg.process_pairwise(cat1, cat2, metric='Euclidean', num_threads=2) assert "for cats first, second" in cl.output
def test_sample_pairs(): nobj = 10000 rng = np.random.RandomState(8675309) x1 = rng.random_sample(nobj) # All from 0..1 y1 = rng.random_sample(nobj) z1 = rng.random_sample(nobj) w1 = rng.random_sample(nobj) use = rng.randint(30, size=nobj).astype(float) w1[use == 0] = 0 g11 = rng.random_sample(nobj) g21 = rng.random_sample(nobj) k1 = rng.random_sample(nobj) x2 = rng.random_sample(nobj) # All from 0..1 y2 = rng.random_sample(nobj) z2 = rng.random_sample(nobj) w2 = rng.random_sample(nobj) use = rng.randint(30, size=nobj).astype(float) w2[use == 0] = 0 g12 = rng.random_sample(nobj) g22 = rng.random_sample(nobj) k2 = rng.random_sample(nobj) # Start with flat coords cat1 = treecorr.Catalog(x=x1, y=y1, w=w1, g1=g11, g2=g21, k=k1, keep_zero_weight=True) cat2 = treecorr.Catalog(x=x2, y=y2, w=w2, g1=g12, g2=g22, k=k2, keep_zero_weight=True) # Note: extend range low enough that some bins have < 100 pairs. nn = treecorr.NNCorrelation(min_sep=0.001, max_sep=0.01, bin_size=0.1, max_top=0) nn.process(cat1, cat2) print('rnom = ',nn.rnom) print('npairs = ',nn.npairs.astype(int)) # Start with a bin near the bottom with < 100 pairs # This only exercises case 1 in the sampleFrom function. b = 1 i1, i2, sep = nn.sample_pairs(100, cat1, cat2, min_sep=nn.left_edges[b], max_sep=nn.right_edges[b]) print('i1 = ',i1) print('i2 = ',i2) print('sep = ',sep) assert nn.npairs[b] <= 100 # i.e. make sure these next tests are what we want to do. assert len(i1) == nn.npairs[b] assert len(i2) == nn.npairs[b] assert len(sep) == nn.npairs[b] actual_sep = ((x1[i1]-x2[i2])**2 + (y1[i1]-y2[i2])**2)**0.5 np.testing.assert_allclose(sep, actual_sep, rtol=0.1) # half bin size with slop. np.testing.assert_array_less(sep, nn.right_edges[b]) np.testing.assert_array_less(nn.left_edges[b], sep) # Next one that still isn't too many pairs, but more than 100 # This exercises cases 1,2 in the sampleFrom function. b = 10 i1, i2, sep = nn.sample_pairs(100, cat1, cat2, min_sep=nn.left_edges[b], max_sep=nn.right_edges[b]) print('i1 = ',i1) print('i2 = ',i2) print('sep = ',sep) assert nn.npairs[b] > 100 assert len(i1) == 100 assert len(i2) == 100 assert len(sep) == 100 actual_sep = ((x1[i1]-x2[i2])**2 + (y1[i1]-y2[i2])**2)**0.5 np.testing.assert_allclose(sep, actual_sep, rtol=0.1) np.testing.assert_array_less(sep, nn.right_edges[b]) np.testing.assert_array_less(nn.left_edges[b], sep) # To exercise case 3, we need to go to larger separations, so the recursion # more often stops before getting to the leaves. # Also switch to 3d coordinates. cat1 = treecorr.Catalog(x=x1, y=y1, z=z1, w=w1, g1=g11, g2=g21, k=k1, keep_zero_weight=True) cat2 = treecorr.Catalog(x=x2, y=y2, z=z2, w=w2, g1=g12, g2=g22, k=k2, keep_zero_weight=True) gg = treecorr.GGCorrelation(min_sep=0.4, nbins=10, bin_size=0.1, max_top=0) gg.process(cat1, cat2) print('rnom = ',gg.rnom) print('npairs = ',gg.npairs.astype(int)) for b in [0,5]: i1, i2, sep = gg.sample_pairs(100, cat1, cat2, min_sep=gg.left_edges[b], max_sep=gg.right_edges[b]) print('len(npairs) = ',len(gg.npairs)) print('npairs = ',gg.npairs) print('i1 = ',i1) print('i2 = ',i2) print('sep = ',sep) assert len(i1) == 100 assert len(i2) == 100 assert len(sep) == 100 actual_sep = ((x1[i1]-x2[i2])**2 + (y1[i1]-y2[i2])**2 + (z1[i1]-z2[i2])**2)**0.5 np.testing.assert_allclose(sep, actual_sep, rtol=0.2) np.testing.assert_array_less(sep, gg.right_edges[b]) np.testing.assert_array_less(gg.left_edges[b], sep) # Check a different metric. # Also ability to generate the field automatically. cat1.clear_cache() # Clears the previously made cat1.field cat2.clear_cache() # and cat2.field b = 3 with CaptureLog() as cl: nk = treecorr.NKCorrelation(min_sep=0.4, max_sep=1.0, bin_size=0.1, max_top=0, logger=cl.logger) i1, i2, sep = nk.sample_pairs(100, cat1, cat2, metric='Arc', min_sep=nk.left_edges[b], max_sep=nk.right_edges[b]) print(cl.output) nk.process(cat1, cat2, metric='Arc') print('len(npairs) = ',len(nk.npairs)) print('npairs = ',nk.npairs) assert "Sampled %d pairs out of a total of %d"%(100, nk.npairs[b]) in cl.output print('i1 = ',i1) print('i2 = ',i2) print('sep = ',sep) assert len(i1) == 100 assert len(i2) == 100 assert len(sep) == 100 r1 = (x1**2 + y1**2 + z1**2)**0.5 r2 = (x2**2 + y2**2 + z2**2)**0.5 xx1 = x1/r1 yy1 = y1/r1 zz1 = z1/r1 xx2 = x2/r2 yy2 = y2/r2 zz2 = z2/r2 chord_sep = ((xx1[i1]-xx2[i2])**2 + (yy1[i1]-yy2[i2])**2 + (zz1[i1]-zz2[i2])**2)**0.5 arc_sep = np.arcsin(chord_sep/2.)*2. print('arc_sep = ',arc_sep) np.testing.assert_allclose(sep, arc_sep, rtol=0.1) np.testing.assert_array_less(sep, nk.right_edges[b]) np.testing.assert_array_less(nk.left_edges[b], sep) # Finally, check spherical coords with non-default units. ra1, dec1 = coord.CelestialCoord.xyz_to_radec(x1,y1,z1) ra2, dec2 = coord.CelestialCoord.xyz_to_radec(x2,y2,z2) cat1 = treecorr.Catalog(ra=ra1, dec=dec1, ra_units='rad', dec_units='rad') cat2 = treecorr.Catalog(ra=ra2, dec=dec2, ra_units='rad', dec_units='rad') nn = treecorr.NNCorrelation(min_sep=1., max_sep=60., nbins=50, sep_units='deg', metric='Arc') nn.process(cat1, cat2) print('rnom = ',nn.rnom) print('npairs = ',nn.npairs.astype(int)) b = 5 n = 50 i1, i2, sep = nn.sample_pairs(n, cat1, cat2, min_sep=nn.left_edges[b], max_sep=nn.right_edges[b]) print('i1 = ',i1) print('i2 = ',i2) print('sep = ',sep) assert nn.npairs[b] > n assert len(i1) == n assert len(i2) == n assert len(sep) == n c1 = [coord.CelestialCoord(r*coord.radians, d*coord.radians) for (r,d) in zip(ra1,dec1)] c2 = [coord.CelestialCoord(r*coord.radians, d*coord.radians) for (r,d) in zip(ra2,dec2)] actual_sep = np.array([c1[i1[k]].distanceTo(c2[i2[k]]) / coord.degrees for k in range(n)]) print('actual_sep = ',actual_sep) np.testing.assert_allclose(sep, actual_sep, rtol=0.1) np.testing.assert_array_less(sep, nn.right_edges[b]) np.testing.assert_array_less(nn.left_edges[b], sep)
def test_nan(): # Test handling of Nan values (w -> 0) nobj = 5000 np.random.seed(8675309) x = np.random.random_sample(nobj) y = np.random.random_sample(nobj) z = np.random.random_sample(nobj) ra = np.random.random_sample(nobj) dec = np.random.random_sample(nobj) r = np.random.random_sample(nobj) w = np.random.random_sample(nobj) wpos = np.random.random_sample(nobj) g1 = np.random.random_sample(nobj) g2 = np.random.random_sample(nobj) k = np.random.random_sample(nobj) # Turn 1% of these values into NaN x[np.random.choice(nobj, nobj//100)] = np.nan y[np.random.choice(nobj, nobj//100)] = np.nan z[np.random.choice(nobj, nobj//100)] = np.nan ra[np.random.choice(nobj, nobj//100)] = np.nan dec[np.random.choice(nobj, nobj//100)] = np.nan r[np.random.choice(nobj, nobj//100)] = np.nan w[np.random.choice(nobj, nobj//100)] = np.nan wpos[np.random.choice(nobj, nobj//100)] = np.nan g1[np.random.choice(nobj, nobj//100)] = np.nan g2[np.random.choice(nobj, nobj//100)] = np.nan k[np.random.choice(nobj, nobj//100)] = np.nan print('x is nan at ',np.where(np.isnan(x))) print('y is nan at ',np.where(np.isnan(y))) print('z is nan at ',np.where(np.isnan(z))) print('ra is nan at ',np.where(np.isnan(ra))) print('dec is nan at ',np.where(np.isnan(dec))) print('w is nan at ',np.where(np.isnan(w))) print('wpos is nan at ',np.where(np.isnan(wpos))) print('g1 is nan at ',np.where(np.isnan(g1))) print('g2 is nan at ',np.where(np.isnan(g2))) print('k is nan at ',np.where(np.isnan(k))) with CaptureLog() as cl: cat1 = treecorr.Catalog(x=x, y=y, z=z, w=w, k=k, logger=cl.logger) assert "NaNs found in x column." in cl.output assert "NaNs found in y column." in cl.output assert "NaNs found in z column." in cl.output assert "NaNs found in k column." in cl.output assert "NaNs found in w column." in cl.output mask = np.isnan(x) | np.isnan(y) | np.isnan(z) | np.isnan(k) | np.isnan(w) good = ~mask assert cat1.ntot == nobj assert cat1.nobj == np.sum(good) np.testing.assert_almost_equal(cat1.x[good], x[good]) np.testing.assert_almost_equal(cat1.y[good], y[good]) np.testing.assert_almost_equal(cat1.z[good], z[good]) np.testing.assert_almost_equal(cat1.w[good], w[good]) np.testing.assert_almost_equal(cat1.k[good], k[good]) np.testing.assert_almost_equal(cat1.w[mask], 0) with CaptureLog() as cl: cat2 = treecorr.Catalog(ra=ra, dec=dec, r=r, w=w, wpos=wpos, g1=g1, g2=g2, ra_units='hours', dec_units='degrees', logger=cl.logger) assert "NaNs found in ra column." in cl.output assert "NaNs found in dec column." in cl.output assert "NaNs found in r column." in cl.output assert "NaNs found in g1 column." in cl.output assert "NaNs found in g2 column." in cl.output assert "NaNs found in w column." in cl.output assert "NaNs found in wpos column." in cl.output mask = np.isnan(ra) | np.isnan(dec) | np.isnan(r) | np.isnan(g1) | np.isnan(g2) | np.isnan(wpos) | np.isnan(w) good = ~mask assert cat2.ntot == nobj assert cat2.nobj == np.sum(good) np.testing.assert_almost_equal(cat2.ra[good], ra[good] * coord.hours / coord.radians) np.testing.assert_almost_equal(cat2.dec[good], dec[good] * coord.degrees / coord.radians) np.testing.assert_almost_equal(cat2.r[good], r[good]) np.testing.assert_almost_equal(cat2.w[good], w[good]) np.testing.assert_almost_equal(cat2.wpos[good], wpos[good]) np.testing.assert_almost_equal(cat2.g1[good], g1[good]) np.testing.assert_almost_equal(cat2.g2[good], g2[good]) np.testing.assert_almost_equal(cat2.w[mask], 0) # If no weight column, it is make automatically to deal with Nans. with CaptureLog() as cl: cat3 = treecorr.Catalog(x=x, y=y, g1=g1, g2=g2, logger=cl.logger) mask = np.isnan(x) | np.isnan(y) | np.isnan(g1) | np.isnan(g2) good = ~mask assert cat3.ntot == nobj assert cat3.nobj == np.sum(good) np.testing.assert_almost_equal(cat3.x[good], x[good]) np.testing.assert_almost_equal(cat3.y[good], y[good]) np.testing.assert_almost_equal(cat3.w[good], 1.) np.testing.assert_almost_equal(cat3.g1[good], g1[good]) np.testing.assert_almost_equal(cat3.g2[good], g2[good]) np.testing.assert_almost_equal(cat3.w[mask], 0)