[Python: track values that are removed from an origin array] #maskedarray #numpy #np #ma

drop_tracker.py

import numpy as np

def drop_tracker(curArr, rmIndxs, dropTracker):
	"""
    Track values that are removed from an array.
   	dropTracker is a MaskedArray that curArr has derived from and slightly changed.
    Requires curArr pts to be closer to their complimentary dropTracker pts more than any other dropTracker pts.
    """
    assert type(dropTracker) is np.ma.MaskedArray
    newArr = np.delete(curArr, rmIndxs)
    dropVals = curArr[rmIndxs]
    for val in dropVals:
        toMask = np.argmin(np.abs(dropTracker-val))
        dropTracker[toMask] = np.ma.masked
    return newArr, dropTracker

test_drop_tracker.py

import unittest
import numpy as np

import dropTracker

class TestUtils(unittest.TestCase):

    def test_drop_tracker_basic(self):
        x = np.linspace(0, 10000, num=101, dtype=int)
        d = np.ma.array(x)
        i = [0, 2, 4, 6, 8]
        for _ in range(5):
            x, d = dropTracker.drop_tracker(x, i, d)
        equal = (x == d.compressed()).all()
        self.assertTrue(equal)

    def test_drop_tracker_adv(self):
        x = np.linspace(0, 100000, num=101, dtype=int)
        d_track = np.ma.array(x) # for testing

        randomizer = np.vectorize( lambda x: x + np.random.choice([-1, 1])*np.random.randint(0, 10) )
        d = np.ma.apply_along_axis(randomizer, 0, x) # d is slightly different from d_track

        x1 = x2 = x.copy()
        for _ in range(10):
            i = np.random.choice(x1.size-1, size=10)
            x1, d_track = dropTracker.drop_tracker(x1, i, d_track)
            x2, d = dropTracker.drop_tracker(x2, i, d)
        equal = (d.mask == d_track.mask).all()
        self.assertTrue(equal, "paramters:\nx = {}\nd_track = {}\nd = {}")