Commit adc4e019 authored by Ryan Gutenkunst's avatar Ryan Gutenkunst
Browse files

Add filter_pops to Spectrum object

parent b17bc6fb
......@@ -889,7 +889,7 @@ def filter_pops(phi, xx, tokeep):
"""
Filter phi to keep only certain populations.
Returns new phi with len(tokeep) fewer populations.
Returns new phi with len(tokeep) populations.
phi: phi corresponding to original populations
xx: Mapping of points in phi to frequencies in population to be removed
......
......@@ -468,6 +468,24 @@ class Spectrum(numpy.ma.masked_array):
else:
return output
def filter_pops(self, tokeep, mask_corners=True):
"""
Filter Spectrum to keep only certain populations.
Returns new Spectrum with len(tokeep) populations.
Note: This is similar in practice to the marginalize operation. But here
populations are numbered from 1, as in the majority of dadi.
tokeep: List of population numbers to keep, numbering from 1.
mask_corners: If True, the typical corners of the resulting fs will be
masked
"""
toremove = list(range(0, self.ndim))
for pop_ii in tokeep:
# Apply -1 factor to account for indexing in marginalize
toremove.remove(pop_ii-1)
return self.marginalize(toremove)
def _counts_per_entry(self):
"""
Counts per population for each entry in the fs.
......
......@@ -38,8 +38,8 @@ class SpectrumTestCase(unittest.TestCase):
return_comments=True)
os.remove(filename)
# Ensure that fs was read correctly.
self.assert_(numpy.allclose(fsout.data, fsin.data))
self.assert_(numpy.all(fsout.mask == fsin.mask))
self.assertTrue(numpy.allclose(fsout.data, fsin.data))
self.assertTrue(numpy.all(fsout.mask == fsin.mask))
self.assertEqual(fsout.folded, fsin.folded)
# Ensure comments were read correctly.
for ii,line in enumerate(commentsin):
......@@ -53,8 +53,8 @@ class SpectrumTestCase(unittest.TestCase):
return_comments=True)
os.remove(filename)
# Ensure that fs was read correctly.
self.assert_(numpy.allclose(fsout.data, fsin.data))
self.assert_(numpy.all(fsout.mask == fsin.mask))
self.assertTrue(numpy.allclose(fsout.data, fsin.data))
self.assertTrue(numpy.all(fsout.mask == fsin.mask))
self.assertEqual(fsout.folded, fsin.folded)
# Ensure comments were read correctly.
for ii,line in enumerate(commentsin):
......@@ -71,8 +71,8 @@ class SpectrumTestCase(unittest.TestCase):
os.remove(filename)
# Ensure that fs was read correctly.
self.assert_(numpy.allclose(fsout.data, fsin.data))
self.assert_(numpy.all(fsout.mask == fsin.mask))
self.assertTrue(numpy.allclose(fsout.data, fsin.data))
self.assertTrue(numpy.all(fsout.mask == fsin.mask))
self.assertEqual(fsout.folded, fsin.folded)
def test_folding(self):
......@@ -87,11 +87,11 @@ class SpectrumTestCase(unittest.TestCase):
self.assertAlmostEqual(fs.sum(), ff.sum(), 6)
self.assertAlmostEqual(fs.data.sum(), ff.data.sum(), 6)
# Ensure that the empty entries are actually empty.
self.assert_(numpy.all(ff.data[::-1] == numpy.tril(ff.data[::-1])))
self.assertTrue(numpy.all(ff.data[::-1] == numpy.tril(ff.data[::-1])))
# This turns out to be the correct result.
correct = numpy.tri(4)[::-1][-3:]*11
self.assert_(numpy.allclose(correct, ff.data))
self.assertTrue(numpy.allclose(correct, ff.data))
def test_ambiguous_folding(self):
"""
......@@ -108,7 +108,7 @@ class SpectrumTestCase(unittest.TestCase):
correct = numpy.zeros((4,4))
correct[0,3] = correct[3,0] = 2
self.assert_(numpy.allclose(correct, ff.data))
self.assertTrue(numpy.allclose(correct, ff.data))
def test_masked_folding(self):
"""
......@@ -123,21 +123,21 @@ class SpectrumTestCase(unittest.TestCase):
ff = fs.fold()
# Ensure that all those are masked.
for entry in [(1,2), (3,4), (1,1)]:
self.assert_(ff.mask[entry])
self.assertTrue(ff.mask[entry])
def test_folded_slices(self):
ns = (3,4)
fs1 = dadi.Spectrum(numpy.random.rand(*ns))
folded1 = fs1.fold()
self.assert_(fs1[:].folded == False)
self.assert_(folded1[:].folded == True)
self.assertTrue(fs1[:].folded == False)
self.assertTrue(folded1[:].folded == True)
self.assert_(fs1[0].folded == False)
self.assert_(folded1[1].folded == True)
self.assertTrue(fs1[0].folded == False)
self.assertTrue(folded1[1].folded == True)
self.assert_(fs1[:,0].folded == False)
self.assert_(folded1[:,1].folded == True)
self.assertTrue(fs1[:,0].folded == False)
self.assertTrue(folded1[:,1].folded == True)
def test_folded_arithmetic(self):
"""
......@@ -175,77 +175,77 @@ class SpectrumTestCase(unittest.TestCase):
result = op(fs1,fs2)
self.assertFalse(result.folded)
self.assert_(numpy.all(result.mask == fs1.mask))
self.assertTrue(numpy.all(result.mask == fs1.mask))
result = op(fs1,2.0)
self.assertFalse(result.folded)
self.assert_(numpy.all(result.mask == fs1.mask))
self.assertTrue(numpy.all(result.mask == fs1.mask))
result = op(2.0,fs2)
self.assertFalse(result.folded)
self.assert_(numpy.all(result.mask == fs2.mask))
self.assertTrue(numpy.all(result.mask == fs2.mask))
result = op(fs1,numpyfloat)
self.assertFalse(result.folded)
self.assert_(numpy.all(result.mask == fs1.mask))
self.assertTrue(numpy.all(result.mask == fs1.mask))
result = op(numpyfloat,fs2)
self.assertFalse(result.folded)
self.assert_(numpy.all(result.mask == fs2.mask))
self.assertTrue(numpy.all(result.mask == fs2.mask))
result = op(fs1,arr)
self.assertFalse(result.folded)
self.assert_(numpy.all(result.mask == fs1.mask))
self.assertTrue(numpy.all(result.mask == fs1.mask))
result = op(arr,fs2)
self.assertFalse(result.folded)
self.assert_(numpy.all(result.mask == fs2.mask))
self.assertTrue(numpy.all(result.mask == fs2.mask))
result = op(fs1,marr)
self.assertFalse(result.folded)
self.assert_(numpy.all(result.mask == fs1.mask))
self.assertTrue(numpy.all(result.mask == fs1.mask))
result = op(marr,fs2)
self.assertFalse(result.folded)
self.assert_(numpy.all(result.mask == fs2.mask))
self.assertTrue(numpy.all(result.mask == fs2.mask))
# Now with folded Spectra
result = op(folded1,folded2)
self.assertTrue(result.folded)
self.assert_(numpy.all(result.mask == folded1.mask))
self.assertTrue(numpy.all(result.mask == folded1.mask))
result = op(folded1,2.0)
self.assertTrue(result.folded)
self.assert_(numpy.all(result.mask == folded1.mask))
self.assertTrue(numpy.all(result.mask == folded1.mask))
result = op(2.0,folded2)
self.assertTrue(result.folded)
self.assert_(numpy.all(result.mask == folded2.mask))
self.assertTrue(numpy.all(result.mask == folded2.mask))
result = op(folded1,numpyfloat)
self.assertTrue(result.folded)
self.assert_(numpy.all(result.mask == folded1.mask))
self.assertTrue(numpy.all(result.mask == folded1.mask))
result = op(numpyfloat,folded2)
self.assertTrue(result.folded)
self.assert_(numpy.all(result.mask == folded2.mask))
self.assertTrue(numpy.all(result.mask == folded2.mask))
result = op(folded1,arr)
self.assertTrue(result.folded)
self.assert_(numpy.all(result.mask == folded1.mask))
self.assertTrue(numpy.all(result.mask == folded1.mask))
result = op(arr,folded2)
self.assertTrue(result.folded)
self.assert_(numpy.all(result.mask == folded2.mask))
self.assertTrue(numpy.all(result.mask == folded2.mask))
result = op(folded1,marr)
self.assertTrue(result.folded)
self.assert_(numpy.all(result.mask == folded1.mask))
self.assertTrue(numpy.all(result.mask == folded1.mask))
result = op(marr,folded2)
self.assertTrue(result.folded)
self.assert_(numpy.all(result.mask == folded2.mask))
self.assertTrue(numpy.all(result.mask == folded2.mask))
# Check that exceptions are properly raised when folding status
# differs
......@@ -268,46 +268,46 @@ class SpectrumTestCase(unittest.TestCase):
# Check that in-place operations preserve folding status.
op(fs1,fs2)
self.assertFalse(fs1.folded)
self.assert_(numpy.all(fs1.mask == fs1origmask))
self.assertTrue(numpy.all(fs1.mask == fs1origmask))
op(fs1,2.0)
self.assertFalse(fs1.folded)
self.assert_(numpy.all(fs1.mask == fs1origmask))
self.assertTrue(numpy.all(fs1.mask == fs1origmask))
op(fs1,numpyfloat)
self.assertFalse(fs1.folded)
self.assert_(numpy.all(fs1.mask == fs1origmask))
self.assertTrue(numpy.all(fs1.mask == fs1origmask))
op(fs1,arr)
self.assertFalse(fs1.folded)
self.assert_(numpy.all(fs1.mask == fs1origmask))
self.assertTrue(numpy.all(fs1.mask == fs1origmask))
op(fs1,marr)
self.assertFalse(fs1.folded)
self.assert_(numpy.all(fs1.mask == fs1origmask))
self.assertTrue(numpy.all(fs1.mask == fs1origmask))
# Now folded Spectra
folded1origmask = folded1.mask.copy()
op(folded1,folded2)
self.assertTrue(folded1.folded)
self.assert_(numpy.all(folded1.mask == folded1origmask))
self.assertTrue(numpy.all(folded1.mask == folded1origmask))
op(folded1,2.0)
self.assertTrue(folded1.folded)
self.assert_(numpy.all(folded1.mask == folded1origmask))
self.assertTrue(numpy.all(folded1.mask == folded1origmask))
op(folded1,numpyfloat)
self.assertTrue(folded1.folded)
self.assert_(numpy.all(folded1.mask == folded1origmask))
self.assertTrue(numpy.all(folded1.mask == folded1origmask))
op(folded1,arr)
self.assertTrue(folded1.folded)
self.assert_(numpy.all(folded1.mask == folded1origmask))
self.assertTrue(numpy.all(folded1.mask == folded1origmask))
op(folded1,marr)
self.assertTrue(folded1.folded)
self.assert_(numpy.all(folded1.mask == folded1origmask))
self.assertTrue(numpy.all(folded1.mask == folded1origmask))
# Check that exceptions are properly raised.
self.assertRaises(ValueError, op, fs1, folded2)
......@@ -356,14 +356,14 @@ class SpectrumTestCase(unittest.TestCase):
manual = dadi.Spectrum(fs.data.sum(axis=1))
# Check that these are equal in the unmasked entries.
self.assert_(numpy.allclose(numpy.where(marg1.mask, 0, marg1.data),
self.assertTrue(numpy.allclose(numpy.where(marg1.mask, 0, marg1.data),
numpy.where(manual.mask, 0, manual.data)))
# Check folded Spectrum objects. I should get the same result if I
# marginalize then fold, as if I fold then marginalize.
mf1 = marg1.fold()
mf2 = folded.marginalize([1])
self.assert_(numpy.allclose(mf1,mf2))
self.assertTrue(numpy.allclose(mf1,mf2))
def test_projection(self):
# Test that projecting a multi-dimensional Spectrum succeeds
......@@ -377,27 +377,27 @@ class SpectrumTestCase(unittest.TestCase):
# equilibrium spectrum
fs = dadi.Spectrum(1./numpy.arange(100))
p = fs.project([17])
self.assert_(numpy.allclose(p[1:-1], 1./numpy.arange(1,len(p)-1)))
self.assertTrue(numpy.allclose(p[1:-1], 1./numpy.arange(1,len(p)-1)))
# Check that masked values are propagated correctly.
fs = dadi.Spectrum(1./numpy.arange(20))
# All values with 3 or fewer observed should be masked.
fs.mask[3] = True
p = fs.project([10])
self.assert_(numpy.all(p.mask[:4]))
self.assertTrue(numpy.all(p.mask[:4]))
# Check that masked values are propagated correctly.
fs = dadi.Spectrum(1./numpy.arange(20))
fs.mask[-3] = True
# All values with 3 or fewer observed should be masked.
p = fs.project([10])
self.assert_(numpy.all(p.mask[-3:]))
self.assertTrue(numpy.all(p.mask[-3:]))
# A more complicated two dimensional projection problem...
fs = dadi.Spectrum(numpy.random.uniform(size=(9,7)))
fs.mask[2,3] = True
p = fs.project([4,4])
self.assert_(numpy.all(p.mask[:3,1:4]))
self.assertTrue(numpy.all(p.mask[:3,1:4]))
# Test that projecting a folded multi-dimensional Spectrum succeeds
# Should get the same result if I fold then project as if I project
......@@ -412,7 +412,25 @@ class SpectrumTestCase(unittest.TestCase):
pf2 = folded.project([3,4,5])
# Check equality
self.assert_(numpy.all(pf1.mask == pf2.mask))
self.assert_(numpy.allclose(pf1.data, pf2.data))
self.assertTrue(numpy.all(pf1.mask == pf2.mask))
self.assertTrue(numpy.allclose(pf1.data, pf2.data))
def test_filter_pops(self):
"""
Test filtering populations in Spectrum.
"""
ns = (7,8,6)
fs = dadi.Spectrum(numpy.random.uniform(size=ns))
folded = fs.fold()
marg1 = fs.filter_pops([1,3])
self.assertEqual(marg1.shape, (ns[0],ns[2]))
marg1 = fs.filter_pops([2])
self.assertEqual(marg1.shape, (ns[1],))
suite = unittest.TestLoader().loadTestsFromTestCase(SpectrumTestCase)
if __name__ == '__main__':
unittest.main()
\ No newline at end of file
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment