How to use the simdkalman.KalmanFilter function in simdkalman

def test_em_algorithm_multi_dimensional_observations(self):
        training_matrix = np.ones((5,10,2))
        training_matrix[1,1,1] = np.nan

        kf = simdkalman.KalmanFilter(
            state_transition = np.eye(2),
            process_noise = 0.1,
            observation_model = np.array([[1,1], [0,1]]),
            observation_noise = 0.1)

        r = kf.em(training_matrix, n_iter=5, verbose=False)

        self.assertSequenceEqual(r.process_noise.shape, (5,2,2))
        A0 = r.process_noise[0,...]
        self.assertMatrixEqual(A0, A0.T, epsilon=1e-8)
        self.assertTrue(min(np.linalg.eig(A0)[0]) > 0)

        B = r.observation_noise
        self.assertSequenceEqual(B.shape, (5,2,2))
        self.assertTrue(min(np.linalg.eig(B[0,...])[0]) > 0)

def test_one_dimensional(self):

        training_matrix = range(10)

        kf = simdkalman.KalmanFilter(
            state_transition = np.eye(2),
            process_noise = 0.1,
            observation_model = np.array([[1,1]]),
            observation_noise = 0.1)

        r = kf.compute(
            training_matrix,
            n_test = 4,
            initial_value = [0,0],
            initial_covariance = 1.0,
            smoothed = True,
            gains = True,
            log_likelihood = True)

        self.assertSequenceEqual(r.predicted.observations.mean.shape, (4,))
        self.assertSequenceEqual(r.smoothed.observations.mean.shape, (10,))

def test_em_algorithm(self):
        training_matrix = np.ones((5,10))

        kf = simdkalman.KalmanFilter(
            state_transition = np.eye(2),
            process_noise = 0.1,
            observation_model = np.array([[1,1]]),
            observation_noise = 0.1)

        r = kf.em(training_matrix, n_iter=5, verbose=False)

        self.assertSequenceEqual(r.process_noise.shape, (5,2,2))
        A0 = r.process_noise[0,...]
        self.assertMatrixEqual(A0, A0.T)
        self.assertTrue(min(np.linalg.eig(A0)[0]) > 0)

        B = r.observation_noise
        self.assertSequenceEqual(B.shape, (5,1,1))
        self.assertTrue(min(list(B)) > 0)

def test_smooth_helper_kalman_filter_2_states(self):
        training_matrix = np.ones((5,10))
        training_matrix[1,1] = np.nan

        kf = simdkalman.KalmanFilter(
            state_transition = np.eye(2),
            process_noise = 0.1,
            observation_model = np.array([[1,1]]),
            observation_noise = 0.1)

        r = kf.smooth(training_matrix)

        self.assertSequenceEqual(r.observations.mean.shape, training_matrix.shape)
        self.assertSequenceEqual(r.states.mean.shape, (5,10,2))
        self.assertSequenceEqual(r.states.cov.shape, (5,10,2,2))

def test_train_and_predict_vectorized_kalman_filter_ema(self):
        training_matrix = np.ones((5,10))

        kf = simdkalman.KalmanFilter(
            state_transition = 1,
            process_noise = 0.1,
            observation_model = 1,
            observation_noise = 0.1)

        r = kf.compute(
            training_matrix,
            n_test = 4,
            initial_covariance = 1.0)

        self.assertSequenceEqual(r.predicted.observations.mean.shape, (5,4))
        self.assertSequenceEqual(r.smoothed.observations.mean.shape, training_matrix.shape)
        self.assertSequenceEqual(r.predicted.states.mean.shape, (5,4,1))
        self.assertSequenceEqual(r.smoothed.states.mean.shape, (5,10,1))
        self.assertSequenceEqual(r.predicted.states.cov.shape, (5,4,1,1))
        self.assertSequenceEqual(r.smoothed.states.cov.shape, (5,10,1,1))

def test_em_algorithm_3_dimensional_observations(self):
        training_matrix = np.ones((5,10,3))
        training_matrix[1,1,1] = np.nan

        kf = simdkalman.KalmanFilter(
            state_transition = np.eye(2),
            process_noise = 0.1,
            observation_model = np.array([[1,1], [0,1], [1,0]]),
            observation_noise = 0.1)

        r = kf.em(training_matrix, n_iter=5, verbose=False)

        self.assertSequenceEqual(r.process_noise.shape, (5,2,2))
        A0 = r.process_noise[0,...]
        self.assertMatrixEqual(A0, A0.T, epsilon=1e-8)
        self.assertTrue(min(np.linalg.eig(A0)[0]) > 0)

        B = r.observation_noise
        self.assertSequenceEqual(B.shape, (5,3,3))
        self.assertTrue(min(np.linalg.eig(B[0,...])[0]) > 0)

def test_predict_helper_ema(self):
        training_matrix = np.ones((5,10))

        kf = simdkalman.KalmanFilter(
            state_transition = 1,
            process_noise = 0.1,
            observation_model = 1,
            observation_noise = 0.1)

        r = kf.predict(training_matrix, n_test = 4)

        self.assertSequenceEqual(r.observations.mean.shape, (5,4))
        self.assertSequenceEqual(r.states.mean.shape, (5,4,1))
        self.assertSequenceEqual(r.states.cov.shape, (5,4,1,1))

def compute_simd():
    kf = simdkalman.KalmanFilter(
        state_transition,
        process_noise,
        observation_model,
        observation_noise)

    r = kf.smooth(data,
        initial_value = initial_values[...,np.newaxis],
        initial_covariance = initial_covariance,
        observations = False)

    return (r.states.mean, r.states.cov)