from __future__ import print_function
import GPinv
import numpy as np
import unittest
import tensorflow as tf

def make_LosMatrix(r,z):
    """
    Constructing a matrix for a cylindrical plasma coordinate.
    A[i,j] element stores a passing length for a shell j by chord i.

    r: radius coordinate
    z: los position.
    """
    n = r.shape[0]
    N = z.shape[0]
    dr = r[1] - r[0]
    A = np.zeros((N,n))
    for i in range(N):
        for j in range(1, n-1):
            # inner radius of the shell j
            r_in  = 0.5*(r[j-1]+r[j])
            # outer radius of the shell j
            r_out = 0.5*(r[j+1]+r[j])
            # if the code passes the shell j
            if np.abs(z[i]) < r_out:
                # if the chord passes both the inner and outer sides
                if np.abs(z[i]) < r_in:
                    A[i,j] = 2.*np.sqrt(r_out**2. - z[i]**2) - \
                             2.*np.sqrt(r_in **2. - z[i]**2)
                # if the chord passes only the outer side
                elif np.abs(z[i]) > r_in:
                    A[i,j] = 2.*np.sqrt(r_out **2. - z[i]**2)

        # inner radius of the shell j
        r_in  = 0.5*(r[n-2]+r[n-1])
        # outer radius of the shell j
        r_out = r[n-1]
        if np.abs(z[i]) < r_out:
            if np.abs(z[i]) < r_in:
                # inner radius of the shell j
                A[i,n-1] = 2.*np.sqrt(r_out**2. - z[i]**2) - \
                         2.*np.sqrt(r_in **2. - z[i]**2)
                # if the chord passes only the outer side
            elif np.abs(z[i]) > r_in:
                A[i,n-1] = 2.*np.sqrt(r_out **2. - z[i]**2)

    return A


class Test_linear_model(unittest.TestCase):
    def test(self):
        """
        Abel inversion for the synthetic data.
        The observation y is generated by the latent function f
        y = A f + e
        with
        A = los-length-matrix
        """
        rng = np.random.RandomState(0)
        n = 30
        N = 40
        # radial position. 0~1
        r = np.linspace(0, 1., n)
        # height of LOS. -0.9~1
        z = np.linspace(-0.9, 0.9, N)
        # observation results at each Z
        y = np.zeros((N, 1))
        # noise amplitude on the observation
        e = 0.1
        # synthetic latent function
        f = np.exp(-(r-0.5)*(r-0.5)/0.1) + np.exp(-(r+0.5)*(r+0.5)/0.1)
        # constructing LOS-matrix
        A = make_LosMatrix(r,z)
        # synthetic signals
        y = np.dot(A, f) + e*rng.randn(N)

        linear_model = GPinv.linear_model.LinearModel(X=r.reshape(-1,1),
                            Y=y.reshape(-1,1), Amat=A,
                            kern=GPinv.kernels.RBF(1))
        linear_model.optimize(disp=False)
        # predict
        Xnew = np.linspace(0,1.,n-1).reshape(-1,1)
        # just check the variance value
        self.assertTrue(np.allclose(linear_model.likelihood.variance.value,
                                    e*e, rtol=0.1))

if __name__ == '__main__':
    unittest.main()

@@ -2819,16 +2819,86 @@
 e value%0A
+        linear_model.predict_f(Xnew)%0A        linear_model.predict_y()%0A
         

test added for predict_f and predict_y